[Mlir-commits] [mlir] [mlir][bufferization] Fix stale BufferOriginAnalysis in BufferDeallocationSimplification (PR #210105)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 16 09:42:05 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Vito Secona (secona)
<details>
<summary>Changes</summary>
Fixes #<!-- -->205228
This change fixes a use-after-free bug in the BufferDeallocationSimplification pipeline caused by the greedy pattern rewriter deleting operation tracked by the BufferOriginAnalysis. BufferViewFlowAnalysis (internally used by BufferOriginAnalysis) creates the dependencies map once during initialization. In this specific case, the folder removed the `memref.cast` ops making later uses a use-after-free bug.
Below is the snippet of running the pass using `-debug-only=greedy-rewriter`. It erases a `memref.cast` and crashes in RemoveDeallocMemrefsContainedInRetained.
```
...
[greedy-rewriter:1] //===-------------------------------------------===//
[greedy-rewriter:1] Processing operation : 'memref.cast'(0x654d43d18670) {
[greedy-rewriter:1] %4 = "memref.cast"(%3) : (memref<?x?xf32, 1>) -> memref<4x4xf32, 1>
[greedy-rewriter:1]
[greedy-rewriter:1] } -> success : operation was folded
[greedy-rewriter:1] //===-------------------------------------------===//
[greedy-rewriter:1] ** Replace : 'memref.cast'(0x654d43d18670)
[greedy-rewriter:1] ** Modified: 'scf.yield'(0x654d43c97bb0)
[greedy-rewriter:1] ** Erase : 'memref.cast'(0x654d43d18670)
...
[greedy-rewriter:1] Processing operation : 'bufferization.dealloc'(0x654d43cbd7d0) {
[greedy-rewriter:1] %5 = "bufferization.dealloc"(%4#<!-- -->0, %1, %3) <{operandSegmentSizes = array<i32: 1, 1, 1>}> : (memref<f32, 1>, i1, memref<4x4xf32, 1>) -> i1
[greedy-rewriter:1]
[greedy-rewriter:1]
[greedy-rewriter:1] * Pattern (anonymous namespace)::RemoveDeallocMemrefsContainedInRetained : 'bufferization.dealloc -> ()' {
```
The fix is by recreating BufferOriginAnalysis on-the-fly inside matchAndRewrite of each rewrite patterns, where it was previously created once in BufferDeallocationSimplificationPass::runOnOperation.
---
Full diff: https://github.com/llvm/llvm-project/pull/210105.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp (+27-30)
- (modified) mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir (+29)
``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
index a465c957d063e..64b8f6d1fbfa8 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
@@ -46,6 +46,14 @@ static Value getViewBase(Value value) {
return value;
}
+static BufferOriginAnalysis getBufferOriginAnalysis(Operation *op) {
+ // dataflow boundary; we want to the origin of all internral buffers.
+ Operation *parent = op->getParentWithTrait<OpTrait::IsIsolatedFromAbove>();
+ assert(parent && "expected dealloc inside an IsIsolatedFromAbove region");
+ BufferOriginAnalysis analysis(parent);
+ return analysis;
+}
+
static LogicalResult updateDeallocIfChanged(DeallocOp deallocOp,
ValueRange memrefs,
ValueRange conditions,
@@ -134,9 +142,8 @@ namespace {
/// given that `%r0` and `%r1` may not alias with `%m0`.
struct RemoveDeallocMemrefsContainedInRetained
: public OpRewritePattern<DeallocOp> {
- RemoveDeallocMemrefsContainedInRetained(MLIRContext *context,
- BufferOriginAnalysis &analysis)
- : OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
+ RemoveDeallocMemrefsContainedInRetained(MLIRContext *context)
+ : OpRewritePattern<DeallocOp>(context) {}
/// The passed 'memref' must not have a may-alias relation to any retained
/// memref, and at least one must-alias relation. If there is no must-aliasing
@@ -145,7 +152,8 @@ struct RemoveDeallocMemrefsContainedInRetained
/// no-alias, then just proceed, if it's must-alias we need to update the
/// updated condition returned by the dealloc operation for that alias.
LogicalResult handleOneMemref(DeallocOp deallocOp, Value memref, Value cond,
- PatternRewriter &rewriter) const {
+ PatternRewriter &rewriter,
+ BufferOriginAnalysis analysis) const {
rewriter.setInsertionPointAfter(deallocOp);
// Check that there is no may-aliasing memref and that at least one memref
@@ -183,6 +191,8 @@ struct RemoveDeallocMemrefsContainedInRetained
LogicalResult matchAndRewrite(DeallocOp deallocOp,
PatternRewriter &rewriter) const override {
+ auto analysis = getBufferOriginAnalysis(deallocOp);
+
// There must not be any duplicates in the retain list anymore because we
// would miss updating one of the result values otherwise.
DenseSet<Value> retained(deallocOp.getRetained().begin(),
@@ -194,13 +204,14 @@ struct RemoveDeallocMemrefsContainedInRetained
for (auto [memref, cond] :
llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
- if (succeeded(handleOneMemref(deallocOp, memref, cond, rewriter)))
+ if (succeeded(
+ handleOneMemref(deallocOp, memref, cond, rewriter, analysis)))
continue;
if (auto extractOp =
memref.getDefiningOp<memref::ExtractStridedMetadataOp>())
if (succeeded(handleOneMemref(deallocOp, extractOp.getOperand(), cond,
- rewriter)))
+ rewriter, analysis)))
continue;
newMemrefs.push_back(memref);
@@ -212,9 +223,6 @@ struct RemoveDeallocMemrefsContainedInRetained
return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
rewriter);
}
-
-private:
- BufferOriginAnalysis &analysis;
};
/// Remove memrefs from the `retained` list which are guaranteed to not alias
@@ -235,12 +243,12 @@ struct RemoveDeallocMemrefsContainedInRetained
/// ```
struct RemoveRetainedMemrefsGuaranteedToNotAlias
: public OpRewritePattern<DeallocOp> {
- RemoveRetainedMemrefsGuaranteedToNotAlias(MLIRContext *context,
- BufferOriginAnalysis &analysis)
- : OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
+ RemoveRetainedMemrefsGuaranteedToNotAlias(MLIRContext *context)
+ : OpRewritePattern<DeallocOp>(context) {}
LogicalResult matchAndRewrite(DeallocOp deallocOp,
PatternRewriter &rewriter) const override {
+ auto analysis = getBufferOriginAnalysis(deallocOp);
SmallVector<Value> newRetainedMemrefs, replacements;
for (auto retainedMemref : deallocOp.getRetained()) {
@@ -270,9 +278,6 @@ struct RemoveRetainedMemrefsGuaranteedToNotAlias
rewriter.replaceOp(deallocOp, replacements);
return success();
}
-
-private:
- BufferOriginAnalysis &analysis;
};
/// Split off memrefs to separate dealloc operations to reduce the number of
@@ -304,12 +309,12 @@ struct RemoveRetainedMemrefsGuaranteedToNotAlias
/// ```
struct SplitDeallocWhenNotAliasingAnyOther
: public OpRewritePattern<DeallocOp> {
- SplitDeallocWhenNotAliasingAnyOther(MLIRContext *context,
- BufferOriginAnalysis &analysis)
- : OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
+ SplitDeallocWhenNotAliasingAnyOther(MLIRContext *context)
+ : OpRewritePattern<DeallocOp>(context) {}
LogicalResult matchAndRewrite(DeallocOp deallocOp,
PatternRewriter &rewriter) const override {
+ auto analysis = getBufferOriginAnalysis(deallocOp);
Location loc = deallocOp.getLoc();
if (deallocOp.getMemrefs().size() <= 1)
return failure();
@@ -359,9 +364,6 @@ struct SplitDeallocWhenNotAliasingAnyOther
rewriter.replaceOp(deallocOp, replacements);
return success();
}
-
-private:
- BufferOriginAnalysis &analysis;
};
/// Check for every retained memref if a must-aliasing memref exists in the
@@ -389,12 +391,12 @@ struct SplitDeallocWhenNotAliasingAnyOther
/// don't have uses anymore.
struct RetainedMemrefAliasingAlwaysDeallocatedMemref
: public OpRewritePattern<DeallocOp> {
- RetainedMemrefAliasingAlwaysDeallocatedMemref(MLIRContext *context,
- BufferOriginAnalysis &analysis)
- : OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
+ RetainedMemrefAliasingAlwaysDeallocatedMemref(MLIRContext *context)
+ : OpRewritePattern<DeallocOp>(context) {}
LogicalResult matchAndRewrite(DeallocOp deallocOp,
PatternRewriter &rewriter) const override {
+ auto analysis = getBufferOriginAnalysis(deallocOp);
BitVector aliasesWithConstTrueMemref(deallocOp.getRetained().size());
SmallVector<Value> newMemrefs, newConditions;
for (auto [memref, cond] :
@@ -441,9 +443,6 @@ struct RetainedMemrefAliasingAlwaysDeallocatedMemref
return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
rewriter);
}
-
-private:
- BufferOriginAnalysis &analysis;
};
} // namespace
@@ -461,13 +460,11 @@ struct BufferDeallocationSimplificationPass
: public bufferization::impl::BufferDeallocationSimplificationPassBase<
BufferDeallocationSimplificationPass> {
void runOnOperation() override {
- BufferOriginAnalysis analysis(getOperation());
RewritePatternSet patterns(&getContext());
patterns.add<RemoveDeallocMemrefsContainedInRetained,
RemoveRetainedMemrefsGuaranteedToNotAlias,
SplitDeallocWhenNotAliasingAnyOther,
- RetainedMemrefAliasingAlwaysDeallocatedMemref>(&getContext(),
- analysis);
+ RetainedMemrefAliasingAlwaysDeallocatedMemref>(&getContext());
populateDeallocOpCanonicalizationPatterns(patterns, &getContext());
// We don't want that the block structure changes invalidating the
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
index b40a17cf800bf..3acc09d23b59e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
@@ -169,3 +169,32 @@ func.func @duplicate_memref(%arg0: memref<5xf32>, %arg1: memref<6xf32>, %c: i1)
// CHECK-LABEL: func @duplicate_memref(
// CHECK: %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<5xf32>) if (%{{.*}}) retain (%{{.*}} : memref<6xf32>)
// CHECK: return %[[r]]
+
+// -----
+
+module {
+ func.func @memref_cast_folding(%arg0: memref<4x4xf32>) -> memref<4x4xf32, 1> {
+ %cfalse = arith.constant false
+ %ctrue = arith.constant true
+
+ %memspacecast = memref.memory_space_cast %arg0 : memref<4x4xf32> to memref<4x4xf32, 1>
+ %cast = memref.cast %memspacecast : memref<4x4xf32, 1> to memref<?x?xf32, 1>
+ %cast_1 = memref.cast %cast : memref<?x?xf32, 1> to memref<4x4xf32, 1>
+
+ %5 = scf.if %cfalse -> (memref<4x4xf32, 1>) {
+ scf.yield %cast_1 : memref<4x4xf32, 1>
+ } else {
+ %7 = bufferization.clone %cast_1 : memref<4x4xf32, 1> to memref<4x4xf32, 1>
+ scf.yield %7 : memref<4x4xf32, 1>
+ }
+
+ %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 : memref<4x4xf32> -> memref<f32>, index, index, index, index, index
+ %base_buffer_5, %offset_6, %sizes_7:2, %strides_8:2 = memref.extract_strided_metadata %memspacecast : memref<4x4xf32, 1> -> memref<f32, 1>, index, index, index, index, index
+
+ %6 = bufferization.dealloc (%base_buffer, %base_buffer_5 : memref<f32>, memref<f32, 1>) if (%cfalse, %ctrue) retain (%5 : memref<4x4xf32, 1>)
+ return %memspacecast : memref<4x4xf32, 1>
+ }
+}
+
+// CHECK-LABEL: func @memref_cast_folding
+// CHECK: %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<f32, 1>) if (%{{.*}}) retain (%{{.*}} : memref<4x4xf32, 1>)
``````````
</details>
https://github.com/llvm/llvm-project/pull/210105
More information about the Mlir-commits
mailing list