Function ternary_dot_neon

nodedb-codec/src/vector_quant/ternary/simd.rs:58–96

(a_hot: &[u8], b_hot: &[u8], dim: usize)

#[cfg(target_arch = "aarch64")]
fn ternary_dot_neon(a_hot: &[u8], b_hot: &[u8], dim: usize) -> i32 {
    use std::arch::aarch64::*;

    let a_trits = unpack_hot(a_hot, dim);
    let b_trits = unpack_hot(b_hot, dim);

    let len = a_trits.len();
    let chunks = len / 16;
    let mut acc: i32;

    unsafe {
        let mut sum = vdupq_n_s32(0i32);
        for i in 0..chunks {
            let a_ptr = a_trits.as_ptr().add(i * 16);
            let b_ptr = b_trits.as_ptr().add(i * 16);
            let a_vec = vld1q_s8(a_ptr);
            let b_vec = vld1q_s8(b_ptr);
            let prod = vmulq_s8(a_vec, b_vec);
            let prod_lo = vmovl_s8(vget_low_s8(prod));
            let prod_hi = vmovl_s8(vget_high_s8(prod));
            let prod32_lo = vmovl_s16(vget_low_s16(prod_lo));
            let prod32_hi = vmovl_s16(vget_high_s16(prod_lo));
            let prod32_lo2 = vmovl_s16(vget_low_s16(prod_hi));
            let prod32_hi2 = vmovl_s16(vget_high_s16(prod_hi));
            sum = vaddq_s32(
                sum,
                vaddq_s32(
                    vaddq_s32(prod32_lo, prod32_hi),
                    vaddq_s32(prod32_lo2, prod32_hi2),
                ),
            );
        }
        acc = vaddvq_s32(sum);
        for i in (chunks * 16)..len {
            acc += a_trits[i] as i32 * b_trits[i] as i32;
        }
    }
    acc
}

type DotFn = fn(&[u8], &[u8], usize) -> i32;