himada-dispatch 0.1.1

use himada_core::HardwareDNA;

/// Softmax: output[i] = exp(input[i] - max) / sum(exp(input - max))
pub fn softmax_scalar(input: &[f64], output: &mut [f64]) {
    let max = input.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    let mut sum = 0.0;
    for i in 0..input.len() {
        let e = (input[i] - max).exp();
        output[i] = e;
        sum += e;
    }
    let inv = 1.0 / sum;
    for v in output.iter_mut() {
        *v *= inv;
    }
}

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

#[cfg(target_arch = "x86_64")]
pub fn softmax_sse(input: &[f64], output: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let n = input.len().min(output.len());

    let mut i = 0;
    // SAFETY: `_mm_set1_pd(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut vmax = unsafe { _mm_set1_pd(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("sse2")` confirms CPU support
        // - `i + 2 <= n` ensures in-bounds access for `_mm_loadu_pd`
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= n {
                let v = _mm_loadu_pd(input.as_ptr().add(i));
                vmax = _mm_max_pd(vmax, v);
                i += 2;
            }
        }
    }
    // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes
    let max_arr: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(vmax) };
    let mut max = max_arr[0].max(max_arr[1]);
    for &v in &input[i..n] {
        if v > max { max = v; }
    }

    let mut sum = 0.0;
    let mut i = 0;
    // SAFETY: `_mm_setzero_pd()` creates a zero-initialized vector; no data is accessed
    let mut vacc = unsafe { _mm_setzero_pd() };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("sse2")` confirms CPU support
        // - `i + 2 <= n` ensures in-bounds access for `_mm_loadu_pd` / `_mm_storeu_pd`
        if is_x86_feature_detected!("sse2") {
            let vmax = _mm_set1_pd(max);
            while i + 2 <= n {
                let v = _mm_loadu_pd(input.as_ptr().add(i));
                let e = _mm_sub_pd(v, vmax);
                // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes
                let ea: [f64; 2] = std::mem::transmute::<_, [f64; 2]>(e);
                let e0 = ea[0].exp();
                let e1 = ea[1].exp();
                _mm_storeu_pd(output.as_mut_ptr().add(i), _mm_set_pd(e1, e0));
                vacc = _mm_add_pd(vacc, _mm_set_pd(e1, e0));
                i += 2;
            }
        }
    }
    // SAFETY: transmuting __m128d to [f64; 2] is safe; both are 16 bytes
    let sum_arr: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(vacc) };
    sum = sum_arr[0] + sum_arr[1];
    for j in i..n {
        let e = (input[j] - max).exp();
        output[j] = e;
        sum += e;
    }

    let inv = 1.0 / sum;
    let mut i = 0;
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("sse2")` confirms CPU support
        // - `i + 2 <= n` ensures in-bounds access for `_mm_loadu_pd` / `_mm_storeu_pd`
        if is_x86_feature_detected!("sse2") {
            let vinv = _mm_set1_pd(inv);
            while i + 2 <= n {
                let v = _mm_loadu_pd(output.as_ptr().add(i));
                _mm_storeu_pd(output.as_mut_ptr().add(i), _mm_mul_pd(v, vinv));
                i += 2;
            }
        }
    }
    for v in output[i..n].iter_mut() {
        *v *= inv;
    }
}

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

#[cfg(not(target_arch = "x86_64"))]
pub fn softmax_sse(_: &[f64], _: &mut [f64]) {}

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

#[cfg(target_arch = "x86_64")]
pub fn softmax_avx2(input: &[f64], output: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let n = input.len().min(output.len());

    let mut i = 0;
    // SAFETY: `_mm256_set1_pd(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut vmax = unsafe { _mm256_set1_pd(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("avx2")` confirms CPU support
        // - `i + 4 <= n` ensures in-bounds access for `_mm256_loadu_pd`
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= n {
                let v = _mm256_loadu_pd(input.as_ptr().add(i));
                vmax = _mm256_max_pd(vmax, v);
                i += 4;
            }
        }
    }
    // SAFETY: transmuting __m256d to [f64; 4] is safe; both are 32 bytes
    let max_arr: [f64; 4] = unsafe { std::mem::transmute::<_, [f64; 4]>(vmax) };
    let mut max = max_arr[0].max(max_arr[1]).max(max_arr[2]).max(max_arr[3]);
    for &v in &input[i..n] {
        if v > max { max = v; }
    }

    let mut i = 0;
    // SAFETY: `_mm256_setzero_pd()` creates a zero-initialized vector; no data is accessed
    let mut vacc = unsafe { _mm256_setzero_pd() };
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("avx2")` confirms CPU support
        // - `i + 4 <= n` ensures in-bounds access for `_mm256_loadu_pd` / `_mm256_storeu_pd`
        if is_x86_feature_detected!("avx2") {
            let vmax = _mm256_set1_pd(max);
            while i + 4 <= n {
                let v = _mm256_loadu_pd(input.as_ptr().add(i));
                let e = _mm256_sub_pd(v, vmax);
                // SAFETY: transmuting __m256d to [f64; 4] is safe; both are 32 bytes
                let ea: [f64; 4] = std::mem::transmute::<_, [f64; 4]>(e);
                let e0 = ea[0].exp();
                let e1 = ea[1].exp();
                let e2 = ea[2].exp();
                let e3 = ea[3].exp();
                _mm256_storeu_pd(output.as_mut_ptr().add(i), _mm256_set_pd(e3, e2, e1, e0));
                vacc = _mm256_add_pd(vacc, _mm256_set_pd(e3, e2, e1, e0));
                i += 4;
            }
        }
    }
    // SAFETY: transmuting __m256d to [f64; 4] is safe; both are 32 bytes
    let sum_arr: [f64; 4] = unsafe { std::mem::transmute::<_, [f64; 4]>(vacc) };
    let mut sum = sum_arr[0] + sum_arr[1] + sum_arr[2] + sum_arr[3];
    for j in i..n {
        let e = (input[j] - max).exp();
        output[j] = e;
        sum += e;
    }

    let inv = 1.0 / sum;
    let mut i = 0;
    unsafe {
        // SAFETY:
        // - `is_x86_feature_detected!("avx2")` confirms CPU support
        // - `i + 4 <= n` ensures in-bounds access for `_mm256_loadu_pd` / `_mm256_storeu_pd`
        if is_x86_feature_detected!("avx2") {
            let vinv = _mm256_set1_pd(inv);
            while i + 4 <= n {
                let v = _mm256_loadu_pd(output.as_ptr().add(i));
                _mm256_storeu_pd(output.as_mut_ptr().add(i), _mm256_mul_pd(v, vinv));
                i += 4;
            }
        }
    }
    for v in output[i..n].iter_mut() {
        *v *= inv;
    }
}

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

#[cfg(not(target_arch = "x86_64"))]
pub fn softmax_avx2(_: &[f64], _: &mut [f64]) {}

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

#[cfg(target_arch = "aarch64")]
pub fn softmax_neon(input: &[f64], output: &mut [f64]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let n = input.len().min(output.len());

    let mut i = 0;
    // SAFETY: `vdupq_n_f64(f64::NEG_INFINITY)` broadcasts a constant; no data is accessed
    let mut vmax = unsafe { vdupq_n_f64(f64::NEG_INFINITY) };
    unsafe {
        // SAFETY:
        // - NEON is available on all aarch64 targets
        // - `i + 2 <= n` ensures in-bounds access for `vld1q_f64`
        while i + 2 <= n {
            let v = vld1q_f64(input.as_ptr().add(i));
            vmax = vmaxq_f64(vmax, v);
            i += 2;
        }
    }
    // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes
    let max_arr: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(vmax) };
    let mut max = max_arr[0].max(max_arr[1]);
    for &v in &input[i..n] {
        if v > max { max = v; }
    }

    let mut i = 0;
    // SAFETY: `vdupq_n_f64(0.0)` creates a zero-initialized vector; no data is accessed
    let mut vacc = unsafe { vdupq_n_f64(0.0) };
    unsafe {
        // SAFETY:
        // - NEON is available on all aarch64 targets
        // - `i + 2 <= n` ensures in-bounds access for `vld1q_f64` / `vst1q_f64`
        let vmax = vdupq_n_f64(max);
        while i + 2 <= n {
            let v = vld1q_f64(input.as_ptr().add(i));
            let e = vsubq_f64(v, vmax);
            // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes
            let ea: [f64; 2] = std::mem::transmute::<_, [f64; 2]>(e);
            let e0 = ea[0].exp();
            let e1 = ea[1].exp();
            // SAFETY: transmuting [f64; 2] back to float64x2_t is safe; both are 16 bytes
            let ve = { let t: [f64; 2] = [e0, e1]; std::mem::transmute::<[f64; 2], float64x2_t>(t) };
            vst1q_f64(output.as_mut_ptr().add(i), ve);
            vacc = vaddq_f64(vacc, ve);
            i += 2;
        }
    }
    // SAFETY: transmuting float64x2_t to [f64; 2] is safe; both are 16 bytes
    let sum_arr: [f64; 2] = unsafe { std::mem::transmute::<_, [f64; 2]>(vacc) };
    let mut sum = sum_arr[0] + sum_arr[1];
    for j in i..n {
        let e = (input[j] - max).exp();
        output[j] = e;
        sum += e;
    }

    let inv = 1.0 / sum;
    let mut i = 0;
    unsafe {
        // SAFETY:
        // - NEON is available on all aarch64 targets
        // - `i + 2 <= n` ensures in-bounds access for `vld1q_f64` / `vst1q_f64`
        let vinv = vdupq_n_f64(inv);
        while i + 2 <= n {
            let v = vld1q_f64(output.as_ptr().add(i));
            vst1q_f64(output.as_mut_ptr().add(i), vmulq_f64(v, vinv));
            i += 2;
        }
    }
    for v in output[i..n].iter_mut() {
        *v *= inv;
    }
}

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

#[cfg(not(target_arch = "aarch64"))]
pub fn softmax_neon(_: &[f64], _: &mut [f64]) {}

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