par_iter_bp_with_buf

Author: RagnarGrootKoerkamp

src/packed_seq.rs

@@ -1,4 +1,5 @@
 use core::cell::RefCell;
+use std::ops::{Deref, DerefMut};
 use traits::Seq;
 use wide::u16x8;
 
@@ -35,6 +36,21 @@ impl RecycledBox {
     }
 }
 
+impl Deref for RecycledBox {
+    type Target = SimdBuf;
+
+    #[inline(always)]
+    fn deref(&self) -> &Self::Target {
+        self.get()
+    }
+}
+impl DerefMut for RecycledBox {
+    #[inline(always)]
+    fn deref_mut(&mut self) -> &mut SimdBuf {
+        self.get_mut()
+    }
+}
+
 impl Drop for RecycledBox {
     #[inline(always)]
     fn drop(&mut self) {
@@ -495,6 +511,72 @@ where
 
     #[inline(always)]
     fn par_iter_bp(self, context: usize) -> PaddedIt<impl ChunkIt<S>> {
+        // Boxed, so it doesn't consume precious registers.
+        // Without this, cur is not always inlined into a register.
+        let mut buf = IT_BUF.with_borrow_mut(|v| RecycledBox(v.pop()));
+        buf.init_if_needed();
+        self.par_iter_bp_with_buf(context, buf)
+    }
+
+    #[inline(always)]
+    fn par_iter_bp_delayed(self, context: usize, delay: Delay) -> PaddedIt<impl ChunkIt<(S, S)>> {
+        self.par_iter_bp_delayed_with_factor(context, delay, 1)
+    }
+
+    /// NOTE: When `self` starts does not start at a byte boundary, the
+    /// 'delayed' character is not guaranteed to be `0`.
+    #[inline(always)]
+    fn par_iter_bp_delayed_2(
+        self,
+        context: usize,
+        delay1: Delay,
+        delay2: Delay,
+    ) -> PaddedIt<impl ChunkIt<(S, S, S)>> {
+        self.par_iter_bp_delayed_2_with_factor(context, delay1, delay2, 1)
+    }
+
+    /// Compares 29 characters at a time.
+    fn cmp_lcp(&self, other: &Self) -> (std::cmp::Ordering, usize) {
+        let mut lcp = 0;
+        let min_len = self.len.min(other.len);
+        for i in (0..min_len).step_by(Self::K64) {
+            let len = (min_len - i).min(Self::K64);
+            let this = self.slice(i..i + len);
+            let other = other.slice(i..i + len);
+            let this_word = this.as_u64();
+            let other_word = other.as_u64();
+            if this_word != other_word {
+                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
+                let eq = this_word ^ other_word;
+                let t = eq.trailing_zeros() as usize / B * B;
+                lcp += t / B;
+                let mask = (Self::CHAR_MASK) << t;
+                return ((this_word & mask).cmp(&(other_word & mask)), lcp);
+            }
+            lcp += len;
+        }
+        (self.len.cmp(&other.len), lcp)
+    }
+
+    #[inline(always)]
+    fn get(&self, index: usize) -> u8 {
+        let offset = self.offset + index;
+        let idx = offset / Self::C8;
+        let offset = offset % Self::C8;
+        (self.seq[idx] >> (B * offset)) & Self::CHAR_MASK as u8
+    }
+}
+
+impl<'s, const B: usize> PackedSeqBase<'s, B>
+where
+    Bits<B>: SupportedBits,
+{
+    #[inline(always)]
+    pub fn par_iter_bp_with_buf<BUF: DerefMut<Target = [S; 8]>>(
+        self,
+        context: usize,
+        mut buf: BUF,
+    ) -> PaddedIt<impl ChunkIt<S> + use<'s, B, BUF>> {
         #[cfg(target_endian = "big")]
         panic!("Big endian architectures are not supported.");
 
@@ -516,11 +598,6 @@ where
         let offsets: [usize; 8] = from_fn(|l| l * bytes_per_chunk);
         let mut cur = S::ZERO;
 
-        // Boxed, so it doesn't consume precious registers.
-        // Without this, cur is not always inlined into a register.
-        let mut buf = IT_BUF.with_borrow_mut(|v| RecycledBox(v.pop()));
-        buf.init_if_needed();
-
         let simd_char_mask: u32x8 = S::splat(Self::CHAR_MASK as u32);
         let simd_b: u32x8 = S::splat(B as u32);
 
@@ -550,9 +627,9 @@ where
                                     )
                                 },
                             );
-                            *buf.get_mut() = transpose(data);
+                            *buf = transpose(data);
                         }
-                        cur = buf.get()[(i % Self::C256) / Self::C32];
+                        cur = buf[(i % Self::C256) / Self::C32];
                     }
                     // Extract the last 2 bits of each character.
                     let chars = cur & simd_char_mask;
@@ -566,59 +643,6 @@ where
         PaddedIt { it, padding }
     }
 
-    #[inline(always)]
-    fn par_iter_bp_delayed(self, context: usize, delay: Delay) -> PaddedIt<impl ChunkIt<(S, S)>> {
-        self.par_iter_bp_delayed_with_factor(context, delay, 1)
-    }
-
-    /// NOTE: When `self` starts does not start at a byte boundary, the
-    /// 'delayed' character is not guaranteed to be `0`.
-    #[inline(always)]
-    fn par_iter_bp_delayed_2(
-        self,
-        context: usize,
-        delay1: Delay,
-        delay2: Delay,
-    ) -> PaddedIt<impl ChunkIt<(S, S, S)>> {
-        self.par_iter_bp_delayed_2_with_factor(context, delay1, delay2, 1)
-    }
-
-    /// Compares 29 characters at a time.
-    fn cmp_lcp(&self, other: &Self) -> (std::cmp::Ordering, usize) {
-        let mut lcp = 0;
-        let min_len = self.len.min(other.len);
-        for i in (0..min_len).step_by(Self::K64) {
-            let len = (min_len - i).min(Self::K64);
-            let this = self.slice(i..i + len);
-            let other = other.slice(i..i + len);
-            let this_word = this.as_u64();
-            let other_word = other.as_u64();
-            if this_word != other_word {
-                // Unfortunately, bases are packed in little endian order, so the default order is reversed.
-                let eq = this_word ^ other_word;
-                let t = eq.trailing_zeros() as usize / B * B;
-                lcp += t / B;
-                let mask = (Self::CHAR_MASK) << t;
-                return ((this_word & mask).cmp(&(other_word & mask)), lcp);
-            }
-            lcp += len;
-        }
-        (self.len.cmp(&other.len), lcp)
-    }
-
-    #[inline(always)]
-    fn get(&self, index: usize) -> u8 {
-        let offset = self.offset + index;
-        let idx = offset / Self::C8;
-        let offset = offset % Self::C8;
-        (self.seq[idx] >> (B * offset)) & Self::CHAR_MASK as u8
-    }
-}
-
-impl<'s, const B: usize> PackedSeqBase<'s, B>
-where
-    Bits<B>: SupportedBits,
-{
     #[inline(always)]
     pub fn par_iter_bp_delayed_with_factor(
         self,

src/test.rs

@@ -1000,6 +1000,34 @@ fn par_iter_bp_bench() {
     }
 }
 
+#[test]
+#[ignore = "This is a benchmark, not a test"]
+fn par_iter_bp_buf_bench() {
+    eprintln!("\nBench PackedSeq::par_iter_bp_buf");
+
+    let mut buf = [S::ZERO; 8];
+
+    for len in [100, 150, 200, 1000, 1_000_000] {
+        // 1Gbp input.
+        let rep = 1_000_000_000 / len;
+        let seq = PackedSeqVec::random(len);
+
+        let start = std::time::Instant::now();
+        for _ in 0..rep {
+            let PaddedIt { it, .. } = seq.as_slice().par_iter_bp_with_buf(1, &mut buf);
+            let mut sum = S::ZERO;
+            for x in it {
+                sum += x;
+            }
+            black_box(&sum);
+        }
+        eprintln!(
+            "Len {len:>7} => {:.03} Gbp/s",
+            start.elapsed().as_secs_f64().recip()
+        );
+    }
+}
+
 #[test]
 #[ignore = "This is a benchmark, not a test"]
 fn par_iter_bp_delayed_bench() {