Skip to content

LowerPipelineToSlots Pass

Multi-buffers a pl.pipeline(N, stage=F) loop by rotating one body through the F slots of a single allocation, instead of replicating the body F times.

Overview

pl.pipeline(N, stage=F) asks for ping-pong buffering. LowerPipelineLoops delivers it by replication: F copies of the body, each with fresh def-vars, so each copy's tiles are distinct MemRefs that MemoryReuse is forbidden to coalesce. That works, but it costs F times the code, a static-or-dynamic remainder dispatch, and the pipeline_membership machinery that keeps the copies apart.

This pass expresses the same intent in the form ptoas already understands. pl.MemRef(name, slots=F) says "one allocation, F uniform slots, this use takes slot k" — see Slots — and PTO codegen lowers exactly that to pto.alloc_multi_tile + pto.multi_tile_get. So the loop keeps one body and each per-stage buffer becomes slot iv % F of a synthesized declaration:

# Before (as this pass sees it)
for i in pl.pipeline(64, stage=2):
    x: pl.Tile[[128], pl.FP32, pl.Mem.Vec] = pl.tile.load(a, [i * 128], [128])
    pl.tile.store(x, [i * 128], out)

# After — one body, bounds and step untouched, kind demoted
for i in pl.range(64):
    x: pl.Tile[[128], pl.FP32, pl.MemRef("pipe_x", slots=2)[i % 2], pl.Mem.Vec] = \
        pl.tile.load(a, [i * 128], [128])
    pl.tile.store(x, [i * 128], out)

No new IR op and no new user-facing switch. The synthesized MemRef is shaped exactly like an author's declaration, so InitMemRef resolves it through the same path and codegen sees no difference between a rotation the author wrote and one the compiler derived.

Because bounds, step and iter_args are untouched, there is no remainder to dispatch — a dynamic trip count needs no special case at all.

Requires: SSAForm, SplitIncoreOrch, IncoreTileOps, TileOps2D, TileMemoryInferred, NormalizedStmtStructure.

Pipeline position: After SkewCrossCorePipeline, immediately before LowerPipelineLoops. Late enough that memory spaces are inferred and the tile structure is final; early enough that InitMemRef has not yet handed the tiles compiler-owned MemRefs.

The two passes are complementary, not alternatives

Both run, in this order. This pass takes the loops it can prove safe and demotes them; every loop it declines keeps ForKind::Pipeline and is replicated by LowerPipelineLoops exactly as before. Nothing loses its ping-pong because this pass exists — matmul L0 stage loops, nested pipelines and unusual loop shapes all keep the replication path.

This mirrors SkewCrossCorePipeline, which handles cross-core pipeline loops the same way and leaves the rest intact.

Gated on memory_planner=PTOAS

Under the default PyPTO planner the pass returns every function untouched, so that path stays byte-identical.

The gate tracks where a region is emitted, not where ptoas can use one. PTO codegen's PlanMultiBufferRegions bails under the PyPTO planner, so a rotation synthesized there would resolve to a runtime address on an ordinary alloc_tile — correct, but with none of the slot analysis this transform exists to trigger.

ptoas itself is not the limitation. Given pto.alloc_multi_tile addr = <constant base>, ptoas 0.55 derives the same per-slot dynamic-event synchronization at --pto-level=level3 as it does at level2 — measured on a prefetch loop, the sync-op sequence is identical after normalizing event ids (two primed events, wait_flag/set_flag keyed on the slot, two drains). ptoas 0.54 does not, which is what the older PTOAS#1106 note describes.

Widening the gate therefore needs work on the PyPTO side rather than upstream: the address allocator must reserve slot_count * slot_size for the region base and codegen must emit that address on the region. That is follow-up work; until it lands, this pass stays scoped to the planner whose codegen path already exists.

API

C++ Python Level
pass::LowerPipelineToSlots() passes.lower_pipeline_to_slots() Function-level 
from pypto import passes
with passes.PassContext([], memory_planner=passes.MemoryPlanner.PTOAS):
    result = passes.lower_pipeline_to_slots()(program)

Behavior

For a ForStmt(kind == ForKind::Pipeline, attrs["pipeline_stages"] == F) with F > 1 that passes every gate below:

  1. Each candidate tile's TileType is rebound onto a fresh pinned MemRef(name, slots=F) with slot_index = iv % F. The definition and all its uses are rewritten in one walk — the IR is in SSA form, so no use precedes its definition.
  2. kind_ becomes ForKind::Sequential and pipeline_stages is stripped. The two always travel together, so the PipelineLoopValid invariant (kind == Pipeline iff pipeline_stages present) holds at every observable state.
  3. Nothing else changes: no statement is added, removed or reordered, and start / stop / step / iter_args are left alone.

F == 1 is a user-written pl.pipeline(stage=1) or the marker a previous LowerPipelineLoops run left behind. Nothing needs multi-buffering, so the (kind, attr) pair stays whole for CanonicalizeIOOrder to scope on.

Which tiles take a slot

Every top-level tile.load result the author has not already bound.

  • Loads only. A load buffer is what must stay private so iteration i+1's prefetch overlaps iteration i's compute; compute intermediates may coalesce. This is the same distinction MemoryReuse already draws with pipeline_load_tiles. Giving every tile F private copies overflows the on-chip budget on real kernels — a stage=4 RMSNorm would need 4 x 67 KB > 188 KB UB. tile.read is not included: it returns a scalar element, not a tile, so there is no buffer to rotate.
  • Top-level. Only direct members of the loop body's SeqStmts; a load nested in an inner loop or if belongs to that region.
  • No loop-invariance filter. Every unbound top-level load qualifies, including one whose arguments never mention the induction variable. Invariance cannot be read off the induction variable: a load addressed through a loop-carried IterArg reads different data each iteration without naming iv, and skipping it would strand it with neither a slot nor a replicated copy once a sibling candidate demotes the loop. Slotting a genuinely invariant load costs nothing over the fallback either — LowerPipelineLoops replicates its buffer F times all the same.
  • A tile the author already bound to a declared allocation stays the author's.

Eligibility

Codegen refuses a region it cannot describe rather than degrading it, because falling back to per-slot alloc_tiles would let ptoas plan the slots on top of each other. So every gate here mirrors a PlanMultiBufferRegions blocker: synthesizing a doubtful region would turn a kernel that compiles today into a compile failure.

Gate Why
F in [2, 16] ptoas' multi_tile_buf slot-count bounds
Memory space is Vec, Mat or Acc The spaces ptoas accepts for a slot
Static valid shape A region declares one static extent for all its slots
Not carried into a phi A tile that is yielded, or used as a nested loop's init_values, makes that phi share its MemRef. Both reach the phi the same way — one through a YieldStmt, the other through IterArg::initValue_. Checked through alias roots: InitMemRef shares one MemRef across a bare a = b tile copy and across a view / in-place result, so yielding an alias carries the original's slot just as yielding it directly would
Not consumed by a view / in-place op Such a result is its source's buffer, so it would land on the same allocation with a different tile_buf type
Constant step > 0, constant start >= 0; for step > 1, a constant stop <= INT64_MAX - (step - 1) See below
No enclosing pipeline loop was declined See below
Slots fit the memory space See below

The decline is per loop, not per tile. The four tile gates above (space, static valid shape, phi, view / in-place) are checked only for a load that actually wants a slot — top-level and not already bound by the author. If any such load trips one of them, the whole loop is declined, even when its siblings are eligible. Dropping just that one load would not work: a surviving sibling still demotes the loop to Sequential, and the blocked load would then reach neither these slots nor LowerPipelineLoops' replication, silently losing the per-stage privacy the pl.pipeline(stage=F) annotation asked for. Only an author-bound tile is skipped without affecting the loop, because declining over it would push its declaration onto the replication path, which rejects it.

Why the slot index is iv % F, (iv - start) % F or ((iv - start) / step) % F. ptoas' slot accounting reads the value pto.multi_tile_get carries and reasons about (symbol + offset) mod N; what it needs is for two accesses to share the identity of that symbol, not for the expression to have a particular shape. So all three qualify.

All three are ((iv - start) / step) % F with the parts that are identities dropped. Which one is emitted is decided by the loop shape. Eligibility is not: any constant step > 0 and constant start >= 0 qualifies.

  • step == 1, start % F == 0: the induction variable is the rotation counter, so the slot is literally iv % F. This is the form these loops have always had.
  • step == 1, start % F != 0: the counter starts at start, so the slot is (iv - start) % F. ptoas sees (symbol + constant) mod N, which its slot matcher peels. pl.pipeline(1, 5, 1, stage=2) selects 0, 1, 0, 1.
  • step > 1: iv takes start + j*step, so the slot is ((iv - start) / step) % F. The division is exact by construction, so this is still a remainder over a single SSA value, never a folded byte offset.

A runtime step or start is still declined: the slot would then depend on a value the analysis cannot bound. A stepped loop (step > 1) is also declined unless stop is a constant no larger than INT64_MAX - (step - 1): the retained scf.for computes iv + step after its last iteration, and that update can overshoot stop by up to step - 1, overflowing the signed index. pl.pipeline(INT64_MAX - 4, INT64_MAX, 3, stage=3) is the shape this closes. step == 1 cannot overshoot, so it keeps a runtime stop. Note that the counter form loses the kDisjoint inference for an (iv + c) % N-style pair, because the constant peel sits above the division and the two sides end up with different symbols; that degrades to kUnknown, which the callers treat conservatively (the dependency is kept), never to a wrong answer.

A negative start is declined in every form: iv % F is a signed remainder, so a negative iv selects a negative slot, and iv - start is not representable for a sufficiently negative start.

Why the slots must fit on chip. The declared slots are pinned: InitMemRef sizes the allocation at F * slot_size and ptoas may not reuse any of it, so this pass is directly accountable for those bytes. A loop with many eligible loads otherwise multiplies its footprint by F and ptoas answers a region it cannot place with a hard overflow error — it does not degrade. The replication path does: MemoryReuse's capacity gate lowers the effective double-buffering depth (F_g = min(depth_g, ⌊C_s / slot_g⌋)) and sheds groups until the space fits, so declining hands the loop to a path that shrinks rather than fails.

The budget is per memory space, summed over the loop's candidates and accumulated across the whole function — a slotted inner loop's region is co-live with its slotted ancestor's. It is seeded with the allocations the author already declared: those are is_pinned_ too, so ptoas cannot reuse them either, and ignoring them would admit a synthesized region that fits on its own while the pair overflows. Capacity comes from Backend::GetMemSize(space); a space with unknown capacity (no backend configured) is left ungated, mirroring MemoryReuse. This bounds only what this pass pins: tiles it does not slot are still planned by ptoas with lifetime reuse, which the pass cannot model.

Why a declined enclosing loop disqualifies everything below it. That loop will be replicated, and its F clones would each select one slot of the same allocation inside one loop body — a shape PTO codegen rejects under this planner (see one slot per iteration).

Generated PTO IR

%pipe_t_mb = pto.alloc_multi_tile valid_row = %c64_index valid_col = %c64_index
           : !pto.multi_tile_buf<!pto.tile_buf<loc=vec, dtype=f32, rows=64, cols=64, ...>, count=2>
scf.for %i = %c0_index to %c4_index step %c1_index {
  %slot = arith.remsi %i, %c2_index : index
  %t = pto.multi_tile_get %pipe_t_mb[%slot] : !pto.multi_tile_buf<..., count=2> -> !pto.tile_buf<...>
  pto.tload ins(...) outs(%t : ...)
  ...
}

The operand remains a logical index, without byte scaling; ptoas can still derive dynamic slot selection and per-slot event ids.

Signed integer and INDEX FloorMod preserve Python's % semantics: a nonzero remainder has the divisor's sign, which arith.remsi alone does not give — it follows the dividend's sign instead. Codegen therefore corrects the remainder, so that -97 % 24 produces 23:

%rem = arith.remsi %i, %c2_index : index
%negative = arith.cmpi slt, %rem, %c0_index : index
%corrected = arith.addi %rem, %c2_index : index
%slot = arith.select %negative, %corrected, %rem : index

A positive constant divisor needs only the negative-remainder check above; a negative constant uses a positive-remainder check. For a dynamic divisor, codegen checks that the remainder is nonzero and its sign differs from the divisor before adding the divisor. Division by zero and signed minimum modulo -1 remain outside the operation's valid inputs.

The correction is skipped when it cannot change the result, which is why the example above shows a bare remsi. With a non-negative dividend and a positive constant divisor, arith.remsi already lands in [0, divisor), the two agree, and the correction would be dead code. Codegen proves the dividend non-negative syntactically, from the shapes an index expression is built out of: a literal; an induction variable whose loop starts at a non-negative literal (a device loop is always ascending, so the variable never goes below its start); that variable minus its own start; and floor divisions and floor remainders of those by a positive constant literal. Sums and products are deliberately excluded: being built from non-negative terms does not make a result non-negative, because the emitted arithmetic wraps — i * 4611686018427387904 is negative at i == 2.

Every form in that proof must also execute at the dtype the proof is made at. EmitArithOperand coerces each operand into its operation's dtype, so an induction variable reaching a narrower operation is a different value: over an INDEX loop from 2147483648, the i an INT32 operation sees is -2147483648. That applies to nested operations too, not just to the variable — an INT32 FloorDiv inside an INDEX FloorMod yields a negative quotient that is then widened back, so proving the subtree at the outer width would skip a correction the result needs. The check is conservative — where it cannot tell, the correction is emitted.

This matters beyond the three instructions saved. A multi-buffer slot operand has to stay in the (symbol + constant) mod N shape that ptoas' slot matcher peels, and a rem/cmp/add/select chain is not that shape. A rotation derived from a stepped or unaligned loop — ((iv - start) / step) % count — is both the case that needs the bare remainder and the case whose dividend is non-negative by construction.

Signed integer and INDEX FloorDiv use arith.floordivsi for the matching floor quotient, preserving a == (a // b) * b + a % b, and fall back to arith.divsi under the same proof — ptoas expands floordivsi into a truncating division plus the same conditional correction. Floating-point and unsigned lowering are unchanged. Floor-division and floor-modulo constant folding leave out-of-range typed literals for runtime evaluation so they cannot replace their emitted bit patterns with calculations on the raw stored integers.

The loop strides by its original step with a single body, and the region carries no addr — ptoas PlanMemory places it, which is what lets it overlap iteration i's load with iteration i-1's compute.