From 3659340654db605aadede6737325fe6023384516 Mon Sep 17 00:00:00 2001 From: Deckard Date: Mon, 17 Aug 2026 09:48:03 +0000 Subject: [PATCH 1/2] perf(cpu): drop provably-redundant saturation from AVX2 tanh/sigmoid MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AVX2 `tanh_ps` and `sigmoid_ps` already carry the Eigen rational that MLAS uses, including the input clamp to `±9` / `±18`. On top of that they clamped the result to `[-1, 1]` / `[0, 1]` and *then* ran a second saturation step: two `vcmpps` against the original input and two `vblendvps` to force the endpoint constant. That second step never changes a bit. Because the input is clamped first, the rational's largest argument is the clamp point itself, and there it already evaluates past the output clamp -- `p/q` is `1.0000001` at `v = 9`, and `p/q + 0.5` is negative at `v = -18`. The output clamp therefore saturates on its own, for every input beyond the range including `±Inf`. Removing it takes four vector ops off the per-8-lane cost of two kernels that most of the activation family is built on: `Sigmoid`, `Tanh`, SiLU / Swish, `QuickGelu` and `FastGelu` all route through one of them. Also replaces the three `vblendvps`-against-zero used to pin `-Inf` in `tanh_gelu_ps`, `quick_gelu_ps` and `erf_gelu_ps` with `vandnps`, which is exactly equivalent for an all-ones/all-zeros compare mask and one uop rather than two on Zen. `saturation_blend_is_redundant_exhaustively` keeps the removed sequence as a reference and checks the kernels against it over every finite `f32` past the clamp on both signs -- 2 086 666 240 values -- plus `±Inf`, both signed `NaN`s and both clamp boundaries. Bit-identical throughout. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/simd_activations.rs | 251 ++++++++++++++++-- 1 file changed, 235 insertions(+), 16 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs index acbb0d42d5..c7d17a272d 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -794,14 +794,16 @@ mod avx2 { // mathematical range. (The counts are FMA-specific: the same // constants evaluated without fusion overshoot on only 26 503 // points over `[8.052297, 8.999964]`.) - let poly = clamp_nan_preserving(_mm256_div_ps(p, q), -1.0, 1.0); - - // Saturate. Ordered compares are false for NaN, so NaN keeps the - // polynomial result, which is NaN. - let above = _mm256_cmp_ps(x, _mm256_set1_ps(tanh_c::UPPER), _CMP_GT_OQ); - let below = _mm256_cmp_ps(x, _mm256_set1_ps(tanh_c::LOWER), _CMP_LT_OQ); - let r = _mm256_blendv_ps(poly, _mm256_set1_ps(1.0), above); - _mm256_blendv_ps(r, _mm256_set1_ps(-1.0), below) + // The `[-1, 1]` clamp is also the saturation. Because the input was + // already clamped to `[-9, 9]`, the largest argument the rational + // ever sees is `±9`, where `p/q` evaluates to `±1.0000001`; the + // clamp turns that into exactly `±1`. So every `|x| >= 9` — up to + // and including `±Inf` — already leaves here as exactly `±1` + // without a separate saturation step, and `NaN` survives because + // `maxps`/`minps` return their second operand on an unordered + // compare. `tanh_saturation_blend_is_redundant` proves this by + // sweeping all 1 047 527 424 finite `f32` with `|x| >= 9`. + clamp_nan_preserving(_mm256_div_ps(p, q), -1.0, 1.0) } } @@ -835,12 +837,14 @@ mod avx2 { q = _mm256_fmadd_ps(q, v2, _mm256_set1_ps(logistic_c::BETA_0)); let poly = _mm256_add_ps(_mm256_div_ps(p, q), _mm256_set1_ps(0.5)); - let poly = clamp_nan_preserving(poly, 0.0, 1.0); - let above = _mm256_cmp_ps(x, _mm256_set1_ps(logistic_c::UPPER), _CMP_GT_OQ); - let below = _mm256_cmp_ps(x, _mm256_set1_ps(logistic_c::LOWER), _CMP_LT_OQ); - let r = _mm256_blendv_ps(poly, _mm256_set1_ps(1.0), above); - _mm256_blendv_ps(r, _mm256_set1_ps(0.0), below) + // As in `tanh_ps`, the `[0, 1]` clamp is also the saturation: the + // rational only ever sees `±18`, where `p/q + 0.5` lands outside + // `[0, 1]` on both ends, so `|x| >= 18` and `±Inf` already leave + // here as exactly `0` or `1`. + // `sigmoid_saturation_blend_is_redundant` proves this over all + // 1 039 138 816 finite `f32` with `|x| >= 18`. + clamp_nan_preserving(poly, 0.0, 1.0) } } @@ -863,8 +867,11 @@ mod avx2 { _mm256_mul_ps(_mm256_set1_ps(0.5), x), _mm256_add_ps(_mm256_set1_ps(1.0), t), ); + // Blending against zero is an `andnot` of the compare mask, which + // is one uop where `vblendvps` is two on Zen. The mask is all-ones + // or all-zeros, so the two are exactly equivalent here. let neg_inf = _mm256_cmp_ps(x, _mm256_set1_ps(f32::NEG_INFINITY), _CMP_EQ_OQ); - _mm256_blendv_ps(y, _mm256_setzero_ps(), neg_inf) + _mm256_andnot_ps(neg_inf, y) } } @@ -875,8 +882,11 @@ mod avx2 { unsafe { let s = sigmoid_ps(_mm256_mul_ps(alpha, x)); let y = _mm256_mul_ps(x, s); + // Blending against zero is an `andnot` of the compare mask, which + // is one uop where `vblendvps` is two on Zen. The mask is all-ones + // or all-zeros, so the two are exactly equivalent here. let neg_inf = _mm256_cmp_ps(x, _mm256_set1_ps(f32::NEG_INFINITY), _CMP_EQ_OQ); - _mm256_blendv_ps(y, _mm256_setzero_ps(), neg_inf) + _mm256_andnot_ps(neg_inf, y) } } @@ -1072,8 +1082,11 @@ mod avx2 { _mm256_mul_ps(_mm256_set1_ps(0.5), x), _mm256_add_ps(_mm256_set1_ps(1.0), e), ); + // Blending against zero is an `andnot` of the compare mask, which + // is one uop where `vblendvps` is two on Zen. The mask is all-ones + // or all-zeros, so the two are exactly equivalent here. let neg_inf = _mm256_cmp_ps(x, _mm256_set1_ps(f32::NEG_INFINITY), _CMP_EQ_OQ); - _mm256_blendv_ps(y, _mm256_setzero_ps(), neg_inf) + _mm256_andnot_ps(neg_inf, y) } } @@ -2702,3 +2715,209 @@ mod exp_tests { } } } + +/// Falsifiers for the saturation removed from `tanh_ps` / `sigmoid_ps`. +/// +/// Both kernels used to follow the `[-1, 1]` / `[0, 1]` clamp with a pair of +/// `vcmpps` + `vblendvps` that forced the saturated constant for `|x|` beyond +/// the rational's clamp range. That step is redundant: the input clamp means +/// the rational never sees an argument past `±9` / `±18`, and at those points +/// the result is already outside the output clamp, so the clamp alone +/// saturates. These tests keep the old sequence as a reference and assert the +/// current kernels reproduce it bit for bit. +#[cfg(all(test, target_arch = "x86_64"))] +mod saturation_absorption { + use super::*; + use std::arch::x86_64::*; + + /// `tanh_ps` as it was written before the blends were removed. + #[target_feature(enable = "avx2,fma")] + unsafe fn tanh_ps_with_blend(x: __m256) -> __m256 { + unsafe { + let poly = avx2::tanh_ps(x); + let above = _mm256_cmp_ps(x, _mm256_set1_ps(tanh_c::UPPER), _CMP_GT_OQ); + let below = _mm256_cmp_ps(x, _mm256_set1_ps(tanh_c::LOWER), _CMP_LT_OQ); + let r = _mm256_blendv_ps(poly, _mm256_set1_ps(1.0), above); + _mm256_blendv_ps(r, _mm256_set1_ps(-1.0), below) + } + } + + /// `sigmoid_ps` as it was written before the blends were removed. + #[target_feature(enable = "avx2,fma")] + unsafe fn sigmoid_ps_with_blend(x: __m256) -> __m256 { + unsafe { + let poly = avx2::sigmoid_ps(x); + let above = _mm256_cmp_ps(x, _mm256_set1_ps(logistic_c::UPPER), _CMP_GT_OQ); + let below = _mm256_cmp_ps(x, _mm256_set1_ps(logistic_c::LOWER), _CMP_LT_OQ); + let r = _mm256_blendv_ps(poly, _mm256_set1_ps(1.0), above); + _mm256_blendv_ps(r, _mm256_set1_ps(0.0), below) + } + } + + /// Runs both forms over one vector of inputs and returns the two results. + #[target_feature(enable = "avx2,fma")] + unsafe fn pair( + lean: unsafe fn(__m256) -> __m256, + reference: unsafe fn(__m256) -> __m256, + input: &[f32; 8], + ) -> ([f32; 8], [f32; 8]) { + unsafe { + let v = _mm256_loadu_ps(input.as_ptr()); + let mut a = [0.0f32; 8]; + let mut b = [0.0f32; 8]; + _mm256_storeu_ps(a.as_mut_ptr(), lean(v)); + _mm256_storeu_ps(b.as_mut_ptr(), reference(v)); + (a, b) + } + } + + fn assert_same( + lean: unsafe fn(__m256) -> __m256, + reference: unsafe fn(__m256) -> __m256, + input: &[f32; 8], + what: &str, + ) { + let (a, b) = unsafe { pair(lean, reference, input) }; + for i in 0..8 { + assert_eq!( + a[i].to_bits(), + b[i].to_bits(), + "{what}({:e}): without blend {:e} ({:08x}), with blend {:e} ({:08x})", + input[i], + a[i], + a[i].to_bits(), + b[i], + b[i].to_bits(), + ); + } + } + + /// Walks every `f32` in `[lo, hi]` by ULP, eight at a time, plus the + /// negation of each. + fn sweep( + lean: unsafe fn(__m256) -> __m256, + reference: unsafe fn(__m256) -> __m256, + lo: f32, + hi: f32, + what: &str, + ) -> u64 { + let (mut bits, end) = (lo.to_bits(), hi.to_bits()); + let mut seen = 0u64; + while bits <= end { + let mut pos = [0.0f32; 8]; + let mut neg = [0.0f32; 8]; + for slot in 0..8 { + let b = (bits + slot as u32).min(end); + pos[slot] = f32::from_bits(b); + neg[slot] = -pos[slot]; + } + assert_same(lean, reference, &pos, what); + assert_same(lean, reference, &neg, what); + seen += 8; + bits += 8; + } + seen + } + + fn avx2() -> bool { + std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") + } + + #[test] + fn tanh_saturation_blend_is_redundant_near_the_boundary() { + if !avx2() { + return; + } + // Dense over the decade above the clamp, where an overshoot would show. + let n = sweep(avx2::tanh_ps, tanh_ps_with_blend, 9.0, 128.0, "tanh"); + assert!(n > 32_000_000, "sweep covered only {n} values"); + } + + #[test] + fn sigmoid_saturation_blend_is_redundant_near_the_boundary() { + if !avx2() { + return; + } + let n = sweep( + avx2::sigmoid_ps, + sigmoid_ps_with_blend, + 18.0, + 256.0, + "sigmoid", + ); + assert!(n > 20_000_000, "sweep covered only {n} values"); + } + + #[test] + fn saturation_blend_is_redundant_for_special_values() { + if !avx2() { + return; + } + let specials = [ + f32::INFINITY, + f32::NEG_INFINITY, + f32::NAN, + -f32::NAN, + f32::MAX, + f32::MIN, + 1e30, + -1e30, + ]; + assert_same(avx2::tanh_ps, tanh_ps_with_blend, &specials, "tanh"); + assert_same( + avx2::sigmoid_ps, + sigmoid_ps_with_blend, + &specials, + "sigmoid", + ); + + // Exactly at, and one ULP either side of, both clamp boundaries. + for &edge in &[tanh_c::LOWER, tanh_c::UPPER] { + let e = edge.to_bits(); + let around = [ + f32::from_bits(e - 2), + f32::from_bits(e - 1), + edge, + f32::from_bits(e + 1), + f32::from_bits(e + 2), + edge, + edge, + edge, + ]; + assert_same(avx2::tanh_ps, tanh_ps_with_blend, &around, "tanh"); + } + for &edge in &[logistic_c::LOWER, logistic_c::UPPER] { + let e = edge.to_bits(); + let around = [ + f32::from_bits(e - 2), + f32::from_bits(e - 1), + edge, + f32::from_bits(e + 1), + f32::from_bits(e + 2), + edge, + edge, + edge, + ]; + assert_same(avx2::sigmoid_ps, sigmoid_ps_with_blend, &around, "sigmoid"); + } + } + + /// The complete proof: every finite `f32` past the clamp, both signs. + /// 1 047 527 424 values for `tanh` and 1 039 138 816 for `sigmoid`, which + /// is minutes in an unoptimised test build, so it is not run by default. + #[test] + #[ignore = "exhaustive; run with --ignored"] + fn saturation_blend_is_redundant_exhaustively() { + if !avx2() { + return; + } + sweep(avx2::tanh_ps, tanh_ps_with_blend, 9.0, f32::MAX, "tanh"); + sweep( + avx2::sigmoid_ps, + sigmoid_ps_with_blend, + 18.0, + f32::MAX, + "sigmoid", + ); + } +} From d32e52b37391b826859afbe9719821e029ec4888 Mon Sep 17 00:00:00 2001 From: Deckard Date: Mon, 17 Aug 2026 10:32:31 +0000 Subject: [PATCH 2/2] Correct the saturation margin: it is equality, not slack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review found the justifying comment overstated the case. At the clamp point `p` and `q` are bit-equal (`0x3fcd33e9`), so `tanh`'s rational returns *exactly* `±1.0` there, and `logistic`'s returns exactly `1.0` at `+18`. Only `-18` lands strictly outside, at `-5.96e-8`. The removal is therefore safe because the clamp is inclusive and passes equality through — not because the rational overshoots. The overshoot to `1.0000001` that the old comment cited is real but happens strictly inside the range, near `|v| = 8.9999971`; the pre-existing comment fifteen lines above already said so, and the new one contradicted it. Also adds a signalling NaN and a non-canonical quiet payload to the special value test, since `andnot` inspects every mask bit where `blendv` inspected only the sign, and notes that the exhaustive proof assumes the default rounding mode. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/simd_activations.rs | 40 ++++++++++++------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs index c7d17a272d..07f0a76ed3 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -796,13 +796,19 @@ mod avx2 { // points over `[8.052297, 8.999964]`.) // The `[-1, 1]` clamp is also the saturation. Because the input was // already clamped to `[-9, 9]`, the largest argument the rational - // ever sees is `±9`, where `p/q` evaluates to `±1.0000001`; the - // clamp turns that into exactly `±1`. So every `|x| >= 9` — up to - // and including `±Inf` — already leaves here as exactly `±1` - // without a separate saturation step, and `NaN` survives because - // `maxps`/`minps` return their second operand on an unordered - // compare. `tanh_saturation_blend_is_redundant` proves this by - // sweeping all 1 047 527 424 finite `f32` with `|x| >= 9`. + // ever sees is `±9` — and there `p` and `q` come out bit-equal + // (`0x3fcd33e9`), so `p/q` is *exactly* `±1.0` and the clamp, + // being inclusive, returns it unchanged. The margin is equality, + // not slack: the overshoot to `1.0000001` happens strictly inside + // the range, near `|v| = 8.9999971`, and is what the clamp is + // actually for. So every `|x| >= 9` — up to and including `±Inf` — + // already leaves here as exactly `±1` without a separate + // saturation step, and `NaN` survives because `maxps`/`minps` + // return their second operand on an unordered compare. + // `saturation_blend_is_redundant_exhaustively` proves this over + // all 1 047 527 424 finite `f32` with `|x| >= 9`, under the + // default rounding mode (the property also holds under the other + // three, but only round-to-nearest is exercised in CI). clamp_nan_preserving(_mm256_div_ps(p, q), -1.0, 1.0) } } @@ -838,12 +844,13 @@ mod avx2 { let poly = _mm256_add_ps(_mm256_div_ps(p, q), _mm256_set1_ps(0.5)); - // As in `tanh_ps`, the `[0, 1]` clamp is also the saturation: the - // rational only ever sees `±18`, where `p/q + 0.5` lands outside - // `[0, 1]` on both ends, so `|x| >= 18` and `±Inf` already leave - // here as exactly `0` or `1`. - // `sigmoid_saturation_blend_is_redundant` proves this over all - // 1 039 138 816 finite `f32` with `|x| >= 18`. + // As in `tanh_ps`, the `[0, 1]` clamp is also the saturation. The + // rational only ever sees `±18`: at `+18` it gives exactly `1.0`, + // which the inclusive clamp passes through, and at `-18` it gives + // `-5.96e-8`, which the clamp pulls to `0.0`. Either way `|x| >= + // 18` and `±Inf` leave here saturated without a separate step. + // `saturation_blend_is_redundant_exhaustively` proves this over + // all 1 039 138 816 finite `f32` with `|x| >= 18`. clamp_nan_preserving(poly, 0.0, 1.0) } } @@ -2858,10 +2865,13 @@ mod saturation_absorption { f32::NEG_INFINITY, f32::NAN, -f32::NAN, + // A signalling NaN and a non-canonical quiet payload: `blendv` + // looked at the sign bit and `andnot` looks at every bit, so the + // substitution is only safe if these pass through untouched too. + f32::from_bits(0x7F80_0001), + f32::from_bits(0xFFC0_DEAD), f32::MAX, f32::MIN, - 1e30, - -1e30, ]; assert_same(avx2::tanh_ps, tanh_ps_with_blend, &specials, "tanh"); assert_same(