diskann_disk/search/pq/
pq_scratch.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5//! Aligned allocator
6
7use diskann::{ANNError, ANNResult};
8
9use diskann_quantization::alloc::{AlignedAllocator, Poly};
10
11#[derive(Debug)]
12/// PQ scratch
13pub struct PQScratch {
14    /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids.
15    /// This is used to store the distance between each chunk in the query vector to each centroid, which is why the length is num of centroids * num of chunks
16    pub aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>,
17
18    /// Aligned dist scratch, must be at least diskann MAX_DEGREE
19    /// This is used to temporarily save the pq distance between query vector to the candidate vectors.
20    pub aligned_dist_scratch: Poly<[f32], AlignedAllocator>,
21
22    /// Aligned pq coord scratch, must be at least [N_CHUNKS * MAX_DEGREE]
23    /// This is used to store the pq coordinates of the candidate vectors.
24    pub aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>,
25
26    /// Query scratch buffer stored as `f32`, sized by the PQ table's logical dimension.
27    /// `set` populates it from a caller-provided `&[f32]`; `PQTable::preprocess_query` can
28    /// then rotate or otherwise preprocess it.
29    pub query_scratch: Vec<f32>,
30}
31
32impl PQScratch {
33    /// Create a new pq scratch.
34    ///
35    /// `dim` is the PQ table's logical dimension (`PQData::get_dim()`); the
36    /// internal `query_scratch` buffer is sized to exactly this many `f32` slots.
37    pub fn new(
38        graph_degree: usize,
39        dim: usize,
40        num_pq_chunks: usize,
41        num_centers: usize,
42    ) -> ANNResult<Self> {
43        let aligned_pq_coord_scratch =
44            Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
45                .map_err(ANNError::log_index_error)?;
46        let aligned_pqtable_dist_scratch =
47            Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128)
48                .map_err(ANNError::log_index_error)?;
49        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
50            .map_err(ANNError::log_index_error)?;
51        let query_scratch = vec![0.0f32; dim];
52
53        Ok(Self {
54            aligned_pqtable_dist_scratch,
55            aligned_dist_scratch,
56            aligned_pq_coord_scratch,
57            query_scratch,
58        })
59    }
60
61    /// Copy the first `dim` elements of `query` into `query_scratch`.
62    ///
63    /// `query` must already be in full-precision `f32` representation; quantized
64    /// inputs (e.g. `MinMaxElement`) should be decoded via `VectorRepr::as_f32`
65    /// at the caller boundary before invoking this method.
66    ///
67    /// Accepts oversized `query` (only the first `dim` elements are used) for
68    /// backwards compatibility with callers that hold alignment-padded buffers.
69    /// Returns `DimensionMismatchError` if `query.len() < query_scratch.len()`.
70    pub fn set(&mut self, query: &[f32]) -> ANNResult<()> {
71        let dim = self.query_scratch.len();
72        if query.len() < dim {
73            return Err(ANNError::log_dimension_mismatch_error(format!(
74                "PQScratch::set: expected query of length >= {dim}, got {}",
75                query.len()
76            )));
77        }
78        self.query_scratch.copy_from_slice(&query[..dim]);
79        Ok(())
80    }
81
82    /// Return the largest number of PQ vectors whose distances can be computed using this
83    /// scratch data structure.
84    pub(crate) fn max_vectors(&self) -> usize {
85        self.aligned_dist_scratch.len()
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use diskann_quantization::num::PowerOfTwo;
92    use rstest::rstest;
93
94    use super::PQScratch;
95
96    #[rstest]
97    #[case(512, 8, 128, 256)] // default test case
98    #[case(59, 16, 37, 41)] // not multiple of 256
99    fn test_pq_scratch(
100        #[case] graph_degree: usize,
101        #[case] dim: usize,
102        #[case] num_pq_chunks: usize,
103        #[case] num_centers: usize,
104    ) {
105        let mut pq_scratch: PQScratch =
106            PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap();
107
108        assert_eq!(
109            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
110            0
111        );
112        assert_eq!(
113            (pq_scratch.aligned_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
114            0
115        );
116        assert_eq!(
117            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
118            0
119        );
120
121        assert_eq!(pq_scratch.max_vectors(), graph_degree);
122
123        // Test set() method
124        let query: Vec<f32> = (1..=dim).map(|i| i as f32).collect();
125        pq_scratch.set(&query).unwrap();
126
127        (0..query.len()).for_each(|i| {
128            assert_eq!(pq_scratch.query_scratch[i], query[i]);
129        });
130    }
131}
diskann_disk/search/pq/pq_scratch.rs

diskann_disk/search/pq/
pq_scratch.rs
