[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 cbec6e6ced12605ff4f27e6e54553a61bece0545 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

---
 .../Dialect/Linalg/Transforms/DataLayoutPropagation.cpp   | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index d1c481e175afb..b90c177244e44 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -1316,10 +1316,11 @@ struct PushDownUnPackThroughPadOp : public OpRewritePattern<tensor::PadOp> {
 
     SmallVector<OpFoldResult> outputSizes;
     AffineExpr d0, d1, d2;
+    AffineExpr sumExpr = 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,
+          rewriter, loc, sumExpr,
           {sourceSizes[i], originalLowPad[i], originalHighPad[i]}));
     }
     Value outputUnPack = tensor::EmptyOp::create(
@@ -1584,10 +1585,7 @@ pushDownExtractSliceOpThroughGenericOp(RewriterBase &rewriter,
       genericOp.getIndexingMapsArray(), genericOp.getIteratorTypesArray(),
       /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp));
   rewriter.cloneRegionBefore(genericOp.getRegion(), newGenericOp.getRegion(),
-                             newGenericOp.getRegion().begin());
-
-  auto extractOp = tensor::ExtractSliceOp::create(
-      rewriter, loc,
+                             newGenericOp.getRegion().begNit: 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.
       newGenericOp.getTiedOpResult(newGenericOp.getDpsInitOperand(0)),
       outputLowPads, newSizes, newStrides);
   Value extractRes = extractOp.getResult();