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::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#[expect(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>::u32x4;
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        // Signed absolute differences occupy unsigned byte lanes (0..=255).
1249        // Squaring them with unsigned dot product preserves the full result.
1250        let diff = acc.arch().vabdq_s8(x, y);
1251        acc.dot_simd(diff, diff)
1252    }
1253
1254    #[inline(always)]
1255    unsafe fn epilogue(
1256        &self,
1257        arch: Neon,
1258        x: *const i8,
1259        y: *const i8,
1260        len: usize,
1261        acc: Self::Accumulator,
1262    ) -> Self::Accumulator {
1263        let scalar = scalar_epilogue(
1264            x,
1265            y,
1266            len.min(Self::SIMDWidth::value() - 1),
1267            0u32,
1268            |acc, x: i8, y: i8| -> u32 {
1269                let c = (x as i32) - (y as i32);
1270                acc + (c * c) as u32
1271            },
1272        );
1273        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
1274    }
1275
1276    // Perform a final reduction.
1277    #[inline(always)]
1278    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1279        x.sum_tree().as_f32_lossy()
1280    }
1281}
1282
1283impl SIMDSchema<i8, i8, Scalar> for L2 {
1284    type SIMDWidth = Const<4>;
1285    type Accumulator = Emulated<i32, 4>;
1286    type Left = Emulated<i8, 4>;
1287    type Right = Emulated<i8, 4>;
1288    type Return = f32;
1289    type Main = Strategy1x1;
1290
1291    #[inline(always)]
1292    fn init(&self, arch: Scalar) -> Self::Accumulator {
1293        Self::Accumulator::default(arch)
1294    }
1295
1296    #[inline(always)]
1297    fn accumulate(
1298        &self,
1299        x: Self::Left,
1300        y: Self::Right,
1301        acc: Self::Accumulator,
1302    ) -> Self::Accumulator {
1303        let x: Self::Accumulator = x.into();
1304        let y: Self::Accumulator = y.into();
1305        let c = x - y;
1306        acc + c * c
1307    }
1308
1309    // Perform a final reduction.
1310    #[inline(always)]
1311    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1312        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
1313    }
1314
1315    #[inline(always)]
1316    unsafe fn epilogue(
1317        &self,
1318        arch: Scalar,
1319        x: *const i8,
1320        y: *const i8,
1321        len: usize,
1322        acc: Self::Accumulator,
1323    ) -> Self::Accumulator {
1324        let mut s: i32 = 0;
1325        for i in 0..len {
1326            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1327            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1328            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1329            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
1330            let d = vx - vy;
1331            s += d * d;
1332        }
1333        acc + Self::Accumulator::from_array(arch, [s, 0, 0, 0])
1334    }
1335}
1336
1337#[cfg(target_arch = "x86_64")]
1338impl SIMDSchema<u8, u8, V4> for L2 {
1339    type SIMDWidth = Const<32>;
1340    type Accumulator = <V4 as Architecture>::i32x16;
1341    type Left = <V4 as Architecture>::u8x32;
1342    type Right = <V4 as Architecture>::u8x32;
1343    type Return = f32;
1344    type Main = Strategy4x1;
1345
1346    #[inline(always)]
1347    fn init(&self, arch: V4) -> Self::Accumulator {
1348        Self::Accumulator::default(arch)
1349    }
1350
1351    #[inline(always)]
1352    fn accumulate(
1353        &self,
1354        x: Self::Left,
1355        y: Self::Right,
1356        acc: Self::Accumulator,
1357    ) -> Self::Accumulator {
1358        diskann_wide::alias!(i16s = <V4>::i16x32);
1359
1360        let x: i16s = x.into();
1361        let y: i16s = y.into();
1362        let c = x - y;
1363        acc.dot_simd(c, c)
1364    }
1365
1366    #[inline(always)]
1367    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1368        x.sum_tree().as_f32_lossy()
1369    }
1370}
1371
1372#[cfg(target_arch = "x86_64")]
1373impl SIMDSchema<u8, u8, V3> for L2 {
1374    type SIMDWidth = Const<16>;
1375    type Accumulator = <V3 as Architecture>::i32x8;
1376    type Left = <V3 as Architecture>::u8x16;
1377    type Right = <V3 as Architecture>::u8x16;
1378    type Return = f32;
1379    type Main = Strategy4x1;
1380
1381    #[inline(always)]
1382    fn init(&self, arch: V3) -> Self::Accumulator {
1383        Self::Accumulator::default(arch)
1384    }
1385
1386    #[inline(always)]
1387    fn accumulate(
1388        &self,
1389        x: Self::Left,
1390        y: Self::Right,
1391        acc: Self::Accumulator,
1392    ) -> Self::Accumulator {
1393        diskann_wide::alias!(i16s = <V3>::i16x16);
1394
1395        let x: i16s = x.into();
1396        let y: i16s = y.into();
1397        let c = x - y;
1398        acc.dot_simd(c, c)
1399    }
1400
1401    // Perform a final reduction.
1402    #[inline(always)]
1403    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1404        x.sum_tree().as_f32_lossy()
1405    }
1406}
1407
1408#[cfg(target_arch = "aarch64")]
1409impl SIMDSchema<u8, u8, Neon> for L2 {
1410    type SIMDWidth = Const<16>;
1411    type Accumulator = <Neon as Architecture>::u32x4;
1412    type Left = diskann_wide::arch::aarch64::u8x16;
1413    type Right = diskann_wide::arch::aarch64::u8x16;
1414    type Return = f32;
1415    type Main = Strategy2x1;
1416
1417    #[inline(always)]
1418    fn init(&self, arch: Neon) -> Self::Accumulator {
1419        Self::Accumulator::default(arch)
1420    }
1421
1422    #[inline(always)]
1423    fn accumulate(
1424        &self,
1425        x: Self::Left,
1426        y: Self::Right,
1427        acc: Self::Accumulator,
1428    ) -> Self::Accumulator {
1429        let diff = acc.arch().vabdq_u8(x, y);
1430        acc.dot_simd(diff, diff)
1431    }
1432
1433    #[inline(always)]
1434    unsafe fn epilogue(
1435        &self,
1436        arch: Neon,
1437        x: *const u8,
1438        y: *const u8,
1439        len: usize,
1440        acc: Self::Accumulator,
1441    ) -> Self::Accumulator {
1442        let scalar = scalar_epilogue(
1443            x,
1444            y,
1445            len.min(Self::SIMDWidth::value() - 1),
1446            0u32,
1447            |acc, x: u8, y: u8| -> u32 {
1448                let c = (x as i32) - (y as i32);
1449                acc + ((c * c) as u32)
1450            },
1451        );
1452        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
1453    }
1454
1455    // Perform a final reduction.
1456    #[inline(always)]
1457    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1458        x.sum_tree().as_f32_lossy()
1459    }
1460}
1461
1462impl SIMDSchema<u8, u8, Scalar> for L2 {
1463    type SIMDWidth = Const<4>;
1464    type Accumulator = Emulated<i32, 4>;
1465    type Left = Emulated<u8, 4>;
1466    type Right = Emulated<u8, 4>;
1467    type Return = f32;
1468    type Main = Strategy1x1;
1469
1470    #[inline(always)]
1471    fn init(&self, arch: Scalar) -> Self::Accumulator {
1472        Self::Accumulator::default(arch)
1473    }
1474
1475    #[inline(always)]
1476    fn accumulate(
1477        &self,
1478        x: Self::Left,
1479        y: Self::Right,
1480        acc: Self::Accumulator,
1481    ) -> Self::Accumulator {
1482        let x: Self::Accumulator = x.into();
1483        let y: Self::Accumulator = y.into();
1484        let c = x - y;
1485        acc + c * c
1486    }
1487
1488    // Perform a final reduction.
1489    #[inline(always)]
1490    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1491        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
1492    }
1493
1494    #[inline(always)]
1495    unsafe fn epilogue(
1496        &self,
1497        arch: Scalar,
1498        x: *const u8,
1499        y: *const u8,
1500        len: usize,
1501        acc: Self::Accumulator,
1502    ) -> Self::Accumulator {
1503        let mut s: i32 = 0;
1504        for i in 0..len {
1505            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1506            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1507            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1508            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
1509            let d = vx - vy;
1510            s += d * d;
1511        }
1512        acc + Self::Accumulator::from_array(arch, [s, 0, 0, 0])
1513    }
1514}
1515
1516// A L2 distance function that defers a final reduction, allowing for a distance
1517// computation to take place across multiple slice pairs.
1518#[derive(Clone, Copy, Debug)]
1519pub struct ResumableL2<A = diskann_wide::arch::Current>
1520where
1521    A: Architecture,
1522    L2: SIMDSchema<f32, f32, A>,
1523{
1524    acc: <L2 as SIMDSchema<f32, f32, A>>::Accumulator,
1525}
1526
1527impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableL2<A>
1528where
1529    A: Architecture,
1530    L2: SIMDSchema<f32, f32, A, Return = f32>,
1531{
1532    type NonResumable = L2;
1533    type FinalReturn = f32;
1534
1535    #[inline(always)]
1536    fn init(arch: A) -> Self {
1537        Self { acc: L2.init(arch) }
1538    }
1539
1540    #[inline(always)]
1541    fn combine_with(&self, other: <L2 as SIMDSchema<f32, f32, A>>::Accumulator) -> Self {
1542        Self {
1543            acc: self.acc + other,
1544        }
1545    }
1546
1547    #[inline(always)]
1548    fn sum(&self) -> f32 {
1549        L2.reduce(self.acc)
1550    }
1551}
1552
1553/////
1554///// IP Implementations
1555/////
1556
1557// A pure IP distance function that provides a final reduction.
1558#[derive(Clone, Copy, Debug, Default)]
1559pub struct IP;
1560
1561#[cfg(target_arch = "x86_64")]
1562impl SIMDSchema<f32, f32, V4> for IP {
1563    type SIMDWidth = Const<8>;
1564    type Accumulator = <V4 as Architecture>::f32x8;
1565    type Left = <V4 as Architecture>::f32x8;
1566    type Right = <V4 as Architecture>::f32x8;
1567    type Return = f32;
1568    type Main = Strategy4x1;
1569
1570    #[inline(always)]
1571    fn init(&self, arch: V4) -> Self::Accumulator {
1572        Self::Accumulator::default(arch)
1573    }
1574
1575    #[inline(always)]
1576    fn accumulate(
1577        &self,
1578        x: Self::Left,
1579        y: Self::Right,
1580        acc: Self::Accumulator,
1581    ) -> Self::Accumulator {
1582        x.mul_add_simd(y, acc)
1583    }
1584
1585    #[inline(always)]
1586    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1587        x.sum_tree()
1588    }
1589}
1590
1591#[cfg(target_arch = "x86_64")]
1592impl SIMDSchema<f32, f32, V3> for IP {
1593    type SIMDWidth = Const<8>;
1594    type Accumulator = <V3 as Architecture>::f32x8;
1595    type Left = <V3 as Architecture>::f32x8;
1596    type Right = <V3 as Architecture>::f32x8;
1597    type Return = f32;
1598    type Main = Strategy4x1;
1599
1600    #[inline(always)]
1601    fn init(&self, arch: V3) -> Self::Accumulator {
1602        Self::Accumulator::default(arch)
1603    }
1604
1605    #[inline(always)]
1606    fn accumulate(
1607        &self,
1608        x: Self::Left,
1609        y: Self::Right,
1610        acc: Self::Accumulator,
1611    ) -> Self::Accumulator {
1612        x.mul_add_simd(y, acc)
1613    }
1614
1615    // Perform a final reduction.
1616    #[inline(always)]
1617    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1618        x.sum_tree()
1619    }
1620}
1621
1622#[cfg(target_arch = "aarch64")]
1623impl SIMDSchema<f32, f32, Neon> for IP {
1624    type SIMDWidth = Const<4>;
1625    type Accumulator = <Neon as Architecture>::f32x4;
1626    type Left = <Neon as Architecture>::f32x4;
1627    type Right = <Neon as Architecture>::f32x4;
1628    type Return = f32;
1629    type Main = Strategy4x1;
1630
1631    #[inline(always)]
1632    fn init(&self, arch: Neon) -> Self::Accumulator {
1633        Self::Accumulator::default(arch)
1634    }
1635
1636    #[inline(always)]
1637    fn accumulate(
1638        &self,
1639        x: Self::Left,
1640        y: Self::Right,
1641        acc: Self::Accumulator,
1642    ) -> Self::Accumulator {
1643        x.mul_add_simd(y, acc)
1644    }
1645
1646    #[inline(always)]
1647    unsafe fn epilogue(
1648        &self,
1649        arch: Neon,
1650        x: *const f32,
1651        y: *const f32,
1652        len: usize,
1653        acc: Self::Accumulator,
1654    ) -> Self::Accumulator {
1655        let scalar = scalar_epilogue(
1656            x,
1657            y,
1658            len.min(Self::SIMDWidth::value() - 1),
1659            0.0f32,
1660            |acc, x: f32, y: f32| -> f32 { x.mul_add(y, acc) },
1661        );
1662        acc + Self::Accumulator::from_array(arch, [scalar, 0.0, 0.0, 0.0])
1663    }
1664
1665    #[inline(always)]
1666    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1667        x.sum_tree()
1668    }
1669}
1670
1671impl SIMDSchema<f32, f32, Scalar> for IP {
1672    type SIMDWidth = Const<4>;
1673    type Accumulator = Emulated<f32, 4>;
1674    type Left = Emulated<f32, 4>;
1675    type Right = Emulated<f32, 4>;
1676    type Return = f32;
1677    type Main = Strategy2x1;
1678
1679    #[inline(always)]
1680    fn init(&self, arch: Scalar) -> Self::Accumulator {
1681        Self::Accumulator::default(arch)
1682    }
1683
1684    #[inline(always)]
1685    fn accumulate(
1686        &self,
1687        x: Self::Left,
1688        y: Self::Right,
1689        acc: Self::Accumulator,
1690    ) -> Self::Accumulator {
1691        x * y + acc
1692    }
1693
1694    // Perform a final reduction.
1695    #[inline(always)]
1696    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1697        x.sum_tree()
1698    }
1699
1700    #[inline(always)]
1701    unsafe fn epilogue(
1702        &self,
1703        arch: Scalar,
1704        x: *const f32,
1705        y: *const f32,
1706        len: usize,
1707        acc: Self::Accumulator,
1708    ) -> Self::Accumulator {
1709        let mut s: f32 = 0.0;
1710        for i in 0..len {
1711            // SAFETY: The range `[x, x.add(len))` is valid for reads.
1712            let vx = unsafe { x.add(i).read_unaligned() };
1713            // SAFETY: The range `[y, y.add(len))` is valid for reads.
1714            let vy = unsafe { y.add(i).read_unaligned() };
1715            s += vx * vy;
1716        }
1717        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
1718    }
1719}
1720
1721#[cfg(target_arch = "x86_64")]
1722impl SIMDSchema<Half, Half, V4> for IP {
1723    type SIMDWidth = Const<8>;
1724    type Accumulator = <V4 as Architecture>::f32x8;
1725    type Left = <V4 as Architecture>::f16x8;
1726    type Right = <V4 as Architecture>::f16x8;
1727    type Return = f32;
1728    type Main = Strategy4x1;
1729
1730    #[inline(always)]
1731    fn init(&self, arch: V4) -> Self::Accumulator {
1732        Self::Accumulator::default(arch)
1733    }
1734
1735    #[inline(always)]
1736    fn accumulate(
1737        &self,
1738        x: Self::Left,
1739        y: Self::Right,
1740        acc: Self::Accumulator,
1741    ) -> Self::Accumulator {
1742        diskann_wide::alias!(f32s = <V4>::f32x8);
1743
1744        let x: f32s = x.into();
1745        let y: f32s = y.into();
1746        x.mul_add_simd(y, acc)
1747    }
1748
1749    #[inline(always)]
1750    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1751        x.sum_tree()
1752    }
1753}
1754
1755#[cfg(target_arch = "x86_64")]
1756impl SIMDSchema<Half, Half, V3> for IP {
1757    type SIMDWidth = Const<8>;
1758    type Accumulator = <V3 as Architecture>::f32x8;
1759    type Left = <V3 as Architecture>::f16x8;
1760    type Right = <V3 as Architecture>::f16x8;
1761    type Return = f32;
1762    type Main = Strategy2x4;
1763
1764    #[inline(always)]
1765    fn init(&self, arch: V3) -> Self::Accumulator {
1766        Self::Accumulator::default(arch)
1767    }
1768
1769    #[inline(always)]
1770    fn accumulate(
1771        &self,
1772        x: Self::Left,
1773        y: Self::Right,
1774        acc: Self::Accumulator,
1775    ) -> Self::Accumulator {
1776        diskann_wide::alias!(f32s = <V3>::f32x8);
1777
1778        let x: f32s = x.into();
1779        let y: f32s = y.into();
1780        x.mul_add_simd(y, acc)
1781    }
1782
1783    // Perform a final reduction.
1784    #[inline(always)]
1785    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1786        x.sum_tree()
1787    }
1788}
1789
1790#[cfg(target_arch = "aarch64")]
1791impl SIMDSchema<Half, Half, Neon> for IP {
1792    type SIMDWidth = Const<4>;
1793    type Accumulator = <Neon as Architecture>::f32x4;
1794    type Left = diskann_wide::arch::aarch64::f16x4;
1795    type Right = diskann_wide::arch::aarch64::f16x4;
1796    type Return = f32;
1797    type Main = Strategy4x1;
1798
1799    #[inline(always)]
1800    fn init(&self, arch: Neon) -> Self::Accumulator {
1801        Self::Accumulator::default(arch)
1802    }
1803
1804    #[inline(always)]
1805    fn accumulate(
1806        &self,
1807        x: Self::Left,
1808        y: Self::Right,
1809        acc: Self::Accumulator,
1810    ) -> Self::Accumulator {
1811        acc.dot_simd(x, y)
1812    }
1813
1814    #[inline(always)]
1815    unsafe fn epilogue(
1816        &self,
1817        arch: Neon,
1818        x: *const Half,
1819        y: *const Half,
1820        len: usize,
1821        acc: Self::Accumulator,
1822    ) -> Self::Accumulator {
1823        diskann_wide::alias!(f32s = <Neon>::f32x4);
1824
1825        let rest = scalar_epilogue(
1826            x,
1827            y,
1828            len.min(Self::SIMDWidth::value() - 1),
1829            f32s::default(arch),
1830            |acc, x: Half, y: Half| -> f32s {
1831                let zero = Half::default();
1832                let x: f32s = Self::Left::from_array(arch, [x, zero, zero, zero]).into();
1833                let y: f32s = Self::Right::from_array(arch, [y, zero, zero, zero]).into();
1834                x.mul_add_simd(y, acc)
1835            },
1836        );
1837        acc + rest
1838    }
1839
1840    #[inline(always)]
1841    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1842        x.sum_tree()
1843    }
1844}
1845
1846impl SIMDSchema<Half, Half, Scalar> for IP {
1847    type SIMDWidth = Const<1>;
1848    type Accumulator = Emulated<f32, 1>;
1849    type Left = Emulated<Half, 1>;
1850    type Right = Emulated<Half, 1>;
1851    type Return = f32;
1852    type Main = Strategy1x1;
1853
1854    #[inline(always)]
1855    fn init(&self, arch: Scalar) -> Self::Accumulator {
1856        Self::Accumulator::default(arch)
1857    }
1858
1859    #[inline(always)]
1860    fn accumulate(
1861        &self,
1862        x: Self::Left,
1863        y: Self::Right,
1864        acc: Self::Accumulator,
1865    ) -> Self::Accumulator {
1866        let x: Self::Accumulator = x.into();
1867        let y: Self::Accumulator = y.into();
1868        x * y + acc
1869    }
1870
1871    #[inline(always)]
1872    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1873        x.to_array()[0]
1874    }
1875}
1876
1877impl<A> SIMDSchema<f32, Half, A> for IP
1878where
1879    A: Architecture,
1880{
1881    type SIMDWidth = Const<8>;
1882    type Accumulator = A::f32x8;
1883    type Left = A::f32x8;
1884    type Right = A::f16x8;
1885    type Return = f32;
1886    type Main = Strategy4x2;
1887
1888    #[inline(always)]
1889    fn init(&self, arch: A) -> Self::Accumulator {
1890        Self::Accumulator::default(arch)
1891    }
1892
1893    #[inline(always)]
1894    fn accumulate(
1895        &self,
1896        x: Self::Left,
1897        y: Self::Right,
1898        acc: Self::Accumulator,
1899    ) -> Self::Accumulator {
1900        let y: A::f32x8 = y.into();
1901        x.mul_add_simd(y, acc)
1902    }
1903
1904    // Perform a final reduction.
1905    #[inline(always)]
1906    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1907        x.sum_tree()
1908    }
1909}
1910
1911#[cfg(target_arch = "x86_64")]
1912impl SIMDSchema<i8, i8, V4> for IP {
1913    type SIMDWidth = Const<32>;
1914    type Accumulator = <V4 as Architecture>::i32x16;
1915    type Left = <V4 as Architecture>::i8x32;
1916    type Right = <V4 as Architecture>::i8x32;
1917    type Return = f32;
1918    type Main = Strategy4x1;
1919
1920    #[inline(always)]
1921    fn init(&self, arch: V4) -> Self::Accumulator {
1922        Self::Accumulator::default(arch)
1923    }
1924
1925    #[inline(always)]
1926    fn accumulate(
1927        &self,
1928        x: Self::Left,
1929        y: Self::Right,
1930        acc: Self::Accumulator,
1931    ) -> Self::Accumulator {
1932        diskann_wide::alias!(i16s = <V4>::i16x32);
1933
1934        let x: i16s = x.into();
1935        let y: i16s = y.into();
1936        acc.dot_simd(x, y)
1937    }
1938
1939    #[inline(always)]
1940    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1941        x.sum_tree().as_f32_lossy()
1942    }
1943}
1944
1945#[cfg(target_arch = "x86_64")]
1946impl SIMDSchema<i8, i8, V3> for IP {
1947    type SIMDWidth = Const<16>;
1948    type Accumulator = <V3 as Architecture>::i32x8;
1949    type Left = <V3 as Architecture>::i8x16;
1950    type Right = <V3 as Architecture>::i8x16;
1951    type Return = f32;
1952    type Main = Strategy4x1;
1953
1954    #[inline(always)]
1955    fn init(&self, arch: V3) -> Self::Accumulator {
1956        Self::Accumulator::default(arch)
1957    }
1958
1959    #[inline(always)]
1960    fn accumulate(
1961        &self,
1962        x: Self::Left,
1963        y: Self::Right,
1964        acc: Self::Accumulator,
1965    ) -> Self::Accumulator {
1966        diskann_wide::alias!(i16s = <V3>::i16x16);
1967
1968        let x: i16s = x.into();
1969        let y: i16s = y.into();
1970        acc.dot_simd(x, y)
1971    }
1972
1973    // Perform a final reduction.
1974    #[inline(always)]
1975    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
1976        x.sum_tree().as_f32_lossy()
1977    }
1978}
1979
1980#[cfg(target_arch = "aarch64")]
1981impl SIMDSchema<i8, i8, Neon> for IP {
1982    type SIMDWidth = Const<16>;
1983    type Accumulator = <Neon as Architecture>::i32x4;
1984    type Left = <Neon as Architecture>::i8x16;
1985    type Right = <Neon as Architecture>::i8x16;
1986    type Return = f32;
1987    type Main = Strategy2x1;
1988
1989    #[inline(always)]
1990    fn init(&self, arch: Neon) -> Self::Accumulator {
1991        Self::Accumulator::default(arch)
1992    }
1993
1994    #[inline(always)]
1995    fn accumulate(
1996        &self,
1997        x: Self::Left,
1998        y: Self::Right,
1999        acc: Self::Accumulator,
2000    ) -> Self::Accumulator {
2001        acc.dot_simd(x, y)
2002    }
2003
2004    #[inline(always)]
2005    unsafe fn epilogue(
2006        &self,
2007        arch: Neon,
2008        x: *const i8,
2009        y: *const i8,
2010        len: usize,
2011        acc: Self::Accumulator,
2012    ) -> Self::Accumulator {
2013        let scalar = scalar_epilogue(
2014            x,
2015            y,
2016            len.min(Self::SIMDWidth::value() - 1),
2017            0i32,
2018            |acc, x: i8, y: i8| -> i32 { acc + (x as i32) * (y as i32) },
2019        );
2020        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
2021    }
2022
2023    #[inline(always)]
2024    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2025        x.sum_tree().as_f32_lossy()
2026    }
2027}
2028
2029impl SIMDSchema<i8, i8, Scalar> for IP {
2030    type SIMDWidth = Const<1>;
2031    type Accumulator = Emulated<i32, 1>;
2032    type Left = Emulated<i8, 1>;
2033    type Right = Emulated<i8, 1>;
2034    type Return = f32;
2035    type Main = Strategy1x1;
2036
2037    #[inline(always)]
2038    fn init(&self, arch: Scalar) -> Self::Accumulator {
2039        Self::Accumulator::default(arch)
2040    }
2041
2042    #[inline(always)]
2043    fn accumulate(
2044        &self,
2045        x: Self::Left,
2046        y: Self::Right,
2047        acc: Self::Accumulator,
2048    ) -> Self::Accumulator {
2049        let x: Self::Accumulator = x.into();
2050        let y: Self::Accumulator = y.into();
2051        x * y + acc
2052    }
2053
2054    // Perform a final reduction.
2055    #[inline(always)]
2056    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2057        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
2058    }
2059
2060    #[inline(always)]
2061    unsafe fn epilogue(
2062        &self,
2063        _arch: Scalar,
2064        _x: *const i8,
2065        _y: *const i8,
2066        _len: usize,
2067        _acc: Self::Accumulator,
2068    ) -> Self::Accumulator {
2069        unreachable!("The SIMD width is 1, so there should be no epilogue")
2070    }
2071}
2072
2073#[cfg(target_arch = "x86_64")]
2074impl SIMDSchema<u8, u8, V4> for IP {
2075    type SIMDWidth = Const<32>;
2076    type Accumulator = <V4 as Architecture>::i32x16;
2077    type Left = <V4 as Architecture>::u8x32;
2078    type Right = <V4 as Architecture>::u8x32;
2079    type Return = f32;
2080    type Main = Strategy4x1;
2081
2082    #[inline(always)]
2083    fn init(&self, arch: V4) -> Self::Accumulator {
2084        Self::Accumulator::default(arch)
2085    }
2086
2087    #[inline(always)]
2088    fn accumulate(
2089        &self,
2090        x: Self::Left,
2091        y: Self::Right,
2092        acc: Self::Accumulator,
2093    ) -> Self::Accumulator {
2094        diskann_wide::alias!(i16s = <V4>::i16x32);
2095
2096        let x: i16s = x.into();
2097        let y: i16s = y.into();
2098        acc.dot_simd(x, y)
2099    }
2100
2101    #[inline(always)]
2102    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2103        x.sum_tree().as_f32_lossy()
2104    }
2105}
2106
2107#[cfg(target_arch = "x86_64")]
2108impl SIMDSchema<u8, u8, V3> for IP {
2109    type SIMDWidth = Const<16>;
2110    type Accumulator = <V3 as Architecture>::i32x8;
2111    type Left = <V3 as Architecture>::u8x16;
2112    type Right = <V3 as Architecture>::u8x16;
2113    type Return = f32;
2114    type Main = Strategy4x1;
2115
2116    #[inline(always)]
2117    fn init(&self, arch: V3) -> Self::Accumulator {
2118        Self::Accumulator::default(arch)
2119    }
2120
2121    #[inline(always)]
2122    fn accumulate(
2123        &self,
2124        x: Self::Left,
2125        y: Self::Right,
2126        acc: Self::Accumulator,
2127    ) -> Self::Accumulator {
2128        diskann_wide::alias!(i16s = <V3>::i16x16);
2129
2130        // NOTE: Promotiving to `i16` rather than `u16` to hit specialized AVX2
2131        // instructions on x86 hardware.
2132        let x: i16s = x.into();
2133        let y: i16s = y.into();
2134        acc.dot_simd(x, y)
2135    }
2136
2137    // Perform a final reduction.
2138    #[inline(always)]
2139    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2140        x.sum_tree().as_f32_lossy()
2141    }
2142}
2143
2144#[cfg(target_arch = "aarch64")]
2145impl SIMDSchema<u8, u8, Neon> for IP {
2146    type SIMDWidth = Const<16>;
2147    type Accumulator = <Neon as Architecture>::u32x4;
2148    type Left = <Neon as Architecture>::u8x16;
2149    type Right = <Neon as Architecture>::u8x16;
2150    type Return = f32;
2151    type Main = Strategy2x1;
2152
2153    #[inline(always)]
2154    fn init(&self, arch: Neon) -> Self::Accumulator {
2155        Self::Accumulator::default(arch)
2156    }
2157
2158    #[inline(always)]
2159    fn accumulate(
2160        &self,
2161        x: Self::Left,
2162        y: Self::Right,
2163        acc: Self::Accumulator,
2164    ) -> Self::Accumulator {
2165        acc.dot_simd(x, y)
2166    }
2167
2168    #[inline(always)]
2169    unsafe fn epilogue(
2170        &self,
2171        arch: Neon,
2172        x: *const u8,
2173        y: *const u8,
2174        len: usize,
2175        acc: Self::Accumulator,
2176    ) -> Self::Accumulator {
2177        let scalar = scalar_epilogue(
2178            x,
2179            y,
2180            len.min(Self::SIMDWidth::value() - 1),
2181            0u32,
2182            |acc, x: u8, y: u8| -> u32 { acc + (x as u32) * (y as u32) },
2183        );
2184        acc + Self::Accumulator::from_array(arch, [scalar, 0, 0, 0])
2185    }
2186
2187    #[inline(always)]
2188    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2189        x.sum_tree().as_f32_lossy()
2190    }
2191}
2192
2193impl SIMDSchema<u8, u8, Scalar> for IP {
2194    type SIMDWidth = Const<1>;
2195    type Accumulator = Emulated<i32, 1>;
2196    type Left = Emulated<u8, 1>;
2197    type Right = Emulated<u8, 1>;
2198    type Return = f32;
2199    type Main = Strategy1x1;
2200
2201    #[inline(always)]
2202    fn init(&self, arch: Scalar) -> Self::Accumulator {
2203        Self::Accumulator::default(arch)
2204    }
2205
2206    #[inline(always)]
2207    fn accumulate(
2208        &self,
2209        x: Self::Left,
2210        y: Self::Right,
2211        acc: Self::Accumulator,
2212    ) -> Self::Accumulator {
2213        let x: Self::Accumulator = x.into();
2214        let y: Self::Accumulator = y.into();
2215        x * y + acc
2216    }
2217
2218    // Perform a final reduction.
2219    #[inline(always)]
2220    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2221        x.to_array().into_iter().sum::<i32>().as_f32_lossy()
2222    }
2223
2224    #[inline(always)]
2225    unsafe fn epilogue(
2226        &self,
2227        _arch: Scalar,
2228        _x: *const u8,
2229        _y: *const u8,
2230        _len: usize,
2231        _acc: Self::Accumulator,
2232    ) -> Self::Accumulator {
2233        unreachable!("The SIMD width is 1, so there should be no epilogue")
2234    }
2235}
2236
2237// An IP distance function that defers a final reduction.
2238#[derive(Clone, Copy, Debug)]
2239pub struct ResumableIP<A = diskann_wide::arch::Current>
2240where
2241    A: Architecture,
2242    IP: SIMDSchema<f32, f32, A>,
2243{
2244    acc: <IP as SIMDSchema<f32, f32, A>>::Accumulator,
2245}
2246
2247impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableIP<A>
2248where
2249    A: Architecture,
2250    IP: SIMDSchema<f32, f32, A, Return = f32>,
2251{
2252    type NonResumable = IP;
2253    type FinalReturn = f32;
2254
2255    #[inline(always)]
2256    fn init(arch: A) -> Self {
2257        Self { acc: IP.init(arch) }
2258    }
2259
2260    #[inline(always)]
2261    fn combine_with(&self, other: <IP as SIMDSchema<f32, f32, A>>::Accumulator) -> Self {
2262        Self {
2263            acc: self.acc + other,
2264        }
2265    }
2266
2267    #[inline(always)]
2268    fn sum(&self) -> f32 {
2269        IP.reduce(self.acc)
2270    }
2271}
2272
2273/////////////////////////////////
2274// Stateless Cosine Similarity //
2275/////////////////////////////////
2276
2277/// Accumulator of partial products for a full cosine distance computation (where
2278/// the norms of both the query and the dataset vector are computed on the fly).
2279#[derive(Debug, Clone, Copy)]
2280pub struct FullCosineAccumulator<T> {
2281    normx: T,
2282    normy: T,
2283    xy: T,
2284}
2285
2286impl<T> FullCosineAccumulator<T>
2287where
2288    T: SIMDVector
2289        + SIMDSumTree
2290        + SIMDMulAdd
2291        + std::ops::Mul<Output = T>
2292        + std::ops::Add<Output = T>,
2293    T::Scalar: LossyF32Conversion,
2294{
2295    #[inline(always)]
2296    pub fn new(arch: T::Arch) -> Self {
2297        // SAFETY: Zero initializing a SIMD vector is safe.
2298        let zero = T::default(arch);
2299        Self {
2300            normx: zero,
2301            normy: zero,
2302            xy: zero,
2303        }
2304    }
2305
2306    #[inline(always)]
2307    pub fn add_with(&self, x: T, y: T) -> Self {
2308        // SAFETY: Arithmetic on valid arguments is valid.
2309        FullCosineAccumulator {
2310            normx: x.mul_add_simd(x, self.normx),
2311            normy: y.mul_add_simd(y, self.normy),
2312            xy: x.mul_add_simd(y, self.xy),
2313        }
2314    }
2315
2316    #[inline(always)]
2317    pub fn add_with_unfused(&self, x: T, y: T) -> Self {
2318        // SAFETY: Arithmetic on valid arguments is valid.
2319        FullCosineAccumulator {
2320            normx: x * x + self.normx,
2321            normy: y * y + self.normy,
2322            xy: x * y + self.xy,
2323        }
2324    }
2325
2326    #[inline(always)]
2327    pub fn sum(&self) -> f32 {
2328        let normx = self.normx.sum_tree().as_f32_lossy();
2329        let normy = self.normy.sum_tree().as_f32_lossy();
2330
2331        // Evaluate the denominator early and use `force_eval`.
2332        // This will allow the long `sqrt` to be overlapped with some other instructions
2333        // rather than waiting at the end of the function.
2334        //
2335        // There is some worry of subnormal numbers, but we're optimizing for the common
2336        // case where norms are reasonable values.
2337        let denominator = normx.sqrt() * normy.sqrt();
2338        let prod = self.xy.sum_tree().as_f32_lossy();
2339
2340        // Force the final products to be completely computed before the range check.
2341        //
2342        // This prevents LLVM from trying to compute `normy` or `prod` *after* the check
2343        // to `normx`, which causes it to spill heavily to the stack.
2344        //
2345        // Unfortunately, this results in a reduction pattern that appears to be slightly
2346        // slower on AMD or Windows.
2347        force_eval(denominator);
2348        force_eval(prod);
2349
2350        // This basically checks if either norm is subnormal and if so, we treat the vector
2351        // as having norm zero.
2352        //
2353        // The reason to do this rather than checking `denominator` directly is to have
2354        // consistent behavior when one vector has a small norm (i.e., always treat it as
2355        // zero) rather than potentially changing behavior when the other vector has a very
2356        // large norm to compensate.
2357        if normx < f32::MIN_POSITIVE || normy < f32::MIN_POSITIVE {
2358            return 0.0;
2359        }
2360
2361        let v = prod / denominator;
2362        (-1.0f32).max(1.0f32.min(v))
2363    }
2364
2365    /// Compute the L2 distance from the partial products rather than the cosine similarity.
2366    #[inline(always)]
2367    pub fn sum_as_l2(&self) -> f32 {
2368        let normx = self.normx.sum_tree().as_f32_lossy();
2369        let normy = self.normy.sum_tree().as_f32_lossy();
2370        let xy = self.xy.sum_tree().as_f32_lossy();
2371        normx + normy - (xy + xy)
2372    }
2373}
2374
2375impl<T> std::ops::Add for FullCosineAccumulator<T>
2376where
2377    T: std::ops::Add<Output = T>,
2378{
2379    type Output = Self;
2380    #[inline(always)]
2381    fn add(self, other: Self) -> Self {
2382        FullCosineAccumulator {
2383            normx: self.normx + other.normx,
2384            normy: self.normy + other.normy,
2385            xy: self.xy + other.xy,
2386        }
2387    }
2388}
2389
2390/// A pure Cosine Similarity function that provides a final reduction.
2391#[derive(Default, Clone, Copy)]
2392pub struct CosineStateless;
2393
2394#[cfg(target_arch = "x86_64")]
2395impl SIMDSchema<f32, f32, V4> for CosineStateless {
2396    type SIMDWidth = Const<16>;
2397    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::f32x16>;
2398    type Left = <V4 as Architecture>::f32x16;
2399    type Right = <V4 as Architecture>::f32x16;
2400    type Return = f32;
2401
2402    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2403    // hefty unroll factor.
2404    type Main = Strategy2x4;
2405
2406    #[inline(always)]
2407    fn init(&self, arch: V4) -> Self::Accumulator {
2408        Self::Accumulator::new(arch)
2409    }
2410
2411    #[inline(always)]
2412    fn accumulate(
2413        &self,
2414        x: Self::Left,
2415        y: Self::Right,
2416        acc: Self::Accumulator,
2417    ) -> Self::Accumulator {
2418        acc.add_with(x, y)
2419    }
2420
2421    // Perform a final reduction.
2422    #[inline(always)]
2423    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2424        acc.sum()
2425    }
2426}
2427
2428#[cfg(target_arch = "x86_64")]
2429impl SIMDSchema<f32, f32, V3> for CosineStateless {
2430    type SIMDWidth = Const<8>;
2431    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::f32x8>;
2432    type Left = <V3 as Architecture>::f32x8;
2433    type Right = <V3 as Architecture>::f32x8;
2434    type Return = f32;
2435
2436    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2437    // hefty unroll factor.
2438    type Main = Strategy2x4;
2439
2440    #[inline(always)]
2441    fn init(&self, arch: V3) -> Self::Accumulator {
2442        Self::Accumulator::new(arch)
2443    }
2444
2445    #[inline(always)]
2446    fn accumulate(
2447        &self,
2448        x: Self::Left,
2449        y: Self::Right,
2450        acc: Self::Accumulator,
2451    ) -> Self::Accumulator {
2452        acc.add_with(x, y)
2453    }
2454
2455    // Perform a final reduction.
2456    #[inline(always)]
2457    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2458        acc.sum()
2459    }
2460}
2461
2462#[cfg(target_arch = "aarch64")]
2463impl SIMDSchema<f32, f32, Neon> for CosineStateless {
2464    type SIMDWidth = Const<4>;
2465    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::f32x4>;
2466    type Left = <Neon as Architecture>::f32x4;
2467    type Right = <Neon as Architecture>::f32x4;
2468    type Return = f32;
2469
2470    // Cosine accumulators are pretty large, so only use 2 parallel accumulator with a
2471    // hefty unroll factor.
2472    type Main = Strategy2x4;
2473
2474    #[inline(always)]
2475    fn init(&self, arch: Neon) -> Self::Accumulator {
2476        Self::Accumulator::new(arch)
2477    }
2478
2479    #[inline(always)]
2480    fn accumulate(
2481        &self,
2482        x: Self::Left,
2483        y: Self::Right,
2484        acc: Self::Accumulator,
2485    ) -> Self::Accumulator {
2486        acc.add_with(x, y)
2487    }
2488
2489    #[inline(always)]
2490    unsafe fn epilogue(
2491        &self,
2492        arch: Neon,
2493        x: *const f32,
2494        y: *const f32,
2495        len: usize,
2496        acc: Self::Accumulator,
2497    ) -> Self::Accumulator {
2498        let mut xx: f32 = 0.0;
2499        let mut yy: f32 = 0.0;
2500        let mut xy: f32 = 0.0;
2501        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
2502            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2503            let vx = unsafe { x.add(i).read_unaligned() };
2504            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2505            let vy = unsafe { y.add(i).read_unaligned() };
2506            xx = vx.mul_add(vx, xx);
2507            yy = vy.mul_add(vy, yy);
2508            xy = vx.mul_add(vy, xy);
2509        }
2510        type V = <Neon as Architecture>::f32x4;
2511        acc + FullCosineAccumulator {
2512            normx: V::from_array(arch, [xx, 0.0, 0.0, 0.0]),
2513            normy: V::from_array(arch, [yy, 0.0, 0.0, 0.0]),
2514            xy: V::from_array(arch, [xy, 0.0, 0.0, 0.0]),
2515        }
2516    }
2517
2518    // Perform a final reduction.
2519    #[inline(always)]
2520    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2521        acc.sum()
2522    }
2523}
2524
2525impl SIMDSchema<f32, f32, Scalar> for CosineStateless {
2526    type SIMDWidth = Const<4>;
2527    type Accumulator = FullCosineAccumulator<Emulated<f32, 4>>;
2528    type Left = Emulated<f32, 4>;
2529    type Right = Emulated<f32, 4>;
2530    type Return = f32;
2531
2532    type Main = Strategy2x1;
2533
2534    #[inline(always)]
2535    fn init(&self, arch: Scalar) -> Self::Accumulator {
2536        Self::Accumulator::new(arch)
2537    }
2538
2539    #[inline(always)]
2540    fn accumulate(
2541        &self,
2542        x: Self::Left,
2543        y: Self::Right,
2544        acc: Self::Accumulator,
2545    ) -> Self::Accumulator {
2546        acc.add_with_unfused(x, y)
2547    }
2548
2549    #[inline(always)]
2550    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2551        acc.sum()
2552    }
2553}
2554
2555#[cfg(target_arch = "x86_64")]
2556impl SIMDSchema<Half, Half, V4> for CosineStateless {
2557    type SIMDWidth = Const<16>;
2558    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::f32x16>;
2559    type Left = <V4 as Architecture>::f16x16;
2560    type Right = <V4 as Architecture>::f16x16;
2561    type Return = f32;
2562    type Main = Strategy2x4;
2563
2564    #[inline(always)]
2565    fn init(&self, arch: V4) -> Self::Accumulator {
2566        Self::Accumulator::new(arch)
2567    }
2568
2569    #[inline(always)]
2570    fn accumulate(
2571        &self,
2572        x: Self::Left,
2573        y: Self::Right,
2574        acc: Self::Accumulator,
2575    ) -> Self::Accumulator {
2576        diskann_wide::alias!(f32s = <V4>::f32x16);
2577
2578        let x: f32s = x.into();
2579        let y: f32s = y.into();
2580        acc.add_with(x, y)
2581    }
2582
2583    #[inline(always)]
2584    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2585        acc.sum()
2586    }
2587}
2588
2589#[cfg(target_arch = "x86_64")]
2590impl SIMDSchema<Half, Half, V3> for CosineStateless {
2591    type SIMDWidth = Const<8>;
2592    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::f32x8>;
2593    type Left = <V3 as Architecture>::f16x8;
2594    type Right = <V3 as Architecture>::f16x8;
2595    type Return = f32;
2596    type Main = Strategy2x4;
2597
2598    #[inline(always)]
2599    fn init(&self, arch: V3) -> Self::Accumulator {
2600        Self::Accumulator::new(arch)
2601    }
2602
2603    #[inline(always)]
2604    fn accumulate(
2605        &self,
2606        x: Self::Left,
2607        y: Self::Right,
2608        acc: Self::Accumulator,
2609    ) -> Self::Accumulator {
2610        diskann_wide::alias!(f32s = <V3>::f32x8);
2611
2612        let x: f32s = x.into();
2613        let y: f32s = y.into();
2614        acc.add_with(x, y)
2615    }
2616
2617    // Perform a final reduction.
2618    #[inline(always)]
2619    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2620        acc.sum()
2621    }
2622}
2623
2624#[cfg(target_arch = "aarch64")]
2625impl SIMDSchema<Half, Half, Neon> for CosineStateless {
2626    type SIMDWidth = Const<4>;
2627    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::f32x4>;
2628    type Left = diskann_wide::arch::aarch64::f16x4;
2629    type Right = diskann_wide::arch::aarch64::f16x4;
2630    type Return = f32;
2631
2632    type Main = Strategy2x4;
2633
2634    #[inline(always)]
2635    fn init(&self, arch: Neon) -> Self::Accumulator {
2636        Self::Accumulator::new(arch)
2637    }
2638
2639    #[inline(always)]
2640    fn accumulate(
2641        &self,
2642        x: Self::Left,
2643        y: Self::Right,
2644        acc: Self::Accumulator,
2645    ) -> Self::Accumulator {
2646        diskann_wide::alias!(f32s = <Neon>::f32x4);
2647
2648        let x: f32s = x.into();
2649        let y: f32s = y.into();
2650        acc.add_with(x, y)
2651    }
2652
2653    #[inline(always)]
2654    unsafe fn epilogue(
2655        &self,
2656        arch: Neon,
2657        x: *const Half,
2658        y: *const Half,
2659        len: usize,
2660        acc: Self::Accumulator,
2661    ) -> Self::Accumulator {
2662        type V = <Neon as Architecture>::f32x4;
2663
2664        let rest = scalar_epilogue(
2665            x,
2666            y,
2667            len.min(Self::SIMDWidth::value() - 1),
2668            FullCosineAccumulator::<V>::new(arch),
2669            |acc, x: Half, y: Half| -> FullCosineAccumulator<V> {
2670                let zero = Half::default();
2671                let x: V = Self::Left::from_array(arch, [x, zero, zero, zero]).into();
2672                let y: V = Self::Right::from_array(arch, [y, zero, zero, zero]).into();
2673                acc.add_with(x, y)
2674            },
2675        );
2676        acc + rest
2677    }
2678
2679    #[inline(always)]
2680    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2681        acc.sum()
2682    }
2683}
2684
2685impl SIMDSchema<Half, Half, Scalar> for CosineStateless {
2686    type SIMDWidth = Const<1>;
2687    type Accumulator = FullCosineAccumulator<Emulated<f32, 1>>;
2688    type Left = Emulated<Half, 1>;
2689    type Right = Emulated<Half, 1>;
2690    type Return = f32;
2691    type Main = Strategy1x1;
2692
2693    #[inline(always)]
2694    fn init(&self, arch: Scalar) -> Self::Accumulator {
2695        Self::Accumulator::new(arch)
2696    }
2697
2698    #[inline(always)]
2699    fn accumulate(
2700        &self,
2701        x: Self::Left,
2702        y: Self::Right,
2703        acc: Self::Accumulator,
2704    ) -> Self::Accumulator {
2705        let x: Emulated<f32, 1> = x.into();
2706        let y: Emulated<f32, 1> = y.into();
2707        acc.add_with_unfused(x, y)
2708    }
2709
2710    #[inline(always)]
2711    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2712        acc.sum()
2713    }
2714}
2715impl<A> SIMDSchema<f32, Half, A> for CosineStateless
2716where
2717    A: Architecture,
2718{
2719    type SIMDWidth = Const<8>;
2720    type Accumulator = FullCosineAccumulator<A::f32x8>;
2721    type Left = A::f32x8;
2722    type Right = A::f16x8;
2723    type Return = f32;
2724    type Main = Strategy2x4;
2725
2726    #[inline(always)]
2727    fn init(&self, arch: A) -> Self::Accumulator {
2728        Self::Accumulator::new(arch)
2729    }
2730
2731    #[inline(always)]
2732    fn accumulate(
2733        &self,
2734        x: Self::Left,
2735        y: Self::Right,
2736        acc: Self::Accumulator,
2737    ) -> Self::Accumulator {
2738        let y: A::f32x8 = y.into();
2739        acc.add_with(x, y)
2740    }
2741
2742    #[inline(always)]
2743    fn reduce(&self, acc: Self::Accumulator) -> Self::Return {
2744        acc.sum()
2745    }
2746}
2747
2748#[cfg(target_arch = "x86_64")]
2749impl SIMDSchema<i8, i8, V4> for CosineStateless {
2750    type SIMDWidth = Const<32>;
2751    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::i32x16>;
2752    type Left = <V4 as Architecture>::i8x32;
2753    type Right = <V4 as Architecture>::i8x32;
2754    type Return = f32;
2755    type Main = Strategy4x1;
2756
2757    #[inline(always)]
2758    fn init(&self, arch: V4) -> Self::Accumulator {
2759        Self::Accumulator::new(arch)
2760    }
2761
2762    #[inline(always)]
2763    fn accumulate(
2764        &self,
2765        x: Self::Left,
2766        y: Self::Right,
2767        acc: Self::Accumulator,
2768    ) -> Self::Accumulator {
2769        diskann_wide::alias!(i16s = <V4>::i16x32);
2770
2771        let x: i16s = x.into();
2772        let y: i16s = y.into();
2773
2774        FullCosineAccumulator {
2775            normx: acc.normx.dot_simd(x, x),
2776            normy: acc.normy.dot_simd(y, y),
2777            xy: acc.xy.dot_simd(x, y),
2778        }
2779    }
2780
2781    // Perform a final reduction.
2782    #[inline(always)]
2783    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2784        x.sum()
2785    }
2786}
2787
2788#[cfg(target_arch = "x86_64")]
2789impl SIMDSchema<i8, i8, V3> for CosineStateless {
2790    type SIMDWidth = Const<16>;
2791    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::i32x8>;
2792    type Left = <V3 as Architecture>::i8x16;
2793    type Right = <V3 as Architecture>::i8x16;
2794    type Return = f32;
2795    type Main = Strategy4x1;
2796
2797    #[inline(always)]
2798    fn init(&self, arch: V3) -> Self::Accumulator {
2799        Self::Accumulator::new(arch)
2800    }
2801
2802    #[inline(always)]
2803    fn accumulate(
2804        &self,
2805        x: Self::Left,
2806        y: Self::Right,
2807        acc: Self::Accumulator,
2808    ) -> Self::Accumulator {
2809        diskann_wide::alias!(i16s = <V3>::i16x16);
2810
2811        let x: i16s = x.into();
2812        let y: i16s = y.into();
2813
2814        FullCosineAccumulator {
2815            normx: acc.normx.dot_simd(x, x),
2816            normy: acc.normy.dot_simd(y, y),
2817            xy: acc.xy.dot_simd(x, y),
2818        }
2819    }
2820
2821    // Perform a final reduction.
2822    #[inline(always)]
2823    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2824        x.sum()
2825    }
2826}
2827
2828#[cfg(target_arch = "aarch64")]
2829impl SIMDSchema<i8, i8, Neon> for CosineStateless {
2830    type SIMDWidth = Const<16>;
2831    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::i32x4>;
2832    type Left = <Neon as Architecture>::i8x16;
2833    type Right = <Neon as Architecture>::i8x16;
2834    type Return = f32;
2835    type Main = Strategy2x1;
2836
2837    #[inline(always)]
2838    fn init(&self, arch: Neon) -> Self::Accumulator {
2839        Self::Accumulator::new(arch)
2840    }
2841
2842    #[inline(always)]
2843    fn accumulate(
2844        &self,
2845        x: Self::Left,
2846        y: Self::Right,
2847        acc: Self::Accumulator,
2848    ) -> Self::Accumulator {
2849        FullCosineAccumulator {
2850            normx: acc.normx.dot_simd(x, x),
2851            normy: acc.normy.dot_simd(y, y),
2852            xy: acc.xy.dot_simd(x, y),
2853        }
2854    }
2855
2856    #[inline(always)]
2857    unsafe fn epilogue(
2858        &self,
2859        arch: Neon,
2860        x: *const i8,
2861        y: *const i8,
2862        len: usize,
2863        acc: Self::Accumulator,
2864    ) -> Self::Accumulator {
2865        let mut xx: i32 = 0;
2866        let mut yy: i32 = 0;
2867        let mut xy: i32 = 0;
2868        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
2869            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2870            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2871            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2872            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
2873            xx += vx * vx;
2874            xy += vx * vy;
2875            yy += vy * vy;
2876        }
2877        type V = <Neon as Architecture>::i32x4;
2878        acc + FullCosineAccumulator {
2879            normx: V::from_array(arch, [xx, 0, 0, 0]),
2880            normy: V::from_array(arch, [yy, 0, 0, 0]),
2881            xy: V::from_array(arch, [xy, 0, 0, 0]),
2882        }
2883    }
2884
2885    #[inline(always)]
2886    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2887        x.sum()
2888    }
2889}
2890
2891impl SIMDSchema<i8, i8, Scalar> for CosineStateless {
2892    type SIMDWidth = Const<4>;
2893    type Accumulator = FullCosineAccumulator<Emulated<i32, 4>>;
2894    type Left = Emulated<i8, 4>;
2895    type Right = Emulated<i8, 4>;
2896    type Return = f32;
2897    type Main = Strategy1x1;
2898
2899    #[inline(always)]
2900    fn init(&self, arch: Scalar) -> Self::Accumulator {
2901        Self::Accumulator::new(arch)
2902    }
2903
2904    #[inline(always)]
2905    fn accumulate(
2906        &self,
2907        x: Self::Left,
2908        y: Self::Right,
2909        acc: Self::Accumulator,
2910    ) -> Self::Accumulator {
2911        let x: Emulated<i32, 4> = x.into();
2912        let y: Emulated<i32, 4> = y.into();
2913        acc.add_with(x, y)
2914    }
2915
2916    // Perform a final reduction.
2917    #[inline(always)]
2918    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2919        x.sum()
2920    }
2921
2922    #[inline(always)]
2923    unsafe fn epilogue(
2924        &self,
2925        arch: Scalar,
2926        x: *const i8,
2927        y: *const i8,
2928        len: usize,
2929        acc: Self::Accumulator,
2930    ) -> Self::Accumulator {
2931        let mut xy: i32 = 0;
2932        let mut xx: i32 = 0;
2933        let mut yy: i32 = 0;
2934
2935        for i in 0..len {
2936            // SAFETY: The range `[x, x.add(len))` is valid for reads.
2937            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2938            // SAFETY: The range `[y, y.add(len))` is valid for reads.
2939            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
2940
2941            xx += vx * vx;
2942            xy += vx * vy;
2943            yy += vy * vy;
2944        }
2945
2946        acc + FullCosineAccumulator {
2947            normx: Emulated::from_array(arch, [xx, 0, 0, 0]),
2948            normy: Emulated::from_array(arch, [yy, 0, 0, 0]),
2949            xy: Emulated::from_array(arch, [xy, 0, 0, 0]),
2950        }
2951    }
2952}
2953
2954#[cfg(target_arch = "x86_64")]
2955impl SIMDSchema<u8, u8, V4> for CosineStateless {
2956    type SIMDWidth = Const<32>;
2957    type Accumulator = FullCosineAccumulator<<V4 as Architecture>::i32x16>;
2958    type Left = <V4 as Architecture>::u8x32;
2959    type Right = <V4 as Architecture>::u8x32;
2960    type Return = f32;
2961    type Main = Strategy4x1;
2962
2963    #[inline(always)]
2964    fn init(&self, arch: V4) -> Self::Accumulator {
2965        Self::Accumulator::new(arch)
2966    }
2967
2968    #[inline(always)]
2969    fn accumulate(
2970        &self,
2971        x: Self::Left,
2972        y: Self::Right,
2973        acc: Self::Accumulator,
2974    ) -> Self::Accumulator {
2975        diskann_wide::alias!(i16s = <V4>::i16x32);
2976
2977        let x: i16s = x.into();
2978        let y: i16s = y.into();
2979
2980        FullCosineAccumulator {
2981            normx: acc.normx.dot_simd(x, x),
2982            normy: acc.normy.dot_simd(y, y),
2983            xy: acc.xy.dot_simd(x, y),
2984        }
2985    }
2986
2987    // Perform a final reduction.
2988    #[inline(always)]
2989    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
2990        x.sum()
2991    }
2992}
2993
2994#[cfg(target_arch = "x86_64")]
2995impl SIMDSchema<u8, u8, V3> for CosineStateless {
2996    type SIMDWidth = Const<16>;
2997    type Accumulator = FullCosineAccumulator<<V3 as Architecture>::i32x8>;
2998    type Left = <V3 as Architecture>::u8x16;
2999    type Right = <V3 as Architecture>::u8x16;
3000    type Return = f32;
3001    type Main = Strategy4x1;
3002
3003    #[inline(always)]
3004    fn init(&self, arch: V3) -> Self::Accumulator {
3005        Self::Accumulator::new(arch)
3006    }
3007
3008    #[inline(always)]
3009    fn accumulate(
3010        &self,
3011        x: Self::Left,
3012        y: Self::Right,
3013        acc: Self::Accumulator,
3014    ) -> Self::Accumulator {
3015        diskann_wide::alias!(i16s = <V3>::i16x16);
3016
3017        let x: i16s = x.into();
3018        let y: i16s = y.into();
3019
3020        FullCosineAccumulator {
3021            normx: acc.normx.dot_simd(x, x),
3022            normy: acc.normy.dot_simd(y, y),
3023            xy: acc.xy.dot_simd(x, y),
3024        }
3025    }
3026
3027    // Perform a final reduction.
3028    #[inline(always)]
3029    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3030        x.sum()
3031    }
3032}
3033
3034#[cfg(target_arch = "aarch64")]
3035impl SIMDSchema<u8, u8, Neon> for CosineStateless {
3036    type SIMDWidth = Const<16>;
3037    type Accumulator = FullCosineAccumulator<<Neon as Architecture>::u32x4>;
3038    type Left = <Neon as Architecture>::u8x16;
3039    type Right = <Neon as Architecture>::u8x16;
3040    type Return = f32;
3041    type Main = Strategy2x1;
3042
3043    #[inline(always)]
3044    fn init(&self, arch: Neon) -> Self::Accumulator {
3045        Self::Accumulator::new(arch)
3046    }
3047
3048    #[inline(always)]
3049    fn accumulate(
3050        &self,
3051        x: Self::Left,
3052        y: Self::Right,
3053        acc: Self::Accumulator,
3054    ) -> Self::Accumulator {
3055        FullCosineAccumulator {
3056            normx: acc.normx.dot_simd(x, x),
3057            normy: acc.normy.dot_simd(y, y),
3058            xy: acc.xy.dot_simd(x, y),
3059        }
3060    }
3061
3062    #[inline(always)]
3063    unsafe fn epilogue(
3064        &self,
3065        arch: Neon,
3066        x: *const u8,
3067        y: *const u8,
3068        len: usize,
3069        acc: Self::Accumulator,
3070    ) -> Self::Accumulator {
3071        let mut xx: u32 = 0;
3072        let mut yy: u32 = 0;
3073        let mut xy: u32 = 0;
3074        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
3075            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3076            let vx: u32 = unsafe { x.add(i).read_unaligned() }.into();
3077            // SAFETY: The range `[y, y.add(len))` is valid for reads.
3078            let vy: u32 = unsafe { y.add(i).read_unaligned() }.into();
3079            xx += vx * vx;
3080            xy += vx * vy;
3081            yy += vy * vy;
3082        }
3083        type V = <Neon as Architecture>::u32x4;
3084        acc + FullCosineAccumulator {
3085            normx: V::from_array(arch, [xx, 0, 0, 0]),
3086            normy: V::from_array(arch, [yy, 0, 0, 0]),
3087            xy: V::from_array(arch, [xy, 0, 0, 0]),
3088        }
3089    }
3090
3091    #[inline(always)]
3092    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3093        x.sum()
3094    }
3095}
3096
3097impl SIMDSchema<u8, u8, Scalar> for CosineStateless {
3098    type SIMDWidth = Const<4>;
3099    type Accumulator = FullCosineAccumulator<Emulated<i32, 4>>;
3100    type Left = Emulated<u8, 4>;
3101    type Right = Emulated<u8, 4>;
3102    type Return = f32;
3103    type Main = Strategy1x1;
3104
3105    #[inline(always)]
3106    fn init(&self, arch: Scalar) -> Self::Accumulator {
3107        Self::Accumulator::new(arch)
3108    }
3109
3110    #[inline(always)]
3111    fn accumulate(
3112        &self,
3113        x: Self::Left,
3114        y: Self::Right,
3115        acc: Self::Accumulator,
3116    ) -> Self::Accumulator {
3117        let x: Emulated<i32, 4> = x.into();
3118        let y: Emulated<i32, 4> = y.into();
3119        acc.add_with(x, y)
3120    }
3121
3122    // Perform a final reduction.
3123    #[inline(always)]
3124    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3125        x.sum()
3126    }
3127
3128    #[inline(always)]
3129    unsafe fn epilogue(
3130        &self,
3131        arch: Scalar,
3132        x: *const u8,
3133        y: *const u8,
3134        len: usize,
3135        acc: Self::Accumulator,
3136    ) -> Self::Accumulator {
3137        let mut xy: i32 = 0;
3138        let mut xx: i32 = 0;
3139        let mut yy: i32 = 0;
3140
3141        for i in 0..len {
3142            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3143            let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
3144            // SAFETY: The range `[y, y.add(len))` is valid for reads.
3145            let vy: i32 = unsafe { y.add(i).read_unaligned() }.into();
3146
3147            xx += vx * vx;
3148            xy += vx * vy;
3149            yy += vy * vy;
3150        }
3151
3152        acc + FullCosineAccumulator {
3153            normx: Emulated::from_array(arch, [xx, 0, 0, 0]),
3154            normy: Emulated::from_array(arch, [yy, 0, 0, 0]),
3155            xy: Emulated::from_array(arch, [xy, 0, 0, 0]),
3156        }
3157    }
3158}
3159
3160/// A resumable cosine similarity computation.
3161#[derive(Debug, Clone, Copy)]
3162pub struct ResumableCosine<A = diskann_wide::arch::Current>(
3163    <CosineStateless as SIMDSchema<f32, f32, A>>::Accumulator,
3164)
3165where
3166    A: Architecture,
3167    CosineStateless: SIMDSchema<f32, f32, A>;
3168
3169impl<A> ResumableSIMDSchema<f32, f32, A> for ResumableCosine<A>
3170where
3171    A: Architecture,
3172    CosineStateless: SIMDSchema<f32, f32, A, Return = f32>,
3173{
3174    type NonResumable = CosineStateless;
3175    type FinalReturn = f32;
3176
3177    #[inline(always)]
3178    fn init(arch: A) -> Self {
3179        Self(CosineStateless.init(arch))
3180    }
3181
3182    #[inline(always)]
3183    fn combine_with(
3184        &self,
3185        other: <CosineStateless as SIMDSchema<f32, f32, A>>::Accumulator,
3186    ) -> Self {
3187        Self(self.0 + other)
3188    }
3189
3190    #[inline(always)]
3191    fn sum(&self) -> f32 {
3192        CosineStateless.reduce(self.0)
3193    }
3194}
3195
3196/////
3197///// L1 Norm Implementations
3198/////
3199
3200// ==================================================================================================
3201// NOTE: L1Norm IS A LOGICAL UNARY OPERATION
3202// --------------------------------------------------------------------------------------------------
3203// Although wired through the generic binary 'SIMDSchema'/'simd_op' infrastructure (which expects
3204// two input slices of equal length), 'L1Norm' conceptually computes: sum_i |x_i|
3205// The right-hand operand is completely ignored and exists ONLY to satisfy the shared execution
3206// machinery (loop tiling, epilogue handling, etc.).
3207// ==================================================================================================
3208
3209// A pure L1 norm function that provides a final reduction.
3210#[derive(Clone, Copy, Debug, Default)]
3211pub struct L1Norm;
3212
3213#[cfg(target_arch = "x86_64")]
3214impl SIMDSchema<f32, f32, V4> for L1Norm {
3215    type SIMDWidth = Const<16>;
3216    type Accumulator = <V4 as Architecture>::f32x16;
3217    type Left = <V4 as Architecture>::f32x16;
3218    type Right = <V4 as Architecture>::f32x16;
3219    type Return = f32;
3220    type Main = Strategy4x1;
3221
3222    #[inline(always)]
3223    fn init(&self, arch: V4) -> Self::Accumulator {
3224        Self::Accumulator::default(arch)
3225    }
3226
3227    #[inline(always)]
3228    fn accumulate(
3229        &self,
3230        x: Self::Left,
3231        _y: Self::Right,
3232        acc: Self::Accumulator,
3233    ) -> Self::Accumulator {
3234        x.abs_simd() + acc
3235    }
3236
3237    // Perform a final reduction.
3238    #[inline(always)]
3239    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3240        x.sum_tree()
3241    }
3242}
3243
3244#[cfg(target_arch = "x86_64")]
3245impl SIMDSchema<f32, f32, V3> for L1Norm {
3246    type SIMDWidth = Const<8>;
3247    type Accumulator = <V3 as Architecture>::f32x8;
3248    type Left = <V3 as Architecture>::f32x8;
3249    type Right = <V3 as Architecture>::f32x8;
3250    type Return = f32;
3251    type Main = Strategy4x1;
3252
3253    #[inline(always)]
3254    fn init(&self, arch: V3) -> Self::Accumulator {
3255        Self::Accumulator::default(arch)
3256    }
3257
3258    #[inline(always)]
3259    fn accumulate(
3260        &self,
3261        x: Self::Left,
3262        _y: Self::Right,
3263        acc: Self::Accumulator,
3264    ) -> Self::Accumulator {
3265        x.abs_simd() + acc
3266    }
3267
3268    // Perform a final reduction.
3269    #[inline(always)]
3270    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3271        x.sum_tree()
3272    }
3273}
3274
3275#[cfg(target_arch = "aarch64")]
3276impl SIMDSchema<f32, f32, Neon> for L1Norm {
3277    type SIMDWidth = Const<4>;
3278    type Accumulator = <Neon as Architecture>::f32x4;
3279    type Left = <Neon as Architecture>::f32x4;
3280    type Right = <Neon as Architecture>::f32x4;
3281    type Return = f32;
3282    type Main = Strategy4x1;
3283
3284    #[inline(always)]
3285    fn init(&self, arch: Neon) -> Self::Accumulator {
3286        Self::Accumulator::default(arch)
3287    }
3288
3289    #[inline(always)]
3290    fn accumulate(
3291        &self,
3292        x: Self::Left,
3293        _y: Self::Right,
3294        acc: Self::Accumulator,
3295    ) -> Self::Accumulator {
3296        x.abs_simd() + acc
3297    }
3298
3299    #[inline(always)]
3300    unsafe fn epilogue(
3301        &self,
3302        arch: Neon,
3303        x: *const f32,
3304        _y: *const f32,
3305        len: usize,
3306        acc: Self::Accumulator,
3307    ) -> Self::Accumulator {
3308        let mut s: f32 = 0.0;
3309        for i in 0..len.min(Self::SIMDWidth::value() - 1) {
3310            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3311            let vx = unsafe { x.add(i).read_unaligned() };
3312            s += vx.abs();
3313        }
3314        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
3315    }
3316
3317    #[inline(always)]
3318    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3319        x.sum_tree()
3320    }
3321}
3322
3323impl SIMDSchema<f32, f32, Scalar> for L1Norm {
3324    type SIMDWidth = Const<4>;
3325    type Accumulator = Emulated<f32, 4>;
3326    type Left = Emulated<f32, 4>;
3327    type Right = Emulated<f32, 4>;
3328    type Return = f32;
3329    type Main = Strategy2x1;
3330
3331    #[inline(always)]
3332    fn init(&self, arch: Scalar) -> Self::Accumulator {
3333        Self::Accumulator::default(arch)
3334    }
3335
3336    #[inline(always)]
3337    fn accumulate(
3338        &self,
3339        x: Self::Left,
3340        _y: Self::Right,
3341        acc: Self::Accumulator,
3342    ) -> Self::Accumulator {
3343        x.abs_simd() + acc
3344    }
3345
3346    // Perform a final reduction.
3347    #[inline(always)]
3348    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3349        x.sum_tree()
3350    }
3351
3352    #[inline(always)]
3353    unsafe fn epilogue(
3354        &self,
3355        arch: Scalar,
3356        x: *const f32,
3357        _y: *const f32,
3358        len: usize,
3359        acc: Self::Accumulator,
3360    ) -> Self::Accumulator {
3361        let mut s: f32 = 0.0;
3362        for i in 0..len {
3363            // SAFETY: The range `[x, x.add(len))` is valid for reads.
3364            let vx = unsafe { x.add(i).read_unaligned() };
3365            s += vx.abs();
3366        }
3367        acc + Self::Accumulator::from_array(arch, [s, 0.0, 0.0, 0.0])
3368    }
3369}
3370
3371#[cfg(target_arch = "x86_64")]
3372impl SIMDSchema<Half, Half, V4> for L1Norm {
3373    type SIMDWidth = Const<8>;
3374    type Accumulator = <V4 as Architecture>::f32x8;
3375    type Left = <V4 as Architecture>::f16x8;
3376    type Right = <V4 as Architecture>::f16x8;
3377    type Return = f32;
3378    type Main = Strategy2x4;
3379
3380    #[inline(always)]
3381    fn init(&self, arch: V4) -> Self::Accumulator {
3382        Self::Accumulator::default(arch)
3383    }
3384
3385    #[inline(always)]
3386    fn accumulate(
3387        &self,
3388        x: Self::Left,
3389        _y: Self::Right,
3390        acc: Self::Accumulator,
3391    ) -> Self::Accumulator {
3392        let x: <V4 as Architecture>::f32x8 = x.into();
3393        x.abs_simd() + acc
3394    }
3395
3396    // Perform a final reduction.
3397    #[inline(always)]
3398    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3399        x.sum_tree()
3400    }
3401}
3402
3403#[cfg(target_arch = "x86_64")]
3404impl SIMDSchema<Half, Half, V3> for L1Norm {
3405    type SIMDWidth = Const<8>;
3406    type Accumulator = <V3 as Architecture>::f32x8;
3407    type Left = <V3 as Architecture>::f16x8;
3408    type Right = <V3 as Architecture>::f16x8;
3409    type Return = f32;
3410    type Main = Strategy2x4;
3411
3412    #[inline(always)]
3413    fn init(&self, arch: V3) -> Self::Accumulator {
3414        Self::Accumulator::default(arch)
3415    }
3416
3417    #[inline(always)]
3418    fn accumulate(
3419        &self,
3420        x: Self::Left,
3421        _y: Self::Right,
3422        acc: Self::Accumulator,
3423    ) -> Self::Accumulator {
3424        let x: <V3 as Architecture>::f32x8 = x.into();
3425        x.abs_simd() + acc
3426    }
3427
3428    // Perform a final reduction.
3429    #[inline(always)]
3430    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3431        x.sum_tree()
3432    }
3433}
3434
3435#[cfg(target_arch = "aarch64")]
3436impl SIMDSchema<Half, Half, Neon> for L1Norm {
3437    type SIMDWidth = Const<4>;
3438    type Accumulator = <Neon as Architecture>::f32x4;
3439    type Left = diskann_wide::arch::aarch64::f16x4;
3440    type Right = diskann_wide::arch::aarch64::f16x4;
3441    type Return = f32;
3442    type Main = Strategy2x4;
3443
3444    #[inline(always)]
3445    fn init(&self, arch: Neon) -> Self::Accumulator {
3446        Self::Accumulator::default(arch)
3447    }
3448
3449    #[inline(always)]
3450    fn accumulate(
3451        &self,
3452        x: Self::Left,
3453        _y: Self::Right,
3454        acc: Self::Accumulator,
3455    ) -> Self::Accumulator {
3456        let x: <Neon as Architecture>::f32x4 = x.into();
3457        x.abs_simd() + acc
3458    }
3459
3460    #[inline(always)]
3461    unsafe fn epilogue(
3462        &self,
3463        arch: Neon,
3464        x: *const Half,
3465        _y: *const Half,
3466        len: usize,
3467        acc: Self::Accumulator,
3468    ) -> Self::Accumulator {
3469        let rest = scalar_epilogue(
3470            x,
3471            x, // unused, but scalar_epilogue requires a right pointer
3472            len.min(Self::SIMDWidth::value() - 1),
3473            Self::Accumulator::default(arch),
3474            |acc, x: Half, _: Half| -> Self::Accumulator {
3475                let zero = Half::default();
3476                let x: Self::Accumulator =
3477                    Self::Left::from_array(arch, [x, zero, zero, zero]).into();
3478                x.abs_simd() + acc
3479            },
3480        );
3481        acc + rest
3482    }
3483
3484    // Perform a final reduction.
3485    #[inline(always)]
3486    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3487        x.sum_tree()
3488    }
3489}
3490
3491impl SIMDSchema<Half, Half, Scalar> for L1Norm {
3492    type SIMDWidth = Const<1>;
3493    type Accumulator = Emulated<f32, 1>;
3494    type Left = Emulated<Half, 1>;
3495    type Right = Emulated<Half, 1>;
3496    type Return = f32;
3497    type Main = Strategy1x1;
3498
3499    #[inline(always)]
3500    fn init(&self, arch: Scalar) -> Self::Accumulator {
3501        Self::Accumulator::default(arch)
3502    }
3503
3504    #[inline(always)]
3505    fn accumulate(
3506        &self,
3507        x: Self::Left,
3508        _y: Self::Right,
3509        acc: Self::Accumulator,
3510    ) -> Self::Accumulator {
3511        let x: Self::Accumulator = x.into();
3512        x.abs_simd() + acc
3513    }
3514
3515    // Perform a final reduction.
3516    #[inline(always)]
3517    fn reduce(&self, x: Self::Accumulator) -> Self::Return {
3518        x.to_array()[0]
3519    }
3520
3521    #[inline(always)]
3522    unsafe fn epilogue(
3523        &self,
3524        _arch: Scalar,
3525        _x: *const Half,
3526        _y: *const Half,
3527        _len: usize,
3528        _acc: Self::Accumulator,
3529    ) -> Self::Accumulator {
3530        unreachable!("The SIMD width is 1, so there should be no epilogue")
3531    }
3532}
3533
3534///////////
3535// Tests //
3536///////////
3537
3538#[cfg(test)]
3539mod tests {
3540    use std::{collections::HashMap, sync::LazyLock};
3541
3542    use approx::assert_relative_eq;
3543    use diskann_wide::{arch::Target1, ARCH};
3544    use half::f16;
3545    use rand::{distr::StandardUniform, rngs::StdRng, Rng, SeedableRng};
3546    use rand_distr;
3547
3548    use super::*;
3549    use crate::{distance::reference, norm::LInfNorm, test_util, unaligned};
3550
3551    ///////////////////////
3552    // Cosine Norm Check //
3553    ///////////////////////
3554
3555    fn cosine_norm_check_impl<A>(arch: A)
3556    where
3557        A: diskann_wide::Architecture,
3558        CosineStateless:
3559            SIMDSchema<f32, f32, A, Return = f32> + SIMDSchema<Half, Half, A, Return = f32>,
3560    {
3561        // Zero - f32
3562        {
3563            let x: [f32; 2] = [0.0, 0.0];
3564            let y: [f32; 2] = [0.0, 1.0];
3565            assert_eq!(
3566                simd_op(&CosineStateless {}, arch, x, x),
3567                0.0,
3568                "when both vectors are zero, similarity should be zero",
3569            );
3570            assert_eq!(
3571                simd_op(&CosineStateless {}, arch, x, y),
3572                0.0,
3573                "when one vector is zero, similarity should be zero",
3574            );
3575            assert_eq!(
3576                simd_op(&CosineStateless {}, arch, y, x),
3577                0.0,
3578                "when one vector is zero, similarity should be zero",
3579            );
3580        }
3581
3582        // Subnormal - f32
3583        {
3584            let x: [f32; 4] = [0.0, 0.0, 2.938736e-39f32, 0.0];
3585            let y: [f32; 4] = [0.0, 0.0, 1.0, 0.0];
3586            assert_eq!(
3587                simd_op(&CosineStateless {}, arch, x, x),
3588                0.0,
3589                "when both vectors are almost zero, similarity should be zero",
3590            );
3591            assert_eq!(
3592                simd_op(&CosineStateless {}, arch, x, y),
3593                0.0,
3594                "when one vector is almost zero, similarity should be zero",
3595            );
3596            assert_eq!(
3597                simd_op(&CosineStateless {}, arch, y, x),
3598                0.0,
3599                "when one vector is almost zero, similarity should be zero",
3600            );
3601        }
3602
3603        // Small - f32
3604        {
3605            let x: [f32; 4] = [0.0, 0.0, 1.0842022e-19f32, 0.0];
3606            let y: [f32; 4] = [0.0, 0.0, 1.0, 0.0];
3607            assert_eq!(
3608                simd_op(&CosineStateless {}, arch, x, x),
3609                1.0,
3610                "cosine-stateless should handle vectors this small",
3611            );
3612            assert_eq!(
3613                simd_op(&CosineStateless {}, arch, x, y),
3614                1.0,
3615                "cosine-stateless should handle vectors this small",
3616            );
3617            assert_eq!(
3618                simd_op(&CosineStateless {}, arch, y, x),
3619                1.0,
3620                "cosine-stateless should handle vectors this small",
3621            );
3622        }
3623
3624        let cvt = diskann_wide::cast_f32_to_f16;
3625
3626        // Zero - f16
3627        {
3628            let x: [Half; 2] = [Half::default(), Half::default()];
3629            let y: [Half; 2] = [Half::default(), cvt(1.0)];
3630            assert_eq!(
3631                simd_op(&CosineStateless {}, arch, x, x),
3632                0.0,
3633                "when both vectors are zero, similarity should be zero",
3634            );
3635            assert_eq!(
3636                simd_op(&CosineStateless {}, arch, x, y),
3637                0.0,
3638                "when one vector is zero, similarity should be zero",
3639            );
3640            assert_eq!(
3641                simd_op(&CosineStateless {}, arch, y, x),
3642                0.0,
3643                "when one vector is zero, similarity should be zero",
3644            );
3645        }
3646
3647        // Subnormal - f16
3648        {
3649            let x: [Half; 4] = [
3650                Half::default(),
3651                Half::default(),
3652                Half::MIN_POSITIVE_SUBNORMAL,
3653                Half::default(),
3654            ];
3655            let y: [Half; 4] = [Half::default(), Half::default(), cvt(1.0), Half::default()];
3656            assert_eq!(
3657                simd_op(&CosineStateless {}, arch, x, x),
3658                1.0,
3659                "when both vectors are almost zero, similarity should be zero",
3660            );
3661            assert_eq!(
3662                simd_op(&CosineStateless {}, arch, x, y),
3663                1.0,
3664                "when one vector is almost zero, similarity should be zero",
3665            );
3666            assert_eq!(
3667                simd_op(&CosineStateless {}, arch, y, x),
3668                1.0,
3669                "when one vector is almost zero, similarity should be zero",
3670            );
3671
3672            // Grab a range of floating point numbers whose squares cover the range of
3673            // our target threshold.
3674            //
3675            // Ensure that all combinations of values within this critical range to not
3676            // result in a misrounding.
3677            let threshold = f32::MIN_POSITIVE;
3678            let bound = 50;
3679            let values = {
3680                let mut down = threshold;
3681                let mut up = threshold;
3682                for _ in 0..bound {
3683                    down = down.next_down();
3684                    up = up.next_up();
3685                }
3686                assert!(down > 0.0);
3687                let min = down.sqrt();
3688                let max = up.sqrt();
3689                let mut v = min;
3690                let mut values = Vec::new();
3691                while v <= max {
3692                    values.push(v);
3693                    v = v.next_up();
3694                }
3695                values
3696            };
3697
3698            let mut lo = 0;
3699            let mut hi = 0;
3700            for i in values.iter() {
3701                for j in values.iter() {
3702                    let s: f32 = simd_op(&CosineStateless {}, arch, [*i], [*j]);
3703                    if i * i < threshold || j * j < threshold {
3704                        lo += 1;
3705                        assert_eq!(s, 0.0, "failed for i = {}, j = {}", i, j);
3706                    } else {
3707                        hi += 1;
3708                        assert_eq!(s, 1.0, "failed for i = {}, j = {}", i, j);
3709                    }
3710                }
3711            }
3712            assert_ne!(lo, 0);
3713            assert_ne!(hi, 0);
3714        }
3715    }
3716
3717    #[test]
3718    fn cosine_norm_check() {
3719        cosine_norm_check_impl::<diskann_wide::arch::Current>(diskann_wide::arch::current());
3720        cosine_norm_check_impl::<diskann_wide::arch::Scalar>(diskann_wide::arch::Scalar::new());
3721    }
3722
3723    #[test]
3724    #[cfg(target_arch = "x86_64")]
3725    fn cosine_norm_check_x86_64() {
3726        if let Some(arch) = V3::new_checked() {
3727            cosine_norm_check_impl::<V3>(arch);
3728        }
3729
3730        if let Some(arch) = V4::new_checked_miri() {
3731            cosine_norm_check_impl::<V4>(arch);
3732        }
3733    }
3734
3735    ////////////
3736    // Schema //
3737    ////////////
3738
3739    // Chunk the left and right hand slices and compute the result using a resumable function.
3740    fn test_resumable<T, L, R, A>(arch: A, x: &[L], y: &[R], chunk_size: usize) -> f32
3741    where
3742        A: Architecture,
3743        T: ResumableSIMDSchema<L, R, A, FinalReturn = f32>,
3744    {
3745        let mut acc = Resumable(<T as ResumableSIMDSchema<L, R, A>>::init(arch));
3746        let iter = std::iter::zip(x.chunks(chunk_size), y.chunks(chunk_size));
3747        for (a, b) in iter {
3748            acc = simd_op(&acc, arch, a, b);
3749        }
3750        acc.0.sum()
3751    }
3752
3753    fn stress_test_with_resumable<
3754        A: Architecture,
3755        O: Default + SIMDSchema<f32, f32, A, Return = f32>,
3756        T: ResumableSIMDSchema<f32, f32, A, NonResumable = O, FinalReturn = f32>,
3757        Rand: Rng,
3758    >(
3759        arch: A,
3760        reference: fn(&[f32], &[f32]) -> f32,
3761        dim: usize,
3762        epsilon: f32,
3763        max_relative: f32,
3764        rng: &mut Rand,
3765    ) {
3766        // Pick chunk sizes that exercise combinations of the unrolled loops.
3767        let chunk_divisors: Vec<usize> = vec![1, 2, 3, 4, 16, 54, 64, 65, 70, 77];
3768        let mut checker = test_util::AdHocChecker::<f32, f32>::new(|a: &[f32], b: &[f32]| {
3769            let expected = reference(a, b);
3770            let got = simd_op(&O::default(), arch, a, b);
3771            println!("dim = {}", dim);
3772            assert_relative_eq!(
3773                expected,
3774                got,
3775                epsilon = epsilon,
3776                max_relative = max_relative,
3777            );
3778
3779            if dim == 0 {
3780                return;
3781            }
3782
3783            for d in &chunk_divisors {
3784                let chunk_size = dim / d + (!dim.is_multiple_of(*d) as usize);
3785                let chunked = test_resumable::<T, f32, f32, _>(arch, a, b, chunk_size);
3786                assert_relative_eq!(chunked, got, epsilon = epsilon, max_relative = max_relative);
3787            }
3788        });
3789
3790        let dist = rand_distr::Normal::new(0.0, 10.0).unwrap();
3791        test_util::test_distance_function(&mut checker, dist, dist, dim, 10, rng);
3792
3793        //-----------//
3794        // Unaligned //
3795        //-----------//
3796
3797        let mut left = unaligned::Buffer::default();
3798        let mut right = unaligned::Buffer::default();
3799        let mut checker = test_util::Checker::<f32, f32, f32>::new(
3800            |a, b| simd_op(&O::default(), arch, a.as_unaligned(), b.as_unaligned()),
3801            |a, b| {
3802                left.copy(a);
3803                right.copy(b);
3804                simd_op(
3805                    &O::default(),
3806                    arch,
3807                    left.as_unaligned(),
3808                    right.as_unaligned(),
3809                )
3810            },
3811            |got, expected| assert_eq!(got, expected),
3812        );
3813        test_util::test_distance_function(&mut checker, dist, dist, dim, 10, rng);
3814    }
3815
3816    #[expect(clippy::too_many_arguments)]
3817    fn stress_test<L, R, DistLeft, DistRight, O, Rand, A>(
3818        arch: A,
3819        reference: fn(&[L], &[R]) -> f32,
3820        left_dist: DistLeft,
3821        right_dist: DistRight,
3822        dim: usize,
3823        epsilon: f32,
3824        max_relative: f32,
3825        rng: &mut Rand,
3826    ) where
3827        L: test_util::CornerCases + bytemuck::Pod,
3828        R: test_util::CornerCases + bytemuck::Pod,
3829        DistLeft: test_util::GenerateRandomArguments<L> + Copy,
3830        DistRight: test_util::GenerateRandomArguments<R> + Copy,
3831        O: Default + SIMDSchema<L, R, A, Return = f32>,
3832        Rand: Rng,
3833        A: Architecture,
3834    {
3835        let mut checker = test_util::Checker::<L, R, f32>::new(
3836            |x: &[L], y: &[R]| simd_op(&O::default(), arch, x, y),
3837            reference,
3838            |got, expected| {
3839                assert_relative_eq!(
3840                    expected,
3841                    got,
3842                    epsilon = epsilon,
3843                    max_relative = max_relative
3844                );
3845            },
3846        );
3847
3848        let trials = if cfg!(miri) { 0 } else { 10 };
3849        test_util::test_distance_function(&mut checker, left_dist, right_dist, dim, trials, rng);
3850
3851        //-----------//
3852        // Unaligned //
3853        //-----------//
3854
3855        let mut left = unaligned::Buffer::default();
3856        let mut right = unaligned::Buffer::default();
3857        let mut checker = test_util::Checker::<L, R, f32>::new(
3858            |a, b| simd_op(&O::default(), arch, a.as_unaligned(), b.as_unaligned()),
3859            |a, b| {
3860                left.copy(a);
3861                right.copy(b);
3862                simd_op(
3863                    &O::default(),
3864                    arch,
3865                    left.as_unaligned(),
3866                    right.as_unaligned(),
3867                )
3868            },
3869            |got, expected| assert_eq!(got, expected),
3870        );
3871        test_util::test_distance_function(&mut checker, left_dist, right_dist, dim, trials, rng);
3872    }
3873
3874    fn stress_test_linf<L, Dist, Rand, A>(
3875        arch: A,
3876        reference: fn(&[L]) -> f32,
3877        dist: Dist,
3878        dim: usize,
3879        epsilon: f32,
3880        max_relative: f32,
3881        rng: &mut Rand,
3882    ) where
3883        L: test_util::CornerCases + Copy,
3884        Dist: Clone + test_util::GenerateRandomArguments<L>,
3885        Rand: Rng,
3886        A: Architecture,
3887        LInfNorm: for<'a> Target1<A, f32, &'a [L]>,
3888    {
3889        let mut checker = test_util::Checker::<L, L, f32>::new(
3890            |x: &[L], _y: &[L]| (LInfNorm).run(arch, x),
3891            |x: &[L], _y: &[L]| reference(x),
3892            |got, expected| {
3893                assert_relative_eq!(
3894                    expected,
3895                    got,
3896                    epsilon = epsilon,
3897                    max_relative = max_relative
3898                );
3899            },
3900        );
3901
3902        println!("checking {dim}");
3903        test_util::test_distance_function(&mut checker, dist.clone(), dist, dim, 10, rng);
3904    }
3905
3906    /////////
3907    // f32 //
3908    /////////
3909
3910    macro_rules! float_test {
3911        ($name:ident,
3912         $impl:ty,
3913         $resumable:ident,
3914         $reference:path,
3915         $eps:literal,
3916         $relative:literal,
3917         $seed:literal,
3918         $upper:literal,
3919         $($arch:tt)*
3920        ) => {
3921            #[test]
3922            fn $name() {
3923                if let Some(arch) = $($arch)* {
3924                    let mut rng = StdRng::seed_from_u64($seed);
3925                    for dim in 0..$upper {
3926                        stress_test_with_resumable::<_, $impl, $resumable<_>, StdRng>(
3927                            arch,
3928                            |l, r| $reference(l, r).into_inner(),
3929                            dim,
3930                            $eps,
3931                            $relative,
3932                            &mut rng,
3933                        );
3934                    }
3935                }
3936            }
3937        }
3938    }
3939
3940    //----//
3941    // L2 //
3942    //----//
3943
3944    float_test!(
3945        test_l2_f32_current,
3946        L2,
3947        ResumableL2,
3948        reference::reference_squared_l2_f32_mathematical,
3949        1e-5,
3950        1e-5,
3951        0xf149c2bcde660128,
3952        64,
3953        Some(diskann_wide::ARCH)
3954    );
3955
3956    float_test!(
3957        test_l2_f32_scalar,
3958        L2,
3959        ResumableL2,
3960        reference::reference_squared_l2_f32_mathematical,
3961        1e-5,
3962        1e-5,
3963        0xf149c2bcde660128,
3964        64,
3965        Some(diskann_wide::arch::Scalar)
3966    );
3967
3968    #[cfg(target_arch = "x86_64")]
3969    float_test!(
3970        test_l2_f32_x86_64_v3,
3971        L2,
3972        ResumableL2,
3973        reference::reference_squared_l2_f32_mathematical,
3974        1e-5,
3975        1e-5,
3976        0xf149c2bcde660128,
3977        256,
3978        V3::new_checked()
3979    );
3980
3981    #[cfg(target_arch = "x86_64")]
3982    float_test!(
3983        test_l2_f32_x86_64_v4,
3984        L2,
3985        ResumableL2,
3986        reference::reference_squared_l2_f32_mathematical,
3987        1e-5,
3988        1e-5,
3989        0xf149c2bcde660128,
3990        256,
3991        V4::new_checked_miri()
3992    );
3993
3994    #[cfg(target_arch = "aarch64")]
3995    float_test!(
3996        test_l2_f32_aarch64_neon,
3997        L2,
3998        ResumableL2,
3999        reference::reference_squared_l2_f32_mathematical,
4000        1e-5,
4001        1e-5,
4002        0xf149c2bcde660128,
4003        256,
4004        Neon::new_checked()
4005    );
4006
4007    //----//
4008    // IP //
4009    //----//
4010
4011    float_test!(
4012        test_ip_f32_current,
4013        IP,
4014        ResumableIP,
4015        reference::reference_innerproduct_f32_mathematical,
4016        2e-4,
4017        1e-3,
4018        0xb4687c17a9ea9866,
4019        64,
4020        Some(diskann_wide::ARCH)
4021    );
4022
4023    float_test!(
4024        test_ip_f32_scalar,
4025        IP,
4026        ResumableIP,
4027        reference::reference_innerproduct_f32_mathematical,
4028        2e-4,
4029        1e-3,
4030        0xb4687c17a9ea9866,
4031        64,
4032        Some(diskann_wide::arch::Scalar)
4033    );
4034
4035    #[cfg(target_arch = "x86_64")]
4036    float_test!(
4037        test_ip_f32_x86_64_v3,
4038        IP,
4039        ResumableIP,
4040        reference::reference_innerproduct_f32_mathematical,
4041        2e-4,
4042        1e-3,
4043        0xb4687c17a9ea9866,
4044        256,
4045        V3::new_checked()
4046    );
4047
4048    #[cfg(target_arch = "x86_64")]
4049    float_test!(
4050        test_ip_f32_x86_64_v4,
4051        IP,
4052        ResumableIP,
4053        reference::reference_innerproduct_f32_mathematical,
4054        2e-4,
4055        1e-3,
4056        0xb4687c17a9ea9866,
4057        256,
4058        V4::new_checked_miri()
4059    );
4060
4061    #[cfg(target_arch = "aarch64")]
4062    float_test!(
4063        test_ip_f32_aarch64_neon,
4064        IP,
4065        ResumableIP,
4066        reference::reference_innerproduct_f32_mathematical,
4067        2e-4,
4068        1e-3,
4069        0xb4687c17a9ea9866,
4070        256,
4071        Neon::new_checked()
4072    );
4073
4074    //--------//
4075    // Cosine //
4076    //--------//
4077
4078    float_test!(
4079        test_cosine_f32_current,
4080        CosineStateless,
4081        ResumableCosine,
4082        reference::reference_cosine_f32_mathematical,
4083        1e-5,
4084        1e-5,
4085        0xe860e9dc65f38bb8,
4086        64,
4087        Some(diskann_wide::ARCH)
4088    );
4089
4090    float_test!(
4091        test_cosine_f32_scalar,
4092        CosineStateless,
4093        ResumableCosine,
4094        reference::reference_cosine_f32_mathematical,
4095        1e-5,
4096        1e-5,
4097        0xe860e9dc65f38bb8,
4098        64,
4099        Some(diskann_wide::arch::Scalar)
4100    );
4101
4102    #[cfg(target_arch = "x86_64")]
4103    float_test!(
4104        test_cosine_f32_x86_64_v3,
4105        CosineStateless,
4106        ResumableCosine,
4107        reference::reference_cosine_f32_mathematical,
4108        1e-5,
4109        1e-5,
4110        0xe860e9dc65f38bb8,
4111        256,
4112        V3::new_checked()
4113    );
4114
4115    #[cfg(target_arch = "x86_64")]
4116    float_test!(
4117        test_cosine_f32_x86_64_v4,
4118        CosineStateless,
4119        ResumableCosine,
4120        reference::reference_cosine_f32_mathematical,
4121        1e-5,
4122        1e-5,
4123        0xe860e9dc65f38bb8,
4124        256,
4125        V4::new_checked_miri()
4126    );
4127
4128    #[cfg(target_arch = "aarch64")]
4129    float_test!(
4130        test_cosine_f32_aarch64_neon,
4131        CosineStateless,
4132        ResumableCosine,
4133        reference::reference_cosine_f32_mathematical,
4134        1e-5,
4135        1e-5,
4136        0xe860e9dc65f38bb8,
4137        256,
4138        Neon::new_checked()
4139    );
4140
4141    /////////
4142    // f16 //
4143    /////////
4144
4145    macro_rules! half_test {
4146        ($name:ident,
4147         $impl:ty,
4148         $reference:path,
4149         $eps:literal,
4150         $relative:literal,
4151         $seed:literal,
4152         $upper:literal,
4153         $($arch:tt)*
4154        ) => {
4155            #[test]
4156            fn $name() {
4157                if let Some(arch) = $($arch)* {
4158                    let mut rng = StdRng::seed_from_u64($seed);
4159                    for dim in 0..$upper {
4160                        stress_test::<
4161                            Half,
4162                            Half,
4163                            rand_distr::Normal<f32>,
4164                            rand_distr::Normal<f32>,
4165                            $impl,
4166                            StdRng,
4167                            _
4168                        >(
4169                            arch,
4170                            |l, r| $reference(l, r).into_inner(),
4171                            rand_distr::Normal::new(0.0, 10.0).unwrap(),
4172                            rand_distr::Normal::new(0.0, 10.0).unwrap(),
4173                            dim,
4174                            $eps,
4175                            $relative,
4176                            &mut rng
4177                        );
4178                    }
4179                }
4180            }
4181        }
4182    }
4183
4184    //----//
4185    // L2 //
4186    //----//
4187
4188    half_test!(
4189        test_l2_f16_current,
4190        L2,
4191        reference::reference_squared_l2_f16_mathematical,
4192        1e-5,
4193        1e-5,
4194        0x87ca6f1051667500,
4195        64,
4196        Some(diskann_wide::ARCH)
4197    );
4198
4199    half_test!(
4200        test_l2_f16_scalar,
4201        L2,
4202        reference::reference_squared_l2_f16_mathematical,
4203        1e-5,
4204        1e-5,
4205        0x87ca6f1051667500,
4206        64,
4207        Some(diskann_wide::arch::Scalar)
4208    );
4209
4210    #[cfg(target_arch = "x86_64")]
4211    half_test!(
4212        test_l2_f16_x86_64_v3,
4213        L2,
4214        reference::reference_squared_l2_f16_mathematical,
4215        1e-5,
4216        1e-5,
4217        0x87ca6f1051667500,
4218        256,
4219        V3::new_checked()
4220    );
4221
4222    #[cfg(target_arch = "x86_64")]
4223    half_test!(
4224        test_l2_f16_x86_64_v4,
4225        L2,
4226        reference::reference_squared_l2_f16_mathematical,
4227        1e-5,
4228        1e-5,
4229        0x87ca6f1051667500,
4230        256,
4231        V4::new_checked_miri()
4232    );
4233
4234    #[cfg(target_arch = "aarch64")]
4235    half_test!(
4236        test_l2_f16_aarch64_neon,
4237        L2,
4238        reference::reference_squared_l2_f16_mathematical,
4239        1e-5,
4240        1e-5,
4241        0x87ca6f1051667500,
4242        256,
4243        Neon::new_checked()
4244    );
4245
4246    //----//
4247    // IP //
4248    //----//
4249
4250    half_test!(
4251        test_ip_f16_current,
4252        IP,
4253        reference::reference_innerproduct_f16_mathematical,
4254        2e-4,
4255        2e-4,
4256        0x5909f5f20307ccbe,
4257        64,
4258        Some(diskann_wide::ARCH)
4259    );
4260
4261    half_test!(
4262        test_ip_f16_scalar,
4263        IP,
4264        reference::reference_innerproduct_f16_mathematical,
4265        2e-4,
4266        2e-4,
4267        0x5909f5f20307ccbe,
4268        64,
4269        Some(diskann_wide::arch::Scalar)
4270    );
4271
4272    #[cfg(target_arch = "x86_64")]
4273    half_test!(
4274        test_ip_f16_x86_64_v3,
4275        IP,
4276        reference::reference_innerproduct_f16_mathematical,
4277        2e-4,
4278        2e-4,
4279        0x5909f5f20307ccbe,
4280        256,
4281        V3::new_checked()
4282    );
4283
4284    #[cfg(target_arch = "x86_64")]
4285    half_test!(
4286        test_ip_f16_x86_64_v4,
4287        IP,
4288        reference::reference_innerproduct_f16_mathematical,
4289        2e-4,
4290        2e-4,
4291        0x5909f5f20307ccbe,
4292        256,
4293        V4::new_checked_miri()
4294    );
4295
4296    #[cfg(target_arch = "aarch64")]
4297    half_test!(
4298        test_ip_f16_aarch64_neon,
4299        IP,
4300        reference::reference_innerproduct_f16_mathematical,
4301        2e-4,
4302        2e-4,
4303        0x5909f5f20307ccbe,
4304        256,
4305        Neon::new_checked()
4306    );
4307
4308    //--------//
4309    // Cosine //
4310    //--------//
4311
4312    half_test!(
4313        test_cosine_f16_current,
4314        CosineStateless,
4315        reference::reference_cosine_f16_mathematical,
4316        1e-5,
4317        1e-5,
4318        0x41dda34655f05ef6,
4319        64,
4320        Some(diskann_wide::ARCH)
4321    );
4322
4323    half_test!(
4324        test_cosine_f16_scalar,
4325        CosineStateless,
4326        reference::reference_cosine_f16_mathematical,
4327        1e-5,
4328        1e-5,
4329        0x41dda34655f05ef6,
4330        64,
4331        Some(diskann_wide::arch::Scalar)
4332    );
4333
4334    #[cfg(target_arch = "x86_64")]
4335    half_test!(
4336        test_cosine_f16_x86_64_v3,
4337        CosineStateless,
4338        reference::reference_cosine_f16_mathematical,
4339        1e-5,
4340        1e-5,
4341        0x41dda34655f05ef6,
4342        256,
4343        V3::new_checked()
4344    );
4345
4346    #[cfg(target_arch = "x86_64")]
4347    half_test!(
4348        test_cosine_f16_x86_64_v4,
4349        CosineStateless,
4350        reference::reference_cosine_f16_mathematical,
4351        1e-5,
4352        1e-5,
4353        0x41dda34655f05ef6,
4354        256,
4355        V4::new_checked_miri()
4356    );
4357
4358    #[cfg(target_arch = "aarch64")]
4359    half_test!(
4360        test_cosine_f16_aarch64_neon,
4361        CosineStateless,
4362        reference::reference_cosine_f16_mathematical,
4363        1e-5,
4364        1e-5,
4365        0x41dda34655f05ef6,
4366        256,
4367        Neon::new_checked()
4368    );
4369
4370    /////////////
4371    // Integer //
4372    /////////////
4373
4374    macro_rules! int_test {
4375        (
4376            $name:ident,
4377            $T:ty,
4378            $impl:ty,
4379            $reference:path,
4380            $seed:literal,
4381            $upper:literal,
4382            { $($arch:tt)* }
4383        ) => {
4384            #[test]
4385            fn $name() {
4386                if let Some(arch) = $($arch)* {
4387                    let mut rng = StdRng::seed_from_u64($seed);
4388                    for dim in 0..$upper {
4389                        stress_test::<$T, $T, _, _, $impl, _, _>(
4390                            arch,
4391                            |l, r| $reference(l, r).into_inner(),
4392                            StandardUniform,
4393                            StandardUniform,
4394                            dim,
4395                            0.0,
4396                            0.0,
4397                            &mut rng,
4398                        )
4399                    }
4400                }
4401            }
4402        }
4403    }
4404
4405    //----//
4406    // U8 //
4407    //----//
4408
4409    int_test!(
4410        test_l2_u8_current,
4411        u8,
4412        L2,
4413        reference::reference_squared_l2_u8_mathematical,
4414        0x945bdc37d8279d4b,
4415        128,
4416        { Some(ARCH) }
4417    );
4418
4419    int_test!(
4420        test_l2_u8_scalar,
4421        u8,
4422        L2,
4423        reference::reference_squared_l2_u8_mathematical,
4424        0x74c86334ab7a51f9,
4425        128,
4426        { Some(diskann_wide::arch::Scalar) }
4427    );
4428
4429    #[cfg(target_arch = "x86_64")]
4430    int_test!(
4431        test_l2_u8_x86_64_v3,
4432        u8,
4433        L2,
4434        reference::reference_squared_l2_u8_mathematical,
4435        0x74c86334ab7a51f9,
4436        256,
4437        { V3::new_checked() }
4438    );
4439
4440    #[cfg(target_arch = "x86_64")]
4441    int_test!(
4442        test_l2_u8_x86_64_v4,
4443        u8,
4444        L2,
4445        reference::reference_squared_l2_u8_mathematical,
4446        0x74c86334ab7a51f9,
4447        320,
4448        { V4::new_checked_miri() }
4449    );
4450
4451    #[cfg(target_arch = "aarch64")]
4452    int_test!(
4453        test_l2_u8_aarch64_neon,
4454        u8,
4455        L2,
4456        reference::reference_squared_l2_u8_mathematical,
4457        0x74c86334ab7a51f9,
4458        320,
4459        { Neon::new_checked() }
4460    );
4461
4462    int_test!(
4463        test_ip_u8_current,
4464        u8,
4465        IP,
4466        reference::reference_innerproduct_u8_mathematical,
4467        0xcbe0342c75085fd5,
4468        64,
4469        { Some(ARCH) }
4470    );
4471
4472    int_test!(
4473        test_ip_u8_scalar,
4474        u8,
4475        IP,
4476        reference::reference_innerproduct_u8_mathematical,
4477        0x888e07fc489e773f,
4478        64,
4479        { Some(diskann_wide::arch::Scalar) }
4480    );
4481
4482    #[cfg(target_arch = "x86_64")]
4483    int_test!(
4484        test_ip_u8_x86_64_v3,
4485        u8,
4486        IP,
4487        reference::reference_innerproduct_u8_mathematical,
4488        0x888e07fc489e773f,
4489        256,
4490        { V3::new_checked() }
4491    );
4492
4493    #[cfg(target_arch = "x86_64")]
4494    int_test!(
4495        test_ip_u8_x86_64_v4,
4496        u8,
4497        IP,
4498        reference::reference_innerproduct_u8_mathematical,
4499        0x888e07fc489e773f,
4500        320,
4501        { V4::new_checked_miri() }
4502    );
4503
4504    #[cfg(target_arch = "aarch64")]
4505    int_test!(
4506        test_ip_u8_aarch64_neon,
4507        u8,
4508        IP,
4509        reference::reference_innerproduct_u8_mathematical,
4510        0x888e07fc489e773f,
4511        320,
4512        { Neon::new_checked() }
4513    );
4514
4515    int_test!(
4516        test_cosine_u8_current,
4517        u8,
4518        CosineStateless,
4519        reference::reference_cosine_u8_mathematical,
4520        0x96867b6aff616b28,
4521        64,
4522        { Some(ARCH) }
4523    );
4524
4525    int_test!(
4526        test_cosine_u8_scalar,
4527        u8,
4528        CosineStateless,
4529        reference::reference_cosine_u8_mathematical,
4530        0xcc258c9391733211,
4531        64,
4532        { Some(diskann_wide::arch::Scalar) }
4533    );
4534
4535    #[cfg(target_arch = "x86_64")]
4536    int_test!(
4537        test_cosine_u8_x86_64_v3,
4538        u8,
4539        CosineStateless,
4540        reference::reference_cosine_u8_mathematical,
4541        0xcc258c9391733211,
4542        256,
4543        { V3::new_checked() }
4544    );
4545
4546    #[cfg(target_arch = "x86_64")]
4547    int_test!(
4548        test_cosine_u8_x86_64_v4,
4549        u8,
4550        CosineStateless,
4551        reference::reference_cosine_u8_mathematical,
4552        0xcc258c9391733211,
4553        320,
4554        { V4::new_checked_miri() }
4555    );
4556
4557    #[cfg(target_arch = "aarch64")]
4558    int_test!(
4559        test_cosine_u8_aarch64_neon,
4560        u8,
4561        CosineStateless,
4562        reference::reference_cosine_u8_mathematical,
4563        0xcc258c9391733211,
4564        320,
4565        { Neon::new_checked() }
4566    );
4567
4568    //----//
4569    // I8 //
4570    //----//
4571
4572    int_test!(
4573        test_l2_i8_current,
4574        i8,
4575        L2,
4576        reference::reference_squared_l2_i8_mathematical,
4577        0xa60136248cd3c2f0,
4578        64,
4579        { Some(ARCH) }
4580    );
4581
4582    int_test!(
4583        test_l2_i8_scalar,
4584        i8,
4585        L2,
4586        reference::reference_squared_l2_i8_mathematical,
4587        0x3e8bada709e176be,
4588        64,
4589        { Some(diskann_wide::arch::Scalar) }
4590    );
4591
4592    #[cfg(target_arch = "x86_64")]
4593    int_test!(
4594        test_l2_i8_x86_64_v3,
4595        i8,
4596        L2,
4597        reference::reference_squared_l2_i8_mathematical,
4598        0x3e8bada709e176be,
4599        256,
4600        { V3::new_checked() }
4601    );
4602
4603    #[cfg(target_arch = "x86_64")]
4604    int_test!(
4605        test_l2_i8_x86_64_v4,
4606        i8,
4607        L2,
4608        reference::reference_squared_l2_i8_mathematical,
4609        0x3e8bada709e176be,
4610        320,
4611        { V4::new_checked_miri() }
4612    );
4613
4614    #[cfg(target_arch = "aarch64")]
4615    int_test!(
4616        test_l2_i8_aarch64_neon,
4617        i8,
4618        L2,
4619        reference::reference_squared_l2_i8_mathematical,
4620        0x3e8bada709e176be,
4621        320,
4622        { Neon::new_checked() }
4623    );
4624
4625    int_test!(
4626        test_ip_i8_current,
4627        i8,
4628        IP,
4629        reference::reference_innerproduct_i8_mathematical,
4630        0xe8306104740509e1,
4631        64,
4632        { Some(ARCH) }
4633    );
4634
4635    int_test!(
4636        test_ip_i8_scalar,
4637        i8,
4638        IP,
4639        reference::reference_innerproduct_i8_mathematical,
4640        0x8a263408c7b31d85,
4641        64,
4642        { Some(diskann_wide::arch::Scalar) }
4643    );
4644
4645    #[cfg(target_arch = "x86_64")]
4646    int_test!(
4647        test_ip_i8_x86_64_v3,
4648        i8,
4649        IP,
4650        reference::reference_innerproduct_i8_mathematical,
4651        0x8a263408c7b31d85,
4652        256,
4653        { V3::new_checked() }
4654    );
4655
4656    #[cfg(target_arch = "x86_64")]
4657    int_test!(
4658        test_ip_i8_x86_64_v4,
4659        i8,
4660        IP,
4661        reference::reference_innerproduct_i8_mathematical,
4662        0x8a263408c7b31d85,
4663        320,
4664        { V4::new_checked_miri() }
4665    );
4666
4667    #[cfg(target_arch = "aarch64")]
4668    int_test!(
4669        test_ip_i8_aarch64_neon,
4670        i8,
4671        IP,
4672        reference::reference_innerproduct_i8_mathematical,
4673        0x8a263408c7b31d85,
4674        320,
4675        { Neon::new_checked() }
4676    );
4677
4678    int_test!(
4679        test_cosine_i8_current,
4680        i8,
4681        CosineStateless,
4682        reference::reference_cosine_i8_mathematical,
4683        0x818c210190701e4b,
4684        64,
4685        { Some(ARCH) }
4686    );
4687
4688    int_test!(
4689        test_cosine_i8_scalar,
4690        i8,
4691        CosineStateless,
4692        reference::reference_cosine_i8_mathematical,
4693        0x2d077bed2629b18e,
4694        64,
4695        { Some(diskann_wide::arch::Scalar) }
4696    );
4697
4698    #[cfg(target_arch = "x86_64")]
4699    int_test!(
4700        test_cosine_i8_x86_64_v3,
4701        i8,
4702        CosineStateless,
4703        reference::reference_cosine_i8_mathematical,
4704        0x2d077bed2629b18e,
4705        256,
4706        { V3::new_checked() }
4707    );
4708
4709    #[cfg(target_arch = "x86_64")]
4710    int_test!(
4711        test_cosine_i8_x86_64_v4,
4712        i8,
4713        CosineStateless,
4714        reference::reference_cosine_i8_mathematical,
4715        0x2d077bed2629b18e,
4716        320,
4717        { V4::new_checked_miri() }
4718    );
4719
4720    #[cfg(target_arch = "aarch64")]
4721    int_test!(
4722        test_cosine_i8_aarch64_neon,
4723        i8,
4724        CosineStateless,
4725        reference::reference_cosine_i8_mathematical,
4726        0x2d077bed2629b18e,
4727        320,
4728        { Neon::new_checked() }
4729    );
4730
4731    //////////
4732    // LInf //
4733    //////////
4734
4735    macro_rules! linf_test {
4736        ($name:ident,
4737         $T:ty,
4738         $reference:path,
4739         $eps:literal,
4740         $relative:literal,
4741         $seed:literal,
4742         $upper:literal,
4743         $($arch:tt)*
4744        ) => {
4745            #[test]
4746            fn $name() {
4747                if let Some(arch) = $($arch)* {
4748                    let mut rng = StdRng::seed_from_u64($seed);
4749                    for dim in 0..$upper {
4750                        stress_test_linf::<$T, _, StdRng, _>(
4751                            arch,
4752                            |l| $reference(l).into_inner(),
4753                            rand_distr::Normal::new(-10.0, 10.0).unwrap(),
4754                            dim,
4755                            $eps,
4756                            $relative,
4757                            &mut rng,
4758                        );
4759                    }
4760                }
4761            }
4762        }
4763    }
4764
4765    linf_test!(
4766        test_linf_f32_scalar,
4767        f32,
4768        reference::reference_linf_f32_mathematical,
4769        1e-6,
4770        1e-6,
4771        0xf149c2bcde660128,
4772        256,
4773        Some(Scalar::new())
4774    );
4775
4776    #[cfg(target_arch = "x86_64")]
4777    linf_test!(
4778        test_linf_f32_v3,
4779        f32,
4780        reference::reference_linf_f32_mathematical,
4781        1e-6,
4782        1e-6,
4783        0xf149c2bcde660128,
4784        256,
4785        V3::new_checked()
4786    );
4787
4788    #[cfg(target_arch = "x86_64")]
4789    linf_test!(
4790        test_linf_f32_v4,
4791        f32,
4792        reference::reference_linf_f32_mathematical,
4793        1e-6,
4794        1e-6,
4795        0xf149c2bcde660128,
4796        256,
4797        V4::new_checked_miri()
4798    );
4799
4800    #[cfg(target_arch = "aarch64")]
4801    linf_test!(
4802        test_linf_f32_neon,
4803        f32,
4804        reference::reference_linf_f32_mathematical,
4805        1e-6,
4806        1e-6,
4807        0xf149c2bcde660128,
4808        256,
4809        Neon::new_checked()
4810    );
4811
4812    linf_test!(
4813        test_linf_f16_scalar,
4814        f16,
4815        reference::reference_linf_f16_mathematical,
4816        1e-6,
4817        1e-6,
4818        0xf149c2bcde660128,
4819        256,
4820        Some(Scalar::new())
4821    );
4822
4823    #[cfg(target_arch = "x86_64")]
4824    linf_test!(
4825        test_linf_f16_v3,
4826        f16,
4827        reference::reference_linf_f16_mathematical,
4828        1e-6,
4829        1e-6,
4830        0xf149c2bcde660128,
4831        256,
4832        V3::new_checked()
4833    );
4834
4835    #[cfg(target_arch = "x86_64")]
4836    linf_test!(
4837        test_linf_f16_v4,
4838        f16,
4839        reference::reference_linf_f16_mathematical,
4840        1e-6,
4841        1e-6,
4842        0xf149c2bcde660128,
4843        256,
4844        V4::new_checked_miri()
4845    );
4846
4847    #[cfg(target_arch = "aarch64")]
4848    linf_test!(
4849        test_linf_f16_neon,
4850        f16,
4851        reference::reference_linf_f16_mathematical,
4852        1e-6,
4853        1e-6,
4854        0xf149c2bcde660128,
4855        256,
4856        Neon::new_checked()
4857    );
4858
4859    ////////////////
4860    // Miri Tests //
4861    ////////////////
4862
4863    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4864    enum DataType {
4865        Float32,
4866        Float16,
4867        UInt8,
4868        Int8,
4869    }
4870
4871    trait AsDataType {
4872        fn as_data_type() -> DataType;
4873    }
4874
4875    impl AsDataType for f32 {
4876        fn as_data_type() -> DataType {
4877            DataType::Float32
4878        }
4879    }
4880
4881    impl AsDataType for f16 {
4882        fn as_data_type() -> DataType {
4883            DataType::Float16
4884        }
4885    }
4886
4887    impl AsDataType for u8 {
4888        fn as_data_type() -> DataType {
4889            DataType::UInt8
4890        }
4891    }
4892
4893    impl AsDataType for i8 {
4894        fn as_data_type() -> DataType {
4895            DataType::Int8
4896        }
4897    }
4898
4899    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4900    enum Arch {
4901        Scalar,
4902        #[expect(non_camel_case_types)]
4903        X86_64_V3,
4904        #[expect(non_camel_case_types)]
4905        X86_64_V4,
4906        Aarch64Neon,
4907    }
4908
4909    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
4910    struct Key {
4911        arch: Arch,
4912        left: DataType,
4913        right: DataType,
4914    }
4915
4916    impl Key {
4917        fn new(arch: Arch, left: DataType, right: DataType) -> Self {
4918            Self { arch, left, right }
4919        }
4920    }
4921
4922    static MIRI_BOUNDS: LazyLock<HashMap<Key, usize>> = LazyLock::new(|| {
4923        use Arch::{Aarch64Neon, Scalar, X86_64_V3, X86_64_V4};
4924        use DataType::{Float16, Float32, Int8, UInt8};
4925
4926        [
4927            (Key::new(Scalar, Float32, Float32), 64),
4928            (Key::new(X86_64_V3, Float32, Float32), 256),
4929            (Key::new(X86_64_V4, Float32, Float32), 256),
4930            (Key::new(Aarch64Neon, Float32, Float32), 128),
4931            (Key::new(Scalar, Float16, Float16), 64),
4932            (Key::new(X86_64_V3, Float16, Float16), 256),
4933            (Key::new(X86_64_V4, Float16, Float16), 256),
4934            (Key::new(Aarch64Neon, Float16, Float16), 128),
4935            (Key::new(Scalar, Float32, Float16), 64),
4936            (Key::new(X86_64_V3, Float32, Float16), 256),
4937            (Key::new(X86_64_V4, Float32, Float16), 256),
4938            (Key::new(Aarch64Neon, Float32, Float16), 128),
4939            (Key::new(Scalar, UInt8, UInt8), 64),
4940            (Key::new(X86_64_V3, UInt8, UInt8), 256),
4941            (Key::new(X86_64_V4, UInt8, UInt8), 320),
4942            (Key::new(Aarch64Neon, UInt8, UInt8), 128),
4943            (Key::new(Scalar, Int8, Int8), 64),
4944            (Key::new(X86_64_V3, Int8, Int8), 256),
4945            (Key::new(X86_64_V4, Int8, Int8), 320),
4946            (Key::new(Aarch64Neon, Int8, Int8), 128),
4947        ]
4948        .into_iter()
4949        .collect()
4950    });
4951
4952    macro_rules! test_bounds {
4953        (
4954            $function:ident,
4955            $left:ty,
4956            $left_ex:expr,
4957            $right:ty,
4958            $right_ex:expr
4959        ) => {
4960            #[test]
4961            fn $function() {
4962                let left: $left = $left_ex;
4963                let right: $right = $right_ex;
4964
4965                let left_type = <$left>::as_data_type();
4966                let right_type = <$right>::as_data_type();
4967
4968                // Scalar
4969                {
4970                    let max = MIRI_BOUNDS[&Key::new(Arch::Scalar, left_type, right_type)];
4971                    for dim in 0..max {
4972                        let left: Vec<$left> = vec![left; dim];
4973                        let right: Vec<$right> = vec![right; dim];
4974
4975                        let arch = diskann_wide::arch::Scalar;
4976                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
4977                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
4978                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
4979
4980                        let left = unaligned::Buffer::new(&left);
4981                        let right = unaligned::Buffer::new(&right);
4982                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
4983                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
4984                        simd_op(
4985                            &CosineStateless,
4986                            arch,
4987                            left.as_unaligned(),
4988                            right.as_unaligned(),
4989                        );
4990                    }
4991                }
4992
4993                #[cfg(target_arch = "x86_64")]
4994                if let Some(arch) = V3::new_checked() {
4995                    let max = MIRI_BOUNDS[&Key::new(Arch::X86_64_V3, left_type, right_type)];
4996                    for dim in 0..max {
4997                        let left: Vec<$left> = vec![left; dim];
4998                        let right: Vec<$right> = vec![right; dim];
4999
5000                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5001                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5002                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5003
5004                        let left = unaligned::Buffer::new(&left);
5005                        let right = unaligned::Buffer::new(&right);
5006                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5007                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5008                        simd_op(
5009                            &CosineStateless,
5010                            arch,
5011                            left.as_unaligned(),
5012                            right.as_unaligned(),
5013                        );
5014                    }
5015                }
5016
5017                #[cfg(target_arch = "x86_64")]
5018                if let Some(arch) = V4::new_checked_miri() {
5019                    let max = MIRI_BOUNDS[&Key::new(Arch::X86_64_V4, left_type, right_type)];
5020                    for dim in 0..max {
5021                        let left: Vec<$left> = vec![left; dim];
5022                        let right: Vec<$right> = vec![right; dim];
5023
5024                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5025                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5026                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5027
5028                        let left = unaligned::Buffer::new(&left);
5029                        let right = unaligned::Buffer::new(&right);
5030                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5031                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5032                        simd_op(
5033                            &CosineStateless,
5034                            arch,
5035                            left.as_unaligned(),
5036                            right.as_unaligned(),
5037                        );
5038                    }
5039                }
5040
5041                #[cfg(target_arch = "aarch64")]
5042                if let Some(arch) = Neon::new_checked() {
5043                    let max = MIRI_BOUNDS[&Key::new(Arch::Aarch64Neon, left_type, right_type)];
5044                    for dim in 0..max {
5045                        let left: Vec<$left> = vec![left; dim];
5046                        let right: Vec<$right> = vec![right; dim];
5047
5048                        simd_op(&L2, arch, left.as_slice(), right.as_slice());
5049                        simd_op(&IP, arch, left.as_slice(), right.as_slice());
5050                        simd_op(&CosineStateless, arch, left.as_slice(), right.as_slice());
5051                        let left = unaligned::Buffer::new(&left);
5052                        let right = unaligned::Buffer::new(&right);
5053                        simd_op(&L2, arch, left.as_unaligned(), right.as_unaligned());
5054                        simd_op(&IP, arch, left.as_unaligned(), right.as_unaligned());
5055                        simd_op(
5056                            &CosineStateless,
5057                            arch,
5058                            left.as_unaligned(),
5059                            right.as_unaligned(),
5060                        );
5061                    }
5062                }
5063            }
5064        };
5065    }
5066
5067    test_bounds!(miri_test_bounds_f32xf32, f32, 1.0f32, f32, 2.0f32);
5068    test_bounds!(
5069        miri_test_bounds_f16xf16,
5070        f16,
5071        diskann_wide::cast_f32_to_f16(1.0f32),
5072        f16,
5073        diskann_wide::cast_f32_to_f16(2.0f32)
5074    );
5075    test_bounds!(
5076        miri_test_bounds_f32xf16,
5077        f32,
5078        1.0f32,
5079        f16,
5080        diskann_wide::cast_f32_to_f16(2.0f32)
5081    );
5082    test_bounds!(miri_test_bounds_u8xu8, u8, 1u8, u8, 1u8);
5083    test_bounds!(miri_test_bounds_i8xi8, i8, 1i8, i8, 1i8);
5084}