oxillama-quant 0.1.3

//! Q5_1 NEON (AArch64) SIMD kernel.
//!
//! Q5_1 block format (24 bytes per 32 weights):
//! - bytes[0..2]: FP16 scale `d` (little-endian)
//! - bytes[2..4]: FP16 minimum `m` (little-endian)
//! - bytes[4..8]: `qh` — 32 high bits (bit 4 of each 5-bit quant), u32 LE
//! - bytes[8..24]: `qs` — 32 × lower 4 bits packed (2 per byte)
//!
//! Weight layout (sequential, not interleaved):
//!   output\[i\]      = d * (qs\[i\].lo4 | ((qh >> i) & 1) << 4) + m  for i in 0..16
//!   output\[i + 16\] = d * (qs\[i\].hi4 | ((qh >> (i+16)) & 1) << 4) + m  for i in 0..16
//!
//! Q5_1 is UNSIGNED: range [0..31], no -16 bias. Affine: d * q + m.

#![cfg(all(feature = "simd-neon", target_arch = "aarch64"))]

use core::arch::aarch64::*;

use crate::error::{QuantError, QuantResult};
use crate::traits::QuantKernel;
use crate::types::QuantTensor;

/// Number of weights per Q5_1 block.
pub const BLOCK_SIZE: usize = 32;
/// Number of bytes per Q5_1 block.
pub const BLOCK_BYTES: usize = 24;

/// NEON-accelerated Q5_1 kernel (AArch64 only).
#[allow(non_camel_case_types)]
pub struct Q5_1Neon;

#[inline(always)]
fn f16_to_f32(bits: u16) -> f32 {
    half::f16::from_bits(bits).to_f32()
}

#[inline(always)]
unsafe fn hsum_f32x4(v: float32x4_t) -> f32 {
    unsafe { vaddvq_f32(v) }
}

/// Expand `qh` into two `[u8; 16]` arrays for the lo and hi halves.
#[inline(always)]
fn expand_qh(qh: u32) -> ([u8; 16], [u8; 16]) {
    let lo: [u8; 16] = core::array::from_fn(|i| ((qh >> i) & 1) as u8);
    let hi: [u8; 16] = core::array::from_fn(|i| ((qh >> (i + 16)) & 1) as u8);
    (lo, hi)
}

/// Dequantize one Q5_1 block using NEON intrinsics.
///
/// Produces 32 f32 values in sequential layout:
///   output[0..16]  = lo-half (from lo nibbles + qh bits 0..15)
///   output[16..32] = hi-half (from hi nibbles + qh bits 16..31)
/// Formula: `d * q5 + m` where q5 ∈ [0..31] (unsigned, no centering).
///
/// # Safety
/// Must be called on AArch64 with NEON. `qs_ptr` must point to 16 valid bytes.
/// `output` must have at least 32 slots.
#[inline]
unsafe fn dequant_block_neon(
    qs_ptr: *const u8,
    qh_lo: &[u8; 16],
    qh_hi: &[u8; 16],
    d: f32,
    m: f32,
    output: &mut [f32],
) {
    let d_vec = unsafe { vdupq_n_f32(d) };
    let m_vec = unsafe { vdupq_n_f32(m) };

    let raw = unsafe { vld1q_u8(qs_ptr) };
    let mask = unsafe { vdupq_n_u8(0x0F) };
    let lo_nib = unsafe { vandq_u8(raw, mask) };
    let hi_nib = unsafe { vshrq_n_u8::<4>(raw) };

    let vqh_lo = unsafe { vld1q_u8(qh_lo.as_ptr()) };
    let vqh_hi = unsafe { vld1q_u8(qh_hi.as_ptr()) };

    // Shift high bit into bit-position 4
    let shift4 = unsafe { vdupq_n_u8(4) };
    let qh_lo_shifted = unsafe { vshlq_u8(vqh_lo, vreinterpretq_s8_u8(shift4)) };
    let qh_hi_shifted = unsafe { vshlq_u8(vqh_hi, vreinterpretq_s8_u8(shift4)) };

    // 5-bit unsigned quants
    let q5_lo = unsafe { vorrq_u8(lo_nib, qh_lo_shifted) };
    let q5_hi = unsafe { vorrq_u8(hi_nib, qh_hi_shifted) };

    // Widen to u16 → u32 → f32, apply affine: m + d * q (no centering bias)
    let q5_lo_u16_low = unsafe { vmovl_u8(vget_low_u8(q5_lo)) };
    let q5_lo_u16_high = unsafe { vmovl_u8(vget_high_u8(q5_lo)) };
    let q5_hi_u16_low = unsafe { vmovl_u8(vget_low_u8(q5_hi)) };
    let q5_hi_u16_high = unsafe { vmovl_u8(vget_high_u8(q5_hi)) };

    let fa = unsafe {
        vfmaq_f32(
            m_vec,
            vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_low))),
            d_vec,
        )
    };
    let fb = unsafe { vfmaq_f32(m_vec, vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_low)), d_vec) };
    let fc = unsafe {
        vfmaq_f32(
            m_vec,
            vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_high))),
            d_vec,
        )
    };
    let fe = unsafe { vfmaq_f32(m_vec, vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_high)), d_vec) };
    let fg = unsafe {
        vfmaq_f32(
            m_vec,
            vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_low))),
            d_vec,
        )
    };
    let fh = unsafe { vfmaq_f32(m_vec, vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_low)), d_vec) };
    let fj = unsafe {
        vfmaq_f32(
            m_vec,
            vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_high))),
            d_vec,
        )
    };
    let fk = unsafe { vfmaq_f32(m_vec, vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_high)), d_vec) };

    unsafe { vst1q_f32(output.as_mut_ptr(), fa) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(4), fb) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(8), fc) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(12), fe) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(16), fg) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(20), fh) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(24), fj) };
    unsafe { vst1q_f32(output.as_mut_ptr().add(28), fk) };
}

/// Compute dot product for one Q5_1 block against a contiguous f32 input slice.
///
/// Mirrors `reference::q5_1::Q5_1Ref::gemv`: computes Σ d*q*inp[i] + m*Σinp[i].
///
/// # Safety
/// Must be called on AArch64 with NEON. `qs_ptr` must point to 16 valid bytes.
/// `input` must have exactly 32 elements.
#[inline]
unsafe fn dot_block_neon(
    qs_ptr: *const u8,
    qh_lo: &[u8; 16],
    qh_hi: &[u8; 16],
    d: f32,
    m: f32,
    input: &[f32],
) -> f32 {
    let d_vec = unsafe { vdupq_n_f32(d) };

    let raw = unsafe { vld1q_u8(qs_ptr) };
    let mask = unsafe { vdupq_n_u8(0x0F) };
    let lo_nib = unsafe { vandq_u8(raw, mask) };
    let hi_nib = unsafe { vshrq_n_u8::<4>(raw) };

    let vqh_lo = unsafe { vld1q_u8(qh_lo.as_ptr()) };
    let vqh_hi = unsafe { vld1q_u8(qh_hi.as_ptr()) };

    let shift4 = unsafe { vdupq_n_u8(4) };
    let qh_lo_shifted = unsafe { vshlq_u8(vqh_lo, vreinterpretq_s8_u8(shift4)) };
    let qh_hi_shifted = unsafe { vshlq_u8(vqh_hi, vreinterpretq_s8_u8(shift4)) };

    let q5_lo = unsafe { vorrq_u8(lo_nib, qh_lo_shifted) };
    let q5_hi = unsafe { vorrq_u8(hi_nib, qh_hi_shifted) };

    let q5_lo_u16_low = unsafe { vmovl_u8(vget_low_u8(q5_lo)) };
    let q5_lo_u16_high = unsafe { vmovl_u8(vget_high_u8(q5_lo)) };
    let q5_hi_u16_low = unsafe { vmovl_u8(vget_low_u8(q5_hi)) };
    let q5_hi_u16_high = unsafe { vmovl_u8(vget_high_u8(q5_hi)) };

    let qf_a = unsafe { vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_low))) };
    let qf_b = unsafe { vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_low)) };
    let qf_c = unsafe { vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_high))) };
    let qf_e = unsafe { vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_high)) };
    let qf_g = unsafe { vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_low))) };
    let qf_h = unsafe { vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_low)) };
    let qf_j = unsafe { vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_high))) };
    let qf_k = unsafe { vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_high)) };

    // Sequential input: [0..4, 4..8, 8..12, 12..16] for lo, [16..20, 20..24, 24..28, 28..32] for hi
    let ip = input.as_ptr();
    let i0 = unsafe { vld1q_f32(ip) };
    let i1 = unsafe { vld1q_f32(ip.add(4)) };
    let i2 = unsafe { vld1q_f32(ip.add(8)) };
    let i3 = unsafe { vld1q_f32(ip.add(12)) };
    let i4 = unsafe { vld1q_f32(ip.add(16)) };
    let i5 = unsafe { vld1q_f32(ip.add(20)) };
    let i6 = unsafe { vld1q_f32(ip.add(24)) };
    let i7 = unsafe { vld1q_f32(ip.add(28)) };

    // Accumulate d * q * inp
    let mut acc = unsafe { vmulq_f32(qf_a, i0) };
    acc = unsafe { vfmaq_f32(acc, qf_b, i1) };
    acc = unsafe { vfmaq_f32(acc, qf_c, i2) };
    acc = unsafe { vfmaq_f32(acc, qf_e, i3) };
    acc = unsafe { vfmaq_f32(acc, qf_g, i4) };
    acc = unsafe { vfmaq_f32(acc, qf_h, i5) };
    acc = unsafe { vfmaq_f32(acc, qf_j, i6) };
    acc = unsafe { vfmaq_f32(acc, qf_k, i7) };

    // Accumulate input sum for m correction: m * Σ inp[i]
    let mut inp_sum = unsafe { vaddq_f32(i0, i1) };
    inp_sum = unsafe { vaddq_f32(inp_sum, vaddq_f32(i2, i3)) };
    inp_sum = unsafe { vaddq_f32(inp_sum, vaddq_f32(i4, i5)) };
    inp_sum = unsafe { vaddq_f32(inp_sum, vaddq_f32(i6, i7)) };

    let _ = d_vec;
    d * unsafe { hsum_f32x4(acc) } + m * unsafe { hsum_f32x4(inp_sum) }
}

impl QuantKernel for Q5_1Neon {
    fn dequant_block(&self, block: &[u8], output: &mut [f32]) -> QuantResult<()> {
        if block.len() < BLOCK_BYTES {
            return Err(QuantError::BufferTooSmall {
                needed: BLOCK_BYTES,
                available: block.len(),
            });
        }
        if output.len() < BLOCK_SIZE {
            return Err(QuantError::BufferTooSmall {
                needed: BLOCK_SIZE,
                available: output.len(),
            });
        }

        let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]]));
        let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]]));
        let qh = u32::from_le_bytes([block[4], block[5], block[6], block[7]]);
        let (qh_lo, qh_hi) = expand_qh(qh);
        unsafe {
            dequant_block_neon(
                block.as_ptr().add(8),
                &qh_lo,
                &qh_hi,
                d,
                m,
                &mut output[..BLOCK_SIZE],
            )
        };
        Ok(())
    }

    fn gemv(
        &self,
        quant_matrix: &QuantTensor,
        input: &[f32],
        output: &mut [f32],
    ) -> QuantResult<()> {
        let n_rows = quant_matrix.shape[0];
        let n_cols = if quant_matrix.shape.len() > 1 {
            quant_matrix.shape[1]
        } else {
            quant_matrix.n_elements() / n_rows
        };

        if input.len() < n_cols {
            return Err(QuantError::DimensionMismatch {
                expected: n_cols,
                got: input.len(),
            });
        }
        if output.len() < n_rows {
            return Err(QuantError::DimensionMismatch {
                expected: n_rows,
                got: output.len(),
            });
        }

        let blocks_per_row = n_cols.div_ceil(BLOCK_SIZE);
        let row_bytes = blocks_per_row * BLOCK_BYTES;

        for (row, out) in output.iter_mut().enumerate().take(n_rows) {
            let row_start = row * row_bytes;
            let mut sum = 0.0f32;

            for blk in 0..blocks_per_row {
                let block_offset = row_start + blk * BLOCK_BYTES;
                let block = &quant_matrix.data[block_offset..block_offset + BLOCK_BYTES];
                let d = f16_to_f32(u16::from_le_bytes([block[0], block[1]]));
                let m = f16_to_f32(u16::from_le_bytes([block[2], block[3]]));
                let qh = u32::from_le_bytes([block[4], block[5], block[6], block[7]]);
                let (qh_lo, qh_hi) = expand_qh(qh);
                let input_offset = blk * BLOCK_SIZE;
                let block_input_end = (input_offset + BLOCK_SIZE).min(n_cols);
                let block_input_len = block_input_end - input_offset;

                if block_input_len == BLOCK_SIZE {
                    sum += unsafe {
                        dot_block_neon(
                            block.as_ptr().add(8),
                            &qh_lo,
                            &qh_hi,
                            d,
                            m,
                            &input[input_offset..input_offset + BLOCK_SIZE],
                        )
                    };
                } else {
                    // Scalar tail (sequential layout, unsigned, affine)
                    let qs = &block[8..24];
                    let inp = &input[input_offset..];
                    let mut input_sum = 0.0f32;
                    for i in 0..block_input_len {
                        let byte = qs[i % 16];
                        let lo_nibble = byte & 0x0F;
                        let hi_nibble = (byte >> 4) & 0x0F;
                        let hi_bit_lo = ((qh >> i) & 1) as u8;
                        let hi_bit_hi = ((qh >> (i + 16)) & 1) as u8;
                        let q0 = (lo_nibble | (hi_bit_lo << 4)) as f32;
                        let q1 = (hi_nibble | (hi_bit_hi << 4)) as f32;
                        if i < 16 {
                            sum += d * q0 * inp[i];
                            input_sum += inp[i];
                        } else {
                            sum += d * q1 * inp[i];
                            input_sum += inp[i];
                        }
                    }
                    sum += m * input_sum;
                }
            }
            *out = sum;
        }

        Ok(())
    }

    fn gemm(
        &self,
        quant_matrix: &QuantTensor,
        input: &[f32],
        output: &mut [f32],
        m: usize,
        n: usize,
        k: usize,
    ) -> QuantResult<()> {
        for row in 0..m {
            let input_row = &input[row * k..(row + 1) * k];
            let output_row = &mut output[row * n..(row + 1) * n];
            self.gemv(quant_matrix, input_row, output_row)?;
        }
        Ok(())
    }

    fn block_size(&self) -> usize {
        BLOCK_SIZE
    }

    fn block_bytes(&self) -> usize {
        BLOCK_BYTES
    }

    fn name(&self) -> &'static str {
        "Q5_1_Neon"
    }

    fn matvec_q8_fused(
        &self,
        weights: &[u8],
        acts_q8: &[u8],
        out: &mut [f32],
        n_rows: usize,
        n_cols: usize,
    ) -> crate::error::QuantResult<()> {
        use crate::error::QuantError;

        if out.len() < n_rows {
            return Err(QuantError::DimensionMismatch {
                expected: n_rows,
                got: out.len(),
            });
        }
        let blocks_per_row = n_cols.div_ceil(BLOCK_SIZE);
        let row_bytes = blocks_per_row * BLOCK_BYTES;
        let q8_block_bytes: usize = 34;

        if weights.len() < n_rows * row_bytes {
            return Err(QuantError::BufferTooSmall {
                needed: n_rows * row_bytes,
                available: weights.len(),
            });
        }
        if acts_q8.len() < blocks_per_row * q8_block_bytes {
            return Err(QuantError::BufferTooSmall {
                needed: blocks_per_row * q8_block_bytes,
                available: acts_q8.len(),
            });
        }

        for (row, out_val) in out.iter_mut().enumerate().take(n_rows) {
            let row_start = row * row_bytes;
            let partial = unsafe {
                fused_q5_1_q8_0_row_neon(
                    &weights[row_start..row_start + row_bytes],
                    acts_q8,
                    blocks_per_row,
                    n_cols,
                )
            };
            *out_val += partial;
        }

        Ok(())
    }
}

/// Fused Q5_1 weight × Q8_0 activation dot product for one matrix row.
///
/// Q5_1 affine formula: `(d_w * q5_unsigned + m_w) * (d_a * q8_act)`.
///
/// # Safety
/// Must be called on AArch64 with NEON. All slice bounds must be pre-validated.
unsafe fn fused_q5_1_q8_0_row_neon(
    weights_row: &[u8],
    acts_q8: &[u8],
    blocks_per_row: usize,
    n_cols: usize,
) -> f32 {
    const Q8_BLOCK_BYTES: usize = 34;
    let mut acc = vdupq_n_f32(0.0f32);

    for blk in 0..blocks_per_row {
        let bo = blk * BLOCK_BYTES;
        let block = &weights_row[bo..bo + BLOCK_BYTES];
        let col_start = blk * BLOCK_SIZE;

        let d_w = f16_to_f32(u16::from_le_bytes([block[0], block[1]]));
        let m_w = f16_to_f32(u16::from_le_bytes([block[2], block[3]]));
        let qh = u32::from_le_bytes([block[4], block[5], block[6], block[7]]);
        let (qh_lo, qh_hi) = expand_qh(qh);

        // Decode Q8_0 activation block
        let ab = blk * Q8_BLOCK_BYTES;
        let a_block = &acts_q8[ab..ab + Q8_BLOCK_BYTES];
        let d_a = f16_to_f32(u16::from_le_bytes([a_block[0], a_block[1]]));

        let col_end = (col_start + BLOCK_SIZE).min(n_cols);
        let avail = col_end - col_start;

        if avail == BLOCK_SIZE {
            // Full block: NEON decode
            let qs_ptr = block.as_ptr().add(8);
            let raw = vld1q_u8(qs_ptr);
            let mask = vdupq_n_u8(0x0F);
            let lo_nib = vandq_u8(raw, mask);
            let hi_nib = vshrq_n_u8::<4>(raw);

            let vqh_lo = vld1q_u8(qh_lo.as_ptr());
            let vqh_hi = vld1q_u8(qh_hi.as_ptr());
            let shift4 = vdupq_n_u8(4);
            let qh_lo_sh = vshlq_u8(vqh_lo, vreinterpretq_s8_u8(shift4));
            let qh_hi_sh = vshlq_u8(vqh_hi, vreinterpretq_s8_u8(shift4));
            let q5_lo = vorrq_u8(lo_nib, qh_lo_sh);
            let q5_hi = vorrq_u8(hi_nib, qh_hi_sh);

            // Widen to u32 → f32 (unsigned, no bias)
            let q5_lo_u16_low = vmovl_u8(vget_low_u8(q5_lo));
            let q5_lo_u16_high = vmovl_u8(vget_high_u8(q5_lo));
            let q5_hi_u16_low = vmovl_u8(vget_low_u8(q5_hi));
            let q5_hi_u16_high = vmovl_u8(vget_high_u8(q5_hi));

            let d_w_vec = vdupq_n_f32(d_w);
            let m_w_vec = vdupq_n_f32(m_w);

            // Affine weight: d_w * q5 + m_w
            let wfa = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_low))),
                d_w_vec,
            );
            let wfb = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_low)),
                d_w_vec,
            );
            let wfc = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_lo_u16_high))),
                d_w_vec,
            );
            let wfe = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_high_u16(q5_lo_u16_high)),
                d_w_vec,
            );
            let wfg = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_low))),
                d_w_vec,
            );
            let wfh = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_low)),
                d_w_vec,
            );
            let wfj = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_u16(vget_low_u16(q5_hi_u16_high))),
                d_w_vec,
            );
            let wfk = vfmaq_f32(
                m_w_vec,
                vcvtq_f32_u32(vmovl_high_u16(q5_hi_u16_high)),
                d_w_vec,
            );

            // Load and scale Q8_0 activation i8 values by d_a
            let q8_ptr = a_block.as_ptr().add(2) as *const i8;
            let aq0 = vld1q_s8(q8_ptr);
            let aq1 = vld1q_s8(q8_ptr.add(16));

            let aq0_lo = vmovl_s8(vget_low_s8(aq0));
            let aq0_hi = vmovl_s8(vget_high_s8(aq0));
            let aq1_lo = vmovl_s8(vget_low_s8(aq1));
            let aq1_hi = vmovl_s8(vget_high_s8(aq1));

            let d_a_vec = vdupq_n_f32(d_a);
            let aa0 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(aq0_lo))), d_a_vec);
            let aa1 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(aq0_lo)), d_a_vec);
            let aa2 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(aq0_hi))), d_a_vec);
            let aa3 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(aq0_hi)), d_a_vec);
            let aa4 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(aq1_lo))), d_a_vec);
            let aa5 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(aq1_lo)), d_a_vec);
            let aa6 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(aq1_hi))), d_a_vec);
            let aa7 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(aq1_hi)), d_a_vec);

            // FMA: acc += (d_w*q5 + m_w) * (d_a * q8_act)
            acc = vfmaq_f32(acc, wfa, aa0);
            acc = vfmaq_f32(acc, wfb, aa1);
            acc = vfmaq_f32(acc, wfc, aa2);
            acc = vfmaq_f32(acc, wfe, aa3);
            acc = vfmaq_f32(acc, wfg, aa4);
            acc = vfmaq_f32(acc, wfh, aa5);
            acc = vfmaq_f32(acc, wfj, aa6);
            acc = vfmaq_f32(acc, wfk, aa7);
        } else {
            // Scalar tail for partial blocks
            let qs = &block[8..24];
            let q8_vals = &a_block[2..];
            let mut block_dot = 0.0f32;
            for i in 0..avail {
                let byte = qs[i % 16];
                let nibble = if i < 16 {
                    byte & 0x0F
                } else {
                    (byte >> 4) & 0x0F
                };
                let hi_bit = ((qh >> i) & 1) as u8;
                let q5 = (nibble | (hi_bit << 4)) as f32;
                let w = d_w * q5 + m_w;
                let a_val = (q8_vals[i] as i8) as f32 * d_a;
                block_dot += w * a_val;
            }
            let lane0 = vgetq_lane_f32::<0>(acc) + block_dot;
            acc = vdupq_n_f32(0.0f32);
            acc = vsetq_lane_f32::<0>(lane0, acc);
        }
    }

    hsum_f32x4(acc)
}

#[cfg(all(test, feature = "simd-neon", target_arch = "aarch64"))]
mod tests {
    use super::*;
    use crate::reference::q5_1::Q5_1Ref;
    use crate::traits::QuantKernel;
    use crate::types::QuantTensor;

    fn make_block(d: f32, m: f32, qh: u32, qs: &[u8; 16]) -> Vec<u8> {
        let mut block = Vec::with_capacity(BLOCK_BYTES);
        block.extend_from_slice(&half::f16::from_f32(d).to_bits().to_le_bytes());
        block.extend_from_slice(&half::f16::from_f32(m).to_bits().to_le_bytes());
        block.extend_from_slice(&qh.to_le_bytes());
        block.extend_from_slice(qs);
        block
    }

    #[test]
    fn test_dequant_zeros() {
        let block = make_block(0.0, 0.0, 0, &[0; 16]);
        let mut out = vec![0.0f32; 32];
        Q5_1Neon.dequant_block(&block, &mut out).expect("dequant");
        for &v in &out {
            assert!(v.abs() < 1e-6, "expected 0, got {v}");
        }
    }

    #[test]
    fn test_dequant_min_only() {
        // d=0, m=3.0, all q=0 → w = 3.0
        let block = make_block(0.0, 3.0, 0, &[0; 16]);
        let mut out = vec![0.0f32; 32];
        Q5_1Neon.dequant_block(&block, &mut out).expect("dequant");
        for &v in &out {
            assert!((v - 3.0).abs() < 1e-4, "expected 3.0, got {v}");
        }
    }

    #[test]
    fn test_dequant_max() {
        // d=1.0, m=0, qh=0xFFFFFFFF, qs=0xFF → q5=31 → w=31
        let block = make_block(1.0, 0.0, 0xFFFF_FFFF, &[0xFF; 16]);
        let mut out_neon = vec![0.0f32; 32];
        let mut out_ref = vec![0.0f32; 32];
        Q5_1Neon.dequant_block(&block, &mut out_neon).expect("neon");
        Q5_1Ref.dequant_block(&block, &mut out_ref).expect("ref");
        let max_err = out_neon
            .iter()
            .zip(out_ref.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_err < 1e-4, "max err {max_err}");
    }

    #[test]
    fn test_dequant_matches_reference() {
        let qh: u32 = 0x5A5A_5A5A;
        let mut qs = [0u8; 16];
        for (i, v) in qs.iter_mut().enumerate() {
            *v = ((i * 9 + 3) & 0xFF) as u8;
        }
        let block = make_block(0.5, 0.25, qh, &qs);
        let mut out_neon = vec![0.0f32; 32];
        let mut out_ref = vec![0.0f32; 32];
        Q5_1Neon.dequant_block(&block, &mut out_neon).expect("neon");
        Q5_1Ref.dequant_block(&block, &mut out_ref).expect("ref");
        let max_err = out_neon
            .iter()
            .zip(out_ref.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(
            max_err < 1e-4,
            "dequant max error {max_err}; neon[0]={} ref[0]={}",
            out_neon[0],
            out_ref[0]
        );
    }

    #[test]
    fn test_gemv_matches_reference() {
        let qh: u32 = 0x5A5A_5A5A;
        let mut qs = [0u8; 16];
        for (i, v) in qs.iter_mut().enumerate() {
            *v = ((i * 9 + 3) & 0xFF) as u8;
        }
        let block = make_block(0.5, 0.25, qh, &qs);
        let n_cols = BLOCK_SIZE;
        let tensor_neon = QuantTensor::new(
            block.clone(),
            vec![1, n_cols],
            oxillama_gguf::GgufTensorType::Q5_1,
        );
        let tensor_ref =
            QuantTensor::new(block, vec![1, n_cols], oxillama_gguf::GgufTensorType::Q5_1);
        let input: Vec<f32> = (0..n_cols).map(|i| (i as f32) * 0.1 - 1.6).collect();
        let mut out_neon = vec![0.0f32; 1];
        let mut out_ref = vec![0.0f32; 1];
        Q5_1Neon
            .gemv(&tensor_neon, &input, &mut out_neon)
            .expect("neon");
        Q5_1Ref
            .gemv(&tensor_ref, &input, &mut out_ref)
            .expect("ref");
        let err = (out_neon[0] - out_ref[0]).abs();
        assert!(
            err < 1e-3,
            "gemv: neon={} ref={} err={}",
            out_neon[0],
            out_ref[0],
            err
        );
    }

    #[test]
    fn fused_q5_1_neon_matches_reference() {
        let qh: u32 = 0x5A5A_5A5A;
        let mut qs_w = [0u8; 16];
        for (i, v) in qs_w.iter_mut().enumerate() {
            *v = ((i * 9 + 3) & 0xFF) as u8;
        }
        let weight_block = make_block(0.5, 0.25, qh, &qs_w);

        // Build Q8_0 activation block (34 bytes)
        let d_a = 0.25f32;
        let mut acts_raw = [0i8; 32];
        for (i, v) in acts_raw.iter_mut().enumerate() {
            *v = ((i as i16 * 5 - 40).clamp(-128, 127)) as i8;
        }
        let mut acts_block = Vec::with_capacity(34);
        acts_block.extend_from_slice(&half::f16::from_f32(d_a).to_bits().to_le_bytes());
        for &v in &acts_raw {
            acts_block.push(v as u8);
        }

        // Reference: dequant weight + dot with scaled activations
        let mut w_dequant = vec![0.0f32; BLOCK_SIZE];
        Q5_1Ref
            .dequant_block(&weight_block, &mut w_dequant)
            .expect("ref dequant");
        let acts_f32: Vec<f32> = acts_raw.iter().map(|&v| v as f32 * d_a).collect();
        let expected: f32 = w_dequant
            .iter()
            .zip(acts_f32.iter())
            .map(|(w, a)| w * a)
            .sum();

        // NEON fused
        let mut out_neon = vec![0.0f32; 1];
        Q5_1Neon
            .matvec_q8_fused(&weight_block, &acts_block, &mut out_neon, 1, BLOCK_SIZE)
            .expect("neon fused");

        let err = (out_neon[0] - expected).abs();
        assert!(
            err < 0.1,
            "fused_q5_1_neon: got={} expected={} err={}",
            out_neon[0],
            expected,
            err
        );

        // Verify against reference Q5_1Ref::matvec_q8_fused
        let mut out_ref = vec![0.0f32; 1];
        Q5_1Ref
            .matvec_q8_fused(&weight_block, &acts_block, &mut out_ref, 1, BLOCK_SIZE)
            .expect("ref fused");
        let ref_err = (out_neon[0] - out_ref[0]).abs();
        assert!(
            ref_err < 0.1,
            "fused: neon={} ref={} err={}",
            out_neon[0],
            out_ref[0],
            ref_err
        );
    }
}