diff --git a/Cargo.toml b/Cargo.toml index 8c8fa2e2..2d82d288 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,10 @@ required-features = ["std"] name = "bf16_rne_exhaustive" required-features = ["std"] +[[example]] +name = "vstanh_exhaustive" +required-features = ["std"] + [[example]] name = "cpu_guard_probe" required-features = ["std"] diff --git a/crates/wasm-simd-parity/src/lib.rs b/crates/wasm-simd-parity/src/lib.rs index c4dc637d..1dc5cd2a 100644 --- a/crates/wasm-simd-parity/src/lib.rs +++ b/crates/wasm-simd-parity/src/lib.rs @@ -52,6 +52,14 @@ fn check_exp_and_round() -> Result<(), u32> { 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; + const EXP_P: [f32; 6] = [ + 1.987_569_150_0e-4, + 1.398_199_950_7e-3, + 8.333_451_907_3e-3, + 4.166_579_589_4e-2, + 1.666_666_545_9e-1, + 5.000_000_120_1e-1, + ]; fn reference(x: f32) -> f32 { if x.is_nan() { return f32::NAN; @@ -60,7 +68,8 @@ fn check_exp_and_round() -> Result<(), u32> { 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 p = EXP_P.iter().skip(1).fold(EXP_P[0], |p, &c| p * r + c); + let poly = (p * (r * r) + r) + 1.0; let (n, post) = if n > 127.0 { (n - 1.0, 2.0) } else if n < -126.0 { diff --git a/examples/vstanh_exhaustive.rs b/examples/vstanh_exhaustive.rs new file mode 100644 index 00000000..3503308f --- /dev/null +++ b/examples/vstanh_exhaustive.rs @@ -0,0 +1,337 @@ +//! Exhaustive accuracy check of `ndarray::hpc::vml::vstanh` over every f32. +//! +//! Every non-negative bit pattern `p` in `0..=0x7fff_ffff` is evaluated +//! together with its negation, so the run covers all 2^32 inputs. The +//! reference is `tanh` in `f64` rounded once to `f32`. Inputs where that +//! single rounding could land on the wrong side of an f32 midpoint are +//! written to a file for an independent high-precision re-check. +//! +//! A second arm runs the previous kernel, `2·sigmoid(2x) − 1`, rebuilt from +//! the public `F32x16` and `simd_exp_f32`, as the baseline. +//! +//! ```text +//! cargo run --release --example vstanh_exhaustive [flip_risk.txt] [worst.txt] +//! ``` + +use std::fmt::Write as _; +use std::time::Instant; + +use ndarray::hpc::vml::vstanh; +use ndarray::simd::{simd_exp_f32, F32x16}; +use ndarray::{Array1, ArrayView1, ArrayViewMut1}; + +/// Signed ordering key: -0.0 and +0.0 one apart, monotone in value. +fn key(x: f32) -> i64 { + let i = x.to_bits() as i32; + if i < 0 { + -((i & i32::MAX) as i64) - 1 + } else { + i as i64 + } +} + +fn ulp(a: f32, b: f32) -> u64 { + if a.is_nan() || b.is_nan() { + return if a.is_nan() && b.is_nan() { 0 } else { u64::MAX }; + } + (key(a) - key(b)).unsigned_abs() +} + +const BANDS: [(&str, f32, f32); 11] = [ + ("zero", 0.0, 0.0), + ("subnormal", f32::from_bits(1), f32::MIN_POSITIVE), + ("min_pos..1e-10", f32::MIN_POSITIVE, 1e-10), + ("1e-10..1e-4", 1e-10, 1e-4), + ("1e-4..0.1", 1e-4, 0.1), + ("0.1..0.625", 0.1, 0.625), + ("0.625..1", 0.625, 1.0), + ("1..4", 1.0, 4.0), + ("4..10", 4.0, 10.0), + ("10..MAX", 10.0, f32::MAX), + ("inf", f32::INFINITY, f32::INFINITY), +]; + +fn band_of(a: f32) -> usize { + if a == 0.0 { + 0 + } else if a.is_infinite() { + 10 + } else if a < f32::MIN_POSITIVE { + 1 + } else { + (2..10).find(|&b| a < BANDS[b].2 || b == 9).unwrap() + } +} + +#[derive(Clone, Copy, Default)] +struct Band { + n: u64, + max_ulp: u64, + worst_x: f32, + worst_got: f32, + worst_want: f32, + max_abs: f64, + max_rel: f64, + zeroed: u64, + hist: [u64; 4], // 0, 1, 2, >2 ulp +} + +#[derive(Clone, Default)] +struct Stats { + bands: [Band; 11], + nan_ok: u64, + nan_bad: u64, + odd_bad: u64, + bound_bad: u64, + mono_bad: u64, + flip_risk: Vec, +} + +fn old_kernel(x: &[f32], out: &mut [f32]) { + for (xc, oc) in x + .as_chunks::<16>() + .0 + .iter() + .zip(out.as_chunks_mut::<16>().0.iter_mut()) + { + let v = F32x16::from_slice(xc); + let two = F32x16::splat(2.0); + let e = simd_exp_f32(-(v * two)); + let s = F32x16::splat(1.0) / (e + F32x16::splat(1.0)); + s.mul_add(two, F32x16::splat(-1.0)).copy_to_slice(oc); + } +} + +fn run(lo: u32, hi: u32, use_old: bool) -> Stats { + const CH: usize = 1 << 16; + let mut st = Stats::default(); + let mut xs = vec![0.0f32; CH]; + let mut nxs = vec![0.0f32; CH]; + let mut ys = vec![0.0f32; CH]; + let mut nys = vec![0.0f32; CH]; + let mut prev: Option = None; + let mut p = lo as u64; + while p < hi as u64 { + let m = (CH as u64).min(hi as u64 - p) as usize; + for i in 0..m { + xs[i] = f32::from_bits((p + i as u64) as u32); + nxs[i] = -xs[i]; + } + if use_old { + for i in m..CH { + xs[i] = 0.0; + nxs[i] = 0.0; + } + old_kernel(&xs, &mut ys); + old_kernel(&nxs, &mut nys); + } else { + vstanh(ArrayView1::from(&xs[..m]), ArrayViewMut1::from(&mut ys[..m])); + vstanh(ArrayView1::from(&nxs[..m]), ArrayViewMut1::from(&mut nys[..m])); + } + for i in 0..m { + let x = xs[i]; + if x.is_nan() { + if ys[i].is_nan() && nys[i].is_nan() { + st.nan_ok += 2; + } else { + st.nan_bad += 2; + } + prev = None; + continue; + } + // odd symmetry, bit for bit + if nys[i].to_bits() != (ys[i].to_bits() ^ 0x8000_0000) { + st.odd_bad += 1; + } + if !(-1.0..=1.0).contains(&ys[i]) || !(-1.0..=1.0).contains(&nys[i]) { + st.bound_bad += 1; + } + if let Some(q) = prev { + if ys[i] < q { + st.mono_bad += 1; + } + } + prev = Some(ys[i]); + let y64 = (x as f64).tanh(); + let want = y64 as f32; + // Flip risk: y64 within 4 f64-ulp of the midpoint between `want` + // and its neighbour away from y64 rounding. + if want.is_finite() && want != 0.0 && x.is_finite() { + let w = want as f64; + let nb = if y64 >= w { + f32::from_bits(want.to_bits() + 1) + } else { + f32::from_bits(want.to_bits() - 1) + } as f64; + let mid = 0.5 * (w + nb); + if (y64 - mid).abs() <= 4.0 * f64::EPSILON * y64.abs() && st.flip_risk.len() < 200_000 { + st.flip_risk.push(x); + } + } + let b = band_of(x); + for (xx, got, want, y64) in [(x, ys[i], want, y64), (-x, nys[i], -want, -y64)] { + let bd = &mut st.bands[b]; + bd.n += 1; + let u = ulp(got, want); + bd.hist[u.min(3) as usize] += 1; + if u > bd.max_ulp || bd.n == 1 { + bd.max_ulp = u; + (bd.worst_x, bd.worst_got, bd.worst_want) = (xx, got, want); + } + let abs = (got as f64 - y64).abs(); + bd.max_abs = bd.max_abs.max(abs); + if y64 != 0.0 { + bd.max_rel = bd.max_rel.max(abs / y64.abs()); + } + if got == 0.0 && xx != 0.0 { + bd.zeroed += 1; + } + } + } + p += m as u64; + } + st +} + +fn merge(mut a: Stats, b: Stats) -> Stats { + for (x, y) in a.bands.iter_mut().zip(b.bands) { + let x_empty = x.n == 0; + x.n += y.n; + if y.n > 0 && (x_empty || y.max_ulp > x.max_ulp) { + (x.max_ulp, x.worst_x, x.worst_got, x.worst_want) = (y.max_ulp, y.worst_x, y.worst_got, y.worst_want); + } + x.max_abs = x.max_abs.max(y.max_abs); + x.max_rel = x.max_rel.max(y.max_rel); + x.zeroed += y.zeroed; + for k in 0..4 { + x.hist[k] += y.hist[k]; + } + } + a.nan_ok += b.nan_ok; + a.nan_bad += b.nan_bad; + a.odd_bad += b.odd_bad; + a.bound_bad += b.bound_bad; + a.mono_bad += b.mono_bad; + a.flip_risk.extend(b.flip_risk); + a +} + +fn sweep(use_old: bool, threads: u32) -> Stats { + let end: u64 = 0x8000_0000; + let step = end / threads as u64; + let parts: Vec = std::thread::scope(|s| { + let hs: Vec<_> = (0..threads as u64) + .map(|t| { + let lo = (t * step) as u32; + let hi = if t + 1 == threads as u64 { end } else { (t + 1) * step }; + s.spawn(move || run(lo, hi as u32, use_old)) + }) + .collect(); + hs.into_iter().map(|h| h.join().unwrap()).collect() + }); + // Monotonicity across thread boundaries. + let mut st = parts.into_iter().reduce(merge).unwrap(); + for t in 1..threads as u64 { + let b = (t * step) as u32; + let (a, c) = (f32::from_bits(b - 1), f32::from_bits(b)); + let mut y = [0.0f32; 2]; + vstanh(ArrayView1::from(&[a, c][..]), ArrayViewMut1::from(&mut y[..])); + if !use_old && y[1] < y[0] { + st.mono_bad += 1; + } + } + st +} + +fn report(name: &str, st: &Stats) -> String { + let mut s = String::new(); + writeln!(s, "== {name} ==").unwrap(); + writeln!( + s, + "{:<16} {:>12} {:>10} {:>11} {:>11} {:>9} {:>10} {:>10} {:>8} {:>8} worst x -> got / want", + "band", "n", "max_ulp", "max_abs", "max_rel", "zeroed", "0ulp", "1ulp", "2ulp", ">2ulp" + ) + .unwrap(); + for (b, (nm, _, _)) in st.bands.iter().zip(BANDS) { + writeln!( + s, + "{:<16} {:>12} {:>10} {:>11.3e} {:>11.3e} {:>9} {:>10} {:>10} {:>8} {:>8} {:e} -> {:e} ({:08x}) / {:e} ({:08x})", + nm, b.n, b.max_ulp, b.max_abs, b.max_rel, b.zeroed, b.hist[0], b.hist[1], b.hist[2], b.hist[3], + b.worst_x, b.worst_got, b.worst_got.to_bits(), b.worst_want, b.worst_want.to_bits() + ) + .unwrap(); + } + writeln!(s, "NaN inputs: {} kept NaN, {} not", st.nan_ok, st.nan_bad).unwrap(); + writeln!(s, "odd-symmetry violations (bitwise): {}", st.odd_bad).unwrap(); + writeln!(s, "outside [-1, 1]: {}", st.bound_bad).unwrap(); + writeln!(s, "monotonicity violations (x >= 0, ordered): {}", st.mono_bad).unwrap(); + writeln!(s, "reference flip-risk inputs (|x| only): {}", st.flip_risk.len()).unwrap(); + s +} + +fn bench() -> String { + let n = 1 << 20; + let x: Array1 = (0..n).map(|i| (i as f32 / n as f32) * 8.0 - 4.0).collect(); + let mut out = Array1::::zeros(n); + let reps = 50; + let t = Instant::now(); + for _ in 0..reps { + vstanh(x.view(), out.view_mut()); + } + let new_ns = t.elapsed().as_nanos() as f64 / (reps * n) as f64; + let xs = x.to_vec(); + let mut o = vec![0.0f32; n]; + let t = Instant::now(); + for _ in 0..reps { + old_kernel(&xs, &mut o); + } + let old_ns = t.elapsed().as_nanos() as f64 / (reps * n) as f64; + let t = Instant::now(); + for _ in 0..reps { + for (a, b) in xs.iter().zip(o.iter_mut()) { + *b = a.tanh(); + } + } + let scalar_ns = t.elapsed().as_nanos() as f64 / (reps * n) as f64; + std::hint::black_box((&out, &o)); + format!( + "throughput over 2^20 f32 in [-4, 4], {reps} reps: vstanh {new_ns:.3} ns/elem, old kernel {old_ns:.3} ns/elem, scalar f32::tanh {scalar_ns:.3} ns/elem\n" + ) +} + +fn main() { + let args: Vec = std::env::args().collect(); + let threads = std::thread::available_parallelism() + .map(|n| n.get() as u32) + .unwrap_or(4); + println!( + "target: avx512f={} avx2={} neon={} threads={threads}", + cfg!(target_feature = "avx512f"), + cfg!(target_feature = "avx2"), + cfg!(target_arch = "aarch64") + ); + let t = Instant::now(); + let new = sweep(false, threads); + print!("{}", report("vstanh (repaired), all 2^32 f32", &new)); + println!("sweep time {:.1} s", t.elapsed().as_secs_f64()); + let old = sweep(true, threads); + print!("{}", report("previous kernel 2*sigmoid(2x)-1, all 2^32 f32", &old)); + print!("{}", bench()); + if let Some(path) = args.get(1) { + let s: String = new + .flip_risk + .iter() + .map(|x| format!("{:08x}\n", x.to_bits())) + .collect(); + std::fs::write(path, s).unwrap(); + } + if let Some(path) = args.get(2) { + let s: String = new + .bands + .iter() + .filter(|b| b.n > 0) + .map(|b| format!("{:08x}\n", b.worst_x.to_bits())) + .collect(); + std::fs::write(path, s).unwrap(); + } +} diff --git a/src/hpc/vml.rs b/src/hpc/vml.rs index 00c77e7a..d5499616 100644 --- a/src/hpc/vml.rs +++ b/src/hpc/vml.rs @@ -232,25 +232,54 @@ fn vspow_slice(a: &[f32], b: &[f32], out: &mut [f32]) { } } +/// 16-lane `tanh`, computed on `a = |x|` and negated where `x < 0`, so the +/// function is exactly odd. Lanes where `x == 0` return `x` itself, which keeps +/// `tanh(-0.0) = -0.0`. (Float compares and `select` only: the NEON `U32x16` +/// has no `&`/`|`, so the sign is not restored by bit masking.) +/// +/// Two branches, chosen per lane: +/// - `|x| < 0.625`: `a + a·z·P(z)` with `z = a²` (the Cephes `tanhf` +/// minimax polynomial). It has no cancellation, and for subnormal or tiny +/// `a` the correction vanishes and the result is `a` itself. +/// - otherwise: `1 − 2/(exp(2a) + 1)`. `exp` overflows to `+inf` for large +/// `a`, which gives exactly `1.0`; a NaN input stays NaN. +/// +/// The earlier `2·sigmoid(2x) − 1` cancelled for small `|x|`: below about +/// `1e-7` the sigmoid rounds to `0.5` and the result was `0`. +#[inline(always)] +fn tanh_f32x16(v: F32x16) -> F32x16 { + let zero = F32x16::splat(0.0); + let a = v.abs(); + let z = a * a; + let p = F32x16::splat(-5.704_988_727_45e-3) + .mul_add(z, F32x16::splat(2.063_908_879_54e-2)) + .mul_add(z, F32x16::splat(-5.373_971_555_31e-2)) + .mul_add(z, F32x16::splat(1.333_144_220_36e-1)) + .mul_add(z, F32x16::splat(-3.333_328_194_22e-1)); + let small = (p * z).mul_add(a, a); + let e = simd_exp_f32(a + a); + let large = F32x16::splat(1.0) - F32x16::splat(2.0) / (e + F32x16::splat(1.0)); + let r = a.simd_lt(F32x16::splat(0.625)).select(small, large); + let r = v.simd_lt(zero).select(-r, r); + v.simd_eq(zero).select(v, r) +} + fn vstanh_slice(x: &[f32], out: &mut [f32]) { let n = x.len().min(out.len()); let chunks = n / 16; - let two = F32x16::splat(2.0); - let one = F32x16::splat(1.0); for i in 0..chunks { let off = i * 16; - let v = F32x16::from_slice(&x[off..off + 16]); - // tanh(x) = 2·sigmoid(2x) - 1 - let two_x = v * two; - // sigmoid(z) = 1/(1+exp(-z)) - 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)); - tanh_v.copy_to_slice(&mut out[off..off + 16]); - } - for i in (chunks * 16)..n { - out[i] = x[i].tanh(); + tanh_f32x16(F32x16::from_slice(&x[off..off + 16])).copy_to_slice(&mut out[off..off + 16]); + } + // The tail goes through the same 16-lane kernel, so a value gives the same + // bits wherever it sits in the array. + let rest = n - chunks * 16; + if rest > 0 { + let mut buf = [0.0f32; 16]; + buf[..rest].copy_from_slice(&x[chunks * 16..n]); + let r = tanh_f32x16(F32x16::from_slice(&buf)); + r.copy_to_slice(&mut buf); + out[chunks * 16..n].copy_from_slice(&buf[..rest]); } } @@ -605,8 +634,12 @@ pub fn vspow(a: ArrayView<'_, f32, D>, b: ArrayView<'_, f32, D>, m .for_each(|o, &x, &y| *o = x.powf(y)); } -/// Element-wise `tanh` via the identity `tanh(x) = 2·sigmoid(2x) − 1`. Hot -/// path reuses the F32x16 sigmoid built from `simd_exp_f32`. +/// Element-wise `tanh`, within 2 ULP of the correctly rounded result over +/// every finite `f32`. `tanh(±0) = ±0`, subnormals map to themselves, +/// `tanh(±inf) = ±1` and NaN stays NaN. +/// +/// Every element, contiguous or strided, goes through the same 16-lane +/// kernel, so a value's result does not depend on its position or layout. /// /// # Panics /// Panics if `x` and `out` shapes differ. @@ -615,7 +648,9 @@ pub fn vstanh(x: ArrayView<'_, f32, D>, mut out: ArrayViewMut<'_, if dispatch_unary_contig(x.view(), &mut out, vstanh_slice) { return; } - Zip::from(&mut out).and(&x).for_each(|o, &v| *o = v.tanh()); + Zip::from(&mut out) + .and(&x) + .for_each(|o, &v| *o = tanh_f32x16(F32x16::splat(v)).to_array()[0]); } /// Element-wise `floor`. Hot path uses `F32x16::floor()` (VRNDSCALEPS on @@ -1123,6 +1158,192 @@ mod tests { } } + /// Distance in units in the last place between two `f32`, with `-0.0` + /// and `+0.0` one apart so a lost sign of zero counts as an error. + fn ulp_distance(a: f32, b: f32) -> u32 { + if a.is_nan() || b.is_nan() { + return if a.is_nan() && b.is_nan() { 0 } else { u32::MAX }; + } + let key = |x: f32| { + let i = x.to_bits() as i32; + if i < 0 { + -((i & i32::MAX) as i64) - 1 + } else { + i as i64 + } + }; + (key(a) - key(b)).unsigned_abs().min(u32::MAX as u64) as u32 + } + + /// Correctly rounded reference: `tanh` in `f64`, rounded once to `f32`. + fn tanh_ref(x: f32) -> f32 { + (x as f64).tanh() as f32 + } + + fn vstanh_vec(xs: &[f32]) -> Vec { + let x = Array1::from(xs.to_vec()); + let mut out = Array1::::zeros(xs.len()); + vstanh(x.view(), out.view_mut()); + out.to_vec() + } + + #[test] + fn test_vstanh_within_2_ulp_across_magnitudes() { + // Log-spaced magnitudes in each band, both signs. The length is not a + // multiple of 16, so the last values take the tail path. + let bands = [ + (1.0e-38f32, 1.0e-10f32), + (1.0e-10, 1.0e-4), + (1.0e-4, 0.1), + (0.1, 0.625), + (0.625, 1.0), + (1.0, 4.0), + (4.0, 10.0), + (10.0, 100.0), + ]; + for (lo, hi) in bands { + let n = 4099; + let mut xs = Vec::with_capacity(2 * n + 1); + for i in 0..n { + let t = i as f64 / (n - 1) as f64; + let v = ((lo as f64).ln() + t * ((hi as f64).ln() - (lo as f64).ln())).exp() as f32; + xs.push(v); + xs.push(-v); + } + xs.push(hi); + assert_ne!(xs.len() % 16, 0, "fixture must reach the tail path"); + let ys = vstanh_vec(&xs); + let (worst, at) = xs + .iter() + .zip(&ys) + .map(|(&x, &y)| (ulp_distance(y, tanh_ref(x)), x)) + .max_by_key(|&(u, _)| u) + .unwrap(); + assert!(worst <= 2, "band [{lo:e}, {hi:e}]: {worst} ulp at x = {at:e}"); + } + } + + #[test] + fn test_vstanh_special_values_in_chunk_and_tail() { + let sub = f32::from_bits(1); + let cases: [(f32, f32); 9] = [ + (0.0, 0.0), + (-0.0, -0.0), + (sub, sub), + (-sub, -sub), + (f32::MIN_POSITIVE, f32::MIN_POSITIVE), + (-f32::MIN_POSITIVE, -f32::MIN_POSITIVE), + (f32::INFINITY, 1.0), + (f32::NEG_INFINITY, -1.0), + (1.0e-8, 1.0e-8), + ]; + // 16 values fill one SIMD chunk; the same 9 again land in the tail. + let mut xs: Vec = cases.iter().map(|c| c.0).collect(); + xs.resize(16, 0.5); + xs.extend(cases.iter().map(|c| c.0)); + let ys = vstanh_vec(&xs); + for (k, &(x, want)) in cases.iter().enumerate() { + for (path, y) in [("chunk", ys[k]), ("tail", ys[16 + k])] { + assert_eq!(y.to_bits(), want.to_bits(), "{path}: vstanh({x:e}) = {y:e}, want {want:e}"); + } + } + assert!(vstanh_vec(&[f32::NAN; 17]).iter().all(|v| v.is_nan())); + } + + #[test] + fn test_vstanh_tail_equals_chunk() { + // The same input gives the same bits in a full chunk and in the tail, + // so a result never depends on where in the array it sits. + let vals: Vec = (0..15).map(|i| (i as f32 - 7.0) * 0.37 + 1.0e-3).collect(); + let mut xs = vals.clone(); + xs.push(0.0); + xs.extend(&vals); + let ys = vstanh_vec(&xs); + for k in 0..vals.len() { + assert_eq!(ys[k].to_bits(), ys[16 + k].to_bits(), "x = {}", vals[k]); + } + } + + #[test] + fn test_vstanh_every_length_agrees_with_a_single_lane() { + // Lengths around the 16-lane chunk: each element must equal the + // one-element result for the same input. + for len in [0usize, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65] { + let xs: Vec = (0..len).map(|i| (i as f32 - 32.0) * 0.173).collect(); + let ys = vstanh_vec(&xs); + assert_eq!(ys.len(), len); + for (&x, &y) in xs.iter().zip(&ys) { + let one = vstanh_vec(&[x])[0]; + assert_eq!(y.to_bits(), one.to_bits(), "len {len}, x = {x}"); + assert!(ulp_distance(y, tanh_ref(x)) <= 2, "len {len}, x = {x}"); + } + } + } + + #[test] + fn test_vstanh_strided_and_2d_views_match_contiguous() { + let data: Vec = (0..96).map(|i| (i as f32 - 48.0) * 0.0911).collect(); + let a = Array2::from_shape_vec((8, 12), data.clone()).unwrap(); + let want: Vec = vstanh_vec(&data).iter().map(|v| v.to_bits()).collect(); + // C-order 2-D (contiguous), its transpose (non-contiguous), a strided column slice. + let mut out = Array2::::zeros((8, 12)); + vstanh(a.view(), out.view_mut()); + assert_eq!(out.iter().map(|v| v.to_bits()).collect::>(), want); + let t = a.t(); + let mut out_t = Array2::::zeros((12, 8)); + vstanh(t, out_t.view_mut()); + for i in 0..8 { + for j in 0..12 { + assert_eq!(out_t[[j, i]].to_bits(), want[i * 12 + j]); + } + } + let col = a.slice(s![.., 5]); + let mut out_c = Array1::::zeros(8); + vstanh(col, out_c.view_mut()); + for i in 0..8 { + assert_eq!(out_c[i].to_bits(), want[i * 12 + 5]); + } + } + + #[test] + fn test_vstanh_branch_threshold_neighbourhood() { + // 2000 consecutive floats either side of 0.625, both signs: the two + // branches must meet within the gate and without a step backwards. + let mut xs = Vec::new(); + let mut lo = 0.625f32; + for _ in 0..2000 { + lo = lo.next_down(); + } + let mut x = lo; + for _ in 0..4001 { + xs.push(x); + x = x.next_up(); + } + let ys = vstanh_vec(&xs); + for (k, (&x, &y)) in xs.iter().zip(&ys).enumerate() { + assert!(ulp_distance(y, tanh_ref(x)) <= 2, "x = {x:e}"); + if k > 0 { + assert!(y >= ys[k - 1], "not monotone at x = {x:e}"); + } + } + let neg: Vec = xs.iter().map(|v| -v).collect(); + for (y, ny) in ys.iter().zip(vstanh_vec(&neg)) { + assert_eq!(ny.to_bits(), y.to_bits() ^ 0x8000_0000); + } + } + + #[test] + fn test_vstanh_saturation_and_extremes() { + let ys = vstanh_vec(&[f32::MAX, -f32::MAX, 9.0, -9.0, 44.0, 89.0, -89.0]); + assert_eq!(ys[0].to_bits(), 1.0f32.to_bits()); + assert_eq!(ys[1].to_bits(), (-1.0f32).to_bits()); + assert!(ulp_distance(ys[2], tanh_ref(9.0)) <= 2); + assert!(ulp_distance(ys[3], tanh_ref(-9.0)) <= 2); + for &y in &ys[4..] { + assert_eq!(y.abs(), 1.0); + } + } + // ----------------------------------------------------------------------- // vsfloor / vsceil / vsround / vsneg / vstrunc // ----------------------------------------------------------------------- diff --git a/src/simd.rs b/src/simd.rs index 3e5c86d0..cf819e51 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -526,10 +526,11 @@ pub fn f32_to_bf16_batch(input: &[f32], output: &mut [u16]) { // SIMD math functions — ndarray additions (not in std::simd) // ============================================================================ -/// Fast exp(x) for F32x16 — Remez polynomial on [-87, 87]. +/// Fast exp(x) for F32x16 — Cephes minimax polynomial after range reduction. /// -/// Max error ~2 ULP in [-10, 10]. Uses the standard range-reduction -/// approach: exp(x) = 2^n * exp(r) where r = x - n*ln(2). +/// Within 1 ULP of the correctly rounded result for normal results +/// (measured on 4 M points over [-87, 88.5]). Uses the standard +/// range-reduction approach: exp(x) = 2^n * exp(r) where r = x - n*ln(2). /// /// Domain: inputs from ln(f32::MAX) up return +Inf, and inputs are clamped /// to [-110, ln(f32::MAX)] before reduction, so -Inf and anything below @@ -580,13 +581,16 @@ pub fn simd_exp_f32(x: F32x16) -> F32x16 { let r = x - n * F32x16::splat(LN2_HI); let r = r - n * F32x16::splat(LN2_LO); - // Polynomial: exp(r) ≈ 1 + r + r²/2 + r³/6 + r⁴/24 + r⁵/120 - let c2 = F32x16::splat(0.5); - let c3 = F32x16::splat(1.0 / 6.0); - let c4 = F32x16::splat(1.0 / 24.0); - let c5 = F32x16::splat(1.0 / 120.0); - - let poly = one + r * (one + r * (c2 + r * (c3 + r * (c4 + r * c5)))); + // Polynomial: exp(r) ≈ 1 + r + r²·P(r) on |r| ≤ ln2/2, with P the Cephes + // `expf` minimax of degree 5. The degree-5 Taylor series used before was + // 39 ULP off at |r| = ln2/2 (its r⁶/720 term is ~2.4e-6 relative). + let p = F32x16::splat(EXP_P[0]); + let p = p * r + F32x16::splat(EXP_P[1]); + let p = p * r + F32x16::splat(EXP_P[2]); + let p = p * r + F32x16::splat(EXP_P[3]); + let p = p * r + F32x16::splat(EXP_P[4]); + let p = p * r + F32x16::splat(EXP_P[5]); + let poly = (p * (r * r) + r) + one; // Reconstruct: exp(x) = 2^n * poly. 2^n is only a normal f32 for n in // [-126, 127], but n reaches 128 near the overflow edge and about -159 @@ -604,6 +608,12 @@ pub fn simd_exp_f32(x: F32x16) -> F32x16 { nan_mask.select(F32x16::splat(f32::NAN), result) } +/// Cephes `expf` coefficients, highest degree first: `exp(r) ≈ 1 + r + r²·P(r)`. +const EXP_P: [f32; 6] = [ + 1.987_569_150_0e-4, 1.398_199_950_7e-3, 8.333_451_907_3e-3, 4.166_579_589_4e-2, 1.666_666_545_9e-1, + 5.000_000_120_1e-1, +]; + /// f32 just above ln(f32::MAX): `exp` of it overflows to +inf. const OVERFLOW_X: f32 = 88.722_84; /// High part of ln 2 (Cephes): 355/512, 9 significant bits. @@ -1764,7 +1774,11 @@ mod tests { let n = (xc * (1.0 / core::f32::consts::LN_2) + 0.5).floor(); let r = xc - n * super::LN2_HI; let r = r - n * super::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 p = super::EXP_P + .iter() + .skip(1) + .fold(super::EXP_P[0], |p, &c| p * r + c); + let poly = (p * (r * r) + r) + 1.0; let (n, post) = if n > 127.0 { (n - 1.0, 2.0) } else if n < -126.0 { @@ -1796,6 +1810,37 @@ mod tests { } } + /// FAILS IF: the polynomial is too short. Measured against a correctly + /// rounded reference (`exp` in `f64`, rounded once) over the normal-result + /// range, the worst case sits at reduced argument |r| = ln2/2. + #[test] + fn simd_exp_f32_is_within_2_ulp_of_the_correctly_rounded_result() { + fn key(x: f32) -> i64 { + let i = x.to_bits() as i32; + if i < 0 { + -((i & i32::MAX) as i64) - 1 + } else { + i as i64 + } + } + let n = 1 << 16; + let xs: Vec = (0..n) + .map(|i| -87.0 + 175.5 * (i as f32 / n as f32)) + .collect(); + let mut worst = (0i64, 0.0f32); + for chunk in xs.chunks_exact(16) { + let got = simd_exp_f32(F32x16::from_slice(chunk)).to_array(); + for (&x, g) in chunk.iter().zip(got) { + let want = (x as f64).exp() as f32; + let u = (key(g) - key(want)).abs(); + if u > worst.0 { + worst = (u, x); + } + } + } + assert!(worst.0 <= 2, "exp: {} ulp at x = {}", worst.0, worst.1); + } + // ── W1a-#9 parity: U64x8/U32x16 andnot + ternlog ──────────────────────── // // The backends are compile-time exclusive, so "all three agree" is proven