1#[cfg(target_arch = "aarch64")]
7use diskann_wide::arch::aarch64::Neon;
8#[cfg(target_arch = "x86_64")]
9use diskann_wide::arch::x86_64::{V3, V4};
10use diskann_wide::{
11 arch::{Dispatched2, FTarget2, Scalar},
12 Architecture,
13};
14use half::f16;
15
16use super::{implementations::Specialize, Cosine, CosineNormalized, InnerProduct, SquaredL2};
17use crate::{distance::Metric, AsUnaligned, UnalignedSlice};
18
19pub trait DistanceProvider<T>: Sized + 'static {
45 fn distance_comparer(metric: Metric, dimension: Option<usize>) -> Distance<Self, T>;
46}
47
48#[derive(Debug)]
49struct Unaligned<T>(std::marker::PhantomData<T>);
50
51impl<T> diskann_wide::lifetime::AddLifetime for Unaligned<T>
52where
53 T: 'static,
54{
55 type Of<'a> = UnalignedSlice<'a, T>;
56}
57
58#[derive(Debug)]
62pub struct Distance<T, U>
63where
64 T: 'static,
65 U: 'static,
66{
67 f: Dispatched2<f32, Unaligned<T>, Unaligned<U>>,
68}
69
70impl<T, U> Distance<T, U>
71where
72 T: 'static,
73 U: 'static,
74{
75 fn over(f: Dispatched2<f32, Unaligned<T>, Unaligned<U>>) -> Self {
76 Self { f }
77 }
78
79 pub fn new(metric: Metric, dim: Option<usize>) -> Self
81 where
82 T: DistanceProvider<U>,
83 {
84 T::distance_comparer(metric, dim)
85 }
86
87 #[inline]
94 pub fn call(&self, x: &[T], y: &[U]) -> f32 {
95 self.call_unaligned(x.as_unaligned(), y.as_unaligned())
96 }
97
98 #[inline]
105 pub fn call_unaligned(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
106 self.f.call(x, y)
107 }
108}
109
110impl<T, U> Clone for Distance<T, U>
111where
112 T: 'static,
113 U: 'static,
114{
115 fn clone(&self) -> Self {
116 *self
117 }
118}
119
120impl<T, U> Copy for Distance<T, U>
121where
122 T: 'static,
123 U: 'static,
124{
125}
126
127impl<T, U> crate::DistanceFunction<&[T], &[U], f32> for Distance<T, U>
128where
129 T: 'static,
130 U: 'static,
131{
132 fn evaluate_similarity(&self, x: &[T], y: &[U]) -> f32 {
133 self.call(x, y)
134 }
135}
136
137impl<T, U> crate::DistanceFunction<UnalignedSlice<'_, T>, UnalignedSlice<'_, U>, f32>
138 for Distance<T, U>
139where
140 T: 'static,
141 U: 'static,
142{
143 fn evaluate_similarity(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
144 self.f.call(x, y)
145 }
146}
147
148macro_rules! provider {
201 ($T:ty, $U:ty) => {
202 impl DistanceProvider<$U> for $T {
203 fn distance_comparer(metric: Metric, dimension: Option<usize>) -> Distance<$T, $U> {
204 diskann_wide::arch::dispatch2_no_features(
209 ArgumentTypes::<$T, $U>::new(),
210 metric,
211 dimension,
212 )
213 }
214 }
215 };
216}
217
218provider!(f32, f32);
219provider!(f16, f16);
220provider!(f32, f16);
221provider!(i8, i8);
222provider!(u8, u8);
223
224macro_rules! spec_list {
229 ($($Ns:literal),* $(,)?) => {
230 spec_list!(@value, $($Ns,)*)
231 };
232 (@value $(,)?) => {
233 Cons::new(Null, Null)
234 };
235 (@value, $N0:literal $(,)?) => {
236 Cons::new(Spec::<$N0>, Null)
237 };
238 (@value, $N0:literal, $N1:literal $(,)?) => {
239 Cons::new(Spec::<$N0>, Spec::<$N1>)
240 };
241 (@value, $N0:literal, $N1:literal, $($Ns:literal),+ $(,)?) => {
242 Cons::new(Spec::<$N0>, spec_list!(@value, $N1, $($Ns,)+))
243 };
244}
245
246struct ArgumentTypes<T: 'static, U: 'static>(std::marker::PhantomData<(T, U)>);
247
248impl<T, U> ArgumentTypes<T, U>
249where
250 T: 'static,
251 U: 'static,
252{
253 fn new() -> Self {
254 Self(std::marker::PhantomData)
255 }
256}
257
258macro_rules! specialize {
259 ($arch:ty, $T:ty, $U:ty, $($Ns:literal),* $(,)?) => {
260 impl diskann_wide::arch::Target2<
261 $arch,
262 Distance<$T, $U>,
263 Metric,
264 Option<usize>,
265 > for ArgumentTypes<$T, $U> {
266 fn run(
267 self,
268 arch: $arch,
269 metric: Metric,
270 dim: Option<usize>,
271 ) -> Distance<$T, $U> {
272 let spec = spec_list!($($Ns),*);
273 match metric {
274 Metric::L2 => spec.specialize(arch, SquaredL2 {}, dim),
275 Metric::Cosine => spec.specialize(arch, Cosine {}, dim),
276 Metric::CosineNormalized => spec.specialize(arch, CosineNormalized {}, dim),
277 Metric::InnerProduct => spec.specialize(arch, InnerProduct {}, dim),
278 }
279 }
280 }
281 };
282 (@integer, $arch:ty, $T:ty, $U:ty, $($Ns:literal),* $(,)?) => {
284 impl diskann_wide::arch::Target2<
285 $arch,
286 Distance<$T, $U>,
287 Metric,
288 Option<usize>,
289 > for ArgumentTypes<$T, $U> {
290 fn run(
291 self,
292 arch: $arch,
293 metric: Metric,
294 dim: Option<usize>,
295 ) -> Distance<$T, $U> {
296 let spec = spec_list!($($Ns),*);
297 match metric {
298 Metric::L2 => spec.specialize(arch, SquaredL2 {}, dim),
299 Metric::Cosine | Metric::CosineNormalized => {
300 spec.specialize(arch, Cosine {}, dim)
301 },
302 Metric::InnerProduct => spec.specialize(arch, InnerProduct {}, dim),
303 }
304 }
305 }
306 };
307}
308
309specialize!(Scalar, f32, f32,);
310specialize!(Scalar, f32, f16,);
311specialize!(Scalar, f16, f16,);
312specialize!(@integer, Scalar, u8, u8,);
313specialize!(@integer, Scalar, i8, i8,);
314
315#[cfg(target_arch = "x86_64")]
316mod x86_64 {
317 use super::*;
318
319 specialize!(V3, f32, f32, 768, 384, 128, 100);
320 specialize!(V4, f32, f32, 768, 384, 128, 100);
321
322 specialize!(V3, f32, f16, 768, 384, 128, 100);
323 specialize!(V4, f32, f16, 768, 384, 128, 100);
324
325 specialize!(V3, f16, f16, 768, 384, 128, 100);
326 specialize!(V4, f16, f16, 768, 384, 128, 100);
327
328 specialize!(@integer, V3, u8, u8, 128);
329 specialize!(@integer, V4, u8, u8, 128);
330
331 specialize!(@integer, V3, i8, i8, 128, 100);
332 specialize!(@integer, V4, i8, i8, 128, 100);
333}
334
335#[cfg(target_arch = "aarch64")]
336mod aarch64 {
337 use super::*;
338
339 specialize!(Neon, f32, f32, 768, 384, 128, 100);
340 specialize!(Neon, f32, f16, 768, 384, 128, 100);
341 specialize!(Neon, f16, f16, 768, 384, 128, 100);
342
343 specialize!(@integer, Neon, u8, u8, 128);
344 specialize!(@integer, Neon, i8, i8, 128, 100);
345}
346
347trait TrySpecialize<A, F, T, U>
350where
351 A: Architecture,
352 T: 'static,
353 U: 'static,
354{
355 fn try_specialize(&self, arch: A, dim: Option<usize>) -> Option<Distance<T, U>>;
356}
357
358struct Spec<const N: usize>;
360
361impl<A, F, const N: usize, T, U> TrySpecialize<A, F, T, U> for Spec<N>
362where
363 A: Architecture,
364 Specialize<N, F>: for<'a, 'b> FTarget2<A, f32, UnalignedSlice<'a, T>, UnalignedSlice<'b, U>>,
365 T: 'static,
366 U: 'static,
367{
368 fn try_specialize(&self, arch: A, dim: Option<usize>) -> Option<Distance<T, U>> {
369 if let Some(d) = dim {
370 if d == N {
371 return Some(Distance::over(
372 arch.dispatch2::<Specialize<N, F>, f32, Unaligned<T>, Unaligned<U>>(),
374 ));
375 }
376 }
377 None
378 }
379}
380
381struct Null;
383
384impl<A, F, T, U> TrySpecialize<A, F, T, U> for Null
385where
386 A: Architecture,
387 T: 'static,
388 U: 'static,
389{
390 fn try_specialize(&self, _arch: A, _dim: Option<usize>) -> Option<Distance<T, U>> {
391 None
392 }
393}
394
395struct Cons<Head, Tail> {
397 head: Head,
398 tail: Tail,
399}
400
401impl<Head, Tail> Cons<Head, Tail> {
402 const fn new(head: Head, tail: Tail) -> Self {
403 Self { head, tail }
404 }
405
406 fn specialize<A, F, T, U>(&self, arch: A, _f: F, dim: Option<usize>) -> Distance<T, U>
409 where
410 A: Architecture,
411 F: for<'a, 'b> FTarget2<A, f32, UnalignedSlice<'a, T>, UnalignedSlice<'b, U>>,
412 Head: TrySpecialize<A, F, T, U>,
413 Tail: TrySpecialize<A, F, T, U>,
414 T: 'static,
415 U: 'static,
416 {
417 if let Some(f) = self.try_specialize(arch, dim) {
418 f
419 } else {
420 Distance::over(arch.dispatch2::<F, f32, Unaligned<T>, Unaligned<U>>())
421 }
422 }
423}
424
425impl<A, Head, Tail, F, T, U> TrySpecialize<A, F, T, U> for Cons<Head, Tail>
427where
428 A: Architecture,
429 Head: TrySpecialize<A, F, T, U>,
430 Tail: TrySpecialize<A, F, T, U>,
431 T: 'static,
432 U: 'static,
433{
434 fn try_specialize(&self, arch: A, dim: Option<usize>) -> Option<Distance<T, U>> {
435 if let Some(f) = self.head.try_specialize(arch, dim) {
436 Some(f)
437 } else {
438 self.tail.try_specialize(arch, dim)
439 }
440 }
441}
442
443#[cfg(test)]
448mod test_unaligned_distance_provider {
449 use approx::assert_relative_eq;
450 use rand::{self, SeedableRng};
451
452 use super::*;
453 use crate::{
454 distance::{reference::ReferenceProvider, Metric},
455 test_util, SimilarityScore,
456 };
457
458 struct EpsilonAndRelative {
460 epsilon: f32,
461 max_relative: f32,
462 }
463
464 fn get_float_bounds(metric: Metric) -> EpsilonAndRelative {
474 match metric {
475 Metric::L2 => EpsilonAndRelative {
476 epsilon: 1e-5,
477 max_relative: 1e-5,
478 },
479 Metric::InnerProduct => EpsilonAndRelative {
480 epsilon: 1e-4,
481 max_relative: 1e-4,
482 },
483 Metric::Cosine => EpsilonAndRelative {
484 epsilon: 1e-4,
485 max_relative: 1e-4,
486 },
487 Metric::CosineNormalized => EpsilonAndRelative {
488 epsilon: 1e-4,
489 max_relative: 1e-4,
490 },
491 }
492 }
493
494 fn get_int_bounds(metric: Metric) -> EpsilonAndRelative {
495 match metric {
496 Metric::Cosine | Metric::CosineNormalized => EpsilonAndRelative {
498 epsilon: 1e-6,
499 max_relative: 1e-6,
500 },
501 Metric::L2 | Metric::InnerProduct => EpsilonAndRelative {
503 epsilon: 0.0,
504 max_relative: 0.0,
505 },
506 }
507 }
508
509 fn do_test<T, Distribution>(
510 under_test: Distance<T, T>,
511 reference: fn(&[T], &[T]) -> SimilarityScore<f32>,
512 bounds: EpsilonAndRelative,
513 dim: usize,
514 distribution: Distribution,
515 ) where
516 T: test_util::CornerCases,
517 Distribution: test_util::GenerateRandomArguments<T> + Clone,
518 {
519 let mut rng = rand::rngs::StdRng::seed_from_u64(0xef0053c);
520
521 let converted = |a: &[T], b: &[T]| -> f32 { reference(a, b).into_inner() };
523
524 let mut checker = test_util::Checker::<T, T, f32>::new(
525 |a, b| under_test.call(a, b),
526 converted,
527 |got: f32, expected: f32| {
528 assert_relative_eq!(
529 got,
530 expected,
531 epsilon = bounds.epsilon,
532 max_relative = bounds.max_relative
533 );
534 },
535 );
536
537 test_util::test_distance_function(
538 &mut checker,
539 distribution.clone(),
540 distribution.clone(),
541 dim,
542 10,
543 &mut rng,
544 );
545 }
546
547 fn all_metrics() -> [Metric; 4] {
548 [
549 Metric::L2,
550 Metric::InnerProduct,
551 Metric::Cosine,
552 Metric::CosineNormalized,
553 ]
554 }
555
556 const MAX_DIM: usize = 256;
558
559 #[test]
560 fn test_unaligned_f32() {
561 let dist = rand_distr::Normal::new(0.0, 1.0).unwrap();
562 for metric in all_metrics() {
563 for dim in 0..MAX_DIM {
564 println!("Metric = {:?}, dim = {}", metric, dim);
565 let unaligned = <f32 as DistanceProvider<f32>>::distance_comparer(metric, None);
566 let simple = <f32 as ReferenceProvider<f32>>::reference_implementation(metric);
567 let bounds = get_float_bounds(metric);
568 do_test(unaligned, simple, bounds, dim, dist);
569 }
570 }
571 }
572
573 #[test]
574 fn test_unaligned_f16() {
575 let dist = rand_distr::Normal::new(0.0, 1.0).unwrap();
576 for metric in all_metrics() {
577 for dim in 0..MAX_DIM {
578 println!("Metric = {:?}, dim = {}", metric, dim);
579 let unaligned = <f16 as DistanceProvider<f16>>::distance_comparer(metric, None);
580 let simple = <f16 as ReferenceProvider<f16>>::reference_implementation(metric);
581 let bounds = get_float_bounds(metric);
582 do_test(unaligned, simple, bounds, dim, dist);
583 }
584 }
585 }
586
587 #[test]
588 fn test_unaligned_u8() {
589 let dist = rand::distr::StandardUniform {};
590 for metric in all_metrics() {
591 for dim in 0..MAX_DIM {
592 println!("Metric = {:?}, dim = {}", metric, dim);
593 let unaligned = <u8 as DistanceProvider<u8>>::distance_comparer(metric, None);
594 let simple = <u8 as ReferenceProvider<u8>>::reference_implementation(metric);
595 let bounds = get_int_bounds(metric);
596 do_test(unaligned, simple, bounds, dim, dist);
597 }
598 }
599 }
600
601 #[test]
602 fn test_unaligned_i8() {
603 let dist = rand::distr::StandardUniform {};
604 for metric in all_metrics() {
605 for dim in 0..MAX_DIM {
606 println!("Metric = {:?}, dim = {}", metric, dim);
607 let unaligned = <i8 as DistanceProvider<i8>>::distance_comparer(metric, None);
608 let simple = <i8 as ReferenceProvider<i8>>::reference_implementation(metric);
609
610 let bounds = get_int_bounds(metric);
611 do_test(unaligned, simple, bounds, dim, dist);
612 }
613 }
614 }
615}
616
617#[cfg(test)]
618mod distance_provider_f32_tests {
619 use approx::assert_abs_diff_eq;
620 use rand::{rngs::StdRng, Rng, SeedableRng};
621
622 use super::*;
623 use crate::{distance::reference, test_util::*};
624
625 #[repr(C, align(32))]
626 pub struct F32Slice112([f32; 112]);
627 #[repr(C, align(32))]
628 pub struct F32Slice104([f32; 104]);
629 #[repr(C, align(32))]
630 pub struct F32Slice128([f32; 128]);
631 #[repr(C, align(32))]
632 pub struct F32Slice256([f32; 256]);
633 #[repr(C, align(32))]
634 pub struct F32Slice4096([f32; 4096]);
635
636 pub fn get_turing_test_data_f32_dim(dim: usize) -> (Vec<f32>, Vec<f32>) {
637 let mut a_slice = vec![0.0f32; dim];
638 let mut b_slice = vec![0.0f32; dim];
639
640 let mut rng = StdRng::seed_from_u64(42);
641 for i in 0..dim {
642 a_slice[i] = rng.random_range(-1.0..1.0);
643 b_slice[i] = rng.random_range(-1.0..1.0);
644 }
645
646 ((a_slice), (b_slice))
647 }
648
649 #[test]
650 fn test_dist_l2_float_turing_104() {
651 let (a_data, b_data) = get_turing_test_data_f32_dim(104);
652 let (a_slice, b_slice) = (
653 F32Slice104(a_data.try_into().unwrap()),
654 F32Slice104(b_data.try_into().unwrap()),
655 );
656
657 let distance: f32 = compare_two_vec::<f32>(104, Metric::L2, &a_slice.0, &b_slice.0);
658
659 assert_abs_diff_eq!(
660 distance as f64,
661 no_vector_compare_f32_as_f64(&a_slice.0, &b_slice.0),
662 epsilon = 1e-4f64
663 );
664 }
665
666 #[test]
667 fn test_dist_l2_float_turing_112() {
668 let (a_data, b_data) = get_turing_test_data_f32_dim(112);
669 let (a_slice, b_slice) = (
670 F32Slice112(a_data.try_into().unwrap()),
671 F32Slice112(b_data.try_into().unwrap()),
672 );
673
674 let distance: f32 = compare_two_vec::<f32>(112, Metric::L2, &a_slice.0, &b_slice.0);
675
676 assert_abs_diff_eq!(
677 distance as f64,
678 no_vector_compare_f32_as_f64(&a_slice.0, &b_slice.0),
679 epsilon = 1e-4f64
680 );
681 }
682
683 #[test]
684 fn test_dist_l2_float_turing_128() {
685 let (a_data, b_data) = get_turing_test_data_f32_dim(128);
686 let (a_slice, b_slice) = (
687 F32Slice128(a_data.try_into().unwrap()),
688 F32Slice128(b_data.try_into().unwrap()),
689 );
690
691 let distance: f32 = compare_two_vec::<f32>(128, Metric::L2, &a_slice.0, &b_slice.0);
692
693 assert_abs_diff_eq!(
694 distance as f64,
695 no_vector_compare_f32_as_f64(&a_slice.0, &b_slice.0),
696 epsilon = 1e-4f64
697 );
698 }
699
700 #[test]
701 fn test_dist_l2_float_turing_256() {
702 let (a_data, b_data) = get_turing_test_data_f32_dim(256);
703 let (a_slice, b_slice) = (
704 F32Slice256(a_data.try_into().unwrap()),
705 F32Slice256(b_data.try_into().unwrap()),
706 );
707
708 let distance: f32 = compare_two_vec::<f32>(256, Metric::L2, &a_slice.0, &b_slice.0);
709
710 assert_abs_diff_eq!(
711 distance as f64,
712 no_vector_compare_f32_as_f64(&a_slice.0, &b_slice.0),
713 epsilon = 1e-3f64
714 );
715 }
716
717 #[test]
718 fn test_dist_l2_float_turing_4096() {
719 let (a_data, b_data) = get_turing_test_data_f32_dim(4096);
720 let (a_slice, b_slice) = (
721 F32Slice4096(a_data.try_into().unwrap()),
722 F32Slice4096(b_data.try_into().unwrap()),
723 );
724
725 let distance: f32 = compare_two_vec::<f32>(4096, Metric::L2, &a_slice.0, &b_slice.0);
726
727 assert_abs_diff_eq!(
728 distance as f64,
729 no_vector_compare_f32_as_f64(&a_slice.0, &b_slice.0),
730 epsilon = 1e-2f64
731 );
732 }
733
734 #[test]
735 fn test_dist_ip_float_turing_112() {
736 let (a_data, b_data) = get_turing_test_data_f32_dim(112);
737 let (a_slice, b_slice) = (
738 F32Slice112(a_data.try_into().unwrap()),
739 F32Slice112(b_data.try_into().unwrap()),
740 );
741
742 let distance: f32 =
743 compare_two_vec::<f32>(112, Metric::InnerProduct, &a_slice.0, &b_slice.0);
744
745 assert_abs_diff_eq!(
746 distance,
747 reference::reference_innerproduct_f32_similarity(&a_slice.0, &b_slice.0).into_inner(),
748 epsilon = 1e-4f32
749 );
750 }
751
752 #[test]
753 fn distance_test() {
754 #[repr(C, align(32))]
755 struct Vector32ByteAligned {
756 v: [f32; 512],
757 }
758
759 let two_vec = Box::new(Vector32ByteAligned {
761 v: [
762 69.02492, 78.84786, 63.125072, 90.90581, 79.2592, 70.81731, 3.0829668, 33.33287,
763 20.777142, 30.147898, 23.681915, 42.553043, 12.602162, 7.3808074, 19.157589,
764 65.6791, 76.44677, 76.89124, 86.40756, 84.70118, 87.86142, 16.126896, 5.1277637,
765 95.11038, 83.946945, 22.735607, 11.548555, 59.51482, 24.84603, 15.573776, 78.27185,
766 71.13179, 38.574017, 80.0228, 13.175261, 62.887978, 15.205181, 18.89392, 96.13162,
767 87.55455, 34.179806, 62.920044, 4.9305916, 54.349373, 21.731495, 14.982187,
768 40.262867, 20.15214, 36.61963, 72.450806, 55.565, 95.5375, 93.73356, 95.36308,
769 66.30762, 58.0397, 18.951357, 67.11702, 43.043316, 30.65622, 99.85361, 2.5889993,
770 27.844774, 39.72441, 46.463238, 71.303764, 90.45308, 36.390602, 63.344395,
771 26.427078, 35.99528, 82.35505, 32.529175, 23.165905, 74.73179, 9.856939, 59.38126,
772 35.714924, 79.81213, 46.704124, 24.47884, 36.01743, 0.46678782, 29.528152,
773 1.8980742, 24.68853, 75.58984, 98.72279, 68.62601, 11.890173, 49.49361, 55.45572,
774 72.71067, 34.107483, 51.357758, 76.400635, 81.32725, 66.45081, 17.848074,
775 62.398876, 94.20444, 2.10886, 17.416393, 64.88253, 29.000723, 62.434315, 53.907238,
776 70.51412, 78.70744, 55.181683, 64.45116, 23.419212, 53.68544, 43.506958, 46.89598,
777 35.905994, 64.51397, 91.95555, 20.322979, 74.80128, 97.548744, 58.312725, 78.81985,
778 31.911612, 14.445949, 49.85094, 70.87396, 40.06766, 7.129991, 78.48008, 75.21636,
779 93.623604, 95.95479, 29.571129, 22.721554, 26.73875, 52.075504, 56.783104,
780 94.65493, 61.778534, 85.72401, 85.369514, 29.922367, 41.410553, 94.12884,
781 80.276855, 55.604828, 54.70947, 74.07216, 44.61955, 31.38113, 68.48596, 34.56782,
782 14.424729, 48.204506, 9.675444, 32.01946, 92.32695, 36.292683, 78.31955, 98.05327,
783 14.343918, 46.017002, 95.90888, 82.63626, 16.873539, 3.698051, 7.8042626,
784 64.194405, 96.71023, 67.93692, 21.618402, 51.92182, 22.834194, 61.56986, 19.749891,
785 55.31206, 38.29552, 67.57593, 67.145836, 38.92673, 94.95708, 72.38746, 90.70901,
786 69.43995, 9.394085, 31.646872, 88.20112, 9.134722, 99.98214, 5.423498, 41.51995,
787 76.94409, 77.373276, 3.2966614, 9.611201, 57.231106, 30.747868, 76.10228, 91.98308,
788 70.893585, 0.9067178, 43.96515, 16.321218, 27.734184, 83.271835, 88.23312,
789 87.16445, 5.556643, 15.627432, 58.547127, 93.6459, 40.539192, 49.124157, 91.13276,
790 57.485855, 8.827019, 4.9690843, 46.511234, 53.91469, 97.71925, 20.135271,
791 23.353004, 70.92099, 93.38748, 87.520134, 51.684677, 29.89813, 9.110392, 65.809204,
792 34.16554, 93.398605, 84.58669, 96.409645, 9.876037, 94.767784, 99.21523, 1.9330144,
793 94.92429, 75.12728, 17.218828, 97.89164, 35.476578, 77.629456, 69.573746,
794 40.200542, 42.117836, 5.861628, 75.45282, 82.73633, 0.98086596, 77.24894,
795 11.248695, 61.070026, 52.692616, 80.5449, 80.76036, 29.270136, 67.60252, 48.782394,
796 95.18851, 83.47162, 52.068756, 46.66002, 90.12216, 15.515327, 33.694042, 96.963036,
797 73.49627, 62.805485, 44.715607, 59.98627, 3.8921833, 37.565327, 29.69184,
798 39.429665, 83.46899, 44.286453, 21.54851, 56.096413, 18.169249, 5.214751,
799 14.691341, 99.779335, 26.32643, 67.69903, 36.41243, 67.27333, 12.157213, 96.18984,
800 2.438283, 78.14289, 0.14715195, 98.769, 53.649532, 21.615898, 39.657497, 95.45616,
801 18.578386, 71.47976, 22.348118, 17.85519, 6.3717127, 62.176777, 22.033644,
802 23.178005, 79.44858, 89.70233, 37.21273, 71.86182, 21.284317, 52.908623, 30.095518,
803 63.64478, 77.55823, 80.04871, 15.133011, 30.439043, 70.16561, 4.4014096, 89.28944,
804 26.29093, 46.827854, 11.764729, 61.887516, 47.774887, 57.19503, 59.444664,
805 28.592825, 98.70386, 1.2497544, 82.28431, 46.76423, 83.746124, 53.032673, 86.53457,
806 99.42168, 90.184, 92.27852, 9.059965, 71.75723, 70.45299, 10.924053, 68.329704,
807 77.27232, 6.677854, 75.63629, 57.370533, 17.09031, 10.554659, 99.56178, 37.53221,
808 72.311104, 75.7565, 65.2042, 36.096478, 64.69502, 38.88497, 64.33723, 84.87812,
809 66.84958, 8.508932, 79.134, 83.431015, 66.72124, 61.801838, 64.30524, 37.194263,
810 77.94725, 89.705185, 23.643505, 19.505919, 48.40264, 43.01083, 21.171177,
811 18.717121, 10.805857, 69.66983, 77.85261, 57.323063, 3.28964, 38.758026, 5.349946,
812 7.46572, 57.485138, 30.822384, 33.9411, 95.53746, 65.57723, 42.1077, 28.591347,
813 11.917269, 5.031073, 31.835615, 19.34116, 85.71027, 87.4516, 1.3798475, 70.70583,
814 51.988052, 45.217144, 14.308596, 54.557167, 86.18323, 79.13666, 76.866745,
815 46.010685, 79.739235, 44.667603, 39.36416, 72.605896, 73.83187, 13.137412,
816 6.7911267, 63.952374, 10.082436, 86.00318, 99.760376, 92.84948, 63.786434,
817 3.4429908, 18.244314, 75.65299, 14.964747, 70.126366, 80.89449, 91.266655,
818 96.58798, 46.439327, 38.253975, 87.31036, 21.093178, 37.19671, 58.28973, 9.75231,
819 12.350321, 25.75115, 87.65073, 53.610504, 36.850048, 18.66356, 94.48941, 83.71898,
820 44.49315, 44.186737, 19.360733, 84.365974, 46.76272, 44.924366, 50.279808,
821 54.868866, 91.33004, 18.683397, 75.13282, 15.070831, 47.04839, 53.780903,
822 26.911152, 74.65651, 57.659935, 25.604189, 37.235474, 65.39667, 53.952206,
823 40.37131, 59.173275, 96.00756, 54.591274, 10.787476, 69.51549, 31.970142,
824 25.408005, 55.972492, 85.01888, 97.48981, 91.006134, 28.98619, 97.151276,
825 34.388496, 47.498177, 11.985874, 64.73775, 33.877014, 13.370312, 34.79146,
826 86.19321, 15.019405, 94.07832, 93.50433, 60.168625, 50.95409, 38.27827, 47.458614,
827 32.83715, 69.54998, 69.0361, 84.1418, 34.270298, 74.23852, 70.707466, 78.59845,
828 9.651399, 24.186779, 58.255756, 53.72362, 92.46477, 97.75528, 20.257462, 30.122698,
829 50.41517, 28.156603, 42.644154,
830 ],
831 });
832
833 let distance: f32 = compare::<f32>(256, Metric::L2, &two_vec.v);
834
835 assert_eq!(distance, 429141.2);
836 }
837
838 fn compare<T>(dim: usize, metric: Metric, v: &[T]) -> f32
839 where
840 T: DistanceProvider<T>,
841 {
842 let distance_comparer = T::distance_comparer(metric, Some(dim));
843 distance_comparer.call(&v[..dim], &v[dim..])
844 }
845
846 pub fn compare_two_vec<T>(dim: usize, metric: Metric, v1: &[T], v2: &[T]) -> f32
847 where
848 T: DistanceProvider<T>,
849 {
850 let distance_comparer = T::distance_comparer(metric, Some(dim));
851 distance_comparer.call(&v1[..dim], &v2[..dim])
852 }
853}
854
855#[cfg(test)]
856mod distance_provider_f16_tests {
857 use approx::assert_abs_diff_eq;
858
859 use super::{distance_provider_f32_tests::get_turing_test_data_f32_dim, *};
860 use crate::{
861 distance::distance_provider::distance_provider_f32_tests::compare_two_vec,
862 test_util::no_vector_compare_f16_as_f64,
863 };
864
865 #[repr(C, align(32))]
866 pub struct F16Slice112([f16; 112]);
867 #[repr(C, align(32))]
868 pub struct F16Slice104([f16; 104]);
869 #[repr(C, align(32))]
870 pub struct F16Slice128([f16; 128]);
871 #[repr(C, align(32))]
872 pub struct F16Slice256([f16; 256]);
873 #[repr(C, align(32))]
874 pub struct F16Slice4096([f16; 4096]);
875
876 fn get_turing_test_data_f16_dim(dim: usize) -> (Vec<f16>, Vec<f16>) {
877 let (a_slice, b_slice) = get_turing_test_data_f32_dim(dim);
878 let a_data = a_slice.iter().map(|x| f16::from_f32(*x)).collect();
879 let b_data = b_slice.iter().map(|x| f16::from_f32(*x)).collect();
880 (a_data, b_data)
881 }
882
883 #[test]
884 fn test_dist_l2_f16_turing_112() {
885 let (a_data, b_data) = get_turing_test_data_f16_dim(112);
887 let (a_slice, b_slice) = (
888 F16Slice112(a_data.try_into().unwrap()),
889 F16Slice112(b_data.try_into().unwrap()),
890 );
891
892 let distance: f32 = compare_two_vec::<f16>(112, Metric::L2, &a_slice.0, &b_slice.0);
893
894 assert_abs_diff_eq!(
896 distance as f64,
897 no_vector_compare_f16_as_f64(&a_slice.0, &b_slice.0),
898 epsilon = 1e-3f64
899 );
900 }
901
902 #[test]
903 fn test_dist_l2_f16_turing_104() {
904 let (a_data, b_data) = get_turing_test_data_f16_dim(104);
906 let (a_slice, b_slice) = (
907 F16Slice104(a_data.try_into().unwrap()),
908 F16Slice104(b_data.try_into().unwrap()),
909 );
910
911 let distance: f32 = compare_two_vec::<f16>(104, Metric::L2, &a_slice.0, &b_slice.0);
912
913 assert_abs_diff_eq!(
915 distance as f64,
916 no_vector_compare_f16_as_f64(&a_slice.0, &b_slice.0),
917 epsilon = 1e-3f64
918 );
919 }
920
921 #[test]
922 fn test_dist_l2_f16_turing_256() {
923 let (a_data, b_data) = get_turing_test_data_f16_dim(256);
925 let (a_slice, b_slice) = (
926 F16Slice256(a_data.try_into().unwrap()),
927 F16Slice256(b_data.try_into().unwrap()),
928 );
929
930 let distance: f32 = compare_two_vec::<f16>(256, Metric::L2, &a_slice.0, &b_slice.0);
931
932 assert_abs_diff_eq!(
934 distance as f64,
935 no_vector_compare_f16_as_f64(&a_slice.0, &b_slice.0),
936 epsilon = 1e-3f64
937 );
938 }
939
940 #[test]
941 fn test_dist_l2_f16_turing_128() {
942 let (a_data, b_data) = get_turing_test_data_f16_dim(128);
944 let (a_slice, b_slice) = (
945 F16Slice128(a_data.try_into().unwrap()),
946 F16Slice128(b_data.try_into().unwrap()),
947 );
948
949 let distance: f32 = compare_two_vec::<f16>(128, Metric::L2, &a_slice.0, &b_slice.0);
950
951 assert_abs_diff_eq!(
953 distance as f64,
954 no_vector_compare_f16_as_f64(&a_slice.0, &b_slice.0),
955 epsilon = 1e-3f64
956 );
957 }
958
959 #[test]
960 fn test_dist_l2_f16_turing_4096() {
961 let (a_data, b_data) = get_turing_test_data_f16_dim(4096);
963 let (a_slice, b_slice) = (
964 F16Slice4096(a_data.try_into().unwrap()),
965 F16Slice4096(b_data.try_into().unwrap()),
966 );
967
968 let distance: f32 = compare_two_vec::<f16>(4096, Metric::L2, &a_slice.0, &b_slice.0);
969
970 assert_abs_diff_eq!(
972 distance as f64,
973 no_vector_compare_f16_as_f64(&a_slice.0, &b_slice.0),
974 epsilon = 1e-2f64
975 );
976 }
977
978 #[test]
979 fn test_dist_l2_f16_produces_nan_distance_for_infinity_vectors() {
980 let a_data = vec![f16::INFINITY; 384];
981 let b_data = vec![f16::INFINITY; 384];
982
983 let distance: f32 = compare_two_vec::<f16>(384, Metric::L2, &a_data, &b_data);
984 assert!(distance.is_nan());
985 }
986}