Skip to content

Commit 92c8317

Browse files
committed
Add [T]::as_simd(_mut)
SIMD-style optimizations are the most common use for `[T]::align_to(_mut)`, but that's `unsafe`. So these are *safe* wrappers around it, now that we have the `Simd` type available, to make it easier to use. ```rust impl [T] { pub fn as_simd<const LANES: usize>(&self) -> (&[T], &[Simd<T, LANES>], &[T]); pub fn as_simd_mut<const LANES: usize>(&mut self) -> (&mut [T], &mut [Simd<T, LANES>], &mut [T]); } ```
1 parent ff23ad3 commit 92c8317

File tree

1 file changed

+83
-0
lines changed

1 file changed

+83
-0
lines changed

library/core/src/slice/mod.rs

+83
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ use crate::option::Option::{None, Some};
1616
use crate::ptr;
1717
use crate::result::Result;
1818
use crate::result::Result::{Err, Ok};
19+
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
20+
use crate::simd::{self, Simd};
1921
use crate::slice;
2022

2123
#[unstable(
@@ -3434,6 +3436,87 @@ impl<T> [T] {
34343436
}
34353437
}
34363438

3439+
/// Split a slice into a prefix, a middle of aligned simd types, and a suffix.
3440+
///
3441+
/// This is a safe wrapper around [`slice::align_to`], so has the same weak
3442+
/// preconditions as that method. Notably, you must not assume any particular
3443+
/// split between the three parts: it's legal for the middle slice to be
3444+
/// empty even if the input slice is longer than `3 * LANES`.
3445+
///
3446+
/// # Examples
3447+
///
3448+
/// ```
3449+
/// #![feature(portable_simd)]
3450+
///
3451+
/// let short = &[1, 2, 3];
3452+
/// let (prefix, middle, suffix) = short.as_simd::<4>();
3453+
/// assert_eq!(middle, []); // Not enough elements for anything in the middle
3454+
///
3455+
/// // They might be split in any possible way between prefix and suffix
3456+
/// let it = prefix.iter().chain(suffix).copied();
3457+
/// assert_eq!(it.collect::<Vec<_>>(), vec![1, 2, 3]);
3458+
///
3459+
/// fn basic_simd_sum(x: &[f32]) -> f32 {
3460+
/// use std::ops::Add;
3461+
/// use std::simd::f32x4;
3462+
/// let (prefix, middle, suffix) = x.as_simd();
3463+
/// let sums = f32x4::from_array([
3464+
/// prefix.iter().copied().sum(),
3465+
/// 0.0,
3466+
/// 0.0,
3467+
/// suffix.iter().copied().sum(),
3468+
/// ]);
3469+
/// let sums = middle.iter().copied().fold(sums, f32x4::add);
3470+
/// sums.horizontal_sum()
3471+
/// }
3472+
///
3473+
/// let numbers: Vec<f32> = (1..101).map(|x| x as _).collect();
3474+
/// assert_eq!(basic_simd_sum(&numbers[1..99]), 4949.0);
3475+
/// ```
3476+
#[unstable(feature = "portable_simd", issue = "86656")]
3477+
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
3478+
pub fn as_simd<const LANES: usize>(&self) -> (&[T], &[Simd<T, LANES>], &[T])
3479+
where
3480+
Simd<T, LANES>: AsRef<[T; LANES]>,
3481+
T: simd::SimdElement,
3482+
simd::LaneCount<LANES>: simd::SupportedLaneCount,
3483+
{
3484+
// These are expected to always match, as vector types are laid out like
3485+
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
3486+
// might as well double-check since it'll optimize away anyhow.
3487+
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());
3488+
3489+
// SAFETY: The simd types have the same layout as arrays, just with
3490+
// potentially-higher alignment, so the de-facto transmutes are sound.
3491+
unsafe { self.align_to() }
3492+
}
3493+
3494+
/// Split a slice into a prefix, a middle of aligned simd types, and a suffix.
3495+
///
3496+
/// This is a safe wrapper around [`slice::align_to`], so has the same weak
3497+
/// preconditions as that method. Notably, you must not assume any particular
3498+
/// split between the three parts: it's legal for the middle slice to be
3499+
/// empty even if the input slice is longer than `3 * LANES`.
3500+
///
3501+
/// This is the mutable version of [`slice::as_simd`]; see that for more.
3502+
#[unstable(feature = "portable_simd", issue = "86656")]
3503+
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
3504+
pub fn as_simd_mut<const LANES: usize>(&mut self) -> (&mut [T], &mut [Simd<T, LANES>], &mut [T])
3505+
where
3506+
Simd<T, LANES>: AsMut<[T; LANES]>,
3507+
T: simd::SimdElement,
3508+
simd::LaneCount<LANES>: simd::SupportedLaneCount,
3509+
{
3510+
// These are expected to always match, as vector types are laid out like
3511+
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
3512+
// might as well double-check since it'll optimize away anyhow.
3513+
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());
3514+
3515+
// SAFETY: The simd types have the same layout as arrays, just with
3516+
// potentially-higher alignment, so the de-facto transmutes are sound.
3517+
unsafe { self.align_to_mut() }
3518+
}
3519+
34373520
/// Checks if the elements of this slice are sorted.
34383521
///
34393522
/// That is, for each element `a` and its following element `b`, `a <= b` must hold. If the

0 commit comments

Comments
 (0)