From 47db8328b5db8acd6d7d1530191b89b283dce6cb Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 14 Oct 2026 16:54:37 -0700 Subject: [PATCH 1/5] Checkpoint. --- diskann-inmem/integration/index/object.rs | 1 - diskann-inmem/integration/index/runner.rs | 50 ++++--- diskann-inmem/src/repr/internal/intrusive.rs | 6 - .../src/repr/internal/quantization/rerank.rs | 16 +-- diskann-inmem/src/repr/mod.rs | 14 +- diskann-inmem/src/repr/quantized/mod.rs | 1 + diskann-inmem/src/repr/quantized/spherical.rs | 135 ++++++++++++++++++ diskann-inmem/src/store/intrusive.rs | 15 +- .../src/spherical/quantizer.rs | 13 +- 9 files changed, 192 insertions(+), 59 deletions(-) create mode 100644 diskann-inmem/src/repr/quantized/spherical.rs diff --git a/diskann-inmem/integration/index/object.rs b/diskann-inmem/integration/index/object.rs index bac89d556b..560a4b4fc1 100644 --- a/diskann-inmem/integration/index/object.rs +++ b/diskann-inmem/integration/index/object.rs @@ -236,5 +236,4 @@ macro_rules! index { } index!({ T } repr::Full where T: repr::FullPrecision + FromSlice + AsDataType); -index!(repr::Spherical); index!(repr::Quantized); diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index b6968eb3c3..5b964bb91f 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -309,15 +309,6 @@ mod quantization { } } } - - impl From for inmem::repr::spherical::Rerank { - fn from(rerank: Rerank) -> Self { - match rerank { - Rerank::None => inmem::repr::spherical::Rerank::None, - Rerank::F16 => inmem::repr::spherical::Rerank::F16, - } - } - } } mod spherical { @@ -641,6 +632,7 @@ impl Test { alloc::GlobalAllocator, spherical::{PreScale, SphericalQuantizer}, }; + use inmem::repr::quantized::Config; use rand::SeedableRng; let DatasetView::F32(data) = data else { @@ -665,24 +657,36 @@ impl Test { )?; // Step 2: Associate it with the target bit-width. - let quantizer = match bits { - spherical::Bits::One => quantizer.as_quantizer::<1>()?, - spherical::Bits::Two => quantizer.as_quantizer::<2>()?, - spherical::Bits::Four => quantizer.as_quantizer::<4>()?, - }; - + let capacity = Capacity::new(data.nrows()); + let max_degree = MaxDegree::new(self.build.config.max_degree().get()); let start_point = rowmajor::Owned::row_vector(Box::from( ::compute_medoid(data), )); + let rerank = quantizer_rerank(rerank, self.data.metric); - // Step 3: Create the config. - let config = inmem::repr::Spherical::config( - quantizer, - Capacity::new(data.nrows()), - MaxDegree::new(self.build.config.max_degree().get()), - start_point, - rerank.into(), - )?; + let config = match bits { + spherical::Bits::One => Config::new( + quantizer.as_quantizer_impl::<1>()?, + capacity, + max_degree, + start_point, + rerank, + )?, + spherical::Bits::Two => Config::new( + quantizer.as_quantizer_impl::<2>()?, + capacity, + max_degree, + start_point, + rerank, + )?, + spherical::Bits::Four => Config::new( + quantizer.as_quantizer_impl::<4>()?, + capacity, + max_degree, + start_point, + rerank, + )?, + }; let index_config = self.build.config.clone(); diff --git a/diskann-inmem/src/repr/internal/intrusive.rs b/diskann-inmem/src/repr/internal/intrusive.rs index eb10c2597f..1d0fa3ed2a 100644 --- a/diskann-inmem/src/repr/internal/intrusive.rs +++ b/diskann-inmem/src/repr/internal/intrusive.rs @@ -80,12 +80,6 @@ impl<'a, D, P> ExpandBeam<'a, D, P> { } } - /// Return the [`crate::epoch::Guard`] of the embedded reader. - #[cfg(feature = "quantization")] - pub(in crate::repr) fn guard(&self) -> &crate::epoch::Guard<'a> { - self.reader.guard() - } - /// Wrap `self` in a `Box`. pub(in crate::repr) fn boxed(self) -> Box { Box::new(self) diff --git a/diskann-inmem/src/repr/internal/quantization/rerank.rs b/diskann-inmem/src/repr/internal/quantization/rerank.rs index 95dfb3d29c..9f133607d7 100644 --- a/diskann-inmem/src/repr/internal/quantization/rerank.rs +++ b/diskann-inmem/src/repr/internal/quantization/rerank.rs @@ -25,10 +25,6 @@ use crate::{ /// Choose how data is reranked. #[derive(Debug, Clone, Copy, PartialEq)] pub(in crate::repr) enum Rerank { - /// No reranking will be performed and no space for higher-precision vectors will be - /// allocated. - None, - /// Use 16-bit floating-point numbers to store the higher-precision representation. /// These will be used automatically during search to rerank candidates. F16, @@ -234,8 +230,12 @@ mod tests { const TEST_DIM: usize = 1; /// Create a reranker with the test dimension for the given configuration and metric. - fn make_store(rerank: Rerank, metric: Metric) -> (Reranker, Store>) { - let (reranker, config) = Reranker::new_with_config(rerank, metric, TEST_DIM).unwrap(); + fn make_store(rerank: Option, metric: Metric) -> (Reranker, Store>) { + let (reranker, config) = if let Some(rerank) = rerank { + Reranker::new_with_config(rerank, metric, TEST_DIM).unwrap(); + } else { + Reranker::disabled() + }; let store = Store::new( store::Layout::new(Capacity::new(10), MaxDegree::new(0), 0), @@ -255,7 +255,7 @@ mod tests { #[test] fn test_disabled_rerank() { - let (reranker, store) = make_store(Rerank::None, Metric::SquaredL2); + let (reranker, store) = make_store(None, Metric::SquaredL2); assert_matches!(reranker, Reranker::None); assert!( store.slots().slots().is_none(), @@ -275,7 +275,7 @@ mod tests { #[test] fn test_reranker_f16() { for metric in Metric::all() { - let (reranker, store) = make_store(Rerank::F16, metric); + let (reranker, store) = make_store(Some(Rerank::F16), metric); let mut map = Reference::new(TEST_DIM); // Insert the following values by logical id: diff --git a/diskann-inmem/src/repr/mod.rs b/diskann-inmem/src/repr/mod.rs index e144103c8d..0209576d3a 100644 --- a/diskann-inmem/src/repr/mod.rs +++ b/diskann-inmem/src/repr/mod.rs @@ -26,13 +26,13 @@ mod internal; pub mod full; pub use full::{Full, FullPrecision}; -#[cfg(feature = "quantization")] -#[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] -pub mod spherical; - -#[cfg(feature = "quantization")] -#[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] -pub use spherical::Spherical; +// #[cfg(feature = "quantization")] +// #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] +// pub mod spherical; +// +// #[cfg(feature = "quantization")] +// #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] +// pub use spherical::Spherical; #[cfg(feature = "quantization")] #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index 172719219d..1e3db109ec 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -30,6 +30,7 @@ use crate::{ // Built-in quantizers pub mod product; pub mod scalar; +pub mod spherical; /// Plug-in for a generalized quantizer compatible with [`Quantized`]. /// diff --git a/diskann-inmem/src/repr/quantized/spherical.rs b/diskann-inmem/src/repr/quantized/spherical.rs new file mode 100644 index 0000000000..d6f6d83b0e --- /dev/null +++ b/diskann-inmem/src/repr/quantized/spherical.rs @@ -0,0 +1,135 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann::{ANNError, ANNResult}; +use diskann_quantization::{ + alloc::{Allocator, GlobalAllocator, ScopedAllocator}, + spherical::iface, +}; + +use crate::{num::Bytes, repr::quantized}; + +impl quantized::Quantizer for T +where + T: iface::Quantizer + 'static, +{ + fn bytes(&self) -> Bytes { + Bytes::new(iface::Quantizer::bytes(self)) + } + fn dim(&self) -> usize { + iface::Quantizer::full_dim(self) + } + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { + iface::Quantizer::compress( + self, + data, + iface::OpaqueMut::new(quantized), + ScopedAllocator::global(), + ) + .map_err(ANNError::new) + } + fn distance<'a, V>(&'a self, visitor: V) -> V::Output + where + V: quantized::DistanceVisitor<'a>, + { + visitor.visit(iface::Quantizer::distance_computer_ref(self)) + } + fn insert_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: quantized::QueryDistanceVisitor<'a>, + { + let query_computer = iface::Quantizer::fused_query_computer( + self, + query, + iface::QueryLayout::SameAsData, + false, + GlobalAllocator, + ScopedAllocator::global(), + ) + .map_err(ANNError::new)?; + + Ok(visitor.visit(query_computer)) + } + fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: quantized::QueryDistanceVisitor<'a>, + { + let query_computer = iface::Quantizer::fused_query_computer( + self, + query, + iface::QueryLayout::FullPrecision, + true, + GlobalAllocator, + ScopedAllocator::global(), + ) + .map_err(ANNError::new)?; + + Ok(visitor.visit(query_computer)) + } +} + +impl quantized::QueryDistance for iface::QueryComputer +where + A: Allocator + std::fmt::Debug + Send + Sync, +{ + fn eval(&self, x: &[u8]) -> ANNResult { + use diskann_vector::PreprocessedDistanceFunction; + self.evaluate_similarity(iface::Opaque::new(x)) + .map_err(ANNError::new) + } +} + +impl quantized::Distance for &dyn iface::DynDistanceComputer { + fn eval(&self, x: &[u8], y: &[u8]) -> ANNResult { + self.evaluate(iface::Opaque::new(x), iface::Opaque::new(y)) + .map_err(ANNError::new) + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + use diskann::graph::test::synthetic::Grid; + use diskann_quantization::spherical::SupportedMetric; + use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; + use diskann_vector::distance::Distance as VDistance; + use hashbrown::HashMap; + + use crate::{ + counters::Counters, + num::{LogicalId, SlotId}, + repr::test::{Reference, test_expand_beam, test_prune}, + }; + + fn train_quantizer( + data: rowmajor::Ref<'_, f32>, + metric: SupportedMetric, + bits: Bits, + ) -> Poly { + use diskann_quantization::{algorithms::transforms, spherical}; + use rand::{SeedableRng, rngs::StdRng}; + + let q = spherical::SphericalQuantizer::train( + data.as_view(), + transforms::TransformKind::Null, + metric, + spherical::PreScale::None, + &mut StdRng::seed_from_u64(0x019823745), + GlobalAllocator, + ) + .unwrap(); + + match bits { + Bits::One => q.as_quantizer::<1>().unwrap(), + Bits::Two => q.as_quantizer::<2>().unwrap(), + Bits::Four => q.as_quantizer::<4>().unwrap(), + } + } +} diff --git a/diskann-inmem/src/store/intrusive.rs b/diskann-inmem/src/store/intrusive.rs index b0957127b9..8704078f4e 100644 --- a/diskann-inmem/src/store/intrusive.rs +++ b/diskann-inmem/src/store/intrusive.rs @@ -158,7 +158,7 @@ impl Intrusive { Reader { buffer: &self.buffer, unpadded: self.unpadded, - guard, + _guard: guard, } } @@ -263,12 +263,7 @@ impl slots::Slots for Intrusive { pub(crate) struct Reader<'a> { buffer: &'a Buffer, unpadded: Bytes, - - #[cfg_attr( - not(feature = "quantization"), - expect(unused, reason = "quantization uses this to share the guard") - )] - guard: epoch::Guard<'a>, + _guard: epoch::Guard<'a>, } impl<'a> Reader<'a> { @@ -402,12 +397,6 @@ impl<'a> Reader<'a> { pub(crate) fn bytes_plus_tag(&self) -> Bytes { self.unpadded } - - /// Return a reference to the contained [`epoch::Guard`]. - #[cfg(feature = "quantization")] - pub(crate) fn guard(&self) -> &epoch::Guard<'a> { - &self.guard - } } /// A [`slots::Exclusive`] for [`Intrusive`]. diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index d29e2ea5c7..47b167a8a8 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -441,6 +441,17 @@ where }) } + /// Construct an [`iface::Impl`] `self`. + pub fn as_quantizer_impl( + self, + ) -> Result, AllocatorError> + where + A: 'static, + iface::Impl: iface::Constructible + iface::Quantizer, + { + iface::Impl::::new(self) + } + /// Construct an [`iface::Quantizer`] trait object from `self`. pub fn as_quantizer( self, @@ -449,7 +460,7 @@ where A: 'static, iface::Impl: iface::Constructible + iface::Quantizer, { - let iface = iface::Impl::::new(self)?; + let iface = self.as_quantizer_impl()?; crate::poly!({ iface::Quantizer }, iface, GlobalAllocator) } } From 2e139258a1a45aeff8f57bdbb622d30b2044bc34 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 14 Oct 2026 17:40:27 -0700 Subject: [PATCH 2/5] Move `spherical` behind `Quantizer`. --- diskann-benchmark/src/index/inmem2.rs | 75 +- diskann-inmem/integration/index/runner.rs | 40 +- .../src/repr/internal/quantization/rerank.rs | 5 +- diskann-inmem/src/repr/mod.rs | 8 - diskann-inmem/src/repr/quantized/mod.rs | 64 +- diskann-inmem/src/repr/quantized/spherical.rs | 256 +++- diskann-inmem/src/repr/spherical.rs | 1034 ----------------- .../src/spherical/quantizer.rs | 6 +- 8 files changed, 338 insertions(+), 1150 deletions(-) delete mode 100644 diskann-inmem/src/repr/spherical.rs diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index 6ba51879d6..59f6961884 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -35,12 +35,7 @@ use diskann_inmem::{ repr::{self, Full, FullPrecision}, Provider, Strategy, }; -use diskann_quantization::{ - self as dq, - alloc::{GlobalAllocator, Poly}, - product::tables::BasicTable, - spherical::iface, -}; +use diskann_quantization::{self as dq, alloc::GlobalAllocator, product::tables::BasicTable}; use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::distance::Metric; use half::f16; @@ -498,8 +493,9 @@ impl Spherical { &self, data: rowmajor::Ref<'_, f32>, metric: Metric, - ) -> anyhow::Result> { + ) -> anyhow::Result { use diskann_quantization::{algorithms::transforms, spherical}; + use inmem::repr::quantized::Backend; use rand::SeedableRng; let mut rng = rand::rngs::StdRng::seed_from_u64(0x7); @@ -518,9 +514,9 @@ impl Spherical { )?; let quantizer = match self.bits { - SphericalBits::One => quantizer.as_quantizer::<1>()?, - SphericalBits::Two => quantizer.as_quantizer::<2>()?, - SphericalBits::Four => quantizer.as_quantizer::<4>()?, + SphericalBits::One => Backend::new(quantizer.as_quantizer_impl::<1>()?), + SphericalBits::Two => Backend::new(quantizer.as_quantizer_impl::<2>()?), + SphericalBits::Four => Backend::new(quantizer.as_quantizer_impl::<4>()?), }; Ok(quantizer) @@ -1044,17 +1040,17 @@ impl Benchmark for SphericalBuild { let num_points = data.nrows(); writeln!(output, "Loaded {num_points} points, dim={dim}")?; - let quantizer = spherical.train(data.as_view(), input.data.distance)?; - // Compute the medoid of the dataset as the single start point. + let capacity = Capacity::new(num_points); + let max_degree = MaxDegree::new(input.build.config.max_degree().get()); let start = StartPointStrategy::Medoid.compute(data.as_view())?; - let config = repr::spherical::Spherical::config( - quantizer, - Capacity::new(num_points), - MaxDegree::new(input.build.config.max_degree().get()), - start, - repr::spherical::Rerank::F16, - )?; + let rerank = repr::quantized::Rerank::F16 { + metric: input.data.distance.into(), + }; + + let backend = spherical.train(data.as_view(), input.data.distance)?; + + let config = repr::Quantized::config(backend, capacity, max_degree, start, rerank)?; let provider = Provider::<_, u32>::new(config)?; let index = Arc::new(DiskANNIndex::new( @@ -1329,7 +1325,10 @@ impl Benchmark for ScalarBuild { checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result<()> { - use repr::{quantized::scalar::Scalar, Quantized}; + use repr::{ + quantized::{scalar::Scalar, Backend}, + Quantized, + }; writeln!(output, "{input}\n")?; @@ -1357,37 +1356,15 @@ impl Benchmark for ScalarBuild { metric: input.data.distance.into(), }; - let config = match scalar.bits { - ScalarBits::One => Quantized::config( - Scalar::<1>::new(quantizer, metric), - capacity, - max_degree, - start, - rerank, - )?, - ScalarBits::Two => Quantized::config( - Scalar::<2>::new(quantizer, metric), - capacity, - max_degree, - start, - rerank, - )?, - ScalarBits::Four => Quantized::config( - Scalar::<4>::new(quantizer, metric), - capacity, - max_degree, - start, - rerank, - )?, - ScalarBits::Eight => Quantized::config( - Scalar::<8>::new(quantizer, metric), - capacity, - max_degree, - start, - rerank, - )?, + let backend = match scalar.bits { + ScalarBits::One => Backend::new(Scalar::<1>::new(quantizer, metric)), + ScalarBits::Two => Backend::new(Scalar::<2>::new(quantizer, metric)), + ScalarBits::Four => Backend::new(Scalar::<4>::new(quantizer, metric)), + ScalarBits::Eight => Backend::new(Scalar::<8>::new(quantizer, metric)), }; + let config = Quantized::config(backend, capacity, max_degree, start, rerank)?; + let provider = Provider::<_, u32>::new(config)?; let index = Arc::new(DiskANNIndex::new( input.build.config.clone(), diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index 5b964bb91f..6f88fe1884 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -754,7 +754,7 @@ impl Test { rerank: quantization::Rerank, ) -> anyhow::Result> { use diskann_quantization::{num::Positive, scalar as qscalar}; - use inmem::repr::quantized::{Quantized, scalar::Scalar}; + use inmem::repr::quantized::{Backend, Quantized, scalar::Scalar}; let DatasetView::F32(data) = data else { anyhow::bail!("scalar quantization only supports f32 data"); @@ -776,40 +776,18 @@ impl Test { inmem::repr::quantized::scalar::Metric::try_from(self.data.metric)? }; + let backend = match bits { + scalar::Bits::One => Backend::new(Scalar::<1>::new(quantizer, metric)), + scalar::Bits::Two => Backend::new(Scalar::<2>::new(quantizer, metric)), + scalar::Bits::Four => Backend::new(Scalar::<4>::new(quantizer, metric)), + scalar::Bits::Eight => Backend::new(Scalar::<8>::new(quantizer, metric)), + }; + let capacity = Capacity::new(data.nrows()); let max_degree = MaxDegree::new(self.build.config.max_degree().get()); let rerank = quantizer_rerank(rerank, self.data.metric); - let config = match bits { - scalar::Bits::One => Quantized::config( - Scalar::<1>::new(quantizer, metric), - capacity, - max_degree, - start_point, - rerank, - )?, - scalar::Bits::Two => Quantized::config( - Scalar::<2>::new(quantizer, metric), - capacity, - max_degree, - start_point, - rerank, - )?, - scalar::Bits::Four => Quantized::config( - Scalar::<4>::new(quantizer, metric), - capacity, - max_degree, - start_point, - rerank, - )?, - scalar::Bits::Eight => Quantized::config( - Scalar::<8>::new(quantizer, metric), - capacity, - max_degree, - start_point, - rerank, - )?, - }; + let config = Quantized::config(backend, capacity, max_degree, start_point, rerank)?; let index_config = self.build.config.clone(); Ok(finish(Provider::new(config)?, index_config)) diff --git a/diskann-inmem/src/repr/internal/quantization/rerank.rs b/diskann-inmem/src/repr/internal/quantization/rerank.rs index 9f133607d7..000c0cdddf 100644 --- a/diskann-inmem/src/repr/internal/quantization/rerank.rs +++ b/diskann-inmem/src/repr/internal/quantization/rerank.rs @@ -33,7 +33,6 @@ pub(in crate::repr) enum Rerank { impl Rerank { fn bytes_for(&self, dim: usize) -> Result { match self { - Self::None => Ok(Bytes::new(0)), Self::F16 => match dim.checked_mul(2) { Some(bytes) => Ok(Bytes::new(bytes)), None => Err(Overflow { kind: *self, dim }), @@ -43,7 +42,6 @@ impl Rerank { fn as_str(&self) -> &'static str { match self { - Self::None => "none", Self::F16 => "f16", } } @@ -80,7 +78,6 @@ impl Reranker { dim: usize, ) -> Result<(Self, Option), Overflow> { match rerank { - Rerank::None => Ok(Self::disabled()), Rerank::F16 => { let distance = Distance::new(metric.as_vector_metric(), Some(dim)); @@ -232,7 +229,7 @@ mod tests { /// Create a reranker with the test dimension for the given configuration and metric. fn make_store(rerank: Option, metric: Metric) -> (Reranker, Store>) { let (reranker, config) = if let Some(rerank) = rerank { - Reranker::new_with_config(rerank, metric, TEST_DIM).unwrap(); + Reranker::new_with_config(rerank, metric, TEST_DIM).unwrap() } else { Reranker::disabled() }; diff --git a/diskann-inmem/src/repr/mod.rs b/diskann-inmem/src/repr/mod.rs index 0209576d3a..e797e893d5 100644 --- a/diskann-inmem/src/repr/mod.rs +++ b/diskann-inmem/src/repr/mod.rs @@ -26,14 +26,6 @@ mod internal; pub mod full; pub use full::{Full, FullPrecision}; -// #[cfg(feature = "quantization")] -// #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] -// pub mod spherical; -// -// #[cfg(feature = "quantization")] -// #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] -// pub use spherical::Spherical; - #[cfg(feature = "quantization")] #[cfg_attr(docsrs, doc(cfg(feature = "quantization")))] pub mod quantized; diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index 1e3db109ec..fa10b0d570 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -109,6 +109,18 @@ pub trait Quantizer: std::fmt::Debug + Send + Sync + 'static { { self.query_distance(query, visitor) } + + /// Construct the [`Backend`] from `self`. + /// + /// This method is not implementable outside of `diskann-inmem`. + fn into_backend(self) -> Backend + where + Self: Sized, + { + Backend { + inner: inner::Quantizer::new(self), + } + } } /// Return the distance between two raw byte slices. @@ -141,6 +153,52 @@ pub trait QueryDistanceVisitor<'a> { D: QueryDistance + 'a; } +/// The opaque backend quantizer for [`Quantized`]. +/// +/// This is left opaque to allow for future optimization and customization and is +/// constructible from all types that implement [`Quantizer`] either through +/// [`Quantizer::into_backend`] or the blanket implementation of [`IntoBackend`]. +/// +/// Since this type is [`IntoBackend`] for itself, it can be used as a type-erased container +/// for multiple [`Quantizer`] implementations, which can simplify quantizer preparation. +pub struct Backend { + inner: inner::Quantizer, +} + +impl Backend { + /// Construct a new [`Backend`] using [`IntoBackend`]. + pub fn new(quantizer: T) -> Self + where + T: IntoBackend, + { + quantizer.into_backend() + } + + fn into_inner(self) -> inner::Quantizer { + self.inner + } +} + +/// Custom conversion to [`Backend`]. +pub trait IntoBackend { + fn into_backend(self) -> Backend; +} + +impl IntoBackend for Backend { + fn into_backend(self) -> Backend { + self + } +} + +impl IntoBackend for Q +where + Q: Quantizer, +{ + fn into_backend(self) -> Backend { + Quantizer::into_backend(self) + } +} + /// Configuration for [`Quantized`]. #[derive(Debug)] pub struct Config { @@ -186,10 +244,10 @@ impl Config { rerank: Rerank, ) -> Result where - T: Quantizer, + T: IntoBackend, { Self::new_inner( - inner::Quantizer::new(quantizer), + quantizer.into_backend().into_inner(), capacity, max_degree, start_points, @@ -383,7 +441,7 @@ impl Quantized { rerank: Rerank, ) -> Result where - T: Quantizer, + T: IntoBackend, { Config::new(quantizer, capacity, max_degree, start_points, rerank) } diff --git a/diskann-inmem/src/repr/quantized/spherical.rs b/diskann-inmem/src/repr/quantized/spherical.rs index d6f6d83b0e..fb10fb9cc9 100644 --- a/diskann-inmem/src/repr/quantized/spherical.rs +++ b/diskann-inmem/src/repr/quantized/spherical.rs @@ -18,9 +18,11 @@ where fn bytes(&self) -> Bytes { Bytes::new(iface::Quantizer::bytes(self)) } + fn dim(&self) -> usize { iface::Quantizer::full_dim(self) } + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { iface::Quantizer::compress( self, @@ -30,12 +32,14 @@ where ) .map_err(ANNError::new) } + fn distance<'a, V>(&'a self, visitor: V) -> V::Output where V: quantized::DistanceVisitor<'a>, { visitor.visit(iface::Quantizer::distance_computer_ref(self)) } + fn insert_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult where V: quantized::QueryDistanceVisitor<'a>, @@ -52,6 +56,7 @@ where Ok(visitor.visit(query_computer)) } + fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult where V: quantized::QueryDistanceVisitor<'a>, @@ -96,27 +101,33 @@ impl quantized::Distance for &dyn iface::DynDistanceComputer { mod tests { use super::*; - use diskann::graph::test::synthetic::Grid; use diskann_quantization::spherical::SupportedMetric; - use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; - use diskann_vector::distance::Distance as VDistance; - use hashbrown::HashMap; - - use crate::{ - counters::Counters, - num::{LogicalId, SlotId}, - repr::test::{Reference, test_expand_beam, test_prune}, + use diskann_utils::{ + assert_contains, + views::rowmajor::{self, Matrix, MatrixMut}, + }; + use diskann_vector::PreprocessedDistanceFunction; + use rand::{ + SeedableRng, + distr::{Distribution, Uniform}, + rngs::StdRng, }; - fn train_quantizer( + use crate::repr::quantized::test::DynQuantizer; + + const DIM: usize = 5; + const TRIALS: usize = 10; + + fn train_quantizer( data: rowmajor::Ref<'_, f32>, metric: SupportedMetric, - bits: Bits, - ) -> Poly { + ) -> iface::Impl + where + iface::Impl: iface::Constructible + iface::Quantizer, + { use diskann_quantization::{algorithms::transforms, spherical}; - use rand::{SeedableRng, rngs::StdRng}; - let q = spherical::SphericalQuantizer::train( + spherical::SphericalQuantizer::train( data.as_view(), transforms::TransformKind::Null, metric, @@ -124,12 +135,219 @@ mod tests { &mut StdRng::seed_from_u64(0x019823745), GlobalAllocator, ) - .unwrap(); + .unwrap() + .as_quantizer_impl::() + .unwrap() + } + + /// Integration tests cover the correctness of the distances. Here, check that the + /// adapter agrees with the spherical API and propagates errors. + fn test_spherical( + quantizer: &dyn DynQuantizer, + reference: &dyn iface::Quantizer, + rng: &mut StdRng, + distribution: Uniform, + ntrials: usize, + ctx: &dyn std::fmt::Display, + ) { + assert_eq!(quantizer.dim(), reference.full_dim(), "{}", ctx); + assert_eq!(quantizer.bytes().value(), reference.bytes(), "{}", ctx); + + let dim = quantizer.dim(); + let bytes = quantizer.bytes().value(); - match bits { - Bits::One => q.as_quantizer::<1>().unwrap(), - Bits::Two => q.as_quantizer::<2>().unwrap(), - Bits::Four => q.as_quantizer::<4>().unwrap(), + let mut x = vec![0.0f32; dim]; + let mut bx = vec![0u8; bytes]; + + let mut y = vec![0.0f32; dim]; + let mut by = vec![0u8; bytes]; + + let mut expected = vec![0u8; bytes]; + + for trial in 0..ntrials { + // Include a zero query to exercise query scaling at zero norm. + if trial == 0 { + x.fill(0.0); + } else { + x.iter_mut().for_each(|x| *x = distribution.sample(rng)); + } + y.iter_mut().for_each(|y| *y = distribution.sample(rng)); + + quantizer.quantize(&mut bx, &x).unwrap(); + quantizer.quantize(&mut by, &y).unwrap(); + + // Check compression. + reference + .compress( + &x, + iface::OpaqueMut::new(&mut expected), + ScopedAllocator::global(), + ) + .unwrap(); + assert_eq!(bx, expected, "trial = {trial} -- {ctx}"); + + // Check distance. + let expected = reference + .distance_computer_ref() + .evaluate(iface::Opaque::new(&bx), iface::Opaque::new(&by)) + .unwrap(); + + let distance = quantizer.distance(); + let from_distance = distance.eval(&bx, &by).unwrap(); + + assert_eq!(from_distance, expected, "trial = {trial} -- {ctx}"); + + for (layout, scale, actual) in [ + ( + iface::QueryLayout::SameAsData, + false, + quantizer.insert_distance(&x).unwrap(), + ), + ( + iface::QueryLayout::FullPrecision, + true, + quantizer.query_distance(&x).unwrap(), + ), + ] { + let expected = reference + .fused_query_computer( + &x, + layout, + scale, + GlobalAllocator, + ScopedAllocator::global(), + ) + .unwrap() + .evaluate_similarity(iface::Opaque::new(&by)) + .unwrap(); + + assert_eq!( + actual.eval(&by).unwrap(), + expected, + "trial = {trial}, layout = {layout} -- {ctx}" + ); + } } + + // Compression Errors + { + // Input size + let x = vec![1.0f32; dim - 1]; + let mut bx = vec![0u8; bytes]; + let err = quantizer.quantize(&mut bx, &x).unwrap_err().to_string(); + assert_contains!(err, "expected source vector to have length 5", "{}", ctx); + + let x = vec![1.0f32; dim + 1]; + let err = quantizer.quantize(&mut bx, &x).unwrap_err().to_string(); + assert_contains!(err, "expected source vector to have length 5", "{}", ctx); + + // Output size + let x = vec![1.0f32; dim]; + let mut bx = vec![0u8; bytes - 1]; + assert!(quantizer.quantize(&mut bx, &x).is_err(), "{}", ctx); + + let mut bx = vec![0u8; bytes + 1]; + assert!(quantizer.quantize(&mut bx, &x).is_err(), "{}", ctx); + + // NaN + let x = vec![f32::NAN; dim]; + let mut bx = vec![0u8; bytes]; + let err = quantizer.quantize(&mut bx, &x).unwrap_err().to_string(); + assert_contains!(err, "NaN"); + } + + // Distance + { + let d = quantizer.distance(); + for length in [bytes - 1, bytes + 1] { + let invalid = vec![0u8; length]; + assert!(d.eval(&bx, &invalid).is_err(), "{ctx}"); + assert!(d.eval(&invalid, &bx).is_err(), "{ctx}"); + } + assert!(d.eval(&bx, &by).is_ok(), "{ctx}"); + } + + // Query Distance + { + for length in [dim - 1, dim + 1] { + let x = vec![1.0f32; length]; + assert!(quantizer.query_distance(&x).is_err(), "{ctx}"); + } + + let x = vec![f32::NAN; dim]; + let err = quantizer.query_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "NaN"); + + let x = vec![0.0f32; dim]; + let d = quantizer.query_distance(&x).unwrap(); + for length in [bytes - 1, bytes + 1] { + let invalid = vec![0u8; length]; + assert!(d.eval(&invalid).is_err(), "{ctx}"); + } + assert!(d.eval(&bx).is_ok(), "{ctx}"); + } + + // Insert Distance + { + for length in [dim - 1, dim + 1] { + let x = vec![1.0f32; length]; + assert!(quantizer.insert_distance(&x).is_err(), "{ctx}"); + } + + let x = vec![f32::NAN; dim]; + let err = quantizer.insert_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "NaN"); + + let x = vec![0.0f32; dim]; + let d = quantizer.insert_distance(&x).unwrap(); + for length in [bytes - 1, bytes + 1] { + let invalid = vec![0u8; length]; + assert!(d.eval(&invalid).is_err(), "{ctx}"); + } + assert!(d.eval(&bx).is_ok(), "{ctx}"); + } + } + + fn run_test() + where + iface::Impl: iface::Constructible + iface::Quantizer, + { + let mut rng = StdRng::seed_from_u64(0x019823745); + let distribution = Uniform::new(-4.0f32, 4.0).unwrap(); + let mut data = rowmajor::Owned::from_element(TRIALS, DIM, 0.0); + data.as_mut_slice() + .iter_mut() + .for_each(|v| *v = distribution.sample(&mut rng)); + + for metric in [ + SupportedMetric::SquaredL2, + SupportedMetric::InnerProduct, + SupportedMetric::Cosine, + ] { + let quantizer = train_quantizer::(data.as_view(), metric); + test_spherical( + &quantizer, + &quantizer, + &mut rng, + distribution, + TRIALS, + &format_args!("metric = {metric:?}, bits = {N}"), + ); + } + } + + #[test] + fn test_one_bit() { + run_test::<1>(); + } + + #[test] + fn test_two_bit() { + run_test::<2>(); + } + + #[test] + fn test_four_bit() { + run_test::<4>(); } } diff --git a/diskann-inmem/src/repr/spherical.rs b/diskann-inmem/src/repr/spherical.rs deleted file mode 100644 index 6c94a6ccdf..0000000000 --- a/diskann-inmem/src/repr/spherical.rs +++ /dev/null @@ -1,1034 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! A quantized store for RabitQ style compressed vectors. - -use std::num::NonZeroUsize; - -use diskann::{ANNError, ANNResult, error::ErrorContext, utils::IntoUsize}; -use diskann_quantization::{ - alloc::{GlobalAllocator, Poly, ScopedAllocator}, - spherical::iface, -}; -use diskann_utils::{ - lazy_format, - views::rowmajor::{self, Matrix}, -}; -use thiserror::Error; - -use crate::{ - counters::LocalCounters, - num::{Bytes, Capacity, IdLimit, MaxDegree}, - prefetch, repr, - store::{ - self, Store, - cons::{self, Cons}, - intrusive::{self, Intrusive}, - optional::Optional, - simple::{self, Simple}, - }, -}; - -/// Choose how data is going to be reranked. -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum Rerank { - /// No reranking will be performed and no space for higher precision vectors will be - /// allocated. - None, - - /// Use 16-bit floating point numbers to store the higher precision representation. - /// These will be used automatically during search to rerank candidates. - F16, -} - -impl Rerank { - fn as_internal_rerank(&self) -> repr::internal::quantization::Rerank { - use repr::internal::quantization::Rerank as IRerank; - match self { - Self::None => IRerank::None, - Self::F16 => IRerank::F16, - } - } -} - -/// The configuration for a [`Spherical`] representation. -#[derive(Debug)] -pub struct Config { - /// The underlying quantizer for the compressed store. - quantizer: Poly, - /// The start points. These must have dimensions equal to `quantizer.full_dim()`. - start_points: rowmajor::Owned, - layout: store::Layout, - store: store::Config, - lookahead: Option, - rerank: repr::internal::quantization::Rerank, -} - -const DEFAULT_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(16).unwrap(); - -impl Config { - /// Create a new [`Config`]. Parameters will be used as described below: - /// - /// * `quantizer`: The [`iface::Quantizer`] used to compress vectors and compute - /// distances among compressed vectors. - /// - /// * `capacity`: The number points to allocate space for. - /// - /// * `max_degree`: The maximum degree of the internal graph. - /// - /// * `start_points`: The points to use as frozen start points in the index. - /// - /// * `rerank`: Whether or not reranking is enabled and if so, the representation of the - /// higher precision vectors. - /// - /// # Errors - /// - /// Errors under the following conditions: - /// - /// * `start_points.ncols() != quantizer.full_dim()`: The dimensionality of the start - /// points must agree with the quantizer. - /// - /// * `start_points.nrows() == 0`: Currently, empty start points are not supported. - /// - /// * The number of start points exceeds `u32::MAX`. - pub fn new( - quantizer: Poly, - capacity: Capacity, - max_degree: MaxDegree, - start_points: rowmajor::Owned, - rerank: Rerank, - ) -> Result { - let quantizer_dim = quantizer.full_dim(); - if quantizer_dim != start_points.ncols() { - return Err(ConfigError::dim_mismatch( - quantizer_dim, - start_points.ncols(), - )); - } - - if start_points.nrows() == 0 { - return Err(ConfigError::empty_start_points()); - } - - let num_start_points: u32 = start_points - .nrows() - .try_into() - .map_err(|_| ConfigError::too_many_start_points(start_points.nrows()))?; - - Ok(Self { - quantizer, - start_points, - layout: store::Layout::new(capacity, max_degree, num_start_points), - store: store::Config::default(), - lookahead: Some(DEFAULT_LOOKAHEAD), - rerank: rerank.as_internal_rerank(), - }) - } - - /// Override the [`store::Config`] for tailoring concurrency details. - pub fn store(mut self, config: store::Config) -> Self { - self.store = config; - self - } - - /// Set the prefetch lookahead. - /// - /// This controls how many iterations ahead in - /// [`diskann::graph::glue::SearchAccessor::expand_beam`] data is prefetched into the CPU - /// cache. Passing `None` disables prefetching. - pub fn prefetch(mut self, lookahead: Option) -> Self { - self.lookahead = lookahead; - self - } - - /// Build the [`Spherical`] from `self`. - pub fn build(self) -> ANNResult { - Spherical::new(self) - } -} - -/// Errors that can occur during the construction of [`Config`]. -#[derive(Debug, Error)] -#[error(transparent)] -pub struct ConfigError { - inner: ConfigErrorInner, -} - -diskann::convert_error!(ConfigError); - -impl ConfigError { - fn dim_mismatch(quantizer: usize, start_points: usize) -> Self { - Self { - inner: ConfigErrorInner::DimMismatch { - quantizer, - start_points, - }, - } - } - - fn empty_start_points() -> Self { - Self { - inner: ConfigErrorInner::EmptyStartPoints, - } - } - - fn too_many_start_points(num_start_points: usize) -> Self { - Self { - inner: ConfigErrorInner::TooManyStartPoints { num_start_points }, - } - } -} - -#[derive(Debug, Error)] -enum ConfigErrorInner { - #[error( - "quantizer configured for dimension {} but given start points have dimension {}", - quantizer, - start_points - )] - DimMismatch { - quantizer: usize, - start_points: usize, - }, - #[error("at least one start point must be provided")] - EmptyStartPoints, - #[error("{} start points exceeds u32::MAX", num_start_points)] - TooManyStartPoints { num_start_points: usize }, -} - -impl repr::RepresentationConfig for Config { - type Representation = Spherical; - - fn build(self) -> ANNResult { - ::build(self) - } -} - -/// Spherically quantized data representation. -#[derive(Debug)] -pub struct Spherical { - store: Store>>, - quantizer: Poly, - // These values come directly from `quantizer`, but are hoisted out to avoid a - // trait-object function call when accessing. - full_dim: usize, - lookahead: Option, - reranker: repr::internal::quantization::Reranker, -} - -impl Spherical { - /// Initialize a [`Config`] for this representation. - /// - /// See also: [`Config::new`]. - /// - /// # Errors - /// - /// Returns the errors described by [`Config::new`]. - pub fn config( - quantizer: Poly, - capacity: Capacity, - max_degree: MaxDegree, - start_points: rowmajor::Owned, - rerank: Rerank, - ) -> Result { - Config::new(quantizer, capacity, max_degree, start_points, rerank) - } - - /// Create a new full-precision representation from `config`. - /// - /// See: [`Config::build`]. - fn new(config: Config) -> ANNResult { - let Config { - quantizer, - start_points, - layout, - store, - lookahead, - rerank, - } = config; - - let full_dim = quantizer.full_dim(); - let (reranker, rerank_config) = repr::internal::quantization::Reranker::new_with_config( - rerank, - quantizer.metric().into(), - full_dim, - )?; - - let slots = cons::Config::new( - Intrusive::config(Bytes::new(quantizer.bytes())), - rerank_config, - ); - - let store = Store::new(layout, store, slots)?; - - let this = Self { - store, - quantizer, - full_dim, - lookahead, - reranker, - }; - - // Initialize start points. - let num_start_points = start_points.nrows(); - for (i, (slot_index, row)) in - std::iter::zip(this.store.frozen(), start_points.rows()).enumerate() - { - #[expect( - clippy::expect_used, - reason = "failing this is an internal, unrecoverable bug" - )] - let mut slot = this - .store - .slot(slot_index) - .expect("internal store should leave frozen points available for writing"); - - this.set(row, slot.data()).with_context(|| { - lazy_format!(move, "on start point {} of {}", i + 1, num_start_points) - })?; - - slot.freeze(); - } - - Ok(this) - } - - /// Return the dimension of the data held within `self`. - pub fn dim(&self) -> usize { - self.full_dim - } - - /// * Attempt to compress `v` into the [`cons::Exclusive::first`] position. - /// * If [`cons::Exclusive::second`] is occupied, use `self.reranker` to store data - /// into that slot. - fn set( - &self, - v: &[f32], - slot: &mut cons::Exclusive, Option>>, - ) -> ANNResult<()> { - self.quantizer - .compress( - v, - iface::OpaqueMut::new(slot.first().as_mut_slice()), - ScopedAllocator::global(), - ) - .map_err(ANNError::new)?; - - self.reranker.store(v, slot.second()); - - Ok(()) - } - - /// Shared entry point for creating [`crate::provider::SearchAccessor`]. - /// - /// See: [`AccessorArgs`]. - fn create_accessor<'a>( - &'a self, - query: &'a [f32], - provider: &'a (dyn std::any::Any + Send + Sync), - counters: LocalCounters<'a>, - args: AccessorArgs, - ) -> ANNResult> { - let AccessorArgs { - layout, - allow_rescale, - rerank_if_enabled, - } = args; - - // Create the query computer. - // - // We do this first because this is one of the most likely things to fail since it - // operates on largely untrusted data. If it does fail, we save the work of acquiring - // epoch guards etc. - let query_computer = self - .quantizer - .fused_query_computer( - query, - layout, - allow_rescale, - GlobalAllocator, - ScopedAllocator::global(), - ) - .map_err(ANNError::new)?; - - // Computer is good - time to make `ExpandBeam` and `PostProcess` (if requested). - let (expand_beam, post_process) = self.store.guard(|cons, guard| { - // TODO: Tailor prefetching to the number of cachelines. - // - // Inlining distance functions will require work in `diskann-quantization`, so - // we can at least optimize prefetching. - let expand_beam = repr::internal::intrusive::ExpandBeam::new( - cons.first().reader(guard), - query_computer, - prefetch::Loop::new(), - self.lookahead, - ); - - let post_process = rerank_if_enabled - .then(|| { - self.reranker - .post_process(query, expand_beam.guard(), cons.second(), &counters) - }) - .flatten(); - - (expand_beam, post_process) - })?; - - Ok(crate::provider::SearchAccessor::new( - self.store.neighbors(), - expand_beam.boxed(), - post_process, - provider, - self.store.frozen(), - counters, - )) - } - - #[cfg(test)] - fn quantizer(&self) -> &dyn iface::Quantizer { - &*self.quantizer - } -} - -/// Arguments to [`Spherical::create_accessor`]. -#[derive(Debug)] -struct AccessorArgs { - /// The [`iface::QueryLayout`] to use when compressing the query. - layout: iface::QueryLayout, - - /// Whether or not the query vector can be rescaled. - /// - /// This only applies to query search. For searches used as part of insertion, rescaling - /// should be disabled to ensure distance are compatible with distances computed during - /// prune. - allow_rescale: bool, - - /// Create a [`repr::PostProcess`] for reranking if enabled in the parent struct. - /// - /// Since insert does not use a reranking step, this can be left as `false` to save - /// some processing and allocations. - rerank_if_enabled: bool, -} - -repr::internal::macros::representation!(Spherical); - -repr::internal::macros::set_guard!( - /// A [`repr::Guard`] for [`Spherical`]. - for<'a> cons::Exclusive, Option>> -); - -impl repr::Set<&[f32]> for Spherical { - type Guard<'a> = Guard<'a>; - - fn set(&self, v: &[f32]) -> ANNResult> { - // Easy check to reject invalid vectors before acquiring an epoch guard. - if v.len() != self.full_dim { - let vlen = v.len(); - let full_dim = self.full_dim; - - return Err(ANNError::message(lazy_format!( - move, - "vector dim {} does not match quantizer dim {}", - vlen, - full_dim - ))); - } - - let mut slot = self - .store - .acquire() - .ok_or_else(|| ANNError::message("could not allocate a new slot"))?; - - self.set(v, slot.data())?; - Ok(Guard::new(slot)) - } -} - -impl repr::Search for Spherical { - type Query<'a> = &'a [f32]; - - fn search_accessor<'a>( - &'a self, - query: &'a [f32], - provider: &'a (dyn std::any::Any + Send + Sync), - counters: LocalCounters<'a>, - ) -> ANNResult> { - self.create_accessor( - query, - provider, - counters, - AccessorArgs { - layout: iface::QueryLayout::FullPrecision, - allow_rescale: true, - rerank_if_enabled: true, - }, - ) - } -} - -impl repr::Insert for Spherical { - fn insert_search_accessor<'a>( - &'a self, - query: Self::Query<'a>, - provider: &'a (dyn std::any::Any + Send + Sync), - counters: LocalCounters<'a>, - ) -> ANNResult> { - self.create_accessor( - query, - provider, - counters, - AccessorArgs { - layout: iface::QueryLayout::SameAsData, - allow_rescale: false, - rerank_if_enabled: false, - }, - ) - } - - fn prune_accessor<'a>( - &'a self, - counters: LocalCounters<'a>, - ) -> ANNResult> { - let distance = self.quantizer.distance_computer_ref(); - let reader = self - .store - .guard(|slots, guard| slots.first().reader(guard))?; - - let prune = repr::internal::intrusive::Prune::new(reader, distance); - Ok(crate::provider::PruneAccessor::new( - prune.boxed(), - self.store.neighbors(), - counters, - )) - } -} - -//----------------// -// Trait Adaptors // -//----------------// - -impl repr::internal::RawQueryDistance for iface::QueryComputer -where - A: diskann_quantization::alloc::Allocator + std::fmt::Debug + Send + Sync, -{ - type Error = ANNError; - - fn eval(&self, x: &[u8]) -> Result { - use diskann_vector::PreprocessedDistanceFunction; - self.evaluate_similarity(iface::Opaque::new(x)) - .map_err(ANNError::new) - } -} - -impl repr::internal::RawDistance for &dyn iface::DynDistanceComputer { - type Error = ANNError; - - fn eval(&self, x: &[u8], y: &[u8]) -> Result { - self.evaluate(iface::Opaque::new(x), iface::Opaque::new(y)) - .map_err(ANNError::new) - } -} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - use diskann::graph::test::synthetic::Grid; - use diskann_quantization::spherical::SupportedMetric; - use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; - use diskann_vector::distance::Distance as VDistance; - use hashbrown::HashMap; - - use crate::{ - counters::Counters, - num::{LogicalId, SlotId}, - repr::test::{Reference, test_expand_beam, test_prune}, - }; - - #[derive(Debug, Clone, Copy)] - enum Bits { - One, - Two, - Four, - } - - fn train_quantizer( - data: rowmajor::Ref<'_, f32>, - metric: SupportedMetric, - bits: Bits, - ) -> Poly { - use diskann_quantization::{algorithms::transforms, spherical}; - use rand::{SeedableRng, rngs::StdRng}; - - let q = spherical::SphericalQuantizer::train( - data.as_view(), - transforms::TransformKind::Null, - metric, - spherical::PreScale::None, - &mut StdRng::seed_from_u64(0x019823745), - GlobalAllocator, - ) - .unwrap(); - - match bits { - Bits::One => q.as_quantizer::<1>().unwrap(), - Bits::Two => q.as_quantizer::<2>().unwrap(), - Bits::Four => q.as_quantizer::<4>().unwrap(), - } - } - - /// See the description in [`make_test_repr`]. - const TEST_LIMIT: IdLimit = IdLimit::new(11); - - // For the spherical quantizer tests, we use the canonical grid layout, but center the - // data around the origin. - // - // This allows cosine distances to return reasonable results as the data is distributed - // around the origin. - // - // To keep computation mostly tractable, we only use a 2d grid with 9 points. So the - // coordinates are as follows: - // - // 0: [-1, -1] - // 1: [-1, 0] - // 2: [-1, +1] - // - // 3: [ 0, -1] - // 4: [ 0, 0] - // 5: [ 0, +1] - // - // 6: [+1, -1] - // 7: [+1, 0] - // 8: [+1, +1] - // - // We put two start points at `[-2, -2]` and `[+2, +2]`. - fn make_test_repr( - metric: SupportedMetric, - bits: Bits, - rerank: Rerank, - fill: bool, - ) -> (Spherical, Reference) { - let grid = Grid::Two; - let mut data = grid.data(3); - let offset = 1.5; - data.as_mut_slice().iter_mut().for_each(|v| *v -= offset); - - let quantizer = train_quantizer(data.as_view(), metric, bits); - - let mut start_points = rowmajor::Owned::from_element(2, data.ncols(), 0.0); - start_points.row_mut(0).fill(-2.0); - start_points.row_mut(1).fill(2.0); - - let config = Spherical::config( - quantizer, - Capacity::new(data.nrows()), - MaxDegree::new(0), - start_points.clone(), - rerank, - ) - .unwrap(); - - let spherical = config.build().unwrap(); - - assert_eq!(repr::Representation::id_limit(&spherical), TEST_LIMIT); - assert_eq!(spherical.dim(), usize::from(grid.dim())); - - let mut reference = Reference::new(grid.dim().into()); - - if fill { - for (i, row) in data.rows().enumerate() { - let guard = repr::Set::set(&spherical, row).unwrap(); - let id = repr::Guard::id(&guard); - - reference.insert(LogicalId(i), SlotId(id), row); - repr::Guard::publish(guard); - } - } - - // Insert frozen points. - for (slot, point) in spherical.store.frozen().zip(start_points.rows()) { - reference.insert(LogicalId(slot.into_usize()), SlotId(slot), point); - } - - (spherical, reference) - } - - /// Here - we don't test the whole `expand_beam` loop. That would be a waste of time and - /// is already tested by other code. - /// - /// Instead we use: - /// - /// * `ExpandBeam::evaluate` to verify that the correct type of distance computer is - /// created. - /// - /// * `Prune` to validate that the correct distance computer is made. - /// - /// * If reranking exists, that reranking works as expected. - fn test_distances( - spherical: &Spherical, - reference: &mut Reference, - metric: SupportedMetric, - rerank: Rerank, - ctx: &dyn std::fmt::Display, - ) { - use repr::internal::{RawDistance, RawQueryDistance}; - - assert_eq!( - repr::Representation::id_limit(spherical), - TEST_LIMIT, - "{ctx}" - ); - - let query = [10.0, -10.0]; - - // Generate a couple of data points in the dataset. - let i0 = LogicalId(2); - let s0 = reference.slot_id_for(i0); - - let i1 = LogicalId(5); - let s1 = reference.slot_id_for(i1); - - let i2 = LogicalId(10); - let s2 = reference.slot_id_for(i2); - - // Perform Deletes // - let i3_deleted = LogicalId(8); - let s3_deleted = reference.slot_id_for(i3_deleted); - - repr::Representation::retire(spherical, s3_deleted.value()).unwrap(); - reference.delete(i3_deleted); - - let i4_deleted = LogicalId(4); - let s4_deleted = reference.slot_id_for(i4_deleted); - - repr::Representation::retire(spherical, s4_deleted.value()).unwrap(); - reference.delete(i4_deleted); - - // Extract Values. - let v0 = &reference[i0]; - let v1 = &reference[i1]; - let v2 = &reference[i2]; - - // Compress the dataset vectors. - let quantizer = spherical.quantizer(); - let bytes = quantizer.bytes(); - let mut b0 = vec![0u8; bytes]; - let mut b1 = vec![0u8; bytes]; - let mut b2 = vec![0u8; bytes]; - - fn _compress(quantizer: &dyn iface::Quantizer, v: &[f32], b: &mut [u8]) { - let alloc = ScopedAllocator::global(); - quantizer - .compress(v, iface::OpaqueMut::new(b), alloc) - .unwrap(); - - // Make sure that compression actually did something. - assert!(!b.iter().all(|i| *i == 0)); - } - - _compress(quantizer, v0, &mut b0); - _compress(quantizer, v1, &mut b1); - _compress(quantizer, v2, &mut b2); - - // Insert - { - let computer = quantizer - .fused_query_computer( - &query, - iface::QueryLayout::SameAsData, - false, - GlobalAllocator, - ScopedAllocator::global(), - ) - .unwrap(); - - let distances = HashMap::from_iter([ - (s0, RawQueryDistance::eval(&computer, &b0).unwrap()), - (s1, RawQueryDistance::eval(&computer, &b1).unwrap()), - (s2, RawQueryDistance::eval(&computer, &b2).unwrap()), - ]); - - let counters = Counters::new(); - let mut sa = repr::Insert::insert_search_accessor( - spherical, - query.as_slice(), - &(), - counters.local(), - ) - .unwrap(); - - assert!( - sa.get_post_process().is_none(), - "insert accessors should not build a post-processor -- {}", - ctx, - ); - - test_expand_beam( - sa.get_expand_beam(), - TEST_LIMIT, - distances, - &[s0, s1, s3_deleted, s2, s4_deleted], - ctx, - ); - } - - // Search - { - let computer = quantizer - .fused_query_computer( - &query, - iface::QueryLayout::FullPrecision, - true, - GlobalAllocator, - ScopedAllocator::global(), - ) - .unwrap(); - - let distances = HashMap::from_iter([ - (s0, RawQueryDistance::eval(&computer, &b0).unwrap()), - (s1, RawQueryDistance::eval(&computer, &b1).unwrap()), - (s2, RawQueryDistance::eval(&computer, &b2).unwrap()), - ]); - - let counters = Counters::new(); - let mut sa = - repr::Search::search_accessor(spherical, query.as_slice(), &(), counters.local()) - .unwrap(); - - test_expand_beam( - sa.get_expand_beam(), - TEST_LIMIT, - distances, - &[s0, s1, s3_deleted, s2, s4_deleted], - ctx, - ); - - if rerank == Rerank::None { - assert!( - sa.get_post_process().is_none(), - "search accessors should not build a post-processor with reranking disabled -- {}", - ctx, - ); - } else { - let post_process = match sa.get_post_process() { - Some(post_process) => post_process, - None => panic!("expected a post processor -- {}", ctx), - }; - - let f = VDistance::::new( - repr::internal::quantization::Metric::from(metric).as_vector_metric(), - None, - ); - - let distances = HashMap::from_iter([ - (s0, f.call(&query, v0)), - (s1, f.call(&query, v1)), - (s2, f.call(&query, v2)), - ]); - - repr::internal::quantization::rerank::test_rerank( - post_process, - distances, - &[s0, s1, s3_deleted, s2, s4_deleted], - ctx, - ); - } - } - - // Prune - { - let computer = quantizer.distance_computer_ref(); - let b00 = RawDistance::eval(&computer, &b0, &b0).unwrap(); - let b01 = RawDistance::eval(&computer, &b0, &b1).unwrap(); - let b02 = RawDistance::eval(&computer, &b0, &b2).unwrap(); - - let b10 = b01; - let b11 = RawDistance::eval(&computer, &b1, &b1).unwrap(); - let b12 = RawDistance::eval(&computer, &b1, &b2).unwrap(); - - let b20 = b02; - let b21 = b12; - let b22 = RawDistance::eval(&computer, &b2, &b2).unwrap(); - - let distances = HashMap::from_iter([ - ((s0, s0), b00), - ((s0, s1), b01), - ((s0, s2), b02), - ((s1, s0), b10), - ((s1, s1), b11), - ((s1, s2), b12), - ((s2, s0), b20), - ((s2, s1), b21), - ((s2, s2), b22), - ]); - - let counters = Counters::new(); - let mut pa = repr::Insert::prune_accessor(spherical, counters.local()).unwrap(); - - test_prune( - pa.get_prune(), - distances, - &[s0, s1, s3_deleted, s2, s4_deleted], - ctx, - ); - } - } - - /// The happy-patch entry point. - /// - /// Note that this method does a grid of the supported parameters. Adding new values - /// to these parameters has a multiplicative effect on runtime. - /// - /// This mainly matters for Miri tests. Currently, the Miri test for this function takes - /// about 30 seconds. If it starts to take much longer, this test should be split into - /// multiple entry points for parallelism. - #[test] - fn test_spherical() { - let metrics = [ - SupportedMetric::SquaredL2, - SupportedMetric::InnerProduct, - SupportedMetric::Cosine, - ]; - - let bits = [Bits::One, Bits::Two, Bits::Four]; - let rerank = [Rerank::None, Rerank::F16]; - - for metric in metrics { - for bits in bits { - for rerank in rerank { - let (spherical, mut reference) = make_test_repr(metric, bits, rerank, true); - - test_distances( - &spherical, - &mut reference, - metric, - rerank, - &format_args!( - "metric = {:?}, bits = {:?}, rerank = {:?}", - metric, bits, rerank - ), - ); - } - } - } - } - - //-------------// - // Error Paths // - //-------------// - - #[test] - fn test_config_dim_mismatch() { - let data = rowmajor::Owned::from_element(2, 5, 1.0f32); - let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - - let start_points = rowmajor::Owned::from_element(1, 6, 0.0f32); // Wrong number of columns - let err = Spherical::config( - quantizer, - Capacity::new(10), - MaxDegree::new(0), - start_points, - Rerank::None, - ) - .unwrap_err(); - - let msg = err.to_string(); - assert_contains!( - msg, - "quantizer configured for dimension 5 but given start points have dimension 6" - ); - } - - #[test] - fn test_empty_start_points() { - let data = rowmajor::Owned::from_element(2, 5, 1.0f32); - let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - - let start_points = rowmajor::Owned::from_element(0, 5, 0.0f32); // Empty - let err = Spherical::config( - quantizer, - Capacity::new(10), - MaxDegree::new(0), - start_points, - Rerank::None, - ) - .unwrap_err(); - - let msg = err.to_string(); - assert_contains!(msg, "at least one start point must be provided",); - } - - #[test] - fn test_build_error_uncompressible_query() { - let data = rowmajor::Owned::from_element(2, 5, 1.0f32); - let quantizer = train_quantizer(data.as_view(), SupportedMetric::SquaredL2, Bits::One); - - let start_points = rowmajor::Owned::from_element(1, 5, f32::INFINITY); // Wrong number of columns - let config = Spherical::config( - quantizer, - Capacity::new(10), - MaxDegree::new(0), - start_points, - Rerank::None, - ) - .unwrap(); - - let err = repr::RepresentationConfig::build(config).unwrap_err(); - let msg = err.to_string(); - assert_contains!( - msg, - "query compression", - "we tried to compress a start point with infinites in it - this should error" - ); - - assert_contains!( - msg, - "1 of 1", - "error message should contain which query errored", - ); - } - - #[test] - fn test_set_capacity_exhaustion() { - let (spherical, _) = - make_test_repr(SupportedMetric::SquaredL2, Bits::One, Rerank::None, true); - - let err = repr::Set::set(&spherical, &[1.0, 2.0]).unwrap_err(); - assert_contains!(err.to_string(), "could not allocate a new slot",); - } - - #[test] - fn test_set_dim_mismatch() { - let (spherical, _) = - make_test_repr(SupportedMetric::SquaredL2, Bits::One, Rerank::None, false); - - let err = repr::Set::set(&spherical, &[1.0, 2.0, 3.0]).unwrap_err(); - assert_contains!( - err.to_string(), - "vector dim 3 does not match quantizer dim 2", - ); - } - - #[test] - fn test_insert_slot_recovery() { - let (spherical, _) = - make_test_repr(SupportedMetric::SquaredL2, Bits::One, Rerank::F16, true); - - // Free up one slot. - repr::Representation::retire(&spherical, 0).unwrap(); - - // Insert something that is incompressible. - let err = repr::Set::set(&spherical, &[f32::INFINITY, f32::INFINITY]).unwrap_err(); - assert_contains!(err.to_string(), "query compression"); - - // If we insert again, this should succeed. - let guard = repr::Set::set(&spherical, &[1.0, 2.0]).unwrap(); - assert_eq!(repr::Guard::id(&guard), 0); - } -} diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index 47b167a8a8..d8fed47fd8 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -387,8 +387,10 @@ where /// Rescale the argument `v` to be in the rough dynamic range of the training dataset. pub fn rescale(&self, v: &mut [f32]) { let norm = FastL2Norm.evaluate(&*v); - let m = self.mean_norm.into_inner() / norm; - v.iter_mut().for_each(|i| *i *= m); + if norm != 0.0 { + let m = self.mean_norm.into_inner() / norm; + v.iter_mut().for_each(|i| *i *= m); + } } /// Private helper function to do common data pre-processing. From 62d5043844304549025b98bccbba043c3146d462 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Tue, 20 Oct 2026 17:54:26 -0700 Subject: [PATCH 3/5] Cleanups. --- diskann-benchmark/src/index/inmem2.rs | 10 +++-- diskann-inmem/integration/index/runner.rs | 35 +++++---------- diskann-inmem/src/repr/quantized/spherical.rs | 45 +++++++++++++++---- 3 files changed, 54 insertions(+), 36 deletions(-) diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index 59f6961884..0b8690451f 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -495,7 +495,7 @@ impl Spherical { metric: Metric, ) -> anyhow::Result { use diskann_quantization::{algorithms::transforms, spherical}; - use inmem::repr::quantized::Backend; + use inmem::repr::quantized::{spherical::Spherical, Backend}; use rand::SeedableRng; let mut rng = rand::rngs::StdRng::seed_from_u64(0x7); @@ -514,9 +514,11 @@ impl Spherical { )?; let quantizer = match self.bits { - SphericalBits::One => Backend::new(quantizer.as_quantizer_impl::<1>()?), - SphericalBits::Two => Backend::new(quantizer.as_quantizer_impl::<2>()?), - SphericalBits::Four => Backend::new(quantizer.as_quantizer_impl::<4>()?), + SphericalBits::One => Backend::new(Spherical::new(quantizer.as_quantizer_impl::<1>()?)), + SphericalBits::Two => Backend::new(Spherical::new(quantizer.as_quantizer_impl::<2>()?)), + SphericalBits::Four => { + Backend::new(Spherical::new(quantizer.as_quantizer_impl::<4>()?)) + } }; Ok(quantizer) diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index 6f88fe1884..d66aea4c07 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -632,7 +632,7 @@ impl Test { alloc::GlobalAllocator, spherical::{PreScale, SphericalQuantizer}, }; - use inmem::repr::quantized::Config; + use inmem::repr::quantized::{Backend, Config, spherical::Spherical}; use rand::SeedableRng; let DatasetView::F32(data) = data else { @@ -664,30 +664,19 @@ impl Test { )); let rerank = quantizer_rerank(rerank, self.data.metric); - let config = match bits { - spherical::Bits::One => Config::new( - quantizer.as_quantizer_impl::<1>()?, - capacity, - max_degree, - start_point, - rerank, - )?, - spherical::Bits::Two => Config::new( - quantizer.as_quantizer_impl::<2>()?, - capacity, - max_degree, - start_point, - rerank, - )?, - spherical::Bits::Four => Config::new( - quantizer.as_quantizer_impl::<4>()?, - capacity, - max_degree, - start_point, - rerank, - )?, + let backend = match bits { + spherical::Bits::One => { + Backend::new(Spherical::new(quantizer.as_quantizer_impl::<1>()?)) + } + spherical::Bits::Two => { + Backend::new(Spherical::new(quantizer.as_quantizer_impl::<2>()?)) + } + spherical::Bits::Four => { + Backend::new(Spherical::new(quantizer.as_quantizer_impl::<4>()?)) + } }; + let config = Config::new(backend, capacity, max_degree, start_point, rerank)?; let index_config = self.build.config.clone(); Ok(finish(Provider::new(config)?, index_config)) diff --git a/diskann-inmem/src/repr/quantized/spherical.rs b/diskann-inmem/src/repr/quantized/spherical.rs index fb10fb9cc9..88aa2c0398 100644 --- a/diskann-inmem/src/repr/quantized/spherical.rs +++ b/diskann-inmem/src/repr/quantized/spherical.rs @@ -11,21 +11,48 @@ use diskann_quantization::{ use crate::{num::Bytes, repr::quantized}; -impl quantized::Quantizer for T +/// A new-type wrapper around [`iface::Quantizer`] for implementing [`quantized::Quantizer`]. +#[derive(Debug)] +pub struct Spherical(T); + +impl Spherical { + pub fn new(quantizer: T) -> Self { + Self(quantizer) + } + + pub fn into_inner(self) -> T { + self.0 + } +} + +impl std::ops::Deref for Spherical { + type Target = T; + fn deref(&self) -> &T { + &self.0 + } +} + +impl std::ops::DerefMut for Spherical { + fn deref_mut(&mut self) -> &mut T { + &mut self.0 + } +} + +impl quantized::Quantizer for Spherical where T: iface::Quantizer + 'static, { fn bytes(&self) -> Bytes { - Bytes::new(iface::Quantizer::bytes(self)) + Bytes::new(iface::Quantizer::bytes(&**self)) } fn dim(&self) -> usize { - iface::Quantizer::full_dim(self) + iface::Quantizer::full_dim(&**self) } fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { iface::Quantizer::compress( - self, + &**self, data, iface::OpaqueMut::new(quantized), ScopedAllocator::global(), @@ -37,7 +64,7 @@ where where V: quantized::DistanceVisitor<'a>, { - visitor.visit(iface::Quantizer::distance_computer_ref(self)) + visitor.visit(iface::Quantizer::distance_computer_ref(&**self)) } fn insert_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult @@ -45,7 +72,7 @@ where V: quantized::QueryDistanceVisitor<'a>, { let query_computer = iface::Quantizer::fused_query_computer( - self, + &**self, query, iface::QueryLayout::SameAsData, false, @@ -62,7 +89,7 @@ where V: quantized::QueryDistanceVisitor<'a>, { let query_computer = iface::Quantizer::fused_query_computer( - self, + &**self, query, iface::QueryLayout::FullPrecision, true, @@ -324,10 +351,10 @@ mod tests { SupportedMetric::InnerProduct, SupportedMetric::Cosine, ] { - let quantizer = train_quantizer::(data.as_view(), metric); + let quantizer = Spherical::new(train_quantizer::(data.as_view(), metric)); test_spherical( &quantizer, - &quantizer, + &*quantizer, &mut rng, distribution, TRIALS, From 025ffc300f1943d61a12a12f5fbc92d32a1284d6 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Tue, 20 Oct 2026 17:55:16 -0700 Subject: [PATCH 4/5] Few more docs. --- diskann-inmem/src/repr/quantized/spherical.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/diskann-inmem/src/repr/quantized/spherical.rs b/diskann-inmem/src/repr/quantized/spherical.rs index 88aa2c0398..19670b60b0 100644 --- a/diskann-inmem/src/repr/quantized/spherical.rs +++ b/diskann-inmem/src/repr/quantized/spherical.rs @@ -16,10 +16,12 @@ use crate::{num::Bytes, repr::quantized}; pub struct Spherical(T); impl Spherical { + /// Construct a new [`Spherical`] new-type wrapper around `quantizer`. pub fn new(quantizer: T) -> Self { Self(quantizer) } + /// Return the wrapped quantizer. pub fn into_inner(self) -> T { self.0 } From 61b1bdd6707fb4396fa3ac72ed5e02541bafab52 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 21 Oct 2026 12:34:27 -0700 Subject: [PATCH 5/5] Small doc cleanups. --- diskann-inmem/src/repr/quantized/mod.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index fa10b0d570..9b887283b4 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -111,8 +111,6 @@ pub trait Quantizer: std::fmt::Debug + Send + Sync + 'static { } /// Construct the [`Backend`] from `self`. - /// - /// This method is not implementable outside of `diskann-inmem`. fn into_backend(self) -> Backend where Self: Sized,