Skip to main content

diskann_vector/distance/
simd.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6#[cfg(target_arch = "x86_64")]
7use diskann_wide::arch::x86_64::{V3, V4};
8
9#[cfg(target_arch = "aarch64")]
10use diskann_wide::arch::aarch64::{algorithms, Neon};
11
12use diskann_wide::{
13    arch::Scalar, Architecture, Const, Constant, Emulated, SIMDAbs, SIMDDotProduct, SIMDMulAdd,
14    SIMDSumTree, SIMDVector,
15};
16
17use crate::{AsUnaligned, Half};
18
19/// A helper trait to allow integer to f32 conversion (which may be lossy).
20pub trait LossyF32Conversion: Copy {
21    fn as_f32_lossy(self) -> f32;
22}
23
24impl LossyF32Conversion for f32 {
25    fn as_f32_lossy(self) -> f32 {
26        self
27    }
28}
29
30impl LossyF32Conversion for i32 {
31    fn as_f32_lossy(self) -> f32 {
32        self as f32
33    }
34}
35
36impl LossyF32Conversion for u32 {
37    fn as_f32_lossy(self) -> f32 {
38        self as f32
39    }
40}
41
42cfg_if::cfg_if! {
43    if #[cfg(miri)] {
44        fn force_eval(_x: f32) {}
45    } else if #[cfg(target_arch = "x86_64")] {
46        use std::arch::asm;
47
48        /// Force the evaluation of the argument, preventing the compiler from reordering the
49        /// computation of `x` behind a condition.
50        ///
51        /// In the context of Cosine similarity, this can help code generation for
52        /// static-dimensional kernels
53        #[inline(always)]
54        fn force_eval(x: f32) {
55            // SAFETY: This function executes no instructions. As such, it satisfies the long
56            // list of requirements for inline assembly.
57            //
58            // See: https://doc.rust-lang.org/reference/inline-assembly.html#rules-for-inline-assembly
59            unsafe {
60                asm!(
61                    // Assembly comment to "use" the argument
62                    "/* {0} */",
63                    // Use an `xmm_reg` since LLVM almost always uses such a register for
64                    // scalar floating point
65                    in(xmm_reg) x,
66                    // Explanation:
67                    // * `nostack`: This function does not touch the stack, so the compiler
68                    //   does not need to worry that the stack gets messed up.
69                    // * `nomem`: This function does not touch memory. The compiler doesn't
70                    //   have to reload any values.
71                    // * `preserves_flags`: This function preserves architectural condition
72                    //   flags. We can make this guarantee because this function literally
73                    //   does nothing.
74                    options(nostack, nomem, preserves_flags)
75                )
76            }
77        }
78    } else {
79        // Fallback implementation.
80        fn force_eval(_x: f32) {}
81    }
82}
83
84/// A utility struct to help with SIMD loading.
85///
86/// The main loop of SIMD kernels consists of various tilings of loads and arithmetic.
87/// Outside of the epilogue, these loads are all full-width vector loads.
88///
89/// To aid in defining different tilings, this struct takes the base pointers for left and
90/// right hand pointers and provides a `load` method to extract full vectors for both
91/// the left and right-hand sides.
92///
93/// This works in conjunction with the [`SIMDSchema`] to help write unrolled loops.
94#[derive(Debug, Clone, Copy)]
95pub struct Loader<Schema, Left, Right, A>
96where
97    Schema: SIMDSchema<Left, Right, A>,
98    A: Architecture,
99{
100    arch: A,
101    schema: Schema,
102    left: *const Left,
103    right: *const Right,
104    len: usize,
105}
106
107impl<Schema, Left, Right, A> Loader<Schema, Left, Right, A>
108where
109    Schema: SIMDSchema<Left, Right, A>,
110    A: Architecture,
111{
112    /// Construct a new loader for the left and right hand pointers.
113    ///
114    /// Requires that the memory ranges `[left, left + len)` and `[right, right + len)` are
115    /// both valid, where `len` is the *number* of the elements of type `T` and `U`.
116    #[inline(always)]
117    fn new(arch: A, schema: Schema, left: *const Left, right: *const Right, len: usize) -> Self {
118        Self {
119            arch,
120            schema,
121            left,
122            right,
123            len,
124        }
125    }
126
127    /// Return the underlying architecture.
128    #[inline(always)]
129    fn arch(&self) -> A {
130        self.arch
131    }
132
133    /// Return the SIMD Schema.
134    #[inline(always)]
135    fn schema(&self) -> Schema {
136        self.schema
137    }
138
139    /// Load full width vectors for the left and right hand memory spans.
140    ///
141    /// This loads a [`SIMDSchema::SIMDWidth`] chunk of data using the following formula:
142    ///
143    /// ```text
144    /// // The number of elements in an unrolled [`MainLoop`].
145    /// let simd_width = Schema::SIMDWidth::value();
146    /// let block_size = simd_width * Schema::Main::BLOCK_SIZE;
147    ///
148    /// load(px + block_size * block + simd_width * offset);
149    /// ```
150    ///
151    /// # Safety
152    ///
153    /// Requires that the following memory addresses are in-bounds (i.e., the highest
154    /// read address is at an offset less than `len`):
155    ///
156    /// ```text
157    /// [
158    ///     px + block_size * block + simd_width * offset,
159    ///     px + block_size * block + simd_width * (offset + 1)
160    /// )
161    ///
162    /// [
163    ///     py + block_size * block + simd_width * offset,
164    ///     py + block_size * block + simd_width * (offset + 1)
165    /// )
166    /// ```
167    ///
168    /// This invariant is checked in debug builds.
169    #[inline(always)]
170    unsafe fn load(&self, block: usize, offset: usize) -> (Schema::Left, Schema::Right) {
171        let stride = Schema::SIMDWidth::value();
172        let block_stride = stride * Schema::Main::BLOCK_SIZE;
173        let offset = block_stride * block + stride * offset;
174
175        debug_assert!(
176            offset + stride <= self.len,
177            "length = {}, offset = {}",
178            self.len,
179            offset
180        );
181
182        (
183            Schema::Left::load_simd(self.arch, self.left.add(offset)),
184            Schema::Right::load_simd(self.arch, self.right.add(offset)),
185        )
186    }
187}
188
189/// A representation of the main unrolled-loop for SIMD kernels.
190pub trait MainLoop {
191    /// The effective number of unrolling (in terms of SIMD vectors) performed by this
192    /// kernel. For example, if `BLOCK_SIZE = 4` and the SIMD width is 8, than each iteration
193    /// of the main loop will process `4 * 8 = 32` elements.
194    ///
195    /// This parameter will be used to compute the number of full-width epilogues that need
196    /// to be executed.
197    const BLOCK_SIZE: usize;
198
199    /// Perform the main unrolled loops of a SIMD kernel. This loop is expected to process
200    /// all elements in the range `[0, trip_count * S::get_simd_width() * Self::BLOCK_SIZE)`
201    /// and return an accumulator consisting of the result.
202    ///
203    /// # Arguments
204    ///
205    /// * `loader`: A SIMD loader to emit loads to the two source spans.
206    /// * `trip_count`: The number of blocks of size `BLOCK_SIZE` to process. A single "trip"
207    ///   will process `S::get_simd_width() * Self::BLOCK_SIZE` elements. So, computation of
208    ///   `trip_count` should be computed as:
209    ///   ```math
210    ///   let trip_count = len / (S::get_simd_width() * <_ as MainLoop>::BLOCK_SIZE);
211    ///   ```
212    /// * `epilogues`: The number of `S::get_simd_width()` vectors remaining after all the
213    ///   main blocks have been processed. This is guaranteed to be less than
214    ///   `Self::BLOCK_SIZE`.
215    ///
216    /// # Safety
217    ///
218    /// All elements in the accessed range must be valid. The memory addresses touched are
219    ///
220    /// ```text
221    /// let block_size = Self::BLOCK_SIZE;
222    /// let simd_width = S::get_simd_width();
223    /// [
224    ///     loader.left,
225    ///     loader.left + trip_count * simd_width + block_size + epilogues * simd_width
226    /// )
227    /// [
228    ///     loader.right,
229    ///     loader.right + trip_count * simd_width + block_size + epilogues * simd_width
230    /// )
231    /// ```
232    ///
233    /// The `loader` will ensure that all accesses are in-bounds in debug builds.
234    unsafe fn main<S, L, R, A>(
235        loader: &Loader<S, L, R, A>,
236        trip_count: usize,
237        epilogues: usize,
238    ) -> S::Accumulator
239    where
240        A: Architecture,
241        S: SIMDSchema<L, R, A>;
242}
243/// An inner loop implementation strategy using 1 parallel instances of the schema
244/// accumulator with a manual inner loop unroll of 1.
245pub struct Strategy1x1;
246
247/// An inner loop implementation strategy using 2 parallel instances of the schema
248/// accumulator with a manual inner loop unroll of 1.
249pub struct Strategy2x1;
250
251/// An inner loop implementation strategy using 4 parallel instances of the schema
252/// accumulator with a manual inner loop unroll of 1.
253pub struct Strategy4x1;
254
255/// An inner loop implementation strategy using 4 parallel instances of the schema
256/// accumulator with a manual inner loop unroll of 2.
257pub struct Strategy4x2;
258
259/// An inner loop implementation strategy using 2 parallel instances of the schema
260/// accumulator with a manual inner loop unroll of 4.
261pub struct Strategy2x4;
262
263impl MainLoop for Strategy1x1 {
264    const BLOCK_SIZE: usize = 1;
265
266    #[inline(always)]
267    unsafe fn main<S, L, R, A>(
268        loader: &Loader<S, L, R, A>,
269        trip_count: usize,
270        _epilogues: usize,
271    ) -> S::Accumulator
272    where
273        A: Architecture,
274        S: SIMDSchema<L, R, A>,
275    {
276        let arch = loader.arch();
277        let schema = loader.schema();
278
279        let mut s0 = schema.init(arch);
280        for i in 0..trip_count {
281            s0 = schema.accumulate_tuple(s0, loader.load(i, 0));
282        }
283
284        s0
285    }
286}
287
288impl MainLoop for Strategy2x1 {
289    const BLOCK_SIZE: usize = 2;
290
291    #[inline(always)]
292    unsafe fn main<S, L, R, A>(
293        loader: &Loader<S, L, R, A>,
294        trip_count: usize,
295        epilogues: usize,
296    ) -> S::Accumulator
297    where
298        A: Architecture,
299        S: SIMDSchema<L, R, A>,
300    {
301        let arch = loader.arch();
302        let schema = loader.schema();
303
304        let mut s0 = schema.init(arch);
305        let mut s1 = schema.init(arch);
306
307        for i in 0..trip_count {
308            s0 = schema.accumulate_tuple(s0, loader.load(i, 0));
309            s1 = schema.accumulate_tuple(s1, loader.load(i, 1));
310        }
311
312        let mut s = schema.combine(s0, s1);
313        if epilogues != 0 {
314            s = schema.accumulate_tuple(s, loader.load(trip_count, 0));
315        }
316
317        s
318    }
319}
320
321impl MainLoop for Strategy4x1 {
322    const BLOCK_SIZE: usize = 4;
323
324    #[inline(always)]
325    unsafe fn main<S, L, R, A>(
326        loader: &Loader<S, L, R, A>,
327        trip_count: usize,
328        epilogues: usize,
329    ) -> S::Accumulator
330    where
331        A: Architecture,
332        S: SIMDSchema<L, R, A>,
333    {
334        let arch = loader.arch();
335        let schema = loader.schema();
336
337        let mut s0 = schema.init(arch);
338        let mut s1 = schema.init(arch);
339        let mut s2 = schema.init(arch);
340        let mut s3 = schema.init(arch);
341
342        for i in 0..trip_count {
343            s0 = schema.accumulate_tuple(s0, loader.load(i, 0));
344            s1 = schema.accumulate_tuple(s1, loader.load(i, 1));
345            s2 = schema.accumulate_tuple(s2, loader.load(i, 2));
346            s3 = schema.accumulate_tuple(s3, loader.load(i, 3));
347        }
348
349        if epilogues >= 1 {
350            s0 = schema.accumulate_tuple(s0, loader.load(trip_count, 0));
351        }
352
353        if epilogues >= 2 {
354            s1 = schema.accumulate_tuple(s1, loader.load(trip_count, 1));
355        }
356
357        if epilogues >= 3 {
358            s2 = schema.accumulate_tuple(s2, loader.load(trip_count, 2));
359        }
360
361        schema.combine(schema.combine(s0, s1), schema.combine(s2, s3))
362    }
363}
364
365impl MainLoop for Strategy4x2 {
366    const BLOCK_SIZE: usize = 4;
367
368    #[inline(always)]
369    unsafe fn main<S, L, R, A>(
370        loader: &Loader<S, L, R, A>,
371        trip_count: usize,
372        epilogues: usize,
373    ) -> S::Accumulator
374    where
375        A: Architecture,
376        S: SIMDSchema<L, R, A>,
377    {
378        let arch = loader.arch();
379        let schema = loader.schema();
380
381        let mut s0 = schema.init(arch);
382        let mut s1 = schema.init(arch);
383        let mut s2 = schema.init(arch);
384        let mut s3 = schema.init(arch);
385
386        for i in 0..(trip_count / 2) {
387            let j = 2 * i;
388            s0 = schema.accumulate_tuple(s0, loader.load(j, 0));
389            s1 = schema.accumulate_tuple(s1, loader.load(j, 1));
390            s2 = schema.accumulate_tuple(s2, loader.load(j, 2));
391            s3 = schema.accumulate_tuple(s3, loader.load(j, 3));
392
393            s0 = schema.accumulate_tuple(s0, loader.load(j, 4));
394            s1 = schema.accumulate_tuple(s1, loader.load(j, 5));
395            s2 = schema.accumulate_tuple(s2, loader.load(j, 6));
396            s3 = schema.accumulate_tuple(s3, loader.load(j, 7));
397        }
398
399        if !trip_count.is_multiple_of(2) {
400            // Will not underflow because `trip_count` is odd.
401            let j = trip_count - 1;
402            s0 = schema.accumulate_tuple(s0, loader.load(j, 0));
403            s1 = schema.accumulate_tuple(s1, loader.load(j, 1));
404            s2 = schema.accumulate_tuple(s2, loader.load(j, 2));
405            s3 = schema.accumulate_tuple(s3, loader.load(j, 3));
406        }
407
408        if epilogues >= 1 {
409            s0 = schema.accumulate_tuple(s0, loader.load(trip_count, 0));
410        }
411
412        if epilogues >= 2 {
413            s1 = schema.accumulate_tuple(s1, loader.load(trip_count, 1));
414        }
415
416        if epilogues >= 3 {
417            s2 = schema.accumulate_tuple(s2, loader.load(trip_count, 2));
418        }
419
420        schema.combine(schema.combine(s0, s1), schema.combine(s2, s3))
421    }
422}
423
424impl MainLoop for Strategy2x4 {
425    const BLOCK_SIZE: usize = 4;
426
427    /// The implementation here has a global unroll of 4, but the unroll factor of the main
428    /// loop is actually 8.
429    ///
430    /// There is a single peeled iteration at the end that handles the last group of 4
431    /// if needed.
432    #[inline(always)]
433    unsafe fn main<S, L, R, A>(
434        loader: &Loader<S, L, R, A>,
435        trip_count: usize,
436        epilogues: usize,
437    ) -> S::Accumulator
438    where
439        A: Architecture,
440        S: SIMDSchema<L, R, A>,
441    {
442        let arch = loader.arch();
443        let schema = loader.schema();
444
445        let mut s0 = schema.init(arch);
446        let mut s1 = schema.init(arch);
447
448        for i in 0..(trip_count / 2) {
449            let j = 2 * i;
450            s0 = schema.accumulate_tuple(s0, loader.load(j, 0));
451            s1 = schema.accumulate_tuple(s1, loader.load(j, 1));
452            s0 = schema.accumulate_tuple(s0, loader.load(j, 2));
453            s1 = schema.accumulate_tuple(s1, loader.load(j, 3));
454
455            s0 = schema.accumulate_tuple(s0, loader.load(j, 4));
456            s1 = schema.accumulate_tuple(s1, loader.load(j, 5));
457            s0 = schema.accumulate_tuple(s0, loader.load(j, 6));
458            s1 = schema.accumulate_tuple(s1, loader.load(j, 7));
459        }
460
461        if !trip_count.is_multiple_of(2) {
462            let j = trip_count - 1;
463            s0 = schema.accumulate_tuple(s0, loader.load(j, 0));
464            s1 = schema.accumulate_tuple(s1, loader.load(j, 1));
465            s0 = schema.accumulate_tuple(s0, loader.load(j, 2));
466            s1 = schema.accumulate_tuple(s1, loader.load(j, 3));
467        }
468
469        if epilogues >= 1 {
470            s0 = schema.accumulate_tuple(s0, loader.load(trip_count, 0));
471        }
472
473        if epilogues >= 2 {
474            s1 = schema.accumulate_tuple(s1, loader.load(trip_count, 1));
475        }
476
477        if epilogues >= 3 {
478            s0 = schema.accumulate_tuple(s0, loader.load(trip_count, 2));
479        }
480
481        schema.combine(s0, s1)
482    }
483}
484
485/// An interface trait for SIMD operations.
486///
487/// Patterns like unrolling, pointer arithmetic, and epilogue handling are common across
488/// many different combinations of left and right hand types for distance computations.
489///
490/// This higher level handling is delegated to functions like `simd_op`, which in turn
491/// uses a `SIMDSchema` to customize the mechanics of loading and accumulation.
492pub trait SIMDSchema<T, U, A: Architecture = diskann_wide::arch::Current>: Copy {
493    /// The desired SIMD read width.
494    /// Reads from the input slice will be use this stride when accessing memory.
495    type SIMDWidth: Constant<Type = usize>;
496
497    /// The type used to represent partial accumulated values.
498    type Accumulator: std::ops::Add<Output = Self::Accumulator> + std::fmt::Debug + Copy;
499
500    /// The type used for the left-hand side.
501    type Left: SIMDVector<Arch = A, Scalar = T, ConstLanes = Self::SIMDWidth>;
502
503    /// The type used for the right-hand side.
504    type Right: SIMDVector<Arch = A, Scalar = U, ConstLanes = Self::SIMDWidth>;
505
506    /// The final return type.
507    /// This is often `f32` for complete distance functions, but need not always be.
508    type Return;
509
510    /// The implementation of the main loop.
511    type Main: MainLoop;
512
513    /// Initialize an empty (identity) accumulator.
514    fn init(&self, arch: A) -> Self::Accumulator;
515
516    /// Perform an accumulation.
517    fn accumulate(
518        &self,
519        x: Self::Left,
520        y: Self::Right,
521        acc: Self::Accumulator,
522    ) -> Self::Accumulator;
523
524    /// Combine two independent accumulators (allows for unrolling).
525    #[inline(always)]
526    fn combine(&self, x: Self::Accumulator, y: Self::Accumulator) -> Self::Accumulator {
527        x + y
528    }
529
530    /// A supplied trait for dealing with non-full-width epilogues.
531    /// Often, masked based loading will do the right thing, but for architectures like AVX2
532    /// that have limited support for masking 8 and 16-bit operations, using a scalar
533    /// fallback may just be better.
534    ///
535    /// This provides a customization point to enable a scalar fallback.
536    ///
537    /// # Safety
538    ///
539    /// * Both pointers `x` and `y` must point to memory.
540    /// * It must be safe to read `len` contiguous items of type `T` starting at `x` and
541    ///   `len` contiguous items of type `U` starting at `y`.
542    ///
543    /// The following guarantee is made:
544    ///
545    /// * No read will be emitted to memory locations at and after `x.add(len)` and
546    ///   `y.add(len)`.
547    #[inline(always)]
548    unsafe fn epilogue(
549        &self,
550        arch: A,
551        x: *const T,
552        y: *const U,
553        len: usize,
554        acc: Self::Accumulator,
555    ) -> Self::Accumulator {
556        // SAFETY: Performing this read is safe by the safety preconditions of `epilogue`.
557        // Guarentee: The load implementation must be correct.
558        let a = Self::Left::load_simd_first(arch, x, len);
559
560        // SAFETY: Performing this read is safe by the safety preconditions of `epilogue`.
561        // Guarentee: The load implementation must be correct.
562        let b = Self::Right::load_simd_first(arch, y, len);
563        self.accumulate(a, b, acc)
564    }
565
566    /// Perform a reduction on the accumulator to yield the final result.
567    ///
568    /// This will be called at the end of distance processing.
569    fn reduce(&self, x: Self::Accumulator) -> Self::Return;
570
571    /// !! Do not extend this function !!
572    ///
573    /// Due to limitations on how associated constants can be used, we need a function
574    /// to access the SIMD width and rely on the compiler to constant propagate the result.
575    #[inline(always)]
576    fn get_simd_width() -> usize {
577        Self::SIMDWidth::value()
578    }
579
580    /// !! Do not extend this function !!
581    ///
582    /// Due to limitations on how associated constants can be used, we need a function
583    /// to access the unroll factor of the main loop and rely on the compiler to constant
584    /// propagate the result.
585    #[inline(always)]
586    fn get_main_bocksize() -> usize {
587        Self::Main::BLOCK_SIZE
588    }
589
590    /// A helper method to access [`Self::accumulate`] in a way that is immediately
591    /// compatible with [`Loader::load`].
592    #[doc(hidden)]
593    #[inline(always)]
594    fn accumulate_tuple(
595        &self,
596        acc: Self::Accumulator,
597        (x, y): (Self::Left, Self::Right),
598    ) -> Self::Accumulator {
599        self.accumulate(x, y, acc)
600    }
601}
602
603/// In some contexts - it can be beneficial to begin a computation on one pair of slices and
604/// then store intermediate state for resumption on another pair of slices.
605///
606/// A good example of this is direct-computation of PQ distances where different chunks need
607/// to be gathered and partially accumulated before the final reduction.
608///
609/// The `ResumableSchema` provides a relatively straight-forward way of achieving this.
610pub trait ResumableSIMDSchema<T, U, A = diskann_wide::arch::Current>: Copy
611where
612    A: Architecture,
613{
614    // The associated type for this function that is non-reentrant.
615    type NonResumable: SIMDSchema<T, U, A> + Default;
616    type FinalReturn;
617
618    fn init(arch: A) -> Self;
619    fn combine_with(&self, other: <Self::NonResumable as SIMDSchema<T, U, A>>::Accumulator)
620        -> Self;
621    fn sum(&self) -> Self::FinalReturn;
622}
623
624#[derive(Debug, Clone, Copy)]
625pub struct Resumable<T>(T);
626
627impl<T> Resumable<T> {
628    pub fn new(val: T) -> Self {
629        Self(val)
630    }
631
632    pub fn consume(self) -> T {
633        self.0
634    }
635}
636
637impl<T, U, R, A> SIMDSchema<T, U, A> for Resumable<R>
638where
639    A: Architecture,
640    R: ResumableSIMDSchema<T, U, A>,
641{
642    type SIMDWidth = <R::NonResumable as SIMDSchema<T, U, A>>::SIMDWidth;
643    type Accumulator = <R::NonResumable as SIMDSchema<T, U, A>>::Accumulator;
644    type Left = <R::NonResumable as SIMDSchema<T, U, A>>::Left;
645    type Right = <R::NonResumable as SIMDSchema<T, U, A>>::Right;
646    type Return = Self;
647    type Main = <R::NonResumable as SIMDSchema<T, U, A>>::Main;
648
649    fn init(&self, arch: A) -> Self::Accumulator {
650        R::NonResumable::default().init(arch)
651    }
652
653    fn accumulate(
654        &self,
655        x: Self::Left,
656        y: Self::Right,
657        acc: Self::Accumulator,
658    ) -> Self::Accumulator {
659        R::NonResumable::default().accumulate(x, y, acc)
660    }
661
662    fn combine(&self, x: Self::Accumulator, y: Self::Accumulator) -> Self::Accumulator {
663        R::NonResumable::default().combine(x, y)
664    }
665
666    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
667        Self(self.0.combine_with(x))
668    }
669}
670
671#[inline(never)]
672#[allow(clippy::panic)]
673fn emit_length_error(xlen: usize, ylen: usize) -> ! {
674    panic!(
675        "lengths must be equal, instead got: xlen = {}, ylen = {}",
676        xlen, ylen
677    )
678}
679
680/// A SIMD executor for binary ops using the provided `SIMDSchema`.
681///
682/// # Panics
683///
684/// Panics if `x.len() != y.len()`.
685#[inline(always)]
686pub fn simd_op<L, R, S, T, U, A>(schema: &S, arch: A, x: T, y: U) -> S::Return
687where
688    A: Architecture,
689    T: AsUnaligned<Element = L>,
690    U: AsUnaligned<Element = R>,
691    S: SIMDSchema<L, R, A>,
692{
693    let x = x.as_unaligned();
694    let y = y.as_unaligned();
695
696    let len = x.len();
697
698    // The two lengths of the vectors must be the same.
699    // Eventually - it will probably be worth looking into various wrapper functions for
700    // `simd_op` that perform this checking, but for now, consider providing two
701    // different-length slices as a hard program bug.
702    //
703    // N.B.: Redirect through `emit_length_error` to keep code generation as clean as
704    // possible.
705    if len != y.len() {
706        emit_length_error(len, y.len());
707    }
708    let px = x.as_ptr();
709    let py = y.as_ptr();
710
711    // N.B.: Due to limitations in Rust's handling of const generics (and outer type
712    // parameters), we cannot just reach into `S` and pull out the constant SIMDWidth.
713    //
714    // Instead, we need to go through a helper function. Since associated functions cannot
715    // be marked as `const`, we cannot require that the extracted width is evaluated at
716    // compile time.
717    //
718    // HOWEVER, compilers are very good at optimizing these kinds of patterns and
719    // recognizing that this value is indeed constant and optimizing accordingly.
720    let simd_width: usize = S::get_simd_width();
721    let unroll: usize = S::get_main_bocksize();
722
723    let trip_count = len / (simd_width * unroll);
724    let epilogues = (len - simd_width * unroll * trip_count) / simd_width;
725
726    // Create a loader that (in debug mode) will check that all of our full-width accesses
727    // are in-bounds.
728    let loader: Loader<S, L, R, A> = Loader::new(arch, *schema, px, py, len);
729
730    // SAFETY: The value of `trip_count`  and `epilogues` so
731    // `[0, trip_count * simd_width * unroll + epilogues * simd_width)` is in-bounds,
732    // satifying the requirements of `main`.
733    let mut s0 = unsafe { <S as SIMDSchema<L, R, A>>::Main::main(&loader, trip_count, epilogues) };
734
735    let remainder = len % simd_width;
736    if remainder != 0 {
737        let i = len - remainder;
738
739        // SAFETY: We have ensured that the lengths of the two inputs are the same.
740        //
741        // Furthermore, preceding computations on the induction variable mean that the
742        // remaining memory must be valid.
743        s0 = unsafe { schema.epilogue(arch, px.add(i), py.add(i), remainder, s0) };
744    }
745
746    schema.reduce(s0)
747}
748
749//----------//
750// Epilogue //
751//----------//
752
753#[cfg(target_arch = "aarch64")]
754#[inline(always)]
755unsafe fn scalar_epilogue<L, R, F, Acc>(
756    left: *const L,
757    right: *const R,
758    len: usize,
759    mut acc: Acc,
760    mut f: F,
761) -> Acc
762where
763    L: Copy,
764    R: Copy,
765    F: FnMut(Acc, L, R) -> Acc,
766{
767    for i in 0..len {
768        // SAFETY: The range `[x, x.add(len))` is valid for reads.
769        let left = unsafe { left.add(i).read_unaligned() };
770        // SAFETY: The range `[y, y.add(len))` is valid for reads.
771        let right = unsafe { right.add(i).read_unaligned() };
772        acc = f(acc, left, right);
773    }
774    acc
775}
776
777/////
778///// L2 Implementations
779/////
780
781// A pure L2 distance function that provides a final reduction.
782#[derive(Debug, Default, Clone, Copy)]
783pub struct L2;
784
785#[cfg(target_arch = "x86_64")]
786impl SIMDSchema<f32, f32, V4> for L2 {
787    type SIMDWidth = Const<8>;
788    type Accumulator = <V4 as Architecture>::f32x8;
789    type Left = <V4 as Architecture>::f32x8;
790    type Right = <V4 as Architecture>::f32x8;
791    type Return = f32;
792    type Main = Strategy4x1;
793
794    #[inline(always)]
795    fn init(&self, arch: V4) -> Self::Accumulator {
796        Self::Accumulator::default(arch)
797    }
798
799    #[inline(always)]
800    fn accumulate(
801        &self,
802        x: Self::Left,
803        y: Self::Right,
804        acc: Self::Accumulator,
805    ) -> Self::Accumulator {
806        let c = x - y;
807        c.mul_add_simd(c, acc)
808    }
809
810    #[inline(always)]
811    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
812        x.sum_tree()
813    }
814}
815
816#[cfg(target_arch = "x86_64")]
817impl SIMDSchema<f32, f32, V3> for L2 {
818    type SIMDWidth = Const<8>;
819    type Accumulator = <V3 as Architecture>::f32x8;
820    type Left = <V3 as Architecture>::f32x8;
821    type Right = <V3 as Architecture>::f32x8;
822    type Return = f32;
823    type Main = Strategy4x1;
824
825    #[inline(always)]
826    fn init(&self, arch: V3) -> Self::Accumulator {
827        Self::Accumulator::default(arch)
828    }
829
830    #[inline(always)]
831    fn accumulate(
832        &self,
833        x: Self::Left,
834        y: Self::Right,
835        acc: Self::Accumulator,
836    ) -> Self::Accumulator {
837        let c = x - y;
838        c.mul_add_simd(c, acc)
839    }
840
841    #[inline(always)]
842    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
843        x.sum_tree()
844    }
845}
846
847#[cfg(target_arch = "aarch64")]
848impl SIMDSchema<f32, f32, Neon> for L2 {
849    type SIMDWidth = Const<4>;
850    type Accumulator = <Neon as Architecture>::f32x4;
851    type Left = <Neon as Architecture>::f32x4;
852    type Right = <Neon as Architecture>::f32x4;
853    type Return = f32;
854    type Main = Strategy4x1;
855
856    #[inline(always)]
857    fn init(&self, arch: Neon) -> Self::Accumulator {
858        Self::Accumulator::default(arch)
859    }
860
861    #[inline(always)]
862    fn accumulate(
863        &self,
864        x: Self::Left,
865        y: Self::Right,
866        acc: Self::Accumulator,
867    ) -> Self::Accumulator {
868        let c = x - y;
869        c.mul_add_simd(c, acc)
870    }
871
872    #[inline(always)]
873    unsafe fn epilogue(
874        &self,
875        arch: Neon,
876        x: *const f32,
877        y: *const f32,
878        len: usize,
879        acc: Self::Accumulator,
880    ) -> Self::Accumulator {
881        let scalar = scalar_epilogue(
882            x,
883            y,
884            len.min(Self::SIMDWidth::value() - 1),
885            0.0f32,
886            |acc, x, y| -> f32 {
887                let c = x - y;
888                c.mul_add(c, acc)
889            },
890        );
891        acc + Self::Accumulator::from_array(arch, [scalar, 0.0, 0.0, 0.0])
892    }
893
894    #[inline(always)]
895    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
896        x.sum_tree()
897    }
898}
899
900impl SIMDSchema<f32, f32, Scalar> for L2 {
901    type SIMDWidth = Const<4>;
902    type Accumulator = Emulated<f32, 4>;
903    type Left = Emulated<f32, 4>;
904    type Right = Emulated<f32, 4>;
905    type Return = f32;
906    type Main = Strategy2x1;
907
908    #[inline(always)]
909    fn init(&self, arch: Scalar) -> Self::Accumulator {
910        Self::Accumulator::default(arch)
911    }
912
913    #[inline(always)]
914    fn accumulate(
915        &self,
916        x: Self::Left,
917        y: Self::Right,
918        acc: Self::Accumulator,
919    ) -> Self::Accumulator {
920        // Don't assume the presence of FMA.
921        let c = x - y;
922        (c * c) + acc
923    }
924
925    #[inline(always)]
926    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
927        x.sum_tree()
928    }
929
930    #[inline(always)]
931    unsafe fn epilogue(
932        &self,
933        arch: Scalar,
934        x: *const f32,
935        y: *const f32,
936        len: usize,
937        acc: Self::Accumulator,
938    ) -> Self::Accumulator {
939        let mut s: f32 = 0.0;
940        for i in 0..len {
941            // SAFETY: The range `[x, x.add(len))` is valid for reads.
942            let vx = unsafe { x.add(i).read_unaligned() };
943            // SAFETY: The range `[y, y.add(len))` is valid for reads.
944            let vy = unsafe { y.add(i).read_unaligned() };
945            let d = vx - vy;
946            s += d * d;
947        }
948        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
949    }
950}
951
952#[cfg(target_arch = "x86_64")]
953impl SIMDSchema<Half, Half, V4> for L2 {
954    type SIMDWidth = Const<8>;
955    type Accumulator = <V4 as Architecture>::f32x8;
956    type Left = <V4 as Architecture>::f16x8;
957    type Right = <V4 as Architecture>::f16x8;
958    type Return = f32;
959    type Main = Strategy2x4;
960
961    #[inline(always)]
962    fn init(&self, arch: V4) -> Self::Accumulator {
963        Self::Accumulator::default(arch)
964    }
965
966    #[inline(always)]
967    fn accumulate(
968        &self,
969        x: Self::Left,
970        y: Self::Right,
971        acc: Self::Accumulator,
972    ) -> Self::Accumulator {
973        diskann_wide::alias!(f32s = <V4>::f32x8);
974
975        let x: f32s = x.into();
976        let y: f32s = y.into();
977
978        let c = x - y;
979        c.mul_add_simd(c, acc)
980    }
981
982    #[inline(always)]
983    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
984        x.sum_tree()
985    }
986}
987
988#[cfg(target_arch = "x86_64")]
989impl SIMDSchema<Half, Half, V3> for L2 {
990    type SIMDWidth = Const<8>;
991    type Accumulator = <V3 as Architecture>::f32x8;
992    type Left = <V3 as Architecture>::f16x8;
993    type Right = <V3 as Architecture>::f16x8;
994    type Return = f32;
995    type Main = Strategy2x4;
996
997    #[inline(always)]
998    fn init(&self, arch: V3) -> Self::Accumulator {
999        Self::Accumulator::default(arch)
1000    }
1001
1002    #[inline(always)]
1003    fn accumulate(
1004        &self,
1005        x: Self::Left,
1006        y: Self::Right,
1007        acc: Self::Accumulator,
1008    ) -> Self::Accumulator {
1009        diskann_wide::alias!(f32s = <V3>::f32x8);
1010
1011        let x: f32s = x.into();
1012        let y: f32s = y.into();
1013
1014        let c = x - y;
1015        c.mul_add_simd(c, acc)
1016    }
1017
1018    // Perform a final reduction.
1019    #[inline(always)]
1020    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1021        x.sum_tree()
1022    }
1023}
1024
1025#[cfg(target_arch = "aarch64")]
1026impl SIMDSchema<Half, Half, Neon> for L2 {
1027    type SIMDWidth = Const<4>;
1028    type Accumulator = <Neon as Architecture>::f32x4;
1029    type Left = diskann_wide::arch::aarch64::f16x4;
1030    type Right = diskann_wide::arch::aarch64::f16x4;
1031    type Return = f32;
1032    type Main = Strategy4x1;
1033
1034    #[inline(always)]
1035    fn init(&self, arch: Neon) -> Self::Accumulator {
1036        Self::Accumulator::default(arch)
1037    }
1038
1039    #[inline(always)]
1040    fn accumulate(
1041        &self,
1042        x: Self::Left,
1043        y: Self::Right,
1044        acc: Self::Accumulator,
1045    ) -> Self::Accumulator {
1046        diskann_wide::alias!(f32s = <Neon>::f32x4);
1047
1048        let x: f32s = x.into();
1049        let y: f32s = y.into();
1050
1051        let c = x - y;
1052        c.mul_add_simd(c, acc)
1053    }
1054
1055    #[inline(always)]
1056    unsafe fn epilogue(
1057        &self,
1058        arch: Neon,
1059        x: *const Half,
1060        y: *const Half,
1061        len: usize,
1062        acc: Self::Accumulator,
1063    ) -> Self::Accumulator {
1064        diskann_wide::alias!(f32s = <Neon>::f32x4);
1065
1066        let rest = scalar_epilogue(
1067            x,
1068            y,
1069            len.min(Self::SIMDWidth::value() - 1),
1070            f32s::default(arch),
1071            |acc, x: Half, y: Half| -> f32s {
1072                let zero = Half::default();
1073                let x: f32s = Self::Left::from_array(arch, [x, zero, zero, zero]).into();
1074                let y: f32s = Self::Right::from_array(arch, [y, zero, zero, zero]).into();
1075                let c: f32s = x - y;
1076                c.mul_add_simd(c, acc)
1077            },
1078        );
1079        acc + rest
1080    }
1081
1082    #[inline(always)]
1083    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1084        x.sum_tree()
1085    }
1086}
1087
1088impl SIMDSchema<Half, Half, Scalar> for L2 {
1089    type SIMDWidth = Const<1>;
1090    type Accumulator = Emulated<f32, 1>;
1091    type Left = Emulated<Half, 1>;
1092    type Right = Emulated<Half, 1>;
1093    type Return = f32;
1094    type Main = Strategy1x1;
1095
1096    #[inline(always)]
1097    fn init(&self, arch: Scalar) -> Self::Accumulator {
1098        Self::Accumulator::default(arch)
1099    }
1100
1101    #[inline(always)]
1102    fn accumulate(
1103        &self,
1104        x: Self::Left,
1105        y: Self::Right,
1106        acc: Self::Accumulator,
1107    ) -> Self::Accumulator {
1108        let x: Self::Accumulator = x.into();
1109        let y: Self::Accumulator = y.into();
1110
1111        let c = x - y;
1112        acc + (c * c)
1113    }
1114
1115    #[inline(always)]
1116    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1117        x.to_array()[0]
1118    }
1119}
1120
1121impl<A> SIMDSchema<f32, Half, A> for L2
1122where
1123    A: Architecture,
1124{
1125    type SIMDWidth = Const<8>;
1126    type Accumulator = A::f32x8;
1127    type Left = A::f32x8;
1128    type Right = A::f16x8;
1129    type Return = f32;
1130    type Main = Strategy4x2;
1131
1132    #[inline(always)]
1133    fn init(&self, arch: A) -> Self::Accumulator {
1134        Self::Accumulator::default(arch)
1135    }
1136
1137    #[inline(always)]
1138    fn accumulate(
1139        &self,
1140        x: Self::Left,
1141        y: Self::Right,
1142        acc: Self::Accumulator,
1143    ) -> Self::Accumulator {
1144        let y: A::f32x8 = y.into();
1145        let c = x - y;
1146        c.mul_add_simd(c, acc)
1147    }
1148
1149    // Perform a final reduction.
1150    #[inline(always)]
1151    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1152        x.sum_tree()
1153    }
1154}
1155
1156#[cfg(target_arch = "x86_64")]
1157impl SIMDSchema<i8, i8, V4> for L2 {
1158    type SIMDWidth = Const<32>;
1159    type Accumulator = <V4 as Architecture>::i32x16;
1160    type Left = <V4 as Architecture>::i8x32;
1161    type Right = <V4 as Architecture>::i8x32;
1162    type Return = f32;
1163    type Main = Strategy4x1;
1164
1165    #[inline(always)]
1166    fn init(&self, arch: V4) -> Self::Accumulator {
1167        Self::Accumulator::default(arch)
1168    }
1169
1170    #[inline(always)]
1171    fn accumulate(
1172        &self,
1173        x: Self::Left,
1174        y: Self::Right,
1175        acc: Self::Accumulator,
1176    ) -> Self::Accumulator {
1177        diskann_wide::alias!(i16s = <V4>::i16x32);
1178
1179        let x: i16s = x.into();
1180        let y: i16s = y.into();
1181        let c = x - y;
1182        acc.dot_simd(c, c)
1183    }
1184
1185    #[inline(always)]
1186    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1187        x.sum_tree().as_f32_lossy()
1188    }
1189}
1190
1191#[cfg(target_arch = "x86_64")]
1192impl SIMDSchema<i8, i8, V3> for L2 {
1193    type SIMDWidth = Const<16>;
1194    type Accumulator = <V3 as Architecture>::i32x8;
1195    type Left = <V3 as Architecture>::i8x16;
1196    type Right = <V3 as Architecture>::i8x16;
1197    type Return = f32;
1198    type Main = Strategy4x1;
1199
1200    #[inline(always)]
1201    fn init(&self, arch: V3) -> Self::Accumulator {
1202        Self::Accumulator::default(arch)
1203    }
1204
1205    #[inline(always)]
1206    fn accumulate(
1207        &self,
1208        x: Self::Left,
1209        y: Self::Right,
1210        acc: Self::Accumulator,
1211    ) -> Self::Accumulator {
1212        diskann_wide::alias!(i16s = <V3>::i16x16);
1213
1214        let x: i16s = x.into();
1215        let y: i16s = y.into();
1216        let c = x - y;
1217        acc.dot_simd(c, c)
1218    }
1219
1220    // Perform a final reduction.
1221    #[inline(always)]
1222    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1223        x.sum_tree().as_f32_lossy()
1224    }
1225}
1226
1227#[cfg(target_arch = "aarch64")]
1228impl SIMDSchema<i8, i8, Neon> for L2 {
1229    type SIMDWidth = Const<16>;
1230    type Accumulator = <Neon as Architecture>::i32x8;
1231    type Left = diskann_wide::arch::aarch64::i8x16;
1232    type Right = diskann_wide::arch::aarch64::i8x16;
1233    type Return = f32;
1234    type Main = Strategy2x1;
1235
1236    #[inline(always)]
1237    fn init(&self, arch: Neon) -> Self::Accumulator {
1238        Self::Accumulator::default(arch)
1239    }
1240
1241    #[inline(always)]
1242    fn accumulate(
1243        &self,
1244        x: Self::Left,
1245        y: Self::Right,
1246        acc: Self::Accumulator,
1247    ) -> Self::Accumulator {
1248        algorithms::squared_euclidean_accum_i8x16(x, y, acc)
1249    }
1250
1251    #[inline(always)]
1252    unsafe fn epilogue(
1253        &self,
1254        arch: Neon,
1255        x: *const i8,
1256        y: *const i8,
1257        len: usize,
1258        acc: Self::Accumulator,
1259    ) -> Self::Accumulator {
1260        let scalar = scalar_epilogue(
1261            x,
1262            y,
1263            len.min(Self::SIMDWidth::value() - 1),
1264            0i32,
1265            |acc, x: i8, y: i8| -> i32 {
1266                let c = (x as i32) - (y as i32);
1267                acc + c * c
1268            },
1269        );
1270        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0, 0, 0, 0, 0])
1271    }
1272
1273    // Perform a final reduction.
1274    #[inline(always)]
1275    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1276        x.sum_tree().as_f32_lossy()
1277    }
1278}
1279
1280impl SIMDSchema<i8, i8, Scalar> for L2 {
1281    type SIMDWidth = Const<4>;
1282    type Accumulator = Emulated<i32, 4>;
1283    type Left = Emulated<i8, 4>;
1284    type Right = Emulated<i8, 4>;
1285    type Return = f32;
1286    type Main = Strategy1x1;
1287
1288    #[inline(always)]
1289    fn init(&self, arch: Scalar) -> Self::Accumulator {
1290        Self::Accumulator::default(arch)
1291    }
1292
1293    #[inline(always)]
1294    fn accumulate(
1295        &self,
1296        x: Self::Left,
1297        y: Self::Right,
1298        acc: Self::Accumulator,
1299    ) -> Self::Accumulator {
1300        let x: Self::Accumulator = x.into();
1301        let y: Self::Accumulator = y.into();
1302        let c = x - y;
1303        acc + c * c
1304    }
1305
1306    // Perform a final reduction.
1307    #[inline(always)]
1308    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1309        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
1310    }
1311
1312    #[inline(always)]
1313    unsafe fn epilogue(
1314        &self,
1315        arch: Scalar,
1316        x: *const i8,
1317        y: *const i8,
1318        len: usize,
1319        acc: Self::Accumulator,
1320    ) -> Self::Accumulator {
1321        let mut s: i32 = 0;
1322        for i in 0..len {
1323            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1324            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1325            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1326            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
1327            let d = vx - vy;
1328            s += d * d;
1329        }
1330        acc + Self::Accumulator::from_array(arch, [s, 0, 0, 0])
1331    }
1332}
1333
1334#[cfg(target_arch = "x86_64")]
1335impl SIMDSchema<u8, u8, V4> for L2 {
1336    type SIMDWidth = Const<32>;
1337    type Accumulator = <V4 as Architecture>::i32x16;
1338    type Left = <V4 as Architecture>::u8x32;
1339    type Right = <V4 as Architecture>::u8x32;
1340    type Return = f32;
1341    type Main = Strategy4x1;
1342
1343    #[inline(always)]
1344    fn init(&self, arch: V4) -> Self::Accumulator {
1345        Self::Accumulator::default(arch)
1346    }
1347
1348    #[inline(always)]
1349    fn accumulate(
1350        &self,
1351        x: Self::Left,
1352        y: Self::Right,
1353        acc: Self::Accumulator,
1354    ) -> Self::Accumulator {
1355        diskann_wide::alias!(i16s = <V4>::i16x32);
1356
1357        let x: i16s = x.into();
1358        let y: i16s = y.into();
1359        let c = x - y;
1360        acc.dot_simd(c, c)
1361    }
1362
1363    #[inline(always)]
1364    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1365        x.sum_tree().as_f32_lossy()
1366    }
1367}
1368
1369#[cfg(target_arch = "x86_64")]
1370impl SIMDSchema<u8, u8, V3> for L2 {
1371    type SIMDWidth = Const<16>;
1372    type Accumulator = <V3 as Architecture>::i32x8;
1373    type Left = <V3 as Architecture>::u8x16;
1374    type Right = <V3 as Architecture>::u8x16;
1375    type Return = f32;
1376    type Main = Strategy4x1;
1377
1378    #[inline(always)]
1379    fn init(&self, arch: V3) -> Self::Accumulator {
1380        Self::Accumulator::default(arch)
1381    }
1382
1383    #[inline(always)]
1384    fn accumulate(
1385        &self,
1386        x: Self::Left,
1387        y: Self::Right,
1388        acc: Self::Accumulator,
1389    ) -> Self::Accumulator {
1390        diskann_wide::alias!(i16s = <V3>::i16x16);
1391
1392        let x: i16s = x.into();
1393        let y: i16s = y.into();
1394        let c = x - y;
1395        acc.dot_simd(c, c)
1396    }
1397
1398    // Perform a final reduction.
1399    #[inline(always)]
1400    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1401        x.sum_tree().as_f32_lossy()
1402    }
1403}
1404
1405#[cfg(target_arch = "aarch64")]
1406impl SIMDSchema<u8, u8, Neon> for L2 {
1407    type SIMDWidth = Const<16>;
1408    type Accumulator = <Neon as Architecture>::u32x8;
1409    type Left = diskann_wide::arch::aarch64::u8x16;
1410    type Right = diskann_wide::arch::aarch64::u8x16;
1411    type Return = f32;
1412    type Main = Strategy2x1;
1413
1414    #[inline(always)]
1415    fn init(&self, arch: Neon) -> Self::Accumulator {
1416        Self::Accumulator::default(arch)
1417    }
1418
1419    #[inline(always)]
1420    fn accumulate(
1421        &self,
1422        x: Self::Left,
1423        y: Self::Right,
1424        acc: Self::Accumulator,
1425    ) -> Self::Accumulator {
1426        algorithms::squared_euclidean_accum_u8x16(x, y, acc)
1427    }
1428
1429    #[inline(always)]
1430    unsafe fn epilogue(
1431        &self,
1432        arch: Neon,
1433        x: *const u8,
1434        y: *const u8,
1435        len: usize,
1436        acc: Self::Accumulator,
1437    ) -> Self::Accumulator {
1438        let scalar = scalar_epilogue(
1439            x,
1440            y,
1441            len.min(Self::SIMDWidth::value() - 1),
1442            0u32,
1443            |acc, x: u8, y: u8| -> u32 {
1444                let c = (x as i32) - (y as i32);
1445                acc + ((c * c) as u32)
1446            },
1447        );
1448        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0, 0, 0, 0, 0])
1449    }
1450
1451    // Perform a final reduction.
1452    #[inline(always)]
1453    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1454        x.sum_tree().as_f32_lossy()
1455    }
1456}
1457
1458impl SIMDSchema<u8, u8, Scalar> for L2 {
1459    type SIMDWidth = Const<4>;
1460    type Accumulator = Emulated<i32, 4>;
1461    type Left = Emulated<u8, 4>;
1462    type Right = Emulated<u8, 4>;
1463    type Return = f32;
1464    type Main = Strategy1x1;
1465
1466    #[inline(always)]
1467    fn init(&self, arch: Scalar) -> Self::Accumulator {
1468        Self::Accumulator::default(arch)
1469    }
1470
1471    #[inline(always)]
1472    fn accumulate(
1473        &self,
1474        x: Self::Left,
1475        y: Self::Right,
1476        acc: Self::Accumulator,
1477    ) -> Self::Accumulator {
1478        let x: Self::Accumulator = x.into();
1479        let y: Self::Accumulator = y.into();
1480        let c = x - y;
1481        acc + c * c
1482    }
1483
1484    // Perform a final reduction.
1485    #[inline(always)]
1486    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1487        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
1488    }
1489
1490    #[inline(always)]
1491    unsafe fn epilogue(
1492        &self,
1493        arch: Scalar,
1494        x: *const u8,
1495        y: *const u8,
1496        len: usize,
1497        acc: Self::Accumulator,
1498    ) -> Self::Accumulator {
1499        let mut s: i32 = 0;
1500        for i in 0..len {
1501            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1502            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1503            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1504            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
1505            let d = vx - vy;
1506            s += d * d;
1507        }
1508        acc + Self::Accumulator::from_array(arch, [s, 0, 0, 0])
1509    }
1510}
1511
1512// A L2 distance function that defers a final reduction, allowing for a distance
1513// computation to take place across multiple slice pairs.
1514#[derive(Clone, Copy, Debug)]
1515pub struct ResumableL2<A = diskann_wide::arch::Current>
1516where
1517    A: Architecture,
1518    L2: SIMDSchema<f32, f32, A>,
1519{
1520    acc: <L2 as SIMDSchema<f32, f32, A>>::Accumulator,
1521}
1522
1523impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableL2<A>
1524where
1525    A: Architecture,
1526    L2: SIMDSchema<f32, f32, A, Return = f32>,
1527{
1528    type NonResumable = L2;
1529    type FinalReturn = f32;
1530
1531    #[inline(always)]
1532    fn init(arch: A) -> Self {
1533        Self { acc: L2.init(arch) }
1534    }
1535
1536    #[inline(always)]
1537    fn combine_with(&self, other: <L2 as SIMDSchema<f32, f32, A>>::Accumulator) -> Self {
1538        Self {
1539            acc: self.acc + other,
1540        }
1541    }
1542
1543    #[inline(always)]
1544    fn sum(&self) -> f32 {
1545        L2.reduce(self.acc)
1546    }
1547}
1548
1549/////
1550///// IP Implementations
1551/////
1552
1553// A pure IP distance function that provides a final reduction.
1554#[derive(Clone, Copy, Debug, Default)]
1555pub struct IP;
1556
1557#[cfg(target_arch = "x86_64")]
1558impl SIMDSchema<f32, f32, V4> for IP {
1559    type SIMDWidth = Const<8>;
1560    type Accumulator = <V4 as Architecture>::f32x8;
1561    type Left = <V4 as Architecture>::f32x8;
1562    type Right = <V4 as Architecture>::f32x8;
1563    type Return = f32;
1564    type Main = Strategy4x1;
1565
1566    #[inline(always)]
1567    fn init(&self, arch: V4) -> Self::Accumulator {
1568        Self::Accumulator::default(arch)
1569    }
1570
1571    #[inline(always)]
1572    fn accumulate(
1573        &self,
1574        x: Self::Left,
1575        y: Self::Right,
1576        acc: Self::Accumulator,
1577    ) -> Self::Accumulator {
1578        x.mul_add_simd(y, acc)
1579    }
1580
1581    #[inline(always)]
1582    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1583        x.sum_tree()
1584    }
1585}
1586
1587#[cfg(target_arch = "x86_64")]
1588impl SIMDSchema<f32, f32, V3> for IP {
1589    type SIMDWidth = Const<8>;
1590    type Accumulator = <V3 as Architecture>::f32x8;
1591    type Left = <V3 as Architecture>::f32x8;
1592    type Right = <V3 as Architecture>::f32x8;
1593    type Return = f32;
1594    type Main = Strategy4x1;
1595
1596    #[inline(always)]
1597    fn init(&self, arch: V3) -> Self::Accumulator {
1598        Self::Accumulator::default(arch)
1599    }
1600
1601    #[inline(always)]
1602    fn accumulate(
1603        &self,
1604        x: Self::Left,
1605        y: Self::Right,
1606        acc: Self::Accumulator,
1607    ) -> Self::Accumulator {
1608        x.mul_add_simd(y, acc)
1609    }
1610
1611    // Perform a final reduction.
1612    #[inline(always)]
1613    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1614        x.sum_tree()
1615    }
1616}
1617
1618#[cfg(target_arch = "aarch64")]
1619impl SIMDSchema<f32, f32, Neon> for IP {
1620    type SIMDWidth = Const<4>;
1621    type Accumulator = <Neon as Architecture>::f32x4;
1622    type Left = <Neon as Architecture>::f32x4;
1623    type Right = <Neon as Architecture>::f32x4;
1624    type Return = f32;
1625    type Main = Strategy4x1;
1626
1627    #[inline(always)]
1628    fn init(&self, arch: Neon) -> Self::Accumulator {
1629        Self::Accumulator::default(arch)
1630    }
1631
1632    #[inline(always)]
1633    fn accumulate(
1634        &self,
1635        x: Self::Left,
1636        y: Self::Right,
1637        acc: Self::Accumulator,
1638    ) -> Self::Accumulator {
1639        x.mul_add_simd(y, acc)
1640    }
1641
1642    #[inline(always)]
1643    unsafe fn epilogue(
1644        &self,
1645        arch: Neon,
1646        x: *const f32,
1647        y: *const f32,
1648        len: usize,
1649        acc: Self::Accumulator,
1650    ) -> Self::Accumulator {
1651        let scalar = scalar_epilogue(
1652            x,
1653            y,
1654            len.min(Self::SIMDWidth::value() - 1),
1655            0.0f32,
1656            |acc, x: f32, y: f32| -> f32 { x.mul_add(y, acc) },
1657        );
1658        acc + Self::Accumulator::from_array(arch, [scalar, 0.0, 0.0, 0.0])
1659    }
1660
1661    #[inline(always)]
1662    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1663        x.sum_tree()
1664    }
1665}
1666
1667impl SIMDSchema<f32, f32, Scalar> for IP {
1668    type SIMDWidth = Const<4>;
1669    type Accumulator = Emulated<f32, 4>;
1670    type Left = Emulated<f32, 4>;
1671    type Right = Emulated<f32, 4>;
1672    type Return = f32;
1673    type Main = Strategy2x1;
1674
1675    #[inline(always)]
1676    fn init(&self, arch: Scalar) -> Self::Accumulator {
1677        Self::Accumulator::default(arch)
1678    }
1679
1680    #[inline(always)]
1681    fn accumulate(
1682        &self,
1683        x: Self::Left,
1684        y: Self::Right,
1685        acc: Self::Accumulator,
1686    ) -> Self::Accumulator {
1687        x * y + acc
1688    }
1689
1690    // Perform a final reduction.
1691    #[inline(always)]
1692    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1693        x.sum_tree()
1694    }
1695
1696    #[inline(always)]
1697    unsafe fn epilogue(
1698        &self,
1699        arch: Scalar,
1700        x: *const f32,
1701        y: *const f32,
1702        len: usize,
1703        acc: Self::Accumulator,
1704    ) -> Self::Accumulator {
1705        let mut s: f32 = 0.0;
1706        for i in 0..len {
1707            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1708            let vx = unsafe { x.add(i).read_unaligned() };
1709            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1710            let vy = unsafe { y.add(i).read_unaligned() };
1711            s += vx * vy;
1712        }
1713        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
1714    }
1715}
1716
1717#[cfg(target_arch = "x86_64")]
1718impl SIMDSchema<Half, Half, V4> for IP {
1719    type SIMDWidth = Const<8>;
1720    type Accumulator = <V4 as Architecture>::f32x8;
1721    type Left = <V4 as Architecture>::f16x8;
1722    type Right = <V4 as Architecture>::f16x8;
1723    type Return = f32;
1724    type Main = Strategy4x1;
1725
1726    #[inline(always)]
1727    fn init(&self, arch: V4) -> Self::Accumulator {
1728        Self::Accumulator::default(arch)
1729    }
1730
1731    #[inline(always)]
1732    fn accumulate(
1733        &self,
1734        x: Self::Left,
1735        y: Self::Right,
1736        acc: Self::Accumulator,
1737    ) -> Self::Accumulator {
1738        diskann_wide::alias!(f32s = <V4>::f32x8);
1739
1740        let x: f32s = x.into();
1741        let y: f32s = y.into();
1742        x.mul_add_simd(y, acc)
1743    }
1744
1745    #[inline(always)]
1746    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1747        x.sum_tree()
1748    }
1749}
1750
1751#[cfg(target_arch = "x86_64")]
1752impl SIMDSchema<Half, Half, V3> for IP {
1753    type SIMDWidth = Const<8>;
1754    type Accumulator = <V3 as Architecture>::f32x8;
1755    type Left = <V3 as Architecture>::f16x8;
1756    type Right = <V3 as Architecture>::f16x8;
1757    type Return = f32;
1758    type Main = Strategy2x4;
1759
1760    #[inline(always)]
1761    fn init(&self, arch: V3) -> Self::Accumulator {
1762        Self::Accumulator::default(arch)
1763    }
1764
1765    #[inline(always)]
1766    fn accumulate(
1767        &self,
1768        x: Self::Left,
1769        y: Self::Right,
1770        acc: Self::Accumulator,
1771    ) -> Self::Accumulator {
1772        diskann_wide::alias!(f32s = <V3>::f32x8);
1773
1774        let x: f32s = x.into();
1775        let y: f32s = y.into();
1776        x.mul_add_simd(y, acc)
1777    }
1778
1779    // Perform a final reduction.
1780    #[inline(always)]
1781    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1782        x.sum_tree()
1783    }
1784}
1785
1786#[cfg(target_arch = "aarch64")]
1787impl SIMDSchema<Half, Half, Neon> for IP {
1788    type SIMDWidth = Const<4>;
1789    type Accumulator = <Neon as Architecture>::f32x4;
1790    type Left = diskann_wide::arch::aarch64::f16x4;
1791    type Right = diskann_wide::arch::aarch64::f16x4;
1792    type Return = f32;
1793    type Main = Strategy4x1;
1794
1795    #[inline(always)]
1796    fn init(&self, arch: Neon) -> Self::Accumulator {
1797        Self::Accumulator::default(arch)
1798    }
1799
1800    #[inline(always)]
1801    fn accumulate(
1802        &self,
1803        x: Self::Left,
1804        y: Self::Right,
1805        acc: Self::Accumulator,
1806    ) -> Self::Accumulator {
1807        diskann_wide::alias!(f32s = <Neon>::f32x4);
1808
1809        let x: f32s = x.into();
1810        let y: f32s = y.into();
1811
1812        x.mul_add_simd(y, acc)
1813    }
1814
1815    #[inline(always)]
1816    unsafe fn epilogue(
1817        &self,
1818        arch: Neon,
1819        x: *const Half,
1820        y: *const Half,
1821        len: usize,
1822        acc: Self::Accumulator,
1823    ) -> Self::Accumulator {
1824        diskann_wide::alias!(f32s = <Neon>::f32x4);
1825
1826        let rest = scalar_epilogue(
1827            x,
1828            y,
1829            len.min(Self::SIMDWidth::value() - 1),
1830            f32s::default(arch),
1831            |acc, x: Half, y: Half| -> f32s {
1832                let zero = Half::default();
1833                let x: f32s = Self::Left::from_array(arch, [x, zero, zero, zero]).into();
1834                let y: f32s = Self::Right::from_array(arch, [y, zero, zero, zero]).into();
1835                x.mul_add_simd(y, acc)
1836            },
1837        );
1838        acc + rest
1839    }
1840
1841    #[inline(always)]
1842    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1843        x.sum_tree()
1844    }
1845}
1846
1847impl SIMDSchema<Half, Half, Scalar> for IP {
1848    type SIMDWidth = Const<1>;
1849    type Accumulator = Emulated<f32, 1>;
1850    type Left = Emulated<Half, 1>;
1851    type Right = Emulated<Half, 1>;
1852    type Return = f32;
1853    type Main = Strategy1x1;
1854
1855    #[inline(always)]
1856    fn init(&self, arch: Scalar) -> Self::Accumulator {
1857        Self::Accumulator::default(arch)
1858    }
1859
1860    #[inline(always)]
1861    fn accumulate(
1862        &self,
1863        x: Self::Left,
1864        y: Self::Right,
1865        acc: Self::Accumulator,
1866    ) -> Self::Accumulator {
1867        let x: Self::Accumulator = x.into();
1868        let y: Self::Accumulator = y.into();
1869        x * y + acc
1870    }
1871
1872    #[inline(always)]
1873    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1874        x.to_array()[0]
1875    }
1876}
1877
1878impl<A> SIMDSchema<f32, Half, A> for IP
1879where
1880    A: Architecture,
1881{
1882    type SIMDWidth = Const<8>;
1883    type Accumulator = A::f32x8;
1884    type Left = A::f32x8;
1885    type Right = A::f16x8;
1886    type Return = f32;
1887    type Main = Strategy4x2;
1888
1889    #[inline(always)]
1890    fn init(&self, arch: A) -> Self::Accumulator {
1891        Self::Accumulator::default(arch)
1892    }
1893
1894    #[inline(always)]
1895    fn accumulate(
1896        &self,
1897        x: Self::Left,
1898        y: Self::Right,
1899        acc: Self::Accumulator,
1900    ) -> Self::Accumulator {
1901        let y: A::f32x8 = y.into();
1902        x.mul_add_simd(y, acc)
1903    }
1904
1905    // Perform a final reduction.
1906    #[inline(always)]
1907    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1908        x.sum_tree()
1909    }
1910}
1911
1912#[cfg(target_arch = "x86_64")]
1913impl SIMDSchema<i8, i8, V4> for IP {
1914    type SIMDWidth = Const<32>;
1915    type Accumulator = <V4 as Architecture>::i32x16;
1916    type Left = <V4 as Architecture>::i8x32;
1917    type Right = <V4 as Architecture>::i8x32;
1918    type Return = f32;
1919    type Main = Strategy4x1;
1920
1921    #[inline(always)]
1922    fn init(&self, arch: V4) -> Self::Accumulator {
1923        Self::Accumulator::default(arch)
1924    }
1925
1926    #[inline(always)]
1927    fn accumulate(
1928        &self,
1929        x: Self::Left,
1930        y: Self::Right,
1931        acc: Self::Accumulator,
1932    ) -> Self::Accumulator {
1933        diskann_wide::alias!(i16s = <V4>::i16x32);
1934
1935        let x: i16s = x.into();
1936        let y: i16s = y.into();
1937        acc.dot_simd(x, y)
1938    }
1939
1940    #[inline(always)]
1941    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1942        x.sum_tree().as_f32_lossy()
1943    }
1944}
1945
1946#[cfg(target_arch = "x86_64")]
1947impl SIMDSchema<i8, i8, V3> for IP {
1948    type SIMDWidth = Const<16>;
1949    type Accumulator = <V3 as Architecture>::i32x8;
1950    type Left = <V3 as Architecture>::i8x16;
1951    type Right = <V3 as Architecture>::i8x16;
1952    type Return = f32;
1953    type Main = Strategy4x1;
1954
1955    #[inline(always)]
1956    fn init(&self, arch: V3) -> Self::Accumulator {
1957        Self::Accumulator::default(arch)
1958    }
1959
1960    #[inline(always)]
1961    fn accumulate(
1962        &self,
1963        x: Self::Left,
1964        y: Self::Right,
1965        acc: Self::Accumulator,
1966    ) -> Self::Accumulator {
1967        diskann_wide::alias!(i16s = <V3>::i16x16);
1968
1969        let x: i16s = x.into();
1970        let y: i16s = y.into();
1971        acc.dot_simd(x, y)
1972    }
1973
1974    // Perform a final reduction.
1975    #[inline(always)]
1976    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1977        x.sum_tree().as_f32_lossy()
1978    }
1979}
1980
1981#[cfg(target_arch = "aarch64")]
1982impl SIMDSchema<i8, i8, Neon> for IP {
1983    type SIMDWidth = Const<16>;
1984    type Accumulator = <Neon as Architecture>::i32x4;
1985    type Left = <Neon as Architecture>::i8x16;
1986    type Right = <Neon as Architecture>::i8x16;
1987    type Return = f32;
1988    type Main = Strategy2x1;
1989
1990    #[inline(always)]
1991    fn init(&self, arch: Neon) -> Self::Accumulator {
1992        Self::Accumulator::default(arch)
1993    }
1994
1995    #[inline(always)]
1996    fn accumulate(
1997        &self,
1998        x: Self::Left,
1999        y: Self::Right,
2000        acc: Self::Accumulator,
2001    ) -> Self::Accumulator {
2002        acc.dot_simd(x, y)
2003    }
2004
2005    #[inline(always)]
2006    unsafe fn epilogue(
2007        &self,
2008        arch: Neon,
2009        x: *const i8,
2010        y: *const i8,
2011        len: usize,
2012        acc: Self::Accumulator,
2013    ) -> Self::Accumulator {
2014        let scalar = scalar_epilogue(
2015            x,
2016            y,
2017            len.min(Self::SIMDWidth::value() - 1),
2018            0i32,
2019            |acc, x: i8, y: i8| -> i32 { acc + (x as i32) * (y as i32) },
2020        );
2021        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
2022    }
2023
2024    #[inline(always)]
2025    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2026        x.sum_tree().as_f32_lossy()
2027    }
2028}
2029
2030impl SIMDSchema<i8, i8, Scalar> for IP {
2031    type SIMDWidth = Const<1>;
2032    type Accumulator = Emulated<i32, 1>;
2033    type Left = Emulated<i8, 1>;
2034    type Right = Emulated<i8, 1>;
2035    type Return = f32;
2036    type Main = Strategy1x1;
2037
2038    #[inline(always)]
2039    fn init(&self, arch: Scalar) -> Self::Accumulator {
2040        Self::Accumulator::default(arch)
2041    }
2042
2043    #[inline(always)]
2044    fn accumulate(
2045        &self,
2046        x: Self::Left,
2047        y: Self::Right,
2048        acc: Self::Accumulator,
2049    ) -> Self::Accumulator {
2050        let x: Self::Accumulator = x.into();
2051        let y: Self::Accumulator = y.into();
2052        x * y + acc
2053    }
2054
2055    // Perform a final reduction.
2056    #[inline(always)]
2057    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2058        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
2059    }
2060
2061    #[inline(always)]
2062    unsafe fn epilogue(
2063        &self,
2064        _arch: Scalar,
2065        _x: *const i8,
2066        _y: *const i8,
2067        _len: usize,
2068        _acc: Self::Accumulator,
2069    ) -> Self::Accumulator {
2070        unreachable!("The SIMD width is 1, so there should be no epilogue")
2071    }
2072}
2073
2074#[cfg(target_arch = "x86_64")]
2075impl SIMDSchema<u8, u8, V4> for IP {
2076    type SIMDWidth = Const<32>;
2077    type Accumulator = <V4 as Architecture>::i32x16;
2078    type Left = <V4 as Architecture>::u8x32;
2079    type Right = <V4 as Architecture>::u8x32;
2080    type Return = f32;
2081    type Main = Strategy4x1;
2082
2083    #[inline(always)]
2084    fn init(&self, arch: V4) -> Self::Accumulator {
2085        Self::Accumulator::default(arch)
2086    }
2087
2088    #[inline(always)]
2089    fn accumulate(
2090        &self,
2091        x: Self::Left,
2092        y: Self::Right,
2093        acc: Self::Accumulator,
2094    ) -> Self::Accumulator {
2095        diskann_wide::alias!(i16s = <V4>::i16x32);
2096
2097        let x: i16s = x.into();
2098        let y: i16s = y.into();
2099        acc.dot_simd(x, y)
2100    }
2101
2102    #[inline(always)]
2103    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2104        x.sum_tree().as_f32_lossy()
2105    }
2106}
2107
2108#[cfg(target_arch = "x86_64")]
2109impl SIMDSchema<u8, u8, V3> for IP {
2110    type SIMDWidth = Const<16>;
2111    type Accumulator = <V3 as Architecture>::i32x8;
2112    type Left = <V3 as Architecture>::u8x16;
2113    type Right = <V3 as Architecture>::u8x16;
2114    type Return = f32;
2115    type Main = Strategy4x1;
2116
2117    #[inline(always)]
2118    fn init(&self, arch: V3) -> Self::Accumulator {
2119        Self::Accumulator::default(arch)
2120    }
2121
2122    #[inline(always)]
2123    fn accumulate(
2124        &self,
2125        x: Self::Left,
2126        y: Self::Right,
2127        acc: Self::Accumulator,
2128    ) -> Self::Accumulator {
2129        diskann_wide::alias!(i16s = <V3>::i16x16);
2130
2131        // NOTE: Promotiving to `i16` rather than `u16` to hit specialized AVX2
2132        // instructions on x86 hardware.
2133        let x: i16s = x.into();
2134        let y: i16s = y.into();
2135        acc.dot_simd(x, y)
2136    }
2137
2138    // Perform a final reduction.
2139    #[inline(always)]
2140    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2141        x.sum_tree().as_f32_lossy()
2142    }
2143}
2144
2145#[cfg(target_arch = "aarch64")]
2146impl SIMDSchema<u8, u8, Neon> for IP {
2147    type SIMDWidth = Const<16>;
2148    type Accumulator = <Neon as Architecture>::u32x4;
2149    type Left = <Neon as Architecture>::u8x16;
2150    type Right = <Neon as Architecture>::u8x16;
2151    type Return = f32;
2152    type Main = Strategy2x1;
2153
2154    #[inline(always)]
2155    fn init(&self, arch: Neon) -> Self::Accumulator {
2156        Self::Accumulator::default(arch)
2157    }
2158
2159    #[inline(always)]
2160    fn accumulate(
2161        &self,
2162        x: Self::Left,
2163        y: Self::Right,
2164        acc: Self::Accumulator,
2165    ) -> Self::Accumulator {
2166        acc.dot_simd(x, y)
2167    }
2168
2169    #[inline(always)]
2170    unsafe fn epilogue(
2171        &self,
2172        arch: Neon,
2173        x: *const u8,
2174        y: *const u8,
2175        len: usize,
2176        acc: Self::Accumulator,
2177    ) -> Self::Accumulator {
2178        let scalar = scalar_epilogue(
2179            x,
2180            y,
2181            len.min(Self::SIMDWidth::value() - 1),
2182            0u32,
2183            |acc, x: u8, y: u8| -> u32 { acc + (x as u32) * (y as u32) },
2184        );
2185        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
2186    }
2187
2188    #[inline(always)]
2189    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2190        x.sum_tree().as_f32_lossy()
2191    }
2192}
2193
2194impl SIMDSchema<u8, u8, Scalar> for IP {
2195    type SIMDWidth = Const<1>;
2196    type Accumulator = Emulated<i32, 1>;
2197    type Left = Emulated<u8, 1>;
2198    type Right = Emulated<u8, 1>;
2199    type Return = f32;
2200    type Main = Strategy1x1;
2201
2202    #[inline(always)]
2203    fn init(&self, arch: Scalar) -> Self::Accumulator {
2204        Self::Accumulator::default(arch)
2205    }
2206
2207    #[inline(always)]
2208    fn accumulate(
2209        &self,
2210        x: Self::Left,
2211        y: Self::Right,
2212        acc: Self::Accumulator,
2213    ) -> Self::Accumulator {
2214        let x: Self::Accumulator = x.into();
2215        let y: Self::Accumulator = y.into();
2216        x * y + acc
2217    }
2218
2219    // Perform a final reduction.
2220    #[inline(always)]
2221    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2222        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
2223    }
2224
2225    #[inline(always)]
2226    unsafe fn epilogue(
2227        &self,
2228        _arch: Scalar,
2229        _x: *const u8,
2230        _y: *const u8,
2231        _len: usize,
2232        _acc: Self::Accumulator,
2233    ) -> Self::Accumulator {
2234        unreachable!("The SIMD width is 1, so there should be no epilogue")
2235    }
2236}
2237
2238// An IP distance function that defers a final reduction.
2239#[derive(Clone, Copy, Debug)]
2240pub struct ResumableIP<A = diskann_wide::arch::Current>
2241where
2242    A: Architecture,
2243    IP: SIMDSchema<f32, f32, A>,
2244{
2245    acc: <IP as SIMDSchema<f32, f32, A>>::Accumulator,
2246}
2247
2248impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableIP<A>
2249where
2250    A: Architecture,
2251    IP: SIMDSchema<f32, f32, A, Return = f32>,
2252{
2253    type NonResumable = IP;
2254    type FinalReturn = f32;
2255
2256    #[inline(always)]
2257    fn init(arch: A) -> Self {
2258        Self { acc: IP.init(arch) }
2259    }
2260
2261    #[inline(always)]
2262    fn combine_with(&self, other: <IP as SIMDSchema<f32, f32, A>>::Accumulator) -> Self {
2263        Self {
2264            acc: self.acc + other,
2265        }
2266    }
2267
2268    #[inline(always)]
2269    fn sum(&self) -> f32 {
2270        IP.reduce(self.acc)
2271    }
2272}
2273
2274/////////////////////////////////
2275// Stateless Cosine Similarity //
2276/////////////////////////////////
2277
2278/// Accumulator of partial products for a full cosine distance computation (where
2279/// the norms of both the query and the dataset vector are computed on the fly).
2280#[derive(Debug, Clone, Copy)]
2281pub struct FullCosineAccumulator<T> {
2282    normx: T,
2283    normy: T,
2284    xy: T,
2285}
2286
2287impl<T> FullCosineAccumulator<T>
2288where
2289    T: SIMDVector
2290        + SIMDSumTree
2291        + SIMDMulAdd
2292        + std::ops::Mul<Output = T>
2293        + std::ops::Add<Output = T>,
2294    T::Scalar: LossyF32Conversion,
2295{
2296    #[inline(always)]
2297    pub fn new(arch: T::Arch) -> Self {
2298        // SAFETY: Zero initializing a SIMD vector is safe.
2299        let zero = T::default(arch);
2300        Self {
2301            normx: zero,
2302            normy: zero,
2303            xy: zero,
2304        }
2305    }
2306
2307    #[inline(always)]
2308    pub fn add_with(&self, x: T, y: T) -> Self {
2309        // SAFETY: Arithmetic on valid arguments is valid.
2310        FullCosineAccumulator {
2311            normx: x.mul_add_simd(x, self.normx),
2312            normy: y.mul_add_simd(y, self.normy),
2313            xy: x.mul_add_simd(y, self.xy),
2314        }
2315    }
2316
2317    #[inline(always)]
2318    pub fn add_with_unfused(&self, x: T, y: T) -> Self {
2319        // SAFETY: Arithmetic on valid arguments is valid.
2320        FullCosineAccumulator {
2321            normx: x * x + self.normx,
2322            normy: y * y + self.normy,
2323            xy: x * y + self.xy,
2324        }
2325    }
2326
2327    #[inline(always)]
2328    pub fn sum(&self) -> f32 {
2329        let normx = self.normx.sum_tree().as_f32_lossy();
2330        let normy = self.normy.sum_tree().as_f32_lossy();
2331
2332        // Evaluate the denominator early and use `force_eval`.
2333        // This will allow the long `sqrt` to be overlapped with some other instructions
2334        // rather than waiting at the end of the function.
2335        //
2336        // There is some worry of subnormal numbers, but we're optimizing for the common
2337        // case where norms are reasonable values.
2338        let denominator = normx.sqrt() * normy.sqrt();
2339        let prod = self.xy.sum_tree().as_f32_lossy();
2340
2341        // Force the final products to be completely computed before the range check.
2342        //
2343        // This prevents LLVM from trying to compute `normy` or `prod` *after* the check
2344        // to `normx`, which causes it to spill heavily to the stack.
2345        //
2346        // Unfortunately, this results in a reduction pattern that appears to be slightly
2347        // slower on AMD or Windows.
2348        force_eval(denominator);
2349        force_eval(prod);
2350
2351        // This basically checks if either norm is subnormal and if so, we treat the vector
2352        // as having norm zero.
2353        //
2354        // The reason to do this rather than checking `denominator` directly is to have
2355        // consistent behavior when one vector has a small norm (i.e., always treat it as
2356        // zero) rather than potentially changing behavior when the other vector has a very
2357        // large norm to compensate.
2358        if normx < f32::MIN_POSITIVE || normy < f32::MIN_POSITIVE {
2359            return 0.0;
2360        }
2361
2362        let v = prod / denominator;
2363        (-1.0f32).max(1.0f32.min(v))
2364    }
2365
2366    /// Compute the L2 distance from the partial products rather than the cosine similarity.
2367    #[inline(always)]
2368    pub fn sum_as_l2(&self) -> f32 {
2369        let normx = self.normx.sum_tree().as_f32_lossy();
2370        let normy = self.normy.sum_tree().as_f32_lossy();
2371        let xy = self.xy.sum_tree().as_f32_lossy();
2372        normx + normy - (xy + xy)
2373    }
2374}
2375
2376impl<T> std::ops::Add for FullCosineAccumulator<T>
2377where
2378    T: std::ops::Add<Output = T>,
2379{
2380    type Output = Self;
2381    #[inline(always)]
2382    fn add(self, other: Self) -> Self {
2383        FullCosineAccumulator {
2384            normx: self.normx + other.normx,
2385            normy: self.normy + other.normy,
2386            xy: self.xy + other.xy,
2387        }
2388    }
2389}
2390
2391/// A pure Cosine Similarity function that provides a final reduction.
2392#[derive(Default, Clone, Copy)]
2393pub struct CosineStateless;
2394
2395#[cfg(target_arch = "x86_64")]
2396impl SIMDSchema<f32, f32, V4> for CosineStateless {
2397    type SIMDWidth = Const<16>;
2398    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::f32x16>;
2399    type Left = <V4 as Architecture>::f32x16;
2400    type Right = <V4 as Architecture>::f32x16;
2401    type Return = f32;
2402
2403    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2404    // hefty unroll factor.
2405    type Main = Strategy2x4;
2406
2407    #[inline(always)]
2408    fn init(&self, arch: V4) -> Self::Accumulator {
2409        Self::Accumulator::new(arch)
2410    }
2411
2412    #[inline(always)]
2413    fn accumulate(
2414        &self,
2415        x: Self::Left,
2416        y: Self::Right,
2417        acc: Self::Accumulator,
2418    ) -> Self::Accumulator {
2419        acc.add_with(x, y)
2420    }
2421
2422    // Perform a final reduction.
2423    #[inline(always)]
2424    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2425        acc.sum()
2426    }
2427}
2428
2429#[cfg(target_arch = "x86_64")]
2430impl SIMDSchema<f32, f32, V3> for CosineStateless {
2431    type SIMDWidth = Const<8>;
2432    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::f32x8>;
2433    type Left = <V3 as Architecture>::f32x8;
2434    type Right = <V3 as Architecture>::f32x8;
2435    type Return = f32;
2436
2437    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2438    // hefty unroll factor.
2439    type Main = Strategy2x4;
2440
2441    #[inline(always)]
2442    fn init(&self, arch: V3) -> Self::Accumulator {
2443        Self::Accumulator::new(arch)
2444    }
2445
2446    #[inline(always)]
2447    fn accumulate(
2448        &self,
2449        x: Self::Left,
2450        y: Self::Right,
2451        acc: Self::Accumulator,
2452    ) -> Self::Accumulator {
2453        acc.add_with(x, y)
2454    }
2455
2456    // Perform a final reduction.
2457    #[inline(always)]
2458    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2459        acc.sum()
2460    }
2461}
2462
2463#[cfg(target_arch = "aarch64")]
2464impl SIMDSchema<f32, f32, Neon> for CosineStateless {
2465    type SIMDWidth = Const<4>;
2466    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::f32x4>;
2467    type Left = <Neon as Architecture>::f32x4;
2468    type Right = <Neon as Architecture>::f32x4;
2469    type Return = f32;
2470
2471    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2472    // hefty unroll factor.
2473    type Main = Strategy2x4;
2474
2475    #[inline(always)]
2476    fn init(&self, arch: Neon) -> Self::Accumulator {
2477        Self::Accumulator::new(arch)
2478    }
2479
2480    #[inline(always)]
2481    fn accumulate(
2482        &self,
2483        x: Self::Left,
2484        y: Self::Right,
2485        acc: Self::Accumulator,
2486    ) -> Self::Accumulator {
2487        acc.add_with(x, y)
2488    }
2489
2490    #[inline(always)]
2491    unsafe fn epilogue(
2492        &self,
2493        arch: Neon,
2494        x: *const f32,
2495        y: *const f32,
2496        len: usize,
2497        acc: Self::Accumulator,
2498    ) -> Self::Accumulator {
2499        let mut xx: f32 = 0.0;
2500        let mut yy: f32 = 0.0;
2501        let mut xy: f32 = 0.0;
2502        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
2503            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2504            let vx = unsafe { x.add(i).read_unaligned() };
2505            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2506            let vy = unsafe { y.add(i).read_unaligned() };
2507            xx = vx.mul_add(vx, xx);
2508            yy = vy.mul_add(vy, yy);
2509            xy = vx.mul_add(vy, xy);
2510        }
2511        type V = <Neon as Architecture>::f32x4;
2512        acc + FullCosineAccumulator {
2513            normx: V::from_array(arch, [xx, 0.0, 0.0, 0.0]),
2514            normy: V::from_array(arch, [yy, 0.0, 0.0, 0.0]),
2515            xy: V::from_array(arch, [xy, 0.0, 0.0, 0.0]),
2516        }
2517    }
2518
2519    // Perform a final reduction.
2520    #[inline(always)]
2521    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2522        acc.sum()
2523    }
2524}
2525
2526impl SIMDSchema<f32, f32, Scalar> for CosineStateless {
2527    type SIMDWidth = Const<4>;
2528    type Accumulator = FullCosineAccumulator<Emulated<f32, 4>>;
2529    type Left = Emulated<f32, 4>;
2530    type Right = Emulated<f32, 4>;
2531    type Return = f32;
2532
2533    type Main = Strategy2x1;
2534
2535    #[inline(always)]
2536    fn init(&self, arch: Scalar) -> Self::Accumulator {
2537        Self::Accumulator::new(arch)
2538    }
2539
2540    #[inline(always)]
2541    fn accumulate(
2542        &self,
2543        x: Self::Left,
2544        y: Self::Right,
2545        acc: Self::Accumulator,
2546    ) -> Self::Accumulator {
2547        acc.add_with_unfused(x, y)
2548    }
2549
2550    #[inline(always)]
2551    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2552        acc.sum()
2553    }
2554}
2555
2556#[cfg(target_arch = "x86_64")]
2557impl SIMDSchema<Half, Half, V4> for CosineStateless {
2558    type SIMDWidth = Const<16>;
2559    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::f32x16>;
2560    type Left = <V4 as Architecture>::f16x16;
2561    type Right = <V4 as Architecture>::f16x16;
2562    type Return = f32;
2563    type Main = Strategy2x4;
2564
2565    #[inline(always)]
2566    fn init(&self, arch: V4) -> Self::Accumulator {
2567        Self::Accumulator::new(arch)
2568    }
2569
2570    #[inline(always)]
2571    fn accumulate(
2572        &self,
2573        x: Self::Left,
2574        y: Self::Right,
2575        acc: Self::Accumulator,
2576    ) -> Self::Accumulator {
2577        diskann_wide::alias!(f32s = <V4>::f32x16);
2578
2579        let x: f32s = x.into();
2580        let y: f32s = y.into();
2581        acc.add_with(x, y)
2582    }
2583
2584    #[inline(always)]
2585    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2586        acc.sum()
2587    }
2588}
2589
2590#[cfg(target_arch = "x86_64")]
2591impl SIMDSchema<Half, Half, V3> for CosineStateless {
2592    type SIMDWidth = Const<8>;
2593    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::f32x8>;
2594    type Left = <V3 as Architecture>::f16x8;
2595    type Right = <V3 as Architecture>::f16x8;
2596    type Return = f32;
2597    type Main = Strategy2x4;
2598
2599    #[inline(always)]
2600    fn init(&self, arch: V3) -> Self::Accumulator {
2601        Self::Accumulator::new(arch)
2602    }
2603
2604    #[inline(always)]
2605    fn accumulate(
2606        &self,
2607        x: Self::Left,
2608        y: Self::Right,
2609        acc: Self::Accumulator,
2610    ) -> Self::Accumulator {
2611        diskann_wide::alias!(f32s = <V3>::f32x8);
2612
2613        let x: f32s = x.into();
2614        let y: f32s = y.into();
2615        acc.add_with(x, y)
2616    }
2617
2618    // Perform a final reduction.
2619    #[inline(always)]
2620    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2621        acc.sum()
2622    }
2623}
2624
2625#[cfg(target_arch = "aarch64")]
2626impl SIMDSchema<Half, Half, Neon> for CosineStateless {
2627    type SIMDWidth = Const<4>;
2628    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::f32x4>;
2629    type Left = diskann_wide::arch::aarch64::f16x4;
2630    type Right = diskann_wide::arch::aarch64::f16x4;
2631    type Return = f32;
2632
2633    type Main = Strategy2x4;
2634
2635    #[inline(always)]
2636    fn init(&self, arch: Neon) -> Self::Accumulator {
2637        Self::Accumulator::new(arch)
2638    }
2639
2640    #[inline(always)]
2641    fn accumulate(
2642        &self,
2643        x: Self::Left,
2644        y: Self::Right,
2645        acc: Self::Accumulator,
2646    ) -> Self::Accumulator {
2647        diskann_wide::alias!(f32s = <Neon>::f32x4);
2648
2649        let x: f32s = x.into();
2650        let y: f32s = y.into();
2651        acc.add_with(x, y)
2652    }
2653
2654    #[inline(always)]
2655    unsafe fn epilogue(
2656        &self,
2657        arch: Neon,
2658        x: *const Half,
2659        y: *const Half,
2660        len: usize,
2661        acc: Self::Accumulator,
2662    ) -> Self::Accumulator {
2663        type V = <Neon as Architecture>::f32x4;
2664
2665        let rest = scalar_epilogue(
2666            x,
2667            y,
2668            len.min(Self::SIMDWidth::value() - 1),
2669            FullCosineAccumulator::<V>::new(arch),
2670            |acc, x: Half, y: Half| -> FullCosineAccumulator<V> {
2671                let zero = Half::default();
2672                let x: V = Self::Left::from_array(arch, [x, zero, zero, zero]).into();
2673                let y: V = Self::Right::from_array(arch, [y, zero, zero, zero]).into();
2674                acc.add_with(x, y)
2675            },
2676        );
2677        acc + rest
2678    }
2679
2680    #[inline(always)]
2681    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2682        acc.sum()
2683    }
2684}
2685
2686impl SIMDSchema<Half, Half, Scalar> for CosineStateless {
2687    type SIMDWidth = Const<1>;
2688    type Accumulator = FullCosineAccumulator<Emulated<f32, 1>>;
2689    type Left = Emulated<Half, 1>;
2690    type Right = Emulated<Half, 1>;
2691    type Return = f32;
2692    type Main = Strategy1x1;
2693
2694    #[inline(always)]
2695    fn init(&self, arch: Scalar) -> Self::Accumulator {
2696        Self::Accumulator::new(arch)
2697    }
2698
2699    #[inline(always)]
2700    fn accumulate(
2701        &self,
2702        x: Self::Left,
2703        y: Self::Right,
2704        acc: Self::Accumulator,
2705    ) -> Self::Accumulator {
2706        let x: Emulated<f32, 1> = x.into();
2707        let y: Emulated<f32, 1> = y.into();
2708        acc.add_with_unfused(x, y)
2709    }
2710
2711    #[inline(always)]
2712    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2713        acc.sum()
2714    }
2715}
2716impl<A> SIMDSchema<f32, Half, A> for CosineStateless
2717where
2718    A: Architecture,
2719{
2720    type SIMDWidth = Const<8>;
2721    type Accumulator = FullCosineAccumulator<A::f32x8>;
2722    type Left = A::f32x8;
2723    type Right = A::f16x8;
2724    type Return = f32;
2725    type Main = Strategy2x4;
2726
2727    #[inline(always)]
2728    fn init(&self, arch: A) -> Self::Accumulator {
2729        Self::Accumulator::new(arch)
2730    }
2731
2732    #[inline(always)]
2733    fn accumulate(
2734        &self,
2735        x: Self::Left,
2736        y: Self::Right,
2737        acc: Self::Accumulator,
2738    ) -> Self::Accumulator {
2739        let y: A::f32x8 = y.into();
2740        acc.add_with(x, y)
2741    }
2742
2743    #[inline(always)]
2744    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2745        acc.sum()
2746    }
2747}
2748
2749#[cfg(target_arch = "x86_64")]
2750impl SIMDSchema<i8, i8, V4> for CosineStateless {
2751    type SIMDWidth = Const<32>;
2752    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::i32x16>;
2753    type Left = <V4 as Architecture>::i8x32;
2754    type Right = <V4 as Architecture>::i8x32;
2755    type Return = f32;
2756    type Main = Strategy4x1;
2757
2758    #[inline(always)]
2759    fn init(&self, arch: V4) -> Self::Accumulator {
2760        Self::Accumulator::new(arch)
2761    }
2762
2763    #[inline(always)]
2764    fn accumulate(
2765        &self,
2766        x: Self::Left,
2767        y: Self::Right,
2768        acc: Self::Accumulator,
2769    ) -> Self::Accumulator {
2770        diskann_wide::alias!(i16s = <V4>::i16x32);
2771
2772        let x: i16s = x.into();
2773        let y: i16s = y.into();
2774
2775        FullCosineAccumulator {
2776            normx: acc.normx.dot_simd(x, x),
2777            normy: acc.normy.dot_simd(y, y),
2778            xy: acc.xy.dot_simd(x, y),
2779        }
2780    }
2781
2782    // Perform a final reduction.
2783    #[inline(always)]
2784    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2785        x.sum()
2786    }
2787}
2788
2789#[cfg(target_arch = "x86_64")]
2790impl SIMDSchema<i8, i8, V3> for CosineStateless {
2791    type SIMDWidth = Const<16>;
2792    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::i32x8>;
2793    type Left = <V3 as Architecture>::i8x16;
2794    type Right = <V3 as Architecture>::i8x16;
2795    type Return = f32;
2796    type Main = Strategy4x1;
2797
2798    #[inline(always)]
2799    fn init(&self, arch: V3) -> Self::Accumulator {
2800        Self::Accumulator::new(arch)
2801    }
2802
2803    #[inline(always)]
2804    fn accumulate(
2805        &self,
2806        x: Self::Left,
2807        y: Self::Right,
2808        acc: Self::Accumulator,
2809    ) -> Self::Accumulator {
2810        diskann_wide::alias!(i16s = <V3>::i16x16);
2811
2812        let x: i16s = x.into();
2813        let y: i16s = y.into();
2814
2815        FullCosineAccumulator {
2816            normx: acc.normx.dot_simd(x, x),
2817            normy: acc.normy.dot_simd(y, y),
2818            xy: acc.xy.dot_simd(x, y),
2819        }
2820    }
2821
2822    // Perform a final reduction.
2823    #[inline(always)]
2824    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2825        x.sum()
2826    }
2827}
2828
2829#[cfg(target_arch = "aarch64")]
2830impl SIMDSchema<i8, i8, Neon> for CosineStateless {
2831    type SIMDWidth = Const<16>;
2832    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::i32x4>;
2833    type Left = <Neon as Architecture>::i8x16;
2834    type Right = <Neon as Architecture>::i8x16;
2835    type Return = f32;
2836    type Main = Strategy2x1;
2837
2838    #[inline(always)]
2839    fn init(&self, arch: Neon) -> Self::Accumulator {
2840        Self::Accumulator::new(arch)
2841    }
2842
2843    #[inline(always)]
2844    fn accumulate(
2845        &self,
2846        x: Self::Left,
2847        y: Self::Right,
2848        acc: Self::Accumulator,
2849    ) -> Self::Accumulator {
2850        FullCosineAccumulator {
2851            normx: acc.normx.dot_simd(x, x),
2852            normy: acc.normy.dot_simd(y, y),
2853            xy: acc.xy.dot_simd(x, y),
2854        }
2855    }
2856
2857    #[inline(always)]
2858    unsafe fn epilogue(
2859        &self,
2860        arch: Neon,
2861        x: *const i8,
2862        y: *const i8,
2863        len: usize,
2864        acc: Self::Accumulator,
2865    ) -> Self::Accumulator {
2866        let mut xx: i32 = 0;
2867        let mut yy: i32 = 0;
2868        let mut xy: i32 = 0;
2869        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
2870            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2871            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2872            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2873            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
2874            xx += vx * vx;
2875            xy += vx * vy;
2876            yy += vy * vy;
2877        }
2878        type V = <Neon as Architecture>::i32x4;
2879        acc + FullCosineAccumulator {
2880            normx: V::from_array(arch, [xx, 0, 0, 0]),
2881            normy: V::from_array(arch, [yy, 0, 0, 0]),
2882            xy: V::from_array(arch, [xy, 0, 0, 0]),
2883        }
2884    }
2885
2886    #[inline(always)]
2887    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2888        x.sum()
2889    }
2890}
2891
2892impl SIMDSchema<i8, i8, Scalar> for CosineStateless {
2893    type SIMDWidth = Const<4>;
2894    type Accumulator = FullCosineAccumulator<Emulated<i32, 4>>;
2895    type Left = Emulated<i8, 4>;
2896    type Right = Emulated<i8, 4>;
2897    type Return = f32;
2898    type Main = Strategy1x1;
2899
2900    #[inline(always)]
2901    fn init(&self, arch: Scalar) -> Self::Accumulator {
2902        Self::Accumulator::new(arch)
2903    }
2904
2905    #[inline(always)]
2906    fn accumulate(
2907        &self,
2908        x: Self::Left,
2909        y: Self::Right,
2910        acc: Self::Accumulator,
2911    ) -> Self::Accumulator {
2912        let x: Emulated<i32, 4> = x.into();
2913        let y: Emulated<i32, 4> = y.into();
2914        acc.add_with(x, y)
2915    }
2916
2917    // Perform a final reduction.
2918    #[inline(always)]
2919    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2920        x.sum()
2921    }
2922
2923    #[inline(always)]
2924    unsafe fn epilogue(
2925        &self,
2926        arch: Scalar,
2927        x: *const i8,
2928        y: *const i8,
2929        len: usize,
2930        acc: Self::Accumulator,
2931    ) -> Self::Accumulator {
2932        let mut xy: i32 = 0;
2933        let mut xx: i32 = 0;
2934        let mut yy: i32 = 0;
2935
2936        for i in 0..len {
2937            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2938            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2939            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2940            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
2941
2942            xx += vx * vx;
2943            xy += vx * vy;
2944            yy += vy * vy;
2945        }
2946
2947        acc + FullCosineAccumulator {
2948            normx: Emulated::from_array(arch, [xx, 0, 0, 0]),
2949            normy: Emulated::from_array(arch, [yy, 0, 0, 0]),
2950            xy: Emulated::from_array(arch, [xy, 0, 0, 0]),
2951        }
2952    }
2953}
2954
2955#[cfg(target_arch = "x86_64")]
2956impl SIMDSchema<u8, u8, V4> for CosineStateless {
2957    type SIMDWidth = Const<32>;
2958    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::i32x16>;
2959    type Left = <V4 as Architecture>::u8x32;
2960    type Right = <V4 as Architecture>::u8x32;
2961    type Return = f32;
2962    type Main = Strategy4x1;
2963
2964    #[inline(always)]
2965    fn init(&self, arch: V4) -> Self::Accumulator {
2966        Self::Accumulator::new(arch)
2967    }
2968
2969    #[inline(always)]
2970    fn accumulate(
2971        &self,
2972        x: Self::Left,
2973        y: Self::Right,
2974        acc: Self::Accumulator,
2975    ) -> Self::Accumulator {
2976        diskann_wide::alias!(i16s = <V4>::i16x32);
2977
2978        let x: i16s = x.into();
2979        let y: i16s = y.into();
2980
2981        FullCosineAccumulator {
2982            normx: acc.normx.dot_simd(x, x),
2983            normy: acc.normy.dot_simd(y, y),
2984            xy: acc.xy.dot_simd(x, y),
2985        }
2986    }
2987
2988    // Perform a final reduction.
2989    #[inline(always)]
2990    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2991        x.sum()
2992    }
2993}
2994
2995#[cfg(target_arch = "x86_64")]
2996impl SIMDSchema<u8, u8, V3> for CosineStateless {
2997    type SIMDWidth = Const<16>;
2998    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::i32x8>;
2999    type Left = <V3 as Architecture>::u8x16;
3000    type Right = <V3 as Architecture>::u8x16;
3001    type Return = f32;
3002    type Main = Strategy4x1;
3003
3004    #[inline(always)]
3005    fn init(&self, arch: V3) -> Self::Accumulator {
3006        Self::Accumulator::new(arch)
3007    }
3008
3009    #[inline(always)]
3010    fn accumulate(
3011        &self,
3012        x: Self::Left,
3013        y: Self::Right,
3014        acc: Self::Accumulator,
3015    ) -> Self::Accumulator {
3016        diskann_wide::alias!(i16s = <V3>::i16x16);
3017
3018        let x: i16s = x.into();
3019        let y: i16s = y.into();
3020
3021        FullCosineAccumulator {
3022            normx: acc.normx.dot_simd(x, x),
3023            normy: acc.normy.dot_simd(y, y),
3024            xy: acc.xy.dot_simd(x, y),
3025        }
3026    }
3027
3028    // Perform a final reduction.
3029    #[inline(always)]
3030    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3031        x.sum()
3032    }
3033}
3034
3035#[cfg(target_arch = "aarch64")]
3036impl SIMDSchema<u8, u8, Neon> for CosineStateless {
3037    type SIMDWidth = Const<16>;
3038    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::u32x4>;
3039    type Left = <Neon as Architecture>::u8x16;
3040    type Right = <Neon as Architecture>::u8x16;
3041    type Return = f32;
3042    type Main = Strategy2x1;
3043
3044    #[inline(always)]
3045    fn init(&self, arch: Neon) -> Self::Accumulator {
3046        Self::Accumulator::new(arch)
3047    }
3048
3049    #[inline(always)]
3050    fn accumulate(
3051        &self,
3052        x: Self::Left,
3053        y: Self::Right,
3054        acc: Self::Accumulator,
3055    ) -> Self::Accumulator {
3056        FullCosineAccumulator {
3057            normx: acc.normx.dot_simd(x, x),
3058            normy: acc.normy.dot_simd(y, y),
3059            xy: acc.xy.dot_simd(x, y),
3060        }
3061    }
3062
3063    #[inline(always)]
3064    unsafe fn epilogue(
3065        &self,
3066        arch: Neon,
3067        x: *const u8,
3068        y: *const u8,
3069        len: usize,
3070        acc: Self::Accumulator,
3071    ) -> Self::Accumulator {
3072        let mut xx: u32 = 0;
3073        let mut yy: u32 = 0;
3074        let mut xy: u32 = 0;
3075        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
3076            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3077            let vx: u32 = unsafe { x.add(i).read_unaligned() }.into();
3078            // SAFETY: The range `[y, y.add(len))` is valid for reads.
3079            let vy: u32 = unsafe { y.add(i).read_unaligned() }.into();
3080            xx += vx * vx;
3081            xy += vx * vy;
3082            yy += vy * vy;
3083        }
3084        type V = <Neon as Architecture>::u32x4;
3085        acc + FullCosineAccumulator {
3086            normx: V::from_array(arch, [xx, 0, 0, 0]),
3087            normy: V::from_array(arch, [yy, 0, 0, 0]),
3088            xy: V::from_array(arch, [xy, 0, 0, 0]),
3089        }
3090    }
3091
3092    #[inline(always)]
3093    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3094        x.sum()
3095    }
3096}
3097
3098impl SIMDSchema<u8, u8, Scalar> for CosineStateless {
3099    type SIMDWidth = Const<4>;
3100    type Accumulator = FullCosineAccumulator<Emulated<i32, 4>>;
3101    type Left = Emulated<u8, 4>;
3102    type Right = Emulated<u8, 4>;
3103    type Return = f32;
3104    type Main = Strategy1x1;
3105
3106    #[inline(always)]
3107    fn init(&self, arch: Scalar) -> Self::Accumulator {
3108        Self::Accumulator::new(arch)
3109    }
3110
3111    #[inline(always)]
3112    fn accumulate(
3113        &self,
3114        x: Self::Left,
3115        y: Self::Right,
3116        acc: Self::Accumulator,
3117    ) -> Self::Accumulator {
3118        let x: Emulated<i32, 4> = x.into();
3119        let y: Emulated<i32, 4> = y.into();
3120        acc.add_with(x, y)
3121    }
3122
3123    // Perform a final reduction.
3124    #[inline(always)]
3125    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3126        x.sum()
3127    }
3128
3129    #[inline(always)]
3130    unsafe fn epilogue(
3131        &self,
3132        arch: Scalar,
3133        x: *const u8,
3134        y: *const u8,
3135        len: usize,
3136        acc: Self::Accumulator,
3137    ) -> Self::Accumulator {
3138        let mut xy: i32 = 0;
3139        let mut xx: i32 = 0;
3140        let mut yy: i32 = 0;
3141
3142        for i in 0..len {
3143            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3144            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
3145            // SAFETY: The range `[y, y.add(len))` is valid for reads.
3146            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
3147
3148            xx += vx * vx;
3149            xy += vx * vy;
3150            yy += vy * vy;
3151        }
3152
3153        acc + FullCosineAccumulator {
3154            normx: Emulated::from_array(arch, [xx, 0, 0, 0]),
3155            normy: Emulated::from_array(arch, [yy, 0, 0, 0]),
3156            xy: Emulated::from_array(arch, [xy, 0, 0, 0]),
3157        }
3158    }
3159}
3160
3161/// A resumable cosine similarity computation.
3162#[derive(Debug, Clone, Copy)]
3163pub struct ResumableCosine<A = diskann_wide::arch::Current>(
3164    <CosineStateless as SIMDSchema<f32, f32, A>>::Accumulator,
3165)
3166where
3167    A: Architecture,
3168    CosineStateless: SIMDSchema<f32, f32, A>;
3169
3170impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableCosine<A>
3171where
3172    A: Architecture,
3173    CosineStateless: SIMDSchema<f32, f32, A, Return = f32>,
3174{
3175    type NonResumable = CosineStateless;
3176    type FinalReturn = f32;
3177
3178    #[inline(always)]
3179    fn init(arch: A) -> Self {
3180        Self(CosineStateless.init(arch))
3181    }
3182
3183    #[inline(always)]
3184    fn combine_with(
3185        &self,
3186        other: <CosineStateless as SIMDSchema<f32, f32, A>>::Accumulator,
3187    ) -> Self {
3188        Self(self.0 + other)
3189    }
3190
3191    #[inline(always)]
3192    fn sum(&self) -> f32 {
3193        CosineStateless.reduce(self.0)
3194    }
3195}
3196
3197/////
3198///// L1 Norm Implementations
3199/////
3200
3201// ==================================================================================================
3202// NOTE: L1Norm IS A LOGICAL UNARY OPERATION
3203// --------------------------------------------------------------------------------------------------
3204// Although wired through the generic binary 'SIMDSchema'/'simd_op' infrastructure (which expects
3205// two input slices of equal length), 'L1Norm' conceptually computes: sum_i |x_i|
3206// The right-hand operand is completely ignored and exists ONLY to satisfy the shared execution
3207// machinery (loop tiling, epilogue handling, etc.).
3208// ==================================================================================================
3209
3210// A pure L1 norm function that provides a final reduction.
3211#[derive(Clone, Copy, Debug, Default)]
3212pub struct L1Norm;
3213
3214#[cfg(target_arch = "x86_64")]
3215impl SIMDSchema<f32, f32, V4> for L1Norm {
3216    type SIMDWidth = Const<16>;
3217    type Accumulator = <V4 as Architecture>::f32x16;
3218    type Left = <V4 as Architecture>::f32x16;
3219    type Right = <V4 as Architecture>::f32x16;
3220    type Return = f32;
3221    type Main = Strategy4x1;
3222
3223    #[inline(always)]
3224    fn init(&self, arch: V4) -> Self::Accumulator {
3225        Self::Accumulator::default(arch)
3226    }
3227
3228    #[inline(always)]
3229    fn accumulate(
3230        &self,
3231        x: Self::Left,
3232        _y: Self::Right,
3233        acc: Self::Accumulator,
3234    ) -> Self::Accumulator {
3235        x.abs_simd() + acc
3236    }
3237
3238    // Perform a final reduction.
3239    #[inline(always)]
3240    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3241        x.sum_tree()
3242    }
3243}
3244
3245#[cfg(target_arch = "x86_64")]
3246impl SIMDSchema<f32, f32, V3> for L1Norm {
3247    type SIMDWidth = Const<8>;
3248    type Accumulator = <V3 as Architecture>::f32x8;
3249    type Left = <V3 as Architecture>::f32x8;
3250    type Right = <V3 as Architecture>::f32x8;
3251    type Return = f32;
3252    type Main = Strategy4x1;
3253
3254    #[inline(always)]
3255    fn init(&self, arch: V3) -> Self::Accumulator {
3256        Self::Accumulator::default(arch)
3257    }
3258
3259    #[inline(always)]
3260    fn accumulate(
3261        &self,
3262        x: Self::Left,
3263        _y: Self::Right,
3264        acc: Self::Accumulator,
3265    ) -> Self::Accumulator {
3266        x.abs_simd() + acc
3267    }
3268
3269    // Perform a final reduction.
3270    #[inline(always)]
3271    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3272        x.sum_tree()
3273    }
3274}
3275
3276#[cfg(target_arch = "aarch64")]
3277impl SIMDSchema<f32, f32, Neon> for L1Norm {
3278    type SIMDWidth = Const<4>;
3279    type Accumulator = <Neon as Architecture>::f32x4;
3280    type Left = <Neon as Architecture>::f32x4;
3281    type Right = <Neon as Architecture>::f32x4;
3282    type Return = f32;
3283    type Main = Strategy4x1;
3284
3285    #[inline(always)]
3286    fn init(&self, arch: Neon) -> Self::Accumulator {
3287        Self::Accumulator::default(arch)
3288    }
3289
3290    #[inline(always)]
3291    fn accumulate(
3292        &self,
3293        x: Self::Left,
3294        _y: Self::Right,
3295        acc: Self::Accumulator,
3296    ) -> Self::Accumulator {
3297        x.abs_simd() + acc
3298    }
3299
3300    #[inline(always)]
3301    unsafe fn epilogue(
3302        &self,
3303        arch: Neon,
3304        x: *const f32,
3305        _y: *const f32,
3306        len: usize,
3307        acc: Self::Accumulator,
3308    ) -> Self::Accumulator {
3309        let mut s: f32 = 0.0;
3310        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
3311            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3312            let vx = unsafe { x.add(i).read_unaligned() };
3313            s += vx.abs();
3314        }
3315        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
3316    }
3317
3318    #[inline(always)]
3319    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3320        x.sum_tree()
3321    }
3322}
3323
3324impl SIMDSchema<f32, f32, Scalar> for L1Norm {
3325    type SIMDWidth = Const<4>;
3326    type Accumulator = Emulated<f32, 4>;
3327    type Left = Emulated<f32, 4>;
3328    type Right = Emulated<f32, 4>;
3329    type Return = f32;
3330    type Main = Strategy2x1;
3331
3332    #[inline(always)]
3333    fn init(&self, arch: Scalar) -> Self::Accumulator {
3334        Self::Accumulator::default(arch)
3335    }
3336
3337    #[inline(always)]
3338    fn accumulate(
3339        &self,
3340        x: Self::Left,
3341        _y: Self::Right,
3342        acc: Self::Accumulator,
3343    ) -> Self::Accumulator {
3344        x.abs_simd() + acc
3345    }
3346
3347    // Perform a final reduction.
3348    #[inline(always)]
3349    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3350        x.sum_tree()
3351    }
3352
3353    #[inline(always)]
3354    unsafe fn epilogue(
3355        &self,
3356        arch: Scalar,
3357        x: *const f32,
3358        _y: *const f32,
3359        len: usize,
3360        acc: Self::Accumulator,
3361    ) -> Self::Accumulator {
3362        let mut s: f32 = 0.0;
3363        for i in 0..len {
3364            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3365            let vx = unsafe { x.add(i).read_unaligned() };
3366            s += vx.abs();
3367        }
3368        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
3369    }
3370}
3371
3372#[cfg(target_arch = "x86_64")]
3373impl SIMDSchema<Half, Half, V4> for L1Norm {
3374    type SIMDWidth = Const<8>;
3375    type Accumulator = <V4 as Architecture>::f32x8;
3376    type Left = <V4 as Architecture>::f16x8;
3377    type Right = <V4 as Architecture>::f16x8;
3378    type Return = f32;
3379    type Main = Strategy2x4;
3380
3381    #[inline(always)]
3382    fn init(&self, arch: V4) -> Self::Accumulator {
3383        Self::Accumulator::default(arch)
3384    }
3385
3386    #[inline(always)]
3387    fn accumulate(
3388        &self,
3389        x: Self::Left,
3390        _y: Self::Right,
3391        acc: Self::Accumulator,
3392    ) -> Self::Accumulator {
3393        let x: <V4 as Architecture>::f32x8 = x.into();
3394        x.abs_simd() + acc
3395    }
3396
3397    // Perform a final reduction.
3398    #[inline(always)]
3399    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3400        x.sum_tree()
3401    }
3402}
3403
3404#[cfg(target_arch = "x86_64")]
3405impl SIMDSchema<Half, Half, V3> for L1Norm {
3406    type SIMDWidth = Const<8>;
3407    type Accumulator = <V3 as Architecture>::f32x8;
3408    type Left = <V3 as Architecture>::f16x8;
3409    type Right = <V3 as Architecture>::f16x8;
3410    type Return = f32;
3411    type Main = Strategy2x4;
3412
3413    #[inline(always)]
3414    fn init(&self, arch: V3) -> Self::Accumulator {
3415        Self::Accumulator::default(arch)
3416    }
3417
3418    #[inline(always)]
3419    fn accumulate(
3420        &self,
3421        x: Self::Left,
3422        _y: Self::Right,
3423        acc: Self::Accumulator,
3424    ) -> Self::Accumulator {
3425        let x: <V3 as Architecture>::f32x8 = x.into();
3426        x.abs_simd() + acc
3427    }
3428
3429    // Perform a final reduction.
3430    #[inline(always)]
3431    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3432        x.sum_tree()
3433    }
3434}
3435
3436#[cfg(target_arch = "aarch64")]
3437impl SIMDSchema<Half, Half, Neon> for L1Norm {
3438    type SIMDWidth = Const<4>;
3439    type Accumulator = <Neon as Architecture>::f32x4;
3440    type Left = diskann_wide::arch::aarch64::f16x4;
3441    type Right = diskann_wide::arch::aarch64::f16x4;
3442    type Return = f32;
3443    type Main = Strategy2x4;
3444
3445    #[inline(always)]
3446    fn init(&self, arch: Neon) -> Self::Accumulator {
3447        Self::Accumulator::default(arch)
3448    }
3449
3450    #[inline(always)]
3451    fn accumulate(
3452        &self,
3453        x: Self::Left,
3454        _y: Self::Right,
3455        acc: Self::Accumulator,
3456    ) -> Self::Accumulator {
3457        let x: <Neon as Architecture>::f32x4 = x.into();
3458        x.abs_simd() + acc
3459    }
3460
3461    #[inline(always)]
3462    unsafe fn epilogue(
3463        &self,
3464        arch: Neon,
3465        x: *const Half,
3466        _y: *const Half,
3467        len: usize,
3468        acc: Self::Accumulator,
3469    ) -> Self::Accumulator {
3470        let rest = scalar_epilogue(
3471            x,
3472            x, // unused, but scalar_epilogue requires a right pointer
3473            len.min(Self::SIMDWidth::value() - 1),
3474            Self::Accumulator::default(arch),
3475            |acc, x: Half, _: Half| -> Self::Accumulator {
3476                let zero = Half::default();
3477                let x: Self::Accumulator =
3478                    Self::Left::from_array(arch, [x, zero, zero, zero]).into();
3479                x.abs_simd() + acc
3480            },
3481        );
3482        acc + rest
3483    }
3484
3485    // Perform a final reduction.
3486    #[inline(always)]
3487    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3488        x.sum_tree()
3489    }
3490}
3491
3492impl SIMDSchema<Half, Half, Scalar> for L1Norm {
3493    type SIMDWidth = Const<1>;
3494    type Accumulator = Emulated<f32, 1>;
3495    type Left = Emulated<Half, 1>;
3496    type Right = Emulated<Half, 1>;
3497    type Return = f32;
3498    type Main = Strategy1x1;
3499
3500    #[inline(always)]
3501    fn init(&self, arch: Scalar) -> Self::Accumulator {
3502        Self::Accumulator::default(arch)
3503    }
3504
3505    #[inline(always)]
3506    fn accumulate(
3507        &self,
3508        x: Self::Left,
3509        _y: Self::Right,
3510        acc: Self::Accumulator,
3511    ) -> Self::Accumulator {
3512        let x: Self::Accumulator = x.into();
3513        x.abs_simd() + acc
3514    }
3515
3516    // Perform a final reduction.
3517    #[inline(always)]
3518    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3519        x.to_array()[0]
3520    }
3521
3522    #[inline(always)]
3523    unsafe fn epilogue(
3524        &self,
3525        _arch: Scalar,
3526        _x: *const Half,
3527        _y: *const Half,
3528        _len: usize,
3529        _acc: Self::Accumulator,
3530    ) -> Self::Accumulator {
3531        unreachable!("The SIMD width is 1, so there should be no epilogue")
3532    }
3533}
3534
3535///////////
3536// Tests //
3537///////////
3538
3539#[cfg(test)]
3540mod tests {
3541    use std::{collections::HashMap, sync::LazyLock};
3542
3543    use approx::assert_relative_eq;
3544    use diskann_wide::{arch::Target1, ARCH};
3545    use half::f16;
3546    use rand::{distr::StandardUniform, rngs::StdRng, Rng, SeedableRng};
3547    use rand_distr;
3548
3549    use super::*;
3550    use crate::{distance::reference, norm::LInfNorm, test_util, unaligned};
3551
3552    ///////////////////////
3553    // Cosine Norm Check //
3554    ///////////////////////
3555
3556    fn cosine_norm_check_impl<A>(arch: A)
3557    where
3558        A: diskann_wide::Architecture,
3559        CosineStateless:
3560            SIMDSchema<f32, f32, A, Return = f32> + SIMDSchema<Half, Half, A, Return = f32>,
3561    {
3562        // Zero - f32
3563        {
3564            let x: [f32; 2] = [0.0, 0.0];
3565            let y: [f32; 2] = [0.0, 1.0];
3566            assert_eq!(
3567                simd_op(&CosineStateless {}, arch, x, x),
3568                0.0,
3569                "when both vectors are zero, similarity should be zero",
3570            );
3571            assert_eq!(
3572                simd_op(&CosineStateless {}, arch, x, y),
3573                0.0,
3574                "when one vector is zero, similarity should be zero",
3575            );
3576            assert_eq!(
3577                simd_op(&CosineStateless {}, arch, y, x),
3578                0.0,
3579                "when one vector is zero, similarity should be zero",
3580            );
3581        }
3582
3583        // Subnormal - f32
3584        {
3585            let x: [f32; 4] = [0.0, 0.0, 2.938736e-39f32, 0.0];
3586            let y: [f32; 4] = [0.0, 0.0, 1.0, 0.0];
3587            assert_eq!(
3588                simd_op(&CosineStateless {}, arch, x, x),
3589                0.0,
3590                "when both vectors are almost zero, similarity should be zero",
3591            );
3592            assert_eq!(
3593                simd_op(&CosineStateless {}, arch, x, y),
3594                0.0,
3595                "when one vector is almost zero, similarity should be zero",
3596            );
3597            assert_eq!(
3598                simd_op(&CosineStateless {}, arch, y, x),
3599                0.0,
3600                "when one vector is almost zero, similarity should be zero",
3601            );
3602        }
3603
3604        // Small - f32
3605        {
3606            let x: [f32; 4] = [0.0, 0.0, 1.0842022e-19f32, 0.0];
3607            let y: [f32; 4] = [0.0, 0.0, 1.0, 0.0];
3608            assert_eq!(
3609                simd_op(&CosineStateless {}, arch, x, x),
3610                1.0,
3611                "cosine-stateless should handle vectors this small",
3612            );
3613            assert_eq!(
3614                simd_op(&CosineStateless {}, arch, x, y),
3615                1.0,
3616                "cosine-stateless should handle vectors this small",
3617            );
3618            assert_eq!(
3619                simd_op(&CosineStateless {}, arch, y, x),
3620                1.0,
3621                "cosine-stateless should handle vectors this small",
3622            );
3623        }
3624
3625        let cvt = diskann_wide::cast_f32_to_f16;
3626
3627        // Zero - f16
3628        {
3629            let x: [Half; 2] = [Half::default(), Half::default()];
3630            let y: [Half; 2] = [Half::default(), cvt(1.0)];
3631            assert_eq!(
3632                simd_op(&CosineStateless {}, arch, x, x),
3633                0.0,
3634                "when both vectors are zero, similarity should be zero",
3635            );
3636            assert_eq!(
3637                simd_op(&CosineStateless {}, arch, x, y),
3638                0.0,
3639                "when one vector is zero, similarity should be zero",
3640            );
3641            assert_eq!(
3642                simd_op(&CosineStateless {}, arch, y, x),
3643                0.0,
3644                "when one vector is zero, similarity should be zero",
3645            );
3646        }
3647
3648        // Subnormal - f16
3649        {
3650            let x: [Half; 4] = [
3651                Half::default(),
3652                Half::default(),
3653                Half::MIN_POSITIVE_SUBNORMAL,
3654                Half::default(),
3655            ];
3656            let y: [Half; 4] = [Half::default(), Half::default(), cvt(1.0), Half::default()];
3657            assert_eq!(
3658                simd_op(&CosineStateless {}, arch, x, x),
3659                1.0,
3660                "when both vectors are almost zero, similarity should be zero",
3661            );
3662            assert_eq!(
3663                simd_op(&CosineStateless {}, arch, x, y),
3664                1.0,
3665                "when one vector is almost zero, similarity should be zero",
3666            );
3667            assert_eq!(
3668                simd_op(&CosineStateless {}, arch, y, x),
3669                1.0,
3670                "when one vector is almost zero, similarity should be zero",
3671            );
3672
3673            // Grab a range of floating point numbers whose squares cover the range of
3674            // our target threshold.
3675            //
3676            // Ensure that all combinations of values within this critical range to not
3677            // result in a misrounding.
3678            let threshold = f32::MIN_POSITIVE;
3679            let bound = 50;
3680            let values = {
3681                let mut down = threshold;
3682                let mut up = threshold;
3683                for _ in 0..bound {
3684                    down = down.next_down();
3685                    up = up.next_up();
3686                }
3687                assert!(down > 0.0);
3688                let min = down.sqrt();
3689                let max = up.sqrt();
3690                let mut v = min;
3691                let mut values = Vec::new();
3692                while v <= max {
3693                    values.push(v);
3694                    v = v.next_up();
3695                }
3696                values
3697            };
3698
3699            let mut lo = 0;
3700            let mut hi = 0;
3701            for i in values.iter() {
3702                for j in values.iter() {
3703                    let s: f32 = simd_op(&CosineStateless {}, arch, [*i], [*j]);
3704                    if i * i < threshold || j * j < threshold {
3705                        lo += 1;
3706                        assert_eq!(s, 0.0, "failed for i = {}, j = {}", i, j);
3707                    } else {
3708                        hi += 1;
3709                        assert_eq!(s, 1.0, "failed for i = {}, j = {}", i, j);
3710                    }
3711                }
3712            }
3713            assert_ne!(lo, 0);
3714            assert_ne!(hi, 0);
3715        }
3716    }
3717
3718    #[test]
3719    fn cosine_norm_check() {
3720        cosine_norm_check_impl::<diskann_wide::arch::Current>(diskann_wide::arch::current());
3721        cosine_norm_check_impl::<diskann_wide::arch::Scalar>(diskann_wide::arch::Scalar::new());
3722    }
3723
3724    #[test]
3725    #[cfg(target_arch = "x86_64")]
3726    fn cosine_norm_check_x86_64() {
3727        if let Some(arch) = V3::new_checked() {
3728            cosine_norm_check_impl::<V3>(arch);
3729        }
3730
3731        if let Some(arch) = V4::new_checked_miri() {
3732            cosine_norm_check_impl::<V4>(arch);
3733        }
3734    }
3735
3736    ////////////
3737    // Schema //
3738    ////////////
3739
3740    // Chunk the left and right hand slices and compute the result using a resumable function.
3741    fn test_resumable<T, L, R, A>(arch: A, x: &[L], y: &[R], chunk_size: usize) -> f32
3742    where
3743        A: Architecture,
3744        T: ResumableSIMDSchema<L, R, A, FinalReturn = f32>,
3745    {
3746        let mut acc = Resumable(<T as ResumableSIMDSchema<L, R, A>>::init(arch));
3747        let iter = std::iter::zip(x.chunks(chunk_size), y.chunks(chunk_size));
3748        for (a, b) in iter {
3749            acc = simd_op(&acc, arch, a, b);
3750        }
3751        acc.0.sum()
3752    }
3753
3754    fn stress_test_with_resumable<
3755        A: Architecture,
3756        O: Default + SIMDSchema<f32, f32, A, Return = f32>,
3757        T: ResumableSIMDSchema<f32, f32, A, NonResumable = O, FinalReturn = f32>,
3758        Rand: Rng,
3759    >(
3760        arch: A,
3761        reference: fn(&[f32], &[f32]) -> f32,
3762        dim: usize,
3763        epsilon: f32,
3764        max_relative: f32,
3765        rng: &mut Rand,
3766    ) {
3767        // Pick chunk sizes that exercise combinations of the unrolled loops.
3768        let chunk_divisors: Vec<usize> = vec![1, 2, 3, 4, 16, 54, 64, 65, 70, 77];
3769        let mut checker = test_util::AdHocChecker::<f32, f32>::new(|a: &[f32], b: &[f32]| {
3770            let expected = reference(a, b);
3771            let got = simd_op(&O::default(), arch, a, b);
3772            println!("dim = {}", dim);
3773            assert_relative_eq!(
3774                expected,
3775                got,
3776                epsilon = epsilon,
3777                max_relative = max_relative,
3778            );
3779
3780            if dim == 0 {
3781                return;
3782            }
3783
3784            for d in &chunk_divisors {
3785                let chunk_size = dim / d + (!dim.is_multiple_of(*d) as usize);
3786                let chunked = test_resumable::<T, f32, f32, _>(arch, a, b, chunk_size);
3787                assert_relative_eq!(chunked, got, epsilon = epsilon, max_relative = max_relative);
3788            }
3789        });
3790
3791        let dist = rand_distr::Normal::new(0.0, 10.0).unwrap();
3792        test_util::test_distance_function(&mut checker, dist, dist, dim, 10, rng);
3793
3794        //-----------//
3795        // Unaligned //
3796        //-----------//
3797
3798        let mut left = unaligned::Buffer::default();
3799        let mut right = unaligned::Buffer::default();
3800        let mut checker = test_util::Checker::<f32, f32, f32>::new(
3801            |a, b| simd_op(&O::default(), arch, a.as_unaligned(), b.as_unaligned()),
3802            |a, b| {
3803                left.copy(a);
3804                right.copy(b);
3805                simd_op(
3806                    &O::default(),
3807                    arch,
3808                    left.as_unaligned(),
3809                    right.as_unaligned(),
3810                )
3811            },
3812            |got, expected| assert_eq!(got, expected),
3813        );
3814        test_util::test_distance_function(&mut checker, dist, dist, dim, 10, rng);
3815    }
3816
3817    #[allow(clippy::too_many_arguments)]
3818    fn stress_test<L, R, DistLeft, DistRight, O, Rand, A>(
3819        arch: A,
3820        reference: fn(&[L], &[R]) -> f32,
3821        left_dist: DistLeft,
3822        right_dist: DistRight,
3823        dim: usize,
3824        epsilon: f32,
3825        max_relative: f32,
3826        rng: &mut Rand,
3827    ) where
3828        L: test_util::CornerCases + bytemuck::Pod,
3829        R: test_util::CornerCases + bytemuck::Pod,
3830        DistLeft: test_util::GenerateRandomArguments<L> + Copy,
3831        DistRight: test_util::GenerateRandomArguments<R> + Copy,
3832        O: Default + SIMDSchema<L, R, A, Return = f32>,
3833        Rand: Rng,
3834        A: Architecture,
3835    {
3836        let mut checker = test_util::Checker::<L, R, f32>::new(
3837            |x: &[L], y: &[R]| simd_op(&O::default(), arch, x, y),
3838            reference,
3839            |got, expected| {
3840                assert_relative_eq!(
3841                    expected,
3842                    got,
3843                    epsilon = epsilon,
3844                    max_relative = max_relative
3845                );
3846            },
3847        );
3848
3849        let trials = if cfg!(miri) { 0 } else { 10 };
3850        test_util::test_distance_function(&mut checker, left_dist, right_dist, dim, trials, rng);
3851
3852        //-----------//
3853        // Unaligned //
3854        //-----------//
3855
3856        let mut left = unaligned::Buffer::default();
3857        let mut right = unaligned::Buffer::default();
3858        let mut checker = test_util::Checker::<L, R, f32>::new(
3859            |a, b| simd_op(&O::default(), arch, a.as_unaligned(), b.as_unaligned()),
3860            |a, b| {
3861                left.copy(a);
3862                right.copy(b);
3863                simd_op(
3864                    &O::default(),
3865                    arch,
3866                    left.as_unaligned(),
3867                    right.as_unaligned(),
3868                )
3869            },
3870            |got, expected| assert_eq!(got, expected),
3871        );
3872        test_util::test_distance_function(&mut checker, left_dist, right_dist, dim, trials, rng);
3873    }
3874
3875    fn stress_test_linf<L, Dist, Rand, A>(
3876        arch: A,
3877        reference: fn(&[L]) -> f32,
3878        dist: Dist,
3879        dim: usize,
3880        epsilon: f32,
3881        max_relative: f32,
3882        rng: &mut Rand,
3883    ) where
3884        L: test_util::CornerCases + Copy,
3885        Dist: Clone + test_util::GenerateRandomArguments<L>,
3886        Rand: Rng,
3887        A: Architecture,
3888        LInfNorm: for<'a> Target1<A, f32, &'a [L]>,
3889    {
3890        let mut checker = test_util::Checker::<L, L, f32>::new(
3891            |x: &[L], _y: &[L]| (LInfNorm).run(arch, x),
3892            |x: &[L], _y: &[L]| reference(x),
3893            |got, expected| {
3894                assert_relative_eq!(
3895                    expected,
3896                    got,
3897                    epsilon = epsilon,
3898                    max_relative = max_relative
3899                );
3900            },
3901        );
3902
3903        println!("checking {dim}");
3904        test_util::test_distance_function(&mut checker, dist.clone(), dist, dim, 10, rng);
3905    }
3906
3907    /////////
3908    // f32 //
3909    /////////
3910
3911    macro_rules! float_test {
3912        ($name:ident,
3913         $impl:ty,
3914         $resumable:ident,
3915         $reference:path,
3916         $eps:literal,
3917         $relative:literal,
3918         $seed:literal,
3919         $upper:literal,
3920         $($arch:tt)*
3921        ) => {
3922            #[test]
3923            fn $name() {
3924                if let Some(arch) = $($arch)* {
3925                    let mut rng = StdRng::seed_from_u64($seed);
3926                    for dim in 0..$upper {
3927                        stress_test_with_resumable::<_, $impl, $resumable<_>, StdRng>(
3928                            arch,
3929                            |l, r| $reference(l, r).into_inner(),
3930                            dim,
3931                            $eps,
3932                            $relative,
3933                            &mut rng,
3934                        );
3935                    }
3936                }
3937            }
3938        }
3939    }
3940
3941    //----//
3942    // L2 //
3943    //----//
3944
3945    float_test!(
3946        test_l2_f32_current,
3947        L2,
3948        ResumableL2,
3949        reference::reference_squared_l2_f32_mathematical,
3950        1e-5,
3951        1e-5,
3952        0xf149c2bcde660128,
3953        64,
3954        Some(diskann_wide::ARCH)
3955    );
3956
3957    float_test!(
3958        test_l2_f32_scalar,
3959        L2,
3960        ResumableL2,
3961        reference::reference_squared_l2_f32_mathematical,
3962        1e-5,
3963        1e-5,
3964        0xf149c2bcde660128,
3965        64,
3966        Some(diskann_wide::arch::Scalar)
3967    );
3968
3969    #[cfg(target_arch = "x86_64")]
3970    float_test!(
3971        test_l2_f32_x86_64_v3,
3972        L2,
3973        ResumableL2,
3974        reference::reference_squared_l2_f32_mathematical,
3975        1e-5,
3976        1e-5,
3977        0xf149c2bcde660128,
3978        256,
3979        V3::new_checked()
3980    );
3981
3982    #[cfg(target_arch = "x86_64")]
3983    float_test!(
3984        test_l2_f32_x86_64_v4,
3985        L2,
3986        ResumableL2,
3987        reference::reference_squared_l2_f32_mathematical,
3988        1e-5,
3989        1e-5,
3990        0xf149c2bcde660128,
3991        256,
3992        V4::new_checked_miri()
3993    );
3994
3995    #[cfg(target_arch = "aarch64")]
3996    float_test!(
3997        test_l2_f32_aarch64_neon,
3998        L2,
3999        ResumableL2,
4000        reference::reference_squared_l2_f32_mathematical,
4001        1e-5,
4002        1e-5,
4003        0xf149c2bcde660128,
4004        256,
4005        Neon::new_checked()
4006    );
4007
4008    //----//
4009    // IP //
4010    //----//
4011
4012    float_test!(
4013        test_ip_f32_current,
4014        IP,
4015        ResumableIP,
4016        reference::reference_innerproduct_f32_mathematical,
4017        2e-4,
4018        1e-3,
4019        0xb4687c17a9ea9866,
4020        64,
4021        Some(diskann_wide::ARCH)
4022    );
4023
4024    float_test!(
4025        test_ip_f32_scalar,
4026        IP,
4027        ResumableIP,
4028        reference::reference_innerproduct_f32_mathematical,
4029        2e-4,
4030        1e-3,
4031        0xb4687c17a9ea9866,
4032        64,
4033        Some(diskann_wide::arch::Scalar)
4034    );
4035
4036    #[cfg(target_arch = "x86_64")]
4037    float_test!(
4038        test_ip_f32_x86_64_v3,
4039        IP,
4040        ResumableIP,
4041        reference::reference_innerproduct_f32_mathematical,
4042        2e-4,
4043        1e-3,
4044        0xb4687c17a9ea9866,
4045        256,
4046        V3::new_checked()
4047    );
4048
4049    #[cfg(target_arch = "x86_64")]
4050    float_test!(
4051        test_ip_f32_x86_64_v4,
4052        IP,
4053        ResumableIP,
4054        reference::reference_innerproduct_f32_mathematical,
4055        2e-4,
4056        1e-3,
4057        0xb4687c17a9ea9866,
4058        256,
4059        V4::new_checked_miri()
4060    );
4061
4062    #[cfg(target_arch = "aarch64")]
4063    float_test!(
4064        test_ip_f32_aarch64_neon,
4065        IP,
4066        ResumableIP,
4067        reference::reference_innerproduct_f32_mathematical,
4068        2e-4,
4069        1e-3,
4070        0xb4687c17a9ea9866,
4071        256,
4072        Neon::new_checked()
4073    );
4074
4075    //--------//
4076    // Cosine //
4077    //--------//
4078
4079    float_test!(
4080        test_cosine_f32_current,
4081        CosineStateless,
4082        ResumableCosine,
4083        reference::reference_cosine_f32_mathematical,
4084        1e-5,
4085        1e-5,
4086        0xe860e9dc65f38bb8,
4087        64,
4088        Some(diskann_wide::ARCH)
4089    );
4090
4091    float_test!(
4092        test_cosine_f32_scalar,
4093        CosineStateless,
4094        ResumableCosine,
4095        reference::reference_cosine_f32_mathematical,
4096        1e-5,
4097        1e-5,
4098        0xe860e9dc65f38bb8,
4099        64,
4100        Some(diskann_wide::arch::Scalar)
4101    );
4102
4103    #[cfg(target_arch = "x86_64")]
4104    float_test!(
4105        test_cosine_f32_x86_64_v3,
4106        CosineStateless,
4107        ResumableCosine,
4108        reference::reference_cosine_f32_mathematical,
4109        1e-5,
4110        1e-5,
4111        0xe860e9dc65f38bb8,
4112        256,
4113        V3::new_checked()
4114    );
4115
4116    #[cfg(target_arch = "x86_64")]
4117    float_test!(
4118        test_cosine_f32_x86_64_v4,
4119        CosineStateless,
4120        ResumableCosine,
4121        reference::reference_cosine_f32_mathematical,
4122        1e-5,
4123        1e-5,
4124        0xe860e9dc65f38bb8,
4125        256,
4126        V4::new_checked_miri()
4127    );
4128
4129    #[cfg(target_arch = "aarch64")]
4130    float_test!(
4131        test_cosine_f32_aarch64_neon,
4132        CosineStateless,
4133        ResumableCosine,
4134        reference::reference_cosine_f32_mathematical,
4135        1e-5,
4136        1e-5,
4137        0xe860e9dc65f38bb8,
4138        256,
4139        Neon::new_checked()
4140    );
4141
4142    /////////
4143    // f16 //
4144    /////////
4145
4146    macro_rules! half_test {
4147        ($name:ident,
4148         $impl:ty,
4149         $reference:path,
4150         $eps:literal,
4151         $relative:literal,
4152         $seed:literal,
4153         $upper:literal,
4154         $($arch:tt)*
4155        ) => {
4156            #[test]
4157            fn $name() {
4158                if let Some(arch) = $($arch)* {
4159                    let mut rng = StdRng::seed_from_u64($seed);
4160                    for dim in 0..$upper {
4161                        stress_test::<
4162                            Half,
4163                            Half,
4164                            rand_distr::Normal<f32>,
4165                            rand_distr::Normal<f32>,
4166                            $impl,
4167                            StdRng,
4168                            _
4169                        >(
4170                            arch,
4171                            |l, r| $reference(l, r).into_inner(),
4172                            rand_distr::Normal::new(0.0, 10.0).unwrap(),
4173                            rand_distr::Normal::new(0.0, 10.0).unwrap(),
4174                            dim,
4175                            $eps,
4176                            $relative,
4177                            &mut rng
4178                        );
4179                    }
4180                }
4181            }
4182        }
4183    }
4184
4185    //----//
4186    // L2 //
4187    //----//
4188
4189    half_test!(
4190        test_l2_f16_current,
4191        L2,
4192        reference::reference_squared_l2_f16_mathematical,
4193        1e-5,
4194        1e-5,
4195        0x87ca6f1051667500,
4196        64,
4197        Some(diskann_wide::ARCH)
4198    );
4199
4200    half_test!(
4201        test_l2_f16_scalar,
4202        L2,
4203        reference::reference_squared_l2_f16_mathematical,
4204        1e-5,
4205        1e-5,
4206        0x87ca6f1051667500,
4207        64,
4208        Some(diskann_wide::arch::Scalar)
4209    );
4210
4211    #[cfg(target_arch = "x86_64")]
4212    half_test!(
4213        test_l2_f16_x86_64_v3,
4214        L2,
4215        reference::reference_squared_l2_f16_mathematical,
4216        1e-5,
4217        1e-5,
4218        0x87ca6f1051667500,
4219        256,
4220        V3::new_checked()
4221    );
4222
4223    #[cfg(target_arch = "x86_64")]
4224    half_test!(
4225        test_l2_f16_x86_64_v4,
4226        L2,
4227        reference::reference_squared_l2_f16_mathematical,
4228        1e-5,
4229        1e-5,
4230        0x87ca6f1051667500,
4231        256,
4232        V4::new_checked_miri()
4233    );
4234
4235    #[cfg(target_arch = "aarch64")]
4236    half_test!(
4237        test_l2_f16_aarch64_neon,
4238        L2,
4239        reference::reference_squared_l2_f16_mathematical,
4240        1e-5,
4241        1e-5,
4242        0x87ca6f1051667500,
4243        256,
4244        Neon::new_checked()
4245    );
4246
4247    //----//
4248    // IP //
4249    //----//
4250
4251    half_test!(
4252        test_ip_f16_current,
4253        IP,
4254        reference::reference_innerproduct_f16_mathematical,
4255        2e-4,
4256        2e-4,
4257        0x5909f5f20307ccbe,
4258        64,
4259        Some(diskann_wide::ARCH)
4260    );
4261
4262    half_test!(
4263        test_ip_f16_scalar,
4264        IP,
4265        reference::reference_innerproduct_f16_mathematical,
4266        2e-4,
4267        2e-4,
4268        0x5909f5f20307ccbe,
4269        64,
4270        Some(diskann_wide::arch::Scalar)
4271    );
4272
4273    #[cfg(target_arch = "x86_64")]
4274    half_test!(
4275        test_ip_f16_x86_64_v3,
4276        IP,
4277        reference::reference_innerproduct_f16_mathematical,
4278        2e-4,
4279        2e-4,
4280        0x5909f5f20307ccbe,
4281        256,
4282        V3::new_checked()
4283    );
4284
4285    #[cfg(target_arch = "x86_64")]
4286    half_test!(
4287        test_ip_f16_x86_64_v4,
4288        IP,
4289        reference::reference_innerproduct_f16_mathematical,
4290        2e-4,
4291        2e-4,
4292        0x5909f5f20307ccbe,
4293        256,
4294        V4::new_checked_miri()
4295    );
4296
4297    #[cfg(target_arch = "aarch64")]
4298    half_test!(
4299        test_ip_f16_aarch64_neon,
4300        IP,
4301        reference::reference_innerproduct_f16_mathematical,
4302        2e-4,
4303        2e-4,
4304        0x5909f5f20307ccbe,
4305        256,
4306        Neon::new_checked()
4307    );
4308
4309    //--------//
4310    // Cosine //
4311    //--------//
4312
4313    half_test!(
4314        test_cosine_f16_current,
4315        CosineStateless,
4316        reference::reference_cosine_f16_mathematical,
4317        1e-5,
4318        1e-5,
4319        0x41dda34655f05ef6,
4320        64,
4321        Some(diskann_wide::ARCH)
4322    );
4323
4324    half_test!(
4325        test_cosine_f16_scalar,
4326        CosineStateless,
4327        reference::reference_cosine_f16_mathematical,
4328        1e-5,
4329        1e-5,
4330        0x41dda34655f05ef6,
4331        64,
4332        Some(diskann_wide::arch::Scalar)
4333    );
4334
4335    #[cfg(target_arch = "x86_64")]
4336    half_test!(
4337        test_cosine_f16_x86_64_v3,
4338        CosineStateless,
4339        reference::reference_cosine_f16_mathematical,
4340        1e-5,
4341        1e-5,
4342        0x41dda34655f05ef6,
4343        256,
4344        V3::new_checked()
4345    );
4346
4347    #[cfg(target_arch = "x86_64")]
4348    half_test!(
4349        test_cosine_f16_x86_64_v4,
4350        CosineStateless,
4351        reference::reference_cosine_f16_mathematical,
4352        1e-5,
4353        1e-5,
4354        0x41dda34655f05ef6,
4355        256,
4356        V4::new_checked_miri()
4357    );
4358
4359    #[cfg(target_arch = "aarch64")]
4360    half_test!(
4361        test_cosine_f16_aarch64_neon,
4362        CosineStateless,
4363        reference::reference_cosine_f16_mathematical,
4364        1e-5,
4365        1e-5,
4366        0x41dda34655f05ef6,
4367        256,
4368        Neon::new_checked()
4369    );
4370
4371    /////////////
4372    // Integer //
4373    /////////////
4374
4375    macro_rules! int_test {
4376        (
4377            $name:ident,
4378            $T:ty,
4379            $impl:ty,
4380            $reference:path,
4381            $seed:literal,
4382            $upper:literal,
4383            { $($arch:tt)* }
4384        ) => {
4385            #[test]
4386            fn $name() {
4387                if let Some(arch) = $($arch)* {
4388                    let mut rng = StdRng::seed_from_u64($seed);
4389                    for dim in 0..$upper {
4390                        stress_test::<$T, $T, _, _, $impl, _, _>(
4391                            arch,
4392                            |l, r| $reference(l, r).into_inner(),
4393                            StandardUniform,
4394                            StandardUniform,
4395                            dim,
4396                            0.0,
4397                            0.0,
4398                            &mut rng,
4399                        )
4400                    }
4401                }
4402            }
4403        }
4404    }
4405
4406    //----//
4407    // U8 //
4408    //----//
4409
4410    int_test!(
4411        test_l2_u8_current,
4412        u8,
4413        L2,
4414        reference::reference_squared_l2_u8_mathematical,
4415        0x945bdc37d8279d4b,
4416        128,
4417        { Some(ARCH) }
4418    );
4419
4420    int_test!(
4421        test_l2_u8_scalar,
4422        u8,
4423        L2,
4424        reference::reference_squared_l2_u8_mathematical,
4425        0x74c86334ab7a51f9,
4426        128,
4427        { Some(diskann_wide::arch::Scalar) }
4428    );
4429
4430    #[cfg(target_arch = "x86_64")]
4431    int_test!(
4432        test_l2_u8_x86_64_v3,
4433        u8,
4434        L2,
4435        reference::reference_squared_l2_u8_mathematical,
4436        0x74c86334ab7a51f9,
4437        256,
4438        { V3::new_checked() }
4439    );
4440
4441    #[cfg(target_arch = "x86_64")]
4442    int_test!(
4443        test_l2_u8_x86_64_v4,
4444        u8,
4445        L2,
4446        reference::reference_squared_l2_u8_mathematical,
4447        0x74c86334ab7a51f9,
4448        320,
4449        { V4::new_checked_miri() }
4450    );
4451
4452    #[cfg(target_arch = "aarch64")]
4453    int_test!(
4454        test_l2_u8_aarch64_neon,
4455        u8,
4456        L2,
4457        reference::reference_squared_l2_u8_mathematical,
4458        0x74c86334ab7a51f9,
4459        320,
4460        { Neon::new_checked() }
4461    );
4462
4463    int_test!(
4464        test_ip_u8_current,
4465        u8,
4466        IP,
4467        reference::reference_innerproduct_u8_mathematical,
4468        0xcbe0342c75085fd5,
4469        64,
4470        { Some(ARCH) }
4471    );
4472
4473    int_test!(
4474        test_ip_u8_scalar,
4475        u8,
4476        IP,
4477        reference::reference_innerproduct_u8_mathematical,
4478        0x888e07fc489e773f,
4479        64,
4480        { Some(diskann_wide::arch::Scalar) }
4481    );
4482
4483    #[cfg(target_arch = "x86_64")]
4484    int_test!(
4485        test_ip_u8_x86_64_v3,
4486        u8,
4487        IP,
4488        reference::reference_innerproduct_u8_mathematical,
4489        0x888e07fc489e773f,
4490        256,
4491        { V3::new_checked() }
4492    );
4493
4494    #[cfg(target_arch = "x86_64")]
4495    int_test!(
4496        test_ip_u8_x86_64_v4,
4497        u8,
4498        IP,
4499        reference::reference_innerproduct_u8_mathematical,
4500        0x888e07fc489e773f,
4501        320,
4502        { V4::new_checked_miri() }
4503    );
4504
4505    #[cfg(target_arch = "aarch64")]
4506    int_test!(
4507        test_ip_u8_aarch64_neon,
4508        u8,
4509        IP,
4510        reference::reference_innerproduct_u8_mathematical,
4511        0x888e07fc489e773f,
4512        320,
4513        { Neon::new_checked() }
4514    );
4515
4516    int_test!(
4517        test_cosine_u8_current,
4518        u8,
4519        CosineStateless,
4520        reference::reference_cosine_u8_mathematical,
4521        0x96867b6aff616b28,
4522        64,
4523        { Some(ARCH) }
4524    );
4525
4526    int_test!(
4527        test_cosine_u8_scalar,
4528        u8,
4529        CosineStateless,
4530        reference::reference_cosine_u8_mathematical,
4531        0xcc258c9391733211,
4532        64,
4533        { Some(diskann_wide::arch::Scalar) }
4534    );
4535
4536    #[cfg(target_arch = "x86_64")]
4537    int_test!(
4538        test_cosine_u8_x86_64_v3,
4539        u8,
4540        CosineStateless,
4541        reference::reference_cosine_u8_mathematical,
4542        0xcc258c9391733211,
4543        256,
4544        { V3::new_checked() }
4545    );
4546
4547    #[cfg(target_arch = "x86_64")]
4548    int_test!(
4549        test_cosine_u8_x86_64_v4,
4550        u8,
4551        CosineStateless,
4552        reference::reference_cosine_u8_mathematical,
4553        0xcc258c9391733211,
4554        320,
4555        { V4::new_checked_miri() }
4556    );
4557
4558    #[cfg(target_arch = "aarch64")]
4559    int_test!(
4560        test_cosine_u8_aarch64_neon,
4561        u8,
4562        CosineStateless,
4563        reference::reference_cosine_u8_mathematical,
4564        0xcc258c9391733211,
4565        320,
4566        { Neon::new_checked() }
4567    );
4568
4569    //----//
4570    // I8 //
4571    //----//
4572
4573    int_test!(
4574        test_l2_i8_current,
4575        i8,
4576        L2,
4577        reference::reference_squared_l2_i8_mathematical,
4578        0xa60136248cd3c2f0,
4579        64,
4580        { Some(ARCH) }
4581    );
4582
4583    int_test!(
4584        test_l2_i8_scalar,
4585        i8,
4586        L2,
4587        reference::reference_squared_l2_i8_mathematical,
4588        0x3e8bada709e176be,
4589        64,
4590        { Some(diskann_wide::arch::Scalar) }
4591    );
4592
4593    #[cfg(target_arch = "x86_64")]
4594    int_test!(
4595        test_l2_i8_x86_64_v3,
4596        i8,
4597        L2,
4598        reference::reference_squared_l2_i8_mathematical,
4599        0x3e8bada709e176be,
4600        256,
4601        { V3::new_checked() }
4602    );
4603
4604    #[cfg(target_arch = "x86_64")]
4605    int_test!(
4606        test_l2_i8_x86_64_v4,
4607        i8,
4608        L2,
4609        reference::reference_squared_l2_i8_mathematical,
4610        0x3e8bada709e176be,
4611        320,
4612        { V4::new_checked_miri() }
4613    );
4614
4615    #[cfg(target_arch = "aarch64")]
4616    int_test!(
4617        test_l2_i8_aarch64_neon,
4618        i8,
4619        L2,
4620        reference::reference_squared_l2_i8_mathematical,
4621        0x3e8bada709e176be,
4622        320,
4623        { Neon::new_checked() }
4624    );
4625
4626    int_test!(
4627        test_ip_i8_current,
4628        i8,
4629        IP,
4630        reference::reference_innerproduct_i8_mathematical,
4631        0xe8306104740509e1,
4632        64,
4633        { Some(ARCH) }
4634    );
4635
4636    int_test!(
4637        test_ip_i8_scalar,
4638        i8,
4639        IP,
4640        reference::reference_innerproduct_i8_mathematical,
4641        0x8a263408c7b31d85,
4642        64,
4643        { Some(diskann_wide::arch::Scalar) }
4644    );
4645
4646    #[cfg(target_arch = "x86_64")]
4647    int_test!(
4648        test_ip_i8_x86_64_v3,
4649        i8,
4650        IP,
4651        reference::reference_innerproduct_i8_mathematical,
4652        0x8a263408c7b31d85,
4653        256,
4654        { V3::new_checked() }
4655    );
4656
4657    #[cfg(target_arch = "x86_64")]
4658    int_test!(
4659        test_ip_i8_x86_64_v4,
4660        i8,
4661        IP,
4662        reference::reference_innerproduct_i8_mathematical,
4663        0x8a263408c7b31d85,
4664        320,
4665        { V4::new_checked_miri() }
4666    );
4667
4668    #[cfg(target_arch = "aarch64")]
4669    int_test!(
4670        test_ip_i8_aarch64_neon,
4671        i8,
4672        IP,
4673        reference::reference_innerproduct_i8_mathematical,
4674        0x8a263408c7b31d85,
4675        320,
4676        { Neon::new_checked() }
4677    );
4678
4679    int_test!(
4680        test_cosine_i8_current,
4681        i8,
4682        CosineStateless,
4683        reference::reference_cosine_i8_mathematical,
4684        0x818c210190701e4b,
4685        64,
4686        { Some(ARCH) }
4687    );
4688
4689    int_test!(
4690        test_cosine_i8_scalar,
4691        i8,
4692        CosineStateless,
4693        reference::reference_cosine_i8_mathematical,
4694        0x2d077bed2629b18e,
4695        64,
4696        { Some(diskann_wide::arch::Scalar) }
4697    );
4698
4699    #[cfg(target_arch = "x86_64")]
4700    int_test!(
4701        test_cosine_i8_x86_64_v3,
4702        i8,
4703        CosineStateless,
4704        reference::reference_cosine_i8_mathematical,
4705        0x2d077bed2629b18e,
4706        256,
4707        { V3::new_checked() }
4708    );
4709
4710    #[cfg(target_arch = "x86_64")]
4711    int_test!(
4712        test_cosine_i8_x86_64_v4,
4713        i8,
4714        CosineStateless,
4715        reference::reference_cosine_i8_mathematical,
4716        0x2d077bed2629b18e,
4717        320,
4718        { V4::new_checked_miri() }
4719    );
4720
4721    #[cfg(target_arch = "aarch64")]
4722    int_test!(
4723        test_cosine_i8_aarch64_neon,
4724        i8,
4725        CosineStateless,
4726        reference::reference_cosine_i8_mathematical,
4727        0x2d077bed2629b18e,
4728        320,
4729        { Neon::new_checked() }
4730    );
4731
4732    //////////
4733    // LInf //
4734    //////////
4735
4736    macro_rules! linf_test {
4737        ($name:ident,
4738         $T:ty,
4739         $reference:path,
4740         $eps:literal,
4741         $relative:literal,
4742         $seed:literal,
4743         $upper:literal,
4744         $($arch:tt)*
4745        ) => {
4746            #[test]
4747            fn $name() {
4748                if let Some(arch) = $($arch)* {
4749                    let mut rng = StdRng::seed_from_u64($seed);
4750                    for dim in 0..$upper {
4751                        stress_test_linf::<$T, _, StdRng, _>(
4752                            arch,
4753                            |l| $reference(l).into_inner(),
4754                            rand_distr::Normal::new(-10.0, 10.0).unwrap(),
4755                            dim,
4756                            $eps,
4757                            $relative,
4758                            &mut rng,
4759                        );
4760                    }
4761                }
4762            }
4763        }
4764    }
4765
4766    linf_test!(
4767        test_linf_f32_scalar,
4768        f32,
4769        reference::reference_linf_f32_mathematical,
4770        1e-6,
4771        1e-6,
4772        0xf149c2bcde660128,
4773        256,
4774        Some(Scalar::new())
4775    );
4776
4777    #[cfg(target_arch = "x86_64")]
4778    linf_test!(
4779        test_linf_f32_v3,
4780        f32,
4781        reference::reference_linf_f32_mathematical,
4782        1e-6,
4783        1e-6,
4784        0xf149c2bcde660128,
4785        256,
4786        V3::new_checked()
4787    );
4788
4789    #[cfg(target_arch = "x86_64")]
4790    linf_test!(
4791        test_linf_f32_v4,
4792        f32,
4793        reference::reference_linf_f32_mathematical,
4794        1e-6,
4795        1e-6,
4796        0xf149c2bcde660128,
4797        256,
4798        V4::new_checked_miri()
4799    );
4800
4801    #[cfg(target_arch = "aarch64")]
4802    linf_test!(
4803        test_linf_f32_neon,
4804        f32,
4805        reference::reference_linf_f32_mathematical,
4806        1e-6,
4807        1e-6,
4808        0xf149c2bcde660128,
4809        256,
4810        Neon::new_checked()
4811    );
4812
4813    linf_test!(
4814        test_linf_f16_scalar,
4815        f16,
4816        reference::reference_linf_f16_mathematical,
4817        1e-6,
4818        1e-6,
4819        0xf149c2bcde660128,
4820        256,
4821        Some(Scalar::new())
4822    );
4823
4824    #[cfg(target_arch = "x86_64")]
4825    linf_test!(
4826        test_linf_f16_v3,
4827        f16,
4828        reference::reference_linf_f16_mathematical,
4829        1e-6,
4830        1e-6,
4831        0xf149c2bcde660128,
4832        256,
4833        V3::new_checked()
4834    );
4835
4836    #[cfg(target_arch = "x86_64")]
4837    linf_test!(
4838        test_linf_f16_v4,
4839        f16,
4840        reference::reference_linf_f16_mathematical,
4841        1e-6,
4842        1e-6,
4843        0xf149c2bcde660128,
4844        256,
4845        V4::new_checked_miri()
4846    );
4847
4848    #[cfg(target_arch = "aarch64")]
4849    linf_test!(
4850        test_linf_f16_neon,
4851        f16,
4852        reference::reference_linf_f16_mathematical,
4853        1e-6,
4854        1e-6,
4855        0xf149c2bcde660128,
4856        256,
4857        Neon::new_checked()
4858    );
4859
4860    ////////////////
4861    // Miri Tests //
4862    ////////////////
4863
4864    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4865    enum DataType {
4866        Float32,
4867        Float16,
4868        UInt8,
4869        Int8,
4870    }
4871
4872    trait AsDataType {
4873        fn as_data_type() -> DataType;
4874    }
4875
4876    impl AsDataType for f32 {
4877        fn as_data_type() -> DataType {
4878            DataType::Float32
4879        }
4880    }
4881
4882    impl AsDataType for f16 {
4883        fn as_data_type() -> DataType {
4884            DataType::Float16
4885        }
4886    }
4887
4888    impl AsDataType for u8 {
4889        fn as_data_type() -> DataType {
4890            DataType::UInt8
4891        }
4892    }
4893
4894    impl AsDataType for i8 {
4895        fn as_data_type() -> DataType {
4896            DataType::Int8
4897        }
4898    }
4899
4900    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4901    enum Arch {
4902        Scalar,
4903        #[expect(non_camel_case_types)]
4904        X86_64_V3,
4905        #[expect(non_camel_case_types)]
4906        X86_64_V4,
4907        Aarch64Neon,
4908    }
4909
4910    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4911    struct Key {
4912        arch: Arch,
4913        left: DataType,
4914        right: DataType,
4915    }
4916
4917    impl Key {
4918        fn new(arch: Arch, left: DataType, right: DataType) -> Self {
4919            Self { arch, left, right }
4920        }
4921    }
4922
4923    static MIRI_BOUNDS: LazyLock<HashMap<Key, usize>> = LazyLock::new(|| {
4924        use Arch::{Aarch64Neon, Scalar, X86_64_V3, X86_64_V4};
4925        use DataType::{Float16, Float32, Int8, UInt8};
4926
4927        [
4928            (Key::new(Scalar, Float32, Float32), 64),
4929            (Key::new(X86_64_V3, Float32, Float32), 256),
4930            (Key::new(X86_64_V4, Float32, Float32), 256),
4931            (Key::new(Aarch64Neon, Float32, Float32), 128),
4932            (Key::new(Scalar, Float16, Float16), 64),
4933            (Key::new(X86_64_V3, Float16, Float16), 256),
4934            (Key::new(X86_64_V4, Float16, Float16), 256),
4935            (Key::new(Aarch64Neon, Float16, Float16), 128),
4936            (Key::new(Scalar, Float32, Float16), 64),
4937            (Key::new(X86_64_V3, Float32, Float16), 256),
4938            (Key::new(X86_64_V4, Float32, Float16), 256),
4939            (Key::new(Aarch64Neon, Float32, Float16), 128),
4940            (Key::new(Scalar, UInt8, UInt8), 64),
4941            (Key::new(X86_64_V3, UInt8, UInt8), 256),
4942            (Key::new(X86_64_V4, UInt8, UInt8), 320),
4943            (Key::new(Aarch64Neon, UInt8, UInt8), 128),
4944            (Key::new(Scalar, Int8, Int8), 64),
4945            (Key::new(X86_64_V3, Int8, Int8), 256),
4946            (Key::new(X86_64_V4, Int8, Int8), 320),
4947            (Key::new(Aarch64Neon, Int8, Int8), 128),
4948        ]
4949        .into_iter()
4950        .collect()
4951    });
4952
4953    macro_rules! test_bounds {
4954        (
4955            $function:ident,
4956            $left:ty,
4957            $left_ex:expr,
4958            $right:ty,
4959            $right_ex:expr
4960        ) => {
4961            #[test]
4962            fn $function() {
4963                let left: $left = $left_ex;
4964                let right: $right = $right_ex;
4965
4966                let left_type = <$left>::as_data_type();
4967                let right_type = <$right>::as_data_type();
4968
4969                // Scalar
4970                {
4971                    let max = MIRI_BOUNDS[&Key::new(Arch::Scalar, left_type, right_type)];
4972                    for dim in 0..max {
4973                        let left: Vec<$left> = vec![left; dim];
4974                        let right: Vec<$right> = vec![right; dim];
4975
4976                        let arch = diskann_wide::arch::Scalar;
4977                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
4978                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
4979                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
4980
4981                        let left = unaligned::Buffer::new(&left);
4982                        let right = unaligned::Buffer::new(&right);
4983                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
4984                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
4985                        simd_op(
4986                            &CosineStateless,
4987                            arch,
4988                            left.as_unaligned(),
4989                            right.as_unaligned(),
4990                        );
4991                    }
4992                }
4993
4994                #[cfg(target_arch = "x86_64")]
4995                if let Some(arch) = V3::new_checked() {
4996                    let max = MIRI_BOUNDS[&Key::new(Arch::X86_64_V3, left_type, right_type)];
4997                    for dim in 0..max {
4998                        let left: Vec<$left> = vec![left; dim];
4999                        let right: Vec<$right> = vec![right; dim];
5000
5001                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5002                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5003                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5004
5005                        let left = unaligned::Buffer::new(&left);
5006                        let right = unaligned::Buffer::new(&right);
5007                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5008                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5009                        simd_op(
5010                            &CosineStateless,
5011                            arch,
5012                            left.as_unaligned(),
5013                            right.as_unaligned(),
5014                        );
5015                    }
5016                }
5017
5018                #[cfg(target_arch = "x86_64")]
5019                if let Some(arch) = V4::new_checked_miri() {
5020                    let max = MIRI_BOUNDS[&Key::new(Arch::X86_64_V4, left_type, right_type)];
5021                    for dim in 0..max {
5022                        let left: Vec<$left> = vec![left; dim];
5023                        let right: Vec<$right> = vec![right; dim];
5024
5025                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5026                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5027                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5028
5029                        let left = unaligned::Buffer::new(&left);
5030                        let right = unaligned::Buffer::new(&right);
5031                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5032                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5033                        simd_op(
5034                            &CosineStateless,
5035                            arch,
5036                            left.as_unaligned(),
5037                            right.as_unaligned(),
5038                        );
5039                    }
5040                }
5041
5042                #[cfg(target_arch = "aarch64")]
5043                if let Some(arch) = Neon::new_checked() {
5044                    let max = MIRI_BOUNDS[&Key::new(Arch::Aarch64Neon, left_type, right_type)];
5045                    for dim in 0..max {
5046                        let left: Vec<$left> = vec![left; dim];
5047                        let right: Vec<$right> = vec![right; dim];
5048
5049                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5050                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5051                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5052                        let left = unaligned::Buffer::new(&left);
5053                        let right = unaligned::Buffer::new(&right);
5054                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5055                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5056                        simd_op(
5057                            &CosineStateless,
5058                            arch,
5059                            left.as_unaligned(),
5060                            right.as_unaligned(),
5061                        );
5062                    }
5063                }
5064            }
5065        };
5066    }
5067
5068    test_bounds!(miri_test_bounds_f32xf32, f32, 1.0f32, f32, 2.0f32);
5069    test_bounds!(
5070        miri_test_bounds_f16xf16,
5071        f16,
5072        diskann_wide::cast_f32_to_f16(1.0f32),
5073        f16,
5074        diskann_wide::cast_f32_to_f16(2.0f32)
5075    );
5076    test_bounds!(
5077        miri_test_bounds_f32xf16,
5078        f32,
5079        1.0f32,
5080        f16,
5081        diskann_wide::cast_f32_to_f16(2.0f32)
5082    );
5083    test_bounds!(miri_test_bounds_u8xu8, u8, 1u8, u8, 1u8);
5084    test_bounds!(miri_test_bounds_i8xi8, i8, 1i8, i8, 1i8);
5085}