gelu: saturate the fused f32 kernel's tanh to exactly -1 - #2582
Open
czoli1976 wants to merge 1 commit into
Open
Conversation
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>
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.
Fixes #2569.
arm64simd_gelu_f32_4n_fusedclamps 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 + tanhtherefore never cancels, andgelu(x)comes out as0.5 * x * 2^-23instead of 0 — an error that grows linearly with |x| and is unbounded.Measured before:
-5.36e-7at x=-9,-5.96e-5at x=-1000,-3.90e-3at x=-65504, against-0from the scalar kernel and from libm.The fix
The lanes that need correcting are exactly the ones the low clamp pinned, so
fcmeqagainst the clamp constant identifies them andbslsubstitutes an exact-1.0before the final combine.v12-v15were 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_negativescovers 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:
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
-9.30does land on exactly-1— but the neighbourhood is noisy under fma rounding (-9.28and-9.32do not), so it would be luck rather than a property.+0where the scalar kernel returns-0, because the combine isfma(0.5x, tanh, 0.5x)anda - ais+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: thefmlarounds once where the reference rounds twice.Tests
tract-linalg 4415 — green.
cargo fmt --allandcargo clippyclean.🍍