Skip to content

Commit

Permalink
Rollup merge of #91479 - scottmcm:slice-as-simd, r=workingjubilee
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
matthiaskrgr authored Dec 15, 2021
2 parents d258e92 + e4c44c5 commit efc49c1
Showing 1 changed file with 119 additions and 0 deletions.
119 changes: 119 additions & 0 deletions library/core/src/slice/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ use crate::option::Option::{None, Some};
use crate::ptr;
use crate::result::Result;
use crate::result::Result::{Err, Ok};
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
use crate::simd::{self, Simd};
use crate::slice;

#[unstable(
Expand Down Expand Up @@ -3512,6 +3514,123 @@ impl<T> [T] {
}
}

/// Split a slice into a prefix, a middle of aligned SIMD types, and a suffix.
///
/// This is a safe wrapper around [`slice::align_to`], so has the same weak
/// postconditions as that method. You're only assured that
/// `self.len() == prefix.len() + middle.len() * LANES + suffix.len()`.
///
/// Notably, all of the following are possible:
/// - `prefix.len() >= LANES`.
/// - `middle.is_empty()` despite `self.len() >= 3 * LANES`.
/// - `suffix.len() >= LANES`.
///
/// That said, this is a safe method, so if you're only writing safe code,
/// then this can at most cause incorrect logic, not unsoundness.
///
/// # Panics
///
/// This will panic if the size of the SIMD type is different from
/// `LANES` times that of the scalar.
///
/// At the time of writing, the trait restrictions on `Simd<T, LANES>` keeps
/// that from ever happening, as only power-of-two numbers of lanes are
/// supported. It's possible that, in the future, those restrictions might
/// be lifted in a way that would make it possible to see panics from this
/// method for something like `LANES == 3`.
///
/// # Examples
///
/// ```
/// #![feature(portable_simd)]
///
/// let short = &[1, 2, 3];
/// let (prefix, middle, suffix) = short.as_simd::<4>();
/// assert_eq!(middle, []); // Not enough elements for anything in the middle
///
/// // They might be split in any possible way between prefix and suffix
/// let it = prefix.iter().chain(suffix).copied();
/// assert_eq!(it.collect::<Vec<_>>(), vec![1, 2, 3]);
///
/// fn basic_simd_sum(x: &[f32]) -> f32 {
/// use std::ops::Add;
/// use std::simd::f32x4;
/// let (prefix, middle, suffix) = x.as_simd();
/// let sums = f32x4::from_array([
/// prefix.iter().copied().sum(),
/// 0.0,
/// 0.0,
/// suffix.iter().copied().sum(),
/// ]);
/// let sums = middle.iter().copied().fold(sums, f32x4::add);
/// sums.horizontal_sum()
/// }
///
/// let numbers: Vec<f32> = (1..101).map(|x| x as _).collect();
/// assert_eq!(basic_simd_sum(&numbers[1..99]), 4949.0);
/// ```
#[unstable(feature = "portable_simd", issue = "86656")]
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
pub fn as_simd<const LANES: usize>(&self) -> (&[T], &[Simd<T, LANES>], &[T])
where
Simd<T, LANES>: AsRef<[T; LANES]>,
T: simd::SimdElement,
simd::LaneCount<LANES>: simd::SupportedLaneCount,
{
// These are expected to always match, as vector types are laid out like
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
// might as well double-check since it'll optimize away anyhow.
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());

// SAFETY: The simd types have the same layout as arrays, just with
// potentially-higher alignment, so the de-facto transmutes are sound.
unsafe { self.align_to() }
}

/// Split a slice into a prefix, a middle of aligned SIMD types, and a suffix.
///
/// This is a safe wrapper around [`slice::align_to_mut`], so has the same weak
/// postconditions as that method. You're only assured that
/// `self.len() == prefix.len() + middle.len() * LANES + suffix.len()`.
///
/// Notably, all of the following are possible:
/// - `prefix.len() >= LANES`.
/// - `middle.is_empty()` despite `self.len() >= 3 * LANES`.
/// - `suffix.len() >= LANES`.
///
/// That said, this is a safe method, so if you're only writing safe code,
/// then this can at most cause incorrect logic, not unsoundness.
///
/// This is the mutable version of [`slice::as_simd`]; see that for examples.
///
/// # Panics
///
/// This will panic if the size of the SIMD type is different from
/// `LANES` times that of the scalar.
///
/// At the time of writing, the trait restrictions on `Simd<T, LANES>` keeps
/// that from ever happening, as only power-of-two numbers of lanes are
/// supported. It's possible that, in the future, those restrictions might
/// be lifted in a way that would make it possible to see panics from this
/// method for something like `LANES == 3`.
#[unstable(feature = "portable_simd", issue = "86656")]
#[cfg(not(miri))] // Miri does not support all SIMD intrinsics
pub fn as_simd_mut<const LANES: usize>(&mut self) -> (&mut [T], &mut [Simd<T, LANES>], &mut [T])
where
Simd<T, LANES>: AsMut<[T; LANES]>,
T: simd::SimdElement,
simd::LaneCount<LANES>: simd::SupportedLaneCount,
{
// These are expected to always match, as vector types are laid out like
// arrays per <https://llvm.org/docs/LangRef.html#vector-type>, but we
// might as well double-check since it'll optimize away anyhow.
assert_eq!(mem::size_of::<Simd<T, LANES>>(), mem::size_of::<[T; LANES]>());

// SAFETY: The simd types have the same layout as arrays, just with
// potentially-higher alignment, so the de-facto transmutes are sound.
unsafe { self.align_to_mut() }
}

/// Checks if the elements of this slice are sorted.
///
/// That is, for each element `a` and its following element `b`, `a <= b` must hold. If the
Expand Down

0 comments on commit efc49c1

Please sign in to comment.