When the model runs out
Most of what XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s GPU compiler decides comes from a cost model: count the FLOPs, count the bytes, pick the option the arithmetic favors. A few choices resist that treatment. A GEMM's algorithm, a Triton kernel's tile sizes, a convolution's algorithm: the option space is wide enough, and the candidates close enough in theory, that no cost model built from constants can call the winner reliably.
When the cost model runs out, the compiler measures.
For exactly those lowerings, the autotuner (xla/backends/autotuner/) compiles every candidate, runs each one on the real device at compile time, and keeps the fastest result that is also correct. The winner gets cached, keyed on the device and a fingerprint of the program, and that cache can be serialized to disk so a second compile of the same program on the same chip skips the race entirely.
How a candidate gets timed
The autotuner treats identical instructions as one job. Every candidate gets a 128-bit fingerprint from its own printed form, duplicates collapse to a single representative, and whatever configuration wins for that representative gets stamped onto every equivalent instruction afterward. Config resolution then runs a three-tier decision: a cache hit returns immediately, a default-config request takes the first backend willing to answer, and only a genuinely new candidate pays for the full race.
The race compiles every supported backend-and-config pair, in parallel when a thread pool is available, then profiles them one mutex at a time: a single lock covers the whole profiling phase, so compilation can overlap but measurement cannot, because the device state being timed is exactly what a concurrent measurement would corrupt. Outputs get clustered by numerical closeness, trusted backends checked first, and the winning cluster has to contain a value XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → actually trusts before anything gets declared a winner.
Among candidates in the winning cluster, the fastest one usually takes it outright. One narrow exception exists: any config within a small time window of the fastest, two microseconds by default, counts as tied, and the tiebreak goes to whichever of those uses the least scratch memory. On fast hardware that window is tight enough that a slower config rarely steals a win on memory alone.
Two sides of measurement
Autotuning is one half of how XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → knows a kernel is fast; the other half is measuring the same fact after the program actually runs. compiled.cost_analysis(), the same call the JAX path meets when it lowers a function ahead of time, is the compiler's own estimate: flops and bytes, computed from the program, no device involved. XProf sits at the opposite end entirely, a trace of what really happened on real hardware, the same instrument the kernel path's timeline reads down to individual bytes moved.
Autotuning sits between the two. It does not estimate the way cost analysis does, and it does not merely observe the way a profiler does after the fact; it runs real candidates on a real device, during compilation, specifically to settle a question neither a static model nor a post-hoc trace can answer alone.
Reproducibility has a price
An autotuned answer is a fact about one chip and one driver, not a fact about the algorithm. The same program compiled on a different chip generation, or even a different driver version on the same chip, can autotune to a different winner. Pinning the autotune cache, shipping it alongside the compiled program instead of recomputing it, is how a production build keeps that variable out of the loop.
Distributed autotuning splits the candidate list across processes and exchanges results through a key-value store, and one deliberate shortcut in that path is worth knowing: the store's key hashes the program's fingerprint and the backend name, not the full compilation config. Two compiles of the same HLO under different debug options can therefore collide in the store and swap results neither run actually produced, a trade XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → accepts on purpose to avoid a slower, less deterministic key.
Readings
- XLA tools ↗ the dump flags and the tooling around autotuning
- XLA GPU architecture ↗ where autotuned decisions sit in the GPU pipeline