A neural network is written as a sequence of operators, and every boundary between two of them is a place where a result has to be written out to slow memory and read back — so compilers fuse adjacent operators into single kernels, but they decide the boundaries first using hand-written rules and then generate code, which locks in choices made without knowing what the code will look like; this paper (Chinese Academy of Sciences and Leeds) analyses the whole network as mathematics before drawing any boundaries, and gets one model down from 662 kernels to a single one.
The argument in one glance
-
Problem
Every kernel boundary is a round trip through slow memory
A network is written as a graph of operators, and each region of that graph becomes a kernel on the accelerator. Between two kernels, whatever the first computed must be written out to global memory and read back by the second, because that is the only place two kernels can communicate. On a processor whose arithmetic vastly outpaces its memory, those round trips dominate — and each launch costs time on its own, so a model executing thousands of kernels is spending most of its life at boundaries.
-
What people did
Fuse operators into larger kernels
The standard remedy is fusion: merge several operators into one kernel so the intermediate values stay in fast on-chip memory. Every serious compiler does this, guided by hand-written rules about which patterns are safe and profitable to combine — an element-wise operation following a matrix multiply, a chain of element-wise operations, and so on.
-
Why it fell short
The boundaries are drawn before anyone knows what the code needs
Fusion happens on the operator graph and code generation happens afterwards, which fixes the partition before the compiler knows what resources the resulting kernel will actually use. The rules also only cover anticipated patterns, so common cases go unhandled — reductions get their own kernel even when the operator feeding them could have kept the data on chip, and two matrix multiplies are never combined even when one could be loading its weights while the other computes.
-
This paper's idea
Analyse the whole network as mathematics, then decide
This paper (Chunwei Xia, Jiacheng Zhao, Qianqi Sun, Zheng Wang, Yuan Wen, Teng Yu, Xiaobing Feng and Huimin Cui) lowers the entire model into tensor expressions — small mathematical statements saying how each output element is computed — and does global dataflow analysis over those. Boundaries are then drawn from what the analysis and the hardware limits say, and the operators inside each region are recombined by algebraic transformation rather than by rule.
Why this problem is worth solving
The representation the paper chooses is what makes everything else possible.
AN
OPERATOR in a neural network graph — a convolution, an attention block, a softmax — is a name
for something a library implements, and a compiler that sees only the name must reason from
whatever the library author documented. A TENSOR EXPRESSION is smaller and more honest: a
mathematical statement saying how each element of an output tensor is computed from elements of
its inputs, including which axes are summed over. Lowering an operator to tensor expressions
decomposes it into pieces with simple, explicit dependence structure — a softmax becomes a
reduction, in which one output element depends on many inputs, followed by an element-wise
operation, in which one output depends on exactly one input. The analysis then works on
mathematics rather than on library calls, so it makes no assumptions about how anything is
implemented and can reason across operators whose combined dependence would be too complicated
to describe at the operator level.
The second thing to be clear about is why kernel boundaries cost what they do.
TWO
kernels on a graphics processor cannot share on-chip memory, because the on-chip storage
belongs to thread blocks that exist only for the duration of a kernel. So anything one kernel
computes and another consumes must go out to global memory and come back — orders of magnitude
slower than the registers and shared memory it was sitting in. There is also a fixed cost to
launching each kernel, which matters when a model runs thousands of them. Fusion removes both,
but it is not free to attempt: a fused kernel needs enough registers and shared memory for
everything inside it, and if the combination exceeds what the hardware provides, occupancy
collapses and the fusion is worse than the split. That trade cannot be evaluated until the code
exists, which is precisely why deciding the boundaries first is the wrong order.
What people did before, and where it broke
-
Fusion by hand-written rule. Established compilers and vendor libraries recognise particular shapes — an element-wise chain after a matrix multiply, sibling operators that can run side by side — and merge those. This is effective for the patterns it covers.
-
The gap. Only anticipated patterns are covered, and models keep producing new ones. The paper finds compilers failing to fuse element-wise memory operations with the matrix multiply feeding them, so those reshapes and permutations become their own kernels.
-
Keep reductions in their own kernel. Because a reduction combines many elements into one, the usual rules treat it as a barrier that fusion cannot cross.
-
The gap. That forces the entire input to the reduction out to global memory first. Doing the reduction partially within each thread block and combining across blocks keeps the data on chip, but requires reasoning the operator-level rules cannot express.
-
Call an optimised library for the heavy operators. Compute-intensive operators are dispatched to a vendor library, which is excellent at those in isolation.
-
The gap. A library call is opaque, so it can never be fused with anything around it. One compiler is measured issuing six separate library calls for the matrix multiplies of a single transformer layer, where the whole layer could have been one kernel.
-
Search for a good schedule within a fixed partition. Automatic schedulers explore tiling, ordering, and unrolling extremely well once the boundaries are given.
-
The gap. They optimise inside the boxes somebody else drew. The paper builds on such a scheduler and shows that changing where the boxes go is worth an average of 3.9 times on top of everything the scheduler already achieves.
How a whole model becomes a few kernels
From an operator graph to a small number of large kernels
-
-
new
-
new
-
new
-
new
-
new
The idea
Do the analysis on mathematics, not on operator names. The reason fusion rules keep failing is that an operator is a label rather than a description — the compiler knows a softmax is happening but not what shape of data movement that implies, so it falls back on whatever somebody wrote down about softmaxes. Lowering everything to tensor expressions removes the label. What remains is a statement of which input elements each output element needs, from which the compiler derives the fine-grained dependence itself. The crucial simplification is that these dependences come in only two kinds. An expression with no summed axis has each output element depending on exactly one input element, which means it can be inlined into whatever produced that input. An expression with a summed axis has each output depending on many inputs, which is a reduction and needs different treatment. Every operator in a model decomposes into some arrangement of those two, so a small analysis covers everything, including combinations no rule anticipated. Choosing a representation whose semantics are explicit enough to reason over, rather than working around an opaque one, is the same argument as the memory-bound fusion work covered elsewhere on this site, which stitches whole subgraphs by deciding where each intermediate value should live.
Draw boundaries from resource limits rather than patterns. With the analysis global, the partition can be posed as a budget problem. Expressions are added to a growing subprogram, and after each addition the compiler asks whether the schedules generated so far would collectively demand more registers and shared memory than the hardware provides — a check made possible by running the schedule optimiser before committing, rather than after. When the budget is exceeded, a new subprogram starts. There is a second constraint doing real work here: fusing across a reduction requires a barrier that synchronises every thread block, and such a barrier only functions if all the blocks are resident at once, which caps the launch dimension. So the partition is bounded by two hardware realities and by nothing else. That is why one model ends up as a single kernel — nothing about it exceeds the budget, so nothing forces a split. Deciding structure from a measured resource constraint rather than a heuristic is the same discipline as the graph transformation work covered elsewhere on this site, which also co-designs the partition with the schedule rather than fixing one first.
Reductions can be fused if you are willing to restructure them. The most consequential single transformation concerns reductions, which conventional rules treat as walls. A reduction following a matrix multiply normally forces the multiply’s entire output through global memory so a second kernel can sum it. The alternative is to reduce in two phases — each thread block sums the portion it holds, keeping that data in on-chip memory where it was already sitting, then the partial results are combined across blocks with atomic operations and a barrier. The full tensor never reaches global memory; only the much smaller partial sums do. Doing this requires knowing precisely which elements each output depends on, which is exactly what the element-wise analysis provides and what operator-level rules cannot. It is the mirror image of the fission work covered elsewhere on this site, which splits operators apart to reduce peak memory where this merges them to reduce memory traffic — the same dial turned in opposite directions for different objectives.
Fusing across a reduction
How a reduction stops being a barrier to fusion
A matrix multiply followed by a reduction, first as two kernels and then as one.
Step 1 of 5
The conventional split, and what it costs
The matrix multiply runs as one kernel and writes its complete output to global memory. A second kernel reads that output back and sums along an axis. The full intermediate tensor makes a round trip it did not need to make.
Rules treat a reduction as a wall because its output depends on many inputs spread across whatever tiling the previous operator used, and reasoning about that at the operator level would require knowing the previous operator's schedule — which the rule does not.
Why it works
- Two dependence kinds cover everything. Every expression either has a summed axis or does not → each output element depends on many inputs or on exactly one → a small analysis handles arbitrary operators, including combinations no hand-written rule anticipated.
- Analysis before partitioning removes the circularity. Generating candidate schedules and then checking their resource demands → boundaries are set by what the code actually needs rather than by pattern matching → one model compiles to a single kernel because nothing in it exceeds the budget.
- Restructuring a reduction keeps its input on chip. Summing within each block and combining with atomics and a barrier → only partial results reach global memory instead of the whole tensor → reductions stop being fusion barriers, which is where several of the largest gains come from.
- Collapsing one-to-one chains is pure algebra. Composing the index mappings of consecutive element-wise expressions into a single function → reshapes, permutations and slices disappear into their producer → the memory operators that compilers were treating as separate kernels cease to exist.
- Bigger kernels unlock optimisations boundaries forbid. Once operators share a kernel → loads for one can be overlapped with arithmetic for another, and buffers can be reused across them → gains that no amount of per-kernel scheduling could reach.
The evaluation
| Property | What was measured | The result | What it tells you |
|---|---|---|---|
| End-to-end inference time | |||
| How many kernels remain | |||
| Traffic to global memory | |||
| Which optimisation contributes what | |||
| Robustness of the baselines |
Check yourself
Check your understanding
1. Why is deciding fusion boundaries before generating code the wrong order?
2. What does lowering operators to tensor expressions buy the analysis?
3. How can a reduction be fused with the operator producing its input, when rules normally treat it as a barrier?
4. The ablation shows the last two optimisations improving transformer models and leaving the others unchanged. What does that indicate?
Sources
Optimizing Deep Learning Inference via Global Analysis and Tensor Expressions. Chunwei Xia, Jiacheng Zhao, Qianqi Sun, Zheng Wang, Yuan Wen, Teng Yu, Xiaobing Feng, Huimin Cui. ASPLOS 2024, Volume 1, pages 286 to 301. DOI 10.1145/3617232.3624858. The system is named Souffle. Full text read from a copy hosted by one of the authors, which appears to be a preprint since its header carries an unfilled placeholder identifier.