Operator Fusion and the Conditions That Block It

Take two elementwise operators in sequence: a bias add and a rectified linear unit. Executed as two kernels, the first reads the input tensor, writes an output of the same size, and the second reads that output and writes again. For a float32 tensor of $N$ elements, the pair reads $8N$ bytes and writes $8N$ bytes, and the arithmetic is roughly $2N$ operations, so one operation is issued per eight bytes moved. That ratio is far below what a compute unit can sustain, which makes the pair bandwidth-bound by a wide margin. Fused into one kernel, the same computation reads $4N$ bytes, writes $4N$ bytes, and keeps the intermediate in registers. The memory traffic halves and the launch count halves, and the result is unchanged because the same operations in the same order are performed, only without the round trip through memory.

That arithmetic is the entire economic argument for fusion, and it is why the transformation dominates the pass list of every graph compiler. The transformation itself is trivial to state: substitute the consumer’s use of the producer’s output with the producer’s expression. What is hard is deciding when the substitution preserves semantics and when it makes the resulting kernel worse.

Vertical fusion and the producer-consumer classification

A fusion group is a set of nodes that will be emitted as one kernel. Vertical fusion merges a producer into its consumer along a data dependency. Whether that is legal depends on what the producer and consumer do, not on their size.

The clean case is elementwise into elementwise. An elementwise operator computes each output element from the corresponding input elements, so its loop nest is a single flat iteration over the output index space, and merging two of them composes the bodies without changing the iteration space. This is the case that every compiler implements first, and it covers the bulk of the node count in a typical convolutional or transformer block: adds, multiplies, activations, casts, and scalar constants.

A reduction as the consumer is also tractable. Fusing a producer into a reduction means computing the producer’s value inside the reduction loop rather than reading a precomputed tensor. The producer’s work is then performed once per reduction step rather than once per output element, which is a real cost multiplier, but the saving on memory traffic usually dominates when the reduction is the outer structure and the producer is cheap.

A reduction as the producer is the case that breaks. If the producer reduces over an axis and the consumer is elementwise, the fused kernel needs the entire reduction result before the consumer can start, which reintroduces exactly the intermediate buffer that fusion was meant to avoid. Compilers handle this by fusing the consumer into the reduction’s epilogue: the final step of the reduction loop applies the elementwise operation to the accumulator before storing. This works when the consumer is elementwise and the reduction is complete in one pass, and it fails when the consumer needs to combine multiple reduction outputs, as a softmax does across a row and a column.

Two properties determine what else is possible. A producer with a single consumer can be fused without duplicating work. A producer with $k$ consumers, fused into each of them, executes its computation $k$ times. That is a profitable trade when the producer is cheap and memory-bound, and a disaster when it is an expensive convolution. This is the central quantity a fusion heuristic has to estimate: recomputation cost against saved traffic.

Horizontal fusion and the shared loop

The second form merges nodes that are independent but have the same iteration space. Two branches of a residual block, or two attention heads, can be computed in one kernel whose loop nest iterates over the union of their outputs, with each thread block assigned to one branch. The saving is launch overhead and, when the branches read disjoint slices of the same tensor, improved cache locality. The cost is register pressure: a kernel holding the working set of two branches needs twice the registers, and a kernel that spills to local memory can be slower than the two launches it replaced.

Horizontal fusion is also the mechanism by which small tensors stop being pathological. A graph with two hundred elementwise nodes over tensors of a few thousand elements each would otherwise issue two hundred launches, each of which occupies the device for less time than the launch takes to enqueue. Merging them into a single kernel with a block-index-to-branch dispatch converts the launch overhead into a table lookup.

What blocks a merge

The blocking conditions fall into three groups, and each has a different remedy.

Semantic blockers. A consumer that needs its input in a layout the producer does not produce cannot be fused without a transpose inside the fused kernel, which is usually worse than a separate transpose that the layout assignment pass can eliminate elsewhere. A consumer that reads a producer’s output at an index unrelated to its own output index, such as a gather, breaks the single-loop-nest assumption: the fused kernel would need an indirect addressing pattern that the vectoriser cannot handle. An operator with internal state across the loop, such as a scan or a cumulative sum, cannot accept an elementwise producer fused into its recurrence without serialising the producer.

Structural blockers. A producer whose output is a graph output must be materialised, because something outside the compiled region reads it. A producer with multiple consumers may be duplicated, as above, or kept. A producer inside a conditional branch cannot be fused with a consumer outside the branch, because the branch may not execute.

Resource blockers. Fusion grows the live range of everything the fused kernel holds. A group that spans more values than the register file can hold spills, and spilling is a memory operation with worse locality than the buffer it replaced. The practical limit is set by the register count per thread and the size of the per-thread working set, which the tiling decides. A group tiled at $128 \times 128$ outputs and assigned to 256 threads gives each thread 64 output elements, so a fused kernel that keeps the whole tile in registers needs 64 float32 registers for accumulators alone, a quarter of the 255 that the instruction set allows per thread, before addressing and temporaries.

/* Fused bias + relu: one pass over memory instead of two. */
void bias_relu(const float *in, const float *bias, float *out, int n, int c) {
    for (int i = 0; i < n; i++) {
        float v = in[i] + bias[i % c];   /* broadcast along the channel axis */
        out[i] = v > 0.0f ? v : 0.0f;    /* relu, no intermediate buffer */
    }
}

Fusion changes what the arithmetic looks like

Merging a producer into a consumer is not always a pure win, because the fused kernel performs the producer’s arithmetic on values that the unfused version would have written to memory and read back. When the producer’s output is consumed more than once inside the consumer, the fused version recomputes it unless the compiler is smart enough to keep it in a register. A compiler that fuses naively turns a producer consumed twice into a doubled computation, and the resulting kernel does more arithmetic than the two it replaced. Whether that is faster depends on whether the kernel is bandwidth-bound or compute-bound, which is exactly the quantity a cost model must estimate.

There is also a numerical consequence. Fusion need not change rounding, since the same operations on the same values give the same result, but it removes the store-and-reload that would otherwise force a value through memory at the graph’s dtype, and it gives the code generator the option of contracting a multiply followed by an add into a fused-multiply-add. A fused kernel that accumulates in float32 and stores in float32 matches the unfused version; one that keeps an intermediate in a register at a wider type, or contracts an add into an FMA, does not. Compilers that must match a reference output keep the accumulator dtype equal to the graph dtype and rely on precision reduction being applied deliberately rather than incidentally.

Where the transformation stops

Fusion cannot cross a barrier that requires global synchronisation. A softmax needs the maximum of a row before it can exponentiate, and the sum before it can divide, and a row may not fit in one thread block’s shared memory. The standard resolution is a multi-pass kernel: one pass computes the row maximum and sum, a second applies them. That is two launches, and no amount of fusion removes the second. The same structure appears in layer normalisation, in any reduction followed by a normalisation, and in attention, whose score matrix must be materialised or tiled depending on its size.

Fusion also cannot cross a graph break. Anything the compiler could not trace into the graph is a barrier by construction, and the fusion region ends at the boundary. The practical fusion ratio of a model, meaning the fraction of nodes absorbed into larger groups, is therefore a property of how much of the model was captured, not only of how good the fusion heuristic is.

The last limitation is the one that keeps the transformation honest: every fusion decision is made without knowing the actual input sizes, unless the compiler has been given them. A group that is optimal at batch size one, where every kernel is launch-bound, can be pessimal at batch size sixty-four, where the same group spills. This is why shape-specialised compilation and dynamic shape handling matter as much as the fusion rules themselves.