Implement common swizzle operations.

crates/core_simd/src/swizzle.rs

@@ -364,4 +364,176 @@ 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 or constrain it by context. 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();
+    /// 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, change signature to
+    //   `pub fn split(self) -> [Simd<T, {LANES / 2}>; 2]`
+    pub fn split<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 or constrain it by context.
+    /// A compile time error will be raised if `LANES + LANES2 != OUTPUT_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(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, change signature to
+    //   `pub fn concat<const LANES2>(self, other: Simd<T, LANES2>) -> Simd<T, {LANES + LANES2}>`
+    pub fn concat<const OUTPUT_LANES: usize, const LANES2: usize>(
+        self,
+        other: Simd<T, LANES2>,
+    ) -> Simd<T, OUTPUT_LANES>
+    where
+        LaneCount<OUTPUT_LANES>: SupportedLaneCount,
+        LaneCount<LANES2>: SupportedLaneCount,
+    {
+        struct Extend;
+        impl<const I: usize, const O: usize> Swizzle<I, O> for Extend {
+            const INDEX: [usize; O] = {
+                assert!(I <= O);
+                let mut index = [0; O];
+                let mut i = 0;
+                while i < I {
+                    index[i] = i;
+                    i += 1;
+                }
+                index
+            };
+        }
+        struct Concat<const A: usize, const B: usize, const Y: usize>;
+        impl<const A: usize, const B: usize, const Y: usize> Swizzle2<Y, Y> for Concat<A, B, Y> {
+            const INDEX: [Which; Y] = {
+                assert!(
+                    A + B == Y,
+                    "concat: OUTPUT_LANES must be the sum of all input lane counts"
+                );
+                let mut retval = [Which::First(0); Y];
+                let mut i = 0;
+                while i < Y {
+                    if i < A {
+                        retval[i] = Which::First(i);
+                    } else {
+                        retval[i] = Which::Second(i - A);
+                    }
+                    i += 1;
+                }
+                retval
+            };
+        }
+        Concat::<LANES, LANES2, OUTPUT_LANES>::swizzle2(
+            Extend::swizzle(self),
+            Extend::swizzle(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,67 @@ 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();
+    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_equal_width() {
+    let x = Simd::from_array([0, 1, 2, 3]);
+    let y = Simd::from_array([4, 5, 6, 7]);
+    let z = x.concat(y);
+    assert_eq!(z.to_array(), [0, 1, 2, 3, 4, 5, 6, 7]);
+}
+
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+#[cfg(feature = "all_lane_counts")]
+fn concat_different_width() {
+    let x = Simd::from_array([0, 1, 2, 3]);
+    let y = Simd::from_array([4, 5]);
+    let z = x.concat(y);
+    assert_eq!(z.to_array(), [0, 1, 2, 3, 4, 5]);
+}
+
+#[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]);
+}