How does XLA decide which operations to fuse on a GPU?

With a priority queue of producers, drained greedily. For every instruction that could be fused into its users, the pass asks a GPU performance model for two estimates, the time with the instruction as its own kernel and the time with it merged into its consumers, and the difference between them is the priority. That priority is not a score: its type is absl::Duration, so it is literally the wall-clock time the merge is estimated to save. Producers with a benefit at or above zero go into an ordered map keyed by the pair of the priority and the instruction unique id; the pass pops the largest key, performs the merge, then recomputes the priorities of the neighbours the merge disturbed and pops again. Two rules sit outside the measurement. A fusible bitcast gets positive infinity, because it is a pure relabelling of bytes and folding it away cannot cost anything. A constant gets negative infinity and is swept up by a separate loop after the queue drains, because folding constants early would throw the cost model off for merges that were going to happen anyway.

Why did XLA not fuse two operations I expected it to fuse?

Because one of thirteen checks said no, and every one of them carries a sentence you can read back. The most common answers on a stock build: the consumer is a bitcast (not fusing into a single bitcast as consumer); both sides contain a reduction whose input-to-output element ratio is at least sixteen (both the producer and the consumer contain a reduce); the merged kernel would exceed the budget of ninety-six operands plus output buffers, or eight unnested reductions, or the shared memory of one block; the merged IR would exceed ten thousand nodes (the fusion would result in an overly large code duplication); or the producer is the root instruction, which cannot be fused into because instructions after the root can still be live. There is also a rule that catches people out: a producer is fused into all of its non-bitcast users or none of them, so a single user that refuses makes the whole producer ineligible. One string is not a refusal at all. triton heroless fusion is not enabled simply means the Triton route declined and the ordinary elemental checks ran instead.

What does priority = time_unfused - time_fused actually measure?

The estimated difference in wall-clock time between running the producer as its own kernel and running it inside its consumers. Three things go into it. One kernel launch disappears, which the model prices at a fixed one microsecond. The producer output stops being written to memory and read back by each consumer, which is the memory traffic the merge deletes. And against those, the producer body is duplicated into every consumer, so its arithmetic is performed once per consumer and each consumer now reads the producer inputs rather than the single intermediate. That last term is why a producer with one consumer is almost always worth fusing while the same producer with three consumers can score below zero. Note the shape of the accumulation: on an incremental update the pass adds the delta from the newly gained consumers to the score already in the queue and subtracts the stored runtimes of the consumers that went away, rather than re-estimating the whole user list.

What is the difference between instruction_fusion.cc and priority_fusion.cc?

The first is backend-agnostic legality and the second is the GPU decision with a cost model. Priority fusion still reaches into the older pass for one shared rule, the in-place-operation check that keeps aliasing correct, and it borrows the shared budgets from the GPU fusibility helpers. What it adds is the ordering. Where a rule-based pass fuses whatever matches whenever it matches, priority fusion measures every candidate first, does the biggest estimated win, and then recomputes what changed. That ordering matters because merges interact: absorbing one producer into a fusion can change how many consumers a different producer has, and therefore whether fusing that producer is profitable at all.

How do I see what the XLA fusion pass decided?

Three environment flags, all off by default, all set through XLA_FLAGS. Setting a dump directory and adding this pass to the per-pass regular expression writes a file named priority_fusion_dump, which holds the module before fusion as text, the GPU device description, and one record per decision: every merge with its producer, consumer and result names, every ineligible producer with its reason string, and every priority update with both estimated runtimes in microseconds. Turning on the fusion visualisation writes the same story as HTML, with captions such as About to fuse, Fused, Rejected with a formatted negative benefit, and Ineligible with a reason, though it needs a graph renderer plugin registered to produce anything. And compiler fuel, set as a budget for the pass named priority-fusion, simply stops the pass after a given number of merges, so a binary search over that number isolates the exact merge that changed your numbers before you go reading any reasoning.

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.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of What Shares a GPU Kernel: priority_fusion.cc, Line by Line — built to be understood, not skimmed.

Questions

How does XLA decide which operations to fuse on a GPU?: With a priority queue of producers, drained greedily. For every instruction that could be fused into its users, the pass asks a GPU performance model for two estimates, the time with the instruction as its own kernel and the time with it merged into its consumers, and the difference between them is the priority. That priority is not a score: its type is absl::Duration, so it is literally the wall-clock time the merge is estimated to save. Producers with a benefit at or above zero go into an ordered map keyed by the pair of the priority and the instruction unique id; the pass pops the largest key, performs the merge, then recomputes the priorities of the neighbours the merge disturbed and pops again. Two rules sit outside the measurement. A fusible bitcast gets positive infinity, because it is a pure relabelling of bytes and folding it away cannot cost anything. A constant gets negative infinity and is swept up by a separate loop after the queue drains, because folding constants early would throw the cost model off for merges that were going to happen anyway.
Why did XLA not fuse two operations I expected it to fuse?: Because one of thirteen checks said no, and every one of them carries a sentence you can read back. The most common answers on a stock build: the consumer is a bitcast (not fusing into a single bitcast as consumer); both sides contain a reduction whose input-to-output element ratio is at least sixteen (both the producer and the consumer contain a reduce); the merged kernel would exceed the budget of ninety-six operands plus output buffers, or eight unnested reductions, or the shared memory of one block; the merged IR would exceed ten thousand nodes (the fusion would result in an overly large code duplication); or the producer is the root instruction, which cannot be fused into because instructions after the root can still be live. There is also a rule that catches people out: a producer is fused into all of its non-bitcast users or none of them, so a single user that refuses makes the whole producer ineligible. One string is not a refusal at all. triton heroless fusion is not enabled simply means the Triton route declined and the ordinary elemental checks ran instead.
What does priority = time_unfused - time_fused actually measure?: The estimated difference in wall-clock time between running the producer as its own kernel and running it inside its consumers. Three things go into it. One kernel launch disappears, which the model prices at a fixed one microsecond. The producer output stops being written to memory and read back by each consumer, which is the memory traffic the merge deletes. And against those, the producer body is duplicated into every consumer, so its arithmetic is performed once per consumer and each consumer now reads the producer inputs rather than the single intermediate. That last term is why a producer with one consumer is almost always worth fusing while the same producer with three consumers can score below zero. Note the shape of the accumulation: on an incremental update the pass adds the delta from the newly gained consumers to the score already in the queue and subtracts the stored runtimes of the consumers that went away, rather than re-estimating the whole user list.
What is the difference between instruction_fusion.cc and priority_fusion.cc?: The first is backend-agnostic legality and the second is the GPU decision with a cost model. Priority fusion still reaches into the older pass for one shared rule, the in-place-operation check that keeps aliasing correct, and it borrows the shared budgets from the GPU fusibility helpers. What it adds is the ordering. Where a rule-based pass fuses whatever matches whenever it matches, priority fusion measures every candidate first, does the biggest estimated win, and then recomputes what changed. That ordering matters because merges interact: absorbing one producer into a fusion can change how many consumers a different producer has, and therefore whether fusing that producer is profitable at all.
How do I see what the XLA fusion pass decided?: Three environment flags, all off by default, all set through XLA_FLAGS. Setting a dump directory and adding this pass to the per-pass regular expression writes a file named priority_fusion_dump, which holds the module before fusion as text, the GPU device description, and one record per decision: every merge with its producer, consumer and result names, every ineligible producer with its reason string, and every priority update with both estimated runtimes in microseconds. Turning on the fusion visualisation writes the same story as HTML, with captions such as About to fuse, Fused, Rejected with a formatted negative benefit, and Ineligible with a reason, though it needs a graph renderer plugin registered to produce anything. And compiler fuel, set as a budget for the pass named priority-fusion, simply stops the pass after a given number of merges, so a binary search over that number isolates the exact merge that changed your numbers before you go reading any reasoning.

What Shares a GPU Kernel: priority_fusion.cc, Line by Line

Questions

Related explainers