[Mlir-commits] [mlir] [mlir][bufferization] Handle arith.select-based deallocs in static memory planner (PR #209106)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jul 13 01:29:04 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-bufferization
@llvm/pr-subscribers-mlir
Author: Krish Gupta (KrxGu)
<details>
<summary>Changes</summary>
The static memory planner currently skips any allocation that doesn't have a direct `memref.dealloc` user. This is overly conservative, after running `ownership-based-buffer-deallocation`, it's common to see patterns like:
`%2 = arith.select %c, %0, %1 : memref<1024xf32>`
`memref.dealloc %2 : memref<1024xf32>`
where both `%0` and `%1` get skipped with `++numSkipNoDealloc` even though their lifetimes are well-defined.
This patch teaches `collectCandidates` to follow `arith.select` chains when looking for potential deallocs. We traverse the use-def graph forward from each alloc, collecting any `memref.dealloc` ops reachable through select results.
Since a single select-based dealloc can conditionally free one of several allocs, we enforce a group constraint: all allocs that share a dealloc via a select must either all go into the arena or all be skipped. Without this, we could end up with a `memref.view` (an arena slice) and a raw alloc being fed into the same select, making the resulting dealloc invalid.
The group constraint is computed with a simple fixpoint iteration , if any member of a group is ineligible, the whole group is dropped.
The lifetime indices (`timeStart`/`timeEnd`) in `buildAllocInfos` are also fixed: the old code did one block scan per candidate (O(n×m)). This replaces it with a single pass upfront using a `DenseMap`, and sets `timeEnd` conservatively to the latest dealloc index across all potential deallocs for an alloc.
Tests added for:
- Single alloc freed via a self-select dealloc
- Two allocs sharing one select-based dealloc (group constraint active)
- The two-select two-dealloc pattern from the design discussion
---
Full diff: https://github.com/llvm/llvm-project/pull/209106.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp (+113-38)
- (modified) mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir (+61)
``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index c2ac40a8427e7..feb7594ac43c9 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/Debug.h"
#include <numeric>
@@ -34,10 +35,12 @@ using namespace mlir;
namespace {
-/// A candidate allocation with its matching deallocation and assigned offset.
+/// A candidate allocation with its matching deallocation(s) and assigned
+/// offset. An alloc may be freed indirectly through arith.select chains,
+/// yielding multiple potential deallocs — all must be in the same block.
struct AllocationCandidate {
memref::AllocOp alloc;
- memref::DeallocOp dealloc;
+ SmallVector<memref::DeallocOp> deallocs;
int64_t offset = 0; // Offset in bytes from arena start (assigned by planner)
int64_t sizeInBytes = 0; // Size in bytes
int64_t alignment = 1; // Required alignment in bytes
@@ -47,18 +50,25 @@ struct AllocationCandidate {
// Helper utilities
//===----------------------------------------------------------------------===//
-/// Finds the unique dealloc operation for a given alloc value.
-/// Returns nullptr if there are zero or multiple deallocs.
-static memref::DeallocOp findUniqueDealloc(Value allocValue) {
- memref::DeallocOp deallocOp = nullptr;
- for (Operation *user : allocValue.getUsers()) {
+/// Collect all dealloc ops that might free the given value, following
+/// arith.select chains. For example:
+/// %0 = memref.alloc()
+/// %2 = arith.select %c, %0, %1
+/// memref.dealloc %2 <- this covers %0 conditionally
+/// `visited` prevents cycles in the use-def graph.
+static void findPotentialDeallocs(Value value,
+ SmallVectorImpl<memref::DeallocOp> &deallocs,
+ SmallPtrSetImpl<Value> &visited) {
+ if (!visited.insert(value).second)
+ return;
+ for (Operation *user : value.getUsers()) {
if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
- if (deallocOp)
- return nullptr; // Multiple deallocs found
- deallocOp = dealloc;
+ deallocs.push_back(dealloc);
+ } else if (auto select = dyn_cast<arith::SelectOp>(user)) {
+ if (isa<MemRefType>(select.getType()))
+ findPotentialDeallocs(select.getResult(), deallocs, visited);
}
}
- return deallocOp;
}
/// Compute the size in bytes for a memref type.
@@ -70,25 +80,31 @@ static int64_t computeSizeInBytes(MemRefType memrefType) {
/// Build lifetime-annotated allocation descriptors from candidates.
/// Returns the arena alignment (LCM of all individual alignments).
+/// Uses a single block scan (O(n+m)) instead of one scan per candidate.
static int64_t buildAllocInfos(
MutableArrayRef<AllocationCandidate> candidates,
SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
+ // Build an op-index map with a single pass over the block.
+ DenseMap<Operation *, int64_t> opIndex;
+ if (!candidates.empty()) {
+ Block *block = candidates.front().alloc->getBlock();
+ int64_t idx = 0;
+ for (Operation &op : *block)
+ opIndex[&op] = idx++;
+ }
+
int64_t arenaAlignment = 1;
for (auto &candidate : candidates) {
bufferization::MemoryPlannerAlloc info;
info.sizeInBytes = candidate.sizeInBytes;
info.alignment = candidate.alignment;
-
- Block *block = candidate.alloc->getBlock();
- int64_t opIdx = 0;
- for (Operation &op : *block) {
- if (&op == candidate.alloc.getOperation())
- info.timeStart = opIdx;
- if (&op == candidate.dealloc.getOperation())
- info.timeEnd = opIdx;
- ++opIdx;
- }
-
+ info.timeStart = opIndex.lookup(candidate.alloc.getOperation());
+ // Conservative: timeEnd = latest dealloc index among all potential
+ // deallocs.
+ int64_t timeEnd = 0;
+ for (memref::DeallocOp d : candidate.deallocs)
+ timeEnd = std::max(timeEnd, opIndex.lookup(d.getOperation()));
+ info.timeEnd = timeEnd;
allocInfos.push_back(info);
arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
}
@@ -96,45 +112,90 @@ static int64_t buildAllocInfos(
}
/// Collect alloc/dealloc pairs eligible for arena placement.
-/// An allocation is eligible if it has a static shape and a unique dealloc
-/// in the same block.
+/// An allocation is eligible if it has a static shape and all of its
+/// potential deallocs (including those reached via arith.select chains)
+/// are in the same block. A group constraint ensures that all allocs
+/// sharing a select-based dealloc are either all eligible or all skipped.
static SmallVector<AllocationCandidate>
collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
llvm::Statistic &numSkipNoDealloc,
llvm::Statistic &numEligible) {
- SmallVector<AllocationCandidate> candidates;
+ // Phase 1: walk allocs, find all potential deallocs via select chains.
+ SmallVector<AllocationCandidate> potentialCandidates;
+ // Track which allocs share each dealloc (for the group constraint).
+ DenseMap<Operation *, SmallVector<Value>> deallocToAllocs;
funcOp->walk([&](memref::AllocOp allocOp) {
MemRefType memrefType = allocOp.getType();
-
- // Skip dynamic shapes
if (!memrefType.hasStaticShape()) {
++numSkipDynamic;
return;
}
- // Find unique dealloc in the same block
- memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
- if (!deallocOp) {
+ SmallVector<memref::DeallocOp> deallocs;
+ SmallPtrSet<Value, 8> visited;
+ findPotentialDeallocs(allocOp.getResult(), deallocs, visited);
+
+ if (deallocs.empty()) {
++numSkipNoDealloc;
return;
}
- if (deallocOp->getBlock() != allocOp->getBlock()) {
+ // All deallocs must be in the same block as the alloc.
+ bool allSameBlock = llvm::all_of(deallocs, [&](memref::DeallocOp d) {
+ return d->getBlock() == allocOp->getBlock();
+ });
+ if (!allSameBlock) {
++numSkipNoDealloc;
return;
}
- // This allocation is eligible
- ++numEligible;
AllocationCandidate candidate;
candidate.alloc = allocOp;
- candidate.dealloc = deallocOp;
+ candidate.deallocs = deallocs;
candidate.sizeInBytes = computeSizeInBytes(memrefType);
candidate.alignment = allocOp.getAlignment().value_or(1);
- candidates.push_back(candidate);
+ potentialCandidates.push_back(candidate);
+
+ for (memref::DeallocOp d : deallocs)
+ deallocToAllocs[d.getOperation()].push_back(allocOp.getResult());
});
+ // Phase 2: enforce group constraint — if a dealloc covers multiple allocs,
+ // all of them must be eligible or none are. Iterate to fixpoint.
+ SmallPtrSet<Value, 16> validAllocs;
+ for (auto &c : potentialCandidates)
+ validAllocs.insert(c.alloc.getResult());
+
+ bool changed = true;
+ while (changed) {
+ changed = false;
+ for (auto &c : potentialCandidates) {
+ if (!validAllocs.contains(c.alloc.getResult()))
+ continue;
+ for (memref::DeallocOp d : c.deallocs) {
+ for (Value peer : deallocToAllocs[d.getOperation()]) {
+ if (!validAllocs.contains(peer)) {
+ validAllocs.erase(c.alloc.getResult());
+ changed = true;
+ break;
+ }
+ }
+ if (!validAllocs.contains(c.alloc.getResult()))
+ break;
+ }
+ }
+ }
+
+ SmallVector<AllocationCandidate> candidates;
+ for (auto &c : potentialCandidates) {
+ if (validAllocs.contains(c.alloc.getResult())) {
+ ++numEligible;
+ candidates.push_back(c);
+ } else {
+ ++numSkipNoDealloc;
+ }
+ }
return candidates;
}
@@ -180,8 +241,14 @@ static FailureOr<Value> createArena(OpBuilder &builder,
}
/// Replace each alloc/dealloc pair with a memref.view into the arena.
+/// Handles select-chained deallocs: a single dealloc may cover multiple allocs,
+/// so we track erased deallocs to avoid double-erase.
static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
Value arenaValue) {
+ SmallPtrSet<Operation *, 8> erasedDeallocs;
+ SmallVector<Operation *> allocsToErase;
+
+ // First replace all alloc results (rewires selects too), collect for erase.
for (auto &candidate : candidates) {
OpBuilder builder(candidate.alloc);
Location loc = candidate.alloc.getLoc();
@@ -191,11 +258,19 @@ static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
arith::ConstantIndexOp::create(builder, loc, candidate.offset);
auto view = memref::ViewOp::create(builder, loc, originalType, arenaValue,
offsetIndex, SmallVector<Value>{});
-
candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
- candidate.alloc.erase();
- candidate.dealloc.erase();
+ allocsToErase.push_back(candidate.alloc.getOperation());
}
+
+ // Erase deallocs first (they may reference alloc results via selects).
+ for (auto &candidate : candidates)
+ for (memref::DeallocOp d : candidate.deallocs)
+ if (erasedDeallocs.insert(d.getOperation()).second)
+ d.erase();
+
+ // Erase allocs last (no users remain after replaceAllUsesWith).
+ for (Operation *allocOp : allocsToErase)
+ allocOp->erase();
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
index a80c0e13adc21..14c8dbf8fd0fc 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -190,3 +190,64 @@ func.func @lcm_alignment() {
memref.dealloc %alloc1 : memref<3xi32>
return
}
+
+// -----
+
+// Test 10: Single alloc freed via arith.select-based dealloc.
+// CHECK-LABEL: func @select_single_alloc
+func.func @select_single_alloc() {
+ %c = arith.constant true
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<4096xi8>
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[V:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<4096xi8> to memref<1024xf32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %alloc = memref.alloc() : memref<1024xf32>
+ %sel = arith.select %c, %alloc, %alloc : memref<1024xf32>
+ memref.dealloc %sel : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 11: Two allocs freed via a shared select-based dealloc.
+// Group constraint: both must be eligible together or neither is.
+// CHECK-LABEL: func @select_shared_dealloc
+func.func @select_shared_dealloc() {
+ %c = arith.constant true
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<8192xi8>
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[V0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<8192xi8> to memref<1024xf32>
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %[[V1:.*]] = memref.view %[[ARENA]][%[[C4096]]][] : memref<8192xi8> to memref<1024xf32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %a = memref.alloc() : memref<1024xf32>
+ %b = memref.alloc() : memref<1024xf32>
+ %sel = arith.select %c, %a, %b : memref<1024xf32>
+ memref.dealloc %sel : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 12: Two allocs, two select-based deallocs (mentor's canonical example).
+// %a freed via dealloc(%sel1) or dealloc(%sel2), %b likewise.
+// CHECK-LABEL: func @select_two_deallocs
+func.func @select_two_deallocs() {
+ %c = arith.constant true
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<8192xi8>
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0]]][] : memref<8192xi8> to memref<1024xf32>
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C4096]]][] : memref<8192xi8> to memref<1024xf32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %a = memref.alloc() : memref<1024xf32>
+ %b = memref.alloc() : memref<1024xf32>
+ %sel1 = arith.select %c, %a, %b : memref<1024xf32>
+ memref.dealloc %sel1 : memref<1024xf32>
+ %sel2 = arith.select %c, %b, %a : memref<1024xf32>
+ memref.dealloc %sel2 : memref<1024xf32>
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/209106
More information about the Mlir-commits
mailing list