Make implementation more scalable by using a helper trait to determine bitmask size. Improve bitmask to int conversion.

Author: calebzulawski

crates/core_simd/src/comparisons.rs

@@ -7,6 +7,7 @@ macro_rules! implement_mask_ops {
             where
                 crate::$vector<LANES>: LanesAtMost32,
                 crate::$inner_ty<LANES>: LanesAtMost32,
+                crate::$mask<LANES>: crate::Mask,
             {
                 /// Test if each lane is equal to the corresponding lane in `other`.
                 #[inline]

crates/core_simd/src/intrinsics.rs

@@ -79,6 +79,9 @@ extern "platform-intrinsic" {
 
     // truncate integer vector to bitmask
     pub(crate) fn simd_bitmask<T, U>(x: T) -> U;
+
+    // select
+    pub(crate) fn simd_select_bitmask<T, U>(m: T, a: U, b: U) -> U;
 }
 
 #[cfg(feature = "std")]

crates/core_simd/src/masks/bitmask.rs

@@ -1,56 +1,103 @@
+use crate::Mask;
+use core::marker::PhantomData;
+
+/// Helper trait for limiting int conversion types
+pub trait ConvertToInt {}
+impl<const LANES: usize> ConvertToInt for crate::SimdI8<LANES> where Self: crate::LanesAtMost32 {}
+impl<const LANES: usize> ConvertToInt for crate::SimdI16<LANES> where Self: crate::LanesAtMost32 {}
+impl<const LANES: usize> ConvertToInt for crate::SimdI32<LANES> where Self: crate::LanesAtMost32 {}
+impl<const LANES: usize> ConvertToInt for crate::SimdI64<LANES> where Self: crate::LanesAtMost32 {}
+impl<const LANES: usize> ConvertToInt for crate::SimdIsize<LANES> where Self: crate::LanesAtMost32 {}
+
 /// A mask where each lane is represented by a single bit.
-#[derive(Copy, Clone, Debug, PartialOrd, PartialEq, Ord, Eq, Hash)]
 #[repr(transparent)]
-pub struct BitMask<const LANES: usize>(u64);
+pub struct BitMask<T: Mask, const LANES: usize>(T::BitMask, PhantomData<[(); LANES]>);
 
-impl<const LANES: usize> BitMask<LANES>
-{
+impl<T: Mask, const LANES: usize> Copy for BitMask<T, LANES> {}
+
+impl<T: Mask, const LANES: usize> Clone for BitMask<T, LANES> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T: Mask, const LANES: usize> PartialEq for BitMask<T, LANES> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.as_ref() == other.0.as_ref()
+    }
+}
+
+impl<T: Mask, const LANES: usize> PartialOrd for BitMask<T, LANES> {
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        self.0.as_ref().partial_cmp(other.0.as_ref())
+    }
+}
+
+impl<T: Mask, const LANES: usize> Eq for BitMask<T, LANES> {}
+
+impl<T: Mask, const LANES: usize> Ord for BitMask<T, LANES> {
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        self.0.as_ref().cmp(other.0.as_ref())
+    }
+}
+
+impl<T: Mask, const LANES: usize> BitMask<T, LANES> {
     #[inline]
     pub fn splat(value: bool) -> Self {
+        let mut mask = T::BitMask::default();
         if value {
-            Self(u64::MAX >> (64 - LANES))
+            mask.as_mut().fill(u8::MAX)
         } else {
-            Self(u64::MIN)
+            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]
     pub unsafe fn test_unchecked(&self, lane: usize) -> bool {
-        (self.0 >> lane) & 0x1 > 0
+        (self.0.as_ref()[lane / 8] >> lane % 8) & 0x1 > 0
     }
 
     #[inline]
     pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) {
-        self.0 ^= ((value ^ self.test_unchecked(lane)) as u64) << lane
+        self.0.as_mut()[lane / 8] ^= ((value ^ self.test_unchecked(lane)) as u8) << (lane % 8)
     }
 
     #[inline]
-    pub fn to_int<V, T>(self) -> V
+    pub fn to_int<V>(self) -> V
     where
-        V: Default + AsMut<[T; LANES]>,
-        T: From<i8>,
+        V: ConvertToInt + Default + core::ops::Not<Output = V>,
     {
-        // TODO this should be an intrinsic sign-extension
-        let mut v = V::default();
-        for i in 0..LANES {
-            let lane = unsafe { self.test_unchecked(i) };
-            v.as_mut()[i] = (-(lane as i8)).into();
+        unsafe {
+            let mask: T::IntBitMask = core::mem::transmute_copy(&self);
+            crate::intrinsics::simd_select_bitmask(mask, !V::default(), V::default())
         }
-        v
     }
 
     #[inline]
     pub unsafe fn from_int_unchecked<V>(value: V) -> Self
     where
         V: crate::LanesAtMost32,
     {
-        let mask: V::BitMask = crate::intrinsics::simd_bitmask(value);
-        Self(mask.into())
+        // TODO remove the transmute when rustc is more flexible
+        assert_eq!(
+            core::mem::size_of::<T::IntBitMask>(),
+            core::mem::size_of::<T::BitMask>()
+        );
+        let mask: T::IntBitMask = crate::intrinsics::simd_bitmask(value);
+        Self(core::mem::transmute_copy(&mask), PhantomData)
     }
 
     #[inline]
-    pub fn to_bitmask(self) -> u64 {
-        self.0
+    pub fn to_bitmask<U: Mask>(self) -> U::BitMask {
+        assert_eq!(
+            core::mem::size_of::<T::BitMask>(),
+            core::mem::size_of::<U::BitMask>()
+        );
+        unsafe { core::mem::transmute_copy(&self.0) }
     }
 
     #[inline]
@@ -64,87 +111,61 @@ impl<const LANES: usize> BitMask<LANES>
     }
 }
 
-impl<const LANES: usize> core::ops::BitAnd for BitMask<LANES>
+impl<T: Mask, const LANES: usize> core::ops::BitAnd for BitMask<T, LANES>
+where
+    T::BitMask: Default + AsRef<[u8]> + AsMut<[u8]>,
 {
     type Output = Self;
     #[inline]
-    fn bitand(self, rhs: Self) -> Self {
-        Self(self.0 & rhs.0)
+    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<const LANES: usize> core::ops::BitAnd<bool> for BitMask<LANES>
+impl<T: Mask, const LANES: usize> core::ops::BitOr for BitMask<T, LANES>
+where
+    T::BitMask: Default + AsRef<[u8]> + AsMut<[u8]>,
 {
     type Output = Self;
     #[inline]
-    fn bitand(self, rhs: bool) -> Self {
-        self & Self::splat(rhs)
-    }
-}
-
-impl<const LANES: usize> core::ops::BitAnd<BitMask<LANES>> for bool
-{
-    type Output = BitMask<LANES>;
-    #[inline]
-    fn bitand(self, rhs: BitMask<LANES>) -> BitMask<LANES> {
-        BitMask::<LANES>::splat(self) & rhs
+    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<const LANES: usize> core::ops::BitOr for BitMask<LANES>
-{
+impl<T: Mask, const LANES: usize> core::ops::BitXor for BitMask<T, LANES> {
     type Output = Self;
     #[inline]
-    fn bitor(self, rhs: Self) -> Self {
-        Self(self.0 | rhs.0)
+    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<const LANES: usize> core::ops::BitXor for BitMask<LANES>
-{
+impl<T: Mask, const LANES: usize> core::ops::Not for BitMask<T, LANES> {
     type Output = Self;
     #[inline]
-    fn bitxor(self, rhs: Self) -> Self::Output {
-        Self(self.0 ^ rhs.0)
-    }
-}
-
-impl<const LANES: usize> core::ops::Not for BitMask<LANES>
-{
-    type Output = BitMask<LANES>;
-    #[inline]
-    fn not(self) -> Self::Output {
-        Self(!self.0) & Self::splat(true)
-    }
-}
-
-impl<const LANES: usize> core::ops::BitAndAssign for BitMask<LANES>
-{
-    #[inline]
-    fn bitand_assign(&mut self, rhs: Self) {
-        self.0 &= rhs.0;
-    }
-}
-
-impl<const LANES: usize> core::ops::BitOrAssign for BitMask<LANES>
-{
-    #[inline]
-    fn bitor_assign(&mut self, rhs: Self) {
-        self.0 |= rhs.0;
-    }
-}
-
-impl<const LANES: usize> core::ops::BitXorAssign for BitMask<LANES>
-{
-    #[inline]
-    fn bitxor_assign(&mut self, rhs: Self) {
-        self.0 ^= rhs.0;
+    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
     }
 }
 
-pub type Mask8<const LANES: usize> = BitMask<LANES>;
-pub type Mask16<const LANES: usize> = BitMask<LANES>;
-pub type Mask32<const LANES: usize> = BitMask<LANES>;
-pub type Mask64<const LANES: usize> = BitMask<LANES>;
-pub type Mask128<const LANES: usize> = BitMask<LANES>;
-pub type MaskSize<const LANES: usize> = BitMask<LANES>;
+pub type Mask8<T, const LANES: usize> = BitMask<T, LANES>;
+pub type Mask16<T, const LANES: usize> = BitMask<T, LANES>;
+pub type Mask32<T, const LANES: usize> = BitMask<T, LANES>;
+pub type Mask64<T, const LANES: usize> = BitMask<T, LANES>;
+pub type MaskSize<T, const LANES: usize> = BitMask<T, LANES>;