Skip to content

gelu: saturate the fused f32 kernel's tanh to exactly -1 - #2582

Open
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:fix/gelu-f32-saturation
Open

gelu: saturate the fused f32 kernel's tanh to exactly -1#2582
czoli1976 wants to merge 1 commit into
sonos:mainfrom
czoli1976:fix/gelu-f32-saturation

Conversation

@czoli1976

Copy link
Copy Markdown
Contributor

Fixes #2569. arm64simd_gelu_f32_4n_fused clamps the pre-tanh argument to [-8.9, 8.9], but at the low bound the Padé polynomial evaluates to -(1 - 2^-23) rather than -1. 1 + tanh therefore never cancels, and gelu(x) comes out as 0.5 * x * 2^-23 instead of 0 — an error that grows linearly with |x| and is unbounded.

Measured before: -5.36e-7 at x=-9, -5.96e-5 at x=-1000, -3.90e-3 at x=-65504, against -0 from the scalar kernel and from libm.

The fix

The lanes that need correcting are exactly the ones the low clamp pinned, so fcmeq against the clamp constant identifies them and bsl substitutes an exact -1.0 before the final combine. v12-v15 were unused, so the masks need no spill. Nine instructions in the unrolled body, three in the tail; nothing else in the polynomial changes.

decays_to_zero_on_large_negatives covers both the unrolled body and the tail at six magnitudes.

Cost

This is not free, and the microbenchmark number looks worse than the real one:

before after
kernel, 1024 elems 630 ns 682 ns +8.2%
f32 FFN block (matmul -> gelu -> matmul) 2.244 ms 2.263 ms +0.85%

Both interleaved within one session rather than compared across runs — on this box a cross-run criterion baseline drifts far enough to invert the sign of a result this size, so I re-ran the two arms alternately and both rounds agreed. Sub-1% on a block where GEMMs dominate seemed a fair price for removing an unbounded error, but if you would rather not pay even that in the hot path I am happy to gate it.

Two things I did not do

  • Moving the clamp instead. Zero added instructions, and -9.30 does land on exactly -1 — but the neighbourhood is noisy under fma rounding (-9.28 and -9.32 do not), so it would be luck rather than a property.
  • Matching the sign of zero. Corrected lanes return +0 where the scalar kernel returns -0, because the combine is fma(0.5x, tanh, 0.5x) and a - a is +0, whereas the reference multiplies by zero. Making that exact costs another nine instructions per iteration. The two compare equal and behave identically downstream, so I left it — say the word if you want it exact. Note the fused kernel was never bit-identical to the scalar one anyway: the fmla rounds once where the reference rounds twice.

Tests

tract-linalg 4415 — green. cargo fmt --all and cargo clippy clean.

🍍

The kernel clamps the pre-tanh argument, but the Pade polynomial lands one ulp
short of -1 at the low bound, so 1 + tanh never cancels and gelu(x) comes out
as 0.5 * x * 2^-23 instead of 0 — an error that grows without bound as x goes
more negative. Substitute an exact -1 on the lanes the low clamp pinned, so
those lanes decay to zero the way the scalar kernel and libm do.

Co-Authored-By: Claude Opus 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.

arm64: fused f32 GELU returns 0.5*x*2^-23 instead of 0 for large negative x

1 participant