Add alignment parameter to simd_masked_{load,store}

Author: sayantn

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1829,10 +1829,11 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
     }
 
     if name == sym::simd_masked_load {
-        // simd_masked_load(mask: <N x i{M}>, pointer: *_ T, values: <N x T>) -> <N x T>
+        // simd_masked_load(mask: <N x i{M}>, pointer: *_ T, values: <N x T>, alignment: u32) -> <N x T>
         // * N: number of elements in the input vectors
         // * T: type of the element to load
         // * M: any integer width is supported, will be truncated to i1
+        // * `alignment`: must be a power of two constant
         // Loads contiguous elements from memory behind `pointer`, but only for
         // those lanes whose `mask` bit is enabled.
         // The memory addresses corresponding to the “off” lanes are not accessed.
@@ -1844,10 +1845,18 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         // The second argument must be a pointer matching the element type
         let pointer_ty = args[1].layout.ty;
 
-        // The last argument is a passthrough vector providing values for disabled lanes
+        // The third argument is a passthrough vector providing values for disabled lanes
         let values_ty = args[2].layout.ty;
         let (values_len, values_elem) = require_simd!(values_ty, SimdThird);
 
+        // The fourth argument is the alignment, must be a power of two integer constant
+        let alignment = bx
+            .const_to_opt_u128(args[3].immediate(), false)
+            .expect("typeck should have ensure that this is a const");
+        if !alignment.is_power_of_two() {
+            return_error!(InvalidMonomorphization::AlignmentNotPowerOfTwo { span, name });
+        }
+
         require_simd!(ret_ty, SimdReturn);
 
         // Of the same length:
@@ -1893,7 +1902,7 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         let mask = vector_mask_to_bitmask(bx, args[0].immediate(), m_elem_bitwidth, mask_len);
 
         // Alignment of T, must be a constant integer value:
-        let alignment = bx.const_i32(bx.align_of(values_elem).bytes() as i32);
+        let alignment = bx.const_i32(alignment as i32);
 
         let llvm_pointer = bx.type_ptr();
 
@@ -1908,10 +1917,11 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
     }
 
     if name == sym::simd_masked_store {
-        // simd_masked_store(mask: <N x i{M}>, pointer: *mut T, values: <N x T>) -> ()
+        // simd_masked_store(mask: <N x i{M}>, pointer: *mut T, values: <N x T>, alignment: u32) -> ()
         // * N: number of elements in the input vectors
         // * T: type of the element to load
         // * M: any integer width is supported, will be truncated to i1
+        // * `alignment`: must be a power of two constant
         // Stores contiguous elements to memory behind `pointer`, but only for
         // those lanes whose `mask` bit is enabled.
         // The memory addresses corresponding to the “off” lanes are not accessed.
@@ -1923,10 +1933,18 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         // The second argument must be a pointer matching the element type
         let pointer_ty = args[1].layout.ty;
 
-        // The last argument specifies the values to store to memory
+        // The third argument specifies the values to store to memory
         let values_ty = args[2].layout.ty;
         let (values_len, values_elem) = require_simd!(values_ty, SimdThird);
 
+        // The fourth argument is the alignment, must be a power of two integer constant
+        let alignment = bx
+            .const_to_opt_u128(args[3].immediate(), false)
+            .expect("typeck should have ensure that this is a const");
+        if !alignment.is_power_of_two() {
+            return_error!(InvalidMonomorphization::AlignmentNotPowerOfTwo { span, name });
+        }
+
         // Of the same length:
         require!(
             values_len == mask_len,
@@ -1965,8 +1983,7 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
 
         let mask = vector_mask_to_bitmask(bx, args[0].immediate(), m_elem_bitwidth, mask_len);
 
-        // Alignment of T, must be a constant integer value:
-        let alignment = bx.const_i32(bx.align_of(values_elem).bytes() as i32);
+        let alignment = bx.const_i32(alignment as i32);
 
         let llvm_pointer = bx.type_ptr();
 

compiler/rustc_codegen_ssa/messages.ftl

@@ -169,6 +169,8 @@ codegen_ssa_invalid_monomorphization_unsupported_symbol = invalid monomorphizati
 
 codegen_ssa_invalid_monomorphization_unsupported_symbol_of_size = invalid monomorphization of `{$name}` intrinsic: unsupported {$symbol} from `{$in_ty}` with element `{$in_elem}` of size `{$size}` to `{$ret_ty}`
 
+codegen_ssa_invalid_monomorphization_non_power_of_two_alignment = invalid monomorphization of `{$name}` intrinsic: `alignment` is not a power of two
+
 codegen_ssa_invalid_windows_subsystem = invalid windows subsystem `{$subsystem}`, only `windows` and `console` are allowed
 
 codegen_ssa_ld64_unimplemented_modifier = `as-needed` modifier not implemented yet for ld64

compiler/rustc_codegen_ssa/src/errors.rs

@@ -1100,6 +1100,13 @@ pub enum InvalidMonomorphization<'tcx> {
         expected_element: Ty<'tcx>,
         vector_type: Ty<'tcx>,
     },
+
+    #[diag(codegen_ssa_invalid_monomorphization_non_power_of_two_alignment, code = E0511)]
+    AlignmentNotPowerOfTwo {
+        #[primary_span]
+        span: Span,
+        name: Symbol,
+    },
 }
 
 pub enum ExpectedPointerMutability {

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -695,8 +695,12 @@ pub(crate) fn check_intrinsic_type(
             (1, 0, vec![param(0), param(0), param(0)], param(0))
         }
         sym::simd_gather => (3, 0, vec![param(0), param(1), param(2)], param(0)),
-        sym::simd_masked_load => (3, 0, vec![param(0), param(1), param(2)], param(2)),
-        sym::simd_masked_store => (3, 0, vec![param(0), param(1), param(2)], tcx.types.unit),
+        sym::simd_masked_load => {
+            (3, 0, vec![param(0), param(1), param(2), tcx.types.u32], param(2))
+        }
+        sym::simd_masked_store => {
+            (3, 0, vec![param(0), param(1), param(2), tcx.types.u32], tcx.types.unit)
+        }
         sym::simd_scatter => (3, 0, vec![param(0), param(1), param(2)], tcx.types.unit),
         sym::simd_insert | sym::simd_insert_dyn => {
             (2, 0, vec![param(0), tcx.types.u32, param(1)], param(0))

library/core/src/intrinsics/simd.rs

@@ -392,13 +392,14 @@ 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).
+/// `alignment` must be a **const** power of two, and specifies the alignment of `ptr`.
+///
+/// Unmasked values in `T` must be readable as if by `<ptr>::read_unaligned` (aligned to `alignment`).
 ///
 /// `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>(mask: V, ptr: U, val: T, alignment: u32) -> T;
 
 /// Writes to a vector of pointers.
 ///
@@ -414,13 +415,14 @@ 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).
+/// `alignment` must be a **const** power of two, and specifies the alignment of `ptr`.
+///
+/// Unmasked values in `T` must be writeable as if by `<ptr>::write_unaligned` (aligned to `alignment`).
 ///
 /// `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>(mask: V, ptr: U, val: T, alignment: u32);
 
 /// Adds two simd vectors elementwise, with saturation.
 ///

library/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(
+                enable.to_int(),
+                ptr,
+                or,
+                const { core::mem::align_of::<T>() as u32 },
+            )
+        }
     }
 
     /// 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(
+                enable.to_int(),
+                ptr,
+                self,
+                const { core::mem::align_of::<T>() as u32 },
+            )
+        }
     }
 
     /// Writes the values in a SIMD vector to potentially discontiguous indices in `slice`.

src/tools/miri/src/intrinsics/simd.rs

@@ -733,7 +733,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 }
             }
             "masked_load" => {
-                let [mask, ptr, default] = check_intrinsic_arg_count(args)?;
+                let [mask, ptr, default, _alignment] = check_intrinsic_arg_count(args)?;
                 let (mask, mask_len) = this.project_to_simd(mask)?;
                 let ptr = this.read_pointer(ptr)?;
                 let (default, default_len) = this.project_to_simd(default)?;
@@ -759,7 +759,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 }
             }
             "masked_store" => {
-                let [mask, ptr, vals] = check_intrinsic_arg_count(args)?;
+                let [mask, ptr, vals, _alignment] = check_intrinsic_arg_count(args)?;
                 let (mask, mask_len) = this.project_to_simd(mask)?;
                 let ptr = this.read_pointer(ptr)?;
                 let (vals, vals_len) = this.project_to_simd(vals)?;

src/tools/miri/tests/pass/intrinsics/portable-simd.rs

@@ -684,21 +684,21 @@ fn simd_masked_loadstore() {
     let buf = [3i32; 3];
     let default = i32x4::splat(0);
     let mask = i32x4::from_array([!0, !0, !0, 0]);
-    let vals = unsafe { intrinsics::simd_masked_load(mask, buf.as_ptr(), default) };
+    let vals = unsafe { intrinsics::simd_masked_load(mask, buf.as_ptr(), default, 4) };
     assert_eq!(vals, i32x4::from_array([3, 3, 3, 0]));
     // Also read in a way that the *first* element is OOB.
     let mask2 = i32x4::from_array([0, !0, !0, !0]);
     let vals =
-        unsafe { intrinsics::simd_masked_load(mask2, buf.as_ptr().wrapping_sub(1), default) };
+        unsafe { intrinsics::simd_masked_load(mask2, buf.as_ptr().wrapping_sub(1), default, 4) };
     assert_eq!(vals, i32x4::from_array([0, 3, 3, 3]));
 
     // The buffer is deliberarely too short, so writing the last element would be UB.
     let mut buf = [42i32; 3];
     let vals = i32x4::from_array([1, 2, 3, 4]);
-    unsafe { intrinsics::simd_masked_store(mask, buf.as_mut_ptr(), vals) };
+    unsafe { intrinsics::simd_masked_store(mask, buf.as_mut_ptr(), vals, 4) };
     assert_eq!(buf, [1, 2, 3]);
     // Also write in a way that the *first* element is OOB.
-    unsafe { intrinsics::simd_masked_store(mask2, buf.as_mut_ptr().wrapping_sub(1), vals) };
+    unsafe { intrinsics::simd_masked_store(mask2, buf.as_mut_ptr().wrapping_sub(1), vals, 4) };
     assert_eq!(buf, [2, 3, 4]);
 }
 

tests/assembly-llvm/simd-intrinsic-mask-load.rs

@@ -35,7 +35,7 @@ pub struct f64x4([f64; 4]);
 pub struct m64x4([i64; 4]);
 
 #[rustc_intrinsic]
-unsafe fn simd_masked_load<M, P, T>(mask: M, pointer: P, values: T) -> T;
+unsafe fn simd_masked_load<M, P, T>(mask: M, pointer: P, values: T, alignment: u32) -> T;
 
 // CHECK-LABEL: load_i8x16
 #[no_mangle]
@@ -56,7 +56,7 @@ pub unsafe extern "C" fn load_i8x16(mask: m8x16, pointer: *const i8) -> i8x16 {
     // x86-avx512-NOT: vpsllw
     // x86-avx512: vpmovb2m k1, xmm0
     // x86-avx512-NEXT: vmovdqu8 xmm0 {k1} {z}, xmmword ptr [rdi]
-    simd_masked_load(mask, pointer, i8x16([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]))
+    simd_masked_load(mask, pointer, i8x16([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), 1)
 }
 
 // CHECK-LABEL: load_f32x8
@@ -68,7 +68,12 @@ pub unsafe extern "C" fn load_f32x8(mask: m32x8, pointer: *const f32) -> f32x8 {
     // x86-avx512-NOT: vpslld
     // x86-avx512: vpmovd2m k1, ymm0
     // x86-avx512-NEXT: vmovups ymm0 {k1} {z}, ymmword ptr [rdi]
-    simd_masked_load(mask, pointer, f32x8([0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32]))
+    simd_masked_load(
+        mask,
+        pointer,
+        f32x8([0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32, 0_f32]),
+        4,
+    )
 }
 
 // CHECK-LABEL: load_f64x4
@@ -79,5 +84,5 @@ pub unsafe extern "C" fn load_f64x4(mask: m64x4, pointer: *const f64) -> f64x4 {
     //
     // x86-avx512-NOT: vpsllq
     // x86-avx512: vpmovq2m k1, ymm0
-    simd_masked_load(mask, pointer, f64x4([0_f64, 0_f64, 0_f64, 0_f64]))
+    simd_masked_load(mask, pointer, f64x4([0_f64, 0_f64, 0_f64, 0_f64]), 8)
 }

tests/assembly-llvm/simd-intrinsic-mask-store.rs

@@ -35,7 +35,7 @@ pub struct f64x4([f64; 4]);
 pub struct m64x4([i64; 4]);
 
 #[rustc_intrinsic]
-unsafe fn simd_masked_store<M, P, T>(mask: M, pointer: P, values: T);
+unsafe fn simd_masked_store<M, P, T>(mask: M, pointer: P, values: T, alignment: u32);
 
 // CHECK-LABEL: store_i8x16
 #[no_mangle]
@@ -54,7 +54,7 @@ pub unsafe extern "C" fn store_i8x16(mask: m8x16, pointer: *mut i8, value: i8x16
     // x86-avx512-NOT: vpsllw
     // x86-avx512: vpmovb2m k1, xmm0
     // x86-avx512-NEXT: vmovdqu8 xmmword ptr [rdi] {k1}, xmm1
-    simd_masked_store(mask, pointer, value)
+    simd_masked_store(mask, pointer, value, 1)
 }
 
 // CHECK-LABEL: store_f32x8
@@ -66,7 +66,7 @@ pub unsafe extern "C" fn store_f32x8(mask: m32x8, pointer: *mut f32, value: f32x
     // x86-avx512-NOT: vpslld
     // x86-avx512: vpmovd2m k1, ymm0
     // x86-avx512-NEXT: vmovups ymmword ptr [rdi] {k1}, ymm1
-    simd_masked_store(mask, pointer, value)
+    simd_masked_store(mask, pointer, value, 4)
 }
 
 // CHECK-LABEL: store_f64x4
@@ -78,5 +78,5 @@ pub unsafe extern "C" fn store_f64x4(mask: m64x4, pointer: *mut f64, value: f64x
     // x86-avx512-NOT: vpsllq
     // x86-avx512: vpmovq2m k1, ymm0
     // x86-avx512-NEXT: vmovupd ymmword ptr [rdi] {k1}, ymm1
-    simd_masked_store(mask, pointer, value)
+    simd_masked_store(mask, pointer, value, 8)
 }