[Mlir-commits] [mlir] [mlir][linalg] Remap linalg.index when packing a structured op (PR #221430)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Sep 5 04:51:57 PDT 2026
https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/221430
>From 0c0ef86af30dd523146c213e740c4c3f52295bbc Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Sat, 5 Sep 2026 11:56:54 +0200
Subject: [PATCH] [mlir][linalg] Remap linalg.index when packing a structured
op
Packing a structured operation splits an original loop dimension `d` into
outer and inner dimensions. The operation's region is moved to the packed
`linalg.generic` unchanged, so `linalg.index d` starts returning only the
outer tile coordinate. The resulting IR verifies, but silently computes
incorrect index values.
Track the inner loop dimension introduced for each packed loop and rewrite
`linalg.index d` to `outer * tileSize + inner`. Unpacked dimensions,
including dimensions whose pack size is zero, remain unchanged.
The fix is implemented in the common `linalg::pack` utility and therefore
also applies to its greedy and block-matmul packing callers.
This addresses the `transform.structured.pack` case reported in #220505.
The other three transformations identified in that issue are not changed.
Tested with `build/bin/llvm-lit mlir/test/Dialect/Linalg` (174/174 passed).
Partially addresses #220505.
---
.../Dialect/Linalg/Transforms/Transforms.cpp | 42 +++++
.../Dialect/Linalg/transform-op-pack.mlir | 167 ++++++++++++++++++
.../Linalg/transform-pack-greedily.mlir | 50 ++++++
3 files changed, 259 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index aa021cbad24a7..8480e655601d0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -218,6 +218,44 @@ struct PackedOperandsDimList {
} // namespace
+/// Rewrite index operations in a packed Linalg op to preserve their values in
+/// the original iteration space. Packing loop dimension `d` by `tileSize`
+/// splits its index into an outer and inner index, so the original index is
+/// `outer * tileSize + inner`.
+static void
+remapPackedIndexOps(RewriterBase &rewriter, LinalgOp packedLinalgOp,
+ ArrayRef<OpFoldResult> packedSizes,
+ ArrayRef<std::optional<int64_t>> packedLoopToInnerLoop) {
+ if (!packedLinalgOp.hasIndexSemantics())
+ return;
+
+ SmallVector<IndexOp> indexOps =
+ llvm::to_vector(packedLinalgOp.getBlock()->getOps<IndexOp>());
+ for (IndexOp indexOp : indexOps) {
+ unsigned dim = indexOp.getDim();
+ assert(dim < packedLoopToInnerLoop.size() && "invalid linalg.index dim");
+ std::optional<int64_t> innerDim = packedLoopToInnerLoop[dim];
+ if (!innerDim)
+ continue;
+
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(indexOp);
+ Value outerIndex =
+ IndexOp::create(rewriter, indexOp.getLoc(), dim).getResult();
+ Value innerIndex =
+ IndexOp::create(rewriter, indexOp.getLoc(), *innerDim).getResult();
+
+ AffineExpr outer, inner, tileSize;
+ bindDims(rewriter.getContext(), outer, inner);
+ bindSymbols(rewriter.getContext(), tileSize);
+ OpFoldResult originalIndex = affine::makeComposedFoldedAffineApply(
+ rewriter, indexOp.getLoc(), outer * tileSize + inner,
+ {outerIndex, innerIndex, packedSizes[dim]});
+ rewriter.replaceOp(indexOp, getValueOrCreateConstantIndexOp(
+ rewriter, indexOp.getLoc(), originalIndex));
+ }
+}
+
FailureOr<LowerPackResult> linalg::lowerPack(RewriterBase &rewriter,
linalg::PackOp packOp,
bool lowerPadLikeWithInsertSlice) {
@@ -501,6 +539,7 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
SmallVector<linalg::UnPackOp> unPackOps;
// Step 1. Pack each dim of the LinalgOp metadata by packedSizes[i].
PackedOperandsDimList listOfPackedOperandsDim;
+ SmallVector<std::optional<int64_t>> packedLoopToInnerLoop(packedSizes.size());
for (int64_t i = 0, e = packedSizes.size(); i < e; ++i) {
std::optional<int64_t> maybeConstant = getConstantIntValue(packedSizes[i]);
// Skip tile sizes explicitly set to 0.
@@ -514,6 +553,7 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
packLinalgMetadataOnce(indexingMaps, iteratorTypes, i);
if (failed(maybePackedDimForEachOperand))
return failure();
+ packedLoopToInnerLoop[i] = iteratorTypes.size() - 1;
packedOperandsDims.packedDimForEachOperand = *maybePackedDimForEachOperand;
LDBG() << "++++ After pack size #" << i << ": " << packedSizes[i];
@@ -582,6 +622,8 @@ FailureOr<PackResult> linalg::pack(RewriterBase &rewriter,
linalg::GenericOp::create(rewriter, linalgOp.getLoc(), inits.getTypes(),
inputs, inits, indexingMaps, iteratorTypes);
packedLinalgOp.getRegion().takeBody(linalgOp->getRegion(0));
+ remapPackedIndexOps(rewriter, packedLinalgOp, packedSizes,
+ packedLoopToInnerLoop);
// Step 4. Propagate packing to all the op results.
for (OpResult result : packedLinalgOp->getResults()) {
diff --git a/mlir/test/Dialect/Linalg/transform-op-pack.mlir b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
index 13a07cc38e10b..03de45cc376e0 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pack.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
@@ -686,3 +686,170 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// CHECK-DAG: #[[$PACK_INDEX_MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 4 + s1)>
+// CHECK-LABEL: func.func @pack_linalg_index
+// CHECK: linalg.generic
+// CHECK: ^bb0
+// CHECK-DAG: %[[OUTER:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[INNER:.*]] = linalg.index 1 : index
+// CHECK: %[[INDEX:.*]] = affine.apply #[[$PACK_INDEX_MAP]]()[%[[OUTER]], %[[INNER]]]
+// CHECK: arith.index_cast %[[INDEX]] : index to i32
+func.func @pack_linalg_index(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32> {
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%input : tensor<8xi32>) outs(%init : tensor<8xi32>) {
+ ^bb0(%in: i32, %out: i32):
+ %index = linalg.index 0 : index
+ %index_i32 = arith.index_cast %index : index to i32
+ %sum = arith.addi %in, %index_i32 : i32
+ linalg.yield %sum : i32
+ } -> tensor<8xi32>
+ return %result : tensor<8xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.pack %generic packed_sizes = [4]
+ : (!transform.any_op) -> (!transform.op<"linalg.generic">)
+ transform.yield
+ }
+}
+
+// -----
+
+#identity = affine_map<(d0) -> (d0)>
+
+// Check index remapping when packing adds a padded iteration. Unpacking drops
+// the extra output element at index 8.
+// CHECK-DAG: #[[$PADDED_PACK_INDEX:.*]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
+// CHECK-LABEL: func.func @pack_linalg_index_padded(
+// CHECK: linalg.pack {{.*}} padding_value{{.*}} : tensor<8xi32> -> tensor<3x3xi32>
+// CHECK: %[[PACKED:.*]] = linalg.generic
+// CHECK: ^bb0
+// CHECK-DAG: %[[OUTER:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[INNER:.*]] = linalg.index 1 : index
+// CHECK: %[[INDEX:.*]] = affine.apply #[[$PADDED_PACK_INDEX]]()[%[[OUTER]], %[[INNER]]]
+// CHECK: arith.index_cast %[[INDEX]] : index to i32
+// CHECK: linalg.unpack %[[PACKED]] {{.*}} : tensor<3x3xi32> -> tensor<8xi32>
+func.func @pack_linalg_index_padded(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32> {
+ %result = linalg.generic {
+ indexing_maps = [#identity, #identity],
+ iterator_types = ["parallel"]}
+ ins(%input : tensor<8xi32>) outs(%init : tensor<8xi32>) {
+ ^bb0(%in: i32, %out: i32):
+ %index = linalg.index 0 : index
+ %index_i32 = arith.index_cast %index : index to i32
+ %sum = arith.addi %in, %index_i32 : i32
+ linalg.yield %sum : i32
+ } -> tensor<8xi32>
+ return %result : tensor<8xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.pack %generic packed_sizes = [3]
+ : (!transform.any_op) -> (!transform.op<"linalg.generic">)
+ transform.yield
+ }
+}
+
+// -----
+
+#identity_3d = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK-DAG: #[[$PACK_INDEX_4:.*]] = affine_map<()[s0, s1] -> (s0 * 4 + s1)>
+// CHECK-DAG: #[[$PACK_INDEX_8:.*]] = affine_map<()[s0, s1] -> (s0 * 8 + s1)>
+// CHECK-LABEL: func.func @pack_multiple_linalg_indices
+// CHECK: linalg.generic
+// CHECK: ^bb0
+// CHECK: %[[OUTER_0:.*]] = linalg.index 0 : index
+// CHECK: %[[INNER_0:.*]] = linalg.index 3 : index
+// CHECK: %[[INDEX_0:.*]] = affine.apply #[[$PACK_INDEX_4]]()[%[[OUTER_0]], %[[INNER_0]]]
+// CHECK: %[[INDEX_1:.*]] = linalg.index 1 : index
+// CHECK: %[[OUTER_2:.*]] = linalg.index 2 : index
+// CHECK: %[[INNER_2:.*]] = linalg.index 4 : index
+// CHECK: %[[INDEX_2:.*]] = affine.apply #[[$PACK_INDEX_8]]()[%[[OUTER_2]], %[[INNER_2]]]
+// CHECK: arith.index_cast %[[INDEX_0]] : index to i32
+// CHECK: arith.index_cast %[[INDEX_1]] : index to i32
+// CHECK: arith.index_cast %[[INDEX_2]] : index to i32
+func.func @pack_multiple_linalg_indices(%input: tensor<8x5x16xi32>,
+ %init: tensor<8x5x16xi32>) -> tensor<8x5x16xi32> {
+ %result = linalg.generic {
+ indexing_maps = [#identity_3d, #identity_3d],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%input : tensor<8x5x16xi32>) outs(%init : tensor<8x5x16xi32>) {
+ ^bb0(%in: i32, %out: i32):
+ %index_0 = linalg.index 0 : index
+ %index_1 = linalg.index 1 : index
+ %index_2 = linalg.index 2 : index
+ %index_0_i32 = arith.index_cast %index_0 : index to i32
+ %index_1_i32 = arith.index_cast %index_1 : index to i32
+ %index_2_i32 = arith.index_cast %index_2 : index to i32
+ %sum_0 = arith.addi %in, %index_0_i32 : i32
+ %sum_1 = arith.addi %sum_0, %index_1_i32 : i32
+ %sum_2 = arith.addi %sum_1, %index_2_i32 : i32
+ linalg.yield %sum_2 : i32
+ } -> tensor<8x5x16xi32>
+ return %result : tensor<8x5x16xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.pack %generic packed_sizes = [4, 0, 8]
+ : (!transform.any_op) -> (!transform.op<"linalg.generic">)
+ transform.yield
+ }
+}
+
+// -----
+
+#identity_dynamic = affine_map<(d0) -> (d0)>
+
+// CHECK-DAG: #[[$DYNAMIC_PACK_INDEX:.*]] = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>
+// CHECK-LABEL: func.func @pack_linalg_index_dynamic
+// CHECK: %[[TILE_SIZE:.*]] = "tile_size"() : () -> index
+// CHECK: linalg.generic
+// CHECK: ^bb0
+// CHECK: %[[OUTER:.*]] = linalg.index 0 : index
+// CHECK: %[[INNER:.*]] = linalg.index 1 : index
+// CHECK: %[[INDEX:.*]] = affine.apply #[[$DYNAMIC_PACK_INDEX]]()[%[[TILE_SIZE]], %[[OUTER]], %[[INNER]]]
+// CHECK: arith.index_cast %[[INDEX]] : index to i32
+func.func @pack_linalg_index_dynamic(%input: tensor<?xi32>,
+ %init: tensor<?xi32>) -> tensor<?xi32> {
+ %tile_size = "tile_size"() : () -> index
+ %result = linalg.generic {
+ indexing_maps = [#identity_dynamic, #identity_dynamic],
+ iterator_types = ["parallel"]}
+ ins(%input : tensor<?xi32>) outs(%init : tensor<?xi32>) {
+ ^bb0(%in: i32, %out: i32):
+ %index = linalg.index 0 : index
+ %index_i32 = arith.index_cast %index : index to i32
+ %sum = arith.addi %in, %index_i32 : i32
+ linalg.yield %sum : i32
+ } -> tensor<?xi32>
+ return %result : tensor<?xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ %tile_size = transform.structured.match ops{["tile_size"]} in %arg0
+ : (!transform.any_op) -> !transform.any_op
+ transform.structured.pack %generic packed_sizes = [%tile_size]
+ : (!transform.any_op, !transform.any_op) -> (!transform.op<"linalg.generic">)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir b/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
index 5812c4db88247..f4ae7fa254bf9 100644
--- a/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
+++ b/mlir/test/Dialect/Linalg/transform-pack-greedily.mlir
@@ -402,3 +402,53 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Interchanging (m, n, k) to (k, m, n) maps the original index 2 to index 0.
+// Packing must then reconstruct k as outer_k * 32 + inner_k. All dimensions
+// are divisible by their tiles: the extra +k term is not neutral under padding.
+// CHECK-DAG: #[[$GREEDY_PACK_K:.*]] = affine_map<()[s0, s1] -> (s0 * 32 + s1)>
+// CHECK-DAG: #[[$GREEDY_INTERCHANGE_K:.*]] = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK-LABEL: func.func @pack_greedily_linalg_index(
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]
+// CHECK-SAME: ins(%{{.*}} : tensor<2x2x32x8xf32>, tensor<2x2x32x16xf32>)
+// CHECK-SAME: outs(%{{.*}} : tensor<2x2x8x16xf32>)
+// CHECK: ^bb0
+// CHECK: %[[K_OUTER:.*]] = linalg.index 0 : index
+// CHECK-NEXT: %[[K_INNER:.*]] = linalg.index 3 : index
+// CHECK-NEXT: %[[PACKED_K:.*]] = affine.apply #[[$GREEDY_PACK_K]]()[%[[K_OUTER]], %[[K_INNER]]]
+// CHECK: %[[K:.*]] = affine.apply #[[$GREEDY_INTERCHANGE_K]](%[[PACKED_K]], %{{.*}}, %{{.*}})
+// CHECK-NEXT: arith.index_cast %[[K]] : index to i32
+func.func @pack_greedily_linalg_index(%A: tensor<16x64xf32>, %B: tensor<32x64xf32>,
+ %C: tensor<32x16xf32>) -> tensor<32x16xf32> {
+ %result = linalg.generic {
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n, m)>],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : tensor<16x64xf32>, tensor<32x64xf32>)
+ outs(%C : tensor<32x16xf32>) {
+ ^bb0(%a: f32, %b: f32, %c: f32):
+ %k = linalg.index 2 : index
+ %k_i32 = arith.index_cast %k : index to i32
+ %k_f32 = arith.sitofp %k_i32 : i32 to f32
+ %product = arith.mulf %a, %b : f32
+ %sum = arith.addf %c, %product : f32
+ %indexed_sum = arith.addf %sum, %k_f32 : f32
+ linalg.yield %indexed_sum : f32
+ } -> tensor<32x16xf32>
+ return %result : tensor<32x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+ : (!transform.any_op) -> !transform.op<"linalg.generic">
+ transform.structured.pack_greedily %generic
+ matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
+ : (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list