Static Memory Planner for MLIR:
A GSoC 2026 Journey

Hi, I am Krish Gupta. Over the summer I worked with the LLVM Foundation as part of Google Summer of Code 2026 on building a static memory planner for MLIR, mentored by Matthias Springer and Javed Absar. This post walks through what we built, why it matters, and what the iterative process of landing it upstream actually looked like.

The Problem #

Modern ML and HPC workloads compiled through MLIR often go through the bufferization pipeline, which converts value-semantic tensor operations into explicit memref.alloc / memref.dealloc pairs. At that point, buffer lifetimes are fully explicit in the IR: the compiler knows exactly when each allocation starts and ends. Yet there was no upstream pass that used this information to eliminate redundant heap allocations.

This matters for accelerator-oriented compilation. On many targets (embedded CPUs, DSPs, custom accelerators) heap allocation is either slow, forbidden, or simply unavailable. Even where it is available, individual small allocations hurt performance through fragmentation and cache pressure. What you want instead is a single static arena whose layout is computed at compile time:

// Before: three separate heap allocations
%a = memref.alloc() : memref<1024xf32>  // 4096 bytes
%b = memref.alloc() : memref<512xf32>   // 2048 bytes
memref.dealloc %a : memref<1024xf32>
memref.dealloc %b : memref<512xf32>

// After: one arena, two views
%arena = memref.alloc() : memref<6144xi8>
%a = memref.view %arena[0][]  : memref<6144xi8> to memref<1024xf32>
%b = memref.view %arena[4096][]: memref<6144xi8> to memref<512xf32>

The arena is memref<Nxi8> so it can hold mixed element types. memref.view reinterprets slices back to their original types with zero overhead. This is the core transformation the planner performs.

How the Pass Works #

The pass is a FunctionOpInterface pass: it runs per function and processes the entire function body in six steps:

The core data structure is AllocationCandidate, which groups a memref.alloc, its set of potential deallocs (there can be more than one, reachable through alias chains), the computed byte offset, size, and alignment. The entire alias reasoning goes through BufferViewFlowAnalysis: a unified analysis that models how buffers flow through arith.select, scf.if/scf.for results, cf branches, and view ops without any op-type-specific code in the pass itself.

PR 1 · The Foundation #

#205125 — landed late June

The first PR introduced the core pass: static-memory-planner-analysis. It handled the straightforward case (allocs and deallocs in the same basic block with a direct one-to-one relationship) and put in place the architecture everything else builds on:

• trivialMemoryPlanner: a pure C++ function that takes a list of (size, alignment, timeStart, timeEnd) descriptors and returns offsets. Being a pure function means it can be unit tested independently and is easy to swap out.
• bestFitMemoryPlanner: a lifetime-aware algorithm that tries to reuse arena slots for allocations whose lifetimes do not overlap.
• Arena modes: allocate (create the arena inside the function) and arg (the arena is passed as a function argument, useful for external memory management).
• Alignment: the arena alignment is the LCM of all individual alignments, ensuring every view is correctly aligned regardless of element type.

The review process surfaced an important point early: upstream contribution culture expects every design decision to be justified and every edge case explicitly tested.

PR 2 · arith.select Deallocs #

#209106

After the first PR landed, we hit the first real-world blocker. The ownership-based-buffer-deallocation pipeline routinely produces patterns like this:

%a = memref.alloc() : memref<1024xf32>
%b = memref.alloc() : memref<1024xf32>
%sel = arith.select %cond, %a, %b : memref<1024xf32>
memref.dealloc %sel : memref<1024xf32>

Here neither %a nor %b has a direct memref.dealloc user. The original pass just incremented a skip counter for both and moved on. The fix was a forward DFS over the use-def graph, following BufferViewFlowOpInterface (which arith.select implements) to find reachable deallocs.

This also required a reverse-alias safety guard: since dealloc %sel may free either %a or %b, erasing it after replacing both with arena views is only safe if all allocs it might free are managed by the arena. The review cycle shaped the final approach: drop the group-constraint fixpoint in favour of the guard, emit hard errors (not silent skips) for missing deallocs, and improve the buildAllocInfos scan from O(n × m) to O(n+m) using a single upfront DenseMap.

PR 3 · scf.if and the Architectural Shift #

#213634

The deallocation pipeline also produces structured-control-flow patterns the DFS could not follow:

%a = memref.alloc() : memref<1024xf32>
%0 = scf.if %c -> memref<1024xf32> {
  scf.yield %a
} else {
  scf.yield %b
}
memref.dealloc %0   // dealloc is on %0, not %a

%a = memref.alloc() : memref<1024xf32>
scf.if %c {
  memref.dealloc %a
}

The mentor pointed to BufferViewFlowAnalysis, a unified alias analysis already used elsewhere in the bufferization pipeline. It models buffer flow through all relevant op types in one graph. We deleted the hand-rolled DFS entirely and replaced it with two calls:

• analysis.resolve(alloc) — forward alias set: every SSA value the alloc may flow into. Walking users of each alias finds all potential deallocs.
• analysis.resolveReverse(dealloc.getMemref()) — reverse alias set: every alloc that flows into the dealloc's operand. Used by the safety guard.

Lifetime anchoring via findAncestorOpInBlock. A dealloc nested inside an scf.if body is in a different block. Block::findAncestorOpInBlock walks up the parent chain and returns the enclosing scf.if op in the plan block. That op's index becomes timeEnd — conservative, but correct.

The reverse-alias guard (stronger). For the scf.if-sharing-a-nested-alloc case:

%a = memref.alloc()           // entry block
%0 = scf.if %c -> memref<1024xf32> {
  memref.dealloc %a
  %b = memref.alloc()         // NESTED -- not arena-managed
  scf.yield %b
} else {
  scf.yield %a
}
memref.dealloc %0   // may free EITHER %a OR %b

resolveReverse(%0) traces back to both %a and %b. Since %b is nested (not in the plan block), the guard fires: %a is conservatively skipped. No miscompile, no leak.

PR 4 · scf.for Test Cases #

#215221

After the BufferViewFlowAnalysis switch, scf.for patterns were handled correctly with no additional pass changes. This PR documented the results:

• Allocs nested inside an scf.for body are skipped (same rule as nested scf.if).
• scf.for with iter_args passing buffers across iterations hits the guard in all three mentor patterns.
• The canonical valid case — entry-block allocs used inside the loop body with deallocs in the entry block — transforms cleanly.

This PR was deliberately tests-only. Having a clear separation between "correct transformation that just needed documentation" and "new design decision" kept the review simple.

PR 5 · Edge and Error Cases #

#216610

The final PR stress-tested the pass with patterns existing tests did not cover. Eight new analysis tests and four new error tests:

• Static and dynamic shapes coexisting in one function (static transforms, dynamic is skipped without interference).
• Both branches of an scf.if deallocating the same alloc (duplicate dealloc ops deduplicated correctly).
• Dealloc at depth 3: findAncestorOpInBlock walks up any depth, not just one level.
• Multi-hop alias chain: %a/%b → %0 → %1 → dealloc. The resolve() BFS follows the full chain.
• Cross-interface chain: arith.select → scf.if → dealloc. Two different interface implementations in one alias path.
• cf.cond_br to a sibling block rejected as unstructured control flow.
• WalkResult::interrupt() points the diagnostic at the problematic alloc, not at subsequent valid ones.

Challenges #

What Was Built #

The final pass lives in StaticMemoryPlannerAnalysis.cpp.

Future Work #

The immediate open question is lifting the entry-block restriction. Right now, only allocs in the function's entry block are planned. Extending this to nested regions requires reasoning about nested lifetime intervals and is a natural next step.

Beyond that, the timeStart/timeEnd metadata infrastructure is in place for more sophisticated planning algorithms: polyhedral lifetime analysis, memory space awareness for heterogeneous targets, and integration with the broader bufferization pipeline for end-to-end static allocation.

Acknowledgements #

Huge thanks to my mentors Matthias Springer and Javed Absar for the patient and precise review feedback, for pointing at BufferViewFlowAnalysis at exactly the right moment, and for treating every PR as a teaching opportunity. The LLVM community's review culture is demanding in the best possible way.

The project GSoC page is at summerofcode.withgoogle.com/programs/2026/projects/XsjxBQ9o. Feel free to reach out with questions or to build on this work at krishgupta2832@gmail.com.