himada-dispatch 0.1.1

use himada_core::HardwareDNA;

// ---------------------------------------------------------------------------
// Reduce sum
// ---------------------------------------------------------------------------

pub fn reduce_sum_scalar(a: &[f64]) -> f64 {
    let mut sum = 0.0;
    for &v in a {
        sum += v;
    }
    sum
}

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

#[cfg(target_arch = "x86_64")]
pub fn reduce_sum_sse(a: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let mut i = 0;
    // SAFETY: `_mm_setzero_pd()` creates a zero-initialized vector; no data is accessed
    let mut acc = unsafe { _mm_setzero_pd() };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("sse2")` confirms CPU support
        // - `i + 2 <= a.len()` ensures in-bounds access for `_mm_loadu_pd`
        // - `_mm_loadu_pd` tolerates unaligned pointers
        // - `a.as_ptr().add(i)` is valid because `i + 2 <= a.len()`
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= a.len() {
                let v = _mm_loadu_pd(a.as_ptr().add(i));
                acc = _mm_add_pd(acc, v);
                i += 2;
            }
        }
    }
    // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes and have the same representation
    let tmp: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(acc) };
    let mut sum = tmp[0] + tmp[1];
    for &v in &a[i..] {
        sum += v;
    }
    sum
}

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

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

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

#[cfg(target_arch = "x86_64")]
pub fn reduce_sum_avx2(a: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let mut i = 0;
    // SAFETY: `_mm256_setzero_pd()` creates a zero-initialized vector; no data is accessed
    let mut acc = unsafe { _mm256_setzero_pd() };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("avx2")` confirms CPU support
        // - `i + 4 <= a.len()` ensures in-bounds access for `_mm256_loadu_pd`
        // - `_mm256_loadu_pd` tolerates unaligned pointers
        // - `a.as_ptr().add(i)` is valid because `i + 4 <= a.len()`
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= a.len() {
                let v = _mm256_loadu_pd(a.as_ptr().add(i));
                acc = _mm256_add_pd(acc, v);
                i += 4;
            }
        }
    }
    // SAFETY: transmuting __m256d to [f64; 4] is safe; both are 32 bytes and have the same representation
    let tmp: [f64; 4] = unsafe { std::mem::transmute::<_, [f64; 4]>(acc) };
    let mut sum = tmp[0] + tmp[1] + tmp[2] + tmp[3];
    for &v in &a[i..] {
        sum += v;
    }
    sum
}

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

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

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

#[cfg(target_arch = "aarch64")]
pub fn reduce_sum_neon(a: &[f64]) -> f64 {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let mut i = 0;
    // SAFETY: `vdupq_n_f64(0.0)` creates a zero-initialized vector; no data is accessed
    let mut acc = unsafe { vdupq_n_f64(0.0) };
    unsafe {
        // SAFETY:
        // - NEON is always available on aarch64
        // - `i + 2 <= a.len()` ensures in-bounds access for `vld1q_f64`
        // - `vld1q_f64` tolerates unaligned pointers on aarch64
        // - `a.as_ptr().add(i)` is valid because `i + 2 <= a.len()`
        while i + 2 <= a.len() {
            let v = vld1q_f64(a.as_ptr().add(i));
            acc = vaddq_f64(acc, v);
            i += 2;
        }
    }
    // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes and have the same representation
    let tmp: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(acc) };
    let mut sum = tmp[0] + tmp[1];
    for &v in &a[i..] {
        sum += v;
    }
    sum
}

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

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

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

// ---------------------------------------------------------------------------
// Reduce max
// ---------------------------------------------------------------------------

pub fn reduce_max_scalar(a: &[f64]) -> f64 {
    let mut max = f64::NEG_INFINITY;
    for &v in a {
        if v > max {
            max = v;
        }
    }
    max
}

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

#[cfg(target_arch = "x86_64")]
pub fn reduce_max_sse(a: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let mut i = 0;
    // SAFETY: `_mm_set1_pd(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut acc = unsafe { _mm_set1_pd(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("sse2")` confirms CPU support
        // - `i + 2 <= a.len()` ensures in-bounds access for `_mm_loadu_pd`
        // - `a.as_ptr().add(i)` is valid because `i + 2 <= a.len()`
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= a.len() {
                let v = _mm_loadu_pd(a.as_ptr().add(i));
                acc = _mm_max_pd(acc, v);
                i += 2;
            }
        }
    }
    // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes and have the same representation
    let tmp: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(acc) };
    let mut max = tmp[0].max(tmp[1]);
    for &v in &a[i..] {
        if v > max { max = v; }
    }
    max
}

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

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

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

#[cfg(target_arch = "x86_64")]
pub fn reduce_max_avx2(a: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let mut i = 0;
    // SAFETY: `_mm256_set1_pd(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut acc = unsafe { _mm256_set1_pd(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("avx2")` confirms CPU support
        // - `i + 4 <= a.len()` ensures in-bounds access for `_mm256_loadu_pd`
        // - `a.as_ptr().add(i)` is valid because `i + 4 <= a.len()`
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= a.len() {
                let v = _mm256_loadu_pd(a.as_ptr().add(i));
                acc = _mm256_max_pd(acc, v);
                i += 4;
            }
        }
    }
    // SAFETY: transmuting __m256d to [f64; 4] is safe; both are 32 bytes and have the same representation
    let tmp: [f64; 4] = unsafe { std::mem::transmute::<_, [f64; 4]>(acc) };
    let mut max = tmp[0].max(tmp[1]).max(tmp[2]).max(tmp[3]);
    for &v in &a[i..] {
        if v > max { max = v; }
    }
    max
}

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

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

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

#[cfg(target_arch = "aarch64")]
pub fn reduce_max_neon(a: &[f64]) -> f64 {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let mut i = 0;
    // SAFETY: `vdupq_n_f64(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut acc = unsafe { vdupq_n_f64(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - NEON is always available on aarch64
        // - `i + 2 <= a.len()` ensures in-bounds access for `vld1q_f64`
        // - `vld1q_f64` tolerates unaligned pointers on aarch64
        // - `a.as_ptr().add(i)` is valid because `i + 2 <= a.len()`
        while i + 2 <= a.len() {
            let v = vld1q_f64(a.as_ptr().add(i));
            acc = vmaxq_f64(acc, v);
            i += 2;
        }
    }
    // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes and have the same representation
    let tmp: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(acc) };
    let mut max = tmp[0].max(tmp[1]);
    for &v in &a[i..] {
        if v > max { max = v; }
    }
    max
}

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

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

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