himada-dispatch 0.1.1

use himada_core::HardwareDNA;

// ---------------------------------------------------------------------------
// cumsum_f64
// ---------------------------------------------------------------------------

pub fn cumsum_f64_scalar(input: &[f64], output: &mut [f64]) {
    let len = input.len().min(output.len());
    if len == 0 { return; }
    let mut acc = 0.0;
    for i in 0..len {
        acc += input[i];
        output[i] = acc;
    }
}

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

// --- NEON f64 ---

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

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 2 <= len` ensures in-bounds access for `vld1q_f64` / `vst1q_f64`
    // - `vld1q_f64` tolerates unaligned pointers on aarch64
    // - `std::mem::transmute` between float64x2_t and [f64; 2] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        while i + 2 <= len {
            let v = vld1q_f64(input.as_ptr().add(i));
            let mut arr: [f64; 2] = std::mem::transmute(v);
            running += arr[0]; arr[0] = running;
            running += arr[1]; arr[1] = running;
            vst1q_f64(output.as_mut_ptr().add(i), std::mem::transmute::<[f64; 2], float64x2_t>(arr));
            i += 2;
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- SSE f64 ---

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

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 2 <= len` ensures in-bounds access for `_mm_loadu_pd` / `_mm_storeu_pd`
    // - `_mm_loadu_pd` / `_mm_storeu_pd` tolerate unaligned pointers
    // - `std::mem::transmute` between __m128d and [f64; 2] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= len {
                let v = _mm_loadu_pd(input.as_ptr().add(i));
                let mut arr: [f64; 2] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                _mm_storeu_pd(output.as_mut_ptr().add(i), std::mem::transmute::<[f64; 2], __m128d>(arr));
                i += 2;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- AVX2 f64 ---

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

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 4 <= len` ensures in-bounds access for `_mm256_loadu_pd` / `_mm256_storeu_pd`
    // - `_mm256_loadu_pd` / `_mm256_storeu_pd` tolerate unaligned pointers
    // - `std::mem::transmute` between __m256d and [f64; 4] is safe (both 32 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= len {
                let v = _mm256_loadu_pd(input.as_ptr().add(i));
                let mut arr: [f64; 4] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                running += arr[2]; arr[2] = running;
                running += arr[3]; arr[3] = running;
                _mm256_storeu_pd(output.as_mut_ptr().add(i), std::mem::transmute::<[f64; 4], __m256d>(arr));
                i += 4;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// ---------------------------------------------------------------------------
// cumsum_f32
// ---------------------------------------------------------------------------

pub fn cumsum_f32_scalar(input: &[f32], output: &mut [f32]) {
    let len = input.len().min(output.len());
    if len == 0 { return; }
    let mut acc = 0.0;
    for i in 0..len {
        acc += input[i];
        output[i] = acc;
    }
}

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

// --- NEON f32 ---

#[cfg(target_arch = "aarch64")]
pub fn cumsum_f32_neon(input: &[f32], output: &mut [f32]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0.0;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 4 <= len` ensures in-bounds access for `vld1q_f32` / `vst1q_f32`
    // - `std::mem::transmute` between float32x4_t and [f32; 4] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        while i + 4 <= len {
            let v = vld1q_f32(input.as_ptr().add(i));
            let mut arr: [f32; 4] = std::mem::transmute(v);
            running += arr[0]; arr[0] = running;
            running += arr[1]; arr[1] = running;
            running += arr[2]; arr[2] = running;
            running += arr[3]; arr[3] = running;
            vst1q_f32(output.as_mut_ptr().add(i), std::mem::transmute::<[f32; 4], float32x4_t>(arr));
            i += 4;
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- SSE f32 ---

#[cfg(target_arch = "x86_64")]
pub fn cumsum_f32_sse(input: &[f32], output: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0.0;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 4 <= len` ensures in-bounds access for `_mm_loadu_ps` / `_mm_storeu_ps`
    // - `std::mem::transmute` between __m128 and [f32; 4] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            while i + 4 <= len {
                let v = _mm_loadu_ps(input.as_ptr().add(i));
                let mut arr: [f32; 4] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                running += arr[2]; arr[2] = running;
                running += arr[3]; arr[3] = running;
                _mm_storeu_ps(output.as_mut_ptr().add(i), std::mem::transmute::<[f32; 4], __m128>(arr));
                i += 4;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- AVX2 f32 ---

#[cfg(target_arch = "x86_64")]
pub fn cumsum_f32_avx2(input: &[f32], output: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0.0;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 8 <= len` ensures in-bounds access for `_mm256_loadu_ps` / `_mm256_storeu_ps`
    // - `std::mem::transmute` between __m256 and [f32; 8] is safe (both 32 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            while i + 8 <= len {
                let v = _mm256_loadu_ps(input.as_ptr().add(i));
                let mut arr: [f32; 8] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                running += arr[2]; arr[2] = running;
                running += arr[3]; arr[3] = running;
                running += arr[4]; arr[4] = running;
                running += arr[5]; arr[5] = running;
                running += arr[6]; arr[6] = running;
                running += arr[7]; arr[7] = running;
                _mm256_storeu_ps(output.as_mut_ptr().add(i), std::mem::transmute::<[f32; 8], __m256>(arr));
                i += 8;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// ---------------------------------------------------------------------------
// cumsum_i32
// ---------------------------------------------------------------------------

pub fn cumsum_i32_scalar(input: &[i32], output: &mut [i32]) {
    let len = input.len().min(output.len());
    if len == 0 { return; }
    let mut acc = 0;
    for i in 0..len {
        acc += input[i];
        output[i] = acc;
    }
}

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

// --- NEON i32 ---

#[cfg(target_arch = "aarch64")]
pub fn cumsum_i32_neon(input: &[i32], output: &mut [i32]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0i32;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 4 <= len` ensures in-bounds access for `vld1q_s32` / `vst1q_s32`
    // - `std::mem::transmute` between int32x4_t and [i32; 4] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        while i + 4 <= len {
            let v = vld1q_s32(input.as_ptr().add(i));
            let mut arr: [i32; 4] = std::mem::transmute(v);
            running += arr[0]; arr[0] = running;
            running += arr[1]; arr[1] = running;
            running += arr[2]; arr[2] = running;
            running += arr[3]; arr[3] = running;
            vst1q_s32(output.as_mut_ptr().add(i), std::mem::transmute::<[i32; 4], int32x4_t>(arr));
            i += 4;
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- SSE i32 ---

#[cfg(target_arch = "x86_64")]
pub fn cumsum_i32_sse(input: &[i32], output: &mut [i32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0i32;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 4 <= len` ensures in-bounds access for `_mm_loadu_si128` / `_mm_storeu_si128`
    // - `std::mem::transmute` between __m128i and [i32; 4] is safe (both 16 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            while i + 4 <= len {
                let v = _mm_loadu_si128(input.as_ptr().add(i) as *const __m128i);
                let mut arr: [i32; 4] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                running += arr[2]; arr[2] = running;
                running += arr[3]; arr[3] = running;
                _mm_storeu_si128(output.as_mut_ptr().add(i) as *mut __m128i, std::mem::transmute::<[i32; 4], __m128i>(arr));
                i += 4;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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

// --- AVX2 i32 ---

#[cfg(target_arch = "x86_64")]
pub fn cumsum_i32_avx2(input: &[i32], output: &mut [i32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(output.len());
    let mut i = 0;
    let mut running = 0i32;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 8 <= len` ensures in-bounds access for `_mm256_loadu_si256` / `_mm256_storeu_si256`
    // - `std::mem::transmute` between __m256i and [i32; 8] is safe (both 32 bytes)
    // - Scalar fallback for remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            while i + 8 <= len {
                let v = _mm256_loadu_si256(input.as_ptr().add(i) as *const __m256i);
                let mut arr: [i32; 8] = std::mem::transmute(v);
                running += arr[0]; arr[0] = running;
                running += arr[1]; arr[1] = running;
                running += arr[2]; arr[2] = running;
                running += arr[3]; arr[3] = running;
                running += arr[4]; arr[4] = running;
                running += arr[5]; arr[5] = running;
                running += arr[6]; arr[6] = running;
                running += arr[7]; arr[7] = running;
                _mm256_storeu_si256(output.as_mut_ptr().add(i) as *mut __m256i, std::mem::transmute::<[i32; 8], __m256i>(arr));
                i += 8;
            }
        }
    }
    for j in i..len {
        running += input[j];
        output[j] = running;
    }
}

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

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

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