XLA, From the Inside
Every JAX and PyTorch/XLA program ends up here: a StableHLO module handed across PJRT, rewritten by about two hundred passes, fused, laid out, assigned to buffers, split across devices, and turned into kernels and a schedule of thunks. The track starts with a dump you can read tonight and ends at the runtime that executes it.
Phase 1 — Read a dump first
One flag writes every stage of the compile to disk. Learn the file names before the passes that produced them.
- LAB·X1 · Read your own dump ↗ — Compile a small JAX function with --xla_dump_to, then read the before and after of the passes that changed it.
- Reading an XLA Dump: What the Compiler Writes, and How to Read It — The directory the compiler writes, read file by file: the naming grammar, the two bookends, the per-pass files, and the two tools that take one of those files without a framework.
- PJRT, the boundary ↗ — The kernels-site chapter that draws the seam this track crosses first.
Phase 2 — The front door
How a framework loads a backend as a plugin, and how a StableHLO module becomes the HLO the rest of the compiler consumes.
- The PJRT Plugin Contract: pjrt_api.cc, Line by Line — pjrt_api.cc line by line: dlopen, one exported symbol, and the two-integer handshake that decides whether a plugin built elsewhere is allowed to run.
- How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line — mlir_to_hlo.cc line by line: the thirteen passes before HLO exists, and the version clamp that lets a framework talk to a plugin built three months ago.
- The PJRT Boundary: One Training Step, Crossing by Crossing — The same seam seen from a frontend: which torch_xla call becomes which PJRT method, one training step at a time.
- LAB·X5 · The mock GPU ↗ — Implement the PJRT C API surface for a fake device in plain C and watch a framework load it.
- xla/pjrt/c/pjrt_c_api.h ↗ — The C ABI itself: one struct of function pointers, a version preamble, and an extension list.
- xla/pjrt/pjrt_executable.h ↗ — What a compiled program is on the PJRT side, loaded or not.
Phase 3 — HLO and the pipeline
The data structure every pass rewrites, and the runner that walks about two hundred of them in order.
- HLO Module Anatomy: What XLA Holds While It Compiles — The data structure every pass rewrites, read before the passes: what each of the three owners holds, what a shape with a layout means, and what the verifier refuses.
- The Loop That Runs Every XLA Pass: hlo_pass_pipeline.cc, Line by Line — hlo_pass_pipeline.cc line by line: the two flags that skip passes, the invariant checkers between them, the dump hook, and the ledger every pass writes a row to.
- Inside XLA: ~200 Passes and the Fusion Decision — The whole compiler journey at design level: the sub-pipelines, the families of passes, and the fusion decision with its cost model.
- LAB·X2 · The TPU pipeline, dumped ↗ — The same dump on a TPU runtime: which passes are TPU-only and what the final module looks like before codegen.
- xla/hlo/ir/hlo_module.h ↗ — The module class: entry computation, config, schedule, and the metadata every pass reads.
- xla/service/hlo_module_config.h ↗ — The knobs a compile carries with it: replica and partition counts, debug options, the device assignment.
Phase 4 — Fusion, layout, and memory
The three decisions that decide the speed of the program, each one read from the pass that makes it.
- Layout Assignment: Where the Copy in Your HLO Dump Comes From — The pass that stamps a physical arrangement onto every shape, and the reason a copy you never wrote is sitting in your dump.
- Buffer Assignment: How Every Value Gets an Address — Where the memory number comes from: live ranges over the chosen schedule, the values that must share, best-fit packing into one temp slab, and the four dump files that say what went where.
- What Shares a GPU Kernel: priority_fusion.cc, Line by Line — priority_fusion.cc line by line: a priority that is a length of time, an ordered map used as a queue, and the thirteen sentences the pass writes down when it refuses a merge.
- Instruction Fusion Legality: instruction_fusion.cc, Line by Line — The base fusion pass line by line: which producer may fuse into which consumer, and every string it uses to say no.
- LAB·X4 · Two loose ends ↗ — Two things the dump shows and the prose skips: a layout copy and a buffer you did not expect, tracked to their passes.
- xla/service/gpu/fusion_pipeline.cc ↗ — The GPU fusion sub-pipeline in 98 lines: the order the fusion passes run in.
- xla/service/gpu/model/gpu_performance_model_base.h ↗ — The cost model priority fusion asks: kernel launch overhead and the bandwidth arithmetic.
- xla/service/copy_insertion.cc ↗ — Where copies are added so that in-place updates and aliasing stay correct.
- Fusion, the chapter ↗ — The kernels-site chapter with the fusion taxonomy and the dumps to read alongside.
Phase 5 — Many devices
One annotated program becomes N per-device programs, and the collectives between them are inserted by a pass.
- The SPMD Partitioner: One Program Across a Device Mesh — The partitioner at design level: sharding propagation, the collectives it inserts, and what the halo exchange costs.
- XLA Collectives: One Guest List, Then One Order — The guest list twice over: the four readings of replica_groups and the switch that turns one into device numbers, then the clique the runtime opens for it and the three places order is made deterministic.
- LAB·X3 · Collectives in the lowering ↗ — Shard a matmul three ways on fake devices and read the all-gather and reduce-scatter the partitioner inserted.
- xla/service/spmd/spmd_partitioner.cc ↗ — The partitioner itself, seven thousand lines; the pages above tell you which functions to open.
- Collectives, the chapter ↗ — The kernels-site chapter on the fabric and the collective as a kernel.
Phase 6 — Backends and the runtime
What the compiled module becomes on a CPU and on a GPU, and the executor that runs the schedule.
- The CPU Thunk Executor: thunk_executor.cc, Line by Line — thunk_executor.cc line by line: how declared buffer use becomes a dependency graph, and how a ready queue and one atomic counter per node run it without a GPU stream.
- The GPU Codegen Path: From a Fused HLO to a Kernel — The two decisions that turn a fused HLO into a kernel: which emitter claims it, and which config the autotuner measures its way to.
- xla/backends/cpu/runtime/thunk.h ↗ — The thunk interface: what every unit of CPU execution promises.
- xla/service/gpu/gpu_compiler.cc ↗ — The GPU compiler driver: every stage name, in order, with its AddPass calls.
- xla/stream_executor/stream_executor.h ↗ — The single-device hardware abstraction under every backend: streams, kernels, memory.
- Codegen, the chapter ↗ — The kernels-site chapter that follows one fusion to PTX and SASS.
- Autotuning, the chapter ↗ — How the autotuner measures variants and what its cache file contains.
Phase 7 — Above the compiler
IFRT wraps per-device buffers into one array so a controller can drive thousands of devices; Pathways is what that buys.
- IFRT Arrays: pjrt_array.cc, Line by Line — pjrt_array.cc line by line: an array is a sharding plus per-device buffers matched by index, and every method has to keep that pairing.
- xla/pjrt/cpu/cpu_client.cc ↗ — The in-process CPU client: the PJRT implementation you can single-step on a laptop.
- IFRT, the chapter ↗ — Why a second interface sits above PJRT and what the proxy adds.
- Pathways, the chapter ↗ — Single-controller execution over thousands of chips, and what it asks of IFRT.