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::{
10    alloc::{aligned_slice, AlignedSlice},
11    num::PowerOfTwo,
12};
13
14#[derive(Debug)]
15/// PQ scratch
16pub struct PQScratch {
17    /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids.
18    /// 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
19    pub aligned_pqtable_dist_scratch: AlignedSlice<f32>,
20
21    /// Aligned dist scratch, must be at least diskann MAX_DEGREE
22    /// This is used to temporarily save the pq distance between query vector to the candidate vectors.
23    pub aligned_dist_scratch: AlignedSlice<f32>,
24
25    /// Aligned pq coord scratch, must be at least [N_CHUNKS * MAX_DEGREE]
26    /// This is used to store the pq coordinates of the candidate vectors.
27    pub aligned_pq_coord_scratch: AlignedSlice<u8>,
28
29    /// Query scratch buffer stored as `f32`. `set` initializes it by copying/converting the
30    /// raw query values; `PQTable.PreprocessQuery` can then rotate or otherwise preprocess it.
31    pub rotated_query: Vec<f32>,
32}
33
34impl PQScratch {
35    /// 128 bytes alignment to optimize for the L2 Adjacent Cache Line Prefetcher.
36    const ALIGNED_ALLOC_128: PowerOfTwo = match PowerOfTwo::new(128) {
37        Ok(v) => v,
38        Err(_) => unreachable!(),
39    };
40
41    /// Create a new pq scratch
42    pub fn new(
43        graph_degree: usize,
44        dim: usize,
45        num_pq_chunks: usize,
46        num_centers: usize,
47    ) -> ANNResult<Self> {
48        let aligned_pq_coord_scratch =
49            aligned_slice(graph_degree * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)
50                .map_err(ANNError::log_index_error)?;
51        let aligned_pqtable_dist_scratch =
52            aligned_slice(num_centers * num_pq_chunks, PQScratch::ALIGNED_ALLOC_128)
53                .map_err(ANNError::log_index_error)?;
54        let aligned_dist_scratch = aligned_slice(graph_degree, PQScratch::ALIGNED_ALLOC_128)
55            .map_err(ANNError::log_index_error)?;
56        let rotated_query = vec![0.0f32; dim];
57
58        Ok(Self {
59            aligned_pqtable_dist_scratch,
60            aligned_dist_scratch,
61            aligned_pq_coord_scratch,
62            rotated_query,
63        })
64    }
65
66    /// Copy `query` into `rotated_query`, converting to `f32`.
67    ///
68    /// `dim` is the element count in the `T` representation. The decompressed
69    /// `f32` length returned by `T::as_f32` may differ (e.g. `MinMaxElement`
70    /// expands to more `f32`s than its raw element count), so the destination
71    /// slice is sized by that actual length.
72    ///
73    /// Returns `DimensionMismatchError` if `dim > query.len()` or the
74    /// decompressed vector does not fit in `rotated_query`.
75    pub fn set<T: VectorRepr>(&mut self, dim: usize, query: &[T]) -> ANNResult<()> {
76        if dim > query.len() {
77            return Err(ANNError::log_dimension_mismatch_error(format!(
78                "PQScratch::set: expected query of length >= {dim}, got {}",
79                query.len()
80            )));
81        }
82        let query = T::as_f32(&query[..dim]).into_ann_result()?;
83        if query.len() > self.rotated_query.len() {
84            return Err(ANNError::log_dimension_mismatch_error(format!(
85                "PQScratch::set: decompressed query of length {} does not fit rotated_query buffer of length {}",
86                query.len(),
87                self.rotated_query.len()
88            )));
89        }
90        self.rotated_query[..query.len()].copy_from_slice(&query);
91        Ok(())
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use rstest::rstest;
98
99    use super::PQScratch;
100
101    #[rstest]
102    #[case(512, 8, 128, 256)] // default test case
103    #[case(59, 16, 37, 41)] // not multiple of 256
104    fn test_pq_scratch(
105        #[case] graph_degree: usize,
106        #[case] dim: usize,
107        #[case] num_pq_chunks: usize,
108        #[case] num_centers: usize,
109    ) {
110        let mut pq_scratch: PQScratch =
111            PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap();
112
113        // Check alignment of the remaining AlignedSlice buffers.
114        assert_eq!(
115            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize)
116                % PQScratch::ALIGNED_ALLOC_128.raw(),
117            0
118        );
119        assert_eq!(
120            (pq_scratch.aligned_dist_scratch.as_ptr() as usize)
121                % PQScratch::ALIGNED_ALLOC_128.raw(),
122            0
123        );
124        assert_eq!(
125            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize)
126                % PQScratch::ALIGNED_ALLOC_128.raw(),
127            0
128        );
129
130        // Test set() method
131        let query: Vec<u8> = (1..=dim).map(|i| i as u8).collect();
132        pq_scratch.set::<u8>(query.len(), &query).unwrap();
133
134        (0..query.len()).for_each(|i| {
135            assert_eq!(pq_scratch.rotated_query[i], query[i] as f32);
136        });
137    }
138}
diskann_disk/search/pq/pq_scratch.rs

diskann_disk/search/pq/
pq_scratch.rs
