Manual DMA and semaphores

Everything the automatic pipeline gives you for free, staging a block into VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → before your kernel body needs it, waiting for the copy, rotating buffers so the next block arrives while the current one is still in use, is itself built from three primitives you can call directly: an ANYThe memory space that tells the compiler not to place a ref at all: a promise that the kernel will move the data itself with a manual DMA.taught in /l/pallas →-space ref that has not been staged anywhere, pltpu.make_async_copy to issue the transfer, and a semaphoreThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici →, requested from scratch_shapesRequests per-invocation workspace (VMEM buffers, semaphores) that belongs to the kernel itself, not to any input or output, and persists across grid steps.taught in /l/pallas →, that the copy signals when it completes.

You reach for this by hand when the access pattern cannot be expressed by the automatic pipeline at all: an irregular gather across a ref the BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → grammar has no vocabulary for, rather than a walk over a fixed grid where each step's block is a simple function of the grid coordinate. If your access pattern already fits the BlockSpec grammar, the automatic pipeline is doing the same overlap you would build by hand, with less code and less room for a mistake, so manual DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → is a fallback for what the grammar cannot say, not a default choice.

Issue the copy, do other work, then wait on the semaphoreThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici → right before you actually need the data. The window between issue and wait is exactly where the transfer overlaps with useful compute, the same overlap the automatic pipeline was already giving you, just under manual control now. Wait too early, right after issuing the copy instead of after doing the other work, and you have written a synchronous copy with extra steps: correct, but with none of the overlap you built this for.

This is not a separate mechanism from the one chapter 09 covers for multi-chip communication. A remote DMAA chip pushes a buffer straight into a neighbor’s memory and signals a semaphore, while its compute keeps working. The native distributed operation.taught in /l/ici → that moves data to a neighboring chip over the ICIThe inter-chip links (4.5e10 bytes per second each way per link on v5e); every collective resolves to hops over these.taught in /l/ici → link is the same make_async_copy and semaphoreThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici → pair, aimed at a different destination: one chip's HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → instead of one part of the same chip's memory hierarchy. Learn the local version here, and the distributed version is the same three primitives with one more argument telling it where to go.

Skip the wait entirely, or wait on the wrong semaphoreThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici →, and the kernel reads whatever happened to already be sitting in that memory before the copy finished landing: not a compile error, since the shapes and types all still check out, but a value that is wrong in a way that depends on timing rather than on logic. That is the price of moving a guarantee the automatic pipeline gave you for free, that a block is fully staged before your kernel body touches it, onto your own bookkeeping instead.

None of this is worth reaching for casually. Every manual DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → is one more place a kernel can be correct in interpret mode and still misbehave on real hardware, because interpret mode has no memory spaces in which to get the timing wrong. Reserve it for the access patterns that genuinely have no BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → expression, and let the automatic pipeline handle everything else.

def dma_kernel(x_hbm_ref, o_ref, scratch_ref, sem):
    copy = pltpu.make_async_copy(x_hbm_ref, scratch_ref, sem)
    copy.start()
    copy.wait()
    o_ref[...] = scratch_ref[...] * 2

def double_via_dma(x):
    return pl.pallas_call(
        dma_kernel,
        in_specs=[pl.BlockSpec(memory_space=MS.ANY)],
        out_specs=pl.BlockSpec(x.shape, lambda: (0, 0)),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        scratch_shapes=[pltpu.VMEM(x.shape, x.dtype), pltpu.SemaphoreType.DMA],
        interpret=True,
    )(x)