[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
Mon Aug 17 14:51:03 PDT 2026
https://github.com/alepot55 created https://github.com/llvm/llvm-project/pull/216853
Fixes #216289.
`--affine-loop-coalescing` silently changes the result of a nested `scf.for`
when the inner loop's body reads the **outer** loop's iteration argument
directly, rather than only through its own:
```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
}
```
After the pass the body reads `arith.addi %arg5, %arg5` -- the two iteration
arguments have collapsed into one. They are not interchangeable: the outer one
is fixed for a whole run of the inner loop, the inner one is updated every
iteration. A 4x4 nest computing `a -> 5a` four times becomes 16 steps of `a -> 2a`.
Cause: the legality test in `mlir::coalesceLoops` compares SSA identity along the
chain (`llvm::equal(outerloop.getRegionIterArgs(), innerLoop.getInitArgs())` plus
a yield check) and never asks what the inner *region* reads from above.
`inlineBlockBefore` then maps the inner block arguments onto the outer loop's
iteration arguments, so a body reading both collapses to one value.
Decline to coalesce when an inner loop's region reads an enclosing loop's
iteration argument other than through its own. Using the outer iteration argument
as the inner loop's *init* operand is a use on the inner loop operation, not
inside its region, so `isProperAncestor` (proper, not `isAncestor`) leaves the
chained nests this transform exists for untouched. The check sits before the
first IR mutation, so it covers `--affine-loop-coalescing`,
`transform.loop.coalesce`, `transform.loop.coalesce_nested` and OpenACC's
`convertACCLoopToSCFFor` at once.
This is **not** the same bug as #216494. That PR fixes `affine::coalesceLoops`
for `affine.for` nests; `LoopCoalescing.cpp` dispatches by op type, so an
`scf.for` nest goes to `SCF/Utils/Utils.cpp` instead and never reaches that
guard. Its condition would also be wrong here: it rejects any loop that carries a
value, which SCF coalescing deliberately supports (`@tensor_loops`).
### Two existing tests change, and the change is forced
`@tensor_loops_first_two` and `@tensor_loops_first_two_2` in
`transform-op-coalesce.mlir` contain the miscompiling pattern by accident: their
innermost body reads `%arg3`, an iteration argument of the *outermost* loop,
while the band actually coalesced is the outer two. Their committed golden output
therefore records the bug. Both tests exist to check *which* loops are selected
into a band, which the body does not influence, so I changed one token in each
(`"use"(%arg3, ...)` -> `"use"(%arg6, ...)`) to read their own loop's iteration
argument, exactly as the sibling `@tensor_loops` already does. No CHECK line
references the `"use"` op.
I verified this is necessary rather than cosmetic: with the patch and the
*unmodified* file, `mlir-opt` reports `error: unexpected error: failed to
coalesce` at lines 177 and 220. That is independent evidence the bug is real and
already reachable from in-tree tests.
### Verified by execution
- Miscompile reproduces on `d4e78d7f5`; the new test fails on unpatched `mlir-opt` and passes with the patch.
- No regression: `Dialect/Affine` 71/71, `Dialect/SCF` 46/46, plus MemRef, Vector, XeGPU, Transforms (368 total).
- I walked every existing iteration-argument coalescing test and confirmed the guard is vacuous
for each (`@tensor_loops`, `@tensor_loops_last_two`, `@trip_one_loops`, `@opsinbetween_nested_loops`,
the three `@noramalized_loops_*`), and that OpenACC's caller builds loops with no init args.
- `git clang-format` reports no changes.
Two things I left out on purpose, happy to add either if you disagree:
- For `coalesce_nested` over an *imperfect* nest, an op between the two loops reading the outer
iteration argument is unsound for the same reason and is not caught here. Widening the rule to
"outer iteration arguments have no use other than the inner loop's init operand" breaks no
existing test, but I could not construct an executed repro for that case.
- For a 3-deep nest where only the inner pair is legal, this now coalesces nothing rather than the
legal pair. #216494 takes the same bail-out-entirely approach.
---
Assisted-by: Claude (Anthropic)
This patch was written with AI assistance, disclosed per the LLVM AI Tool Use Policy.
Everything reported above as verified was verified by building and running.
>From 8dd0ff61ced42deb8935bbf8fe3d90d7a6161116 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] [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 c158e624002bd..6d64f90206448 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -953,6 +953,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 d08d2bf79c781..4b21603eff112 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -462,3 +462,27 @@ func.func @no_coalesce_zero_step(%lb: index, %ub: index) {
}
return
}
+
+// -----
+
+// 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>
More information about the Mlir-commits
mailing list