[Mlir-commits] [mlir] feat(linalg): enable lowering/decomposing scalable pack ops (PR #200216)
Ege Beysel
llvmlistbot at llvm.org
Mon Jun 1 05:22:14 PDT 2026
https://github.com/egebeysel updated https://github.com/llvm/llvm-project/pull/200216
>From 81789c0ef4a9e90da6995c98cb6a9df5caa3b947 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Tue, 28 Apr 2026 13:58:34 +0200
Subject: [PATCH 1/2] feat(linalg): enable lowering/decomposing scalable pack
ops
Enables lowering/decomposing `linalg.pack` ops with dynamic inner tiles to a sequence of `tensor.pad` -> `tensor.expand_shape` -> `linalg.transpose` ops.
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
.../Dialect/Linalg/Transforms/Transforms.cpp | 16 ++--
.../Dialect/Linalg/transform-lower-pack.mlir | 80 +++++++++++++++++++
2 files changed, 89 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 260e36fb47f04..aa1be54a93ed4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -225,14 +225,8 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
if (!packOp.hasPureTensorSemantics())
return failure();
- // 1. Filter out NYI cases.
auto packedTensorType =
cast<RankedTensorType>(packOp->getResultTypes().front());
- if (llvm::any_of(packOp.getStaticInnerTiles(), ShapedType::isDynamic)) {
- return rewriter.notifyMatchFailure(
- packOp,
- "non-static shape NYI, needs a more powerful tensor.expand_shape op");
- }
Location loc = packOp->getLoc();
OpBuilder::InsertionGuard g(rewriter);
@@ -249,6 +243,14 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
SmallVector<int64_t> stripMinedShape(packedTensorType.getShape());
applyPermutationToVector(stripMinedShape, packedToStripMinedShapePerm);
+ // Also compute the mixed (static+dynamic) strip-mined sizes for the
+ // expand_shape output. This is needed to support dynamic inner tile sizes,
+ // since the 3-arg ExpandShapeOp builder cannot infer shapes when multiple
+ // dynamic dims appear in a single reassociation group.
+ SmallVector<OpFoldResult> stripMinedMixedSizes =
+ tensor::getMixedSizes(rewriter, loc, packOp.getDest());
+ applyPermutationToVector(stripMinedMixedSizes, packedToStripMinedShapePerm);
+
// 4. Pad the source of packOp to a shape we can expand into stripMinedShape.
SmallVector<OpFoldResult> lows(packOp.getSourceRank(),
rewriter.getIndexAttr(0));
@@ -331,7 +333,7 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
RankedTensorType::Builder(packedTensorType).setShape(stripMinedShape);
auto reshapeOp = tensor::ExpandShapeOp::create(
rewriter, loc, expandShapeResultType, padOp.getResult(),
- packingMetadata.reassociations);
+ packingMetadata.reassociations, stripMinedMixedSizes);
// 6. Transpose stripMinedShape to packedShape.
SmallVector<int64_t> transpPerm =
diff --git a/mlir/test/Dialect/Linalg/transform-lower-pack.mlir b/mlir/test/Dialect/Linalg/transform-lower-pack.mlir
index b6fe67a9ae1f3..8f5a810ff428f 100644
--- a/mlir/test/Dialect/Linalg/transform-lower-pack.mlir
+++ b/mlir/test/Dialect/Linalg/transform-lower-pack.mlir
@@ -692,6 +692,86 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1 - 256)>
+// CHECK: func.func @pack_dynamic_inner_tile(
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<256x128xf32>,
+// CHECK-SAME: %[[TILE_N:[a-zA-Z0-9]+]]: index
+func.func @pack_dynamic_inner_tile(%src: tensor<256x128xf32>, %tile_n: index) -> tensor<?x128x?xf32> {
+ // CHECK: %[[TILED_D0:.+]] = arith.ceildivui %{{.+}}, %[[TILE_N]] : index
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[TILED_D0]], %[[TILE_N]]) : tensor<?x128x?xf32>
+ // CHECK: %[[DIM0:.+]] = tensor.dim %[[EMPTY]], %{{.+}}
+ // CHECK: %[[DIM2:.+]] = tensor.dim %[[EMPTY]], %{{.+}}
+ // CHECK: %[[HIGH:.+]] = affine.apply #[[MAP]]()[%[[TILE_N]], %[[DIM0]]]
+ // CHECK: %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[%[[HIGH]], 0]
+ // CHECK: : tensor<256x128xf32> to tensor<?x128xf32>
+ // CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]] {{\[}}[0, 1], [2]]
+ // CHECK-SAME: output_shape [%[[DIM0]], %[[DIM2]], 128]
+ // CHECK-SAME: : tensor<?x128xf32> into tensor<?x?x128xf32>
+ // CHECK: %[[TRANSP:.+]] = linalg.transpose
+ // CHECK-SAME: ins(%[[EXPAND]] : tensor<?x?x128xf32>)
+ // CHECK-SAME: outs(%[[EMPTY]] : tensor<?x128x?xf32>)
+ // CHECK-SAME: permutation = [0, 2, 1]
+ // CHECK: return %[[TRANSP]]
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f32
+ %c256 = arith.constant 256 : index
+ %tiled_d0 = arith.ceildivui %c256, %tile_n : index
+ %init_pack = tensor.empty(%tiled_d0, %tile_n) : tensor<?x128x?xf32>
+ %pack = linalg.pack %src padding_value(%cst : f32)
+ inner_dims_pos = [0] inner_tiles = [%tile_n] into %init_pack
+ : tensor<256x128xf32> -> tensor<?x128x?xf32>
+ return %pack : tensor<?x128x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %module_op
+ : (!transform.any_op) -> !transform.op<"linalg.pack">
+ transform.structured.lower_pack %pack : (!transform.op<"linalg.pack">)
+ -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_dynamic_inner_tile_with_outer_dims_perm(
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<64x128xf32>,
+// CHECK-SAME: %[[TILE_N:[a-zA-Z0-9]+]]: index
+func.func @pack_dynamic_inner_tile_with_outer_dims_perm(%src: tensor<64x128xf32>, %tile_n: index) -> tensor<?x64x?xf32> {
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%{{.+}}, %[[TILE_N]]) : tensor<?x64x?xf32>
+ // CHECK: %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[0, %{{.+}}]
+ // CHECK: : tensor<64x128xf32> to tensor<64x?xf32>
+ // CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]] {{\[}}[0], [1, 2]]
+ // CHECK-SAME: output_shape [64, %{{.+}}, %{{.+}}]
+ // CHECK-SAME: : tensor<64x?xf32> into tensor<64x?x?xf32>
+ // CHECK: %[[TRANSP:.+]] = linalg.transpose
+ // CHECK-SAME: ins(%[[EXPAND]] : tensor<64x?x?xf32>)
+ // CHECK-SAME: outs(%[[EMPTY]] : tensor<?x64x?xf32>)
+ // CHECK-SAME: permutation = [1, 0, 2]
+ // CHECK: return %[[TRANSP]]
+ %cst = arith.constant 0.0 : f32
+ %c128 = arith.constant 128 : index
+ %tiled_d1 = arith.ceildivui %c128, %tile_n : index
+ %init_pack = tensor.empty(%tiled_d1, %tile_n) : tensor<?x64x?xf32>
+ %pack = linalg.pack %src padding_value(%cst : f32)
+ outer_dims_perm = [1, 0] inner_dims_pos = [1] inner_tiles = [%tile_n] into %init_pack
+ : tensor<64x128xf32> -> tensor<?x64x?xf32>
+ return %pack : tensor<?x64x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+ %pack = transform.structured.match ops{["linalg.pack"]} in %module_op
+ : (!transform.any_op) -> !transform.op<"linalg.pack">
+ transform.structured.lower_pack %pack : (!transform.op<"linalg.pack">)
+ -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: @unpack_with_outer_dims_perm
// CHECK-SAME: %[[ARG0:.*]]: tensor<32x64xf32>, %[[ARG1:.*]]: tensor<2x4x32x8xf32>
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8x2x32xf32>
>From dc4fcec560c71cf519eec752953df9d8d9b64ed4 Mon Sep 17 00:00:00 2001
From: Ege Beysel <beyselege at gmail.com>
Date: Mon, 1 Jun 2026 14:21:54 +0200
Subject: [PATCH 2/2] address comment
Signed-off-by: Ege Beysel <beyselege at gmail.com>
---
mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index aa1be54a93ed4..f2d491ac442e6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -245,8 +245,9 @@ FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
// Also compute the mixed (static+dynamic) strip-mined sizes for the
// expand_shape output. This is needed to support dynamic inner tile sizes,
- // since the 3-arg ExpandShapeOp builder cannot infer shapes when multiple
- // dynamic dims appear in a single reassociation group.
+ // since the shapes cannot be inferred automatically when multiple dynamic
+ // dims appear in a single reassociation group during ExpandShapeOp
+ // construction.
SmallVector<OpFoldResult> stripMinedMixedSizes =
tensor::getMixedSizes(rewriter, loc, packOp.getDest());
applyPermutationToVector(stripMinedMixedSizes, packedToStripMinedShapePerm);
More information about the Mlir-commits
mailing list