Rollup merge of #147355 - sayantn:masked-loads, r=RalfJung,bjorn3

Add alignment parameter to `simd_masked_{load,store}`

This PR adds an alignment parameter in `simd_masked_load` and `simd_masked_store`, in the form of a const-generic enum `core::intrinsics::simd::SimdAlign`. This represents the alignment of the `ptr` argument in these intrinsics as follows

 - `SimdAlign::Unaligned` - `ptr` is unaligned/1-byte aligned
 - `SimdAlign::Element` - `ptr` is aligned to the element type of the SIMD vector (default behavior in the old signature)
 - `SimdAlign::Vector` - `ptr` is aligned to the SIMD vector type

The main motive for this is stdarch - most vector loads are either fully aligned (to the vector size) or unaligned (byte-aligned), so the previous signature doesn't cut it.

Now, stdarch will mostly use `SimdAlign::Unaligned` and `SimdAlign::Vector`, whereas portable-simd will use `SimdAlign::Element`.

 - [x] `cg_llvm`
 - [x] `cg_clif`
 - [x] `miri`/`const_eval`

## Alternatives

Using a const-generic/"const" `u32` parameter as alignment (and we error during codegen if this argument is not a power of two). This, although more flexible than this, has a few drawbacks

 - If we use an const-generic argument, then portable-simd somehow needs to pass `align_of::<T>()` as the alignment, which isn't possible without GCE
 - "const" function parameters are just an ugly hack, and a pain to deal with in non-LLVM backends

We can remedy the problem with the const-generic `u32` parameter by adding a special rule for the element alignment case (e.g. `0` can mean "use the alignment of the element type), but I feel like this is not as expressive as the enum approach, although I am open to suggestions

cc `@workingjubilee` `@RalfJung` `@BoxyUwU`

Author: Zalathar

tests/fail/intrinsics/simd_masked_load_element_misaligned.rs

@@ -0,0 +1,17 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::*;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let buf = [0u32; 5];
+        //~v ERROR: accessing memory with alignment 
+        simd_masked_load::<_, _, _, { SimdAlign::Element }>(
+            i32x4::splat(-1),
+            // This is not i32-aligned
+            buf.as_ptr().byte_offset(1),
+            i32x4::splat(0),
+        );
+    }
+}

tests/fail/intrinsics/simd_masked_load_element_misaligned.stderr

@@ -0,0 +1,18 @@
+error: Undefined Behavior: accessing memory with alignment ALIGN, but alignment ALIGN is required
+  --> tests/fail/intrinsics/simd_masked_load_element_misaligned.rs:LL:CC
+   |
+LL | /         simd_masked_load::<_, _, _, { SimdAlign::Element }>(
+LL | |             i32x4::splat(-1),
+LL | |             // This is not i32-aligned
+LL | |             buf.as_ptr().byte_offset(1),
+LL | |             i32x4::splat(0),
+LL | |         );
+   | |_________^ Undefined Behavior occurred here
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

tests/fail/intrinsics/simd_masked_load_vector_misaligned.rs

@@ -0,0 +1,17 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::*;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let buf = Simd::<i32, 8>::splat(0);
+        //~v ERROR: accessing memory with alignment
+        simd_masked_load::<_, _, _, { SimdAlign::Vector }>(
+            i32x4::splat(-1),
+            // This is i32-aligned but not i32x4-aligned.
+            buf.as_array()[1..].as_ptr(),
+            i32x4::splat(0),
+        );
+    }
+}

tests/fail/intrinsics/simd_masked_load_vector_misaligned.stderr

@@ -0,0 +1,18 @@
+error: Undefined Behavior: accessing memory with alignment ALIGN, but alignment ALIGN is required
+  --> tests/fail/intrinsics/simd_masked_load_vector_misaligned.rs:LL:CC
+   |
+LL | /         simd_masked_load::<_, _, _, { SimdAlign::Vector }>(
+LL | |             i32x4::splat(-1),
+LL | |             // This is i32-aligned but not i32x4-aligned.
+LL | |             buf.as_array()[1..].as_ptr(),
+LL | |             i32x4::splat(0),
+LL | |         );
+   | |_________^ Undefined Behavior occurred here
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

tests/fail/intrinsics/simd_masked_store_element_misaligned.rs

@@ -0,0 +1,17 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::*;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let mut buf = [0u32; 5];
+        //~v ERROR: accessing memory with alignment
+        simd_masked_store::<_, _, _, { SimdAlign::Element }>(
+            i32x4::splat(-1),
+            // This is not i32-aligned
+            buf.as_mut_ptr().byte_offset(1),
+            i32x4::splat(0),
+        );
+    }
+}

tests/fail/intrinsics/simd_masked_store_element_misaligned.stderr

@@ -0,0 +1,18 @@
+error: Undefined Behavior: accessing memory with alignment ALIGN, but alignment ALIGN is required
+  --> tests/fail/intrinsics/simd_masked_store_element_misaligned.rs:LL:CC
+   |
+LL | /         simd_masked_store::<_, _, _, { SimdAlign::Element }>(
+LL | |             i32x4::splat(-1),
+LL | |             // This is not i32-aligned
+LL | |             buf.as_mut_ptr().byte_offset(1),
+LL | |             i32x4::splat(0),
+LL | |         );
+   | |_________^ Undefined Behavior occurred here
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

tests/fail/intrinsics/simd_masked_store_vector_misaligned.rs

@@ -0,0 +1,17 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::*;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let mut buf = Simd::<i32, 8>::splat(0);
+        //~v ERROR: accessing memory with alignment
+        simd_masked_store::<_, _, _, { SimdAlign::Vector }>(
+            i32x4::splat(-1),
+            // This is i32-aligned but not i32x4-aligned.
+            buf.as_mut_array()[1..].as_mut_ptr(),
+            i32x4::splat(0),
+        );
+    }
+}

tests/fail/intrinsics/simd_masked_store_vector_misaligned.stderr

@@ -0,0 +1,18 @@
+error: Undefined Behavior: accessing memory with alignment ALIGN, but alignment ALIGN is required
+  --> tests/fail/intrinsics/simd_masked_store_vector_misaligned.rs:LL:CC
+   |
+LL | /         simd_masked_store::<_, _, _, { SimdAlign::Vector }>(
+LL | |             i32x4::splat(-1),
+LL | |             // This is i32-aligned but not i32x4-aligned.
+LL | |             buf.as_mut_array()[1..].as_mut_ptr(),
+LL | |             i32x4::splat(0),
+LL | |         );
+   | |_________^ Undefined Behavior occurred here
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

tests/pass/intrinsics/portable-simd.rs

@@ -936,26 +936,93 @@ fn simd_float_intrinsics() {
 }
 
 fn simd_masked_loadstore() {
+    use intrinsics::*;
+
     // The buffer is deliberarely too short, so reading the last element would be UB.
     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 { simd_masked_load::<_, _, _, { SimdAlign::Element }>(mask, buf.as_ptr(), default) };
     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) };
+    let vals = unsafe {
+        simd_masked_load::<_, _, _, { SimdAlign::Element }>(
+            mask2,
+            buf.as_ptr().wrapping_sub(1),
+            default,
+        )
+    };
     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 { simd_masked_store::<_, _, _, { SimdAlign::Element }>(mask, buf.as_mut_ptr(), vals) };
     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 {
+        simd_masked_store::<_, _, _, { SimdAlign::Element }>(
+            mask2,
+            buf.as_mut_ptr().wrapping_sub(1),
+            vals,
+        )
+    };
     assert_eq!(buf, [2, 3, 4]);
+
+    // we use a purposely misaliged buffer to make sure Miri doesn't error in this case
+    let buf = [0x03030303_i32; 5];
+    let default = i32x4::splat(0);
+    let mask = i32x4::splat(!0);
+    let vals = unsafe {
+        simd_masked_load::<_, _, _, { SimdAlign::Unaligned }>(
+            mask,
+            buf.as_ptr().byte_offset(1), // this is guaranteed to be unaligned
+            default,
+        )
+    };
+    assert_eq!(vals, i32x4::splat(0x03030303));
+
+    let mut buf = [0i32; 5];
+    let mask = i32x4::splat(!0);
+    unsafe {
+        simd_masked_store::<_, _, _, { SimdAlign::Unaligned }>(
+            mask,
+            buf.as_mut_ptr().byte_offset(1), // this is guaranteed to be unaligned
+            vals,
+        )
+    };
+    assert_eq!(
+        buf,
+        [
+            i32::from_ne_bytes([0, 3, 3, 3]),
+            0x03030303,
+            0x03030303,
+            0x03030303,
+            i32::from_ne_bytes([3, 0, 0, 0])
+        ]
+    );
+
+    // `repr(simd)` types like `Simd<T, N>` have the correct alignment for vectors
+    let buf = i32x4::splat(3);
+    let default = i32x4::splat(0);
+    let mask = i32x4::splat(!0);
+    let vals = unsafe {
+        simd_masked_load::<_, _, _, { SimdAlign::Vector }>(
+            mask,
+            &raw const buf as *const i32,
+            default,
+        )
+    };
+    assert_eq!(vals, buf);
+
+    let mut buf = i32x4::splat(0);
+    let mask = i32x4::splat(!0);
+    unsafe {
+        simd_masked_store::<_, _, _, { SimdAlign::Vector }>(mask, &raw mut buf as *mut i32, vals)
+    };
+    assert_eq!(buf, vals);
 }
 
 fn simd_ops_non_pow2() {