Skip to content

Commit 977f26f

Browse files
committed
Add some common shuffles
1 parent e73985f commit 977f26f

File tree

2 files changed

+103
-0
lines changed

2 files changed

+103
-0
lines changed

crates/core_simd/src/permute.rs

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,89 @@ macro_rules! impl_shuffle_lane {
1313
pub fn shuffle<const IDX: [u32; $n]>(self, second: Self) -> Self {
1414
unsafe { crate::intrinsics::$fn(self, second, IDX) }
1515
}
16+
17+
/// Reverse the order of the lanes in the vector.
18+
#[inline]
19+
pub fn reverse(self) -> Self {
20+
const fn idx() -> [u32; $n] {
21+
let mut idx = [0u32; $n];
22+
let mut i = 0;
23+
while i < $n {
24+
idx[i] = ($n - i - 1) as u32;
25+
i += 1;
26+
}
27+
idx
28+
}
29+
self.shuffle::<{ idx() }>(self)
30+
}
31+
32+
/// Interleave two vectors.
33+
///
34+
/// The even lanes of the first result contain the lower half of `self`, and the odd
35+
/// lanes contain the lower half of `other`.
36+
///
37+
/// The even lanes of the second result contain the upper half of `self`, and the odd
38+
/// lanes contain the upper half of `other`.
39+
#[inline]
40+
pub fn interleave(self, other: Self) -> (Self, Self) {
41+
const fn lo() -> [u32; $n] {
42+
let mut idx = [0u32; $n];
43+
let mut i = 0;
44+
while i < $n {
45+
let offset = i / 2;
46+
idx[i] = if i % 2 == 0 {
47+
offset
48+
} else {
49+
$n + offset
50+
} as u32;
51+
i += 1;
52+
}
53+
idx
54+
}
55+
const fn hi() -> [u32; $n] {
56+
let mut idx = [0u32; $n];
57+
let mut i = 0;
58+
while i < $n {
59+
let offset = ($n + i) / 2;
60+
idx[i] = if i % 2 == 0 {
61+
offset
62+
} else {
63+
$n + offset
64+
} as u32;
65+
i += 1;
66+
}
67+
idx
68+
}
69+
(self.shuffle::<{ lo() }>(other), self.shuffle::<{ hi() }>(other))
70+
}
71+
72+
/// Deinterleave two vectors.
73+
///
74+
/// The first result contains the even lanes of `self` and `other` concatenated.
75+
///
76+
/// The second result contains the odd lanes of `self` and `other` concatenated.
77+
#[inline]
78+
pub fn deinterleave(self, other: Self) -> (Self, Self) {
79+
const fn even() -> [u32; $n] {
80+
let mut idx = [0u32; $n];
81+
let mut i = 0;
82+
while i < $n {
83+
idx[i] = 2 * i as u32;
84+
i += 1;
85+
}
86+
idx
87+
}
88+
const fn odd() -> [u32; $n] {
89+
let mut idx = [0u32; $n];
90+
let mut i = 0;
91+
while i < $n {
92+
idx[i] = 1 + 2 * i as u32;
93+
i += 1;
94+
}
95+
idx
96+
}
97+
(self.shuffle::<{ even() }>(other), self.shuffle::<{ odd() }>(other))
98+
}
1699
}
17100
}
18101
}

crates/core_simd/tests/permute.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,23 @@ fn simple_shuffle() {
1313
let b = a;
1414
assert_eq!(a.shuffle::<{ [3, 1, 4, 6] }>(b).to_array(), [9, 4, 2, 1]);
1515
}
16+
17+
#[test]
18+
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
19+
fn reverse() {
20+
let a = SimdU32::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
21+
assert_eq!(a.reverse().to_array(), [7, 6, 5, 4, 3, 2, 1, 0]);
22+
}
23+
24+
#[test]
25+
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
26+
fn interleave() {
27+
let a = SimdU32::from_array([0, 1, 2, 3, 4, 5, 6, 7]);
28+
let b = SimdU32::from_array([8, 9, 10, 11, 12, 13, 14, 15]);
29+
let (lo, hi) = a.interleave(b);
30+
assert_eq!(lo.to_array(), [0, 8, 1, 9, 2, 10, 3, 11]);
31+
assert_eq!(hi.to_array(), [4, 12, 5, 13, 6, 14, 7, 15]);
32+
let (even, odd) = lo.deinterleave(hi);
33+
assert_eq!(even, a);
34+
assert_eq!(odd, b);
35+
}

0 commit comments

Comments
 (0)