himada-dispatch 0.1.1

Adaptive SIMD dispatch for Himada — auto-selects fastest kernel at runtime
use himada_core::HardwareDNA;

// --- Scalar ---

pub fn dot_scalar(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

pub fn dot_scalar_supported(_: &HardwareDNA) -> bool {
    true
}

// --- x86 SSE ---

#[cfg(target_arch = "x86_64")]
pub fn dot_sse(a: &[f64], b: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    let len = a.len().min(b.len());
    let mut sum = 0.0;
    let mut i = 0;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` ensures the CPU supports SSE2 intrinsics
    // - `len >= 2` ensures at least one full 2×f64 vector can be processed
    // - `_mm_loadu_pd` works on unaligned pointers, valid for a.len() elements
    // - `a.as_ptr().add(i)` and `b.as_ptr().add(i)` are in-bounds because `i + 2 <= len`
    // - `_mm_add_pd` / `_mm_mul_pd` are safe SIMD arithmetic operations
    // - `transmute` from __m128d to [f64; 2] is safe because both are 16 bytes
    // - remainder is handled by scalar code after the block
    unsafe {
        if is_x86_feature_detected!("sse2") && len >= 2 {
            let mut acc = _mm_setzero_pd();
            while i + 2 <= len {
                let va = _mm_loadu_pd(a.as_ptr().add(i));
                let vb = _mm_loadu_pd(b.as_ptr().add(i));
                acc = _mm_add_pd(acc, _mm_mul_pd(va, vb));
                i += 2;
            }
            let tmp: [f64; 2] = std::mem::transmute::<_, [f64; 2]>(acc);
            sum += tmp[0] + tmp[1];
        }
    }

    for j in i..len {
        sum += a[j] * b[j];
    }
    sum
}

#[cfg(target_arch = "x86_64")]
pub fn dot_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn dot_sse(_: &[f64], _: &[f64]) -> f64 { f64::NAN }

#[cfg(not(target_arch = "x86_64"))]
pub fn dot_sse_supported(_: &HardwareDNA) -> bool { false }

// --- x86 AVX2 ---

#[cfg(target_arch = "x86_64")]
pub fn dot_avx2(a: &[f64], b: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    let len = a.len().min(b.len());
    let mut sum = 0.0;
    let mut i = 0;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` ensures the CPU supports AVX2 intrinsics
    // - `len >= 4` ensures at least one full 4×f64 vector can be processed
    // - `_mm256_loadu_pd` works on unaligned pointers, valid for a.len() elements
    // - `a.as_ptr().add(i)` and `b.as_ptr().add(i)` are in-bounds because `i + 4 <= len`
    // - `_mm256_add_pd` / `_mm256_mul_pd` are safe SIMD arithmetic operations
    // - `transmute` from __m256d to [f64; 4] is safe because both are 32 bytes
    // - remainder is handled by scalar code after the block
    unsafe {
        if is_x86_feature_detected!("avx2") && len >= 4 {
            let mut acc = _mm256_setzero_pd();
            while i + 4 <= len {
                let va = _mm256_loadu_pd(a.as_ptr().add(i));
                let vb = _mm256_loadu_pd(b.as_ptr().add(i));
                acc = _mm256_add_pd(acc, _mm256_mul_pd(va, vb));
                i += 4;
            }
            let tmp: [f64; 4] = std::mem::transmute::<_, [f64; 4]>(acc);
            sum += tmp[0] + tmp[1] + tmp[2] + tmp[3];
        }
    }

    for j in i..len {
        sum += a[j] * b[j];
    }
    sum
}

#[cfg(target_arch = "x86_64")]
pub fn dot_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn dot_avx2(_: &[f64], _: &[f64]) -> f64 { f64::NAN }

#[cfg(not(target_arch = "x86_64"))]
pub fn dot_avx2_supported(_: &HardwareDNA) -> bool { false }

// --- ARM NEON ---

#[cfg(target_arch = "aarch64")]
pub fn dot_neon(a: &[f64], b: &[f64]) -> f64 {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;

    let len = a.len().min(b.len());
    let mut sum = 0.0;
    let mut i = 0;

    // SAFETY:
    // - NEON is always available on aarch64, no runtime check needed
    // - `len >= 2` ensures at least one full 2×f64 vector can be processed
    // - `vld1q_f64` works on unaligned pointers (aarch64 allows unaligned access)
    // - `a.as_ptr().add(i)` and `b.as_ptr().add(i)` are in-bounds because `i + 2 <= len`
    // - `vaddq_f64` / `vmulq_f64` are safe SIMD arithmetic operations
    // - `transmute` from float64x2_t to [f64; 2] is safe because both are 16 bytes
    // - remainder is handled by scalar code after the block
    unsafe {
        if len >= 2 {
            let mut acc = vdupq_n_f64(0.0);
            while i + 2 <= len {
                let va = vld1q_f64(a.as_ptr().add(i));
                let vb = vld1q_f64(b.as_ptr().add(i));
                acc = vaddq_f64(acc, vmulq_f64(va, vb));
                i += 2;
            }
            let tmp: [f64; 2] = std::mem::transmute::<_, [f64; 2]>(acc);
            sum += tmp[0] + tmp[1];
        }
    }

    for j in i..len {
        sum += a[j] * b[j];
    }
    sum
}

#[cfg(target_arch = "aarch64")]
pub fn dot_neon_supported(_: &HardwareDNA) -> bool {
    true
}

#[cfg(not(target_arch = "aarch64"))]
pub fn dot_neon(_: &[f64], _: &[f64]) -> f64 { f64::NAN }

#[cfg(not(target_arch = "aarch64"))]
pub fn dot_neon_supported(_: &HardwareDNA) -> bool { false }