[Mlir-commits] [mlir] [mlir][bufferization] Handle arith.select-based deallocs in static memory planner (PR #209106)
Krish Gupta
llvmlistbot at llvm.org
Sun Jul 19 14:57:51 PDT 2026
https://github.com/KrxGu updated https://github.com/llvm/llvm-project/pull/209106
>From fbe630086b549c739d2c2686804a5ec9afd0b13c Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 13 Jul 2026 13:40:50 +0530
Subject: [PATCH 1/2] [mlir][bufferization] Handle arith.select-based deallocs
in static memory planner
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Allocs freed indirectly via arith.select chains were previously skipped.
This adds forward select-chain traversal so patterns like:
%2 = arith.select %c, %0, %1
memref.dealloc %2
are now handled correctly.
A group constraint ensures that all allocs sharing a select-based
dealloc are either all placed in the arena or all skipped — putting
one alloc in while leaving its peer out would break the dealloc.
Also fixes the O(n*m) block scan in buildAllocInfos by doing a single
upfront pass with a DenseMap index.
Tests added for single-alloc select, shared select-dealloc, and the
two-select two-dealloc pattern.
---
.../StaticMemoryPlannerAnalysis.cpp | 120 ++++++++++++------
.../static-memory-planner-analysis.mlir | 61 +++++++++
2 files changed, 143 insertions(+), 38 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index c2ac40a8427e7..770a46ff3de14 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,59 @@ 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;
+ // Walk allocs, find all potential deallocs via select chains.
+ SmallVector<AllocationCandidate> potentialCandidates;
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);
});
+ // Note: arith.select requires both operands to have the same type, so if
+ // one alloc in a select group has a static shape, all others must too.
+ // A mixed-eligibility group (some eligible, some not) is therefore
+ // impossible and no group-constraint fixpoint is needed.
+ SmallVector<AllocationCandidate> candidates;
+ for (auto &c : potentialCandidates) {
+ ++numEligible;
+ candidates.push_back(c);
+ }
return candidates;
}
@@ -180,8 +210,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 +227,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
+}
>From 30fd90d09690c96a63076c1df77d100448213f74 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 20 Jul 2026 03:27:35 +0530
Subject: [PATCH 2/2] [mlir][bufferization] Address review feedback on
arith.select dealloc pass
- Use BufferViewFlowOpInterface instead of hardcoded arith::SelectOp check,
so any future op implementing the interface is handled automatically
- Error out (not skip) when no dealloc or cross-block dealloc is found;
the only valid skip is dynamic shapes
- Remove redundant potentialCandidates vector copy in collectCandidates
- Replace erasedDeallocs with deallocsToErase set in rewriteAllocations
- Remove outdated comment about group constraint
- Add error tests for missing-dealloc and cross-block cases
---
.../StaticMemoryPlannerAnalysis.cpp | 95 ++++++++++---------
.../static-memory-planner-analysis.mlir | 40 ++------
.../static-memory-planner-errors.mlir | 24 +++++
3 files changed, 78 insertions(+), 81 deletions(-)
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 770a46ff3de14..f1db0d5ee55a4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -13,6 +13,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -50,12 +51,16 @@ struct AllocationCandidate {
// Helper utilities
//===----------------------------------------------------------------------===//
-/// Collect all dealloc ops that might free the given value, following
-/// arith.select chains. For example:
+/// Collect all dealloc ops that might free the given value, following ops
+/// that implement BufferViewFlowOpInterface (e.g. arith.select). 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) {
@@ -64,9 +69,11 @@ static void findPotentialDeallocs(Value value,
for (Operation *user : value.getUsers()) {
if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
deallocs.push_back(dealloc);
- } else if (auto select = dyn_cast<arith::SelectOp>(user)) {
- if (isa<MemRefType>(select.getType()))
- findPotentialDeallocs(select.getResult(), deallocs, visited);
+ } 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);
}
}
}
@@ -112,22 +119,20 @@ static int64_t buildAllocInfos(
}
/// Collect alloc/dealloc pairs eligible for arena placement.
-/// 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>
+/// 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.
+static LogicalResult
collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
- llvm::Statistic &numSkipNoDealloc,
- llvm::Statistic &numEligible) {
- // Walk allocs, find all potential deallocs via select chains.
- SmallVector<AllocationCandidate> potentialCandidates;
-
- funcOp->walk([&](memref::AllocOp allocOp) {
+ llvm::Statistic &numEligible,
+ SmallVector<AllocationCandidate> &candidates) {
+ bool walkFailed = false;
+ funcOp->walk([&](memref::AllocOp allocOp) -> WalkResult {
MemRefType memrefType = allocOp.getType();
if (!memrefType.hasStaticShape()) {
++numSkipDynamic;
- return;
+ return WalkResult::advance();
}
SmallVector<memref::DeallocOp> deallocs;
@@ -135,37 +140,32 @@ collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
findPotentialDeallocs(allocOp.getResult(), deallocs, visited);
if (deallocs.empty()) {
- ++numSkipNoDealloc;
- return;
+ allocOp.emitError("no dealloc found; run the deallocation pipeline "
+ "before this pass");
+ walkFailed = true;
+ return WalkResult::interrupt();
}
- // 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;
+ 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");
+ walkFailed = true;
+ return WalkResult::interrupt();
+ }
}
+ ++numEligible;
AllocationCandidate candidate;
candidate.alloc = allocOp;
candidate.deallocs = deallocs;
candidate.sizeInBytes = computeSizeInBytes(memrefType);
candidate.alignment = allocOp.getAlignment().value_or(1);
- potentialCandidates.push_back(candidate);
+ candidates.push_back(candidate);
+ return WalkResult::advance();
});
- // Note: arith.select requires both operands to have the same type, so if
- // one alloc in a select group has a static shape, all others must too.
- // A mixed-eligibility group (some eligible, some not) is therefore
- // impossible and no group-constraint fixpoint is needed.
- SmallVector<AllocationCandidate> candidates;
- for (auto &c : potentialCandidates) {
- ++numEligible;
- candidates.push_back(c);
- }
- return candidates;
+ return failure(walkFailed);
}
/// Create or obtain the arena buffer based on the arena mode.
@@ -210,14 +210,12 @@ 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;
+ SmallPtrSet<Operation *, 8> deallocsToErase;
SmallVector<Operation *> allocsToErase;
- // First replace all alloc results (rewires selects too), collect for erase.
+ // Replace all alloc results with views (rewires selects too).
for (auto &candidate : candidates) {
OpBuilder builder(candidate.alloc);
Location loc = candidate.alloc.getLoc();
@@ -229,13 +227,14 @@ static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
offsetIndex, SmallVector<Value>{});
candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
allocsToErase.push_back(candidate.alloc.getOperation());
+
+ for (memref::DeallocOp d : candidate.deallocs)
+ deallocsToErase.insert(d.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();
+ for (Operation *d : deallocsToErase)
+ d->erase();
// Erase allocs last (no users remain after replaceAllUsesWith).
for (Operation *allocOp : allocsToErase)
@@ -270,8 +269,10 @@ void StaticMemoryPlannerAnalysisPass::runOnOperation() {
}
// Step 1: Collect eligible allocation candidates.
- SmallVector<AllocationCandidate> candidates =
- collectCandidates(funcOp, numSkipDynamic, numSkipNoDealloc, numEligible);
+ SmallVector<AllocationCandidate> candidates;
+ if (failed(
+ collectCandidates(funcOp, numSkipDynamic, numEligible, candidates)))
+ return signalPassFailure();
if (candidates.empty())
return;
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 14c8dbf8fd0fc..af900be29a22e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -85,35 +85,7 @@ func.func @dynamic_shape_skipped(%n: index) {
// -----
-// Test 5: No dealloc - should be skipped
-// CHECK-LABEL: func @no_dealloc_skipped
-func.func @no_dealloc_skipped() {
- // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1024xf32>
- // CHECK-NOT: memref.subview
- %alloc = memref.alloc() : memref<1024xf32>
- return
-}
-
-// -----
-
-// Test 6: Dealloc in different block - should be skipped
-// CHECK-LABEL: func @different_block_skipped
-func.func @different_block_skipped(%cond: i1) {
- // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1024xf32>
- // CHECK: scf.if
- // CHECK: memref.dealloc %[[ALLOC]]
- // CHECK-NOT: memref.subview
- %alloc = memref.alloc() : memref<1024xf32>
- scf.if %cond {
- memref.dealloc %alloc : memref<1024xf32>
- scf.yield
- }
- return
-}
-
-// -----
-
-// Test 7: Multiple allocations with sequential offsets
+// Test 5: Multiple allocations with sequential offsets
// CHECK-LABEL: func @multiple_sequential
func.func @multiple_sequential() {
// Arena: 1024*4 + 512*4 + 2048*4 = 14336 bytes
@@ -140,7 +112,7 @@ func.func @multiple_sequential() {
// -----
-// Test 8: Alignment requirements with padding
+// Test 6: Alignment requirements with padding
// CHECK-LABEL: func @alignment_padding
func.func @alignment_padding() {
// Arena: 256*4 + 128*4 + 64*4 = 1792 bytes, alignment = lcm(128,64,128) = 128
@@ -167,7 +139,7 @@ func.func @alignment_padding() {
// -----
-// Test 9: LCM arena alignment (alignment=4, alignment=16 → lcm=16).
+// Test 7: LCM arena alignment (alignment=4, alignment=16 → lcm=16).
// For power-of-2 alignments lcm equals max, but lcm is the correct
// general formula. Arena must be aligned to 16 so that all views are
// correctly aligned regardless of their individual requirements.
@@ -193,7 +165,7 @@ func.func @lcm_alignment() {
// -----
-// Test 10: Single alloc freed via arith.select-based dealloc.
+// Test 8: Single alloc freed via arith.select-based dealloc.
// CHECK-LABEL: func @select_single_alloc
func.func @select_single_alloc() {
%c = arith.constant true
@@ -210,7 +182,7 @@ func.func @select_single_alloc() {
// -----
-// Test 11: Two allocs freed via a shared select-based dealloc.
+// Test 9: 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() {
@@ -231,7 +203,7 @@ func.func @select_shared_dealloc() {
// -----
-// Test 12: Two allocs, two select-based deallocs (mentor's canonical example).
+// Test 10: 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() {
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
new file mode 100644
index 0000000000000..86af0649c67ba
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(static-memory-planner-analysis))" \
+// RUN: -split-input-file -verify-diagnostics
+
+// -----
+
+// Test 1: Alloc with no dealloc should be an error (not silently skipped).
+func.func @error_no_dealloc() {
+ // expected-error @+1 {{no dealloc found; run the deallocation pipeline before this pass}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ return
+}
+
+// -----
+
+// 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}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ scf.if %cond {
+ memref.dealloc %alloc : memref<1024xf32>
+ scf.yield
+ }
+ return
+}
More information about the Mlir-commits
mailing list