oxillama-quant 0.1.3

//! Q3_K NEON (AArch64) SIMD kernel.
//!
//! Q3_K super-block format (110 bytes per 256 weights):
//! - bytes[0..32]:   hmask  — 1 bit per weight; if set → subtract 0; if clear → subtract 4
//! - bytes[32..96]:  qs     — lower 2 bits of each 3-bit quant (4 per byte via shifts 0,2,4,6)
//! - bytes[96..108]: scales — 16 × 6-bit signed sub-block scales, packed into 12 bytes
//! - bytes[108..110]: FP16 super-block scale `d`
//!
//! Symmetric format (no minimum offset). Weight formula:
//!   `w = d * scale_i * (q_lo | (hmask_bit ? 0 : -4))`
//! where q_lo ∈ [0..3], giving effective range [-4..3].

#![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 Q3_K super-block.
pub const BLOCK_SIZE: usize = 256;
/// Bytes per Q3_K super-block.
pub const BLOCK_BYTES: usize = 110;

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

#[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) }
}

/// Decode 16 signed 6-bit sub-block scales from the 12-byte packed representation.
///
/// Mirrors `reference::q3_k::decode_scales` exactly.
/// Returns unsigned 6-bit values minus 32 → signed range [-32..31].
pub fn unpack_scales(scales_raw: &[u8; 12]) -> [i8; 16] {
    let mut sc = [0u32; 16];

    for j in 0..4 {
        sc[j] = (scales_raw[j] & 0x3F) as u32;
    }
    for j in 0..4 {
        sc[4 + j] = (scales_raw[4 + j] & 0x3F) as u32;
    }
    for j in 0..4 {
        let lo = (scales_raw[8 + j] & 0x0F) as u32;
        let hi = ((scales_raw[j] >> 6) & 0x03) as u32;
        sc[8 + j] = lo | (hi << 4);
    }
    for j in 0..4 {
        let lo = ((scales_raw[8 + j] >> 4) & 0x0F) as u32;
        let hi = ((scales_raw[4 + j] >> 6) & 0x03) as u32;
        sc[12 + j] = lo | (hi << 4);
    }

    let mut result = [0i8; 16];
    for i in 0..16 {
        result[i] = (sc[i] as i32 - 32) as i8;
    }
    result
}

/// Extract 2-bit fields from 16 bytes, selecting field by shift (0, 2, 4, 6).
#[inline(always)]
unsafe fn extract_2bit(raw: uint8x16_t, shift: u32) -> uint8x16_t {
    let mask = unsafe { vdupq_n_u8(0x03) };
    let shifted = match shift {
        0 => raw,
        2 => unsafe { vshrq_n_u8::<2>(raw) },
        4 => unsafe { vshrq_n_u8::<4>(raw) },
        _ => unsafe { vshrq_n_u8::<6>(raw) },
    };
    unsafe { vandq_u8(shifted, mask) }
}

/// Compute per-element correction from hmask bytes: 0 if bit set, 4 if bit clear.
///
/// `m_bit` is the single-bit selector (a power of 2 ≤ 128).
/// Returns a u8x16 with values in {0, 4}.
#[inline(always)]
unsafe fn hmask_correction(hmask_chunk: uint8x16_t, m_bit: u8) -> uint8x16_t {
    let m_vec = unsafe { vdupq_n_u8(m_bit) };
    let four = unsafe { vdupq_n_u8(4) };
    let zero = unsafe { vdupq_n_u8(0) };
    // Mask the bit: if (hmask & m_bit) != 0 → bit present
    let masked = unsafe { vandq_u8(hmask_chunk, m_vec) };
    // If masked == m_bit → correction 0, else correction 4
    // Use vcgtq_u8: masked > 0 → all-ones lane, else all-zeros.
    // Then select: 0 where bit set, 4 where bit clear.
    let is_set = unsafe { vcgtq_u8(masked, vdupq_n_u8(0)) };
    unsafe { vbslq_u8(is_set, zero, four) }
}

/// Dequantize 16 weights from a Q3_K sub-block using NEON.
///
/// `qs_raw` = 16-byte chunk of qs at appropriate base.
/// `hmask_chunk` = corresponding 16-byte slice of hmask.
/// Returns nothing; writes to `out_ptr`.
///
/// # Safety
/// `out_ptr` must have at least 16 valid f32 slots.
#[inline]
unsafe fn dequant_16_weights(
    qs_raw: uint8x16_t,
    hmask_chunk: uint8x16_t,
    shift: u32,
    m_bit: u8,
    dl: f32,
    out_ptr: *mut f32,
) {
    let q_bytes = unsafe { extract_2bit(qs_raw, shift) };
    let corr = unsafe { hmask_correction(hmask_chunk, m_bit) };

    let dl_vec = unsafe { vdupq_n_f32(dl) };

    // q_signed = q_lo - correction (4 when hmask bit clear → subtracts 4 for signed range)
    // Widen u8 to i16, subtract correction (u8→i16 too, then sub)
    let q_lo_u16 = unsafe { vmovl_u8(vget_low_u8(q_bytes)) };
    let q_hi_u16 = unsafe { vmovl_u8(vget_high_u8(q_bytes)) };
    let c_lo_u16 = unsafe { vmovl_u8(vget_low_u8(corr)) };
    let c_hi_u16 = unsafe { vmovl_u8(vget_high_u8(corr)) };

    let qs_lo = unsafe { vreinterpretq_s16_u16(vsubq_u16(q_lo_u16, c_lo_u16)) };
    let qs_hi = unsafe { vreinterpretq_s16_u16(vsubq_u16(q_hi_u16, c_hi_u16)) };

    // Widen to s32 → f32, scale by dl
    let q0 = unsafe { vcvtq_f32_s32(vmovl_s16(vget_low_s16(qs_lo))) };
    let q1 = unsafe { vcvtq_f32_s32(vmovl_high_s16(qs_lo)) };
    let q2 = unsafe { vcvtq_f32_s32(vmovl_s16(vget_low_s16(qs_hi))) };
    let q3 = unsafe { vcvtq_f32_s32(vmovl_high_s16(qs_hi)) };

    let w0 = unsafe { vmulq_f32(dl_vec, q0) };
    let w1 = unsafe { vmulq_f32(dl_vec, q1) };
    let w2 = unsafe { vmulq_f32(dl_vec, q2) };
    let w3 = unsafe { vmulq_f32(dl_vec, q3) };

    unsafe { vst1q_f32(out_ptr, w0) };
    unsafe { vst1q_f32(out_ptr.add(4), w1) };
    unsafe { vst1q_f32(out_ptr.add(8), w2) };
    unsafe { vst1q_f32(out_ptr.add(12), w3) };
}

/// Dot product: 16 Q3_K weights × 16 f32 inputs.
///
/// # Safety
/// `inp_ptr` must point to 16 valid f32 values.
#[inline]
unsafe fn dot_16_weights(
    qs_raw: uint8x16_t,
    hmask_chunk: uint8x16_t,
    shift: u32,
    m_bit: u8,
    dl: f32,
    inp_ptr: *const f32,
) -> f32 {
    let q_bytes = unsafe { extract_2bit(qs_raw, shift) };
    let corr = unsafe { hmask_correction(hmask_chunk, m_bit) };

    let dl_vec = unsafe { vdupq_n_f32(dl) };

    let q_lo_u16 = unsafe { vmovl_u8(vget_low_u8(q_bytes)) };
    let q_hi_u16 = unsafe { vmovl_u8(vget_high_u8(q_bytes)) };
    let c_lo_u16 = unsafe { vmovl_u8(vget_low_u8(corr)) };
    let c_hi_u16 = unsafe { vmovl_u8(vget_high_u8(corr)) };

    let qs_lo = unsafe { vreinterpretq_s16_u16(vsubq_u16(q_lo_u16, c_lo_u16)) };
    let qs_hi = unsafe { vreinterpretq_s16_u16(vsubq_u16(q_hi_u16, c_hi_u16)) };

    let q0 = unsafe { vcvtq_f32_s32(vmovl_s16(vget_low_s16(qs_lo))) };
    let q1 = unsafe { vcvtq_f32_s32(vmovl_high_s16(qs_lo)) };
    let q2 = unsafe { vcvtq_f32_s32(vmovl_s16(vget_low_s16(qs_hi))) };
    let q3 = unsafe { vcvtq_f32_s32(vmovl_high_s16(qs_hi)) };

    let w0 = unsafe { vmulq_f32(dl_vec, q0) };
    let w1 = unsafe { vmulq_f32(dl_vec, q1) };
    let w2 = unsafe { vmulq_f32(dl_vec, q2) };
    let w3 = unsafe { vmulq_f32(dl_vec, q3) };

    let i0 = unsafe { vld1q_f32(inp_ptr) };
    let i1 = unsafe { vld1q_f32(inp_ptr.add(4)) };
    let i2 = unsafe { vld1q_f32(inp_ptr.add(8)) };
    let i3 = unsafe { vld1q_f32(inp_ptr.add(12)) };

    let mut acc = unsafe { vmulq_f32(w0, i0) };
    acc = unsafe { vfmaq_f32(acc, w1, i1) };
    acc = unsafe { vfmaq_f32(acc, w2, i2) };
    acc = unsafe { vfmaq_f32(acc, w3, i3) };

    unsafe { hsum_f32x4(acc) }
}

impl QuantKernel for Q3_KNeon {
    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 hmask = &block[0..32];
        let qs = &block[32..96];
        let scales_raw: &[u8; 12] = block[96..108].try_into().unwrap_or(&[0u8; 12]);
        let d = f16_to_f32(u16::from_le_bytes([block[108], block[109]]));
        let sc = unpack_scales(scales_raw);

        let hmask_lo = unsafe { vld1q_u8(hmask.as_ptr()) };
        let hmask_hi = unsafe { vld1q_u8(hmask.as_ptr().add(16)) };

        let mut is = 0usize;
        let mut out_off = 0usize;

        for group in 0..2usize {
            let qs_base = group * 32;
            let raw_a = unsafe { vld1q_u8(qs.as_ptr().add(qs_base)) };
            let raw_b = unsafe { vld1q_u8(qs.as_ptr().add(qs_base + 16)) };

            for shift_idx in 0..4u32 {
                let shift = shift_idx * 2;
                let bit_pos = (group as u32) * 4 + shift_idx;
                let m_bit: u8 = 1u8 << bit_pos;

                let dl_a = d * sc[is] as f32;
                is += 1;
                unsafe {
                    dequant_16_weights(
                        raw_a,
                        hmask_lo,
                        shift,
                        m_bit,
                        dl_a,
                        output.as_mut_ptr().add(out_off),
                    )
                };
                out_off += 16;

                let dl_b = d * sc[is] as f32;
                is += 1;
                unsafe {
                    dequant_16_weights(
                        raw_b,
                        hmask_hi,
                        shift,
                        m_bit,
                        dl_b,
                        output.as_mut_ptr().add(out_off),
                    )
                };
                out_off += 16;
            }
        }

        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 hmask = &block[0..32];
                let qs = &block[32..96];
                let scales_raw: &[u8; 12] = block[96..108].try_into().unwrap_or(&[0u8; 12]);
                let d = f16_to_f32(u16::from_le_bytes([block[108], block[109]]));
                let sc = unpack_scales(scales_raw);
                let input_offset = blk * BLOCK_SIZE;
                let cols_in_block = (n_cols - input_offset).min(BLOCK_SIZE);

                let hmask_lo = unsafe { vld1q_u8(hmask.as_ptr()) };
                let hmask_hi = unsafe { vld1q_u8(hmask.as_ptr().add(16)) };

                if cols_in_block == BLOCK_SIZE {
                    let mut is = 0usize;
                    let mut w_off = input_offset;

                    for group in 0..2usize {
                        let qs_base = group * 32;
                        let raw_a = unsafe { vld1q_u8(qs.as_ptr().add(qs_base)) };
                        let raw_b = unsafe { vld1q_u8(qs.as_ptr().add(qs_base + 16)) };

                        for shift_idx in 0..4u32 {
                            let shift = shift_idx * 2;
                            let bit_pos = (group as u32) * 4 + shift_idx;
                            let m_bit: u8 = 1u8 << bit_pos;

                            let dl_a = d * sc[is] as f32;
                            is += 1;
                            sum += unsafe {
                                dot_16_weights(
                                    raw_a,
                                    hmask_lo,
                                    shift,
                                    m_bit,
                                    dl_a,
                                    input.as_ptr().add(w_off),
                                )
                            };
                            w_off += 16;

                            let dl_b = d * sc[is] as f32;
                            is += 1;
                            sum += unsafe {
                                dot_16_weights(
                                    raw_b,
                                    hmask_hi,
                                    shift,
                                    m_bit,
                                    dl_b,
                                    input.as_ptr().add(w_off),
                                )
                            };
                            w_off += 16;
                        }
                    }
                } else {
                    // Scalar tail for partial blocks
                    let inp = &input[input_offset..];
                    let mut is = 0usize;
                    let mut in_off = 0usize;
                    let mut m_bit: u8 = 1;

                    for group in 0..2 {
                        let qs_base = group * 32;
                        for shift in (0..8usize).step_by(2) {
                            for n in 0..2 {
                                let dl = d * sc[is] as f32;
                                is += 1;
                                for l in 0..16 {
                                    if in_off + l < cols_in_block {
                                        let qs_idx = qs_base + n * 16 + l;
                                        let q_lo = ((qs[qs_idx] >> shift) & 3) as i32;
                                        let sub =
                                            if hmask[n * 16 + l] & m_bit != 0 { 0 } else { 4 };
                                        sum += dl * (q_lo - sub) as f32 * inp[in_off + l];
                                    }
                                }
                                in_off += 16;
                            }
                            m_bit = m_bit.wrapping_shl(1);
                        }
                    }
                }
            }
            *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(())
    }

    /// Fused Q3_K weight × Q8_0 activation GEMV using NEON integer dot products.
    ///
    /// Each Q3_K super-block (256 weights, 110 bytes) maps to 8 Q8_0 activation blocks
    /// (32 weights each, 34 bytes each).  The 256 weights are organized as 16 sub-blocks
    /// of 16 weights.  Two consecutive sub-blocks share one Q8_0 block of 32 activations.
    ///
    /// Fused formula per sub-block (symmetric format, no minimum):
    ///   contribution = dl * Σ(q3_i * q8_i) * d_a
    /// where q3_i = q_lo_i - correction_i (correction = 4 if hmask bit clear, else 0).
    /// No intermediate f32 scratch buffer is allocated.
    fn matvec_q8_fused(
        &self,
        weights: &[u8],
        acts_q8: &[u8],
        out: &mut [f32],
        n_rows: usize,
        n_cols: usize,
    ) -> QuantResult<()> {
        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_blocks_per_row = blocks_per_row * 8;
        let acts_needed = q8_blocks_per_row * Q8_0_BLOCK_BYTES;

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

        for (row, out_val) in out.iter_mut().enumerate().take(n_rows) {
            let row_start = row * row_bytes;
            // SAFETY: bounds checked above; AArch64 always has NEON.
            let row_sum = unsafe {
                fused_q3k_q8_0_row_neon(
                    &weights[row_start..row_start + row_bytes],
                    acts_q8,
                    blocks_per_row,
                    n_cols,
                )
            };
            *out_val += row_sum;
        }

        Ok(())
    }

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

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

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

/// Q8_0 block bytes for fused GEMV.
const Q8_0_BLOCK_BYTES: usize = 34;

/// Fused Q3_K weight × Q8_0 activation dot product for one row using NEON integer arithmetic.
///
/// Q3_K layout (per block): hmask[0..32], qs[32..96], scales[96..108], d[108..110].
/// 256 weights = 16 sub-blocks × 16 weights.  Two consecutive sub-blocks map to one Q8_0 block.
/// Sub-block `sb` (0..16) → Q8_0 block `blk*8 + sb/2`.
///
/// For each sub-block with signed scale `dl = d * sc[sb]`:
///   contribution = dl * Σ(q3_i * q8_i) * d_a
/// where q3_i = q_lo_i - correction_i (correction = 4 if hmask bit clear).
///
/// # Safety
/// - `row_data.len() == blocks_per_row * BLOCK_BYTES`
/// - `acts_q8.len() >= blocks_per_row * 8 * Q8_0_BLOCK_BYTES`
/// - Must be called on AArch64 with NEON
unsafe fn fused_q3k_q8_0_row_neon(
    row_data: &[u8],
    acts_q8: &[u8],
    blocks_per_row: usize,
    n_cols: usize,
) -> f32 {
    let mut row_sum = 0.0f32;

    for blk in 0..blocks_per_row {
        let bo = blk * BLOCK_BYTES;
        // SAFETY: row_data.len() == blocks_per_row * BLOCK_BYTES; blk < blocks_per_row.
        let block = &row_data[bo..bo + BLOCK_BYTES];

        let hmask = &block[0..32];
        let qs = &block[32..96];
        let scales_raw: &[u8; 12] = block[96..108].try_into().unwrap_or(&[0u8; 12]);
        let d = f16_to_f32(u16::from_le_bytes([block[108], block[109]]));
        let sc = unpack_scales(scales_raw);

        let input_offset = blk * BLOCK_SIZE;
        let cols_in_block = (n_cols - input_offset).min(BLOCK_SIZE);

        let hmask_lo = unsafe { vld1q_u8(hmask.as_ptr()) };
        let hmask_hi = unsafe { vld1q_u8(hmask.as_ptr().add(16)) };

        if cols_in_block < BLOCK_SIZE {
            // Scalar tail path — partial block.
            let mut is = 0usize;
            let mut col_off = 0usize;
            let mut m_bit: u8 = 1;

            for group in 0..2usize {
                let qs_base = group * 32;
                for _shift_idx in 0..4u32 {
                    let shift = _shift_idx * 2;
                    // Sub-block A (hmask[0..16])
                    {
                        let dl = d * sc[is] as f32;
                        is += 1;
                        let q8_blk_idx = blk * 8 + col_off / 32;
                        let q8_lane_base = col_off % 32;
                        let a_start = q8_blk_idx * Q8_0_BLOCK_BYTES;
                        let a_block = &acts_q8[a_start..a_start + Q8_0_BLOCK_BYTES];
                        let d_a = f16_to_f32(u16::from_le_bytes([a_block[0], a_block[1]]));
                        let q8_vals = &a_block[2..];
                        let mut dot = 0i32;
                        for l in 0..16 {
                            if col_off + l < cols_in_block {
                                let q_lo = ((qs[qs_base + l] >> shift) & 3) as i32;
                                let sub = if hmask[l] & m_bit != 0 { 0 } else { 4 };
                                let q3 = q_lo - sub;
                                let q_a = q8_vals[q8_lane_base + l] as i8 as i32;
                                dot += q3 * q_a;
                            }
                        }
                        row_sum += dl * dot as f32 * d_a;
                        col_off += 16;
                    }
                    // Sub-block B (hmask[16..32])
                    {
                        let dl = d * sc[is] as f32;
                        is += 1;
                        let q8_blk_idx = blk * 8 + col_off / 32;
                        let q8_lane_base = col_off % 32;
                        let a_start = q8_blk_idx * Q8_0_BLOCK_BYTES;
                        let a_block = &acts_q8[a_start..a_start + Q8_0_BLOCK_BYTES];
                        let d_a = f16_to_f32(u16::from_le_bytes([a_block[0], a_block[1]]));
                        let q8_vals = &a_block[2..];
                        let mut dot = 0i32;
                        for l in 0..16 {
                            if col_off + l < cols_in_block {
                                let q_lo = ((qs[qs_base + 16 + l] >> shift) & 3) as i32;
                                let sub = if hmask[16 + l] & m_bit != 0 { 0 } else { 4 };
                                let q3 = q_lo - sub;
                                let q_a = q8_vals[q8_lane_base + l] as i8 as i32;
                                dot += q3 * q_a;
                            }
                        }
                        row_sum += dl * dot as f32 * d_a;
                        col_off += 16;
                    }
                    m_bit = m_bit.wrapping_shl(1);
                }
            }
        } else {
            // Fast path: full 256-weight block.
            // Process 16 sub-blocks of 16 weights, two per Q8_0 block.
            let mut is = 0usize;
            let mut col_off = 0usize;

            for group in 0..2usize {
                let qs_base = group * 32;
                // Pre-load 32 qs bytes (two 16-byte halves).
                // SAFETY: qs_base + 32 <= 64; qs.len() == 64.
                let raw_a = unsafe { vld1q_u8(qs.as_ptr().add(qs_base)) };
                let raw_b = unsafe { vld1q_u8(qs.as_ptr().add(qs_base + 16)) };

                for shift_idx in 0..4u32 {
                    let shift = shift_idx * 2;
                    let bit_pos = (group as u32) * 4 + shift_idx;
                    // SAFETY: bit_pos 0..7.
                    let m_bit: u8 = 1u8 << bit_pos;

                    // SAFETY: extract_2bit valid on AArch64.
                    let q2_a = unsafe { extract_2bit(raw_a, shift) }; // u8x16, values 0..3
                    let q2_b = unsafe { extract_2bit(raw_b, shift) }; // u8x16, values 0..3
                                                                      // SAFETY: hmask_correction valid on AArch64.
                    let corr_a = unsafe { hmask_correction(hmask_lo, m_bit) }; // u8x16, {0,4}
                    let corr_b = unsafe { hmask_correction(hmask_hi, m_bit) }; // u8x16, {0,4}

                    // Sub-block A: lanes 0..16 of Q8_0 block at col_off/32
                    {
                        let dl = d * sc[is] as f32;
                        is += 1;
                        let q8_blk_idx = blk * 8 + col_off / 32;
                        let q8_lane_base = col_off % 32;
                        let a_start = q8_blk_idx * Q8_0_BLOCK_BYTES;
                        // SAFETY: acts_q8.len() >= blocks_per_row*8*Q8_0_BLOCK_BYTES.
                        let a_block = &acts_q8[a_start..a_start + Q8_0_BLOCK_BYTES];
                        let d_a = f16_to_f32(u16::from_le_bytes([a_block[0], a_block[1]]));

                        // Load 16 Q8_0 i8 activations at lane base.
                        // SAFETY: a_block has 34 bytes; 2 + q8_lane_base + 16 <= 34.
                        let qa_ptr = a_block.as_ptr().add(2 + q8_lane_base) as *const i8;
                        let qa_i8 = unsafe { vld1q_s8(qa_ptr) };

                        // q3 = q2 - correction (signed, range -4..3)
                        // Widen u8 to i16, subtract correction (also widened), then i16 × i16 → i32.
                        let q2_lo_u16 = unsafe { vmovl_u8(vget_low_u8(q2_a)) };
                        let q2_hi_u16 = unsafe { vmovl_u8(vget_high_u8(q2_a)) };
                        let c_lo_u16 = unsafe { vmovl_u8(vget_low_u8(corr_a)) };
                        let c_hi_u16 = unsafe { vmovl_u8(vget_high_u8(corr_a)) };

                        // q3 as signed i16 = (u16 q2) - (u16 correction), safe since both ≤ 4
                        let q3_lo_i16 =
                            unsafe { vreinterpretq_s16_u16(vsubq_u16(q2_lo_u16, c_lo_u16)) };
                        let q3_hi_i16 =
                            unsafe { vreinterpretq_s16_u16(vsubq_u16(q2_hi_u16, c_hi_u16)) };

                        let qa_lo_i16 = unsafe { vmovl_s8(vget_low_s8(qa_i8)) };
                        let qa_hi_i16 = unsafe { vmovl_s8(vget_high_s8(qa_i8)) };

                        // dot = Σ(q3 * q8) as i32
                        let p0 =
                            unsafe { vmull_s16(vget_low_s16(q3_lo_i16), vget_low_s16(qa_lo_i16)) };
                        let p1 = unsafe {
                            vmlal_s16(p0, vget_high_s16(q3_lo_i16), vget_high_s16(qa_lo_i16))
                        };
                        let p2 = unsafe {
                            vmlal_s16(p1, vget_low_s16(q3_hi_i16), vget_low_s16(qa_hi_i16))
                        };
                        let p3 = unsafe {
                            vmlal_s16(p2, vget_high_s16(q3_hi_i16), vget_high_s16(qa_hi_i16))
                        };
                        let dot = unsafe { vaddvq_s32(p3) };

                        row_sum += dl * dot as f32 * d_a;
                        col_off += 16;
                    }

                    // Sub-block B: lanes 16..32 of Q8_0 block at col_off/32
                    {
                        let dl = d * sc[is] as f32;
                        is += 1;
                        let q8_blk_idx = blk * 8 + col_off / 32;
                        let q8_lane_base = col_off % 32;
                        let a_start = q8_blk_idx * Q8_0_BLOCK_BYTES;
                        // SAFETY: acts_q8.len() >= blocks_per_row*8*Q8_0_BLOCK_BYTES.
                        let a_block = &acts_q8[a_start..a_start + Q8_0_BLOCK_BYTES];
                        let d_a = f16_to_f32(u16::from_le_bytes([a_block[0], a_block[1]]));

                        let qa_ptr = a_block.as_ptr().add(2 + q8_lane_base) as *const i8;
                        let qa_i8 = unsafe { vld1q_s8(qa_ptr) };

                        let q2_lo_u16 = unsafe { vmovl_u8(vget_low_u8(q2_b)) };
                        let q2_hi_u16 = unsafe { vmovl_u8(vget_high_u8(q2_b)) };
                        let c_lo_u16 = unsafe { vmovl_u8(vget_low_u8(corr_b)) };
                        let c_hi_u16 = unsafe { vmovl_u8(vget_high_u8(corr_b)) };

                        let q3_lo_i16 =
                            unsafe { vreinterpretq_s16_u16(vsubq_u16(q2_lo_u16, c_lo_u16)) };
                        let q3_hi_i16 =
                            unsafe { vreinterpretq_s16_u16(vsubq_u16(q2_hi_u16, c_hi_u16)) };

                        let qa_lo_i16 = unsafe { vmovl_s8(vget_low_s8(qa_i8)) };
                        let qa_hi_i16 = unsafe { vmovl_s8(vget_high_s8(qa_i8)) };

                        let p0 =
                            unsafe { vmull_s16(vget_low_s16(q3_lo_i16), vget_low_s16(qa_lo_i16)) };
                        let p1 = unsafe {
                            vmlal_s16(p0, vget_high_s16(q3_lo_i16), vget_high_s16(qa_lo_i16))
                        };
                        let p2 = unsafe {
                            vmlal_s16(p1, vget_low_s16(q3_hi_i16), vget_low_s16(qa_hi_i16))
                        };
                        let p3 = unsafe {
                            vmlal_s16(p2, vget_high_s16(q3_hi_i16), vget_high_s16(qa_hi_i16))
                        };
                        let dot = unsafe { vaddvq_s32(p3) };

                        row_sum += dl * dot as f32 * d_a;
                        col_off += 16;
                    }
                }
            }
        }
    }

    row_sum
}

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

    fn make_block(d: f32, scales: &[u8; 12], hmask: &[u8; 32], qs: &[u8; 64]) -> Vec<u8> {
        let mut block = Vec::with_capacity(BLOCK_BYTES);
        block.extend_from_slice(hmask);
        block.extend_from_slice(qs);
        block.extend_from_slice(scales);
        block.extend_from_slice(&half::f16::from_f32(d).to_bits().to_le_bytes());
        block
    }

    fn all_one_scales() -> [u8; 12] {
        // All 16 sub-block scales = signed +1 (raw = 33)
        [
            0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0xA1, 0x11, 0x11, 0x11, 0x11,
        ]
    }

    #[test]
    fn test_unpack_scales_all_one() {
        let scales = all_one_scales();
        let decoded = unpack_scales(&scales);
        for (i, &s) in decoded.iter().enumerate() {
            assert_eq!(s, 1, "scale[{i}] = {s}, expected 1");
        }
    }

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

    #[test]
    fn test_dequant_matches_reference() {
        let mut hmask = [0u8; 32];
        let mut qs = [0u8; 64];
        for (i, h) in hmask.iter_mut().enumerate() {
            *h = ((i * 7 + 3) & 0xFF) as u8;
        }
        for (i, q) in qs.iter_mut().enumerate() {
            *q = ((i * 11 + 5) & 0xFF) as u8;
        }
        let scales = all_one_scales();
        let block = make_block(0.5, &scales, &hmask, &qs);
        let mut out_neon = vec![0.0f32; 256];
        let mut out_ref = vec![0.0f32; 256];
        Q3_KNeon.dequant_block(&block, &mut out_neon).expect("neon");
        Q3KRef.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_dequant_hmask_set_q3() {
        // hmask all set → subtract 0, all qs=0xFF → q_lo=3, scale=+1, d=1.0 → weight=3.0
        let hmask = [0xFFu8; 32];
        let qs = [0xFFu8; 64];
        let scales = all_one_scales();
        let block = make_block(1.0, &scales, &hmask, &qs);
        let mut out_neon = vec![0.0f32; 256];
        let mut out_ref = vec![0.0f32; 256];
        Q3_KNeon.dequant_block(&block, &mut out_neon).expect("neon");
        Q3KRef.dequant_block(&block, &mut out_ref).expect("ref");
        for i in 0..256 {
            assert!(
                (out_neon[i] - out_ref[i]).abs() < 1e-4,
                "weight[{i}]: neon={} ref={}",
                out_neon[i],
                out_ref[i]
            );
        }
    }

    #[test]
    fn test_gemv_matches_reference() {
        let mut hmask = [0u8; 32];
        let mut qs = [0u8; 64];
        for (i, h) in hmask.iter_mut().enumerate() {
            *h = ((i * 7 + 3) & 0xFF) as u8;
        }
        for (i, q) in qs.iter_mut().enumerate() {
            *q = ((i * 11 + 5) & 0xFF) as u8;
        }
        let scales = all_one_scales();
        let block = make_block(0.5, &scales, &hmask, &qs);

        let input: Vec<f32> = (0..256).map(|i| (i as f32) * 0.01 - 1.28).collect();

        let tensor_neon = QuantTensor::new(
            block.clone(),
            vec![1, 256],
            oxillama_gguf::GgufTensorType::Q3K,
        );
        let tensor_ref = QuantTensor::new(block, vec![1, 256], oxillama_gguf::GgufTensorType::Q3K);

        let mut out_neon = vec![0.0f32; 1];
        let mut out_ref = vec![0.0f32; 1];
        Q3_KNeon
            .gemv(&tensor_neon, &input, &mut out_neon)
            .expect("neon");
        Q3KRef.gemv(&tensor_ref, &input, &mut out_ref).expect("ref");

        let err = (out_neon[0] - out_ref[0]).abs();
        assert!(
            err < 0.1,
            "gemv: neon={} ref={} err={}",
            out_neon[0],
            out_ref[0],
            err
        );
    }

    // ── matvec_q8_fused Q3_K (NEON) ──────────────────────────────────────

    fn make_q8_0_block(scale: f32, values: &[i8; 32]) -> Vec<u8> {
        let mut block = Vec::with_capacity(34);
        block.extend_from_slice(&half::f16::from_f32(scale).to_bits().to_le_bytes());
        for &v in values {
            block.push(v as u8);
        }
        block
    }

    fn make_q8_acts(n_q8_blocks: usize, scale: f32, values: &[i8; 32]) -> Vec<u8> {
        make_q8_0_block(scale, values).repeat(n_q8_blocks)
    }

    #[test]
    fn test_q3k_neon_fused_matches_reference() {
        // One Q3_K super-block (256 weights) needs 8 Q8_0 activation blocks.
        let mut scales_raw = [0u8; 12];
        for (i, s) in scales_raw.iter_mut().enumerate() {
            *s = ((i * 11 + 3) & 0x3F) as u8;
        }
        let mut hmask = [0u8; 32];
        for (i, h) in hmask.iter_mut().enumerate() {
            *h = ((i * 13 + 7) & 0xFF) as u8;
        }
        let mut qs = [0u8; 64];
        for (i, q) in qs.iter_mut().enumerate() {
            *q = ((i * 7 + 3) & 0xFF) as u8;
        }

        let w_block = make_block(0.5, &scales_raw, &hmask, &qs);
        let act_vals: [i8; 32] = [
            1, -2, 3, -4, 5, -6, 7, -8, 9, -10, 11, -12, 13, -14, 15, -16, 0, 1, -1, 2, -2, 3, -3,
            4, -4, 5, -5, 6, -6, 7, -7, 8,
        ];
        let acts = make_q8_acts(8, 0.1, &act_vals);

        let mut out_neon = vec![0.0f32; 1];
        let mut out_ref = vec![0.0f32; 1];

        Q3_KNeon
            .matvec_q8_fused(&w_block, &acts, &mut out_neon, 1, 256)
            .expect("neon fused q3k single block");
        Q3KRef
            .matvec_q8_fused(&w_block, &acts, &mut out_ref, 1, 256)
            .expect("ref fused q3k single block");

        let err = (out_neon[0] - out_ref[0]).abs();
        assert!(
            err < 1e-3,
            "q3k_neon_fused_matches_reference: neon={} ref={} err={}",
            out_neon[0],
            out_ref[0],
            err
        );
    }

    #[test]
    fn test_q3k_neon_fused_multi_row() {
        // 3 rows × 512 cols = 2 Q3_K super-blocks per row → 16 Q8_0 act blocks.
        let n_rows = 3usize;
        let n_cols = 512usize;
        let blocks_per_row = 2usize;
        let q8_blocks_per_row = blocks_per_row * 8; // 16

        let mut all_weights = Vec::new();
        for r in 0..n_rows {
            for b in 0..blocks_per_row {
                let mut scales_raw = [0u8; 12];
                for (i, s) in scales_raw.iter_mut().enumerate() {
                    *s = ((r * 11 + b * 7 + i * 5 + 3) & 0x3F) as u8;
                }
                let mut hmask = [0u8; 32];
                for (i, h) in hmask.iter_mut().enumerate() {
                    *h = ((r * 13 + b * 17 + i * 7) & 0xFF) as u8;
                }
                let mut qs = [0u8; 64];
                for (i, q) in qs.iter_mut().enumerate() {
                    *q = ((r * 5 + b * 23 + i * 11) & 0xFF) as u8;
                }
                all_weights.extend(make_block(0.5 + r as f32 * 0.1, &scales_raw, &hmask, &qs));
            }
        }

        let act_vals: [i8; 32] = [
            2, -3, 5, -7, 1, -1, 4, -4, 6, -6, 3, -3, 2, -2, 1, -1, 8, -8, 7, -7, 6, -6, 5, -5, 4,
            -4, 3, -3, 2, -2, 1, -1,
        ];
        let acts = make_q8_acts(q8_blocks_per_row, 0.05, &act_vals);

        let mut out_neon = vec![0.0f32; n_rows];
        let mut out_ref = vec![0.0f32; n_rows];

        Q3_KNeon
            .matvec_q8_fused(&all_weights, &acts, &mut out_neon, n_rows, n_cols)
            .expect("neon fused q3k multi-row");
        Q3KRef
            .matvec_q8_fused(&all_weights, &acts, &mut out_ref, n_rows, n_cols)
            .expect("ref fused q3k multi-row");

        for i in 0..n_rows {
            let err = (out_neon[i] - out_ref[i]).abs();
            assert!(
                err < 1e-3,
                "q3k neon fused multi-row row {i}: neon={} ref={} err={}",
                out_neon[i],
                out_ref[i],
                err
            );
        }
    }
}