clawhdf5/crates/clawhdf5-accel/src/neon.rs

//! ARM NEON SIMD implementations.
//! NEON is always available on aarch64.

#![cfg(target_arch = "aarch64")]

use std::arch::aarch64::*;

/// NEON dot product for f32 slices.
///
/// # Safety
/// Caller must ensure aarch64 target (NEON always available).
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
#[target_feature(enable = "neon")]
pub unsafe fn dot_product(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let len = a.len();
    let mut i = 0;
    let mut acc0 = vdupq_n_f32(0.0);
    let mut acc1 = vdupq_n_f32(0.0);

    // Process 8 elements per iteration (2x4 unrolled)
    while i + 8 <= len {
        // SAFETY: Caller guarantees NEON/FP16 available per the # Safety contract on this fn.
        unsafe {
            let va0 = vld1q_f32(a.as_ptr().add(i));
            let vb0 = vld1q_f32(b.as_ptr().add(i));
            acc0 = vfmaq_f32(acc0, va0, vb0);

            let va1 = vld1q_f32(a.as_ptr().add(i + 4));
            let vb1 = vld1q_f32(b.as_ptr().add(i + 4));
            acc1 = vfmaq_f32(acc1, va1, vb1);
        }
        i += 8;
    }

    // Process remaining 4-element chunk
    if i + 4 <= len {
        // SAFETY: Caller guarantees NEON/FP16 available per the # Safety contract on this fn.
        unsafe {
            let va = vld1q_f32(a.as_ptr().add(i));
            let vb = vld1q_f32(b.as_ptr().add(i));
            acc0 = vfmaq_f32(acc0, va, vb);
        }
        i += 4;
    }

    let mut sum = vaddvq_f32(vaddq_f32(acc0, acc1));

    // Scalar tail
    while i < len {
        sum += a[i] * b[i];
        i += 1;
    }

    sum
}

/// NEON cosine similarity — fused single pass with 3 accumulators.
///
/// # Safety
/// Caller must ensure aarch64 target.
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
#[target_feature(enable = "neon")]
pub unsafe fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let len = a.len();
    let mut i = 0;

    let mut dot_acc = vdupq_n_f32(0.0);
    let mut norm_a_acc = vdupq_n_f32(0.0);
    let mut norm_b_acc = vdupq_n_f32(0.0);

    while i + 4 <= len {
        // SAFETY: Caller guarantees NEON/FP16 available per the # Safety contract on this fn.
        unsafe {
            let va = vld1q_f32(a.as_ptr().add(i));
            let vb = vld1q_f32(b.as_ptr().add(i));
            dot_acc = vfmaq_f32(dot_acc, va, vb);
            norm_a_acc = vfmaq_f32(norm_a_acc, va, va);
            norm_b_acc = vfmaq_f32(norm_b_acc, vb, vb);
        }
        i += 4;
    }

    let mut dot = vaddvq_f32(dot_acc);
    let mut norm_a = vaddvq_f32(norm_a_acc);
    let mut norm_b = vaddvq_f32(norm_b_acc);

    while i < len {
        dot += a[i] * b[i];
        norm_a += a[i] * a[i];
        norm_b += b[i] * b[i];
        i += 1;
    }

    let denom = (norm_a * norm_b).sqrt();
    if denom < f32::EPSILON {
        0.0
    } else {
        dot / denom
    }
}

/// NEON L2 distance.
///
/// # Safety
/// Caller must ensure aarch64 target.
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
#[target_feature(enable = "neon")]
pub unsafe fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());
    let len = a.len();
    let mut i = 0;
    let mut acc = vdupq_n_f32(0.0);

    while i + 4 <= len {
        // SAFETY: Caller guarantees NEON/FP16 available per the # Safety contract on this fn.
        unsafe {
            let va = vld1q_f32(a.as_ptr().add(i));
            let vb = vld1q_f32(b.as_ptr().add(i));
            let diff = vsubq_f32(va, vb);
            acc = vfmaq_f32(acc, diff, diff);
        }
        i += 4;
    }

    let mut sum = vaddvq_f32(acc);

    while i < len {
        let d = a[i] - b[i];
        sum += d * d;
        i += 1;
    }

    sum.sqrt()
}

/// NEON f16 to f32 batch conversion.
///
/// Note: Hardware vcvt_f32_f16 requires nightly (stdarch_neon_f16).
/// On stable Rust, we delegate to the scalar implementation.
/// The NEON module still provides the function for API uniformity.
pub fn f16_to_f32_batch(input: &[u16], output: &mut [f32]) {
    // Delegate to scalar — hardware f16 intrinsics are unstable on aarch64.
    crate::scalar::f16_to_f32_batch(input, output);
}

/// Fletcher32 checksum (scalar implementation, no NEON intrinsics used).
///
/// This function uses no NEON intrinsics despite living in the neon module.
/// It is safe to call without feature detection.
pub fn checksum_fletcher32(data: &[u8]) -> u32 {
    let mut sum1: u32 = 0xFFFF;
    let mut sum2: u32 = 0xFFFF;

    let mut i = 0;
    // Process in blocks of 360 words (720 bytes) to avoid overflow before modulo
    // 360 * 65535 fits in u32
    while i + 1 < data.len() {
        let remaining_words = (data.len() - i) / 2;
        let block_words = remaining_words.min(360);

        for _ in 0..block_words {
            let word = ((data[i] as u32) << 8) | (data[i + 1] as u32);
            sum1 += word;
            sum2 += sum1;
            i += 2;
        }

        sum1 %= 65535;
        sum2 %= 65535;
    }

    // Handle trailing byte
    if i < data.len() {
        let word = (data[i] as u32) << 8;
        sum1 = (sum1 + word) % 65535;
        sum2 = (sum2 + sum1) % 65535;
    }

    (sum2 << 16) | sum1
}

/// NEON dot product of two `i8` slices, widened to `i32`, for any aarch64 CPU.
///
/// `vmull_s8` multiplies eight lanes into `i16` — even `-128 * -128` is 16 384,
/// inside `i16` — and `vpadalq_s16` adds adjacent pairs of those into `i32`
/// accumulators, so nothing can overflow before the final horizontal sum.
///
/// CPUs with the ARMv8.2 dot-product extension should use
/// [`dot_i8_dotprod`], which does the multiply and the accumulate in one
/// instruction.
///
/// # Safety
/// Caller must ensure aarch64 target (NEON always available).
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
#[target_feature(enable = "neon")]
pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
    assert_eq!(a.len(), b.len());
    let len = a.len();
    let mut i = 0;
    let mut acc0 = vdupq_n_s32(0);
    let mut acc1 = vdupq_n_s32(0);

    while i + 16 <= len {
        // SAFETY: NEON is available per the # Safety contract, and both
        // 16-byte loads start at an index checked against `len` above.
        unsafe {
            let va = vld1q_s8(a.as_ptr().add(i));
            let vb = vld1q_s8(b.as_ptr().add(i));
            acc0 = vpadalq_s16(acc0, vmull_s8(vget_low_s8(va), vget_low_s8(vb)));
            acc1 = vpadalq_s16(acc1, vmull_high_s8(va, vb));
        }
        i += 16;
    }

    let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
    while i < len {
        sum += i32::from(a[i]) * i32::from(b[i]);
        i += 1;
    }
    sum
}

/// One `SDOT`: for each of the four `i32` lanes of `acc`, add the dot
/// product of the corresponding four `i8` pairs from `a` and `b`.
///
/// Written as inline assembly because the `vdotq_s32` intrinsic is still
/// behind the unstable `stdarch_neon_dotprod` feature; inline assembly is
/// stable on aarch64.
///
/// # Safety
/// Caller must ensure the CPU supports the `dotprod` extension.
#[inline]
#[target_feature(enable = "neon,dotprod")]
unsafe fn sdot(acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
    let mut acc = acc;
    // SAFETY: `dotprod` is enabled for this function and the caller
    // guarantees the CPU supports it. The instruction reads only its three
    // vector registers and touches no memory.
    unsafe {
        std::arch::asm!(
            "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
            acc = inout(vreg) acc,
            a = in(vreg) a,
            b = in(vreg) b,
            options(pure, nomem, nostack),
        );
    }
    acc
}

/// NEON dot product of two `i8` slices using the ARMv8.2 dot-product
/// extension (`SDOT`): sixteen multiply-accumulates per instruction, straight
/// into `i32` lanes.
///
/// Present on the cores this crate actually runs on — Cortex-A76 and later
/// (Raspberry Pi 5, current Android phones), Neoverse-N1 (Graviton2, Ampere
/// Altra), and every Apple Silicon generation.
///
/// # Safety
/// Caller must verify `is_aarch64_feature_detected!("dotprod")`.
// SAFETY: caller has verified the dotprod extension at runtime.
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn dot_i8_dotprod(a: &[i8], b: &[i8]) -> i32 {
    assert_eq!(a.len(), b.len());
    let len = a.len();
    let mut i = 0;
    let mut acc0 = vdupq_n_s32(0);
    let mut acc1 = vdupq_n_s32(0);

    // Two independent accumulators so consecutive SDOTs are not serialised on
    // one register.
    while i + 32 <= len {
        // SAFETY: dotprod is available per the # Safety contract, and every
        // 16-byte load starts at an index checked against `len` above.
        unsafe {
            acc0 = sdot(
                acc0,
                vld1q_s8(a.as_ptr().add(i)),
                vld1q_s8(b.as_ptr().add(i)),
            );
            acc1 = sdot(
                acc1,
                vld1q_s8(a.as_ptr().add(i + 16)),
                vld1q_s8(b.as_ptr().add(i + 16)),
            );
        }
        i += 32;
    }
    if i + 16 <= len {
        // SAFETY: as above; the load is bounds-checked by this condition.
        unsafe {
            acc0 = sdot(
                acc0,
                vld1q_s8(a.as_ptr().add(i)),
                vld1q_s8(b.as_ptr().add(i)),
            );
        }
        i += 16;
    }

    let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
    while i < len {
        sum += i32::from(a[i]) * i32::from(b[i]);
        i += 1;
    }
    sum
}