first_true_[mask]

sarah-quinones · Oct 12, 2024 · d9dea94 · d9dea94
1 parent a8e0222
commit d9dea94
Showing 1 changed file with 67 additions and 2 deletions.
diff --git a/pulp/src/lib.rs b/pulp/src/lib.rs
@@ -135,13 +135,13 @@ impl<F: NullaryFnOnce> WithSimd for F {
 }
 
 pub trait Simd: Seal + Debug + Copy + Send + Sync + 'static {
-    type m32s: Debug + Copy + Send + Sync + Zeroable + 'static;
+    type m32s: Debug + Copy + Send + Sync + Zeroable + NoUninit + 'static;
     type f32s: Debug + Copy + Send + Sync + Pod + 'static;
     type c32s: Debug + Copy + Send + Sync + Pod + 'static;
     type i32s: Debug + Copy + Send + Sync + Pod + 'static;
     type u32s: Debug + Copy + Send + Sync + Pod + 'static;
 
-    type m64s: Debug + Copy + Send + Sync + Zeroable + 'static;
+    type m64s: Debug + Copy + Send + Sync + Zeroable + NoUninit + 'static;
     type f64s: Debug + Copy + Send + Sync + Pod + 'static;
     type c64s: Debug + Copy + Send + Sync + Pod + 'static;
     type i64s: Debug + Copy + Send + Sync + Pod + 'static;
@@ -985,6 +985,53 @@ pub trait Simd: Seal + Debug + Copy + Send + Sync + 'static {
         self.partial_store_last_f64s(bytemuck::cast_slice_mut(slice), cast(values))
     }
 
+    #[inline(always)]
+    fn first_true_m32s(self, mask: Self::m32s) -> usize {
+        if const { size_of::<Self::m32s>() == size_of::<Self::u32s>() } {
+            let mask: Self::u32s = bytemuck::cast(mask);
+            let slice = bytemuck::cast_slice::<Self::u32s, u32>(core::slice::from_ref(&mask));
+            let mut i = 0;
+            for &x in slice.iter() {
+                if x != 0 {
+                    break;
+                }
+                i += 1;
+            }
+            i
+        } else if const { size_of::<Self::m32s>() == size_of::<u8>() } {
+            let mask: u8 = bytemuck::cast(mask);
+            mask.leading_zeros() as usize
+        } else if const { size_of::<Self::m32s>() == size_of::<u16>() } {
+            let mask: u16 = bytemuck::cast(mask);
+            mask.leading_zeros() as usize
+        } else {
+            panic!()
+        }
+    }
+    #[inline(always)]
+    fn first_true_m64s(self, mask: Self::m64s) -> usize {
+        if const { size_of::<Self::m64s>() == size_of::<Self::u64s>() } {
+            let mask: Self::u64s = bytemuck::cast(mask);
+            let slice = bytemuck::cast_slice::<Self::u64s, u64>(core::slice::from_ref(&mask));
+            let mut i = 0;
+            for &x in slice.iter() {
+                if x != 0 {
+                    break;
+                }
+                i += 1;
+            }
+            i
+        } else if const { size_of::<Self::m64s>() == size_of::<u8>() } {
+            let mask: u8 = bytemuck::cast(mask);
+            mask.leading_zeros() as usize
+        } else if const { size_of::<Self::m64s>() == size_of::<u16>() } {
+            let mask: u16 = bytemuck::cast(mask);
+            mask.leading_zeros() as usize
+        } else {
+            panic!()
+        }
+    }
+
     #[inline(always)]
     fn tail_mask_f64s(self, len: usize) -> Self::m64s {
         let iota: Self::u64s = const {
@@ -2385,6 +2432,24 @@ impl Simd for Scalar {
         let im = re;
         Complex { re, im }
     }
+
+    #[inline(always)]
+    fn first_true_m32s(self, mask: Self::m32s) -> usize {
+        if mask {
+            0
+        } else {
+            1
+        }
+    }
+
+    #[inline(always)]
+    fn first_true_m64s(self, mask: Self::m64s) -> usize {
+        if mask {
+            0
+        } else {
+            1
+        }
+    }
 }
 
 #[derive(Copy, Clone)]