Add alignment parameter to simd_masked_{load,store}

Author: sayantn

core/src/intrinsics/simd.rs

@@ -2,6 +2,8 @@
 //!
 //! In this module, a "vector" is any `repr(simd)` type.
 
+use crate::marker::ConstParamTy;
+
 /// Inserts an element into a vector, returning the updated vector.
 ///
 /// `T` must be a vector with element type `U`, and `idx` must be `const`.
@@ -377,6 +379,19 @@ pub unsafe fn simd_gather<T, U, V>(val: T, ptr: U, mask: V) -> T;
 #[rustc_nounwind]
 pub unsafe fn simd_scatter<T, U, V>(val: T, ptr: U, mask: V);
 
+/// A type for alignment options for SIMD masked load/store intrinsics.
+#[derive(Debug, ConstParamTy, PartialEq, Eq)]
+pub enum SimdAlign {
+    // These values must match the compiler's `SimdAlign` defined in
+    // `rustc_middle/src/ty/consts/int.rs`!
+    /// No alignment requirements on the pointer
+    Unaligned = 0,
+    /// The pointer must be aligned to the element type of the SIMD vector
+    Element = 1,
+    /// The pointer must be aligned to the SIMD vector type
+    Vector = 2,
+}
+
 /// Reads a vector of pointers.
 ///
 /// `T` must be a vector.
@@ -392,13 +407,12 @@ pub unsafe fn simd_scatter<T, U, V>(val: T, ptr: U, mask: V);
 /// `val`.
 ///
 /// # Safety
-/// Unmasked values in `T` must be readable as if by `<ptr>::read` (e.g. aligned to the element
-/// type).
+/// `ptr` must be aligned according to the `ALIGN` parameter, see [`SimdAlign`] for details.
 ///
 /// `mask` must only contain `0` or `!0` values.
 #[rustc_intrinsic]
 #[rustc_nounwind]
-pub unsafe fn simd_masked_load<V, U, T>(mask: V, ptr: U, val: T) -> T;
+pub unsafe fn simd_masked_load<V, U, T, const ALIGN: SimdAlign>(mask: V, ptr: U, val: T) -> T;
 
 /// Writes to a vector of pointers.
 ///
@@ -414,13 +428,12 @@ pub unsafe fn simd_masked_load<V, U, T>(mask: V, ptr: U, val: T) -> T;
 /// Otherwise if the corresponding value in `mask` is `0`, do nothing.
 ///
 /// # Safety
-/// Unmasked values in `T` must be writeable as if by `<ptr>::write` (e.g. aligned to the element
-/// type).
+/// `ptr` must be aligned according to the `ALIGN` parameter, see [`SimdAlign`] for details.
 ///
 /// `mask` must only contain `0` or `!0` values.
 #[rustc_intrinsic]
 #[rustc_nounwind]
-pub unsafe fn simd_masked_store<V, U, T>(mask: V, ptr: U, val: T);
+pub unsafe fn simd_masked_store<V, U, T, const ALIGN: SimdAlign>(mask: V, ptr: U, val: T);
 
 /// Adds two simd vectors elementwise, with saturation.
 ///

portable-simd/crates/core_simd/src/vector.rs

@@ -474,7 +474,14 @@ where
         or: Self,
     ) -> Self {
         // SAFETY: The safety of reading elements through `ptr` is ensured by the caller.
-        unsafe { core::intrinsics::simd::simd_masked_load(enable.to_int(), ptr, or) }
+        unsafe {
+            core::intrinsics::simd::simd_masked_load::<
+                _,
+                _,
+                _,
+                { core::intrinsics::simd::SimdAlign::Element },
+            >(enable.to_int(), ptr, or)
+        }
     }
 
     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
@@ -723,7 +730,14 @@ where
     #[inline]
     pub unsafe fn store_select_ptr(self, ptr: *mut T, enable: Mask<<T as SimdElement>::Mask, N>) {
         // SAFETY: The safety of writing elements through `ptr` is ensured by the caller.
-        unsafe { core::intrinsics::simd::simd_masked_store(enable.to_int(), ptr, self) }
+        unsafe {
+            core::intrinsics::simd::simd_masked_store::<
+                _,
+                _,
+                _,
+                { core::intrinsics::simd::SimdAlign::Element },
+            >(enable.to_int(), ptr, self)
+        }
     }
 
     /// Writes the values in a SIMD vector to potentially discontiguous indices in `slice`.