diff --git a/crates/core_simd/src/masks.rs b/crates/core_simd/src/masks.rs index b1f98d9eb4e..e1cd7930450 100644 --- a/crates/core_simd/src/masks.rs +++ b/crates/core_simd/src/masks.rs @@ -12,8 +12,10 @@ )] mod mask_impl; -use crate::simd::intrinsics; -use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount}; +mod to_bitmask; +pub use to_bitmask::ToBitMask; + +use crate::simd::{intrinsics, LaneCount, Simd, SimdElement, SupportedLaneCount}; use core::cmp::Ordering; use core::{fmt, mem}; @@ -225,22 +227,6 @@ where } } - /// Convert this mask to a bitmask, with one bit set per lane. - #[cfg(feature = "generic_const_exprs")] - #[inline] - #[must_use = "method returns a new array and does not mutate the original value"] - pub fn to_bitmask(self) -> [u8; LaneCount::::BITMASK_LEN] { - self.0.to_bitmask() - } - - /// Convert a bitmask to a mask. - #[cfg(feature = "generic_const_exprs")] - #[inline] - #[must_use = "method returns a new mask and does not mutate the original value"] - pub fn from_bitmask(bitmask: [u8; LaneCount::::BITMASK_LEN]) -> Self { - Self(mask_impl::Mask::from_bitmask(bitmask)) - } - /// 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"] diff --git a/crates/core_simd/src/masks/bitmask.rs b/crates/core_simd/src/masks/bitmask.rs index b7f8b2c236e..ec4dd357ee9 100644 --- a/crates/core_simd/src/masks/bitmask.rs +++ b/crates/core_simd/src/masks/bitmask.rs @@ -1,7 +1,7 @@ #![allow(unused_imports)] use super::MaskElement; use crate::simd::intrinsics; -use crate::simd::{LaneCount, Simd, SupportedLaneCount}; +use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask}; use core::marker::PhantomData; /// A mask where each lane is represented by a single bit. @@ -115,20 +115,22 @@ where 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(self) -> [u8; LaneCount::::BITMASK_LEN] { - // Safety: these are the same type and we are laundering the generic + pub fn to_bitmask_integer(self) -> U + where + super::Mask: ToBitMask, + { + // Safety: these are the same types 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(bitmask: [u8; LaneCount::::BITMASK_LEN]) -> Self { - // Safety: these are the same type and we are laundering the generic - Self(unsafe { core::mem::transmute_copy(&bitmask) }, PhantomData) + pub fn from_bitmask_integer(bitmask: U) -> Self + where + super::Mask: ToBitMask, + { + // Safety: these are the same types + unsafe { Self(core::mem::transmute_copy(&bitmask), PhantomData) } } #[inline] diff --git a/crates/core_simd/src/masks/full_masks.rs b/crates/core_simd/src/masks/full_masks.rs index 02b5593c8f4..8bbdf637de8 100644 --- a/crates/core_simd/src/masks/full_masks.rs +++ b/crates/core_simd/src/masks/full_masks.rs @@ -2,7 +2,7 @@ use super::MaskElement; use crate::simd::intrinsics; -use crate::simd::{LaneCount, Simd, SupportedLaneCount}; +use crate::simd::{LaneCount, Simd, SupportedLaneCount, ToBitMask}; #[repr(transparent)] pub struct Mask(Simd) @@ -66,6 +66,23 @@ where } } +// Used for bitmask bit order workaround +pub(crate) trait ReverseBits { + fn reverse_bits(self) -> Self; +} + +macro_rules! impl_reverse_bits { + { $($int:ty),* } => { + $( + impl ReverseBits for $int { + fn reverse_bits(self) -> Self { <$int>::reverse_bits(self) } + } + )* + } +} + +impl_reverse_bits! { u8, u16, u32, u64 } + impl Mask where T: MaskElement, @@ -110,41 +127,36 @@ where 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(self) -> [u8; LaneCount::::BITMASK_LEN] { - unsafe { - let mut bitmask: [u8; LaneCount::::BITMASK_LEN] = - intrinsics::simd_bitmask(self.0); - - // There is a bug where LLVM appears to implement this operation with the wrong - // bit order. - // TODO fix this in a better way - if cfg!(target_endian = "big") { - for x in bitmask.as_mut() { - *x = x.reverse_bits(); - } - } + pub(crate) fn to_bitmask_integer(self) -> U + where + super::Mask: ToBitMask, + { + // 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() + } else { 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(mut bitmask: [u8; LaneCount::::BITMASK_LEN]) -> Self { - unsafe { - // There is a bug where LLVM appears to implement this operation with the wrong - // bit order. - // TODO fix this in a better way - if cfg!(target_endian = "big") { - for x in bitmask.as_mut() { - *x = x.reverse_bits(); - } - } + pub(crate) fn from_bitmask_integer(bitmask: U) -> Self + where + super::Mask: ToBitMask, + { + // LLVM assumes bit order should match endianness + let bitmask = if cfg!(target_endian = "big") { + bitmask.reverse_bits() + } 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(), diff --git a/crates/core_simd/src/masks/to_bitmask.rs b/crates/core_simd/src/masks/to_bitmask.rs new file mode 100644 index 00000000000..1c2037764c1 --- /dev/null +++ b/crates/core_simd/src/masks/to_bitmask.rs @@ -0,0 +1,57 @@ +use super::{mask_impl, Mask, MaskElement}; +use crate::simd::{LaneCount, SupportedLaneCount}; + +mod sealed { + pub trait Sealed {} +} +pub use sealed::Sealed; + +impl Sealed for Mask +where + T: MaskElement, + LaneCount: SupportedLaneCount, +{ +} + +/// Converts masks to and from integer bitmasks. +/// +/// Each bit of the bitmask corresponds to a mask lane, starting with the LSB. +/// +/// # Safety +/// This trait is `unsafe` and sealed, since the `BitMask` type must match the number of lanes in +/// the mask. +pub unsafe 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; +} + +macro_rules! impl_integer_intrinsic { + { $(unsafe impl ToBitMask for Mask<_, $lanes:literal>)* } => { + $( + unsafe impl ToBitMask for Mask { + 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! { + unsafe impl ToBitMask for Mask<_, 8> + unsafe impl ToBitMask for Mask<_, 16> + unsafe impl ToBitMask for Mask<_, 32> + unsafe impl ToBitMask for Mask<_, 64> +} diff --git a/crates/core_simd/tests/masks.rs b/crates/core_simd/tests/masks.rs index 6a8ecd33a73..3aec36ca7b7 100644 --- a/crates/core_simd/tests/masks.rs +++ b/crates/core_simd/tests/masks.rs @@ -68,16 +68,16 @@ macro_rules! test_mask_api { assert_eq!(core_simd::Mask::<$type, 8>::from_int(int), mask); } - #[cfg(feature = "generic_const_exprs")] #[test] fn roundtrip_bitmask_conversion() { + use core_simd::ToBitMask; let values = [ true, false, false, true, false, false, true, false, true, true, false, false, false, false, false, true, ]; let mask = core_simd::Mask::<$type, 16>::from_array(values); let bitmask = mask.to_bitmask(); - assert_eq!(bitmask, [0b01001001, 0b10000011]); + assert_eq!(bitmask, 0b1000001101001001); assert_eq!(core_simd::Mask::<$type, 16>::from_bitmask(bitmask), mask); } }