Optimizing UchenML inference for a browser tab

Model training is an incremental process — start small, scale up. For Dots, I started with an 8×8 field, got a moderately “fun” model, then moved to 16×16. That made the model roughly 4× slower, and I started worrying about what kind of demo it would make. Apprehensive as I was, I decided to take a detour from training and focus on lower-level optimizations.

I ended up with a 6.45× performance improvement — 21.98 ms down to 3.41 ms per forward pass, one evaluation of the model. Most gains came from the following three optimizations:

• switching the convolution algorithm from im2col to Winograd F(2,3);
• adding kernel fusion, used mostly to fuse ReLU activations into convolutions and fully-connected layers;
• adding kernel elision, to remove unnecessary copies.

WASM runtime constraints

Running in the browser imposes some constraints that are less common when targeting modern hardware directly:

1. Single threaded. Workers with copied buffers are possible; the added overhead has not been worth it so far. It might be worth looking into once I have larger models.
2. No FP16 arithmetic. The WASM half-precision proposal that adds f16 SIMD lanes is still at phase 2 and no engine ships it, so compute is 32-bit throughout. Parameters still ship as FP16 and are unpacked on load, which halves the download but not the traffic once the model is running — and memory bandwidth, not arithmetic, is what this model runs out of first.
3. 4-lane SIMD. SIMD support in WASM is limited to 128-bit registers — four floats at a time — even on AVX2/AVX-512 hardware with wider ones. Native builds can use the wider registers, courtesy of Google Highway — but so far I am mostly working on WASM and ARM Neon so there’s likely some performance left on the table on wider instruction sets.
4. No GPU. WebGPU now ships by default in Chrome, Safari and Firefox, and WebNN is at Candidate Recommendation, but adopting either would be a significant departure from the current codebase.

Specialized algorithm — Winograd F(2,3) replacing im2col: 21.98 → 3.75 ms

The slowdown was mostly in the convolutions. UchenML ran all of them through im2col, a common general-purpose algorithm — any filter size, stride or padding — and the obvious default. The models I was building did not need that flexibility: eleven of Dots’ seventeen convolutions are dense 3×3 at unit stride and unit padding, and the remaining six are 1×1, which is a matrix multiply already. The common specialization for the 3×3 case is Winograd, introduced for CNNs by Lavin and Gray in Fast Algorithms for Convolutional Neural Networks.

Winograd transforms the input and the filter into another domain, multiplies there, and transforms the result back. For 3×3 filters it produces a 2×2 output patch from a 4×4 input tile with 16 multiplies where the direct form needs 36. The filter half of that transform depends only on the weights, so it runs once, when the model is built. The three transform matrices that do this are not something you derive by hand — Lavin’s wincnn generates them symbolically for F(2,3) and its larger relatives.

Winograd was designed to cut multiplies, and it does: 2.25× fewer. Both algorithms also stage the input before they can multiply anything — im2col writes every input value nine times, once per filter tap, while Winograd writes it four times, once per overlapping tile: 2.53 MiB of staging writes against 1.125 MiB for the same 24×24 layer at 128 channels.

UchenML picks the algorithm from the shape of the convolution — both are still there. The forward pass drops from 21.98 ms to 3.75 ms — 5.86×.

A second look at that number

5.86× is more than the algorithm can account for. Fewer multiplies is 2.25×; less staging is 2.25×; neither, nor both, multiplies out to 5.86×.

Turns out my im2col did not have enough scratch. It tiles its GEMM through a scratch buffer with a 128 KiB target. At 128 input channels a single 3×3 patch is 4.6 KiB, so a tile holds 28 output pixels — and a 16×16 layer therefore takes ten passes, re-reading the entire 576 KiB weight matrix on every one of them. Raising that budget to 8 MiB and changing nothing else takes the im2col build from 21.58 ms to 11.60 ms, against Winograd’s 3.72 ms in the same run. Still significantly slower.

The ladder further down still says 21.98 → 3.75 ms, because that is what actually shipped and what the other rungs were measured against. But “Winograd is 5.86× faster than im2col” is not a statement about the two algorithms.

What the specialization costs

Winograd runs faster — but the cost is not zero.

1. A second convolution implementation. im2col is still needed and is not going away — and with the separate depthwise kernel that makes three. More code, more maintenance, a bigger binary.

2. The model gets 72% bigger. Winograd stores the transformed 4×4 values rather than the 3×3 taps, so every dense 3×3 convolution stores sixteen parameters per channel pair where im2col stores nine:

That is 1,154,048 extra parameters — at FP16, 5.52 MB of weights on the wire against 3.21 MB, a 2.3 MB increase. A browser tab downloads the model once and answers moves for the rest of the session, so 2.3 MB for 5.86× on every move is a trade I will take. The weights could be repacked to 3×3 for transfer and unpacked at load, but that is lossy for the reason below, and I have not built it.

3. The trained layer is no longer a 3×3 convolution. The sixteen values start as a linear transform of the nine taps, but training updates all sixteen independently and never projects back. Training in this format is not equivalent to training in the 3×3 format.

For my project, these costs are worth it. They allowed me to grow the play field 9×, while keeping the latency in the range I am comfortable with.

Reducing the data round trip — kernel fusion: 3.75 → 3.42 ms

A model can be viewed as a pipeline of layers, each one taking the output of the previous and producing its own. Broadly they come in two kinds. Parametric layers do most of the actual “learning” — convolution detects spatial patterns, linear mixes them up, attention identifies relationships between values; they own the weights and do real arithmetic for every value they produce. Activations are fixed functions that make the relationship between input and output non-linear (hence the other name, “nonlinearities”); the common ones have no parameters and transform a single value at a time.

A naive implementation runs each layer separately, storing its output to memory and reading it back for the next. For an activation that is a full round trip through memory to do one instruction’s worth of work per value. A fused kernel does both operations in one pass, so the values never leave the CPU:

Whether that costs anything depends on the shape. One convolution output here is 24×24×128 floats — 288 KiB, bigger than the L1 cache on any current core — and there are eleven of them per forward pass, so the extra pass spills out to L2.

This is one of the jobs a graph optimizer in a framework like PyTorch does — it replaces separate Conv and ReLU nodes with a single fused node.

UchenML uses the C++ compiler as the graph optimizer. At compile time, C++ magic (argument-dependent lookup plus a requires clause) picks the fused function instead of the two separate ones.

Linear and convolution layers (Winograd and im2col alike) have separate implementations fused with Clamp. Relu is a special case of Clamp with a minimum of 0 and a maximum of FLT_MAX. The fused kernel does the multiply and the clamp in one pass instead of two.

template <typename D, bool Clamp, bool HasTail>
HWY_ATTR void ColumnMajorGemvKernel(D d, std::span<float> output,
...
    for (size_t i = 0; i < full; i += lanes) {
      const V acc = hn::LoadU(d, output.data() + i);
      const V w = hn::LoadU(d, col + i);
      V r = hn::MulAdd(w, scalar, acc);
// Only compiled in for Clamp=true instantiations
      if constexpr (Clamp) {
        r = hn::Min(hn::Max(r, min_vec), max_vec);
      }
      hn::StoreU(r, d, output.data() + i);
    }
...
}

Clamp is a template parameter, not a runtime flag: the unfused build contains no branch and no Min/Max at all. Fusing ReLU into the convolution and linear layers gains 9%. I actually managed to break the fusion once during the refactor, and the benchmark immediately noticed.

What the fusion costs

Cheaper than the algorithm swap, but not free either.

1. The kernel gets copied, once per activation. Clamp is a template parameter, which is exactly why the unfused build carries no branch — and also means every activation variant emits its own complete copy of the kernel, with the activation welded into the innermost loop. Crossed with the tail-handling parameter, the matrix-vector kernel goes from two instantiations to four the moment fusion is switched on. For this model, that adds 1,540 bytes of code.

2. The rules are hand-written, one per layer pair. Nothing here will discover Conv → BatchNorm → ReLU on its own — you get exactly the fusions someone (you) wrote down. And the fused path is picked by a requires guard, so a miss is silent: no error, the framework just runs the clamp as a second pass. Still correct, only slower, and only the benchmark notices.

Doing nothing when there’s nothing to do — kernel elision: 3.42 → 3.41 ms

Some layers exist only for the developer’s convenience. Identity does nothing at all; it fills a slot where a real layer might otherwise go — the pass-through branch of a residual block, for instance. Reshape says the tensor shape changes — a 3D tensor for convolution becoming a 1D vector for a fully-connected layer — without anything in the underlying data changing.

Identity is dropped from the model type. Piping a real layer onto a model whose last layer is Identity does not append — it replaces the Identity. A mid-chain | Identity is a straight no-op, which matters for composability: an explicit Identity never breaks a downstream fusion rule that keys on “the last layer.”

Reshape survives, but emits nothing. A residual block hands the next block a flat buffer where it wants a conv-shaped one. The obvious fix is a reinterpreting copy per block; instead the reshape declares the following layer’s scratch buffer as its own and hands it over, so the copy never happens — the shape change is a change of type, and types are free.

Standing alone at the end of a model the reshape does pay one copy — then it has nothing to borrow scratch from.

Optimizing memory: reuse scratch

The fourth optimization was a core part of the design from the start: no heap allocations during inference. UchenML pre-allocates a scratch arena and uses it for the whole forward pass. Because the model is known at compile time — the shapes of every input and output, and the scratch each layer needs, such as im2col buffers — the size of that arena is computable before the program runs.

As mentioned above, a model is a pipeline of layers. While a layer runs, the only buffers that need to be live are its own scratch and the scratch of the previous layer — which holds its input. Everything earlier is dead, and its space can be reused.

Every serious runtime arrives at that observation somewhere: ONNX Runtime, TensorRT and TFLite all analyze the graph to work out which buffers can overlap, then hand out slices of one arena. The difference here is when the planning happens. There is no graph to analyse at load time, because the model is a constexpr variable, so the arena is sized while the program is compiled. ContextForInfer is two variants, one over even-indexed layers and one over odd:

// uchen/inferrence_context.h
typename Evens::type_t evens_;   // std::variant<CT<0>, CT<2>, CT<4>, …>
typename Odds::type_t  odds_;    // std::variant<CT<1>, CT<3>, CT<5>, …>

Along a straight chain, peak scratch is max(even layers) + max(odd layers), not the sum. A Parallel block counts as one layer whose scratch is the sum of its branches, each of them recursing into the same scheme. It is a classic ping-pong buffer with two differences that matter: the sizes are computed by the type system from the model expression, and the “allocation” is a variant emplace inside a lazily-initialized getter. Scratch types are trivially constructible and have no destructors, so switching between them is cheap. No allocator is involved, and a mis-sized buffer is a compile error rather than memory corruption. The default Model::operator() still heap-allocates one context per call, but an overload takes a caller-held context instead — and then the inference path allocates nothing at all. A deeper net costs more time, not more scratch.

Two layer kinds are not a straight pipeline, and they treat the arena in opposite ways. Parallel — what a residual fork and the seven-head fan-out are built from — gives every branch its own nested ContextForInfer and keeps them in a tuple, because the Join at the end consumes all the branch outputs together. Inside a branch the ping-pong resumes; across branches it stops, so a fork costs the sum of its branches. Scan, the recurrent (RNN) layer, goes the other way: it holds one cell context and reuses it for every step, so inference over a sequence is flat in the length of that sequence rather than proportional to it.

The main catch is that every layer output is a view into the arena, not a value, so a result is only valid until the next call on that context — hold a policy vector across two inferences and you are reading the second one’s activations. The type system does not catch that. The context is also single-occupancy, so every thread needs its own arena — easier to design in than to retrofit.

Backward pass

The backward pass needs the intermediate activations to compute gradients, so it cannot reuse the same arena. Training gets its own context type, and the whole difference between the two is where the scratch lives.

Both implement the same memory::Context<M, I> interface, so the layers never know which one they are running under. ContextForInfer keeps two std::variants and hands out a pointer into whichever one matches the layer’s parity, constructing it on first use — a variant holds one alternative at a time, which is exactly what makes the ping-pong work. ContextForGradientCalculation keeps a std::tuple over every layer index instead, and its getter is the whole story:

// uchen/training/training_pass.h
return [=]() { return &std::get<Idx>(*tuple); };

No parity, no construction on demand, no sharing: every layer’s scratch is its own tuple member and all of them stay live for the entire pass. Peak scratch goes from max(even) + max(odd) to the sum over every layer. A few scratch types even change shape for the backward pass — ConcreteTypeForGradient is the identity for most layers, but Fork/Parallel and the RNN scan keep per-branch and per-step state that only the gradient path reads. Scan is the sharpest case: the single reused cell context becomes a std::list holding one backward pass per step, so the flat cost inference enjoys turns into one proportional to sequence length. The reverse walk runs from the last step to the first, so every step’s forward context has to still be there when its turn comes — that is simply what backpropagation through time costs.

It is also the one place that still allocates on the heap, twice per step: a list node, and the unique_ptr each BackwardPass holds for its own context. The compile-time arena cannot help here. Every other size in the model comes from the type, but the number of steps is a property of the input, so there is nothing to size ahead of time; node-based storage then keeps each step’s record at a stable address as the sequence grows. The no-heap rule is an inference-path constraint and training runs offline, so this is affordable — but it is where the compile-time story runs out.

The asymmetry is deliberate. Training runs offline on a machine with 128 GiB of RAM, where keeping every activation is cheap and the alternative — recomputing them — is not. Inference runs in a browser tab, where it is the other way round. Same model type, same layer code, two context types.

A fun challenge

At this point I am fairly happy with the performance of UchenML. I am looking into deploying it to more targets — Cloudflare Workers, Raspberry Pi Picos — which pretty much guarantees more performance challenges.

That is the appeal, and why I enjoyed this more than I expected to. A browser tab gives you one thread, no FP16, four SIMD lanes — no easy way to throw more hardware at it. There is no bigger instance to move to, no accelerator to offload to, no framework flag that makes the problem someone else’s.

In the meantime, how about a game of Dots?

Companion pieces: SIMD Softmax, NEON to WASM, with Google Highway on the kernels underneath every number above, and The Model Is a Type on why the model is a C++ type in the first place.