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

[llvmlistbot at llvm.org]

Author: Stefan Schuermans
Date: 2026-06-15T11:20:20Z
New Revision: c33a47e6863b3294e693e540df8134a486eafef7

URL: https://github.com/llvm/llvm-project/commit/c33a47e6863b3294e693e540df8134a486eafef7
DIFF: https://github.com/llvm/llvm-project/commit/c33a47e6863b3294e693e540df8134a486eafef7.diff

LOG: [mlir][CSE] fix use-after-free (#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 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>

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/CSE.cpp

Removed: 
    


################################################################################
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,