[Mlir-commits] [mlir] [mlir][linalg] Remap linalg.index in splitReduction (PR #223281)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Sep 13 14:45:16 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Nitin John Raj (njayinthehouse)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/223281.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp (+64)
- (modified) mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir (+52-1)
- (modified) mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir (+131)
``````````diff
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
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/223281
More information about the Mlir-commits
mailing list