Dumps on demand, and where XLA stops

Everything up to this chapter has been about reading dumps XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → already decided to write. This last piece is about producing more of them on demand, and about knowing which decisions are XLA's own versus decisions it handed off to somebody else's compiler entirely, which matters once you start comparing behavior across chips.

Every pass XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → runs can be written to disk. Setting XLA_FLAGS=--xla_dump_to=DIR before your process starts writes out every intermediate stage, from the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → XLA received through every optimization pass to the final compiled HLO, one file per stage, into DIR. For a kernel behaving strangely, this is the only way to see which specific pass changed something, rather than comparing just the two endpoints from .lower() and .compile().

The dump from a real program is large, often hundreds of files per compilation, so --xla_dump_hlo_pass_re narrows it down to passes whose name matches a regular expression. Pointed at the 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 → pass, it isolates exactly the step where your op graph collapsed into the fusion ops described earlier in this chapter, before and after, so you can watch one specific decision get made instead of reading through the entire pipeline.

The same compiler runs on GPU, and the source lives in the same open repository, but XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →:GPU is not a from-scratch compiler the way XLA:TPU is. For a large share of the work, it dispatches instead of compiling directly. Big matmuls go to cuBLAS. Attention and convolution go to cuDNN. Generated fusionsSeveral 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 → increasingly go to Triton. Collectives go to NCCL. XLA:TPU, by contrast, compiles everything itself, down through 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 → and into low-level TPU instructions, without handing any of it to a separate vendor library.

Every path through this stack ends in a box neither compiler opens itself. On TPU that box is LLOThe TPU’s near-assembly, closed inside libtpu. The readable world ends one layer above, at Mosaic.taught in /l/vliw-bundles-and-llo → and libtpu. On GPU it's ptxas turning PTX into SASS, the actual instructions the streaming multiprocessorsOne of a GPU’s many cores; each keeps many warps resident and switches among them to hide memory latency.taught in /l/tpu → run. The difference worth knowing is that PTX is a documented, stable instruction set you can read one layer deeper than the equivalent point on TPU, where the instructions inside libtpu are not public.

Knowing where XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → stops and a vendor library starts changes how you debug a slow kernel. On GPU, a slowdown inside a matmul might be a cuBLAS kernel selection issue that XLA never chose, and no amount of staring at the HLO dump shows you what happened inside cuBLAS. On TPU, everything from the 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 → decision down to the instruction stream is XLA's own compiler, so the dump is the whole story: there's no separate vendor pass between the compiled HLO and what runs on the chip, other than LLOThe TPU’s near-assembly, closed inside libtpu. The readable world ends one layer above, at Mosaic.taught in /l/vliw-bundles-and-llo → and libtpu turning it into instructions.