[Mlir-commits] [mlir] [mlir][SCF] Fix use-after-free in coalesceLoops when inner loop yields its induction var (PR #217510)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Aman Singh (amanyagami)

<details>
<summary>Changes</summary>

coalesceLoops rewrites yielded values that reference the inner loop's
iter_args (about to be inlined and destroyed) to their replacement
values, but never checked whether a yielded value was the inner
loop's induction variable itself — also a block argument about to be
destroyed. When the inner loop's `scf.yield` yields its own
induction variable, this left a dangling `Value` reference that
crashed `RewriterBase::replaceOp` with a use-after-free.

Fix by also rewriting a yielded induction variable to its delinearized
replacement value, mirroring the existing iter_arg case.

Verified: reverting this fix reproduces the reported segfault; with
the fix, `mlir-opt --affine-loop-coalescing` on the reported
reproducer succeeds, and `mlir/test/Dialect/Affine/loop-coalescing.mlir`
passes.

Fixes #<!-- -->216903

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---
Full diff: https://github.com/llvm/llvm-project/pull/217510.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+7) 
- (modified) mlir/test/Dialect/Affine/loop-coalescing.mlir (+27) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c158e624002bd..36cddd5cf8bd7 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1002,6 +1002,13 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
     auto yieldedVals = llvm::to_vector(innerTerminator->getOperands());
     assert(llvm::equal(outerLoop.getRegionIterArgs(), innerLoop.getInitArgs()));
     for (Value &yieldedVal : yieldedVals) {
+      // The yielded value may be the induction variable of the inner loop,
+      // which is about to be inlined and whose block argument is about to
+      // be destroyed. Use its replacement value instead.
+      if (yieldedVal == innerLoop.getInductionVar()) {
+        yieldedVal = delinearizeIvs[i];
+        continue;
+      }
       // The yielded value may be an iteration argument of the inner loop
       // which is about to be inlined.
       auto iter = llvm::find(innerLoop.getRegionIterArgs(), yieldedVal);
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index d08d2bf79c781..b8a68b9c5c1f9 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -462,3 +462,30 @@ func.func @no_coalesce_zero_step(%lb: index, %ub: index) {
   }
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/216903:
+// coalescing must not crash when the inner loop yields its own induction
+// variable. The yielded induction variable has to be rewritten to its
+// linearized replacement before the inner loop's induction variable block
+// argument is destroyed by inlining.
+
+// CHECK-LABEL: @inner_loop_yields_induction_var
+func.func @inner_loop_yields_induction_var() -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (index) {
+  // CHECK-NOT: scf.for
+  // CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[IV]]
+  // CHECK: scf.yield %[[DELIN]]#1 : index
+  // CHECK: }
+  %r = scf.for %i = %c0 to %c4 step %c1 iter_args(%a = %c0) -> (index) {
+    %s = scf.for %j = %c0 to %c4 step %c1 iter_args(%b = %a) -> (index) {
+      scf.yield %j : index
+    }
+    scf.yield %s : index
+  }
+  return %r : index
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/217510