LowerPipelineLoops Pass

Lowers pl.pipeline(N, stage=F) at the tile level: replicates the loop body F times per outer iteration to enable ping-pong buffering, while keeping the outer loop sequential.

Overview

This is one of two lowerings for pl.pipeline. LowerPipelineToSlots runs immediately before this pass and, under memory_planner=PTOAS, multi-buffers what it can by rotating a single body through the slots of one allocation. It demotes every loop it takes, so this pass sees — and replicates — only the loops it declined, plus every loop under the default PyPTO planner. Replication remains the general path.

pl.unroll(N) fully expands a loop into N body copies at slot #1 (before SSA). Users reach for this not because they want N copies but because they need distinct tile MemRefs — MemoryReuse would otherwise coalesce sequentially-live tiles into a single buffer, defeating ping-pong execution.

pl.pipeline(N, stage=F) is the user-facing surface for that targeted knob: replicate the body F times (typically 2–4) at the tile level, leaving an outer loop of N/F iterations. Each clone gets fresh def-vars (SSA preserved) and operates on independent tiles.

Fresh SSA vars are not by themselves enough to keep the clones in separate buffers: the F clones are sequential in program order, so their per-clone tiles have disjoint program-order lifetimes — exactly the condition under which MemoryReuse would coalesce them into one buffer (defeating the ping-pong). To make stage separation explicit, this pass tags each stage's tile-producing Calls in clone k with a pipeline_membership attr recording (group, stage=k) (see include/pypto/ir/transforms/utils/attrs.h). A tile inside a nested pipeline carries one membership pair per enclosing replicated region, so nested same-core pipelines stay separated at every level. (Historically this separation was an unreliable side effect of CanonicalizeIOOrder clustering sibling-clone loads to induce lifetime overlap; the explicit attr makes it a hard constraint independent of statement ordering.)

Cube accumulators are the one exception by default — they are not tagged. A pipeline stage double-buffers the operands it loads: their loads overlap the previous stage's compute, so stages' operand buffers are genuinely co-live and must stay separate. An accumulator is written by the serialized cube, so the default drain-before-next schedule needs one L0C buffer regardless of operand pipeline depth. A loop carrying pipeline_double_buffer_c=true is the explicit exception: its cube Acc definitions are initially tagged with the full source stage so CanonicalizeIOOrder can form depth-two MMSS chunks; that pass then rotates the Acc membership modulo two before MemoryReuse, producing two L0C stage residues in each fully replicated group even when F > 2. The automatic AutoTile recognizer requires a trip count divisible by F, avoiding a separately allocated tail group. Enclosing pipelines still do not tag that accumulator. The producer-op test remains important: a non-MAD data movement targeting Acc (for example tile.extract(..., target_memory=Acc)) is a genuine per-stage buffer and keeps the ordinary full-stage membership.

MemoryReuse consumes the attr with role-aware granularity — forbidding all cross-stage reuse (depth = F) would need F full copies of every intermediate and overflows the on-chip budget on real kernels (e.g. stage=4 RMSNorm needs 4 × 67 KB > 188 KB UB). Only load buffers generally need per-stage privacy so iteration i+1's prefetch overlaps iteration i's compute. So the legacy rule is: block a cross-stage buffer share iff the two tiles are same-group / different-stage and at least one is a load (tile.load / tile.read), with the L0 matmul spaces exempt entirely. The default path is the capacity gate (#1475): it separates operand L0 spaces (Left/Right/Bias) per stage up to the affordable depth. Acc normally stays untagged and coalesces to one slot; the explicit dbC path described above supplies a modulo-two membership and therefore exactly two slots.

Because the tag is a generic op-call attr, it is serialized through the python printer / parser (attrs={"pipeline_membership": "..."}) so it survives the print→parse round-trip the test harness runs after every pass.

Internally, pl.pipeline(...) emits ForStmt(kind=ForKind::Pipeline, attrs={"pipeline_stages": F}). The bidirectional structural invariant kind == Pipeline ⇔ pipeline_stages attr present (verified by PipelineLoopValid) holds at every observable IR state — the kind and the attr always travel together. LowerPipelineLoops triggers when F > 1; it replicates the body and downgrades the attr to 1 as the post-lowering marker, keeping ForKind::Pipeline for the downstream CanonicalizeIOOrder pass to scope on. Re-running LowerPipelineLoops is a no-op (it sees factor == 1 and skips).

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

Pipeline position: After SkewCrossCorePipeline (and NormalizeReturnOrder), before CanonicalizeIOOrder and InitMemRef. Cross-core (cube/vector) pipeline loops are skewed to ForKind::Sequential by the upstream skew pass, so by here only same-core pipeline loops (GM→L1, L1→L0, nested matmul stage loops) remain ForKind::Pipeline for this pass to replicate. Late enough that all tile-structural decisions are made; early enough that CanonicalizeIOOrder / InitMemRef / MemoryReuse see distinct tile vars per clone.

API

C++	Python	Level
pass::LowerPipelineLoops()	passes.lower_pipeline_loops()	Function-level

from pypto import passes
result = passes.lower_pipeline_loops()(program)

DSL Syntax

# Replicate the body 4 times per outer iteration; outer loop runs 16 iters with stride 4.
for i in pl.pipeline(64, stage=4):
    tile_x = pl.tile.load(input_a, [i * 128], [128])
    pl.tile.store(tile_x, [i * 128], output)

pl.pipeline accepts the same positional args as pl.range — (stop), (start, stop), or (start, stop, step) — plus the required stage= kwarg. init_values= is supported.

Behavior

For ForStmt(kind=Pipeline, attrs={"pipeline_stages": F}, start, stop, step, body) with F > 1:

• Main loop: stride F*step, body is a SeqStmts of F clones, kind still ForKind::Pipeline (marker), attr downgraded to 1. The (kind, attr) pair stays together so PipelineLoopValid holds and the IR survives print/parse round-trip (renders as pl.pipeline(..., stage=1)).
• Cloning: each clone uses DeepClone(body, {loop_var → new_var + k * step}, clone_def_vars=true). Fresh def-vars per clone keep SSA intact and give MemoryReuse distinct tile identities to work with.

stage=1 is a no-op trigger: the pass leaves the loop intact (kind and attr stay) and only recurses into the body to lower nested pipelines. CanonicalizeIOOrder then scopes on the kept marker, applies the IO reorder, and demotes/strips on exit. The same path handles user-written pl.pipeline(stage=1) and the post-lowering output of a factor>1 invocation — both ask for IO reorder without further replication, so re-running LowerPipelineLoops is naturally idempotent.

Two lowering modes — static vs dynamic — differ only in how the main loop's stop and the remainder are computed.

Static bounds — all of start, stop, step are compile-time integers

With trip count T = (stop - start) / step:

• Main loop stops at start + (T // F) * F * step.
• If T % F != 0, a bare SeqStmts of T % F cloned bodies at offsets start + (T // F) * F * step + j * step (for j ∈ [0, T%F)) is flattened directly into the outer scope. No runtime dispatch and no wrapper are needed — the remainder count is known.
• When the source loop has iter_args, trailing AssignStmts bind the outer loop's return_vars to the tail's final yielded expressions so downstream references stay valid.

Dynamic bounds — start and/or stop are runtime Exprs (step still static, positive)

• Compute the total trip count as trip_iters = ceil_div(stop - start, step). When step == 1 this collapses to stop - start and the pass emits the shorter form.
• Let main_iters = trip_iters / factor (floor-div) and materialize main_end = start + main_iters * (factor * step) as a fresh SSA AssignStmt (named unroll_main_end).
• Main loop is for i in range(start, main_end, F*step).
• Materialize rem_iters = trip_iters - main_iters * factor as a fresh SSA AssignStmt (named unroll_rem). When step == 1 this is equivalent to stop - main_end, and the pass emits that shorter form. The remainder is dispatched through a cascaded IfStmt chain:

if rem_iters == 1:    <1 clone>
else if rem_iters == 2: <2 clones>
else if rem_iters == 3: <3 clones>
# ...
else if rem_iters == F-1: <F-1 clones>
# rem_iters == 0 falls through — no tail work.

Each branch body is a bare SeqStmts of k cloned bodies (followed by a YieldStmt when the source loop has iter_args). The enclosing IfStmt carries return_vars; at the outermost level these are the original outer loop's return_vars, at inner levels they are fresh vars fed upward via successive YieldStmts. SSA stays clean: each branch is self-contained; no conditionally-defined var escapes its IfStmt.

Constraints

Constraint	Reason
step must be a compile-time integer constant	Main loop's stride and per-clone offsets both require stage * step as an integer
Dynamic bounds require step > 0	The dynamic trip-count formula assumes positive step; negative-step ranges must use static bounds
stage= is only accepted on pl.pipeline()	Scoped feature; pl.range() / pl.parallel() / pl.unroll() have different semantics

Examples

Static — trip count known (N=10, F=4)

# Before
for i in pl.pipeline(0, 10, 1, stage=4):
    tile_x = pl.tile.load(input_a, [i * 128], [128])
    pl.tile.store(tile_x, [i * 128], output)

# After: main loop covers [0, 8) with kind=Pipeline (marker), attr downgraded to
# stage=1; bare tail clones flattened into the outer scope.
for i in pl.pipeline(0, 8, 4, stage=1):
    tile_x_0 = pl.tile.load(input_a, [i * 128], [128]); pl.tile.store(tile_x_0, [i * 128], output)
    tile_x_1 = pl.tile.load(input_a, [(i + 1) * 128], [128]); pl.tile.store(tile_x_1, [(i + 1) * 128], output)
    tile_x_2 = pl.tile.load(input_a, [(i + 2) * 128], [128]); pl.tile.store(tile_x_2, [(i + 2) * 128], output)
    tile_x_3 = pl.tile.load(input_a, [(i + 3) * 128], [128]); pl.tile.store(tile_x_3, [(i + 3) * 128], output)

tile_x_4 = pl.tile.load(input_a, [8 * 128], [128]); pl.tile.store(tile_x_4, [8 * 128], output)
tile_x_5 = pl.tile.load(input_a, [9 * 128], [128]); pl.tile.store(tile_x_5, [9 * 128], output)

Dynamic — runtime stop n

# Before
for i in pl.pipeline(0, n, 1, stage=4):
    tile_x = pl.tile.load(input_a, [i * 128], [128])
    pl.tile.store(tile_x, [i * 128], output)

# After
unroll_main_end: pl.Scalar[pl.INDEX] = ((n - 0) // 4) * 4 + 0
for i in pl.pipeline(0, unroll_main_end, 4, stage=1):  # post-lowering marker
    <4 clones as above>

unroll_rem: pl.Scalar[pl.INDEX] = n - unroll_main_end
if unroll_rem == 1:
    tile_x_t0 = pl.tile.load(input_a, [unroll_main_end * 128], [128])
    pl.tile.store(tile_x_t0, [unroll_main_end * 128], output)
else:
    if unroll_rem == 2:
        <2 clones at offsets unroll_main_end + 0, unroll_main_end + 1>
    else:
        if unroll_rem == 3:
            <3 clones at offsets unroll_main_end + 0, +1, +2>

After this pass, CanonicalizeIOOrder runs scoped to the pipeline loop's body, clusters loads at the top and stores at the bottom, and demotes the outer loop's kind_ to Sequential — making the cloned input tiles co-live so MemoryReuse keeps them in distinct buffers. Ping-pong buffering applies to both the bulk main loop and the tail clones.

Related

• CanonicalizeIOOrder — the IO-order canonicalization pass that runs next, scoped to pipeline bodies
• UnrollLoops — full-unroll pass at slot #1, kept as the primary pl.unroll(N) lowering