Layouts

Layout is the part of a compiled program that has nothing to do with what you computed and everything to do with how it sits in memory, and the compiled dump is the only place it's visible. Get comfortable reading layout annotations early, because a good share of the copies you'll end up chasing down for performance turn out to be layout disagreements rather than anything about the math itself.

Every tensor in the compiled dump carries a layout annotation like {1,0}, and it's worth learning to read on sight. The numbers list dimensions from fastest varying to slowest varying, minor to major. {1,0} on a rank-2 tensor means dimension 1, the columns, is minor and laid out contiguously in memory, while dimension 0, the rows, is major, stepping between rows. {0,1} flips that, making rows the contiguous dimension instead. The same logical shape can carry either layout, and the compiler picks based on what's cheapest for the ops sitting around it.

When a tensor's layout has to change between one op and the next, whether because two ops disagree on what's cheap or because a shape crosses a fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → boundary, XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → inserts a copy op to physically rearrange the data. Copies cost a full pass over the tensor, so a compiled dump full of copies is a compiled dump full of wasted 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 → bandwidth. Scanning the dump for copy ops is a fast way to find where the compiler disagreed with itself, or where your code forced a rearrangement it never needed.

This is also why a transpose is sometimes free and sometimes not. If the op consuming the transposed tensor can read it with the new layout annotation, treating a different dimension as minor, the transpose costs nothing: it disappears at compile time into a layout choice. If the consumer needs a specific physical layout it can't bend on, a matmul feeding the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → expects its operands in a fixed orientation, the transpose becomes a real copy. The op graph looks identical either way. Only the compiled dump tells you which case you're actually in.

Layout interacts directly with how a dtype gets tiled into registers on the chip. f32 tiles as (8,128), bf16 as (16,128), int8 as (32,128), and the minor-to-major ordering in the layout annotation decides which of a tensor's dimensions actually line up with that tile shape. Get the layout wrong relative to the tile and the compiler inserts a copy to fix it before the op that needs the aligned shape can run.