[Mlir-commits] [mlir] [mlir][CSE] fix use-after-free (PR #203849)

[Stefan Schuermans]

https://github.com/schuermans-roofline created https://github.com/llvm/llvm-project/pull/203849

There was a use-after free introduced accidentally by e79e056ee982 [mlir][CSE] Remove the opsToErase container and immediately delete dead ops. (#203702) This change fixes it.

Locating the use-after-free has been assisted by Claude Code. The implementation of the fix has been done without it.

>From e2fa29a2f499aa2e5af30c0634b0daa14aabe13b Mon Sep 17 00:00:00 2001
From: Stefan Schuermans <schuermans at roofline.ai>
Date: Mon, 15 Jun 2026 10:45:51 +0200
Subject: [PATCH] [mlir][CSE] fix use-after-free

There was a use-after free introduced accidentally by
e79e056ee982 [mlir][CSE] Remove the opsToErase container and immediately delete dead ops. (#203702)
This change fixes it.

Locating the use-after-free has been assisted by Claude Code. The implementation of the fix has been done without it.

Signed-off-by: Stefan Schuermans <schuermans at roofline.ai>
---
 mlir/lib/Transforms/Utils/CSE.cpp | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/CSE.cpp b/mlir/lib/Transforms/Utils/CSE.cpp
index eed77f20dd2da..a5cc2376f4ba8 100644
--- a/mlir/lib/Transforms/Utils/CSE.cpp
+++ b/mlir/lib/Transforms/Utils/CSE.cpp
@@ -126,14 +126,15 @@ void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                      Operation *existing,
                                      bool hasSSADominance) {
   // If we find one then replace all uses of the current operation with the
-  // existing one and mark it for deletion. We can only replace an operand in
-  // an operation if it has not been visited yet.
+  // existing one and mark it for deletion at the bottom of the function. We
+  // can only replace an operand in an operation if it has not been visited yet.
+  bool opIsDead = false;
   if (hasSSADominance) {
     // If the region has SSA dominance, then we are guaranteed to have not
     // visited any use of the current operation.
     // Replace all uses, but do not remove the operation yet.
     rewriter.replaceAllOpUsesWith(op, existing->getResults());
-    eraseDeadOp(op);
+    opIsDead = true;
   } else {
     // When the region does not have SSA dominance, we need to check if we
     // have visited a use before replacing any use.
@@ -152,8 +153,7 @@ void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                wasVisited);
 
     // There may be some remaining uses of the operation.
-    if (op->use_empty())
-      eraseDeadOp(op);
+    opIsDead = op->use_empty();
   }
 
   // If the existing operation has an unknown location and the current
@@ -162,6 +162,11 @@ void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
   if (isa<UnknownLoc>(existing->getLoc()) && !isa<UnknownLoc>(op->getLoc()))
     existing->setLoc(op->getLoc());
 
+  // Deletion of the dead current op can only be done here, and not before,
+  // because the its location is still read above.
+  if (opIsDead)
+    eraseDeadOp(op);
+
   ++numCSE;
 }