The Shader That Repeated Itself

A Posterior on Any GPU ended on a promise: no model falls over on a GPU that can hold it. The toy models kept that promise easily. Then a real one — a three-regime Gaussian mixture, the kind of model that decides whether a market is trending, mean-reverting, or just volatile — walked up to the same 2013 Kepler and made it think for an uncomfortably long time before the first posterior ever came back. The GPU wasn’t slow. The shader was enormous.

3-regime mixture, 30 obs
  emitted GLSL shader:  837,132 bytes

Eight hundred kilobytes of generated GLSL for a model with six free parameters. That is not a big model. It is a big shader, and the difference between those two facts is the whole story.

A DAG, flattened into a tree

The custom-likelihood synthesizer walks an Nx.Defn expression and emits a GLSL scalar expression for the per-observation log-density. That expression is a graph — a mixture’s log-sum-exp shares a softmax denominator, and every regime’s term divides by it. But the emitter writes the graph out as a tree. Every time a shared node is reached by a different parent, it is emitted again, in full, subtree and all. The softmax denominator log(1 + exp(min(w1,10)) + exp(min(w2,10))) didn’t appear once and get reused; it appeared, expanded, thousands of times. The shader was mostly the same handful of subexpressions, copied until the file was a megabyte.

This is invisible on a normal-priors model — nothing is shared, so nothing repeats. It only bites when the likelihood has structure that reuses intermediate results, which is exactly what mixtures, hierarchies, and anything with a normalizing term do. The models you actually reach for.

Say it once

The fix is the oldest trick in a compiler’s book — common-subexpression elimination — applied at the last possible moment, on the emitted string, inside the per-observation loop. Walk the body, find balanced subexpressions that occur more than once, and hoist each into a local computed once per iteration:

double _c0 = log(1.0lf + exp_d(min(q[5],10.0)) + exp_d(min(q[6],10.0)));
double _c1 = obs_j - q[0];
// ... the body now references _c0, _c1, ... instead of re-deriving them

Shortest repeats are hoisted first, so a later, larger binding naturally references the smaller ones already introduced — creation order is dependency order, no topological sort required. Every binding and every use share the one obs-loop scope, so it is a safe, local transform. About sixty lines of Elixir, no change to the math.

The race

The repository ships the benchmark that measures it — bench/cse_race.exs, the same three-regime mixture rendered twice, once with config :exmc, glsl_cse: false and once with it on:

=== CSE race: 3-regime mixture, 30 obs ===
mode     GLSL bytes   SPIR-V bytes
CSE off  837,132      2,999,572
CSE on    94,030        223,280
         8.9x smaller  13.4x smaller

glslang -V (generate the SPIR-V the driver ingests):
  RTX 3060 Ti / Linux      1088ms -> 204ms    5.3x faster
  GT 650M / FreeBSD (2013)  1623ms -> 255ms    6.4x faster

posterior:   bit-identical  (|diff| = 0.0)

Eight-point-nine times smaller source, thirteen times smaller binary, five to six times faster to compile — and the sampled mu_trend mean is identical to the last decimal in both modes. Note which machine gains the most: the thirteen-year-old GT 650M on FreeBSD, where the slow CPU pays the steepest tax for a megabyte of shader, goes from 1.6 seconds to a quarter-second per compile. The oldest hardware in the fleet is exactly where saying it once matters most. CSE doesn’t approximate anything; it just refuses to say the same thing twice.

The honest part: this is a compile-time win

It is tempting to claim the shader now runs thirteen times faster. It doesn’t. Measured on-device, the sampling throughput barely moved — because the SPIR-V optimizer and the GPU driver already perform common-subexpression elimination at the instruction level. The GPU was never actually recomputing that exp() thousands of times per step; it was computing it once. The redundancy was a compile-time problem the whole way down: eight hundred kilobytes of text for glslangValidator to parse, three megabytes of SPIR-V to generate and hand to the driver, all to describe a computation the driver would then collapse anyway.

Which is precisely why it mattered. The thing that made the Kepler pause before its first posterior was not the arithmetic; it was compiling a megabyte of shader that expanded into three megabytes of SPIR-V, per model, before a single leapfrog step ran. Cut the shader to ninety kilobytes and the first-posterior compile stops being something you wait through. The GPU that could hold the model can now also compile it in a hurry. Correct was never in question — CSE made tractable the part that 0.3.0 left slow.

The change lives in the custom-synthesis emitter (borodark/exmc), toggled with config :exmc, glsl_cse: false if you ever want to see the megabyte for yourself. The GPU backend it renders for is borodark/nx_vulkan, and the release that put a NUTS sampler on any Vulkan driver in the first place is written up in A Posterior on Any GPU.

A shader is just a program you compile a lot. The fastest way to compile a program is to make it smaller — and the smallest program is the one that never says the same thing twice.