simd_shuffle: require index argument to be a vector

qinheping · Sep 14, 2024 · 79d937d · 79d937d
1 parent 34e22e8
commit 79d937d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/core/src/intrinsics/simd.rs b/core/src/intrinsics/simd.rs
@@ -232,7 +232,7 @@ extern "rust-intrinsic" {
     ///
     /// `T` must be a vector.
     ///
-    /// `U` must be a **const** array or vector of `u32`s. This means it must either refer to a named
+    /// `U` must be a **const** vector of `u32`s. This means it must either refer to a named
     /// const or be given as an inline const expression (`const { ... }`).
     ///
     /// `V` must be a vector with the same element type as `T` and the same length as `U`.

diff --git a/portable-simd/crates/core_simd/src/swizzle.rs b/portable-simd/crates/core_simd/src/swizzle.rs
@@ -85,7 +85,7 @@ pub trait Swizzle<const N: usize> {
         LaneCount<N>: SupportedLaneCount,
         LaneCount<M>: SupportedLaneCount,
     {
-        // Safety: `vector` is a vector, and the index is a const array of u32.
+        // Safety: `vector` is a vector, and the index is a const vector of u32.
         unsafe {
             core::intrinsics::simd::simd_shuffle(
                 vector,
@@ -103,7 +103,11 @@ pub trait Swizzle<const N: usize> {
                         output[i] = index as u32;
                         i += 1;
                     }
-                    output
+
+                    // The index list needs to be returned as a vector.
+                    #[repr(simd)]
+                    struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
+                    SimdShuffleIdx(output)
                 },
             )
         }
@@ -121,7 +125,7 @@ pub trait Swizzle<const N: usize> {
         LaneCount<N>: SupportedLaneCount,
         LaneCount<M>: SupportedLaneCount,
     {
-        // Safety: `first` and `second` are vectors, and the index is a const array of u32.
+        // Safety: `first` and `second` are vectors, and the index is a const vector of u32.
         unsafe {
             core::intrinsics::simd::simd_shuffle(
                 first,
@@ -139,7 +143,11 @@ pub trait Swizzle<const N: usize> {
                         output[i] = index as u32;
                         i += 1;
                     }
-                    output
+
+                    // The index list needs to be returned as a vector.
+                    #[repr(simd)]
+                    struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
+                    SimdShuffleIdx(output)
                 },
             )
         }