Skip to main content

diskann_vector/distance/
distance_provider.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6#[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
19/// Return a function pointer-like [`Distance`] to compute the requested metric.
20///
21/// If `dimension` is provided, then the returned function may **only** be used on
22/// slices with length `dimension`. Calling the returned function with a different sized
23/// slice **may** panic.
24///
25/// If `dimension` is not provided, then the returned function will work for all sizes.
26///
27/// The functions returned by `distance_comparer` do not have strict alignment
28/// requirements, though aligning your data *may* yield better memory performance.
29///
30/// # Metric Semantics
31///
32/// The values computed by the returned functions may be modified from the true mathematical
33/// definition of the metric to ensure that values closer to `-infinity` imply more similar.
34///
35/// * `L2`: Computes the squared L2 distance between vectors.
36/// * `InnerProduct`: Returns the **negative** inner-product.
37/// * `Cosine`: Returns `1 - cosine-similarity` and will work on un-normalized vectors.
38/// * `CosineNormalized`: Returns `1 - cosinesimilarity` with the hint that the provided
39///   vectors have norm 1. This allows for potentially more-efficient implementations but the
40///   results may be incorrect if called with unnormalized data.
41///
42///   When provided with integer arguments (for which normalization does not make sense), this
43///   behaves as if `Cosine` was provided.
44pub 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/// A function pointer-like type for computing distances between `&[T]` and `&[U]`.
59///
60/// See: [`DistanceProvider`].
61#[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    /// Create a new [`Distance`] using the [`DistanceProvider`] implementation.
80    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    /// Compute the distance between `x` and `y`.
88    ///
89    /// The actual distance computed depends on the metric supplied to [`DistanceProvider`].
90    ///
91    /// Additionally, if a dimension were given to [`DistanceProvider`], this function may
92    /// panic if provided with slices with a length not equal to this dimension.
93    #[inline]
94    pub fn call(&self, x: &[T], y: &[U]) -> f32 {
95        self.call_unaligned(x.as_unaligned(), y.as_unaligned())
96    }
97
98    /// Compute the distance between `x` and `y`.
99    ///
100    /// The actual distance computed depends on the metric supplied to [`DistanceProvider`].
101    ///
102    /// Additionally, if a dimension were given to [`DistanceProvider`], this function may
103    /// panic if provided with slices with a length not equal to this dimension.
104    #[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
148////////////////////
149// Implementation //
150////////////////////
151
152// Implementation Notes
153//
154// Our implementation of `DistanceProvider` dispatches across:
155//
156// * Data Types
157// * Metric
158// * Dimensions
159// * Runtime Micro-architecture
160//
161// This is a combinatorial explosing of potentially compiled kernels. To get a handle on the
162// sheer number of compiled functions, we manually control the dimensional specialization on
163// a case-by-case basis.
164//
165// This is facilitated by the `specialize!` macro, which accepts a list of dimensions and
166// instantiates the necessary machinery.
167//
168// To explain the machiner a little, a [`Cons`] compile-time list is constructed. This type
169// might look like
170//
171// * `Cons<Spec<100>, Cons<Spec<64>, Spec<32>>>`: To specialize dimensions 100, 64, and 32.
172// * `Cons<Spec<100>, Null>`: To specialize just dimension 100.
173// * `Cons<Null, Null>`: To specialize no dimensions.
174//
175// The `TrySpecialize` trait is then used specialize a kernel `F` for an architecture `A`
176// with implementations
177//
178// * `Spec<N>`: Check if the requested dimension is equal to `N` and if so, return the
179//   specialized method.
180// * `Null`: Never specialize.
181// * `Cons<Head, Tail>`: Try to specialize using `Head` returning if successful. Otherwise,
182//   return the specialization of `Tail`.
183//
184//   This definition is what allows nested `Cons` structures to specialize multiple dimensions.
185//
186//   The `Cons` list also compiles a generic-dimensional fallback if none of the
187//   specializations match.
188//
189// The overall flow is
190//
191// 1. Enter the `DistanceProvider` implementation.
192//
193// 2. First dispatch across micro-architecture using `ArgumentTypes` to hold the data types.
194//    `ArgumentTypes` implements `Target2` to facilitate this dispatch.
195//
196// 3. The implementations of `Target2` for `ArgumentTypes` are performed by the `specialize!`
197//    macro, which creates a `Cons` list of requested specializations, switches across
198//    metrics and invokes `Cons:specialize` on the requested metric.
199
200macro_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                // Use the `no-features` variant because we do not care if the target gets
205                // compiled for higher micro-architecture levels.
206                //
207                // It's the returned kernel that matters.
208                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
224/////////////////////////
225// Specialization List //
226/////////////////////////
227
228macro_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 types redirect `CosineNormalized` to `Cosine`.
283    (@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
347/// Specialize a distance function `F` for the dimension `dim` if possible. Otherwise,
348/// return `None`.
349trait 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
358/// Specialize a distance function for the requested dimensionality.
359struct 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                    // NOTE: This line here is what actually compiles the specialized kernel.
373                    arch.dispatch2::<Specialize<N, F>, f32, Unaligned<T>, Unaligned<U>>(),
374                ));
375            }
376        }
377        None
378    }
379}
380
381/// Don't specialize at all.
382struct 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
395/// A recursive compile-time list for building a list of specializations.
396struct 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    /// Try to specialize `F`. If no such specialization is available, return a fallback
407    /// implementation.
408    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
425// Try `Head` and then `Tail`.
426impl<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///////////
444// Tests //
445///////////
446
447#[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    // Comparison Bounds
459    struct EpsilonAndRelative {
460        epsilon: f32,
461        max_relative: f32,
462    }
463
464    /// For now - these are rough bounds selected heuristically.
465    /// Eventually (once we have implementations using compensated arithmetic), we should
466    /// empirically derive bounds based on a combination of
467    ///
468    /// 1. Input Distribution
469    /// 2. Distance Function
470    /// 3. Dimensionality
471    ///
472    /// To ensure that these bounds are tight.
473    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            // Allow for some error when handling the normalization at the end.
497            Metric::Cosine | Metric::CosineNormalized => EpsilonAndRelative {
498                epsilon: 1e-6,
499                max_relative: 1e-6,
500            },
501            // These should be exact.
502            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        // Unwrap the SimilarityScore for the reference implementation.
522        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    /// The maximum dimension used for unaligned behavior checking with simple distances.
557    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        // two vectors are allocated in the contiguous heap memory
760        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        // two vectors are allocated in the contiguous heap memory
886        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        // Note the variance between the full 32 bit precision and the 16 bit precision
895        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        // two vectors are allocated in the contiguous heap memory
905        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        // Note the variance between the full 32 bit precision and the 16 bit precision
914        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        // two vectors are allocated in the contiguous heap memory
924        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        // Note the variance between the full 32 bit precision and the 16 bit precision
933        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        // two vectors are allocated in the contiguous heap memory
943        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        // Note the variance between the full 32 bit precision and the 16 bit precision
952        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        // two vectors are allocated in the contiguous heap memory
962        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        // Note the variance between the full 32 bit precision and the 16 bit precision
971        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}