Merge pull request #299 from rust-lang/interleave-one

Fix interleave/deinterleave for vectors with only one lane
This commit is contained in:
Caleb Zulawski 2022-08-03 23:23:19 -04:00 committed by GitHub
commit 7145dc5a41
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 41 additions and 45 deletions

View File

@ -265,16 +265,13 @@ where
/// Interleave two vectors.
///
/// Produces two vectors with lanes taken alternately from `self` and `other`.
/// The resulting vectors contain lanes taken alternatively from `self` and `other`, first
/// filling the first result, and then the second.
///
/// 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`.
/// The reverse of this operation is [`Simd::deinterleave`].
///
/// ```
/// #![feature(portable_simd)]
/// # #![feature(portable_simd)]
/// # use core::simd::Simd;
/// let a = Simd::from_array([0, 1, 2, 3]);
/// let b = Simd::from_array([4, 5, 6, 7]);
@ -285,29 +282,17 @@ 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 lo<const LANES: usize>() -> [Which; LANES] {
const fn interleave<const LANES: usize>(high: bool) -> [Which; LANES] {
let mut idx = [Which::First(0); LANES];
let mut i = 0;
while i < LANES {
let offset = i / 2;
idx[i] = if i % 2 == 0 {
Which::First(offset)
// 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(offset)
};
i += 1;
}
idx
}
const fn hi<const LANES: usize>() -> [Which; LANES] {
let mut idx = [Which::First(0); LANES];
let mut i = 0;
while i < LANES {
let offset = (LANES + i) / 2;
idx[i] = if i % 2 == 0 {
Which::First(offset)
} else {
Which::Second(offset)
Which::Second(src_index % LANES)
};
i += 1;
}
@ -318,11 +303,11 @@ where
struct Hi;
impl<const LANES: usize> Swizzle2<LANES, LANES> for Lo {
const INDEX: [Which; LANES] = lo::<LANES>();
const INDEX: [Which; LANES] = interleave::<LANES>(false);
}
impl<const LANES: usize> Swizzle2<LANES, LANES> for Hi {
const INDEX: [Which; LANES] = hi::<LANES>();
const INDEX: [Which; LANES] = interleave::<LANES>(true);
}
(Lo::swizzle2(self, other), Hi::swizzle2(self, other))
@ -336,8 +321,10 @@ where
/// The second result takes every other lane of `self` and then `other`, starting with
/// the second lane.
///
/// The reverse of this operation is [`Simd::interleave`].
///
/// ```
/// #![feature(portable_simd)]
/// # #![feature(portable_simd)]
/// # use core::simd::Simd;
/// let a = Simd::from_array([0, 4, 1, 5]);
/// let b = Simd::from_array([2, 6, 3, 7]);
@ -348,22 +335,17 @@ 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 even<const LANES: usize>() -> [Which; LANES] {
const fn deinterleave<const LANES: usize>(second: bool) -> [Which; LANES] {
let mut idx = [Which::First(0); LANES];
let mut i = 0;
while i < LANES / 2 {
idx[i] = Which::First(2 * i);
idx[i + LANES / 2] = Which::Second(2 * i);
i += 1;
}
idx
}
const fn odd<const LANES: usize>() -> [Which; LANES] {
let mut idx = [Which::First(0); LANES];
let mut i = 0;
while i < LANES / 2 {
idx[i] = Which::First(2 * i + 1);
idx[i + LANES / 2] = Which::Second(2 * i + 1);
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)
};
i += 1;
}
idx
@ -373,11 +355,11 @@ where
struct Odd;
impl<const LANES: usize> Swizzle2<LANES, LANES> for Even {
const INDEX: [Which; LANES] = even::<LANES>();
const INDEX: [Which; LANES] = deinterleave::<LANES>(false);
}
impl<const LANES: usize> Swizzle2<LANES, LANES> for Odd {
const INDEX: [Which; LANES] = odd::<LANES>();
const INDEX: [Which; LANES] = deinterleave::<LANES>(true);
}
(Even::swizzle2(self, other), Odd::swizzle2(self, other))

View File

@ -60,3 +60,17 @@ fn interleave() {
assert_eq!(even, a);
assert_eq!(odd, b);
}
// portable-simd#298
#[test]
#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
fn interleave_one() {
let a = Simd::from_array([0]);
let b = Simd::from_array([1]);
let (lo, hi) = a.interleave(b);
assert_eq!(lo.to_array(), [0]);
assert_eq!(hi.to_array(), [1]);
let (even, odd) = lo.deinterleave(hi);
assert_eq!(even, a);
assert_eq!(odd, b);
}