geographdb-core 0.5.4

//! CPU-native SIMD kernels for small geometric ML primitives.
//!
//! These kernels are intentionally low-level and private to the crate.  They
//! provide portable, feature-detected SIMD implementations of attention
//! primitives (softmax) and RoPE with a scalar fallback on every platform.
//!
//! Note on GEMV: the hot `vec_matmul` / `mat_t_vec` paths in
//! `algorithms::attention` use small row-major matrices that auto-vectorize
//! extremely well when the scalar loops are inlined.  Hand-written AVX2 kernels
//! with runtime dispatch are *slower* here because they cannot be inlined into
//! the caller, so we intentionally keep GEMV inline in `attention.rs`.
//!
//! The dispatch style mirrors `spatial::simd`: runtime feature detection, no
//! `unsafe` in the public entry points, and a scalar path that is always
//! correct.

/// Stable softmax computed in place.
///
/// After the call `scores[i] = exp(scores[i] - max) / sum_j exp(scores[j] - max)`.
pub fn softmax_in_place(scores: &mut [f32]) {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        if is_avx2_supported() {
            unsafe { softmax_avx2(scores) };
            return;
        }
    }

    softmax_scalar(scores);
}

/// Apply standard RoPE rotation to a single embedding vector in place.
///
/// `position` is the 1-based sequence distance from the target.  `embed_dim`
/// must be even.  Kept scalar because per-pair trigonometry dominates and the
/// operation is not on the critical path relative to GEMV.
pub fn apply_rope_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    assert_eq!(
        vec.len(),
        embed_dim,
        "apply_rope_in_place: dimension mismatch"
    );
    assert!(
        embed_dim % 2 == 0,
        "apply_rope_in_place: embed_dim must be even"
    );

    let ln_theta = theta.ln();
    let pos_f = position as f32;
    let embed_dim_f = embed_dim as f32;

    for pair in 0..embed_dim / 2 {
        let d0 = pair * 2;
        let d1 = d0 + 1;
        let freq = (-2.0f32 * pair as f32 * ln_theta / embed_dim_f).exp();
        let angle = pos_f * freq;
        let (cos_a, sin_a) = (angle.cos(), angle.sin());
        let v0 = vec[d0];
        let v1 = vec[d1];
        vec[d0] = v0 * cos_a - v1 * sin_a;
        vec[d1] = v0 * sin_a + v1 * cos_a;
    }
}

/// Apply the inverse RoPE rotation to a single embedding vector in place.
pub fn apply_rope_inv_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    assert_eq!(
        vec.len(),
        embed_dim,
        "apply_rope_inv_in_place: dimension mismatch"
    );
    assert!(
        embed_dim % 2 == 0,
        "apply_rope_inv_in_place: embed_dim must be even"
    );

    let ln_theta = theta.ln();
    let pos_f = position as f32;
    let embed_dim_f = embed_dim as f32;

    for pair in 0..embed_dim / 2 {
        let d0 = pair * 2;
        let d1 = d0 + 1;
        let freq = (-2.0f32 * pair as f32 * ln_theta / embed_dim_f).exp();
        let angle = pos_f * freq;
        let (cos_a, sin_a) = (angle.cos(), angle.sin());
        let v0 = vec[d0];
        let v1 = vec[d1];
        vec[d0] = v0 * cos_a + v1 * sin_a;
        vec[d1] = -v0 * sin_a + v1 * cos_a;
    }
}

// ---------------------------------------------------------------------------
// Scalar fallbacks
// ---------------------------------------------------------------------------

fn softmax_scalar(scores: &mut [f32]) {
    let max = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0f32;
    for s in scores.iter_mut() {
        *s = (*s - max).exp();
        sum += *s;
    }
    let inv_sum = 1.0 / sum;
    for s in scores.iter_mut() {
        *s *= inv_sum;
    }
}

// ---------------------------------------------------------------------------
// x86/x86_64 runtime feature detection
// ---------------------------------------------------------------------------

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn is_avx2_supported() -> bool {
    use std::sync::atomic::{AtomicI8, Ordering};
    static CACHED: AtomicI8 = AtomicI8::new(-1);

    let cached = CACHED.load(Ordering::Relaxed);
    if cached >= 0 {
        cached != 0
    } else {
        let supported = std::arch::x86_64::__cpuid(7).ebx & (1 << 5) != 0;
        CACHED.store(if supported { 1 } else { 0 }, Ordering::Relaxed);
        supported
    }
}

// ---------------------------------------------------------------------------
// AVX2 implementations
// ---------------------------------------------------------------------------

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn softmax_avx2(scores: &mut [f32]) {
    use std::arch::x86_64::*;

    if scores.is_empty() {
        return;
    }

    // Find max.
    let mut max_vec = _mm256_set1_ps(f32::NEG_INFINITY);
    let mut i = 0;
    while i + 8 <= scores.len() {
        let v = _mm256_loadu_ps(scores.as_ptr().add(i));
        max_vec = _mm256_max_ps(max_vec, v);
        i += 8;
    }
    let mut max = hsum256_ps(max_vec);
    while i < scores.len() {
        max = max.max(*scores.as_ptr().add(i));
        i += 1;
    }

    // Compute exp(x - max) and sum.
    let max_broadcast = _mm256_set1_ps(max);
    let mut sum_vec = _mm256_setzero_ps();
    i = 0;
    while i + 8 <= scores.len() {
        let v = _mm256_loadu_ps(scores.as_ptr().add(i));
        let shifted = _mm256_sub_ps(v, max_broadcast);
        let exps = exp256_ps(shifted);
        _mm256_storeu_ps(scores.as_mut_ptr().add(i), exps);
        sum_vec = _mm256_add_ps(sum_vec, exps);
        i += 8;
    }
    let mut sum = hsum256_ps(sum_vec);
    while i < scores.len() {
        let e = (*scores.as_ptr().add(i) - max).exp();
        *scores.as_mut_ptr().add(i) = e;
        sum += e;
        i += 1;
    }

    // Normalize.
    let inv_sum = 1.0 / sum;
    let inv_vec = _mm256_set1_ps(inv_sum);
    i = 0;
    while i + 8 <= scores.len() {
        let v = _mm256_loadu_ps(scores.as_ptr().add(i));
        _mm256_storeu_ps(scores.as_mut_ptr().add(i), _mm256_mul_ps(v, inv_vec));
        i += 8;
    }
    while i < scores.len() {
        *scores.as_mut_ptr().add(i) *= inv_sum;
        i += 1;
    }
}

// ---------------------------------------------------------------------------
// AVX2 helpers
// ---------------------------------------------------------------------------

/// Horizontal sum of eight f32 lanes in an AVX2 register.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn hsum256_ps(v: std::arch::x86_64::__m256) -> f32 {
    use std::arch::x86_64::*;
    let hi = _mm256_extractf128_ps(v, 1);
    let lo = _mm256_castps256_ps128(v);
    let sum128 = _mm_add_ps(lo, hi);
    let shuf = _mm_movehdup_ps(sum128);
    let sums = _mm_add_ps(sum128, shuf);
    let shuf2 = _mm_movehl_ps(shuf, sums);
    let sums2 = _mm_add_ss(sums, shuf2);
    _mm_cvtss_f32(sums2)
}

/// Vectorized `exp` for AVX2.
///
/// Computes `exp(x)` using the identity `exp(x) = 2^k * exp(r)` where
/// `k = round(x / ln 2)` and `r = x - k * ln 2`.  The scalar `exp(r)` is
/// approximated by a degree-5 Taylor polynomial on `[-ln2/2, ln2/2]`.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn exp256_ps(x: std::arch::x86_64::__m256) -> std::arch::x86_64::__m256 {
    use std::arch::x86_64::*;

    let min_x = _mm256_set1_ps(-87.336_55_f32);
    let max_x = _mm256_set1_ps(87.336_55_f32);
    let x = _mm256_max_ps(min_x, _mm256_min_ps(max_x, x));

    let inv_ln2 = _mm256_set1_ps(std::f32::consts::LOG2_E);
    let ln2 = _mm256_set1_ps(std::f32::consts::LN_2);
    let half = _mm256_set1_ps(0.5f32);

    let fx = _mm256_mul_ps(x, inv_ln2);
    let fx = _mm256_add_ps(fx, half);
    let k = _mm256_floor_ps(fx);
    let r = _mm256_sub_ps(x, _mm256_mul_ps(k, ln2));

    let c2 = _mm256_set1_ps(0.5_f32);
    let c3 = _mm256_set1_ps(0.166_666_67_f32);
    let c4 = _mm256_set1_ps(0.041_666_668_f32);
    let c5 = _mm256_set1_ps(0.008_333_334_f32);

    let r2 = _mm256_mul_ps(r, r);
    let r3 = _mm256_mul_ps(r2, r);
    let r4 = _mm256_mul_ps(r2, r2);

    let poly = _mm256_add_ps(
        _mm256_add_ps(_mm256_set1_ps(1.0f32), r),
        _mm256_add_ps(
            _mm256_mul_ps(c2, r2),
            _mm256_add_ps(
                _mm256_mul_ps(c3, r3),
                _mm256_add_ps(
                    _mm256_mul_ps(c4, r4),
                    _mm256_mul_ps(c5, _mm256_mul_ps(r4, r)),
                ),
            ),
        ),
    );

    let ki = _mm256_cvtps_epi32(k);
    let biased = _mm256_add_epi32(ki, _mm256_set1_epi32(127));
    let two_k = _mm256_slli_epi32::<23>(biased);
    _mm256_mul_ps(poly, _mm256_castsi256_ps(two_k))
}

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

    #[test]
    fn softmax_sums_to_one() {
        let mut scores = vec![1.0f32, 2.0, 3.0, 0.0];
        softmax_in_place(&mut scores);
        let sum: f32 = scores.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
    }

    #[test]
    fn softmax_matches_scalar_reference() {
        let values = vec![0.1f32, -0.5, 2.3, 1.0, -1.0, 0.0, 3.0];
        let mut scalar = values.clone();
        softmax_scalar(&mut scalar);
        let mut simd = values.clone();
        softmax_in_place(&mut simd);
        for (a, b) in scalar.iter().zip(simd.iter()) {
            assert!((a - b).abs() < 1e-5, "softmax mismatch: {} vs {}", a, b);
        }
    }

    #[test]
    fn rope_inverse_recover() {
        let theta = 10000.0f32;
        let original = (0..16).map(|i| i as f32).collect::<Vec<_>>();
        let mut v = original.clone();
        apply_rope_in_place(&mut v, 16, 3, theta);
        apply_rope_inv_in_place(&mut v, 16, 3, theta);
        for (a, b) in original.iter().zip(v.iter()) {
            assert!(
                (a - b).abs() < 1e-5,
                "RoPE inverse mismatch: {} vs {}",
                a,
                b
            );
        }
    }
}