Why did XLA not fuse my two operations?

The pass asks six questions about the edge and stops at the first no. Is the consumer fusible at all, which excludes a parameter, a while loop, a conditional, a call and a domain. Is the operand fusible. Is the producer in the set of instructions a pre-pass banned from duplication, which happens when it has two or more large inputs, duplicating it would grow memory traffic, and at least one of its consumers will not take it. Is the producer the computation root, which can never be fused away because its value has to exist when the computation returns. Does the cost condition fire, which it does when fusing would duplicate the producer and either the pass may not duplicate at all or the producer opcode is on the expensive list, unless the producer is a broadcast or a widening convert. And finally, is the fusion safe with respect to a consumer that writes into one of its own inputs. Every no is a string, and asking the compiler for the fusion visualisation writes those strings into a dump as "Not fusing |producer| into |consumer| as" followed by the reason.

What makes an HLO instruction expensive for fusion in XLA?

A hand-written list, not a cost model. InstructionFusion::IsExpensive is a switch on the opcode with 134 cases and no default label, which is exactly the number of opcodes XLA defines, so every operation in the language has been classified by hand and adding a new opcode breaks the build until somebody decides. Fifty-four are cheap: element-at-a-time arithmetic, comparisons and selects, and everything that only moves or reinterprets bytes such as reshape, transpose, slice, pad, concatenate and get-tuple-element. Seventy-two are expensive: dot, convolution, the transcendentals, reductions, sort, gather and scatter, plus everything with a side effect or a control-flow body, where the classification is really about duplication rather than speed. Six depend on the element type and are cheap for reals but expensive for complex. Divide and remainder are the interesting pair: expensive in general, but cheap when the element type is integral and the divisor matches an effective scalar constant or a broadcast of one, because that compiles down to a multiply and a couple of shifts.

Why does XLA duplicate a broadcast into every consumer?

Because it is one of exactly two opcodes exempted from the duplication rule by IsAlwaysDuplicable, and the comment above it gives the principle: these should be nodes that get cheaper the more they are duplicated. A broadcast writes a large tensor made out of a small one. Materialise it once and every consumer reads a full tensor from memory; compute it inside each consumer and the cost is an index calculation. The other exemption is a convert whose input is smaller in bytes than its output, so a half-precision tensor widened to single precision. Fusing that into the consumer means the consumer reads the small tensor and widens it in a register, which is strictly less traffic even though the arithmetic now happens more than once. A narrowing convert gets no exemption. Both exemptions apply even when the pass was constructed with duplication disallowed, which is what the TODO comment in the source is being wry about.

What is multi-output fusion, and why does the base pass refuse it?

Regular fusion copies the producer into the consumer, so the producer can be deleted when nothing else uses it. Multi-output fusion copies nothing: it puts producer and consumer in one kernel that returns both values as a tuple, so the producer keeps its identity and its other users keep reading it. That is how a producer with several consumers gets fused without being duplicated. The base class declares two hooks for it in the header and both default to a refusal explaining that multi-output fusion is not supported by this pass, so in the base class every refused edge picks up that phrase as the second half of its combined message. A subclass that implements it overrides both hooks, and then one extra check applies: merging two nodes merges their dependencies, so if anything the producer feeds is also something the consumer depends on, the merged node depends on itself. The cycle check tries the reachability map first and falls back to a forward depth-first search, because the map was built before the loop started rewriting the graph.

When is fusing a slice into a dynamic-update-slice unsafe?

A dynamic-update-slice is compiled to write into its input buffer instead of allocating a new one, which is the only sane way to change one block of a large tensor. That is safe as long as nothing else needs the old contents. So when the producer being fused reads the same buffer the update writes, ShouldFuseInPlaceOp has to prove the accesses line up. One pattern is recognised: a slice or dynamic-slice reading exactly the block the update writes, at the same indices, with unit strides. Then the read happens before the write to those exact addresses and both can live in one kernel. Everything else is refused, and the refusals are specific: shapes that differ between the slice and the update, indices that differ or cannot be proven equal, more than one update inside the consumer, a producer holding more than one non-elementwise operation, and a recursive check that finds another non-elementwise user that would end up sharing the buffer. A producer whose single non-elementwise operation is neither a slice nor a dynamic slice falls through to a refusal naming it an unrecognised in-place pair, which is the honest answer when the function cannot prove safety.

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.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Instruction Fusion Legality: instruction_fusion.cc, Line by Line — built to be understood, not skimmed.

Questions

Why did XLA not fuse my two operations?: The pass asks six questions about the edge and stops at the first no. Is the consumer fusible at all, which excludes a parameter, a while loop, a conditional, a call and a domain. Is the operand fusible. Is the producer in the set of instructions a pre-pass banned from duplication, which happens when it has two or more large inputs, duplicating it would grow memory traffic, and at least one of its consumers will not take it. Is the producer the computation root, which can never be fused away because its value has to exist when the computation returns. Does the cost condition fire, which it does when fusing would duplicate the producer and either the pass may not duplicate at all or the producer opcode is on the expensive list, unless the producer is a broadcast or a widening convert. And finally, is the fusion safe with respect to a consumer that writes into one of its own inputs. Every no is a string, and asking the compiler for the fusion visualisation writes those strings into a dump as "Not fusing |producer| into |consumer| as" followed by the reason.
What makes an HLO instruction expensive for fusion in XLA?: A hand-written list, not a cost model. InstructionFusion::IsExpensive is a switch on the opcode with 134 cases and no default label, which is exactly the number of opcodes XLA defines, so every operation in the language has been classified by hand and adding a new opcode breaks the build until somebody decides. Fifty-four are cheap: element-at-a-time arithmetic, comparisons and selects, and everything that only moves or reinterprets bytes such as reshape, transpose, slice, pad, concatenate and get-tuple-element. Seventy-two are expensive: dot, convolution, the transcendentals, reductions, sort, gather and scatter, plus everything with a side effect or a control-flow body, where the classification is really about duplication rather than speed. Six depend on the element type and are cheap for reals but expensive for complex. Divide and remainder are the interesting pair: expensive in general, but cheap when the element type is integral and the divisor matches an effective scalar constant or a broadcast of one, because that compiles down to a multiply and a couple of shifts.
Why does XLA duplicate a broadcast into every consumer?: Because it is one of exactly two opcodes exempted from the duplication rule by IsAlwaysDuplicable, and the comment above it gives the principle: these should be nodes that get cheaper the more they are duplicated. A broadcast writes a large tensor made out of a small one. Materialise it once and every consumer reads a full tensor from memory; compute it inside each consumer and the cost is an index calculation. The other exemption is a convert whose input is smaller in bytes than its output, so a half-precision tensor widened to single precision. Fusing that into the consumer means the consumer reads the small tensor and widens it in a register, which is strictly less traffic even though the arithmetic now happens more than once. A narrowing convert gets no exemption. Both exemptions apply even when the pass was constructed with duplication disallowed, which is what the TODO comment in the source is being wry about.
What is multi-output fusion, and why does the base pass refuse it?: Regular fusion copies the producer into the consumer, so the producer can be deleted when nothing else uses it. Multi-output fusion copies nothing: it puts producer and consumer in one kernel that returns both values as a tuple, so the producer keeps its identity and its other users keep reading it. That is how a producer with several consumers gets fused without being duplicated. The base class declares two hooks for it in the header and both default to a refusal explaining that multi-output fusion is not supported by this pass, so in the base class every refused edge picks up that phrase as the second half of its combined message. A subclass that implements it overrides both hooks, and then one extra check applies: merging two nodes merges their dependencies, so if anything the producer feeds is also something the consumer depends on, the merged node depends on itself. The cycle check tries the reachability map first and falls back to a forward depth-first search, because the map was built before the loop started rewriting the graph.
When is fusing a slice into a dynamic-update-slice unsafe?: A dynamic-update-slice is compiled to write into its input buffer instead of allocating a new one, which is the only sane way to change one block of a large tensor. That is safe as long as nothing else needs the old contents. So when the producer being fused reads the same buffer the update writes, ShouldFuseInPlaceOp has to prove the accesses line up. One pattern is recognised: a slice or dynamic-slice reading exactly the block the update writes, at the same indices, with unit strides. Then the read happens before the write to those exact addresses and both can live in one kernel. Everything else is refused, and the refusals are specific: shapes that differ between the slice and the update, indices that differ or cannot be proven equal, more than one update inside the consumer, a producer holding more than one non-elementwise operation, and a recursive check that finds another non-elementwise user that would end up sharing the buffer. A producer whose single non-elementwise operation is neither a slice nor a dynamic slice falls through to a refusal naming it an unrecognised in-place pair, which is the honest answer when the function cannot prove safety.

Instruction Fusion Legality: instruction_fusion.cc, Line by Line

Questions

Related explainers