Auto merge of #115933 - oli-obk:simd_shuffle_const, r=workingjubilee

Prototype using const generic for simd_shuffle IDX array

cc rust-lang/rust#85229

r? `@workingjubilee` on the design

TLDR: there is now a `fn simd_shuffle_generic<T, U, const IDX: &'static [u32]>(x: T, y: T) -> U;` intrinsic that allows replacing

```rust
simd_shuffle(a, b, const { stuff })
```

with

```rust
simd_shuffle_generic::<_, _, {&stuff}>(a, b)
```

which makes the compiler implementations much simpler, if we manage to at some point eliminate `simd_shuffle`.

There are some issues with this today though (can't do math without bubbling it up in the generic arguments). With this change, we can start porting the simple cases and get better data on the others.

Author: bors

src/shims/intrinsics/mod.rs

@@ -60,7 +60,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
 
         // The rest jumps to `ret` immediately.
-        this.emulate_intrinsic_by_name(intrinsic_name, args, dest)?;
+        this.emulate_intrinsic_by_name(intrinsic_name, instance.args, args, dest)?;
 
         trace!("{:?}", this.dump_place(dest));
         this.go_to_block(ret);
@@ -71,6 +71,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     fn emulate_intrinsic_by_name(
         &mut self,
         intrinsic_name: &str,
+        generic_args: ty::GenericArgsRef<'tcx>,
         args: &[OpTy<'tcx, Provenance>],
         dest: &PlaceTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx> {
@@ -80,7 +81,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             return this.emulate_atomic_intrinsic(name, args, dest);
         }
         if let Some(name) = intrinsic_name.strip_prefix("simd_") {
-            return this.emulate_simd_intrinsic(name, args, dest);
+            return this.emulate_simd_intrinsic(name, generic_args, args, dest);
         }
 
         match intrinsic_name {

src/shims/intrinsics/simd.rs

@@ -12,6 +12,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     fn emulate_simd_intrinsic(
         &mut self,
         intrinsic_name: &str,
+        generic_args: ty::GenericArgsRef<'tcx>,
         args: &[OpTy<'tcx, Provenance>],
         dest: &PlaceTy<'tcx, Provenance>,
     ) -> InterpResult<'tcx> {
@@ -488,6 +489,38 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     this.write_immediate(*val, &dest)?;
                 }
             }
+            "shuffle_generic" => {
+                let [left, right] = check_arg_count(args)?;
+                let (left, left_len) = this.operand_to_simd(left)?;
+                let (right, right_len) = this.operand_to_simd(right)?;
+                let (dest, dest_len) = this.place_to_simd(dest)?;
+
+                let index = generic_args[2].expect_const().eval(*this.tcx, this.param_env(), Some(this.tcx.span)).unwrap().unwrap_branch();
+                let index_len = index.len();
+
+                assert_eq!(left_len, right_len);
+                assert_eq!(index_len as u64, dest_len);
+
+                for i in 0..dest_len {
+                    let src_index: u64 = index[i as usize].unwrap_leaf()
+                        .try_to_u32().unwrap()
+                        .into();
+                    let dest = this.project_index(&dest, i)?;
+
+                    let val = if src_index < left_len {
+                        this.read_immediate(&this.project_index(&left, src_index)?)?
+                    } else if src_index < left_len.checked_add(right_len).unwrap() {
+                        let right_idx = src_index.checked_sub(left_len).unwrap();
+                        this.read_immediate(&this.project_index(&right, right_idx)?)?
+                    } else {
+                        span_bug!(
+                            this.cur_span(),
+                            "simd_shuffle index {src_index} is out of bounds for 2 vectors of size {left_len}",
+                        );
+                    };
+                    this.write_immediate(*val, &dest)?;
+                }
+            }
             "shuffle" => {
                 let [left, right, index] = check_arg_count(args)?;
                 let (left, left_len) = this.operand_to_simd(left)?;

tests/pass/portable-simd.rs

@@ -1,5 +1,6 @@
 //@compile-flags: -Zmiri-strict-provenance
-#![feature(portable_simd, platform_intrinsics)]
+#![feature(portable_simd, platform_intrinsics, adt_const_params, inline_const)]
+#![allow(incomplete_features)]
 use std::simd::*;
 
 extern "platform-intrinsic" {
@@ -390,6 +391,8 @@ fn simd_intrinsics() {
         fn simd_reduce_any<T>(x: T) -> bool;
         fn simd_reduce_all<T>(x: T) -> bool;
         fn simd_select<M, T>(m: M, yes: T, no: T) -> T;
+        fn simd_shuffle_generic<T, U, const IDX: &'static [u32]>(x: T, y: T) -> U;
+        fn simd_shuffle<T, IDX, U>(x: T, y: T, idx: IDX) -> U;
     }
     unsafe {
         // Make sure simd_eq returns all-1 for `true`
@@ -413,6 +416,22 @@ fn simd_intrinsics() {
             simd_select(i8x4::from_array([0, -1, -1, 0]), b, a),
             i32x4::from_array([10, 2, 10, 10])
         );
+        assert_eq!(
+            simd_shuffle_generic::<_, i32x4, {&[3, 1, 0, 2]}>(a, b),
+            a,
+        );
+        assert_eq!(
+            simd_shuffle::<_, _, i32x4>(a, b, const {[3, 1, 0, 2]}),
+            a,
+        );
+        assert_eq!(
+            simd_shuffle_generic::<_, i32x4, {&[7, 5, 4, 6]}>(a, b),
+            i32x4::from_array([4, 2, 1, 10]),
+        );
+        assert_eq!(
+            simd_shuffle::<_, _, i32x4>(a, b, const {[7, 5, 4, 6]}),
+            i32x4::from_array([4, 2, 1, 10]),
+        );
     }
 }