inferencelayer 0.2.4

//! Deterministic vectorized transcendentals for the CPU encoder: `exp`, `erf` (A&S 7.1.26 — the
//! SAME approximation the GPU's `ACT_FNS` uses), and the activation/softmax passes built on them.
//!
//! Two invariants, both load-bearing:
//!
//! 1. **The NEON lanes and the scalar twin produce IDENTICAL bits.** Every formula is written as
//!    separate multiplies and adds — never `mul_add`, never a fused intrinsic — in one fixed
//!    order, so the only difference between the paths is how many lanes evaluate at once. The
//!    scalar twin is not a fallback approximation; it IS the same function. (This is why these
//!    don't call `libm`: exact-`erff`/`exp` differ per platform's libm, while a fixed polynomial
//!    is the same everywhere.)
//! 2. **Reduction order is fixed by construction.** `exp_sub_sum` accumulates in four lanes and
//!    folds them in one documented order on both paths, so the CPU encoder stays bit-reproducible
//!    run to run — it is the oracle the GPU kernels are gated against.
//!
//! Accuracy: `exp` is the Cephes single-precision polynomial (~2 ulp over the clamped range);
//! `erf` is Abramowitz & Stegun 7.1.26 (|ε| ≤ 1.5e-7 absolute). Both sit far inside the 1e-4
//! parity gates. The old scalar-`libm` activation pass measured 6.1 ns/elem on an M4 Max — a
//! 12-layer erf-GELU encode spends whole milliseconds there; these run ~4-6× faster and, unlike
//! libm, identically on every host.

/// Clamp bounds for `exp`: below → ~1.2e-38 (min normal), above → the scale step saturates to
/// +inf. exp(+88.72) truly is ~3.4e38 ≈ f32::MAX, so the +inf edge only exists AT the clamp,
/// where every consumer (silu, tanh, softmax) has the correct limit anyway.
const EXP_LO: f32 = -87.336_54;
const EXP_HI: f32 = 88.722_83;
const LOG2E: f32 = std::f32::consts::LOG2_E;
/// ln2 split for Cody–Waite range reduction (hi + lo = ln2 to extra precision).
const LN2_HI: f32 = 0.693_359_4;
const LN2_LO: f32 = -2.121_944_4e-4;
// Cephes expf polynomial: exp(r) ≈ 1 + r + r²·P(r), r ∈ [-½ln2, ½ln2].
const C0: f32 = 5e-1;
const C1: f32 = 1.666_666_5e-1;
const C2: f32 = 4.166_579_6e-2;
const C3: f32 = 8.333_452e-3;
const C4: f32 = 1.398_2e-3;
const C5: f32 = 1.987_569_1e-4;
// A&S 7.1.26 erf coefficients (the GPU's ACT_FNS uses the same five).
const ERF_P: f32 = 0.327_591_1;
const ERF_A1: f32 = 0.254_829_6;
const ERF_A2: f32 = -0.284_496_72;
const ERF_A3: f32 = 1.421_413_8;
const ERF_A4: f32 = -1.453_152_1;
const ERF_A5: f32 = 1.061_405_4;

/// Scalar `exp`, the twin of the NEON lane. Mul+add only, fixed order.
#[inline(always)]
pub fn exp_f32(x: f32) -> f32 {
    let x = x.max(EXP_LO).min(EXP_HI);
    let n = (x * LOG2E + 0.5).floor();
    let r = (x - n * LN2_HI) - n * LN2_LO;
    let mut p = C5;
    p = p * r + C4;
    p = p * r + C3;
    p = p * r + C2;
    p = p * r + C1;
    p = p * r + C0;
    let y = ((r * r) * p + r) + 1.0;
    y * f32::from_bits((((n as i32) + 127) << 23) as u32)
}

/// Scalar `erf`, the twin of the NEON lane (A&S 7.1.26, odd extension via the sign).
#[inline(always)]
pub fn erf_f32(x: f32) -> f32 {
    let ax = x.abs();
    let t = 1.0 / (1.0 + ERF_P * ax);
    let mut p = ERF_A5;
    p = p * t + ERF_A4;
    p = p * t + ERF_A3;
    p = p * t + ERF_A2;
    p = p * t + ERF_A1;
    let e = exp_f32(-(ax * ax));
    let y = 1.0 - (p * t) * e;
    if x < 0.0 { -y } else { y }
}

#[inline(always)]
fn tanh_via_exp(x: f32) -> f32 {
    // tanh(x) = 1 − 2/(exp(2x)+1); exp's clamp gives the correct ±1 limits.
    1.0 - 2.0 / (exp_f32(x + x) + 1.0)
}

/// The five activation kinds the encoders use — mirrors `encoder_weights::Act` semantics; the
/// caller maps its enum to avoid a dependency cycle.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum ActKind {
    GeluErf,
    GeluTanh,
    Silu,
    Tanh,
    Relu,
}

const GELU_TANH_C: f32 = 0.797_884_6; // sqrt(2/π)
const FRAC_1_SQRT_2: f32 = std::f32::consts::FRAC_1_SQRT_2;

#[inline(always)]
fn act_one(v: f32, act: ActKind) -> f32 {
    match act {
        ActKind::GeluErf => 0.5 * v * (1.0 + erf_f32(v * FRAC_1_SQRT_2)),
        ActKind::GeluTanh => {
            let inner = GELU_TANH_C * (v + 0.044715 * (v * v * v));
            0.5 * v * (1.0 + tanh_via_exp(inner))
        }
        ActKind::Silu => v / (1.0 + exp_f32(-v)),
        ActKind::Tanh => tanh_via_exp(v),
        ActKind::Relu => v.max(0.0),
    }
}

/// Elementwise activation in place. On aarch64 the polynomial runs four lanes wide; every other
/// host runs the scalar twin — same bits either way.
pub fn act_slice(xs: &mut [f32], act: ActKind) {
    #[cfg(target_arch = "aarch64")]
    {
        act_slice_neon(xs, act);
        return;
    }
    #[cfg(not(target_arch = "aarch64"))]
    for v in xs {
        *v = act_one(*v, act);
    }
}

/// `xs[i] = exp(xs[i] − mx)`, returning the sum — the softmax middle. `-inf` scores (masked
/// keys) map to EXACTLY 0.0, preserving the "all keys masked ⇒ denom == 0" contract the ragged
/// attention relies on. The sum folds four fixed accumulator lanes as `(s0+s1) + (s2+s3)`, tail
/// elements joining lane `i % 4` — the same order scalar and NEON.
pub fn exp_sub_sum(xs: &mut [f32], mx: f32) -> f32 {
    #[cfg(target_arch = "aarch64")]
    {
        return exp_sub_sum_neon(xs, mx);
    }
    #[cfg(not(target_arch = "aarch64"))]
    {
        let mut acc = [0f32; 4];
        for (i, v) in xs.iter_mut().enumerate() {
            let e = if *v == f32::NEG_INFINITY {
                0.0
            } else {
                exp_f32(*v - mx)
            };
            *v = e;
            acc[i % 4] += e;
        }
        (acc[0] + acc[1]) + (acc[2] + acc[3])
    }
}

#[cfg(target_arch = "aarch64")]
mod neon {
    use super::*;
    use std::arch::aarch64::*;

    /// Four-lane `exp`, the NEON body of [`exp_f32`] — the identical op sequence.
    ///
    /// # Safety
    /// NEON is baseline on aarch64; no memory is touched.
    #[inline(always)]
    pub(super) unsafe fn exp4(x: float32x4_t) -> float32x4_t {
        unsafe {
            let x = vminq_f32(vmaxq_f32(x, vdupq_n_f32(EXP_LO)), vdupq_n_f32(EXP_HI));
            // n = floor(x·log2e + 0.5)
            let z = vaddq_f32(vmulq_f32(x, vdupq_n_f32(LOG2E)), vdupq_n_f32(0.5));
            let n = vrndmq_f32(z);
            let r = vsubq_f32(
                vsubq_f32(x, vmulq_f32(n, vdupq_n_f32(LN2_HI))),
                vmulq_f32(n, vdupq_n_f32(LN2_LO)),
            );
            let mut p = vdupq_n_f32(C5);
            p = vaddq_f32(vmulq_f32(p, r), vdupq_n_f32(C4));
            p = vaddq_f32(vmulq_f32(p, r), vdupq_n_f32(C3));
            p = vaddq_f32(vmulq_f32(p, r), vdupq_n_f32(C2));
            p = vaddq_f32(vmulq_f32(p, r), vdupq_n_f32(C1));
            p = vaddq_f32(vmulq_f32(p, r), vdupq_n_f32(C0));
            let y = vaddq_f32(
                vaddq_f32(vmulq_f32(vmulq_f32(r, r), p), r),
                vdupq_n_f32(1.0),
            );
            // 2^n via exponent bits.
            let bits = vshlq_n_s32::<23>(vaddq_s32(vcvtq_s32_f32(n), vdupq_n_s32(127)));
            vmulq_f32(y, vreinterpretq_f32_s32(bits))
        }
    }

    /// Four-lane `erf` (A&S 7.1.26), twin of [`erf_f32`].
    ///
    /// # Safety
    /// NEON is baseline on aarch64; no memory is touched.
    #[inline(always)]
    pub(super) unsafe fn erf4(x: float32x4_t) -> float32x4_t {
        unsafe {
            let ax = vabsq_f32(x);
            let t = vdivq_f32(
                vdupq_n_f32(1.0),
                vaddq_f32(vdupq_n_f32(1.0), vmulq_f32(vdupq_n_f32(ERF_P), ax)),
            );
            let mut p = vdupq_n_f32(ERF_A5);
            p = vaddq_f32(vmulq_f32(p, t), vdupq_n_f32(ERF_A4));
            p = vaddq_f32(vmulq_f32(p, t), vdupq_n_f32(ERF_A3));
            p = vaddq_f32(vmulq_f32(p, t), vdupq_n_f32(ERF_A2));
            p = vaddq_f32(vmulq_f32(p, t), vdupq_n_f32(ERF_A1));
            let e = exp4(vnegq_f32(vmulq_f32(ax, ax)));
            let y = vsubq_f32(vdupq_n_f32(1.0), vmulq_f32(vmulq_f32(p, t), e));
            // Odd extension: copy x's sign onto y.
            let sign = vandq_u32(vreinterpretq_u32_f32(x), vdupq_n_u32(0x8000_0000));
            vreinterpretq_f32_u32(vorrq_u32(
                vandq_u32(vreinterpretq_u32_f32(y), vdupq_n_u32(0x7FFF_FFFF)),
                sign,
            ))
        }
    }

    #[inline(always)]
    unsafe fn tanh4(x: float32x4_t) -> float32x4_t {
        unsafe {
            let e = exp4(vaddq_f32(x, x));
            vsubq_f32(
                vdupq_n_f32(1.0),
                vdivq_f32(vdupq_n_f32(2.0), vaddq_f32(e, vdupq_n_f32(1.0))),
            )
        }
    }

    #[inline(always)]
    unsafe fn act4(v: float32x4_t, act: ActKind) -> float32x4_t {
        unsafe {
            match act {
                ActKind::GeluErf => {
                    let er = erf4(vmulq_f32(v, vdupq_n_f32(FRAC_1_SQRT_2)));
                    vmulq_f32(
                        vmulq_f32(vdupq_n_f32(0.5), v),
                        vaddq_f32(vdupq_n_f32(1.0), er),
                    )
                }
                ActKind::GeluTanh => {
                    let v3 = vmulq_f32(vmulq_f32(v, v), v);
                    let inner = vmulq_f32(
                        vdupq_n_f32(GELU_TANH_C),
                        vaddq_f32(v, vmulq_f32(vdupq_n_f32(0.044715), v3)),
                    );
                    vmulq_f32(
                        vmulq_f32(vdupq_n_f32(0.5), v),
                        vaddq_f32(vdupq_n_f32(1.0), tanh4(inner)),
                    )
                }
                ActKind::Silu => vdivq_f32(v, vaddq_f32(vdupq_n_f32(1.0), exp4(vnegq_f32(v)))),
                ActKind::Tanh => tanh4(v),
                ActKind::Relu => vmaxq_f32(v, vdupq_n_f32(0.0)),
            }
        }
    }

    pub(super) fn act_slice_neon(xs: &mut [f32], act: ActKind) {
        let n4 = xs.len() / 4 * 4;
        // SAFETY: in-bounds 4-lane loads/stores below n4; NEON baseline on aarch64.
        unsafe {
            let p = xs.as_mut_ptr();
            let mut i = 0;
            while i < n4 {
                vst1q_f32(p.add(i), act4(vld1q_f32(p.add(i)), act));
                i += 4;
            }
        }
        for v in &mut xs[n4..] {
            *v = act_one(*v, act);
        }
    }

    pub(super) fn exp_sub_sum_neon(xs: &mut [f32], mx: f32) -> f32 {
        let n4 = xs.len() / 4 * 4;
        let mut acc = [0f32; 4];
        // SAFETY: in-bounds 4-lane loads/stores below n4; NEON baseline on aarch64.
        unsafe {
            let p = xs.as_mut_ptr();
            let mxv = vdupq_n_f32(mx);
            let ninf = vdupq_n_f32(f32::NEG_INFINITY);
            let mut accv = vdupq_n_f32(0.0);
            let mut i = 0;
            while i < n4 {
                let v = vld1q_f32(p.add(i));
                let e = exp4(vsubq_f32(v, mxv));
                // Masked keys (score == −inf) contribute EXACTLY zero.
                let e = vbslq_f32(vceqq_f32(v, ninf), vdupq_n_f32(0.0), e);
                vst1q_f32(p.add(i), e);
                accv = vaddq_f32(accv, e);
                i += 4;
            }
            vst1q_f32(acc.as_mut_ptr(), accv);
        }
        for (i, v) in xs[n4..].iter_mut().enumerate() {
            let e = if *v == f32::NEG_INFINITY {
                0.0
            } else {
                exp_f32(*v - mx)
            };
            *v = e;
            acc[i % 4] += e;
        }
        (acc[0] + acc[1]) + (acc[2] + acc[3])
    }
}
#[cfg(target_arch = "aarch64")]
use neon::{act_slice_neon, exp_sub_sum_neon};

#[cfg(test)]
mod tests {
    use super::*;

    /// The polynomial exp must sit within ~2 ulp of libm over the whole clamped range.
    #[test]
    fn exp_matches_libm_within_tolerance() {
        let mut worst = 0f64;
        for i in -87_000..88_000 {
            let x = i as f32 * 1e-3;
            let (got, want) = (exp_f32(x) as f64, (x as f64).exp());
            let rel = ((got - want) / want).abs();
            worst = worst.max(rel);
        }
        assert!(worst < 3e-7, "worst rel err {worst:.3e}");
    }

    /// A&S 7.1.26 promises |ε| ≤ 1.5e-7 in exact arithmetic; evaluated in f32 (Horner + the
    /// polynomial exp) the measured worst is 4.1e-7 on this grid. Gate at 1e-6 — still 100×
    /// inside the 1e-4 parity band, and the GPU's ACT_FNS runs the SAME f32 formula.
    #[test]
    fn erf_matches_libm_within_its_published_bound() {
        let mut worst = 0f64;
        for i in -6_000..6_000 {
            let x = i as f32 * 1e-3;
            let d = (erf_f32(x) as f64 - libm::erf(x as f64)).abs();
            worst = worst.max(d);
        }
        assert!(worst < 1e-6, "worst abs err {worst:.3e}");
    }

    /// The NEON lanes and the scalar twin must be BITWISE identical — they are one function.
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn neon_and_scalar_twins_are_bitwise_identical() {
        for act in [
            ActKind::GeluErf,
            ActKind::GeluTanh,
            ActKind::Silu,
            ActKind::Tanh,
            ActKind::Relu,
        ] {
            let xs: Vec<f32> = (0..1001).map(|i| (i as f32 - 500.0) * 0.037).collect();
            let mut a = xs.clone();
            let mut b = xs.clone();
            act_slice(&mut a, act); // NEON (len 1001 also exercises the scalar tail)
            for v in &mut b {
                *v = act_one(*v, act);
            }
            assert!(
                a.iter().zip(&b).all(|(x, y)| x.to_bits() == y.to_bits()),
                "twin drift"
            );
        }
    }

    /// Masked (−inf) softmax lanes stay exactly zero, and the all-masked row keeps denom == 0.
    #[test]
    fn exp_sub_sum_keeps_masked_lanes_exactly_zero() {
        let mut xs = vec![0.5f32, f32::NEG_INFINITY, -0.25, 1.0, f32::NEG_INFINITY];
        let s = exp_sub_sum(&mut xs, 1.0);
        assert_eq!(xs[1], 0.0);
        assert_eq!(xs[4], 0.0);
        assert!((s - (xs[0] + xs[2] + xs[3])).abs() < 1e-6);
        let mut all = vec![f32::NEG_INFINITY; 7];
        assert_eq!(exp_sub_sum(&mut all, f32::NEG_INFINITY), 0.0);
        assert!(all.iter().all(|&v| v == 0.0));
    }
}