use super::compact::QuantizationKind;
use crate::error::SearchError;
const SCALE: f32 = 127.0;
#[derive(Debug)]
pub(crate) enum QuantizedStorage {
BFloat16(Vec<u16>),
Float16(Vec<u16>),
Float8E4M3(Vec<u8>),
Int8(Vec<i8>),
Binary(Vec<u64>),
}
impl QuantizedStorage {
pub(super) fn with_capacity(
kind: QuantizationKind,
elements: usize,
words: usize,
) -> Result<Self, SearchError> {
match kind {
QuantizationKind::BFloat16 => {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| SearchError::AllocationFailed)?;
Ok(Self::BFloat16(values))
}
QuantizationKind::Float16 => {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| SearchError::AllocationFailed)?;
Ok(Self::Float16(values))
}
QuantizationKind::Float8E4M3 => {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| SearchError::AllocationFailed)?;
Ok(Self::Float8E4M3(values))
}
QuantizationKind::Int8 => {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| SearchError::AllocationFailed)?;
Ok(Self::Int8(values))
}
QuantizationKind::Binary => {
let mut values = Vec::new();
values
.try_reserve_exact(words)
.map_err(|_| SearchError::AllocationFailed)?;
Ok(Self::Binary(values))
}
}
}
pub(super) fn push_vector(
&mut self,
vector: &[f32],
scale: f32,
words_per_vector: usize,
) -> Result<f32, SearchError> {
match self {
Self::BFloat16(values) => {
values.extend(vector.iter().map(|value| to_bf16(*value * scale)));
}
Self::Float16(values) => {
values.extend(vector.iter().map(|value| to_f16(*value * scale)));
}
Self::Float8E4M3(values) => {
values.extend(vector.iter().map(|value| to_f8_e4m3(*value * scale)));
}
Self::Int8(values) => {
values.extend(
vector
.iter()
.map(|value| quantize_component(*value * scale)),
);
}
Self::Binary(values) => {
for word in 0..words_per_vector {
values.push(binary_word(vector, word));
}
return Ok(1.0);
}
}
let start = self
.len()
.checked_sub(vector.len())
.ok_or(SearchError::CapacityOverflow)?;
Ok((start..self.len())
.map(|index| {
let value = self.value(index);
value * value
})
.sum())
}
pub(crate) fn len(&self) -> usize {
match self {
Self::BFloat16(values) | Self::Float16(values) => values.len(),
Self::Float8E4M3(values) => values.len(),
Self::Int8(values) => values.len(),
Self::Binary(values) => values.len(),
}
}
pub(super) fn value(&self, index: usize) -> f32 {
match self {
Self::BFloat16(values) => from_bf16(values[index]),
Self::Float16(values) => from_f16(values[index]),
Self::Float8E4M3(values) => from_f8_e4m3(values[index]),
Self::Int8(values) => f32::from(values[index]),
Self::Binary(_) => unreachable!("binary values use packed distance"),
}
}
pub(super) fn estimated_bytes(&self) -> usize {
match self {
Self::BFloat16(values) | Self::Float16(values) => {
values.capacity().saturating_mul(std::mem::size_of::<u16>())
}
Self::Float8E4M3(values) => values.capacity(),
Self::Int8(values) => values.capacity(),
Self::Binary(values) => values.capacity().saturating_mul(std::mem::size_of::<u64>()),
}
}
}
pub(super) fn binary_word(vector: &[f32], word_index: usize) -> u64 {
let start = word_index * 64;
vector
.iter()
.skip(start)
.take(64)
.enumerate()
.fold(0_u64, |word, (bit, value)| {
if *value >= 0.0 {
word | (1_u64 << bit)
} else {
word
}
})
}
#[allow(clippy::cast_possible_truncation)]
pub(super) fn quantize_component(value: f32) -> i8 {
(value * SCALE).round().clamp(-SCALE, SCALE) as i8
}
fn to_bf16(value: f32) -> u16 {
let bits = value.to_bits();
let rounding = 0x7fff_u32 + ((bits >> 16) & 1);
u16::try_from(bits.wrapping_add(rounding) >> 16).expect("upper f32 bits fit u16")
}
fn from_bf16(value: u16) -> f32 {
f32::from_bits(u32::from(value) << 16)
}
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::cast_sign_loss
)]
fn to_f16(value: f32) -> u16 {
let bits = value.to_bits();
let sign = ((bits >> 16) & 0x8000) as u16;
let exponent = ((bits >> 23) & 0xff) as i32;
let mantissa = bits & 0x7f_ffff;
if exponent == 0xff {
return sign | if mantissa == 0 { 0x7c00 } else { 0x7e00 };
}
let half_exponent = exponent - 127 + 15;
if half_exponent >= 0x1f {
return sign | 0x7c00;
}
if half_exponent <= 0 {
if half_exponent < -10 {
return sign;
}
let mantissa = mantissa | 0x80_0000;
let shift = 14 - half_exponent;
let rounded = (mantissa + (1_u32 << (shift - 1))) >> shift;
return sign | rounded as u16;
}
let rounded = mantissa + 0x1000;
if rounded & 0x80_0000 != 0 {
let next = half_exponent + 1;
return sign
| if next >= 0x1f {
0x7c00
} else {
(next as u16) << 10
};
}
sign | ((half_exponent as u16) << 10) | ((rounded >> 13) as u16)
}
fn from_f16(value: u16) -> f32 {
let sign = u32::from(value & 0x8000) << 16;
let exponent = u32::from((value >> 10) & 0x1f);
let mantissa = u32::from(value & 0x03ff);
let bits = match exponent {
0 if mantissa == 0 => sign,
0 => {
let mut mantissa = mantissa;
let mut exponent = 113_u32;
while mantissa & 0x0400 == 0 {
mantissa <<= 1;
exponent -= 1;
}
sign | (exponent << 23) | ((mantissa & 0x03ff) << 13)
}
0x1f => sign | 0x7f80_0000 | (mantissa << 13),
_ => sign | ((exponent + 112) << 23) | (mantissa << 13),
};
f32::from_bits(bits)
}
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn to_f8_e4m3(value: f32) -> u8 {
if value.is_nan() {
return 0x7f;
}
let sign = if value.is_sign_negative() { 0x80 } else { 0 };
let absolute = value.abs().min(240.0);
if absolute < 2_f32.powi(-9) {
return sign;
}
let exponent = absolute.log2().floor() as i32;
if exponent < -6 {
let mantissa = (absolute / 2_f32.powi(-9)).round().clamp(0.0, 7.0) as u8;
return sign | mantissa;
}
let mut biased = exponent + 7;
let base = 2_f32.powi(exponent);
let mut mantissa = ((absolute / base - 1.0) * 8.0).round() as i32;
if mantissa == 8 {
biased += 1;
mantissa = 0;
}
if biased >= 15 {
return sign | 0x77;
}
sign | ((biased as u8) << 3) | mantissa.clamp(0, 7) as u8
}
fn from_f8_e4m3(value: u8) -> f32 {
let sign = if value & 0x80 == 0 { 1.0 } else { -1.0 };
let exponent = (value >> 3) & 0x0f;
let mantissa = value & 0x07;
if exponent == 0 {
return sign * f32::from(mantissa) * 2_f32.powi(-9);
}
if exponent == 0x0f {
return f32::NAN;
}
sign * (1.0 + f32::from(mantissa) / 8.0) * 2_f32.powi(i32::from(exponent) - 7)
}