[Mlir-commits] [mlir] [MLIR] Use dynamic dim in PushDownUnPackThroughPadOp (PR #171122)
guan jian
llvmlistbot at llvm.org
Sun Dec 14 02:42:28 PST 2025
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/4] [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/4] 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/4] 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(
>From 1c83542f4992c06c9910f2d90d7dec6e9d1654c7 Mon Sep 17 00:00:00 2001
From: rez5427 <guanjian at stu.cdut.edu.cn>
Date: Sun, 14 Dec 2025 18:42:16 +0800
Subject: [PATCH 4/4] Use createDestinationTensor
---
.../Transforms/DataLayoutPropagation.cpp | 20 ++++---------------
.../Linalg/data-layout-propagation.mlir | 16 +++++----------
2 files changed, 9 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index 01024dbe9bec8..a1499824fde15 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -1309,22 +1309,10 @@ struct PushDownUnPackThroughPadOp : public OpRewritePattern<tensor::PadOp> {
paddingVal, padOp.getNofold());
// Inject the linalg.unpack right after the packed padOp.
- 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);
- AffineExpr sumExpr = d0 + d1 + d2;
- for (size_t i = 0; i < sourceSizes.size(); ++i) {
- outputSizes.push_back(affine::makeComposedFoldedAffineApply(
- rewriter, loc, sumExpr,
- {sourceSizes[i], originalLowPad[i], originalHighPad[i]}));
- }
- Value outputUnPack = tensor::EmptyOp::create(
- rewriter, loc, outputSizes, padOp.getResultType().getElementType());
+ // Compute the unpacked output size directly from the padded packed tensor.
+ Value outputUnPack = linalg::UnPackOp::createDestinationTensor(
+ rewriter, loc, newPadOp.getResult(), unpackOp.getMixedTiles(),
+ innerDimsPos, outerDimsPerm);
Value replacement = linalg::UnPackOp::create(
rewriter, loc, newPadOp.getResult(), outputUnPack, innerDimsPos,
diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index 30a7d9c1be6bc..6121b69a3ecd8 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1651,22 +1651,16 @@ func.func @test_dynamic_unpack_pad(%arg0: tensor<?x4x8xf32>, %dim: index) -> ten
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: %[[DIM_VAL:.+]] = tensor.dim %[[PAD]], %[[C0]]
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM_VAL]])
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[PAD]]
// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [8]
-// CHECK-SAME: into %[[EMPTY1]]
-// CHECK: return %[[UNPACK1]]
+// CHECK-SAME: into %[[EMPTY]]
+// CHECK: return %[[UNPACK]]
More information about the Mlir-commits
mailing list