IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

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​