diff options
Diffstat (limited to '')
4 files changed, 1257 insertions, 0 deletions
diff --git a/library/portable-simd/crates/core_simd/src/masks.rs b/library/portable-simd/crates/core_simd/src/masks.rs new file mode 100644 index 000000000..c36c336d8 --- /dev/null +++ b/library/portable-simd/crates/core_simd/src/masks.rs @@ -0,0 +1,595 @@ +//! Types and traits associated with masking lanes of vectors. +//! Types representing +#![allow(non_camel_case_types)] + +#[cfg_attr( + not(all(target_arch = "x86_64", target_feature = "avx512f")), + path = "masks/full_masks.rs" +)] +#[cfg_attr( + all(target_arch = "x86_64", target_feature = "avx512f"), + path = "masks/bitmask.rs" +)] +mod mask_impl; + +mod to_bitmask; +pub use to_bitmask::ToBitMask; + +#[cfg(feature = "generic_const_exprs")] +pub use to_bitmask::{bitmask_len, ToBitMaskArray}; + +use crate::simd::{intrinsics, LaneCount, Simd, SimdElement, SimdPartialEq, SupportedLaneCount}; +use core::cmp::Ordering; +use core::{fmt, mem}; + +mod sealed { + use super::*; + + /// Not only does this seal the `MaskElement` trait, but these functions prevent other traits + /// from bleeding into the parent bounds. + /// + /// For example, `eq` could be provided by requiring `MaskElement: PartialEq`, but that would + /// prevent us from ever removing that bound, or from implementing `MaskElement` on + /// non-`PartialEq` types in the future. + pub trait Sealed { + fn valid<const LANES: usize>(values: Simd<Self, LANES>) -> bool + where + LaneCount<LANES>: SupportedLaneCount, + Self: SimdElement; + + fn eq(self, other: Self) -> bool; + + const TRUE: Self; + + const FALSE: Self; + } +} +use sealed::Sealed; + +/// Marker trait for types that may be used as SIMD mask elements. +/// +/// # Safety +/// Type must be a signed integer. +pub unsafe trait MaskElement: SimdElement + Sealed {} + +macro_rules! impl_element { + { $ty:ty } => { + impl Sealed for $ty { + fn valid<const LANES: usize>(value: Simd<Self, LANES>) -> bool + where + LaneCount<LANES>: SupportedLaneCount, + { + (value.simd_eq(Simd::splat(0 as _)) | value.simd_eq(Simd::splat(-1 as _))).all() + } + + fn eq(self, other: Self) -> bool { self == other } + + const TRUE: Self = -1; + const FALSE: Self = 0; + } + + // Safety: this is a valid mask element type + unsafe impl MaskElement for $ty {} + } +} + +impl_element! { i8 } +impl_element! { i16 } +impl_element! { i32 } +impl_element! { i64 } +impl_element! { isize } + +/// A SIMD vector mask for `LANES` elements of width specified by `Element`. +/// +/// Masks represent boolean inclusion/exclusion on a per-lane basis. +/// +/// The layout of this type is unspecified. +#[repr(transparent)] +pub struct Mask<T, const LANES: usize>(mask_impl::Mask<T, LANES>) +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount; + +impl<T, const LANES: usize> Copy for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +impl<T, const LANES: usize> Clone for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn clone(&self) -> Self { + *self + } +} + +impl<T, const LANES: usize> Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + /// Construct a mask by setting all lanes to the given value. + pub fn splat(value: bool) -> Self { + Self(mask_impl::Mask::splat(value)) + } + + /// Converts an array of bools to a SIMD mask. + pub fn from_array(array: [bool; LANES]) -> Self { + // SAFETY: Rust's bool has a layout of 1 byte (u8) with a value of + // true: 0b_0000_0001 + // false: 0b_0000_0000 + // Thus, an array of bools is also a valid array of bytes: [u8; N] + // This would be hypothetically valid as an "in-place" transmute, + // but these are "dependently-sized" types, so copy elision it is! + unsafe { + let bytes: [u8; LANES] = mem::transmute_copy(&array); + let bools: Simd<i8, LANES> = + intrinsics::simd_ne(Simd::from_array(bytes), Simd::splat(0u8)); + Mask::from_int_unchecked(intrinsics::simd_cast(bools)) + } + } + + /// Converts a SIMD mask to an array of bools. + pub fn to_array(self) -> [bool; LANES] { + // This follows mostly the same logic as from_array. + // SAFETY: Rust's bool has a layout of 1 byte (u8) with a value of + // true: 0b_0000_0001 + // false: 0b_0000_0000 + // Thus, an array of bools is also a valid array of bytes: [u8; N] + // Since our masks are equal to integers where all bits are set, + // we can simply convert them to i8s, and then bitand them by the + // bitpattern for Rust's "true" bool. + // This would be hypothetically valid as an "in-place" transmute, + // but these are "dependently-sized" types, so copy elision it is! + unsafe { + let mut bytes: Simd<i8, LANES> = intrinsics::simd_cast(self.to_int()); + bytes &= Simd::splat(1i8); + mem::transmute_copy(&bytes) + } + } + + /// Converts a vector of integers to a mask, where 0 represents `false` and -1 + /// represents `true`. + /// + /// # Safety + /// All lanes must be either 0 or -1. + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub unsafe fn from_int_unchecked(value: Simd<T, LANES>) -> Self { + // Safety: the caller must confirm this invariant + unsafe { Self(mask_impl::Mask::from_int_unchecked(value)) } + } + + /// Converts a vector of integers to a mask, where 0 represents `false` and -1 + /// represents `true`. + /// + /// # Panics + /// Panics if any lane is not 0 or -1. + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn from_int(value: Simd<T, LANES>) -> Self { + assert!(T::valid(value), "all values must be either 0 or -1",); + // Safety: the validity has been checked + unsafe { Self::from_int_unchecked(value) } + } + + /// Converts the mask to a vector of integers, where 0 represents `false` and -1 + /// represents `true`. + #[inline] + #[must_use = "method returns a new vector and does not mutate the original value"] + pub fn to_int(self) -> Simd<T, LANES> { + self.0.to_int() + } + + /// Converts the mask to a mask of any other lane size. + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn cast<U: MaskElement>(self) -> Mask<U, LANES> { + Mask(self.0.convert()) + } + + /// Tests the value of the specified lane. + /// + /// # Safety + /// `lane` must be less than `LANES`. + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub unsafe fn test_unchecked(&self, lane: usize) -> bool { + // Safety: the caller must confirm this invariant + unsafe { self.0.test_unchecked(lane) } + } + + /// Tests the value of the specified lane. + /// + /// # Panics + /// Panics if `lane` is greater than or equal to the number of lanes in the vector. + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn test(&self, lane: usize) -> bool { + assert!(lane < LANES, "lane index out of range"); + // Safety: the lane index has been checked + unsafe { self.test_unchecked(lane) } + } + + /// Sets the value of the specified lane. + /// + /// # Safety + /// `lane` must be less than `LANES`. + #[inline] + pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) { + // Safety: the caller must confirm this invariant + unsafe { + self.0.set_unchecked(lane, value); + } + } + + /// Sets the value of the specified lane. + /// + /// # Panics + /// Panics if `lane` is greater than or equal to the number of lanes in the vector. + #[inline] + pub fn set(&mut self, lane: usize, value: bool) { + assert!(lane < LANES, "lane index out of range"); + // Safety: the lane index has been checked + unsafe { + self.set_unchecked(lane, value); + } + } + + /// Returns true if any lane is set, or false otherwise. + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn any(self) -> bool { + self.0.any() + } + + /// Returns true if all lanes are set, or false otherwise. + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn all(self) -> bool { + self.0.all() + } +} + +// vector/array conversion +impl<T, const LANES: usize> From<[bool; LANES]> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn from(array: [bool; LANES]) -> Self { + Self::from_array(array) + } +} + +impl<T, const LANES: usize> From<Mask<T, LANES>> for [bool; LANES] +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn from(vector: Mask<T, LANES>) -> Self { + vector.to_array() + } +} + +impl<T, const LANES: usize> Default for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a defaulted mask with all lanes set to false (0)"] + fn default() -> Self { + Self::splat(false) + } +} + +impl<T, const LANES: usize> PartialEq for Mask<T, LANES> +where + T: MaskElement + PartialEq, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + fn eq(&self, other: &Self) -> bool { + self.0 == other.0 + } +} + +impl<T, const LANES: usize> PartialOrd for Mask<T, LANES> +where + T: MaskElement + PartialOrd, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a new Ordering and does not mutate the original value"] + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + self.0.partial_cmp(&other.0) + } +} + +impl<T, const LANES: usize> fmt::Debug for Mask<T, LANES> +where + T: MaskElement + fmt::Debug, + LaneCount<LANES>: SupportedLaneCount, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries((0..LANES).map(|lane| self.test(lane))) + .finish() + } +} + +impl<T, const LANES: usize> core::ops::BitAnd for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitand(self, rhs: Self) -> Self { + Self(self.0 & rhs.0) + } +} + +impl<T, const LANES: usize> core::ops::BitAnd<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitand(self, rhs: bool) -> Self { + self & Self::splat(rhs) + } +} + +impl<T, const LANES: usize> core::ops::BitAnd<Mask<T, LANES>> for bool +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Mask<T, LANES>; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitand(self, rhs: Mask<T, LANES>) -> Mask<T, LANES> { + Mask::splat(self) & rhs + } +} + +impl<T, const LANES: usize> core::ops::BitOr for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitor(self, rhs: Self) -> Self { + Self(self.0 | rhs.0) + } +} + +impl<T, const LANES: usize> core::ops::BitOr<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitor(self, rhs: bool) -> Self { + self | Self::splat(rhs) + } +} + +impl<T, const LANES: usize> core::ops::BitOr<Mask<T, LANES>> for bool +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Mask<T, LANES>; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitor(self, rhs: Mask<T, LANES>) -> Mask<T, LANES> { + Mask::splat(self) | rhs + } +} + +impl<T, const LANES: usize> core::ops::BitXor for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitxor(self, rhs: Self) -> Self::Output { + Self(self.0 ^ rhs.0) + } +} + +impl<T, const LANES: usize> core::ops::BitXor<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitxor(self, rhs: bool) -> Self::Output { + self ^ Self::splat(rhs) + } +} + +impl<T, const LANES: usize> core::ops::BitXor<Mask<T, LANES>> for bool +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Mask<T, LANES>; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitxor(self, rhs: Mask<T, LANES>) -> Self::Output { + Mask::splat(self) ^ rhs + } +} + +impl<T, const LANES: usize> core::ops::Not for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Mask<T, LANES>; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn not(self) -> Self::Output { + Self(!self.0) + } +} + +impl<T, const LANES: usize> core::ops::BitAndAssign for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitand_assign(&mut self, rhs: Self) { + self.0 = self.0 & rhs.0; + } +} + +impl<T, const LANES: usize> core::ops::BitAndAssign<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitand_assign(&mut self, rhs: bool) { + *self &= Self::splat(rhs); + } +} + +impl<T, const LANES: usize> core::ops::BitOrAssign for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitor_assign(&mut self, rhs: Self) { + self.0 = self.0 | rhs.0; + } +} + +impl<T, const LANES: usize> core::ops::BitOrAssign<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitor_assign(&mut self, rhs: bool) { + *self |= Self::splat(rhs); + } +} + +impl<T, const LANES: usize> core::ops::BitXorAssign for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitxor_assign(&mut self, rhs: Self) { + self.0 = self.0 ^ rhs.0; + } +} + +impl<T, const LANES: usize> core::ops::BitXorAssign<bool> for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + fn bitxor_assign(&mut self, rhs: bool) { + *self ^= Self::splat(rhs); + } +} + +/// A mask for SIMD vectors with eight elements of 8 bits. +pub type mask8x8 = Mask<i8, 8>; + +/// A mask for SIMD vectors with 16 elements of 8 bits. +pub type mask8x16 = Mask<i8, 16>; + +/// A mask for SIMD vectors with 32 elements of 8 bits. +pub type mask8x32 = Mask<i8, 32>; + +/// A mask for SIMD vectors with 64 elements of 8 bits. +pub type mask8x64 = Mask<i8, 64>; + +/// A mask for SIMD vectors with four elements of 16 bits. +pub type mask16x4 = Mask<i16, 4>; + +/// A mask for SIMD vectors with eight elements of 16 bits. +pub type mask16x8 = Mask<i16, 8>; + +/// A mask for SIMD vectors with 16 elements of 16 bits. +pub type mask16x16 = Mask<i16, 16>; + +/// A mask for SIMD vectors with 32 elements of 16 bits. +pub type mask16x32 = Mask<i16, 32>; + +/// A mask for SIMD vectors with two elements of 32 bits. +pub type mask32x2 = Mask<i32, 2>; + +/// A mask for SIMD vectors with four elements of 32 bits. +pub type mask32x4 = Mask<i32, 4>; + +/// A mask for SIMD vectors with eight elements of 32 bits. +pub type mask32x8 = Mask<i32, 8>; + +/// A mask for SIMD vectors with 16 elements of 32 bits. +pub type mask32x16 = Mask<i32, 16>; + +/// A mask for SIMD vectors with two elements of 64 bits. +pub type mask64x2 = Mask<i64, 2>; + +/// A mask for SIMD vectors with four elements of 64 bits. +pub type mask64x4 = Mask<i64, 4>; + +/// A mask for SIMD vectors with eight elements of 64 bits. +pub type mask64x8 = Mask<i64, 8>; + +/// A mask for SIMD vectors with two elements of pointer width. +pub type masksizex2 = Mask<isize, 2>; + +/// A mask for SIMD vectors with four elements of pointer width. +pub type masksizex4 = Mask<isize, 4>; + +/// A mask for SIMD vectors with eight elements of pointer width. +pub type masksizex8 = Mask<isize, 8>; + +macro_rules! impl_from { + { $from:ty => $($to:ty),* } => { + $( + impl<const LANES: usize> From<Mask<$from, LANES>> for Mask<$to, LANES> + where + LaneCount<LANES>: SupportedLaneCount, + { + fn from(value: Mask<$from, LANES>) -> Self { + value.cast() + } + } + )* + } +} +impl_from! { i8 => i16, i32, i64, isize } +impl_from! { i16 => i32, i64, isize, i8 } +impl_from! { i32 => i64, isize, i8, i16 } +impl_from! { i64 => isize, i8, i16, i32 } +impl_from! { isize => i8, i16, i32, i64 } diff --git a/library/portable-simd/crates/core_simd/src/masks/bitmask.rs b/library/portable-simd/crates/core_simd/src/masks/bitmask.rs new file mode 100644 index 000000000..365ecc0a3 --- /dev/null +++ b/library/portable-simd/crates/core_simd/src/masks/bitmask.rs @@ -0,0 +1,246 @@ +#![allow(unused_imports)] +use super::MaskElement; +use crate::simd::intrinsics; +use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask}; +use core::marker::PhantomData; + +/// A mask where each lane is represented by a single bit. +#[repr(transparent)] +pub struct Mask<T, const LANES: usize>( + <LaneCount<LANES> as SupportedLaneCount>::BitMask, + PhantomData<T>, +) +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount; + +impl<T, const LANES: usize> Copy for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +impl<T, const LANES: usize> Clone for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn clone(&self) -> Self { + *self + } +} + +impl<T, const LANES: usize> PartialEq for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn eq(&self, other: &Self) -> bool { + self.0.as_ref() == other.0.as_ref() + } +} + +impl<T, const LANES: usize> PartialOrd for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> { + self.0.as_ref().partial_cmp(other.0.as_ref()) + } +} + +impl<T, const LANES: usize> Eq for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +impl<T, const LANES: usize> Ord for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + self.0.as_ref().cmp(other.0.as_ref()) + } +} + +impl<T, const LANES: usize> Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn splat(value: bool) -> Self { + let mut mask = <LaneCount<LANES> as SupportedLaneCount>::BitMask::default(); + if value { + mask.as_mut().fill(u8::MAX) + } else { + mask.as_mut().fill(u8::MIN) + } + if LANES % 8 > 0 { + *mask.as_mut().last_mut().unwrap() &= u8::MAX >> (8 - LANES % 8); + } + Self(mask, PhantomData) + } + + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub unsafe fn test_unchecked(&self, lane: usize) -> bool { + (self.0.as_ref()[lane / 8] >> (lane % 8)) & 0x1 > 0 + } + + #[inline] + pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) { + unsafe { + self.0.as_mut()[lane / 8] ^= ((value ^ self.test_unchecked(lane)) as u8) << (lane % 8) + } + } + + #[inline] + #[must_use = "method returns a new vector and does not mutate the original value"] + pub fn to_int(self) -> Simd<T, LANES> { + unsafe { + intrinsics::simd_select_bitmask(self.0, Simd::splat(T::TRUE), Simd::splat(T::FALSE)) + } + } + + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub unsafe fn from_int_unchecked(value: Simd<T, LANES>) -> Self { + unsafe { Self(intrinsics::simd_bitmask(value), PhantomData) } + } + + #[cfg(feature = "generic_const_exprs")] + #[inline] + #[must_use = "method returns a new array and does not mutate the original value"] + pub fn to_bitmask_array<const N: usize>(self) -> [u8; N] { + assert!(core::mem::size_of::<Self>() == N); + + // Safety: converting an integer to an array of bytes of the same size is safe + unsafe { core::mem::transmute_copy(&self.0) } + } + + #[cfg(feature = "generic_const_exprs")] + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn from_bitmask_array<const N: usize>(bitmask: [u8; N]) -> Self { + assert!(core::mem::size_of::<Self>() == N); + + // Safety: converting an array of bytes to an integer of the same size is safe + Self(unsafe { core::mem::transmute_copy(&bitmask) }, PhantomData) + } + + #[inline] + pub fn to_bitmask_integer<U>(self) -> U + where + super::Mask<T, LANES>: ToBitMask<BitMask = U>, + { + // Safety: these are the same types + unsafe { core::mem::transmute_copy(&self.0) } + } + + #[inline] + pub fn from_bitmask_integer<U>(bitmask: U) -> Self + where + super::Mask<T, LANES>: ToBitMask<BitMask = U>, + { + // Safety: these are the same types + unsafe { Self(core::mem::transmute_copy(&bitmask), PhantomData) } + } + + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn convert<U>(self) -> Mask<U, LANES> + where + U: MaskElement, + { + // Safety: bitmask layout does not depend on the element width + unsafe { core::mem::transmute_copy(&self) } + } + + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn any(self) -> bool { + self != Self::splat(false) + } + + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn all(self) -> bool { + self == Self::splat(true) + } +} + +impl<T, const LANES: usize> core::ops::BitAnd for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, + <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitand(mut self, rhs: Self) -> Self { + for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) { + *l &= r; + } + self + } +} + +impl<T, const LANES: usize> core::ops::BitOr for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, + <LaneCount<LANES> as SupportedLaneCount>::BitMask: AsRef<[u8]> + AsMut<[u8]>, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitor(mut self, rhs: Self) -> Self { + for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) { + *l |= r; + } + self + } +} + +impl<T, const LANES: usize> core::ops::BitXor for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitxor(mut self, rhs: Self) -> Self::Output { + for (l, r) in self.0.as_mut().iter_mut().zip(rhs.0.as_ref().iter()) { + *l ^= r; + } + self + } +} + +impl<T, const LANES: usize> core::ops::Not for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn not(mut self) -> Self::Output { + for x in self.0.as_mut() { + *x = !*x; + } + if LANES % 8 > 0 { + *self.0.as_mut().last_mut().unwrap() &= u8::MAX >> (8 - LANES % 8); + } + self + } +} diff --git a/library/portable-simd/crates/core_simd/src/masks/full_masks.rs b/library/portable-simd/crates/core_simd/src/masks/full_masks.rs new file mode 100644 index 000000000..adf0fcbea --- /dev/null +++ b/library/portable-simd/crates/core_simd/src/masks/full_masks.rs @@ -0,0 +1,323 @@ +//! Masks that take up full SIMD vector registers. + +use super::MaskElement; +use crate::simd::intrinsics; +use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask}; + +#[cfg(feature = "generic_const_exprs")] +use crate::simd::ToBitMaskArray; + +#[repr(transparent)] +pub struct Mask<T, const LANES: usize>(Simd<T, LANES>) +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount; + +impl<T, const LANES: usize> Copy for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +impl<T, const LANES: usize> Clone for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn clone(&self) -> Self { + *self + } +} + +impl<T, const LANES: usize> PartialEq for Mask<T, LANES> +where + T: MaskElement + PartialEq, + LaneCount<LANES>: SupportedLaneCount, +{ + fn eq(&self, other: &Self) -> bool { + self.0.eq(&other.0) + } +} + +impl<T, const LANES: usize> PartialOrd for Mask<T, LANES> +where + T: MaskElement + PartialOrd, + LaneCount<LANES>: SupportedLaneCount, +{ + fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> { + self.0.partial_cmp(&other.0) + } +} + +impl<T, const LANES: usize> Eq for Mask<T, LANES> +where + T: MaskElement + Eq, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +impl<T, const LANES: usize> Ord for Mask<T, LANES> +where + T: MaskElement + Ord, + LaneCount<LANES>: SupportedLaneCount, +{ + fn cmp(&self, other: &Self) -> core::cmp::Ordering { + self.0.cmp(&other.0) + } +} + +// Used for bitmask bit order workaround +pub(crate) trait ReverseBits { + // Reverse the least significant `n` bits of `self`. + // (Remaining bits must be 0.) + fn reverse_bits(self, n: usize) -> Self; +} + +macro_rules! impl_reverse_bits { + { $($int:ty),* } => { + $( + impl ReverseBits for $int { + #[inline(always)] + fn reverse_bits(self, n: usize) -> Self { + let rev = <$int>::reverse_bits(self); + let bitsize = core::mem::size_of::<$int>() * 8; + if n < bitsize { + // Shift things back to the right + rev >> (bitsize - n) + } else { + rev + } + } + } + )* + } +} + +impl_reverse_bits! { u8, u16, u32, u64 } + +impl<T, const LANES: usize> Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn splat(value: bool) -> Self { + Self(Simd::splat(if value { T::TRUE } else { T::FALSE })) + } + + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub unsafe fn test_unchecked(&self, lane: usize) -> bool { + T::eq(self.0[lane], T::TRUE) + } + + #[inline] + pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) { + self.0[lane] = if value { T::TRUE } else { T::FALSE } + } + + #[inline] + #[must_use = "method returns a new vector and does not mutate the original value"] + pub fn to_int(self) -> Simd<T, LANES> { + self.0 + } + + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub unsafe fn from_int_unchecked(value: Simd<T, LANES>) -> Self { + Self(value) + } + + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn convert<U>(self) -> Mask<U, LANES> + where + U: MaskElement, + { + // Safety: masks are simply integer vectors of 0 and -1, and we can cast the element type. + unsafe { Mask(intrinsics::simd_cast(self.0)) } + } + + #[cfg(feature = "generic_const_exprs")] + #[inline] + #[must_use = "method returns a new array and does not mutate the original value"] + pub fn to_bitmask_array<const N: usize>(self) -> [u8; N] + where + super::Mask<T, LANES>: ToBitMaskArray, + [(); <super::Mask<T, LANES> as ToBitMaskArray>::BYTES]: Sized, + { + assert_eq!(<super::Mask<T, LANES> as ToBitMaskArray>::BYTES, N); + + // Safety: N is the correct bitmask size + unsafe { + // Compute the bitmask + let bitmask: [u8; <super::Mask<T, LANES> as ToBitMaskArray>::BYTES] = + intrinsics::simd_bitmask(self.0); + + // Transmute to the return type, previously asserted to be the same size + let mut bitmask: [u8; N] = core::mem::transmute_copy(&bitmask); + + // LLVM assumes bit order should match endianness + if cfg!(target_endian = "big") { + for x in bitmask.as_mut() { + *x = x.reverse_bits(); + } + }; + + bitmask + } + } + + #[cfg(feature = "generic_const_exprs")] + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + pub fn from_bitmask_array<const N: usize>(mut bitmask: [u8; N]) -> Self + where + super::Mask<T, LANES>: ToBitMaskArray, + [(); <super::Mask<T, LANES> as ToBitMaskArray>::BYTES]: Sized, + { + assert_eq!(<super::Mask<T, LANES> as ToBitMaskArray>::BYTES, N); + + // Safety: N is the correct bitmask size + unsafe { + // LLVM assumes bit order should match endianness + if cfg!(target_endian = "big") { + for x in bitmask.as_mut() { + *x = x.reverse_bits(); + } + } + + // Transmute to the bitmask type, previously asserted to be the same size + let bitmask: [u8; <super::Mask<T, LANES> as ToBitMaskArray>::BYTES] = + core::mem::transmute_copy(&bitmask); + + // Compute the regular mask + Self::from_int_unchecked(intrinsics::simd_select_bitmask( + bitmask, + Self::splat(true).to_int(), + Self::splat(false).to_int(), + )) + } + } + + #[inline] + pub(crate) fn to_bitmask_integer<U: ReverseBits>(self) -> U + where + super::Mask<T, LANES>: ToBitMask<BitMask = U>, + { + // Safety: U is required to be the appropriate bitmask type + let bitmask: U = unsafe { intrinsics::simd_bitmask(self.0) }; + + // LLVM assumes bit order should match endianness + if cfg!(target_endian = "big") { + bitmask.reverse_bits(LANES) + } else { + bitmask + } + } + + #[inline] + pub(crate) fn from_bitmask_integer<U: ReverseBits>(bitmask: U) -> Self + where + super::Mask<T, LANES>: ToBitMask<BitMask = U>, + { + // LLVM assumes bit order should match endianness + let bitmask = if cfg!(target_endian = "big") { + bitmask.reverse_bits(LANES) + } else { + bitmask + }; + + // Safety: U is required to be the appropriate bitmask type + unsafe { + Self::from_int_unchecked(intrinsics::simd_select_bitmask( + bitmask, + Self::splat(true).to_int(), + Self::splat(false).to_int(), + )) + } + } + + #[inline] + #[must_use = "method returns a new bool and does not mutate the original value"] + pub fn any(self) -> bool { + // Safety: use `self` as an integer vector + unsafe { intrinsics::simd_reduce_any(self.to_int()) } + } + + #[inline] + #[must_use = "method returns a new vector and does not mutate the original value"] + pub fn all(self) -> bool { + // Safety: use `self` as an integer vector + unsafe { intrinsics::simd_reduce_all(self.to_int()) } + } +} + +impl<T, const LANES: usize> core::convert::From<Mask<T, LANES>> for Simd<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + fn from(value: Mask<T, LANES>) -> Self { + value.0 + } +} + +impl<T, const LANES: usize> core::ops::BitAnd for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitand(self, rhs: Self) -> Self { + // Safety: `self` is an integer vector + unsafe { Self(intrinsics::simd_and(self.0, rhs.0)) } + } +} + +impl<T, const LANES: usize> core::ops::BitOr for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitor(self, rhs: Self) -> Self { + // Safety: `self` is an integer vector + unsafe { Self(intrinsics::simd_or(self.0, rhs.0)) } + } +} + +impl<T, const LANES: usize> core::ops::BitXor for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn bitxor(self, rhs: Self) -> Self { + // Safety: `self` is an integer vector + unsafe { Self(intrinsics::simd_xor(self.0, rhs.0)) } + } +} + +impl<T, const LANES: usize> core::ops::Not for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ + type Output = Self; + #[inline] + #[must_use = "method returns a new mask and does not mutate the original value"] + fn not(self) -> Self::Output { + Self::splat(true) ^ self + } +} diff --git a/library/portable-simd/crates/core_simd/src/masks/to_bitmask.rs b/library/portable-simd/crates/core_simd/src/masks/to_bitmask.rs new file mode 100644 index 000000000..65d3ce9be --- /dev/null +++ b/library/portable-simd/crates/core_simd/src/masks/to_bitmask.rs @@ -0,0 +1,93 @@ +use super::{mask_impl, Mask, MaskElement}; +use crate::simd::{LaneCount, SupportedLaneCount}; + +mod sealed { + pub trait Sealed {} +} +pub use sealed::Sealed; + +impl<T, const LANES: usize> Sealed for Mask<T, LANES> +where + T: MaskElement, + LaneCount<LANES>: SupportedLaneCount, +{ +} + +/// Converts masks to and from integer bitmasks. +/// +/// Each bit of the bitmask corresponds to a mask lane, starting with the LSB. +pub trait ToBitMask: Sealed { + /// The integer bitmask type. + type BitMask; + + /// Converts a mask to a bitmask. + fn to_bitmask(self) -> Self::BitMask; + + /// Converts a bitmask to a mask. + fn from_bitmask(bitmask: Self::BitMask) -> Self; +} + +/// Converts masks to and from byte array bitmasks. +/// +/// Each bit of the bitmask corresponds to a mask lane, starting with the LSB of the first byte. +#[cfg(feature = "generic_const_exprs")] +pub trait ToBitMaskArray: Sealed { + /// The length of the bitmask array. + const BYTES: usize; + + /// Converts a mask to a bitmask. + fn to_bitmask_array(self) -> [u8; Self::BYTES]; + + /// Converts a bitmask to a mask. + fn from_bitmask_array(bitmask: [u8; Self::BYTES]) -> Self; +} + +macro_rules! impl_integer_intrinsic { + { $(impl ToBitMask<BitMask=$int:ty> for Mask<_, $lanes:literal>)* } => { + $( + impl<T: MaskElement> ToBitMask for Mask<T, $lanes> { + type BitMask = $int; + + fn to_bitmask(self) -> $int { + self.0.to_bitmask_integer() + } + + fn from_bitmask(bitmask: $int) -> Self { + Self(mask_impl::Mask::from_bitmask_integer(bitmask)) + } + } + )* + } +} + +impl_integer_intrinsic! { + impl ToBitMask<BitMask=u8> for Mask<_, 1> + impl ToBitMask<BitMask=u8> for Mask<_, 2> + impl ToBitMask<BitMask=u8> for Mask<_, 4> + impl ToBitMask<BitMask=u8> for Mask<_, 8> + impl ToBitMask<BitMask=u16> for Mask<_, 16> + impl ToBitMask<BitMask=u32> for Mask<_, 32> + impl ToBitMask<BitMask=u64> for Mask<_, 64> +} + +/// Returns the minimum numnber of bytes in a bitmask with `lanes` lanes. +#[cfg(feature = "generic_const_exprs")] +pub const fn bitmask_len(lanes: usize) -> usize { + (lanes + 7) / 8 +} + +#[cfg(feature = "generic_const_exprs")] +impl<T: MaskElement, const LANES: usize> ToBitMaskArray for Mask<T, LANES> +where + LaneCount<LANES>: SupportedLaneCount, +{ + const BYTES: usize = bitmask_len(LANES); + + fn to_bitmask_array(self) -> [u8; Self::BYTES] { + self.0.to_bitmask_array() + } + + fn from_bitmask_array(bitmask: [u8; Self::BYTES]) -> Self { + Mask(mask_impl::Mask::from_bitmask_array(bitmask)) + } +} |