Skip to content

Commit

Permalink
Merge pull request rust-lang#98 from rust-lang/feature/common-shuffles
Browse files Browse the repository at this point in the history
Add some common shuffles
  • Loading branch information
workingjubilee authored Apr 21, 2021
2 parents b2e25bc + 7028a58 commit 2fa62b9
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 0 deletions.
109 changes: 109 additions & 0 deletions crates/core_simd/src/permute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,115 @@ macro_rules! impl_shuffle_lane {
pub fn shuffle<const IDX: [u32; $n]>(self, second: Self) -> Self {
unsafe { crate::intrinsics::$fn(self, second, IDX) }
}

/// Reverse the order of the lanes in the vector.
#[inline]
pub fn reverse(self) -> Self {
const fn idx() -> [u32; $n] {
let mut idx = [0u32; $n];
let mut i = 0;
while i < $n {
idx[i] = ($n - i - 1) as u32;
i += 1;
}
idx
}
self.shuffle::<{ idx() }>(self)
}

/// Interleave two vectors.
///
/// Produces two vectors with lanes taken alternately from `self` and `other`.
///
/// The first result contains the first `LANES / 2` lanes from `self` and `other`,
/// alternating, starting with the first lane of `self`.
///
/// The second result contains the last `LANES / 2` lanes from `self` and `other`,
/// alternating, starting with the lane `LANES / 2` from the start of `self`.
///
/// This particular permutation is efficient on many architectures.
///
/// ```
/// # use core_simd::SimdU32;
/// let a = SimdU32::from_array([0, 1, 2, 3]);
/// let b = SimdU32::from_array([4, 5, 6, 7]);
/// let (x, y) = a.interleave(b);
/// assert_eq!(x.to_array(), [0, 4, 1, 5]);
/// assert_eq!(y.to_array(), [2, 6, 3, 7]);
/// ```
#[inline]
pub fn interleave(self, other: Self) -> (Self, Self) {
const fn lo() -> [u32; $n] {
let mut idx = [0u32; $n];
let mut i = 0;
while i < $n {
let offset = i / 2;
idx[i] = if i % 2 == 0 {
offset
} else {
$n + offset
} as u32;
i += 1;
}
idx
}
const fn hi() -> [u32; $n] {
let mut idx = [0u32; $n];
let mut i = 0;
while i < $n {
let offset = ($n + i) / 2;
idx[i] = if i % 2 == 0 {
offset
} else {
$n + offset
} as u32;
i += 1;
}
idx
}
(self.shuffle::<{ lo() }>(other), self.shuffle::<{ hi() }>(other))
}

/// Deinterleave two vectors.
///
/// The first result takes every other lane of `self` and then `other`, starting with
/// the first lane.
///
/// The second result takes every other lane of `self` and then `other`, starting with
/// the second lane.
///
/// This particular permutation is efficient on many architectures.
///
/// ```
/// # use core_simd::SimdU32;
/// let a = SimdU32::from_array([0, 4, 1, 5]);
/// let b = SimdU32::from_array([2, 6, 3, 7]);
/// let (x, y) = a.deinterleave(b);
/// assert_eq!(x.to_array(), [0, 1, 2, 3]);
/// assert_eq!(y.to_array(), [4, 5, 6, 7]);
/// ```
#[inline]
pub fn deinterleave(self, other: Self) -> (Self, Self) {
const fn even() -> [u32; $n] {
let mut idx = [0u32; $n];
let mut i = 0;
while i < $n {
idx[i] = 2 * i as u32;
i += 1;
}
idx
}
const fn odd() -> [u32; $n] {
let mut idx = [0u32; $n];
let mut i = 0;
while i < $n {
idx[i] = 1 + 2 * i as u32;
i += 1;
}
idx
}
(self.shuffle::<{ even() }>(other), self.shuffle::<{ odd() }>(other))
}
}
}
}
Expand Down
20 changes: 20 additions & 0 deletions crates/core_simd/tests/permute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,23 @@ fn simple_shuffle() {
let b = a;
assert_eq!(a.shuffle::<{ [3, 1, 4, 6] }>(b).to_array(), [9, 4, 2, 1]);
}

#[test]
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
fn reverse() {
let a = SimdU32::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
assert_eq!(a.reverse().to_array(), [7, 6, 5, 4, 3, 2, 1, 0]);
}

#[test]
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
fn interleave() {
let a = SimdU32::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
let b = SimdU32::from_array([8, 9, 10, 11, 12, 13, 14, 15]);
let (lo, hi) = a.interleave(b);
assert_eq!(lo.to_array(), [0, 8, 1, 9, 2, 10, 3, 11]);
assert_eq!(hi.to_array(), [4, 12, 5, 13, 6, 14, 7, 15]);
let (even, odd) = lo.deinterleave(hi);
assert_eq!(even, a);
assert_eq!(odd, b);
}

0 comments on commit 2fa62b9

Please sign in to comment.