[Mlir-commits] [mlir] [mlir][affine] Fix crash coalescing loops whose inner results are used (PR #233320)

[llvmlistbot at llvm.org]

https://github.com/Naaz30 created https://github.com/llvm/llvm-project/pull/233320

Replace uses of the erased inner loop's results with its inits and bail out when inner iter_args are not yielded unchanged because replacing them with their inits would change the result.

Fixes #231902.

>From 09c05fc9a75ea39204c2e180d78b525326fe2248 Mon Sep 17 00:00:00 2001
From: Naazni Yahya <naazniyahya at gmail.com>
Date: Sun, 18 Oct 2026 19:30:31 -0400
Subject: [PATCH] [mlir][affine] Fix crash coalescing loops whose inner results
 are used

---
 mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp   | 25 +++++--
 mlir/test/Dialect/Affine/loop-coalescing.mlir | 71 +++++++++++++++++++
 2 files changed, 90 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index b1efda47945304..dc6d2846b00fd7 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1652,6 +1652,16 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
         loop.getConstantLowerBound() != 0)
       return failure();
   }
+
+  // The iter_args of the inner loops are replaced by their initial values
+  // below, which is only correct if they are yielded unchanged.
+  for (AffineForOp loop : loops.drop_front()) {
+    if (loop.getNumIterOperands() == 0)
+      continue;
+    auto yield = cast<AffineYieldOp>(loop.getBody()->getTerminator());
+    if (!llvm::equal(yield.getOperands(), loop.getRegionIterArgs()))
+      return failure();
+  }
   SmallVector<Value, 4> upperBoundSymbols;
   SmallVector<Value, 4> ubOperands(ub.getOperands().begin(),
                                    ub.getOperands().end());
@@ -1745,12 +1755,15 @@ 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();
+  // Go from outer to inner so that an inner loop's init
+  // that refers to an enclosing loop's iter_arg is updated first.
+  for (AffineForOp loop : loops.drop_front()) {
+    for (auto [iter, init] :
+         llvm::zip_equal(loop.getRegionIterArgs(), loop.getInits()))
+      iter.replaceAllUsesWith(init);
   }
+
+  secondOutermostLoop->replaceAllUsesWith(secondOutermostLoop.getInits());
   secondOutermostLoop.erase();
   return success();
 }
@@ -2863,4 +2876,4 @@ int64_t mlir::affine::numEnclosingInvariantLoops(OpOperand &operand) {
     count++;
   }
   return count;
-}
+}
\ No newline at end of file
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index b8a68b9c5c1f98..a0ffe927b6c9ef 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -489,3 +489,74 @@ func.func @inner_loop_yields_induction_var() -> index {
   }
   return %r : index
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/231902
+// Yielded values of inner loops that are removed must be retained to
+// replace later uses of the loop results.
+
+// CHECK-LABEL: func @inner_iter_arg_passthrough_result
+// CHECK-SAME:    (%[[INIT:.*]]: i32)
+// CHECK:         %[[RES:.*]] = affine.for %{{.*}} = 0 to %{{.*}} iter_args(%{{.*}} = %[[INIT]]) -> (i32) {
+// CHECK-NOT:       affine.for
+// CHECK:           "use"(%[[INIT]]) : (i32) -> ()
+// CHECK:           affine.yield %[[INIT]] : i32
+// CHECK:         }
+// CHECK:         return %[[RES]] : i32
+func.func @inner_iter_arg_passthrough_result(%init: i32) -> i32 {
+  %outer = affine.for %i = 0 to 4 iter_args(%a = %init) -> (i32) {
+    %inner = affine.for %j = 0 to 4 iter_args(%b = %init) -> (i32) {
+      "use"(%b) : (i32) -> ()
+      affine.yield %b : i32
+    }
+    affine.yield %inner : i32
+  }
+  return %outer : i32
+}
+
+// -----
+
+// Inner loops that carry a value across iterations (here a reduction) must
+// not be coalesced: replacing the iter_arg with its initial value would
+// change the result.
+
+// CHECK-LABEL: func @inner_loop_carried_value_not_coalesced
+// CHECK:         affine.for %{{.*}} = 0 to 4 iter_args
+// CHECK:           affine.for %{{.*}} = 0 to 4 iter_args
+// CHECK:             arith.addi
+func.func @inner_loop_carried_value_not_coalesced(%init: i32, %v: i32) -> i32 {
+  %outer = affine.for %i = 0 to 4 iter_args(%outer_arg = %init) -> (i32) {
+    %inner = affine.for %j = 0 to 4 iter_args(%inner_arg = %init) -> (i32) {
+      %next = arith.addi %inner_arg, %v : i32
+      affine.yield %next : i32
+    }
+    affine.yield %inner : i32
+  }
+  return %outer : i32
+}
+
+// -----
+
+// Pass-through iter_args of every inner loop in a 3-deep nest are replaced by
+// their initial values, including uses in the innermost body of the middle
+// loop's iter_arg.
+
+// CHECK-LABEL: func @three_level_passthrough_iter_args
+// CHECK-SAME:    (%[[X:.*]]: i32)
+// CHECK:         affine.for %{{.*}} = 0 to %{{.*}} {
+// CHECK-NOT:       affine.for
+// CHECK:           "use"(%[[X]], %[[X]]) : (i32, i32) -> ()
+// CHECK:         }
+func.func @three_level_passthrough_iter_args(%x: i32) {
+  affine.for %i = 0 to 4 {
+    %r = affine.for %j = 0 to 4 iter_args(%a = %x) -> (i32) {
+      %s = affine.for %k = 0 to 4 iter_args(%b = %x) -> (i32) {
+        "use"(%a, %b) : (i32, i32) -> ()
+        affine.yield %b : i32
+      }
+      affine.yield %a : i32
+    }
+  }
+  return
+}