extract more splats

Author: RagnarGrootKoerkamp

src/collect.rs

@@ -169,8 +169,10 @@ impl<I: ChunkIt<u32x8>> CollectAndDedup for PaddedIt<I> {
                     break;
                 }
 
+                // FIXME: IS this one slow?
                 let mut m = [u32x8::ZERO; 8];
                 let mut i = 0;
+                let eight = S::splat(8);
                 it.for_each(
                     #[inline(always)]
                     |x| {
@@ -218,7 +220,7 @@ impl<I: ChunkIt<u32x8>> CollectAndDedup for PaddedIt<I> {
                                 }
                                 old[j] = lane;
                             }
-                            offsets += u32x8::splat(8);
+                            offsets += eight;
                         }
                         i += 1;
                     },

src/intrinsics/dedup.rs

@@ -1,8 +1,13 @@
 use crate::S;
-use crate::minimizers::SKIPPED;
+use crate::minimizers::SIMD_SKIPPED;
 use core::mem::transmute;
 use packed_seq::L;
 
+#[cfg(target_feature = "neon")]
+const OFFSET: S = unsafe { std::mem::transmute([0x03_02_01_00; 8]) };
+#[cfg(target_feature = "neon")]
+const MASK: S = unsafe { std::mem::transmute([0x04_04_04_04; 8]) };
+
 /// 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)]
@@ -87,7 +92,7 @@ pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
         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));
+            m |= new.cmp_eq(SIMD_SKIPPED);
         }
         let m = _mm256_movemask_ps(transmute(m)) as usize;
         let numberofnewvalues = L - m.count_ones() as usize;
@@ -189,7 +194,7 @@ pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
 
         let mut dup = prec.cmp_eq(new);
         if SKIP_MAX {
-            dup |= new.cmp_eq(S::splat(SKIPPED));
+            dup |= new.cmp_eq(SIMD_SKIPPED);
         }
         // emulate movemask
         let (d1, d2): (u32x4, u32x4) = transmute(dup);
@@ -201,7 +206,7 @@ pub unsafe fn append_unique_vals<const SKIP_MAX: bool>(
 
         let numberofnewvalues = L - m.count_ones() as usize;
         let key = UNIQSHUF[m];
-        let idx = key * S::splat(0x04_04_04_04) + S::splat(0x03_02_01_00);
+        let idx = key * MASK + OFFSET;
         let (i1, i2) = transmute(idx);
         let t = transmute(vals);
         let r1 = vqtbl2q_u8(t, i1);
@@ -273,7 +278,7 @@ pub unsafe fn append_unique_vals_2(
 
         let numberofnewvalues = L - m.count_ones() as usize;
         let key = UNIQSHUF[m];
-        let idx = key * S::splat(0x04_04_04_04) + S::splat(0x03_02_01_00);
+        let idx = key * MASK + OFFSET;
         let (i1, i2) = transmute(idx);
         let t = transmute(vals);
         let r1 = vqtbl2q_u8(t, i1);

src/minimizers.rs

@@ -16,6 +16,7 @@ use seq_hash::KmerHasher;
 use wide::u32x8;
 
 pub const SKIPPED: u32 = u32::MAX - 1;
+pub(crate) const SIMD_SKIPPED: u32x8 = unsafe { std::mem::transmute([SKIPPED; 8]) };
 
 /// Minimizer position of a single window.
 pub fn one_minimizer<'s>(seq: impl Seq<'s>, hasher: &impl KmerHasher) -> usize {
@@ -202,6 +203,6 @@ pub fn canonical_minimizers_skip_ambiguous_windows<'s>(
             let hash = hash_mapper((a, rh));
             let canonical = canonical_mapper((a, rc));
             let (lmin, rmin) = sliding_min_mapper(hash);
-            ambi.blend(u32x8::splat(SKIPPED), canonical.blend(lmin, rmin))
+            ambi.blend(SIMD_SKIPPED, canonical.blend(lmin, rmin))
         })
 }

src/sliding_min.rs

@@ -235,24 +235,24 @@ pub fn sliding_min_mapper_simd<const LEFT: bool>(
     let pos_mask = S::splat(0x0000_ffff);
     let max_pos = S::splat((1 << 16) - 1);
     let mut pos = S::splat(0);
-    let one = S::splat(1);
     // Sliding min is over w+k-1 characters, so chunks overlap w+k-2.
     // Thus, the true length of each lane is len-(k+w-2).
     //
     // The k-mer starting at position 0 is done after processing the char at
     // position k-1, so we compensate for that as well.
     let mut pos_offset: S = from_fn(|l| (l * len.saturating_sub(w - 1)) as u32).into();
+    let delta = S::splat((1 << 16) - 2 - w as u32);
 
     #[inline(always)]
     move |val| {
         // Make sure the position does not interfere with the hash value.
         if pos == max_pos {
             // Slow case extracted to a function to have better inlining here.
-            reset_positions_offsets(w, &mut pos, &mut prefix_min, &mut pos_offset, ring_buf);
+            reset_positions_offsets(delta, &mut pos, &mut prefix_min, &mut pos_offset, ring_buf);
         }
         // slightly faster than assigning S::splat(u32::MAX)
         let elem = (if LEFT { val } else { !val } & val_mask) | pos;
-        pos += one;
+        pos += S::ONE;
         ring_buf.push(elem);
         prefix_min = simd_min::<LEFT>(prefix_min, elem);
         // After a chunk has been filled, compute suffix minima.
@@ -284,13 +284,12 @@ fn suffix_minima<const LEFT: bool>(
 }
 
 fn reset_positions_offsets(
-    w: usize,
+    delta: S,
     pos: &mut S,
     prefix_min: &mut S,
     pos_offset: &mut S,
     ring_buf: &mut RingBuf<S>,
 ) {
-    let delta = S::splat((1 << 16) - 2 - w as u32);
     *pos -= delta;
     *prefix_min -= delta;
     *pos_offset += delta;
@@ -318,7 +317,6 @@ pub fn sliding_lr_min_mapper_simd(
     let max_pos = S::splat((1 << 16) - 1);
     let mut pos = S::splat(0);
     let mut pos_offset: S = from_fn(|l| (l * len.saturating_sub(w - 1)) as u32).into();
-    let one = S::splat(1);
     let delta = S::splat((1 << 16) - 2 - w as u32);
 
     #[inline(always)]
@@ -338,7 +336,7 @@ pub fn sliding_lr_min_mapper_simd(
         let lelem = (val & val_mask) | pos;
         let relem = (!val & val_mask) | pos;
         let elem = (lelem, relem);
-        pos += one;
+        pos += S::ONE;
         ring_buf.push(elem);
         prefix_lr_min = simd_lr_min(prefix_lr_min, elem);
         // After a chunk has been filled, compute suffix minima.