fix(vml, simd): sign of zero, exp range/±inf, round ties-to-even; exp bit-identical on every backend incl. wasm

crates/wasm-simd-parity/src/lib.rs

@@ -35,9 +35,85 @@ pub extern "C" fn selfcheck() -> u32 {
     if let Err(code) = check_i32x16_compare() {
         return code;
     }
+    if let Err(code) = check_exp_and_round() {
+        return code;
+    }
     0
 }
 
+/// `simd_exp_f32` and `F32x16::round` under wasm. exp must equal, bit for
+/// bit, the same unfused op sequence in scalar f32; the native unit test
+/// `simd_exp_f32_matches_the_unfused_scalar_sequence_bit_for_bit` pins the
+/// native backends to that sequence, so together they give wasm == native.
+/// `round` must break ties to even on every backend.
+fn check_exp_and_round() -> Result<(), u32> {
+    // Constants of `simd_exp_f32` (private in ndarray::simd), restated.
+    const OVERFLOW_X: f32 = 88.722_84;
+    const LN2_HI: f32 = 0.693_359_4;
+    const LN2_LO: f32 = -2.121_944_4e-4;
+    const TWO_POW_NEG_64: f32 = 5.421_011e-20;
+    fn reference(x: f32) -> f32 {
+        if x.is_nan() {
+            return f32::NAN;
+        }
+        let xc = x.max(-110.0).min(OVERFLOW_X);
+        let n = (xc * (1.0 / core::f32::consts::LN_2) + 0.5).floor();
+        let r = xc - n * LN2_HI;
+        let r = r - n * LN2_LO;
+        let poly = 1.0 + r * (1.0 + r * (0.5 + r * (1.0 / 6.0 + r * (1.0 / 24.0 + r * (1.0 / 120.0)))));
+        let (n, post) = if n > 127.0 {
+            (n - 1.0, 2.0)
+        } else if n < -126.0 {
+            (n + 64.0, TWO_POW_NEG_64)
+        } else {
+            (n, 1.0)
+        };
+        let scale = f32::from_bits((((n as i32).clamp(-126, 127) + 127) as u32) << 23);
+        poly * scale * post
+    }
+    // -110 .. 88.72 in steps of 0.01, then ties of x/ln2 + 0.5 and specials.
+    let mut xs: Vec<f32> = (0..=19872).map(|j| -110.0 + j as f32 * 0.01).collect();
+    xs.extend((-150..=128).map(|k| (k as f32 + 0.5) * core::f32::consts::LN_2));
+    xs.extend([f32::INFINITY, f32::NEG_INFINITY, f32::NAN, OVERFLOW_X, -0.0, 0.0, f32::MIN_POSITIVE]);
+    while xs.len() % 16 != 0 {
+        xs.push(1.0);
+    }
+    for chunk in xs.chunks_exact(16) {
+        let got = ndarray::simd::simd_exp_f32(F32x16::from_slice(chunk)).to_array();
+        for (&x, g) in chunk.iter().zip(got) {
+            let want = reference(x);
+            if g.to_bits() != want.to_bits() && !(g.is_nan() && want.is_nan()) {
+                return Err(70);
+            }
+        }
+    }
+    let ties: [f32; 16] = [
+        0.5,
+        1.5,
+        2.5,
+        -0.5,
+        -1.5,
+        -2.5,
+        3.5,
+        4.5,
+        -3.5,
+        0.25,
+        0.75,
+        -0.0,
+        1e7 + 0.5,
+        8388607.5,
+        -8388607.5,
+        0.0,
+    ];
+    let r = F32x16::from_array(ties).round().to_array();
+    for k in 0..16 {
+        if r[k].to_bits() != ties[k].round_ties_even().to_bits() {
+            return Err(71);
+        }
+    }
+    Ok(())
+}
+
 /// `U32x16` ARX triple (Add / BitXor / rotate_left) — the ChaCha20/BLAKE lane.
 fn check_u32x16() -> Result<(), u32> {
     let a_arr: [u32; 16] = [

src/hpc/gpt2/inference.rs

@@ -453,7 +453,9 @@ fn gelu_simd(x: &mut [f32]) {
         // tanh approximation via exp: tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
         let two_inner = inner + inner;
         let exp_2x = crate::simd::simd_exp_f32(two_inner);
-        let tanh_v = (exp_2x - one) / (exp_2x + one);
+        // 1 - 2/(e+1) equals (e-1)/(e+1) but stays 1 when exp overflows to
+        // +inf, where the quotient form gives inf/inf = NaN.
+        let tanh_v = one - (one + one) / (exp_2x + one);
         let result = v * half * (one + tanh_v);
         result.copy_to_slice(&mut x[off..off + 16]);
     }
@@ -522,6 +524,20 @@ fn matmul_vec_simd(input: &[f32], weight: &[f32], bias: &[f32], output: &mut [f3
 mod tests {
     use super::*;
 
+    /// FAILS IF: large inputs, where exp(2·inner) overflows to +inf, come
+    /// out NaN from the SIMD body instead of x (positive) or ~0 (negative).
+    #[test]
+    fn gelu_stays_finite_where_exp_overflows() {
+        for v in [12.0f32, 20.0, 1e3, -12.0, -20.0, -1e3] {
+            let mut x = [v; 17];
+            gelu_simd(&mut x);
+            for (i, &g) in x.iter().enumerate() {
+                let want = if v > 0.0 { v } else { 0.0 };
+                assert!((g - want).abs() <= 1e-4 * v.abs(), "gelu({v})[{i}] = {g}");
+            }
+        }
+    }
+
     #[test]
     fn test_layer_norm_identity() {
         let mut x = vec![1.0, 2.0, 3.0, 4.0];

src/hpc/models/layers.rs

@@ -111,7 +111,9 @@ pub fn gelu(x: &mut [f32]) {
         let inner = sqrt_2_over_pi * (v + coeff * v3);
         let two_inner = inner + inner;
         let exp_2x = crate::simd::simd_exp_f32(two_inner);
-        let tanh_v = (exp_2x - one) / (exp_2x + one);
+        // 1 - 2/(e+1) equals (e-1)/(e+1) but stays 1 when exp overflows to
+        // +inf, where the quotient form gives inf/inf = NaN.
+        let tanh_v = one - (one + one) / (exp_2x + one);
         let result = v * half * (one + tanh_v);
         result.copy_to_slice(&mut x[off..off + 16]);
     }
@@ -283,6 +285,20 @@ pub fn rope_apply(q: &mut [f32], k: &mut [f32], head_dim: usize, position: usize
 mod tests {
     use super::*;
 
+    /// FAILS IF: large inputs, where exp(2·inner) overflows to +inf, come
+    /// out NaN from the SIMD body instead of x (positive) or ~0 (negative).
+    #[test]
+    fn gelu_stays_finite_where_exp_overflows() {
+        for v in [12.0f32, 20.0, 1e3, -12.0, -20.0, -1e3] {
+            let mut x = [v; 17];
+            gelu(&mut x);
+            for (i, &g) in x.iter().enumerate() {
+                let want = if v > 0.0 { v } else { 0.0 };
+                assert!((g - want).abs() <= 1e-4 * v.abs(), "gelu({v})[{i}] = {g}");
+            }
+        }
+    }
+
     #[test]
     fn test_layer_norm_zero_mean() {
         let mut x = vec![1.0, 2.0, 3.0, 4.0];

src/hpc/vml.rs

@@ -243,7 +243,7 @@ fn vstanh_slice(x: &[f32], out: &mut [f32]) {
         // tanh(x) = 2·sigmoid(2x) - 1
         let two_x = v * two;
         // sigmoid(z) = 1/(1+exp(-z))
-        let neg_two_x = F32x16::splat(0.0) - two_x;
+        let neg_two_x = -two_x;
         let exp_neg = simd_exp_f32(neg_two_x);
         let sigmoid = one / (exp_neg + one);
         let tanh_v = sigmoid.mul_add(two, F32x16::splat(-1.0));
@@ -273,11 +273,9 @@ fn vsceil_slice(x: &[f32], out: &mut [f32]) {
     for i in 0..chunks {
         let off = i * 16;
         let v = F32x16::from_slice(&x[off..off + 16]);
-        // ceil = -floor(-x)
-        let neg = F32x16::splat(0.0) - v;
-        let floored = neg.floor();
-        let ceiled = F32x16::splat(0.0) - floored;
-        ceiled.copy_to_slice(&mut out[off..off + 16]);
+        // ceil = -floor(-x). `Neg` flips the sign bit; `0.0 - v` would turn
+        // ceil(-0.5) and ceil(-0.0) into +0.0.
+        (-(-v).floor()).copy_to_slice(&mut out[off..off + 16]);
     }
     for i in (chunks * 16)..n {
         out[i] = x[i].ceil();
@@ -300,11 +298,11 @@ fn vsround_slice(x: &[f32], out: &mut [f32]) {
 fn vsneg_slice(x: &[f32], out: &mut [f32]) {
     let n = x.len().min(out.len());
     let chunks = n / 16;
-    let zero = F32x16::splat(0.0);
     for i in 0..chunks {
         let off = i * 16;
         let v = F32x16::from_slice(&x[off..off + 16]);
-        (zero - v).copy_to_slice(&mut out[off..off + 16]);
+        // `Neg` flips the sign bit, so -(+0.0) is -0.0; `0.0 - v` gives +0.0.
+        (-v).copy_to_slice(&mut out[off..off + 16]);
     }
     for i in (chunks * 16)..n {
         out[i] = -x[i];
@@ -317,14 +315,12 @@ fn vstrunc_slice(x: &[f32], out: &mut [f32]) {
     for i in 0..chunks {
         let off = i * 16;
         let v = F32x16::from_slice(&x[off..off + 16]);
-        let abs_v = v.abs();
-        let floored = abs_v.floor();
-        let zero = F32x16::splat(0.0);
-        let mask = v.simd_lt(zero); // true where x < 0
-        let pos_result = floored;
-        let neg_result = zero - floored;
-        let result = mask.select(neg_result, pos_result);
-        result.copy_to_slice(&mut out[off..off + 16]);
+        // Negative lanes take -floor(-x), the rest floor(x). -0.0 is not
+        // < 0, so it takes floor(-0.0) = -0.0 and keeps its sign.
+        let negative = v.simd_lt(F32x16::splat(0.0));
+        negative
+            .select(-(-v).floor(), v.floor())
+            .copy_to_slice(&mut out[off..off + 16]);
     }
     for i in (chunks * 16)..n {
         out[i] = x[i].trunc();
@@ -688,7 +684,7 @@ pub fn vstrunc<D: Dimension>(x: ArrayView<'_, f32, D>, mut out: ArrayViewMut<'_,
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{arr1, arr2, s, Array1, Array2};
+    use crate::{arr1, arr2, s, Array1, Array2, ArrayView1, ArrayViewMut1};
 
     // -----------------------------------------------------------------------
     // vsexp
@@ -1215,4 +1211,112 @@ mod tests {
         vstrunc(strided, out.view_mut());
         assert_eq!(out.as_slice().unwrap(), &[1.0, -1.0, 0.0, 0.0]);
     }
+
+    // -----------------------------------------------------------------------
+    // IEEE edge cases through the 16-lane SIMD body, not just the scalar tail.
+    // Every length from 0 to 40 puts each value both in a SIMD chunk and in
+    // the tail, so a body that disagrees with its own tail cannot pass.
+    // -----------------------------------------------------------------------
+
+    const EDGES: [f32; 12] = [
+        0.0,
+        -0.0,
+        0.5,
+        -0.5,
+        1.5,
+        -1.5,
+        -0.25,
+        f32::INFINITY,
+        f32::NEG_INFINITY,
+        f32::MIN_POSITIVE,
+        -f32::MIN_POSITIVE,
+        f32::NAN,
+    ];
+
+    /// Checks `kernel` against `scalar` bit for bit (NaN only as NaN) for
+    /// every length 0..=40, cycling through `EDGES`.
+    fn assert_bitwise_like_scalar(
+        name: &str, kernel: fn(ArrayView1<'_, f32>, ArrayViewMut1<'_, f32>), scalar: fn(f32) -> f32,
+    ) {
+        for n in 0..=40usize {
+            let x = Array1::from_iter((0..n).map(|i| EDGES[i % EDGES.len()]));
+            let mut out = Array1::<f32>::zeros(n);
+            kernel(x.view(), out.view_mut());
+            for i in 0..n {
+                let (got, want) = (out[i], scalar(x[i]));
+                if want.is_nan() {
+                    assert!(got.is_nan(), "{name}({}) n={n} i={i}: got {got}", x[i]);
+                } else {
+                    assert_eq!(
+                        got.to_bits(),
+                        want.to_bits(),
+                        "{name}({}) n={n} i={i}: got {got:?}, want {want:?}",
+                        x[i]
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn vsneg_flips_the_sign_of_zero_in_the_simd_body() {
+        assert_bitwise_like_scalar("vsneg", vsneg, |v| -v);
+    }
+
+    #[test]
+    fn vsceil_keeps_negative_zero_in_the_simd_body() {
+        assert_bitwise_like_scalar("vsceil", vsceil, f32::ceil);
+    }
+
+    #[test]
+    fn vstrunc_keeps_negative_zero_in_the_simd_body() {
+        assert_bitwise_like_scalar("vstrunc", vstrunc, f32::trunc);
+    }
+
+    /// FAILS IF: the SIMD body breaks ties differently from the scalar tail
+    /// (`round_ties_even`). On the AVX2, scalar and nightly arms
+    /// `F32x16::round` used `f32::round`, which rounds ties away from zero.
+    #[test]
+    fn vsround_breaks_ties_to_even_in_the_simd_body() {
+        assert_bitwise_like_scalar("vsround", vsround, f32::round_ties_even);
+    }
+
+    #[test]
+    fn vsexp_maps_infinities_like_the_scalar_tail() {
+        for n in [1usize, 16, 17, 32] {
+            for (v, want) in [(f32::INFINITY, f32::INFINITY), (f32::NEG_INFINITY, 0.0)] {
+                let x = Array1::from_elem(n, v);
+                let mut out = Array1::<f32>::zeros(n);
+                vsexp(x.view(), out.view_mut());
+                for i in 0..n {
+                    assert_eq!(out[i].to_bits(), want.to_bits(), "vsexp({v}) n={n} i={i}: got {:?}", out[i]);
+                }
+            }
+        }
+    }
+
+    /// FAILS IF: any finite input in [-110, 88.72] is off by more than the
+    /// polynomial's own error. Measured against f64 `exp`. Results below the
+    /// normal range may also be off by one subnormal step.
+    #[test]
+    fn vsexp_is_accurate_up_to_the_overflow_and_down_into_subnormals() {
+        let xs: Vec<f32> = (0..=19872)
+            .map(|i| -110.0 + i as f32 * 0.01)
+            .filter(|&x| x <= 88.72)
+            .collect();
+        let x = Array1::from(xs);
+        let mut out = Array1::<f32>::zeros(x.len());
+        vsexp(x.view(), out.view_mut());
+        for (i, &v) in x.iter().enumerate() {
+            let want = f64::from(v).exp();
+            let got = f64::from(out[i]);
+            if want >= f64::from(f32::MIN_POSITIVE) {
+                assert!(((got - want) / want).abs() < 4e-6, "vsexp({v}): got {got:e}, want {want:e}");
+            } else {
+                // the polynomial's relative error plus one subnormal step
+                let step = f64::from(f32::from_bits(1));
+                assert!((got - want).abs() <= 4e-6 * want + step, "vsexp({v}): got {got:e}, want {want:e}");
+            }
+        }
+    }
 }