zeppelin-embed 0.6.0

//! Four-bit Extended-RaBitQ row encoding.
//!
//! The conceptual codebook is the normalized half-integer grid
//! `{-7.5, -6.5, ..., -0.5, 0.5, ..., 6.5, 7.5}^d`. It is fixed by the bit
//! width: there is no learned codebook, training set, or first-query setup.
//! Random rotation is optional in Extended-RaBitQ and is off by default in v1;
//! the recall harness must validate that identity-rotation assumption for each
//! embedding model.
//!
//! For input `v`, Algorithm 1 selects `y` maximizing `<y,v>/||y||`. At scale
//! zero each signed magnitude is `0.5`. Coordinate `i` then crosses its seven
//! rounding boundaries at `t = level / |v_i|`, `level=1..=7`. Sorting all
//! critical values, applying equal thresholds as one group, and retaining the
//! best normalized projection enumerates every realizable rounded direction
//! exactly in `O(2^B d log d)` time.
//!
//! The stored RaBitQ correction is `c = ||v||^2 / <y,v>`. The paper writes the
//! estimator with normalized `y`; because `||y||` occurs in numerator and
//! denominator it cancels, leaving `c * <y,q_hat>`. The companion norm is kept
//! for reconstruction/error bounds and future distance bookkeeping.

use crate::kernels::{MAX_DOT_I8_DIMENSION, dot_bit4_prepared, score_bit4_prepared_batch};

use super::QuantError;

mod controlled;
pub(crate) use controlled::{
    Bit4ControlError, Bit4Scratch, est_dot_bit4_controlled, prepare_bit4_query_controlled,
    quantize_bit4_controlled,
};

const CODES_PER_BYTE: usize = 2;
const MAX_MAGNITUDE_LEVEL: u8 = 7;

/// Per-row scalars required by the four-bit unbiased dot estimator.
#[derive(Clone, Copy, Debug, PartialEq)]
#[repr(C)]
pub struct Bit4Factors {
    scale: f32,
    normalized_norm: f32,
    normalized_correction: f32,
}

impl Bit4Factors {
    /// Reconstructs the permanent three-f32 persisted factor record.
    #[must_use]
    pub const fn from_persisted(
        scale: f32,
        normalized_norm: f32,
        normalized_correction: f32,
    ) -> Self {
        Self {
            scale,
            normalized_norm,
            normalized_correction,
        }
    }

    /// Returns the permanent little-endian record fields in declaration order.
    #[must_use]
    pub const fn persisted_fields(self) -> [f32; 3] {
        [self.scale, self.normalized_norm, self.normalized_correction]
    }

    /// Returns the original row's Euclidean norm.
    ///
    /// A finite `f32` row scale and finite normalized `f32` norm are expanded
    /// in `f64`, so the logical norm remains usable above `f32::MAX` without
    /// enlarging the three-scalar factor record.
    #[must_use]
    pub const fn norm(self) -> f64 {
        self.scale as f64 * self.normalized_norm as f64
    }

    /// Returns `||v||^2 / <y, v>` for the selected grid direction.
    ///
    /// The stored normalized correction is multiplied by the finite row scale
    /// in `f64`, avoiding overflow in persisted metadata.
    #[must_use]
    pub const fn correction(self) -> f64 {
        self.scale as f64 * self.normalized_correction as f64
    }

    /// Returns the L2 error bound for norm-scaled code reconstruction.
    ///
    /// Every selected grid coordinate is sign-aligned with the input, hence
    /// the direction cosine is non-negative and the deterministic bound is
    /// `sqrt(2) ||v||`. The exact grid search normally gives a tighter error;
    /// retaining this conservative bound avoids a fourth per-row scalar.
    #[must_use]
    pub fn reconstruction_error_bound(self) -> f64 {
        std::f64::consts::SQRT_2 * self.norm()
    }

    pub(crate) const fn scoring_parts(self) -> (f32, f32) {
        (self.scale, self.normalized_correction)
    }
}

/// Query-side signed-byte representation prepared once and reused per row.
///
/// Coordinates are symmetrically quantized to `[-127, 127]` with seeded
/// stochastic rounding, so their reconstruction is unbiased in expectation.
/// The private kernel representation groups even then odd coordinates within
/// each 32-value block, matching the packed row's high/low nibble vectors.
#[derive(Clone, Debug, PartialEq)]
pub struct Bit4Query {
    /// Even coordinates followed by odd coordinates in each 32-value block.
    codes: Vec<i8>,
    code_sum: i32,
    scale_half: f64,
}

impl Bit4Query {
    /// Returns the query coordinate count.
    #[must_use]
    pub fn len(&self) -> usize {
        self.codes.len()
    }

    /// Returns whether the query has zero coordinates.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.codes.is_empty()
    }

    pub(crate) fn resident_bytes(&self) -> usize {
        self.codes.capacity()
    }

    /// Returns the prepared code bytes, pre-interleave sum, and half-scale
    /// exactly as consumed by the scoring kernel.
    #[cfg(any(test, feature = "test-seams"))]
    #[doc(hidden)]
    #[must_use]
    pub fn observation_parts(&self) -> (&[i8], i32, f64) {
        (&self.codes, self.code_sum, self.scale_half)
    }

    pub(crate) fn kernel_parts(&self) -> (&[i8], i32, f64) {
        (&self.codes, self.code_sum, self.scale_half)
    }
}

#[derive(Clone, Copy, Debug)]
struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    const fn new(seed: u64) -> Self {
        Self { state: seed }
    }

    fn next_open_unit_f64(&mut self) -> f64 {
        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut value = self.state;
        value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        value ^= value >> 31;
        ((value >> 11) as f64 + 0.5) * (1.0 / 9_007_199_254_740_992.0)
    }
}

#[derive(Clone, Copy, Debug)]
struct CriticalValue {
    threshold: f64,
    coordinate: usize,
    level: u8,
    magnitude: f64,
}

/// Quantizes one row to four-bit Extended-RaBitQ codes.
///
/// Two coordinates occupy each byte in coordinate order, high nibble first.
/// Unsigned code `u` represents `u - 7.5`. For an odd dimension, the unused
/// low nibble in the final byte is canonical zero padding.
///
/// # Errors
///
/// Returns [`QuantError::EmptyVector`] for an empty input,
/// [`QuantError::DimensionTooLarge`] above the i8 scoring-kernel limit,
/// [`QuantError::NonFinite`] for NaN or infinity, and
/// [`QuantError::OutputLength`] unless `out.len() == ceil(v.len() / 2)`.
/// Validation completes before `out` is modified.
pub fn quantize_bit4(v: &[f32], out: &mut [u8]) -> Result<Bit4Factors, QuantError> {
    validate_input(v, out.len())?;
    let mut scratch = Bit4Scratch::compatibility(v.len());
    match quantize_bit4_controlled(v, out, &mut scratch, &mut |_| Ok::<_, ()>(())) {
        Ok(factors) => Ok(factors),
        Err(Bit4ControlError::Quant(error)) => Err(error),
        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
        Err(Bit4ControlError::Control(())) => Err(QuantError::Allocation),
    }
}

/// Reconstructs the norm-scaled selected four-bit code direction.
///
/// The result is `||v|| * y / ||y||`. Serving uses [`est_dot_bit4`] directly;
/// reconstruction exists to specify and verify the codec's geometric error.
///
/// # Errors
///
/// Returns the code-length and padding errors documented by
/// [`est_dot_bit4`]. The output length defines the logical dimension of the
/// final partial byte.
pub fn dequantize_bit4(
    codes: &[u8],
    factors: Bit4Factors,
    out: &mut [f32],
) -> Result<(), QuantError> {
    let unpacked = unpack(codes, out.len())?;
    if factors.scale == 0.0 {
        out.fill(0.0);
        return Ok(());
    }
    let doubled_code_norm = unpacked
        .iter()
        .map(|&value| {
            let value = f64::from(value);
            value * value
        })
        .sum::<f64>()
        .sqrt();
    for (output, &value) in out.iter_mut().zip(&unpacked) {
        *output = finite_f32(factors.norm() * f64::from(value) / doubled_code_norm);
    }
    Ok(())
}

/// Prepares one query for repeated four-bit row scoring.
///
/// The seed controls reproducible stochastic rounding only. Random rotation is
/// not performed by this v1 API. A zero query receives scale zero and all-zero
/// signed codes.
///
/// # Errors
///
/// Returns [`QuantError::EmptyVector`], [`QuantError::DimensionTooLarge`], or
/// [`QuantError::NonFinite`] under the same input policy as
/// [`quantize_bit4`].
pub fn prepare_bit4_query(q: &[f32], seed: u64) -> Result<Bit4Query, QuantError> {
    #[cfg(any(test, feature = "test-seams"))]
    super::QUERY_PREPARATIONS.with(|calls| {
        if let Some(calls) = calls.borrow_mut().as_mut() {
            calls.bit4.push((q.len(), seed));
        }
    });
    match prepare_bit4_query_controlled(q, seed, &mut |_| Ok::<_, ()>(())) {
        Ok((query, _)) => Ok(query),
        Err(Bit4ControlError::Quant(error)) => Err(error),
        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
        Err(Bit4ControlError::Control(())) => Err(QuantError::Allocation),
    }
}

/// Estimates a dot product directly from one packed four-bit row.
///
/// The native packed kernel scores unsigned nibbles directly, then applies
/// `dot(z, 2*u - 15) = 2*dot(z, u) - 15*sum(z)` once per row. The query's
/// block-interleaved representation pairs directly with high- and low-nibble
/// vectors, so scoring needs no expanded row, ZIP interleave, or per-vector
/// grid conversion. If the prepared query reconstructs as `scale*z`, the
/// integer kernel gives `<y,q_hat> = scale * dot(2*y,z) / 2`; the stored
/// correction then yields the Extended-RaBitQ estimate.
///
/// For an odd dimension, the unused low nibble must be canonical zero padding.
///
/// # Errors
///
/// Returns [`QuantError::CodeLength`] for a dimension mismatch and
/// [`QuantError::NonZeroPadding`] for non-canonical trailing fields.
pub fn est_dot_bit4(
    query: &Bit4Query,
    codes: &[u8],
    factors: Bit4Factors,
) -> Result<f32, QuantError> {
    match est_dot_bit4_controlled(query, codes, factors, &mut |_| {
        Ok::<_, std::convert::Infallible>(())
    }) {
        Ok(value) => Ok(value),
        Err(Bit4ControlError::Quant(error)) => Err(error),
        Err(Bit4ControlError::Memory) => Err(QuantError::Allocation),
        Err(Bit4ControlError::Control(never)) => match never {},
    }
}

/// Estimates dot products for contiguous packed four-bit rows in one dispatch.
///
/// `factors` and `out` contain one entry per row. Row validation is performed
/// before entering the dispatched microkernel, and the scalar table remains
/// the bit-exact behavioral oracle on machines without the selected SIMD tier.
///
/// # Errors
///
/// Returns [`QuantError::OutputLength`] for a factor/output count mismatch,
/// [`QuantError::CodeLength`] for a batch byte-length mismatch, or
/// [`QuantError::NonZeroPadding`] for a non-canonical odd-dimension row.
pub fn est_dot_bit4_batch(
    query: &Bit4Query,
    codes: &[u8],
    factors: &[Bit4Factors],
    out: &mut [f32],
) -> Result<(), QuantError> {
    if out.len() != factors.len() {
        return Err(QuantError::OutputLength {
            expected: factors.len(),
            actual: out.len(),
        });
    }
    let row_bytes = query.codes.len().div_ceil(CODES_PER_BYTE);
    let expected = row_bytes.saturating_mul(factors.len());
    if codes.len() != expected {
        return Err(QuantError::CodeLength {
            expected,
            actual: codes.len(),
        });
    }
    if !query.codes.len().is_multiple_of(CODES_PER_BYTE) {
        for row in codes.chunks_exact(row_bytes) {
            validate_padding(row, query.codes.len())?;
        }
    }
    score_bit4_prepared_batch(
        &query.codes,
        query.code_sum,
        query.scale_half,
        codes,
        query.codes.len(),
        factors,
        out,
    );
    Ok(())
}

fn validate_input(v: &[f32], output_len: usize) -> Result<(), QuantError> {
    validate_vector(v)?;
    let expected = v.len().div_ceil(CODES_PER_BYTE);
    if output_len != expected {
        return Err(QuantError::OutputLength {
            expected,
            actual: output_len,
        });
    }
    Ok(())
}

fn validate_vector_shape(v: &[f32]) -> Result<(), QuantError> {
    if v.is_empty() {
        return Err(QuantError::EmptyVector);
    }
    if v.len() > MAX_DOT_I8_DIMENSION {
        return Err(QuantError::DimensionTooLarge {
            actual: v.len(),
            maximum: MAX_DOT_I8_DIMENSION,
        });
    }
    Ok(())
}

fn validate_vector(v: &[f32]) -> Result<(), QuantError> {
    validate_vector_shape(v)?;
    if let Some((index, _)) = v.iter().enumerate().find(|(_, value)| !value.is_finite()) {
        return Err(QuantError::NonFinite { index });
    }
    Ok(())
}

fn unpack(codes: &[u8], dimension: usize) -> Result<Vec<i8>, QuantError> {
    validate_code(codes, dimension)?;

    let mut unpacked = Vec::with_capacity(dimension);
    for &byte in codes {
        for shift in [4_u32, 0] {
            if unpacked.len() == dimension {
                break;
            }
            let unsigned = (byte >> shift) & 0b1111;
            unpacked.push((i16::from(unsigned) * 2 - 15) as i8);
        }
    }
    Ok(unpacked)
}

fn validate_code(codes: &[u8], dimension: usize) -> Result<(), QuantError> {
    let expected = dimension.div_ceil(CODES_PER_BYTE);
    if codes.len() != expected {
        return Err(QuantError::CodeLength {
            expected,
            actual: codes.len(),
        });
    }
    validate_padding(codes, dimension)?;
    Ok(())
}

fn validate_padding(codes: &[u8], dimension: usize) -> Result<(), QuantError> {
    if dimension.is_multiple_of(CODES_PER_BYTE) {
        return Ok(());
    }
    let mask = 0x0f;
    if let Some(&byte) = codes.last()
        && byte & mask != 0
    {
        return Err(QuantError::NonZeroPadding { byte, mask });
    }
    Ok(())
}

fn normalized_score_squared(numerator: f64, norm_squared: f64) -> f64 {
    numerator * numerator / norm_squared
}

fn finite_f32(value: f64) -> f32 {
    value.clamp(f64::from(f32::MIN), f64::from(f32::MAX)) as f32
}

fn apply_event(event: &CriticalValue, numerator: &mut f64, norm_squared: &mut f64) {
    *numerator += event.magnitude;
    *norm_squared += 2.0 * f64::from(event.level);
}