Reading layout decisions

Reading a module for its layout decisions is not about tracing every op in order. Three signatures carry most of the information you need, and once you know to look for them, a wall of MLIR turns into a short account of how the compiler actually intends to run your kernel. One signature sits around reductions, one sits on the grid attributes themselves, and one sits inside every operand's BlockSpec. Learn all three and the module stops reading like raw output and starts reading like a decision record.

A reduction drops a dimension logically, but the tiled representation underneath cannot drop it the same way: the hardware still holds a full register grid. Look for a shape_cast right before a sum or max reduction and a broadcast right after. Together they are keepdims made physical. The cast reshapes the input into the tile the reduction expects, and the broadcast puts a size-one dimension back so downstream ops see a shape that still lines up with sublanes and lanes. Skim past these two ops and you will miss that a reduction happened at all.

Two attributes on the grid tell you how the pipeline actually runs, not how you asked it to run. dimension_semantics marks each grid axis parallel or arbitrary, and arbitrary marks the axis 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 → cannot reorder, usually because one iteration depends on state a previous one left in VMEM. window_params carries the block shape and index math for every operand, per grid step. Read the two together and you are reading the compiled schedule as data: which axes the pipeline can freely reorder, and which one it cannot, expressed as attributes rather than as prose in a design doc.

The 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 → you write in a BlockSpec does not survive as a Python closure. It compiles into a transform function attached to each operand, one that takes grid indices and returns a memory offset for that block. When a kernel misbehaves at a block boundary, this is the function to trace. It is the compiled answer to which slice of 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 → a given grid step reads, and a mismatch there usually explains the fault before you get anywhere near the compute ops. Read it alongside window_params and you have the full picture: what shape each block is, and where each grid step goes to fetch it.