diskann_quantization/multi_vector/distance/
factory.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Factory + concrete `MaxSimKernel<T>` impls for the multi-vector distance
7//! API. BYOTE entry point — see [`build_max_sim`].
8
9use std::num::NonZeroUsize;
10
11use diskann_utils::views::rowmajor::Matrix;
12use diskann_vector::PureDistanceFunction;
13use diskann_vector::distance::InnerProduct;
14use diskann_wide::Architecture;
15use diskann_wide::arch::Scalar;
16#[cfg(target_arch = "aarch64")]
17use diskann_wide::arch::aarch64::Neon;
18#[cfg(target_arch = "x86_64")]
19use diskann_wide::arch::x86_64::{V3, V4};
20use thiserror::Error;
21
22use super::fallback::FallbackKernel;
23use super::isa::{MaxSimIsa, NotSupported};
24use super::kernel::{Erase, MaxSimKernel};
25use super::max_sim::MaxSimError;
26use crate::matrix_kernels as mk;
27use crate::multi_vector::distance::QueryMatRef;
28use crate::multi_vector::{BlockTransposed, Mat, MatRef, Overflow, Standard};
29
30// ─────────────────────────────────────────────────────────────────────────
31//  Prepared<A, Q, NR> — concrete kernel for the arch-dispatched paths.
32// ─────────────────────────────────────────────────────────────────────────
33
34// ZST for selecting the packing for the documents.
35//
36// TODO: Some of these decisions should eventually ve moved lower into `matrix_kernels`.
37#[derive(Debug, Clone, Copy)]
38struct Pack<const NR: usize>;
39
40#[derive(Debug)]
41struct Prepared<A, Q, const NR: usize> {
42    arch: A,
43    prepared: Q,
44    _packing: Pack<NR>,
45}
46
47impl<A, const GROUP: usize, const NR: usize> MaxSimKernel<f32>
48    for Prepared<A, BlockTransposed<f32, GROUP>, NR>
49where
50    A: Architecture,
51    for<'a> mk::maxsim::packed_f32_x_unpacked_f32::Driver<'a, A, GROUP, NR>: mk::Drive,
52{
53    fn nrows(&self) -> usize {
54        self.prepared.nrows()
55    }
56
57    fn compute_max_sim(
58        &self,
59        doc: MatRef<'_, Standard<f32>>,
60        scores: &mut [f32],
61    ) -> Result<(), MaxSimError> {
62        if scores.len() != self.nrows() {
63            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
64        }
65
66        if doc.vector_dim() != self.prepared.ncols() {
67            return Err(MaxSimError::UnequalDim(
68                doc.vector_dim(),
69                self.prepared.ncols(),
70            ));
71        }
72
73        let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
74            scores.fill(if doc.num_vectors() == 0 {
75                f32::MAX
76            } else {
77                0.0
78            });
79            return Ok(());
80        };
81
82        let Some(a) = mk::blocks::packed::View::from_block_transposed(self.prepared.as_view())
83        else {
84            return Ok(());
85        };
86
87        let Some(b) = mk::blocks::unpacked::View::from_matrix_view(doc.as_matrix_view()) else {
88            scores.fill(f32::MAX);
89            return Ok(());
90        };
91
92        // SAFETY: The dimension check establishes that `a.k() == b.k() == k`.
93        // The length check establishes that `scores` occupies exactly the
94        // packed blocks in `a`.
95        let mut driver = unsafe {
96            mk::maxsim::packed_f32_x_unpacked_f32::Driver::new(
97                self.arch,
98                a,
99                b,
100                scores,
101                k,
102                mk::Cache::detect(),
103            )
104        };
105
106        mk::Drive::drive(&mut driver);
107
108        scores.iter_mut().for_each(|s| *s = -*s);
109
110        Ok(())
111    }
112}
113
114impl<A, const GROUP: usize, const NR: usize> MaxSimKernel<half::f16>
115    for Prepared<A, BlockTransposed<f32, GROUP>, NR>
116where
117    A: Architecture,
118    for<'a> mk::maxsim::packed_f32_x_unpacked_f16::Driver<'a, A, GROUP, NR>: mk::Drive,
119{
120    fn nrows(&self) -> usize {
121        self.prepared.nrows()
122    }
123
124    fn compute_max_sim(
125        &self,
126        doc: MatRef<'_, Standard<half::f16>>,
127        scores: &mut [f32],
128    ) -> Result<(), MaxSimError> {
129        if scores.len() != self.nrows() {
130            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
131        }
132
133        if doc.vector_dim() != self.prepared.ncols() {
134            return Err(MaxSimError::UnequalDim(
135                doc.vector_dim(),
136                self.prepared.ncols(),
137            ));
138        }
139
140        let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
141            scores.fill(if doc.num_vectors() == 0 {
142                f32::MAX
143            } else {
144                0.0
145            });
146            return Ok(());
147        };
148
149        let Some(a) = mk::blocks::packed::View::from_block_transposed(self.prepared.as_view())
150        else {
151            return Ok(());
152        };
153
154        let Some(b) = mk::blocks::unpacked::View::from_matrix_view(doc.as_matrix_view()) else {
155            scores.fill(f32::MAX);
156            return Ok(());
157        };
158
159        // SAFETY: The dimension check establishes that `a.k() == b.k() == k`.
160        // The length check establishes that `scores` occupies exactly the
161        // packed blocks in `a`.
162        let mut driver = unsafe {
163            mk::maxsim::packed_f32_x_unpacked_f16::Driver::new(
164                self.arch,
165                a,
166                b,
167                scores,
168                k,
169                mk::Cache::detect(),
170            )
171        };
172
173        mk::Drive::drive(&mut driver);
174
175        scores.iter_mut().for_each(|s| *s = -*s);
176
177        Ok(())
178    }
179}
180
181impl<A, Q, const GROUP: usize, const NR: usize, const PACK: usize> MaxSimKernel<i8>
182    for Prepared<A, BlockTransposed<Q, GROUP, PACK>, NR>
183where
184    Q: Copy + Send + Sync + std::fmt::Debug,
185    A: mk::maxsim::packed_i8_x_unpacked_i8::PrepareB<AElem = Q> + Architecture,
186    for<'a> mk::maxsim::packed_i8_x_unpacked_i8::Driver<'a, A, GROUP, NR, PACK>: mk::Drive,
187{
188    fn nrows(&self) -> usize {
189        self.prepared.nrows()
190    }
191
192    fn compute_max_sim(
193        &self,
194        doc: MatRef<'_, Standard<i8>>,
195        scores: &mut [i32],
196    ) -> Result<(), MaxSimError> {
197        if scores.len() != self.nrows() {
198            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
199        }
200
201        if doc.vector_dim() != self.prepared.ncols() {
202            return Err(MaxSimError::UnequalDim(
203                doc.vector_dim(),
204                self.prepared.ncols(),
205            ));
206        }
207
208        let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
209            scores.fill(if doc.num_vectors() == 0 { i32::MAX } else { 0 });
210            return Ok(());
211        };
212
213        let Some(a) = mk::blocks::packed::View::from_block_transposed(self.prepared.as_view())
214        else {
215            return Ok(());
216        };
217
218        let Some(b) = mk::blocks::unpacked::View::from_matrix_view(doc.as_matrix_view()) else {
219            scores.fill(i32::MAX);
220            return Ok(());
221        };
222
223        // SAFETY: The dimension check establishes that `a.k() == b.k() == k`.
224        // The length check establishes that `scores` occupies exactly the
225        // packed blocks in `a`.
226        let mut driver = unsafe {
227            mk::maxsim::packed_i8_x_unpacked_i8::Driver::new(
228                self.arch,
229                a,
230                b,
231                scores,
232                k,
233                mk::Cache::detect(),
234            )
235        };
236
237        mk::Drive::drive(&mut driver);
238
239        scores.iter_mut().for_each(|s| *s = -*s);
240
241        Ok(())
242    }
243}
244
245// ─────────────────────────────────────────────────────────────────────────
246//  ReferenceKernel<T> — double loop over the single-vector inner product.
247// ─────────────────────────────────────────────────────────────────────────
248
249/// Reference MaxSim implementation, selected by [`MaxSimElement::build`].
250///
251/// May assume `scores` and `doc` were validated by
252/// [`MaxSimKernel::compute_max_sim`].
253type ReferenceFn<T> =
254    fn(QueryMatRef<'_, Standard<T>>, MatRef<'_, Standard<T>>, &mut [<T as MaxSimElement>::Score]);
255
256/// Unoptimized kernel backing [`MaxSimIsa::Reference`].
257struct ReferenceKernel<T: MaxSimElement> {
258    query: Mat<Standard<T>>,
259    run: ReferenceFn<T>,
260}
261
262impl<T: MaxSimElement> std::fmt::Debug for ReferenceKernel<T> {
263    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
264        f.debug_struct("ReferenceKernel")
265            .field("nrows", &self.query.num_vectors())
266            .finish()
267    }
268}
269
270impl<T: MaxSimElement> ReferenceKernel<T> {
271    fn new(query: MatRef<'_, Standard<T>>, run: ReferenceFn<T>) -> Self {
272        Self {
273            query: query.to_owned(),
274            run,
275        }
276    }
277}
278
279impl<T: MaxSimElement> MaxSimKernel<T> for ReferenceKernel<T> {
280    fn nrows(&self) -> usize {
281        self.query.num_vectors()
282    }
283
284    fn compute_max_sim(
285        &self,
286        doc: MatRef<'_, Standard<T>>,
287        scores: &mut [T::Score],
288    ) -> Result<(), MaxSimError> {
289        if scores.len() != self.nrows() {
290            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
291        }
292        if doc.vector_dim() != self.query.vector_dim() {
293            return Err(MaxSimError::UnequalDim(
294                doc.vector_dim(),
295                self.query.vector_dim(),
296            ));
297        }
298        (self.run)(self.query.as_view().into(), doc, scores);
299        Ok(())
300    }
301}
302
303/// [`ReferenceKernel`] implementation for element types scored in `f32`.
304fn reference_scores<T: Copy>(
305    query: QueryMatRef<'_, Standard<T>>,
306    doc: MatRef<'_, Standard<T>>,
307    scores: &mut [f32],
308) where
309    InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
310{
311    FallbackKernel::max_sim_kernel(query, doc, |i, score| scores[i] = score);
312}
313
314/// [`ReferenceKernel`] implementation for `i8`, which scores in exact `i32`.
315fn reference_scores_i8(
316    query: QueryMatRef<'_, Standard<i8>>,
317    doc: MatRef<'_, Standard<i8>>,
318    scores: &mut [i32],
319) {
320    FallbackKernel::max_sim_kernel_i8(query, doc, |i, score| scores[i] = score);
321}
322
323// ─────────────────────────────────────────────────────────────────────────
324//  BuildAndErase<E> — Target1 impls used by `dispatch1_no_features` (Auto).
325// ─────────────────────────────────────────────────────────────────────────
326
327struct BuildAndErase<E>(E);
328
329// ───── f32 Target1 impls ─────
330
331impl<E: Erase<f32>> diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<f32>>>
332    for BuildAndErase<E>
333{
334    fn run(self, arch: Scalar, query: MatRef<'_, Standard<f32>>) -> E::Output {
335        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
336        self.0.erase(Prepared {
337            arch,
338            prepared,
339            _packing: Pack::<2>,
340        })
341    }
342}
343
344#[cfg(target_arch = "x86_64")]
345impl<E: Erase<f32>> diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<f32>>>
346    for BuildAndErase<E>
347{
348    fn run(self, arch: V3, query: MatRef<'_, Standard<f32>>) -> E::Output {
349        let prepared = BlockTransposed::<f32, 16>::from_matrix_view(query.as_matrix_view());
350        self.0.erase(Prepared {
351            arch,
352            prepared,
353            _packing: Pack::<6>,
354        })
355    }
356}
357
358#[cfg(target_arch = "x86_64")]
359impl<E: Erase<f32>> diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<f32>>>
360    for BuildAndErase<E>
361{
362    fn run(self, arch: V4, query: MatRef<'_, Standard<f32>>) -> E::Output {
363        let prepared = BlockTransposed::<f32, 32>::from_matrix_view(query.as_matrix_view());
364        self.0.erase(Prepared {
365            arch,
366            prepared,
367            _packing: Pack::<6>,
368        })
369    }
370}
371
372#[cfg(target_arch = "aarch64")]
373impl<E: Erase<f32>> diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<f32>>>
374    for BuildAndErase<E>
375{
376    fn run(self, arch: Neon, query: MatRef<'_, Standard<f32>>) -> E::Output {
377        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
378        self.0.erase(Prepared {
379            arch,
380            prepared,
381            _packing: Pack::<6>,
382        })
383    }
384}
385
386// ───── f16 Target1 impls ─────
387
388impl<E: Erase<half::f16>>
389    diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<half::f16>>>
390    for BuildAndErase<E>
391{
392    fn run(self, arch: Scalar, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
393        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(
394            query
395                .as_matrix_view()
396                .map(|v| diskann_wide::cast_f16_to_f32(*v))
397                .as_view(),
398        );
399        self.0.erase(Prepared {
400            arch,
401            prepared,
402            _packing: Pack::<2>,
403        })
404    }
405}
406
407#[cfg(target_arch = "x86_64")]
408impl<E: Erase<half::f16>>
409    diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<half::f16>>>
410    for BuildAndErase<E>
411{
412    fn run(self, arch: V3, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
413        let prepared = BlockTransposed::<f32, 16>::from_matrix_view(
414            query
415                .as_matrix_view()
416                .map(|v| diskann_wide::cast_f16_to_f32(*v))
417                .as_view(),
418        );
419        self.0.erase(Prepared {
420            arch,
421            prepared,
422            _packing: Pack::<6>,
423        })
424    }
425}
426
427#[cfg(target_arch = "x86_64")]
428impl<E: Erase<half::f16>>
429    diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<half::f16>>>
430    for BuildAndErase<E>
431{
432    fn run(self, arch: V4, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
433        let prepared = BlockTransposed::<f32, 32>::from_matrix_view(
434            query
435                .as_matrix_view()
436                .map(|v| diskann_wide::cast_f16_to_f32(*v))
437                .as_view(),
438        );
439        self.0.erase(Prepared {
440            arch,
441            prepared,
442            _packing: Pack::<6>,
443        })
444    }
445}
446
447#[cfg(target_arch = "aarch64")]
448impl<E: Erase<half::f16>>
449    diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<half::f16>>>
450    for BuildAndErase<E>
451{
452    fn run(self, arch: Neon, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
453        // Neon dispatches to Scalar (no Neon-specific kernel).
454        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(
455            query
456                .as_matrix_view()
457                .map(|v| diskann_wide::cast_f16_to_f32(*v))
458                .as_view(),
459        );
460        self.0.erase(Prepared {
461            arch,
462            prepared,
463            _packing: Pack::<6>,
464        })
465    }
466}
467
468// ───── i8 Target1 impls ─────
469
470impl<E: Erase<i8>> diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<i8>>>
471    for BuildAndErase<E>
472{
473    fn run(self, arch: Scalar, query: MatRef<'_, Standard<i8>>) -> E::Output {
474        let prepared = BlockTransposed::<i8, 8, 2>::from_matrix_view(query.as_matrix_view());
475        self.0.erase(Prepared {
476            arch,
477            prepared,
478            _packing: Pack::<2>,
479        })
480    }
481}
482
483#[cfg(target_arch = "x86_64")]
484impl<E: Erase<i8>> diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<i8>>>
485    for BuildAndErase<E>
486{
487    fn run(self, arch: V3, query: MatRef<'_, Standard<i8>>) -> E::Output {
488        // `PrepareB for V3` expects the query widened to `i16`.
489        let prepared = BlockTransposed::<i16, 16, 2>::from_matrix_view(
490            query.as_matrix_view().map(|v| i16::from(*v)).as_view(),
491        );
492        self.0.erase(Prepared {
493            arch,
494            prepared,
495            _packing: Pack::<6>,
496        })
497    }
498}
499
500#[cfg(target_arch = "x86_64")]
501impl<E: Erase<i8>> diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<i8>>>
502    for BuildAndErase<E>
503{
504    fn run(self, arch: V4, query: MatRef<'_, Standard<i8>>) -> E::Output {
505        // `PrepareB for V4` expects the query packed as `x + 128` in `u8`.
506        let prepared = BlockTransposed::<u8, 32, 4>::from_matrix_view(
507            query.as_matrix_view().map(|v| (*v as u8) ^ 0x80).as_view(),
508        );
509        self.0.erase(Prepared {
510            arch,
511            prepared,
512            _packing: Pack::<6>,
513        })
514    }
515}
516
517#[cfg(target_arch = "aarch64")]
518impl<E: Erase<i8>> diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<i8>>>
519    for BuildAndErase<E>
520{
521    fn run(self, arch: Neon, query: MatRef<'_, Standard<i8>>) -> E::Output {
522        let prepared = BlockTransposed::<i8, 8, 4>::from_matrix_view(query.as_matrix_view());
523        self.0.erase(Prepared {
524            arch,
525            prepared,
526            _packing: Pack::<6>,
527        })
528    }
529}
530
531// ─────────────────────────────────────────────────────────────────────────
532//  MaxSimElement — sealed trait gating accepted element types.
533// ─────────────────────────────────────────────────────────────────────────
534
535mod sealed {
536    pub trait Sealed {}
537}
538
539/// Scalar element types accepted by [`build_max_sim`].
540///
541/// Sealed: external crates cannot add impls. Quantized representations
542/// (PQ, SQ, packed sub-byte) are intentionally excluded — they need
543/// codebook/scale state that [`MatRef<'_, Standard<Self>>`] can't carry.
544pub trait MaxSimElement: sealed::Sealed + Sized + Copy + Send + Sync + 'static {
545    /// Score produced per query row: `f32` for floating-point elements, `i32`
546    /// for `i8` where the inner product is exact.
547    type Score: Copy + Default + PartialEq + std::fmt::Debug;
548
549    /// Score written for every query row when the document set is empty.
550    const NO_MATCH: Self::Score;
551
552    /// Build the concrete kernel for this element type and hand it to
553    /// `erase.erase(...)`.
554    ///
555    /// # Errors
556    ///
557    /// Returns [`BuildMaxSimError::NotSupported`] when the requested ISA cannot
558    /// run on this build (e.g. AVX-512 unavailable; aarch64 on x86_64), and
559    /// [`BuildMaxSimError::DimTooLarge`] when an `i8` query has more than
560    /// `131_071` dimensions, beyond which `i32` scores could overflow.
561    fn build<E: Erase<Self>>(
562        isa: MaxSimIsa,
563        query: MatRef<'_, Standard<Self>>,
564        erase: E,
565    ) -> Result<E::Output, BuildMaxSimError>;
566}
567
568impl sealed::Sealed for f32 {}
569impl sealed::Sealed for half::f16 {}
570impl sealed::Sealed for i8 {}
571
572impl MaxSimElement for f32 {
573    type Score = f32;
574    const NO_MATCH: f32 = f32::MAX;
575
576    fn build<E: Erase<f32>>(
577        isa: MaxSimIsa,
578        query: MatRef<'_, Standard<f32>>,
579        erase: E,
580    ) -> Result<E::Output, BuildMaxSimError> {
581        match isa {
582            MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
583                BuildAndErase(erase),
584                query,
585            )),
586            MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
587            #[cfg(target_arch = "x86_64")]
588            MaxSimIsa::X86_64_V3 => {
589                let arch = V3::new_checked().ok_or(NotSupported {
590                    isa,
591                    reason: "AVX2/FMA unavailable on this CPU",
592                })?;
593                Ok(arch.run1(BuildAndErase(erase), query))
594            }
595            #[cfg(target_arch = "x86_64")]
596            MaxSimIsa::X86_64_V4 => {
597                let arch = V4::new_checked().ok_or(NotSupported {
598                    isa,
599                    reason: "AVX-512 unavailable on this CPU",
600                })?;
601                Ok(arch.run1(BuildAndErase(erase), query))
602            }
603            #[cfg(not(target_arch = "x86_64"))]
604            MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
605                isa,
606                reason: "x86_64 target only",
607            }
608            .into()),
609            #[cfg(target_arch = "aarch64")]
610            MaxSimIsa::Neon => {
611                let arch = Neon::new_checked().ok_or(NotSupported {
612                    isa,
613                    reason: "Neon unavailable on this CPU",
614                })?;
615                Ok(arch.run1(BuildAndErase(erase), query))
616            }
617            #[cfg(not(target_arch = "aarch64"))]
618            MaxSimIsa::Neon => Err(NotSupported {
619                isa,
620                reason: "aarch64 target only",
621            }
622            .into()),
623            MaxSimIsa::Reference => {
624                Ok(erase.erase(ReferenceKernel::new(query, reference_scores::<f32>)))
625            }
626        }
627    }
628}
629
630impl MaxSimElement for half::f16 {
631    type Score = f32;
632    const NO_MATCH: f32 = f32::MAX;
633
634    fn build<E: Erase<half::f16>>(
635        isa: MaxSimIsa,
636        query: MatRef<'_, Standard<half::f16>>,
637        erase: E,
638    ) -> Result<E::Output, BuildMaxSimError> {
639        match isa {
640            MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
641                BuildAndErase(erase),
642                query,
643            )),
644            MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
645            #[cfg(target_arch = "x86_64")]
646            MaxSimIsa::X86_64_V3 => {
647                let arch = V3::new_checked().ok_or(NotSupported {
648                    isa,
649                    reason: "AVX2/FMA unavailable on this CPU",
650                })?;
651                Ok(arch.run1(BuildAndErase(erase), query))
652            }
653            #[cfg(target_arch = "x86_64")]
654            MaxSimIsa::X86_64_V4 => {
655                let arch = V4::new_checked().ok_or(NotSupported {
656                    isa,
657                    reason: "AVX-512 unavailable on this CPU",
658                })?;
659                Ok(arch.run1(BuildAndErase(erase), query))
660            }
661            #[cfg(not(target_arch = "x86_64"))]
662            MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
663                isa,
664                reason: "x86_64 target only",
665            }
666            .into()),
667            #[cfg(target_arch = "aarch64")]
668            MaxSimIsa::Neon => {
669                let arch = Neon::new_checked().ok_or(NotSupported {
670                    isa,
671                    reason: "Neon unavailable on this CPU",
672                })?;
673                Ok(arch.run1(BuildAndErase(erase), query))
674            }
675            #[cfg(not(target_arch = "aarch64"))]
676            MaxSimIsa::Neon => Err(NotSupported {
677                isa,
678                reason: "aarch64 target only",
679            }
680            .into()),
681            MaxSimIsa::Reference => {
682                Ok(erase.erase(ReferenceKernel::new(query, reference_scores::<half::f16>)))
683            }
684        }
685    }
686}
687
688/// Largest `i8` dimension for which every inner product fits in `i32`.
689const MAX_I8_DIM: usize = (i32::MAX / (128 * 128)) as usize;
690
691impl MaxSimElement for i8 {
692    type Score = i32;
693    const NO_MATCH: i32 = i32::MAX;
694
695    fn build<E: Erase<i8>>(
696        isa: MaxSimIsa,
697        query: MatRef<'_, Standard<i8>>,
698        erase: E,
699    ) -> Result<E::Output, BuildMaxSimError> {
700        if query.vector_dim() > MAX_I8_DIM {
701            return Err(BuildMaxSimError::DimTooLarge(
702                query.vector_dim(),
703                MAX_I8_DIM,
704            ));
705        }
706
707        match isa {
708            MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
709                BuildAndErase(erase),
710                query,
711            )),
712            MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
713            #[cfg(target_arch = "x86_64")]
714            MaxSimIsa::X86_64_V3 => {
715                let arch = V3::new_checked().ok_or(NotSupported {
716                    isa,
717                    reason: "AVX2/FMA unavailable on this CPU",
718                })?;
719                Ok(arch.run1(BuildAndErase(erase), query))
720            }
721            #[cfg(target_arch = "x86_64")]
722            MaxSimIsa::X86_64_V4 => {
723                let arch = V4::new_checked().ok_or(NotSupported {
724                    isa,
725                    reason: "AVX-512 unavailable on this CPU",
726                })?;
727                Ok(arch.run1(BuildAndErase(erase), query))
728            }
729            #[cfg(not(target_arch = "x86_64"))]
730            MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
731                isa,
732                reason: "x86_64 target only",
733            }
734            .into()),
735            #[cfg(target_arch = "aarch64")]
736            MaxSimIsa::Neon => {
737                let arch = Neon::new_checked().ok_or(NotSupported {
738                    isa,
739                    reason: "Neon unavailable on this CPU",
740                })?;
741                Ok(arch.run1(BuildAndErase(erase), query))
742            }
743            #[cfg(not(target_arch = "aarch64"))]
744            MaxSimIsa::Neon => Err(NotSupported {
745                isa,
746                reason: "aarch64 target only",
747            }
748            .into()),
749            MaxSimIsa::Reference => {
750                Ok(erase.erase(ReferenceKernel::new(query, reference_scores_i8)))
751            }
752        }
753    }
754}
755
756// ─────────────────────────────────────────────────────────────────────────
757//  Factory entry point.
758// ─────────────────────────────────────────────────────────────────────────
759
760/// Failure to prepare a MaxSim query.
761#[derive(Debug, Clone, Copy, Error)]
762#[non_exhaustive]
763pub enum BuildMaxSimError {
764    #[error(transparent)]
765    NotSupported(#[from] NotSupported),
766    #[error("query-vector dim {0} exceeds the maximum of {1}")]
767    DimTooLarge(usize, usize),
768    #[error(transparent)]
769    Overflow(#[from] Overflow),
770}
771
772/// Build a multi-vector MaxSim kernel for any [`MaxSimElement`] type.
773///
774/// Thin wrapper over [`MaxSimElement::build`] so callers don't have to name
775/// the trait at the call site.
776///
777/// # Errors
778///
779/// Returns [`BuildMaxSimError::NotSupported`] when the requested ISA cannot run
780/// on this build, and [`BuildMaxSimError::DimTooLarge`] when an `i8` query has
781/// more than `131_071` dimensions.
782pub fn build_max_sim<T: MaxSimElement, E: Erase<T>>(
783    isa: MaxSimIsa,
784    query: MatRef<'_, Standard<T>>,
785    erase: E,
786) -> Result<E::Output, BuildMaxSimError> {
787    T::build(isa, query, erase)
788}
789
790#[cfg(test)]
791mod tests {
792    use super::*;
793    use crate::multi_vector::{BoxErase, Chamfer, MaxSim, QueryMatRef};
794    use diskann_vector::DistanceFunctionMut;
795
796    /// Local helper trait — picks a sane test value of `T` from an `f32`
797    /// so every element type shares the same data generator.
798    trait FromF32 {
799        fn from_f32(v: f32) -> Self;
800    }
801
802    impl FromF32 for f32 {
803        fn from_f32(v: f32) -> Self {
804            v
805        }
806    }
807
808    impl FromF32 for half::f16 {
809        fn from_f32(v: f32) -> Self {
810            diskann_wide::cast_f32_to_f16(v)
811        }
812    }
813
814    impl FromF32 for i8 {
815        fn from_f32(v: f32) -> Self {
816            v as i8
817        }
818    }
819
820    /// Projects a kernel score onto the `f32` distance the fallback path
821    /// produces, so every element type shares the same assertions.
822    trait ScoreAsF32: MaxSimElement {
823        fn score_as_f32(score: Self::Score) -> f32;
824    }
825
826    impl ScoreAsF32 for f32 {
827        fn score_as_f32(score: f32) -> f32 {
828            score
829        }
830    }
831
832    impl ScoreAsF32 for half::f16 {
833        fn score_as_f32(score: f32) -> f32 {
834            score
835        }
836    }
837
838    impl ScoreAsF32 for i8 {
839        fn score_as_f32(score: i32) -> f32 {
840            if score == Self::NO_MATCH {
841                f32::MAX
842            } else {
843                score as f32
844            }
845        }
846    }
847
848    fn scores_buffer<T: MaxSimElement>(len: usize) -> Vec<T::Score> {
849        vec![T::Score::default(); len]
850    }
851
852    fn make_mat<T: Copy>(data: &[T], nrows: usize, ncols: usize) -> MatRef<'_, Standard<T>> {
853        MatRef::new(Standard::new(nrows, ncols).unwrap(), data).unwrap()
854    }
855
856    fn make_test_data<T: FromF32>(len: usize, ceil: usize, shift: usize) -> Vec<T> {
857        (0..len)
858            .map(|v| T::from_f32(((v + shift) % ceil) as f32))
859            .collect()
860    }
861
862    /// Shapes for the `chamfer_matches_fallback` / `max_sim_matches_fallback`
863    /// agreement checks: `(num_queries, num_docs, dim)`.
864    ///
865    /// Targets the factory wiring (query setup, score writeback) above the
866    /// kernel layer; exhaustive panel/remainder coverage is pinned in
867    /// `kernels::tiled_reduce::tests`.
868    const TEST_CASES: &[(usize, usize, usize)] = &[
869        (0, 1, 1),   // Empty query
870        (1, 0, 1),   // Empty docs
871        (1, 1, 0),   // Empty dim
872        (1, 0, 0),   // Empty docs and dim
873        (1, 1, 4),   // Degenerate
874        (5, 3, 5),   // Prime k; nq > 1 and nd > 1 exercise per-row writeback
875        (17, 4, 64), // A-panel remainder crossing both Scalar and V3 panel widths
876        (16, 6, 32), // B-remainder ≠ 1 (V3 b_remainder = 2)
877    ];
878
879    fn check_chamfer_matches<T>(tol: f32, label: &str)
880    where
881        T: ScoreAsF32 + FromF32,
882        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
883    {
884        for &(nq, nd, dim) in TEST_CASES {
885            let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
886            let doc_data = make_test_data::<T>(nd * dim, dim, dim);
887
888            let query = make_mat(&query_data, nq, dim);
889            let doc = make_mat(&doc_data, nd, dim);
890
891            let expected = Chamfer::evaluate(QueryMatRef::from(query), doc);
892
893            let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
894            let mut scores = scores_buffer::<T>(nq);
895            kernel.compute_max_sim(doc, &mut scores).unwrap();
896            let actual: f32 = scores.iter().map(|&s| T::score_as_f32(s)).sum();
897
898            assert!(
899                (actual - expected).abs() < tol,
900                "{label}Chamfer mismatch for ({nq},{nd},{dim}): actual={actual}, expected={expected}",
901            );
902        }
903    }
904
905    fn check_max_sim_matches<T>(tol: f32, label: &str)
906    where
907        T: ScoreAsF32 + FromF32,
908        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
909    {
910        for &(nq, nd, dim) in TEST_CASES {
911            let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
912            let doc_data = make_test_data::<T>(nd * dim, dim, dim);
913
914            let query = make_mat(&query_data, nq, dim);
915            let doc = make_mat(&doc_data, nd, dim);
916
917            let mut expected_scores = vec![0.0f32; nq];
918            let _ = MaxSim::new(&mut expected_scores).evaluate(QueryMatRef::from(query), doc);
919
920            let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
921            let mut actual_scores = scores_buffer::<T>(nq);
922            kernel.compute_max_sim(doc, &mut actual_scores).unwrap();
923
924            for i in 0..nq {
925                let actual = T::score_as_f32(actual_scores[i]);
926                assert!(
927                    (actual - expected_scores[i]).abs() < tol,
928                    "{label}MaxSim[{i}] mismatch for ({nq},{nd},{dim}): actual={actual}, expected={}",
929                    expected_scores[i],
930                );
931            }
932        }
933    }
934
935    fn check_i8_isas(nq: usize, nd: usize, dim: usize, query_data: &[i8], doc_data: &[i8]) {
936        let query = make_mat(query_data, nq, dim);
937        let doc = make_mat(doc_data, nd, dim);
938
939        let mut expected = vec![0.0f32; nq];
940        let _ = MaxSim::new(&mut expected).evaluate(QueryMatRef::from(query), doc);
941
942        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
943            let kernel = build_max_sim::<i8, _>(isa, query, BoxErase).unwrap();
944            let mut scores = scores_buffer::<i8>(nq);
945            kernel.compute_max_sim(doc, &mut scores).unwrap();
946
947            for i in 0..nq {
948                let actual = <i8 as ScoreAsF32>::score_as_f32(scores[i]);
949                assert!(
950                    (actual - expected[i]).abs() < 1e-10,
951                    "i8 {isa} MaxSim[{i}] mismatch for ({nq},{nd},{dim}): \
952                     actual={actual}, expected={}",
953                    expected[i],
954                );
955            }
956        }
957    }
958
959    #[test]
960    fn i8_isas_match_oracle() {
961        for &(nq, nd, dim) in TEST_CASES {
962            let query_data = make_test_data::<i8>(nq * dim, dim, dim / 2);
963            let doc_data = make_test_data::<i8>(nd * dim, dim, dim);
964            check_i8_isas(nq, nd, dim, &query_data, &doc_data);
965        }
966
967        let full_range: &[(usize, usize, usize)] = if cfg!(miri) {
968            &[(5, 4, 64)]
969        } else {
970            &[(33, 13, 64), (70, 1000, 131)]
971        };
972        for &(nq, nd, dim) in full_range {
973            let query_data: Vec<i8> = (0..nq * dim).map(|v| (37 * v) as i8).collect();
974            let doc_data: Vec<i8> = (0..nd * dim).map(|v| (91 * v + 5) as i8).collect();
975            check_i8_isas(nq, nd, dim, &query_data, &doc_data);
976        }
977    }
978
979    #[test]
980    fn dimensions_f32() {
981        let data = vec![1.0f32; 5 * 8];
982        let query = make_mat(&data, 5, 8);
983        let kernel = build_max_sim::<f32, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
984        assert_eq!(kernel.nrows(), 5);
985    }
986
987    #[test]
988    fn dimensions_f16() {
989        let data = vec![diskann_wide::cast_f32_to_f16(1.0); 5 * 8];
990        let query = make_mat(data.as_slice(), 5, 8);
991        let kernel = build_max_sim::<half::f16, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
992        assert_eq!(kernel.nrows(), 5);
993    }
994
995    #[test]
996    fn dimensions_i8() {
997        let data = vec![1i8; 5 * 8];
998        let query = make_mat(&data, 5, 8);
999        let kernel = build_max_sim::<i8, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
1000        assert_eq!(kernel.nrows(), 5);
1001    }
1002
1003    #[test]
1004    fn i8_rejects_dim_too_large() {
1005        let data = vec![0i8; 131_072];
1006        let query = make_mat(&data, 1, 131_072);
1007
1008        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
1009            let err = build_max_sim::<i8, _>(isa, query, BoxErase).err();
1010            assert!(
1011                matches!(err, Some(BuildMaxSimError::DimTooLarge(131_072, 131_071))),
1012                "{isa:?}: expected DimTooLarge(131_072, 131_071), got {err:?}",
1013            );
1014        }
1015    }
1016
1017    #[test]
1018    #[cfg(not(miri))]
1019    fn i8_max_dim_is_exact() {
1020        let data = vec![i8::MIN; 131_071];
1021        let query = make_mat(&data, 1, 131_071);
1022
1023        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
1024            let kernel = build_max_sim::<i8, _>(isa, query, BoxErase).unwrap();
1025            let mut scores = scores_buffer::<i8>(1);
1026            kernel.compute_max_sim(query, &mut scores).unwrap();
1027            assert_eq!(scores[0], -131_071 * 16_384, "{isa:?}");
1028        }
1029    }
1030
1031    fn check_size_mismatch<T>(label: &str)
1032    where
1033        T: MaxSimElement + FromF32,
1034    {
1035        let query_data = make_test_data::<T>(3 * 4, 4, 0);
1036        let doc_data = make_test_data::<T>(2 * 4, 4, 1);
1037        let query = make_mat(&query_data, 3, 4);
1038        let doc = make_mat(&doc_data, 2, 4);
1039
1040        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
1041            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
1042
1043            let mut too_short = scores_buffer::<T>(2);
1044            match kernel.compute_max_sim(doc, &mut too_short) {
1045                Err(MaxSimError::InvalidBufferLength(2, 3)) => {}
1046                other => {
1047                    panic!("{label}({isa:?}) expected InvalidBufferLength(2, 3), got {other:?}",)
1048                }
1049            }
1050
1051            let mut too_long = scores_buffer::<T>(4);
1052            match kernel.compute_max_sim(doc, &mut too_long) {
1053                Err(MaxSimError::InvalidBufferLength(4, 3)) => {}
1054                other => {
1055                    panic!("{label}({isa:?}) expected InvalidBufferLength(4, 3), got {other:?}",)
1056                }
1057            }
1058        }
1059    }
1060
1061    fn check_zero_docs_fills_sentinel<T>(label: &str)
1062    where
1063        T: MaxSimElement + FromF32,
1064    {
1065        let query_data = make_test_data::<T>(3 * 4, 4, 0);
1066        let doc_data: Vec<T> = Vec::new();
1067        let query = make_mat(&query_data, 3, 4);
1068        let doc = make_mat(doc_data.as_slice(), 0, 4);
1069
1070        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
1071            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
1072            let mut scores = scores_buffer::<T>(3);
1073            kernel.compute_max_sim(doc, &mut scores).unwrap();
1074            for (i, &s) in scores.iter().enumerate() {
1075                assert_eq!(
1076                    s,
1077                    T::NO_MATCH,
1078                    "{label}({isa:?}) zero-doc slot {i} should be the NO_MATCH sentinel",
1079                );
1080            }
1081        }
1082    }
1083
1084    fn check_zero_query<T>(label: &str)
1085    where
1086        T: MaxSimElement + FromF32,
1087    {
1088        let query_data: Vec<T> = Vec::new();
1089        let doc_data = make_test_data::<T>(2 * 4, 4, 0);
1090        let query = make_mat(query_data.as_slice(), 0, 4);
1091        let doc = make_mat(&doc_data, 2, 4);
1092
1093        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
1094            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
1095            assert_eq!(
1096                kernel.nrows(),
1097                0,
1098                "{label}({isa:?}) empty query should yield nrows=0",
1099            );
1100            let mut scores = scores_buffer::<T>(0);
1101            kernel
1102                .compute_max_sim(doc, &mut scores)
1103                .unwrap_or_else(|e| panic!("{label}({isa:?}) expected Ok, got {e:?}"));
1104        }
1105    }
1106
1107    macro_rules! test_matches_fallback {
1108        ($mod_name:ident, $ty:ty, $tol:expr, $label:literal) => {
1109            mod $mod_name {
1110                use super::*;
1111
1112                #[test]
1113                fn chamfer_matches_fallback() {
1114                    check_chamfer_matches::<$ty>($tol, $label);
1115                }
1116
1117                #[test]
1118                fn max_sim_matches_fallback() {
1119                    check_max_sim_matches::<$ty>($tol, $label);
1120                }
1121
1122                #[test]
1123                fn errors_on_size_mismatch() {
1124                    check_size_mismatch::<$ty>($label);
1125                }
1126
1127                #[test]
1128                fn zero_docs_fills_sentinel() {
1129                    check_zero_docs_fills_sentinel::<$ty>($label);
1130                }
1131
1132                #[test]
1133                fn zero_query_returns_ok() {
1134                    check_zero_query::<$ty>($label);
1135                }
1136            }
1137        };
1138    }
1139
1140    test_matches_fallback!(f32, f32, 1e-10, "f32 ");
1141    test_matches_fallback!(f16, half::f16, 1e-10, "f16 ");
1142    test_matches_fallback!(i8, i8, 1e-10, "i8 ");
1143}
diskann_quantization/multi_vector/distance/factory.rs

diskann_quantization/multi_vector/distance/
factory.rs
