Implement simd_shuffle in const_eval

Author: sayantn

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -485,6 +485,39 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 }
                 self.copy_op(&self.project_index(&input, index)?, dest)?;
             }
+            sym::simd_shuffle => {
+                let (x, input_len) = self.project_to_simd(&args[0])?;
+                let (y, _) = self.project_to_simd(&args[1])?;
+                let (shuffle, shuffle_len) = self.project_to_simd(&args[2])?;
+                let (dest, dest_len) = self.project_to_simd(dest)?;
+                assert_eq!(
+                    shuffle_len, dest_len,
+                    "Return vector length must match shuffle vector length"
+                );
+                assert_eq!(
+                    shuffle.layout.ty.sequence_element_type(*self.tcx).kind(),
+                    &ty::Uint(ty::UintTy::U32),
+                    "Shuffle vector must be a vector of `u32`s"
+                );
+
+                for i in 0..dest_len {
+                    let idx = self.project_index(&shuffle, i)?;
+                    let idx = u64::from(self.read_scalar(&idx)?.to_u32()?);
+                    if idx >= 2 * input_len {
+                        throw_ub_format!(
+                            "`simd_shuffle` index {idx} is out-of-bounds of vector with length {input_len}"
+                        );
+                    }
+
+                    let value = if idx < input_len {
+                        self.project_index(&x, idx)?
+                    } else {
+                        self.project_index(&y, idx - input_len)?
+                    };
+                    let place = self.project_index(&dest, i)?;
+                    self.copy_op(&value, &place)?;
+                }
+            }
             sym::black_box => {
                 // These just return their argument
                 self.copy_op(&args[0], dest)?;

library/core/src/intrinsics/simd.rs

@@ -330,7 +330,7 @@ pub unsafe fn simd_ge<T, U>(x: T, y: T) -> U;
 /// of `xy`.
 #[rustc_intrinsic]
 #[rustc_nounwind]
-pub unsafe fn simd_shuffle<T, U, V>(x: T, y: T, idx: U) -> V;
+pub const unsafe fn simd_shuffle<T, U, V>(x: T, y: T, idx: U) -> V;
 
 /// Reads a vector of pointers.
 ///

tests/ui/consts/const-eval/simd/shuffle.rs

@@ -0,0 +1,62 @@
+//@ run-pass
+
+// these tests are copied from `tests/ui/simd/intrinsic/generic-elements-pass.rs`
+#![feature(repr_simd, intrinsics, core_intrinsics)]
+
+#[path = "../../../../auxiliary/minisimd.rs"]
+mod minisimd;
+use std::intrinsics::simd::simd_shuffle;
+
+use minisimd::*;
+
+#[repr(simd)]
+struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
+
+macro_rules! all_eq {
+    ($a: expr, $b: expr) => {{
+        let a = $a;
+        let b = $b;
+        // type inference works better with the concrete type on the
+        // left, but humans work better with the expected on the
+        // right.
+        assert!(b == a, "{:?} != {:?}", a, b);
+    }};
+}
+
+macro_rules! simd_shuffle {
+    ($a: expr, $b: expr, $shuffle: expr) => {
+        const { simd_shuffle($a, $b, const { SimdShuffleIdx($shuffle) } ) }
+    };
+}
+
+fn main() {
+    const X2: i32x2 = i32x2::from_array([20, 21]);
+    const X4: i32x4 = i32x4::from_array([40, 41, 42, 43]);
+    const X8: i32x8 = i32x8::from_array([80, 81, 82, 83, 84, 85, 86, 87]);
+
+    const Y2: i32x2 = i32x2::from_array([120, 121]);
+    const Y4: i32x4 = i32x4::from_array([140, 141, 142, 143]);
+    const Y8: i32x8 = i32x8::from_array([180, 181, 182, 183, 184, 185, 186, 187]);
+    unsafe {
+        all_eq!(simd_shuffle!(X2, Y2, [3, 0]), i32x2::from_array([121, 20]));
+        all_eq!(simd_shuffle!(X2, Y2, [3, 0, 1, 2]), i32x4::from_array([121, 20, 21, 120]));
+        all_eq!(
+            simd_shuffle!(X2, Y2, [3, 0, 1, 2, 1, 2, 3, 0]),
+            i32x8::from_array([121, 20, 21, 120, 21, 120, 121, 20])
+        );
+
+        all_eq!(simd_shuffle!(X4, Y4, [7, 2]), i32x2::from_array([143, 42]));
+        all_eq!(simd_shuffle!(X4, Y4, [7, 2, 5, 0]), i32x4::from_array([143, 42, 141, 40]));
+        all_eq!(
+            simd_shuffle!(X4, Y4, [7, 2, 5, 0, 3, 6, 4, 1]),
+            i32x8::from_array([143, 42, 141, 40, 43, 142, 140, 41])
+        );
+
+        all_eq!(simd_shuffle!(X8, Y8, [11, 5]), i32x2::from_array([183, 85]));
+        all_eq!(simd_shuffle!(X8, Y8, [11, 5, 15, 0]), i32x4::from_array([183, 85, 187, 80]));
+        all_eq!(
+            simd_shuffle!(X8, Y8, [11, 5, 15, 0, 3, 8, 12, 1]),
+            i32x8::from_array([183, 85, 187, 80, 83, 180, 184, 81])
+        );
+    }
+}