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