diskann_quantization/minmax/multi/
max_sim.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Distance implementations for MinMax quantized multi-vectors.
7
8use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
9
10use super::super::vectors::{DataRef, MinMaxIP};
11use super::meta::MinMaxMeta;
12use crate::bits::{Representation, Unsigned};
13use crate::distances::{self, UnequalLengths};
14use crate::multi_vector::distance::QueryMatRef;
15use crate::multi_vector::{Chamfer, MatRef, MaxSim};
16
17//////////////////
18// MinMaxKernel //
19//////////////////
20
21/// Kernel for computing [`MaxSim`] and [`Chamfer`] distance using MinMax quantized vectors.
22///
23/// Uses a simple double-iteration strategy and computes pairwise inner-products between
24/// query vectors and document vectors using [`MinMaxIP`].
25pub struct MinMaxKernel;
26
27impl MinMaxKernel {
28    /// Core kernel for computing per-query-vector max similarities using MinMax.
29    ///
30    /// For each query vector, computes the maximum similarity (min distance using
31    /// MinMax inner product) to any document vector, then calls `f(index, score)` with the result.
32    ///
33    /// # Arguments
34    ///
35    /// * `query` - The query MinMax multi-vector (wrapped as [`QueryMatRef`])
36    /// * `doc` - The document MinMax multi-vector
37    /// * `f` - Callback invoked with `(query_index, min_distance)` for each query vector
38    #[inline(always)]
39    pub(crate) fn max_sim_kernel<const NBITS: usize, const MBITS: usize, F>(
40        query: QueryMatRef<'_, MinMaxMeta<NBITS>>,
41        doc: MatRef<'_, MinMaxMeta<MBITS>>,
42        mut f: F,
43    ) -> Result<(), UnequalLengths>
44    where
45        Unsigned: Representation<NBITS> + Representation<MBITS>,
46        distances::InnerProduct: for<'x, 'y> PureDistanceFunction<
47                crate::bits::BitSlice<'x, NBITS, Unsigned>,
48                crate::bits::BitSlice<'y, MBITS, Unsigned>,
49                distances::MathematicalResult<u32>,
50            >,
51        F: FnMut(usize, f32),
52    {
53        for (i, q_ref) in query.rows().enumerate() {
54            // Find min distance (IP returns negated, so min = max similarity)
55            let mut min_distance = f32::MAX;
56
57            for d_ref in doc.rows() {
58                // Use MinMaxIP to compute negated inner product as distance
59                let dist = <MinMaxIP as PureDistanceFunction<
60                    DataRef<'_, NBITS>,
61                    DataRef<'_, MBITS>,
62                    distances::Result<f32>,
63                >>::evaluate(q_ref, d_ref)?;
64
65                min_distance = min_distance.min(dist);
66            }
67
68            f(i, min_distance);
69        }
70
71        Ok(())
72    }
73}
74
75////////////
76// MaxSim //
77////////////
78
79impl<const NBITS: usize, const MBITS: usize>
80    DistanceFunctionMut<QueryMatRef<'_, MinMaxMeta<NBITS>>, MatRef<'_, MinMaxMeta<MBITS>>>
81    for MaxSim<'_>
82where
83    Unsigned: Representation<NBITS> + Representation<MBITS>,
84    distances::InnerProduct: for<'x, 'y> PureDistanceFunction<
85            crate::bits::BitSlice<'x, NBITS, Unsigned>,
86            crate::bits::BitSlice<'y, MBITS, Unsigned>,
87            distances::MathematicalResult<u32>,
88        >,
89{
90    #[inline(always)]
91    fn evaluate(
92        &mut self,
93        query: QueryMatRef<'_, MinMaxMeta<NBITS>>,
94        doc: MatRef<'_, MinMaxMeta<MBITS>>,
95    ) {
96        assert!(
97            self.size() == query.num_vectors(),
98            "scores buffer not right size : {} != {}",
99            self.size(),
100            query.num_vectors()
101        );
102
103        let _ = MinMaxKernel::max_sim_kernel(query, doc, |i, score| {
104            // SAFETY: We asserted that self.size() == query.num_vectors(),
105            // and i < query.num_vectors() due to the kernel loop bound.
106            let _ = self.set(i, score);
107        });
108    }
109}
110
111/////////////
112// Chamfer //
113/////////////
114
115impl<const NBITS: usize, const MBITS: usize>
116    PureDistanceFunction<QueryMatRef<'_, MinMaxMeta<NBITS>>, MatRef<'_, MinMaxMeta<MBITS>>, f32>
117    for Chamfer
118where
119    Unsigned: Representation<NBITS> + Representation<MBITS>,
120    distances::InnerProduct: for<'a, 'b> PureDistanceFunction<
121            crate::bits::BitSlice<'a, NBITS, Unsigned>,
122            crate::bits::BitSlice<'b, MBITS, Unsigned>,
123            distances::MathematicalResult<u32>,
124        >,
125{
126    #[inline(always)]
127    fn evaluate(
128        query: QueryMatRef<'_, MinMaxMeta<NBITS>>,
129        doc: MatRef<'_, MinMaxMeta<MBITS>>,
130    ) -> f32 {
131        let mut sum = 0.0f32;
132
133        let _ = MinMaxKernel::max_sim_kernel(query, doc, |_i, score| {
134            sum += score;
135        });
136
137        sum
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::CompressInto;
145    use crate::algorithms::Transform;
146    use crate::algorithms::transforms::NullTransform;
147    use crate::bits::{Representation, Unsigned};
148    use crate::minmax::{Data, MinMaxQuantizer};
149    use crate::multi_vector::{Defaulted, Mat, Standard};
150    use crate::num::Positive;
151    use diskann_utils::ReborrowMut;
152    use std::num::NonZeroUsize;
153
154    macro_rules! expand_to_bitrates {
155        ($name:ident, $func:ident) => {
156            #[test]
157            fn $name() {
158                // Homogeneous
159                $func::<1, 1>();
160                $func::<2, 2>();
161                $func::<4, 4>();
162                $func::<8, 8>();
163                // Heterogeneous
164                $func::<8, 4>();
165                $func::<8, 2>();
166                $func::<8, 1>();
167            }
168        };
169    }
170
171    /// Test cases: (num_queries, num_docs, dim)
172    const TEST_CASES: &[(usize, usize, usize)] = &[
173        (1, 1, 4),   // Single query, single doc
174        (1, 5, 8),   // Single query, multiple docs
175        (5, 1, 8),   // Multiple queries, single doc
176        (3, 4, 16),  // General case
177        (7, 7, 32),  // Square case
178        (2, 3, 128), // Larger dimension
179    ];
180
181    fn make_quantizer(dim: usize) -> MinMaxQuantizer {
182        MinMaxQuantizer::new(
183            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
184            Positive::new(1.0).unwrap(),
185        )
186    }
187
188    fn generate_input_mat(n: usize, dim: usize, offset: usize) -> Vec<f32> {
189        (0..n * dim)
190            .map(|idx| {
191                let i = idx / dim;
192                let j = idx % dim;
193                ((i + offset) * dim + j) as f32 * 0.1
194            })
195            .collect()
196    }
197
198    fn compress_mat<const NBITS: usize>(
199        quantizer: &MinMaxQuantizer,
200        input: &[f32],
201        n: usize,
202        dim: usize,
203    ) -> Mat<MinMaxMeta<NBITS>>
204    where
205        Unsigned: Representation<NBITS>,
206    {
207        let input_mat = MatRef::new(Standard::<f32>::new(n, dim).unwrap(), input).unwrap();
208        let mut output: Mat<MinMaxMeta<NBITS>> =
209            Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap();
210        quantizer
211            .compress_into(input_mat, output.reborrow_mut())
212            .unwrap();
213        output
214    }
215
216    /// Naive max-sim for one query vector: min distance to any doc vector.
217    fn naive_max_sim_single<const NBITS: usize, const MBITS: usize>(
218        query: DataRef<'_, NBITS>,
219        doc: &MatRef<'_, MinMaxMeta<MBITS>>,
220    ) -> f32
221    where
222        Unsigned: Representation<NBITS> + Representation<MBITS>,
223        distances::InnerProduct: for<'x, 'y> PureDistanceFunction<
224                crate::bits::BitSlice<'x, NBITS, Unsigned>,
225                crate::bits::BitSlice<'y, MBITS, Unsigned>,
226                distances::MathematicalResult<u32>,
227            >,
228    {
229        doc.rows()
230            .map(|d| {
231                <MinMaxIP as PureDistanceFunction<
232                    DataRef<'_, NBITS>,
233                    DataRef<'_, MBITS>,
234                    distances::Result<f32>,
235                >>::evaluate(query, d)
236                .unwrap()
237            })
238            .fold(f32::MAX, f32::min)
239    }
240
241    fn test_matches_naive<const NBITS: usize, const MBITS: usize>()
242    where
243        Unsigned: Representation<NBITS> + Representation<MBITS>,
244        distances::InnerProduct: for<'x, 'y> PureDistanceFunction<
245                crate::bits::BitSlice<'x, NBITS, Unsigned>,
246                crate::bits::BitSlice<'y, MBITS, Unsigned>,
247                distances::MathematicalResult<u32>,
248            >,
249    {
250        for &(nq, nd, dim) in TEST_CASES {
251            let quantizer = make_quantizer(dim);
252
253            let query_data = generate_input_mat(nq, dim, 0);
254            let doc_data = generate_input_mat(nd, dim, nq);
255
256            let query_mat = compress_mat::<NBITS>(&quantizer, &query_data, nq, dim);
257            let doc_mat = compress_mat::<MBITS>(&quantizer, &doc_data, nd, dim);
258
259            let query: QueryMatRef<_> = query_mat.as_view().into();
260            let doc = doc_mat.as_view();
261
262            // Test MaxSim matches naive
263            let expected: Vec<f32> = query
264                .rows()
265                .map(|q| naive_max_sim_single(q, &doc))
266                .collect();
267
268            let mut scores = vec![0.0f32; nq];
269            MaxSim::new(&mut scores).evaluate(query, doc);
270
271            for (i, (&got, &exp)) in scores.iter().zip(expected.iter()).enumerate() {
272                assert!(
273                    (got - exp).abs() < 1e-5,
274                    "({NBITS},{MBITS}) ({nq},{nd},{dim}) MaxSim[{i}]: {got} != {exp}"
275                );
276            }
277
278            // Test kernel matches MaxSim
279            let mut kernel_scores = vec![0.0f32; nq];
280            MinMaxKernel::max_sim_kernel(query, doc, |i, s| kernel_scores[i] = s).unwrap();
281            assert_eq!(
282                scores, kernel_scores,
283                "({NBITS},{MBITS}) ({nq},{nd},{dim}) kernel mismatch"
284            );
285
286            // Test Chamfer equals sum of MaxSim
287            let chamfer = Chamfer::evaluate(query, doc);
288            let sum: f32 = scores.iter().sum();
289            assert!(
290                (chamfer - sum).abs() < 1e-4,
291                "({NBITS},{MBITS}) ({nq},{nd},{dim}) Chamfer {chamfer} != sum {sum}"
292            );
293        }
294    }
295
296    expand_to_bitrates!(matches_naive, test_matches_naive);
297
298    #[test]
299    #[should_panic(expected = "scores buffer not right size")]
300    fn max_sim_panics_on_size_mismatch() {
301        let dim = 4;
302        let row_bytes = Data::<8>::canonical_bytes(dim);
303        let query_data = vec![0u8; 2 * row_bytes];
304        let doc_data = vec![0u8; 3 * row_bytes];
305
306        let query: QueryMatRef<_> = MatRef::new(MinMaxMeta::<8>::new(2, dim), &query_data)
307            .unwrap()
308            .into();
309        let doc = MatRef::new(MinMaxMeta::<8>::new(3, dim), &doc_data).unwrap();
310
311        let mut scores = vec![0.0f32; 5]; // Wrong size
312        MaxSim::new(&mut scores).evaluate(query, doc);
313    }
314}
diskann_quantization/minmax/multi/max_sim.rs

diskann_quantization/minmax/multi/
max_sim.rs
