Mojo module

pipeline_body

Builder for declarative pipeline body specifications.

Provides a PipelineBody context manager with role-typed methods and combinators that produce List[OpDesc] for use with the pipeline scheduling framework.

Role-typed methods (load/store/frag/compute) communicate programmer intent; actual roles are stamped by annotate_ops() via the TargetCostModel.

Example — single-buffer matmul:

with PipelineBody() as b:
    b.load(LOAD_DRAM, ch=0)
    b.store(STORE_SMEM, ch=0)
    b.barrier()
    b.fan[num_k_tiles](LOAD_FRAG, ch=0)
    b.fan[num_k_tiles](COMPUTE)
    return b.done()

Example — ping-pong half:

with PipelineBody() as b:
    b.load(LOAD_A, ch=0, stage=os, sub=1, k=k_special)
    b.load(LOAD_A, ch=0, stage=s,  sub=0, k=k_off)
    b.load(LOAD_B, ch=1, stage=s,  sub=0, k=k_off)
    b.load(LOAD_B, ch=1, stage=s,  sub=1, k=k_off)
    b.fan[2](MMA_LOAD_A, ch=0, stage=s)
    b.fan[2](MMA_LOAD_B, ch=1, stage=s)
    b.grid[2, 2](MMA)
    return b.done()

Structs

PipelineBody: Builder for declarative pipeline body specifications.

Structs​

Structs