Skip to content

Commit 872e341

Browse files
authored
Rollup merge of rust-lang#91479 - scottmcm:slice-as-simd, r=workingjubilee
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]); } ``` They're `cfg`'d out for miri because the `simd` module as a whole is unavailable there.
2 parents d594910 + e4c44c5 commit 872e341

File tree

1 file changed

+119
-0
lines changed

1 file changed

+119
-0
lines changed

library/core/src/slice/mod.rs

+119
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(
@@ -3512,6 +3514,123 @@ impl<T> [T] {
35123514
}
35133515
}
35143516

3517+
/// Split a slice into a prefix, a middle of aligned SIMD types, and a suffix.
3518+
///
3519+
/// This is a safe wrapper around [`slice::align_to`], so has the same weak
3520+
/// postconditions as that method. You're only assured that
3521+
/// `self.len() == prefix.len() + middle.len() * LANES + suffix.len()`.
3522+
///
3523+
/// Notably, all of the following are possible:
3524+
/// - `prefix.len() >= LANES`.
3525+
/// - `middle.is_empty()` despite `self.len() >= 3 * LANES`.
3526+
/// - `suffix.len() >= LANES`.
3527+
///
3528+
/// That said, this is a safe method, so if you're only writing safe code,
3529+
/// then this can at most cause incorrect logic, not unsoundness.
3530+
///
3531+
/// # Panics
3532+
///
3533+
/// This will panic if the size of the SIMD type is different from
3534+
/// `LANES` times that of the scalar.
3535+
///
3536+
/// At the time of writing, the trait restrictions on `Simd<T, LANES>` keeps
3537+
/// that from ever happening, as only power-of-two numbers of lanes are
3538+
/// supported. It's possible that, in the future, those restrictions might
3539+
/// be lifted in a way that would make it possible to see panics from this
3540+
/// method for something like `LANES == 3`.
3541+
///
3542+
/// # Examples
3543+
///
3544+
/// ```
3545+
/// #![feature(portable_simd)]
3546+
///
3547+
/// let short = &[1, 2, 3];
3548+
/// let (prefix, middle, suffix) = short.as_simd::<4>();
3549+
/// assert_eq!(middle, []); // Not enough elements for anything in the middle
3550+
///
3551+
/// // They might be split in any possible way between prefix and suffix
3552+
/// let it = prefix.iter().chain(suffix).copied();
3553+
/// assert_eq!(it.collect::<Vec<_>>(), vec![1, 2, 3]);
3554+
///
3555+
/// fn basic_simd_sum(x: &[f32]) -> f32 {
3556+
/// use std::ops::Add;
3557+
/// use std::simd::f32x4;
3558+
/// let (prefix, middle, suffix) = x.as_simd();
3559+
/// let sums = f32x4::from_array([
3560+
/// prefix.iter().copied().sum(),
3561+
/// 0.0,
3562+
/// 0.0,
3563+
/// suffix.iter().copied().sum(),
3564+
/// ]);
3565+
/// let sums = middle.iter().copied().fold(sums, f32x4::add);
3566+
/// sums.horizontal_sum()
3567+
/// }
3568+
///
3569+
/// let numbers: Vec<f32> = (1..101).map(|x| x as _).collect();
3570+
/// assert_eq!(basic_simd_sum(&numbers[1..99]), 4949.0);
3571+
/// ```
3572+
#[unstable(feature = "portable_simd", issue = "86656")]
3573+
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
3574+
pub fn as_simd<const LANES: usize>(&self) -> (&[T], &[Simd<T, LANES>], &[T])
3575+
where
3576+
Simd<T, LANES>: AsRef<[T; LANES]>,
3577+
T: simd::SimdElement,
3578+
simd::LaneCount<LANES>: simd::SupportedLaneCount,
3579+
{
3580+
// These are expected to always match, as vector types are laid out like
3581+
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
3582+
// might as well double-check since it'll optimize away anyhow.
3583+
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());
3584+
3585+
// SAFETY: The simd types have the same layout as arrays, just with
3586+
// potentially-higher alignment, so the de-facto transmutes are sound.
3587+
unsafe { self.align_to() }
3588+
}
3589+
3590+
/// Split a slice into a prefix, a middle of aligned SIMD types, and a suffix.
3591+
///
3592+
/// This is a safe wrapper around [`slice::align_to_mut`], so has the same weak
3593+
/// postconditions as that method. You're only assured that
3594+
/// `self.len() == prefix.len() + middle.len() * LANES + suffix.len()`.
3595+
///
3596+
/// Notably, all of the following are possible:
3597+
/// - `prefix.len() >= LANES`.
3598+
/// - `middle.is_empty()` despite `self.len() >= 3 * LANES`.
3599+
/// - `suffix.len() >= LANES`.
3600+
///
3601+
/// That said, this is a safe method, so if you're only writing safe code,
3602+
/// then this can at most cause incorrect logic, not unsoundness.
3603+
///
3604+
/// This is the mutable version of [`slice::as_simd`]; see that for examples.
3605+
///
3606+
/// # Panics
3607+
///
3608+
/// This will panic if the size of the SIMD type is different from
3609+
/// `LANES` times that of the scalar.
3610+
///
3611+
/// At the time of writing, the trait restrictions on `Simd<T, LANES>` keeps
3612+
/// that from ever happening, as only power-of-two numbers of lanes are
3613+
/// supported. It's possible that, in the future, those restrictions might
3614+
/// be lifted in a way that would make it possible to see panics from this
3615+
/// method for something like `LANES == 3`.
3616+
#[unstable(feature = "portable_simd", issue = "86656")]
3617+
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
3618+
pub fn as_simd_mut<const LANES: usize>(&mut self) -> (&mut [T], &mut [Simd<T, LANES>], &mut [T])
3619+
where
3620+
Simd<T, LANES>: AsMut<[T; LANES]>,
3621+
T: simd::SimdElement,
3622+
simd::LaneCount<LANES>: simd::SupportedLaneCount,
3623+
{
3624+
// These are expected to always match, as vector types are laid out like
3625+
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
3626+
// might as well double-check since it'll optimize away anyhow.
3627+
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());
3628+
3629+
// SAFETY: The simd types have the same layout as arrays, just with
3630+
// potentially-higher alignment, so the de-facto transmutes are sound.
3631+
unsafe { self.align_to_mut() }
3632+
}
3633+
35153634
/// Checks if the elements of this slice are sorted.
35163635
///
35173636
/// That is, for each element `a` and its following element `b`, `a <= b` must hold. If the

0 commit comments

Comments
 (0)