[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 1/3] [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,

>From a90078ae36646d1e4fab9b42a5db01075bfd4639 Mon Sep 17 00:00:00 2001
From: rez5427 <guanjian at stu.cdut.edu.cn>
Date: Sun, 14 Dec 2025 10:10:19 +0800
Subject: [PATCH 2/3] Add test case

---
 .../Linalg/data-layout-propagation.mlir       | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index 7a16bc0a4faee..30a7d9c1be6bc 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1634,3 +1634,39 @@ func.func @push_extract_through_generic_secondextract(%arg0: tensor<128x128xf32>
 // CHECK-SAME:        ins(%[[PAD]], %[[ARG0]]
 // CHECK:           %[[EXTRACT2:.+]] =  tensor.extract_slice %[[GENERIC]]
 // CHECK:           scf.yield %[[EXTRACT2]]
+
+// -----
+
+func.func @test_dynamic_unpack_pad(%arg0: tensor<?x4x8xf32>, %dim: index) -> tensor<?x32xf32> {
+  %dest = tensor.empty(%dim) : tensor<?x32xf32>
+  %unpack = linalg.unpack %arg0 inner_dims_pos = [1] inner_tiles = [8]
+            into %dest : tensor<?x4x8xf32> -> tensor<?x32xf32>
+
+  %c0 = arith.constant 0.0 : f32
+  %pad = tensor.pad %unpack low[2, 0] high[3, 0] {
+    ^bb0(%arg1: index, %arg2: index):
+      tensor.yield %c0 : f32
+  } : tensor<?x32xf32> to tensor<?x32xf32>
+
+  return %pad : tensor<?x32xf32>
+}
+
+// CHECK-DAG:  #[[$MAP:.+]] = affine_map<()[s0] -> (s0 + 5)>
+// CHECK-LABEL: func.func @test_dynamic_unpack_pad
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:    %[[DIM:[a-zA-Z0-9]+]]
+// CHECK-DAG:     %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG:     %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK:         %[[EMPTY0:.+]] = tensor.empty(%[[DIM]])
+// CHECK:         %[[UNPACK0:.+]] = linalg.unpack %[[ARG0]]
+// CHECK-SAME:      inner_dims_pos = [1] inner_tiles = [8]
+// CHECK-SAME:      into %[[EMPTY0]]
+// CHECK:         %[[PAD:.+]] = tensor.pad %[[ARG0]] low[2, 0, 0] high[3, 0, 0]
+// CHECK:           tensor.yield %[[CST]]
+// CHECK:         %[[DIM_VAL:.+]] = tensor.dim %[[UNPACK0]], %[[C0]]
+// CHECK:         %[[NEW_DIM:.+]] = affine.apply #[[$MAP]]()[%[[DIM_VAL]]]
+// CHECK:         %[[EMPTY1:.+]] = tensor.empty(%[[NEW_DIM]])
+// CHECK:         %[[UNPACK1:.+]] = linalg.unpack %[[PAD]]
+// CHECK-SAME:      inner_dims_pos = [1] inner_tiles = [8]
+// CHECK-SAME:      into %[[EMPTY1]]
+// CHECK:         return %[[UNPACK1]]

>From 5bf6778fe308b41ce6d329b482d2ebd1e2c215c7 Mon Sep 17 00:00:00 2001
From: rez5427 <guanjian at stu.cdut.edu.cn>
Date: Sun, 14 Dec 2025 16:44:23 +0800
Subject: [PATCH 3/3] Use sumExpr

---
 mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index d1c481e175afb..01024dbe9bec8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -1317,9 +1317,10 @@ struct PushDownUnPackThroughPadOp : public OpRewritePattern<tensor::PadOp> {
     SmallVector<OpFoldResult> outputSizes;
     AffineExpr d0, d1, d2;
     bindDims(rewriter.getContext(), d0, d1, d2);
+    AffineExpr sumExpr = d0 + d1 + d2;
     for (size_t i = 0; i < sourceSizes.size(); ++i) {
       outputSizes.push_back(affine::makeComposedFoldedAffineApply(
-          rewriter, loc, d0 + d1 + d2,
+          rewriter, loc, sumExpr,
           {sourceSizes[i], originalLowPad[i], originalHighPad[i]}));
     }
     Value outputUnPack = tensor::EmptyOp::create(