[Mlir-commits] [mlir] [mlir][scf] Do not coalesce when an inner loop reads an outer iter_arg (PR #216853)
Alessandro Potenza
llvmlistbot at llvm.org
Sat Aug 29 23:56:53 PDT 2026
https://github.com/alepot55 updated https://github.com/llvm/llvm-project/pull/216853
>From 1464ecbe9ca0aced5a544c22a304f80875a5b3ec Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Mon, 17 Aug 2026 22:59:39 +0200
Subject: [PATCH 1/2] [mlir][scf] Do not coalesce when an inner loop reads an
outer iter_arg
`mlir-opt --affine-loop-coalescing` miscompiles this nest:
```mlir
func.func @coalesce(%init: i64, %lb: index, %ub: index, %st: index) -> i64 {
%r = scf.for %i = %lb to %ub step %st iter_args(%a = %init) -> (i64) {
%s = scf.for %j = %lb to %ub step %st iter_args(%b = %a) -> (i64) {
%t = arith.addi %b, %a : i64
scf.yield %t : i64
}
scf.yield %s : i64
}
return %r : i64
}
```
The body reads `%a`, the iteration argument of the outer loop, in addition
to `%b`, the one of the inner loop. `coalesceLoops` inlines the body of the
inner loop into the outer one and maps the iteration arguments of the inner
loop onto those of the outer loop, so both `%b` and `%a` become the single
iteration argument of the coalesced loop and the body collapses to
`arith.addi %arg, %arg`. That is not the same program: `%a` is fixed for a
whole run of the inner loop while `%b` is updated on every iteration, so 4x4
steps of `a -> 5a` turn into 16 steps of `a -> 2a`.
The existing legality check in `coalescePerfectlyNestedSCFForLoops` only
compares SSA identity of the iteration arguments and the yields along the
chain. It never asks what else the region of the inner loop reads from the
enclosing loop, so it accepts this nest.
Check up front that no loop in the band has its iteration arguments read
inside the region of the loop below it, and bail out otherwise. The check
sits in `coalesceLoops`, before any IR is modified, so it covers all the
entry points (`--affine-loop-coalescing`, `transform.loop.coalesce` and
`transform.loop.coalesce_nested`) and the transformation does not stop half
way. Using an iteration argument of the outer loop as an init argument of
the inner loop is a use on the inner loop operation itself, not inside its
region, so the chained nests that are the point of this transformation are
still coalesced.
`@tensor_loops_first_two` and `@tensor_loops_first_two_2` in
transform-op-coalesce.mlir contained the same pattern: their innermost body
reads `%arg3`, an iteration argument of the outermost loop, while the band
being coalesced is the outer two loops. Their expected output therefore
recorded the miscompilation. Both tests exist to check which loops get
picked for the band, which the body does not influence, so they now read the
iteration argument of their own loop like `@tensor_loops` already does.
`@tensor_loops_last_two` keeps reading `%arg3` because there the band is the
inner two loops, which leaves the outermost iteration argument alone.
This is the scf.for counterpart of the affine.for problem reported in #216455.
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 17 +++++++++++++
mlir/test/Dialect/Affine/loop-coalescing.mlir | 24 +++++++++++++++++++
.../Dialect/SCF/transform-op-coalesce.mlir | 4 ++--
3 files changed, 43 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 490972a837e96..b10ce6c5cbe4b 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1006,6 +1006,23 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
}
}
}
+
+ // Bail out if the region of an inner loop reads an iteration argument of an
+ // enclosing loop other than through its own iteration arguments. Coalescing
+ // maps the iteration arguments of every loop in the band onto the ones of the
+ // outermost loop, which turns such a read into a read of the value carried by
+ // the coalesced loop. That value is updated on every iteration, whereas the
+ // one the inner loop reads is fixed for a whole run of that loop.
+ for (unsigned i = 1, e = loops.size(); i < e; ++i) {
+ scf::ForOp innerLoop = loops[i];
+ for (BlockArgument iterArg : loops[i - 1].getRegionIterArgs()) {
+ if (llvm::any_of(iterArg.getUsers(), [&](Operation *user) {
+ return innerLoop->isProperAncestor(user);
+ }))
+ return failure();
+ }
+ }
+
// 1. Make sure all loops iterate from 0 to upperBound with step 1. This
// allows the following code to assume upperBound is the number of iterations.
for (auto loop : loops) {
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index b8a68b9c5c1f9..f928056935965 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -489,3 +489,27 @@ func.func @inner_loop_yields_induction_var() -> index {
}
return %r : index
}
+
+// -----
+
+// Verify that coalescing is not attempted when the body of the inner loop reads
+// an iteration argument of the outer loop directly. That value is fixed for a
+// whole run of the inner loop, whereas the value carried by the coalesced loop
+// is updated on every iteration.
+
+// CHECK-LABEL: @no_coalesce_outer_iter_arg_read_in_inner_loop
+// CHECK-SAME: %[[INIT:[A-Za-z0-9]+]]: i64
+func.func @no_coalesce_outer_iter_arg_read_in_inner_loop(%init: i64, %lb: index,
+ %ub: index, %step: index) -> i64 {
+ // CHECK: scf.for %{{.*}} iter_args(%[[OUTER:.*]] = %[[INIT]]) -> (i64)
+ %0 = scf.for %i = %lb to %ub step %step iter_args(%outer = %init) -> (i64) {
+ // CHECK: scf.for %{{.*}} iter_args(%[[INNER:.*]] = %[[OUTER]]) -> (i64)
+ %1 = scf.for %j = %lb to %ub step %step iter_args(%inner = %outer) -> (i64) {
+ // CHECK: arith.addi %[[INNER]], %[[OUTER]] : i64
+ %2 = arith.addi %inner, %outer : i64
+ scf.yield %2 : i64
+ }
+ scf.yield %1 : i64
+ }
+ return %0 : i64
+}
diff --git a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
index 467020e331295..ed6cdb0c5b688 100644
--- a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
@@ -161,7 +161,7 @@ func.func @tensor_loops_first_two(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf3
%0:2 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg2 = %arg0, %arg3 = %arg1) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%1:2 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg4 = %arg2, %arg5 = %arg3) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%2:2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg6 = %arg5, %arg7 = %arg4) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
- %3:2 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
+ %3:2 = "use"(%arg6, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
scf.yield %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
}
scf.yield %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xf32>
@@ -204,7 +204,7 @@ func.func @tensor_loops_first_two_2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?x
%0:2 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg2 = %arg0, %arg3 = %arg1) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%1:2 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg4 = %arg2, %arg5 = %arg3) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%2:2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg6 = %arg4, %arg7 = %arg5) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
- %3:2 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
+ %3:2 = "use"(%arg6, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
scf.yield %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
}
scf.yield %2#1, %2#0 : tensor<?x?xf32>, tensor<?x?xf32>
>From 3f1267228aaf695ffc7c298bbb6586c34d872d47 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Tue, 18 Aug 2026 01:35:49 +0200
Subject: [PATCH 2/2] Also decline when an op between the loops reads the outer
iter_arg
The first guard only looked at uses inside the inner loop's region, so an
imperfect nest whose in-between op reads the outer iteration argument was
still coalesced. Verified: transform.loop.coalesce_nested merged such a nest
and the read then saw the value carried by the coalesced loop.
Widen the condition to any use other than the inner loop's init operand,
which covers both shapes, and add the negative test.
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 17 +++++-----
.../Dialect/SCF/transform-op-coalesce.mlir | 32 +++++++++++++++++++
2 files changed, 41 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index b10ce6c5cbe4b..c990f2a352418 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1007,18 +1007,19 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
}
}
- // Bail out if the region of an inner loop reads an iteration argument of an
- // enclosing loop other than through its own iteration arguments. Coalescing
- // maps the iteration arguments of every loop in the band onto the ones of the
+ // Bail out if an iteration argument of an enclosing loop is read anywhere
+ // other than as the init operand of the loop nested in it. Coalescing maps
+ // the iteration arguments of every loop in the band onto the ones of the
// outermost loop, which turns such a read into a read of the value carried by
// the coalesced loop. That value is updated on every iteration, whereas the
- // one the inner loop reads is fixed for a whole run of that loop.
+ // one being read is fixed for a whole run of the inner loop. This covers both
+ // a read inside the inner loop's region and, in an imperfect nest, one from
+ // an operation sitting between the two loops.
for (unsigned i = 1, e = loops.size(); i < e; ++i) {
- scf::ForOp innerLoop = loops[i];
+ Operation *innerLoop = loops[i].getOperation();
for (BlockArgument iterArg : loops[i - 1].getRegionIterArgs()) {
- if (llvm::any_of(iterArg.getUsers(), [&](Operation *user) {
- return innerLoop->isProperAncestor(user);
- }))
+ if (llvm::any_of(iterArg.getUsers(),
+ [&](Operation *user) { return user != innerLoop; }))
return failure();
}
}
diff --git a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
index ed6cdb0c5b688..a9c19aaa9f12b 100644
--- a/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-coalesce.mlir
@@ -414,3 +414,35 @@ module attributes {transform.with_named_sequence} {
// CHECK: scf.yield %[[UPDATED]]
// CHECK: }
// CHECK: return %[[RESULT]]
+
+// -----
+
+// An operation sitting between the two loops of an imperfect nest and reading
+// the outer loop's iteration argument is unsound to coalesce for the same
+// reason as a read from inside the inner loop: after the merge it would see a
+// value updated on every iteration of the coalesced loop.
+
+func.func @no_coalesce_outer_iter_arg_read_between_loops(%init: f32) -> f32 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c3 = arith.constant 3 : index
+ %result = scf.for %i = %c0 to %c3 step %c1 iter_args(%outer = %init) -> (f32) {
+ %read = "use"(%outer) : (f32) -> f32
+ %inner_result = scf.for %j = %c0 to %c3 step %c1 iter_args(%inner = %outer) -> (f32) {
+ %updated = "use"(%inner, %read) : (f32, f32) -> f32
+ scf.yield %updated : f32
+ }
+ scf.yield %inner_result : f32
+ } {coalesce_nested}
+ return %result : f32
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce_nested} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
+ // expected-error @below {{failed to coalesce}}
+ %2 = transform.loop.coalesce_nested %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list