1#[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
19pub 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 #[inline(always)]
54 fn force_eval(x: f32) {
55 unsafe {
60 asm!(
61 "/* {0} */",
63 in(xmm_reg) x,
66 options(nostack, nomem, preserves_flags)
75 )
76 }
77 }
78 } else {
79 fn force_eval(_x: f32) {}
81 }
82}
83
84#[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 #[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 #[inline(always)]
129 fn arch(&self) -> A {
130 self.arch
131 }
132
133 #[inline(always)]
135 fn schema(&self) -> Schema {
136 self.schema
137 }
138
139 #[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
189pub trait MainLoop {
191 const BLOCK_SIZE: usize;
198
199 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}
243pub struct Strategy1x1;
246
247pub struct Strategy2x1;
250
251pub struct Strategy4x1;
254
255pub struct Strategy4x2;
258
259pub 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 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 #[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
485pub trait SIMDSchema<T, U, A: Architecture = diskann_wide::arch::Current>: Copy {
493 type SIMDWidth: Constant<Type = usize>;
496
497 type Accumulator: std::ops::Add<Output = Self::Accumulator> + std::fmt::Debug + Copy;
499
500 type Left: SIMDVector<Arch = A, Scalar = T, ConstLanes = Self::SIMDWidth>;
502
503 type Right: SIMDVector<Arch = A, Scalar = U, ConstLanes = Self::SIMDWidth>;
505
506 type Return;
509
510 type Main: MainLoop;
512
513 fn init(&self, arch: A) -> Self::Accumulator;
515
516 fn accumulate(
518 &self,
519 x: Self::Left,
520 y: Self::Right,
521 acc: Self::Accumulator,
522 ) -> Self::Accumulator;
523
524 #[inline(always)]
526 fn combine(&self, x: Self::Accumulator, y: Self::Accumulator) -> Self::Accumulator {
527 x + y
528 }
529
530 #[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 let a = Self::Left::load_simd_first(arch, x, len);
559
560 let b = Self::Right::load_simd_first(arch, y, len);
563 self.accumulate(a, b, acc)
564 }
565
566 fn reduce(&self, x: Self::Accumulator) -> Self::Return;
570
571 #[inline(always)]
576 fn get_simd_width() -> usize {
577 Self::SIMDWidth::value()
578 }
579
580 #[inline(always)]
586 fn get_main_bocksize() -> usize {
587 Self::Main::BLOCK_SIZE
588 }
589
590 #[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
603pub trait ResumableSIMDSchema<T, U, A = diskann_wide::arch::Current>: Copy
611where
612 A: Architecture,
613{
614 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#[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 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 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 let loader: Loader<S, L, R, A> = Loader::new(arch, *schema, px, py, len);
729
730 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 s0 = unsafe { schema.epilogue(arch, px.add(i), py.add(i), remainder, s0) };
744 }
745
746 schema.reduce(s0)
747}
748
749#[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 let left = unsafe { left.add(i).read_unaligned() };
770 let right = unsafe { right.add(i).read_unaligned() };
772 acc = f(acc, left, right);
773 }
774 acc
775}
776
777#[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 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 let vx = unsafe { x.add(i).read_unaligned() };
943 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 #[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 #[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 #[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 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 #[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 #[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 let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1328 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 #[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 #[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 #[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 let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
1507 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#[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#[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 #[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 #[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 let vx = unsafe { x.add(i).read_unaligned() };
1713 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 #[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 #[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 #[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 #[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 let x: i16s = x.into();
2133 let y: i16s = y.into();
2134 acc.dot_simd(x, y)
2135 }
2136
2137 #[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 #[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#[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#[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 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 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 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 let denominator = normx.sqrt() * normy.sqrt();
2338 let prod = self.xy.sum_tree().as_f32_lossy();
2339
2340 force_eval(denominator);
2348 force_eval(prod);
2349
2350 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 #[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#[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 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 #[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 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 #[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 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 let vx = unsafe { x.add(i).read_unaligned() };
2504 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 #[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 #[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 #[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 #[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 let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2871 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 #[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 let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
2938 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 #[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 #[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 let vx: u32 = unsafe { x.add(i).read_unaligned() }.into();
3077 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 #[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 let vx: i32 = unsafe { x.add(i).read_unaligned() }.into();
3144 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#[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#[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 #[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 #[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 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 #[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 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 #[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 #[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, 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 #[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 #[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#[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 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 {
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 {
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 {
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 {
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 {
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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 #[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 {
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}