lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp - llvm-project/mlir - Git at Google

 //===- 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.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Transforms memref.alloc/memref.dealloc pairs into a single arena allocation
 // with memref.view. Delegates offset computation to planning algorithms in
 // StaticMemoryPlanning.h.
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include <numeric>

 #define DEBUG_TYPE "static-memory-planner"

 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 {

 /// A candidate allocation with its matching deallocation(s) and assigned
 /// offset. An alloc may be freed indirectly through arith.select chains,
 /// yielding multiple potential deallocs — all must be in the same block.
 struct AllocationCandidate {
   memref::AllocOp alloc;
   SmallVector<memref::DeallocOp> deallocs;
   int64_t offset = 0; // Offset in bytes from arena start (assigned by planner)
   int64_t sizeInBytes = 0; // Size in bytes
   int64_t alignment = 1;   // Required alignment in bytes
 };

 //===----------------------------------------------------------------------===//
 // Helper utilities
 //===----------------------------------------------------------------------===//

 /// Collect all dealloc ops that might free the given alloc value. Instead of a
 /// bespoke traversal, this uses the shared `BufferViewFlowAnalysis`, which
 /// already models all the ways a buffer can flow to a dealloc:
 ///   - `arith.select`     (via BufferViewFlowOpInterface)
 ///   - `scf.if`/`scf.for` (via RegionBranchOpInterface region/result wiring)
 ///   - `cf.br`/`cf.cond_br` (via BranchOpInterface block arguments)
 ///   - `memref.view`/subview (via ViewLikeOpInterface)
 /// `analysis.resolve(alloc)` returns the forward alias set (the alloc plus
 /// every value it may flow into); a dealloc on any of those aliases frees the
 /// alloc. For example:
 ///   %0 = memref.alloc()
 ///   %2 = arith.select %c, %0, %1
 ///   memref.dealloc %2      <- covers %0 conditionally (via alias set)
 ///   %3 = scf.if %c { yield %0 } else { yield %1 }
 ///   memref.dealloc %3      <- also covers %0 conditionally
 static void collectDeallocs(Value alloc, const BufferViewFlowAnalysis &analysis,
                             SmallVectorImpl<memref::DeallocOp> &deallocs) {
   for (Value alias : analysis.resolve(alloc))
     for (Operation *user : alias.getUsers())
       if (auto dealloc = dyn_cast<memref::DeallocOp>(user))
         deallocs.push_back(dealloc);
 }

 /// Return the set of allocation ops whose buffer may be freed by `dealloc`,
 /// i.e. the terminal `memref.alloc` sources that flow into the dealloc operand.
 /// Uses the reverse alias set so that a dealloc reached through a `scf.if`
 /// result or `arith.select` is attributed to every alloc it may free.
 static SmallVector<memref::AllocOp>
 findFreedAllocs(memref::DeallocOp dealloc,
                 const BufferViewFlowAnalysis &analysis) {
   SmallVector<memref::AllocOp> allocs;
   for (Value source : analysis.resolveReverse(dealloc.getMemref()))
     if (auto allocOp = source.getDefiningOp<memref::AllocOp>())
       allocs.push_back(allocOp);
   return allocs;
 }

 /// Compute the size in bytes for a memref type.
 static int64_t computeSizeInBytes(MemRefType memrefType) {
   int64_t numElements = memrefType.getNumElements();
   unsigned elementSizeInBits = memrefType.getElementTypeBitWidth();
   return (numElements * elementSizeInBits + 7) / 8; // Round up to bytes
 }

 /// Build lifetime-annotated allocation descriptors from candidates.
 /// Returns the arena alignment (LCM of all individual alignments).
 /// Uses a single block scan (O(n+m)) instead of one scan per candidate.
 static int64_t buildAllocInfos(
     MutableArrayRef<AllocationCandidate> candidates,
     SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
   // Build an op-index map with a single pass over the plan block.
   DenseMap<Operation *, int64_t> opIndex;
   Block *planBlock = nullptr;
   if (!candidates.empty()) {
     planBlock = candidates.front().alloc->getBlock();
     int64_t idx = 0;
     for (Operation &op : *planBlock)
       opIndex[&op] = idx++;
   }

   int64_t arenaAlignment = 1;
   for (auto &candidate : candidates) {
     bufferization::MemoryPlannerAlloc info;
     info.sizeInBytes = candidate.sizeInBytes;
     info.alignment = candidate.alignment;
     info.timeStart = opIndex.lookup(candidate.alloc.getOperation());
     // Conservative: timeEnd = latest dealloc position among all potential
     // deallocs. A dealloc may be nested (e.g. inside an scf.if body); its
     // lifetime contribution is bounded by the enclosing op in the plan block.
     int64_t timeEnd = info.timeStart;
     for (memref::DeallocOp d : candidate.deallocs) {
       Operation *anchor = planBlock->findAncestorOpInBlock(*d.getOperation());
       timeEnd = std::max(timeEnd, opIndex.lookup(anchor));
     }
     info.timeEnd = timeEnd;
     allocInfos.push_back(info);
     arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
   }
   return arenaAlignment;
 }

 /// Collect alloc/dealloc groups eligible for arena placement.
 ///
 /// Eligibility uses the shared `BufferViewFlowAnalysis` so that buffers flowing
 /// through `arith.select`, `scf.if`/`scf.for` results, `cf` branches, or view
 /// ops are handled uniformly. An allocation is eligible when:
 ///   - it has a static shape (dynamic shapes are silently skipped), and
 ///   - it lives directly in the function's entry block (allocs nested in a
 ///     region are skipped for now), and
 ///   - every dealloc that may free it is anchored in that same entry block --
 ///     either directly, or nested inside an op of that block (e.g. an
 ///     `scf.if` body), which conservatively bounds the lifetime.
 /// A dealloc that escapes the entry block entirely (e.g. lives in a sibling
 /// `cf` block) is reported as an error, as is an alloc with no dealloc.
 ///
 /// To keep the rewrite safe, a dealloc is only accepted if *all* allocs it may
 /// free (per the reverse alias set) are themselves candidates in this block;
 /// otherwise erasing it during the rewrite could leak or double-free a buffer
 /// that is not managed by the arena.
 static LogicalResult
 collectCandidates(FunctionOpInterface funcOp,
                   const BufferViewFlowAnalysis &analysis,
                   llvm::Statistic &numSkipDynamic,
                   llvm::Statistic &numSkipNested, llvm::Statistic &numEligible,
                   SmallVector<AllocationCandidate> &candidates) {
   // All candidates are planned relative to the function's entry block.
   if (funcOp.getFunctionBody().empty())
     return success();
   Block *planBlock = &funcOp.getFunctionBody().front();

   bool walkFailed = false;
   funcOp->walk([&](memref::AllocOp allocOp) -> WalkResult {
     MemRefType memrefType = allocOp.getType();
     if (!memrefType.hasStaticShape()) {
       ++numSkipDynamic;
       return WalkResult::advance();
     }

     // Only plan allocs that live directly in the entry block. Allocs nested in
     // a region (loop/conditional body) are skipped for now.
     if (allocOp->getBlock() != planBlock) {
       ++numSkipNested;
       return WalkResult::advance();
     }

     SmallVector<memref::DeallocOp> deallocs;
     collectDeallocs(allocOp.getResult(), analysis, deallocs);

     if (deallocs.empty()) {
       allocOp.emitError("no dealloc found; run the deallocation pipeline "
                         "before this pass");
       walkFailed = true;
       return WalkResult::interrupt();
     }

     for (memref::DeallocOp d : deallocs) {
       // The dealloc must be anchored in the plan block (directly or via an
       // enclosing op such as an scf.if). A dealloc in a sibling block escapes.
       if (!planBlock->findAncestorOpInBlock(*d.getOperation())) {
         allocOp.emitError("unstructured control flow is not supported");
         walkFailed = true;
         return WalkResult::interrupt();
       }
       // Every alloc that this dealloc may free must also be an entry-block
       // candidate; otherwise erasing it during the rewrite is unsafe.
       for (memref::AllocOp freed : findFreedAllocs(d, analysis)) {
         if (freed->getBlock() != planBlock) {
           ++numSkipNested;
           return WalkResult::advance();
         }
       }
     }

     ++numEligible;
     AllocationCandidate candidate;
     candidate.alloc = allocOp;
     candidate.deallocs = deallocs;
     candidate.sizeInBytes = computeSizeInBytes(memrefType);
     candidate.alignment = allocOp.getAlignment().value_or(1);
     candidates.push_back(candidate);
     return WalkResult::advance();
   });

   return failure(walkFailed);
 }

 /// 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 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) {
   SmallPtrSet<Operation *, 8> deallocsToErase;
   SmallVector<Operation *> allocsToErase;

   // Replace all alloc results with views (rewires selects too).
   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());
     allocsToErase.push_back(candidate.alloc.getOperation());

     for (memref::DeallocOp d : candidate.deallocs)
       deallocsToErase.insert(d.getOperation());
   }

   // Erase deallocs first (they may reference alloc results via selects).
   for (Operation *d : deallocsToErase)
     d->erase();

   // Erase allocs last (no users remain after replaceAllUsesWith).
   for (Operation *allocOp : allocsToErase)
     allocOp->erase();
 }

 //===----------------------------------------------------------------------===//
 // StaticMemoryPlannerAnalysisPass
 //===----------------------------------------------------------------------===//

 struct StaticMemoryPlannerAnalysisPass
     : public bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
           StaticMemoryPlannerAnalysisPass> {
 public:
   using Base = bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
       StaticMemoryPlannerAnalysisPass>;
   using Base::Base;

   void runOnOperation() override;
 };

 void StaticMemoryPlannerAnalysisPass::runOnOperation() {
   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. The buffer view-flow
   // analysis models how buffers flow through selects, scf.if results, branches,
   // and view ops so we can find deallocs and freed allocs uniformly.
   BufferViewFlowAnalysis analysis(funcOp);
   SmallVector<AllocationCandidate> candidates;
   if (failed(collectCandidates(funcOp, analysis, numSkipDynamic, numSkipNested,
                                numEligible, candidates)))
     return signalPassFailure();

   if (candidates.empty())
     return;

   // Step 2: Build allocation descriptors with lifetime info.
   SmallVector<bufferization::MemoryPlannerAlloc> allocInfos;
   int64_t arenaAlignment = buildAllocInfos(candidates, allocInfos);

   // 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;
   }

   // 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
	//===- 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.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	//
	// Transforms memref.alloc/memref.dealloc pairs into a single arena allocation
	// with memref.view. Delegates offset computation to planning algorithms in
	// StaticMemoryPlanning.h.
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/Dialect/Arith/IR/Arith.h"
	#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
	#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
	#include "mlir/Dialect/Bufferization/Transforms/StaticMemoryPlanning.h"
	#include "mlir/Dialect/MemRef/IR/MemRef.h"
	#include "mlir/IR/Builders.h"
	#include "mlir/Interfaces/FunctionInterfaces.h"
	#include "llvm/ADT/SmallPtrSet.h"
	#include "llvm/Support/Debug.h"
	#include <numeric>

	#define DEBUG_TYPE "static-memory-planner"

	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 {

	/// A candidate allocation with its matching deallocation(s) and assigned
	/// offset. An alloc may be freed indirectly through arith.select chains,
	/// yielding multiple potential deallocs — all must be in the same block.
	struct AllocationCandidate {
	memref::AllocOp alloc;
	SmallVector<memref::DeallocOp> deallocs;
	int64_t offset = 0; // Offset in bytes from arena start (assigned by planner)
	int64_t sizeInBytes = 0; // Size in bytes
	int64_t alignment = 1; // Required alignment in bytes
	};

	//===----------------------------------------------------------------------===//
	// Helper utilities
	//===----------------------------------------------------------------------===//

	/// Collect all dealloc ops that might free the given alloc value. Instead of a
	/// bespoke traversal, this uses the shared `BufferViewFlowAnalysis`, which
	/// already models all the ways a buffer can flow to a dealloc:
	/// - `arith.select` (via BufferViewFlowOpInterface)
	/// - `scf.if`/`scf.for` (via RegionBranchOpInterface region/result wiring)
	/// - `cf.br`/`cf.cond_br` (via BranchOpInterface block arguments)
	/// - `memref.view`/subview (via ViewLikeOpInterface)
	/// `analysis.resolve(alloc)` returns the forward alias set (the alloc plus
	/// every value it may flow into); a dealloc on any of those aliases frees the
	/// alloc. For example:
	/// %0 = memref.alloc()
	/// %2 = arith.select %c, %0, %1
	/// memref.dealloc %2 <- covers %0 conditionally (via alias set)
	/// %3 = scf.if %c { yield %0 } else { yield %1 }
	/// memref.dealloc %3 <- also covers %0 conditionally
	static void collectDeallocs(Value alloc, const BufferViewFlowAnalysis &analysis,
	SmallVectorImpl<memref::DeallocOp> &deallocs) {
	for (Value alias : analysis.resolve(alloc))
	for (Operation *user : alias.getUsers())
	if (auto dealloc = dyn_cast<memref::DeallocOp>(user))
	deallocs.push_back(dealloc);
	}

	/// Return the set of allocation ops whose buffer may be freed by `dealloc`,
	/// i.e. the terminal `memref.alloc` sources that flow into the dealloc operand.
	/// Uses the reverse alias set so that a dealloc reached through a `scf.if`
	/// result or `arith.select` is attributed to every alloc it may free.
	static SmallVector<memref::AllocOp>
	findFreedAllocs(memref::DeallocOp dealloc,
	const BufferViewFlowAnalysis &analysis) {
	SmallVector<memref::AllocOp> allocs;
	for (Value source : analysis.resolveReverse(dealloc.getMemref()))
	if (auto allocOp = source.getDefiningOp<memref::AllocOp>())
	allocs.push_back(allocOp);
	return allocs;
	}

	/// Compute the size in bytes for a memref type.
	static int64_t computeSizeInBytes(MemRefType memrefType) {
	int64_t numElements = memrefType.getNumElements();
	unsigned elementSizeInBits = memrefType.getElementTypeBitWidth();
	return (numElements * elementSizeInBits + 7) / 8; // Round up to bytes
	}

	/// Build lifetime-annotated allocation descriptors from candidates.
	/// Returns the arena alignment (LCM of all individual alignments).
	/// Uses a single block scan (O(n+m)) instead of one scan per candidate.
	static int64_t buildAllocInfos(
	MutableArrayRef<AllocationCandidate> candidates,
	SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
	// Build an op-index map with a single pass over the plan block.
	DenseMap<Operation *, int64_t> opIndex;
	Block *planBlock = nullptr;
	if (!candidates.empty()) {
	planBlock = candidates.front().alloc->getBlock();
	int64_t idx = 0;
	for (Operation &op : *planBlock)
	opIndex[&op] = idx++;
	}

	int64_t arenaAlignment = 1;
	for (auto &candidate : candidates) {
	bufferization::MemoryPlannerAlloc info;
	info.sizeInBytes = candidate.sizeInBytes;
	info.alignment = candidate.alignment;
	info.timeStart = opIndex.lookup(candidate.alloc.getOperation());
	// Conservative: timeEnd = latest dealloc position among all potential
	// deallocs. A dealloc may be nested (e.g. inside an scf.if body); its
	// lifetime contribution is bounded by the enclosing op in the plan block.
	int64_t timeEnd = info.timeStart;
	for (memref::DeallocOp d : candidate.deallocs) {
	Operation anchor = planBlock->findAncestorOpInBlock(d.getOperation());
	timeEnd = std::max(timeEnd, opIndex.lookup(anchor));
	}
	info.timeEnd = timeEnd;
	allocInfos.push_back(info);
	arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
	}
	return arenaAlignment;
	}

	/// Collect alloc/dealloc groups eligible for arena placement.
	///
	/// Eligibility uses the shared `BufferViewFlowAnalysis` so that buffers flowing
	/// through `arith.select`, `scf.if`/`scf.for` results, `cf` branches, or view
	/// ops are handled uniformly. An allocation is eligible when:
	/// - it has a static shape (dynamic shapes are silently skipped), and
	/// - it lives directly in the function's entry block (allocs nested in a
	/// region are skipped for now), and
	/// - every dealloc that may free it is anchored in that same entry block --
	/// either directly, or nested inside an op of that block (e.g. an
	/// `scf.if` body), which conservatively bounds the lifetime.
	/// A dealloc that escapes the entry block entirely (e.g. lives in a sibling
	/// `cf` block) is reported as an error, as is an alloc with no dealloc.
	///
	/// To keep the rewrite safe, a dealloc is only accepted if all allocs it may
	/// free (per the reverse alias set) are themselves candidates in this block;
	/// otherwise erasing it during the rewrite could leak or double-free a buffer
	/// that is not managed by the arena.
	static LogicalResult
	collectCandidates(FunctionOpInterface funcOp,
	const BufferViewFlowAnalysis &analysis,
	llvm::Statistic &numSkipDynamic,
	llvm::Statistic &numSkipNested, llvm::Statistic &numEligible,
	SmallVector<AllocationCandidate> &candidates) {
	// All candidates are planned relative to the function's entry block.
	if (funcOp.getFunctionBody().empty())
	return success();
	Block *planBlock = &funcOp.getFunctionBody().front();

	bool walkFailed = false;
	funcOp->walk([&](memref::AllocOp allocOp) -> WalkResult {
	MemRefType memrefType = allocOp.getType();
	if (!memrefType.hasStaticShape()) {
	++numSkipDynamic;
	return WalkResult::advance();
	}

	// Only plan allocs that live directly in the entry block. Allocs nested in
	// a region (loop/conditional body) are skipped for now.
	if (allocOp->getBlock() != planBlock) {
	++numSkipNested;
	return WalkResult::advance();
	}

	SmallVector<memref::DeallocOp> deallocs;
	collectDeallocs(allocOp.getResult(), analysis, deallocs);

	if (deallocs.empty()) {
	allocOp.emitError("no dealloc found; run the deallocation pipeline "
	"before this pass");
	walkFailed = true;
	return WalkResult::interrupt();
	}

	for (memref::DeallocOp d : deallocs) {
	// The dealloc must be anchored in the plan block (directly or via an
	// enclosing op such as an scf.if). A dealloc in a sibling block escapes.
	if (!planBlock->findAncestorOpInBlock(*d.getOperation())) {
	allocOp.emitError("unstructured control flow is not supported");
	walkFailed = true;
	return WalkResult::interrupt();
	}
	// Every alloc that this dealloc may free must also be an entry-block
	// candidate; otherwise erasing it during the rewrite is unsafe.
	for (memref::AllocOp freed : findFreedAllocs(d, analysis)) {
	if (freed->getBlock() != planBlock) {
	++numSkipNested;
	return WalkResult::advance();
	}
	}
	}

	++numEligible;
	AllocationCandidate candidate;
	candidate.alloc = allocOp;
	candidate.deallocs = deallocs;
	candidate.sizeInBytes = computeSizeInBytes(memrefType);
	candidate.alignment = allocOp.getAlignment().value_or(1);
	candidates.push_back(candidate);
	return WalkResult::advance();
	});

	return failure(walkFailed);
	}

	/// 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 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) {
	SmallPtrSet<Operation *, 8> deallocsToErase;
	SmallVector<Operation *> allocsToErase;

	// Replace all alloc results with views (rewires selects too).
	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());
	allocsToErase.push_back(candidate.alloc.getOperation());

	for (memref::DeallocOp d : candidate.deallocs)
	deallocsToErase.insert(d.getOperation());
	}

	// Erase deallocs first (they may reference alloc results via selects).
	for (Operation *d : deallocsToErase)
	d->erase();

	// Erase allocs last (no users remain after replaceAllUsesWith).
	for (Operation *allocOp : allocsToErase)
	allocOp->erase();
	}

	//===----------------------------------------------------------------------===//
	// StaticMemoryPlannerAnalysisPass
	//===----------------------------------------------------------------------===//

	struct StaticMemoryPlannerAnalysisPass
	: public bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
	StaticMemoryPlannerAnalysisPass> {
	public:
	using Base = bufferization::impl::StaticMemoryPlannerAnalysisPassBase<
	StaticMemoryPlannerAnalysisPass>;
	using Base::Base;

	void runOnOperation() override;
	};

	void StaticMemoryPlannerAnalysisPass::runOnOperation() {
	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. The buffer view-flow
	// analysis models how buffers flow through selects, scf.if results, branches,
	// and view ops so we can find deallocs and freed allocs uniformly.
	BufferViewFlowAnalysis analysis(funcOp);
	SmallVector<AllocationCandidate> candidates;
	if (failed(collectCandidates(funcOp, analysis, numSkipDynamic, numSkipNested,
	numEligible, candidates)))
	return signalPassFailure();

	if (candidates.empty())
	return;

	// Step 2: Build allocation descriptors with lifetime info.
	SmallVector<bufferization::MemoryPlannerAlloc> allocInfos;
	int64_t arenaAlignment = buildAllocInfos(candidates, allocInfos);

	// 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;
	}

	// 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