1use diskann_utils::Reborrow;
8use diskann_vector::distance::InnerProduct;
9use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
10use diskann_wide::Architecture;
11use diskann_wide::arch::Scalar;
12#[cfg(target_arch = "aarch64")]
13use diskann_wide::arch::aarch64::Neon;
14#[cfg(target_arch = "x86_64")]
15use diskann_wide::arch::x86_64::{V3, V4};
16
17use super::isa::{MaxSimIsa, NotSupported};
18use super::kernel::{Erase, MaxSimKernel};
19use super::kernels::f16::F16Entry;
20use super::kernels::f32::F32Kernel;
21use super::max_sim::{MaxSim, MaxSimError};
22use crate::multi_vector::distance::QueryMatRef;
23use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, Standard};
24
25#[derive(Debug)]
30struct Prepared<A, Q> {
31 arch: A,
32 prepared: Q,
33}
34
35impl<A, const GROUP: usize> MaxSimKernel<f32> for Prepared<A, BlockTransposed<f32, GROUP>>
36where
37 A: Architecture,
38 F32Kernel<GROUP>: for<'a> diskann_wide::arch::Target3<
39 A,
40 (),
41 BlockTransposedRef<'a, f32, GROUP>,
42 MatRef<'a, Standard<f32>>,
43 &'a mut [f32],
44 >,
45{
46 fn nrows(&self) -> usize {
47 self.prepared.nrows()
48 }
49
50 fn compute_max_sim(
51 &self,
52 doc: MatRef<'_, Standard<f32>>,
53 scores: &mut [f32],
54 ) -> Result<(), MaxSimError> {
55 if scores.len() != self.nrows() {
56 return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
57 }
58 if doc.num_vectors() == 0 {
59 scores.fill(f32::MAX);
60 return Ok(());
61 }
62 let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()];
63 self.arch.run3(
64 F32Kernel::<GROUP>,
65 self.prepared.reborrow(),
66 doc,
67 &mut scratch,
68 );
69 for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) {
70 *dst = -src;
71 }
72 Ok(())
73 }
74}
75
76impl<A, const GROUP: usize> MaxSimKernel<half::f16>
77 for Prepared<A, BlockTransposed<half::f16, GROUP>>
78where
79 A: Architecture,
80 F16Entry<GROUP>: for<'a> diskann_wide::arch::Target3<
81 A,
82 (),
83 BlockTransposedRef<'a, half::f16, GROUP>,
84 MatRef<'a, Standard<half::f16>>,
85 &'a mut [f32],
86 >,
87{
88 fn nrows(&self) -> usize {
89 self.prepared.nrows()
90 }
91
92 fn compute_max_sim(
93 &self,
94 doc: MatRef<'_, Standard<half::f16>>,
95 scores: &mut [f32],
96 ) -> Result<(), MaxSimError> {
97 if scores.len() != self.nrows() {
98 return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
99 }
100 if doc.num_vectors() == 0 {
101 scores.fill(f32::MAX);
102 return Ok(());
103 }
104 let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()];
105 self.arch.run3(
106 F16Entry::<GROUP>,
107 self.prepared.reborrow(),
108 doc,
109 &mut scratch,
110 );
111 for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) {
112 *dst = -src;
113 }
114 Ok(())
115 }
116}
117
118struct ReferenceKernel<T: Copy> {
123 query: Mat<Standard<T>>,
124}
125
126impl<T: Copy + std::fmt::Debug> std::fmt::Debug for ReferenceKernel<T> {
127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128 f.debug_struct("ReferenceKernel")
129 .field("nrows", &self.query.num_vectors())
130 .finish()
131 }
132}
133
134impl<T: Copy> ReferenceKernel<T> {
135 fn new(query: MatRef<'_, Standard<T>>) -> Self {
136 Self {
137 query: query.to_owned(),
138 }
139 }
140}
141
142impl<T> MaxSimKernel<T> for ReferenceKernel<T>
143where
144 T: Copy + Send + Sync + std::fmt::Debug + 'static,
145 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
146{
147 fn nrows(&self) -> usize {
148 self.query.num_vectors()
149 }
150
151 fn compute_max_sim(
152 &self,
153 doc: MatRef<'_, Standard<T>>,
154 scores: &mut [f32],
155 ) -> Result<(), MaxSimError> {
156 if scores.len() != self.nrows() {
157 return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows()));
158 }
159 if doc.num_vectors() == 0 {
160 scores.fill(f32::MAX);
161 return Ok(());
162 }
163 let query: QueryMatRef<'_, Standard<T>> = self.query.as_view().into();
164 let mut max_sim = MaxSim::new(scores);
165 max_sim.evaluate(query, doc)
166 }
167}
168
169struct BuildAndErase<E>(E);
174
175impl<E: Erase<f32>> diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<f32>>>
178 for BuildAndErase<E>
179{
180 fn run(self, arch: Scalar, query: MatRef<'_, Standard<f32>>) -> E::Output {
181 let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
182 self.0.erase(Prepared { arch, prepared })
183 }
184}
185
186#[cfg(target_arch = "x86_64")]
187impl<E: Erase<f32>> diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<f32>>>
188 for BuildAndErase<E>
189{
190 fn run(self, arch: V3, query: MatRef<'_, Standard<f32>>) -> E::Output {
191 let prepared = BlockTransposed::<f32, 16>::from_matrix_view(query.as_matrix_view());
192 self.0.erase(Prepared { arch, prepared })
193 }
194}
195
196#[cfg(target_arch = "x86_64")]
197impl<E: Erase<f32>> diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<f32>>>
198 for BuildAndErase<E>
199{
200 fn run(self, arch: V4, query: MatRef<'_, Standard<f32>>) -> E::Output {
201 let arch = arch.retarget();
203 let prepared = BlockTransposed::<f32, 16>::from_matrix_view(query.as_matrix_view());
204 self.0.erase(Prepared { arch, prepared })
205 }
206}
207
208#[cfg(target_arch = "aarch64")]
209impl<E: Erase<f32>> diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<f32>>>
210 for BuildAndErase<E>
211{
212 fn run(self, arch: Neon, query: MatRef<'_, Standard<f32>>) -> E::Output {
213 let arch = arch.retarget();
215 let prepared = BlockTransposed::<f32, 8>::from_matrix_view(query.as_matrix_view());
216 self.0.erase(Prepared { arch, prepared })
217 }
218}
219
220impl<E: Erase<half::f16>>
223 diskann_wide::arch::Target1<Scalar, E::Output, MatRef<'_, Standard<half::f16>>>
224 for BuildAndErase<E>
225{
226 fn run(self, arch: Scalar, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
227 let prepared = BlockTransposed::<half::f16, 8>::from_matrix_view(query.as_matrix_view());
228 self.0.erase(Prepared { arch, prepared })
229 }
230}
231
232#[cfg(target_arch = "x86_64")]
233impl<E: Erase<half::f16>>
234 diskann_wide::arch::Target1<V3, E::Output, MatRef<'_, Standard<half::f16>>>
235 for BuildAndErase<E>
236{
237 fn run(self, arch: V3, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
238 let prepared = BlockTransposed::<half::f16, 16>::from_matrix_view(query.as_matrix_view());
239 self.0.erase(Prepared { arch, prepared })
240 }
241}
242
243#[cfg(target_arch = "x86_64")]
244impl<E: Erase<half::f16>>
245 diskann_wide::arch::Target1<V4, E::Output, MatRef<'_, Standard<half::f16>>>
246 for BuildAndErase<E>
247{
248 fn run(self, arch: V4, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
249 let arch = arch.retarget();
251 let prepared = BlockTransposed::<half::f16, 16>::from_matrix_view(query.as_matrix_view());
252 self.0.erase(Prepared { arch, prepared })
253 }
254}
255
256#[cfg(target_arch = "aarch64")]
257impl<E: Erase<half::f16>>
258 diskann_wide::arch::Target1<Neon, E::Output, MatRef<'_, Standard<half::f16>>>
259 for BuildAndErase<E>
260{
261 fn run(self, arch: Neon, query: MatRef<'_, Standard<half::f16>>) -> E::Output {
262 let arch = arch.retarget();
264 let prepared = BlockTransposed::<half::f16, 8>::from_matrix_view(query.as_matrix_view());
265 self.0.erase(Prepared { arch, prepared })
266 }
267}
268
269mod sealed {
274 pub trait Sealed {}
275}
276
277pub trait MaxSimElement: sealed::Sealed + Sized + Copy + Send + Sync + 'static {
283 fn build<E: Erase<Self>>(
291 isa: MaxSimIsa,
292 query: MatRef<'_, Standard<Self>>,
293 erase: E,
294 ) -> Result<E::Output, NotSupported>;
295}
296
297impl sealed::Sealed for f32 {}
298impl sealed::Sealed for half::f16 {}
299
300impl MaxSimElement for f32 {
301 fn build<E: Erase<f32>>(
302 isa: MaxSimIsa,
303 query: MatRef<'_, Standard<f32>>,
304 erase: E,
305 ) -> Result<E::Output, NotSupported> {
306 match isa {
307 MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
308 BuildAndErase(erase),
309 query,
310 )),
311 MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
312 #[cfg(target_arch = "x86_64")]
313 MaxSimIsa::X86_64_V3 => {
314 let arch = V3::new_checked().ok_or(NotSupported {
315 isa,
316 reason: "AVX2/FMA unavailable on this CPU",
317 })?;
318 Ok(arch.run1(BuildAndErase(erase), query))
319 }
320 #[cfg(target_arch = "x86_64")]
321 MaxSimIsa::X86_64_V4 => {
322 let arch = V4::new_checked().ok_or(NotSupported {
323 isa,
324 reason: "AVX-512 unavailable on this CPU",
325 })?;
326 Ok(arch.run1(BuildAndErase(erase), query))
327 }
328 #[cfg(not(target_arch = "x86_64"))]
329 MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
330 isa,
331 reason: "x86_64 target only",
332 }),
333 #[cfg(target_arch = "aarch64")]
334 MaxSimIsa::Neon => {
335 let arch = Neon::new_checked().ok_or(NotSupported {
336 isa,
337 reason: "Neon unavailable on this CPU",
338 })?;
339 Ok(arch.run1(BuildAndErase(erase), query))
340 }
341 #[cfg(not(target_arch = "aarch64"))]
342 MaxSimIsa::Neon => Err(NotSupported {
343 isa,
344 reason: "aarch64 target only",
345 }),
346 MaxSimIsa::Reference => Ok(erase.erase(ReferenceKernel::<f32>::new(query))),
347 }
348 }
349}
350
351impl MaxSimElement for half::f16 {
352 fn build<E: Erase<half::f16>>(
353 isa: MaxSimIsa,
354 query: MatRef<'_, Standard<half::f16>>,
355 erase: E,
356 ) -> Result<E::Output, NotSupported> {
357 match isa {
358 MaxSimIsa::Auto => Ok(diskann_wide::arch::dispatch1_no_features(
359 BuildAndErase(erase),
360 query,
361 )),
362 MaxSimIsa::Scalar => Ok(Scalar::new().run1(BuildAndErase(erase), query)),
363 #[cfg(target_arch = "x86_64")]
364 MaxSimIsa::X86_64_V3 => {
365 let arch = V3::new_checked().ok_or(NotSupported {
366 isa,
367 reason: "AVX2/FMA unavailable on this CPU",
368 })?;
369 Ok(arch.run1(BuildAndErase(erase), query))
370 }
371 #[cfg(target_arch = "x86_64")]
372 MaxSimIsa::X86_64_V4 => {
373 let arch = V4::new_checked().ok_or(NotSupported {
374 isa,
375 reason: "AVX-512 unavailable on this CPU",
376 })?;
377 Ok(arch.run1(BuildAndErase(erase), query))
378 }
379 #[cfg(not(target_arch = "x86_64"))]
380 MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported {
381 isa,
382 reason: "x86_64 target only",
383 }),
384 #[cfg(target_arch = "aarch64")]
385 MaxSimIsa::Neon => {
386 let arch = Neon::new_checked().ok_or(NotSupported {
387 isa,
388 reason: "Neon unavailable on this CPU",
389 })?;
390 Ok(arch.run1(BuildAndErase(erase), query))
391 }
392 #[cfg(not(target_arch = "aarch64"))]
393 MaxSimIsa::Neon => Err(NotSupported {
394 isa,
395 reason: "aarch64 target only",
396 }),
397 MaxSimIsa::Reference => Ok(erase.erase(ReferenceKernel::<half::f16>::new(query))),
398 }
399 }
400}
401
402pub fn build_max_sim<T: MaxSimElement, E: Erase<T>>(
415 isa: MaxSimIsa,
416 query: MatRef<'_, Standard<T>>,
417 erase: E,
418) -> Result<E::Output, NotSupported> {
419 T::build(isa, query, erase)
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use crate::multi_vector::{BoxErase, Chamfer, MaxSim, QueryMatRef};
426
427 trait FromF32 {
431 fn from_f32(v: f32) -> Self;
432 }
433
434 impl FromF32 for f32 {
435 fn from_f32(v: f32) -> Self {
436 v
437 }
438 }
439
440 impl FromF32 for half::f16 {
441 fn from_f32(v: f32) -> Self {
442 diskann_wide::cast_f32_to_f16(v)
443 }
444 }
445
446 fn make_mat<T: Copy>(data: &[T], nrows: usize, ncols: usize) -> MatRef<'_, Standard<T>> {
447 MatRef::new(Standard::new(nrows, ncols).unwrap(), data).unwrap()
448 }
449
450 fn make_test_data<T: FromF32>(len: usize, ceil: usize, shift: usize) -> Vec<T> {
451 (0..len)
452 .map(|v| T::from_f32(((v + shift) % ceil) as f32))
453 .collect()
454 }
455
456 const TEST_CASES: &[(usize, usize, usize)] = &[
463 (1, 1, 4), (5, 3, 5), (17, 4, 64), (16, 6, 32), ];
468
469 fn check_chamfer_matches<T>(tol: f32, label: &str)
470 where
471 T: MaxSimElement + FromF32,
472 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
473 {
474 for &(nq, nd, dim) in TEST_CASES {
475 let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
476 let doc_data = make_test_data::<T>(nd * dim, dim, dim);
477
478 let query = make_mat(&query_data, nq, dim);
479 let doc = make_mat(&doc_data, nd, dim);
480
481 let expected = Chamfer::evaluate(QueryMatRef::from(query), doc);
482
483 let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
484 let mut scores = vec![0.0f32; nq];
485 kernel.compute_max_sim(doc, &mut scores).unwrap();
486 let actual: f32 = scores.iter().sum();
487
488 assert!(
489 (actual - expected).abs() < tol,
490 "{label}Chamfer mismatch for ({nq},{nd},{dim}): actual={actual}, expected={expected}",
491 );
492 }
493 }
494
495 fn check_max_sim_matches<T>(tol: f32, label: &str)
496 where
497 T: MaxSimElement + FromF32,
498 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
499 {
500 for &(nq, nd, dim) in TEST_CASES {
501 let query_data = make_test_data::<T>(nq * dim, dim, dim / 2);
502 let doc_data = make_test_data::<T>(nd * dim, dim, dim);
503
504 let query = make_mat(&query_data, nq, dim);
505 let doc = make_mat(&doc_data, nd, dim);
506
507 let mut expected_scores = vec![0.0f32; nq];
508 let _ = MaxSim::new(&mut expected_scores).evaluate(QueryMatRef::from(query), doc);
509
510 let kernel = build_max_sim::<T, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
511 let mut actual_scores = vec![0.0f32; nq];
512 kernel.compute_max_sim(doc, &mut actual_scores).unwrap();
513
514 for i in 0..nq {
515 assert!(
516 (actual_scores[i] - expected_scores[i]).abs() < tol,
517 "{label}MaxSim[{i}] mismatch for ({nq},{nd},{dim}): actual={}, expected={}",
518 actual_scores[i],
519 expected_scores[i],
520 );
521 }
522 }
523 }
524
525 #[test]
526 fn dimensions_f32() {
527 let data = vec![1.0f32; 5 * 8];
528 let query = make_mat(&data, 5, 8);
529 let kernel = build_max_sim::<f32, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
530 assert_eq!(kernel.nrows(), 5);
531 }
532
533 #[test]
534 fn dimensions_f16() {
535 let data = vec![diskann_wide::cast_f32_to_f16(1.0); 5 * 8];
536 let query = make_mat(data.as_slice(), 5, 8);
537 let kernel = build_max_sim::<half::f16, _>(MaxSimIsa::Auto, query, BoxErase).unwrap();
538 assert_eq!(kernel.nrows(), 5);
539 }
540
541 fn check_size_mismatch<T>(label: &str)
542 where
543 T: MaxSimElement + FromF32,
544 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
545 {
546 let query_data = make_test_data::<T>(3 * 4, 4, 0);
547 let doc_data = make_test_data::<T>(2 * 4, 4, 1);
548 let query = make_mat(&query_data, 3, 4);
549 let doc = make_mat(&doc_data, 2, 4);
550
551 for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
552 let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
553
554 let mut too_short = vec![0.0f32; 2];
555 match kernel.compute_max_sim(doc, &mut too_short) {
556 Err(MaxSimError::InvalidBufferLength(2, 3)) => {}
557 other => {
558 panic!("{label}({isa:?}) expected InvalidBufferLength(2, 3), got {other:?}",)
559 }
560 }
561
562 let mut too_long = vec![0.0f32; 4];
563 match kernel.compute_max_sim(doc, &mut too_long) {
564 Err(MaxSimError::InvalidBufferLength(4, 3)) => {}
565 other => {
566 panic!("{label}({isa:?}) expected InvalidBufferLength(4, 3), got {other:?}",)
567 }
568 }
569 }
570 }
571
572 fn check_zero_docs_fills_sentinel<T>(label: &str)
573 where
574 T: MaxSimElement + FromF32,
575 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
576 {
577 let query_data = make_test_data::<T>(3 * 4, 4, 0);
578 let doc_data: Vec<T> = Vec::new();
579 let query = make_mat(&query_data, 3, 4);
580 let doc = make_mat(doc_data.as_slice(), 0, 4);
581
582 for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
583 let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
584 let mut scores = vec![0.0f32; 3];
585 kernel.compute_max_sim(doc, &mut scores).unwrap();
586 for (i, &s) in scores.iter().enumerate() {
587 assert_eq!(
588 s,
589 f32::MAX,
590 "{label}({isa:?}) zero-doc slot {i} should be f32::MAX sentinel",
591 );
592 }
593 }
594 }
595
596 fn check_zero_query<T>(label: &str)
597 where
598 T: MaxSimElement + FromF32,
599 InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
600 {
601 let query_data: Vec<T> = Vec::new();
602 let doc_data = make_test_data::<T>(2 * 4, 4, 0);
603 let query = make_mat(query_data.as_slice(), 0, 4);
604 let doc = make_mat(&doc_data, 2, 4);
605
606 for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] {
607 let kernel = build_max_sim::<T, _>(isa, query, BoxErase).unwrap();
608 assert_eq!(
609 kernel.nrows(),
610 0,
611 "{label}({isa:?}) empty query should yield nrows=0",
612 );
613 let mut scores: Vec<f32> = Vec::new();
614 kernel
615 .compute_max_sim(doc, &mut scores)
616 .unwrap_or_else(|e| panic!("{label}({isa:?}) expected Ok, got {e:?}"));
617 }
618 }
619
620 macro_rules! test_matches_fallback {
621 ($mod_name:ident, $ty:ty, $tol:expr, $label:literal) => {
622 mod $mod_name {
623 use super::*;
624
625 #[test]
626 fn chamfer_matches_fallback() {
627 check_chamfer_matches::<$ty>($tol, $label);
628 }
629
630 #[test]
631 fn max_sim_matches_fallback() {
632 check_max_sim_matches::<$ty>($tol, $label);
633 }
634
635 #[test]
636 fn errors_on_size_mismatch() {
637 check_size_mismatch::<$ty>($label);
638 }
639
640 #[test]
641 fn zero_docs_fills_sentinel() {
642 check_zero_docs_fills_sentinel::<$ty>($label);
643 }
644
645 #[test]
646 fn zero_query_returns_ok() {
647 check_zero_query::<$ty>($label);
648 }
649 }
650 };
651 }
652
653 test_matches_fallback!(f32, f32, 1e-10, "f32 ");
654 test_matches_fallback!(f16, half::f16, 1e-10, "f16 ");
655}