Skip to main content

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::distance::InnerProduct;
13use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
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};
20
21use super::isa::{MaxSimIsa, NotSupported};
22use super::kernel::{Erase, MaxSimKernel};
23use super::max_sim::{MaxSim, MaxSimError};
24use crate::matrix_kernels as mk;
25use crate::multi_vector::distance::QueryMatRef;
26use crate::multi_vector::{BlockTransposed, Mat, MatRef, Standard};
27
28// ─────────────────────────────────────────────────────────────────────────
29//  Prepared<A, Q, NR> — concrete kernel for the arch-dispatched paths.
30// ─────────────────────────────────────────────────────────────────────────
31
32// ZST for selecting the packing for the documents.
33//
34// TODO: Some of these decisions should eventually ve moved lower into `matrix_kernels`.
35#[derive(Debug, Clone, Copy)]
36struct Pack<const NR: usize>;
37
38#[derive(Debug)]
39struct Prepared<A, Q, const NR: usize> {
40    arch: A,
41    prepared: Q,
42    _packing: Pack<NR>,
43}
44
45impl<A, const GROUP: usize, const NR: usize> MaxSimKernel<f32>
46    for Prepared<A, BlockTransposed<f32, GROUP>, NR>
47where
48    A: Architecture,
49    for<'a> mk::maxsim::packed_f32_x_unpacked_f32::Driver<'a, A, GROUP, NR>: mk::Drive,
50{
51    fn nrows(&self) -> usize {
52        self.prepared.nrows()
53    }
54
55    fn compute_max_sim(
56        &self,
57        doc: MatRef<'_, Standard<f32>>,
58        scores: &mut [f32],
59    ) -> Result<(), MaxSimError> {
60        if scores.len() != self.nrows() {
61            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
62        }
63
64        if doc.vector_dim() != self.prepared.ncols() {
65            return Err(MaxSimError::UnequalDim(
66                doc.vector_dim(),
67                self.prepared.ncols(),
68            ));
69        }
70
71        let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
72            scores.fill(if doc.num_vectors() == 0 {
73                f32::MAX
74            } else {
75                0.0
76            });
77            return Ok(());
78        };
79
80        let Some(a) = mk::blocks::packed::View::from_block_transposed(self.prepared.as_view())
81        else {
82            return Ok(());
83        };
84
85        let Some(b) = mk::blocks::unpacked::View::from_matrix_view(doc.as_matrix_view()) else {
86            scores.fill(f32::MAX);
87            return Ok(());
88        };
89
90        // SAFETY: The dimension check establishes that `a.k() == b.k() == k`.
91        // The length check establishes that `scores` occupies exactly the
92        // packed blocks in `a`.
93        let mut driver = unsafe {
94            mk::maxsim::packed_f32_x_unpacked_f32::Driver::new(
95                self.arch,
96                a,
97                b,
98                scores,
99                k,
100                mk::Cache::detect(),
101            )
102        };
103
104        mk::Drive::drive(&mut driver);
105
106        scores.iter_mut().for_each(|s| *s = -*s);
107
108        Ok(())
109    }
110}
111
112impl<A, const GROUP: usize, const NR: usize> MaxSimKernel<half::f16>
113    for Prepared<A, BlockTransposed<f32, GROUP>, NR>
114where
115    A: Architecture,
116    for<'a> mk::maxsim::packed_f32_x_unpacked_f16::Driver<'a, A, GROUP, NR>: mk::Drive,
117{
118    fn nrows(&self) -> usize {
119        self.prepared.nrows()
120    }
121
122    fn compute_max_sim(
123        &self,
124        doc: MatRef<'_, Standard<half::f16>>,
125        scores: &mut [f32],
126    ) -> Result<(), MaxSimError> {
127        if scores.len() != self.nrows() {
128            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
129        }
130
131        if doc.vector_dim() != self.prepared.ncols() {
132            return Err(MaxSimError::UnequalDim(
133                doc.vector_dim(),
134                self.prepared.ncols(),
135            ));
136        }
137
138        let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
139            scores.fill(if doc.num_vectors() == 0 {
140                f32::MAX
141            } else {
142                0.0
143            });
144            return Ok(());
145        };
146
147        let Some(a) = mk::blocks::packed::View::from_block_transposed(self.prepared.as_view())
148        else {
149            return Ok(());
150        };
151
152        let Some(b) = mk::blocks::unpacked::View::from_matrix_view(doc.as_matrix_view()) else {
153            scores.fill(f32::MAX);
154            return Ok(());
155        };
156
157        // SAFETY: The dimension check establishes that `a.k() == b.k() == k`.
158        // The length check establishes that `scores` occupies exactly the
159        // packed blocks in `a`.
160        let mut driver = unsafe {
161            mk::maxsim::packed_f32_x_unpacked_f16::Driver::new(
162                self.arch,
163                a,
164                b,
165                scores,
166                k,
167                mk::Cache::detect(),
168            )
169        };
170
171        mk::Drive::drive(&mut driver);
172
173        scores.iter_mut().for_each(|s| *s = -*s);
174
175        Ok(())
176    }
177}
178
179// ─────────────────────────────────────────────────────────────────────────
180//  ReferenceKernel<T> — non-SIMD fallback that wraps MaxSim::evaluate.
181// ─────────────────────────────────────────────────────────────────────────
182
183struct ReferenceKernel<T: Copy> {
184    query: Mat<Standard<T>>,
185}
186
187impl<T: Copy + std::fmt::Debug> std::fmt::Debug for ReferenceKernel<T> {
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        f.debug_struct("ReferenceKernel")
190            .field("nrows", &self.query.num_vectors())
191            .finish()
192    }
193}
194
195impl<T: Copy> ReferenceKernel<T> {
196    fn new(query: MatRef<'_, Standard<T>>) -> Self {
197        Self {
198            query: query.to_owned(),
199        }
200    }
201}
202
203impl<T> MaxSimKernel<T> for ReferenceKernel<T>
204where
205    T: Copy + Send + Sync + std::fmt::Debug + 'static,
206    InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
207{
208    fn nrows(&self) -> usize {
209        self.query.num_vectors()
210    }
211
212    fn compute_max_sim(
213        &self,
214        doc: MatRef<'_, Standard<T>>,
215        scores: &mut [f32],
216    ) -> Result<(), MaxSimError> {
217        if scores.len() != self.nrows() {
218            return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
219        }
220        if doc.vector_dim() != self.query.vector_dim() {
221            return Err(MaxSimError::UnequalDim(
222                doc.vector_dim(),
223                self.query.vector_dim(),
224            ));
225        }
226        if doc.num_vectors() == 0 {
227            scores.fill(f32::MAX);
228            return Ok(());
229        }
230        let query: QueryMatRef<'_, Standard<T>> = self.query.as_view().into();
231        let mut max_sim = MaxSim::new(scores);
232        max_sim.evaluate(query, doc)
233    }
234}
235
236// ─────────────────────────────────────────────────────────────────────────
237//  BuildAndErase<E> — Target1 impls used by `dispatch1_no_features` (Auto).
238// ─────────────────────────────────────────────────────────────────────────
239
240struct BuildAndErase<E>(E);
241
242// ───── f32 Target1 impls ─────
243
244impl<E: Erase<f32>> diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<f32>>>
245    for BuildAndErase<E>
246{
247    fn run(self, arch: Scalar, query: MatRef<'_, Standard<f32>>) -> E::Output {
248        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
249        self.0.erase(Prepared {
250            arch,
251            prepared,
252            _packing: Pack::<2>,
253        })
254    }
255}
256
257#[cfg(target_arch = "x86_64")]
258impl<E: Erase<f32>> diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<f32>>>
259    for BuildAndErase<E>
260{
261    fn run(self, arch: V3, query: MatRef<'_, Standard<f32>>) -> E::Output {
262        let prepared = BlockTransposed::<f32, 16>::from_matrix_view(query.as_matrix_view());
263        self.0.erase(Prepared {
264            arch,
265            prepared,
266            _packing: Pack::<6>,
267        })
268    }
269}
270
271#[cfg(target_arch = "x86_64")]
272impl<E: Erase<f32>> diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<f32>>>
273    for BuildAndErase<E>
274{
275    fn run(self, arch: V4, query: MatRef<'_, Standard<f32>>) -> E::Output {
276        let prepared = BlockTransposed::<f32, 32>::from_matrix_view(query.as_matrix_view());
277        self.0.erase(Prepared {
278            arch,
279            prepared,
280            _packing: Pack::<6>,
281        })
282    }
283}
284
285#[cfg(target_arch = "aarch64")]
286impl<E: Erase<f32>> diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<f32>>>
287    for BuildAndErase<E>
288{
289    fn run(self, arch: Neon, query: MatRef<'_, Standard<f32>>) -> E::Output {
290        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
291        self.0.erase(Prepared {
292            arch,
293            prepared,
294            _packing: Pack::<6>,
295        })
296    }
297}
298
299// ───── f16 Target1 impls ─────
300
301impl<E: Erase<half::f16>>
302    diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<half::f16>>>
303    for BuildAndErase<E>
304{
305    fn run(self, arch: Scalar, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
306        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(
307            query
308                .as_matrix_view()
309                .map(|v| diskann_wide::cast_f16_to_f32(*v))
310                .as_view(),
311        );
312        self.0.erase(Prepared {
313            arch,
314            prepared,
315            _packing: Pack::<2>,
316        })
317    }
318}
319
320#[cfg(target_arch = "x86_64")]
321impl<E: Erase<half::f16>>
322    diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<half::f16>>>
323    for BuildAndErase<E>
324{
325    fn run(self, arch: V3, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
326        let prepared = BlockTransposed::<f32, 16>::from_matrix_view(
327            query
328                .as_matrix_view()
329                .map(|v| diskann_wide::cast_f16_to_f32(*v))
330                .as_view(),
331        );
332        self.0.erase(Prepared {
333            arch,
334            prepared,
335            _packing: Pack::<6>,
336        })
337    }
338}
339
340#[cfg(target_arch = "x86_64")]
341impl<E: Erase<half::f16>>
342    diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<half::f16>>>
343    for BuildAndErase<E>
344{
345    fn run(self, arch: V4, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
346        let prepared = BlockTransposed::<f32, 32>::from_matrix_view(
347            query
348                .as_matrix_view()
349                .map(|v| diskann_wide::cast_f16_to_f32(*v))
350                .as_view(),
351        );
352        self.0.erase(Prepared {
353            arch,
354            prepared,
355            _packing: Pack::<6>,
356        })
357    }
358}
359
360#[cfg(target_arch = "aarch64")]
361impl<E: Erase<half::f16>>
362    diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<half::f16>>>
363    for BuildAndErase<E>
364{
365    fn run(self, arch: Neon, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
366        // Neon dispatches to Scalar (no Neon-specific kernel).
367        let prepared = BlockTransposed::<f32, 8>::from_matrix_view(
368            query
369                .as_matrix_view()
370                .map(|v| diskann_wide::cast_f16_to_f32(*v))
371                .as_view(),
372        );
373        self.0.erase(Prepared {
374            arch,
375            prepared,
376            _packing: Pack::<6>,
377        })
378    }
379}
380
381// ─────────────────────────────────────────────────────────────────────────
382//  MaxSimElement — sealed trait gating accepted element types.
383// ─────────────────────────────────────────────────────────────────────────
384
385mod sealed {
386    pub trait Sealed {}
387}
388
389/// Scalar element types accepted by [`build_max_sim`].
390///
391/// Sealed: external crates cannot add impls. Quantized representations
392/// (PQ, SQ, packed sub-byte) are intentionally excluded — they need
393/// codebook/scale state that [`MatRef<'_, Standard<Self>>`] can't carry.
394pub trait MaxSimElement: sealed::Sealed + Sized + Copy + Send + Sync + 'static {
395    /// Build the concrete kernel for this element type and hand it to
396    /// `erase.erase(...)`.
397    ///
398    /// # Errors
399    ///
400    /// Returns [`NotSupported`] when the requested ISA cannot run on this
401    /// build (e.g. AVX-512 unavailable; aarch64 on x86_64).
402    fn build<E: Erase<Self>>(
403        isa: MaxSimIsa,
404        query: MatRef<'_, Standard<Self>>,
405        erase: E,
406    ) -> Result<E::Output, NotSupported>;
407}
408
409impl sealed::Sealed for f32 {}
410impl sealed::Sealed for half::f16 {}
411
412impl MaxSimElement for f32 {
413    fn build<E: Erase<f32>>(
414        isa: MaxSimIsa,
415        query: MatRef<'_, Standard<f32>>,
416        erase: E,
417    ) -> Result<E::Output, NotSupported> {
418        match isa {
419            MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
420                BuildAndErase(erase),
421                query,
422            )),
423            MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
424            #[cfg(target_arch = "x86_64")]
425            MaxSimIsa::X86_64_V3 => {
426                let arch = V3::new_checked().ok_or(NotSupported {
427                    isa,
428                    reason: "AVX2/FMA unavailable on this CPU",
429                })?;
430                Ok(arch.run1(BuildAndErase(erase), query))
431            }
432            #[cfg(target_arch = "x86_64")]
433            MaxSimIsa::X86_64_V4 => {
434                let arch = V4::new_checked().ok_or(NotSupported {
435                    isa,
436                    reason: "AVX-512 unavailable on this CPU",
437                })?;
438                Ok(arch.run1(BuildAndErase(erase), query))
439            }
440            #[cfg(not(target_arch = "x86_64"))]
441            MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
442                isa,
443                reason: "x86_64 target only",
444            }),
445            #[cfg(target_arch = "aarch64")]
446            MaxSimIsa::Neon => {
447                let arch = Neon::new_checked().ok_or(NotSupported {
448                    isa,
449                    reason: "Neon unavailable on this CPU",
450                })?;
451                Ok(arch.run1(BuildAndErase(erase), query))
452            }
453            #[cfg(not(target_arch = "aarch64"))]
454            MaxSimIsa::Neon => Err(NotSupported {
455                isa,
456                reason: "aarch64 target only",
457            }),
458            MaxSimIsa::Reference => Ok(erase.erase(ReferenceKernel::<f32>::new(query))),
459        }
460    }
461}
462
463impl MaxSimElement for half::f16 {
464    fn build<E: Erase<half::f16>>(
465        isa: MaxSimIsa,
466        query: MatRef<'_, Standard<half::f16>>,
467        erase: E,
468    ) -> Result<E::Output, NotSupported> {
469        match isa {
470            MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
471                BuildAndErase(erase),
472                query,
473            )),
474            MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
475            #[cfg(target_arch = "x86_64")]
476            MaxSimIsa::X86_64_V3 => {
477                let arch = V3::new_checked().ok_or(NotSupported {
478                    isa,
479                    reason: "AVX2/FMA unavailable on this CPU",
480                })?;
481                Ok(arch.run1(BuildAndErase(erase), query))
482            }
483            #[cfg(target_arch = "x86_64")]
484            MaxSimIsa::X86_64_V4 => {
485                let arch = V4::new_checked().ok_or(NotSupported {
486                    isa,
487                    reason: "AVX-512 unavailable on this CPU",
488                })?;
489                Ok(arch.run1(BuildAndErase(erase), query))
490            }
491            #[cfg(not(target_arch = "x86_64"))]
492            MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
493                isa,
494                reason: "x86_64 target only",
495            }),
496            #[cfg(target_arch = "aarch64")]
497            MaxSimIsa::Neon => {
498                let arch = Neon::new_checked().ok_or(NotSupported {
499                    isa,
500                    reason: "Neon unavailable on this CPU",
501                })?;
502                Ok(arch.run1(BuildAndErase(erase), query))
503            }
504            #[cfg(not(target_arch = "aarch64"))]
505            MaxSimIsa::Neon => Err(NotSupported {
506                isa,
507                reason: "aarch64 target only",
508            }),
509            MaxSimIsa::Reference => Ok(erase.erase(ReferenceKernel::<half::f16>::new(query))),
510        }
511    }
512}
513
514// ─────────────────────────────────────────────────────────────────────────
515//  Factory entry point.
516// ─────────────────────────────────────────────────────────────────────────
517
518/// Build a multi-vector MaxSim kernel for any [`MaxSimElement`] type.
519///
520/// Thin wrapper over [`MaxSimElement::build`] so callers don't have to name
521/// the trait at the call site.
522///
523/// # Errors
524///
525/// Returns [`NotSupported`] when the requested ISA cannot run on this build.
526pub fn build_max_sim<T: MaxSimElement, E: Erase<T>>(
527    isa: MaxSimIsa,
528    query: MatRef<'_, Standard<T>>,
529    erase: E,
530) -> Result<E::Output, NotSupported> {
531    T::build(isa, query, erase)
532}
533
534#[cfg(test)]
535mod tests {
536    use super::*;
537    use crate::multi_vector::{BoxErase, Chamfer, MaxSim, QueryMatRef};
538
539    /// Local helper trait — picks a sane test value of `T` from an `f32`
540    /// so both `f32` and `half::f16` parameterizations share the same data
541    /// generator.
542    trait FromF32 {
543        fn from_f32(v: f32) -> Self;
544    }
545
546    impl FromF32 for f32 {
547        fn from_f32(v: f32) -> Self {
548            v
549        }
550    }
551
552    impl FromF32 for half::f16 {
553        fn from_f32(v: f32) -> Self {
554            diskann_wide::cast_f32_to_f16(v)
555        }
556    }
557
558    fn make_mat<T: Copy>(data: &[T], nrows: usize, ncols: usize) -> MatRef<'_, Standard<T>> {
559        MatRef::new(Standard::new(nrows, ncols).unwrap(), data).unwrap()
560    }
561
562    fn make_test_data<T: FromF32>(len: usize, ceil: usize, shift: usize) -> Vec<T> {
563        (0..len)
564            .map(|v| T::from_f32(((v + shift) % ceil) as f32))
565            .collect()
566    }
567
568    /// Shapes for the `chamfer_matches_fallback` / `max_sim_matches_fallback`
569    /// agreement checks: `(num_queries, num_docs, dim)`.
570    ///
571    /// Targets the factory wiring (query setup, score writeback) above the
572    /// kernel layer; exhaustive panel/remainder coverage is pinned in
573    /// `kernels::tiled_reduce::tests`.
574    const TEST_CASES: &[(usize, usize, usize)] = &[
575        (0, 1, 1),   // Empty query
576        (1, 0, 1),   // Empty docs
577        (1, 1, 0),   // Empty dim
578        (1, 0, 0),   // Empty docs and dim
579        (1, 1, 4),   // Degenerate
580        (5, 3, 5),   // Prime k; nq > 1 and nd > 1 exercise per-row writeback
581        (17, 4, 64), // A-panel remainder crossing both Scalar and V3 panel widths
582        (16, 6, 32), // B-remainder ≠ 1 (V3 b_remainder = 2)
583    ];
584
585    fn check_chamfer_matches<T>(tol: f32, label: &str)
586    where
587        T: MaxSimElement + FromF32,
588        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
589    {
590        for &(nq, nd, dim) in TEST_CASES {
591            let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
592            let doc_data = make_test_data::<T>(nd * dim, dim, dim);
593
594            let query = make_mat(&query_data, nq, dim);
595            let doc = make_mat(&doc_data, nd, dim);
596
597            let expected = Chamfer::evaluate(QueryMatRef::from(query), doc);
598
599            let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
600            let mut scores = vec![0.0f32; nq];
601            kernel.compute_max_sim(doc, &mut scores).unwrap();
602            let actual: f32 = scores.iter().sum();
603
604            assert!(
605                (actual - expected).abs() < tol,
606                "{label}Chamfer mismatch for ({nq},{nd},{dim}): actual={actual}, expected={expected}",
607            );
608        }
609    }
610
611    fn check_max_sim_matches<T>(tol: f32, label: &str)
612    where
613        T: MaxSimElement + FromF32,
614        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
615    {
616        for &(nq, nd, dim) in TEST_CASES {
617            let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
618            let doc_data = make_test_data::<T>(nd * dim, dim, dim);
619
620            let query = make_mat(&query_data, nq, dim);
621            let doc = make_mat(&doc_data, nd, dim);
622
623            let mut expected_scores = vec![0.0f32; nq];
624            let _ = MaxSim::new(&mut expected_scores).evaluate(QueryMatRef::from(query), doc);
625
626            let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
627            let mut actual_scores = vec![0.0f32; nq];
628            kernel.compute_max_sim(doc, &mut actual_scores).unwrap();
629
630            for i in 0..nq {
631                assert!(
632                    (actual_scores[i] - expected_scores[i]).abs() < tol,
633                    "{label}MaxSim[{i}] mismatch for ({nq},{nd},{dim}): actual={}, expected={}",
634                    actual_scores[i],
635                    expected_scores[i],
636                );
637            }
638        }
639    }
640
641    #[test]
642    fn dimensions_f32() {
643        let data = vec![1.0f32; 5 * 8];
644        let query = make_mat(&data, 5, 8);
645        let kernel = build_max_sim::<f32, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
646        assert_eq!(kernel.nrows(), 5);
647    }
648
649    #[test]
650    fn dimensions_f16() {
651        let data = vec![diskann_wide::cast_f32_to_f16(1.0); 5 * 8];
652        let query = make_mat(data.as_slice(), 5, 8);
653        let kernel = build_max_sim::<half::f16, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
654        assert_eq!(kernel.nrows(), 5);
655    }
656
657    fn check_size_mismatch<T>(label: &str)
658    where
659        T: MaxSimElement + FromF32,
660        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
661    {
662        let query_data = make_test_data::<T>(3 * 4, 4, 0);
663        let doc_data = make_test_data::<T>(2 * 4, 4, 1);
664        let query = make_mat(&query_data, 3, 4);
665        let doc = make_mat(&doc_data, 2, 4);
666
667        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
668            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
669
670            let mut too_short = vec![0.0f32; 2];
671            match kernel.compute_max_sim(doc, &mut too_short) {
672                Err(MaxSimError::InvalidBufferLength(2, 3)) => {}
673                other => {
674                    panic!("{label}({isa:?}) expected InvalidBufferLength(2, 3), got {other:?}",)
675                }
676            }
677
678            let mut too_long = vec![0.0f32; 4];
679            match kernel.compute_max_sim(doc, &mut too_long) {
680                Err(MaxSimError::InvalidBufferLength(4, 3)) => {}
681                other => {
682                    panic!("{label}({isa:?}) expected InvalidBufferLength(4, 3), got {other:?}",)
683                }
684            }
685        }
686    }
687
688    fn check_zero_docs_fills_sentinel<T>(label: &str)
689    where
690        T: MaxSimElement + FromF32,
691        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
692    {
693        let query_data = make_test_data::<T>(3 * 4, 4, 0);
694        let doc_data: Vec<T> = Vec::new();
695        let query = make_mat(&query_data, 3, 4);
696        let doc = make_mat(doc_data.as_slice(), 0, 4);
697
698        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
699            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
700            let mut scores = vec![0.0f32; 3];
701            kernel.compute_max_sim(doc, &mut scores).unwrap();
702            for (i, &s) in scores.iter().enumerate() {
703                assert_eq!(
704                    s,
705                    f32::MAX,
706                    "{label}({isa:?}) zero-doc slot {i} should be f32::MAX sentinel",
707                );
708            }
709        }
710    }
711
712    fn check_zero_query<T>(label: &str)
713    where
714        T: MaxSimElement + FromF32,
715        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
716    {
717        let query_data: Vec<T> = Vec::new();
718        let doc_data = make_test_data::<T>(2 * 4, 4, 0);
719        let query = make_mat(query_data.as_slice(), 0, 4);
720        let doc = make_mat(&doc_data, 2, 4);
721
722        for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
723            let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
724            assert_eq!(
725                kernel.nrows(),
726                0,
727                "{label}({isa:?}) empty query should yield nrows=0",
728            );
729            let mut scores: Vec<f32> = Vec::new();
730            kernel
731                .compute_max_sim(doc, &mut scores)
732                .unwrap_or_else(|e| panic!("{label}({isa:?}) expected Ok, got {e:?}"));
733        }
734    }
735
736    macro_rules! test_matches_fallback {
737        ($mod_name:ident, $ty:ty, $tol:expr, $label:literal) => {
738            mod $mod_name {
739                use super::*;
740
741                #[test]
742                fn chamfer_matches_fallback() {
743                    check_chamfer_matches::<$ty>($tol, $label);
744                }
745
746                #[test]
747                fn max_sim_matches_fallback() {
748                    check_max_sim_matches::<$ty>($tol, $label);
749                }
750
751                #[test]
752                fn errors_on_size_mismatch() {
753                    check_size_mismatch::<$ty>($label);
754                }
755
756                #[test]
757                fn zero_docs_fills_sentinel() {
758                    check_zero_docs_fills_sentinel::<$ty>($label);
759                }
760
761                #[test]
762                fn zero_query_returns_ok() {
763                    check_zero_query::<$ty>($label);
764                }
765            }
766        };
767    }
768
769    test_matches_fallback!(f32, f32, 1e-10, "f32 ");
770    test_matches_fallback!(f16, half::f16, 1e-10, "f16 ");
771}