[Mlir-commits] [mlir] [mlir][bufferization] Fix use-after-free in ownership-based buffer deallocation (PR #184118)

[Matthias Springer]

================
@@ -717,8 +717,18 @@ BufferDeallocation::handleInterface(RegionBranchOpInterface op) {
   int counter = op->getNumResults();
   unsigned numMemrefResults = llvm::count_if(op->getResults(), isMemref);
   SmallVector<Type> ownershipResults(numMemrefResults, builder.getI1Type());
+  // Save the old result values before appendOpResults erases the op. The
+  // liveness analysis holds references to these values and they may be queried
+  // later (e.g., from handleInterface(BranchOpInterface) in the same block).
+  SmallVector<Value> oldResults(op->getResults());
   RegionBranchOpInterface newOp = appendOpResults(op, ownershipResults);
 
+  // Register the replacement of each old result with the corresponding new
+  // result so that stale liveness entries can be translated on demand.
+  for (auto [oldResult, newResult] :
+       llvm::zip(oldResults, newOp->getResults().take_front(oldResults.size())))
+    state.mapValue(oldResult, newResult);
----------------
matthias-springer wrote:

This should probably be done in `appendOpResults`. I suspect that this issue also reproducible from other places such as `BufferDeallocation::handleInterface(CallOpInterface op)`.

I'm also wondering if this mapping is the right way to fix the issue. Maybe we should maintain a set of all live/dead values and update the set whenever an op in created/replaced/erased.

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