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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: guan jian (rez5427)

<details>
<summary>Changes</summary>

In some case original code will crash.

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


1 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp (+15-3) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index 3bb5f8af821c0..8d8465c84580f 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -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,
+          {sourceSizes[i], originalLowPad[i], originalHighPad[i]}))
+    }
+    Value outputUnPack = tensor::EmptyOp::create(
+        rewriter, loc, outputSizes, padOp.getResultType().getElementType());
 
     Value replacement = linalg::UnPackOp::create(
         rewriter, loc, newPadOp.getResult(), outputUnPack, innerDimsPos,

``````````

</details>


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