SIMD Softmax, Neon to WASM, with Google Highway

SIMD Softmax, NEON to WASM, with Google Highway — UchenML 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: softmax(xi)=exi∑jnexj 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::VFromDD>; const size_t lanes = hn::Lanes(d); const V neg_inf = hn::Set(d, -std::numeric_limitsfloat>::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 — 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, a 256-bit one under AVX2’s ScalableTag and a 512-bit one under...

SIMD Softmax, Neon to WASM, with Google Highway

Related Articles

Incident with Github.com

(no title)

Scientists reverse brain aging, with a nasal spray

Malware infects Android-based automotive head unit firmware

NHS apologises and admits Palantir have access to identifiable patient data