[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 03:04:36 PDT 2026
https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/221430
## Problem
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.
## Approach
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.
## Scope
This addresses the `transform.structured.pack` case reported in #220505. The other three transformations identified in that issue are not changed.
## Testing
- `mlir/test/Dialect/Linalg`: 174/174 passed
- `git clang-format --diff`: clean
Partially addresses #220505.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01LHzHeLy5ALk92REMAnY3kb
>From 5460530c68436a58bf87a32696c33abd5e0d80ca 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.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
Claude-Session: https://claude.ai/code/session_01LHzHeLy5ALk92REMAnY3kb
---
.../Dialect/Linalg/Transforms/Transforms.cpp | 42 +++++++++++++++++++
.../Dialect/Linalg/transform-op-pack.mlir | 36 ++++++++++++++++
2 files changed, 78 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..0c9005193a8fa 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pack.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
@@ -686,3 +686,39 @@ 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
+ }
+}
More information about the Mlir-commits
mailing list