[Mlir-commits] [mlir] [mlir][bufferization] Add static memory planner pass for compile-time buffer allocation (PR #205125)
Krish Gupta
llvmlistbot at llvm.org
Mon Jun 29 01:56:34 PDT 2026
https://github.com/KrxGu updated https://github.com/llvm/llvm-project/pull/205125
>From f5111c17c2f6c050e25811c3168e8f2541e084ec Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Sat, 28 Mar 2026 10:24:15 +0530
Subject: [PATCH 01/13] [mlir][bufferization] add same-block alloc lifetime
analysis for static buffer planning
Adds an analysis-only pass that finds same-block memref.alloc/dealloc pairs
eligible for static memory reuse. For each eligible alloc it computes a
conservative alias-aware lifetime interval using BufferViewFlowAnalysis,
collects size/alignment metadata, and emits op remarks. Ineligible allocs
get a skip reason (dynamic shape, nested in loop/conditional, no unique
same-block dealloc, escaping alias). Aggregate counts are tracked as pass
statistics.
No IR mutations. Intended to run after the deallocation pipeline
(ownership-based-buffer-deallocation + bufferization-lower-deallocations).
This is the first upstream step for a buffer reuse pass discussed in the
MLIR discourse thread (RFC: GSoC buffer reuse pass for non-overlapping
allocations after lower-deallocations).
Adds six lit/FileCheck tests covering the main eligibility paths.
Signed-off-by: KrxGu <krishom70 at gmail.com>
---
.../Bufferization/Transforms/Passes.td | 62 ++++
.../Bufferization/Transforms/CMakeLists.txt | 1 +
.../StaticMemoryPlannerAnalysis.cpp | 272 ++++++++++++++++++
.../static-memory-planner-analysis.mlir | 107 +++++++
4 files changed, 442 insertions(+)
create mode 100644 mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index cd28bd6cf73a5..3c3b40137f00e 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -184,6 +184,68 @@ def OptimizeAllocationLivenessPass
let dependentDialects = ["mlir::memref::MemRefDialect"];
}
+def StaticMemoryPlannerAnalysisPass
+ : Pass<"static-memory-planner-analysis", "func::FuncOp"> {
+ let summary = "Identifies same-block alloc/dealloc pairs eligible for "
+ "static memory planning";
+ let description = [{
+ This analysis-only pass identifies `memref.alloc` / `memref.dealloc` pairs
+ in the same basic block that are candidates for static memory planning and
+ buffer reuse.
+
+ For each allocation the pass checks a conservative eligibility envelope:
+ - Static memref shape.
+ - Unique same-block dealloc (found via `MemoryEffects::Free`).
+ - Not nested inside any loop or conditional region (e.g., `scf.for`,
+ `scf.while`, `scf.forall`, `scf.parallel`, `scf.if`).
+ - No cross-block alias or escaping use (alias set resolved via
+ `BufferViewFlowAnalysis`).
+
+ Eligible allocations receive an alias-aware lifetime interval computed
+ using `BufferViewFlowAnalysis::resolve()`. The interval spans from the
+ alloc op's block-local index to the maximum of the dealloc index and the
+ last use index of any derived alias within the block. Allocation size in
+ bytes (for integer and floating-point element types) and alignment are
+ also reported where available.
+
+ Ineligible allocations receive a skip reason. Both eligibility results
+ and skip reasons are reported via op remarks. Aggregate counts are
+ reported as pass statistics.
+
+ No IR modifications are made by this pass.
+
+ This pass is expected to run after the deallocation pipeline
+ (`ownership-based-buffer-deallocation` followed by
+ `bufferization-lower-deallocations`).
+
+ Example:
+ ```mlir
+ func.func @example() {
+ // Remark: static-memory-planner: eligible: size=4096 bytes, interval=[0,2]
+ %0 = memref.alloc() : memref<1024xf32>
+ "some.use"(%0) : (memref<1024xf32>) -> ()
+ memref.dealloc %0 : memref<1024xf32>
+ return
+ }
+ ```
+ }];
+
+ let statistics =
+ [Statistic<"numEligible", "num-eligible",
+ "Number of alloc/dealloc pairs eligible for static planning">,
+ Statistic<"numSkipDynamic", "num-skipped-dynamic",
+ "Number of allocations skipped: dynamic shape">,
+ Statistic<
+ "numSkipNested", "num-skipped-nested",
+ "Number of allocations skipped: nested in loop or conditional">,
+ Statistic<"numSkipNoDealloc", "num-skipped-no-dealloc",
+ "Number of allocations skipped: no unique same-block dealloc">,
+ Statistic<
+ "numSkipEscaping", "num-skipped-escaping",
+ "Number of allocations skipped: cross-block alias or escaping use">,
+ ];
+}
+
def LowerDeallocationsPass : Pass<"bufferization-lower-deallocations"> {
let summary = "Lowers `bufferization.dealloc` operations to `memref.dealloc`"
"operations";
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
index 7c38621be1bb5..5df9f19a5e30a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_dialect_library(MLIRBufferizationTransforms
OwnershipBasedBufferDeallocation.cpp
TensorCopyInsertion.cpp
OptimizeAllocationLiveness.cpp
+ StaticMemoryPlannerAnalysis.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
new file mode 100644
index 0000000000000..28038be76cb4d
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -0,0 +1,272 @@
+//===- StaticMemoryPlannerAnalysis.cpp - Analysis for static memory -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements an analysis-only pass that discovers same-block
+// memref.alloc / memref.dealloc pairs eligible for static memory planning.
+//
+// For each eligible allocation the pass:
+// - Computes a conservative alias-aware lifetime interval using
+// BufferViewFlowAnalysis.
+// - Collects metadata (static size in bytes, alignment, memory space).
+// - Emits the results as op remarks on the alloc op.
+//
+// Ineligible allocations also receive a remark describing the skip reason.
+//
+// This pass is the first upstream step for the static memory planner project.
+// It is intentionally analysis-only (no IR mutations) and covers the simplest
+// structured case: same-block, static-shape, non-escaping allocations not
+// nested inside loops or conditionals.
+//
+// Expected pipeline position: after ownership-based-buffer-deallocation and
+// bufferization-lower-deallocations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "static-memory-planner-analysis"
+
+namespace mlir {
+namespace bufferization {
+#define GEN_PASS_DEF_STATICMEMORYPLANNERANALYSISPASS
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
+} // namespace bufferization
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper utilities
+//===----------------------------------------------------------------------===//
+
+/// Returns true if `op` is nested inside any loop or conditional region,
+/// i.e., any ancestor op (up to but not including the nearest
+/// IsIsolatedFromAbove boundary) is a loop or conditional.
+static bool isNestedInLoopOrConditional(Operation *op) {
+ Operation *parent = op->getParentOp();
+ while (parent) {
+ if (isa<scf::ForOp, scf::ForallOp, scf::WhileOp, scf::ParallelOp,
+ scf::IfOp>(parent))
+ return true;
+ // Do not cross function/isolated-region boundaries.
+ if (parent->hasTrait<OpTrait::IsIsolatedFromAbove>())
+ break;
+ parent = parent->getParentOp();
+ }
+ return false;
+}
+
+/// Returns the unique user of `value` that carries a MemoryEffects::Free
+/// effect, or nullptr when there are zero or multiple such users.
+static Operation *findUniqueFreeSideEffectUser(Value value) {
+ Operation *freeUser = nullptr;
+ for (Operation *user : value.getUsers()) {
+ auto memEffectOp = dyn_cast<MemoryEffectOpInterface>(user);
+ if (!memEffectOp)
+ continue;
+ SmallVector<MemoryEffects::EffectInstance, 2> effects;
+ memEffectOp.getEffects(effects);
+ for (const auto &effect : effects) {
+ if (isa<MemoryEffects::Free>(effect.getEffect())) {
+ if (freeUser)
+ return nullptr; // Multiple free users — not uniquely deallocated.
+ freeUser = user;
+ }
+ }
+ }
+ return freeUser;
+}
+
+/// Builds a block-local operation index map: op → position in block order.
+static DenseMap<Operation *, unsigned> buildOpIndexMap(Block *block) {
+ DenseMap<Operation *, unsigned> indexMap;
+ unsigned idx = 0;
+ for (Operation &op : *block)
+ indexMap[&op] = idx++;
+ return indexMap;
+}
+
+/// Returns the static allocation size in bytes for `type`, or -1 if it cannot
+/// be determined (e.g., non-integer/float element types such as index).
+static int64_t getStaticSizeBytes(MemRefType type) {
+ Type elemType = type.getElementType();
+ if (!elemType.isIntOrFloat())
+ return -1;
+ int64_t numElems = type.getNumElements();
+ unsigned elemBits = type.getElementTypeBitWidth();
+ return (numElems * static_cast<int64_t>(elemBits) + 7) / 8;
+}
+
+//===----------------------------------------------------------------------===//
+// StaticMemoryPlannerAnalysisPass
+//===----------------------------------------------------------------------===//
+
+struct StaticMemoryPlannerAnalysisPass
+ : public bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
+ StaticMemoryPlannerAnalysisPass> {
+public:
+ StaticMemoryPlannerAnalysisPass() = default;
+
+ void runOnOperation() override {
+ func::FuncOp func = getOperation();
+
+ if (func.isExternal())
+ return;
+
+ // Build alias/view-flow analysis once for the entire function.
+ // BufferViewFlowAnalysis::resolve(v) gives the transitive closure of all
+ // values derived from v (subviews, expands, casts, etc.).
+ BufferViewFlowAnalysis aliasAnalysis(func);
+
+ // Lazily-populated per-block operation index maps.
+ // Keyed by Block*; maps each Op* to its zero-based position in the block.
+ DenseMap<Block *, DenseMap<Operation *, unsigned>> blockIndexMaps;
+
+ // Walk every memref.alloc in the function and classify it.
+ func.walk([&](memref::AllocOp allocOp) {
+ auto memrefType = allocOp.getType();
+
+ //----------------------------------------------------------------
+ // Eligibility check 1: static shape.
+ //----------------------------------------------------------------
+ if (!memrefType.hasStaticShape()) {
+ ++numSkipDynamic;
+ (void)allocOp.emitRemark("static-memory-planner: skip: dynamic shape");
+ return;
+ }
+
+ //----------------------------------------------------------------
+ // Eligibility check 2: not nested inside a loop or conditional.
+ //----------------------------------------------------------------
+ if (isNestedInLoopOrConditional(allocOp)) {
+ ++numSkipNested;
+ (void)allocOp.emitRemark(
+ "static-memory-planner: skip: nested in loop or conditional");
+ return;
+ }
+
+ //----------------------------------------------------------------
+ // Eligibility check 3: unique same-block dealloc.
+ //----------------------------------------------------------------
+ Value allocResult = allocOp.getResult();
+ Operation *deallocOp = findUniqueFreeSideEffectUser(allocResult);
+ if (!deallocOp || deallocOp->getBlock() != allocOp->getBlock()) {
+ ++numSkipNoDealloc;
+ (void)allocOp.emitRemark(
+ "static-memory-planner: skip: no unique same-block dealloc");
+ return;
+ }
+
+ Block *block = allocOp->getBlock();
+
+ //----------------------------------------------------------------
+ // Eligibility check 4: no cross-block alias or escaping use.
+ // Resolve the full alias set and verify every user is in the same
+ // block (or is the dealloc itself). A use in func.return also
+ // counts as escaping.
+ //----------------------------------------------------------------
+ const BufferViewFlowAnalysis::ValueSetT &aliases =
+ aliasAnalysis.resolve(allocResult);
+ bool escapes = false;
+ for (Value alias : aliases) {
+ for (Operation *user : alias.getUsers()) {
+ if (user == deallocOp)
+ continue;
+ if (user->getBlock() != block || isa<func::ReturnOp>(user)) {
+ escapes = true;
+ break;
+ }
+ }
+ if (escapes)
+ break;
+ }
+ if (escapes) {
+ ++numSkipEscaping;
+ (void)allocOp.emitRemark(
+ "static-memory-planner: skip: escaping or cross-block alias");
+ return;
+ }
+
+ //----------------------------------------------------------------
+ // Compute alias-aware lifetime interval.
+ // Start = block-local index of the alloc op.
+ // End = max(dealloc index, last use index of any alias in block).
+ //----------------------------------------------------------------
+ auto &indexMap = blockIndexMaps.try_emplace(block).first->second;
+ if (indexMap.empty())
+ indexMap = buildOpIndexMap(block);
+
+ unsigned allocIdx = indexMap.lookup(allocOp);
+ unsigned deallocIdx = indexMap.lookup(deallocOp);
+ unsigned endIdx = deallocIdx;
+
+ for (Value alias : aliases) {
+ for (Operation *user : alias.getUsers()) {
+ if (user == deallocOp)
+ continue;
+ // Lift the user to the ancestor op that lives directly in `block`
+ // (handles users inside nested regions attached to block-level ops).
+ if (Operation *ancestor = block->findAncestorOpInBlock(*user)) {
+ auto it = indexMap.find(ancestor);
+ if (it != indexMap.end() && it->second > endIdx)
+ endIdx = it->second;
+ }
+ }
+ }
+
+ //----------------------------------------------------------------
+ // Collect metadata.
+ //----------------------------------------------------------------
+ int64_t sizeBytes = getStaticSizeBytes(memrefType);
+ std::optional<uint64_t> alignment = allocOp.getAlignment();
+
+ //----------------------------------------------------------------
+ // Emit eligibility remark and update statistics.
+ //----------------------------------------------------------------
+ ++numEligible;
+
+ LDBG() << "eligible: " << allocOp << " size=" << sizeBytes
+ << " interval=[" << allocIdx << "," << endIdx << "]";
+
+ std::string msg = "static-memory-planner: eligible";
+ if (sizeBytes >= 0)
+ msg += ": size=" + std::to_string(sizeBytes) + " bytes";
+ else
+ msg += ": size=unknown";
+
+ if (alignment)
+ msg += ", align=" + std::to_string(*alignment);
+
+ msg += ", interval=[" + std::to_string(allocIdx) + "," +
+ std::to_string(endIdx) + "]";
+
+ (void)allocOp.emitRemark(msg);
+ });
+
+ LDBG() << "[" << func.getName()
+ << "] summary: eligible=" << (unsigned)numEligible
+ << " skip-dynamic=" << (unsigned)numSkipDynamic
+ << " skip-nested=" << (unsigned)numSkipNested
+ << " skip-no-dealloc=" << (unsigned)numSkipNoDealloc
+ << " skip-escaping=" << (unsigned)numSkipEscaping;
+ }
+};
+
+} // end anonymous namespace
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
new file mode 100644
index 0000000000000..230bc854ee66f
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -0,0 +1,107 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics \
+// RUN: --static-memory-planner-analysis
+
+
+// -----
+
+// Test 1: Two sequential non-overlapping alloc/dealloc pairs — both eligible.
+//
+// Block op indices (zero-based):
+// 0: %alloc0 = memref.alloc
+// 1: memref.dealloc %alloc0
+// 2: %alloc1 = memref.alloc
+// 3: memref.dealloc %alloc1
+// 4: return
+//
+// alloc0 interval = [0, 1] — dealloc is lastUse, no other users.
+// alloc1 interval = [2, 3] — dealloc is lastUse, no other users.
+// The two intervals are non-overlapping: a future planner can reuse the same
+// static region for both buffers.
+
+func.func @simple_sequential() {
+ // expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,1]}}
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ // expected-remark @below {{static-memory-planner: eligible: size=2048 bytes, interval=[2,3]}}
+ %alloc1 = memref.alloc() : memref<512xf32>
+ memref.dealloc %alloc1 : memref<512xf32>
+ return
+}
+
+// -----
+
+// Test 2: Dynamic-shape allocation — skipped.
+// The alloc has a runtime dimension (%n), so the pass cannot compute a static
+// size or reason about it conservatively for static planning.
+
+func.func @dynamic_shape_skipped(%n: index) {
+ // expected-remark @below {{static-memory-planner: skip: dynamic shape}}
+ %alloc = memref.alloc(%n) : memref<?xf32>
+ return
+}
+
+// -----
+
+// Test 3: No same-block dealloc — skipped.
+// The alloc is unmatched (no memref.dealloc in any block), so the pass
+// cannot establish a lifetime interval.
+
+func.func @no_dealloc_skipped() {
+ // expected-remark @below {{static-memory-planner: skip: no unique same-block dealloc}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 4: Alloc inside scf.if (conditional) — skipped.
+// Allocations nested inside conditionals are excluded from v1 scope because
+// their liveness depends on runtime predicate evaluation.
+
+func.func @conditional_alloc_skipped(%cond: i1) {
+ scf.if %cond {
+ // expected-remark @below {{static-memory-planner: skip: nested in loop or conditional}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc : memref<1024xf32>
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// Test 5: Alloc inside scf.for (loop) — skipped.
+// Allocations inside loop bodies may execute a dynamic number of times;
+// static planning requires reasoning outside loop nests.
+
+func.func @loop_alloc_skipped(%lb: index, %ub: index, %step: index) {
+ scf.for %i = %lb to %ub step %step {
+ // expected-remark @below {{static-memory-planner: skip: nested in loop or conditional}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc : memref<1024xf32>
+ scf.yield
+ }
+ return
+}
+
+// -----
+
+// Test 6: expand_shape alias — alloc remains eligible and the alias is tracked.
+//
+// BufferViewFlowAnalysis::resolve(%alloc) finds both %alloc and %view.
+// Users of %alloc (non-dealloc): expand_shape at idx 1.
+// Users of %view: none.
+// endIdx = max(deallocIdx=2, lastAliasUse=1) = 2.
+// Interval = [0, 2].
+//
+// This test verifies that the presence of a derived alias does not cause
+// misclassification or a crash.
+
+func.func @view_alias_tracked() {
+ // expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,2]}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ %view = memref.expand_shape %alloc [[0, 1]] output_shape [2, 512]
+ : memref<1024xf32> into memref<2x512xf32>
+ memref.dealloc %alloc : memref<1024xf32>
+ return
+}
>From 9e6c33fe10614fae5c765b2cd6f885e035710731 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Sat, 28 Mar 2026 11:07:20 +0530
Subject: [PATCH 02/13] clean up comments and add cross-block escape test
Rewrote function-level and inline comments to be more concise.
Added test 7 covering the cross-block alias / escaping skip path,
which was previously exercised by the implementation but not tested.
Signed-off-by: KrxGu <krishom70 at gmail.com>
---
.../StaticMemoryPlannerAnalysis.cpp | 81 +++++--------------
.../static-memory-planner-analysis.mlir | 41 ++++++----
2 files changed, 47 insertions(+), 75 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 28038be76cb4d..3509decfdd492 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -6,24 +6,13 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements an analysis-only pass that discovers same-block
-// memref.alloc / memref.dealloc pairs eligible for static memory planning.
+// Walk each block in a function and find memref.alloc / memref.dealloc pairs
+// that can be reasoned about statically. For each eligible alloc compute a
+// conservative alias-aware lifetime interval via BufferViewFlowAnalysis and
+// report it as an op remark. Ineligible allocs get a skip reason instead.
//
-// For each eligible allocation the pass:
-// - Computes a conservative alias-aware lifetime interval using
-// BufferViewFlowAnalysis.
-// - Collects metadata (static size in bytes, alignment, memory space).
-// - Emits the results as op remarks on the alloc op.
-//
-// Ineligible allocations also receive a remark describing the skip reason.
-//
-// This pass is the first upstream step for the static memory planner project.
-// It is intentionally analysis-only (no IR mutations) and covers the simplest
-// structured case: same-block, static-shape, non-escaping allocations not
-// nested inside loops or conditionals.
-//
-// Expected pipeline position: after ownership-based-buffer-deallocation and
-// bufferization-lower-deallocations.
+// Meant to run after ownership-based-buffer-deallocation followed by
+// bufferization-lower-deallocations, once all pairs are explicit in the IR.
//
//===----------------------------------------------------------------------===//
@@ -56,9 +45,8 @@ namespace {
// Helper utilities
//===----------------------------------------------------------------------===//
-/// Returns true if `op` is nested inside any loop or conditional region,
-/// i.e., any ancestor op (up to but not including the nearest
-/// IsIsolatedFromAbove boundary) is a loop or conditional.
+/// Returns true if `op` lives inside a loop or conditional body. Stops
+/// walking at function/isolated-region boundaries.
static bool isNestedInLoopOrConditional(Operation *op) {
Operation *parent = op->getParentOp();
while (parent) {
@@ -73,8 +61,8 @@ static bool isNestedInLoopOrConditional(Operation *op) {
return false;
}
-/// Returns the unique user of `value` that carries a MemoryEffects::Free
-/// effect, or nullptr when there are zero or multiple such users.
+/// Returns the single user of `value` with a MemoryEffects::Free effect, or
+/// nullptr if there are zero or more than one.
static Operation *findUniqueFreeSideEffectUser(Value value) {
Operation *freeUser = nullptr;
for (Operation *user : value.getUsers()) {
@@ -86,7 +74,7 @@ static Operation *findUniqueFreeSideEffectUser(Value value) {
for (const auto &effect : effects) {
if (isa<MemoryEffects::Free>(effect.getEffect())) {
if (freeUser)
- return nullptr; // Multiple free users — not uniquely deallocated.
+ return nullptr; // more than one free user, not uniquely deallocated
freeUser = user;
}
}
@@ -94,7 +82,7 @@ static Operation *findUniqueFreeSideEffectUser(Value value) {
return freeUser;
}
-/// Builds a block-local operation index map: op → position in block order.
+/// Numbers each op in `block` by its position (0-based).
static DenseMap<Operation *, unsigned> buildOpIndexMap(Block *block) {
DenseMap<Operation *, unsigned> indexMap;
unsigned idx = 0;
@@ -103,8 +91,7 @@ static DenseMap<Operation *, unsigned> buildOpIndexMap(Block *block) {
return indexMap;
}
-/// Returns the static allocation size in bytes for `type`, or -1 if it cannot
-/// be determined (e.g., non-integer/float element types such as index).
+/// Returns the size of `type` in bytes, or -1 for non int/float element types.
static int64_t getStaticSizeBytes(MemRefType type) {
Type elemType = type.getElementType();
if (!elemType.isIntOrFloat())
@@ -130,31 +117,23 @@ struct StaticMemoryPlannerAnalysisPass
if (func.isExternal())
return;
- // Build alias/view-flow analysis once for the entire function.
- // BufferViewFlowAnalysis::resolve(v) gives the transitive closure of all
- // values derived from v (subviews, expands, casts, etc.).
+ // Build alias analysis once; we call resolve() per alloc below.
BufferViewFlowAnalysis aliasAnalysis(func);
- // Lazily-populated per-block operation index maps.
- // Keyed by Block*; maps each Op* to its zero-based position in the block.
+ // Op index maps, built on demand per block.
DenseMap<Block *, DenseMap<Operation *, unsigned>> blockIndexMaps;
- // Walk every memref.alloc in the function and classify it.
func.walk([&](memref::AllocOp allocOp) {
auto memrefType = allocOp.getType();
- //----------------------------------------------------------------
- // Eligibility check 1: static shape.
- //----------------------------------------------------------------
+ // Skip dynamic shapes; size is not known at compile time.
if (!memrefType.hasStaticShape()) {
++numSkipDynamic;
(void)allocOp.emitRemark("static-memory-planner: skip: dynamic shape");
return;
}
- //----------------------------------------------------------------
- // Eligibility check 2: not nested inside a loop or conditional.
- //----------------------------------------------------------------
+ // Skip allocs inside loops or conditionals.
if (isNestedInLoopOrConditional(allocOp)) {
++numSkipNested;
(void)allocOp.emitRemark(
@@ -162,9 +141,7 @@ struct StaticMemoryPlannerAnalysisPass
return;
}
- //----------------------------------------------------------------
- // Eligibility check 3: unique same-block dealloc.
- //----------------------------------------------------------------
+ // Need exactly one dealloc in the same block to form a pair.
Value allocResult = allocOp.getResult();
Operation *deallocOp = findUniqueFreeSideEffectUser(allocResult);
if (!deallocOp || deallocOp->getBlock() != allocOp->getBlock()) {
@@ -176,12 +153,7 @@ struct StaticMemoryPlannerAnalysisPass
Block *block = allocOp->getBlock();
- //----------------------------------------------------------------
- // Eligibility check 4: no cross-block alias or escaping use.
- // Resolve the full alias set and verify every user is in the same
- // block (or is the dealloc itself). A use in func.return also
- // counts as escaping.
- //----------------------------------------------------------------
+ // Skip if any alias escapes or is used in a different block.
const BufferViewFlowAnalysis::ValueSetT &aliases =
aliasAnalysis.resolve(allocResult);
bool escapes = false;
@@ -204,11 +176,7 @@ struct StaticMemoryPlannerAnalysisPass
return;
}
- //----------------------------------------------------------------
- // Compute alias-aware lifetime interval.
- // Start = block-local index of the alloc op.
- // End = max(dealloc index, last use index of any alias in block).
- //----------------------------------------------------------------
+ // Interval: [allocIdx, max(deallocIdx, last alias use in block)].
auto &indexMap = blockIndexMaps.try_emplace(block).first->second;
if (indexMap.empty())
indexMap = buildOpIndexMap(block);
@@ -221,8 +189,7 @@ struct StaticMemoryPlannerAnalysisPass
for (Operation *user : alias.getUsers()) {
if (user == deallocOp)
continue;
- // Lift the user to the ancestor op that lives directly in `block`
- // (handles users inside nested regions attached to block-level ops).
+ // Users in nested regions: lift to the ancestor in this block.
if (Operation *ancestor = block->findAncestorOpInBlock(*user)) {
auto it = indexMap.find(ancestor);
if (it != indexMap.end() && it->second > endIdx)
@@ -231,15 +198,9 @@ struct StaticMemoryPlannerAnalysisPass
}
}
- //----------------------------------------------------------------
- // Collect metadata.
- //----------------------------------------------------------------
int64_t sizeBytes = getStaticSizeBytes(memrefType);
std::optional<uint64_t> alignment = allocOp.getAlignment();
- //----------------------------------------------------------------
- // Emit eligibility remark and update statistics.
- //----------------------------------------------------------------
++numEligible;
LDBG() << "eligible: " << allocOp << " size=" << sizeBytes
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 230bc854ee66f..932da38747e72 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -13,10 +13,8 @@
// 3: memref.dealloc %alloc1
// 4: return
//
-// alloc0 interval = [0, 1] — dealloc is lastUse, no other users.
-// alloc1 interval = [2, 3] — dealloc is lastUse, no other users.
-// The two intervals are non-overlapping: a future planner can reuse the same
-// static region for both buffers.
+// alloc0 interval = [0, 1] — dealloc is the last use, no other users.
+// alloc1 interval = [2, 3] — same. The two intervals don't overlap.
func.func @simple_sequential() {
// expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,1]}}
@@ -31,8 +29,7 @@ func.func @simple_sequential() {
// -----
// Test 2: Dynamic-shape allocation — skipped.
-// The alloc has a runtime dimension (%n), so the pass cannot compute a static
-// size or reason about it conservatively for static planning.
+// The alloc has a runtime dimension (%n); size is unknown at compile time.
func.func @dynamic_shape_skipped(%n: index) {
// expected-remark @below {{static-memory-planner: skip: dynamic shape}}
@@ -43,8 +40,7 @@ func.func @dynamic_shape_skipped(%n: index) {
// -----
// Test 3: No same-block dealloc — skipped.
-// The alloc is unmatched (no memref.dealloc in any block), so the pass
-// cannot establish a lifetime interval.
+// No memref.dealloc anywhere, so we can't form a pair.
func.func @no_dealloc_skipped() {
// expected-remark @below {{static-memory-planner: skip: no unique same-block dealloc}}
@@ -55,8 +51,7 @@ func.func @no_dealloc_skipped() {
// -----
// Test 4: Alloc inside scf.if (conditional) — skipped.
-// Allocations nested inside conditionals are excluded from v1 scope because
-// their liveness depends on runtime predicate evaluation.
+// Liveness depends on a runtime predicate, so we skip it.
func.func @conditional_alloc_skipped(%cond: i1) {
scf.if %cond {
@@ -71,8 +66,7 @@ func.func @conditional_alloc_skipped(%cond: i1) {
// -----
// Test 5: Alloc inside scf.for (loop) — skipped.
-// Allocations inside loop bodies may execute a dynamic number of times;
-// static planning requires reasoning outside loop nests.
+// Iteration count is not known statically; skip the alloc.
func.func @loop_alloc_skipped(%lb: index, %ub: index, %step: index) {
scf.for %i = %lb to %ub step %step {
@@ -94,9 +88,6 @@ func.func @loop_alloc_skipped(%lb: index, %ub: index, %step: index) {
// endIdx = max(deallocIdx=2, lastAliasUse=1) = 2.
// Interval = [0, 2].
//
-// This test verifies that the presence of a derived alias does not cause
-// misclassification or a crash.
-
func.func @view_alias_tracked() {
// expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,2]}}
%alloc = memref.alloc() : memref<1024xf32>
@@ -105,3 +96,23 @@ func.func @view_alias_tracked() {
memref.dealloc %alloc : memref<1024xf32>
return
}
+
+// -----
+
+// Test 7: Cross-block use — skipped.
+// %alloc has a unique same-block dealloc, so it passes the dealloc check.
+// However it is also used inside an scf.if body (a different block), which
+// the alias/escape analysis detects as a cross-block use.
+
+func.func @cross_block_use_skipped(%cond: i1) {
+ // expected-remark @below {{static-memory-planner: skip: escaping or cross-block alias}}
+ %alloc = memref.alloc() : memref<1024xf32>
+ scf.if %cond {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f32
+ memref.store %cst, %alloc[%c0] : memref<1024xf32>
+ scf.yield
+ }
+ memref.dealloc %alloc : memref<1024xf32>
+ return
+}
>From 647b82db11af3a81630ca48e2af033845527cb6b Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 18 May 2026 13:02:30 +0530
Subject: [PATCH 03/13] [mlir][bufferization] Simplify static memory planner to
basic lifetime analysis
Stripped down the initial implementation to focus on fundamentals.
The pass now just identifies eligible alloc/dealloc pairs in the same
block and checks which ones could potentially reuse memory based on
their order in the IR.
Removed pool grouping, offset assignment, and interval packing logic.
That complexity belongs in a follow-up once the basic approach is solid.
Signed-off-by: KrxGu <krishom70 at gmail.com>
---
.../StaticMemoryPlannerAnalysis.cpp | 235 +++++++-----------
.../static-memory-planner-analysis.mlir | 101 +++-----
2 files changed, 112 insertions(+), 224 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 3509decfdd492..33f689ffc195f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -6,27 +6,21 @@
//
//===----------------------------------------------------------------------===//
//
-// Walk each block in a function and find memref.alloc / memref.dealloc pairs
-// that can be reasoned about statically. For each eligible alloc compute a
-// conservative alias-aware lifetime interval via BufferViewFlowAnalysis and
-// report it as an op remark. Ineligible allocs get a skip reason instead.
+// Discovers same-block memref.alloc/memref.dealloc pairs and analyzes their
+// lifetime relationships to identify opportunities for memory reuse.
//
-// Meant to run after ownership-based-buffer-deallocation followed by
-// bufferization-lower-deallocations, once all pairs are explicit in the IR.
+// This pass performs basic structural analysis without computing actual memory
+// layouts. It reports which allocations are eligible for static planning and
+// whether pairs of allocations have non-overlapping lifetimes (can reuse).
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/Block.h"
#include "mlir/IR/Operation.h"
-#include "mlir/IR/Value.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/Support/DebugLog.h"
+#include "llvm/Support/Debug.h"
#define DEBUG_TYPE "static-memory-planner-analysis"
@@ -42,63 +36,56 @@ using namespace mlir;
namespace {
//===----------------------------------------------------------------------===//
-// Helper utilities
+// Data structures
//===----------------------------------------------------------------------===//
-/// Returns true if `op` lives inside a loop or conditional body. Stops
-/// walking at function/isolated-region boundaries.
-static bool isNestedInLoopOrConditional(Operation *op) {
- Operation *parent = op->getParentOp();
- while (parent) {
- if (isa<scf::ForOp, scf::ForallOp, scf::WhileOp, scf::ParallelOp,
- scf::IfOp>(parent))
- return true;
- // Do not cross function/isolated-region boundaries.
- if (parent->hasTrait<OpTrait::IsIsolatedFromAbove>())
- break;
- parent = parent->getParentOp();
- }
- return false;
-}
+/// A candidate allocation with its matching deallocation.
+struct AllocationCandidate {
+ memref::AllocOp alloc;
+ memref::DeallocOp dealloc;
+};
-/// Returns the single user of `value` with a MemoryEffects::Free effect, or
-/// nullptr if there are zero or more than one.
-static Operation *findUniqueFreeSideEffectUser(Value value) {
- Operation *freeUser = nullptr;
- for (Operation *user : value.getUsers()) {
- auto memEffectOp = dyn_cast<MemoryEffectOpInterface>(user);
- if (!memEffectOp)
- continue;
- SmallVector<MemoryEffects::EffectInstance, 2> effects;
- memEffectOp.getEffects(effects);
- for (const auto &effect : effects) {
- if (isa<MemoryEffects::Free>(effect.getEffect())) {
- if (freeUser)
- return nullptr; // more than one free user, not uniquely deallocated
- freeUser = user;
- }
+//===----------------------------------------------------------------------===//
+// 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()) {
+ if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
+ if (deallocOp)
+ return nullptr; // Multiple deallocs found
+ deallocOp = dealloc;
}
}
- return freeUser;
+ return deallocOp;
}
-/// Numbers each op in `block` by its position (0-based).
-static DenseMap<Operation *, unsigned> buildOpIndexMap(Block *block) {
- DenseMap<Operation *, unsigned> indexMap;
- unsigned idx = 0;
- for (Operation &op : *block)
- indexMap[&op] = idx++;
- return indexMap;
-}
+/// Checks if two allocation candidates have non-overlapping lifetimes.
+/// Returns true if the first's dealloc is strictly before the second's alloc,
+/// or vice versa.
+static bool canReuseMemory(const AllocationCandidate &first,
+ const AllocationCandidate &second) {
+ Operation *firstDealloc = first.dealloc;
+ Operation *firstAlloc = first.alloc;
+ Operation *secondDealloc = second.dealloc;
+ Operation *secondAlloc = second.alloc;
+
+ // Check if both are in the same block
+ if (firstAlloc->getBlock() != secondAlloc->getBlock())
+ return false;
+
+ // Check if first ends before second starts
+ if (firstDealloc->isBeforeInBlock(secondAlloc))
+ return true;
-/// Returns the size of `type` in bytes, or -1 for non int/float element types.
-static int64_t getStaticSizeBytes(MemRefType type) {
- Type elemType = type.getElementType();
- if (!elemType.isIntOrFloat())
- return -1;
- int64_t numElems = type.getNumElements();
- unsigned elemBits = type.getElementTypeBitWidth();
- return (numElems * static_cast<int64_t>(elemBits) + 7) / 8;
+ // Check if second ends before first starts
+ if (secondDealloc->isBeforeInBlock(firstAlloc))
+ return true;
+
+ return false;
}
//===----------------------------------------------------------------------===//
@@ -117,116 +104,60 @@ struct StaticMemoryPlannerAnalysisPass
if (func.isExternal())
return;
- // Build alias analysis once; we call resolve() per alloc below.
- BufferViewFlowAnalysis aliasAnalysis(func);
-
- // Op index maps, built on demand per block.
- DenseMap<Block *, DenseMap<Operation *, unsigned>> blockIndexMaps;
+ // Collect eligible allocation candidates
+ SmallVector<AllocationCandidate> candidates;
func.walk([&](memref::AllocOp allocOp) {
- auto memrefType = allocOp.getType();
+ MemRefType memrefType = allocOp.getType();
- // Skip dynamic shapes; size is not known at compile time.
+ // Skip dynamic shapes
if (!memrefType.hasStaticShape()) {
++numSkipDynamic;
- (void)allocOp.emitRemark("static-memory-planner: skip: dynamic shape");
+ allocOp.emitRemark("static-memory-planner: skip: dynamic shape");
return;
}
- // Skip allocs inside loops or conditionals.
- if (isNestedInLoopOrConditional(allocOp)) {
- ++numSkipNested;
- (void)allocOp.emitRemark(
- "static-memory-planner: skip: nested in loop or conditional");
+ // Find unique dealloc in the same block
+ memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
+ if (!deallocOp) {
+ ++numSkipNoDealloc;
+ allocOp.emitRemark(
+ "static-memory-planner: skip: no unique dealloc");
return;
}
- // Need exactly one dealloc in the same block to form a pair.
- Value allocResult = allocOp.getResult();
- Operation *deallocOp = findUniqueFreeSideEffectUser(allocResult);
- if (!deallocOp || deallocOp->getBlock() != allocOp->getBlock()) {
+ if (deallocOp->getBlock() != allocOp->getBlock()) {
++numSkipNoDealloc;
- (void)allocOp.emitRemark(
- "static-memory-planner: skip: no unique same-block dealloc");
+ allocOp.emitRemark(
+ "static-memory-planner: skip: dealloc in different block");
return;
}
- Block *block = allocOp->getBlock();
-
- // Skip if any alias escapes or is used in a different block.
- const BufferViewFlowAnalysis::ValueSetT &aliases =
- aliasAnalysis.resolve(allocResult);
- bool escapes = false;
- for (Value alias : aliases) {
- for (Operation *user : alias.getUsers()) {
- if (user == deallocOp)
- continue;
- if (user->getBlock() != block || isa<func::ReturnOp>(user)) {
- escapes = true;
- break;
- }
- }
- if (escapes)
- break;
- }
- if (escapes) {
- ++numSkipEscaping;
- (void)allocOp.emitRemark(
- "static-memory-planner: skip: escaping or cross-block alias");
- return;
- }
+ // This allocation is eligible
+ ++numEligible;
+ allocOp.emitRemark("static-memory-planner: eligible");
+ candidates.push_back({allocOp, deallocOp});
+ });
- // Interval: [allocIdx, max(deallocIdx, last alias use in block)].
- auto &indexMap = blockIndexMaps.try_emplace(block).first->second;
- if (indexMap.empty())
- indexMap = buildOpIndexMap(block);
-
- unsigned allocIdx = indexMap.lookup(allocOp);
- unsigned deallocIdx = indexMap.lookup(deallocOp);
- unsigned endIdx = deallocIdx;
-
- for (Value alias : aliases) {
- for (Operation *user : alias.getUsers()) {
- if (user == deallocOp)
- continue;
- // Users in nested regions: lift to the ancestor in this block.
- if (Operation *ancestor = block->findAncestorOpInBlock(*user)) {
- auto it = indexMap.find(ancestor);
- if (it != indexMap.end() && it->second > endIdx)
- endIdx = it->second;
- }
+ // Analyze reuse opportunities between pairs of candidates
+ unsigned numReusable = 0;
+ for (size_t i = 0; i < candidates.size(); ++i) {
+ for (size_t j = i + 1; j < candidates.size(); ++j) {
+ if (canReuseMemory(candidates[i], candidates[j])) {
+ ++numReusable;
+ LLVM_DEBUG(llvm::dbgs()
+ << "[static-memory-planner] reuse opportunity: alloc "
+ << i << " and alloc " << j << "\n");
}
}
+ }
- int64_t sizeBytes = getStaticSizeBytes(memrefType);
- std::optional<uint64_t> alignment = allocOp.getAlignment();
-
- ++numEligible;
-
- LDBG() << "eligible: " << allocOp << " size=" << sizeBytes
- << " interval=[" << allocIdx << "," << endIdx << "]";
-
- std::string msg = "static-memory-planner: eligible";
- if (sizeBytes >= 0)
- msg += ": size=" + std::to_string(sizeBytes) + " bytes";
- else
- msg += ": size=unknown";
-
- if (alignment)
- msg += ", align=" + std::to_string(*alignment);
-
- msg += ", interval=[" + std::to_string(allocIdx) + "," +
- std::to_string(endIdx) + "]";
-
- (void)allocOp.emitRemark(msg);
- });
-
- LDBG() << "[" << func.getName()
- << "] summary: eligible=" << (unsigned)numEligible
- << " skip-dynamic=" << (unsigned)numSkipDynamic
- << " skip-nested=" << (unsigned)numSkipNested
- << " skip-no-dealloc=" << (unsigned)numSkipNoDealloc
- << " skip-escaping=" << (unsigned)numSkipEscaping;
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] summary for "
+ << func.getName() << ": eligible="
+ << (unsigned)numEligible << " skip-dynamic="
+ << (unsigned)numSkipDynamic << " skip-no-dealloc="
+ << (unsigned)numSkipNoDealloc << " reusable-pairs="
+ << numReusable << "\n");
}
};
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 932da38747e72..a60098ab84f35 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -1,26 +1,14 @@
// RUN: mlir-opt %s -split-input-file -verify-diagnostics \
// RUN: --static-memory-planner-analysis
-
// -----
-// Test 1: Two sequential non-overlapping alloc/dealloc pairs — both eligible.
-//
-// Block op indices (zero-based):
-// 0: %alloc0 = memref.alloc
-// 1: memref.dealloc %alloc0
-// 2: %alloc1 = memref.alloc
-// 3: memref.dealloc %alloc1
-// 4: return
-//
-// alloc0 interval = [0, 1] — dealloc is the last use, no other users.
-// alloc1 interval = [2, 3] — same. The two intervals don't overlap.
-
+// Test 1: Simple sequential alloc/dealloc pairs
func.func @simple_sequential() {
- // expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,1]}}
+ // expected-remark @below {{static-memory-planner: eligible}}
%alloc0 = memref.alloc() : memref<1024xf32>
memref.dealloc %alloc0 : memref<1024xf32>
- // expected-remark @below {{static-memory-planner: eligible: size=2048 bytes, interval=[2,3]}}
+ // expected-remark @below {{static-memory-planner: eligible}}
%alloc1 = memref.alloc() : memref<512xf32>
memref.dealloc %alloc1 : memref<512xf32>
return
@@ -28,9 +16,7 @@ func.func @simple_sequential() {
// -----
-// Test 2: Dynamic-shape allocation — skipped.
-// The alloc has a runtime dimension (%n); size is unknown at compile time.
-
+// Test 2: Dynamic shape - should be skipped
func.func @dynamic_shape_skipped(%n: index) {
// expected-remark @below {{static-memory-planner: skip: dynamic shape}}
%alloc = memref.alloc(%n) : memref<?xf32>
@@ -39,24 +25,20 @@ func.func @dynamic_shape_skipped(%n: index) {
// -----
-// Test 3: No same-block dealloc — skipped.
-// No memref.dealloc anywhere, so we can't form a pair.
-
+// Test 3: No dealloc - should be skipped
func.func @no_dealloc_skipped() {
- // expected-remark @below {{static-memory-planner: skip: no unique same-block dealloc}}
+ // expected-remark @below {{static-memory-planner: skip: no unique dealloc}}
%alloc = memref.alloc() : memref<1024xf32>
return
}
// -----
-// Test 4: Alloc inside scf.if (conditional) — skipped.
-// Liveness depends on a runtime predicate, so we skip it.
-
-func.func @conditional_alloc_skipped(%cond: i1) {
+// Test 4: Dealloc in different block - should be skipped
+func.func @different_block_skipped(%cond: i1) {
+ // expected-remark @below {{static-memory-planner: skip: dealloc in different block}}
+ %alloc = memref.alloc() : memref<1024xf32>
scf.if %cond {
- // expected-remark @below {{static-memory-planner: skip: nested in loop or conditional}}
- %alloc = memref.alloc() : memref<1024xf32>
memref.dealloc %alloc : memref<1024xf32>
scf.yield
}
@@ -65,54 +47,29 @@ func.func @conditional_alloc_skipped(%cond: i1) {
// -----
-// Test 5: Alloc inside scf.for (loop) — skipped.
-// Iteration count is not known statically; skip the alloc.
-
-func.func @loop_alloc_skipped(%lb: index, %ub: index, %step: index) {
- scf.for %i = %lb to %ub step %step {
- // expected-remark @below {{static-memory-planner: skip: nested in loop or conditional}}
- %alloc = memref.alloc() : memref<1024xf32>
- memref.dealloc %alloc : memref<1024xf32>
- scf.yield
- }
+// Test 5: Overlapping lifetimes
+func.func @overlapping_lifetimes() {
+ // expected-remark @below {{static-memory-planner: eligible}}
+ %alloc0 = memref.alloc() : memref<512xf32>
+ // expected-remark @below {{static-memory-planner: eligible}}
+ %alloc1 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc1 : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<512xf32>
return
}
// -----
-// Test 6: expand_shape alias — alloc remains eligible and the alias is tracked.
-//
-// BufferViewFlowAnalysis::resolve(%alloc) finds both %alloc and %view.
-// Users of %alloc (non-dealloc): expand_shape at idx 1.
-// Users of %view: none.
-// endIdx = max(deallocIdx=2, lastAliasUse=1) = 2.
-// Interval = [0, 2].
-//
-func.func @view_alias_tracked() {
- // expected-remark @below {{static-memory-planner: eligible: size=4096 bytes, interval=[0,2]}}
- %alloc = memref.alloc() : memref<1024xf32>
- %view = memref.expand_shape %alloc [[0, 1]] output_shape [2, 512]
- : memref<1024xf32> into memref<2x512xf32>
- memref.dealloc %alloc : memref<1024xf32>
- return
-}
-
-// -----
-
-// Test 7: Cross-block use — skipped.
-// %alloc has a unique same-block dealloc, so it passes the dealloc check.
-// However it is also used inside an scf.if body (a different block), which
-// the alias/escape analysis detects as a cross-block use.
-
-func.func @cross_block_use_skipped(%cond: i1) {
- // expected-remark @below {{static-memory-planner: skip: escaping or cross-block alias}}
- %alloc = memref.alloc() : memref<1024xf32>
- scf.if %cond {
- %c0 = arith.constant 0 : index
- %cst = arith.constant 0.0 : f32
- memref.store %cst, %alloc[%c0] : memref<1024xf32>
- scf.yield
- }
- memref.dealloc %alloc : memref<1024xf32>
+// Test 6: Multiple allocations with non-overlapping lifetimes
+func.func @multiple_reusable() {
+ // expected-remark @below {{static-memory-planner: eligible}}
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ // expected-remark @below {{static-memory-planner: eligible}}
+ %alloc1 = memref.alloc() : memref<512xf32>
+ memref.dealloc %alloc1 : memref<512xf32>
+ // expected-remark @below {{static-memory-planner: eligible}}
+ %alloc2 = memref.alloc() : memref<2048xf32>
+ memref.dealloc %alloc2 : memref<2048xf32>
return
}
>From e92aeb20f7293708f7db178c5e0328ec0b503171 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Sun, 24 May 2026 12:47:13 +0530
Subject: [PATCH 04/13] [mlir][bufferization] Make static memory planner
generic to work on any operation
Address review feedback: removed func::FuncOp restriction so the pass
can run on any operation (e.g., gpu.launch, async.execute). Updated
test to use pass-pipeline notation to explicitly schedule on func.func.
This makes the pass more flexible for users who want to apply static
memory planning to other operation types.
Signed-off-by: KrxGu <krishom70 at gmail.com>
---
.../mlir/Dialect/Bufferization/Transforms/Passes.td | 2 +-
.../Transforms/StaticMemoryPlannerAnalysis.cpp | 10 +++-------
.../Transforms/static-memory-planner-analysis.mlir | 4 ++--
3 files changed, 6 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 3c3b40137f00e..54c9d6ad14442 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -185,7 +185,7 @@ def OptimizeAllocationLivenessPass
}
def StaticMemoryPlannerAnalysisPass
- : Pass<"static-memory-planner-analysis", "func::FuncOp"> {
+ : Pass<"static-memory-planner-analysis"> {
let summary = "Identifies same-block alloc/dealloc pairs eligible for "
"static memory planning";
let description = [{
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 33f689ffc195f..98803bcd12992 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -99,15 +99,12 @@ struct StaticMemoryPlannerAnalysisPass
StaticMemoryPlannerAnalysisPass() = default;
void runOnOperation() override {
- func::FuncOp func = getOperation();
-
- if (func.isExternal())
- return;
+ Operation *op = getOperation();
// Collect eligible allocation candidates
SmallVector<AllocationCandidate> candidates;
- func.walk([&](memref::AllocOp allocOp) {
+ op->walk([&](memref::AllocOp allocOp) {
MemRefType memrefType = allocOp.getType();
// Skip dynamic shapes
@@ -152,8 +149,7 @@ struct StaticMemoryPlannerAnalysisPass
}
}
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] summary for "
- << func.getName() << ": eligible="
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] summary: eligible="
<< (unsigned)numEligible << " skip-dynamic="
<< (unsigned)numSkipDynamic << " skip-no-dealloc="
<< (unsigned)numSkipNoDealloc << " reusable-pairs="
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 a60098ab84f35..a8c7c728b8423 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -split-input-file -verify-diagnostics \
-// RUN: --static-memory-planner-analysis
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(static-memory-planner-analysis))" \
+// RUN: -split-input-file -verify-diagnostics
// -----
>From 06a8d9b138ca05d7158f045b1dc79bb321241428 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 1 Jun 2026 16:56:14 +0530
Subject: [PATCH 05/13] [mlir][bufferization] Implement e2e IR transformation
for static memory planner
This adds the complete transformation pass that converts multiple
memref.alloc/dealloc pairs into a single arena with subviews.
The offset assignment is intentionally simple (just sequential) - this
establishes the e2e pipeline so we can add smarter bin-packing later.
Tests verify arena sizing, sequential offsets, and that dynamic shapes
or missing deallocations are correctly skipped.
---
.../StaticMemoryPlannerAnalysis.cpp | 151 ++++++++++--------
.../static-memory-planner-analysis.mlir | 50 ++++--
2 files changed, 116 insertions(+), 85 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 98803bcd12992..dbf2661fb8b50 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -1,4 +1,4 @@
-//===- StaticMemoryPlannerAnalysis.cpp - Analysis for static memory -------===//
+//===- StaticMemoryPlannerAnalysis.cpp - Static memory planning -----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,23 +6,21 @@
//
//===----------------------------------------------------------------------===//
//
-// Discovers same-block memref.alloc/memref.dealloc pairs and analyzes their
-// lifetime relationships to identify opportunities for memory reuse.
-//
-// This pass performs basic structural analysis without computing actual memory
-// layouts. It reports which allocations are eligible for static planning and
-// whether pairs of allocations have non-overlapping lifetimes (can reuse).
+// Transforms memref.alloc/memref.dealloc pairs into a single arena allocation
+// with subviews. Uses simple sequential offset assignment where each allocation
+// gets its own space without overlap (baseline algorithm for e2e testing).
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/Support/Debug.h"
-#define DEBUG_TYPE "static-memory-planner-analysis"
+#define DEBUG_TYPE "static-memory-planner"
namespace mlir {
namespace bufferization {
@@ -31,18 +29,18 @@ namespace bufferization {
} // namespace bufferization
} // namespace mlir
-using namespace mlir;
-
namespace {
//===----------------------------------------------------------------------===//
// Data structures
//===----------------------------------------------------------------------===//
-/// A candidate allocation with its matching deallocation.
+/// A candidate allocation with its matching deallocation and assigned offset.
struct AllocationCandidate {
- memref::AllocOp alloc;
- memref::DeallocOp dealloc;
+ mlir::memref::AllocOp alloc;
+ mlir::memref::DeallocOp dealloc;
+ int64_t offset = 0; // Offset in elements from arena start
+ int64_t sizeInElements = 0; // Size in elements
};
//===----------------------------------------------------------------------===//
@@ -51,10 +49,10 @@ struct AllocationCandidate {
/// 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()) {
- if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
+static mlir::memref::DeallocOp findUniqueDealloc(mlir::Value allocValue) {
+ mlir::memref::DeallocOp deallocOp = nullptr;
+ for (mlir::Operation *user : allocValue.getUsers()) {
+ if (auto dealloc = mlir::dyn_cast<mlir::memref::DeallocOp>(user)) {
if (deallocOp)
return nullptr; // Multiple deallocs found
deallocOp = dealloc;
@@ -63,29 +61,12 @@ static memref::DeallocOp findUniqueDealloc(Value allocValue) {
return deallocOp;
}
-/// Checks if two allocation candidates have non-overlapping lifetimes.
-/// Returns true if the first's dealloc is strictly before the second's alloc,
-/// or vice versa.
-static bool canReuseMemory(const AllocationCandidate &first,
- const AllocationCandidate &second) {
- Operation *firstDealloc = first.dealloc;
- Operation *firstAlloc = first.alloc;
- Operation *secondDealloc = second.dealloc;
- Operation *secondAlloc = second.alloc;
-
- // Check if both are in the same block
- if (firstAlloc->getBlock() != secondAlloc->getBlock())
- return false;
-
- // Check if first ends before second starts
- if (firstDealloc->isBeforeInBlock(secondAlloc))
- return true;
-
- // Check if second ends before first starts
- if (secondDealloc->isBeforeInBlock(firstAlloc))
- return true;
-
- return false;
+/// Compute the number of elements in a static-shape memref.
+static int64_t computeSizeInElements(mlir::MemRefType memrefType) {
+ int64_t size = 1;
+ for (int64_t dim : memrefType.getShape())
+ size *= dim;
+ return size;
}
//===----------------------------------------------------------------------===//
@@ -93,67 +74,97 @@ static bool canReuseMemory(const AllocationCandidate &first,
//===----------------------------------------------------------------------===//
struct StaticMemoryPlannerAnalysisPass
- : public bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
+ : public mlir::bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
StaticMemoryPlannerAnalysisPass> {
public:
StaticMemoryPlannerAnalysisPass() = default;
void runOnOperation() override {
- Operation *op = getOperation();
+ mlir::Operation *op = getOperation();
- // Collect eligible allocation candidates
- SmallVector<AllocationCandidate> candidates;
+ // Step 1: Collect eligible allocation candidates
+ llvm::SmallVector<AllocationCandidate> candidates;
- op->walk([&](memref::AllocOp allocOp) {
- MemRefType memrefType = allocOp.getType();
+ op->walk([&](mlir::memref::AllocOp allocOp) {
+ mlir::MemRefType memrefType = allocOp.getType();
// Skip dynamic shapes
if (!memrefType.hasStaticShape()) {
++numSkipDynamic;
- allocOp.emitRemark("static-memory-planner: skip: dynamic shape");
return;
}
// Find unique dealloc in the same block
- memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
+ mlir::memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
if (!deallocOp) {
++numSkipNoDealloc;
- allocOp.emitRemark(
- "static-memory-planner: skip: no unique dealloc");
return;
}
if (deallocOp->getBlock() != allocOp->getBlock()) {
++numSkipNoDealloc;
- allocOp.emitRemark(
- "static-memory-planner: skip: dealloc in different block");
return;
}
// This allocation is eligible
++numEligible;
- allocOp.emitRemark("static-memory-planner: eligible");
- candidates.push_back({allocOp, deallocOp});
+ AllocationCandidate candidate;
+ candidate.alloc = allocOp;
+ candidate.dealloc = deallocOp;
+ candidate.sizeInElements = computeSizeInElements(memrefType);
+ candidates.push_back(candidate);
});
- // Analyze reuse opportunities between pairs of candidates
- unsigned numReusable = 0;
- for (size_t i = 0; i < candidates.size(); ++i) {
- for (size_t j = i + 1; j < candidates.size(); ++j) {
- if (canReuseMemory(candidates[i], candidates[j])) {
- ++numReusable;
- LLVM_DEBUG(llvm::dbgs()
- << "[static-memory-planner] reuse opportunity: alloc "
- << i << " and alloc " << j << "\n");
- }
- }
+ if (candidates.empty())
+ return;
+
+ // Step 2: Compute simple sequential offsets (no overlap optimization)
+ int64_t totalSize = 0;
+ for (auto &candidate : candidates) {
+ candidate.offset = totalSize;
+ totalSize += candidate.sizeInElements;
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] offset="
+ << candidate.offset
+ << " size=" << candidate.sizeInElements << "\n");
}
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] summary: eligible="
- << (unsigned)numEligible << " skip-dynamic="
- << (unsigned)numSkipDynamic << " skip-no-dealloc="
- << (unsigned)numSkipNoDealloc << " reusable-pairs="
- << numReusable << "\n");
+ // Step 3: Find the first allocation's location to place arena
+ mlir::Operation *firstAlloc = candidates.front().alloc;
+ mlir::OpBuilder builder(firstAlloc);
+
+ // Get element type from first allocation (assume all same type for now)
+ mlir::Type elementType = candidates.front().alloc.getType().getElementType();
+
+ // Step 4: Create arena allocation
+ auto arenaType = mlir::MemRefType::get({totalSize}, elementType);
+ auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
+
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena: size="
+ << totalSize << " elements\n");
+
+ // Step 5: Replace each alloc with a subview and remove deallocs
+ for (auto &candidate : candidates) {
+ mlir::OpBuilder subviewBuilder(candidate.alloc);
+
+ // Create subview into arena
+ llvm::SmallVector<mlir::OpFoldResult> offsets, sizes, strides;
+
+ // Single offset into flat arena
+ offsets.push_back(subviewBuilder.getIndexAttr(candidate.offset));
+ sizes.push_back(subviewBuilder.getIndexAttr(candidate.sizeInElements));
+ strides.push_back(subviewBuilder.getIndexAttr(1));
+
+ auto subview = mlir::memref::SubViewOp::create(
+ subviewBuilder, candidate.alloc.getLoc(), arenaAlloc.getResult(),
+ offsets, sizes, strides);
+
+ // Replace all uses of the original alloc
+ candidate.alloc.getResult().replaceAllUsesWith(subview.getResult());
+
+ // Remove the original alloc and dealloc
+ candidate.alloc.erase();
+ candidate.dealloc.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 a8c7c728b8423..1cde192103017 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -1,14 +1,18 @@
// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(static-memory-planner-analysis))" \
-// RUN: -split-input-file -verify-diagnostics
+// RUN: -split-input-file | FileCheck %s
// -----
// Test 1: Simple sequential alloc/dealloc pairs
+// CHECK-LABEL: func @simple_sequential
func.func @simple_sequential() {
- // expected-remark @below {{static-memory-planner: eligible}}
+ // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<1536xf32>
+ // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1]>>
+ // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1], offset: 1024>>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<1024xf32>
memref.dealloc %alloc0 : memref<1024xf32>
- // expected-remark @below {{static-memory-planner: eligible}}
%alloc1 = memref.alloc() : memref<512xf32>
memref.dealloc %alloc1 : memref<512xf32>
return
@@ -16,9 +20,11 @@ func.func @simple_sequential() {
// -----
-// Test 2: Dynamic shape - should be skipped
+// Test 2: Dynamic shape - should be skipped (no transformation)
+// CHECK-LABEL: func @dynamic_shape_skipped
func.func @dynamic_shape_skipped(%n: index) {
- // expected-remark @below {{static-memory-planner: skip: dynamic shape}}
+ // CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) : memref<?xf32>
+ // CHECK-NOT: memref.subview
%alloc = memref.alloc(%n) : memref<?xf32>
return
}
@@ -26,8 +32,10 @@ func.func @dynamic_shape_skipped(%n: index) {
// -----
// Test 3: No dealloc - should be skipped
+// CHECK-LABEL: func @no_dealloc_skipped
func.func @no_dealloc_skipped() {
- // expected-remark @below {{static-memory-planner: skip: no unique dealloc}}
+ // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1024xf32>
+ // CHECK-NOT: memref.subview
%alloc = memref.alloc() : memref<1024xf32>
return
}
@@ -35,8 +43,12 @@ func.func @no_dealloc_skipped() {
// -----
// Test 4: Dealloc in different block - should be skipped
+// CHECK-LABEL: func @different_block_skipped
func.func @different_block_skipped(%cond: i1) {
- // expected-remark @below {{static-memory-planner: skip: dealloc in different block}}
+ // 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>
@@ -47,11 +59,15 @@ func.func @different_block_skipped(%cond: i1) {
// -----
-// Test 5: Overlapping lifetimes
+// Test 5: Overlapping lifetimes (both eligible, sequential offsets)
+// CHECK-LABEL: func @overlapping_lifetimes
func.func @overlapping_lifetimes() {
- // expected-remark @below {{static-memory-planner: eligible}}
+ // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<1536xf32>
+ // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1]>>
+ // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][512] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1], offset: 512>>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<512xf32>
- // expected-remark @below {{static-memory-planner: eligible}}
%alloc1 = memref.alloc() : memref<1024xf32>
memref.dealloc %alloc1 : memref<1024xf32>
memref.dealloc %alloc0 : memref<512xf32>
@@ -60,15 +76,19 @@ func.func @overlapping_lifetimes() {
// -----
-// Test 6: Multiple allocations with non-overlapping lifetimes
-func.func @multiple_reusable() {
- // expected-remark @below {{static-memory-planner: eligible}}
+// Test 6: Multiple allocations with sequential offsets
+// CHECK-LABEL: func @multiple_sequential
+func.func @multiple_sequential() {
+ // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<3584xf32>
+ // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<3584xf32> to memref<1024xf32, strided<[1]>>
+ // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<3584xf32> to memref<512xf32, strided<[1], offset: 1024>>
+ // CHECK-NEXT: %[[SUBVIEW2:.*]] = memref.subview %[[ARENA]][1536] [2048] [1] : memref<3584xf32> to memref<2048xf32, strided<[1], offset: 1536>>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<1024xf32>
memref.dealloc %alloc0 : memref<1024xf32>
- // expected-remark @below {{static-memory-planner: eligible}}
%alloc1 = memref.alloc() : memref<512xf32>
memref.dealloc %alloc1 : memref<512xf32>
- // expected-remark @below {{static-memory-planner: eligible}}
%alloc2 = memref.alloc() : memref<2048xf32>
memref.dealloc %alloc2 : memref<2048xf32>
return
>From 8d006922581e7f435815bfbf71839e6059eba2ad Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 15 Jun 2026 14:37:55 +0530
Subject: [PATCH 06/13] [mlir][bufferization] Add alignment support to static
memory planner
Track alignment requirements from memref.alloc operations and ensure
offsets are properly padded to meet alignment constraints. The arena
allocation receives the maximum alignment of all transformed allocations.
Changes:
- Add alignment field to AllocationCandidate structure
- Compute sizes in bytes to handle alignment padding correctly
- Implement alignOffset() helper to pad offsets to alignment boundaries
- Set arena alignment attribute to maximum required alignment
- Add test demonstrating alignment padding with 64 and 128-byte requirements
This ensures correctness for SIMD and other alignment-sensitive operations.
---
.../StaticMemoryPlannerAnalysis.cpp | 56 +++++++++++++++----
.../static-memory-planner-analysis.mlir | 31 +++++++++-
2 files changed, 72 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index dbf2661fb8b50..9c1898aa96cc3 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -39,8 +39,9 @@ namespace {
struct AllocationCandidate {
mlir::memref::AllocOp alloc;
mlir::memref::DeallocOp dealloc;
- int64_t offset = 0; // Offset in elements from arena start
- int64_t sizeInElements = 0; // Size in elements
+ int64_t offset = 0; // Offset in bytes from arena start
+ int64_t sizeInBytes = 0; // Size in bytes
+ int64_t alignment = 1; // Required alignment in bytes
};
//===----------------------------------------------------------------------===//
@@ -69,6 +70,21 @@ static int64_t computeSizeInElements(mlir::MemRefType memrefType) {
return size;
}
+/// Compute the size in bytes for a memref type.
+static int64_t computeSizeInBytes(mlir::MemRefType memrefType) {
+ int64_t numElements = computeSizeInElements(memrefType);
+ unsigned elementSizeInBits = memrefType.getElementTypeBitWidth();
+ return (numElements * elementSizeInBits + 7) / 8; // Round up to bytes
+}
+
+/// Align an offset to the specified alignment.
+/// Returns the smallest value >= offset that is a multiple of alignment.
+static int64_t alignOffset(int64_t offset, int64_t alignment) {
+ if (alignment <= 1)
+ return offset;
+ return (offset + alignment - 1) / alignment * alignment;
+}
+
//===----------------------------------------------------------------------===//
// StaticMemoryPlannerAnalysisPass
//===----------------------------------------------------------------------===//
@@ -111,21 +127,28 @@ struct StaticMemoryPlannerAnalysisPass
AllocationCandidate candidate;
candidate.alloc = allocOp;
candidate.dealloc = deallocOp;
- candidate.sizeInElements = computeSizeInElements(memrefType);
+ candidate.sizeInBytes = computeSizeInBytes(memrefType);
+ // Extract alignment requirement (default to 1 if not specified)
+ candidate.alignment = allocOp.getAlignment().value_or(1);
candidates.push_back(candidate);
});
if (candidates.empty())
return;
- // Step 2: Compute simple sequential offsets (no overlap optimization)
+ // Step 2: Compute simple sequential offsets with alignment padding
int64_t totalSize = 0;
+ int64_t maxAlignment = 1;
for (auto &candidate : candidates) {
+ // Align current offset to this allocation's requirement
+ totalSize = alignOffset(totalSize, candidate.alignment);
candidate.offset = totalSize;
- totalSize += candidate.sizeInElements;
+ totalSize += candidate.sizeInBytes;
+ maxAlignment = std::max(maxAlignment, candidate.alignment);
LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] offset="
<< candidate.offset
- << " size=" << candidate.sizeInElements << "\n");
+ << " size=" << candidate.sizeInBytes
+ << " alignment=" << candidate.alignment << "\n");
}
// Step 3: Find the first allocation's location to place arena
@@ -134,13 +157,19 @@ struct StaticMemoryPlannerAnalysisPass
// Get element type from first allocation (assume all same type for now)
mlir::Type elementType = candidates.front().alloc.getType().getElementType();
+ unsigned elementSizeInBits = elementType.getIntOrFloatBitWidth();
+ unsigned elementSizeInBytes = (elementSizeInBits + 7) / 8;
- // Step 4: Create arena allocation
- auto arenaType = mlir::MemRefType::get({totalSize}, elementType);
+ // Step 4: Create arena allocation with maximum alignment
+ // Convert total size from bytes to elements
+ int64_t arenaSizeInElements = (totalSize + elementSizeInBytes - 1) / elementSizeInBytes;
+ auto arenaType = mlir::MemRefType::get({arenaSizeInElements}, elementType);
auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
+ arenaAlloc.setAlignmentAttr(builder.getI64IntegerAttr(maxAlignment));
LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena: size="
- << totalSize << " elements\n");
+ << arenaSizeInElements << " elements (" << totalSize
+ << " bytes), alignment=" << maxAlignment << " bytes\n");
// Step 5: Replace each alloc with a subview and remove deallocs
for (auto &candidate : candidates) {
@@ -149,9 +178,12 @@ struct StaticMemoryPlannerAnalysisPass
// Create subview into arena
llvm::SmallVector<mlir::OpFoldResult> offsets, sizes, strides;
- // Single offset into flat arena
- offsets.push_back(subviewBuilder.getIndexAttr(candidate.offset));
- sizes.push_back(subviewBuilder.getIndexAttr(candidate.sizeInElements));
+ // Convert byte offsets/sizes back to element indices for the arena
+ int64_t offsetInElements = candidate.offset / elementSizeInBytes;
+ int64_t sizeInElements = candidate.sizeInBytes / elementSizeInBytes;
+
+ offsets.push_back(subviewBuilder.getIndexAttr(offsetInElements));
+ sizes.push_back(subviewBuilder.getIndexAttr(sizeInElements));
strides.push_back(subviewBuilder.getIndexAttr(1));
auto subview = mlir::memref::SubViewOp::create(
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 1cde192103017..f793fd42d9947 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -6,7 +6,7 @@
// Test 1: Simple sequential alloc/dealloc pairs
// CHECK-LABEL: func @simple_sequential
func.func @simple_sequential() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<1536xf32>
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<1536xf32>
// CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1]>>
// CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1], offset: 1024>>
// CHECK-NOT: memref.alloc
@@ -62,7 +62,7 @@ func.func @different_block_skipped(%cond: i1) {
// Test 5: Overlapping lifetimes (both eligible, sequential offsets)
// CHECK-LABEL: func @overlapping_lifetimes
func.func @overlapping_lifetimes() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<1536xf32>
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<1536xf32>
// CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1]>>
// CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][512] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1], offset: 512>>
// CHECK-NOT: memref.alloc
@@ -79,7 +79,7 @@ func.func @overlapping_lifetimes() {
// Test 6: Multiple allocations with sequential offsets
// CHECK-LABEL: func @multiple_sequential
func.func @multiple_sequential() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() : memref<3584xf32>
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<3584xf32>
// CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<3584xf32> to memref<1024xf32, strided<[1]>>
// CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<3584xf32> to memref<512xf32, strided<[1], offset: 1024>>
// CHECK-NEXT: %[[SUBVIEW2:.*]] = memref.subview %[[ARENA]][1536] [2048] [1] : memref<3584xf32> to memref<2048xf32, strided<[1], offset: 1536>>
@@ -93,3 +93,28 @@ func.func @multiple_sequential() {
memref.dealloc %alloc2 : memref<2048xf32>
return
}
+
+// -----
+
+// Test 7: Alignment requirements with padding
+// CHECK-LABEL: func @alignment_padding
+func.func @alignment_padding() {
+ // Arena has max alignment (128 bytes)
+ // Total: 256*4 + 128*4 + 64*4 = 1792 bytes = 448 f32 elements
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 128 : i64} : memref<448xf32>
+ // First alloc: 256 f32, alignment=128, offset=0 bytes (0 elements)
+ // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [256] [1] : memref<448xf32> to memref<256xf32, strided<[1]>>
+ // Second alloc: 128 f32, alignment=64, offset=1024 bytes (256 elements, 64-aligned)
+ // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][256] [128] [1] : memref<448xf32> to memref<128xf32, strided<[1], offset: 256>>
+ // Third alloc: 64 f32, alignment=128, offset=1536 bytes (384 elements, 128-aligned)
+ // CHECK-NEXT: %[[SUBVIEW2:.*]] = memref.subview %[[ARENA]][384] [64] [1] : memref<448xf32> to memref<64xf32, strided<[1], offset: 384>>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %alloc0 = memref.alloc() {alignment = 128 : i64} : memref<256xf32>
+ memref.dealloc %alloc0 : memref<256xf32>
+ %alloc1 = memref.alloc() {alignment = 64 : i64} : memref<128xf32>
+ memref.dealloc %alloc1 : memref<128xf32>
+ %alloc2 = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
+ memref.dealloc %alloc2 : memref<64xf32>
+ return
+}
>From ac3e313c30444f8bebbfb3e0aa29c0a4910aa67f Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 15 Jun 2026 14:40:39 +0530
Subject: [PATCH 07/13] [mlir][bufferization] Extract memory planning into pure
function
Separate memory planning logic from IR transformation by introducing
trivialMemoryPlanner() - a pure function that computes buffer offsets
without depending on MLIR operations.
Changes:
- Add Alloc structure for allocation-independent planning
- Implement trivialMemoryPlanner(arenaAlignment, allocs) -> offsets
- Refactor runOnOperation() to use the planning function
- Planning logic is now testable independently of MLIR
This architecture enables plugging in different allocation strategies
(firstFit, bestFit) without modifying IR transformation code.
---
.../StaticMemoryPlannerAnalysis.cpp | 73 +++++++++++++++----
1 file changed, 59 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 9c1898aa96cc3..f9cd01bbd2c0b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -35,11 +35,19 @@ namespace {
// Data structures
//===----------------------------------------------------------------------===//
+/// Allocation info for memory planning (independent of MLIR).
+/// This can be used with pure planning algorithms.
+struct Alloc {
+ int64_t sizeInBytes = 0; // Size in bytes
+ int64_t alignment = 1; // Required alignment in bytes
+ // Note: time_start and time_end will be added later for lifetime-aware planning
+};
+
/// A candidate allocation with its matching deallocation and assigned offset.
struct AllocationCandidate {
mlir::memref::AllocOp alloc;
mlir::memref::DeallocOp dealloc;
- int64_t offset = 0; // Offset in bytes from arena start
+ 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
};
@@ -85,6 +93,32 @@ static int64_t alignOffset(int64_t offset, int64_t alignment) {
return (offset + alignment - 1) / alignment * alignment;
}
+//===----------------------------------------------------------------------===//
+// Memory Planning Algorithms
+//===----------------------------------------------------------------------===//
+
+/// Simple sequential memory planner (baseline algorithm).
+/// Allocates each buffer one after another with proper alignment padding.
+/// Returns offsets in bytes for each allocation.
+static llvm::SmallVector<int64_t>
+trivialMemoryPlanner(int64_t arenaAlignment,
+ llvm::ArrayRef<Alloc> allocs) {
+ llvm::SmallVector<int64_t> offsets;
+ int64_t currentOffset = 0;
+
+ for (const auto &alloc : allocs) {
+ // Ensure offset respects both arena alignment and this allocation's alignment
+ // The comment from mentor: (arenaAlignment + offset) % alloc.alignment == 0
+ // This means: if arena starts at an arenaAlignment boundary,
+ // then offset must make the final address properly aligned for alloc.alignment
+ currentOffset = alignOffset(currentOffset, alloc.alignment);
+ offsets.push_back(currentOffset);
+ currentOffset += alloc.sizeInBytes;
+ }
+
+ return offsets;
+}
+
//===----------------------------------------------------------------------===//
// StaticMemoryPlannerAnalysisPass
//===----------------------------------------------------------------------===//
@@ -136,22 +170,33 @@ struct StaticMemoryPlannerAnalysisPass
if (candidates.empty())
return;
- // Step 2: Compute simple sequential offsets with alignment padding
- int64_t totalSize = 0;
+ // Step 2: Prepare allocation info for planner
+ llvm::SmallVector<Alloc> allocInfos;
int64_t maxAlignment = 1;
- for (auto &candidate : candidates) {
- // Align current offset to this allocation's requirement
- totalSize = alignOffset(totalSize, candidate.alignment);
- candidate.offset = totalSize;
- totalSize += candidate.sizeInBytes;
+ for (const auto &candidate : candidates) {
+ Alloc allocInfo;
+ allocInfo.sizeInBytes = candidate.sizeInBytes;
+ allocInfo.alignment = candidate.alignment;
+ allocInfos.push_back(allocInfo);
maxAlignment = std::max(maxAlignment, candidate.alignment);
+ }
+
+ // Step 3: Run the planning algorithm
+ llvm::SmallVector<int64_t> offsets =
+ trivialMemoryPlanner(maxAlignment, allocInfos);
+
+ // Assign computed offsets back to candidates
+ int64_t totalSize = 0;
+ for (size_t i = 0; i < candidates.size(); ++i) {
+ candidates[i].offset = offsets[i];
+ totalSize = std::max(totalSize, offsets[i] + candidates[i].sizeInBytes);
LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] offset="
- << candidate.offset
- << " size=" << candidate.sizeInBytes
- << " alignment=" << candidate.alignment << "\n");
+ << candidates[i].offset
+ << " size=" << candidates[i].sizeInBytes
+ << " alignment=" << candidates[i].alignment << "\n");
}
- // Step 3: Find the first allocation's location to place arena
+ // Step 4: Find the first allocation's location to place arena
mlir::Operation *firstAlloc = candidates.front().alloc;
mlir::OpBuilder builder(firstAlloc);
@@ -160,7 +205,7 @@ struct StaticMemoryPlannerAnalysisPass
unsigned elementSizeInBits = elementType.getIntOrFloatBitWidth();
unsigned elementSizeInBytes = (elementSizeInBits + 7) / 8;
- // Step 4: Create arena allocation with maximum alignment
+ // Step 5: Create arena allocation with maximum alignment
// Convert total size from bytes to elements
int64_t arenaSizeInElements = (totalSize + elementSizeInBytes - 1) / elementSizeInBytes;
auto arenaType = mlir::MemRefType::get({arenaSizeInElements}, elementType);
@@ -171,7 +216,7 @@ struct StaticMemoryPlannerAnalysisPass
<< arenaSizeInElements << " elements (" << totalSize
<< " bytes), alignment=" << maxAlignment << " bytes\n");
- // Step 5: Replace each alloc with a subview and remove deallocs
+ // Step 6: Replace each alloc with a subview and remove deallocs
for (auto &candidate : candidates) {
mlir::OpBuilder subviewBuilder(candidate.alloc);
>From 0d2a04c31c0669b1eb87ca1504bbe5ddac60caa0 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 15 Jun 2026 14:51:59 +0530
Subject: [PATCH 08/13] [mlir][bufferization] Convert arena to i8 byte buffer
with memref.view
Change the arena from typed (e.g., memref<Nxf32>) to a generic i8 byte buffer
(memref<Nxi8>). This allows a single arena to hold allocations of different
element types (f32, i64, i16, etc.).
Use memref.view to create typed views into the i8 arena at computed byte
offsets. This is the standard MLIR pattern for type-agnostic memory buffers.
Changes:
- Arena is now memref<totalSizexi8> instead of element-typed
- Use memref.view instead of memref.subview + reinterpret_cast
- Byte offsets passed directly to memref.view via arith.constant
- Update all tests to reflect i8 arena + view pattern
Example transformation:
Before: memref.alloc() : memref<1024xf32>
After: %arena = memref.alloc() : memref<4096xi8>
%c0 = arith.constant 0 : index
%view = memref.view %arena[%c0][] : memref<4096xi8> to memref<1024xf32>
---
.../StaticMemoryPlannerAnalysis.cpp | 49 ++++++++--------
.../static-memory-planner-analysis.mlir | 57 ++++++++++++-------
2 files changed, 61 insertions(+), 45 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index f9cd01bbd2c0b..03cd9814d4f4f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -13,6 +13,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
@@ -199,44 +200,40 @@ struct StaticMemoryPlannerAnalysisPass
// Step 4: Find the first allocation's location to place arena
mlir::Operation *firstAlloc = candidates.front().alloc;
mlir::OpBuilder builder(firstAlloc);
-
- // Get element type from first allocation (assume all same type for now)
- mlir::Type elementType = candidates.front().alloc.getType().getElementType();
- unsigned elementSizeInBits = elementType.getIntOrFloatBitWidth();
- unsigned elementSizeInBytes = (elementSizeInBits + 7) / 8;
- // Step 5: Create arena allocation with maximum alignment
- // Convert total size from bytes to elements
- int64_t arenaSizeInElements = (totalSize + elementSizeInBytes - 1) / elementSizeInBytes;
- auto arenaType = mlir::MemRefType::get({arenaSizeInElements}, elementType);
+ // Step 5: Create arena allocation as i8 byte buffer
+ // This allows the same arena to hold multiple data types (f32, i64, etc.)
+ auto i8Type = builder.getI8Type();
+ auto arenaType = mlir::MemRefType::get({totalSize}, i8Type);
auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
arenaAlloc.setAlignmentAttr(builder.getI64IntegerAttr(maxAlignment));
LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena: size="
- << arenaSizeInElements << " elements (" << totalSize
- << " bytes), alignment=" << maxAlignment << " bytes\n");
+ << totalSize << " bytes, alignment="
+ << maxAlignment << " bytes\n");
- // Step 6: Replace each alloc with a subview and remove deallocs
+ // Step 6: Replace each alloc with memref.view directly on arena
for (auto &candidate : candidates) {
- mlir::OpBuilder subviewBuilder(candidate.alloc);
+ mlir::OpBuilder viewBuilder(candidate.alloc);
+ mlir::Location loc = candidate.alloc.getLoc();
- // Create subview into arena
- llvm::SmallVector<mlir::OpFoldResult> offsets, sizes, strides;
+ // Get the original memref type that we need to recreate
+ mlir::MemRefType originalType = candidate.alloc.getType();
- // Convert byte offsets/sizes back to element indices for the arena
- int64_t offsetInElements = candidate.offset / elementSizeInBytes;
- int64_t sizeInElements = candidate.sizeInBytes / elementSizeInBytes;
+ // Create a constant for the byte offset into the arena
+ mlir::Value offsetIndex = mlir::arith::ConstantIndexOp::create(
+ viewBuilder, loc, candidate.offset);
- offsets.push_back(subviewBuilder.getIndexAttr(offsetInElements));
- sizes.push_back(subviewBuilder.getIndexAttr(sizeInElements));
- strides.push_back(subviewBuilder.getIndexAttr(1));
+ // Use memref.view to create a typed view directly on the i8 arena
+ // memref.view: memref<...xi8>, offset -> memref<shape x type>
+ llvm::SmallVector<mlir::Value> dynamicSizes; // Empty for static shapes
- auto subview = mlir::memref::SubViewOp::create(
- subviewBuilder, candidate.alloc.getLoc(), arenaAlloc.getResult(),
- offsets, sizes, strides);
+ auto view = mlir::memref::ViewOp::create(
+ viewBuilder, loc, originalType, arenaAlloc.getResult(),
+ offsetIndex, dynamicSizes);
- // Replace all uses of the original alloc
- candidate.alloc.getResult().replaceAllUsesWith(subview.getResult());
+ // Replace all uses of the original alloc with the viewed memref
+ candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
// Remove the original alloc and dealloc
candidate.alloc.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 f793fd42d9947..3648558740b6e 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -6,9 +6,14 @@
// Test 1: Simple sequential alloc/dealloc pairs
// CHECK-LABEL: func @simple_sequential
func.func @simple_sequential() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<1536xf32>
- // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1]>>
- // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1], offset: 1024>>
+ // Arena is i8 buffer: 1024*4 + 512*4 = 6144 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<6144xi8>
+ // First allocation at offset 0
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<6144xi8> to memref<1024xf32>
+ // Second allocation at offset 4096 bytes (1024 * 4)
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C4096]]][] : memref<6144xi8> to memref<512xf32>
// CHECK-NOT: memref.alloc
// CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<1024xf32>
@@ -62,9 +67,14 @@ func.func @different_block_skipped(%cond: i1) {
// Test 5: Overlapping lifetimes (both eligible, sequential offsets)
// CHECK-LABEL: func @overlapping_lifetimes
func.func @overlapping_lifetimes() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<1536xf32>
- // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [512] [1] : memref<1536xf32> to memref<512xf32, strided<[1]>>
- // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][512] [1024] [1] : memref<1536xf32> to memref<1024xf32, strided<[1], offset: 512>>
+ // Arena: 512*4 + 1024*4 = 6144 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<6144xi8>
+ // First allocation at offset 0
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<6144xi8> to memref<512xf32>
+ // Second allocation at offset 2048 bytes (512 * 4)
+ // CHECK-NEXT: %[[C2048:.*]] = arith.constant 2048 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C2048]]][] : memref<6144xi8> to memref<1024xf32>
// CHECK-NOT: memref.alloc
// CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<512xf32>
@@ -79,10 +89,17 @@ func.func @overlapping_lifetimes() {
// Test 6: Multiple allocations with sequential offsets
// CHECK-LABEL: func @multiple_sequential
func.func @multiple_sequential() {
- // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<3584xf32>
- // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [1024] [1] : memref<3584xf32> to memref<1024xf32, strided<[1]>>
- // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][1024] [512] [1] : memref<3584xf32> to memref<512xf32, strided<[1], offset: 1024>>
- // CHECK-NEXT: %[[SUBVIEW2:.*]] = memref.subview %[[ARENA]][1536] [2048] [1] : memref<3584xf32> to memref<2048xf32, strided<[1], offset: 1536>>
+ // Arena: 1024*4 + 512*4 + 2048*4 = 14336 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<14336xi8>
+ // First at offset 0
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<14336xi8> to memref<1024xf32>
+ // Second at offset 4096 bytes (1024 * 4)
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C4096]]][] : memref<14336xi8> to memref<512xf32>
+ // Third at offset 6144 bytes (1024*4 + 512*4)
+ // CHECK-NEXT: %[[C6144:.*]] = arith.constant 6144 : index
+ // CHECK-NEXT: %[[VIEW2:.*]] = memref.view %[[ARENA]][%[[C6144]]][] : memref<14336xi8> to memref<2048xf32>
// CHECK-NOT: memref.alloc
// CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() : memref<1024xf32>
@@ -99,15 +116,17 @@ func.func @multiple_sequential() {
// Test 7: Alignment requirements with padding
// CHECK-LABEL: func @alignment_padding
func.func @alignment_padding() {
- // Arena has max alignment (128 bytes)
- // Total: 256*4 + 128*4 + 64*4 = 1792 bytes = 448 f32 elements
- // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 128 : i64} : memref<448xf32>
- // First alloc: 256 f32, alignment=128, offset=0 bytes (0 elements)
- // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[ARENA]][0] [256] [1] : memref<448xf32> to memref<256xf32, strided<[1]>>
- // Second alloc: 128 f32, alignment=64, offset=1024 bytes (256 elements, 64-aligned)
- // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[ARENA]][256] [128] [1] : memref<448xf32> to memref<128xf32, strided<[1], offset: 256>>
- // Third alloc: 64 f32, alignment=128, offset=1536 bytes (384 elements, 128-aligned)
- // CHECK-NEXT: %[[SUBVIEW2:.*]] = memref.subview %[[ARENA]][384] [64] [1] : memref<448xf32> to memref<64xf32, strided<[1], offset: 384>>
+ // Arena has max alignment (128 bytes), total: 256*4 + 128*4 + 64*4 = 1792 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 128 : i64} : memref<1792xi8>
+ // First alloc: 256 f32, alignment=128, offset=0 bytes (128-aligned)
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<1792xi8> to memref<256xf32>
+ // Second alloc: 128 f32, alignment=64, offset=1024 bytes (64-aligned)
+ // CHECK-NEXT: %[[C1024:.*]] = arith.constant 1024 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C1024]]][] : memref<1792xi8> to memref<128xf32>
+ // Third alloc: 64 f32, alignment=128, offset=1536 bytes (128-aligned)
+ // CHECK-NEXT: %[[C1536:.*]] = arith.constant 1536 : index
+ // CHECK-NEXT: %[[VIEW2:.*]] = memref.view %[[ARENA]][%[[C1536]]][] : memref<1792xi8> to memref<64xf32>
// CHECK-NOT: memref.alloc
// CHECK-NOT: memref.dealloc
%alloc0 = memref.alloc() {alignment = 128 : i64} : memref<256xf32>
>From f30ab188cc144527ba2fadc621ef668640b9cc98 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 15 Jun 2026 15:14:33 +0530
Subject: [PATCH 09/13] [mlir][bufferization] Add arena-mode pass option
(allocate vs arg)
Add arena-mode pass option to control how the shared arena is obtained:
- 'allocate' (default): Creates arena via memref.alloc within the function
- 'arg': Uses function's first argument as the pre-allocated arena
The 'arg' mode is useful when the arena is pre-allocated externally and
passed to the function, enabling use cases like pre-allocated scratch
buffers or memory pools.
In 'arg' mode, the pass validates that:
1. The context is a function operation
2. The function has at least one argument
3. The first argument is memref<...xi8>
If validation fails, the pass emits an error and fails gracefully.
Changes:
- Add arena-mode option to Passes.td with default 'allocate'
- Update pass description to reflect transformation behavior
- Implement conditional arena acquisition based on mode
- Add tests for arg mode with error validation
---
.../Bufferization/Transforms/Passes.td | 49 +++++++++++------
.../StaticMemoryPlannerAnalysis.cpp | 55 +++++++++++++++----
.../static-memory-planner-arena-arg.mlir | 34 ++++++++++++
3 files changed, 108 insertions(+), 30 deletions(-)
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 54c9d6ad14442..3e6931051b49e 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -186,12 +186,13 @@ def OptimizeAllocationLivenessPass
def StaticMemoryPlannerAnalysisPass
: Pass<"static-memory-planner-analysis"> {
- let summary = "Identifies same-block alloc/dealloc pairs eligible for "
- "static memory planning";
+ let summary = "Transforms same-block alloc/dealloc pairs into static arena "
+ "allocation with computed offsets";
let description = [{
- This analysis-only pass identifies `memref.alloc` / `memref.dealloc` pairs
- in the same basic block that are candidates for static memory planning and
- buffer reuse.
+ This pass identifies and transforms `memref.alloc` / `memref.dealloc` pairs
+ in the same basic block into a single static arena allocation. All eligible
+ allocations are packed into a shared i8 byte buffer with aligned offsets,
+ and accessed via `memref.view` operations.
For each allocation the pass checks a conservative eligibility envelope:
- Static memref shape.
@@ -201,35 +202,47 @@ def StaticMemoryPlannerAnalysisPass
- No cross-block alias or escaping use (alias set resolved via
`BufferViewFlowAnalysis`).
- Eligible allocations receive an alias-aware lifetime interval computed
- using `BufferViewFlowAnalysis::resolve()`. The interval spans from the
- alloc op's block-local index to the maximum of the dealloc index and the
- last use index of any derived alias within the block. Allocation size in
- bytes (for integer and floating-point element types) and alignment are
- also reported where available.
+ Eligible allocations are packed into a single arena using a trivial
+ sequential allocation strategy with alignment padding. The arena is an
+ i8 byte buffer (`memref<Nxi8>`) that can hold allocations of different
+ element types. Each original allocation is replaced with a `memref.view`
+ operation that creates a typed view into the arena at the computed offset.
- Ineligible allocations receive a skip reason. Both eligibility results
- and skip reasons are reported via op remarks. Aggregate counts are
- reported as pass statistics.
-
- No IR modifications are made by this pass.
+ Ineligible allocations are skipped and retain their original
+ alloc/dealloc operations. Skip reasons are reported via op remarks.
+ Aggregate counts are reported as pass statistics.
This pass is expected to run after the deallocation pipeline
(`ownership-based-buffer-deallocation` followed by
`bufferization-lower-deallocations`).
- Example:
+ Example transformation:
```mlir
+ // Before:
func.func @example() {
- // Remark: static-memory-planner: eligible: size=4096 bytes, interval=[0,2]
%0 = memref.alloc() : memref<1024xf32>
"some.use"(%0) : (memref<1024xf32>) -> ()
memref.dealloc %0 : memref<1024xf32>
return
}
+
+ // After:
+ func.func @example() {
+ %arena = memref.alloc() {alignment = 1 : i64} : memref<4096xi8>
+ %c0 = arith.constant 0 : index
+ %0 = memref.view %arena[%c0][] : memref<4096xi8> to memref<1024xf32>
+ "some.use"(%0) : (memref<1024xf32>) -> ()
+ return
+ }
```
}];
+ let options =
+ [Option<"arenaMode", "arena-mode", "std::string",
+ /*default=*/"\"allocate\"",
+ "Arena allocation mode: 'allocate' creates arena via AllocOp, "
+ "'arg' extracts arena from function arguments">];
+
let statistics =
[Statistic<"numEligible", "num-eligible",
"Number of alloc/dealloc pairs eligible for static planning">,
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 03cd9814d4f4f..017a76d957b43 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -128,7 +128,9 @@ struct StaticMemoryPlannerAnalysisPass
: public mlir::bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
StaticMemoryPlannerAnalysisPass> {
public:
- StaticMemoryPlannerAnalysisPass() = default;
+ using Base = mlir::bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
+ StaticMemoryPlannerAnalysisPass>;
+ using Base::Base;
void runOnOperation() override {
mlir::Operation *op = getOperation();
@@ -197,20 +199,49 @@ struct StaticMemoryPlannerAnalysisPass
<< " alignment=" << candidates[i].alignment << "\n");
}
- // Step 4: Find the first allocation's location to place arena
+ // Step 4: Obtain arena based on arena mode
mlir::Operation *firstAlloc = candidates.front().alloc;
mlir::OpBuilder builder(firstAlloc);
+ mlir::Value arenaValue;
- // Step 5: Create arena allocation as i8 byte buffer
- // This allows the same arena to hold multiple data types (f32, i64, etc.)
- auto i8Type = builder.getI8Type();
- auto arenaType = mlir::MemRefType::get({totalSize}, i8Type);
- auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
- arenaAlloc.setAlignmentAttr(builder.getI64IntegerAttr(maxAlignment));
+ if (arenaMode == "allocate") {
+ // Step 5a: Create arena via AllocOp (default mode)
+ // Arena is i8 byte buffer to support multiple data types (f32, i64, etc.)
+ auto i8Type = builder.getI8Type();
+ auto arenaType = mlir::MemRefType::get({totalSize}, i8Type);
+ auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
+ arenaAlloc.setAlignmentAttr(builder.getI64IntegerAttr(maxAlignment));
+ arenaValue = arenaAlloc.getResult();
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena: size="
- << totalSize << " bytes, alignment="
- << maxAlignment << " bytes\n");
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena via AllocOp: size="
+ << totalSize << " bytes, alignment="
+ << maxAlignment << " bytes\n");
+ } else if (arenaMode == "arg") {
+ // Step 5b: Extract arena from function arguments (arg mode)
+ // Assumes first argument is an i8 buffer of sufficient size
+ auto funcOp = llvm::dyn_cast<mlir::func::FuncOp>(op);
+ if (!funcOp) {
+ op->emitError("arena-mode=arg requires function context");
+ return signalPassFailure();
+ }
+
+ if (funcOp.getNumArguments() == 0) {
+ funcOp.emitError("arena-mode=arg requires at least one function argument");
+ return signalPassFailure();
+ }
+
+ arenaValue = funcOp.getArgument(0);
+ auto arenaType = llvm::dyn_cast<mlir::MemRefType>(arenaValue.getType());
+ if (!arenaType || !arenaType.getElementType().isInteger(8)) {
+ funcOp.emitError("arena-mode=arg requires first argument to be memref<...xi8>");
+ return signalPassFailure();
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] using arena from function arg 0\n");
+ } else {
+ op->emitError("invalid arena-mode: '" + arenaMode + "' (must be 'allocate' or 'arg')");
+ return signalPassFailure();
+ }
// Step 6: Replace each alloc with memref.view directly on arena
for (auto &candidate : candidates) {
@@ -229,7 +260,7 @@ struct StaticMemoryPlannerAnalysisPass
llvm::SmallVector<mlir::Value> dynamicSizes; // Empty for static shapes
auto view = mlir::memref::ViewOp::create(
- viewBuilder, loc, originalType, arenaAlloc.getResult(),
+ viewBuilder, loc, originalType, arenaValue,
offsetIndex, dynamicSizes);
// Replace all uses of the original alloc with the viewed memref
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
new file mode 100644
index 0000000000000..b58935d351142
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
@@ -0,0 +1,34 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(static-memory-planner-analysis{arena-mode=arg}))" \
+// RUN: -split-input-file -verify-diagnostics
+
+// -----
+
+// Test 1: Arena from function argument
+func.func @arena_from_arg(%arena: memref<8192xi8>) {
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+
+ %alloc1 = memref.alloc() : memref<512xf32>
+ memref.dealloc %alloc1 : memref<512xf32>
+ return
+}
+
+// -----
+
+// Test 2: Error when no function argument
+// expected-error @+1 {{arena-mode=arg requires at least one function argument}}
+func.func @error_no_args() {
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 3: Error when first argument is not i8 memref
+// expected-error @+1 {{arena-mode=arg requires first argument to be memref<...xi8>}}
+func.func @error_wrong_type(%arena: memref<8192xf32>) {
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ return
+}
>From 8b10fdf6574c30bc85867512ae285ad80796351d Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 15 Jun 2026 15:22:36 +0530
Subject: [PATCH 10/13] [mlir][bufferization] Add error for memref return types
Add validation to reject functions with memref return types, as static
memory planning is incompatible with returning memrefs. In allocate mode,
the arena is freed at function exit, making returned memrefs invalid. In
arg mode, returning a memref from the input arena violates typical memory
ownership patterns.
When a function has memref return types, the pass:
1. Emits a clear error message
2. Fails gracefully without transforming the function
3. Preserves the original IR
This prevents silent bugs where returned memrefs would point to freed or
external memory.
Changes:
- Add return type validation at start of runOnOperation()
- Check all result types for MemRefType
- Emit descriptive error and signal pass failure
- Add test case verifying error emission
---
.../Transforms/StaticMemoryPlannerAnalysis.cpp | 12 ++++++++++++
.../Transforms/static-memory-planner-arena-arg.mlir | 11 +++++++++++
2 files changed, 23 insertions(+)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 017a76d957b43..3d007a79a8045 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -135,6 +135,18 @@ struct StaticMemoryPlannerAnalysisPass
void runOnOperation() override {
mlir::Operation *op = getOperation();
+ // Step 0: Check for memref return types (not supported)
+ if (auto funcOp = llvm::dyn_cast<mlir::func::FuncOp>(op)) {
+ mlir::FunctionType funcType = funcOp.getFunctionType();
+ for (mlir::Type resultType : funcType.getResults()) {
+ if (llvm::isa<mlir::MemRefType>(resultType)) {
+ funcOp.emitError("static-memory-planner does not support functions "
+ "with memref return types");
+ return signalPassFailure();
+ }
+ }
+ }
+
// Step 1: Collect eligible allocation candidates
llvm::SmallVector<AllocationCandidate> candidates;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
index b58935d351142..1ad7e1a008e20 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-arena-arg.mlir
@@ -32,3 +32,14 @@ func.func @error_wrong_type(%arena: memref<8192xf32>) {
memref.dealloc %alloc0 : memref<1024xf32>
return
}
+
+// -----
+
+// Test 4: Error when function returns memref
+// expected-error @+1 {{static-memory-planner does not support functions with memref return types}}
+func.func @error_memref_return(%arena: memref<8192xi8>) -> memref<1024xf32> {
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ %result = memref.alloc() : memref<1024xf32>
+ return %result : memref<1024xf32>
+}
>From f4d2dd82064062788b496702c3cb5e1d62d03a11 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Tue, 23 Jun 2026 12:35:18 +0530
Subject: [PATCH 11/13] [mlir][bufferization] Fix code formatting
---
.../Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 3d007a79a8045..0442acdbc2276 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -197,9 +197,9 @@ struct StaticMemoryPlannerAnalysisPass
}
// Step 3: Run the planning algorithm
- llvm::SmallVector<int64_t> offsets =
+ llvm::SmallVector<int64_t> offsets =
trivialMemoryPlanner(maxAlignment, allocInfos);
-
+
// Assign computed offsets back to candidates
int64_t totalSize = 0;
for (size_t i = 0; i < candidates.size(); ++i) {
>From 4290417bab4e7b85b5775e4c8b8776c7e69e3325 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 29 Jun 2026 10:44:35 +0530
Subject: [PATCH 12/13] [mlir][bufferization] Address review feedback on static
memory planner
- Use InterfacePass<..., FunctionOpInterface> so the pass only runs on
function-like ops; cast to FunctionOpInterface in runOnOperation
- Update pass description: mention memref.alloc/dealloc explicitly,
drop outdated scf.for list and BufferViewFlowAnalysis references,
adopt pluggable-planner wording from javedabsar
- Add time_start/time_end to Alloc struct for future lifetime-aware
planning
- Use getNumElements(), llvm::alignTo, std::lcm (LCM is correct for
arena alignment; std::max was wrong)
- Add precondition doc + NDEBUG assert in trivialMemoryPlanner
- Check BaseMemRefType for return type guard to cover unranked memrefs
- Add 1D rank check for arena arg memref
- Use AllocOp::create overload that takes alignment directly
- Reuse OpBuilder via setInsertionPoint instead of creating a new one
- Drop mlir:: namespace prefixes (using namespace mlir)
---
.../Bufferization/Transforms/Passes.td | 109 +++++----
.../StaticMemoryPlannerAnalysis.cpp | 207 ++++++++----------
2 files changed, 146 insertions(+), 170 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 3e6931051b49e..c7f6cae571e40 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -146,10 +146,10 @@ def OwnershipBasedBufferDeallocationPass
"earlier deallocations.">,
];
- let dependentDialects = [
- "mlir::bufferization::BufferizationDialect", "mlir::arith::ArithDialect",
- "mlir::memref::MemRefDialect", "mlir::scf::SCFDialect"
- ];
+ let dependentDialects = ["mlir::bufferization::BufferizationDialect",
+ "mlir::arith::ArithDialect",
+ "mlir::memref::MemRefDialect",
+ "mlir::scf::SCFDialect"];
}
def BufferDeallocationSimplificationPass
@@ -163,10 +163,9 @@ def BufferDeallocationSimplificationPass
some memref isn't deallocated twice (double free).
}];
- let dependentDialects = [
- "mlir::bufferization::BufferizationDialect", "mlir::arith::ArithDialect",
- "mlir::memref::MemRefDialect"
- ];
+ let dependentDialects = ["mlir::bufferization::BufferizationDialect",
+ "mlir::arith::ArithDialect",
+ "mlir::memref::MemRefDialect"];
}
def OptimizeAllocationLivenessPass
@@ -185,28 +184,26 @@ def OptimizeAllocationLivenessPass
}
def StaticMemoryPlannerAnalysisPass
- : Pass<"static-memory-planner-analysis"> {
+ : InterfacePass<"static-memory-planner-analysis", "FunctionOpInterface"> {
let summary = "Transforms same-block alloc/dealloc pairs into static arena "
"allocation with computed offsets";
let description = [{
This pass identifies and transforms `memref.alloc` / `memref.dealloc` pairs
- in the same basic block into a single static arena allocation. All eligible
- allocations are packed into a shared i8 byte buffer with aligned offsets,
- and accessed via `memref.view` operations.
+ in the same basic block into a single static arena allocation.
- For each allocation the pass checks a conservative eligibility envelope:
+ For each `memref.alloc` the pass checks a conservative eligibility
+ envelope:
- Static memref shape.
- - Unique same-block dealloc (found via `MemoryEffects::Free`).
- - Not nested inside any loop or conditional region (e.g., `scf.for`,
- `scf.while`, `scf.forall`, `scf.parallel`, `scf.if`).
- - No cross-block alias or escaping use (alias set resolved via
- `BufferViewFlowAnalysis`).
-
- Eligible allocations are packed into a single arena using a trivial
- sequential allocation strategy with alignment padding. The arena is an
- i8 byte buffer (`memref<Nxi8>`) that can hold allocations of different
- element types. Each original allocation is replaced with a `memref.view`
- operation that creates a typed view into the arena at the computed offset.
+ - Unique same-block `memref.dealloc`.
+ - Allocations in nested blocks are ignored for now.
+
+ Eligible allocations are packed into a single arena. Currently, we use a
+ trivial sequential allocation strategy with alignment padding. But the
+ interface will also allow any packing algorithm to be plugged in as long as
+ it respects the interface. The arena is an i8 byte buffer
+ (`memref<Nxi8>`) that can hold allocations of different element types.
+ Each original allocation is replaced with a `memref.view` operation that
+ creates a typed view into the arena at the computed offset.
Ineligible allocations are skipped and retain their original
alloc/dealloc operations. Skip reasons are reported via op remarks.
@@ -237,11 +234,11 @@ def StaticMemoryPlannerAnalysisPass
```
}];
- let options =
- [Option<"arenaMode", "arena-mode", "std::string",
- /*default=*/"\"allocate\"",
- "Arena allocation mode: 'allocate' creates arena via AllocOp, "
- "'arg' extracts arena from function arguments">];
+ let options = [Option<
+ "arenaMode", "arena-mode", "std::string",
+ /*default=*/"\"allocate\"",
+ "Arena allocation mode: 'allocate' creates arena via AllocOp, "
+ "'arg' extracts arena from function arguments">];
let statistics =
[Statistic<"numEligible", "num-eligible",
@@ -273,10 +270,8 @@ def LowerDeallocationsPass : Pass<"bufferization-lower-deallocations"> {
library functions to avoid code-size blow-up.
}];
- let dependentDialects = [
- "arith::ArithDialect", "memref::MemRefDialect", "scf::SCFDialect",
- "func::FuncDialect"
- ];
+ let dependentDialects = ["arith::ArithDialect", "memref::MemRefDialect",
+ "scf::SCFDialect", "func::FuncDialect"];
}
def BufferHoistingPass : Pass<"buffer-hoisting", "func::FuncOp"> {
@@ -334,7 +329,8 @@ def BufferResultsToOutParamsPass
Option<"hoistDynamicAllocs", "hoist-dynamic-allocs", "bool",
/*default=*/"false", "Hoist dynamic allocations to call sites.">,
Option<"modifyPublicFunctions", "modify-public-functions", "bool",
- /*default=*/"false", "Modify function signatures of public "
+ /*default=*/"false",
+ "Modify function signatures of public "
"functions.">,
];
let dependentDialects = ["memref::MemRefDialect"];
@@ -351,10 +347,11 @@ def DropEquivalentBufferResultsPass
Note: If a bbArg buffer is not returned directly but casted to beforehand,
the buffer is still considered equivalent.
}];
- let options = [
- Option<"modifyPublicFunctions", "modify-public-functions", "bool",
- /*default=*/"false", "Modify function signatures of public "
- "functions.">,
+ let options = [Option<"modifyPublicFunctions", "modify-public-functions",
+ "bool",
+ /*default=*/"false",
+ "Modify function signatures of public "
+ "functions.">,
];
let dependentDialects = ["memref::MemRefDialect"];
}
@@ -539,7 +536,8 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> {
/*default=*/"false",
"Test only: Annotate IR with RaW conflicts. Requires "
"test-analysis-only.">,
- Option<"unknownTypeConversion", "unknown-type-conversion", "LayoutMapOption",
+ Option<"unknownTypeConversion", "unknown-type-conversion",
+ "LayoutMapOption",
/*default=*/"LayoutMapOption::FullyDynamicLayoutMap",
"Controls layout maps for non-inferrable memref types.",
layoutMapClValues.values>,
@@ -548,18 +546,16 @@ def OneShotBufferizePass : Pass<"one-shot-bufferize", "ModuleOp"> {
"Sets the alignment of newly allocated buffers.">,
];
- let statistics = [
- Statistic<"numBufferAlloc", "num-buffer-alloc",
- "Number of buffer allocations">,
- Statistic<"numTensorInPlace", "num-tensor-in-place",
- "Number of in-place tensor OpOperands">,
- Statistic<"numTensorOutOfPlace", "num-tensor-out-of-place",
- "Number of out-of-place tensor OpOperands">,
+ let statistics = [Statistic<"numBufferAlloc", "num-buffer-alloc",
+ "Number of buffer allocations">,
+ Statistic<"numTensorInPlace", "num-tensor-in-place",
+ "Number of in-place tensor OpOperands">,
+ Statistic<"numTensorOutOfPlace", "num-tensor-out-of-place",
+ "Number of out-of-place tensor OpOperands">,
];
- let dependentDialects = [
- "bufferization::BufferizationDialect", "memref::MemRefDialect"
- ];
+ let dependentDialects = ["bufferization::BufferizationDialect",
+ "memref::MemRefDialect"];
}
def PromoteBuffersToStackPass
@@ -573,13 +569,14 @@ def PromoteBuffersToStackPass
shaped buffers that are limited by the rank of the tensor can be
converted. They are only transformed if they are considered to be small.
}];
- let options = [
- Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
- /*default=*/"1024",
- "Maximal size in bytes to promote allocations to stack.">,
- Option<"maxRankOfAllocatedMemRef", "max-rank-of-allocated-memref", "unsigned",
- /*default=*/"1",
- "Maximal memref rank to promote dynamic buffers.">,
+ let options =
+ [Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
+ /*default=*/"1024",
+ "Maximal size in bytes to promote allocations to stack.">,
+ Option<"maxRankOfAllocatedMemRef", "max-rank-of-allocated-memref",
+ "unsigned",
+ /*default=*/"1",
+ "Maximal memref rank to promote dynamic buffers.">,
];
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 0442acdbc2276..48b6e97c809ef 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -7,19 +7,19 @@
//===----------------------------------------------------------------------===//
//
// Transforms memref.alloc/memref.dealloc pairs into a single arena allocation
-// with subviews. Uses simple sequential offset assignment where each allocation
-// gets its own space without overlap (baseline algorithm for e2e testing).
+// with memref.view. Uses simple sequential offset assignment where each
+// allocation gets its own space without overlap (baseline algorithm).
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
-#include "mlir/IR/Operation.h"
-#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include <numeric>
#define DEBUG_TYPE "static-memory-planner"
@@ -30,6 +30,8 @@ namespace bufferization {
} // namespace bufferization
} // namespace mlir
+using namespace mlir;
+
namespace {
//===----------------------------------------------------------------------===//
@@ -39,18 +41,19 @@ namespace {
/// Allocation info for memory planning (independent of MLIR).
/// This can be used with pure planning algorithms.
struct Alloc {
- int64_t sizeInBytes = 0; // Size in bytes
- int64_t alignment = 1; // Required alignment in bytes
- // Note: time_start and time_end will be added later for lifetime-aware planning
+ int64_t sizeInBytes = 0; // Size in bytes
+ int64_t alignment = 1; // Required alignment in bytes
+ int64_t timeStart = 0; // Operation index when allocation starts
+ int64_t timeEnd = 0; // Operation index when allocation ends (dealloc)
};
/// A candidate allocation with its matching deallocation and assigned offset.
struct AllocationCandidate {
- mlir::memref::AllocOp alloc;
- mlir::memref::DeallocOp dealloc;
- 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
+ memref::AllocOp alloc;
+ memref::DeallocOp dealloc;
+ 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
};
//===----------------------------------------------------------------------===//
@@ -59,10 +62,10 @@ struct AllocationCandidate {
/// Finds the unique dealloc operation for a given alloc value.
/// Returns nullptr if there are zero or multiple deallocs.
-static mlir::memref::DeallocOp findUniqueDealloc(mlir::Value allocValue) {
- mlir::memref::DeallocOp deallocOp = nullptr;
- for (mlir::Operation *user : allocValue.getUsers()) {
- if (auto dealloc = mlir::dyn_cast<mlir::memref::DeallocOp>(user)) {
+static memref::DeallocOp findUniqueDealloc(Value allocValue) {
+ memref::DeallocOp deallocOp = nullptr;
+ for (Operation *user : allocValue.getUsers()) {
+ if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
if (deallocOp)
return nullptr; // Multiple deallocs found
deallocOp = dealloc;
@@ -71,17 +74,9 @@ static mlir::memref::DeallocOp findUniqueDealloc(mlir::Value allocValue) {
return deallocOp;
}
-/// Compute the number of elements in a static-shape memref.
-static int64_t computeSizeInElements(mlir::MemRefType memrefType) {
- int64_t size = 1;
- for (int64_t dim : memrefType.getShape())
- size *= dim;
- return size;
-}
-
/// Compute the size in bytes for a memref type.
-static int64_t computeSizeInBytes(mlir::MemRefType memrefType) {
- int64_t numElements = computeSizeInElements(memrefType);
+static int64_t computeSizeInBytes(MemRefType memrefType) {
+ int64_t numElements = memrefType.getNumElements();
unsigned elementSizeInBits = memrefType.getElementTypeBitWidth();
return (numElements * elementSizeInBits + 7) / 8; // Round up to bytes
}
@@ -89,9 +84,7 @@ static int64_t computeSizeInBytes(mlir::MemRefType memrefType) {
/// Align an offset to the specified alignment.
/// Returns the smallest value >= offset that is a multiple of alignment.
static int64_t alignOffset(int64_t offset, int64_t alignment) {
- if (alignment <= 1)
- return offset;
- return (offset + alignment - 1) / alignment * alignment;
+ return llvm::alignTo(offset, alignment);
}
//===----------------------------------------------------------------------===//
@@ -99,24 +92,24 @@ static int64_t alignOffset(int64_t offset, int64_t alignment) {
//===----------------------------------------------------------------------===//
/// Simple sequential memory planner (baseline algorithm).
+/// arenaAlignment must be a multiple (LCM) of all alloc.alignment values.
/// Allocates each buffer one after another with proper alignment padding.
/// Returns offsets in bytes for each allocation.
-static llvm::SmallVector<int64_t>
-trivialMemoryPlanner(int64_t arenaAlignment,
- llvm::ArrayRef<Alloc> allocs) {
- llvm::SmallVector<int64_t> offsets;
+static SmallVector<int64_t> trivialMemoryPlanner(int64_t arenaAlignment,
+ ArrayRef<Alloc> allocs) {
+ SmallVector<int64_t> offsets;
int64_t currentOffset = 0;
-
+
for (const auto &alloc : allocs) {
- // Ensure offset respects both arena alignment and this allocation's alignment
- // The comment from mentor: (arenaAlignment + offset) % alloc.alignment == 0
- // This means: if arena starts at an arenaAlignment boundary,
- // then offset must make the final address properly aligned for alloc.alignment
currentOffset = alignOffset(currentOffset, alloc.alignment);
+#ifndef NDEBUG
+ assert((arenaAlignment + currentOffset) % alloc.alignment == 0 &&
+ "invalid alignment");
+#endif
offsets.push_back(currentOffset);
currentOffset += alloc.sizeInBytes;
}
-
+
return offsets;
}
@@ -125,33 +118,30 @@ trivialMemoryPlanner(int64_t arenaAlignment,
//===----------------------------------------------------------------------===//
struct StaticMemoryPlannerAnalysisPass
- : public mlir::bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
+ : public bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
StaticMemoryPlannerAnalysisPass> {
public:
- using Base = mlir::bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
+ using Base = bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
StaticMemoryPlannerAnalysisPass>;
using Base::Base;
void runOnOperation() override {
- mlir::Operation *op = getOperation();
+ auto funcOp = llvm::cast<FunctionOpInterface>(getOperation());
// Step 0: Check for memref return types (not supported)
- if (auto funcOp = llvm::dyn_cast<mlir::func::FuncOp>(op)) {
- mlir::FunctionType funcType = funcOp.getFunctionType();
- for (mlir::Type resultType : funcType.getResults()) {
- if (llvm::isa<mlir::MemRefType>(resultType)) {
- funcOp.emitError("static-memory-planner does not support functions "
- "with memref return types");
- return signalPassFailure();
- }
+ for (Type resultType : funcOp.getResultTypes()) {
+ if (isa<BaseMemRefType>(resultType)) {
+ funcOp->emitError("static-memory-planner does not support functions "
+ "with memref return types");
+ return signalPassFailure();
}
}
// Step 1: Collect eligible allocation candidates
- llvm::SmallVector<AllocationCandidate> candidates;
+ SmallVector<AllocationCandidate> candidates;
- op->walk([&](mlir::memref::AllocOp allocOp) {
- mlir::MemRefType memrefType = allocOp.getType();
+ funcOp->walk([&](memref::AllocOp allocOp) {
+ MemRefType memrefType = allocOp.getType();
// Skip dynamic shapes
if (!memrefType.hasStaticShape()) {
@@ -160,7 +150,7 @@ struct StaticMemoryPlannerAnalysisPass
}
// Find unique dealloc in the same block
- mlir::memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
+ memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
if (!deallocOp) {
++numSkipNoDealloc;
return;
@@ -177,7 +167,6 @@ struct StaticMemoryPlannerAnalysisPass
candidate.alloc = allocOp;
candidate.dealloc = deallocOp;
candidate.sizeInBytes = computeSizeInBytes(memrefType);
- // Extract alignment requirement (default to 1 if not specified)
candidate.alignment = allocOp.getAlignment().value_or(1);
candidates.push_back(candidate);
});
@@ -186,99 +175,89 @@ struct StaticMemoryPlannerAnalysisPass
return;
// Step 2: Prepare allocation info for planner
- llvm::SmallVector<Alloc> allocInfos;
- int64_t maxAlignment = 1;
+ SmallVector<Alloc> allocInfos;
+ int64_t arenaAlignment = 1;
for (const auto &candidate : candidates) {
Alloc allocInfo;
allocInfo.sizeInBytes = candidate.sizeInBytes;
allocInfo.alignment = candidate.alignment;
allocInfos.push_back(allocInfo);
- maxAlignment = std::max(maxAlignment, candidate.alignment);
+ arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
}
// Step 3: Run the planning algorithm
- llvm::SmallVector<int64_t> offsets =
- trivialMemoryPlanner(maxAlignment, allocInfos);
+ SmallVector<int64_t> offsets =
+ trivialMemoryPlanner(arenaAlignment, allocInfos);
// Assign computed offsets back to candidates
int64_t totalSize = 0;
for (size_t i = 0; i < candidates.size(); ++i) {
candidates[i].offset = offsets[i];
totalSize = std::max(totalSize, offsets[i] + candidates[i].sizeInBytes);
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] offset="
- << candidates[i].offset
- << " size=" << candidates[i].sizeInBytes
- << " alignment=" << candidates[i].alignment << "\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << "[static-memory-planner] offset=" << candidates[i].offset
+ << " size=" << candidates[i].sizeInBytes
+ << " alignment=" << candidates[i].alignment << "\n");
}
// Step 4: Obtain arena based on arena mode
- mlir::Operation *firstAlloc = candidates.front().alloc;
- mlir::OpBuilder builder(firstAlloc);
- mlir::Value arenaValue;
-
+ Operation *firstAlloc = candidates.front().alloc;
+ OpBuilder builder(firstAlloc);
+ Value arenaValue;
+
if (arenaMode == "allocate") {
- // Step 5a: Create arena via AllocOp (default mode)
// Arena is i8 byte buffer to support multiple data types (f32, i64, etc.)
- auto i8Type = builder.getI8Type();
- auto arenaType = mlir::MemRefType::get({totalSize}, i8Type);
- auto arenaAlloc = mlir::memref::AllocOp::create(builder, firstAlloc->getLoc(), arenaType);
- arenaAlloc.setAlignmentAttr(builder.getI64IntegerAttr(maxAlignment));
+ auto arenaType = MemRefType::get({totalSize}, builder.getI8Type());
+ auto arenaAlloc = memref::AllocOp::create(
+ builder, firstAlloc->getLoc(), arenaType, ValueRange{},
+ builder.getI64IntegerAttr(arenaAlignment));
arenaValue = arenaAlloc.getResult();
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] created arena via AllocOp: size="
- << totalSize << " bytes, alignment="
- << maxAlignment << " bytes\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << "[static-memory-planner] created arena via AllocOp: size="
+ << totalSize << " bytes, alignment=" << arenaAlignment
+ << " bytes\n");
} else if (arenaMode == "arg") {
- // Step 5b: Extract arena from function arguments (arg mode)
- // Assumes first argument is an i8 buffer of sufficient size
- auto funcOp = llvm::dyn_cast<mlir::func::FuncOp>(op);
- if (!funcOp) {
- op->emitError("arena-mode=arg requires function context");
- return signalPassFailure();
- }
-
if (funcOp.getNumArguments() == 0) {
- funcOp.emitError("arena-mode=arg requires at least one function argument");
+ funcOp->emitError(
+ "arena-mode=arg requires at least one function argument");
return signalPassFailure();
}
-
+
arenaValue = funcOp.getArgument(0);
- auto arenaType = llvm::dyn_cast<mlir::MemRefType>(arenaValue.getType());
- if (!arenaType || !arenaType.getElementType().isInteger(8)) {
- funcOp.emitError("arena-mode=arg requires first argument to be memref<...xi8>");
+ auto arenaType = dyn_cast<MemRefType>(arenaValue.getType());
+ if (!arenaType || !arenaType.getElementType().isInteger(8) ||
+ arenaType.getRank() != 1) {
+ funcOp->emitError(
+ "arena-mode=arg requires first argument to be memref<...xi8>");
return signalPassFailure();
}
- LLVM_DEBUG(llvm::dbgs() << "[static-memory-planner] using arena from function arg 0\n");
+ LLVM_DEBUG(
+ llvm::dbgs()
+ << "[static-memory-planner] using arena from function arg 0\n");
} else {
- op->emitError("invalid arena-mode: '" + arenaMode + "' (must be 'allocate' or 'arg')");
+ funcOp->emitError("invalid arena-mode: '" + arenaMode +
+ "' (must be 'allocate' or 'arg')");
return signalPassFailure();
}
- // Step 6: Replace each alloc with memref.view directly on arena
+ // Step 5: Replace each alloc with memref.view directly on arena
for (auto &candidate : candidates) {
- mlir::OpBuilder viewBuilder(candidate.alloc);
- mlir::Location loc = candidate.alloc.getLoc();
-
- // Get the original memref type that we need to recreate
- mlir::MemRefType originalType = candidate.alloc.getType();
-
+ builder.setInsertionPoint(candidate.alloc);
+ Location loc = candidate.alloc.getLoc();
+
+ MemRefType originalType = candidate.alloc.getType();
+
// Create a constant for the byte offset into the arena
- mlir::Value offsetIndex = mlir::arith::ConstantIndexOp::create(
- viewBuilder, loc, candidate.offset);
-
- // Use memref.view to create a typed view directly on the i8 arena
- // memref.view: memref<...xi8>, offset -> memref<shape x type>
- llvm::SmallVector<mlir::Value> dynamicSizes; // Empty for static shapes
-
- auto view = mlir::memref::ViewOp::create(
- viewBuilder, loc, originalType, arenaValue,
- offsetIndex, dynamicSizes);
-
- // Replace all uses of the original alloc with the viewed memref
+ Value offsetIndex =
+ arith::ConstantIndexOp::create(builder, loc, candidate.offset);
+
+ // Use memref.view to create a typed view into the i8 arena
+ auto view = memref::ViewOp::create(builder, loc, originalType, arenaValue,
+ offsetIndex, SmallVector<Value>{});
+
candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
-
- // Remove the original alloc and dealloc
candidate.alloc.erase();
candidate.dealloc.erase();
}
>From 195ebad4a785b6cddc02a675364d0ebdb223d4b6 Mon Sep 17 00:00:00 2001
From: KrxGu <krishom70 at gmail.com>
Date: Mon, 29 Jun 2026 14:26:11 +0530
Subject: [PATCH 13/13] [mlir][bufferization] Restructure tests per review
feedback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- Add Test 2: non-sequential pairs (alloc alloc dealloc dealloc)
- Add Test 3: interleaved pairs (alloc alloc dealloc alloc dealloc dealloc)
- Renumber original Tests 2-7 to Tests 4-8
- Add Test 9: LCM alignment (alignment=4, alignment=16 → lcm=16)
- Fix Test 1 comment wording to 'sequential alloc and dealloc pairs'
---
.../static-memory-planner-analysis.mlir | 111 +++++++++++++-----
1 file changed, 82 insertions(+), 29 deletions(-)
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 3648558740b6e..a80c0e13adc21 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -3,7 +3,7 @@
// -----
-// Test 1: Simple sequential alloc/dealloc pairs
+// Test 1: Sequential alloc and dealloc pairs.
// CHECK-LABEL: func @simple_sequential
func.func @simple_sequential() {
// Arena is i8 buffer: 1024*4 + 512*4 = 6144 bytes
@@ -25,7 +25,56 @@ func.func @simple_sequential() {
// -----
-// Test 2: Dynamic shape - should be skipped (no transformation)
+// Test 2: Non-sequential pairs (alloc alloc dealloc dealloc).
+// CHECK-LABEL: func @non_sequential_pairs
+func.func @non_sequential_pairs() {
+ // Arena: 1024*4 + 512*4 = 6144 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<6144xi8>
+ // First allocation at offset 0
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<6144xi8> to memref<1024xf32>
+ // Second allocation at offset 4096 bytes (1024 * 4)
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C4096]]][] : memref<6144xi8> to memref<512xf32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %alloc0 = memref.alloc() : memref<1024xf32>
+ %alloc1 = memref.alloc() : memref<512xf32>
+ memref.dealloc %alloc0 : memref<1024xf32>
+ memref.dealloc %alloc1 : memref<512xf32>
+ return
+}
+
+// -----
+
+// Test 3: Interleaved pairs (alloc alloc dealloc alloc dealloc dealloc).
+// CHECK-LABEL: func @interleaved_pairs
+func.func @interleaved_pairs() {
+ // Arena: 512*4 + 256*4 + 128*4 = 3584 bytes
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<3584xi8>
+ // First at offset 0
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<3584xi8> to memref<512xf32>
+ // Second at offset 2048 bytes (512 * 4)
+ // CHECK-NEXT: %[[C2048:.*]] = arith.constant 2048 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C2048]]][] : memref<3584xi8> to memref<256xf32>
+ // Third at offset 3072 bytes (512*4 + 256*4)
+ // CHECK-NEXT: %[[C3072:.*]] = arith.constant 3072 : index
+ // CHECK-NEXT: %[[VIEW2:.*]] = memref.view %[[ARENA]][%[[C3072]]][] : memref<3584xi8> to memref<128xf32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %alloc0 = memref.alloc() : memref<512xf32>
+ %alloc1 = memref.alloc() : memref<256xf32>
+ memref.dealloc %alloc0 : memref<512xf32>
+ %alloc2 = memref.alloc() : memref<128xf32>
+ memref.dealloc %alloc1 : memref<256xf32>
+ memref.dealloc %alloc2 : memref<128xf32>
+ return
+}
+
+// -----
+
+// Test 4: Dynamic shape - should be skipped (no transformation)
// CHECK-LABEL: func @dynamic_shape_skipped
func.func @dynamic_shape_skipped(%n: index) {
// CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) : memref<?xf32>
@@ -36,7 +85,7 @@ func.func @dynamic_shape_skipped(%n: index) {
// -----
-// Test 3: No dealloc - should be skipped
+// Test 5: No dealloc - should be skipped
// CHECK-LABEL: func @no_dealloc_skipped
func.func @no_dealloc_skipped() {
// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1024xf32>
@@ -47,7 +96,7 @@ func.func @no_dealloc_skipped() {
// -----
-// Test 4: Dealloc in different block - should be skipped
+// 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>
@@ -64,29 +113,7 @@ func.func @different_block_skipped(%cond: i1) {
// -----
-// Test 5: Overlapping lifetimes (both eligible, sequential offsets)
-// CHECK-LABEL: func @overlapping_lifetimes
-func.func @overlapping_lifetimes() {
- // Arena: 512*4 + 1024*4 = 6144 bytes
- // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<6144xi8>
- // First allocation at offset 0
- // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
- // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<6144xi8> to memref<512xf32>
- // Second allocation at offset 2048 bytes (512 * 4)
- // CHECK-NEXT: %[[C2048:.*]] = arith.constant 2048 : index
- // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C2048]]][] : memref<6144xi8> to memref<1024xf32>
- // CHECK-NOT: memref.alloc
- // CHECK-NOT: memref.dealloc
- %alloc0 = memref.alloc() : memref<512xf32>
- %alloc1 = memref.alloc() : memref<1024xf32>
- memref.dealloc %alloc1 : memref<1024xf32>
- memref.dealloc %alloc0 : memref<512xf32>
- return
-}
-
-// -----
-
-// Test 6: Multiple allocations with sequential offsets
+// Test 7: Multiple allocations with sequential offsets
// CHECK-LABEL: func @multiple_sequential
func.func @multiple_sequential() {
// Arena: 1024*4 + 512*4 + 2048*4 = 14336 bytes
@@ -113,10 +140,10 @@ func.func @multiple_sequential() {
// -----
-// Test 7: Alignment requirements with padding
+// Test 8: Alignment requirements with padding
// CHECK-LABEL: func @alignment_padding
func.func @alignment_padding() {
- // Arena has max alignment (128 bytes), total: 256*4 + 128*4 + 64*4 = 1792 bytes
+ // Arena: 256*4 + 128*4 + 64*4 = 1792 bytes, alignment = lcm(128,64,128) = 128
// CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 128 : i64} : memref<1792xi8>
// First alloc: 256 f32, alignment=128, offset=0 bytes (128-aligned)
// CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
@@ -137,3 +164,29 @@ func.func @alignment_padding() {
memref.dealloc %alloc2 : memref<64xf32>
return
}
+
+// -----
+
+// Test 9: 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.
+// CHECK-LABEL: func @lcm_alignment
+func.func @lcm_alignment() {
+ // Arena: 3*4 + 3*4 = 24 bytes, but second alloc needs 16-byte offset
+ // (alignTo(12, 16) = 16), so total = 28 bytes, alignment = lcm(4,16) = 16
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 16 : i64} : memref<28xi8>
+ // First at offset 0 (alignment=4, 0 % 4 == 0)
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %[[VIEW0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<28xi8> to memref<3xi32>
+ // Second at offset 16 (alignment=16, 16 % 16 == 0)
+ // CHECK-NEXT: %[[C16:.*]] = arith.constant 16 : index
+ // CHECK-NEXT: %[[VIEW1:.*]] = memref.view %[[ARENA]][%[[C16]]][] : memref<28xi8> to memref<3xi32>
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.dealloc
+ %alloc0 = memref.alloc() {alignment = 4 : i64} : memref<3xi32>
+ memref.dealloc %alloc0 : memref<3xi32>
+ %alloc1 = memref.alloc() {alignment = 16 : i64} : memref<3xi32>
+ memref.dealloc %alloc1 : memref<3xi32>
+ return
+}
More information about the Mlir-commits
mailing list