diff --git a/crates/core_simd/src/masks/bitmask.rs b/crates/core_simd/src/masks/bitmask.rs index e27b2689606..7bf2add2036 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. @@ -116,13 +116,20 @@ where } #[inline] - pub unsafe fn to_bitmask_integer(self) -> U { + pub fn to_bitmask_integer(self) -> U + where + super::Mask: ToBitMask, + { + // Safety: these are the same types unsafe { core::mem::transmute_copy(&self.0) } } - // Safety: U must be the integer with the exact number of bits required to hold the bitmask for #[inline] - pub unsafe fn from_bitmask_integer(bitmask: U) -> Self { + 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) } } diff --git a/crates/core_simd/src/masks/full_masks.rs b/crates/core_simd/src/masks/full_masks.rs index 90af486a887..848997a0792 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,16 +127,34 @@ where } #[inline] - pub unsafe fn to_bitmask_integer(self) -> U { - // Safety: caller must only return bitmask types - unsafe { intrinsics::simd_bitmask(self.0) } + 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 + } } - // Safety: U must be the integer with the exact number of bits required to hold the bitmask for - // this mask #[inline] - pub unsafe fn from_bitmask_integer(bitmask: U) -> Self { - // Safety: caller must only pass bitmask types + 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, diff --git a/crates/core_simd/src/masks/to_bitmask.rs b/crates/core_simd/src/masks/to_bitmask.rs index 86143f2331f..1c2037764c1 100644 --- a/crates/core_simd/src/masks/to_bitmask.rs +++ b/crates/core_simd/src/masks/to_bitmask.rs @@ -1,9 +1,26 @@ 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. -pub trait ToBitMask { +/// +/// # 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; @@ -14,32 +31,18 @@ pub trait ToBitMask { 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. -pub trait ToBitMaskArray { - /// 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 { { $(unsafe impl ToBitMask for Mask<_, $lanes:literal>)* } => { $( - impl ToBitMask for Mask { + unsafe impl ToBitMask for Mask { type BitMask = $int; fn to_bitmask(self) -> $int { - unsafe { self.0.to_bitmask_integer() } + self.0.to_bitmask_integer() } fn from_bitmask(bitmask: $int) -> Self { - unsafe { Self(mask_impl::Mask::from_bitmask_integer(bitmask)) } + Self(mask_impl::Mask::from_bitmask_integer(bitmask)) } } )*