[Mlir-commits] [mlir] [mlir][Affine] Fix loadCSE replacing a load with one already scheduled for erasure (PR #217511)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Aman Singh (amanyagami)

<details>
<summary>Changes</summary>

affineScalarReplace defers erasing CSE'd-away loads until after the
whole walk completes. Within `loadCSE`'s candidate search for a
dominating, equivalent load to replace a given load with, a candidate
that had itself already been queued for erasure earlier in the same
walk was never excluded. This allowed a later load to be replaced
with (and its uses pointed at) an operation that was about to be
destroyed, crashing with "operation destroyed but still has uses".

Fix by skipping candidates already present in `loadOpsToErase`.

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

Fixes #<!-- -->216916

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+6) 
- (modified) mlir/test/Dialect/Affine/scalrep.mlir (+42) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7043083298615..4e0de45ecb8a3 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1000,6 +1000,12 @@ static void loadCSE(AffineReadOpInterface loadA,
     if (loadB.getValue().getType() != loadA.getValue().getType())
       continue;
 
+    // loadB should not already be scheduled for erasure. Otherwise, loadA
+    // would end up being replaced with (and then erased along with) an
+    // operation that no longer exists.
+    if (llvm::is_contained(loadOpsToErase, loadB.getOperation()))
+      continue;
+
     loadCandidates.push_back(loadB);
   }
 
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index fb6eef941790a..065a2d143d2b7 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -1034,3 +1034,45 @@ func.func @vector_store_dead_elim_same_type(%arg0: memref<20x1xi64>) {
   affine.vector_store %cst2, %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64>
   return
 }
+
+// Regression test for https://github.com/llvm/llvm-project/issues/216916:
+// loadCSE must not replace a load with a dominating, equivalent load that
+// has itself already been scheduled for erasure earlier in the same walk.
+// Doing so used to leave a load's uses pointing at an operation that was
+// about to be destroyed, which crashed with "operation destroyed but still
+// has uses". %a and %b below read the same, never-overwritten location
+// (%m[3]; the store only touches %m[0]) and dominate one another, so they
+// are redundant and must be CSE'd down to a single load without crashing.
+// %c reads the same location from inside the loop body: it does not
+// dominate (nor is dominated by) %a/%b in the same way, so loadCSE
+// conservatively leaves it as a separate load — this test only asserts
+// that doing so does not crash.
+memref.global "private" @gv : memref<4xi64> = dense<[1, 2, 3, 4]>
+
+// CHECK-LABEL: func @load_cse_no_replace_with_erased_load
+func.func @load_cse_no_replace_with_erased_load() -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c7 = arith.constant 7 : i64
+  %m = memref.get_global @gv : memref<4xi64>
+  // CHECK: %[[A:.*]] = affine.load
+  %a = affine.load %m[3] : memref<4xi64>
+  affine.store %c7, %m[0] : memref<4xi64>
+  // CHECK-NOT: affine.load
+  // CHECK: scf.for
+  // CHECK: affine.load
+  %b = affine.load %m[3] : memref<4xi64>
+  %s = scf.for %i = %c0 to %c2 step %c1 iter_args(%acc = %c0) -> (index) {
+    %c = affine.load %m[3] : memref<4xi64>
+    %x = index.castu %c : i64 to index
+    %y = index.add %acc, %x
+    scf.yield %y : index
+  }
+  %p = index.castu %a : i64 to index
+  %q = index.castu %b : i64 to index
+  %r = index.add %s, %p
+  %t = index.add %r, %q
+  // CHECK: return
+  return %t : index
+}

``````````

</details>


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