[Mlir-commits] [mlir] [mlir][bufferization] Add static memory planner pass for compile-time buffer allocation (PR #205125)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 22 09:05:30 PDT 2026
github-actions[bot] wrote:
<!--LLVM CODE FORMAT COMMENT: {clang-format}-->
:warning: C/C++ code formatter, clang-format found issues in your code. :warning:
<details>
<summary>
You can test this locally with the following command:
</summary>
``````````bash
git-clang-format --diff origin/main HEAD --extensions cpp -- mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp --diff_from_common_commit
``````````
:warning:
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing `origin/main` to the base branch/commit you want to compare against.
:warning:
</details>
<details>
<summary>
View the diff from clang-format here.
</summary>
``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 3d007a79a..f43d65536 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -12,8 +12,8 @@
//
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#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"
@@ -39,18 +39,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
+ // 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 (assigned by planner)
- int64_t sizeInBytes = 0; // Size in bytes
- int64_t alignment = 1; // Required alignment in bytes
+ 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
};
//===----------------------------------------------------------------------===//
@@ -102,21 +103,21 @@ static int64_t alignOffset(int64_t offset, int64_t alignment) {
/// 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) {
+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
+ // 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;
}
@@ -160,7 +161,8 @@ public:
}
// Find unique dealloc in the same block
- mlir::memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
+ mlir::memref::DeallocOp deallocOp =
+ findUniqueDealloc(allocOp.getResult());
if (!deallocOp) {
++numSkipNoDealloc;
return;
@@ -197,37 +199,39 @@ public:
}
// 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) {
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;
-
+
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);
+ 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 via AllocOp: 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
@@ -236,22 +240,27 @@ public:
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>");
+ 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')");
+ op->emitError("invalid arena-mode: '" + arenaMode +
+ "' (must be 'allocate' or 'arg')");
return signalPassFailure();
}
@@ -259,25 +268,25 @@ public:
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();
-
+
// 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);
+
+ auto view =
+ mlir::memref::ViewOp::create(viewBuilder, loc, originalType,
+ arenaValue, offsetIndex, dynamicSizes);
// 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();
candidate.dealloc.erase();
``````````
</details>
https://github.com/llvm/llvm-project/pull/205125
More information about the Mlir-commits
mailing list