himada-dispatch 0.1.0

use himada_core::HardwareDNA;

// ---------------------------------------------------------------------------
// layer_norm_f64
// ---------------------------------------------------------------------------

pub fn layer_norm_f64_scalar(input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) {
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }
    let mut sum = 0.0;
    for &v in input.iter().take(len) { sum += v; }
    let mean = sum / len as f64;
    let mut var = 0.0;
    for &v in input.iter().take(len) { let d = v - mean; var += d * d; }
    let var = var / len as f64;
    let inv_std = 1.0 / (var + epsilon).sqrt();
    for i in 0..len {
        output[i] = gamma[i] * (input[i] - mean) * inv_std + beta[i];
    }
}

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

// --- NEON f64 ---

#[cfg(target_arch = "aarch64")]
pub fn layer_norm_f64_neon(input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `vld1q_f64` / `vst1q_f64` tolerate unaligned pointers on aarch64
    // - `vdupq_n_f64`, `vaddq_f64`, `vsubq_f64`, `vmulq_f64`, `vaddvq_f64` are pure compute
    // - `std::mem::transmute` between SIMD and array types is safe (same size)
    // - Scalar fallback for remainder elements
    unsafe {
        let mut i = 0;
        let mut vsum = vdupq_n_f64(0.0);
        while i + 2 <= len {
            vsum = vaddq_f64(vsum, vld1q_f64(input.as_ptr().add(i)));
            i += 2;
        }
        let mut sum = vaddvq_f64(vsum);
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f64;
        let vmean = vdupq_n_f64(mean);

        let mut i = 0;
        let mut vvar = vdupq_n_f64(0.0);
        while i + 2 <= len {
            let d = vsubq_f64(vld1q_f64(input.as_ptr().add(i)), vmean);
            vvar = vaddq_f64(vvar, vmulq_f64(d, d));
            i += 2;
        }
        let mut var = vaddvq_f64(vvar);
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f64;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = vdupq_n_f64(inv_std);

        let mut i = 0;
        while i + 2 <= len {
            let vin = vld1q_f64(input.as_ptr().add(i));
            let vg = vld1q_f64(gamma.as_ptr().add(i));
            let vb = vld1q_f64(beta.as_ptr().add(i));
            let vnorm = vmulq_f64(vsubq_f64(vin, vmean), vinv_std);
            vst1q_f64(output.as_mut_ptr().add(i), vaddq_f64(vmulq_f64(vg, vnorm), vb));
            i += 2;
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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

// --- SSE f64 ---

#[cfg(target_arch = "x86_64")]
pub fn layer_norm_f64_sse(input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support where checked
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `_mm_loadu_pd` / `_mm_storeu_pd` tolerate unaligned pointers
    // - `_mm_setzero_pd`, `_mm_set1_pd`, `_mm_add_pd`, `_mm_sub_pd`, `_mm_mul_pd` are pure compute
    // - `std::mem::transmute` between __m128d and [f64; 2] is safe (same size)
    // - Scalar fallback for remainder elements and when SSE2 is not detected
    unsafe {
        let mut i = 0;
        let mut vsum = _mm_setzero_pd();
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= len {
                vsum = _mm_add_pd(vsum, _mm_loadu_pd(input.as_ptr().add(i)));
                i += 2;
            }
        }
        let tmp: [f64; 2] = std::mem::transmute(vsum);
        let mut sum = tmp[0] + tmp[1];
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f64;
        let vmean = _mm_set1_pd(mean);

        let mut i = 0;
        let mut vvar = _mm_setzero_pd();
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= len {
                let d = _mm_sub_pd(_mm_loadu_pd(input.as_ptr().add(i)), vmean);
                vvar = _mm_add_pd(vvar, _mm_mul_pd(d, d));
                i += 2;
            }
        }
        let tmp: [f64; 2] = std::mem::transmute(vvar);
        let mut var = tmp[0] + tmp[1];
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f64;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = _mm_set1_pd(inv_std);

        let mut i = 0;
        if is_x86_feature_detected!("sse2") {
            while i + 2 <= len {
                let vin = _mm_loadu_pd(input.as_ptr().add(i));
                let vg = _mm_loadu_pd(gamma.as_ptr().add(i));
                let vb = _mm_loadu_pd(beta.as_ptr().add(i));
                let vnorm = _mm_mul_pd(_mm_sub_pd(vin, vmean), vinv_std);
                _mm_storeu_pd(output.as_mut_ptr().add(i), _mm_add_pd(_mm_mul_pd(vg, vnorm), vb));
                i += 2;
            }
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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

// --- AVX2 f64 ---

#[cfg(target_arch = "x86_64")]
pub fn layer_norm_f64_avx2(input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support where checked
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `_mm256_loadu_pd` / `_mm256_storeu_pd` tolerate unaligned pointers
    // - `_mm256_setzero_pd`, `_mm256_set1_pd`, `_mm256_add_pd`, `_mm256_sub_pd`, `_mm256_mul_pd` are pure compute
    // - `std::mem::transmute` between __m256d and [f64; 4] is safe (same size)
    // - Scalar fallback for remainder elements
    unsafe {
        let mut i = 0;
        let mut vsum = _mm256_setzero_pd();
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= len {
                vsum = _mm256_add_pd(vsum, _mm256_loadu_pd(input.as_ptr().add(i)));
                i += 4;
            }
        }
        let tmp: [f64; 4] = std::mem::transmute(vsum);
        let mut sum = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f64;
        let vmean = _mm256_set1_pd(mean);

        let mut i = 0;
        let mut vvar = _mm256_setzero_pd();
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= len {
                let d = _mm256_sub_pd(_mm256_loadu_pd(input.as_ptr().add(i)), vmean);
                vvar = _mm256_add_pd(vvar, _mm256_mul_pd(d, d));
                i += 4;
            }
        }
        let tmp: [f64; 4] = std::mem::transmute(vvar);
        let mut var = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f64;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = _mm256_set1_pd(inv_std);

        let mut i = 0;
        if is_x86_feature_detected!("avx2") {
            while i + 4 <= len {
                let vin = _mm256_loadu_pd(input.as_ptr().add(i));
                let vg = _mm256_loadu_pd(gamma.as_ptr().add(i));
                let vb = _mm256_loadu_pd(beta.as_ptr().add(i));
                let vnorm = _mm256_mul_pd(_mm256_sub_pd(vin, vmean), vinv_std);
                _mm256_storeu_pd(output.as_mut_ptr().add(i), _mm256_add_pd(_mm256_mul_pd(vg, vnorm), vb));
                i += 4;
            }
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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

// ---------------------------------------------------------------------------
// layer_norm_f32
// ---------------------------------------------------------------------------

pub fn layer_norm_f32_scalar(input: &[f32], gamma: &[f32], beta: &[f32], epsilon: f32, output: &mut [f32]) {
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }
    let mut sum = 0.0;
    for &v in input.iter().take(len) { sum += v; }
    let mean = sum / len as f32;
    let mut var = 0.0;
    for &v in input.iter().take(len) { let d = v - mean; var += d * d; }
    let var = var / len as f32;
    let inv_std = 1.0 / (var + epsilon).sqrt();
    for i in 0..len {
        output[i] = gamma[i] * (input[i] - mean) * inv_std + beta[i];
    }
}

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

// --- NEON f32 ---

#[cfg(target_arch = "aarch64")]
pub fn layer_norm_f32_neon(input: &[f32], gamma: &[f32], beta: &[f32], epsilon: f32, output: &mut [f32]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `vld1q_f32` / `vst1q_f32` tolerate unaligned pointers on aarch64
    // - `vdupq_n_f32`, `vaddq_f32`, `vsubq_f32`, `vmulq_f32` are pure compute
    // - `std::mem::transmute` between float32x4_t and [f32; 4] is safe (same size)
    // - Scalar fallback for remainder elements
    unsafe {
        let mut i = 0;
        let mut vsum = vdupq_n_f32(0.0);
        while i + 4 <= len {
            vsum = vaddq_f32(vsum, vld1q_f32(input.as_ptr().add(i)));
            i += 4;
        }
        let tmp: [f32; 4] = std::mem::transmute(vsum);
        let mut sum = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f32;
        let vmean = vdupq_n_f32(mean);

        let mut i = 0;
        let mut vvar = vdupq_n_f32(0.0);
        while i + 4 <= len {
            let d = vsubq_f32(vld1q_f32(input.as_ptr().add(i)), vmean);
            vvar = vaddq_f32(vvar, vmulq_f32(d, d));
            i += 4;
        }
        let tmp: [f32; 4] = std::mem::transmute(vvar);
        let mut var = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f32;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = vdupq_n_f32(inv_std);

        let mut i = 0;
        while i + 4 <= len {
            let vin = vld1q_f32(input.as_ptr().add(i));
            let vg = vld1q_f32(gamma.as_ptr().add(i));
            let vb = vld1q_f32(beta.as_ptr().add(i));
            let vnorm = vmulq_f32(vsubq_f32(vin, vmean), vinv_std);
            vst1q_f32(output.as_mut_ptr().add(i), vaddq_f32(vmulq_f32(vg, vnorm), vb));
            i += 4;
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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

// --- SSE f32 ---

#[cfg(target_arch = "x86_64")]
pub fn layer_norm_f32_sse(input: &[f32], gamma: &[f32], beta: &[f32], epsilon: f32, output: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support where checked
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `_mm_loadu_ps` / `_mm_storeu_ps` tolerate unaligned pointers
    // - `_mm_setzero_ps`, `_mm_set1_ps`, `_mm_add_ps`, `_mm_sub_ps`, `_mm_mul_ps` are pure compute
    // - `std::mem::transmute` between __m128 and [f32; 4] is safe (same size)
    // - Scalar fallback for remainder elements
    unsafe {
        let mut i = 0;
        let mut vsum = _mm_setzero_ps();
        if is_x86_feature_detected!("sse2") {
            while i + 4 <= len {
                vsum = _mm_add_ps(vsum, _mm_loadu_ps(input.as_ptr().add(i)));
                i += 4;
            }
        }
        let tmp: [f32; 4] = std::mem::transmute(vsum);
        let mut sum = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f32;
        let vmean = _mm_set1_ps(mean);

        let mut i = 0;
        let mut vvar = _mm_setzero_ps();
        if is_x86_feature_detected!("sse2") {
            while i + 4 <= len {
                let d = _mm_sub_ps(_mm_loadu_ps(input.as_ptr().add(i)), vmean);
                vvar = _mm_add_ps(vvar, _mm_mul_ps(d, d));
                i += 4;
            }
        }
        let tmp: [f32; 4] = std::mem::transmute(vvar);
        let mut var = tmp[0] + tmp[1] + tmp[2] + tmp[3];
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f32;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = _mm_set1_ps(inv_std);

        let mut i = 0;
        if is_x86_feature_detected!("sse2") {
            while i + 4 <= len {
                let vin = _mm_loadu_ps(input.as_ptr().add(i));
                let vg = _mm_loadu_ps(gamma.as_ptr().add(i));
                let vb = _mm_loadu_ps(beta.as_ptr().add(i));
                let vnorm = _mm_mul_ps(_mm_sub_ps(vin, vmean), vinv_std);
                _mm_storeu_ps(output.as_mut_ptr().add(i), _mm_add_ps(_mm_mul_ps(vg, vnorm), vb));
                i += 4;
            }
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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

// --- AVX2 f32 ---

#[cfg(target_arch = "x86_64")]
pub fn layer_norm_f32_avx2(input: &[f32], gamma: &[f32], beta: &[f32], epsilon: f32, output: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let len = input.len().min(gamma.len()).min(beta.len()).min(output.len());
    if len == 0 { return; }

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support where checked
    // - All pointer arithmetic is bounded by `len` (min of all slices)
    // - `_mm256_loadu_ps` / `_mm256_storeu_ps` tolerate unaligned pointers
    // - `_mm256_setzero_ps`, `_mm256_set1_ps`, `_mm256_add_ps`, `_mm256_sub_ps`, `_mm256_mul_ps` are pure compute
    // - `std::mem::transmute` between __m256 and [f32; 8] is safe (same size)
    // - Scalar fallback for remainder elements
    unsafe {
        let mut i = 0;
        let mut vsum = _mm256_setzero_ps();
        if is_x86_feature_detected!("avx2") {
            while i + 8 <= len {
                vsum = _mm256_add_ps(vsum, _mm256_loadu_ps(input.as_ptr().add(i)));
                i += 8;
            }
        }
        let tmp: [f32; 8] = std::mem::transmute(vsum);
        let mut sum: f32 = tmp.iter().sum();
        for &v in input.iter().skip(i).take(len.saturating_sub(i)) { sum += v; }
        let mean = sum / len as f32;
        let vmean = _mm256_set1_ps(mean);

        let mut i = 0;
        let mut vvar = _mm256_setzero_ps();
        if is_x86_feature_detected!("avx2") {
            while i + 8 <= len {
                let d = _mm256_sub_ps(_mm256_loadu_ps(input.as_ptr().add(i)), vmean);
                vvar = _mm256_add_ps(vvar, _mm256_mul_ps(d, d));
                i += 8;
            }
        }
        let tmp: [f32; 8] = std::mem::transmute(vvar);
        let mut var: f32 = tmp.iter().sum();
        for (_, &v) in input.iter().enumerate().skip(i).take(len.saturating_sub(i)) {
            let d = v - mean; var += d * d;
        }
        let var = var / len as f32;
        let inv_std = 1.0 / (var + epsilon).sqrt();
        let vinv_std = _mm256_set1_ps(inv_std);

        let mut i = 0;
        if is_x86_feature_detected!("avx2") {
            while i + 8 <= len {
                let vin = _mm256_loadu_ps(input.as_ptr().add(i));
                let vg = _mm256_loadu_ps(gamma.as_ptr().add(i));
                let vb = _mm256_loadu_ps(beta.as_ptr().add(i));
                let vnorm = _mm256_mul_ps(_mm256_sub_ps(vin, vmean), vinv_std);
                _mm256_storeu_ps(output.as_mut_ptr().add(i), _mm256_add_ps(_mm256_mul_ps(vg, vnorm), vb));
                i += 8;
            }
        }
        for j in i..len {
            output[j] = gamma[j] * (input[j] - mean) * inv_std + beta[j];
        }
    }
}

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

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

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