Layout and memory

Every shape carries a layout, and until you have seen one written out, it is easy to assume shape settles memory order by itself, the way it quietly does in plain NumPy. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → makes no such assumption. Look back at any dump from earlier chapters and the notation was already there, sitting right after the shape: f32[64,64]{1,0}. That trailing {1,0} is the layout, a minor-to-major ordering of dimensions, and it is a decision the compiler makes explicitly and separately from the shape.

Read {1,0} as an instruction: dimension 1 is minor, meaning it varies fastest in memory, and dimension 0 is major. For a 64x64 array of 4-byte floats laid out this way, stepping to the next column costs 4 bytes and stepping to the next row costs 256 bytes, the ordinary row-major layout most code assumes by habit. Flip the order to {0,1} and the strides flip with it: rows become contiguous, columns become the expensive direction. Same shape, same values, a completely different walk through memory.

Some instructions arrive with an opinion about layout already fixed. A dot or a convolution wants specific operand layouts to hit its fast path; an entry parameter's layout is whatever the caller promised. LayoutAssignment (xla/service/layout_assignment.h) is the pass that starts from these fixed points and propagates outward, deciding a layout for every remaining shape in the module by following the dataflow from instructions that demand one toward instructions that do not care.

A layout is a promise the whole module must keep.

Propagation does not always land two neighbors on the same page. When an operand's chosen layout does not match what its consumer needs, XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → cannot just reinterpret the bytes and hope; a layout is a promise the whole module must keep, so the only honest fix is to insert a transpose or a copy that actually rearranges the data into the layout the next instruction was promised. Every copy you see in a dump is a broken agreement being repaired on the spot.

Layout assignment finishes before buffer assignment ever starts, and that ordering is load-bearing: once every shape has a fixed layout, BufferAssignment (xla/service/buffer_assignment.h) can treat each value as a fixed number of bytes and ask where those bytes should live on the device, using liveness and aliasing analysis to decide which values can share a physical allocation because they are never alive at the same time.

This is where donation from the JAX path finally cashes out. Marking an argument for donation with donate_argnums does not do anything to memory by itself; it tells the compiler that once this function runs, nobody upstream still needs that input. Buffer assignment is what turns that promise into an actual input-output alias, letting the output reuse the exact allocation the donated input already occupied instead of paying for a second one.

Buffer assignment's liveness analysis is not perfect, and it does not try to be. Some interference between values only becomes visible after allocations are chosen, and resolving that interference is a separate job: copy-insertion, the pass whose handiwork you already saw in chapter 4's dump. It exists precisely because buffer assignment cannot fix the conflicts it creates on its own.

That is the honest reason this whole stretch of the compiler runs as an ordered pipeline rather than one pass doing everything at once. Layout has to be fixed before buffer assignment can count bytes meaningfully; buffer assignment's allocation choices have to exist before copy-insertion knows which conflicts still need patching. Each pass leaves work for the next one on purpose, and the order is not arbitrary.