From 7ddd8ef8e6569e4574d5dc626b4e21934fb2fb70 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Tue, 13 Oct 2026 13:36:03 -0700 Subject: [PATCH 1/6] Extract quantized infrastructure. --- diskann-inmem/src/repr/internal/quantization/rerank.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/diskann-inmem/src/repr/internal/quantization/rerank.rs b/diskann-inmem/src/repr/internal/quantization/rerank.rs index 190c59247..b1890365f 100644 --- a/diskann-inmem/src/repr/internal/quantization/rerank.rs +++ b/diskann-inmem/src/repr/internal/quantization/rerank.rs @@ -224,15 +224,14 @@ mod tests { use std::assert_matches; - use diskann::utils::IntoUsize; - use diskann_utils::assert_contains; - use crate::{ counters::Counters, num::{Capacity, LogicalId, MaxDegree, SlotId}, repr::test::Reference, store::Store, }; + use diskann::utils::IntoUsize; + use diskann_utils::assert_contains; const TEST_DIM: usize = 1; From 93e3a19f17a9e493f30e0ce819e2def001ce77aa Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Tue, 13 Oct 2026 17:34:44 -0700 Subject: [PATCH 2/6] Preliminary scalar quantization support. --- diskann-inmem/integration/index/runner.rs | 154 +++++++++- .../graph/scalar/eight-bit-l2-baseline.json | 125 ++++++++ .../jsons/graph/scalar/eight-bit-l2.json | 53 ++++ .../graph/scalar/four-bit-ip-baseline.json | 125 ++++++++ .../jsons/graph/scalar/four-bit-ip.json | 53 ++++ .../graph/scalar/one-bit-l2-baseline.json | 248 ++++++++++++++++ .../jsons/graph/scalar/one-bit-l2.json | 98 ++++++ .../graph/scalar/two-bit-cosine-baseline.json | 125 ++++++++ .../jsons/graph/scalar/two-bit-cosine.json | 54 ++++ diskann-inmem/integration/main.rs | 44 +++ diskann-inmem/src/repr/quantized/mod.rs | 1 + diskann-inmem/src/repr/quantized/scalar.rs | 279 ++++++++++++++++++ diskann-quantization/src/num.rs | 2 + 13 files changed, 1360 insertions(+), 1 deletion(-) create mode 100644 diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/four-bit-ip.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/one-bit-l2.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json create mode 100644 diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json create mode 100644 diskann-inmem/src/repr/quantized/scalar.rs diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index 2089917bc..fc4602054 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -135,6 +135,19 @@ mod dto { } } + pub(super) mod scalar { + use super::*; + + #[derive(Debug, Serialize, Deserialize)] + #[serde(rename_all = "kebab-case")] + pub(in crate::index::runner) enum Bits { + One, + Two, + Four, + Eight, + } + } + #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub(super) enum Representation { @@ -149,6 +162,11 @@ mod dto { chunks: NonZeroUsize, rerank: quantization::Rerank, }, + Scalar { + bits: scalar::Bits, + override_to_cosine_normalized: bool, + rerank: quantization::Rerank, + }, } #[derive(Debug, Serialize, Deserialize)] @@ -331,6 +349,38 @@ mod spherical { } } +mod scalar { + use super::*; + + #[derive(Debug, Clone, Copy)] + pub(super) enum Bits { + One, + Two, + Four, + Eight, + } + + impl Bits { + pub(super) fn from_raw(raw: dto::scalar::Bits) -> Self { + match raw { + dto::scalar::Bits::One => Self::One, + dto::scalar::Bits::Two => Self::Two, + dto::scalar::Bits::Four => Self::Four, + dto::scalar::Bits::Eight => Self::Eight, + } + } + + pub(super) fn as_raw(&self) -> dto::scalar::Bits { + match self { + Self::One => dto::scalar::Bits::One, + Self::Two => dto::scalar::Bits::Two, + Self::Four => dto::scalar::Bits::Four, + Self::Eight => dto::scalar::Bits::Eight, + } + } + } +} + #[derive(Debug)] enum Representation { FullPrecision { @@ -344,6 +394,13 @@ enum Representation { chunks: NonZeroUsize, rerank: quantization::Rerank, }, + // NOTE: We allow overriding to `cosine_normalized` since that is basically L2 in + // a trench coat. We can at least still get coverage. + Scalar { + bits: scalar::Bits, + override_to_cosine_normalized: bool, + rerank: quantization::Rerank, + }, } impl Representation { @@ -358,6 +415,15 @@ impl Representation { chunks, rerank: quantization::Rerank::from_raw(rerank), }, + dto::Representation::Scalar { + bits, + override_to_cosine_normalized, + rerank, + } => Self::Scalar { + bits: scalar::Bits::from_raw(bits), + override_to_cosine_normalized, + rerank: quantization::Rerank::from_raw(rerank), + }, } } @@ -374,6 +440,15 @@ impl Representation { chunks: *chunks, rerank: rerank.as_raw(), }, + Self::Scalar { + bits, + override_to_cosine_normalized, + rerank, + } => dto::Representation::Scalar { + bits: bits.as_raw(), + override_to_cosine_normalized: *override_to_cosine_normalized, + rerank: rerank.as_raw(), + }, } } } @@ -547,6 +622,11 @@ impl Test { Representation::Product { chunks, rerank } => { self.create_product(data, *chunks, *rerank) } + Representation::Scalar { + bits, + override_to_cosine_normalized, + rerank, + } => self.create_scalar(data, *bits, *override_to_cosine_normalized, *rerank), } } @@ -661,6 +741,75 @@ impl Test { let index_config = self.build.config.clone(); Ok(finish(Provider::new(config)?, index_config)) } + + fn create_scalar( + &self, + data: DatasetView<'_>, + bits: scalar::Bits, + override_to_cosine_normalized: bool, + rerank: quantization::Rerank, + ) -> anyhow::Result> { + use diskann_quantization::{num::Positive, scalar as qscalar}; + use inmem::repr::quantized::{Quantized, scalar::Scalar}; + + let DatasetView::F32(data) = data else { + anyhow::bail!("scalar quantization only supports f32 data"); + }; + + // Step 1: Train a generic quantizer. + let std = Positive::::new(1.5).expect("1.5 is positive"); + let trainer = qscalar::train::ScalarQuantizationParameters::new(std); + let quantizer = trainer.train(data); + + let start_point = rowmajor::Owned::row_vector(Box::from( + ::compute_medoid(data), + )); + + // Step 2: Create the config. + let metric = if override_to_cosine_normalized { + inmem::repr::quantized::scalar::Metric::CosineNormalized + } else { + inmem::repr::quantized::scalar::Metric::try_from(self.data.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 index_config = self.build.config.clone(); + Ok(finish(Provider::new(config)?, index_config)) + } } fn finish(provider: DP, config: diskann::graph::Config) -> Arc @@ -763,7 +912,9 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision { Representation::FullPrecision { .. } => { // We match all valid data-types } - Representation::Spherical { .. } | Representation::Product { .. } => { + Representation::Spherical { .. } + | Representation::Product { .. } + | Representation::Scalar { .. } => { let data_type = input.data.data_type; // Ensure that the data type if `f32`. if data_type != DataType::F32 { @@ -792,6 +943,7 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision { Representation::FullPrecision { data_type } => data_type, Representation::Spherical { .. } => input.data.data_type, Representation::Product { .. } => input.data.data_type, + Representation::Scalar { .. } => input.data.data_type, }; // Load the data and perform any necessary data conversions. diff --git a/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json new file mode 100644 index 000000000..df05e5f86 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json @@ -0,0 +1,125 @@ +[ + { + "input": { + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "/yfcc/yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "/yfcc/groundtruth.bin", + "metric": "l2", + "preprocess": [], + "queries": "/yfcc/yfcc_query_100.fbin" + }, + "representation": { + "scalar": { + "bits": "eight", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": 1, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + }, + "type": "integration-test" + }, + "results": { + "build": { + "append_neighbors": 100311, + "distance": 3063537, + "get_neighbors": 357204, + "get_vector": 2824455, + "query_distance": 2259492, + "set_neighbors": 24670, + "set_vector": 10000 + }, + "knn": [ + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5439, + "get_vector": 45194, + "query_distance": 45194, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 45194, + "hops": 5439 + }, + "recall": { + "average": 0.902, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5790, + "get_vector": 49411, + "query_distance": 49411, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 49411, + "hops": 5790 + }, + "recall": { + "average": 0.904, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 10649, + "get_vector": 74516, + "query_distance": 74516, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 74516, + "hops": 10649 + }, + "recall": { + "average": 0.908, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + } + ] + } + } +] \ No newline at end of file diff --git a/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2.json b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2.json new file mode 100644 index 000000000..bf3c2dfb9 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2.json @@ -0,0 +1,53 @@ +{ + "search_directories": [ + "yfcc" + ], + "output_directory": null, + "jobs": [ + { + "type": "integration-test", + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "groundtruth.bin", + "metric": "l2", + "queries": "yfcc_query_100.fbin", + "preprocess": [] + }, + "representation": { + "scalar": { + "bits": "eight", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": null, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + } + } + ] +} diff --git a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json new file mode 100644 index 000000000..98311d17d --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json @@ -0,0 +1,125 @@ +[ + { + "input": { + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "/yfcc/yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "/yfcc/groundtruth_ip.bin", + "metric": "inner-product", + "preprocess": [], + "queries": "/yfcc/yfcc_query_100.fbin" + }, + "representation": { + "scalar": { + "bits": "four", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": 1, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + }, + "type": "integration-test" + }, + "results": { + "build": { + "append_neighbors": 128219, + "distance": 5554287, + "get_neighbors": 380443, + "get_vector": 3110524, + "query_distance": 2193419, + "set_neighbors": 41661, + "set_vector": 10000 + }, + "knn": [ + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5219, + "get_vector": 49096, + "query_distance": 49096, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 49096, + "hops": 5219 + }, + "recall": { + "average": 0.764, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5250, + "get_vector": 49283, + "query_distance": 49283, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 49283, + "hops": 5250 + }, + "recall": { + "average": 0.764, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 10227, + "get_vector": 76057, + "query_distance": 76057, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 76057, + "hops": 10227 + }, + "recall": { + "average": 0.764, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + } + ] + } + } +] \ No newline at end of file diff --git a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip.json b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip.json new file mode 100644 index 000000000..93fe21b07 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip.json @@ -0,0 +1,53 @@ +{ + "search_directories": [ + "yfcc" + ], + "output_directory": null, + "jobs": [ + { + "type": "integration-test", + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "groundtruth_ip.bin", + "metric": "inner-product", + "queries": "yfcc_query_100.fbin", + "preprocess": [] + }, + "representation": { + "scalar": { + "bits": "four", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": null, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + } + } + ] +} diff --git a/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json new file mode 100644 index 000000000..a4e231ca1 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json @@ -0,0 +1,248 @@ +[ + { + "input": { + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "/yfcc/yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "/yfcc/groundtruth.bin", + "metric": "l2", + "preprocess": [], + "queries": "/yfcc/yfcc_query_100.fbin" + }, + "representation": { + "scalar": { + "bits": "one", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": 1, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + }, + "type": "integration-test" + }, + "results": { + "build": { + "append_neighbors": 125092, + "distance": 4903759, + "get_neighbors": 401105, + "get_vector": 3459145, + "query_distance": 2651349, + "set_neighbors": 35323, + "set_vector": 10000 + }, + "knn": [ + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5479, + "get_vector": 52625, + "query_distance": 52625, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 52625, + "hops": 5479 + }, + "recall": { + "average": 0.455, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5871, + "get_vector": 56863, + "query_distance": 56863, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 56863, + "hops": 5871 + }, + "recall": { + "average": 0.456, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 10739, + "get_vector": 88317, + "query_distance": 88317, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 88317, + "hops": 10739 + }, + "recall": { + "average": 0.458, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + } + ] + } + }, + { + "input": { + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "/yfcc/yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "/yfcc/groundtruth.bin", + "metric": "l2", + "preprocess": [], + "queries": "/yfcc/yfcc_query_100.fbin" + }, + "representation": { + "scalar": { + "bits": "one", + "override_to_cosine_normalized": false, + "rerank": "f16" + } + }, + "search": { + "knn": [ + { + "beam_width": 1, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + }, + "type": "integration-test" + }, + "results": { + "build": { + "append_neighbors": 125092, + "distance": 4903759, + "get_neighbors": 401105, + "get_vector": 3459145, + "query_distance": 2651349, + "set_neighbors": 35323, + "set_vector": 10000 + }, + "knn": [ + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5479, + "get_vector": 57725, + "query_distance": 57725, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 52625, + "hops": 5479 + }, + "recall": { + "average": 0.795, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5871, + "get_vector": 61963, + "query_distance": 61963, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 56863, + "hops": 5871 + }, + "recall": { + "average": 0.793, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 10739, + "get_vector": 98417, + "query_distance": 98417, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 88317, + "hops": 10739 + }, + "recall": { + "average": 0.891, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + } + ] + } + } +] \ No newline at end of file diff --git a/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2.json b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2.json new file mode 100644 index 000000000..f11a11356 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2.json @@ -0,0 +1,98 @@ +{ + "search_directories": [ + "yfcc" + ], + "output_directory": null, + "jobs": [ + { + "type": "integration-test", + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "groundtruth.bin", + "metric": "l2", + "queries": "yfcc_query_100.fbin", + "preprocess": [] + }, + "representation": { + "scalar": { + "bits": "one", + "override_to_cosine_normalized": false, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": null, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + } + }, + { + "type": "integration-test", + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "groundtruth.bin", + "metric": "l2", + "queries": "yfcc_query_100.fbin", + "preprocess": [] + }, + "representation": { + "scalar": { + "bits": "one", + "override_to_cosine_normalized": false, + "rerank": "f16" + } + }, + "search": { + "knn": [ + { + "beam_width": null, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + } + } + ] +} diff --git a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json new file mode 100644 index 000000000..0b2d7f568 --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json @@ -0,0 +1,125 @@ +[ + { + "input": { + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "/yfcc/yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "/yfcc/groundtruth.bin", + "metric": "l2", + "preprocess": [], + "queries": "/yfcc/yfcc_query_100.fbin" + }, + "representation": { + "scalar": { + "bits": "two", + "override_to_cosine_normalized": true, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": 1, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + }, + "type": "integration-test" + }, + "results": { + "build": { + "append_neighbors": 107423, + "distance": 3639726, + "get_neighbors": 368584, + "get_vector": 3001191, + "query_distance": 2360497, + "set_neighbors": 28073, + "set_vector": 10000 + }, + "knn": [ + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5475, + "get_vector": 47483, + "query_distance": 47483, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 47483, + "hops": 5475 + }, + "recall": { + "average": 0.65, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 5802, + "get_vector": 51303, + "query_distance": 51303, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 51303, + "hops": 5802 + }, + "recall": { + "average": 0.652, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + }, + { + "counters": { + "append_neighbors": 0, + "distance": 0, + "get_neighbors": 10617, + "get_vector": 77819, + "query_distance": 77819, + "set_neighbors": 0, + "set_vector": 0 + }, + "misc": { + "cmps": 77819, + "hops": 10617 + }, + "recall": { + "average": 0.659, + "num_queries": 100, + "recall_k": 10, + "recall_n": 10 + } + } + ] + } + } +] \ No newline at end of file diff --git a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json new file mode 100644 index 000000000..ec4fd80ea --- /dev/null +++ b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json @@ -0,0 +1,54 @@ +{ + "search_directories": [ + "yfcc" + ], + "output_directory": null, + "jobs": [ + { + "type": "integration-test", + "content": { + "build": { + "alpha": 1.2000000476837158, + "l_build": 20, + "max_degree": 20, + "pruned_degree": 16 + }, + "data": { + "data": "yfcc_10k.fbin", + "data_type": "f32", + "groundtruth": "groundtruth.b + in", + "metric": "l2", + "queries": "yfcc_query_100.fbin", + "preprocess": [] + }, + "representation": { + "scalar": { + "bits": "two", + "override_to_cosine_normalized": true, + "rerank": "none" + } + }, + "search": { + "knn": [ + { + "beam_width": null, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 50 + }, + { + "beam_width": 3, + "knn": 10, + "search_l": 100 + } + ] + } + } + } + ] +} diff --git a/diskann-inmem/integration/main.rs b/diskann-inmem/integration/main.rs index 6ba14b956..946854ac6 100644 --- a/diskann-inmem/integration/main.rs +++ b/diskann-inmem/integration/main.rs @@ -377,4 +377,48 @@ mod tests { "graph/product/cosine-baseline.json", ); } + + //---------------------// + // Scalar Quantization // + //---------------------// + + #[test] + #[cfg(not(miri))] + fn graph_scalar_one() { + run_regression_example( + "graph/scalar/one-bit-l2.json", + "checks.json", + "graph/scalar/one-bit-l2-baseline.json", + ); + } + + #[test] + #[cfg(not(any(miri, coverage)))] + fn graph_scalar_two() { + run_regression_example( + "graph/scalar/two-bit-cosine.json", + "checks.json", + "graph/scalar/two-bit-cosine-baseline.json", + ); + } + + #[test] + #[cfg(not(any(miri, coverage)))] + fn graph_scalar_four() { + run_regression_example( + "graph/scalar/four-bit-ip.json", + "checks.json", + "graph/scalar/four-bit-ip-baseline.json", + ); + } + + #[test] + #[cfg(not(any(miri, coverage)))] + fn graph_scalar_eight() { + run_regression_example( + "graph/scalar/eight-bit-l2.json", + "checks.json", + "graph/scalar/eight-bit-l2-baseline.json", + ); + } } diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index c9a6a6969..9b6e2a45c 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -29,6 +29,7 @@ use crate::{ // Built-in quantizers pub mod product; +pub mod scalar; /// Plug-in for a generalized quantizer compatible with [`Quantized`]. /// diff --git a/diskann-inmem/src/repr/quantized/scalar.rs b/diskann-inmem/src/repr/quantized/scalar.rs new file mode 100644 index 000000000..b9a820342 --- /dev/null +++ b/diskann-inmem/src/repr/quantized/scalar.rs @@ -0,0 +1,279 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann::{ANNError, ANNResult}; +use diskann_quantization::{ + CompressInto, + bits::{Representation, Unsigned}, + distances as quant_distances, + scalar::{ + self, CompensatedCosineNormalized as CCosineNormalized, CompensatedIP as CInnerProduct, + CompensatedSquaredL2 as CSquaredL2, + }, +}; +use thiserror::Error; + +use crate::{num::Bytes, repr::quantized}; + +// Help bring the ridiculous names under control. +type VectorRef<'a, const N: usize> = scalar::CompensatedVectorRef<'a, N>; +type VectorMut<'a, const N: usize> = scalar::MutCompensatedVectorRef<'a, N>; +type Vector = scalar::CompensatedVector; + +/// A scalar quantizer specialized for `N` bit computation. +#[derive(Debug)] +pub struct Scalar { + quantizer: scalar::ScalarQuantizer, + metric: Metric, +} + +impl Scalar { + /// Construct a new [`Scalar`] quantizer. + pub fn new(quantizer: scalar::ScalarQuantizer, metric: Metric) -> Self { + Self { quantizer, metric } + } + + // The following methods are used to implement `Quantizer` in their full generic glory. + + fn bytes(&self) -> Bytes + where + Unsigned: Representation, + { + Bytes::new(Vector::::canonical_bytes(self.dim())) + } + + fn dim(&self) -> usize { + self.quantizer.dim() + } + + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> + where + Unsigned: Representation, + scalar::ScalarQuantizer: for<'a> CompressInto<&'a [f32], VectorMut<'a, N>>, + { + let quantized = + VectorMut::from_canonical_back_mut(quantized, self.dim()).map_err(ANNError::new)?; + self.quantizer + .compress_into(data, quantized) + .map_err(ANNError::new)?; + Ok(()) + } + + fn distance<'a, V>(&'a self, visitor: V) -> V::Output + where + V: quantized::DistanceVisitor<'a>, + Distance: quantized::Distance, + Distance: quantized::Distance, + Distance: quantized::Distance, + { + use diskann_quantization::AsFunctor; + + let dim = self.dim(); + match self.metric { + Metric::SquaredL2 => visitor.visit(Distance::::new( + self.quantizer.as_functor(), + dim, + )), + Metric::InnerProduct => visitor.visit(Distance::::new( + self.quantizer.as_functor(), + dim, + )), + Metric::CosineNormalized => visitor.visit(Distance::::new( + self.quantizer.as_functor(), + dim, + )), + } + } + + fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: quantized::QueryDistanceVisitor<'a>, + Unsigned: Representation, + scalar::ScalarQuantizer: for<'b> CompressInto<&'b [f32], VectorMut<'b, N>>, + QueryDistance: quantized::QueryDistance, + QueryDistance: quantized::QueryDistance, + QueryDistance: quantized::QueryDistance, + { + use diskann_quantization::AsFunctor; + use diskann_utils::ReborrowMut; + + let mut q = Vector::new_boxed(self.dim()); + + self.quantizer + .compress_into(query, q.reborrow_mut()) + .map_err(ANNError::new)?; + + let visited = match self.metric { + Metric::SquaredL2 => visitor.visit(QueryDistance::::new( + q, + self.quantizer.as_functor(), + )), + Metric::InnerProduct => visitor.visit(QueryDistance::::new( + q, + self.quantizer.as_functor(), + )), + Metric::CosineNormalized => visitor.visit(QueryDistance::::new( + q, + self.quantizer.as_functor(), + )), + }; + + Ok(visited) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Metric { + SquaredL2, + InnerProduct, + CosineNormalized, +} + +impl TryFrom for Metric { + type Error = CosineNotSupported; + fn try_from(m: diskann_vector::distance::Metric) -> Result { + use diskann_vector::distance::Metric as VMetric; + + match m { + VMetric::L2 => Ok(Self::SquaredL2), + VMetric::InnerProduct => Ok(Self::InnerProduct), + VMetric::CosineNormalized => Ok(Self::CosineNormalized), + VMetric::Cosine => Err(CosineNotSupported), + } + } +} + +#[derive(Debug, Error)] +#[error("Metric::Cosine is not supported")] +#[non_exhaustive] +pub struct CosineNotSupported; + +diskann::convert_error!(CosineNotSupported); + +/// A macro is used to stamp out implementations instead of the crazy generic bounds that +/// we *could* use because it's (1) easier to see what exactly is implemented via rustdoc +/// and (2) gives us a bit of flexibility to changes details on the backend. +macro_rules! quantizer { + ($N:literal) => { + impl quantized::Quantizer for Scalar<$N> { + fn bytes(&self) -> Bytes { + >::bytes(self) + } + fn dim(&self) -> usize { + >::dim(self) + } + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { + >::quantize(self, quantized, data) + } + fn distance<'a, V>(&'a self, visitor: V) -> V::Output + where + V: quantized::DistanceVisitor<'a>, + { + >::distance(self, visitor) + } + fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: quantized::QueryDistanceVisitor<'a>, + { + >::query_distance(self, query, visitor) + } + } + }; + ($($Ns:literal),+ $(,)?) => { + $(quantizer!($Ns);)+ + } +} + +quantizer!(1, 2, 3, 4, 5, 6, 7, 8); + +//----------------// +// Query Distance // +//----------------// + +#[derive(Debug)] +struct QueryDistance +where + Unsigned: Representation, +{ + query: Vector, + inner: T, +} + +impl QueryDistance +where + Unsigned: Representation, +{ + fn new(query: Vector, inner: T) -> Self { + Self { query, inner } + } +} + +impl quantized::QueryDistance for QueryDistance +where + Unsigned: Representation, + T: for<'a> diskann_vector::DistanceFunction< + VectorRef<'a, N>, + VectorRef<'a, N>, + quant_distances::Result, + > + Send + + Sync + + std::fmt::Debug, +{ + #[inline] + fn eval(&self, x: &[u8]) -> ANNResult { + use diskann_utils::Reborrow; + + let x = VectorRef::::from_canonical_back(x, self.query.len()).map_err(ANNError::new)?; + self.inner + .evaluate_similarity(self.query.reborrow(), x) + .map_err(ANNError::new) + } +} + +//----------// +// Distance // +//----------// + +#[derive(Debug)] +struct Distance { + inner: T, + dim: usize, +} + +impl Distance { + fn new(inner: T, dim: usize) -> Self { + Self { inner, dim } + } +} + +impl quantized::Distance for Distance +where + Unsigned: Representation, + T: for<'a> diskann_vector::DistanceFunction< + VectorRef<'a, N>, + VectorRef<'a, N>, + quant_distances::Result, + > + Send + + Sync + + std::fmt::Debug, +{ + #[inline] + fn eval(&self, x: &[u8], y: &[u8]) -> ANNResult { + let x = VectorRef::::from_canonical_back(x, self.dim).map_err(ANNError::new)?; + let y = VectorRef::::from_canonical_back(y, self.dim).map_err(ANNError::new)?; + self.inner.evaluate_similarity(x, y).map_err(ANNError::new) + } +} + +/////////// +// Tests // +/////////// + +#[cfg(test)] +mod tests { + use super::*; + + +} diff --git a/diskann-quantization/src/num.rs b/diskann-quantization/src/num.rs index 8fe974a25..ba2ba7a68 100644 --- a/diskann-quantization/src/num.rs +++ b/diskann-quantization/src/num.rs @@ -50,6 +50,8 @@ where // SAFETY: 1.0 is positive. pub const POSITIVE_ONE_F32: Positive = unsafe { Positive::new_unchecked(1.0) }; +// SAFETY: 1.0 is positive. +pub const POSITIVE_ONE_F64: Positive = unsafe { Positive::new_unchecked(1.0) }; #[derive(Debug, Clone, Copy, PartialEq)] #[repr(transparent)] From 102c9162775a9c4c9831f83a3ccf61dd77428b78 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 14 Oct 2026 11:18:49 -0700 Subject: [PATCH 3/6] Scalar quantization! --- diskann-benchmark/src/index/inmem2.rs | 302 +++++++++++++++- .../jsons/graph/scalar/two-bit-cosine.json | 3 +- diskann-inmem/src/neighbors.rs | 3 +- diskann-inmem/src/repr/full.rs | 12 +- .../src/repr/internal/quantization/rerank.rs | 7 +- diskann-inmem/src/repr/quantized/mod.rs | 84 ++++- diskann-inmem/src/repr/quantized/product.rs | 31 +- diskann-inmem/src/repr/quantized/scalar.rs | 333 +++++++++++++++++- diskann-inmem/src/repr/spherical.rs | 6 +- .../src/distance/distance_provider.rs | 14 +- 10 files changed, 728 insertions(+), 67 deletions(-) diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index 3bb3f43a7..2d77b4f98 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -30,11 +30,13 @@ use diskann_benchmark_runner::{ Benchmark, Checker, Checkpoint, Input, Registry, }; use diskann_inmem::{ + self as inmem, num::{Capacity, MaxDegree}, repr::{self, Full, FullPrecision}, Provider, Strategy, }; use diskann_quantization::{ + self as dq, alloc::{GlobalAllocator, Poly}, product::tables::BasicTable, spherical::iface, @@ -60,6 +62,7 @@ pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> registry.register("inmem2-spherical", SphericalBuild)?; registry.register("inmem2-pq", ProductBuild)?; + registry.register("inmem2-scalar", ScalarBuild)?; registry.register("inmem2-f32-stream", StreamingBenchmark::::new())?; Ok(()) @@ -113,6 +116,7 @@ mod dto { pub(super) enum Quantization { Spherical(Spherical), Product(Product), + Scalar(Scalar), } #[derive(Debug, Serialize, Deserialize)] @@ -135,6 +139,20 @@ mod dto { pub(super) seed: u64, } + #[derive(Debug, Serialize, Deserialize)] + pub(super) enum ScalarBits { + One, + Two, + Four, + Eight, + } + + #[derive(Debug, Serialize, Deserialize)] + pub(super) struct Scalar { + pub(super) bits: ScalarBits, + pub(super) standard_deviations: f64, + } + //-----------// // Streaming // //-----------// @@ -369,20 +387,24 @@ enum Quantization { None, Spherical(Spherical), Product(Product), + Scalar(Scalar), } impl Quantization { - fn from_raw(raw: Option) -> Self { - if let Some(raw) = raw { + fn from_raw(raw: Option) -> anyhow::Result { + let this = if let Some(raw) = raw { match raw { dto::Quantization::Spherical(spherical) => { Self::Spherical(Spherical::from_raw(spherical)) } dto::Quantization::Product(product) => Self::Product(Product::from_raw(product)), + dto::Quantization::Scalar(scalar) => Self::Scalar(Scalar::from_raw(scalar)?), } } else { Self::None - } + }; + + Ok(this) } fn as_spherical(&self) -> Option<&Spherical> { @@ -398,6 +420,13 @@ impl Quantization { _ => None, } } + + fn as_scalar(&self) -> Option<&Scalar> { + match self { + Self::Scalar(scalar) => Some(scalar), + _ => None, + } + } } impl std::fmt::Display for Quantization { @@ -414,6 +443,11 @@ impl std::fmt::Display for Quantization { kv.push("product", product); write!(f, "{}", kv) } + Self::Scalar(scalar) => { + let mut kv = KeyValue::new(); + kv.push("scalar", scalar); + write!(f, "{}", kv) + } } } } @@ -578,6 +612,77 @@ impl std::fmt::Display for Product { } } +#[derive(Debug)] +enum ScalarBits { + One, + Two, + Four, + Eight, +} + +impl ScalarBits { + fn from_raw(raw: dto::ScalarBits) -> Self { + match raw { + dto::ScalarBits::One => Self::One, + dto::ScalarBits::Two => Self::Two, + dto::ScalarBits::Four => Self::Four, + dto::ScalarBits::Eight => Self::Eight, + } + } +} + +impl std::fmt::Display for ScalarBits { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::One => write!(f, "1"), + Self::Two => write!(f, "2"), + Self::Four => write!(f, "4"), + Self::Eight => write!(f, "8"), + } + } +} + +#[derive(Debug)] +struct Scalar { + bits: ScalarBits, + standard_deviations: dq::num::Positive, +} + +impl Scalar { + fn from_raw(raw: dto::Scalar) -> anyhow::Result { + let dto::Scalar { bits, standard_deviations } = raw; + + Ok(Self { + bits: ScalarBits::from_raw(bits), + standard_deviations: dq::num::Positive::new(standard_deviations)?, + }) + } + + fn train( + &self, + data: rowmajor::Ref<'_, f32>, + ) -> anyhow::Result { + let trainer = dq::scalar::train::ScalarQuantizationParameters::new(self.standard_deviations); + let quantizer = trainer.train(data); + Ok(quantizer) + } +} + +impl std::fmt::Display for Scalar { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let Self { bits, standard_deviations } = self; + + let std = standard_deviations.into_inner(); + + let mut kv = KeyValue::new(); + kv.push("bits", bits); + kv.push("standard_deviations", &std); + + write!(f, "{}", kv) + } +} + + //-------------// // StaticBuild // //-------------// @@ -607,7 +712,7 @@ impl StaticBuild { let data = Data::from_raw(data, checker.as_deref_mut())?; let build = BuildParams::from_raw(build, data.distance)?; let search = KnnSearch::from_raw(search, checker)?; - let quantization = Quantization::from_raw(quantization); + let quantization = Quantization::from_raw(quantization)?; Ok(Self { data, @@ -724,7 +829,7 @@ impl Build { impl Benchmark for Build where - T: diskann_inmem::repr::FullPrecision + diskann::graph::SampleableForStart + AsDataType, + T: repr::FullPrecision + diskann::graph::SampleableForStart + AsDataType, { type Input = StaticBuild; type Output = (); @@ -1160,6 +1265,191 @@ impl Benchmark for ProductBuild { } } +//--------// +// Scalar // +//-------// + +#[derive(Debug)] +struct ScalarBuild; + +impl Benchmark for ScalarBuild { + type Input = StaticBuild; + type Output = (); + + fn try_match(&self, input: &StaticBuild, context: &MatchContext) -> Score { + let mut score = context.success(0); + + let DispatchParams { + data_type, + quantization, + distance, + } = input.dispatch_params(); + + if !matches!(quantization, Quantization::Scalar(_)) { + score.fail(2000, &"needed scalar-quantization"); + } + + if distance == Metric::Cosine { + score.fail( + 500, + &format_args!( + "{} is not supported for scalar quantization- use \"cosine-normalized\" instead", + Metric::Cosine + ), + ); + } + + if !f32::is_match(data_type) { + score.fail( + 1000, + &format_args!( + "expected data-type {}, instead got {}", + Quote(f32::DATA_TYPE), + Quote(data_type) + ), + ) + } + + score + } + + fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "scalar-quantized index build-and-search",)?; + + Ok(()) + } + + fn run( + &self, + input: &StaticBuild, + checkpoint: Checkpoint<'_>, + mut output: &mut dyn Output, + ) -> anyhow::Result<()> { + use repr::{Quantized, quantized::scalar::Scalar}; + + writeln!(output, "{input}\n")?; + + let scalar = input.quantization.as_scalar().unwrap(); + + // Load data. + let data: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.data.data), + )?); + + let dim = data.ncols(); + let num_points = data.nrows(); + writeln!(output, "Loaded {num_points} points, dim={dim}")?; + + let quantizer = scalar.train(data.as_view())?; + + // Compute the medoid of the dataset as the single start point. + + let metric = repr::quantized::scalar::Metric::try_from(input.data.distance)?; + + 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 rerank = repr::quantized::Rerank::F16 { + 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 provider = Provider::<_, u32>::new(config)?; + let index = Arc::new(DiskANNIndex::new( + input.build.config.clone(), + provider, + None, + )); + + // Build via SingleInsert. + let rt = benchmark_core::tokio::runtime(input.build.num_threads.get())?; + let builder = build_core::graph::SingleInsert::new( + index.clone(), + data, + Strategy, + build_core::ids::Identity::::new(), + ); + + let build_results = build_core::build_tracked( + builder, + build_core::Parallelism::dynamic(diskann::utils::ONE, input.build.num_threads), + &rt, + Some(&ProgressMeter::new(output)), + )?; + + let total_build_time = build_results.end_to_end_latency(); + writeln!( + output, + "\nBuild complete in {:.2}s", + total_build_time.as_seconds() + )?; + checkpoint.checkpoint(&total_build_time)?; + + // Search. + let queries: Arc> = Arc::new(datafiles::load_dataset( + datafiles::BinFile(&input.search.queries), + )?); + let max_k = input.search.maximum_recall_k(); + let groundtruth = datafiles::load_groundtruth( + datafiles::BinFile(&input.search.groundtruth), + Some(max_k), + )?; + + writeln!(output, "Loaded {} queries\n", queries.nrows())?; + + let knn = benchmark_core::search::graph::KNN::new( + index, + queries, + benchmark_core::search::graph::Strategy::broadcast(Strategy), + )?; + + let results = _knn( + &knn, + &groundtruth, + input.search.reps, + &input.search.num_threads, + &input.search.runs, + )?; + + let results = AggregatedSearchResults::Topk(results); + + writeln!(output, "{}", results)?; + + Ok(()) + } +} + + fn _knn( runner: &dyn crate::index::search::knn::Knn, groundtruth: &dyn benchmark_core::recall::Rows, @@ -1637,7 +1927,7 @@ where // // Also, this is *way* more expensive than it needs to be because each delete creates // and then destroys an EBR guard. - let ctx = diskann_inmem::Context; + let ctx = inmem::Context; for id in ids.clone() { self.runtime .block_on(self.index.provider().delete(&ctx, &(id as u32)))?; diff --git a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json index ec4fd80ea..f3729d16d 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json +++ b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine.json @@ -16,8 +16,7 @@ "data": { "data": "yfcc_10k.fbin", "data_type": "f32", - "groundtruth": "groundtruth.b - in", + "groundtruth": "groundtruth.bin", "metric": "l2", "queries": "yfcc_query_100.fbin", "preprocess": [] diff --git a/diskann-inmem/src/neighbors.rs b/diskann-inmem/src/neighbors.rs index f8ad42b0f..b67e7b049 100644 --- a/diskann-inmem/src/neighbors.rs +++ b/diskann-inmem/src/neighbors.rs @@ -498,7 +498,7 @@ mod tests { let n = Neighbors::new(IdLimit::new(4), 4).unwrap(); let lock = n.lock(0).unwrap(); - assert_eq!(lock.as_slice(), &[]); + assert!(lock.as_slice().is_empty()); lock.append(&[1, 2, 3]).unwrap(); let mut out = AdjacencyList::with_capacity(4); @@ -582,7 +582,6 @@ mod tests { assert_eq!(lock.capacity(), neighbors.max_degree().value()); assert_eq!(lock.len(), 0); assert!(lock.is_empty()); - assert_eq!(lock.as_slice(), &[]); } // Verify out-of-bounds accesses error. diff --git a/diskann-inmem/src/repr/full.rs b/diskann-inmem/src/repr/full.rs index 0c265a0e7..27c0faf41 100644 --- a/diskann-inmem/src/repr/full.rs +++ b/diskann-inmem/src/repr/full.rs @@ -43,7 +43,8 @@ use diskann_vector::{ UnalignedSlice, conversion::SliceCast, distance::{ - Cosine, CosineNormalized, DistanceProvider, InnerProduct, Metric, Specialize, SquaredL2, + Cosine, CosineNormalized, Distance as VDistance, InnerProduct, Metric, Specialize, + SquaredL2, }, }; use diskann_wide::{ @@ -753,8 +754,7 @@ impl FullPrecisionImpl for i8 { ) -> ANNResult> { let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?; - let distance = - >::distance_comparer(full.metric(), Some(full.dim())); + let distance = VDistance::::new(full.metric(), Some(full.dim())); let output: Box = into .into_expand_beam(distance, prefetch::Loop::new()) @@ -766,8 +766,7 @@ impl FullPrecisionImpl for i8 { fn make_prune<'a>(full: &'a Full) -> ANNResult> { let reader = full.reader()?; - let distance = - >::distance_comparer(full.metric(), Some(full.dim())); + let distance = VDistance::::new(full.metric(), Some(full.dim())); let output: Box = repr::internal::intrusive::Prune::new(reader, distance).boxed(); @@ -831,6 +830,7 @@ mod tests { use diskann::neighbor::Neighbor; use diskann_utils::lazy_format; + use diskann_vector::distance::DistanceProvider; use hashbrown::{HashMap, HashSet}; use rand::{Rng, SeedableRng, rngs::StdRng}; @@ -1246,7 +1246,7 @@ mod tests { id }; - let distance = >::distance_comparer(metric, None); + let distance = VDistance::::new(metric, None); let expected = distance.call(&start_point, &query); // Expand Beam - both `evaluate` and `expand_beam` share the same distance computer, diff --git a/diskann-inmem/src/repr/internal/quantization/rerank.rs b/diskann-inmem/src/repr/internal/quantization/rerank.rs index b1890365f..5bc602528 100644 --- a/diskann-inmem/src/repr/internal/quantization/rerank.rs +++ b/diskann-inmem/src/repr/internal/quantization/rerank.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann_vector::distance::{Distance, DistanceProvider}; +use diskann_vector::distance::Distance; use half::f16; use thiserror::Error; @@ -86,10 +86,7 @@ impl Reranker { match rerank { Rerank::None => Ok(Self::disabled()), Rerank::F16 => { - let distance = >::distance_comparer( - metric.as_vector_metric(), - Some(dim), - ); + let distance = Distance::new(metric.as_vector_metric(), Some(dim)); let this = Self::F16(distance); let config = Simple::config(rerank.bytes_for(dim)?); diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index 9b6e2a45c..9614d7cd0 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -735,6 +735,72 @@ mod inner { } } +/////////////////////// +// Shared Test Utils // +/////////////////////// + +#[cfg(test)] +mod test { + use super::*; + + #[derive(Debug)] + pub(super) struct Boxit; + + impl<'a> DistanceVisitor<'a> for Boxit { + type Output = Box; + fn visit(self, distance: D) -> Self::Output + where + D: Distance + 'a, + { + Box::new(distance) + } + } + + impl<'a> QueryDistanceVisitor<'a> for Boxit { + type Output = Box; + fn visit(self, distance: D) -> Self::Output + where + D: QueryDistance + 'a, + { + Box::new(distance) + } + } + + /// A dyn-compatible wrapper for [`super::Quantizer`]. + pub(super) trait DynQuantizer { + fn dim(&self) -> usize; + fn bytes(&self) -> Bytes; + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()>; + fn distance<'a>(&'a self) -> Box; + fn query_distance<'a>(&'a self, query: &'a [f32]) + -> ANNResult>; + } + + impl DynQuantizer for Q + where + Q: Quantizer, + { + fn dim(&self) -> usize { + Quantizer::dim(self) + } + fn bytes(&self) -> Bytes { + Quantizer::bytes(self) + } + fn quantize(&self, quantized: &mut [u8], data: &[f32]) -> ANNResult<()> { + Quantizer::quantize(self, quantized, data) + } + fn distance<'a>(&'a self) -> Box { + Quantizer::distance(self, Boxit) + } + fn query_distance<'a>( + &'a self, + query: &'a [f32], + ) -> ANNResult> { + Quantizer::query_distance(self, query, Boxit) + } + } +} + /////////// // Tests // /////////// @@ -747,7 +813,7 @@ mod tests { use diskann::graph::test::synthetic::Grid; use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; - use diskann_vector::distance::DistanceProvider; + use diskann_vector::distance::{Distance as VDistance, Metric as VMetric}; use hashbrown::HashMap; use crate::{ @@ -755,18 +821,19 @@ mod tests { num::{LogicalId, SlotId}, repr::test::{Reference, test_expand_beam, test_prune}, }; + ///////////////////// + // Quantized Tests // + ///////////////////// /// A simple test quantizer that tries its best to convert `f32` to `i8`. #[derive(Debug, Clone, Copy)] struct Round { dim: usize, - distance: diskann_vector::distance::Distance, + distance: VDistance, } impl Round { fn new(dim: usize, metric: Metric) -> Self { - use diskann_vector::distance::{DistanceProvider, Metric as VMetric}; - let vmetric = match metric { Metric::SquaredL2 => VMetric::L2, Metric::InnerProduct => VMetric::InnerProduct, @@ -775,7 +842,7 @@ mod tests { Self { dim, - distance: >::distance_comparer(vmetric, None), + distance: VDistance::new(vmetric, None), } } } @@ -895,7 +962,7 @@ mod tests { let quantized = config.build().unwrap(); assert_eq!(repr::Representation::id_limit(&quantized), TEST_LIMIT); - assert_eq!(quantized.dim(), grid.dim().into()); + assert_eq!(quantized.dim(), usize::from(grid.dim())); let mut reference = Reference::new(grid.dim().into()); @@ -972,10 +1039,7 @@ mod tests { let v2 = &reference[i2]; // The Round quantizer is exact for these integral values. - let f = >::distance_comparer( - metric.as_internal_metric().as_vector_metric(), - None, - ); + let f = VDistance::::new(metric.as_internal_metric().as_vector_metric(), None); let distances = HashMap::from_iter([ (s0, f.call(&query, v0)), diff --git a/diskann-inmem/src/repr/quantized/product.rs b/diskann-inmem/src/repr/quantized/product.rs index a819a11e9..a8b394df8 100644 --- a/diskann-inmem/src/repr/quantized/product.rs +++ b/diskann-inmem/src/repr/quantized/product.rs @@ -357,6 +357,9 @@ mod tests { use diskann::graph::test::synthetic::Grid; use diskann_utils::assert_contains; + use diskann_vector::distance::Distance as VDistance; + + use crate::repr::quantized::test::Boxit; fn train( data: rowmajor::Ref<'_, f32>, @@ -391,29 +394,6 @@ mod tests { Product::new(basic, metric, NonZeroUsize::new(1)).unwrap() } - #[derive(Debug)] - struct Boxit; - - impl<'a> quantized::DistanceVisitor<'a> for Boxit { - type Output = Box; - fn visit(self, distance: D) -> Self::Output - where - D: quantized::Distance + 'a, - { - Box::new(distance) - } - } - - impl<'a> quantized::QueryDistanceVisitor<'a> for Boxit { - type Output = Box; - fn visit(self, distance: D) -> Self::Output - where - D: quantized::QueryDistance + 'a, - { - Box::new(distance) - } - } - // 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 @@ -502,13 +482,10 @@ mod tests { /// Happy path tests. #[test] fn test_product() { - use diskann_vector::distance::DistanceProvider; - for metric in Metric::all() { let (product, data) = make_test(*metric); - let reference = - >::distance_comparer(metric.as_vector_metric(), None); + let reference = VDistance::::new(metric.as_vector_metric(), None); // Quantize the data. let mut quantized = rowmajor::Owned::from_element( diff --git a/diskann-inmem/src/repr/quantized/scalar.rs b/diskann-inmem/src/repr/quantized/scalar.rs index b9a820342..43c3c6acb 100644 --- a/diskann-inmem/src/repr/quantized/scalar.rs +++ b/diskann-inmem/src/repr/quantized/scalar.rs @@ -3,6 +3,8 @@ * Licensed under the MIT license. */ +//! Support for [`diskann_quantization::scalar`]. + use diskann::{ANNError, ANNResult}; use diskann_quantization::{ CompressInto, @@ -35,6 +37,16 @@ impl Scalar { Self { quantizer, metric } } + /// Return the configured metric. + pub fn metric(&self) -> Metric { + self.metric + } + + /// Return the number of bits used for each dimension. + pub fn nbits(&self) -> usize { + N + } + // The following methods are used to implement `Quantizer` in their full generic glory. fn bytes(&self) -> Bytes @@ -53,11 +65,14 @@ impl Scalar { Unsigned: Representation, scalar::ScalarQuantizer: for<'a> CompressInto<&'a [f32], VectorMut<'a, N>>, { + InvalidDim::check(self.dim(), data.len())?; + let quantized = VectorMut::from_canonical_back_mut(quantized, self.dim()).map_err(ANNError::new)?; self.quantizer .compress_into(data, quantized) .map_err(ANNError::new)?; + Ok(()) } @@ -99,6 +114,8 @@ impl Scalar { use diskann_quantization::AsFunctor; use diskann_utils::ReborrowMut; + InvalidDim::check(self.dim(), query.len())?; + let mut q = Vector::new_boxed(self.dim()); self.quantizer @@ -124,6 +141,7 @@ impl Scalar { } } +/// Metrics supported by the [`Scalar`] quantizer. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Metric { SquaredL2, @@ -131,6 +149,22 @@ pub enum Metric { CosineNormalized, } +impl Metric { + fn as_str(&self) -> &'static str { + match self { + Self::SquaredL2 => "squared-l2", + Self::InnerProduct => "inner-product", + Self::CosineNormalized => "cosine-normalized", + } + } +} + +impl std::fmt::Display for Metric { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + impl TryFrom for Metric { type Error = CosineNotSupported; fn try_from(m: diskann_vector::distance::Metric) -> Result { @@ -145,13 +179,44 @@ impl TryFrom for Metric { } } +impl From for diskann_vector::distance::Metric { + fn from(m: Metric) -> Self { + match m { + Metric::SquaredL2 => Self::L2, + Metric::InnerProduct => Self::InnerProduct, + Metric::CosineNormalized => Self::CosineNormalized, + } + } +} + +/// Cosine is not supported. #[derive(Debug, Error)] -#[error("Metric::Cosine is not supported")] +#[error("Metric::Cosine is not supported for scalar quantization")] #[non_exhaustive] pub struct CosineNotSupported; diskann::convert_error!(CosineNotSupported); +/// TODO: This works around `diskann-quantization` panicking instead of erroring. +#[derive(Debug, Error)] +#[error("scalar quantizer expected dim {} - instead got {}", expected, got)] +struct InvalidDim { + expected: usize, + got: usize, +} + +impl InvalidDim { + fn check(expected: usize, got: usize) -> Result<(), Self> { + if expected != got { + Err(Self { expected, got }) + } else { + Ok(()) + } + } +} + +diskann::convert_error!(InvalidDim); + /// A macro is used to stamp out implementations instead of the crazy generic bounds that /// we *could* use because it's (1) easier to see what exactly is implemented via rustdoc /// and (2) gives us a bit of flexibility to changes details on the backend. @@ -173,7 +238,11 @@ macro_rules! quantizer { { >::distance(self, visitor) } - fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + fn query_distance<'a, V>( + &'a self, + query: &'a [f32], + visitor: V, + ) -> ANNResult where V: quantized::QueryDistanceVisitor<'a>, { @@ -183,7 +252,7 @@ macro_rules! quantizer { }; ($($Ns:literal),+ $(,)?) => { $(quantizer!($Ns);)+ - } + }; } quantizer!(1, 2, 3, 4, 5, 6, 7, 8); @@ -275,5 +344,263 @@ where mod tests { use super::*; + use diskann_utils::assert_contains; + use rand::{ + SeedableRng, + distr::{Distribution, Uniform}, + rngs::StdRng, + }; + + use crate::repr::quantized::test::DynQuantizer; + + const DIM: usize = 5; + const TRIALS: usize = 10; + + #[derive(Debug, Clone, Copy)] + struct Bits(usize); + + impl Bits { + fn value(&self, bucket: usize) -> f32 { + assert!(bucket < self.buckets()); + self.step() * (bucket as f32) - self.shift() + } + + fn buckets(&self) -> usize { + 2usize.pow(self.0.try_into().unwrap()) + } + + fn shift(&self) -> f32 { + self.half_step() * ((self.buckets() - 1) as f32) + } + + fn half_step(&self) -> f32 { + self.step() / 2.0 + } + + fn step(&self) -> f32 { + 1.0 + } + } + + /// For the most part, we rely on the integration tests to make sure that distances + /// actually make sense. + /// + /// These tests ensure that the [`Distance`] and [`QueryDistance`] implementations agree + /// and that errors are thrown as needed. + fn test_scalar( + quantizer: &dyn DynQuantizer, + bits: usize, + rng: &mut StdRng, + ntrials: usize, + ctx: &dyn std::fmt::Display, + ) { + let bits = Bits(bits); + let buckets = bits.buckets(); + let dist = Uniform::new(0usize, buckets).unwrap(); + + let dim = quantizer.dim(); + let bytes = quantizer.bytes().value(); + + 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]; + + for trial in 0..ntrials { + x.iter_mut().for_each(|x| *x = bits.value(dist.sample(rng))); + y.iter_mut().for_each(|y| *y = bits.value(dist.sample(rng))); + + quantizer.quantize(&mut bx, &x).unwrap(); + quantizer.quantize(&mut by, &y).unwrap(); + + // Check distance + let distance = quantizer.distance(); + let from_distance = distance.eval(&bx, &by).unwrap(); + + let query_distance = quantizer.query_distance(&x).unwrap(); + let from_query_distance = query_distance.eval(&by).unwrap(); + + assert_eq!( + from_distance, from_query_distance, + "x = {:?}, bx = {:?}, y = {:?}, by = {:?}, trial = {} -- {}", + x, bx, y, by, trial, ctx, + ); + } + + // Compression Errors + { + // Input size + let x = vec![0.0f32; dim - 1]; + let mut bx = vec![0u8; bytes]; + let err = quantizer.quantize(&mut bx, &x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + let x = vec![0.0f32; dim + 1]; + let err = quantizer.quantize(&mut bx, &x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + // Output size + let x = vec![0.0f32; dim]; + let mut bx = vec![0u8; bytes - 1]; + assert!(quantizer.quantize(&mut bx, &x).is_err()); + + let mut bx = vec![0u8; bytes + 1]; + assert!(quantizer.quantize(&mut bx, &x).is_err()); + + // 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(); + let bx = vec![0u8; bytes]; + let by = vec![0u8; bytes - 1]; + let bz = vec![0u8; bytes + 1]; + + assert!(d.eval(&bx, &by).is_err()); + assert!(d.eval(&bx, &bz).is_err()); + + assert!(d.eval(&by, &bx).is_err()); + assert!(d.eval(&bz, &bx).is_err()); + + assert!(d.eval(&bx, &bx).is_ok()); + } + + // Query Distance + { + let x = vec![0.0f32; dim - 1]; + let err = quantizer.query_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + let x = vec![0.0f32; dim + 1]; + let err = quantizer.query_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + 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(); + + let b = vec![0u8; bytes - 1]; + assert!(d.eval(&b).is_err()); + + let b = vec![0u8; bytes + 1]; + assert!(d.eval(&b).is_err()); + + let b = vec![0u8; bytes]; + assert!(d.eval(&b).is_ok()); + } + } + + /// This quantizer is set up to encode data in the range `[-1, 1)`. + /// + /// As such, its shift is centered around 1.0 and it can encode data at a granularity of + /// `1/2^n` where `n` is the number of bits. + fn test_quantizer(bits: Bits) -> scalar::ScalarQuantizer { + scalar::ScalarQuantizer::new( + (bits.buckets() as f32) / 2.0, + vec![-bits.shift(); DIM], + None, + ) + } + fn run_test() + where + Scalar: quantized::Quantizer, + { + let bits = Bits(N); + let mut rng = StdRng::seed_from_u64(0xcbd17799cff736a7); + + test_scalar( + &Scalar::::new(test_quantizer(bits), Metric::SquaredL2), + N, + &mut rng, + TRIALS, + &"squared-l2", + ); + + test_scalar( + &Scalar::::new(test_quantizer(bits), Metric::InnerProduct), + N, + &mut rng, + TRIALS, + &"inner-product", + ); + + test_scalar( + &Scalar::::new(test_quantizer(bits), Metric::CosineNormalized), + N, + &mut rng, + TRIALS, + &"cosine-normalized", + ); + } + + #[test] + fn test_one_bit() { + run_test::<1>(); + } + + #[test] + fn test_two_bit() { + run_test::<2>(); + } + + #[test] + fn test_three_bit() { + run_test::<3>(); + } + + #[test] + fn test_four_bit() { + run_test::<4>(); + } + + #[test] + fn test_five_bit() { + run_test::<5>(); + } + + #[test] + fn test_six_bit() { + run_test::<6>(); + } + + #[test] + fn test_seven_bit() { + run_test::<7>(); + } + + #[test] + fn test_eight_bit() { + run_test::<8>(); + } + + //--------// + // Metric // + //--------// + + #[test] + fn test_metric() { + use diskann_vector::distance::{Metric as VMetric}; + + // VMetric -> Metric + assert_eq!(Metric::try_from(VMetric::L2).unwrap(), Metric::SquaredL2); + assert_eq!(Metric::try_from(VMetric::InnerProduct).unwrap(), Metric::InnerProduct); + assert_eq!(Metric::try_from(VMetric::CosineNormalized).unwrap(), Metric::CosineNormalized); + let err = Metric::try_from(VMetric::Cosine).unwrap_err(); + assert_contains!(err.to_string(), "Metric::Cosine is not supported for scalar quantization"); + + // Metric -> VMetric + assert_eq!(VMetric::from(Metric::SquaredL2), VMetric::L2); + assert_eq!(VMetric::from(Metric::InnerProduct), VMetric::InnerProduct); + assert_eq!(VMetric::from(Metric::CosineNormalized), VMetric::CosineNormalized); + } } diff --git a/diskann-inmem/src/repr/spherical.rs b/diskann-inmem/src/repr/spherical.rs index cbd6de062..6c94a6ccd 100644 --- a/diskann-inmem/src/repr/spherical.rs +++ b/diskann-inmem/src/repr/spherical.rs @@ -542,7 +542,7 @@ mod tests { use diskann::graph::test::synthetic::Grid; use diskann_quantization::spherical::SupportedMetric; use diskann_utils::{assert_contains, views::rowmajor::MatrixMut}; - use diskann_vector::distance::DistanceProvider; + use diskann_vector::distance::Distance as VDistance; use hashbrown::HashMap; use crate::{ @@ -637,7 +637,7 @@ mod tests { let spherical = config.build().unwrap(); assert_eq!(repr::Representation::id_limit(&spherical), TEST_LIMIT); - assert_eq!(spherical.dim(), grid.dim().into()); + assert_eq!(spherical.dim(), usize::from(grid.dim())); let mut reference = Reference::new(grid.dim().into()); @@ -821,7 +821,7 @@ mod tests { None => panic!("expected a post processor -- {}", ctx), }; - let f = >::distance_comparer( + let f = VDistance::::new( repr::internal::quantization::Metric::from(metric).as_vector_metric(), None, ); diff --git a/diskann-vector/src/distance/distance_provider.rs b/diskann-vector/src/distance/distance_provider.rs index 60cd02079..e6dd1b48f 100644 --- a/diskann-vector/src/distance/distance_provider.rs +++ b/diskann-vector/src/distance/distance_provider.rs @@ -72,10 +72,18 @@ where T: 'static, U: 'static, { - fn new(f: Dispatched2, Unaligned>) -> Self { + fn over(f: Dispatched2, Unaligned>) -> Self { Self { f } } + /// Create a new [`Distance`] using the [`DistanceProvider`] implementation. + pub fn new(metric: Metric, dim: Option) -> Self + where + T: DistanceProvider, + { + T::distance_comparer(metric, dim) + } + /// Compute the distance between `x` and `y`. /// /// The actual distance computed depends on the metric supplied to [`DistanceProvider`]. @@ -360,7 +368,7 @@ where fn try_specialize(&self, arch: A, dim: Option) -> Option> { if let Some(d) = dim { if d == N { - return Some(Distance::new( + return Some(Distance::over( // NOTE: This line here is what actually compiles the specialized kernel. arch.dispatch2::, f32, Unaligned, Unaligned>(), )); @@ -409,7 +417,7 @@ impl Cons { if let Some(f) = self.try_specialize(arch, dim) { f } else { - Distance::new(arch.dispatch2::, Unaligned>()) + Distance::over(arch.dispatch2::, Unaligned>()) } } } From ef149c687216ce6ca1ab1af76e4c2233231763d7 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 14 Oct 2026 16:17:38 -0700 Subject: [PATCH 4/6] Polish up tests. --- diskann-benchmark/example/inmem2.json | 44 ++++ diskann-benchmark/src/index/inmem2.rs | 33 +-- diskann-inmem/integration/index/runner.rs | 2 +- .../graph/scalar/four-bit-ip-baseline.json | 36 +-- diskann-inmem/src/repr/quantized/mod.rs | 173 +++++++++++---- diskann-inmem/src/repr/quantized/scalar.rs | 208 ++++++++++++++++-- 6 files changed, 402 insertions(+), 94 deletions(-) diff --git a/diskann-benchmark/example/inmem2.json b/diskann-benchmark/example/inmem2.json index 5a67781a5..10cd6d8d1 100644 --- a/diskann-benchmark/example/inmem2.json +++ b/diskann-benchmark/example/inmem2.json @@ -130,6 +130,50 @@ ] } } + }, + { + "type": "inmem2", + "content": { + "build": { + "alpha": 1.2, + "l_build": 100, + "max_degree": 32, + "num_threads": 4, + "pruned_degree": 28 + }, + "data": { + "data": "disk_index_siftsmall_learn_256pts_data.fbin", + "data_type": "float32", + "distance": "squared_l2" + }, + "quantization": { + "scalar": { + "standard_deviations": 1.5, + "bits": "eight" + } + }, + "search": { + "groundtruth": "disk_index_10pts_idx_uint32_exact_ground_truth.bin", + "num_threads": [ + 4 + ], + "queries": "disk_index_sample_query_10pts.fbin", + "reps": 3, + "runs": [ + { + "recall_k": 10, + "search_l": [ + 10, + 20, + 30, + 40, + 50 + ], + "search_n": 10 + } + ] + } + } } ] } diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs index 2d77b4f98..6ba51879d 100644 --- a/diskann-benchmark/src/index/inmem2.rs +++ b/diskann-benchmark/src/index/inmem2.rs @@ -140,6 +140,7 @@ mod dto { } #[derive(Debug, Serialize, Deserialize)] + #[serde(rename_all = "kebab-case")] pub(super) enum ScalarBits { One, Two, @@ -650,7 +651,10 @@ struct Scalar { impl Scalar { fn from_raw(raw: dto::Scalar) -> anyhow::Result { - let dto::Scalar { bits, standard_deviations } = raw; + let dto::Scalar { + bits, + standard_deviations, + } = raw; Ok(Self { bits: ScalarBits::from_raw(bits), @@ -658,11 +662,9 @@ impl Scalar { }) } - fn train( - &self, - data: rowmajor::Ref<'_, f32>, - ) -> anyhow::Result { - let trainer = dq::scalar::train::ScalarQuantizationParameters::new(self.standard_deviations); + fn train(&self, data: rowmajor::Ref<'_, f32>) -> anyhow::Result { + let trainer = + dq::scalar::train::ScalarQuantizationParameters::new(self.standard_deviations); let quantizer = trainer.train(data); Ok(quantizer) } @@ -670,7 +672,10 @@ impl Scalar { impl std::fmt::Display for Scalar { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let Self { bits, standard_deviations } = self; + let Self { + bits, + standard_deviations, + } = self; let std = standard_deviations.into_inner(); @@ -682,7 +687,6 @@ impl std::fmt::Display for Scalar { } } - //-------------// // StaticBuild // //-------------// @@ -1325,7 +1329,7 @@ impl Benchmark for ScalarBuild { checkpoint: Checkpoint<'_>, mut output: &mut dyn Output, ) -> anyhow::Result<()> { - use repr::{Quantized, quantized::scalar::Scalar}; + use repr::{quantized::scalar::Scalar, Quantized}; writeln!(output, "{input}\n")?; @@ -1344,7 +1348,7 @@ impl Benchmark for ScalarBuild { // Compute the medoid of the dataset as the single start point. - let metric = repr::quantized::scalar::Metric::try_from(input.data.distance)?; + let metric = repr::quantized::scalar::Metric::try_from(input.data.distance)?; let capacity = Capacity::new(num_points); let max_degree = MaxDegree::new(input.build.config.max_degree().get()); @@ -1359,28 +1363,28 @@ impl Benchmark for ScalarBuild { capacity, max_degree, start, - rerank + rerank, )?, ScalarBits::Two => Quantized::config( Scalar::<2>::new(quantizer, metric), capacity, max_degree, start, - rerank + rerank, )?, ScalarBits::Four => Quantized::config( Scalar::<4>::new(quantizer, metric), capacity, max_degree, start, - rerank + rerank, )?, ScalarBits::Eight => Quantized::config( Scalar::<8>::new(quantizer, metric), capacity, max_degree, start, - rerank + rerank, )?, }; @@ -1449,7 +1453,6 @@ impl Benchmark for ScalarBuild { } } - fn _knn( runner: &dyn crate::index::search::knn::Knn, groundtruth: &dyn benchmark_core::recall::Rows, diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs index fc4602054..b6968eb3c 100644 --- a/diskann-inmem/integration/index/runner.rs +++ b/diskann-inmem/integration/index/runner.rs @@ -757,7 +757,7 @@ impl Test { }; // Step 1: Train a generic quantizer. - let std = Positive::::new(1.5).expect("1.5 is positive"); + let std = Positive::::new(1.5)?; let trainer = qscalar::train::ScalarQuantizationParameters::new(std); let quantizer = trainer.train(data); diff --git a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json index 98311d17d..81d7d4133 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json @@ -60,18 +60,18 @@ "counters": { "append_neighbors": 0, "distance": 0, - "get_neighbors": 5219, - "get_vector": 49096, - "query_distance": 49096, + "get_neighbors": 5208, + "get_vector": 49130, + "query_distance": 49130, "set_neighbors": 0, "set_vector": 0 }, "misc": { - "cmps": 49096, - "hops": 5219 + "cmps": 49130, + "hops": 5208 }, "recall": { - "average": 0.764, + "average": 0.769, "num_queries": 100, "recall_k": 10, "recall_n": 10 @@ -81,18 +81,18 @@ "counters": { "append_neighbors": 0, "distance": 0, - "get_neighbors": 5250, - "get_vector": 49283, - "query_distance": 49283, + "get_neighbors": 5240, + "get_vector": 49346, + "query_distance": 49346, "set_neighbors": 0, "set_vector": 0 }, "misc": { - "cmps": 49283, - "hops": 5250 + "cmps": 49346, + "hops": 5240 }, "recall": { - "average": 0.764, + "average": 0.769, "num_queries": 100, "recall_k": 10, "recall_n": 10 @@ -102,18 +102,18 @@ "counters": { "append_neighbors": 0, "distance": 0, - "get_neighbors": 10227, - "get_vector": 76057, - "query_distance": 76057, + "get_neighbors": 10228, + "get_vector": 75966, + "query_distance": 75966, "set_neighbors": 0, "set_vector": 0 }, "misc": { - "cmps": 76057, - "hops": 10227 + "cmps": 75966, + "hops": 10228 }, "recall": { - "average": 0.764, + "average": 0.769, "num_queries": 100, "recall_k": 10, "recall_n": 10 diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index 9614d7cd0..fc0cb0f2d 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -34,11 +34,35 @@ pub mod scalar; /// Plug-in for a generalized quantizer compatible with [`Quantized`]. /// /// This is applicable for quantizers that have a single, non-split, flat representation for -/// quantized vectors with a fixed number of bytes and the same distance computation for insertion -/// searches and standard serving searches. +/// quantized vectors with a fixed number of bytes. /// -/// Distance functors provided by [`Self::distance`] and [`Self::query_distance`] will be -/// inlined into a higher-level kernel, allowing implementations to specialize distances. +/// Three different distance types are allowed: +/// +/// * [`Self::query_distance`]: The distance computation used for general queries over an +/// index. Implementations may assume that a query is fixed for the duration of this search +/// and the distances used will not participate in any stage of index construction. +/// +/// Queries may come from a different distributions than the standard data corpus. +/// +/// * [`Self::insert_distance`]: Distance computations used for the search phase of an +/// insertion. This is like [`Self::query_distance`] but with some important differences: +/// +/// * The item being inserted can be assumed to belong to the corpus of data upon which +/// the quantizer is trained. +/// +/// * Distance yielded from this computer may be mixed with those returned by +/// [`Self::distance`] for index construction and should be compatible. +/// +/// * [`Self::distance`]: General distance computations between quantized items. +/// Unlike the [`QueryDistance`] style methods, the left and right-hand sides of the +/// arguments to the returned [`Distance`] can both change. +/// +/// Distance values returned from this method should be compatible with those yielded +/// by [`Self::insert_distance`]. +/// +/// Distance functors provided by [`Self::distance`], [`Self::insert_distance`], and +/// [`Self::query_distance`] will be inlined into a higher-level kernel, allowing +/// implementations to specialize distances. pub trait Quantizer: std::fmt::Debug + Send + Sync + 'static { /// The number of bytes required for each quantized vector. /// @@ -71,6 +95,19 @@ pub trait Quantizer: std::fmt::Debug + Send + Sync + 'static { fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult where V: QueryDistanceVisitor<'a>; + + /// Return the distance used during the search phase of insert. Implementations may use + /// this to provide distance computations that are more compatible with [`Self::distance`]. + /// + /// Implementations may assume that `query.len() == self.dim()`. + /// + /// This defaults to using [`Self::query_distance`]. + fn insert_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: QueryDistanceVisitor<'a>, + { + self.query_distance(query, visitor) + } } /// Return the distance between two raw byte slices. @@ -430,7 +467,7 @@ impl Quantized { counters: LocalCounters<'a>, args: AccessorArgs, ) -> ANNResult> { - let AccessorArgs { rerank_if_enabled } = args; + let AccessorArgs { insert } = args; // NOTE: `ProcessInto` panics on a query dimension mismatch. // @@ -449,6 +486,7 @@ impl Quantized { let (expand_beam, post_process) = self.store.guard(|cons, guard| -> ANNResult<_> { // This will result in some wasted work if `expand_beam` errors, but that should // be a relatively rare occurrence. + let rerank_if_enabled = !insert; let post_process = rerank_if_enabled .then(|| { self.reranker @@ -456,9 +494,18 @@ impl Quantized { }) .flatten(); - let expand_beam = - self.quantizer - .expand_beam(query, cons.first().reader(guard), self.lookahead)?; + let kind = if insert { + inner::ExpandBeamKind::Insert + } else { + inner::ExpandBeamKind::Query + }; + + let expand_beam = self.quantizer.expand_beam( + query, + cons.first().reader(guard), + self.lookahead, + kind, + )?; Ok((expand_beam, post_process)) })??; @@ -476,7 +523,7 @@ impl Quantized { #[derive(Debug)] struct AccessorArgs { - rerank_if_enabled: bool, + insert: bool, } repr::internal::macros::representation!(Quantized); @@ -521,14 +568,7 @@ impl repr::Search for Quantized { provider: &'a (dyn std::any::Any + Send + Sync), counters: LocalCounters<'a>, ) -> ANNResult> { - self.create_accessor( - query, - provider, - counters, - AccessorArgs { - rerank_if_enabled: true, - }, - ) + self.create_accessor(query, provider, counters, AccessorArgs { insert: false }) } } @@ -539,14 +579,7 @@ impl repr::Insert for Quantized { provider: &'a (dyn std::any::Any + Send + Sync), counters: LocalCounters<'a>, ) -> ANNResult> { - self.create_accessor( - query, - provider, - counters, - AccessorArgs { - rerank_if_enabled: false, - }, - ) + self.create_accessor(query, provider, counters, AccessorArgs { insert: true }) } fn prune_accessor<'a>( @@ -578,6 +611,12 @@ mod inner { use crate::{num::Bytes, prefetch, repr, store::intrusive}; + #[derive(Debug)] + pub(super) enum ExpandBeamKind { + Insert, + Query, + } + #[derive(Debug)] pub(super) struct Quantizer { quantizer: Box, @@ -619,8 +658,9 @@ mod inner { query: &'a [f32], reader: intrusive::Reader<'a>, lookahead: Option, + kind: ExpandBeamKind, ) -> ANNResult> { - self.quantizer.expand_beam(query, reader, lookahead) + self.quantizer.expand_beam(query, reader, lookahead, kind) } pub(super) fn prune<'a>( @@ -640,6 +680,7 @@ mod inner { query: &'a [f32], reader: intrusive::Reader<'a>, lookahead: Option, + kind: ExpandBeamKind, ) -> ANNResult>; fn prune<'a>(&'a self, reader: intrusive::Reader<'a>) -> Box; @@ -725,9 +766,13 @@ mod inner { query: &'a [f32], reader: intrusive::Reader<'a>, lookahead: Option, + kind: ExpandBeamKind, ) -> ANNResult> { - self.0 - .query_distance(query, QueryDistanceVisitor { reader, lookahead }) + let visitor = QueryDistanceVisitor { reader, lookahead }; + match kind { + ExpandBeamKind::Insert => self.0.insert_distance(query, visitor), + ExpandBeamKind::Query => self.0.query_distance(query, visitor), + } } fn prune<'a>(&'a self, reader: intrusive::Reader<'a>) -> Box { self.0.distance(DistanceVisitor { reader }) @@ -774,6 +819,10 @@ mod test { fn distance<'a>(&'a self) -> Box; fn query_distance<'a>(&'a self, query: &'a [f32]) -> ANNResult>; + fn insert_distance<'a>( + &'a self, + query: &'a [f32], + ) -> ANNResult>; } impl DynQuantizer for Q @@ -798,6 +847,12 @@ mod test { ) -> ANNResult> { Quantizer::query_distance(self, query, Boxit) } + fn insert_distance<'a>( + &'a self, + query: &'a [f32], + ) -> ANNResult> { + Quantizer::insert_distance(self, query, Boxit) + } } } @@ -826,6 +881,13 @@ mod tests { ///////////////////// /// A simple test quantizer that tries its best to convert `f32` to `i8`. + /// + /// To test that [`Quantizer::query_distance`] and [`Quantizer::insert_distance`] are + /// called correctly, we assume that queries are all shifted by `[5.0, 5.0]` and need + /// to be recentered around zero. + /// + /// In other words, a query may come in as `[5.0, 5.0]` but for purposes of all distance + /// computations, it should be treated as `[0.0, 0.0]`. #[derive(Debug, Clone, Copy)] struct Round { dim: usize, @@ -833,6 +895,8 @@ mod tests { } impl Round { + const QUERY_OFFSET: f32 = 5.0; + fn new(dim: usize, metric: Metric) -> Self { let vmetric = match metric { Metric::SquaredL2 => VMetric::L2, @@ -875,6 +939,19 @@ mod tests { visitor.visit(RoundDistance(*self)) } fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + where + V: QueryDistanceVisitor<'a>, + { + let distance = RoundQueryDistance { + round: *self, + query: query + .iter() + .map(|v| (*v - Self::QUERY_OFFSET) as i8) + .collect(), + }; + Ok(visitor.visit(distance)) + } + fn insert_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult where V: QueryDistanceVisitor<'a>, { @@ -1008,7 +1085,8 @@ mod tests { "{ctx}" ); - let query = [10.0, -10.0]; + let insert_vector = [10.0, -10.0]; + let query_vector = insert_vector.map(|q| q + Round::QUERY_OFFSET); // Generate a couple of data points in the dataset. let i0 = LogicalId(2); @@ -1041,10 +1119,21 @@ mod tests { // The Round quantizer is exact for these integral values. let f = VDistance::::new(metric.as_internal_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)), + // True distances between the insert vector and the dataset vectors. + let insert_distances = HashMap::from_iter([ + (s0, f.call(&insert_vector, v0)), + (s1, f.call(&insert_vector, v1)), + (s2, f.call(&insert_vector, v2)), + ]); + + // Distances between the un-shifted query vector and the dataset vectors. + let query_distances = insert_distances.clone(); + + // THe true distances between the shifted query vector and the dataset vectors. + let true_query_distances = HashMap::from_iter([ + (s0, f.call(&query_vector, v0)), + (s1, f.call(&query_vector, v1)), + (s2, f.call(&query_vector, v2)), ]); // Insert @@ -1052,7 +1141,7 @@ mod tests { let counters = Counters::new(); let mut sa = repr::Insert::insert_search_accessor( quantized, - query.as_slice(), + insert_vector.as_slice(), &(), counters.local(), ) @@ -1067,7 +1156,7 @@ mod tests { test_expand_beam( sa.get_expand_beam(), TEST_LIMIT, - distances.clone(), + insert_distances, &[s0, s1, s3_deleted, s2, s4_deleted], ctx, ); @@ -1076,14 +1165,18 @@ mod tests { // Search { let counters = Counters::new(); - let mut sa = - repr::Search::search_accessor(quantized, query.as_slice(), &(), counters.local()) - .unwrap(); + let mut sa = repr::Search::search_accessor( + quantized, + query_vector.as_slice(), + &(), + counters.local(), + ) + .unwrap(); test_expand_beam( sa.get_expand_beam(), TEST_LIMIT, - distances.clone(), + query_distances, &[s0, s1, s3_deleted, s2, s4_deleted], ctx, ); @@ -1102,7 +1195,7 @@ mod tests { repr::internal::quantization::rerank::test_rerank( post_process, - distances, + true_query_distances, &[s0, s1, s3_deleted, s2, s4_deleted], ctx, ); diff --git a/diskann-inmem/src/repr/quantized/scalar.rs b/diskann-inmem/src/repr/quantized/scalar.rs index 43c3c6acb..299ae0b02 100644 --- a/diskann-inmem/src/repr/quantized/scalar.rs +++ b/diskann-inmem/src/repr/quantized/scalar.rs @@ -15,9 +15,13 @@ use diskann_quantization::{ CompensatedSquaredL2 as CSquaredL2, }, }; +use diskann_vector::{Norm, norm::FastL2Norm}; use thiserror::Error; -use crate::{num::Bytes, repr::quantized}; +use crate::{ + num::Bytes, + repr::{internal::Calf, quantized}, +}; // Help bring the ridiculous names under control. type VectorRef<'a, const N: usize> = scalar::CompensatedVectorRef<'a, N>; @@ -102,7 +106,12 @@ impl Scalar { } } - fn query_distance<'a, V>(&'a self, query: &'a [f32], visitor: V) -> ANNResult + fn query_distance<'a, V>( + &'a self, + query: &'a [f32], + visitor: V, + kind: QueryDistanceKind, + ) -> ANNResult where V: quantized::QueryDistanceVisitor<'a>, Unsigned: Representation, @@ -118,8 +127,28 @@ impl Scalar { let mut q = Vector::new_boxed(self.dim()); + // TODO: `diskann-quantization` needs to expose an asymmetric distance API. + // + // We really shouldn't have to handle this externally. + let query = (|| match kind { + QueryDistanceKind::Query => { + if matches!(self.metric, Metric::InnerProduct) + && let Some(mean_norm) = self.quantizer.mean_norm() + { + let norm = (FastL2Norm).evaluate(query); + if norm != 0.0 { + let scale = mean_norm / norm; + let rescaled = query.iter().map(|q| q * scale).collect(); + return Calf::Owned(rescaled); + } + } + Calf::Borrowed(query) + } + QueryDistanceKind::Insert => Calf::Borrowed(query), + })(); + self.quantizer - .compress_into(query, q.reborrow_mut()) + .compress_into(&*query, q.reborrow_mut()) .map_err(ANNError::new)?; let visited = match self.metric { @@ -141,6 +170,12 @@ impl Scalar { } } +#[derive(Debug, Clone, Copy)] +enum QueryDistanceKind { + Insert, + Query, +} + /// Metrics supported by the [`Scalar`] quantizer. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Metric { @@ -246,7 +281,17 @@ macro_rules! quantizer { where V: quantized::QueryDistanceVisitor<'a>, { - >::query_distance(self, query, visitor) + >::query_distance(self, query, visitor, QueryDistanceKind::Query) + } + fn insert_distance<'a, V>( + &'a self, + query: &'a [f32], + visitor: V, + ) -> ANNResult + where + V: quantized::QueryDistanceVisitor<'a>, + { + >::query_distance(self, query, visitor, QueryDistanceKind::Insert) } } }; @@ -355,6 +400,7 @@ mod tests { const DIM: usize = 5; const TRIALS: usize = 10; + const SCALE: f32 = 10.0; #[derive(Debug, Clone, Copy)] struct Bits(usize); @@ -390,6 +436,7 @@ mod tests { fn test_scalar( quantizer: &dyn DynQuantizer, bits: usize, + metric: Metric, rng: &mut StdRng, ntrials: usize, ctx: &dyn std::fmt::Display, @@ -407,8 +454,16 @@ mod tests { let mut y = vec![0.0f32; dim]; let mut by = vec![0u8; bytes]; + let mut x_scaled = vec![0.0f32; dim]; + for trial in 0..ntrials { - x.iter_mut().for_each(|x| *x = bits.value(dist.sample(rng))); + // One the first trial, use `0` for `x` to make sure we don't try to scale a + // zeroed vector. + if trial == 0 { + x.fill(0.0); + } else { + x.iter_mut().for_each(|x| *x = bits.value(dist.sample(rng))); + } y.iter_mut().for_each(|y| *y = bits.value(dist.sample(rng))); quantizer.quantize(&mut bx, &x).unwrap(); @@ -418,14 +473,40 @@ mod tests { let distance = quantizer.distance(); let from_distance = distance.eval(&bx, &by).unwrap(); - let query_distance = quantizer.query_distance(&x).unwrap(); - let from_query_distance = query_distance.eval(&by).unwrap(); + let insert_distance = quantizer.insert_distance(&x).unwrap(); + let from_insert_distance = insert_distance.eval(&by).unwrap(); assert_eq!( - from_distance, from_query_distance, + from_distance, from_insert_distance, "x = {:?}, bx = {:?}, y = {:?}, by = {:?}, trial = {} -- {}", x, bx, y, by, trial, ctx, ); + + // Check scaling. + // + // This mimics the scaling logic. + match metric { + Metric::InnerProduct => { + let norm = (FastL2Norm).evaluate(&*x); + let norm = if norm == 0.0 { 1.0 } else { norm }; + let scale = SCALE / norm; + std::iter::zip(x_scaled.iter_mut(), x.iter()) + .for_each(|(xs, x)| *xs = (*x) * scale) + } + Metric::SquaredL2 | Metric::CosineNormalized => x_scaled.copy_from_slice(&x), + }; + + let insert_distance = quantizer.insert_distance(&x_scaled).unwrap(); + let from_insert_distance = insert_distance.eval(&by).unwrap(); + + let query_distance = quantizer.query_distance(&x).unwrap(); + let from_query_distance = query_distance.eval(&by).unwrap(); + + assert_eq!( + from_insert_distance, from_query_distance, + "x = {:?}, bx = {:?}, x_scaled = {:?}, y = {:?}, by = {:?}, trial = {} -- {}", + x, bx, x_scaled, y, by, trial, ctx, + ); } // Compression Errors @@ -497,6 +578,33 @@ mod tests { let b = vec![0u8; bytes]; assert!(d.eval(&b).is_ok()); } + + // Insert Distance + { + let x = vec![0.0f32; dim - 1]; + let err = quantizer.insert_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + let x = vec![0.0f32; dim + 1]; + let err = quantizer.insert_distance(&x).unwrap_err().to_string(); + assert_contains!(err, "expected dim 5"); + + 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(); + + let b = vec![0u8; bytes - 1]; + assert!(d.eval(&b).is_err()); + + let b = vec![0u8; bytes + 1]; + assert!(d.eval(&b).is_err()); + + let b = vec![0u8; bytes]; + assert!(d.eval(&b).is_ok()); + } } /// This quantizer is set up to encode data in the range `[-1, 1)`. @@ -507,7 +615,7 @@ mod tests { scalar::ScalarQuantizer::new( (bits.buckets() as f32) / 2.0, vec![-bits.shift(); DIM], - None, + Some(SCALE), ) } @@ -518,25 +626,40 @@ mod tests { let bits = Bits(N); let mut rng = StdRng::seed_from_u64(0xcbd17799cff736a7); + let quantizer = Scalar::::new(test_quantizer(bits), Metric::SquaredL2); + assert_eq!(quantizer.nbits(), N); + assert_eq!(quantizer.metric(), Metric::SquaredL2); + test_scalar( - &Scalar::::new(test_quantizer(bits), Metric::SquaredL2), - N, + &quantizer, + quantizer.nbits(), + quantizer.metric(), &mut rng, TRIALS, &"squared-l2", ); + let quantizer = Scalar::::new(test_quantizer(bits), Metric::InnerProduct); + assert_eq!(quantizer.nbits(), N); + assert_eq!(quantizer.metric(), Metric::InnerProduct); + test_scalar( - &Scalar::::new(test_quantizer(bits), Metric::InnerProduct), - N, + &quantizer, + quantizer.nbits(), + quantizer.metric(), &mut rng, TRIALS, &"inner-product", ); + let quantizer = Scalar::::new(test_quantizer(bits), Metric::CosineNormalized); + assert_eq!(quantizer.nbits(), N); + assert_eq!(quantizer.metric(), Metric::CosineNormalized); + test_scalar( - &Scalar::::new(test_quantizer(bits), Metric::CosineNormalized), - N, + &quantizer, + quantizer.nbits(), + quantizer.metric(), &mut rng, TRIALS, &"cosine-normalized", @@ -583,24 +706,69 @@ mod tests { run_test::<8>(); } + #[test] + fn test_no_scaling() { + let bits = Bits(8); + + let quantizer = Scalar::<8>::new( + scalar::ScalarQuantizer::new( + (bits.buckets() as f32) / 2.0, + vec![-bits.shift(); DIM], + None, + ), + Metric::InnerProduct, + ); + + let dyn_quantizer: &dyn DynQuantizer = &quantizer; + + let dim = dyn_quantizer.dim(); + let bytes = dyn_quantizer.bytes().value(); + + let x = vec![-1.0, -0.5, -0.5, 0.5, 1.0]; + let y = vec![-1.0f32; dim]; + let mut by = vec![0u8; bytes]; + + dyn_quantizer.quantize(&mut by, &y).unwrap(); + + let insert_distance = dyn_quantizer.insert_distance(&x).unwrap(); + let from_insert_distance = insert_distance.eval(&by).unwrap(); + + let query_distance = dyn_quantizer.query_distance(&x).unwrap(); + let from_query_distance = query_distance.eval(&by).unwrap(); + + assert_eq!(from_insert_distance, from_query_distance); + } + //--------// // Metric // //--------// #[test] fn test_metric() { - use diskann_vector::distance::{Metric as VMetric}; + use diskann_vector::distance::Metric as VMetric; // VMetric -> Metric assert_eq!(Metric::try_from(VMetric::L2).unwrap(), Metric::SquaredL2); - assert_eq!(Metric::try_from(VMetric::InnerProduct).unwrap(), Metric::InnerProduct); - assert_eq!(Metric::try_from(VMetric::CosineNormalized).unwrap(), Metric::CosineNormalized); + assert_eq!( + Metric::try_from(VMetric::InnerProduct).unwrap(), + Metric::InnerProduct + ); + assert_eq!( + Metric::try_from(VMetric::CosineNormalized).unwrap(), + Metric::CosineNormalized + ); let err = Metric::try_from(VMetric::Cosine).unwrap_err(); - assert_contains!(err.to_string(), "Metric::Cosine is not supported for scalar quantization"); + assert_contains!( + err.to_string(), + "Metric::Cosine is not supported for scalar quantization" + ); // Metric -> VMetric assert_eq!(VMetric::from(Metric::SquaredL2), VMetric::L2); assert_eq!(VMetric::from(Metric::InnerProduct), VMetric::InnerProduct); - assert_eq!(VMetric::from(Metric::CosineNormalized), VMetric::CosineNormalized); + assert_eq!( + VMetric::from(Metric::CosineNormalized), + VMetric::CosineNormalized + ); } } From eeca7f939d40bde6822e3bbadc2158887760e34d Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 14 Oct 2026 16:20:21 -0700 Subject: [PATCH 5/6] Remove superflous addition. --- diskann-quantization/src/num.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/diskann-quantization/src/num.rs b/diskann-quantization/src/num.rs index ba2ba7a68..8fe974a25 100644 --- a/diskann-quantization/src/num.rs +++ b/diskann-quantization/src/num.rs @@ -50,8 +50,6 @@ where // SAFETY: 1.0 is positive. pub const POSITIVE_ONE_F32: Positive = unsafe { Positive::new_unchecked(1.0) }; -// SAFETY: 1.0 is positive. -pub const POSITIVE_ONE_F64: Positive = unsafe { Positive::new_unchecked(1.0) }; #[derive(Debug, Clone, Copy, PartialEq)] #[repr(transparent)] From 70b10412d7959e32d5b9718fe899428d8f1c6211 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Wed, 21 Oct 2026 09:38:42 -0700 Subject: [PATCH 6/6] Undo some unnecessary shuffles. --- diskann-inmem/src/repr/internal/quantization/rerank.rs | 5 +++-- diskann-inmem/src/repr/quantized/mod.rs | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/diskann-inmem/src/repr/internal/quantization/rerank.rs b/diskann-inmem/src/repr/internal/quantization/rerank.rs index 5bc602528..95dfb3d29 100644 --- a/diskann-inmem/src/repr/internal/quantization/rerank.rs +++ b/diskann-inmem/src/repr/internal/quantization/rerank.rs @@ -221,14 +221,15 @@ mod tests { use std::assert_matches; + use diskann::utils::IntoUsize; + use diskann_utils::assert_contains; + use crate::{ counters::Counters, num::{Capacity, LogicalId, MaxDegree, SlotId}, repr::test::Reference, store::Store, }; - use diskann::utils::IntoUsize; - use diskann_utils::assert_contains; const TEST_DIM: usize = 1; diff --git a/diskann-inmem/src/repr/quantized/mod.rs b/diskann-inmem/src/repr/quantized/mod.rs index fc0cb0f2d..172719219 100644 --- a/diskann-inmem/src/repr/quantized/mod.rs +++ b/diskann-inmem/src/repr/quantized/mod.rs @@ -876,6 +876,7 @@ mod tests { num::{LogicalId, SlotId}, repr::test::{Reference, test_expand_beam, test_prune}, }; + ///////////////////// // Quantized Tests // /////////////////////