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

[Stefan Schuermans]

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

>From e3193b4867c822e2490b405e2ba26873f716a7b2 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 by moving the copy of the location (the "use") to before the deletion of the op (the "free").

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 | 58 +++++++++++++++----------------
 1 file changed, 29 insertions(+), 29 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/CSE.cpp b/mlir/lib/Transforms/Utils/CSE.cpp
index eed77f20dd2da..5af76bc99e956 100644
--- a/mlir/lib/Transforms/Utils/CSE.cpp
+++ b/mlir/lib/Transforms/Utils/CSE.cpp
@@ -125,44 +125,44 @@ class CSEDriver {
 void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                      Operation *existing,
                                      bool hasSSADominance) {
+  // If the existing operation has an unknown location and the current
+  // operation doesn't, then set the existing op's location to that of the
+  // current op.
+  if (isa<UnknownLoc>(existing->getLoc()) && !isa<UnknownLoc>(op->getLoc()))
+    existing->setLoc(op->getLoc());
+
+  ++numCSE;
+
   // 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 delete it.
   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);
-  } else {
-    // When the region does not have SSA dominance, we need to check if we
-    // have visited a use before replacing any use.
-    auto wasVisited = [&](OpOperand &operand) {
-      return !knownValues.count(operand.getOwner());
-    };
-    if (auto *rewriteListener =
-            dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
-      for (Value v : op->getResults())
-        if (all_of(v.getUses(), wasVisited))
-          rewriteListener->notifyOperationReplaced(op, existing);
-
-    // Replace all uses, but do not remove the operation yet. This does not
-    // notify the listener because the original op is not erased.
-    rewriter.replaceUsesWithIf(op->getResults(), existing->getResults(),
-                               wasVisited);
-
-    // There may be some remaining uses of the operation.
-    if (op->use_empty())
-      eraseDeadOp(op);
+    return;
   }
 
-  // If the existing operation has an unknown location and the current
-  // operation doesn't, then set the existing op's location to that of the
-  // current op.
-  if (isa<UnknownLoc>(existing->getLoc()) && !isa<UnknownLoc>(op->getLoc()))
-    existing->setLoc(op->getLoc());
-
-  ++numCSE;
+  // When the region does not have SSA dominance, we need to check if we
+  // have visited a use before replacing any use.
+  auto wasVisited = [&](OpOperand &operand) {
+    return !knownValues.count(operand.getOwner());
+  };
+  if (auto *rewriteListener =
+          dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
+    for (Value v : op->getResults())
+      if (all_of(v.getUses(), wasVisited))
+        rewriteListener->notifyOperationReplaced(op, existing);
+
+  // Replace all uses, but do not remove the operation yet. This does not
+  // notify the listener because the original op is not erased.
+  rewriter.replaceUsesWithIf(op->getResults(), existing->getResults(),
+                             wasVisited);
+
+  // There may be some remaining uses of the operation.
+  if (op->use_empty())
+    eraseDeadOp(op);
 }
 
 bool CSEDriver::hasOtherSideEffectingOpInBetween(Operation *fromOp,