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::{error::IntoANNResult, utils::VectorRepr, 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`. `set` initializes it by copying/converting the
27    /// raw query values; `PQTable.PreprocessQuery` can then rotate or otherwise preprocess it.
28    pub rotated_query: Vec<f32>,
29}
30
31impl PQScratch {
32    /// Create a new pq scratch
33    pub fn new(
34        graph_degree: usize,
35        dim: usize,
36        num_pq_chunks: usize,
37        num_centers: usize,
38    ) -> ANNResult<Self> {
39        let aligned_pq_coord_scratch =
40            Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
41                .map_err(ANNError::log_index_error)?;
42        let aligned_pqtable_dist_scratch =
43            Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128)
44                .map_err(ANNError::log_index_error)?;
45        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
46            .map_err(ANNError::log_index_error)?;
47        let rotated_query = vec![0.0f32; dim];
48
49        Ok(Self {
50            aligned_pqtable_dist_scratch,
51            aligned_dist_scratch,
52            aligned_pq_coord_scratch,
53            rotated_query,
54        })
55    }
56
57    /// Copy `query` into `rotated_query`, converting to `f32`.
58    ///
59    /// `dim` is the element count in the `T` representation. The decompressed
60    /// `f32` length returned by `T::as_f32` may differ (e.g. `MinMaxElement`
61    /// expands to more `f32`s than its raw element count), so the destination
62    /// slice is sized by that actual length.
63    ///
64    /// Returns `DimensionMismatchError` if `dim > query.len()` or the
65    /// decompressed vector does not fit in `rotated_query`.
66    pub fn set<T: VectorRepr>(&mut self, dim: usize, query: &[T]) -> ANNResult<()> {
67        if dim > query.len() {
68            return Err(ANNError::log_dimension_mismatch_error(format!(
69                "PQScratch::set: expected query of length >= {dim}, got {}",
70                query.len()
71            )));
72        }
73        let query = T::as_f32(&query[..dim]).into_ann_result()?;
74        if query.len() > self.rotated_query.len() {
75            return Err(ANNError::log_dimension_mismatch_error(format!(
76                "PQScratch::set: decompressed query of length {} does not fit rotated_query buffer of length {}",
77                query.len(),
78                self.rotated_query.len()
79            )));
80        }
81        self.rotated_query[..query.len()].copy_from_slice(&query);
82        Ok(())
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use diskann_quantization::num::PowerOfTwo;
89    use rstest::rstest;
90
91    use super::PQScratch;
92
93    #[rstest]
94    #[case(512, 8, 128, 256)] // default test case
95    #[case(59, 16, 37, 41)] // not multiple of 256
96    fn test_pq_scratch(
97        #[case] graph_degree: usize,
98        #[case] dim: usize,
99        #[case] num_pq_chunks: usize,
100        #[case] num_centers: usize,
101    ) {
102        let mut pq_scratch: PQScratch =
103            PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap();
104
105        assert_eq!(
106            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
107            0
108        );
109        assert_eq!(
110            (pq_scratch.aligned_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
111            0
112        );
113        assert_eq!(
114            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
115            0
116        );
117
118        // Test set() method
119        let query: Vec<u8> = (1..=dim).map(|i| i as u8).collect();
120        pq_scratch.set::<u8>(query.len(), &query).unwrap();
121
122        (0..query.len()).for_each(|i| {
123            assert_eq!(pq_scratch.rotated_query[i], query[i] as f32);
124        });
125    }
126}
diskann_disk/search/pq/pq_scratch.rs

diskann_disk/search/pq/
pq_scratch.rs
