[diskann-wide] Add more 64-bit neon types

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

@@ -6,8 +6,8 @@
 use half::f16;
 
 use crate::{
-    Emulated, SIMDAbs, SIMDMask, SIMDMinMax, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd, SIMDSelect,
-    SIMDSumTree, SIMDVector, constant::Const, helpers,
+    AsSIMD, Emulated, SIMDAbs, SIMDMask, SIMDMinMax, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd,
+    SIMDSelect, SIMDSumTree, SIMDVector, constant::Const, helpers,
 };
 
 // AArch64 masks
@@ -38,26 +38,36 @@ macros::aarch64_define_fma!(f32x4, vfmaq_f32);
 impl SIMDMinMax for f32x4 {
     #[inline(always)]
     fn min_simd(self, rhs: Self) -> Self {
-        // SAFETY: `vminnmq_f32` requires "neon", implied by the `Neon` architecture.
-        Self(unsafe { vminnmq_f32(self.0, rhs.0) })
+        self.min_simd_standard(rhs)
     }
 
     #[inline(always)]
     fn min_simd_standard(self, rhs: Self) -> Self {
-        // SAFETY: `vminnmq_f32` requires "neon", implied by the `Neon` architecture.
-        Self(unsafe { vminnmq_f32(self.0, rhs.0) })
+        if cfg!(miri) {
+            self.emulated()
+                .min_simd_standard(rhs.emulated())
+                .as_simd(self.arch())
+        } else {
+            // SAFETY: `vminnmq_f32` requires "neon", implied by the `Neon` architecture.
+            Self(unsafe { vminnmq_f32(self.0, rhs.0) })
+        }
     }
 
     #[inline(always)]
     fn max_simd(self, rhs: Self) -> Self {
-        // SAFETY: `vmaxnmq_f32` requires "neon", implied by the `Neon` architecture.
-        Self(unsafe { vmaxnmq_f32(self.0, rhs.0) })
+        self.max_simd_standard(rhs)
     }
 
     #[inline(always)]
     fn max_simd_standard(self, rhs: Self) -> Self {
-        // SAFETY: `vmaxnmq_f32` requires "neon", implied by the `Neon` architecture.
-        Self(unsafe { vmaxnmq_f32(self.0, rhs.0) })
+        if cfg!(miri) {
+            self.emulated()
+                .max_simd_standard(rhs.emulated())
+                .as_simd(self.arch())
+        } else {
+            // SAFETY: `vmaxnmq_f32` requires "neon", implied by the `Neon` architecture.
+            Self(unsafe { vmaxnmq_f32(self.0, rhs.0) })
+        }
     }
 }
 

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

@@ -0,0 +1,103 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+
+use crate::{
+    Emulated, SIMDAbs, SIMDMask, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd, SIMDVector,
+    constant::Const, helpers,
+};
+
+// AArch64 masks
+use super::{
+    Neon, internal,
+    macros::{self, AArchLoadStore, AArchSplat},
+    masks::mask16x4,
+};
+
+// AArch64 intrinsics
+use std::arch::aarch64::*;
+
+///////////////////
+// 16-bit signed //
+///////////////////
+
+macros::aarch64_define_register!(i16x4, int16x4_t, mask16x4, i16, 4, Neon);
+macros::aarch64_define_splat!(i16x4, vmov_n_s16);
+macros::aarch64_define_loadstore!(i16x4, vld1_s16, internal::load_first::i16x4, vst1_s16, 4);
+
+helpers::unsafe_map_binary_op!(i16x4, std::ops::Add, add, vadd_s16, "neon");
+helpers::unsafe_map_binary_op!(i16x4, std::ops::Sub, sub, vsub_s16, "neon");
+helpers::unsafe_map_binary_op!(i16x4, std::ops::Mul, mul, vmul_s16, "neon");
+helpers::unsafe_map_unary_op!(i16x4, SIMDAbs, abs_simd, vabs_s16, "neon");
+macros::aarch64_define_fma!(i16x4, vmla_s16);
+
+macros::aarch64_define_cmp!(
+    i16x4,
+    vceq_s16,
+    (vmvn_u16),
+    vclt_s16,
+    vcle_s16,
+    vcgt_s16,
+    vcge_s16
+);
+macros::aarch64_define_bitops!(
+    i16x4,
+    vmvn_s16,
+    vand_s16,
+    vorr_s16,
+    veor_s16,
+    (
+        vshl_s16,
+        16,
+        vneg_s16,
+        vmin_u16,
+        vreinterpret_s16_u16,
+        vreinterpret_u16_s16
+    ),
+    (u16, i16, vmov_n_s16),
+);
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{arch::aarch64::test_neon, reference::ReferenceScalarOps, test_utils};
+
+    #[test]
+    fn miri_test_load() {
+        if let Some(arch) = test_neon() {
+            test_utils::test_load_simd::<i16, 4, i16x4>(arch);
+        }
+    }
+
+    #[test]
+    fn miri_test_store() {
+        if let Some(arch) = test_neon() {
+            test_utils::test_store_simd::<i16, 4, i16x4>(arch);
+        }
+    }
+
+    // constructors
+    #[test]
+    fn test_constructors() {
+        if let Some(arch) = test_neon() {
+            test_utils::ops::test_splat::<i16, 4, i16x4>(arch);
+        }
+    }
+
+    // Ops
+    test_utils::ops::test_add!(i16x4, 0x3017fd73c99cc633, test_neon());
+    test_utils::ops::test_sub!(i16x4, 0xfc627f10b5f8db8a, test_neon());
+    test_utils::ops::test_mul!(i16x4, 0x0f4caa80eceaa523, test_neon());
+    test_utils::ops::test_fma!(i16x4, 0xb8f702ba85375041, test_neon());
+    test_utils::ops::test_abs!(i16x4, 0xb8f702ba85375041, test_neon());
+
+    test_utils::ops::test_cmp!(i16x4, 0x941757bd5cc641a1, test_neon());
+
+    // Bit ops
+    test_utils::ops::test_bitops!(i16x4, 0xd62d8de09f82ed4e, test_neon());
+}

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

@@ -0,0 +1,140 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT license.
+ */
+
+use crate::{
+    Emulated, SIMDAbs, SIMDCast, SIMDMask, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd, SIMDSelect,
+    SIMDSumTree, SIMDVector, constant::Const, helpers,
+};
+
+// AArch64 masks
+use super::{
+    Neon, f32x2, internal,
+    macros::{self, AArchLoadStore, AArchSplat},
+    masks::mask32x2,
+};
+
+// AArch64 intrinsics
+use std::arch::aarch64::*;
+
+///////////////////
+// 32-bit signed //
+///////////////////
+
+macros::aarch64_define_register!(i32x2, int32x2_t, mask32x2, i32, 2, Neon);
+macros::aarch64_define_splat!(i32x2, vmov_n_s32);
+macros::aarch64_define_loadstore!(i32x2, vld1_s32, internal::load_first::i32x2, vst1_s32, 2);
+
+helpers::unsafe_map_binary_op!(i32x2, std::ops::Add, add, vadd_s32, "neon");
+helpers::unsafe_map_binary_op!(i32x2, std::ops::Sub, sub, vsub_s32, "neon");
+helpers::unsafe_map_binary_op!(i32x2, std::ops::Mul, mul, vmul_s32, "neon");
+helpers::unsafe_map_unary_op!(i32x2, SIMDAbs, abs_simd, vabs_s32, "neon");
+macros::aarch64_define_fma!(i32x2, vmla_s32);
+
+macros::aarch64_define_cmp!(
+    i32x2,
+    vceq_s32,
+    (vmvn_u32),
+    vclt_s32,
+    vcle_s32,
+    vcgt_s32,
+    vcge_s32
+);
+macros::aarch64_define_bitops!(
+    i32x2,
+    vmvn_s32,
+    vand_s32,
+    vorr_s32,
+    veor_s32,
+    (
+        vshl_s32,
+        32,
+        vneg_s32,
+        vmin_u32,
+        vreinterpret_s32_u32,
+        vreinterpret_u32_s32
+    ),
+    (u32, i32, vmov_n_s32),
+);
+
+impl SIMDSumTree for i32x2 {
+    #[inline(always)]
+    fn sum_tree(self) -> i32 {
+        if cfg!(miri) {
+            self.emulated().sum_tree()
+        } else {
+            // SAFETY: Allowed by the `Neon` architecture.
+            unsafe { vaddv_s32(self.0) }
+        }
+    }
+}
+
+impl SIMDSelect<i32x2> for mask32x2 {
+    #[inline(always)]
+    fn select(self, x: i32x2, y: i32x2) -> i32x2 {
+        // SAFETY: Allowed by the `Neon` architecture.
+        i32x2(unsafe { vbsl_s32(self.0, x.0, y.0) })
+    }
+}
+
+//-------------//
+// Conversions //
+//-------------//
+
+helpers::unsafe_map_cast!(
+    i32x2 => (f32, f32x2),
+    vcvt_f32_s32,
+    "neon"
+);
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{arch::aarch64::test_neon, reference::ReferenceScalarOps, test_utils};
+
+    #[test]
+    fn miri_test_load() {
+        if let Some(arch) = test_neon() {
+            test_utils::test_load_simd::<i32, 2, i32x2>(arch);
+        }
+    }
+
+    #[test]
+    fn miri_test_store() {
+        if let Some(arch) = test_neon() {
+            test_utils::test_store_simd::<i32, 2, i32x2>(arch);
+        }
+    }
+
+    // constructors
+    #[test]
+    fn test_constructors() {
+        if let Some(arch) = test_neon() {
+            test_utils::ops::test_splat::<i32, 2, i32x2>(arch);
+        }
+    }
+
+    // Ops
+    test_utils::ops::test_add!(i32x2, 0x3017fd73c99cc633, test_neon());
+    test_utils::ops::test_sub!(i32x2, 0xfc627f10b5f8db8a, test_neon());
+    test_utils::ops::test_mul!(i32x2, 0x0f4caa80eceaa523, test_neon());
+    test_utils::ops::test_fma!(i32x2, 0xb8f702ba85375041, test_neon());
+    test_utils::ops::test_abs!(i32x2, 0xb8f702ba85375041, test_neon());
+
+    test_utils::ops::test_cmp!(i32x2, 0x941757bd5cc641a1, test_neon());
+
+    // Bit ops
+    test_utils::ops::test_bitops!(i32x2, 0xd62d8de09f82ed4e, test_neon());
+    test_utils::ops::test_select!(i32x2, 0xd62d8de09f82ed4e, test_neon());
+
+    // Reductions
+    test_utils::ops::test_sumtree!(i32x2, 0xb9ac82ab23a855da, test_neon());
+
+    // Conversions
+    test_utils::ops::test_cast!(i32x2 => f32x2, 0xba8fe343fc9dbeff, test_neon());
+}

diskann-wide/src/arch/aarch64/internal/load_first.rs

@@ -51,6 +51,50 @@ pub(in crate::arch::aarch64) unsafe fn u16x4(_: Neon, ptr: *const u16, first: us
     unsafe { vcreate_u16(load_first_of_8_bytes(ptr.cast::<u8>(), 2 * first)) }
 }
 
+/// Load the first `first` elements from `ptr` into a `int16x4_t` register.
+///
+/// # Safety
+///
+/// The caller must ensure `[ptr, ptr + first)` is readable. The presence of `Neon`
+/// enables the use of "neon" intrinsics.
+#[inline(always)]
+pub(in crate::arch::aarch64) unsafe fn i16x4(_: Neon, ptr: *const i16, first: usize) -> int16x4_t {
+    // SAFETY: Pointer access inherited from caller. `Neon` enables "neon" intrinsics.
+    unsafe { vcreate_s16(load_first_of_8_bytes(ptr.cast::<u8>(), 2 * first)) }
+}
+
+/// Load the first `first` elements from `ptr` into a `uint32x2_t` register.
+///
+/// # Safety
+///
+/// The caller must ensure `[ptr, ptr + first)` is readable. The presence of `Neon`
+/// enables the use of "neon" intrinsics.
+#[inline(always)]
+pub(in crate::arch::aarch64) unsafe fn u32x2(
+    arch: Neon,
+    ptr: *const u32,
+    first: usize,
+) -> uint32x2_t {
+    // SAFETY: Pointer access inherited from caller. `Neon` enables "neon" intrinsics.
+    unsafe { load_first_32x2(arch, ptr, first) }
+}
+
+/// Load the first `first` elements from `ptr` into a `int32x2_t` register.
+///
+/// # Safety
+///
+/// The caller must ensure `[ptr, ptr + first)` is readable. The presence of `Neon`
+/// enables the use of "neon" intrinsics.
+#[inline(always)]
+pub(in crate::arch::aarch64) unsafe fn i32x2(
+    arch: Neon,
+    ptr: *const i32,
+    first: usize,
+) -> int32x2_t {
+    // SAFETY: Pointer access inherited from caller. `Neon` enables "neon" intrinsics.
+    unsafe { vreinterpret_s32_u32(load_first_32x2(arch, ptr.cast::<u32>(), first)) }
+}
+
 /// Load the first `first` elements from `ptr` into a `float32x2_t` register.
 ///
 /// # Safety
@@ -359,7 +403,7 @@ unsafe fn load_first_32x4(_: Neon, ptr: *const u32, first: usize) -> uint32x4_t
             vld1q_u32(ptr)
         } else if first == 3 {
             let lo = vld1_u32(ptr);
-            let hi = vld1_lane_u32(ptr.add(2), vcreate_u32(0), 0);
+            let hi = vset_lane_u32(ptr.add(2).read_unaligned(), vcreate_u32(0), 0);
             vcombine_u32(lo, hi)
         } else if first == 2 {
             vcombine_u32(vld1_u32(ptr), vcreate_u32(0))

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

@@ -34,9 +34,15 @@ pub use u8x8_::u8x8;
 pub mod u8x16_;
 pub use u8x16_::u8x16;
 
+pub mod u16x4_;
+pub use u16x4_::u16x4;
+
 pub mod u16x8_;
 pub use u16x8_::u16x8;
 
+pub mod u32x2_;
+pub use u32x2_::u32x2;
+
 pub mod u32x4_;
 pub use u32x4_::u32x4;
 
@@ -50,9 +56,15 @@ pub use i8x8_::i8x8;
 pub mod i8x16_;
 pub use i8x16_::i8x16;
 
+pub mod i16x4_;
+pub use i16x4_::i16x4;
+
 pub mod i16x8_;
 pub use i16x8_::i16x8;
 
+pub mod i32x2_;
+pub use i32x2_::i32x2;
+
 pub mod i32x4_;
 pub use i32x4_::i32x4;