What is the XLA HLO pass pipeline?

XLA does not run one big optimizer — it runs an ordered pipeline of small passes, each with a single job, and pipelines nest inside pipelines. A program that JAX has traced into a jaxpr and lowered to StableHLO enters XLA as an HLO module: a directed graph of instructions like dot, broadcast, add, and reduce. The GPU backend alone registers on the order of two hundred passes, grouped into roughly a dozen named sub-pipelines (optimization, spmd-partitioner, layout assignment, fusion, post-fusion optimization, scheduling, and so on). Each pass is a function from an HLO module to a slightly better one, and after any pass that changes the graph, the runner re-runs invariant checkers to prove the graph is still well-formed — the one rule being that "invariant checkers must not change the graph."

What is operation fusion in XLA and why does it matter?

Fusion merges several adjacent operations into a single GPU kernel. It matters because elementwise tensor ops are almost never limited by arithmetic — they are limited by memory bandwidth, how fast the chip can move bytes to and from HBM (the large global device memory). Run three ops as three kernels and every intermediate result is written to HBM and read straight back, which is pure tax. Fuse them into one kernel and the intermediates never leave the chip: they live in registers, get consumed the instant they are produced, and only the original inputs and final output touch HBM. For the workloads XLA targets, fusion alone is often the difference between memory-bound and compute-bound, which is why it is the single most important pass in the pipeline.

How does XLA’s priority fusion decide what to fuse?

Priority fusion refuses to guess. For a candidate merge it asks a GPU performance model to estimate two runtimes — how long the operations take run separately (time_unfused) and how long they take fused (time_fused) — and the merge’s priority is the difference, the wall-clock time it would save. It keeps every fusable edge in a priority queue keyed by that score, pops the highest one, performs the merge, and recomputes the neighbours’ priorities before popping again: greedy, biggest win first. A merge whose priority is at or below zero is dropped, because "if the priority is negative, it’s not helpful to perform fusion." Two ops jump the queue by rule: bitcasts get infinite priority ("bitcasts should always be fused first, since they are no-ops"), and constants are fused last so they don’t distort the cost model.

When does XLA refuse to fuse two operations?

A positive priority is necessary but not sufficient — a legality check can veto a merge before the cost model votes, returning a fusion decision that is forbidden with a reason string. You can never fuse into the module’s root instruction. You cannot fuse two reduces together ("both the producer and the consumer contain a reduce"), because the cost model lacks the tiling analysis for a reduce feeding a reduce. You cannot exceed the kernel parameter budget: XLA caps operands-plus-outputs per fusion at 96, because "in a fusion with many operands, each GPU thread likely has to do a lot of work, and so possibly uses a lot of registers, thus limiting occupancy." And it forbids duplicating an expensive producer into many consumers, since recomputing it in each one would "lead to an increase in memory bandwidth use" — the legality-layer version of a negative priority.

Why does XLA recompile when my input shape changes?

Because the whole pipeline compiles one specific HLO module — specific tensor sizes and dtypes. The cache key is the shape, so a jitted function called with batches of length 31, 32, and 33 runs all ~200 passes three times and pays the compile latency three times. The fix is to stop the shapes from varying: pad or bucket inputs to a few fixed sizes so the compiled binary is reused. Truly dynamic sizes need shape polymorphism, where XLA compiles once against symbolic dimensions — more flexible, but it gives the cost model less certainty, so the fused kernels can come out slower. Most production JAX code lands in the middle: a handful of bucketed shapes, each compiled once, each fused to the hilt.

Inside XLA: ~200 Passes and the Fusion Decision

A traced program enters XLA as a graph of abstract operations and leaves as a handful of GPU kernels. In between, a couple hundred passes rewrite it — and one of them, fusion, decides which operations share a kernel by putting a stopwatch on every merge. We walk the ordered pass pipeline, the memory round-trip fusion deletes, the priority cost model (time_unfused − time_fused) and its greedy queue, the vetoes that refuse a merge, buffer assignment, codegen, and the recompilation you pay for all of it — drawn, computed, and animated.

System design · AI / ML. The source ↗

A free, interactive, animated visual explainer of Inside XLA: ~200 Passes and the Fusion Decision — built to be understood, not skimmed.

Questions

What is the XLA HLO pass pipeline?: XLA does not run one big optimizer — it runs an ordered pipeline of small passes, each with a single job, and pipelines nest inside pipelines. A program that JAX has traced into a jaxpr and lowered to StableHLO enters XLA as an HLO module: a directed graph of instructions like dot, broadcast, add, and reduce. The GPU backend alone registers on the order of two hundred passes, grouped into roughly a dozen named sub-pipelines (optimization, spmd-partitioner, layout assignment, fusion, post-fusion optimization, scheduling, and so on). Each pass is a function from an HLO module to a slightly better one, and after any pass that changes the graph, the runner re-runs invariant checkers to prove the graph is still well-formed — the one rule being that "invariant checkers must not change the graph."
What is operation fusion in XLA and why does it matter?: Fusion merges several adjacent operations into a single GPU kernel. It matters because elementwise tensor ops are almost never limited by arithmetic — they are limited by memory bandwidth, how fast the chip can move bytes to and from HBM (the large global device memory). Run three ops as three kernels and every intermediate result is written to HBM and read straight back, which is pure tax. Fuse them into one kernel and the intermediates never leave the chip: they live in registers, get consumed the instant they are produced, and only the original inputs and final output touch HBM. For the workloads XLA targets, fusion alone is often the difference between memory-bound and compute-bound, which is why it is the single most important pass in the pipeline.
How does XLA’s priority fusion decide what to fuse?: Priority fusion refuses to guess. For a candidate merge it asks a GPU performance model to estimate two runtimes — how long the operations take run separately (time_unfused) and how long they take fused (time_fused) — and the merge’s priority is the difference, the wall-clock time it would save. It keeps every fusable edge in a priority queue keyed by that score, pops the highest one, performs the merge, and recomputes the neighbours’ priorities before popping again: greedy, biggest win first. A merge whose priority is at or below zero is dropped, because "if the priority is negative, it’s not helpful to perform fusion." Two ops jump the queue by rule: bitcasts get infinite priority ("bitcasts should always be fused first, since they are no-ops"), and constants are fused last so they don’t distort the cost model.
When does XLA refuse to fuse two operations?: A positive priority is necessary but not sufficient — a legality check can veto a merge before the cost model votes, returning a fusion decision that is forbidden with a reason string. You can never fuse into the module’s root instruction. You cannot fuse two reduces together ("both the producer and the consumer contain a reduce"), because the cost model lacks the tiling analysis for a reduce feeding a reduce. You cannot exceed the kernel parameter budget: XLA caps operands-plus-outputs per fusion at 96, because "in a fusion with many operands, each GPU thread likely has to do a lot of work, and so possibly uses a lot of registers, thus limiting occupancy." And it forbids duplicating an expensive producer into many consumers, since recomputing it in each one would "lead to an increase in memory bandwidth use" — the legality-layer version of a negative priority.
Why does XLA recompile when my input shape changes?: Because the whole pipeline compiles one specific HLO module — specific tensor sizes and dtypes. The cache key is the shape, so a jitted function called with batches of length 31, 32, and 33 runs all ~200 passes three times and pays the compile latency three times. The fix is to stop the shapes from varying: pad or bucket inputs to a few fixed sizes so the compiled binary is reused. Truly dynamic sizes need shape polymorphism, where XLA compiles once against symbolic dimensions — more flexible, but it gives the cost model less certainty, so the fused kernels can come out slower. Most production JAX code lands in the middle: a handful of bucketed shapes, each compiled once, each fused to the hilt.

Inside XLA: ~200 Passes and the Fusion Decision

Questions

Related explainers