zeppelin_embed/quant/
int8.rs

1//! Per-vector affine signed-byte quantization.
2//!
3//! Each row stores one signed byte per coordinate plus an eight-byte
4//! `(scale, offset)` tail. A code reconstructs as `scale * code + offset`,
5//! with the row minimum and maximum mapped across `[-127, 127]`. The mapping
6//! is training-free: every row derives its own two scalars and depends on no
7//! corpus calibration or learned state.
8//!
9//! A future quality option may replace extrema with clipped quantiles. Apache
10//! Lucene's [`ScalarQuantizer`](https://lucene.apache.org/core/10_4_0/core/org/apache/lucene/util/quantization/ScalarQuantizer.html)
11//! is the reference for confidence-interval quantiles. That policy is not part
12//! of this v1 codec because clipping changes the deterministic error contract.
13
14use crate::kernels::{MAX_DOT_I8_DIMENSION, dot_i8};
15
16use super::QuantError;
17
18/// Borrowed signed-byte row and its per-vector affine reconstruction factors.
19#[derive(Clone, Copy, Debug, PartialEq)]
20pub struct Int8Vec<'a> {
21    /// One signed code per source coordinate.
22    pub codes: &'a [i8],
23    /// Reconstruction step between adjacent signed codes.
24    pub scale: f32,
25    /// Reconstructed value at code zero.
26    pub offset: f32,
27}
28
29/// Query-side signed-byte representation prepared once and reused per row.
30///
31/// The query uses a symmetric zero-centered map, so a per-row affine candidate
32/// can be scored from one native i8 dot product plus the candidate offset times
33/// the precomputed query-code sum. No query quantization or expansion occurs
34/// inside the row loop.
35#[derive(Clone, Debug, PartialEq)]
36pub struct Int8Query {
37    codes: Vec<i8>,
38    scale: f64,
39    code_sum: i32,
40}
41
42impl Int8Query {
43    /// Returns the query coordinate count.
44    #[must_use]
45    pub fn len(&self) -> usize {
46        self.codes.len()
47    }
48
49    /// Returns whether the query has zero coordinates.
50    #[must_use]
51    pub fn is_empty(&self) -> bool {
52        self.codes.is_empty()
53    }
54
55    /// Returns the prepared code bytes, scale, and code sum exactly as
56    /// consumed by the affine dot-product scorer.
57    #[cfg(any(test, feature = "test-seams"))]
58    #[doc(hidden)]
59    #[must_use]
60    pub fn observation_parts(&self) -> (&[i8], f64, i32) {
61        (&self.codes, self.scale, self.code_sum)
62    }
63
64    #[cfg(test)]
65    pub(crate) fn codes(&self) -> &[i8] {
66        &self.codes
67    }
68
69    pub(crate) fn score_integer_dot(
70        &self,
71        integer_dot: i32,
72        row_scale: f32,
73        row_offset: f32,
74    ) -> f32 {
75        (self.scale
76            * (f64::from(row_scale) * f64::from(integer_dot)
77                + f64::from(row_offset) * f64::from(self.code_sum))) as f32
78    }
79}
80
81/// Quantizes one finite row to signed bytes using a per-vector affine map.
82///
83/// A non-constant row reconstructs coordinate `i` as
84/// `scale * out[i] + offset`, where `scale = (max - min) / 254`. A constant
85/// row uses `scale = 0`, zero codes, and the constant as its offset, so the
86/// edge case is exact. Positive and negative zero are accepted and reconstruct
87/// to numerical zero.
88///
89/// # Errors
90///
91/// Returns [`QuantError::EmptyVector`] for an empty input,
92/// [`QuantError::DimensionTooLarge`] above the signed-byte kernel limit,
93/// [`QuantError::NonFinite`] for the first NaN or infinity, and
94/// [`QuantError::OutputLength`] unless `out.len() == v.len()`. Validation
95/// completes before `out` is modified.
96pub fn quantize_int8(v: &[f32], out: &mut [i8]) -> Result<(f32, f32), QuantError> {
97    validate_vector(v)?;
98    if out.len() != v.len() {
99        return Err(QuantError::OutputLength {
100            expected: v.len(),
101            actual: out.len(),
102        });
103    }
104
105    let (minimum, maximum) = extrema(v);
106    if minimum == maximum {
107        out.fill(0);
108        return Ok((0.0, minimum));
109    }
110
111    let range = f64::from(maximum) - f64::from(minimum);
112    let mut scale = (range / 254.0) as f32;
113    if scale == 0.0 {
114        scale = f32::from_bits(1);
115    }
116    let offset = ((f64::from(maximum) + f64::from(minimum)) * 0.5) as f32;
117    let scale_f64 = f64::from(scale);
118    let offset_f64 = f64::from(offset);
119    for (&value, code) in v.iter().zip(out.iter_mut()) {
120        let mapped = ((f64::from(value) - offset_f64) / scale_f64)
121            .round()
122            .clamp(-127.0, 127.0);
123        *code = mapped as i8;
124    }
125    Ok((scale, offset))
126}
127
128/// Reconstructs one affine signed-byte row into caller-owned storage.
129///
130/// # Errors
131///
132/// Returns [`QuantError::OutputLength`] when the output and code lengths
133/// differ.
134pub fn dequantize_int8(encoded: Int8Vec<'_>, out: &mut [f32]) -> Result<(), QuantError> {
135    if out.len() != encoded.codes.len() {
136        return Err(QuantError::OutputLength {
137            expected: encoded.codes.len(),
138            actual: out.len(),
139        });
140    }
141    for (&code, value) in encoded.codes.iter().zip(out.iter_mut()) {
142        *value = (f64::from(encoded.scale) * f64::from(code) + f64::from(encoded.offset)) as f32;
143    }
144    Ok(())
145}
146
147/// Prepares one full-precision query for repeated affine Int8 row scoring.
148///
149/// Coordinates are symmetrically rounded into `[-127, 127]`. This work and
150/// the signed-code sum happen once per query; [`dot_int8_query`] performs no
151/// allocation and no per-row query conversion.
152///
153/// # Errors
154///
155/// Returns [`QuantError::EmptyVector`], [`QuantError::DimensionTooLarge`], or
156/// [`QuantError::NonFinite`] under the same policy as [`quantize_int8`].
157pub fn prepare_int8_query(q: &[f32]) -> Result<Int8Query, QuantError> {
158    validate_vector(q)?;
159    #[cfg(any(test, feature = "test-seams"))]
160    super::QUERY_PREPARATIONS.with(|calls| {
161        if let Some(calls) = calls.borrow_mut().as_mut() {
162            calls.int8.push(q.len());
163        }
164    });
165    let maximum = q.iter().map(|value| value.abs()).fold(0.0_f32, f32::max);
166    if maximum == 0.0 {
167        return Ok(Int8Query {
168            codes: vec![0_i8; q.len()],
169            scale: 0.0,
170            code_sum: 0,
171        });
172    }
173
174    let scale = f64::from(maximum) / 127.0;
175    let mut code_sum = 0_i32;
176    let codes = q
177        .iter()
178        .map(|&value| {
179            let code = (f64::from(value) / scale).round().clamp(-127.0, 127.0) as i8;
180            code_sum += i32::from(code);
181            code
182        })
183        .collect();
184    Ok(Int8Query {
185        codes,
186        scale,
187        code_sum,
188    })
189}
190
191/// Estimates a dot product against one affine Int8 row.
192///
193/// If `q_hat = query_scale * query_codes` and
194/// `v_hat = row_scale * row_codes + row_offset`, expansion gives
195/// `query_scale * (row_scale * dot(query_codes, row_codes) +
196/// row_offset * sum(query_codes))`. The native task-03 i8 dot kernel supplies
197/// the only dimension-wide operation.
198///
199/// # Errors
200///
201/// Returns [`QuantError::CodeLength`] when the prepared query and stored row
202/// dimensions differ.
203pub fn dot_int8_query(query: &Int8Query, row: Int8Vec<'_>) -> Result<f32, QuantError> {
204    if row.codes.len() != query.codes.len() {
205        return Err(QuantError::CodeLength {
206            expected: query.codes.len(),
207            actual: row.codes.len(),
208        });
209    }
210    let integer_dot = dot_i8(&query.codes, row.codes);
211    Ok(query.score_integer_dot(integer_dot, row.scale, row.offset))
212}
213
214fn validate_vector(v: &[f32]) -> Result<(), QuantError> {
215    if v.is_empty() {
216        return Err(QuantError::EmptyVector);
217    }
218    if v.len() > MAX_DOT_I8_DIMENSION {
219        return Err(QuantError::DimensionTooLarge {
220            actual: v.len(),
221            maximum: MAX_DOT_I8_DIMENSION,
222        });
223    }
224    if let Some((index, _)) = v.iter().enumerate().find(|(_, value)| !value.is_finite()) {
225        return Err(QuantError::NonFinite { index });
226    }
227    Ok(())
228}
229
230fn extrema(v: &[f32]) -> (f32, f32) {
231    v.iter().copied().fold(
232        (f32::INFINITY, f32::NEG_INFINITY),
233        |(minimum, maximum), value| (minimum.min(value), maximum.max(value)),
234    )
235}
zeppelin_embed/quant/int8.rs

zeppelin_embed/quant/
int8.rs
