MemoryReuse Pass¶

Uses dependency analysis to identify memory reuse opportunities and removes redundant alloc operations.

Overview¶

This pass analyzes variable lifetimes and dependencies to enable memory sharing. Variables with non-overlapping lifetimes in the same memory space can share MemRef objects, reducing memory footprint.

After applying MemRef sharing, the pass also removes redundant tile.alloc statements for MemRefs that are no longer referenced by any TileType variable.

Key insights:

Variables that don't overlap in lifetime can reuse memory
Only variables in the same memory space can share MemRef
Lifetime is determined by def-use analysis
After sharing, MemRefs that become unreferenced are cleaned up along with their alloc statements

When to use: Run after MaterializeSemanticAliases and before AllocateMemoryAddr. Reduces memory allocation overhead. This pass does the opportunistic lifetime coalescing only; the semantics-required must-alias retarget (loop-carry / in-place — this pass's former "Step 0") now runs in MaterializeSemanticAliases, so MemoryReuse can be skipped independently (e.g. memory_planner=PTOAS, where ptoas owns lifetime reuse).

API¶

C++	Python	Level
`pass::MemoryReuse()`	`passes.memory_reuse()`	Function-level

Factory function:

Pass MemoryReuse();

Python usage:

from pypto.pypto_core import passes

reuse_pass = passes.memory_reuse()
program_optimized = reuse_pass(program)

Algorithm¶

Lifetime Analysis: Walk the full IR tree (including nested control flow bodies) to compute variable lifetimes via def-use analysis. Variables defined outside a loop but used inside have their lifetime extended to the end of the loop (loop-aware extension)
Interference Check: Identify variables with overlapping lifetimes
MemRef Sharing (global first-fit-decreasing packing, IdentifyReuseOpportunities): Within each memory space, intervals are packed largest-first; every later interval joins the first existing buffer all of whose members it can share with (non-overlapping lifetime + hazard/no-alias safe — see can_share). A buffer's allocated size is fixed by its first (largest) member, so admitting a smaller member afterwards costs nothing — and a later, larger interval can now host an earlier, smaller one. (The previous definition-order greedy had a one-directional size gate source.size >= target.size, so two lifetime-disjoint tiles whose small one was defined first were never coalesced.) The representative each member is rebased onto is the buffer's largest member; its tile.alloc dominates the whole function because InitMemRef hoists every alloc to the function-body head, so a representative defined after some of its members is safe. Because the packer no longer runs in program order, every pairwise gate (hazard, no-alias) is checked in both directions.
Loop-carry re-alignment (AlignLoopCarriesToInitMutator): Sharing (step 3) only retypes AssignStmt-defined vars (producers/init); loop-carried iter_arg/return_var nodes are excluded from the lifetime/sharing maps and keep their original MemRef. This step walks ForStmts top-down and retypes each loop's iter_arg/return_var to its (now-reused) initValue MemRef, seeding var_remap_ before recursing so a nested loop observes the corrected outer iter_arg as its init. Without it, a reused nested pipelined matmul_acc accumulator splits across two Acc buffers and step 5 emits invalid acc→acc tile.move ops that ptoas rejects on Ascend 910B (#1352)
Accumulator if-phi coalescing (TopDownRetargeter::CoalesceAccumulatorIfPhis): LowerPipelineLoops peels a stage-2 K-loop into if-phis whose live branch is an in-place matmul_acc (on the accumulator buffer) and whose dead if k==0 branch is a fresh matmul seed on a different Acc buffer. Left alone, step 6 reconciles them with an acc→acc tile.move — a second co-live L0C buffer (overflow) that ptoas also rejects (no legal Acc→Acc tmov). This step identifies the in-place-accumulator branch by its reuses_input producer and retargets the other branch's seed onto the accumulator buffer, so both branches share it and no move is emitted (matching mad_acc's shared-%dst semantics). Scoped to Acc; the retarget is mandatory (a declined retarget is an INTERNAL_CHECK, never a move — no legal Acc→Acc move exists). It bypasses the global dead-at-assign liveness check (which would false-decline on the legitimate post-if phi consumer), but only after verifying the two preconditions branch exclusivity actually needs: (a) the seed producer is a Call lexically inside the branch (a pre-if value yielded through the branch runs unconditionally and would clobber the accumulator the sibling in-place branch reads), and (b) a branch-scoped IsTargetDeadAtAssign (bounded to stop at the enclosing if) finds no same-branch tail read of the accumulator buffer. When either fails, the phi is left to step 6 instead of being coalesced
Yield fixup: Fix MemRef mismatches in control flow return variables:
ForStmt: Ensure all 4 loop-carry variables (initValue, iter_arg, yield value, return_var) share the same MemRef. Inserts tile.move before yield if MemRefs differ
IfStmt: Patch return_vars to match yield value's MemRef
Identity-copy buffer normalization (NormalizeIdentityCopyBuffersMutator): after step 5 retargets an accumulator if-phi, a downstream bare-Var SSA identity copy of the (now-moved) return_var can still carry the pre-coalesce buffer (e.g. c: …mem_acc_17 = c_phi after c_phi moved to mem_acc_5). An x = y copy (value a bare Var, not a Call) is a pure rename and must alias y's buffer, so this single forward pass retypes such a copy's LHS to the RHS's MemRef and substitutes the LHS's downstream uses. No-op when no mismatch exists
Remove redundant allocs: Collect all MemRefs still referenced by TileType variables, then remove tile.alloc statements whose MemRef is no longer in use

Reuse conditions:

Non-overlapping lifetimes (no interference). Two variables do NOT overlap when prev.last_use <= curr.def (i.e., the source's last use can be at the same statement as the target's definition, since inputs are read before outputs are written within a single statement).
Same memory space
A buffer is sized to its largest member; because packing is largest-first, every member admitted later is no larger than the representative, so no explicit byte-size check is needed (and the reuse direction is no longer constrained to "earlier-and-larger only")
No-alias guard (op-semantic): the op that defines the reusing variable may forbid its output from sharing a buffer with one or more of its input operands, because the hardware reads those inputs while writing the output — an in-place write would corrupt the op mid-flight. Three sources feed one per-output forbidden-input set (ForbidAliasCollector):
not_inplace_safe() — the op cannot run with src == dst, so its output must not alias any input operand.
forbid_output_alias(i) — the op is in-place-safe w.r.t. its value operands but reads a specific operand while writing its output, so the output must not alias that one operand's buffer.
widening tile.cast (handled directly in ForbidAliasCollector) — when the output dtype is wider than the input, the cast cannot run in place: element i is read at i*in_bytes but written at i*out_bytes, so the write cursor outruns the read cursor and clobbers not-yet-converted input. Narrowing / same-width casts stay in-place-safe (preserving the cross-dtype reuse below).

MemoryReuse refuses to place the output on a forbidden operand's physical buffer, resolving each operand through both reuse-map coalescing and VIEW inheritance (a reshape/slice shares its source's MemRef base) — so a forbidden operand reached indirectly (its owning tile reused onto another buffer, or occupied via a view) is still caught.

Ops that currently declare a no-alias constraint:

Op(s)	Constraint	Why the output cannot alias the input
`tile.recip`, `tile.rsqrt`	`not_inplace_safe`	high-precision path reads the input and a tmp scratch while writing the output
`tile.row_sum` / `row_max` / `row_min`	`not_inplace_safe`	`TROW*` reads the full input row + tmp scratch while writing the reduced `[M, 1]` output
`tile.mrgsort_format1`	`not_inplace_safe`	merge-sort intrinsic requires `src != dst`
`tile.fmod`, `tile.fmods`	`not_inplace_safe`	`TFMOD`/`TFMODS` compute `a - trunc(a/b)*b` by overwriting `dst = a/b` first, then re-reading the original `src0` (`a`) for the final subtraction; when `dst == src0` that subtraction sees the already-clobbered quotient and yields `0` for every element
`tile.transpose`	`not_inplace_safe`	`pto.ttrans` is not in-place safe: the a2a3 unaligned scalar path writes `dst` directly from `src` (no tmp staging), so `dst == src` corrupts the data mid-write. The output always gets a fresh buffer (also enforced in InitMemRef, which never inherits the input's buffer for it).
`tile.sel`	`forbid_output_alias(0)` (mask), `(3)` (tmp)	`TSEL` reads the predicate mask + tmp scratch while writing `dst`
`tile.{row,col}_expand{,_mul,_add,_sub,_div}`	`forbid_output_alias(1)` (broadcast vector)	the row/col vector (arg 1) is re-read for every output row/col, so an output aliasing it is overwritten after the first row/col
`tile.cast` (widening only)	output ≠ input buffer (conditional, in `ForbidAliasCollector`)	wider output's write cursor outruns the read cursor (see above)

Pipeline-stage guard (capacity-gated): pl.pipeline(stage=F) replicates a loop body F times for ping-pong, and LowerPipelineLoops tags each clone's tile-producing Call with a pipeline_membership (group, stage) (see 26-lower_pipeline_loops.md). The F clones run concurrently under the scheduler, so their program-order-disjoint lifetimes are not a safe reuse signal — collapsing concurrent clones onto one buffer injects a false write-after-read that serializes the stages (the #1475 cube-matmul-operand collapse). MemoryReuse therefore keeps concurrent clones in distinct buffers in every memory space — including the L0 matmul spaces (Left/Right/Acc/Bias), and regardless of whether a tile is a load or a tile.move result — up to the max-affordable double-buffering depth F_g = min(depth_g, ⌊C_s / slot_g⌋): a clone at stage k lands in residue ordinal(k) mod F_g (the dense stage ordinal, so sparse stage IDs like {0, 2} can't collide via raw 2 mod 2 == 0 mod 2), so concurrent clones never share (full ping-pong when it fits, maximal spread when the space is tight). Whether the separation fits is decided by the exact per-space allocator footprint (SpaceFootprint, shared with AllocateMemoryAddr — parity by construction), not an estimate. When a space still overflows at every group's affordable depth, a graceful cross-group shed lowers one group's depth by a residue and re-packs (choosing the group by the MaxRelief heuristic — free the most bytes first, ties by lowest group id); if the shed exhausts, the pass re-packs from scratch in a legacy fallback (force_legacy), so it never overflows where the legacy packing would have fit. A space whose capacity is unknown (no backend configured) uses the legacy predicate, so the capacity-gated path is never worse than legacy. When the gate reduces a group below its requested stage= depth (or a space hits the legacy fallback), MemoryReuse emits a diagnostic through the unified channel — a PH-MR-001 perf hint (or a warning for the fallback) naming the requested vs achieved depth and the fix (shrink the per-stage tile to ≤ C_s / stage, or lower stage= to what fits) — so a capacity-forced serialization is never silent. Once the reuse decision is made, MemoryReuse strips the now-consumed pipeline_membership attr so it does not ride downstream into later passes or codegen.

No shape / dtype / TileView compatibility gate: tiles that share a physical MemRef may carry different shapes, dtypes, or TileView attributes. PTO codegen binds a per-variable alloc_tile to each tile, so each alias declares the shared base with its own static shape / dtype / layout / valid_shape. This permits, for example:

cross-dtype reuse — a BF16 tile reusing a dead FP32 tile's buffer (e.g. across tile.cast);
tile.fillpad output reusing its input, and two fillpad outputs with different pad sharing one buffer;
N-D tiles with divergent valid_shape sharing a buffer (each keeps its own valid_shape on its own alloc_tile);
L0 cube-input Left / Right sub-tiles of differing shape sharing one slot (e.g. fused-attention QK Right [k, SEQ] reused by PV Right [k', HEAD], halving peak L0B — issue #1595).

Earlier revisions gated reuse on an AreTileTypesCompatible shape / dtype / view match (with a narrow L0 byte-reuse exception); that gate has been removed. Correctness for read-while-write ops is now handled precisely by the no-alias guard above rather than by a coarse whole-tile match.

Alloc cleanup:

After MemRef sharing, some MemRef objects become unreferenced (their variables now point to a different shared MemRef). The pass traverses the surrounding SeqStmts and removes any tile.alloc AssignStmt whose LHS MemRef pointer is not in the set of still-used MemRefs.

Ascend910B load + tpop_from_aic hazard¶

On Ascend910B AIV functions with a non-None SplitMode, a writer that consumes both a tile.load result (or a legal-view descendant of one) and a tile.tpop_from_aic value must not place its output in the same physical buffer as that load result. Allowing the writer's output to in-place-reuse the load buffer produces silently wrong results on this hardware.

MemoryReuse owns every buffer-coalescing decision, so it prevents the hazardous sharing from ever forming rather than relying on a later split. When the guard is active, the reuse decision is blocked exactly when:

the writer's defining op consumes a tile.tpop_from_aic value, and
the buffer member it would reuse in place (whose last use is the writer's def statement) is load-derived.

The guard is gated by BackendHandler::RequiresSplitLoadTpopWorkaround() (true only for Ascend910B) and the function being split-AIV; on every other backend / function kind the inputs are empty and reuse behaviour is unchanged. The writer is still free to reuse any non-load buffer — only the load + tpop in-place combination is rejected. (This guard previously lived in a dedicated LegalizePTOBufferReuse pass that split the buffer after the fact; it now folds into MemoryReuse.)

Example¶

Before (after InitMemRef):

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
mem_vec_2: MemRefType = tile.alloc(Vec, -1, 16384, 2)
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.add(tile_a, ...)
# tile_a last use ↑
tile_c: Tile[[64, 64], FP32, memref=mem_vec_2] = tile.load(...)
# ]

After (tile_c reuses mem_vec_0 from tile_a, alloc for mem_vec_2 removed):

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
# mem_vec_2 alloc removed — no longer referenced
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.add(tile_a, ...)
tile_c: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
# tile_c now shares mem_vec_0 with tile_a
# ]

Producer-Consumer Reuse¶

When a variable's last use is at the same statement that defines a new variable (producer-consumer relationship), the new variable can reuse the old variable's memory because inputs are read before outputs are written:

# Before:
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.create(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.muls(tile_a, 0.0)
# tile_a.last_use == tile_b.def → reuse allowed

# After:
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.create(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.muls(tile_a, 0.0)
# tile_b reuses mem_vec_0

Overlapping Lifetimes (No Reuse)¶

When a variable is still alive after another variable's definition (last_use > def), their lifetimes truly overlap and they cannot share memory:

# SeqStmts [
mem_vec_0: MemRefType = tile.alloc(Vec, -1, 16384, 0)
mem_vec_1: MemRefType = tile.alloc(Vec, -1, 16384, 1)
tile_a: Tile[[64, 64], FP32, memref=mem_vec_0] = tile.load(...)
tile_b: Tile[[64, 64], FP32, memref=mem_vec_1] = tile.load(...)
# tile_a.last_use > tile_b.def → tile_a still live when tile_b is defined
# ]

Implementation¶

Header: include/pypto/ir/transforms/passes.h

Pass MemoryReuse();

Implementation: src/ir/transforms/memory_reuse_pass.cpp

LifetimeAnalyzer walks the full IR tree to compute variable lifetimes (including nested control flow)
ComputeLifetimes builds MemRef sharing groups and lifetime intervals
IdentifyReuseOpportunities finds reuse candidates
ApplyMemRefSharing updates MemRef pointers via MemRefSharingMutator
TopDownRetargeter::CoalesceAccumulatorIfPhis coalesces peeled loop-carried accumulator if-phis by retargeting the dead-branch seed onto the in-place accumulator buffer, so YieldFixupMutator never emits an illegal acc→acc tile.move (see Algorithm step 5)
YieldFixupMutator fixes ForStmt/IfStmt yield/return_var MemRef mismatches after reuse (inserts tile.move when needed)
NormalizeIdentityCopyBuffersMutator reconciles bare-Var SSA identity copies whose LHS/RHS buffers diverged after accumulator if-phi coalescing (see Algorithm step 7)
UsedMemRefCollector gathers still-referenced MemRef pointers after sharing
RemoveUnusedAllocStatements filters out redundant tile.alloc statements from SeqStmts

Python binding: python/bindings/modules/passes.cpp

passes.def("memory_reuse", &pass::MemoryReuse, "Memory reuse optimization");

Tests: tests/ut/ir/transforms/test_memory_reuse.py

Tests non-overlapping lifetime reuse with MemRef sharing
Tests producer-consumer reuse (last_use == def at same statement)
Tests overlapping lifetime no-reuse
Tests memory space separation
Tests byte-size compatibility
Tests cross-dtype / cross-TileView reuse (now permitted: BF16↔FP32, fillpad output↔input, divergent valid_shape)
Tests the no-alias guard (TestForbidOutputAlias + TestInplaceOps), one case per constraint above:
tile.recip / tile.rsqrt / tile.row_sum — output must not alias input (not_inplace_safe)
tile.sel — output must not alias the mask / tmp (forbid_output_alias)
tile.col_expand_mul — output must not alias the broadcast vector
widening tile.cast — output must not alias the (narrower) input
a forbidden operand reached through a VIEW is still honored (physical-buffer resolution)
Tests view operation MemRef sharing preservation
Tests redundant alloc statement removal
Tests control flow lifetime analysis (nested IfStmt in ForStmt, branch variable sharing)