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

[guan jian]

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

>From 2db60279a33f829905e4b5f6ca0d57a678b5a528 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..d1c481e175afb 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,