[diskann-wide] Add u64x8

diskann-wide/src/arch/aarch64/double.rs

@@ -28,6 +28,7 @@ doubled::double_mask!(4, mask64x2);
 doubled::double_mask!(64, Doubled<mask8x16>);
 doubled::double_mask!(32, Doubled<mask16x8>);
 doubled::double_mask!(16, Doubled<mask32x4>);
+doubled::double_mask!(8, Doubled<mask64x2>);
 
 macro_rules! double_alias {
     ($type:ident, $scalar:ty, $lanes:literal, $subtype:ty) => {
@@ -64,6 +65,7 @@ double_alias!(f16x16, f16, 16, f16x8);
 
 double_alias!(u8x64, u8, 64, u8x32);
 double_alias!(u32x16, u32, 16, u32x8);
+double_alias!(u64x8, u64, 8, u64x4);
 
 double_alias!(i8x64, i8, 64, i8x32);
 double_alias!(i16x32, i16, 32, i16x16);
@@ -295,6 +297,16 @@ mod tests {
         test_utils::ops::test_splitjoin!(u64x4 => u64x2, 0x2e301b7e12090d5c, test_neon());
     }
 
+    mod test_u64x8 {
+        use super::*;
+        standard_tests!(u64x8, u64, 8);
+
+        // Bit ops
+        test_utils::ops::test_bitops!(u64x8, 0xc4491a44af4aa58e, test_neon());
+        test_utils::ops::test_sumtree!(u64x8, 0x529c27f62ea171ec, test_neon());
+        test_utils::ops::test_splitjoin!(u64x8 => u64x4, 0x2e301b7e12090d5c, test_neon());
+    }
+
     // i8s
     mod test_i8x32 {
         use super::*;

diskann-wide/src/arch/aarch64/mod.rs

@@ -83,6 +83,7 @@ pub use double::u32x8;
 pub use double::u32x16;
 
 pub use double::u64x4;
+pub use double::u64x8;
 
 // Internal helpers
 mod macros;

diskann-wide/src/arch/emulated/mod.rs

@@ -42,6 +42,7 @@ pub type u32x16 = Emulated<u32, 16>;
 
 pub type u64x2 = Emulated<u64, 2>;
 pub type u64x4 = Emulated<u64, 4>;
+pub type u64x8 = Emulated<u64, 8>;
 
 /// A safe architecture that is guaranteed to be compatible with the machine that
 /// a Rust program was compiled for.

diskann-wide/src/arch/mod.rs

@@ -555,6 +555,7 @@ pub trait Architecture: sealed::Sealed {
     type mask_u32x16: SIMDMask + SIMDSelect<Self::u32x16>;
     type mask_u64x2: SIMDMask;
     type mask_u64x4: SIMDMask;
+    type mask_u64x8: SIMDMask;
 
     /////////////////
     //-- vectors --//
@@ -699,6 +700,12 @@ pub trait Architecture: sealed::Sealed {
         + SIMDUnsigned
         + SIMDSumTree
     );
+    vector!(
+        u64x8: <Self, u64, 8, mask_u64x8>
+        + SplitJoin<Halved = Self::u64x4>
+        + SIMDUnsigned
+        + SIMDSumTree
+    );
 
     //---------//
     // Methods //
@@ -1209,6 +1216,7 @@ macro_rules! maskdef {
 
             mask_u64x2 = u64x2,
             mask_u64x4 = u64x4,
+            mask_u64x8 = u64x8,
         );
     };
 }
@@ -1245,6 +1253,7 @@ macro_rules! typedef {
 
             u64x2,
             u64x4,
+            u64x8,
         );
     };
     ($repr:ident) => {

diskann-wide/src/arch/x86_64/v3/masks.rs

@@ -671,6 +671,7 @@ impl From<mask64x4> for BitMask<4, V3> {
 // Native Masks
 doubled::double_mask!(64, mask8x32);
 doubled::double_mask!(16, mask32x8);
+doubled::double_mask!(8, mask64x4);
 
 // Bit Mask
 doubled::double_mask!(32, BitMask<16, V3>);
@@ -986,6 +987,7 @@ mod test_masks {
     // Double
     test_simdmask!(Doubled<mask8x32>, 64, nop);
     test_simdmask!(Doubled<mask32x8>, 16, nop);
+    test_simdmask!(Doubled<mask64x4>, 8, nop);
 
     // Type alias to work around limitations in `test_simdmask`.
     type BitMask16V3 = BitMask<16, V3>;

diskann-wide/src/arch/x86_64/v3/mod.rs

@@ -84,6 +84,9 @@ pub use u64x2_::u64x2;
 pub mod u64x4_;
 pub use u64x4_::u64x4;
 
+pub mod u64x8_;
+pub use u64x8_::u64x8;
+
 // Masks
 pub mod masks;
 

diskann-wide/src/arch/x86_64/v3/u64x8_.rs

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use super::u64x4;
+use crate::doubled;
+
+/////////////////////
+// 64-bit unsigned //
+/////////////////////
+
+doubled::double_vector!(u64, 8, u64x4);
+doubled::double_scalar_shift!(Doubled<u64x4>);
+
+#[allow(non_camel_case_types)]
+pub type u64x8 = doubled::Doubled<u64x4>;
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::{arch::x86_64::V3, reference::ReferenceScalarOps, test_utils};
+
+    #[test]
+    fn miri_test_load() {
+        if let Some(arch) = V3::new_checked_uncached() {
+            test_utils::test_load_simd::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    #[test]
+    fn miri_test_store() {
+        if let Some(arch) = V3::new_checked_uncached() {
+            test_utils::test_store_simd::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    #[test]
+    fn test_constructors() {
+        if let Some(arch) = V3::new_checked_uncached() {
+            test_utils::ops::test_splat::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    test_utils::ops::test_add!(u64x8, 0xeaee2fd0398fe357, V3::new_checked_uncached());
+    test_utils::ops::test_sub!(u64x8, 0x40af040b0c2c1e28, V3::new_checked_uncached());
+    test_utils::ops::test_mul!(u64x8, 0x68f68933a29c5ea9, V3::new_checked_uncached());
+    test_utils::ops::test_fma!(u64x8, 0x31bc9d25e91e6744, V3::new_checked_uncached());
+
+    test_utils::ops::test_cmp!(u64x8, 0x0beda0dd5141ec40, V3::new_checked_uncached());
+    test_utils::ops::test_splitjoin!(u64x8 => u64x4, 0xb151fcd6141b10c9, V3::new_checked_uncached());
+
+    test_utils::ops::test_sumtree!(u64x8, 0x529c27f62ea171ec, V3::new_checked_uncached());
+
+    test_utils::ops::test_bitops!(u64x8, 0xb1ac2e16327a8d5e, V3::new_checked_uncached());
+}

diskann-wide/src/arch/x86_64/v4/mod.rs

@@ -157,6 +157,9 @@ maybe_miri!(u64x2_, u64x2, u64, 2);
 pub mod u64x4_;
 maybe_miri!(u64x4_, u64x4, u64, 4);
 
+pub mod u64x8_;
+maybe_miri!(u64x8_, u64x8, u64, 8);
+
 // Conversions between intrinsics
 mod conversion;
 

diskann-wide/src/arch/x86_64/v4/u64x8_.rs

@@ -0,0 +1,126 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::arch::x86_64::*;
+
+use crate::{
+    BitMask,
+    arch::x86_64::{
+        V4,
+        common::AllOnes,
+        macros::{self, X86Default, X86LoadStore, X86Splat},
+        v4::u64x4_::u64x4,
+    },
+    constant::Const,
+    helpers,
+    traits::{SIMDMask, SIMDMulAdd, SIMDPopcount, SIMDSumTree, SIMDVector},
+};
+
+/////////////////////
+// 64-bit unsigned //
+/////////////////////
+
+macros::x86_define_register!(u64x8, __m512i, BitMask<8, V4>, u64, 8, V4);
+macros::x86_define_splat!(u64x8 as i64, _mm512_set1_epi64, "avx512f");
+macros::x86_define_default!(u64x8, _mm512_setzero_si512, "avx512f");
+macros::x86_splitjoin!(__m512i, u64x8, u64x4);
+
+helpers::unsafe_map_binary_op!(u64x8, std::ops::Add, add, _mm512_add_epi64, "avx512f");
+helpers::unsafe_map_binary_op!(u64x8, std::ops::Sub, sub, _mm512_sub_epi64, "avx512f");
+helpers::unsafe_map_unary_op!(
+    u64x8,
+    SIMDPopcount,
+    popcount_simd,
+    _mm512_popcnt_epi64,
+    "avx512vpopcntdq"
+);
+helpers::unsafe_map_binary_op!(u64x8, std::ops::Mul, mul, _mm512_mullo_epi64, "avx512dq");
+
+helpers::unsafe_map_binary_op!(u64x8, std::ops::BitAnd, bitand, _mm512_and_si512, "avx512f");
+helpers::unsafe_map_binary_op!(u64x8, std::ops::BitOr, bitor, _mm512_or_si512, "avx512f");
+helpers::unsafe_map_binary_op!(u64x8, std::ops::BitXor, bitxor, _mm512_xor_si512, "avx512f");
+helpers::unsafe_map_binary_op!(u64x8, std::ops::Shr, shr, _mm512_srlv_epi64, "avx512f");
+helpers::unsafe_map_binary_op!(u64x8, std::ops::Shl, shl, _mm512_sllv_epi64, "avx512f");
+helpers::scalar_shift_by_splat!(u64x8, u64);
+
+impl std::ops::Not for u64x8 {
+    type Output = Self;
+
+    #[inline(always)]
+    fn not(self) -> Self {
+        self ^ Self::from_underlying(self.arch(), <Self as SIMDVector>::Underlying::all_ones())
+    }
+}
+
+impl SIMDMulAdd for u64x8 {
+    #[inline(always)]
+    fn mul_add_simd(self, rhs: Self, accumulator: Self) -> Self {
+        self * rhs + accumulator
+    }
+}
+
+macros::x86_avx512_load_store!(
+    u64x8,
+    _mm512_loadu_epi64,
+    _mm512_maskz_loadu_epi64,
+    _mm512_storeu_epi64,
+    _mm512_mask_storeu_epi64,
+    i64,
+    "avx512f"
+);
+
+macros::x86_avx512_int_comparisons!(u64x8, _mm512_cmp_epu64_mask, "avx512f");
+
+impl SIMDSumTree for u64x8 {
+    #[inline(always)]
+    fn sum_tree(self) -> u64 {
+        // SAFETY: `_mm512_reduce_add_epi64` requires AVX-512F, implied by V4.
+        unsafe { _mm512_reduce_add_epi64(self.0) as u64 }
+    }
+}
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod test_x86_u64 {
+    use super::*;
+    use crate::{reference::ReferenceScalarOps, test_utils};
+
+    #[test]
+    fn miri_test_load() {
+        if let Some(arch) = V4::new_checked_uncached() {
+            test_utils::test_load_simd::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    #[test]
+    fn miri_test_store() {
+        if let Some(arch) = V4::new_checked_uncached() {
+            test_utils::test_store_simd::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    #[test]
+    fn test_constructors() {
+        if let Some(arch) = V4::new_checked_uncached() {
+            test_utils::ops::test_splat::<u64, 8, u64x8>(arch);
+        }
+    }
+
+    test_utils::ops::test_add!(u64x8, 0xeaee2fd0398fe357, V4::new_checked_uncached());
+    test_utils::ops::test_sub!(u64x8, 0x40af040b0c2c1e28, V4::new_checked_uncached());
+    test_utils::ops::test_mul!(u64x8, 0x68f68933a29c5ea9, V4::new_checked_uncached());
+    test_utils::ops::test_fma!(u64x8, 0x31bc9d25e91e6744, V4::new_checked_uncached());
+
+    test_utils::ops::test_cmp!(u64x8, 0x0beda0dd5141ec40, V4::new_checked_uncached());
+    test_utils::ops::test_splitjoin!(u64x8 => u64x4, 0xb151fcd6141b10c9, V4::new_checked_uncached());
+
+    test_utils::ops::test_sumtree!(u64x8, 0x529c27f62ea171ec, V4::new_checked_uncached());
+
+    test_utils::ops::test_bitops!(u64x8, 0xb1ac2e16327a8d5e, V4::new_checked_uncached());
+    test_utils::ops::test_popcount!(u64x8, 0xf23de3226c0141be, V4::new_checked_uncached());
+}