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

[guan jian]

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

In some case original code will crash.

>From 586fe5cda2934d27cc498cba6da1e8bed03dbee9 Mon Sep 17 00:00:00 2001
From: rez5427 <guanjian at stu.cdut.edu.cn>
Date: Mon, 8 Dec 2025 21:32:33 +0800
Subject: [PATCH] [MLIR] Use dynamic dim in PushDownUnPackThroughPadOp

---
 .../Transforms/DataLayoutPropagation.cpp       | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

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,