Mosaic → LLO

Pallas does not execute your kernel; it lowers it to Mosaic, an MLIR dialect for TPU kernels. Mosaic maps your block operations onto the chip's real units: matmuls onto the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu →'s systolic array, elementwise work onto the VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu →'s lanes, and everything onto the (8, 128) register tiling. Below Mosaic sits LLOThe TPU’s near-assembly, closed inside libtpu. The readable world ends one layer above, at Mosaic.taught in /l/vliw-bundles-and-llo →, the TPU's near-assembly, where operations become VLIW bundles the sequencer issues.

The two halves have different visibility, and knowing where the line falls matters. Mosaic's front half is open: the tpu dialect and much of its lowering live in the jax repo under jaxlib/mosaic, and any kernel prints its own module with pallas_call(..., debug=True). LLOThe TPU’s near-assembly, closed inside libtpu. The readable world ends one layer above, at Mosaic.taught in /l/vliw-bundles-and-llo → is not: it ships closed inside libtpu, and the public pipeline ends where Mosaic hands the module over. So the honest division of labor is: you write Pallas, you read Mosaic, and you never see LLO.

Reading Mosaic is worth the hour it takes, because everything you wrote is still recognizable, just compiled. Your 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 → index mapsThe 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 → become literal functions in the module. Your pl.when becomes a real scf.if with both branches present. Your jnp.dot becomes tpu.matmul with the dimension numbers carried down intact. The GYM·05 instrument holds three track kernels open at this layer with hover sync, the same way GYM·04 does for the layers above.

You will also read Mosaic's complaints. When a block shape violates the lattice, when a dynamic slice defeats vectorization, when 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 → overflows: the error surfaces from this layer, in this layer's vocabulary of layouts and vector shapes. The reason this track teaches the machine (stage 0) before the kernel language (stage 1) is exactly so those errors read as information rather than noise.

%6 = tpu.matmul %4, %5, %cst {dimension_numbers = #tpu.dot_dimension_numbers<[1], [0], [0], [1], [0, 0, 1, 1], [], []>} : vector<256x256xbf16>, vector<256x256xbf16>, vector<256x256xf32> -> vector<256x256xf32>
%7 = arith.truncf %6 : vector<256x256xf32> to vector<256x256xbf16>

// later in the same module: lambda i, j, kk: (kk, j), compiled
func.func @transform_1(%arg0: i32, %arg1: i32, %arg2: i32) -> (i32, i32) {
  %c0_i32 = arith.constant 0 : i32
  return %arg2, %arg1 : i32, i32
}

One thing this layer explains: why TPU kernel scheduling is tractable at all. TPU grid steps are a sequential pipeline, not thousands of independent thread blocks, so questions like "will this transfer hide under this compute" have answers you can reason about from constants. The whole rooflineThe floor model: latency is at least the larger of FLOPs over peak compute and bytes over bandwidth. Predict first, measure second.taught in /l/tpu → discipline of stage 0 works because the hardware below is this predictable.