How does Pallas lower a kernel to a TPU?

Through one registered MLIR lowering rule for pallas_call on the tpu platform, and that rule does three things. It builds a Mosaic module by walking the kernel jaxpr equation by equation, it optionally prints the jaxpr and the module when you pass debug=True, and it hands the module to a custom call. The walk itself is a dictionary lookup per equation: a rule table keyed first by core type and then by primitive, with the rule called on the already-lowered operand values. Everything the kernel body does becomes ops in the tpu, vector, arith, math, scf and memref dialects, and the surrounding grid becomes function arguments plus window parameters that a later pipeline emitter reads. The module never becomes HLO. It is serialised into one HLO custom call whose target is the string tpu_custom_call, and the TPU compiler picks it up from there.

Which JAX primitives can I use inside a Pallas TPU kernel?

At this commit, 135 primitives have a Mosaic TPU lowering rule, and they are not all available on every core. Registration takes a set of core types, defaulting to the tensor core alone. Seventy-one of the 135 are tensor-core only, which includes dot_general, transpose, iota, gather, top_k, the DMA start and wait pair and the whole random-number surface. Fifty-nine are registered on all three core types, tensor core plus both sparse-core subcores, and those are the ones you would guess: add, multiply, divide, the comparisons, select, the reductions over booleans, control flow, and the semaphore operations. Five sit in between, on the tensor core and the sparse-core vector subcore only: sqrt, rsqrt, tanh, log and round. Anything with no rule for the core you are compiling for raises NotImplementedError naming the primitive and the core type.

What does "Unimplemented primitive in Pallas TPU lowering" mean?

It means the equation walker reached a jaxpr equation whose primitive is not a key in the rule table for the core type you are compiling for, and there is nothing left to try. The message names both the primitive and the core type, and it asks you to file an issue, because the fix is a new rule rather than anything you can do at the call site. In practice you usually get there by calling a jax.numpy function that expands into a primitive nobody has needed on a TPU yet, so the useful first move is to print the kernel jaxpr with debug=True and read which primitive the equation actually is; the name in the error is the primitive name, not the jax.numpy function you wrote. It is worth separating this from the other failure that looks similar: a rule that exists but refuses your dtypes raises from inside the rule, and with JAX_PALLAS_VERBOSE_ERRORS set that exception is re-raised with the failing jaxpr equation attached.

How do I make a Pallas ref live in VMEM instead of HBM?

On a TPU tensor core you already do, unless you asked otherwise. A BlockSpec has a memory_space field that defaults to None, and the code that canonicalises a BlockSpec into a block mapping replaces that None with the logical space DEFAULT. The TPU backend then maps DEFAULT to VMEM for the tensor core and the sparse-core vector subcore, and to SMEM for the sparse-core scalar subcore, and the resulting enum is formatted straight into an MLIR attribute, so the memref type your kernel argument gets carries the literal text vmem. Asking for pl.ANY takes the other branch: it stays ANY, the memref carries any, and the lowering emits empty window parameters for that operand, which means no copy is generated for you and the block shape must equal the array shape with a trivial index map. That is the trade: a windowed VMEM operand is moved for you, an ANY operand stays where it is and your kernel body has to DMA it.

Why does dot_general fail on unsigned integers in a Pallas TPU kernel?

Because the type converter throws signedness away and the matmul rule refuses rather than compute the wrong answer. Every integer dtype becomes a signless MLIR integer of the same width on the way in, with a comment in the source saying the conversion makes Mosaic interpret all unsigned types as signed. For most operations that is harmless or handled: division and remainder branch on the JAX-level dtype and emit the unsigned instruction, and the comparisons pick an unsigned predicate for unsigned operands. Matrix multiply has no such branch, so the rule checks its input avals first and raises NotImplementedError for any unsigned one, telling you that dot_general interprets all integer inputs as signed and suggesting you cast to a signed type before the dot. So unsigned arithmetic inside a kernel is fine and an unsigned matmul is not, which is a surprising place to meet the limit.

Why does Pallas require TPU block shapes divisible by 8 and 128?

Because those are the sublane and lane counts of the vector unit, and the check runs before any lowering happens. A separate validation pass walks every block mapping and, for a block of rank two or more, requires the last dimension to equal the array dimension or be a multiple of 128, and the second-to-last to equal the array dimension or be a multiple of 8. Rank-one blocks get a different rule with three escape hatches: equal to the array dimension, a multiple of sublanes times lanes, or a power of two at least as large as the lane count times the packing factor for the dtype. Blocks in SMEM with a trivial window are skipped entirely, and so are semaphores. The error names the block shape, the array shape, the index map and the memory space, and links the BlockSpec documentation, which makes it one of the more legible failures in the stack.

Pallas on a TPU: From a Jaxpr to the Mosaic Dialect

A Pallas kernel is a Python function, but by the time a TPU sees it the Python is long gone: what the backend receives is a jaxpr, and what it emits is an MLIR module in the Mosaic dialect wrapped inside one HLO custom call. This reads that translation at a pinned commit of jax-ml/jax. The type mapping first, where a ref becomes a memref carrying a memory-space attribute and a value becomes a vector, and the four-step chain that turns a BlockSpec with no memory space into the literal string vmem. Then the window parameters that decide whether a block is copied in for you or left in HBM for your own DMA, and the two allocators behind run_scoped, one calling alloca and one calling a semaphore allocator. Then the table: 135 primitives carry a TPU rule at this commit, computed from the real registrations and filterable by core type and by the MLIR dialect each rule builds into. Then the two doors to the MXU, one of them a single tpu.matmul and the other five primitives that name an accumulator and an MXU index. It closes on refusal: the dtypes the lowering will not take, the block shapes it rejects before lowering starts, and the one message that means no rule exists at all.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Pallas on a TPU: From a Jaxpr to the Mosaic Dialect — built to be understood, not skimmed.

Questions

How does Pallas lower a kernel to a TPU?: Through one registered MLIR lowering rule for pallas_call on the tpu platform, and that rule does three things. It builds a Mosaic module by walking the kernel jaxpr equation by equation, it optionally prints the jaxpr and the module when you pass debug=True, and it hands the module to a custom call. The walk itself is a dictionary lookup per equation: a rule table keyed first by core type and then by primitive, with the rule called on the already-lowered operand values. Everything the kernel body does becomes ops in the tpu, vector, arith, math, scf and memref dialects, and the surrounding grid becomes function arguments plus window parameters that a later pipeline emitter reads. The module never becomes HLO. It is serialised into one HLO custom call whose target is the string tpu_custom_call, and the TPU compiler picks it up from there.
Which JAX primitives can I use inside a Pallas TPU kernel?: At this commit, 135 primitives have a Mosaic TPU lowering rule, and they are not all available on every core. Registration takes a set of core types, defaulting to the tensor core alone. Seventy-one of the 135 are tensor-core only, which includes dot_general, transpose, iota, gather, top_k, the DMA start and wait pair and the whole random-number surface. Fifty-nine are registered on all three core types, tensor core plus both sparse-core subcores, and those are the ones you would guess: add, multiply, divide, the comparisons, select, the reductions over booleans, control flow, and the semaphore operations. Five sit in between, on the tensor core and the sparse-core vector subcore only: sqrt, rsqrt, tanh, log and round. Anything with no rule for the core you are compiling for raises NotImplementedError naming the primitive and the core type.
What does "Unimplemented primitive in Pallas TPU lowering" mean?: It means the equation walker reached a jaxpr equation whose primitive is not a key in the rule table for the core type you are compiling for, and there is nothing left to try. The message names both the primitive and the core type, and it asks you to file an issue, because the fix is a new rule rather than anything you can do at the call site. In practice you usually get there by calling a jax.numpy function that expands into a primitive nobody has needed on a TPU yet, so the useful first move is to print the kernel jaxpr with debug=True and read which primitive the equation actually is; the name in the error is the primitive name, not the jax.numpy function you wrote. It is worth separating this from the other failure that looks similar: a rule that exists but refuses your dtypes raises from inside the rule, and with JAX_PALLAS_VERBOSE_ERRORS set that exception is re-raised with the failing jaxpr equation attached.
How do I make a Pallas ref live in VMEM instead of HBM?: On a TPU tensor core you already do, unless you asked otherwise. A BlockSpec has a memory_space field that defaults to None, and the code that canonicalises a BlockSpec into a block mapping replaces that None with the logical space DEFAULT. The TPU backend then maps DEFAULT to VMEM for the tensor core and the sparse-core vector subcore, and to SMEM for the sparse-core scalar subcore, and the resulting enum is formatted straight into an MLIR attribute, so the memref type your kernel argument gets carries the literal text vmem. Asking for pl.ANY takes the other branch: it stays ANY, the memref carries any, and the lowering emits empty window parameters for that operand, which means no copy is generated for you and the block shape must equal the array shape with a trivial index map. That is the trade: a windowed VMEM operand is moved for you, an ANY operand stays where it is and your kernel body has to DMA it.
Why does dot_general fail on unsigned integers in a Pallas TPU kernel?: Because the type converter throws signedness away and the matmul rule refuses rather than compute the wrong answer. Every integer dtype becomes a signless MLIR integer of the same width on the way in, with a comment in the source saying the conversion makes Mosaic interpret all unsigned types as signed. For most operations that is harmless or handled: division and remainder branch on the JAX-level dtype and emit the unsigned instruction, and the comparisons pick an unsigned predicate for unsigned operands. Matrix multiply has no such branch, so the rule checks its input avals first and raises NotImplementedError for any unsigned one, telling you that dot_general interprets all integer inputs as signed and suggesting you cast to a signed type before the dot. So unsigned arithmetic inside a kernel is fine and an unsigned matmul is not, which is a surprising place to meet the limit.
Why does Pallas require TPU block shapes divisible by 8 and 128?: Because those are the sublane and lane counts of the vector unit, and the check runs before any lowering happens. A separate validation pass walks every block mapping and, for a block of rank two or more, requires the last dimension to equal the array dimension or be a multiple of 128, and the second-to-last to equal the array dimension or be a multiple of 8. Rank-one blocks get a different rule with three escape hatches: equal to the array dimension, a multiple of sublanes times lanes, or a power of two at least as large as the lane count times the packing factor for the dtype. Blocks in SMEM with a trivial window are skipped entirely, and so are semaphores. The error names the block shape, the array shape, the index map and the memory space, and links the BlockSpec documentation, which makes it one of the more legible failures in the stack.

Pallas on a TPU: From a Jaxpr to the Mosaic Dialect

Questions

Related explainers