[Mlir-commits] [mlir] 948b2b5 - [mlir][Linalg] Enable lowering/decomposing scalable pack ops (#200216)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 1 12:08:47 PDT 2026
Author: Ege Beysel
Date: 2026-06-01T21:08:41+02:00
New Revision: 948b2b594af8101f5062c66aa8e23334fed4e2c0
URL: https://github.com/llvm/llvm-project/commit/948b2b594af8101f5062c66aa8e23334fed4e2c0
DIFF: https://github.com/llvm/llvm-project/commit/948b2b594af8101f5062c66aa8e23334fed4e2c0.diff
LOG: [mlir][Linalg] Enable lowering/decomposing scalable pack ops (#200216)
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>
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
mlir/test/Dialect/Linalg/transform-lower-pack.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 260e36fb47f04..f2d491ac442e6 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,15 @@ 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 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);
+
// 4. Pad the source of packOp to a shape we can expand into stripMinedShape.
SmallVector<OpFoldResult> lows(packOp.getSourceRank(),
rewriter.getIndexAttr(0));
@@ -331,7 +334,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>
More information about the Mlir-commits
mailing list