XLA

The compiler and runtime under JAX and PyTorch/XLA: HLO, the pass pipeline, fusion, layout, buffers, SPMD, and the backends.

Explainers

• The Loop That Runs Every XLA Pass: hlo_pass_pipeline.cc, Line by Line — A compiler pass is small and well behaved: it takes a program, makes one kind of improvement, and says whether it found anything. Running a few hundred of them back to back is a different job, and this is the 334-line C++ file that does it, read end to end. The two flags that skip passes, their four-shape grammar, and the two occurrence counters they resolve against. The invariant checkers between passes, why they only run after a pass that changed something, and the single rule that makes them safe. The debug mode that hashes the whole module twice to catch a pass that lied about its own return value. The dump condition with a special case hiding inside it, where the catch-all pattern writes fewer files than a narrow one. And the written ledger every pass appends a row to, opened before the skip decision so its numbering survives a bisect. The signature exhibit is the loop itself, ported line for line over a seven-pass pipeline, so you can watch a flag reshape a run.
• The PJRT Plugin Contract: pjrt_api.cc, Line by Line — A vendor ships a shared library for an accelerator the framework has never been compiled against, and it runs your model anyway. The entire arrangement is one exported C symbol and two integers, and this is the 209-line file where the two sides meet. We read all of it: the global map from a device-type string to a function table, the lowercase that is the whole matching rule, the write-once registration that refuses to let two packages claim one name, and the loader itself, which dlopens a library, asks for a single symbol called GetPjrtApi, and never closes the file again because the table lives inside it. Then the handshake. Where the plugin gets its version number from, why major must match exactly and minor only has a floor, the environment variable that switches the rule from generous to strict, and the value that reads as encouraging while doing the opposite. The signature exhibit runs the real comparison order: set the version the plugin was built against, set the environment variable, and read back the exact error string absl::StrCat assembles, trailing-parenthesis quirk included.
• HLO Module Anatomy: What XLA Holds While It Compiles — An HLO dump is an object graph printed. This walks that graph in the XLA source at a pinned commit: an HloModule that owns a vector of computations and points at one of them as the entry, an HloComputation whose parameters and return value are both just instructions in its own body, and an HloInstruction that owns nothing beyond an opcode, a shape and a list of pointers out to operands and back from users. Then the two things that sit beside the graph and decide most of what it costs: the shape that carries a layout, where minor_to_major turns an index into an address and a disagreement turns into a copy you never wrote; and the schedule, an optional total order added late to hold peak memory down. It closes on the verifier that runs between all two hundred passes, in the order it runs its checks, and the refusals worth knowing by name.
• Reading an XLA Dump: What the Compiler Writes, and How to Read It — Someone asks you to attach the HLO dump. You set one flag, point a directory at it, and get back a pile of files with names like module_0000.jit_step.0007.simplification.after_algsimp.before_reshape-mover.txt. This reads that directory end to end at a pinned commit of openxla/xla: the four inference rules one flag triggers before anything is written, the filename grammar and what every segment of it means, the two named bookends of a compile and the four buffer-assignment reports that ride with the second one, the per-pass files and the rule that quietly drops any pass that changed nothing, and the flags that make the dump larger, smaller, or unreadable. Then the two tools that take a dumped file and do something with it without a framework anywhere in sight: hlo-opt, which re-runs named passes and prints any stage of the compile, and run_hlo_module, which executes the module and checks it against the interpreter. The signature exhibit assembles the actual directory listing from the real inference rules, one flag at a time.
• How a StableHLO Module Becomes HLO: mlir_to_hlo.cc, Line by Line — Your framework does not hand the compiler a graph. It hands it a versioned document written in a dialect called StableHLO, and 434 lines of C++ decide what that document is allowed to say. We read all of them. The includes, which split into the two vocabularies this file sits between. The thirteen MLIR passes a module walks through before anything called HLO exists: seven from a Shardy pipeline it pulls in unconditionally, six added here by name, ending with the one that copies captured constants into control-flow regions because XLA has no implicit capture. The Shardy-loses-to-GSPMD fallback at the top of the function, and the block sixty lines later that silently reads the flag it just switched off. The parser, and the one error message that names three different version problems because it genuinely cannot tell them apart. Then the write side, which is the longer half: the version clamp that keeps the lower of what a plugin asked for and what this build can produce, the allow-list that refuses to serialize any operation nobody has promised to read, and the twelve-week compatibility window written as a requirement rather than a number. The signature exhibit runs the real branches: pick a direction, flip the flags, and every pass your module walks through appears in order, at the line that adds it.
• Layout Assignment: Where the Copy in Your HLO Dump Comes From — A copy instruction shows up in your dump between two lines you wrote next to each other, and the only thing that changed across it is the little list in braces. This reads the pass that put it there, at a pinned commit of openxla/xla. What a layout is, which is one permutation of dimension numbers read from the fastest-moving end. The three kinds of constraint the pass collects, and why exactly one of them can be settled by inserting a copy. The switch statement that answers, for all 134 opcodes, whether an operation may hold one arrangement out and a different one in, and the 89 that may not, which is what lets a constraint travel across most of a graph for free. The deque with two ends, the three propagators hanging off it, and the cap of two rewrites that stops it oscillating. The derivation at a transpose, whose stated goal is to make the transpose move no bytes at all, and the copy that lands one edge above it as a result. Then the priority arithmetic that settles a collision, the row-major default handed to whatever nobody argued about, the tuple rule that sends a constraint further upstream instead, and the moment inside AssignLayouts where the copy is finally created. The signature exhibit runs the real propagation loop over a four-instruction module, one pop at a time, with every layout on screen produced by the same arithmetic the pass uses.
• Buffer Assignment: How Every Value Gets an Address — A compiled XLA program never allocates. It receives a handful of slabs from the runtime and every kernel was compiled with its offsets inside those slabs already baked in. This walks the pass that picks them, at a pinned commit of openxla/xla: the values that must share memory because the language says so, the schedule that turns a dependency graph into a clock so that a live range can be two integers, the closed-interval interference test and the one endpoint case it makes an exception for, the colour that is a memory space rather than a graph colouring, the four kinds of value that never reach the packing, and the biggest-first best-fit heap that packs everything left. Then the nine ordered reasons an allocation refuses a buffer, the in-place update whose copy you never see until something reads the old value, and the four files a dump writes with the answer in them. The signature exhibit packs a real schedule with the real rules and lets you drag a live range until the sharing breaks.
• What Shares a GPU Kernel: priority_fusion.cc, Line by Line — A GPU kernel launch costs about a microsecond of nothing happening, plus a round trip through memory for its inputs and another for its outputs. So the largest single decision an ML compiler makes is which operations get to share a kernel, and on the GPU backend that decision is 1,419 lines of C++ we read end to end. The type choice the whole pass turns on: a priority is an absl::Duration, the wall-clock time this merge is estimated to save, which is what lets an ordered map double as a priority queue and two infinities put bitcasts first and constants last with no special case in the ordering. The correction to the usual summary: the queue holds producers, not edges, and a producer is scored against every one of its consumers at once, so the cost of being duplicated into three kernels is always in the number. The incremental machinery that keeps the pass from being quadratic, where a re-score is the old score plus the delta from the new consumers minus the runtimes of the departed ones. The Triton path that is tried first and the elemental checks that run when it declines. The thirteen refusal strings, quoted verbatim, that are what you actually read in a fusion dump. And the three flags that let you watch it decide on your own model, including the compiler fuel that bisects to the exact merge that changed your numbers. The signature exhibit runs the real ordering rules over a small graph: press pop, watch a producer that scored below zero turn profitable because a neighbour was absorbed somewhere else.
• Instruction Fusion Legality: instruction_fusion.cc, Line by Line — Merging two operations into one kernel deletes a write and a read of a whole tensor, which on this hardware is the difference that matters. So the compiler wants to fuse everything, and this 1,248-line C++ file is where it works out what it may. The unit of decision is one graph edge: a producer, a consumer that takes it as an operand, and one boolean. Read end to end: the hand-written classification of all 134 opcodes into cheap and expensive, in a switch with no default case; the two opcodes that may always be copied and why copying them lowers memory traffic; the global pre-pass that bans a producer from duplication unless every consumer will swallow it; the reverse-post-order queue and the thirty-line comment about the duplicate clone it exists to prevent; and the quarter of the file that is about correctness rather than speed, where a slice read out of a buffer meets an update written back into it. The signature exhibit is the six-gate decision ladder ported from the file, run over one small module, so you can watch an edge get refused and see which predicate did it.
• IFRT Arrays: pjrt_array.cc, Line by Line — A jax.Array spread over 512 chips is one Python object, and underneath it is a list of ordinary single-device buffers plus a description of how they tile the whole. This 696-line C++ file is the class that holds those two things together, and the walk reads all of it. The validator that runs before any array exists, which matches buffers to devices strictly by position and compares only the devices this process can address. The six construction paths, one of which quietly skips the validator. The three copy semantics that the header describes as different and that one function implements identically, with the TODO admitting it. Disassembly, which hands back one array per shard without moving a byte. The read back to the host, which refuses anything but a single shard unless the array is replicated. And a copy path where the same invariant is re-derived against a different device list. The signature exhibit runs the validator itself: pick a sharding, perturb one buffer, and watch which of the five checks refuses it and with which string.
• The CPU Thunk Executor: thunk_executor.cc, Line by Line — A GPU gets a driver queue to submit work to. A CPU gets nothing, so the compiled program arrives as a flat list of runtime actions called thunks, and something has to work out what may run at the same time. This 841-line C++ file is that something, read end to end. Each unit declares which buffers it reads and which it writes, and every dependency in the program is derived from those declarations and nothing else, by a conflict test where read-after-read is free. Then a transitive reduction deletes the edges other edges already imply and hands each node a priority equal to how much it unblocks. At run time it is one atomic counter per node and a ready queue: a thunk finishes, decrements its successors, and any counter that hits zero names work that can start now. We read the two edge kinds and why sharing a collective communicator is a weaker constraint than sharing memory; the three thresholds that quietly make a small program single-threaded; the recursive halving that spreads forty ready thunks over four threads without forty queue pushes; and the completion count over sink nodes that makes success and failure end the same way. The signature exhibit is a real ten-thunk graph you step one loop iteration at a time, with all three ready-queue classes.
• The GPU Codegen Path: From a Fused HLO to a Kernel — When the optimizer stops, your model is a few hundred fusions and a handful of library calls, and none of them is machine code yet. This reads the two decisions that finish the job, at a pinned commit of openxla/xla. First the kind: a string an upstream pass stamped on the backend config, or, when there is no stamp, ten questions asked about the fusion roots in a fixed order until one answers yes, with a tenth that always does. Then the class the switch builds from that kind, including the two extra questions the loop arm asks before it settles, one of which turns a copy into a memcpy and skips the kernel entirely. Then the three families behind those classes: six emitters that write a module in an MLIR dialect XLA defined for the purpose and lower it through forty-one passes to LLVM; the Triton path, which builds one module and hands it to another compiler; and the library paths, where cuBLAS arrives as a custom call and cuDNN arrives as a fusion, for reasons that follow from what each library accepts. Then the autotuner: the only pass that compiles and runs real kernels while it is still compiling, the four ways it can get a config, the sentences it uses to refuse, the two-microsecond window that lets a thriftier kernel win over a faster one, the clustering that treats correctness as agreement rather than truth, and the two-tier cache with a key caveat the source states out loud. The signature exhibit runs the real dispatch function over seven candidate fusions and shows which question fires and which class gets built.
• XLA Collectives: One Guest List, Then One Order — A collective is the one instruction in a compiled program that cannot be executed by looking at its own operands, because it needs values from machines that do not share memory with this one. This reads the whole subsystem at a pinned commit of openxla/xla, from both sides of the line. On the compiler side: the list of integers an instruction carries, the two optional fields that decide whether those integers are replica ids, partition ids or positions in a flat enumeration of the device grid, and the switch with four arms that multiplies one written group out against a mesh into the sets of global device numbers the runtime will actually open communicators for. Then the flattened id and the device number, which walk the same grid in opposite directions and disagree on any mesh with more than one of each. On the runtime side: ranks, communicators, cliques and the factory that makes them, all defined without naming a vendor; the future that means launched rather than finished, and means two different things on GPU and CPU; the rendezvous that agrees on membership before any bytes move; and the three separate places where the answer to how do we avoid a deadlock is the same answer. The signature exhibit runs the real group arithmetic and the real verifier checks over a mesh you dial.