SIMD Softmax, NEON to WASM, with Google Highway

Overview

Pivoting to focus on WASM meant deprioritizing GPU support in UchenML — but I still wanted it to run fast. And “fast” meant two different machines: the model runs in WASM, while training should use everything my Mac Studio has.

SIMD was the obvious place to start, and I was apprehensive about it. It usually means several parallel implementations of the same kernel, one per architecture. I wanted to write my own kernels — just not one copy per target.

That is when I decided to give Google Highway a try. There are many “portable SIMD” libraries and I am generally skeptical of them — they usually offer high level primitives and may not support all the architectures I am interested in. Highway was curious because it mostly achieves portability at compile time, which is well aligned with the UchenML philosophy of offloading everything I can to the compiler.

Google Highway

Google Highway is a portable SIMD library for C++: you write a kernel once and it compiles to NEON, AVX2, AVX-512, SVE or WASM SIMD. What makes it interesting is what it refuses to do. It does not raise the level of abstraction — there is no vector class with overloaded operators, no expression templates, and no autovectorizer to argue with. It gives every operation the hardware already has a consistent name, and then gets out of the way.

On top of that baseline it adds some extra functionality: transcendentals such as Exp and Log, which are essential for softmax, and a set of algorithm-style utilities.

Softmax in Highway

Softmax is one of the basic operations in machine learning, used for classification and an essential part of the attention mechanism. It turns a vector of logits into a probability distribution: $$\mathrm{softmax}\left(x_i\right)=\frac{e^{x_i}}{\sum _j^n{e}^{x_j}}$$

Evaluating that directly overflows, so the kernel does it in three passes over one contiguous span of floats — which is exactly a column of a column-major matrix, so that layout is this kernel called once per column. Row-major is the awkward case: its columns are strided by the row length, and walking one costs a cache line per element. That needs a different kernel — one that tiles columns to a cache line — and its own post.

Here it is as it ships in UchenML, verbatim. hn is the conventional alias for the per-target namespace, namespace hn = hwy::HWY_NAMESPACE;:

template <class D>
void ContiguousRun(D d, const float* HWY_RESTRICT in, float* HWY_RESTRICT out,
                   size_t n) {
  using V = hn::VFromD<D>;
  const size_t lanes = hn::Lanes(d);
  const V neg_inf = hn::Set(d, -std::numeric_limits<float>::infinity());

  V acc_max = neg_inf;
  size_t i = 0;
  for (; i + lanes <= n; i += lanes) {
    acc_max = hn::Max(acc_max, hn::LoadU(d, in + i));
  }
  if (i < n) {
    acc_max = hn::Max(acc_max, hn::LoadNOr(neg_inf, d, in + i, n - i));
  }
  const V vmax = hn::Set(d, hn::ReduceMax(d, acc_max));

  V acc_sum = hn::Zero(d);
  i = 0;
  for (; i + lanes <= n; i += lanes) {
    const V e = hn::Exp(d, hn::Sub(hn::LoadU(d, in + i), vmax));
    acc_sum = hn::Add(acc_sum, e);
    hn::StoreU(e, d, out + i);
  }
  if (i < n) {
    const auto mask = hn::FirstN(d, n - i);
    const V v = hn::LoadNOr(vmax, d, in + i, n - i);
    const V e = hn::IfThenElseZero(mask, hn::Exp(d, hn::Sub(v, vmax)));
    acc_sum = hn::Add(acc_sum, e);
    hn::StoreN(e, d, out + i, n - i);
  }

  const V inv = hn::Set(d, 1.0f / hn::ReduceSum(d, acc_sum));
  i = 0;
  for (; i + lanes <= n; i += lanes) {
    hn::StoreU(hn::Mul(hn::LoadU(d, out + i), inv), d, out + i);
  }
  if (i < n) {
    hn::StoreN(hn::Mul(hn::LoadN(d, out + i, n - i), inv), d, out + i, n - i);
  }
}

The three passes run in the order the formula implies, each a main loop over whole vectors and a tail for the leftovers. One contiguous run is also the everyday case: a classifier’s logit vector over its label set is exactly this shape, and training one evaluates a softmax per example.

The first pass exists only to keep the arithmetic in range. exp outgrows a float just past 88, so a single logit of 100 comes back as infinity, the sum is infinity too, and infinity over infinity is NaN — the whole vector is garbage. Subtracting the largest logit first pulls everything down to at most exp(0) == 1, and the answer is unchanged, because the same constant cancels between the top and the bottom of the fraction. acc_max accumulates lane-wise, ReduceMax collapses it, and Set broadcasts the result into vmax — which also fixes the tail’s fill value at -inf, the one that loses every Max it meets.

The second exponentiates in - vmax, adds into acc_sum and stores into out, so the third never calls Exp again.

The third divides by ReduceSum — once, as a reciprocal and a multiply, rather than a divide per element.

Build and run it without installing anything: the whole benchmark is on Compiler Explorer — more below.

Line by line

template <class D> — D is the tag type: zero-sized, carrying the element type and the lane count. The width lives in the type, so this one source text is a 128-bit kernel under CappedTag<float, 4>, a 256-bit one under AVX2’s ScalableTag and a 512-bit one under AVX-512.

ScalableTag is the default — the target’s native width. CappedTag and FixedTag are for when that is wrong:

• Algorithms correct at only one width: a 4×4 transpose, interleaved complex pairs, anything treating the vector as a structure rather than as independent lanes.
• You profiled and a narrower vector won. AVX-512 clock throttling of the Skylake era is the stock example.

D d — by convention the tag is the first argument to nearly every operation; ops that already take a vector or a mask deduce it and need none. It occupies no storage, so passing it is free.

HWY_RESTRICT — portable __restrict. in and out must not alias: the second pass writes out and the third reads it back.

using V = hn::VFromD<D>; — the target’s vector type, named by the tag rather than by you: float32x4_t on NEON, __m256 on AVX2, v128_t on WASM. Masks get their own alias, hn::MFromD<D>.

hn::Lanes(d) — elements per vector; for float, 4 on NEON, SSE and WASM, 8 on AVX2, 16 on AVX-512. Its compile-time counterpart is hn::MaxLanes(d), which is what you need where a constant is required — array extents, constexpr arithmetic. On fixed-width targets the two agree; on SVE and RVV only MaxLanes is a constant, because the real width belongs to the machine the binary lands on.

hn::Set(d, x) / hn::Zero(d) — broadcast a scalar into every lane. That is where neg_inf, vmax and inv come from, and where acc_sum starts.

hn::LoadU / hn::StoreU — unaligned load and store, and where the API is easiest to appreciate. That one call is vld1q_f32 on NEON, wasm_v128_load on WASM, and on x86 a choice between _mm_loadu_ps, _mm256_loadu_ps and _mm512_loadu_ps. The element type is baked into the intrinsic name too (ps — packed single precision); LoadU is just LoadU, for double as well. Pure magic!

Load and Store without the U require vector-aligned addresses and are worth it only if you own the allocation. Note the argument order: loads take the tag first, stores take the value first — LoadU(d, p) against StoreU(v, d, p).

hn::Max, hn::Add, hn::Sub, hn::Mul — basic arithmetic, every lane at once. Fused multiply-add is MulAdd(a, b, c), one instruction where the target has FMA and a multiply plus an add where it does not — one of the places the library quietly stops you writing the same kernel twice.

hn::LoadN / hn::LoadNOr / hn::StoreN — the tail, for the n % lanes elements the main loop cannot take. They touch exactly the count you pass and carry a guarantee that matters: they never read or write past it. LoadN zero-fills the inactive lanes, LoadNOr fills them with a value you supply, and the max pass needs the second — zeros would win the Max over a run of negative logits, and the max-subtraction that keeps Exp from overflowing would stop protecting it.

The obvious-looking alternative, MaskedLoadOr with BlendedStore, reads and writes a whole vector and blends. On every target where HWY_MEM_OPS_MIGHT_FAULT is 1 — NEON, x86 before AVX-512, and sanitizer builds — the hardware really does touch the lanes the mask disables, which off the end of an array is a fault or an ASAN report.

hn::FirstN(d, k) / hn::IfThenElseZero(m, v) — masks. FirstN builds one with the first k lanes set; IfThenElseZero keeps those lanes and zeroes the rest, which the exp tail needs because its inactive lanes feed a horizontal sum and exp(0) is 1, not 0. A mask is its own type rather than a vector of 0/-1, because that is what AVX-512 and SVE have in hardware. The rest of the family is IfThenElse, IfThenZeroElse and IfNegativeThenElse.

hn::ReduceMax / hn::ReduceSum — the only cross-lane operations here: collapse a vector to one scalar. Both run once per pass rather than once per iteration, because the loops accumulate into a vector — acc_max, acc_sum — and reduce only at the end.

hn::Exp — not a hardware instruction. No mainstream SIMD ISA has a vector exponential; Highway ships one in hwy/contrib/math, a lane-wise polynomial built from the same portable ops. It is why the middle pass does not drop to scalar.

The loops themselves are boilerplate, and hwy/contrib/algo/transform-inl.h exists to remove it. Foreach(d, in, n, no, func) walks a read-only buffer and fills its remainder with LoadNOr(no, …) — the -inf of the max pass is just an argument there, and the whole tail branch goes away. Transform(d, inout, n, func) does the same for an in-place map, which is the scale pass exactly.

The middle pass is the one that resists: it reads in, writes a different buffer, and accumulates a sum on the side, while Transform and its Transform1 / Transform2 siblings are all in-place on inout. Its remainder would also arrive through LoadN, so the dead lanes would hold exp(0 - vmax) — nonzero — and the callback never sees the count it would need to mask them. Two passes out of three, then, which is why this kernel still writes its loops out.

Everything in hwy/contrib arrives through an -inl.h header — hwy/contrib/math/math-inl.h here. The suffix is the convention for anything that has to be compiled once per target rather than once per translation unit, and it is why those includes sit inside the target namespace rather than at the top of the file with the others.

Static and dynamic dispatch

Everything above is a template over a tag, so something still has to pick the target. HWY_STATIC_DISPATCH(Kernel)(...) compiles exactly one — the best enabled at build time — and calls it directly. HWY_DYNAMIC_DISPATCH compiles the kernel once per supported target and chooses at run time, which is what a binary redistributed to unknown x86 machines wants.

Dynamic costs more than a macro swap. It works by self-inclusion: HWY_TARGET_INCLUDE names the current file, hwy/foreach_target.h re-includes it once per target, the kernels sit between HWY_BEFORE_NAMESPACE() and HWY_AFTER_NAMESPACE(), and anything that must exist exactly once — main, the HWY_EXPORT table — goes behind #if HWY_ONCE. So the file has to be able to include itself by path, which is precisely what a Compiler Explorer paste cannot do; that is why the link below is static-only. The Highway quick reference has the full recipe.

Which one you want is a deployment question, not a kernel question. WASM has a single target and fixes its feature set at compile time; aarch64 has NEON in the baseline; x86 is where run-time selection earns its keep, because one binary may have to run on an old SSE4 machine and a new AVX-512 one. Either way the kernel is untouched — both modes call the same function template.

Run the benchmark on Compiler Explorer

The standalone benchmark depends on nothing but Highway and google_benchmark, which means Compiler Explorer can build and run it: open it on Compiler Explorer. CE’s workers are x86, so it builds at AVX2 and eight float lanes. The assembly pane is the more instructive half — that is where the tag either vanishes or does not.

It carries two variants of each kernel from the same source text: CappedTag<float, 4> and ScalableTag<float>. On a 128-bit target those are the same type and compile to the same machine code, so they have to tie; on CE they are four lanes against eight, which turns the pair into a width comparison, bounded above by 2×. The strided/* rows are the row-major kernel, which tiles columns to a cache line and is not shown here; the ragged/* pair is that kernel on 250 columns, where the last tile does not fill:

The eight-lane arm wins every pair, which is the expected direction. Read the CPU column, not the wall column: CPU time counts only the cycles this process actually ran, so it is the kernel’s cost, while wall time also counts everything the machine did instead. items_per_second is items over CPU time for that reason.

On CE that ratio runs 1.6 to 2.8, and it shows: the contiguous pair reports capped4 → native as 2.37×, above the 2× that doubling the lane count can give. A single repetition means there is no cv column to argue with either. Use CE to confirm the target and read the assembly; take ratios on a machine you control.

One trap, and the reason the link passes -maes: Highway gates its AVX2 target behind its SSE4 target, and SSE4 requires __PCLMUL__ && __AES__. -march=haswell supplies PCLMUL but not AES, so without -maes you quietly get a narrower target than you asked for. The full incantation is -std=c++20 -O3 -march=haswell -maes. This is why the benchmark prints what it actually got before it prints a single measurement:

hwy static target: NEON
  capped4: 4 lanes x 2 tuples x 4 = 32 cols/tile
  native:  4 lanes x 2 tuples x 4 = 32 cols/tile

If those lines surprise you, fix the build before reading any timings. Verify the target you actually got, every time: -march=haswell without -maes here, and emscripten without -msimd128 below, both compile and pass their tests while quietly handing you a narrower kernel than you asked for.

One kernel for training and WASM inference

UchenML splits its tree in two. Runtime code — anything an embedder calls during a forward pass — allocates nothing on the heap, because it runs inside emscripten’s allocator in a WASM bundle where allocation is expensive and fragmentation is unpredictable. Training code runs on the host and uses std::vector freely.

Softmax is a runtime kernel: stack-only, and shared by both. The forward pass in a browser tab and the forward pass inside the training loop on the Mac Studio are the same source, compiled twice — no “fast path for training, portable path for deployment” pair to keep in sync.

One caveat if you ship to WASM, the same species of mistake as the -maes one. Highway selects its HWY_WASM target only when __wasm_simd128__ is defined, which means passing -msimd128. Without it everything still compiles and every test still passes — Highway falls back to EMU128, a scalar emulation of a 128-bit vector — and you get none of this. UchenML’s own --config=wasm had it wrong. Print your target; do not assume it.

What else was on the table

Let the compiler do it. Autovectorization is the option that costs nothing to try, and for simple loops it works — the linear layer two years ago needed no help. The problem is that it is a heuristic with no contract. Whether a loop vectorizes depends on the compiler, its version, the flags, whether the pointers are restrict-qualified and whether the trip count is known, and none of that is visible in the source. A refactor that looks harmless turns it off, and nothing fails. It is also overconstrained from the other side: to keep the vectorizer interested the loop has to stay in the shape it recognizes, which is a real restriction on how the rest of the code can be written. For softmax specifically it does not get off the ground — the middle pass calls std::exp, and the sum reduction needs reassociation the compiler is not allowed to do without -ffast-math.

std::simd. The standard answer, eventually. It is C++26, implementations are partial, and UchenML has to build under emscripten today. Beyond the timeline, the standard library is a common denominator by construction: no partial load carrying a documented never-faults guarantee like LoadN, and a permute vocabulary far short of what kernels reach for. It also only solves the source-portability half — there is no equivalent of compiling one source for several targets and choosing at run time.

A linear algebra library. Would have worked, and would have made this someone else’s project.

Verdict

Two years ago the post on optimizing the linear layer concluded that the compiler already vectorized that layer well enough that competing with it was an unnecessary exercise. Softmax is the counter-example, and the exp in its inner loop is the whole reason. Measured against plain scalar code, the Highway kernel is 2.7× on the contiguous axis.

Highway delivered what I wanted from it: one kernel source, no #ifdef, no architecture-specific branch, and generated code that is the intrinsic I would have written by hand.