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

[Aman Singh]

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

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)

>From 66c7d6ac9408ed8323926040079d6149970e792d Mon Sep 17 00:00:00 2001
From: amanyagami <2amansingh2 at gmail.com>
Date: Wed, 19 Aug 2026 19:45:29 -0700
Subject: [PATCH] [mlir][Affine] Fix loadCSE replacing a load with one already
 scheduled for erasure

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.

Fixes https://github.com/llvm/llvm-project/issues/216916
---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp |  6 ++++
 mlir/test/Dialect/Affine/scalrep.mlir   | 42 +++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

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
+}