Add support for masked loads & stores

Author: farnoy

crates/core_simd/src/intrinsics.rs

@@ -107,6 +107,15 @@ extern "platform-intrinsic" {
     /// like gather, but more spicy, as it writes instead of reads
     pub(crate) fn simd_scatter<T, U, V>(val: T, ptr: U, mask: V);
 
+    /// like a loop of reads offset from the same pointer
+    /// val: vector of values to select if a lane is masked
+    /// ptr: vector of pointers to read from
+    /// mask: a "wide" mask of integers, selects as if simd_select(mask, read(ptr), val)
+    /// note, the LLVM intrinsic accepts a mask vector of `<N x i1>`
+    pub(crate) fn simd_masked_load<T, U, V>(val: T, ptr: U, mask: V) -> T;
+    /// like masked_load, but more spicy, as it writes instead of reads
+    pub(crate) fn simd_masked_store<T, U, V>(val: T, ptr: U, mask: V);
+
     // {s,u}add.sat
     pub(crate) fn simd_saturating_add<T>(x: T, y: T) -> T;
 

crates/core_simd/src/vector.rs

@@ -311,6 +311,42 @@ where
         unsafe { self.store(slice.as_mut_ptr().cast()) }
     }
 
+    #[must_use]
+    #[inline]
+    pub fn masked_load_or(slice: &[T], or: Self) -> Self {
+        Self::masked_load_select(slice, Mask::splat(true), or)
+    }
+
+    #[must_use]
+    #[inline]
+    pub fn masked_load_select(slice: &[T], enable: Mask<isize, N>, or: Self) -> Self {
+        let ptr = slice.as_ptr();
+        let idxs = Simd::<usize, N>::from_slice(&[
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
+            24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
+            46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
+        ]);
+        let enable: Mask<isize, N> = enable & idxs.simd_lt(Simd::splat(slice.len()));
+        unsafe { Self::masked_load_select_ptr(ptr, enable, or) }
+    }
+
+    #[must_use]
+    #[inline]
+    pub unsafe fn masked_load_select_unchecked(
+        slice: &[T],
+        enable: Mask<isize, N>,
+        or: Self,
+    ) -> Self {
+        let ptr = slice.as_ptr();
+        unsafe { Self::masked_load_select_ptr(ptr, enable, or) }
+    }
+
+    #[must_use]
+    #[inline]
+    pub unsafe fn masked_load_select_ptr(ptr: *const T, enable: Mask<isize, N>, or: Self) -> Self {
+        unsafe { intrinsics::simd_masked_load(or, ptr, enable.to_int()) }
+    }
+
     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
     /// If an index is out-of-bounds, the element is instead selected from the `or` vector.
     ///
@@ -489,6 +525,29 @@ where
         unsafe { intrinsics::simd_gather(or, source, enable.to_int()) }
     }
 
+    #[inline]
+    pub fn masked_store(self, slice: &mut [T], enable: Mask<isize, N>) {
+        let ptr = slice.as_mut_ptr();
+        let idxs = Simd::<usize, N>::from_slice(&[
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
+            24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
+            46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
+        ]);
+        let enable: Mask<isize, N> = enable & idxs.simd_lt(Simd::splat(slice.len()));
+        unsafe { self.masked_store_ptr(ptr, enable) }
+    }
+
+    #[inline]
+    pub unsafe fn masked_store_unchecked(self, slice: &mut [T], enable: Mask<isize, N>) {
+        let ptr = slice.as_mut_ptr();
+        unsafe { self.masked_store_ptr(ptr, enable) }
+    }
+
+    #[inline]
+    pub unsafe fn masked_store_ptr(self, ptr: *mut T, enable: Mask<isize, N>) {
+        unsafe { intrinsics::simd_masked_store(self, ptr, enable.to_int()) }
+    }
+
     /// Writes the values in a SIMD vector to potentially discontiguous indices in `slice`.
     /// If an index is out-of-bounds, the write is suppressed without panicking.
     /// If two elements in the scattered vector would write to the same index