feox-ann 0.1.0

/// SIMD dot product over two `f32` slices, using NEON on aarch64 and AVX2+FMA
/// on x86_64 when available, with a scalar fallback. Slices are truncated to
/// the shorter length.
pub fn dot(left: &[f32], right: &[f32]) -> f32 {
    #[cfg(target_arch = "aarch64")]
    {
        return unsafe { neon::dot(left, right) };
    }

    #[cfg(target_arch = "x86_64")]
    {
        if avx2::available() {
            return unsafe { avx2::dot(left, right) };
        }
    }

    #[allow(unreachable_code)]
    dot_scalar(left, right)
}

#[allow(dead_code)]
pub(crate) fn dot_scalar(left: &[f32], right: &[f32]) -> f32 {
    let mut i = 0;
    let mut a = 0.0_f32;
    let mut b = 0.0_f32;
    let mut c = 0.0_f32;
    let mut d = 0.0_f32;
    while i + 4 <= left.len() {
        a += left[i] * right[i];
        b += left[i + 1] * right[i + 1];
        c += left[i + 2] * right[i + 2];
        d += left[i + 3] * right[i + 3];
        i += 4;
    }
    let mut sum = a + b + c + d;
    while i < left.len() {
        sum += left[i] * right[i];
        i += 1;
    }
    sum
}

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

    #[target_feature(enable = "neon")]
    pub(crate) unsafe fn dot(left: &[f32], right: &[f32]) -> f32 {
        let len = left.len().min(right.len());
        let a = left.as_ptr();
        let b = right.as_ptr();
        let mut acc0 = vdupq_n_f32(0.0);
        let mut acc1 = vdupq_n_f32(0.0);
        let mut acc2 = vdupq_n_f32(0.0);
        let mut acc3 = vdupq_n_f32(0.0);
        let mut i = 0;
        while i + 16 <= len {
            acc0 = vfmaq_f32(acc0, vld1q_f32(a.add(i)), vld1q_f32(b.add(i)));
            acc1 = vfmaq_f32(acc1, vld1q_f32(a.add(i + 4)), vld1q_f32(b.add(i + 4)));
            acc2 = vfmaq_f32(acc2, vld1q_f32(a.add(i + 8)), vld1q_f32(b.add(i + 8)));
            acc3 = vfmaq_f32(acc3, vld1q_f32(a.add(i + 12)), vld1q_f32(b.add(i + 12)));
            i += 16;
        }
        while i + 4 <= len {
            acc0 = vfmaq_f32(acc0, vld1q_f32(a.add(i)), vld1q_f32(b.add(i)));
            i += 4;
        }
        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3)));
        while i < len {
            sum += *a.add(i) * *b.add(i);
            i += 1;
        }
        sum
    }
}

#[cfg(target_arch = "x86_64")]
mod avx2 {
    use std::arch::x86_64::*;
    use std::sync::OnceLock;

    pub(crate) fn available() -> bool {
        static AVAILABLE: OnceLock<bool> = OnceLock::new();
        *AVAILABLE.get_or_init(|| {
            std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma")
        })
    }

    #[target_feature(enable = "avx2,fma")]
    pub(crate) unsafe fn dot(left: &[f32], right: &[f32]) -> f32 {
        let len = left.len().min(right.len());
        let a = left.as_ptr();
        let b = right.as_ptr();
        let mut acc0 = _mm256_setzero_ps();
        let mut acc1 = _mm256_setzero_ps();
        let mut acc2 = _mm256_setzero_ps();
        let mut acc3 = _mm256_setzero_ps();
        let mut i = 0;
        while i + 32 <= len {
            acc0 = _mm256_fmadd_ps(_mm256_loadu_ps(a.add(i)), _mm256_loadu_ps(b.add(i)), acc0);
            acc1 = _mm256_fmadd_ps(
                _mm256_loadu_ps(a.add(i + 8)),
                _mm256_loadu_ps(b.add(i + 8)),
                acc1,
            );
            acc2 = _mm256_fmadd_ps(
                _mm256_loadu_ps(a.add(i + 16)),
                _mm256_loadu_ps(b.add(i + 16)),
                acc2,
            );
            acc3 = _mm256_fmadd_ps(
                _mm256_loadu_ps(a.add(i + 24)),
                _mm256_loadu_ps(b.add(i + 24)),
                acc3,
            );
            i += 32;
        }
        while i + 8 <= len {
            acc0 = _mm256_fmadd_ps(_mm256_loadu_ps(a.add(i)), _mm256_loadu_ps(b.add(i)), acc0);
            i += 8;
        }
        let reduced = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
        let low = _mm256_castps256_ps128(reduced);
        let high = _mm256_extractf128_ps(reduced, 1);
        let sum128 = _mm_add_ps(low, high);
        let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
        let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0b01));
        let mut sum = _mm_cvtss_f32(sum32);
        while i < len {
            sum += *a.add(i) * *b.add(i);
            i += 1;
        }
        sum
    }
}

pub(crate) fn normalize(vector: &[f32]) -> Vec<f32> {
    let norm = dot(vector, vector).sqrt();
    if norm == 0.0 {
        return vector.to_vec();
    }
    vector.iter().map(|value| value / norm).collect()
}

pub(crate) fn normalize_in_place(vector: &mut [f32]) {
    let norm = dot(vector, vector).sqrt();
    if norm == 0.0 {
        return;
    }
    for value in vector {
        *value /= norm;
    }
}

pub(crate) fn stable_hash(bytes: &[u8]) -> u64 {
    let mut hash = 0xcbf29ce484222325_u64;
    for byte in bytes {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(0x100000001b3);
    }
    hash
}

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

    #[test]
    fn simd_dot_matches_scalar() {
        for len in [0_usize, 1, 3, 4, 7, 8, 15, 16, 31, 32, 100, 128, 384, 1536] {
            let left: Vec<f32> = (0..len).map(|i| (i as f32 * 0.37).sin()).collect();
            let right: Vec<f32> = (0..len).map(|i| (i as f32 * 0.73).cos()).collect();
            let fast = dot(&left, &right);
            let scalar = dot_scalar(&left, &right);
            assert!(
                (fast - scalar).abs() <= scalar.abs().max(1.0) * 1e-5,
                "len {len}: simd {fast} vs scalar {scalar}"
            );
        }
    }
}