What does pltpu.emit_pipeline actually do?

It writes a software pipeline around your kernel body. Given a grid, a set of input BlockSpecs and a set of output BlockSpecs, it allocates a small stack of VMEM buffers per ref, emits a prologue that starts the first transfers, emits a fori_loop whose body issues the copy for a future block, waits for the current block, calls your kernel, issues the output copy for the block just written and waits for the previous one, and emits an epilogue that waits for the final send. Your body never sees any of it: it is handed one block per ref, already in VMEM, and writes one block out. The point of the arrangement is that the memory transfers run underneath the arithmetic instead of beside it, so the core is not idle while a block is in flight.

How many buffers does a Pallas TPU pipeline use by default?

Two, for every input and every output. The number comes from the BlockSpec: if the spec carries a pipeline_mode, its buffer_count is used, and otherwise the default of two is applied when the buffered ref is built. There is one exception in the other direction: a ref whose BlockSpec covers the entire array has no windowing to do, so if no pipeline_mode was given it drops to a single buffer and is copied in once before the loop and out once after, synchronously, instead of participating in the schedule at all. The pipeline as a whole sizes its prologue and its lookahead from the largest buffer count across all specs, floored at two.

Can I use more than two buffers for a Pallas output?

No. The buffered ref refuses it in its constructor with a NotImplementedError saying buffer counts above two are not supported for output buffered refs. Inputs have no such cap. The reason is written as a comment on the output wait: the scheduler never records which grid indices a given output copy was issued for, and instead relies on the fact that the wait always fires on the iteration immediately after the copy, so the previous grid indices are always the correct thing to wait on. Deeper output buffering would break that assumption, and supporting it properly would mean saving the indices for each outstanding copy.

What happens if a Pallas input uses only one buffer?

It still works and it produces correct results, but it stops overlapping. With one buffer the prologue is skipped, because the prefetch loop returns early as soon as the stage index reaches the buffer count. The copy inside the loop then carries a special clause that forces it to fire on the first step regardless of the change predicate, precisely so the wait on the next line has something to wait for. From then on the copy and the wait for the same block land on the same iteration, every iteration, so the core issues the transfer and immediately blocks on it. What you get is a correct pipeline with the overlap removed, in exchange for half the VMEM.

Why does my Pallas TPU kernel not copy a block on some iterations?

Because none of the copies are unconditional. Every one is guarded by a predicate computed from your BlockSpec index maps. The wait for an input fires only when the block index changed since the previous grid position, or on the very first step. The copy fires only when the fetch position is about to change and the pipeline has not run out of blocks to fetch. The output copy fires only when the next step maps to a different block, or on the last step. So an index map that ignores a grid axis, which is the usual shape for the accumulation axis of a matmul, produces a schedule where the block is fetched once and then reused for every iteration of that axis, with no copies in between. That is the intended behaviour and it is the main reason index maps are worth getting right.

Why does an emit_pipeline schedule stop fetching near the end of the grid?

A predicate named out_of_fetch turns true once the step number reaches the total step count minus the buffer count plus one. With four steps and two buffers that is the last step, and with three buffers it is the last two. The bound exists because the loop prefetches ahead, so continuing to fetch on those iterations would read past the end of the grid. It also means the tail of a short grid is less overlapped than the middle, which is a real effect on kernels whose grid is only a handful of steps long: the deeper the buffering, the more of the grid is prologue and tail rather than steady state.

The Mosaic TPU Pipeline, Line by Line

A Pallas kernel on a TPU computes from VMEM, a scratchpad of on-chip memory small enough that the array you passed in does not fit. So every block has to be copied in from HBM before the core can touch it, and copied back out afterwards, and if the core waits for either of those the matrix unit sits idle for most of the schedule. One file decides when each copy is issued, when the core is allowed to wait for it, and how many copies of a block VMEM holds at once. This walks that file at a pinned commit of jax-ml/jax: the buffered ref and its four independent cursors, the semaphore per slot, the four boolean predicates that are the entire schedule, the prologue that primes the buffers stage by stage, the loop body whose line ordering is the only reason anything overlaps, and the epilogue that waits for the last send nobody else waited for. The signature exhibit runs the real predicates over a grid you dial and prints the resulting schedule, slot by slot.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Mosaic TPU Pipeline, Line by Line — built to be understood, not skimmed.

Questions

What does pltpu.emit_pipeline actually do?: It writes a software pipeline around your kernel body. Given a grid, a set of input BlockSpecs and a set of output BlockSpecs, it allocates a small stack of VMEM buffers per ref, emits a prologue that starts the first transfers, emits a fori_loop whose body issues the copy for a future block, waits for the current block, calls your kernel, issues the output copy for the block just written and waits for the previous one, and emits an epilogue that waits for the final send. Your body never sees any of it: it is handed one block per ref, already in VMEM, and writes one block out. The point of the arrangement is that the memory transfers run underneath the arithmetic instead of beside it, so the core is not idle while a block is in flight.
How many buffers does a Pallas TPU pipeline use by default?: Two, for every input and every output. The number comes from the BlockSpec: if the spec carries a pipeline_mode, its buffer_count is used, and otherwise the default of two is applied when the buffered ref is built. There is one exception in the other direction: a ref whose BlockSpec covers the entire array has no windowing to do, so if no pipeline_mode was given it drops to a single buffer and is copied in once before the loop and out once after, synchronously, instead of participating in the schedule at all. The pipeline as a whole sizes its prologue and its lookahead from the largest buffer count across all specs, floored at two.
Can I use more than two buffers for a Pallas output?: No. The buffered ref refuses it in its constructor with a NotImplementedError saying buffer counts above two are not supported for output buffered refs. Inputs have no such cap. The reason is written as a comment on the output wait: the scheduler never records which grid indices a given output copy was issued for, and instead relies on the fact that the wait always fires on the iteration immediately after the copy, so the previous grid indices are always the correct thing to wait on. Deeper output buffering would break that assumption, and supporting it properly would mean saving the indices for each outstanding copy.
What happens if a Pallas input uses only one buffer?: It still works and it produces correct results, but it stops overlapping. With one buffer the prologue is skipped, because the prefetch loop returns early as soon as the stage index reaches the buffer count. The copy inside the loop then carries a special clause that forces it to fire on the first step regardless of the change predicate, precisely so the wait on the next line has something to wait for. From then on the copy and the wait for the same block land on the same iteration, every iteration, so the core issues the transfer and immediately blocks on it. What you get is a correct pipeline with the overlap removed, in exchange for half the VMEM.
Why does my Pallas TPU kernel not copy a block on some iterations?: Because none of the copies are unconditional. Every one is guarded by a predicate computed from your BlockSpec index maps. The wait for an input fires only when the block index changed since the previous grid position, or on the very first step. The copy fires only when the fetch position is about to change and the pipeline has not run out of blocks to fetch. The output copy fires only when the next step maps to a different block, or on the last step. So an index map that ignores a grid axis, which is the usual shape for the accumulation axis of a matmul, produces a schedule where the block is fetched once and then reused for every iteration of that axis, with no copies in between. That is the intended behaviour and it is the main reason index maps are worth getting right.
Why does an emit_pipeline schedule stop fetching near the end of the grid?: A predicate named out_of_fetch turns true once the step number reaches the total step count minus the buffer count plus one. With four steps and two buffers that is the last step, and with three buffers it is the last two. The bound exists because the loop prefetches ahead, so continuing to fetch on those iterations would read past the end of the grid. It also means the tail of a short grid is less overlapped than the middle, which is a real effect on kernels whose grid is only a handful of steps long: the deeper the buffering, the more of the grid is prologue and tail rather than steady state.

The Mosaic TPU Pipeline, Line by Line

Questions

Related explainers