[Mlir-commits] [mlir] 839dc4f - [mlir][bufferization] Fix use-after-free in ownership-based buffer deallocation (#184118)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Mar 5 04:28:54 PST 2026
Author: Mehdi Amini
Date: 2026-03-05T13:28:50+01:00
New Revision: 839dc4f7cfff5d240cc9274696efb056dd3847cd
URL: https://github.com/llvm/llvm-project/commit/839dc4f7cfff5d240cc9274696efb056dd3847cd
DIFF: https://github.com/llvm/llvm-project/commit/839dc4f7cfff5d240cc9274696efb056dd3847cd.diff
LOG: [mlir][bufferization] Fix use-after-free in ownership-based buffer deallocation (#184118)
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
Assisted-by: Claude Code
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h
mlir/lib/Dialect/Bufferization/IR/BufferDeallocationOpInterface.cpp
mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
Removed:
################################################################################
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..e625f172a3bf3 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -668,6 +668,11 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
Operation *BufferDeallocation::appendOpResults(Operation *op,
ArrayRef<Type> types) {
SmallVector<Type> newTypes(op->getResultTypes());
+ // 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());
+
newTypes.append(types.begin(), types.end());
auto *newOp = Operation::create(op->getLoc(), op->getName(), newTypes,
op->getOperands(), op->getAttrDictionary(),
@@ -681,6 +686,12 @@ Operation *BufferDeallocation::appendOpResults(Operation *op,
op->replaceAllUsesWith(newOp->getResults().take_front(op->getNumResults()));
op->erase();
+ // 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);
+
return newOp;
}
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
+}
More information about the Mlir-commits
mailing list