the path · 0/15
start the path

the kernel path · Mosaic → LLO · lesson 02 of 3

The tiling vocabulary, precisely

A vector register is a physical grid of 8 sublanes by 128 lanes, and every vector type in a lowered module is that fact showing through.

the goal Read the trailing tile pair on any vector type and derive it from the dtype: (8, 128) for f32, (16, 128) for bf16, (32, 128) for int8.

mastery work · this chapter0/2
  2. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
§ 01

The tiling vocabulary, precisely

A vector register on a TPU holds a fixed grid: 8 sublanes by 128 lanes. That grid is physical, not a convention, and it is the reason every vector<...> type printed in a lowered module has the shape it does. Sublanes run down the register, lanes run across it, and every value in a compiled kernel gets tiled into that grid before any op touches it. This section names the rule precisely: what the tile is, how dtype changes it, and where it shows up when a kernel breaks.

the boundary, drawn exactly: everything left of the line prints with debug=True; right of it is closed inside libtpu
pallas_call your kernel body kernel jaxpr traced over refs Mosaic module tpu dialect · openin jaxlib/mosaic debug=True prints both the last readable layer LLO libtpu · closed

Change the dtype and the tile changes with it. Register width in bytes stays fixed, so narrower elements pack more rows into the same 8x128 footprint. f32 keeps the base tile at (8, 128). bf16 packs two elements per lane slot and doubles the sublane count to (16, 128). int8 packs four and lands at (32, 128). Rows and lanes are not free parameters you choose. They are consequences of how many bytes each element takes.

Open a lowered module and this rule stops being abstract. Every vector<...> type carries a tile-shaped trailing dimension pair that matches whichever dtype produced it, 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 → checks that pairing on every op. Break it, say by bitcasting an f32 value into an int8 layout without adjusting the tile, and the compiler does not fail quietly. It refuses the rewrite and quotes the exact packing rule back at you in the error text.

The captured module keeps both shapes side by side: the logical array shape you wrote in JAX, and the physical vector type 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 → produced from it. Once you know the packing rule, that pairing stops looking arbitrary. A bf16[128, 256] operand does not lower into a register grid that matches its logical axes one for one. It lowers into a tile count built from (16, 128). The 128-lane width itself is not incidental: it matches the 128x128 systolic array on v5e, the same MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → that consumes these tiles once they are packed.

before you move on

Check yourself

01 Why does bf16 tile as (16, 128) when f32 tiles as (8, 128)?

Register width in bytes is fixed, so narrower elements pack more rows into the same footprint: half-width bf16 doubles the sublane count.

02 What two shapes does the captured module keep side by side, and what does the pairing teach?

The logical array shape you wrote in JAX and the physical vector type Mosaic derived from it. With the packing rule known, the pairing reads as a derivation rather than noise.

assigned

Readings