[Mlir-commits] [mlir] [mlir][bufferization] Add best-fit algorithm to static memory planner (PR #207403)
Javed Absar
llvmlistbot at llvm.org
Fri Jul 3 09:40:47 PDT 2026
https://github.com/javedabsar1 updated https://github.com/llvm/llvm-project/pull/207403
>From 2511e24db6d65443df39400795bccd9f41de9535 Mon Sep 17 00:00:00 2001
From: mabsar <mabsar at qti.qualcommm.com>
Date: Fri, 3 Jul 2026 07:20:15 -0700
Subject: [PATCH] [mlir][bufferization] Add best-fit algorithm to static memory
planner
Introduces an algorithm selection option to the static memory planner pass
and add a best-fit algorithm that reuses memory from expired allocations
by finding the smallest suitable gap. The planning algorithms are factored
into a separate StaticMemoryPlanning.{h,cpp} to keep them independent of
MLIR IR and easily testable.
Signed-off-by: mabsar <mabsar at qti.qualcommm.com>
---
.../Bufferization/IR/BufferizationEnums.td | 8 +
.../Bufferization/Transforms/Passes.td | 31 +-
.../Transforms/StaticMemoryPlanning.h | 50 +++
.../Bufferization/Transforms/CMakeLists.txt | 1 +
.../StaticMemoryPlannerAnalysis.cpp | 345 +++++++++---------
.../Transforms/StaticMemoryPlanning.cpp | 109 ++++++
.../static-memory-planner-best-fit.mlir | 118 ++++++
7 files changed, 483 insertions(+), 179 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h
create mode 100644 mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlanning.cpp
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-best-fit.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td
index bafa84645e57b..a8cad26e9341b 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td
@@ -24,4 +24,12 @@ def LayoutMapOption : I32EnumAttr<"LayoutMapOption",
let cppNamespace = "::mlir::bufferization";
}
+def MemoryPlannerAlgorithm : I32EnumAttr<"MemoryPlannerAlgorithm",
+ "memory planning algorithm", [
+ I32EnumAttrCase<"Trivial", 0, "trivial">,
+ I32EnumAttrCase<"BestFit", 1, "best-fit">
+]> {
+ let cppNamespace = "::mlir::bufferization";
+}
+
#endif // BUFFERIZATION_ENUMS
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index c7f6cae571e40..8408315dda607 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -197,13 +197,11 @@ def StaticMemoryPlannerAnalysisPass
- 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.
+ Eligible allocations are packed into a single arena using a configurable
+ planning algorithm (see the `algorithm` option). 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.
@@ -234,11 +232,24 @@ def StaticMemoryPlannerAnalysisPass
```
}];
- let options = [Option<
- "arenaMode", "arena-mode", "std::string",
+ let options = [
+ Option<"arenaMode", "arena-mode", "std::string",
/*default=*/"\"allocate\"",
"Arena allocation mode: 'allocate' creates arena via AllocOp, "
- "'arg' extracts arena from function arguments">];
+ "'arg' extracts arena from function arguments">,
+ Option<"algorithm", "algorithm",
+ "::mlir::bufferization::MemoryPlannerAlgorithm",
+ /*default=*/"::mlir::bufferization::MemoryPlannerAlgorithm::Trivial",
+ "Memory planning algorithm to use.",
+ [{::llvm::cl::values(
+ clEnumValN(::mlir::bufferization::MemoryPlannerAlgorithm::Trivial,
+ "trivial",
+ "Sequential packing without lifetime overlap"),
+ clEnumValN(::mlir::bufferization::MemoryPlannerAlgorithm::BestFit,
+ "best-fit",
+ "Best-fit packing with lifetime-aware gap reuse")
+ )}]>,
+ ];
let statistics =
[Statistic<"numEligible", "num-eligible",
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h
new file mode 100644
index 0000000000000..626e384ce4642
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h
@@ -0,0 +1,50 @@
+//===- StaticMemoryPlanning.h - Memory planning algorithms ------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Pure memory planning algorithms for static arena allocation. These operate
+// on abstract allocation descriptors (size, alignment, lifetime) and produce
+// byte offsets. They are independent of MLIR IR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_STATICMEMORYPLANNING_H
+#define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_STATICMEMORYPLANNING_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include <cstdint>
+
+namespace mlir {
+namespace bufferization {
+
+/// Descriptor for a single allocation to be placed by the memory planner.
+struct MemoryPlannerAlloc {
+ int64_t sizeInBytes = 0;
+ int64_t alignment = 1;
+ int64_t timeStart = 0; // Operation index when allocation becomes live
+ int64_t timeEnd = 0; // Operation index when allocation is freed
+};
+
+/// Sequential packing without lifetime overlap. Each allocation is placed
+/// immediately after the previous one (with alignment padding). Ignores
+/// lifetimes entirely.
+llvm::SmallVector<int64_t>
+trivialMemoryPlanner(int64_t arenaAlignment,
+ llvm::ArrayRef<MemoryPlannerAlloc> allocs);
+
+/// Best-fit packing with lifetime-aware gap reuse. Processes allocations in
+/// time order and places each one in the smallest gap left by expired
+/// allocations. Falls back to extending the arena if no gap fits.
+llvm::SmallVector<int64_t>
+bestFitMemoryPlanner(int64_t arenaAlignment,
+ llvm::ArrayRef<MemoryPlannerAlloc> allocs);
+
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_STATICMEMORYPLANNING_H
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
index 5df9f19a5e30a..006fcd1ce0ec7 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
+ StaticMemoryPlanning.cpp
StaticMemoryPlannerAnalysis.cpp
ADDITIONAL_HEADER_DIRS
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 48b6e97c809ef..892df0787fb5d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -7,18 +7,18 @@
//===----------------------------------------------------------------------===//
//
// Transforms memref.alloc/memref.dealloc pairs into a single arena allocation
-// with memref.view. Uses simple sequential offset assignment where each
-// allocation gets its own space without overlap (baseline algorithm).
+// with memref.view. Delegates offset computation to planning algorithms in
+// StaticMemoryPlanning.h.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/Support/Debug.h"
-#include "llvm/Support/MathExtras.h"
#include <numeric>
#define DEBUG_TYPE "static-memory-planner"
@@ -34,19 +34,6 @@ using namespace mlir;
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
- 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 {
memref::AllocOp alloc;
@@ -81,36 +68,135 @@ static int64_t computeSizeInBytes(MemRefType memrefType) {
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) {
- return llvm::alignTo(offset, alignment);
+/// Build lifetime-annotated allocation descriptors from candidates.
+/// Returns the arena alignment (LCM of all individual alignments).
+static int64_t
+buildAllocInfos(MutableArrayRef<AllocationCandidate> candidates,
+ SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
+ int64_t arenaAlignment = 1;
+ for (auto &candidate : candidates) {
+ bufferization::MemoryPlannerAlloc info;
+ info.sizeInBytes = candidate.sizeInBytes;
+ info.alignment = candidate.alignment;
+
+ Block *block = candidate.alloc->getBlock();
+ int64_t opIdx = 0;
+ for (Operation &op : *block) {
+ if (&op == candidate.alloc.getOperation())
+ info.timeStart = opIdx;
+ if (&op == candidate.dealloc.getOperation())
+ info.timeEnd = opIdx;
+ ++opIdx;
+ }
+
+ allocInfos.push_back(info);
+ arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
+ }
+ return arenaAlignment;
}
-//===----------------------------------------------------------------------===//
-// Memory Planning Algorithms
-//===----------------------------------------------------------------------===//
+/// Collect alloc/dealloc pairs eligible for arena placement.
+/// An allocation is eligible if it has a static shape and a unique dealloc
+/// in the same block.
+static SmallVector<AllocationCandidate>
+collectCandidates(FunctionOpInterface funcOp,
+ llvm::Statistic &numSkipDynamic,
+ llvm::Statistic &numSkipNoDealloc,
+ llvm::Statistic &numEligible) {
+ SmallVector<AllocationCandidate> candidates;
+
+ funcOp->walk([&](memref::AllocOp allocOp) {
+ MemRefType memrefType = allocOp.getType();
+
+ // Skip dynamic shapes
+ if (!memrefType.hasStaticShape()) {
+ ++numSkipDynamic;
+ return;
+ }
-/// 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 SmallVector<int64_t> trivialMemoryPlanner(int64_t arenaAlignment,
- ArrayRef<Alloc> allocs) {
- SmallVector<int64_t> offsets;
- int64_t currentOffset = 0;
-
- for (const auto &alloc : allocs) {
- currentOffset = alignOffset(currentOffset, alloc.alignment);
-#ifndef NDEBUG
- assert((arenaAlignment + currentOffset) % alloc.alignment == 0 &&
- "invalid alignment");
-#endif
- offsets.push_back(currentOffset);
- currentOffset += alloc.sizeInBytes;
+ // Find unique dealloc in the same block
+ memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
+ if (!deallocOp) {
+ ++numSkipNoDealloc;
+ return;
+ }
+
+ if (deallocOp->getBlock() != allocOp->getBlock()) {
+ ++numSkipNoDealloc;
+ return;
+ }
+
+ // This allocation is eligible
+ ++numEligible;
+ AllocationCandidate candidate;
+ candidate.alloc = allocOp;
+ candidate.dealloc = deallocOp;
+ candidate.sizeInBytes = computeSizeInBytes(memrefType);
+ candidate.alignment = allocOp.getAlignment().value_or(1);
+ candidates.push_back(candidate);
+ });
+
+ return candidates;
+}
+
+/// Create or obtain the arena buffer based on the arena mode.
+/// Returns failure if the mode is invalid or preconditions aren't met.
+static FailureOr<Value> createArena(OpBuilder &builder,
+ FunctionOpInterface funcOp,
+ StringRef arenaMode, int64_t totalSize,
+ int64_t arenaAlignment) {
+ Location loc = funcOp->getLoc();
+
+ if (arenaMode == "allocate") {
+ auto arenaType = MemRefType::get({totalSize}, builder.getI8Type());
+ auto arenaAlloc = memref::AllocOp::create(
+ builder, loc, arenaType, ValueRange{},
+ builder.getI64IntegerAttr(arenaAlignment));
+ LLVM_DEBUG(llvm::dbgs()
+ << "[static-memory-planner] created arena via AllocOp: size="
+ << totalSize << " bytes, alignment=" << arenaAlignment
+ << " bytes\n");
+ return arenaAlloc.getResult();
+ }
+
+ if (arenaMode == "arg") {
+ if (funcOp.getNumArguments() == 0)
+ return funcOp->emitError(
+ "arena-mode=arg requires at least one function argument");
+
+ Value arenaValue = funcOp.getArgument(0);
+ auto arenaType = dyn_cast<MemRefType>(arenaValue.getType());
+ if (!arenaType || !arenaType.getElementType().isInteger(8) ||
+ arenaType.getRank() != 1)
+ return funcOp->emitError(
+ "arena-mode=arg requires first argument to be memref<...xi8>");
+
+ LLVM_DEBUG(llvm::dbgs()
+ << "[static-memory-planner] using arena from function arg 0\n");
+ return arenaValue;
}
- return offsets;
+ return funcOp->emitError("invalid arena-mode: '" + arenaMode +
+ "' (must be 'allocate' or 'arg')");
+}
+
+/// Replace each alloc/dealloc pair with a memref.view into the arena.
+static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
+ Value arenaValue) {
+ for (auto &candidate : candidates) {
+ OpBuilder builder(candidate.alloc);
+ Location loc = candidate.alloc.getLoc();
+ MemRefType originalType = candidate.alloc.getType();
+
+ Value offsetIndex =
+ arith::ConstantIndexOp::create(builder, loc, candidate.offset);
+ auto view = memref::ViewOp::create(builder, loc, originalType, arenaValue,
+ offsetIndex, SmallVector<Value>{});
+
+ candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
+ candidate.alloc.erase();
+ candidate.dealloc.erase();
+ }
}
//===----------------------------------------------------------------------===//
@@ -125,143 +211,64 @@ struct StaticMemoryPlannerAnalysisPass
StaticMemoryPlannerAnalysisPass>;
using Base::Base;
- void runOnOperation() override {
- auto funcOp = llvm::cast<FunctionOpInterface>(getOperation());
-
- // Step 0: Check for memref return types (not supported)
- 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
- SmallVector<AllocationCandidate> candidates;
-
- funcOp->walk([&](memref::AllocOp allocOp) {
- MemRefType memrefType = allocOp.getType();
-
- // Skip dynamic shapes
- if (!memrefType.hasStaticShape()) {
- ++numSkipDynamic;
- return;
- }
-
- // Find unique dealloc in the same block
- memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
- if (!deallocOp) {
- ++numSkipNoDealloc;
- return;
- }
-
- if (deallocOp->getBlock() != allocOp->getBlock()) {
- ++numSkipNoDealloc;
- return;
- }
-
- // This allocation is eligible
- ++numEligible;
- AllocationCandidate candidate;
- candidate.alloc = allocOp;
- candidate.dealloc = deallocOp;
- candidate.sizeInBytes = computeSizeInBytes(memrefType);
- candidate.alignment = allocOp.getAlignment().value_or(1);
- candidates.push_back(candidate);
- });
-
- if (candidates.empty())
- return;
-
- // Step 2: Prepare allocation info for planner
- 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);
- arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
- }
+ void runOnOperation() override;
+};
- // Step 3: Run the planning algorithm
- 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");
- }
+void StaticMemoryPlannerAnalysisPass::runOnOperation() {
+ auto funcOp = llvm::cast<FunctionOpInterface>(getOperation());
- // Step 4: Obtain arena based on arena mode
- Operation *firstAlloc = candidates.front().alloc;
- OpBuilder builder(firstAlloc);
- Value arenaValue;
-
- if (arenaMode == "allocate") {
- // Arena is i8 byte buffer to support multiple data types (f32, i64, etc.)
- 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=" << arenaAlignment
- << " bytes\n");
- } else if (arenaMode == "arg") {
- if (funcOp.getNumArguments() == 0) {
- funcOp->emitError(
- "arena-mode=arg requires at least one function argument");
- return signalPassFailure();
- }
-
- arenaValue = funcOp.getArgument(0);
- 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");
- } else {
- funcOp->emitError("invalid arena-mode: '" + arenaMode +
- "' (must be 'allocate' or 'arg')");
+ // Step 0: Check for memref return types (not supported)
+ 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 5: Replace each alloc with memref.view directly on arena
- for (auto &candidate : candidates) {
- builder.setInsertionPoint(candidate.alloc);
- Location loc = candidate.alloc.getLoc();
+ // Step 1: Collect eligible allocation candidates.
+ SmallVector<AllocationCandidate> candidates =
+ collectCandidates(funcOp, numSkipDynamic, numSkipNoDealloc, numEligible);
- MemRefType originalType = candidate.alloc.getType();
+ if (candidates.empty())
+ return;
- // Create a constant for the byte offset into the arena
- Value offsetIndex =
- arith::ConstantIndexOp::create(builder, loc, candidate.offset);
+ // Step 2: Build allocation descriptors with lifetime info.
+ SmallVector<bufferization::MemoryPlannerAlloc> allocInfos;
+ int64_t arenaAlignment = buildAllocInfos(candidates, allocInfos);
- // Use memref.view to create a typed view into the i8 arena
- auto view = memref::ViewOp::create(builder, loc, originalType, arenaValue,
- offsetIndex, SmallVector<Value>{});
+ // Step 3: Run the planning algorithm.
+ SmallVector<int64_t> offsets;
+ switch (algorithm) {
+ case bufferization::MemoryPlannerAlgorithm::Trivial:
+ offsets = bufferization::trivialMemoryPlanner(arenaAlignment, allocInfos);
+ break;
+ case bufferization::MemoryPlannerAlgorithm::BestFit:
+ offsets = bufferization::bestFitMemoryPlanner(arenaAlignment, allocInfos);
+ break;
+ }
- candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
- candidate.alloc.erase();
- candidate.dealloc.erase();
- }
+ // Step 4: Compute total arena size and assign offsets.
+ 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");
}
-};
+
+ // Step 5: Obtain arena based on arena mode.
+ Operation *firstAlloc = candidates.front().alloc;
+ OpBuilder builder(firstAlloc);
+ FailureOr<Value> arenaValue =
+ createArena(builder, funcOp, arenaMode, totalSize, arenaAlignment);
+ if (failed(arenaValue))
+ return signalPassFailure();
+
+ // Step 6: Replace each alloc with memref.view into the arena.
+ rewriteAllocations(candidates, *arenaValue);
+}
} // end anonymous namespace
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlanning.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlanning.cpp
new file mode 100644
index 0000000000000..eadb1307d1b21
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlanning.cpp
@@ -0,0 +1,109 @@
+//===- StaticMemoryPlanning.cpp - Memory planning algorithms --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
+#include "llvm/Support/MathExtras.h"
+#include <numeric>
+
+using namespace mlir::bufferization;
+
+/// Align an offset to the specified alignment.
+static int64_t alignOffset(int64_t offset, int64_t alignment) {
+ return llvm::alignTo(offset, alignment);
+}
+
+llvm::SmallVector<int64_t> mlir::bufferization::trivialMemoryPlanner(
+ int64_t arenaAlignment, llvm::ArrayRef<MemoryPlannerAlloc> allocs) {
+ llvm::SmallVector<int64_t> offsets;
+ int64_t currentOffset = 0;
+
+ for (const auto &alloc : allocs) {
+ currentOffset = alignOffset(currentOffset, alloc.alignment);
+ assert((arenaAlignment + currentOffset) % alloc.alignment == 0 &&
+ "invalid alignment");
+ offsets.push_back(currentOffset);
+ currentOffset += alloc.sizeInBytes;
+ }
+
+ return offsets;
+}
+
+llvm::SmallVector<int64_t> mlir::bufferization::bestFitMemoryPlanner(
+ int64_t arenaAlignment, llvm::ArrayRef<MemoryPlannerAlloc> allocs) {
+ struct Placement {
+ int64_t offset;
+ int64_t size;
+ int64_t timeEnd;
+ };
+
+ // Process allocations in order of start time.
+ llvm::SmallVector<unsigned> order(allocs.size());
+ std::iota(order.begin(), order.end(), 0);
+ llvm::sort(order, [&](unsigned a, unsigned b) {
+ return allocs[a].timeStart < allocs[b].timeStart;
+ });
+
+ llvm::SmallVector<Placement> placements;
+ llvm::SmallVector<int64_t> offsets(allocs.size(), 0);
+
+ for (unsigned idx : order) {
+ const MemoryPlannerAlloc &alloc = allocs[idx];
+
+ // Collect intervals that are still live at this allocation's start time.
+ llvm::SmallVector<std::pair<int64_t, int64_t>> occupied;
+ for (const auto &p : placements) {
+ if (p.timeEnd > alloc.timeStart)
+ occupied.push_back({p.offset, p.offset + p.size});
+ }
+ llvm::sort(occupied);
+
+ // Find the best (smallest) gap that fits this allocation.
+ int64_t bestOffset = -1;
+ int64_t bestGapSize = INT64_MAX;
+
+ // Arena high-water mark from all placements so far.
+ int64_t arenaEnd = 0;
+ for (const auto &p : placements)
+ arenaEnd = std::max(arenaEnd, p.offset + p.size);
+
+ int64_t gapStart = 0;
+ for (const auto &[occStart, occEnd] : occupied) {
+ int64_t alignedStart = alignOffset(gapStart, alloc.alignment);
+ int64_t gapEnd = occStart;
+ if (alignedStart + alloc.sizeInBytes <= gapEnd) {
+ int64_t gapSize = gapEnd - alignedStart;
+ if (gapSize < bestGapSize) {
+ bestGapSize = gapSize;
+ bestOffset = alignedStart;
+ }
+ }
+ gapStart = std::max(gapStart, occEnd);
+ }
+
+ // Check the trailing gap (between last occupied and arena end).
+ int64_t alignedTrailing = alignOffset(gapStart, alloc.alignment);
+ if (alignedTrailing + alloc.sizeInBytes <= arenaEnd) {
+ int64_t trailingSize = arenaEnd - alignedTrailing;
+ if (trailingSize < bestGapSize) {
+ bestGapSize = trailingSize;
+ bestOffset = alignedTrailing;
+ }
+ }
+
+ // If no existing gap worked, append at the end.
+ if (bestOffset < 0)
+ bestOffset = alignedTrailing;
+
+ assert((arenaAlignment + bestOffset) % alloc.alignment == 0 &&
+ "invalid alignment");
+ offsets[idx] = bestOffset;
+ placements.push_back({bestOffset, alloc.sizeInBytes, alloc.timeEnd});
+ }
+
+ return offsets;
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-best-fit.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-best-fit.mlir
new file mode 100644
index 0000000000000..6c261880140a5
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-best-fit.mlir
@@ -0,0 +1,118 @@
+// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(static-memory-planner-analysis{algorithm=best-fit}))" \
+// RUN: -split-input-file | FileCheck %s
+
+// -----
+
+// Test 1: Non-overlapping lifetimes reuse the same memory.
+// With trivial packing this would be 8192 bytes; best-fit reuses the space.
+// CHECK-LABEL: func @reuse_non_overlapping
+func.func @reuse_non_overlapping() {
+ // Arena should be 4096 bytes (1024 * 4), not 8192.
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<4096xi8>
+ // First allocation at offset 0
+ // CHECK-NEXT: %[[C0_0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0_0]]][] : memref<4096xi8> to memref<1024xf32>
+ // Second allocation also at offset 0 (reuses freed space)
+ // CHECK-NEXT: %[[C0_1:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0_1]]][] : memref<4096xi8> to memref<1024xf32>
+ %0 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %0 : memref<1024xf32>
+ %1 = memref.alloc() : memref<1024xf32>
+ memref.dealloc %1 : memref<1024xf32>
+ return
+}
+
+// -----
+
+// Test 2: Overlapping lifetimes cannot reuse memory.
+// CHECK-LABEL: func @no_reuse_overlapping
+func.func @no_reuse_overlapping() {
+ // Both are live at the same time, so arena = 4096 + 2048 = 6144 bytes.
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<6144xi8>
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0]]][] : memref<6144xi8> to memref<1024xf32>
+ // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C4096]]][] : memref<6144xi8> to memref<512xf32>
+ %0 = memref.alloc() : memref<1024xf32>
+ %1 = memref.alloc() : memref<512xf32>
+ memref.dealloc %0 : memref<1024xf32>
+ memref.dealloc %1 : memref<512xf32>
+ return
+}
+
+// -----
+
+// Test 3: Best-fit picks the smallest suitable gap.
+// Layout: A(4096) at 0, B(1024) at 4096, C(4096) at 5120, D(1024) at 9216.
+// B and D are freed while A and C are still live, creating two gaps:
+// [4096, 5120) = 1024 bytes (B's slot)
+// [9216, 10240) = 1024 bytes (D's slot)
+// Then we free A, creating gap [0, 4096) = 4096 bytes.
+// Now allocate E(512 bytes). Gaps: [0,4096)=4096, [4096,5120)=1024, [9216,10240)=1024.
+// Best-fit should pick one of the 1024-byte gaps (smallest fit for 512).
+// CHECK-LABEL: func @best_fit_smallest_gap
+func.func @best_fit_smallest_gap() {
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<10240xi8>
+ // A at offset 0
+ // CHECK-NEXT: %{{.*}} = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]
+ // B at offset 4096
+ // CHECK-NEXT: %{{.*}} = arith.constant 4096 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]
+ // C at offset 5120
+ // CHECK-NEXT: %{{.*}} = arith.constant 5120 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]
+ // D at offset 9216
+ // CHECK-NEXT: %{{.*}} = arith.constant 9216 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]
+ // E at offset 9216 (best-fit picks 1024-byte trailing gap over 5120-byte gap)
+ // CHECK-NEXT: %{{.*}} = arith.constant 9216 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]
+ %a = memref.alloc() : memref<1024xf32>
+ %b = memref.alloc() : memref<256xf32>
+ %c = memref.alloc() : memref<1024xf32>
+ %d = memref.alloc() : memref<256xf32>
+ memref.dealloc %b : memref<256xf32>
+ memref.dealloc %d : memref<256xf32>
+ memref.dealloc %a : memref<1024xf32>
+ %e = memref.alloc() : memref<128xf32>
+ memref.dealloc %c : memref<1024xf32>
+ memref.dealloc %e : memref<128xf32>
+ return
+}
+
+// -----
+
+// Test 4: Alignment padding can disqualify a gap.
+// A(128, align 128) at 0 (live), B(56, align 1) at 128, C(128, align 128) at 256 (live).
+// B is freed => gap [128, 256) = 128 bytes.
+// D(64, align 128): aligned start = 128, fits in gap.
+// E(64, align 256): next 256-aligned offset in [128,256) is 256 = gap end, doesn't fit.
+// E must go past the arena high-water mark.
+// CHECK-LABEL: func @best_fit_alignment
+func.func @best_fit_alignment() {
+ // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 256 : i64} : memref<576xi8>
+ // CHECK-NEXT: %{{.*}} = arith.constant 0 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]{{.*}} to memref<128xi8>
+ // CHECK-NEXT: %{{.*}} = arith.constant 128 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]{{.*}} to memref<56xi8>
+ // CHECK-NEXT: %{{.*}} = arith.constant 256 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]{{.*}} to memref<128xi8>
+ // D fits in gap at offset 128
+ // CHECK-NEXT: %{{.*}} = arith.constant 128 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]{{.*}} to memref<64xi8>
+ // E cannot fit in gap (alignment 256), placed at offset 512
+ // CHECK-NEXT: %{{.*}} = arith.constant 512 : index
+ // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]]{{.*}} to memref<64xi8>
+ %a = memref.alloc() {alignment = 128 : i64} : memref<128xi8>
+ %b = memref.alloc() : memref<56xi8>
+ %c = memref.alloc() {alignment = 128 : i64} : memref<128xi8>
+ memref.dealloc %b : memref<56xi8>
+ %d = memref.alloc() {alignment = 128 : i64} : memref<64xi8>
+ %e = memref.alloc() {alignment = 256 : i64} : memref<64xi8>
+ memref.dealloc %a : memref<128xi8>
+ memref.dealloc %c : memref<128xi8>
+ memref.dealloc %d : memref<64xi8>
+ memref.dealloc %e : memref<64xi8>
+ return
+}
More information about the Mlir-commits
mailing list