omendb_core/compression/

scalar.rs

//! Scalar Quantization (SQ8) for OmenDB
//!
//! Compresses f32 vectors to u8 (4x compression, ~97% recall, 2-3x faster than FP32).
//!
//! # Algorithm
//!
//! Uniform min/max scaling (single scale/offset for all dimensions):
//! - Train: Compute global min, max from sample vectors
//! - Quantize: u8[d] = round((f32[d] - offset) / scale)
//! - Distance: Integer SIMD dot product with float reconstruction
//!
//! # Performance (768D, Apple M3 Max)
//!
//! - 4x compression (f32 → u8)
//! - 2-3x faster than FP32 (integer SIMD)
//! - ~97% recall (vs 99%+ for per-dimension quantization)

use serde::{Deserialize, Serialize};

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

/// Trained scalar quantization parameters (uniform quantization)
///
/// Uses a single scale/offset for all dimensions, enabling integer SIMD
/// for 2-3x speedup over FP32.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScalarParams {
    /// Global scale factor: (max - min) / 255
    pub scale: f32,
    /// Global offset (minimum value)
    pub offset: f32,
    /// Number of dimensions
    pub dimensions: usize,
}

/// Precomputed data for a quantized vector
#[derive(Debug, Clone)]
pub struct QuantizedVector {
    /// Quantized values (u8)
    pub data: Vec<u8>,
    /// Precomputed: sum of quantized values (Σ data[i])
    pub sum: i32,
    /// Precomputed: squared norm of dequantized vector
    pub norm_sq: f32,
}

/// Precomputed query data for fast integer SIMD distance
#[derive(Debug, Clone)]
pub struct QueryPrep {
    /// Quantized query values (u8 for SIMD dot product)
    pub quantized: Vec<u8>,
    /// Query squared norm: ||q||²
    pub norm_sq: f32,
    /// Sum of quantized query values
    pub sum: i32,
}

impl ScalarParams {
    /// Create uninitialized params (for lazy training)
    ///
    /// Uses identity mapping until trained.
    #[must_use]
    pub fn uninitialized(dimensions: usize) -> Self {
        Self {
            scale: 1.0 / 255.0,
            offset: 0.0,
            dimensions,
        }
    }

    /// Train scalar quantization from sample vectors
    ///
    /// Uses 1st and 99th percentiles to handle outliers.
    ///
    /// # Errors
    /// Returns error if vectors is empty or vectors have inconsistent dimensions.
    pub fn train(vectors: &[&[f32]]) -> Result<Self, &'static str> {
        Self::train_with_percentiles(vectors, 0.01, 0.99)
    }

    /// Train with custom percentile bounds
    pub fn train_with_percentiles(
        vectors: &[&[f32]],
        lower_percentile: f32,
        upper_percentile: f32,
    ) -> Result<Self, &'static str> {
        if vectors.is_empty() {
            return Err("Need at least one vector to train");
        }
        let dimensions = vectors[0].len();
        if !vectors.iter().all(|v| v.len() == dimensions) {
            return Err("All vectors must have same dimensions");
        }

        // Collect ALL values across all vectors and dimensions
        let mut all_values: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
        all_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let n = all_values.len();
        let lower_idx = ((n as f32 * lower_percentile) as usize).min(n - 1);
        let upper_idx = ((n as f32 * upper_percentile) as usize).min(n - 1);

        let min_val = all_values[lower_idx];
        let max_val = all_values[upper_idx];

        let range = max_val - min_val;
        let (offset, scale) = if range < 1e-7 {
            (min_val - 0.5, 1.0 / 255.0)
        } else {
            (min_val, range / 255.0)
        };

        Ok(Self {
            scale,
            offset,
            dimensions,
        })
    }

    /// Quantize a vector to u8 with precomputed metadata
    #[must_use]
    pub fn quantize(&self, vector: &[f32]) -> QuantizedVector {
        debug_assert_eq!(vector.len(), self.dimensions);

        let inv_scale = 1.0 / self.scale;
        let data: Vec<u8> = vector
            .iter()
            .map(|&v| ((v - self.offset) * inv_scale).clamp(0.0, 255.0).round() as u8)
            .collect();

        let sum: i32 = data.iter().map(|&x| x as i32).sum();

        // Compute dequantized norm
        let norm_sq: f32 = data
            .iter()
            .map(|&x| {
                let dequant = x as f32 * self.scale + self.offset;
                dequant * dequant
            })
            .sum();

        QuantizedVector { data, sum, norm_sq }
    }

    /// Quantize a vector, returning only the u8 data (for storage)
    #[must_use]
    pub fn quantize_to_bytes(&self, vector: &[f32]) -> Vec<u8> {
        debug_assert_eq!(vector.len(), self.dimensions);

        let inv_scale = 1.0 / self.scale;
        vector
            .iter()
            .map(|&v| ((v - self.offset) * inv_scale).clamp(0.0, 255.0).round() as u8)
            .collect()
    }

    /// Dequantize a u8 vector back to f32 (approximate)
    #[must_use]
    pub fn dequantize(&self, quantized: &[u8]) -> Vec<f32> {
        quantized
            .iter()
            .map(|&q| q as f32 * self.scale + self.offset)
            .collect()
    }

    /// Compute squared norm of dequantized vector: ||dequant(q)||²
    #[must_use]
    pub fn dequantized_norm_squared(&self, quantized: &[u8]) -> f32 {
        quantized
            .iter()
            .map(|&q| {
                let dequant = q as f32 * self.scale + self.offset;
                dequant * dequant
            })
            .sum()
    }

    /// Compute sum of quantized values (for distance computation)
    #[must_use]
    pub fn quantized_sum(&self, quantized: &[u8]) -> i32 {
        quantized.iter().map(|&x| x as i32).sum()
    }

    /// Prepare query for fast integer SIMD distance computation
    #[must_use]
    pub fn prepare_query(&self, query: &[f32]) -> QueryPrep {
        debug_assert_eq!(query.len(), self.dimensions);

        let inv_scale = 1.0 / self.scale;
        let quantized: Vec<u8> = query
            .iter()
            .map(|&v| ((v - self.offset) * inv_scale).clamp(0.0, 255.0).round() as u8)
            .collect();

        let norm_sq: f32 = query.iter().map(|x| x * x).sum();
        let sum: i32 = quantized.iter().map(|&x| x as i32).sum();

        QueryPrep {
            quantized,
            norm_sq,
            sum,
        }
    }

    /// Compute L2² distance using integer SIMD
    ///
    /// Uses the identity: ||q - v||² = ||q||² + ||v||² - 2⟨q,v⟩
    /// The dot product is computed in integer domain for speed.
    #[inline(always)]
    #[must_use]
    pub fn distance_l2_squared(&self, query_prep: &QueryPrep, vec: &QuantizedVector) -> f32 {
        // Integer dot product (SIMD accelerated) - uses u8×u8→u32
        let int_dot = self.int_dot_product(&query_prep.quantized, &vec.data);

        // Reconstruct actual dot product: scale² × int_dot + corrections
        // dot(q, v) = scale² × Σ q_int[i] × v_int[i]
        //           + scale × offset × (Σ q_int[i] + Σ v_int[i])
        //           + offset² × dim
        let scale_sq = self.scale * self.scale;
        let dot = scale_sq * int_dot as f32
            + self.scale * self.offset * (query_prep.sum + vec.sum) as f32
            + self.offset * self.offset * self.dimensions as f32;

        // L2² = ||q||² + ||v||² - 2⟨q,v⟩
        query_prep.norm_sq + vec.norm_sq - 2.0 * dot
    }

    /// Compute L2² distance from raw bytes (for storage integration)
    ///
    /// Slightly slower than `distance_l2_squared` since it computes sum and norm on the fly.
    #[inline(always)]
    #[must_use]
    pub fn distance_l2_squared_raw(
        &self,
        query_prep: &QueryPrep,
        vec_data: &[u8],
        vec_sum: i32,
        vec_norm_sq: f32,
    ) -> f32 {
        let int_dot = self.int_dot_product(&query_prep.quantized, vec_data);

        let scale_sq = self.scale * self.scale;
        let dot = scale_sq * int_dot as f32
            + self.scale * self.offset * (query_prep.sum + vec_sum) as f32
            + self.offset * self.offset * self.dimensions as f32;

        query_prep.norm_sq + vec_norm_sq - 2.0 * dot
    }

    /// Integer dot product with SIMD acceleration (u8 × u8 → u32)
    #[inline(always)]
    fn int_dot_product(&self, query: &[u8], vec: &[u8]) -> u32 {
        debug_assert_eq!(query.len(), vec.len());

        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return unsafe { self.int_dot_product_avx2(query, vec) };
            }
            Self::int_dot_product_scalar(query, vec)
        }

        #[cfg(target_arch = "aarch64")]
        {
            unsafe { self.int_dot_product_neon(query, vec) }
        }

        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
        {
            Self::int_dot_product_scalar(query, vec)
        }
    }

    #[inline]
    #[allow(dead_code)]
    fn int_dot_product_scalar(query: &[u8], vec: &[u8]) -> u32 {
        query
            .iter()
            .zip(vec.iter())
            .map(|(&q, &v)| q as u32 * v as u32)
            .sum()
    }

    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    #[allow(clippy::unused_self)]
    unsafe fn int_dot_product_avx2(&self, query: &[u8], vec: &[u8]) -> u32 {
        let mut sum = _mm256_setzero_si256();
        let mut i = 0;

        while i + 32 <= query.len() {
            let q = _mm256_loadu_si256(query.as_ptr().add(i).cast());
            let v = _mm256_loadu_si256(vec.as_ptr().add(i).cast());

            let q_lo = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(q, 0));
            let q_hi = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(q, 1));
            let v_lo = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(v, 0));
            let v_hi = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(v, 1));

            let prod_lo = _mm256_madd_epi16(q_lo, v_lo);
            let prod_hi = _mm256_madd_epi16(q_hi, v_hi);
            sum = _mm256_add_epi32(sum, prod_lo);
            sum = _mm256_add_epi32(sum, prod_hi);

            i += 32;
        }

        while i + 16 <= query.len() {
            let q = _mm256_cvtepu8_epi16(_mm_loadu_si128(query.as_ptr().add(i).cast()));
            let v = _mm256_cvtepu8_epi16(_mm_loadu_si128(vec.as_ptr().add(i).cast()));
            let prod = _mm256_madd_epi16(q, v);
            sum = _mm256_add_epi32(sum, prod);
            i += 16;
        }

        let sum128 = _mm_add_epi32(
            _mm256_extracti128_si256(sum, 0),
            _mm256_extracti128_si256(sum, 1),
        );
        let sum64 = _mm_add_epi32(sum128, _mm_srli_si128(sum128, 8));
        let sum32 = _mm_add_epi32(sum64, _mm_srli_si128(sum64, 4));
        let mut result = _mm_cvtsi128_si32(sum32) as u32;

        for j in i..query.len() {
            result += query[j] as u32 * vec[j] as u32;
        }

        result
    }

    #[cfg(target_arch = "aarch64")]
    #[inline(always)]
    #[allow(clippy::unused_self)]
    unsafe fn int_dot_product_neon(&self, query: &[u8], vec: &[u8]) -> u32 {
        use std::arch::aarch64::{
            vaddq_u32, vaddvq_u32, vdupq_n_u32, vget_low_u8, vld1q_u8, vmull_high_u8, vmull_u8,
            vpadalq_u16,
        };

        // Use 4 accumulators to hide latency and increase ILP
        let mut sum0 = vdupq_n_u32(0);
        let mut sum1 = vdupq_n_u32(0);
        let mut sum2 = vdupq_n_u32(0);
        let mut sum3 = vdupq_n_u32(0);
        let mut i = 0;

        // Process 64 elements per iteration (4x unrolling)
        while i + 64 <= query.len() {
            let q0 = vld1q_u8(query.as_ptr().add(i));
            let v0 = vld1q_u8(vec.as_ptr().add(i));
            let prod0_lo = vmull_u8(vget_low_u8(q0), vget_low_u8(v0));
            let prod0_hi = vmull_high_u8(q0, v0);
            sum0 = vpadalq_u16(sum0, prod0_lo);
            sum0 = vpadalq_u16(sum0, prod0_hi);

            let q1 = vld1q_u8(query.as_ptr().add(i + 16));
            let v1 = vld1q_u8(vec.as_ptr().add(i + 16));
            let prod1_lo = vmull_u8(vget_low_u8(q1), vget_low_u8(v1));
            let prod1_hi = vmull_high_u8(q1, v1);
            sum1 = vpadalq_u16(sum1, prod1_lo);
            sum1 = vpadalq_u16(sum1, prod1_hi);

            let q2 = vld1q_u8(query.as_ptr().add(i + 32));
            let v2 = vld1q_u8(vec.as_ptr().add(i + 32));
            let prod2_lo = vmull_u8(vget_low_u8(q2), vget_low_u8(v2));
            let prod2_hi = vmull_high_u8(q2, v2);
            sum2 = vpadalq_u16(sum2, prod2_lo);
            sum2 = vpadalq_u16(sum2, prod2_hi);

            let q3 = vld1q_u8(query.as_ptr().add(i + 48));
            let v3 = vld1q_u8(vec.as_ptr().add(i + 48));
            let prod3_lo = vmull_u8(vget_low_u8(q3), vget_low_u8(v3));
            let prod3_hi = vmull_high_u8(q3, v3);
            sum3 = vpadalq_u16(sum3, prod3_lo);
            sum3 = vpadalq_u16(sum3, prod3_hi);

            i += 64;
        }

        while i + 16 <= query.len() {
            let q = vld1q_u8(query.as_ptr().add(i));
            let v = vld1q_u8(vec.as_ptr().add(i));
            let prod_lo = vmull_u8(vget_low_u8(q), vget_low_u8(v));
            let prod_hi = vmull_high_u8(q, v);
            sum0 = vpadalq_u16(sum0, prod_lo);
            sum0 = vpadalq_u16(sum0, prod_hi);
            i += 16;
        }

        let sum01 = vaddq_u32(sum0, sum1);
        let sum23 = vaddq_u32(sum2, sum3);
        let sum_all = vaddq_u32(sum01, sum23);
        let mut result = vaddvq_u32(sum_all);

        for j in i..query.len() {
            result += query[j] as u32 * vec[j] as u32;
        }

        result
    }
}

/// Compute symmetric L2² distance between two quantized vectors
#[inline]
#[must_use]
pub fn symmetric_l2_squared_u8(a: &[u8], b: &[u8]) -> u32 {
    a.iter()
        .zip(b.iter())
        .map(|(&x, &y)| {
            let diff = (i16::from(x) - i16::from(y)) as i32;
            (diff * diff) as u32
        })
        .sum()
}

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

    #[test]
    fn test_train_and_quantize() {
        let vectors: Vec<Vec<f32>> = vec![
            vec![0.0, 0.5, 1.0, 0.3],
            vec![0.1, 0.6, 0.9, 0.4],
            vec![0.2, 0.4, 0.8, 0.5],
        ];
        let refs: Vec<&[f32]> = vectors.iter().map(Vec::as_slice).collect();

        let params = ScalarParams::train(&refs).unwrap();

        let quantized = params.quantize(&vectors[0]);
        assert_eq!(quantized.data.len(), 4);
        assert!(quantized.sum > 0);
        assert!(quantized.norm_sq > 0.0);
    }

    #[test]
    fn test_distance_accuracy() {
        use rand::Rng;

        let dim = 128;
        let n_vectors = 100;
        let mut rng = rand::thread_rng();

        // Generate normalized vectors (common in embeddings)
        let vectors: Vec<Vec<f32>> = (0..n_vectors)
            .map(|_| {
                let v: Vec<f32> = (0..dim).map(|_| rng.gen_range(-1.0..1.0)).collect();
                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
                v.iter().map(|x| x / norm).collect()
            })
            .collect();

        let refs: Vec<&[f32]> = vectors.iter().map(|v| v.as_slice()).collect();
        let params = ScalarParams::train(&refs).unwrap();

        let quantized: Vec<_> = vectors.iter().map(|v| params.quantize(v)).collect();

        let query = &vectors[0];
        let query_prep = params.prepare_query(query);

        let mut max_rel_error = 0.0f32;

        for (i, (orig, quant)) in vectors.iter().zip(quantized.iter()).enumerate() {
            if i == 0 {
                continue;
            }

            let true_dist: f32 = query
                .iter()
                .zip(orig.iter())
                .map(|(a, b)| (a - b).powi(2))
                .sum();

            let quant_dist = params.distance_l2_squared(&query_prep, quant);

            let rel_error = (true_dist - quant_dist).abs() / true_dist.max(1e-6);
            max_rel_error = max_rel_error.max(rel_error);
        }

        println!(
            "SQ8 max relative distance error: {:.2}%",
            max_rel_error * 100.0
        );
        assert!(
            max_rel_error < 0.15,
            "Distance error too large: {max_rel_error:.4}"
        );
    }

    #[test]
    fn test_int_dot_product() {
        let vectors: Vec<Vec<f32>> = vec![vec![0.5; 768], vec![0.3; 768]];
        let refs: Vec<&[f32]> = vectors.iter().map(Vec::as_slice).collect();

        let params = ScalarParams::train(&refs).unwrap();
        let query_prep = params.prepare_query(&vectors[0]);
        let quantized = params.quantize(&vectors[1]);

        let dist = params.distance_l2_squared(&query_prep, &quantized);
        assert!(dist >= 0.0);
        assert!(!dist.is_nan());
    }

    #[test]
    fn test_dequantize_roundtrip() {
        let vectors: Vec<Vec<f32>> = vec![
            vec![0.0, 0.5, 1.0],
            vec![0.1, 0.6, 0.9],
            vec![0.2, 0.4, 0.8],
        ];
        let refs: Vec<&[f32]> = vectors.iter().map(Vec::as_slice).collect();

        let params = ScalarParams::train(&refs).unwrap();
        let quantized = params.quantize(&vectors[0]);
        let dequantized = params.dequantize(&quantized.data);

        for (orig, deq) in vectors[0].iter().zip(dequantized.iter()) {
            assert!(
                (orig - deq).abs() < 0.05,
                "Roundtrip error too large: {} vs {}",
                orig,
                deq
            );
        }
    }

    #[test]
    fn test_symmetric_distance() {
        let a: Vec<u8> = vec![0, 100, 200, 255];
        let b: Vec<u8> = vec![0, 100, 200, 255];
        let dist = symmetric_l2_squared_u8(&a, &b);
        assert_eq!(dist, 0);

        let c: Vec<u8> = vec![10, 110, 210, 245];
        let dist2 = symmetric_l2_squared_u8(&a, &c);
        assert!(dist2 > 0);
    }

    #[test]
    fn test_compression_ratio() {
        let dims = 768;
        let original_size = dims * 4; // f32 = 4 bytes
        let quantized_size = dims; // u8 = 1 byte

        assert_eq!(original_size / quantized_size, 4);
    }
}