Implement common swizzle operations.

crates/core_simd/src/swizzle.rs

@@ -364,4 +364,153 @@ where
 
         (Even::swizzle2(self, other), Odd::swizzle2(self, other))
     }
+
+    /// Splits a vector into its two halves.
+    ///
+    /// Due to limitations in const generics, the length of the resulting vector cannot be inferred
+    /// from the input vectors. You must specify it explicitly. A compile-time error will be raised
+    /// if `HALF_LANES * 2 != LANES`.
+    ///
+    /// ```
+    /// # #![feature(portable_simd)]
+    /// # #[cfg(feature = "as_crate")] use core_simd::simd::Simd;
+    /// # #[cfg(not(feature = "as_crate"))] use core::simd::Simd;
+    /// let x = Simd::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
+    /// let [y, z] = x.split_to::<4>();
+    /// assert_eq!(y.to_array(), [0, 1, 2, 3]);
+    /// assert_eq!(z.to_array(), [4, 5, 6, 7]);
+    /// ```
+    #[inline]
+    #[must_use = "method returns a new vector and does not mutate the original inputs"]
+    // TODO: when `generic_const_exprs` supports it, provide
+    //   `pub fn split(self) -> [Simd<T, {LANES / 2}>; 2]`
+    // and deprecate `split_to`.
+    pub fn split_to<const HALF_LANES: usize>(self) -> [Simd<T, HALF_LANES>; 2]
+    where
+        LaneCount<HALF_LANES>: SupportedLaneCount,
+    {
+        const fn slice_index<const LEN: usize>(hi_half: bool, lanes: usize) -> [usize; LEN] {
+            assert!(
+                LEN * 2 == lanes,
+                "x.split_to::<N>() must provide N=x.lanes()/2"
+            );
+            let offset = if hi_half { LEN } else { 0 };
+            let mut index = [0; LEN];
+            let mut i = 0;
+            while i < LEN {
+                index[i] = i + offset;
+                i += 1;
+            }
+            index
+        }
+        struct Split<const HI_HALF: bool>;
+        impl<const HI_HALF: bool, const LEN: usize, const LANES: usize> Swizzle<LANES, LEN>
+            for Split<HI_HALF>
+        {
+            const INDEX: [usize; LEN] = slice_index::<LEN>(HI_HALF, LANES);
+        }
+        [Split::<false>::swizzle(self), Split::<true>::swizzle(self)]
+    }
+
+    /// Concatenates two vectors of equal length.
+    ///
+    /// Due to limitations in const generics, the length of the resulting vector cannot be inferred
+    /// from the input vectors. You must specify it explicitly. A compile time error will be raised
+    /// if `LANES * 2 != DOUBLE_LANES`
+    ///
+    /// ```
+    /// # #![feature(portable_simd)]
+    /// # #[cfg(feature = "as_crate")] use core_simd::simd::Simd;
+    /// # #[cfg(not(feature = "as_crate"))] use core::simd::Simd;
+    /// let x = Simd::from_array([0, 1, 2, 3]);
+    /// let y = Simd::from_array([4, 5, 6, 7]);
+    /// let z = x.concat_to::<8>(y);
+    /// assert_eq!(z.to_array(), [0, 1, 2, 3, 4, 5, 6, 7]);
+    /// ```
+    ///
+    /// Will be rejected at compile time if `LANES * 2 != DOUBLE_LANES`.
+    #[inline]
+    #[must_use = "method returns a new vector and does not mutate the original inputs"]
+    // TODO: when `generic_const_exprs` supports it, provide
+    //   `pub fn concat(self, other: Self) -> Simd<T, {LANES * 2}>`
+    // and deprecate `concat_to`.
+    pub fn concat_to<const DOUBLE_LANES: usize>(self, other: Self) -> Simd<T, DOUBLE_LANES>
+    where
+        LaneCount<DOUBLE_LANES>: SupportedLaneCount,
+    {
+        const fn concat_index<const DOUBLE_LANES: usize>(lanes: usize) -> [Which; DOUBLE_LANES] {
+            assert!(lanes * 2 == DOUBLE_LANES);
+            let mut index = [Which::First(0); DOUBLE_LANES];
+            let mut i = 0;
+            while i < lanes {
+                index[i] = Which::First(i);
+                index[i + lanes] = Which::Second(i);
+                i += 1;
+            }
+            index
+        }
+        struct Concat;
+        impl<const LANES: usize, const DOUBLE_LANES: usize> Swizzle2<LANES, DOUBLE_LANES> for Concat {
+            const INDEX: [Which; DOUBLE_LANES] = concat_index::<DOUBLE_LANES>(LANES);
+        }
+        Concat::swizzle2(self, other)
+    }
+
+    /// For each lane `i`, swaps it with lane `i ^ SWAP_MASK`.
+    ///
+    /// This is a powerful swizzle operation that can implement many common patterns as special cases.
+    /// For power-of-2 swap masks, this produces the [butterfly shuffles](https://en.wikipedia.org/wiki/Butterfly_network)
+    /// that are often useful for horizontal reductions.
+    /// 
+    /// A similar operation (operating on bits instead of lanes) is known as `grev` in the RISC-V
+    /// Bitmanip specification.
+    ///
+    /// ```
+    /// # #![feature(portable_simd)]
+    /// # #[cfg(feature = "as_crate")] use core_simd::simd::Simd;
+    /// # #[cfg(not(feature = "as_crate"))] use core::simd::Simd;
+    /// let x = Simd::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
+    /// // Swap adjacent lanes:
+    /// assert_eq!(x.general_reverse::<1>().to_array(), [1, 0, 3, 2, 5, 4, 7, 6]);
+    /// // Swap lanes separated by distance 2:
+    /// assert_eq!(x.general_reverse::<2>().to_array(), [2, 3, 0, 1, 6, 7, 4, 5]);
+    /// // Swap lanes separated by distance 4:
+    /// assert_eq!(x.general_reverse::<4>().to_array(), [4, 5, 6, 7, 0, 1, 2, 3]);
+    /// // Reverse lanes, within each 4-lane group:
+    /// assert_eq!(x.general_reverse::<3>().to_array(), [3, 2, 1, 0, 7, 6, 5, 4]);
+    /// ```
+    ///
+    /// Commonly useful for horizontal reductions, for example:
+    ///
+    /// ```
+    /// # #![feature(portable_simd)]
+    /// # #[cfg(feature = "as_crate")] use core_simd::simd::Simd;
+    /// # #[cfg(not(feature = "as_crate"))] use core::simd::Simd;
+    /// let x = Simd::from_array([0u32, 1, 2, 3, 4, 5, 6, 7]);
+    /// let x = x + x.general_reverse::<1>();
+    /// let x = x + x.general_reverse::<2>();
+    /// let x = x + x.general_reverse::<4>();
+    /// assert_eq!(x.to_array(), [28, 28, 28, 28, 28, 28, 28, 28]);
+    /// ```
+    #[inline]
+    #[must_use = "method returns a new vector and does not mutate the original inputs"]
+    #[doc(alias = "grev")]
+    #[doc(alias = "butterfly")]
+    #[doc(alias = "bfly")]
+    pub fn general_reverse<const SWAP_MASK: usize>(self) -> Self {
+        const fn general_reverse_index<const LANES: usize>(swap_mask: usize) -> [usize; LANES] {
+            let mut index = [0; LANES];
+            let mut i = 0;
+            while i < LANES {
+                index[i] = i ^ swap_mask;
+                i += 1;
+            }
+            index
+        }
+        struct GeneralReverse<const DISTANCE: usize>;
+        impl<const LANES: usize, const DISTANCE: usize> Swizzle<LANES, LANES> for GeneralReverse<DISTANCE> {
+            const INDEX: [usize; LANES] = general_reverse_index::<LANES>(DISTANCE);
+        }
+        GeneralReverse::<SWAP_MASK>::swizzle(self)
+    }
 }

crates/core_simd/tests/swizzle.rs

@@ -74,3 +74,57 @@ fn interleave_one() {
     assert_eq!(even, a);
     assert_eq!(odd, b);
 }
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn slice() {
+    let a = Simd::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
+    let [lo, hi] = a.split_to::<4>();
+    assert_eq!(lo.to_array(), [0, 1, 2, 3]);
+    assert_eq!(hi.to_array(), [4, 5, 6, 7]);
+}
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn concat() {
+    let x = Simd::from_array([0, 1, 2, 3]);
+    let y = Simd::from_array([4, 5, 6, 7]);
+    let z = x.concat_to::<8>(y);
+    assert_eq!(z.to_array(), [0, 1, 2, 3, 4, 5, 6, 7]);
+}
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn general_reverse() {
+    let x = Simd::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
+    // Swap adjacent lanes:
+    assert_eq!(
+        x.general_reverse::<1>().to_array(),
+        [1, 0, 3, 2, 5, 4, 7, 6]
+    );
+    // Swap lanes separated by distance 2:
+    assert_eq!(
+        x.general_reverse::<2>().to_array(),
+        [2, 3, 0, 1, 6, 7, 4, 5]
+    );
+    // Swap lanes separated by distance 4:
+    assert_eq!(
+        x.general_reverse::<4>().to_array(),
+        [4, 5, 6, 7, 0, 1, 2, 3]
+    );
+    // Reverse lanes, within each 4-lane group:
+    assert_eq!(
+        x.general_reverse::<3>().to_array(),
+        [3, 2, 1, 0, 7, 6, 5, 4]
+    );
+}
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn general_reverse_sum() {
+    let x = Simd::from_array([0u32, 1, 2, 3, 4, 5, 6, 7]);
+    let x = x + x.general_reverse::<1>();
+    let x = x + x.general_reverse::<2>();
+    let x = x + x.general_reverse::<4>();
+    assert_eq!(x.to_array(), [28, 28, 28, 28, 28, 28, 28, 28]);
+}