Skip to content

linalg/wasm: fused simd128 f32 kernels for gelu, silu and erf - #2591

Open
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:wasm-f32-activations
Open

linalg/wasm: fused simd128 f32 kernels for gelu, silu and erf#2591
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:wasm-f32-activations

Conversation

@czoli1976

Copy link
Copy Markdown
Contributor

On plain +simd128 builds (stable toolchain, no relaxed-simd) the gelu_f32, silu_f32 and erf_f32 slots fall back to generic kernels that go through libm (tanh, exp) per element; this adds fused single-pass f32x4 kernels for those three, reusing the scalar kernels' coefficients, and wires them in wasm::plug.

Numbers

Through ElementWiseImpl on 64k elements, wasmtime on x86-64, values in [-3, 3) refilled from a source buffer every iteration (iterated in-place application decays gelu/silu into subnormals and poisons the timing):

kernel generic simd128 speedup
gelu 25.6 ns/elem 3.8 ns/elem 6.7x
silu 6.9 ns/elem 2.6 ns/elem 2.6x
erf 1.65 ns/elem 1.44 ns/elem 1.14x

gelu and silu win big because their generic kernels call libm per element. erf's generic kernel is a pure polynomial that LLVM already auto-vectorizes, so the fused kernel only shaves the powi(16)/store schedule.

What is deliberately absent

No simd128 sigmoid or tanh. I wrote them, measured them, and deleted them: the generic polynomial loops auto-vectorize under +simd128 and beat the handwritten kernels (sigmoid 1.43 vs 2.54 ns/elem, tanh 1.30 vs 2.14 — the auto-vectorized loop unrolls wider and hides the f32x4.div latency). Only the relaxed-simd FMA variants in act.rs beat the generic baseline, and those already exist and keep their override.

The gelu tail

The kernel clamps the pre-tanh argument to [-8.9, 8.9] and the Padé polynomial lands one ulp short of -1 at the low bound, the same unbounded-error tail #2582 fixes on arm64. Rather than ship the bug and fix it later, the lanes the low clamp pinned substitute an exact -1.0 (f32x4_eq + v128_bitselect, two extra ops); gelu_saturates_to_zero_below_the_tanh_clamp covers it out to ±1e6.

Tests

Frame tests for all three kernels plus the tail test, 489 green on wasm32-wasip1 +simd128 under wasmtime; the +relaxed-simd configuration still installs its sigmoid/tanh overrides on top and its act suite is green; native and wasm32-unknown-unknown (no simd) builds are unaffected (linalg::wasm is feature-gated out). cargo fmt / clippy clean.

🍍

On plain +simd128 builds the gelu, silu and erf slots fall back to generic
kernels that go through libm per element; give them fused single-pass
f32x4 kernels reusing the scalar coefficients, with gelu's tanh saturated
to exactly -1 on low-clamped lanes so it decays to zero like the scalar
path. Sigmoid and tanh keep the generic polynomial kernels, which LLVM
auto-vectorizes into faster code than a handwritten simd128 loop.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant