[Mlir-commits] [mlir] [MLIR] Remove redundant dealloc memrefs with known true ownership (PR #228339)
Jiaqi He
llvmlistbot at llvm.org
Thu Oct 1 23:16:42 PDT 2026
https://github.com/heturing updated https://github.com/llvm/llvm-project/pull/228339
>From 35be76b93d5fe0d160822efc48f5fc5bcf504997 Mon Sep 17 00:00:00 2001
From: Jiaqi He <heturing at gmail.com>
Date: Thu, 1 Oct 2026 23:07:39 -0600
Subject: [PATCH] [MLIR] Remove redundant dealloc memrefs with known true
ownership
---
mlir/docs/OwnershipBasedBufferDeallocation.md | 3 +-
.../BufferDeallocationSimplification.cpp | 76 +++++++++++++++----
.../buffer-deallocation-simplification.mlir | 47 ++++++++++++
3 files changed, 110 insertions(+), 16 deletions(-)
diff --git a/mlir/docs/OwnershipBasedBufferDeallocation.md b/mlir/docs/OwnershipBasedBufferDeallocation.md
index f5fa01c4c49cc..b093c07d53733 100644
--- a/mlir/docs/OwnershipBasedBufferDeallocation.md
+++ b/mlir/docs/OwnershipBasedBufferDeallocation.md
@@ -375,7 +375,8 @@ The pass applies patterns for the following simplifications:
efficient lowering for this new `bufferization.dealloc` operation.
* Remove values from the 'memref' operand list when it is guaranteed to alias
with at least one value in the 'retained' list and may not alias any other
- value in the 'retain' list.
+ value in the 'retain' list whose updated ownership condition cannot be
+ statically resolved to true.
## Lower Deallocations Pass
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
index 06e90333fe469..aa94d8f86ad0a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocationSimplification.cpp
@@ -106,8 +106,9 @@ namespace {
/// Remove values from the `memref` operand list that are also present in the
/// `retained` list (or a guaranteed alias of it) because they will never
/// actually be deallocated. However, we also need to be certain about which
-/// other memrefs in the `retained` list can alias, i.e., there must not by any
-/// may-aliasing memref. This is necessary because the `dealloc` operation is
+/// other memrefs in the `retained` list can alias, i.e., there must not be any
+/// may-aliasing memref its updated ownership condition can be statically
+/// resolved to true. This is necessary because the `dealloc` operation is
/// defined to return one `i1` value per memref in the `retained` list which
/// represents the disjunction of the condition values corresponding to all
/// aliasing values in the `memref` list. In particular, this means that if
@@ -139,24 +140,30 @@ struct RemoveDeallocMemrefsContainedInRetained
: OpRewritePattern<DeallocOp>(context), analysis(analysis) {}
/// 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
- /// memref in the retain list, we cannot simply remove the memref as there
- /// could be situations in which it actually has to be deallocated. If it's
- /// 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 {
+ /// memref whose updated ownership condition cannot be resolved to true, and
+ /// at least one must-alias relation. If there is no must-aliasing memref in
+ /// the retain list, we cannot simply remove the memref as there could be
+ /// situations in which it actually has to be deallocated. If it's 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 DenseSet<Value> &retainedWithKnownTrueOwnership) const {
rewriter.setInsertionPointAfter(deallocOp);
- // Check that there is no may-aliasing memref and that at least one memref
+ // Check that every may-aliasing retained memref has an updated ownership
+ // condition known to be true and that at least one memref
// in the retain list aliases (because otherwise it might have to be
// deallocated in some situations and can thus not be dropped).
bool atLeastOneMustAlias = false;
for (Value retained : deallocOp.getRetained()) {
std::optional<bool> analysisResult =
analysis.isSameAllocation(retained, memref);
- if (!analysisResult.has_value())
- return failure();
+ if (!analysisResult.has_value()) {
+ if (!retainedWithKnownTrueOwnership.contains(retained))
+ return failure();
+ }
if (analysisResult == true)
atLeastOneMustAlias = true;
}
@@ -181,6 +188,39 @@ struct RemoveDeallocMemrefsContainedInRetained
return success();
}
+ /// Precompute sets of retained that have known true owernships. Since
+ /// matchAndRewrite call handleOneMemref twice on different sets of memrefs,
+ /// we precompute retained sets for both of them.
+ std::pair<DenseSet<Value>, DenseSet<Value>>
+ findRetainedWithKnownTrueOwnership(DeallocOp deallocOp) const {
+ DenseSet<Value> retainedWithKnownTrueOwnershipFromMemref;
+ DenseSet<Value> retainedWithKnownTrueOwnershipFromSource;
+
+ for (auto [memref, cond] :
+ llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
+
+ if (!matchPattern(cond, m_One()))
+ continue;
+
+ auto extractOp = memref.getDefiningOp<memref::ExtractStridedMetadataOp>();
+ Value source = extractOp ? extractOp.getOperand() : Value();
+ for (Value retained : deallocOp.getRetained()) {
+ std::optional<bool> analysisResult =
+ analysis.isSameAllocation(retained, memref);
+ if (analysisResult == true)
+ retainedWithKnownTrueOwnershipFromMemref.insert(retained);
+ if (source) {
+ analysisResult = analysis.isSameAllocation(retained, source);
+ if (analysisResult == true)
+ retainedWithKnownTrueOwnershipFromSource.insert(retained);
+ }
+ }
+ }
+
+ return {retainedWithKnownTrueOwnershipFromMemref,
+ retainedWithKnownTrueOwnershipFromSource};
+ }
+
LogicalResult matchAndRewrite(DeallocOp deallocOp,
PatternRewriter &rewriter) const override {
// There must not be any duplicates in the retain list anymore because we
@@ -190,17 +230,23 @@ struct RemoveDeallocMemrefsContainedInRetained
if (retained.size() != deallocOp.getRetained().size())
return failure();
+ auto [retainedWithKnownTrueOwnershipFromMemref,
+ retainedWithKnownTrueOwnershipFromSource] =
+ findRetainedWithKnownTrueOwnership(deallocOp);
+
SmallVector<Value> newMemrefs, newConditions;
for (auto [memref, cond] :
llvm::zip(deallocOp.getMemrefs(), deallocOp.getConditions())) {
- if (succeeded(handleOneMemref(deallocOp, memref, cond, rewriter)))
+ if (succeeded(handleOneMemref(deallocOp, memref, cond, rewriter,
+ retainedWithKnownTrueOwnershipFromMemref)))
continue;
if (auto extractOp =
memref.getDefiningOp<memref::ExtractStridedMetadataOp>())
- if (succeeded(handleOneMemref(deallocOp, extractOp.getOperand(), cond,
- rewriter)))
+ if (succeeded(handleOneMemref(
+ deallocOp, extractOp.getOperand(), cond, rewriter,
+ retainedWithKnownTrueOwnershipFromSource)))
continue;
newMemrefs.push_back(memref);
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
index 3acc09d23b59e..02b614884d6d5 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation-simplification.mlir
@@ -198,3 +198,50 @@ module {
// CHECK-LABEL: func @memref_cast_folding
// CHECK: %[[r:.*]] = bufferization.dealloc (%{{.*}} : memref<f32, 1>) if (%{{.*}}) retain (%{{.*}} : memref<4x4xf32, 1>)
+
+// -----
+
+func.func @partial_ownership(%a: memref<4xi32>, %b: memref<4xi32>, %pick: i1, %owned: i1) -> (memref<4xi32>, i1, i1, i1) {
+ %true = arith.constant true
+ %fresh = memref.alloc() : memref<4xi32>
+ %selected = scf.if %pick -> (memref<4xi32>) {
+ scf.yield %fresh : memref<4xi32>
+ } else {
+ scf.yield %a : memref<4xi32>
+ }
+ %r:3 = bufferization.dealloc(%a, %b, %selected : memref<4xi32>, memref<4xi32>, memref<4xi32>)
+ if (%true, %true, %owned)
+ retain (%a, %b, %fresh : memref<4xi32>, memref<4xi32>, memref<4xi32>)
+ return %fresh, %r#0, %r#1, %r#2 : memref<4xi32>, i1, i1, i1
+}
+
+// CHECK-LABEL: func.func @partial_ownership
+// CHECK: %[[TRUE:.*]] = arith.constant true
+// CHECK: %[[SELECTED:.*]] = scf.if
+// CHECK: %[[R:.*]]:3 = bufferization.dealloc (%[[SELECTED]] : memref<4xi32>) if (%{{.*}}) retain (%{{.*}}, %{{.*}}, %{{.*}} : memref<4xi32>, memref<4xi32>, memref<4xi32>)
+// CHECK: return %{{.*}}, %[[TRUE]], %[[TRUE]], %[[R]]#2
+
+// -----
+
+func.func @partial_ownership_generalized(%a: memref<4xi32>, %b: memref<4xi32>, %pick: i1, %owned: i1) -> (memref<4xi32>, i1, i1, i1) {
+ %true = arith.constant true
+ %fresh = memref.alloc() : memref<4xi32>
+ %c = memref.cast %b : memref<4xi32> to memref<?xi32>
+
+ %selected = scf.if %pick -> (memref<4xi32>) {
+ scf.yield %fresh : memref<4xi32>
+ } else {
+ scf.yield %a : memref<4xi32>
+ }
+
+ %r:3 = bufferization.dealloc(%a, %c, %selected : memref<4xi32>, memref<?xi32>, memref<4xi32>)
+ if (%true, %true, %owned)
+ retain (%a, %b, %fresh : memref<4xi32>, memref<4xi32>, memref<4xi32>)
+ return %fresh, %r#0, %r#1, %r#2 : memref<4xi32>, i1, i1, i1
+}
+
+// CHECK-LABEL: func.func @partial_ownership_generalized
+// CHECK: %[[TRUE:.*]] = arith.constant true
+// CHECK: %[[SELECTED:.*]] = scf.if
+// CHECK: %[[R:.*]]:3 = bufferization.dealloc (%[[SELECTED]] : memref<4xi32>) if (%{{.*}}) retain (%{{.*}}, %{{.*}}, %{{.*}} : memref<4xi32>, memref<4xi32>, memref<4xi32>)
+// CHECK: return %{{.*}}, %[[TRUE]], %[[TRUE]], %[[R]]#2
More information about the Mlir-commits
mailing list