linalg/wasm: fused simd128 f32 kernels for gelu, silu and erf - #2591
Open
czoli1976 wants to merge 1 commit into
Open
linalg/wasm: fused simd128 f32 kernels for gelu, silu and erf#2591czoli1976 wants to merge 1 commit into
czoli1976 wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On plain
+simd128builds (stable toolchain, no relaxed-simd) thegelu_f32,silu_f32anderf_f32slots 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 inwasm::plug.Numbers
Through
ElementWiseImplon 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):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
+simd128and 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 thef32x4.divlatency). Only the relaxed-simd FMA variants inact.rsbeat 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_clampcovers it out to ±1e6.Tests
Frame tests for all three kernels plus the tail test, 489 green on
wasm32-wasip1 +simd128under wasmtime; the+relaxed-simdconfiguration still installs its sigmoid/tanh overrides on top and itsactsuite is green; native andwasm32-unknown-unknown(no simd) builds are unaffected (linalg::wasmis feature-gated out).cargo fmt/clippyclean.🍍