collect & dedup: add generic to skip u32::MAX-1 during collecting

Author: RagnarGrootKoerkamp

bench/src/bin/paper.rs

@@ -147,7 +147,7 @@ fn plot() {
                     .collect_into(v2);
                 v2.clear();
                 minimizers::canonical_minimizers_seq_simd(packed_seq, &can_hasher, w, &mut cache)
-                    .collect_and_dedup_into(v2);
+                    .collect_and_dedup_into::<false>(v2);
                 v2.clear();
             }
 

src/collect.rs

@@ -8,7 +8,7 @@ use std::{
     mem::transmute,
 };
 
-use crate::S;
+use crate::{S, minimizers::SKIPPED};
 use packed_seq::{ChunkIt, L, PaddedIt, intrinsics::transpose};
 use wide::u32x8;
 
@@ -76,9 +76,9 @@ pub fn collect_and_dedup_with_index_into_scalar(
 
 pub trait CollectAndDedup: Sized {
     /// Convenience wrapper around `collect_and_dedup_into`.
-    fn collect_and_dedup<const SUPER: bool>(self) -> Vec<u32> {
+    fn collect_and_dedup<const SKIP_MAX: bool>(self) -> Vec<u32> {
         let mut v = vec![];
-        self.collect_and_dedup_into(&mut v);
+        self.collect_and_dedup_into::<SKIP_MAX>(&mut v);
         v
     }
 
@@ -95,8 +95,8 @@ pub trait CollectAndDedup: Sized {
     ///
     /// The output is simply the deduplicated input values.
     #[inline(always)]
-    fn collect_and_dedup_into(self, out_vec: &mut Vec<u32>) {
-        self.collect_and_dedup_into_impl::<false>(out_vec, &mut vec![]);
+    fn collect_and_dedup_into<const SKIP_MAX: bool>(self, out_vec: &mut Vec<u32>) {
+        self.collect_and_dedup_into_impl::<false, SKIP_MAX>(out_vec, &mut vec![]);
     }
 
     /// Collect a SIMD-iterator into a single vector, and duplicate adjacent equal elements.
@@ -105,15 +105,15 @@ pub trait CollectAndDedup: Sized {
     /// The deduplicated input values are written in `out_vec` and the index of the stream it first appeared, i.e., the start of its super-k-mer, is written in `idx_vec`.
     #[inline(always)]
     fn collect_and_dedup_with_index_into(self, out_vec: &mut Vec<u32>, idx_vec: &mut Vec<u32>) {
-        self.collect_and_dedup_into_impl::<true>(out_vec, idx_vec);
+        self.collect_and_dedup_into_impl::<true, false>(out_vec, idx_vec);
     }
 
     /// Collect a SIMD-iterator into a single vector, and duplicate adjacent equal elements.
     /// Works by taking 8 elements from each stream, and then transposing the SIMD-matrix before writing out the results.
     ///
     /// By default (when `SUPER` is false), the deduplicated input values are written in `out_vec`.
     /// When `SUPER` is true, the index of the stream in which the input value first appeared, i.e., the start of its super-k-mer, is additionale written in `idx_vec`.
-    fn collect_and_dedup_into_impl<const SUPER: bool>(
+    fn collect_and_dedup_into_impl<const SUPER: bool, const SKIP_MAX: bool>(
         self,
         out_vec: &mut Vec<u32>,
         idx_vec: &mut Vec<u32>,
@@ -126,7 +126,7 @@ thread_local! {
 
 impl<I: ChunkIt<u32x8>> CollectAndDedup for PaddedIt<I> {
     #[inline(always)]
-    fn collect_and_dedup_into_impl<const SUPER: bool>(
+    fn collect_and_dedup_into_impl<const SUPER: bool, const SKIP_MAX: bool>(
         self,
         out_vec: &mut Vec<u32>,
         idx_vec: &mut Vec<u32>,
@@ -207,7 +207,7 @@ impl<I: ChunkIt<u32x8>> CollectAndDedup for PaddedIt<I> {
                                             &mut write_idx[j],
                                         );
                                     } else {
-                                        crate::intrinsics::append_unique_vals(
+                                        crate::intrinsics::append_unique_vals::<SKIP_MAX>(
                                             old[j],
                                             lane,
                                             lane,
@@ -237,7 +237,7 @@ impl<I: ChunkIt<u32x8>> CollectAndDedup for PaddedIt<I> {
                 for j in 0..8 {
                     let lane = t[j].as_array_ref();
                     for (p, x) in lane.iter().take(k).enumerate() {
-                        if v[j].last() != Some(x) {
+                        if v[j].last() != Some(x) && (!SKIP_MAX || *x != SKIPPED) {
                             v[j].push(*x);
                             if SUPER {
                                 v2[j].push(

src/intrinsics/dedup.rs

@@ -1,19 +1,26 @@
 use crate::S;
+use crate::minimizers::SKIPPED;
 use core::mem::transmute;
 use packed_seq::L;
 
 /// Dedup adjacent `new` values (starting with the last element of `old`).
 /// If an element is different from the preceding element, append the corresponding element of `vals` to `v[write_idx]`.
 #[inline(always)]
 #[cfg(not(any(target_feature = "avx2", target_feature = "neon")))]
-pub unsafe fn append_unique_vals(old: S, new: S, vals: S, v: &mut [u32], write_idx: &mut usize) {
+pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
+    old: S,
+    new: S,
+    vals: S,
+    v: &mut [u32],
+    write_idx: &mut usize,
+) {
     unsafe {
         let old = old.to_array();
         let new = new.to_array();
         let vals = vals.to_array();
         let mut prec = old[7];
         for (i, &curr) in new.iter().enumerate() {
-            if curr != prec {
+            if curr != prec && cur != SKIPPED {
                 v.as_mut_ptr().add(*write_idx).write(vals[i]);
                 *write_idx += 1;
                 prec = curr;
@@ -26,7 +33,7 @@ pub unsafe fn append_unique_vals(old: S, new: S, vals: S, v: &mut [u32], write_i
 /// If an element is different from the preceding element, append the corresponding element of `vals` to `v[write_idx]` and `vals2` to `v2[write_idx]`.
 #[inline(always)]
 #[cfg(not(any(target_feature = "avx2", target_feature = "neon")))]
-pub unsafe fn append_unique_vals_2(
+pub unsafe fn append_unique_vals_2<const SKIP_MAX: bool>(
     old: S,
     new: S,
     vals: S,
@@ -60,19 +67,29 @@ pub unsafe fn append_unique_vals_2(
 /// <https://github.com/lemire/Code-used-on-Daniel-Lemire-s-blog/blob/edfd0e8b809d9a57527a7990c4bb44b9d1d05a69/2017/04/10/removeduplicates.cpp>
 #[cfg(target_feature = "avx2")]
 #[inline(always)]
-pub unsafe fn append_unique_vals(old: S, new: S, vals: S, v: &mut [u32], write_idx: &mut usize) {
+pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
+    old: S,
+    new: S,
+    vals: S,
+    v: &mut [u32],
+    write_idx: &mut usize,
+) {
     unsafe {
         use core::arch::x86_64::*;
 
         let old = transmute(old);
-        let new = transmute(new);
         let vals = transmute(vals);
 
-        let recon = _mm256_blend_epi32(old, new, 0b01111111);
+        let recon = _mm256_blend_epi32(old, transmute(new), 0b01111111);
         let movebyone_mask = _mm256_set_epi32(6, 5, 4, 3, 2, 1, 0, 7); // rotate shuffle
-        let vec_tmp = _mm256_permutevar8x32_epi32(recon, movebyone_mask);
+        let vec_tmp: S = transmute(_mm256_permutevar8x32_epi32(recon, movebyone_mask));
 
-        let m = _mm256_movemask_ps(transmute(_mm256_cmpeq_epi32(vec_tmp, new))) as usize;
+        let mut m = vec_tmp.cmp_eq(new);
+        if SKIP_MAX {
+            // skip everything equal to prev, or equal to MAX.
+            m |= new.cmp_eq(S::splat(SKIPPED));
+        }
+        let m = _mm256_movemask_ps(transmute(m)) as usize;
         let numberofnewvalues = L - m.count_ones() as usize;
         let key = transmute(UNIQSHUF[m]);
         let val = _mm256_permutevar8x32_epi32(vals, key);
@@ -129,7 +146,13 @@ pub unsafe fn append_unique_vals_2(
 /// <https://github.com/lemire/Code-used-on-Daniel-Lemire-s-blog/blob/edfd0e8b809d9a57527a7990c4bb44b9d1d05a69/2017/04/10/removeduplicates.cpp>
 #[inline(always)]
 #[cfg(target_feature = "neon")]
-pub unsafe fn append_unique_vals(old: S, new: S, vals: S, v: &mut [u32], write_idx: &mut usize) {
+pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
+    old: S,
+    new: S,
+    vals: S,
+    v: &mut [u32],
+    write_idx: &mut usize,
+) {
     unsafe {
         use core::arch::aarch64::{vaddvq_u32, vqtbl2q_u8, vst1_u32_x4};
         use wide::u32x4;
@@ -164,7 +187,11 @@ pub unsafe fn append_unique_vals(old: S, new: S, vals: S, v: &mut [u32], write_i
         let r2 = vqtbl2q_u8(t, i2);
         let prec: S = transmute((r1, r2));
 
-        let dup = prec.cmp_eq(new);
+        let mut dup = prec.cmp_eq(new);
+        if SKIP_MAX {
+            dup |= new.cmp_eq(S::splat(SKIPPED));
+        }
+        // emulate movemask
         let (d1, d2): (u32x4, u32x4) = transmute(dup);
         let pow1 = u32x4::new([1, 2, 4, 8]);
         let pow2 = u32x4::new([16, 32, 64, 128]);
@@ -555,7 +582,7 @@ mod test {
             let start = Instant::now();
             for new in chunks {
                 unsafe {
-                    append_unique_vals(old, new, new, &mut v2, &mut write_idx);
+                    append_unique_vals::<false>(old, new, new, &mut v2, &mut write_idx);
                 }
                 old = new;
             }

src/lib.rs

@@ -303,11 +303,10 @@ impl<'h, const CANONICAL: bool, H: KmerHasher> Builder<'h, CANONICAL, H, ()> {
                 canonical_minimizers_seq_scalar(seq, hasher, self.w, cache),
                 min_pos,
             ),
-            (true, false) => {
-                minimizers_seq_simd(seq, hasher, self.w, cache).collect_and_dedup_into(min_pos)
-            }
+            (true, false) => minimizers_seq_simd(seq, hasher, self.w, cache)
+                .collect_and_dedup_into::<false>(min_pos),
             (true, true) => canonical_minimizers_seq_simd(seq, hasher, self.w, cache)
-                .collect_and_dedup_into(min_pos),
+                .collect_and_dedup_into::<false>(min_pos),
         });
         Output {
             k: self.k,

src/minimizers.rs

@@ -15,6 +15,8 @@ use packed_seq::{Advance, ChunkIt, Delay, PaddedIt, Seq};
 use seq_hash::KmerHasher;
 use wide::u32x8;
 
+pub const SKIPPED: u32 = u32::MAX - 1;
+
 /// Minimizer position of a single window.
 pub fn one_minimizer<'s>(seq: impl Seq<'s>, hasher: &impl KmerHasher) -> usize {
     hasher

src/test.rs

@@ -341,6 +341,50 @@ fn collect_and_dedup_with_index_scalar() {
 }
 
 #[test]
+fn collect_and_dedup_skip_max() {
+    let x = u32::MAX - 1;
+    let v = [0, 1, 1, x, 2, 3, x, x, 4].map(S::splat);
+
+    let mut out = vec![];
+    PaddedIt {
+        it: v.iter().copied(),
+        padding: 0,
+    }
+    .collect_and_dedup_into::<false>(&mut out);
+    assert!(
+        out.starts_with(&[0, 1, x, 2, 3, x, 4, 0, 1]),
+        "out: {out:?}"
+    );
+
+    let mut out = vec![];
+    PaddedIt {
+        it: v.iter().copied(),
+        padding: 0,
+    }
+    .collect_and_dedup_into::<true>(&mut out);
+    assert!(out.starts_with(&[0, 1, 2, 3, 4, 0, 1]), "out: {out:?}");
+
+    let v = [1, x, x, x, x, x, x, 2, x, x, x, x].map(S::splat);
+
+    let mut out = vec![];
+    PaddedIt {
+        it: v.iter().copied(),
+        padding: 0,
+    }
+    .collect_and_dedup_into::<false>(&mut out);
+    assert!(out.starts_with(&[1, x, 2, x, 1, x]), "out: {out:?}");
+
+    let mut out = vec![];
+    PaddedIt {
+        it: v.iter().copied(),
+        padding: 0,
+    }
+    .collect_and_dedup_into::<true>(&mut out);
+    assert!(out.starts_with(&[1, 2, 1, 2]), "out: {out:?}");
+}
+
+#[test]
+#[allow(unused)]
 fn readme_example() {
     use packed_seq::{PackedSeqVec, SeqVec};