[mlir][bufferization] Handle scf.if deallocs in static memory planner (#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. GitOrigin-RevId: 02960166f1a6814c0efc1c22417b5ab0ff37801e
diff --git a/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/include/mlir/Dialect/Bufferization/Transforms/Passes.td index 8408315..e37b8a0 100644 --- a/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -194,8 +194,14 @@ 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/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp index f1db0d5..6d5a9f0 100644 --- a/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp +++ b/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,41 @@ // 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 +101,13 @@ 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 +117,14 @@ 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 +132,35 @@ 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, - llvm::Statistic &numEligible, +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 +169,15 @@ 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 +187,21 @@ } for (memref::DeallocOp d : deallocs) { - if (d->getBlock() != allocOp->getBlock()) { - allocOp.emitError("dealloc is in a different block than the alloc; " - "run the deallocation pipeline before this pass"); + // 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("unstructured control flow is not supported"); 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 +317,13 @@ } } - // 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/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir b/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir index af900be..5bf7b27 100644 --- a/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir +++ b/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -223,3 +223,102 @@ 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 %b : memref<1024xf32> + } else { + scf.yield %a : 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/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir b/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir index 86af064..cf7388c 100644 --- a/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir +++ b/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
@@ -12,13 +12,13 @@ // ----- -// 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 via unstructured +// control flow (cf.br) should be an error. +func.func @error_escaping_dealloc(%cond: i1) { + // expected-error @+1 {{unstructured control flow is not supported}} %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 }