[Mlir-commits] [mlir] [mlir][bufferization] Handle scf.if deallocs in static memory planner (PR #213634)
Krish Gupta
llvmlistbot at llvm.org
Mon Aug 3 02:46:15 PDT 2026
https://github.com/KrxGu created https://github.com/llvm/llvm-project/pull/213634
Extends the static memory planner (#209106) to handle two scf.if patterns that previously errored or were silently missed.
**What changed**
Replaced the hand-rolled `BufferViewFlowOpInterface` DFS with the shared `BufferViewFlowAnalysis`. This covers arith.select, scf.if/for results, cf branches, and view ops in one place — no new interface needed.
Two new cases are handled:
1. Alloc flows through an `scf.if` result; `dealloc` is on that result. `resolve()` finds the alias and picks up the dealloc.
2. Alloc is in the entry block; `dealloc` is inside an `scf.if` body. `findAncestorOpInBlock` anchors the lifetime to the enclosing `scf.if` — conservative but correct.
A reverse-alias guard (`resolveReverse`) handles the unsafe case where a dealloc may also free a *nested* alloc not managed by the arena. That alloc is conservatively skipped rather than miscompiled.
**Test changes**
- Tests 11–14 added: scf.if nested dealloc, scf.if result alias, nested alloc skip, shared-dealloc conservative skip.
- Error test 2 updated: scf.if-nested dealloc is now valid, replaced with a `cf.br` sibling-block escaping case.
>From ef49a4ac07747b1ce512a77b6b69a6bf7cf9b703 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 3 Aug 2026 15:14:37 +0530
Subject: [PATCH] [mlir][bufferization] Handle scf.if deallocs in static memory
planner
Switch dealloc discovery from a hand-rolled BufferViewFlowOpInterface
DFS to the shared BufferViewFlowAnalysis, which handles arith.select,
scf.if/for results, cf branches, and view ops uniformly.
Two new patterns now work:
- Alloc flows through an scf.if result; dealloc is on that result.
- Dealloc is inside an scf.if body; lifetime anchored conservatively
to the enclosing scf.if via findAncestorOpInBlock.
A reverse-alias guard (resolveReverse) skips allocs whose dealloc also
frees a nested alloc not managed by the arena, avoiding unsafe erasure.
The cross-block error test is updated: scf.if-nested deallocs are now
valid, replaced with a cf.br sibling-block escape case. Four new
FileCheck tests cover the new patterns and conservative skips.
Extends #209106.
---
.../Bufferization/Transforms/Passes.td | 10 +-
.../StaticMemoryPlannerAnalysis.cpp | 144 ++++++++++++------
.../static-memory-planner-analysis.mlir | 99 ++++++++++++
.../static-memory-planner-errors.mlir | 14 +-
4 files changed, 213 insertions(+), 54 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 8408315dda607..e37b8a029aac6 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -194,8 +194,14 @@ def StaticMemoryPlannerAnalysisPass
For each `memref.alloc` the pass checks a conservative eligibility
envelope:
- Static memref shape.
- - Unique same-block `memref.dealloc`.
- - Allocations in nested blocks are ignored for now.
+ - The allocation lives directly in the function's entry block
+ (allocations nested in a region are skipped for now).
+ - Every `memref.dealloc` that may free the buffer is anchored in the entry
+ block. Deallocs reached indirectly through `arith.select`, `scf.if`
+ results, `cf` branches, or view ops are handled via the shared
+ `BufferViewFlowAnalysis`; a dealloc nested inside an entry-block op
+ (e.g. an `scf.if` body) is accepted and bounds the lifetime
+ conservatively.
Eligible allocations are packed into a single arena using a configurable
planning algorithm (see the `algorithm` option). The arena is an i8 byte
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index f1db0d5ee55a4..6087e2ff98ee6 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -13,7 +13,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -51,31 +51,42 @@ struct AllocationCandidate {
// Helper utilities
//===----------------------------------------------------------------------===//
-/// Collect all dealloc ops that might free the given value, following ops
-/// that implement BufferViewFlowOpInterface (e.g. arith.select). For example:
+/// Collect all dealloc ops that might free the given alloc value. Instead of a
+/// bespoke traversal, this uses the shared `BufferViewFlowAnalysis`, which
+/// already models all the ways a buffer can flow to a dealloc:
+/// - `arith.select` (via BufferViewFlowOpInterface)
+/// - `scf.if`/`scf.for` (via RegionBranchOpInterface region/result wiring)
+/// - `cf.br`/`cf.cond_br` (via BranchOpInterface block arguments)
+/// - `memref.view`/subview (via ViewLikeOpInterface)
+/// `analysis.resolve(alloc)` returns the forward alias set (the alloc plus every
+/// value it may flow into); a dealloc on any of those aliases frees the alloc.
+/// 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.
-///
-/// TODO: This relies on BufferViewFlowOpInterface external models being
-/// registered for the ops in the IR (e.g. via
-/// arith::registerBufferViewFlowOpInterfaceExternalModels).
-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)) {
- deallocs.push_back(dealloc);
- } else if (dyn_cast<bufferization::BufferViewFlowOpInterface>(user)) {
- // Follow any op that propagates buffer values to its results.
- for (Value result : user->getResults())
- if (isa<MemRefType>(result.getType()))
- findPotentialDeallocs(result, deallocs, visited);
- }
- }
+/// memref.dealloc %2 <- covers %0 conditionally (via alias set)
+/// %3 = scf.if %c { yield %0 } else { yield %1 }
+/// memref.dealloc %3 <- also covers %0 conditionally
+static void collectDeallocs(Value alloc,
+ const BufferViewFlowAnalysis &analysis,
+ SmallVectorImpl<memref::DeallocOp> &deallocs) {
+ for (Value alias : analysis.resolve(alloc))
+ for (Operation *user : alias.getUsers())
+ if (auto dealloc = dyn_cast<memref::DeallocOp>(user))
+ deallocs.push_back(dealloc);
+}
+
+/// Return the set of allocation ops whose buffer may be freed by `dealloc`,
+/// i.e. the terminal `memref.alloc` sources that flow into the dealloc operand.
+/// Uses the reverse alias set so that a dealloc reached through a `scf.if`
+/// result or `arith.select` is attributed to every alloc it may free.
+static SmallVector<memref::AllocOp>
+findFreedAllocs(memref::DeallocOp dealloc,
+ const BufferViewFlowAnalysis &analysis) {
+ SmallVector<memref::AllocOp> allocs;
+ for (Value source : analysis.resolveReverse(dealloc.getMemref()))
+ if (auto allocOp = source.getDefiningOp<memref::AllocOp>())
+ allocs.push_back(allocOp);
+ return allocs;
}
/// Compute the size in bytes for a memref type.
@@ -91,12 +102,13 @@ static int64_t computeSizeInBytes(MemRefType memrefType) {
static int64_t buildAllocInfos(
MutableArrayRef<AllocationCandidate> candidates,
SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
- // Build an op-index map with a single pass over the block.
+ // Build an op-index map with a single pass over the plan block.
DenseMap<Operation *, int64_t> opIndex;
+ Block *planBlock = nullptr;
if (!candidates.empty()) {
- Block *block = candidates.front().alloc->getBlock();
+ planBlock = candidates.front().alloc->getBlock();
int64_t idx = 0;
- for (Operation &op : *block)
+ for (Operation &op : *planBlock)
opIndex[&op] = idx++;
}
@@ -106,11 +118,14 @@ static int64_t buildAllocInfos(
info.sizeInBytes = candidate.sizeInBytes;
info.alignment = candidate.alignment;
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()));
+ // Conservative: timeEnd = latest dealloc position among all potential
+ // deallocs. A dealloc may be nested (e.g. inside an scf.if body); its
+ // lifetime contribution is bounded by the enclosing op in the plan block.
+ int64_t timeEnd = info.timeStart;
+ for (memref::DeallocOp d : candidate.deallocs) {
+ Operation *anchor = planBlock->findAncestorOpInBlock(*d.getOperation());
+ timeEnd = std::max(timeEnd, opIndex.lookup(anchor));
+ }
info.timeEnd = timeEnd;
allocInfos.push_back(info);
arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
@@ -118,15 +133,35 @@ static int64_t buildAllocInfos(
return arenaAlignment;
}
-/// Collect alloc/dealloc pairs eligible for arena placement.
-/// An allocation is eligible if it has a static shape and its deallocs
-/// (including those reached via BufferViewFlowOpInterface chains)
-/// are in the same block. Allocations with dynamic shapes are skipped.
-/// Missing deallocs or cross-block deallocs are reported as errors.
+/// Collect alloc/dealloc groups eligible for arena placement.
+///
+/// Eligibility uses the shared `BufferViewFlowAnalysis` so that buffers flowing
+/// through `arith.select`, `scf.if`/`scf.for` results, `cf` branches, or view
+/// ops are handled uniformly. An allocation is eligible when:
+/// - it has a static shape (dynamic shapes are silently skipped), and
+/// - it lives directly in the function's entry block (allocs nested in a
+/// region are skipped for now), and
+/// - every dealloc that may free it is anchored in that same entry block --
+/// either directly, or nested inside an op of that block (e.g. an
+/// `scf.if` body), which conservatively bounds the lifetime.
+/// A dealloc that escapes the entry block entirely (e.g. lives in a sibling
+/// `cf` block) is reported as an error, as is an alloc with no dealloc.
+///
+/// To keep the rewrite safe, a dealloc is only accepted if *all* allocs it may
+/// free (per the reverse alias set) are themselves candidates in this block;
+/// otherwise erasing it during the rewrite could leak or double-free a buffer
+/// that is not managed by the arena.
static LogicalResult
-collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
+collectCandidates(FunctionOpInterface funcOp,
+ const BufferViewFlowAnalysis &analysis,
+ llvm::Statistic &numSkipDynamic, llvm::Statistic &numSkipNested,
llvm::Statistic &numEligible,
SmallVector<AllocationCandidate> &candidates) {
+ // All candidates are planned relative to the function's entry block.
+ if (funcOp.getFunctionBody().empty())
+ return success();
+ Block *planBlock = &funcOp.getFunctionBody().front();
+
bool walkFailed = false;
funcOp->walk([&](memref::AllocOp allocOp) -> WalkResult {
MemRefType memrefType = allocOp.getType();
@@ -135,9 +170,15 @@ collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
return WalkResult::advance();
}
+ // Only plan allocs that live directly in the entry block. Allocs nested in
+ // a region (loop/conditional body) are skipped for now.
+ if (allocOp->getBlock() != planBlock) {
+ ++numSkipNested;
+ return WalkResult::advance();
+ }
+
SmallVector<memref::DeallocOp> deallocs;
- SmallPtrSet<Value, 8> visited;
- findPotentialDeallocs(allocOp.getResult(), deallocs, visited);
+ collectDeallocs(allocOp.getResult(), analysis, deallocs);
if (deallocs.empty()) {
allocOp.emitError("no dealloc found; run the deallocation pipeline "
@@ -147,12 +188,22 @@ collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
}
for (memref::DeallocOp d : deallocs) {
- if (d->getBlock() != allocOp->getBlock()) {
- allocOp.emitError("dealloc is in a different block than the alloc; "
+ // The dealloc must be anchored in the plan block (directly or via an
+ // enclosing op such as an scf.if). A dealloc in a sibling block escapes.
+ if (!planBlock->findAncestorOpInBlock(*d.getOperation())) {
+ allocOp.emitError("dealloc is not reachable from the alloc's block; "
"run the deallocation pipeline before this pass");
walkFailed = true;
return WalkResult::interrupt();
}
+ // Every alloc that this dealloc may free must also be an entry-block
+ // candidate; otherwise erasing it during the rewrite is unsafe.
+ for (memref::AllocOp freed : findFreedAllocs(d, analysis)) {
+ if (freed->getBlock() != planBlock) {
+ ++numSkipNested;
+ return WalkResult::advance();
+ }
+ }
}
++numEligible;
@@ -268,10 +319,13 @@ void StaticMemoryPlannerAnalysisPass::runOnOperation() {
}
}
- // Step 1: Collect eligible allocation candidates.
+ // Step 1: Collect eligible allocation candidates. The buffer view-flow
+ // analysis models how buffers flow through selects, scf.if results, branches,
+ // and view ops so we can find deallocs and freed allocs uniformly.
+ BufferViewFlowAnalysis analysis(funcOp);
SmallVector<AllocationCandidate> candidates;
- if (failed(
- collectCandidates(funcOp, numSkipDynamic, numEligible, candidates)))
+ if (failed(collectCandidates(funcOp, analysis, numSkipDynamic, numSkipNested,
+ numEligible, candidates)))
return signalPassFailure();
if (candidates.empty())
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 af900be29a22e..d72aea5b19cd0 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -223,3 +223,102 @@ func.func @select_two_deallocs() {
memref.dealloc %sel2 : memref<1024xf32>
return
}
+
+// -----
+
+// Test 11: Deallocs nested inside scf.if bodies (mentor case_2).
+// Both allocs live in the entry block; each dealloc is anchored by the
+// enclosing scf.if, so both are eligible via the buffer view-flow analysis.
+// CHECK-LABEL: func @scf_if_nested_deallocs
+func.func @scf_if_nested_deallocs(%c: i1, %d: i1) {
+ // 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>
+ scf.if %c {
+ memref.dealloc %a : memref<1024xf32>
+ }
+ scf.if %d {
+ memref.dealloc %b : memref<1024xf32>
+ }
+ return
+}
+
+// -----
+
+// Test 12: Allocs flow through scf.if results, then deallocated (mentor case_1).
+// The analysis follows the scf.if result aliases back to %a and %b, so both
+// are planned and the yielded views are rewired automatically.
+// CHECK-LABEL: func @scf_if_result_aliases
+func.func @scf_if_result_aliases(%c: i1) {
+ // 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
+ // CHECK: scf.if
+ // CHECK: scf.yield %[[V0]]
+ // CHECK: scf.yield %[[V1]]
+ %a = memref.alloc() : memref<1024xf32>
+ %b = memref.alloc() : memref<1024xf32>
+ %0 = scf.if %c -> memref<1024xf32> {
+ scf.yield %a : memref<1024xf32>
+ } else {
+ scf.yield %b : memref<1024xf32>
+ }
+ %1 = scf.if %c -> memref<1024xf32> {
+ scf.yield %a : memref<1024xf32>
+ } else {
+ scf.yield %b : memref<1024xf32>
+ }
+ memref.dealloc %0 : memref<1024xf32>
+ memref.dealloc %1 : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 13: Alloc nested inside an scf.if body is left untouched (not planned).
+// Only entry-block allocs are planned; the nested %b keeps its alloc/dealloc.
+// CHECK-LABEL: func @scf_if_nested_alloc_skipped
+func.func @scf_if_nested_alloc_skipped(%c: i1) {
+ // CHECK-NOT: memref.view
+ // CHECK: scf.if
+ // CHECK: memref.alloc() : memref<1024xf32>
+ // CHECK: memref.dealloc
+ scf.if %c {
+ %b = memref.alloc() : memref<1024xf32>
+ memref.dealloc %b : memref<1024xf32>
+ }
+ return
+}
+
+// -----
+
+// Test 14: A dealloc that may free both an entry-block alloc and a nested
+// alloc (mentor case_3) is conservatively skipped: erasing it would be unsafe
+// for the buffer that is not managed by the arena.
+// CHECK-LABEL: func @scf_if_shared_nested_dealloc_skipped
+func.func @scf_if_shared_nested_dealloc_skipped(%c: i1) {
+ // CHECK: memref.alloc() : memref<1024xf32>
+ // CHECK-NOT: memref.view
+ // CHECK: scf.if
+ // CHECK: memref.dealloc
+ %a = memref.alloc() : memref<1024xf32>
+ %0 = scf.if %c -> memref<1024xf32> {
+ memref.dealloc %a : memref<1024xf32>
+ %b = memref.alloc() : memref<1024xf32>
+ scf.yield %b : memref<1024xf32>
+ } else {
+ scf.yield %a : memref<1024xf32>
+ }
+ memref.dealloc %0 : memref<1024xf32>
+ return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
index 86af0649c67ba..f3dc012b7b840 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
@@ -12,13 +12,13 @@ func.func @error_no_dealloc() {
// -----
-// Test 2: Alloc whose dealloc is in a different block should be an error.
-func.func @error_cross_block_dealloc(%cond: i1) {
- // expected-error @+1 {{dealloc is in a different block than the alloc; run the deallocation pipeline before this pass}}
+// Test 2: Alloc whose dealloc escapes to a sibling block (not reachable from
+// the alloc's block, even via region control flow) should be an error.
+func.func @error_escaping_dealloc(%cond: i1) {
+ // expected-error @+1 {{dealloc is not reachable from the alloc's block; run the deallocation pipeline before this pass}}
%alloc = memref.alloc() : memref<1024xf32>
- scf.if %cond {
- memref.dealloc %alloc : memref<1024xf32>
- scf.yield
- }
+ cf.br ^bb1
+^bb1:
+ memref.dealloc %alloc : memref<1024xf32>
return
}
More information about the Mlir-commits
mailing list