[Mlir-commits] [mlir] [MLIR] Use dynamic dim in PushDownUnPackThroughPadOp (PR #171122)

[llvmlistbot at llvm.org]

================
@@ -1309,9 +1309,21 @@ struct PushDownUnPackThroughPadOp : public OpRewritePattern<tensor::PadOp> {
                                           paddingVal, padOp.getNofold());
 
     // Inject the linalg.unpack right after the packed padOp.
-    Value outputUnPack =
-        tensor::EmptyOp::create(rewriter, loc, padOp.getResultType().getShape(),
-                                padOp.getResultType().getElementType());
+    SmallVector<OpFoldResult> sourceSizes =
+        tensor::getMixedSizes(rewriter, loc, unpackOp);
+    SmallVector<OpFoldResult> originalLowPad = padOp.getMixedLowPad();
+    SmallVector<OpFoldResult> originalHighPad = padOp.getMixedHighPad();
+
+    SmallVector<OpFoldResult> outputSizes;
+    AffineExpr d0, d1, d2;
+    bindDims(rewriter.getContext(), d0, d1, d2);
+    for (size_t i = 0; i < sourceSizes.size(); ++i) {
+      outputSizes.push_back(affine::makeComposedFoldedAffineApply(
+          rewriter, loc, d0 + d1 + d2,
----------------
MaheshRavishankar wrote:

Nit: could you create this affine expression separately. This will have an overhead of creating the same affine expression everytime, checking if it matches with existing affine expr (since affine exprs are uniqued) and then discard it. This wont be automatically hoisted AFAIK.

https://github.com/llvm/llvm-project/pull/171122