From tracing to machine code

Tracing a kernel body against refs produces a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → like any other, with two differences you can see at a glance: the binders are memrefs rather than arrays, and loads and stores appear as their own terms. The path had you print one in LAB·1.4 and read it in chapter 07. The claim worth adding here is about what is absent from it.

Nothing in that jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → mentions a TPU. No tile shape, no memory space, no pipeline. It is a body plus its memory traffic, stated in the same primitives an ordinary JAX program uses, which is what makes the same source reachable by three different compilers.

The jaxpr names no chip. Every backend decision happens below it.

MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → consumes mostly standard-dialect MLIR and emits LLOThe TPU’s near-assembly, closed inside libtpu. The readable world ends one layer above, at Mosaic.taught in /l/vliw-bundles-and-llo → to be compiled for the TPU. Pallas gets there by translating JAX primitives into MLIR, mostly the vector and arith dialects, and by converting the BlockSpecs into the pipeline schedules the grid lesson at /s/pallas/grid-and-pipeline read out of the module's attributes.

Read one of this repo's captures and you find two vocabularies side by side. The algebra is standard: a row max becomes vector.multi_reduction, the exponential becomes math.exp, the divide becomes arith.divf, the bf16 round-trip becomes arith.extf and arith.truncf. Only what is specific to this machine wears the tpu prefix, tpu.matmul and tpu.vector_store among them.

The split is worth internalizing because it tells you where portability ends. Anything expressed in vector and arith is describing computation any vector machine could do. Anything in the tpu dialect is naming a unit that exists on this chip.

%1 = arith.extf %0 : vector<256x512xbf16> to vector<256x512xf32>
%cst = arith.constant dense<0xFF800000> : vector<256xf32>
%2 = vector.multi_reduction <maximumf>, %1, %cst [1] : vector<256x512xf32> to vector<256xf32>
%3 = vector.shape_cast %2 : vector<256xf32> to vector<256x1xf32>
%4 = vector.broadcast %3 : vector<256x1xf32> to vector<256x512xf32>
%5 = arith.subf %1, %4 : vector<256x512xf32>
%6 = math.exp %5 : vector<256x512xf32>

module @softmax_kernel {
  func.func @main(%arg0: memref<256x512xbf16, #tpu.memory_space<vmem>>, %arg1: memref<256x512xbf16, #tpu.memory_space<vmem>>) attributes {dimension_semantics = [], scalar_prefetch = 0 : i64, scratch_operands = 0 : i64} {
    %c0 = arith.constant 0 : index
    %c0_0 = arith.constant 0 : index
    %0 = vector.load %arg0[%c0, %c0_0] : memref<256x512xbf16, #tpu.memory_space<vmem>>, vector<256x512xbf16>
    %1 = arith.extf %0 : vector<256x512xbf16> to vector<256x512xf32>
    %cst = arith.constant dense<0xFF800000> : vector<256xf32>
    %2 = vector.multi_reduction <maximumf>, %1, %cst [1] : vector<256x512xf32> to vector<256xf32>
    %3 = vector.shape_cast %2 : vector<256xf32> to vector<256x1xf32>
    %4 = vector.broadcast %3 : vector<256x1xf32> to vector<256x512xf32>
    %5 = arith.subf %1, %4 : vector<256x512xf32>
    %6 = math.exp %5 : vector<256x512xf32>
    %cst_1 = arith.constant dense<0.000000e+00> : vector<256xf32>
    %7 = vector.multi_reduction <add>, %6, %cst_1 [1] : vector<256x512xf32> to vector<256xf32>
    %8 = vector.shape_cast %7 : vector<256xf32> to vector<256x1xf32>
    %9 = vector.broadcast %8 : vector<256x1xf32> to vector<256x512xf32>
    %10 = arith.divf %6, %9 : vector<256x512xf32>
    %11 = arith.truncf %10 : vector<256x512xf32> to vector<256x512xbf16>
    %c0_2 = arith.constant 0 : index
    %c0_3 = arith.constant 0 : index
    %12 = vector.load %arg1[%c0_2, %c0_3] : memref<256x512xbf16, #tpu.memory_space<vmem>>, vector<256x512xbf16>
    tpu.vector_store %arg1[%c0_2, %c0_3], %11 {strides = array<i32>} : memref<256x512xbf16, #tpu.memory_space<vmem>>, vector<256x512xbf16>,
    return
  }
}

Pallas was designed with Triton as a target, so the original GPU lowering was straightforward: JAX dot products became Triton dot products, unary primitives became their Triton equivalents, and Triton atomics arrived through new Pallas primitives. The GPU path now goes to MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → GPU instead, and the design document carries a footnote saying the Pallas-to-Triton lowering path is officially deprecated and is discussed for historical reasons.

Read the deprecated section anyway, for one paragraph in it. Triton has no notion of a BlockSpec and addresses memory with pointers rather than indices, so lowering x_ref[3, 2] on a (4, 5) ref meant computing the row-major pointer by hand, 5 * 3 + 2 * 1, and lowering a slice like x_ref[4, :] meant producing a whole array of pointers.

That arithmetic is what the block index abstraction removed. The blockspec lesson spent its first section on the fact that an index mapThe BlockSpec function that returns block coordinates (not element offsets) for each grid step; Pallas multiplies by the block shape to find elements.taught in /l/pallas → returns block coordinates and Pallas does the multiplication; this is the layer where you can see who used to do it and what they had to say instead.

Because a kernel is JAX primitives plus a few Pallas ones, a pallas_call can be lowered to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → directly, implemented as a lax.scan over the grid, and compiled by XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → like any other program. That is the whole of interpret mode. It is not a simulator written for Pallas; it is your kernel expressed as a loop that XLA already knows how to run, on any supported platform including CPU.

Three consequences follow, and the design document states all three. Ordinary debugging works, jax.debug.print included. The numerics come from XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →, which the document calls more reliable and better tested, and which is used to verify the Triton and MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → compilers. And the ordering is the scan's, which the document notes could in principle be perturbed to simulate the parallel reads and writes a GPU performs.

What a scan cannot have is a memory hierarchy. 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 →, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → and 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 → collapse into whatever plain JAX does, so a lattice violation and a VMEM overflow, the two exhibits the museum keeps at /mistakes/kernels, are invisible here by construction. A kernel that passes in interpret mode and fails to lower on hardware almost always failed the second question, not the first.

Batching a kernel from the outside has a natural default. vmap of a pallas_call augments the call with an extra grid dimension for the new batch axis and rewrites the BlockSpecs to index along it, and jax.custom_vmap is there for when that default is not the batched kernel you wanted.

Differentiation is where the layering strains, and the design document is specific about why rather than vague. jax.grad decomposes into jvp, partial_eval and transpose, and most of JAX's existing machinery applies. Then the honest sentence: automatic differentiation of kernels can result in a performance hit due to how memory access is transposed. A kernel with overlapping-and-parallel reads and disjoint-but-parallel writes transposes into one with overlapping-but-parallel writes, which are slow when done atomically, and disjoint-and-parallel reads.

The document names the missing capability too, which is more useful than the warning alone: emitting a good transposed kernel would mean reordering loops and changing the vectorization, and Pallas has no program representation amenable to that. Elementwise kernels transpose fine. For everything else the recommendation is jax.custom_vjp, which is exactly what stage 3 does when it wires a hand-derived backward onto flash attention at /s/kernels.

Two more transformations sit on the document's speculative list rather than in the shipping surface. checkify could plumb error codes out of a kernel for out-of-bounds access or NaNs, and custom_partitioning could make a kernel automatically partitionable. Both are written as things one could imagine, and reading them that way keeps the map honest.