[Mlir-commits] [mlir] [mlir][linalg] Remap linalg.index in splitReduction (PR #223281)
Nitin John Raj
llvmlistbot at llvm.org
Sun Sep 13 14:44:19 PDT 2026
https://github.com/njayinthehouse created https://github.com/llvm/llvm-project/pull/223281
split_reduction now rewrites the indices for `linalg.index` ops in the loop body. Prior to this change, the rewritten IR passed the verifier but computed the wrong index values.
Specifically:
1. If the index dimension is the reduction dimension, we calculate the new index from the new parallel and reduction indices.
2. If the index is at or higher than the insertion dimension and not the reduction dimension, we shift it up by one
This partially addressed issue #220505
Assisted-by: Claude Opus 5
>From 39f43c1bf256b8af58fcf751933794168bfd1b3e Mon Sep 17 00:00:00 2001
From: Nitin John Raj <nitinjohnraj at gmail.com>
Date: Sat, 12 Sep 2026 00:41:08 -0700
Subject: [PATCH] [mlir][linalg] Remap linalg.index in splitReduction
splitReduction and splitReductionByScaling move the body of the original
op into a new generic op with an extra parallel loop, but left the
`linalg.index` ops in the body untouched. The resulting IR verifies but
computes wrong index values: loops after the inserted one are off by one
position, and the split reduction loop `k` now reads only one of the two
loops it was split into.
Rewrite the `linalg.index` ops after inlining the body: shift the loops
at or after the inserted loop up by one, and rebuild `k` as
`outer * innerSize + inner` from the new parallel and reduction loops.
Part of #220505.
Assisted-by: Claude Opus 5
---
.../Linalg/Transforms/SplitReduction.cpp | 64 +++++++++
...ansform-op-split-reduction-by-scaling.mlir | 53 ++++++-
.../Linalg/transform-op-split-reduction.mlir | 131 ++++++++++++++++++
3 files changed, 247 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp b/mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
index b8f862001bccf..dd626fc83ec49 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
@@ -15,7 +15,9 @@
#include <utility>
#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
@@ -27,6 +29,54 @@
using namespace mlir;
using namespace mlir::linalg;
+/// Rewrite the `linalg.index` ops of `genericOp`, whose body was moved from
+/// the original op, to the loops of `genericOp`. Splitting inserted a new
+/// parallel loop at `insertedDim`, so every original loop `d >= insertedDim`
+/// is now loop `d + 1`, and replaced the original reduction loop
+/// `reductionDim` by the pair (new parallel loop, remaining reduction loop),
+/// with the parallel loop outermost iff `parallelIsOuter` and the inner loop
+/// of extent `innerSize`. The original reduction index is
+/// `outerIndex * innerSize + innerIndex`.
+static void remapSplitReductionIndexOps(RewriterBase &b, LinalgOp genericOp,
+ unsigned reductionDim,
+ unsigned insertedDim,
+ bool parallelIsOuter,
+ int64_t innerSize) {
+ if (!genericOp.hasIndexSemantics())
+ return;
+
+ unsigned newReductionDim =
+ reductionDim < insertedDim ? reductionDim : reductionDim + 1;
+ unsigned outerDim = parallelIsOuter ? insertedDim : newReductionDim;
+ unsigned innerDim = parallelIsOuter ? newReductionDim : insertedDim;
+
+ // Collect first: the loop inserts new `linalg.index` ops into this block.
+ SmallVector<IndexOp> indexOps =
+ llvm::to_vector(genericOp.getBlock()->getOps<IndexOp>());
+ for (IndexOp indexOp : indexOps) {
+ unsigned dim = indexOp.getDim();
+ if (dim < insertedDim && dim != reductionDim)
+ continue;
+
+ OpBuilder::InsertionGuard guard(b);
+ b.setInsertionPoint(indexOp);
+ Location loc = indexOp.getLoc();
+ if (dim != reductionDim) {
+ // A loop was inserted at `insertedDim`, so later loops shift up by one.
+ b.replaceOp(indexOp, IndexOp::create(b, loc, dim + 1).getResult());
+ continue;
+ }
+
+ Value outerIndex = IndexOp::create(b, loc, outerDim).getResult();
+ Value innerIndex = IndexOp::create(b, loc, innerDim).getResult();
+ AffineExpr d0, d1;
+ bindDims(b.getContext(), d0, d1);
+ OpFoldResult newIndex = affine::makeComposedFoldedAffineApply(
+ b, loc, d0 * innerSize + d1, {outerIndex, innerIndex});
+ b.replaceOp(indexOp, getValueOrCreateConstantIndexOp(b, loc, newIndex));
+ }
+}
+
FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
RewriterBase &b, LinalgOp op,
const ControlSplitReductionFn &controlSplitReductionFn, bool useAlloc) {
@@ -174,6 +224,13 @@ FailureOr<SplitReductionResult> mlir::linalg::splitReduction(
b.inlineRegionBefore(op->getRegion(0), genericOp.getRegion(),
genericOp.getRegion().begin());
+ // The body still refers to the loops of `op`; rewrite its `linalg.index`
+ // ops to the loops of `genericOp`.
+ remapSplitReductionIndexOps(
+ b, genericOp, reductionDim, insertSplitDimension,
+ /*parallelIsOuter=*/!control.innerParallel,
+ /*innerSize=*/control.innerParallel ? ratio : reductionDimSize / ratio);
+
// Then create a new reduction that only reduce the newly added dimension
// from the previous op.
unsigned intermRank = newOutputShape.size();
@@ -372,6 +429,13 @@ FailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
genericOp.getRegion().front().insertArgument(reductionDimPos,
b.getIntegerType(1), loc);
+ // The body still refers to the loops of `op`; rewrite its `linalg.index`
+ // ops to the loops of `genericOp`.
+ remapSplitReductionIndexOps(b, genericOp, reductionDimPos,
+ /*insertedDim=*/reductionDimPos,
+ /*parallelIsOuter=*/true,
+ /*innerSize=*/splitFactor);
+
// Step 5. Create new reduction ops that only reduce the newly added
// dimensions from the previous op.
// For now assume outputs are 1-1 with reduction ops.
diff --git a/mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir b/mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir
index b21a4a78369e4..495d160132bcc 100644
--- a/mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt --transform-interpreter %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --transform-interpreter %s | FileCheck %s
// CHECK-LABEL: func.func @matmul_split
func.func @matmul_split(%A : tensor<?x256xf32>, %B: tensor<256x32xf32>, %C: tensor<?x32xf32>) -> tensor<?x32xf32> {
@@ -27,3 +27,54 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// The reduction loop `k` (position 2) is split into a parallel loop at
+// position 2 and the remaining reduction loop at position 3, so the parallel
+// loop `l` shifts from position 3 to 4. In the body, `linalg.index 2` (`k`)
+// must become `p * 4 + r`, `linalg.index 3` (`l`) must become
+// `linalg.index 4`, and `linalg.index 0` (`i`) must stay as it is.
+func.func @generic_split_index(%A: tensor<2x8xi32>, %B: tensor<8x2x3xi32>, %C: tensor<2x2x3xi32>) -> tensor<2x2x3xi32> {
+ %r = linalg.generic {indexing_maps = [affine_map<(i, j, k, l) -> (i, k)>,
+ affine_map<(i, j, k, l) -> (k, j, l)>,
+ affine_map<(i, j, k, l) -> (i, j, l)>],
+ iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
+ ins(%A, %B : tensor<2x8xi32>, tensor<8x2x3xi32>) outs(%C : tensor<2x2x3xi32>) {
+ ^bb0(%a: i32, %b: i32, %acc: i32):
+ %i = linalg.index 0 : index
+ %k = linalg.index 2 : index
+ %l = linalg.index 3 : index
+ %ik = arith.addi %i, %k : index
+ %ikl = arith.addi %ik, %l : index
+ %c = arith.index_cast %ikl : index to i32
+ %s = arith.addi %acc, %c : i32
+ linalg.yield %s : i32
+ } -> tensor<2x2x3xi32>
+ return %r : tensor<2x2x3xi32>
+}
+
+// CHECK-DAG: #[[$MAP_K:.*]] = affine_map<()[s0, s1] -> (s0 * 4 + s1)>
+// CHECK-LABEL: func.func @generic_split_index
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "parallel"]
+// CHECK: ^bb0
+// CHECK-DAG: %[[I:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[P:.*]] = linalg.index 2 : index
+// CHECK-DAG: %[[R:.*]] = linalg.index 3 : index
+// CHECK: %[[K:.*]] = affine.apply #[[$MAP_K]]()[%[[P]], %[[R]]]
+// CHECK: %[[L:.*]] = linalg.index 4 : index
+// CHECK: %[[IK:.*]] = arith.addi %[[I]], %[[K]] : index
+// CHECK: arith.addi %[[IK]], %[[L]] : index
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel"]
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:4 = transform.structured.split_reduction %0
+ <split_factor = 4, insert_split_dimension = 2, use_scaling_algorithm>
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir b/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
index 82111f911aa24..44ded93f137f2 100644
--- a/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir
@@ -499,3 +499,134 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// The reduction loop `k` (position 1) is split into a parallel loop inserted
+// at position 0 and the remaining reduction loop, now at position 2. In the
+// body, `linalg.index 1` (`k`) must become `p * 2 + r` and `linalg.index 0`
+// (`i`) must shift to `linalg.index 1`.
+func.func @generic_split_index_shifted(%in: tensor<2x8xi32>, %out: tensor<2xi32>) -> tensor<2xi32> {
+ %r = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%in : tensor<2x8xi32>) outs(%out : tensor<2xi32>) {
+ ^bb0(%a: i32, %acc: i32):
+ %i = linalg.index 0 : index
+ %k = linalg.index 1 : index
+ %ik = arith.addi %i, %k : index
+ %c = arith.index_cast %ik : index to i32
+ %s = arith.addi %acc, %c : i32
+ linalg.yield %s : i32
+ } -> tensor<2xi32>
+ return %r : tensor<2xi32>
+}
+
+// CHECK-DAG: #[[$MAP_K:.*]] = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
+// CHECK-LABEL: func @generic_split_index_shifted
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK: ^bb0
+// CHECK-DAG: %[[I:.*]] = linalg.index 1 : index
+// CHECK-DAG: %[[P:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[R:.*]] = linalg.index 2 : index
+// CHECK: %[[K:.*]] = affine.apply #[[$MAP_K]]()[%[[P]], %[[R]]]
+// CHECK: arith.addi %[[I]], %[[K]] : index
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["reduction", "parallel"]
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:4 = transform.structured.split_reduction %0 <split_factor = 4, insert_split_dimension = 0>
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Same op, but the parallel loop is inserted at position 1, after `i`. Then
+// `linalg.index 0` (`i`) must stay as it is, and `k` becomes `p * 2 + r` with
+// `p` at position 1 and `r` at position 2.
+func.func @generic_split_index_unshifted(%in: tensor<2x8xi32>, %out: tensor<2xi32>) -> tensor<2xi32> {
+ %r = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%in : tensor<2x8xi32>) outs(%out : tensor<2xi32>) {
+ ^bb0(%a: i32, %acc: i32):
+ %i = linalg.index 0 : index
+ %k = linalg.index 1 : index
+ %ik = arith.addi %i, %k : index
+ %c = arith.index_cast %ik : index to i32
+ %s = arith.addi %acc, %c : i32
+ linalg.yield %s : i32
+ } -> tensor<2xi32>
+ return %r : tensor<2xi32>
+}
+
+// CHECK-DAG: #[[$MAP_K:.*]] = affine_map<()[s0, s1] -> (s0 * 2 + s1)>
+// CHECK-LABEL: func @generic_split_index_unshifted
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK: ^bb0
+// CHECK-DAG: %[[I:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[P:.*]] = linalg.index 1 : index
+// CHECK-DAG: %[[R:.*]] = linalg.index 2 : index
+// CHECK: %[[K:.*]] = affine.apply #[[$MAP_K]]()[%[[P]], %[[R]]]
+// CHECK: arith.addi %[[I]], %[[K]] : index
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "reduction"]
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:4 = transform.structured.split_reduction %0 <split_factor = 4, insert_split_dimension = 1>
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// With `inner_parallel`, the reduction loop `k` (position 1) keeps its position
+// with extent 8 / 4 = 2 and the parallel loop of extent 4 is inserted right
+// after it, at position 2. In the body, `linalg.index 1` (`k`) must become
+// `r * 4 + p` and `linalg.index 0` (`i`) must stay as it is.
+func.func @generic_split_index_inner_parallel(%in: tensor<2x8xi32>, %out: tensor<2xi32>) -> tensor<2xi32> {
+ %r = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%in : tensor<2x8xi32>) outs(%out : tensor<2xi32>) {
+ ^bb0(%a: i32, %acc: i32):
+ %i = linalg.index 0 : index
+ %k = linalg.index 1 : index
+ %ik = arith.addi %i, %k : index
+ %c = arith.index_cast %ik : index to i32
+ %s = arith.addi %acc, %c : i32
+ linalg.yield %s : i32
+ } -> tensor<2xi32>
+ return %r : tensor<2xi32>
+}
+
+// CHECK-DAG: #[[$MAP_K:.*]] = affine_map<()[s0, s1] -> (s0 * 4 + s1)>
+// CHECK-LABEL: func @generic_split_index_inner_parallel
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"]
+// CHECK: ^bb0
+// CHECK-DAG: %[[I:.*]] = linalg.index 0 : index
+// CHECK-DAG: %[[R:.*]] = linalg.index 1 : index
+// CHECK-DAG: %[[P:.*]] = linalg.index 2 : index
+// CHECK: %[[K:.*]] = affine.apply #[[$MAP_K]]()[%[[R]], %[[P]]]
+// CHECK: arith.addi %[[I]], %[[K]] : index
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "reduction"]
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1:4 = transform.structured.split_reduction %0 <split_factor = 4, insert_split_dimension = 1, inner_parallel>
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list