Accept [u32; N] arrays as simd_shuffle indices

.github/workflows/main.yml

@@ -151,6 +151,7 @@ jobs:
       run: |
         sed -i 's/jit./#jit./' config.txt
         sed -i 's/aot.powi_libcall_signature/#aot.powi_libcall_signature/' config.txt
+        sed -i 's/aot.simd-shuffle-array-index/#aot.simd-shuffle-array-index/' config.txt
 
     - name: Test
       env:

build_system/tests.rs

@@ -87,6 +87,11 @@ const BASE_SYSROOT_SUITE: &[TestCase] = &[
         &[],
     ),
     TestCase::build_bin_and_run("aot.float-minmax-pass", "example/float-minmax-pass.rs", &[]),
+    TestCase::build_bin_and_run(
+        "aot.simd-shuffle-array-index",
+        "example/simd-shuffle-array-index.rs",
+        &[],
+    ),
     TestCase::custom("aot.powi_libcall_signature", &|runner| {
         let mut cmd = runner.rustc_command(["example/powi-libcall-signature.rs"]);
         let output = cmd.output().unwrap();

config.txt

@@ -19,6 +19,7 @@ aot.dst_field_align
 aot.subslice-patterns-const-eval
 aot.track-caller-attribute
 aot.float-minmax-pass
+aot.simd-shuffle-array-index
 aot.powi_libcall_signature
 aot.issue-72793
 aot.issue-59326

example/simd-shuffle-array-index.rs

@@ -0,0 +1,24 @@
+// Copied from pathfinder_simd-style usage.
+// run-pass
+
+#![feature(repr_simd, core_intrinsics)]
+#![allow(internal_features, non_camel_case_types)]
+
+use std::intrinsics::simd::simd_shuffle;
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct f32x4([f32; 4]);
+
+impl f32x4 {
+    fn into_array(self) -> [f32; 4] {
+        unsafe { std::mem::transmute(self) }
+    }
+}
+
+fn main() {
+    let x = f32x4([1.0, 2.0, 3.0, 4.0]);
+    const IDX: [u32; 4] = [3, 3, 3, 3];
+    let r: f32x4 = unsafe { simd_shuffle(x, x, IDX) };
+    assert_eq!(r.into_array(), [4.0, 4.0, 4.0, 4.0]);
+}

src/intrinsics/simd.rs

@@ -18,6 +18,32 @@ fn report_simd_type_validation_error(
     fx.bcx.ins().trap(TrapCode::user(1 /* unreachable */).unwrap());
 }
 
+fn simd_shuffle_index_len<'tcx>(
+    fx: &mut FunctionCx<'_, '_, 'tcx>,
+    span: Span,
+    idx_ty: Ty<'tcx>,
+) -> Option<u64> {
+    if idx_ty.is_simd() {
+        let (count, elem_ty) = idx_ty.simd_size_and_type(fx.tcx);
+        if matches!(elem_ty.kind(), ty::Uint(ty::UintTy::U32)) {
+            return Some(count);
+        }
+    } else if let ty::Array(elem, len) = idx_ty.kind() {
+        if matches!(elem.kind(), ty::Uint(ty::UintTy::U32)) {
+            return Some(
+                len.try_to_target_usize(fx.tcx).expect("expected monomorphic const in codegen"),
+            );
+        }
+    }
+
+    fx.tcx.dcx().span_err(
+        span,
+        format!("simd_shuffle index must be a SIMD vector of `u32` or `[u32; N]`, got `{idx_ty}`",),
+    );
+    fx.bcx.ins().trap(TrapCode::user(1 /* unreachable */).unwrap());
+    None
+}
+
 pub(super) fn codegen_simd_intrinsic_call<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
     intrinsic: Symbol,
@@ -179,20 +205,12 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
                 return;
             }
 
-            // Make sure this is actually a SIMD vector.
+            // The index must be a const SIMD vector of `u32` or a const `[u32; N]` array.
+            // pathfinder_simd passes the latter.
             let idx_ty = fx.monomorphize(idx.node.ty(fx.mir, fx.tcx));
-            if !idx_ty.is_simd()
-                || !matches!(idx_ty.simd_size_and_type(fx.tcx).1.kind(), ty::Uint(ty::UintTy::U32))
-            {
-                fx.tcx.dcx().span_err(
-                    span,
-                    format!("simd_shuffle index must be a SIMD vector of `u32`, got `{}`", idx_ty),
-                );
-                // Prevent verifier error
-                fx.bcx.ins().trap(TrapCode::user(1 /* unreachable */).unwrap());
+            let Some(index_len) = simd_shuffle_index_len(fx, span, idx_ty) else {
                 return;
             };
-            let n: u16 = idx_ty.simd_size_and_type(fx.tcx).0.try_into().unwrap();
 
             assert_eq!(x.layout(), y.layout());
             let layout = x.layout();
@@ -201,7 +219,7 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
             let (ret_lane_count, ret_lane_ty) = ret.layout().ty.simd_size_and_type(fx.tcx);
 
             assert_eq!(lane_ty, ret_lane_ty);
-            assert_eq!(u64::from(n), ret_lane_count);
+            assert_eq!(index_len, ret_lane_count);
 
             let total_len = lane_count * 2;