himada-dispatch 0.1.0

use himada_core::HardwareDNA;

// ---------------------------------------------------------------------------
// dot_clamp_reduce: c[i] = clamp(a[i] * b[i], 0, 1), then sum all c[i]
// ---------------------------------------------------------------------------

pub fn dot_clamp_reduce_scalar(a: &[f64], b: &[f64]) -> f64 {
    let mut sum = 0.0;
    let len = a.len().min(b.len());
    for i in 0..len {
        let v = a[i] * b[i];
        sum += if v < 0.0 { 0.0 } else if v > 1.0 { 1.0 } else { v };
    }
    sum
}

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

#[cfg(target_arch = "aarch64")]
pub fn dot_clamp_reduce_neon(a: &[f64], b: &[f64]) -> f64 {
    use std::arch::aarch64::*;
    let len = a.len().min(b.len());
    let mut sum = 0.0;
    let mut i = 0;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `len >= 2` checks before entering SIMD loop
    // - `i + 2 <= len` ensures in-bounds access for `vld1q_f64`
    // - `vdupq_n_f64`, `vmaxq_f64`, `vminq_f64`, `vaddq_f64` are pure compute intrinsics
    // - `std::mem::transmute(acc)` is safe: float64x2_t and [f64; 2] are both 16 bytes
    // - remainder is handled by scalar code after the block
    unsafe {
        if len >= 2 {
            let mut acc = vdupq_n_f64(0.0);
            let zero = vdupq_n_f64(0.0);
            let one = vdupq_n_f64(1.0);
            while i + 2 <= len {
                let va = vld1q_f64(a.as_ptr().add(i));
                let vb = vld1q_f64(b.as_ptr().add(i));
                let prod = vmulq_f64(va, vb);
                let clamped = vmaxq_f64(vminq_f64(prod, one), zero);
                acc = vaddq_f64(acc, clamped);
                i += 2;
            }
            let tmp: [f64; 2] = std::mem::transmute(acc);
            sum += tmp[0] + tmp[1];
        }
    }
    for j in i..len {
        let v = a[j] * b[j];
        sum += if v < 0.0 { 0.0 } else if v > 1.0 { 1.0 } else { v };
    }
    sum
}

#[cfg(not(target_arch = "aarch64"))]
pub fn dot_clamp_reduce_neon(_: &[f64], _: &[f64]) -> f64 { f64::NAN }
#[cfg(not(target_arch = "aarch64"))]
pub fn dot_clamp_reduce_neon_supported(_: &HardwareDNA) -> bool { false }

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

#[cfg(target_arch = "x86_64")]
pub fn dot_clamp_reduce_sse(a: &[f64], b: &[f64]) -> f64 {
    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")` confirms CPU support
    // - `len >= 2` checks before entering SIMD loop
    // - `i + 2 <= len` ensures in-bounds access for `_mm_loadu_pd`
    // - `_mm_setzero_pd`, `_mm_set1_pd`, `_mm_min_pd`, `_mm_max_pd`, `_mm_add_pd` are pure compute
    // - `std::mem::transmute(acc)` is safe: __m128d and [f64; 2] are both 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();
            let zero = _mm_setzero_pd();
            let one = _mm_set1_pd(1.0);
            while i + 2 <= len {
                let va = _mm_loadu_pd(a.as_ptr().add(i));
                let vb = _mm_loadu_pd(b.as_ptr().add(i));
                let prod = _mm_mul_pd(va, vb);
                let clamped = _mm_min_pd(_mm_max_pd(prod, zero), one);
                acc = _mm_add_pd(acc, clamped);
                i += 2;
            }
            let tmp: [f64; 2] = std::mem::transmute(acc);
            sum += tmp[0] + tmp[1];
        }
    }
    for j in i..len {
        let v = a[j] * b[j];
        sum += if v < 0.0 { 0.0 } else if v > 1.0 { 1.0 } else { v };
    }
    sum
}

#[cfg(target_arch = "x86_64")]
pub fn dot_clamp_reduce_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}
#[cfg(not(target_arch = "x86_64"))]
pub fn dot_clamp_reduce_sse(_: &[f64], _: &[f64]) -> f64 { f64::NAN }
#[cfg(not(target_arch = "x86_64"))]
pub fn dot_clamp_reduce_sse_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn dot_clamp_reduce_avx2(a: &[f64], b: &[f64]) -> f64 {
    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")` confirms CPU support
    // - `len >= 4` checks before entering SIMD loop
    // - `i + 4 <= len` ensures in-bounds access for `_mm256_loadu_pd`
    // - `_mm256_setzero_pd`, `_mm256_set1_pd`, `_mm256_min_pd`, `_mm256_max_pd`, `_mm256_add_pd` are pure compute
    // - `std::mem::transmute(acc)` is safe: __m256d and [f64; 4] are both 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();
            let zero = _mm256_setzero_pd();
            let one = _mm256_set1_pd(1.0);
            while i + 4 <= len {
                let va = _mm256_loadu_pd(a.as_ptr().add(i));
                let vb = _mm256_loadu_pd(b.as_ptr().add(i));
                let prod = _mm256_mul_pd(va, vb);
                let clamped = _mm256_min_pd(_mm256_max_pd(prod, zero), one);
                acc = _mm256_add_pd(acc, clamped);
                i += 4;
            }
            let tmp: [f64; 4] = std::mem::transmute(acc);
            sum += tmp[0] + tmp[1] + tmp[2] + tmp[3];
        }
    }
    for j in i..len {
        let v = a[j] * b[j];
        sum += if v < 0.0 { 0.0 } else if v > 1.0 { 1.0 } else { v };
    }
    sum
}

#[cfg(target_arch = "x86_64")]
pub fn dot_clamp_reduce_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}
#[cfg(not(target_arch = "x86_64"))]
pub fn dot_clamp_reduce_avx2(_: &[f64], _: &[f64]) -> f64 { f64::NAN }
#[cfg(not(target_arch = "x86_64"))]
pub fn dot_clamp_reduce_avx2_supported(_: &HardwareDNA) -> bool { false }

// ---------------------------------------------------------------------------
// matmul_bias_relu: c = relu(matmul(a, w) + bias)
// ---------------------------------------------------------------------------

pub fn matmul_bias_relu_scalar(a: &[f64], w: &[f64], bias: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) {
    for i in 0..m {
        for j in 0..n {
            let mut sum = 0.0;
            for t in 0..k {
                sum += a[i * k + t] * w[t * n + j];
            }
            sum += bias[j];
            c[i * n + j] = if sum > 0.0 { sum } else { 0.0 };
        }
    }
}

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

#[cfg(target_arch = "aarch64")]
pub fn matmul_bias_relu_neon(a: &[f64], w: &[f64], bias: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) {
    use std::arch::aarch64::*;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - All indices are bounds-checked: `j + 1 < n` or `j < n`
    // - Pointer arithmetic via `.add()` is valid because loop bounds are within slice lengths
    // - `vld1q_dup_f64`, `vld1q_f64`, `vld1_f64`, `vdupq_n_f64` are safe SIMD loads
    // - `vfmaq_f64`, `vaddq_f64`, `vmaxq_f64` are pure compute intrinsics
    // - `vcombine_f64(vw0_lo, vdup_n_f64(0.0))` pads with zero for the odd-column case
    // - `vst1q_f64` / direct store of transmuted lane are safe (bounds checked via j < n)
    unsafe {
        for i in 0..m {
            for j in (0..n).step_by(2) {
                let mut vacc0 = vdupq_n_f64(0.0);
                if j + 1 < n {
                    for t in 0..k {
                        let va = vld1q_dup_f64(a.as_ptr().add(i * k + t));
                        let vw0 = vld1q_f64(w.as_ptr().add(t * n + j));
                        vacc0 = vfmaq_f64(vacc0, va, vw0);
                    }
                    vacc0 = vaddq_f64(vacc0, vld1q_f64(bias.as_ptr().add(j)));
                    let zero = vdupq_n_f64(0.0);
                    vacc0 = vmaxq_f64(vacc0, zero);
                    vst1q_f64(c.as_mut_ptr().add(i * n + j), vacc0);
                } else {
                    for t in 0..k {
                        let va = vld1q_dup_f64(a.as_ptr().add(i * k + t));
                        let vw0_lo = vld1_f64(w.as_ptr().add(t * n + j));
                        vacc0 = vfmaq_f64(vacc0, va, vcombine_f64(vw0_lo, vdup_n_f64(0.0)));
                    }
                    let b = vld1q_f64(bias.as_ptr().add(j));
                    vacc0 = vaddq_f64(vacc0, b);
                    let zero = vdupq_n_f64(0.0);
                    vacc0 = vmaxq_f64(vacc0, zero);
                    // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes
                    let tmp: [f64; 2] = std::mem::transmute(vacc0);
                    c[i * n + j] = tmp[0];
                }
            }
        }
    }
}

#[cfg(target_arch = "aarch64")]
pub fn matmul_bias_relu_neon_supported(_: &HardwareDNA) -> bool { true }
#[cfg(not(target_arch = "aarch64"))]
pub fn matmul_bias_relu_neon(_: &[f64], _: &[f64], _: &[f64], _: &mut [f64], _: usize, _: usize, _: usize) {}
#[cfg(not(target_arch = "aarch64"))]
pub fn matmul_bias_relu_neon_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn matmul_bias_relu_sse(a: &[f64], w: &[f64], bias: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) {
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` would be checked by callers; intrinsic functions are unsafe
    // - All indices are bounds-checked: loops respect `m`, `k`, `n` dimensions
    // - `_mm_loadu_pd`, `_mm_load_sd`, `_mm_storeu_pd` tolerate unaligned pointers
    // - `_mm_setzero_pd`, `_mm_set1_pd`, `_mm_unpacklo_pd` are pure compute intrinsics
    // - Scalar fallback handles the odd-column remainder
    unsafe {
        for i in 0..m {
            for j in (0..n).step_by(2) {
                let mut vacc0 = _mm_setzero_pd();
                if j + 1 < n {
                    for t in 0..k {
                        let va = _mm_set1_pd(a[i * k + t]);
                        let vw = _mm_loadu_pd(w.as_ptr().add(t * n + j));
                        vacc0 = _mm_add_pd(vacc0, _mm_mul_pd(va, vw));
                    }
                    vacc0 = _mm_add_pd(vacc0, _mm_loadu_pd(bias.as_ptr().add(j)));
                    let zero = _mm_setzero_pd();
                    vacc0 = _mm_max_pd(vacc0, zero);
                    _mm_storeu_pd(c.as_mut_ptr().add(i * n + j), vacc0);
                } else {
                    for t in 0..k {
                        let va = _mm_set1_pd(a[i * k + t]);
                        let vw_lo = _mm_load_sd(w.as_ptr().add(t * n + j) as *const f64);
                        let vw = _mm_unpacklo_pd(vw_lo, _mm_setzero_pd());
                        vacc0 = _mm_add_pd(vacc0, _mm_mul_pd(va, vw));
                    }
                    let b = _mm_loadu_pd(bias.as_ptr().add(j));
                    vacc0 = _mm_add_pd(vacc0, b);
                    let zero = _mm_setzero_pd();
                    vacc0 = _mm_max_pd(vacc0, zero);
                    // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes
                    let tmp: [f64; 2] = std::mem::transmute(vacc0);
                    c[i * n + j] = tmp[0];
                }
            }
        }
    }
}

#[cfg(target_arch = "x86_64")]
pub fn matmul_bias_relu_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}
#[cfg(not(target_arch = "x86_64"))]
pub fn matmul_bias_relu_sse(_: &[f64], _: &[f64], _: &[f64], _: &mut [f64], _: usize, _: usize, _: usize) {}
#[cfg(not(target_arch = "x86_64"))]
pub fn matmul_bias_relu_sse_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn matmul_bias_relu_avx2(a: &[f64], w: &[f64], bias: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) {
    use std::arch::x86_64::*;

    // SAFETY:
    // - All indices are bounds-checked: loops respect `m`, `k`, `n` dimensions
    // - `_mm256_loadu_pd`, `_mm_loadu_pd`, `_mm256_storeu_pd` tolerate unaligned pointers
    // - `_mm256_setzero_pd`, `_mm256_set1_pd`, `_mm256_castpd128_pd256` are pure compute intrinsics
    // - Scalar fallback (via two/one element path) handles the remainder
    // - `std::mem::transmute(vacc0)` is safe: __m256d and [f64; 4] are both 32 bytes
    unsafe {
        for i in 0..m {
            for j in (0..n).step_by(4) {
                let mut vacc0 = _mm256_setzero_pd();
                let remaining = n - j;
                let full = remaining >= 4;
                let two = remaining >= 2;
                if full {
                    for t in 0..k {
                        let va = _mm256_set1_pd(a[i * k + t]);
                        let vw0 = _mm256_loadu_pd(w.as_ptr().add(t * n + j));
                        vacc0 = _mm256_add_pd(vacc0, _mm256_mul_pd(va, vw0));
                    }
                    vacc0 = _mm256_add_pd(vacc0, _mm256_loadu_pd(bias.as_ptr().add(j)));
                    let zero = _mm256_setzero_pd();
                    vacc0 = _mm256_max_pd(vacc0, zero);
                    _mm256_storeu_pd(c.as_mut_ptr().add(i * n + j), vacc0);
                } else if two {
                    for t in 0..k {
                        let va = _mm256_set1_pd(a[i * k + t]);
                        let vw0_128 = _mm_loadu_pd(w.as_ptr().add(t * n + j));
                        let vw0 = _mm256_castpd128_pd256(vw0_128);
                        vacc0 = _mm256_add_pd(vacc0, _mm256_mul_pd(va, vw0));
                    }
                    let b128 = _mm_loadu_pd(bias.as_ptr().add(j));
                    let b = _mm256_castpd128_pd256(b128);
                    vacc0 = _mm256_add_pd(vacc0, b);
                    let zero = _mm256_setzero_pd();
                    vacc0 = _mm256_max_pd(vacc0, zero);
                    let tmp: [f64; 4] = std::mem::transmute(vacc0);
                    c[i * n + j] = tmp[0];
                    c[i * n + j + 1] = tmp[1];
                }
            }
        }
    }
}

#[cfg(target_arch = "x86_64")]
pub fn matmul_bias_relu_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}
#[cfg(not(target_arch = "x86_64"))]
pub fn matmul_bias_relu_avx2(_: &[f64], _: &[f64], _: &[f64], _: &mut [f64], _: usize, _: usize, _: usize) {}
#[cfg(not(target_arch = "x86_64"))]
pub fn matmul_bias_relu_avx2_supported(_: &HardwareDNA) -> bool { false }

// ---------------------------------------------------------------------------
// General fusion helpers
// ---------------------------------------------------------------------------

use crate::{Dispatch, DotKernel, ReduceKernel, ClampKernel, HwdnaError};

/// A fused operation that does dot product + clamp + reduce in one pass.
pub struct FusedDotClampReduce {
    pub dot: Dispatch<DotKernel>,
    pub clamp: Dispatch<ClampKernel>,
    pub reduce: Dispatch<ReduceKernel>,
}

impl FusedDotClampReduce {
    pub fn new(dot: Dispatch<DotKernel>, clamp: Dispatch<ClampKernel>, reduce: Dispatch<ReduceKernel>) -> Self {
        Self { dot, clamp, reduce }
    }

    pub fn compute(&mut self, a: &[f64], b: &[f64], lo: f64, hi: f64) -> Result<f64, HwdnaError> {
        let n = a.len().min(b.len());
        let mut sum = 0.0f64;
        for i in 0..n {
            let d = a[i] * b[i];
            sum += d.clamp(lo, hi);
        }
        Ok(sum)
    }
}