From 06a1070f1ffe377981dec550b1ebfb679a1b973f Mon Sep 17 00:00:00 2001 From: Caleb Zulawski Date: Sun, 1 Oct 2023 12:31:39 -0400 Subject: [PATCH] Simplify Swizzle trait and condense all swizzles into this trait --- crates/core_simd/examples/matrix_inversion.rs | 9 +- crates/core_simd/src/lib.rs | 1 + crates/core_simd/src/swizzle.rs | 349 ++++++++---------- crates/core_simd/src/vector.rs | 4 +- crates/core_simd/tests/swizzle.rs | 4 +- 5 files changed, 173 insertions(+), 194 deletions(-) diff --git a/crates/core_simd/examples/matrix_inversion.rs b/crates/core_simd/examples/matrix_inversion.rs index 39f530f68f5..979873d885d 100644 --- a/crates/core_simd/examples/matrix_inversion.rs +++ b/crates/core_simd/examples/matrix_inversion.rs @@ -3,7 +3,6 @@ // Run this code with `cargo test --example matrix_inversion` #![feature(array_chunks, portable_simd)] use core_simd::simd::*; -use Which::*; // Gotta define our own 4x4 matrix since Rust doesn't ship multidim arrays yet :^) #[derive(Copy, Clone, Debug, PartialEq, PartialOrd)] @@ -164,10 +163,10 @@ pub fn simd_inv4x4(m: Matrix4x4) -> Option { let m_2 = f32x4::from_array(m[2]); let m_3 = f32x4::from_array(m[3]); - const SHUFFLE01: [Which; 4] = [First(0), First(1), Second(0), Second(1)]; - const SHUFFLE02: [Which; 4] = [First(0), First(2), Second(0), Second(2)]; - const SHUFFLE13: [Which; 4] = [First(1), First(3), Second(1), Second(3)]; - const SHUFFLE23: [Which; 4] = [First(2), First(3), Second(2), Second(3)]; + const SHUFFLE01: [usize; 4] = [0, 1, 4, 5]; + const SHUFFLE02: [usize; 4] = [0, 2, 4, 6]; + const SHUFFLE13: [usize; 4] = [1, 3, 5, 7]; + const SHUFFLE23: [usize; 4] = [2, 3, 6, 7]; let tmp = simd_swizzle!(m_0, m_1, SHUFFLE01); let row1 = simd_swizzle!(m_2, m_3, SHUFFLE01); diff --git a/crates/core_simd/src/lib.rs b/crates/core_simd/src/lib.rs index 2d68e4cce85..74ed0136003 100644 --- a/crates/core_simd/src/lib.rs +++ b/crates/core_simd/src/lib.rs @@ -5,6 +5,7 @@ const_mut_refs, convert_float_to_int, decl_macro, + inline_const, intra_doc_pointers, platform_intrinsics, repr_simd, diff --git a/crates/core_simd/src/swizzle.rs b/crates/core_simd/src/swizzle.rs index 68f20516cf5..fb257e34cf9 100644 --- a/crates/core_simd/src/swizzle.rs +++ b/crates/core_simd/src/swizzle.rs @@ -1,17 +1,15 @@ use crate::simd::intrinsics; -use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount}; +use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount}; -/// Constructs a new SIMD vector by copying elements from selected lanes in other vectors. +/// Constructs a new SIMD vector by copying elements from selected elements in other vectors. /// -/// When swizzling one vector, lanes are selected by a `const` array of `usize`, -/// like [`Swizzle`]. +/// When swizzling one vector, elements are selected like [`Swizzle::swizzle`]. /// -/// When swizzling two vectors, lanes are selected by a `const` array of [`Which`], -/// like [`Swizzle2`]. +/// When swizzling two vectors, elements are selected like [`Swizzle::concat_swizzle`]. /// /// # Examples /// -/// With a single SIMD vector, the const array specifies lane indices in that vector: +/// With a single SIMD vector, the const array specifies element indices in that vector: /// ``` /// # #![feature(portable_simd)] /// # use core::simd::{u32x2, u32x4, simd_swizzle}; @@ -21,25 +19,27 @@ use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount}; /// let r: u32x4 = simd_swizzle!(v, [3, 0, 1, 2]); /// assert_eq!(r.to_array(), [13, 10, 11, 12]); /// -/// // Changing the number of lanes +/// // Changing the number of elements /// let r: u32x2 = simd_swizzle!(v, [3, 1]); /// assert_eq!(r.to_array(), [13, 11]); /// ``` /// -/// With two input SIMD vectors, the const array uses `Which` to specify the source of each index: +/// With two input SIMD vectors, the const array specifies element indices in the concatenation of +/// those vectors: /// ``` /// # #![feature(portable_simd)] -/// # use core::simd::{u32x2, u32x4, simd_swizzle, Which}; -/// use Which::{First, Second}; +/// # #[cfg(feature = "as_crate")] use core_simd::simd; +/// # #[cfg(not(feature = "as_crate"))] use core::simd; +/// # use simd::{u32x2, u32x4, simd_swizzle}; /// let a = u32x4::from_array([0, 1, 2, 3]); /// let b = u32x4::from_array([4, 5, 6, 7]); /// /// // Keeping the same size -/// let r: u32x4 = simd_swizzle!(a, b, [First(0), First(1), Second(2), Second(3)]); +/// let r: u32x4 = simd_swizzle!(a, b, [0, 1, 6, 7]); /// assert_eq!(r.to_array(), [0, 1, 6, 7]); /// -/// // Changing the number of lanes -/// let r: u32x2 = simd_swizzle!(a, b, [First(0), Second(0)]); +/// // Changing the number of elements +/// let r: u32x2 = simd_swizzle!(a, b, [0, 4]); /// assert_eq!(r.to_array(), [0, 4]); /// ``` #[allow(unused_macros)] @@ -50,7 +50,7 @@ pub macro simd_swizzle { { use $crate::simd::Swizzle; struct Impl; - impl Swizzle for Impl { + impl Swizzle<{$index.len()}> for Impl { const INDEX: [usize; {$index.len()}] = $index; } Impl::swizzle($vector) @@ -60,127 +60,117 @@ pub macro simd_swizzle { $first:expr, $second:expr, $index:expr $(,)? ) => { { - use $crate::simd::{Which, Swizzle2}; + use $crate::simd::Swizzle; struct Impl; - impl Swizzle2 for Impl { - const INDEX: [Which; {$index.len()}] = $index; + impl Swizzle<{$index.len()}> for Impl { + const INDEX: [usize; {$index.len()}] = $index; } - Impl::swizzle2($first, $second) + Impl::concat_swizzle($first, $second) } } } -/// Specifies a lane index into one of two SIMD vectors. -/// -/// This is an input type for [Swizzle2] and helper macros like [simd_swizzle]. -#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub enum Which { - /// Index of a lane in the first input SIMD vector. - First(usize), - /// Index of a lane in the second input SIMD vector. - Second(usize), -} - /// Create a vector from the elements of another vector. -pub trait Swizzle { - /// Map from the lanes of the input vector to the output vector. - const INDEX: [usize; OUTPUT_LANES]; +pub trait Swizzle { + /// Map from the elements of the input vector to the output vector. + const INDEX: [usize; N]; - /// Create a new vector from the lanes of `vector`. + /// Create a new vector from the elements of `vector`. /// /// Lane `i` of the output is `vector[Self::INDEX[i]]`. #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] - fn swizzle(vector: Simd) -> Simd + fn swizzle(vector: Simd) -> Simd where T: SimdElement, - LaneCount: SupportedLaneCount, - LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, { - // Safety: `vector` is a vector, and `INDEX_IMPL` is a const array of u32. - unsafe { intrinsics::simd_shuffle(vector, vector, Self::INDEX_IMPL) } + // Safety: `vector` is a vector, and the index is a const array of u32. + unsafe { + intrinsics::simd_shuffle( + vector, + vector, + const { + let mut output = [0; N]; + let mut i = 0; + while i < N { + let index = Self::INDEX[i]; + assert!(index as u32 as usize == index); + assert!(index < M, "source element index exceeds input vector length"); + output[i] = index as u32; + i += 1; + } + output + }, + ) + } } -} -/// Create a vector from the elements of two other vectors. -pub trait Swizzle2 { - /// Map from the lanes of the input vectors to the output vector - const INDEX: [Which; OUTPUT_LANES]; - - /// Create a new vector from the lanes of `first` and `second`. + /// Create a new vector from the elements of `first` and `second`. /// - /// Lane `i` is `first[j]` when `Self::INDEX[i]` is `First(j)`, or `second[j]` when it is - /// `Second(j)`. + /// Lane `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of + /// `first` and `second`. #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] - fn swizzle2( - first: Simd, - second: Simd, - ) -> Simd + fn concat_swizzle(first: Simd, second: Simd) -> Simd where T: SimdElement, - LaneCount: SupportedLaneCount, - LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, { - // Safety: `first` and `second` are vectors, and `INDEX_IMPL` is a const array of u32. - unsafe { intrinsics::simd_shuffle(first, second, Self::INDEX_IMPL) } - } -} - -/// The `simd_shuffle` intrinsic expects `u32`, so do error checking and conversion here. -/// This trait hides `INDEX_IMPL` from the public API. -trait SwizzleImpl { - const INDEX_IMPL: [u32; OUTPUT_LANES]; -} - -impl SwizzleImpl - for T -where - T: Swizzle + ?Sized, -{ - const INDEX_IMPL: [u32; OUTPUT_LANES] = { - let mut output = [0; OUTPUT_LANES]; - let mut i = 0; - while i < OUTPUT_LANES { - let index = Self::INDEX[i]; - assert!(index as u32 as usize == index); - assert!(index < INPUT_LANES, "source lane exceeds input lane count",); - output[i] = index as u32; - i += 1; + // Safety: `first` and `second` are vectors, and the index is a const array of u32. + unsafe { + intrinsics::simd_shuffle( + first, + second, + const { + let mut output = [0; N]; + let mut i = 0; + while i < N { + let index = Self::INDEX[i]; + assert!(index as u32 as usize == index); + assert!(index < 2 * M, "source element index exceeds input vector length"); + output[i] = index as u32; + i += 1; + } + output + }, + ) } - output - }; -} - -/// The `simd_shuffle` intrinsic expects `u32`, so do error checking and conversion here. -/// This trait hides `INDEX_IMPL` from the public API. -trait Swizzle2Impl { - const INDEX_IMPL: [u32; OUTPUT_LANES]; -} + } -impl Swizzle2Impl - for T -where - T: Swizzle2 + ?Sized, -{ - const INDEX_IMPL: [u32; OUTPUT_LANES] = { - let mut output = [0; OUTPUT_LANES]; - let mut i = 0; - while i < OUTPUT_LANES { - let (offset, index) = match Self::INDEX[i] { - Which::First(index) => (false, index), - Which::Second(index) => (true, index), - }; - assert!(index < INPUT_LANES, "source lane exceeds input lane count",); + /// Create a new mask from the elements of `first` and `second`. + /// + /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of + /// `first` and `second`. + #[inline] + #[must_use = "method returns a new mask and does not mutate the original inputs"] + fn swizzle_mask(vector: Mask) -> Mask + where + T: MaskElement, + LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, + { + // SAFETY: all elements of this mask come from another mask + unsafe { Mask::from_int_unchecked(Self::swizzle(vector.to_int())) } + } - // lanes are indexed by the first vector, then second vector - let index = if offset { index + INPUT_LANES } else { index }; - assert!(index as u32 as usize == index); - output[i] = index as u32; - i += 1; - } - output - }; + /// Create a new mask from the elements of `first` and `second`. + /// + /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of + /// `first` and `second`. + #[inline] + #[must_use = "method returns a new mask and does not mutate the original inputs"] + fn concat_swizzle_mask(first: Mask, second: Mask) -> Mask + where + T: MaskElement, + LaneCount: SupportedLaneCount, + LaneCount: SupportedLaneCount, + { + // SAFETY: all elements of this mask come from another mask + unsafe { Mask::from_int_unchecked(Self::concat_swizzle(first.to_int(), second.to_int())) } + } } impl Simd @@ -188,24 +178,22 @@ where T: SimdElement, LaneCount: SupportedLaneCount, { - /// Reverse the order of the lanes in the vector. + /// Reverse the order of the elements in the vector. #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] pub fn reverse(self) -> Self { - const fn reverse_index() -> [usize; LANES] { - let mut index = [0; LANES]; - let mut i = 0; - while i < LANES { - index[i] = LANES - i - 1; - i += 1; - } - index - } - struct Reverse; - impl Swizzle for Reverse { - const INDEX: [usize; LANES] = reverse_index::(); + impl Swizzle for Reverse { + const INDEX: [usize; N] = const { + let mut index = [0; N]; + let mut i = 0; + while i < N { + index[i] = N - i - 1; + i += 1; + } + index + }; } Reverse::swizzle(self) @@ -217,21 +205,19 @@ where #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] pub fn rotate_lanes_left(self) -> Self { - const fn rotate_index() -> [usize; LANES] { - let offset = OFFSET % LANES; - let mut index = [0; LANES]; - let mut i = 0; - while i < LANES { - index[i] = (i + offset) % LANES; - i += 1; - } - index - } - struct Rotate; - impl Swizzle for Rotate { - const INDEX: [usize; LANES] = rotate_index::(); + impl Swizzle for Rotate { + const INDEX: [usize; N] = const { + let offset = OFFSET % N; + let mut index = [0; N]; + let mut i = 0; + while i < N { + index[i] = (i + offset) % N; + i += 1; + } + index + }; } Rotate::::swizzle(self) @@ -243,21 +229,19 @@ where #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] pub fn rotate_lanes_right(self) -> Self { - const fn rotate_index() -> [usize; LANES] { - let offset = LANES - OFFSET % LANES; - let mut index = [0; LANES]; - let mut i = 0; - while i < LANES { - index[i] = (i + offset) % LANES; - i += 1; - } - index - } - struct Rotate; - impl Swizzle for Rotate { - const INDEX: [usize; LANES] = rotate_index::(); + impl Swizzle for Rotate { + const INDEX: [usize; N] = const { + let offset = N - OFFSET % N; + let mut index = [0; N]; + let mut i = 0; + while i < N { + index[i] = (i + offset) % N; + i += 1; + } + index + }; } Rotate::::swizzle(self) @@ -265,7 +249,7 @@ where /// Interleave two vectors. /// - /// The resulting vectors contain lanes taken alternatively from `self` and `other`, first + /// The resulting vectors contain elements taken alternatively from `self` and `other`, first /// filling the first result, and then the second. /// /// The reverse of this operation is [`Simd::deinterleave`]. @@ -282,18 +266,13 @@ where #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] pub fn interleave(self, other: Self) -> (Self, Self) { - const fn interleave(high: bool) -> [Which; LANES] { - let mut idx = [Which::First(0); LANES]; + const fn interleave(high: bool) -> [usize; N] { + let mut idx = [0; N]; let mut i = 0; - while i < LANES { - // Treat the source as a concatenated vector - let dst_index = if high { i + LANES } else { i }; - let src_index = dst_index / 2 + (dst_index % 2) * LANES; - idx[i] = if src_index < LANES { - Which::First(src_index) - } else { - Which::Second(src_index % LANES) - }; + while i < N { + let dst_index = if high { i + N } else { i }; + let src_index = dst_index / 2 + (dst_index % 2) * N; + idx[i] = src_index; i += 1; } idx @@ -302,24 +281,27 @@ where struct Lo; struct Hi; - impl Swizzle2 for Lo { - const INDEX: [Which; LANES] = interleave::(false); + impl Swizzle for Lo { + const INDEX: [usize; N] = interleave::(false); } - impl Swizzle2 for Hi { - const INDEX: [Which; LANES] = interleave::(true); + impl Swizzle for Hi { + const INDEX: [usize; N] = interleave::(true); } - (Lo::swizzle2(self, other), Hi::swizzle2(self, other)) + ( + Lo::concat_swizzle(self, other), + Hi::concat_swizzle(self, other), + ) } /// Deinterleave two vectors. /// - /// The first result takes every other lane of `self` and then `other`, starting with - /// the first lane. + /// The first result takes every other element of `self` and then `other`, starting with + /// the first element. /// - /// The second result takes every other lane of `self` and then `other`, starting with - /// the second lane. + /// The second result takes every other element of `self` and then `other`, starting with + /// the second element. /// /// The reverse of this operation is [`Simd::interleave`]. /// @@ -335,17 +317,11 @@ where #[inline] #[must_use = "method returns a new vector and does not mutate the original inputs"] pub fn deinterleave(self, other: Self) -> (Self, Self) { - const fn deinterleave(second: bool) -> [Which; LANES] { - let mut idx = [Which::First(0); LANES]; + const fn deinterleave(second: bool) -> [usize; N] { + let mut idx = [0; N]; let mut i = 0; - while i < LANES { - // Treat the source as a concatenated vector - let src_index = i * 2 + second as usize; - idx[i] = if src_index < LANES { - Which::First(src_index) - } else { - Which::Second(src_index % LANES) - }; + while i < N { + idx[i] = i * 2 + second as usize; i += 1; } idx @@ -354,14 +330,17 @@ where struct Even; struct Odd; - impl Swizzle2 for Even { - const INDEX: [Which; LANES] = deinterleave::(false); + impl Swizzle for Even { + const INDEX: [usize; N] = deinterleave::(false); } - impl Swizzle2 for Odd { - const INDEX: [Which; LANES] = deinterleave::(true); + impl Swizzle for Odd { + const INDEX: [usize; N] = deinterleave::(true); } - (Even::swizzle2(self, other), Odd::swizzle2(self, other)) + ( + Even::concat_swizzle(self, other), + Odd::concat_swizzle(self, other), + ) } } diff --git a/crates/core_simd/src/vector.rs b/crates/core_simd/src/vector.rs index 9aa7bacfce9..92b1b61d92c 100644 --- a/crates/core_simd/src/vector.rs +++ b/crates/core_simd/src/vector.rs @@ -142,10 +142,10 @@ where // This is preferred over `[value; N]`, since it's explicitly a splat: // https://github.com/rust-lang/rust/issues/97804 struct Splat; - impl Swizzle<1, N> for Splat { + impl Swizzle for Splat { const INDEX: [usize; N] = [0; N]; } - Splat::swizzle(Simd::::from([value])) + Splat::swizzle::(Simd::::from([value])) } /// Returns an array reference containing the entire SIMD vector. diff --git a/crates/core_simd/tests/swizzle.rs b/crates/core_simd/tests/swizzle.rs index 8cd7c33e823..46aaf748ad8 100644 --- a/crates/core_simd/tests/swizzle.rs +++ b/crates/core_simd/tests/swizzle.rs @@ -11,10 +11,10 @@ wasm_bindgen_test_configure!(run_in_browser); #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)] fn swizzle() { struct Index; - impl Swizzle<4, 4> for Index { + impl Swizzle<4> for Index { const INDEX: [usize; 4] = [2, 1, 3, 0]; } - impl Swizzle<4, 2> for Index { + impl Swizzle<2> for Index { const INDEX: [usize; 2] = [1, 1]; }