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

Mon Mar 2 05:24:19 PST 2026

llvmbot wrote:




@llvm/pr-subscribers-mlir-bufferization

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When `handleInterface(RegionBranchOpInterface)` processes an op such as `scf.for`, it calls `appendOpResults` to clone the op with extra ownership result types and erase the original. The `Liveness` analysis is computed once before the transformation begins and may still reference the old (now-freed) result values.

If the same block contains a `BranchOpInterface` terminator (e.g., `cf.br`) after the structured loop, `handleInterface(BranchOpInterface)` calls `getMemrefsToRetain`, which iterates `liveness.getLiveOut()`. That set may contain stale `Value` objects pointing to the erased op's results. Calling `isMemref()` on such a value dereferences freed memory, triggering a crash.

Fix by adding a `valueMapping` map to `DeallocationState`. Before erasing the old op in `handleInterface(RegionBranchOpInterface)`, record the old-to-new result mapping via `state.mapValue`. The `getLiveMemrefsIn` and `getMemrefsToRetain` helpers translate stale liveness values through this map before calling `isMemref`, so they always operate on live pointers.

Fixes #119863

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


4 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h (+13) 
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp (+35-5) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+10) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir (+25) 


``````````diff

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
index 5d33817c7d9fc..c961c2724d2a9 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
@@ -186,6 +186,13 @@ class DeallocationState {
   /// attributes on the function operation.
   SymbolTableCollection *getSymbolTable() { return &symbolTable; }
 
+  /// Register that 'oldValue' has been replaced by 'newValue'. When the
+  /// liveness analysis is consulted after an op has been replaced (e.g., via
+  /// appendOpResults), the cached liveness may still refer to the old value.
+  /// This mapping is used to translate stale values to their replacements
+  /// before checking whether a value is a MemRef.
+  void mapValue(Value oldValue, Value newValue);
+
 private:
   // Symbol cache to lookup functions from call operations to check attributes
   // on the function operation.
@@ -203,6 +210,12 @@ class DeallocationState {
   // The underlying liveness analysis to compute fine grained information about
   // alloc and dealloc positions.
   Liveness liveness;
+
+  // Maps values that have been replaced (e.g., when an op is cloned with extra
+  // results via appendOpResults) to their replacements. The liveness analysis
+  // is computed once and may contain stale values after IR modifications; this
+  // map is used to translate them before accessing their types.
+  DenseMap<Value, Value> valueMapping;
 };
 
 namespace deallocation_impl {
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
index a8703edfec7e2..ff4cad71110b6 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
@@ -124,10 +124,21 @@ void DeallocationState::dropMemrefToDeallocate(Value memref, Block *block) {
   llvm::erase(memrefsToDeallocatePerBlock[block], memref);
 }
 
+void DeallocationState::mapValue(Value oldValue, Value newValue) {
+  valueMapping[oldValue] = newValue;
+}
+
 void DeallocationState::getLiveMemrefsIn(Block *block,
                                          SmallVectorImpl<Value> &memrefs) {
-  SmallVector<Value> liveMemrefs(
-      llvm::make_filter_range(liveness.getLiveIn(block), isMemref));
+  SmallVector<Value> liveMemrefs;
+  for (Value val : liveness.getLiveIn(block)) {
+    // Translate any value that was replaced (e.g., by appendOpResults) to its
+    // current equivalent before checking whether it is a MemRef.
+    if (Value mapped = valueMapping.lookup(val))
+      val = mapped;
+    if (isMemref(val))
+      liveMemrefs.push_back(val);
+  }
   llvm::sort(liveMemrefs, ValueComparator());
   memrefs.append(liveMemrefs);
 }
@@ -167,13 +178,32 @@ void DeallocationState::getMemrefsToRetain(
     toRetain.push_back(operand);
   }
 
+  // Translate any value replaced during the transformation (e.g., when an op
+  // was cloned with extra results via appendOpResults) before checking whether
+  // it is a MemRef. The liveness analysis is computed once and may contain
+  // stale values after IR modifications.
+  auto translateValue = [&](Value val) -> Value {
+    if (Value mapped = valueMapping.lookup(val))
+      return mapped;
+    return val;
+  };
+
   SmallPtrSet<Value, 16> liveOut;
-  for (auto val : liveness.getLiveOut(fromBlock))
+  for (auto val : liveness.getLiveOut(fromBlock)) {
+    val = translateValue(val);
     if (isMemref(val))
       liveOut.insert(val);
+  }
 
-  if (toBlock)
-    llvm::set_intersect(liveOut, liveness.getLiveIn(toBlock));
+  if (toBlock) {
+    SmallPtrSet<Value, 16> liveIn;
+    for (auto val : liveness.getLiveIn(toBlock)) {
+      val = translateValue(val);
+      if (isMemref(val))
+        liveIn.insert(val);
+    }
+    llvm::set_intersect(liveOut, liveIn);
+  }
 
   // liveOut has non-deterministic order because it was constructed by iterating
   // over a hash-set.
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 6081e515d4e3a..8e6b6e806be8b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -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);
+
   for (auto result : llvm::make_filter_range(newOp->getResults(), isMemref)) {
     state.updateOwnership(result, newOp->getResult(counter++));
     state.addMemrefToDeallocate(result, newOp->getBlock());
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
index 8e14990502143..3ce87f2fead21 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
@@ -593,3 +593,28 @@ func.func @blocks_not_preordered_by_dominance() {
 //  CHECK-NEXT:   [[ALLOC]] = memref.alloc()
 //  CHECK-NEXT:   cf.br [[BB2]]
 //  CHECK-NEXT: }
+
+// -----
+
+// Regression test: a block that contains both a structured loop
+// (RegionBranchOpInterface) with iter_args and an explicit cf.br
+// (BranchOpInterface) must not crash. The pass used to use-after-free when
+// accessing the liveness info after the loop op was replaced by
+// appendOpResults.
+// https://github.com/llvm/llvm-project/issues/119863
+
+// CHECK-LABEL: func @region_branch_op_followed_by_cf_br
+func.func @region_branch_op_followed_by_cf_br(%arg0: f32) -> f32 {
+  cf.br ^bb1
+^bb1:
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%acc = %arg0) -> f32 {
+    %1 = arith.addf %acc, %acc : f32
+    scf.yield %1 : f32
+  }
+  cf.br ^bb2
+^bb2:
+  return %0 : f32
+}

``````````

</details>


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