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

[Aman Singh]

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

>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 1/2] [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
+}

>From a1077595e9a68b0cee2e9e25b220b6cf08717398 Mon Sep 17 00:00:00 2001
From: amanyagami <aman at noahlabs.ai>
Date: Sat, 5 Sep 2026 02:03:20 -0700
Subject: [PATCH 2/2] Address review: tighten comments

Take the suggested wording in loadCSE, and drop the regression-test comment
block in scalrep.mlir -- the test name and CHECK lines carry the intent.

Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp |  3 +--
 mlir/test/Dialect/Affine/scalrep.mlir   | 12 ------------
 2 files changed, 1 insertion(+), 14 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 4e0de45ecb8a3..303b16a74e92c 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1001,8 +1001,7 @@ static void loadCSE(AffineReadOpInterface loadA,
       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.
+    // would end up being replaced with an operation that no longer exists.
     if (llvm::is_contained(loadOpsToErase, loadB.getOperation()))
       continue;
 
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 065a2d143d2b7..187fea10fdace 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -1035,18 +1035,6 @@ func.func @vector_store_dead_elim_same_type(%arg0: memref<20x1xi64>) {
   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