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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Stefan Schuermans (schuermans-roofline)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/203849.diff


1 Files Affected:

- (modified) mlir/lib/Transforms/Utils/CSE.cpp (+10-5) 


``````````diff
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;
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/203849