Compute It Twice: When CSE Lost the Race

A fusion compiler for Nx on Vulkan learns to fold whole neural-network layers into one on-GPU schedule — and then the one optimization every compiler textbook endorses loses the race on both GPUs it was measured on.

The edge we were chasing

EXLA’s structural advantage over an eager tensor backend is not a faster matmul. It is whole-graph compilation. Where an eager backend dispatches tanh(add(dot(x, W), b)) as three separate GPU launches — each with its own kernel, its own intermediate buffer, its own round trip through the driver — a compiler traces the whole expression and emits as few kernels as the graph allows. The elementwise tail folds into one shader. The intermediates never touch host memory. The dispatch count collapses.

VulkanoBackend — the pure-Rust Vulkan tensor backend this site has been documenting for a season — is a good eager backend. It had no compiler. This is the story of building one, and of the single optimization inside it that measurement talked us out of shipping.

The compiler is an Nx.Defn.Compiler. It traces a defn to an expression DAG and asks a simple question of every node: can this fuse? A same-shape elementwise chain becomes one generated GLSL shader, compiled once, cached by source hash, dispatched in a single call. On a ten-op chain that path runs 3.62× the eager per-op backend — not because any kernel got faster, but because nine dispatches and nine intermediate buffers became one of each.

Boundaries, and the schedule between them

A chain of elementwise ops is the easy case. Real graphs have boundaries — a matmul, a convolution, a reduction — ops that are not elementwise and cannot melt into the shader around them. The interesting question is what to do with the elementwise regions between the boundaries.

The answer is a stage schedule. Each boundary becomes a stage that writes a GPU-resident buffer; each maximal elementwise region between boundaries becomes one generated shader whose inputs may be earlier stages’ buffers. relu(x @ W + b) is a matmul stage followed by a single fused max(dot + b, 0) stage — the bias-add and the ReLU and the broadcast all collapse into the epilogue, and the matmul result never leaves the card. A two-layer MLP forward pass is four stages. Nothing falls back to the interpreter; the whole layer compiles.

Over a run of increments the boundary vocabulary grew until most of a neural network fit through it:

Two of those deserve a footnote, because they draw the line the whole design turns on. reshape is a view: a row-major relabel of bytes that are already contiguous, so its stage is no dispatch at all — the planner aliases the input buffer and moves on. transpose looks similar and is the opposite: it permutes axes, it moves data, so it must materialise a new buffer with a real dispatch. Reshape is free. Transpose is not. Hold onto that distinction; it is the entire point of the post.

The obvious optimization

Softmax is the graph that exposes the question. Written out, it shares a subexpression with itself:

defn softmax(x) do
  n = Nx.exp(x - Nx.reduce_max(x, axes: [1], keep_axes: true))
  n / Nx.sum(n, axes: [1], keep_axes: true)
end

The numerator n = exp(x - max(x)) appears twice: once as the thing we divide, once inside the sum we divide by. Trace it and n is one node in the DAG with two parents — the divide and the sum. A compiler that inlines naively will emit the exp into both consumers and compute it twice.

Every compiler course has a name for the fix. Common-subexpression elimination: compute the shared value once, into a temporary, and let both consumers read it. It is on the short list of optimizations so foundational nobody argues about them. So we built it — a hoisting pass that materialises a boundary-crossing shared subexpression into its own stage. With it, softmax plans as four stages instead of three:

It worked, exactly as advertised. The exp ran once. The DAG was honoured. Every textbook nodded. We reached for the benchmark to measure the win, and there wasn’t one.

The race

The bench runs full softmax at nine shapes and, for each, compares the compiled graph with CSE on against the same graph with CSE off — the pre-hoisting behaviour, forced in-process with an environment flag. The column that matters is on/off: above 1, hoisting helped; below 1, it hurt. Every correctness error was zero either way; the two graphs compute the same softmax. The question was only which was faster.

On the RTX 3060 Ti — a current-generation Ampere discrete card, on Linux — hoisting the shared numerator never won:

The 2012 GeForce GT 650M on FreeBSD — a fourteen-year-old mobile Kepler, the other end of the fleet — told the same story with the same shape:

Two GPU generations, two operating systems, nine years of silicon between the cards, and they agreed to two decimal places: common subexpression elimination ranges from harmful to neutral, and never — not at one shape, not on one card — pays off. The worst case was 0.72×. The best case was a tie.

Why the textbook is wrong here

The instinct comes from the CPU, where it is correct. On a CPU an exp is a handful of cycles you would rather not spend twice, and the temporary that saves the second one lives in a register or L1. Recompute is expensive; the bookkeeping to avoid it is nearly free. CSE is a strict win because the thing it removes costs more than the thing it adds.

A GPU inverts both sides of that ledger. The arithmetic is nearly free: a compute shader that already touches every element for the subtract can exp it in the same pass for the price of one more instruction across thousands of lanes running in parallel. What is not free is the thing CSE adds. Materialising the shared value means a whole extra dispatch — command buffer, descriptor set, queue submission, fence — and it means the value makes a full round trip through global memory: the producer writes it out, the two consumers read it back. On a bandwidth-and-dispatch-bound machine, that round trip is the expensive operation, and the redundant exp it was meant to save is the cheap one.

So CSE on a GPU is not the free win it is on a CPU. It is a trade: recompute, against a dispatch plus a memory round trip. On the hardware we have, the recompute is the bargain. The transpose/reshape distinction from three sections ago is the same fact wearing different clothes — the whole cost model of this compiler is dispatches and bytes moved, not arithmetic performed. An optimization that trades the cheap axis for the expensive one runs backwards.

The bonus bug the race caught

The fleet earns its keep by turning up things no single run would. While racing an unrelated pair of increments — f64 fusion and the transpose boundary — a bare 512×512 f64 x @ W came back at 0.56× the eager path. That is one matmul dispatch either way; the compiled path had no business being slower. The kernels were fine. The executor was not: resolving a parameter to its device buffer went through the generic backend transfer, which for an already-GPU-resident tensor still round-tripped to_binary → from_binary — device to host to device — before every stage that read it. The fused graph was paying to download and re-upload its own weights. Binding the resident buffer directly fixed it for every multi-stage path, not just the two being raced.

It is the same lesson as the concatenate post from the spring: the cost that hides is the transfer nobody wrote down, paid by a caller, measured by no benchmark that only times the kernel. You do not find those by reading the shader. You find them by racing the whole thing on a machine that will bill you for every copy.

What shipped

CSE ships in the compiler, off by default, behind NXV_CSE=1 for the rare graph whose shared subexpression is expensive enough — not cheap softmax arithmetic — to earn its dispatch back. There is no device-class gate, because gating implies a class where it wins, and the fleet found none. Two device generations that disagree about almost everything agreed about this, so the honest default is the one the measurements chose, not the one the textbook assumed.

One neighbouring optimization is on by default, and the distinction is worth the last paragraph. When a graph has two outputs — {mean, variance}, softmax stats — and both read the same subexpression, the planner still materialises it once and lets both outputs share the buffer. That looks like the same idea and is not: the shared value was going to be a stage buffer anyway, because it is one of the outputs. Reusing a buffer that already exists costs nothing. Manufacturing a buffer that did not have to exist costs a dispatch. The whole judgment is which of those two you are actually doing, and only the benchmark can tell you, because on paper they are the same optimization.

The compiler now folds whole dense layers, CNN layers, classifier heads, softmax, layernorm, multi-output graphs and transposed-weight matmuls into on-GPU stage schedules, in f32 and f64, across a 2012 Kepler and a 2021 Ampere — 863 doctests and 361 tests green on both. The marquee feature is the fusion. The lesson is the one optimization we built, measured, and left switched off: on a CPU you eliminate redundant work, and on a GPU you learn to cherish it. A fusion compiler’s job is to delete dispatches, not arithmetic.

Raw data (softmax CSE race, Kepler + Ampere): bench_results/CSE_SOFTMAX_RACE.md. Script: examples/cse_softmax_bench.exs. The compiler lives in lib/nx_vulkan/compiler.ex and lib/nx_vulkan/codegen.ex. Set NXV_CSE=1 and race it on your own card.