Rollup merge of #147534 - sayantn:simd-funnel-shifts, r=RalfJung

Implement SIMD funnel shifts in const-eval/Miri

Split off from rust-lang/rust#147520 with just this change for easier review

r? ```@RalfJung```

compiler/rustc_const_eval/src/interpret/intrinsics/simd.rs

@@ -3,7 +3,7 @@ use rustc_abi::{BackendRepr, Endian};
 use rustc_apfloat::ieee::{Double, Half, Quad, Single};
 use rustc_apfloat::{Float, Round};
 use rustc_middle::mir::interpret::{InterpErrorKind, Pointer, UndefinedBehaviorInfo};
-use rustc_middle::ty::{FloatTy, SimdAlign};
+use rustc_middle::ty::{FloatTy, ScalarInt, SimdAlign};
 use rustc_middle::{bug, err_ub_format, mir, span_bug, throw_unsup_format, ty};
 use rustc_span::{Symbol, sym};
 use tracing::trace;
@@ -744,6 +744,58 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                     self.write_scalar(val, &dest)?;
                 }
             }
+            sym::simd_funnel_shl | sym::simd_funnel_shr => {
+                let (left, _) = self.project_to_simd(&args[0])?;
+                let (right, _) = self.project_to_simd(&args[1])?;
+                let (shift, _) = self.project_to_simd(&args[2])?;
+                let (dest, _) = self.project_to_simd(&dest)?;
+
+                let (len, elem_ty) = args[0].layout.ty.simd_size_and_type(*self.tcx);
+                let (elem_size, _signed) = elem_ty.int_size_and_signed(*self.tcx);
+                let elem_size_bits = u128::from(elem_size.bits());
+
+                let is_left = intrinsic_name == sym::simd_funnel_shl;
+
+                for i in 0..len {
+                    let left =
+                        self.read_scalar(&self.project_index(&left, i)?)?.to_bits(elem_size)?;
+                    let right =
+                        self.read_scalar(&self.project_index(&right, i)?)?.to_bits(elem_size)?;
+                    let shift_bits =
+                        self.read_scalar(&self.project_index(&shift, i)?)?.to_bits(elem_size)?;
+
+                    if shift_bits >= elem_size_bits {
+                        throw_ub_format!(
+                            "overflowing shift by {shift_bits} in `{intrinsic_name}` in lane {i}"
+                        );
+                    }
+                    let inv_shift_bits = u32::try_from(elem_size_bits - shift_bits).unwrap();
+
+                    // A funnel shift left by S can be implemented as `(x << S) | y.unbounded_shr(SIZE - S)`.
+                    // The `unbounded_shr` is needed because otherwise if `S = 0`, it would be `x | y`
+                    // when it should be `x`.
+                    //
+                    // This selects the least-significant `SIZE - S` bits of `x`, followed by the `S` most
+                    // significant bits of `y`. As `left` and `right` both occupy the lower `SIZE` bits,
+                    // we can treat the lower `SIZE` bits as an integer of the right width and use
+                    // the same implementation, but on a zero-extended `x` and `y`. This works because
+                    // `x << S` just pushes the `SIZE-S` MSBs out, and `y >> (SIZE - S)` shifts in
+                    // zeros, as it is zero-extended. To the lower `SIZE` bits, this looks just like a
+                    // funnel shift left.
+                    //
+                    // Note that the `unbounded_sh{l,r}`s are needed only in case we are using this on
+                    // `u128xN` and `inv_shift_bits == 128`.
+                    let result_bits = if is_left {
+                        (left << shift_bits) | right.unbounded_shr(inv_shift_bits)
+                    } else {
+                        left.unbounded_shl(inv_shift_bits) | (right >> shift_bits)
+                    };
+                    let (result, _overflow) = ScalarInt::truncate_from_uint(result_bits, elem_size);
+
+                    let dest = self.project_index(&dest, i)?;
+                    self.write_scalar(result, &dest)?;
+                }
+            }
 
             // Unsupported intrinsic: skip the return_to_block below.
             _ => return interp_ok(false),

src/tools/miri/tests/fail/intrinsics/simd-funnel_shl-too-far.rs

@@ -0,0 +1,12 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::simd_funnel_shl;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let x = i32x2::from_array([1, 1]);
+        let y = i32x2::from_array([100, 0]);
+        simd_funnel_shl(x, x, y); //~ERROR: overflowing shift by 100 in `simd_funnel_shl` in lane 0
+    }
+}

src/tools/miri/tests/fail/intrinsics/simd-funnel_shl-too-far.stderr

@@ -0,0 +1,13 @@
+error: Undefined Behavior: overflowing shift by 100 in `simd_funnel_shl` in lane 0
+  --> tests/fail/intrinsics/simd-funnel_shl-too-far.rs:LL:CC
+   |
+LL |         simd_funnel_shl(x, x, y);
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^ 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
+

src/tools/miri/tests/fail/intrinsics/simd-funnel_shr-too-far.rs

@@ -0,0 +1,12 @@
+#![feature(core_intrinsics, portable_simd)]
+
+use std::intrinsics::simd::simd_funnel_shr;
+use std::simd::*;
+
+fn main() {
+    unsafe {
+        let x = i32x2::from_array([1, 1]);
+        let y = i32x2::from_array([20, 40]);
+        simd_funnel_shr(x, x, y); //~ERROR: overflowing shift by 40 in `simd_funnel_shr` in lane 1
+    }
+}

src/tools/miri/tests/fail/intrinsics/simd-funnel_shr-too-far.stderr

@@ -0,0 +1,13 @@
+error: Undefined Behavior: overflowing shift by 40 in `simd_funnel_shr` in lane 1
+  --> tests/fail/intrinsics/simd-funnel_shr-too-far.rs:LL:CC
+   |
+LL |         simd_funnel_shr(x, x, y);
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^ 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
+

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

@@ -62,7 +62,7 @@ impl<T: Copy, const N: usize> PackedSimd<T, N> {
 #[rustc_nounwind]
 pub unsafe fn simd_shuffle_const_generic<T, U, const IDX: &'static [u32]>(x: T, y: T) -> U;
 
-pub fn simd_ops_f16() {
+fn simd_ops_f16() {
     use intrinsics::*;
 
     // small hack to make type inference better
@@ -273,7 +273,7 @@ fn simd_ops_f64() {
     assert_eq!(f64x2::from_array([f64::NAN, 0.0]).reduce_min(), 0.0);
 }
 
-pub fn simd_ops_f128() {
+fn simd_ops_f128() {
     use intrinsics::*;
 
     // small hack to make type inference better
@@ -454,6 +454,18 @@ fn simd_ops_i32() {
             0x3fffffffu32 as i32
         ])
     );
+
+    // these values are taken from the doctests of `u32::funnel_shl` and `u32::funnel_shr`
+    let c = u32x4::splat(0x010000b3);
+    let d = u32x4::splat(0x2fe78e45);
+
+    unsafe {
+        assert_eq!(intrinsics::simd_funnel_shl(c, d, u32x4::splat(0)), c);
+        assert_eq!(intrinsics::simd_funnel_shl(c, d, u32x4::splat(8)), u32x4::splat(0x0000b32f));
+
+        assert_eq!(intrinsics::simd_funnel_shr(c, d, u32x4::splat(0)), d);
+        assert_eq!(intrinsics::simd_funnel_shr(c, d, u32x4::splat(8)), u32x4::splat(0xb32fe78e));
+    }
 }
 
 fn simd_mask() {