[Mlir-commits] [mlir] [mlir][affine] Do not coalesce loops that carry a value (PR #216494)
Hamza Qureshi
llvmlistbot at llvm.org
Sat Aug 15 07:33:34 PDT 2026
https://github.com/hamzaqureshi5 created https://github.com/llvm/llvm-project/pull/216494
Fixes #216455
`mlir-opt --affine-loop-coalescing` crashes on this input:
```mlir
func.func @f(%seed: i64) {
affine.for %i = 0 to 4 {
affine.for %j = 0 to 4 {
%r = affine.for %k = 0 to 4 iter_args(%acc = %seed) -> (i64) {
%t = arith.addi %acc, %acc : i64
affine.yield %t : i64
}
}
}
return
}
```
```
Assertion `use_empty() && "Cannot destroy a value that still has uses!"' failed.
```
`coalesceLoops` moves the body of the innermost loop out and erases every loop but the outermost one, replacing their iteration arguments by their initial values. There are two problems with that.
Only the second-outermost loop had its iteration arguments replaced. Erasing it also destroys the loops nested inside it, so in a deeper nest their iteration arguments were still used by the operations that had just been moved out, which is the crash above.
Replacing an iteration argument by its initial value is also only correct when the loop yields that argument unchanged, so that it is really loop invariant. When the loop does carry a value the accumulation is silently dropped:
```mlir
affine.for %i = 0 to 4 {
%r = affine.for %k = 0 to 4 iter_args(%acc = %seed) -> (i64) {
memref.store %acc, %out[%k] : memref<4xi64>
%t = arith.addi %acc, %acc : i64
affine.yield %t : i64
}
}
```
The stores should write `%seed`, `2 x %seed`, `4 x %seed`, `8 x %seed`. On main the coalesced loop writes `%seed` every iteration. This is a two loop nest, so it goes through the existing handling added in #169514 and does not need the deeper nest above.
This checks up front that every loop that will be erased yields its iteration arguments unchanged, and bails out otherwise. The replacement is then valid for all of the erased loops, which fixes both the crash and the dropped accumulation. Loops that only forward their iteration argument, like the existing `@inner_loop_has_iter_args` test, are still coalesced.
The check is done before any IR is modified so the transformation does not stop half way.
## Testing
Added two cases to `mlir/test/Dialect/Affine/loop-coalescing.mlir`: the reported nest, and a two loop nest whose carried value is observed by a store.
Verified on an assertions build that the reproducer aborts before the change and succeeds after, that the two loop nest stores the iteration argument instead of the initial value after the change, that both new tests fail without it, and that `mlir/test/Dialect/Affine` and `mlir/test/Dialect/SCF` stay green (the SCF transform op calls the same entry point).
cc @linuxlonelyeagle, who added the existing iteration argument handling in #169514.
>From e6f8c50e84c5edc4c7fc465695c80680c65e118d Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Sat, 15 Aug 2026 19:27:30 +0500
Subject: [PATCH] [mlir][affine] Do not coalesce loops that carry a value
coalesceLoops moves the body of the innermost loop out and erases all the
loops but the outermost one, replacing their iteration arguments by their
initial values. That replacement is only correct when a loop yields its
iteration arguments unchanged, so that the argument is really loop invariant.
Two things went wrong. Only the second-outermost loop had its iteration
arguments replaced, so in a deeper nest the arguments of the loops below it
were left used by the operations that had just been moved out, and erasing
them tripped the "value still has uses" assertion. And for a loop that does
carry a value, replacing the argument by the initial value silently dropped
the accumulation.
Check up front that every loop that will be erased yields its iteration
arguments unchanged, and bail out otherwise. The replacement is then valid for
all of them.
Fixes #216455
---
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 26 ++++++++--
mlir/test/Dialect/Affine/loop-coalescing.mlir | 49 +++++++++++++++++++
2 files changed, 70 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 90bc57e950cf1..f775299667371 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1624,6 +1624,17 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
loop.getConstantLowerBound() != 0)
return failure();
}
+ // All the loops but the outermost one are erased, and their iteration
+ // arguments are replaced by their initial values. That is only valid when a
+ // loop does not actually carry a value across iterations, i.e. it yields its
+ // iteration arguments unchanged.
+ for (AffineForOp loop : loops.drop_front()) {
+ auto yieldOp = cast<AffineYieldOp>(loop.getBody()->getTerminator());
+ for (auto [iter, yielded] :
+ llvm::zip_equal(loop.getRegionIterArgs(), yieldOp.getOperands()))
+ if (iter != yielded)
+ return failure();
+ }
SmallVector<Value, 4> upperBoundSymbols;
SmallVector<Value, 4> ubOperands(ub.getOperands().begin(),
ub.getOperands().end());
@@ -1717,11 +1728,16 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
outermost.getBody()->getOperations().splice(
Block::iterator(secondOutermostLoop.getOperation()),
innermost.getBody()->getOperations());
- for (auto [iter, init] :
- llvm::zip_equal(secondOutermostLoop.getRegionIterArgs(),
- secondOutermostLoop.getInits())) {
- iter.replaceAllUsesWith(init);
- iter.dropAllUses();
+ // Erasing the second-outermost loop also destroys every loop nested in it,
+ // so the iteration arguments of all of them may still be used by the
+ // operations that were just moved out. They are loop invariant (checked
+ // above), so replace them by their initial values.
+ for (AffineForOp loop : loops.drop_front()) {
+ for (auto [iter, init] :
+ llvm::zip_equal(loop.getRegionIterArgs(), loop.getInits())) {
+ iter.replaceAllUsesWith(init);
+ iter.dropAllUses();
+ }
}
secondOutermostLoop.erase();
return success();
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index d08d2bf79c781..695026e878596 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -462,3 +462,52 @@ func.func @no_coalesce_zero_step(%lb: index, %ub: index) {
}
return
}
+
+// -----
+
+// Verify that coalescing is not attempted when an inner loop really carries a
+// value, i.e. it does not yield its iteration argument unchanged. Erasing such
+// a loop would leave its iteration argument used by the operations moved out
+// of it, and replacing that argument by the initial value would drop the
+// accumulation. Coalescing always introduces affine.apply ops, so their
+// absence shows the nest was left alone.
+
+// CHECK-LABEL: @no_coalesce_carried_value
+// CHECK-NOT: affine.apply
+// CHECK: affine.for
+// CHECK: affine.for
+// CHECK: affine.for {{.*}} iter_args
+// CHECK: arith.addi
+func.func @no_coalesce_carried_value(%seed: i64, %out: memref<4xi64>) {
+ affine.for %i = 0 to 4 {
+ affine.for %j = 0 to 4 {
+ %r = affine.for %k = 0 to 4 iter_args(%acc = %seed) -> (i64) {
+ %t = arith.addi %acc, %acc : i64
+ memref.store %t, %out[%k] : memref<4xi64>
+ affine.yield %t : i64
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// Same for a two loop nest, where the carried value is also observed by a
+// store: coalescing it would store the initial value on every iteration.
+
+// CHECK-LABEL: @no_coalesce_observed_carried_value
+// CHECK-NOT: affine.apply
+// CHECK: affine.for
+// CHECK: affine.for {{.*}} iter_args
+// CHECK: memref.store
+func.func @no_coalesce_observed_carried_value(%seed: i64, %out: memref<4xi64>) {
+ affine.for %i = 0 to 4 {
+ %r = affine.for %k = 0 to 4 iter_args(%acc = %seed) -> (i64) {
+ memref.store %acc, %out[%k] : memref<4xi64>
+ %t = arith.addi %acc, %acc : i64
+ affine.yield %t : i64
+ }
+ }
+ return
+}
More information about the Mlir-commits
mailing list