lattice_embed/simd/

int4.rs

//! INT4 packed-vector quantization and approximate dot products.
//!
//! Nibble layout and offset correction are shared by scalar and NEON paths.
//!
//! See docs/simd.md for the packed format and corrected dot-product derivation.

#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
use super::simd_config;

/// **Unstable**: INT4 quantization internals; scale/bias scheme may change.
///
/// Quantization parameters for INT4 conversion.
///
/// Uses symmetric unsigned quantization: the float range [-max_abs, max_abs]
/// is mapped to the integer range [0, 15].
#[derive(Debug, Clone, Copy)]
pub struct Int4Params {
    /// **Unstable**: scale factor; formula may change with quantization scheme update.
    pub scale: f32,
    /// **Unstable**: maximum absolute value; field may be removed.
    pub max_abs: f32,
}

impl Int4Params {
    /// **Unstable**: quantization parameter computation; may be folded into `Int4Vector::from_f32`.
    pub fn from_vector(vector: &[f32]) -> Self {
        let max_abs = max_abs_finite(vector);

        // Epsilon guard: avoid division by near-zero
        let scale = if max_abs > 1e-10 {
            15.0 / (2.0 * max_abs)
        } else {
            1.0
        };

        Self { scale, max_abs }
    }
}

fn max_abs_finite(vector: &[f32]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        if simd_config().avx2_enabled {
            // SAFETY: AVX2 was detected at runtime; the kernel bounds every load.
            return unsafe { max_abs_finite_avx2(vector) };
        }
    }
    #[cfg(target_arch = "aarch64")]
    {
        if simd_config().neon_enabled {
            // SAFETY: NEON was detected at runtime; the kernel bounds every load.
            return unsafe { max_abs_finite_neon(vector) };
        }
    }
    max_abs_finite_scalar(vector)
}

fn max_abs_finite_scalar(vector: &[f32]) -> f32 {
    vector
        .iter()
        .filter(|value| value.is_finite())
        .map(|value| value.abs())
        .fold(0.0, f32::max)
}

#[cfg(test)]
thread_local! {
    static INT4_MAX_ABS_SIMD_HITS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn max_abs_finite_avx2(vector: &[f32]) -> f32 {
    #[cfg(test)]
    INT4_MAX_ABS_SIMD_HITS.with(|hits| hits.set(hits.get() + 1));

    let chunks = vector.len() / 8;
    let sign = _mm256_set1_ps(-0.0);
    let inf = _mm256_set1_ps(f32::INFINITY);
    let mut maximum = _mm256_setzero_ps();

    for i in 0..chunks {
        let input = _mm256_loadu_ps(vector.as_ptr().add(i * 8));
        let abs = _mm256_andnot_ps(sign, input);
        let finite = _mm256_cmp_ps(abs, inf, _CMP_LT_OQ);
        maximum = _mm256_max_ps(maximum, _mm256_and_ps(abs, finite));
    }

    let mut lanes = [0.0f32; 8];
    _mm256_storeu_ps(lanes.as_mut_ptr(), maximum);
    let mut max_abs = lanes.into_iter().fold(0.0f32, f32::max);
    for &value in &vector[chunks * 8..] {
        if value.is_finite() {
            max_abs = max_abs.max(value.abs());
        }
    }
    max_abs
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn max_abs_finite_neon(vector: &[f32]) -> f32 {
    #[cfg(test)]
    INT4_MAX_ABS_SIMD_HITS.with(|hits| hits.set(hits.get() + 1));

    let chunks = vector.len() / 4;
    let inf = vdupq_n_f32(f32::INFINITY);
    let zero = vdupq_n_f32(0.0);
    let mut maximum = zero;

    for i in 0..chunks {
        let input = vld1q_f32(vector.as_ptr().add(i * 4));
        let abs = vabsq_f32(input);
        let finite = vcltq_f32(abs, inf);
        maximum = vmaxq_f32(maximum, vbslq_f32(finite, abs, zero));
    }

    let mut max_abs = vmaxvq_f32(maximum);
    for &value in &vector[chunks * 4..] {
        if value.is_finite() {
            max_abs = max_abs.max(value.abs());
        }
    }
    max_abs
}

/// **Unstable**: INT4 quantization format is under active design; struct layout may change.
///
/// Quantized INT4 vector with packed nibble storage.
#[derive(Debug, Clone)]
pub struct Int4Vector {
    /// **Unstable**: packed nibble data; bit packing scheme may change.
    pub data: Vec<u8>,
    /// **Unstable**: number of original dimensions.
    pub dims: usize,
    /// **Unstable**: quantization parameters; may be separated from the vector.
    pub params: Int4Params,
    /// **Unstable**: L2 norm; may be removed or moved.
    pub norm: f32,
}

impl Int4Vector {
    /// **Unstable**: quantization format; nibble packing may change.
    ///
    /// Each pair of consecutive dimensions is packed into one byte:
    /// - High nibble (bits 7..4) = even-indexed value
    /// - Low nibble (bits 3..0) = odd-indexed value
    pub fn from_f32(vector: &[f32]) -> Self {
        let params = Int4Params::from_vector(vector);
        let dims = vector.len();

        // Compute L2 norm
        let mut norm_sq = 0.0f32;
        for &v in vector {
            if v.is_finite() {
                norm_sq += v * v;
            }
        }
        let norm = norm_sq.sqrt();

        let data = quantize_int4(vector, params);

        Self {
            data,
            dims,
            params,
            norm,
        }
    }

    /// **Unstable**: dequantizes packed INT4 data, or returns empty for a malformed buffer.
    ///
    /// See [`docs/simd.md`](../../docs/simd.md#int4-vectors) for format and precision bounds.
    pub fn to_f32(&self) -> Vec<f32> {
        let required_bytes = self.dims.div_ceil(2);
        if self.data.len() < required_bytes {
            return Vec::new();
        }

        let scale = if self.params.scale.is_finite() && self.params.scale != 0.0 {
            self.params.scale
        } else {
            1.0
        };

        let mut result = Vec::with_capacity(self.dims);
        for i in 0..self.dims {
            let byte_idx = i / 2;
            let q = if i % 2 == 0 {
                (self.data[byte_idx] >> 4) & 0x0F
            } else {
                self.data[byte_idx] & 0x0F
            };
            result.push(q as f32 / scale - self.params.max_abs);
        }
        result
    }

    /// **Unstable**: INT4 dot product approximation; formula may change.
    ///
    /// Returns the dequantized dot product suitable for cosine distance computation.
    #[inline]
    pub fn dot_product(&self, other: &Int4Vector) -> f32 {
        dot_product_int4(self, other)
    }

    /// **Unstable**: INT4 cosine similarity approximation; delegates to `dot_product`.
    #[inline]
    pub fn cosine_similarity(&self, other: &Int4Vector) -> f32 {
        let denom = self.norm * other.norm;
        if denom == 0.0 || !denom.is_finite() {
            return 0.0;
        }
        self.dot_product(other) / denom
    }

    /// **Unstable**: complement of `cosine_similarity`; definition may evolve.
    #[inline]
    pub fn cosine_distance(&self, other: &Int4Vector) -> f32 {
        1.0 - self.cosine_similarity(other)
    }
}

fn quantize_int4(vector: &[f32], params: Int4Params) -> Vec<u8> {
    #[cfg(target_arch = "x86_64")]
    {
        if simd_config().avx2_enabled {
            // SAFETY: AVX2 was detected at runtime; the kernel bounds every load and store.
            return unsafe { quantize_int4_avx2(vector, params) };
        }
    }
    #[cfg(target_arch = "aarch64")]
    {
        if simd_config().neon_enabled {
            // SAFETY: NEON was detected at runtime; the kernel bounds every load and store.
            return unsafe { quantize_int4_neon(vector, params) };
        }
    }
    quantize_int4_scalar(vector, params)
}

fn quantize_int4_scalar(vector: &[f32], params: Int4Params) -> Vec<u8> {
    let mut data = vec![0u8; vector.len().div_ceil(2)];
    quantize_int4_scalar_tail(vector, params, &mut data, 0);
    data
}

fn quantize_int4_scalar_tail(vector: &[f32], params: Int4Params, data: &mut [u8], start: usize) {
    for (i, &value) in vector.iter().enumerate().skip(start) {
        let quantized = quantize_int4_value(value, params);
        if i % 2 == 0 {
            data[i / 2] |= quantized << 4;
        } else {
            data[i / 2] |= quantized;
        }
    }
}

#[inline]
fn quantize_int4_value(value: f32, params: Int4Params) -> u8 {
    let finite_value = if value.is_finite() { value } else { 0.0 };
    ((finite_value + params.max_abs) * params.scale)
        .round()
        .clamp(0.0, 15.0) as u8
}

#[cfg(test)]
thread_local! {
    static INT4_QUANTIZE_SIMD_HITS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn quantize_int4_avx2(vector: &[f32], params: Int4Params) -> Vec<u8> {
    #[cfg(test)]
    INT4_QUANTIZE_SIMD_HITS.with(|hits| hits.set(hits.get() + 1));

    let mut data = vec![0u8; vector.len().div_ceil(2)];
    let chunks = vector.len() / 8;
    let sign = _mm256_set1_ps(-0.0);
    let inf = _mm256_set1_ps(f32::INFINITY);
    let max_abs = _mm256_set1_ps(params.max_abs);
    let scale = _mm256_set1_ps(params.scale);
    let zero = _mm256_setzero_ps();
    let high = _mm256_set1_ps(15.0);
    let half = _mm256_set1_ps(0.5);
    let one = _mm256_set1_epi32(1);

    for i in 0..chunks {
        let base = i * 8;
        let input = _mm256_loadu_ps(vector.as_ptr().add(base));
        let abs = _mm256_andnot_ps(sign, input);
        let finite = _mm256_cmp_ps(abs, inf, _CMP_LT_OQ);
        let values = _mm256_and_ps(input, finite);
        let scaled = _mm256_mul_ps(_mm256_add_ps(values, max_abs), scale);
        let clamped = _mm256_min_ps(_mm256_max_ps(scaled, zero), high);
        let truncated = _mm256_cvttps_epi32(clamped);
        let fraction = _mm256_sub_ps(clamped, _mm256_cvtepi32_ps(truncated));
        let round_up = _mm256_castps_si256(_mm256_cmp_ps(fraction, half, _CMP_GE_OQ));
        let rounded = _mm256_add_epi32(truncated, _mm256_and_si256(round_up, one));
        let mut lanes = [0i32; 8];
        _mm256_storeu_si256(lanes.as_mut_ptr().cast::<__m256i>(), rounded);
        for pair in 0..4 {
            data[base / 2 + pair] = ((lanes[pair * 2] as u8) << 4) | lanes[pair * 2 + 1] as u8;
        }
    }

    quantize_int4_scalar_tail(vector, params, &mut data, chunks * 8);
    data
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn quantize_int4_neon(vector: &[f32], params: Int4Params) -> Vec<u8> {
    #[cfg(test)]
    INT4_QUANTIZE_SIMD_HITS.with(|hits| hits.set(hits.get() + 1));

    let mut data = vec![0u8; vector.len().div_ceil(2)];
    let chunks = vector.len() / 4;
    let inf = vdupq_n_f32(f32::INFINITY);
    let zero = vdupq_n_f32(0.0);
    let max_abs = vdupq_n_f32(params.max_abs);
    let scale = vdupq_n_f32(params.scale);
    let high = vdupq_n_f32(15.0);

    for i in 0..chunks {
        let base = i * 4;
        let input = vld1q_f32(vector.as_ptr().add(base));
        let finite = vcaltq_f32(input, inf);
        let values = vbslq_f32(finite, input, zero);
        let scaled = vmulq_f32(vaddq_f32(values, max_abs), scale);
        let clamped = vminq_f32(vmaxq_f32(scaled, zero), high);
        let rounded = vcvtaq_s32_f32(clamped);
        let mut lanes = [0i32; 4];
        vst1q_s32(lanes.as_mut_ptr(), rounded);
        data[base / 2] = ((lanes[0] as u8) << 4) | lanes[1] as u8;
        data[base / 2 + 1] = ((lanes[2] as u8) << 4) | lanes[3] as u8;
    }

    quantize_int4_scalar_tail(vector, params, &mut data, chunks * 4);
    data
}

/// **Unstable**: dequantized INT4 dot product; dispatch may change.
#[inline]
pub fn dot_product_int4(a: &Int4Vector, b: &Int4Vector) -> f32 {
    if a.dims != b.dims {
        return 0.0;
    }

    let scale_a = a.params.scale;
    let scale_b = b.params.scale;
    if scale_a == 0.0 || scale_b == 0.0 || !scale_a.is_finite() || !scale_b.is_finite() {
        return 0.0;
    }

    let packed_len = a.dims.div_ceil(2);
    if a.data.len() < packed_len || b.data.len() < packed_len {
        return 0.0;
    }

    #[cfg(target_arch = "aarch64")]
    {
        let config = simd_config();
        if config.neon_enabled {
            // SAFETY: aarch64 NEON is available by config, the packed data length guard
            // above prevents out-of-bounds loads, and the callee handles odd dimensions
            // without reading the padding nibble as a real dimension.
            let (raw_dot, sum_a, sum_b) =
                unsafe { dot_product_int4_neon_unrolled(&a.data, &b.data, a.dims) };
            return finish_int4_dot(raw_dot, sum_a, sum_b, a, b);
        }
    }

    let (raw_dot, sum_a, sum_b) = dot_product_int4_packed_scalar(&a.data, &b.data, a.dims);
    finish_int4_dot(raw_dot, sum_a, sum_b, a, b)
}

#[inline]
fn finish_int4_dot(raw_dot: i32, sum_a: i32, sum_b: i32, a: &Int4Vector, b: &Int4Vector) -> f32 {
    let raw_dot = raw_dot as f32;
    let sum_a = sum_a as f32;
    let sum_b = sum_b as f32;
    let scale_a = a.params.scale;
    let scale_b = b.params.scale;

    raw_dot / (scale_a * scale_b)
        - (b.params.max_abs * sum_a / scale_a)
        - (a.params.max_abs * sum_b / scale_b)
        + (a.dims as f32 * a.params.max_abs * b.params.max_abs)
}

#[inline]
fn dot_product_int4_packed_scalar(a: &[u8], b: &[u8], dims: usize) -> (i32, i32, i32) {
    let full_bytes = dims / 2;
    let mut raw_dot = 0i32;
    let mut sum_a = 0i32;
    let mut sum_b = 0i32;

    for i in 0..full_bytes {
        let av = a[i];
        let bv = b[i];
        let ah = ((av >> 4) & 0x0f) as i32;
        let al = (av & 0x0f) as i32;
        let bh = ((bv >> 4) & 0x0f) as i32;
        let bl = (bv & 0x0f) as i32;
        raw_dot += ah * bh + al * bl;
        sum_a += ah + al;
        sum_b += bh + bl;
    }

    if dims % 2 == 1 {
        let av = a[full_bytes];
        let bv = b[full_bytes];
        let ah = ((av >> 4) & 0x0f) as i32;
        let bh = ((bv >> 4) & 0x0f) as i32;
        raw_dot += ah * bh;
        sum_a += ah;
        sum_b += bh;
    }

    (raw_dot, sum_a, sum_b)
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
#[inline]
unsafe fn dot_product_int4_neon_unrolled(a: &[u8], b: &[u8], dims: usize) -> (i32, i32, i32) {
    debug_assert!(a.len() >= dims.div_ceil(2));
    debug_assert!(b.len() >= dims.div_ceil(2));

    const BLOCK_BYTES: usize = 16;
    const UNROLL: usize = 4;
    const CHUNK_BYTES: usize = BLOCK_BYTES * UNROLL;

    // Only bytes containing two valid dimensions are processed in SIMD.
    // If dims is odd, the final high nibble is handled separately and the low
    // padding nibble is ignored to preserve current to_f32 semantics.
    let full_bytes = dims / 2;
    let chunks = full_bytes / CHUNK_BYTES;

    let mut raw0 = vdupq_n_u32(0);
    let mut raw1 = vdupq_n_u32(0);
    let mut raw2 = vdupq_n_u32(0);
    let mut raw3 = vdupq_n_u32(0);
    let mut sum_a = vdupq_n_u32(0);
    let mut sum_b = vdupq_n_u32(0);
    let mask = vdupq_n_u8(0x0f);

    macro_rules! accumulate_block {
        ($base:expr, $raw:ident) => {{
            let a_bytes = vld1q_u8(a.as_ptr().add($base));
            let b_bytes = vld1q_u8(b.as_ptr().add($base));

            let a_hi = vshrq_n_u8::<4>(a_bytes);
            let b_hi = vshrq_n_u8::<4>(b_bytes);
            let a_lo = vandq_u8(a_bytes, mask);
            let b_lo = vandq_u8(b_bytes, mask);

            $raw = vpadalq_u16($raw, vmull_u8(vget_low_u8(a_hi), vget_low_u8(b_hi)));
            $raw = vpadalq_u16($raw, vmull_u8(vget_high_u8(a_hi), vget_high_u8(b_hi)));
            $raw = vpadalq_u16($raw, vmull_u8(vget_low_u8(a_lo), vget_low_u8(b_lo)));
            $raw = vpadalq_u16($raw, vmull_u8(vget_high_u8(a_lo), vget_high_u8(b_lo)));

            sum_a = vpadalq_u16(sum_a, vpaddlq_u8(a_hi));
            sum_a = vpadalq_u16(sum_a, vpaddlq_u8(a_lo));
            sum_b = vpadalq_u16(sum_b, vpaddlq_u8(b_hi));
            sum_b = vpadalq_u16(sum_b, vpaddlq_u8(b_lo));
        }};
    }

    for i in 0..chunks {
        let base = i * CHUNK_BYTES;
        accumulate_block!(base, raw0);
        accumulate_block!(base + BLOCK_BYTES, raw1);
        accumulate_block!(base + BLOCK_BYTES * 2, raw2);
        accumulate_block!(base + BLOCK_BYTES * 3, raw3);
    }

    let raw_vec = vaddq_u32(vaddq_u32(raw0, raw1), vaddq_u32(raw2, raw3));
    let mut raw_total = (vgetq_lane_u32::<0>(raw_vec)
        + vgetq_lane_u32::<1>(raw_vec)
        + vgetq_lane_u32::<2>(raw_vec)
        + vgetq_lane_u32::<3>(raw_vec)) as i32;
    let mut sum_a_total = (vgetq_lane_u32::<0>(sum_a)
        + vgetq_lane_u32::<1>(sum_a)
        + vgetq_lane_u32::<2>(sum_a)
        + vgetq_lane_u32::<3>(sum_a)) as i32;
    let mut sum_b_total = (vgetq_lane_u32::<0>(sum_b)
        + vgetq_lane_u32::<1>(sum_b)
        + vgetq_lane_u32::<2>(sum_b)
        + vgetq_lane_u32::<3>(sum_b)) as i32;

    let remainder_start = chunks * CHUNK_BYTES;
    for byte_idx in remainder_start..full_bytes {
        let av = *a.get_unchecked(byte_idx);
        let bv = *b.get_unchecked(byte_idx);
        let ah = ((av >> 4) & 0x0f) as i32;
        let al = (av & 0x0f) as i32;
        let bh = ((bv >> 4) & 0x0f) as i32;
        let bl = (bv & 0x0f) as i32;

        raw_total += ah * bh + al * bl;
        sum_a_total += ah + al;
        sum_b_total += bh + bl;
    }

    if dims % 2 == 1 {
        let av = *a.get_unchecked(full_bytes);
        let bv = *b.get_unchecked(full_bytes);
        let ah = ((av >> 4) & 0x0f) as i32;
        let bh = ((bv >> 4) & 0x0f) as i32;

        raw_total += ah * bh;
        sum_a_total += ah;
        sum_b_total += bh;
    }

    (raw_total, sum_a_total, sum_b_total)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn generate_vector(dim: usize, seed: u64) -> Vec<f32> {
        let mut state = seed ^ ((dim as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
        (0..dim)
            .map(|i| {
                state = state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407)
                    .wrapping_add(i as u64);
                let unit = ((state >> 32) as u32) as f32 / u32::MAX as f32;
                unit * 2.0 - 1.0
            })
            .collect()
    }

    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
    #[test]
    fn test_int4_quantize_explicit_simd_matches_scalar_and_is_dispatched() {
        #[cfg(target_arch = "x86_64")]
        if !std::arch::is_x86_feature_detected!("avx2") {
            return;
        }

        let params = Int4Params {
            scale: 7.5,
            max_abs: 1.0,
        };
        let tie_params = Int4Params {
            scale: 1.0,
            max_abs: 7.5,
        };
        let tie_input = [-7.5, -1.0, 1.0, 7.5, -1.0, 1.0, 0.0, -0.0];
        let scalar_ties = quantize_int4_scalar(&tie_input, tie_params);
        assert_eq!(scalar_ties, [0x07, 0x9f, 0x79, 0x88]);
        #[cfg(target_arch = "aarch64")]
        // SAFETY: baseline aarch64 provides NEON; the kernel bounds every access.
        let simd_ties = unsafe { quantize_int4_neon(&tie_input, tie_params) };
        #[cfg(target_arch = "x86_64")]
        // SAFETY: AVX2 was detected above; the kernel bounds every access.
        let simd_ties = unsafe { quantize_int4_avx2(&tie_input, tie_params) };
        assert_eq!(
            simd_ties, scalar_ties,
            "explicit SIMD must preserve ties-away rounding"
        );

        for dim in [0usize, 1, 3, 4, 7, 8, 9, 31, 32, 33, 383, 384, 385] {
            let mut input = generate_vector(dim, 900 + dim as u64);
            if dim > 0 {
                input[0] = f32::NAN;
            }
            if dim > 1 {
                input[1] = f32::INFINITY;
            }
            if dim > 2 {
                input[2] = f32::NEG_INFINITY;
            }
            if dim > 3 {
                let boundary = -1.0 + 0.5 / params.scale;
                input[3] = f32::from_bits(boundary.to_bits() - 1);
            }
            if dim > 4 {
                let boundary = -1.0 + 0.5 / params.scale;
                input[4] = f32::from_bits(boundary.to_bits() + 1);
            }

            let scalar = quantize_int4_scalar(&input, params);
            #[cfg(target_arch = "aarch64")]
            // SAFETY: baseline aarch64 provides NEON; the kernel bounds every access.
            let simd = unsafe { quantize_int4_neon(&input, params) };
            #[cfg(target_arch = "x86_64")]
            // SAFETY: AVX2 was detected above; the kernel bounds every access.
            let simd = unsafe { quantize_int4_avx2(&input, params) };
            assert_eq!(simd, scalar, "explicit SIMD mismatch at dim={dim}");
        }

        let input = generate_vector(385, 1_063);
        let before = INT4_QUANTIZE_SIMD_HITS.with(std::cell::Cell::get);
        let quantized = Int4Vector::from_f32(&input);
        let after = INT4_QUANTIZE_SIMD_HITS.with(std::cell::Cell::get);
        assert_eq!(
            after,
            before + 1,
            "Int4Vector::from_f32 did not execute its explicit SIMD quantizer"
        );
        assert_eq!(
            quantized.data,
            quantize_int4_scalar(&input, quantized.params)
        );
    }

    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
    #[test]
    fn test_int4_max_abs_explicit_simd_matches_scalar() {
        #[cfg(target_arch = "x86_64")]
        if !std::arch::is_x86_feature_detected!("avx2") {
            return;
        }

        for dim in [0usize, 1, 3, 4, 7, 8, 9, 31, 32, 33, 383, 384, 385] {
            let mut input = generate_vector(dim, 1_100 + dim as u64);
            if dim > 0 {
                input[0] = f32::NAN;
            }
            if dim > 1 {
                input[1] = f32::INFINITY;
            }
            if dim > 2 {
                input[2] = f32::NEG_INFINITY;
            }

            let scalar = max_abs_finite_scalar(&input);
            #[cfg(target_arch = "aarch64")]
            // SAFETY: baseline aarch64 provides NEON; the kernel bounds every access.
            let simd = unsafe { max_abs_finite_neon(&input) };
            #[cfg(target_arch = "x86_64")]
            // SAFETY: AVX2 was detected above; the kernel bounds every access.
            let simd = unsafe { max_abs_finite_avx2(&input) };
            assert_eq!(simd, scalar, "finite max-abs mismatch at dim={dim}");
        }

        let input = generate_vector(385, 1_063);
        let before = INT4_MAX_ABS_SIMD_HITS.with(std::cell::Cell::get);
        let params = Int4Params::from_vector(&input);
        let after = INT4_MAX_ABS_SIMD_HITS.with(std::cell::Cell::get);
        assert_eq!(
            after,
            before + 1,
            "Int4Params::from_vector did not execute its explicit SIMD reducer"
        );
        assert_eq!(params.max_abs, max_abs_finite_scalar(&input));
    }

    #[test]
    fn test_int4_roundtrip_accuracy() {
        let original = generate_vector(384, 42);
        let quantized = Int4Vector::from_f32(&original);
        let dequantized = quantized.to_f32();

        assert_eq!(dequantized.len(), original.len());

        // INT4 has only 16 levels, so error is larger than INT8.
        // Max error should be within 1/15 of the range.
        let max_abs = original
            .iter()
            .filter(|v| v.is_finite())
            .map(|v| v.abs())
            .fold(0.0f32, f32::max);
        let expected_max_error = 2.0 * max_abs / 15.0;

        for (i, (orig, deq)) in original.iter().zip(dequantized.iter()).enumerate() {
            let error = (orig - deq).abs();
            assert!(
                error <= expected_max_error + 1e-5,
                "INT4 roundtrip error too large at index {i}: orig={orig}, deq={deq}, error={error}, max_allowed={expected_max_error}"
            );
        }
    }

    #[test]
    fn test_int4_packing_correctness() {
        // Verify nibble packing: even index -> high nibble, odd -> low
        let v = vec![0.5, -0.5, 0.0, 1.0]; // 4 values -> 2 packed bytes
        let q = Int4Vector::from_f32(&v);
        assert_eq!(q.data.len(), 2);
        assert_eq!(q.dims, 4);

        // Verify roundtrip preserves approximate values
        let deq = q.to_f32();
        assert_eq!(deq.len(), 4);
        // 0.5 should map to roughly the right region
        assert!((deq[0] - 0.5).abs() < 0.15, "deq[0]={}", deq[0]);
        assert!((deq[1] - (-0.5)).abs() < 0.15, "deq[1]={}", deq[1]);
    }

    #[test]
    fn test_int4_odd_dimensions() {
        // Odd number of dimensions: last nibble has a padding zero
        let v = generate_vector(383, 77);
        let q = Int4Vector::from_f32(&v);
        assert_eq!(q.data.len(), 192); // ceil(383/2) = 192
        assert_eq!(q.dims, 383);

        let deq = q.to_f32();
        assert_eq!(deq.len(), 383);
    }

    #[test]
    fn test_int4_zero_vector() {
        let v = vec![0.0; 384];
        let q = Int4Vector::from_f32(&v);
        let deq = q.to_f32();
        for &val in &deq {
            assert!(
                val.abs() < 1e-5,
                "Zero vector should dequantize to near-zero"
            );
        }
    }

    #[test]
    fn test_int4_dot_product_vs_f32() {
        // Use correlated vectors so the true dot product is large relative to noise.
        // For uncorrelated random vectors, the expected dot product is ~0 while
        // quantization noise is O(dims * step^2), so relative error is unbounded.
        let a = generate_vector(384, 101);
        let b: Vec<f32> = a
            .iter()
            .enumerate()
            .map(|(i, &x)| x + 0.2 * (i as f32 * 0.3).sin())
            .collect();

        // f32 reference
        let f32_dot: f32 = a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum();

        let qa = Int4Vector::from_f32(&a);
        let qb = Int4Vector::from_f32(&b);
        let int4_dot = qa.dot_product(&qb);

        // INT4 has 16 levels; for correlated vectors the relative error should be
        // within ~15% (quantization step = 2*max_abs/15 per component).
        let rel_error = (f32_dot - int4_dot).abs() / f32_dot.abs().max(1.0);
        assert!(
            rel_error < 0.15,
            "INT4 dot product relative error too large: f32={f32_dot}, int4={int4_dot}, rel_error={rel_error}"
        );
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_packed_scalar_matches_neon_exact() {
        // Directly compare dot_product_int4_packed_scalar against
        // dot_product_int4_neon_unrolled on integer tuples. Both return
        // (raw_dot, sum_a, sum_b) as i32 — integer domain, exact equality expected.
        // This is the executing parity proof for the non-aarch64 fallback kernel.
        for dim in [1usize, 3, 31, 127, 383, 384] {
            let a_f32 = generate_vector(dim, 500 + dim as u64);
            let b_f32 = generate_vector(dim, 600 + dim as u64);
            let qa = Int4Vector::from_f32(&a_f32);
            let qb = Int4Vector::from_f32(&b_f32);

            let scalar_result = dot_product_int4_packed_scalar(&qa.data, &qb.data, dim);
            // SAFETY: aarch64 always has NEON; data slices are correctly sized by
            // Int4Vector::from_f32 (len = dims.div_ceil(2)).
            let neon_result = unsafe { dot_product_int4_neon_unrolled(&qa.data, &qb.data, dim) };

            assert_eq!(
                scalar_result, neon_result,
                "packed_scalar vs NEON integer mismatch at dim={dim}: scalar={scalar_result:?}, neon={neon_result:?}"
            );
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_int4_neon_matches_dequantized_scalar() {
        for dim in [1, 2, 31, 64, 127, 384, 768] {
            let a = generate_vector(dim, 501);
            let b = generate_vector(dim, 777);
            let qa = Int4Vector::from_f32(&a);
            let qb = Int4Vector::from_f32(&b);

            let a_deq = qa.to_f32();
            let b_deq = qb.to_f32();
            let expected: f32 = a_deq.iter().zip(b_deq.iter()).map(|(&x, &y)| x * y).sum();
            let got = qa.dot_product(&qb);

            assert!(
                (expected - got).abs() < 1e-4,
                "INT4 NEON mismatch for dim={dim}: expected={expected}, got={got}"
            );
        }
    }

    #[test]
    fn test_int4_cosine_similarity() {
        let a = generate_vector(384, 301);
        let b = generate_vector(384, 302);

        let qa = Int4Vector::from_f32(&a);
        let qb = Int4Vector::from_f32(&b);
        let int4_cos = qa.cosine_similarity(&qb);

        // Compute f32 reference cosine
        let dot: f32 = a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum();
        let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
        let f32_cos = dot / (norm_a * norm_b);

        assert!(
            (f32_cos - int4_cos).abs() < 0.1,
            "INT4 cosine too far from f32: f32={f32_cos}, int4={int4_cos}"
        );
    }

    #[test]
    fn test_int4_memory_savings() {
        let v = generate_vector(384, 999);
        let q = Int4Vector::from_f32(&v);

        // f32: 384 * 4 = 1536 bytes
        // INT4: ceil(384/2) = 192 bytes = 8x compression
        assert_eq!(q.data.len(), 192);
        assert_eq!(v.len() * 4, 1536);
    }

    #[test]
    fn test_int4_nan_inf_handling() {
        let v = vec![
            1.0,
            f32::NAN,
            f32::INFINITY,
            f32::NEG_INFINITY,
            -1.0,
            0.5,
            0.0,
            -0.3,
        ];
        let q = Int4Vector::from_f32(&v);
        let deq = q.to_f32();
        assert_eq!(deq.len(), 8);
        // NaN and Inf should be treated as 0
        // The dequantized value for the "0" slot should be near -max_abs + something,
        // but the key invariant is no panics and finite output.
        for &val in &deq {
            assert!(val.is_finite(), "Dequantized value should be finite");
        }
    }

    // --- Issue #211 regression tests -------------------------------------------------

    #[test]
    fn test_int4_to_f32_short_data_returns_empty() {
        // dims=128 requires 64 bytes; supply only 4.
        let q = Int4Vector {
            dims: 128,
            data: vec![0xFFu8; 4],
            params: Int4Params {
                scale: 7.5,
                max_abs: 1.0,
            },
            norm: 1.0,
        };
        let result = q.to_f32();
        assert!(
            result.is_empty(),
            "to_f32 on malformed Int4Vector must return empty Vec"
        );
    }

    #[test]
    fn test_int4_to_f32_exact_length_works() {
        // Exactly the right number of bytes — must succeed and not index OOB.
        let v: Vec<f32> = (0..128).map(|i| (i as f32) / 64.0 - 1.0).collect();
        let q = Int4Vector::from_f32(&v);
        let deq = q.to_f32();
        assert_eq!(deq.len(), 128);
        for &val in &deq {
            assert!(val.is_finite());
        }
    }
}