From d5d2bd9e22307887eadf820907936e4703447791 Mon Sep 17 00:00:00 2001 From: Ragnar Groot Koerkamp Date: Sat, 27 Sep 2025 01:15:26 +0200 Subject: [PATCH 1/2] Add NEON version for ambiguous-base detection --- src/packed_seq.rs | 52 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/src/packed_seq.rs b/src/packed_seq.rs index d686307..b3ff586 100644 --- a/src/packed_seq.rs +++ b/src/packed_seq.rs @@ -1,4 +1,5 @@ use traits::Seq; +use wide::u16x8; use crate::{intrinsics::transpose, padded_it::ChunkIt}; @@ -268,6 +269,20 @@ pub(crate) fn read_slice(seq: &[u8], idx: usize) -> u32x8 { } } +/// Read up to 16 bytes starting at idx. +#[allow(unused)] +#[inline(always)] +pub(crate) fn read_slice_16(seq: &[u8], idx: usize) -> u16x8 { + // assert!(idx <= seq.len()); + let mut result = [0u8; 16]; + let num_bytes = 16.min(seq.len().saturating_sub(idx)); + unsafe { + let src = seq.as_ptr().add(idx); + std::ptr::copy_nonoverlapping(src, result.as_mut_ptr(), num_bytes); + std::mem::transmute(result) + } +} + impl<'s, const B: usize> Seq<'s> for PackedSeqBase<'s, B> where Bits: SupportedBits, @@ -904,7 +919,6 @@ where #[allow(unused)] let mut last = unaligned; - // TODO: Vectorization for B=1? if B == 2 { #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] { @@ -950,7 +964,6 @@ where } } if B == 1 { - // FIXME: Add NEON version. #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { last = unaligned + len; @@ -992,6 +1005,41 @@ where } } } + + #[cfg(target_feature = "neon")] + { + last = unaligned + len; + self.len = len; + + for i in (unaligned..last).step_by(16) { + use std::mem::transmute as t; + + use wide::CmpEq; + type S = wide::i8x16; + let chars: S = unsafe { t(read_slice_16(seq, i)) }; + let upper_mask = !(b'a' - b'A'); + // make everything upper case + let chars = chars & S::splat(upper_mask as i8); + let lossy_encoded = chars & S::splat(6); + let table = unsafe { S::from(t::<_, S>(*b"AxCxTxGxxxxxxxxx")) }; + let lookup: S = + unsafe { t(std::arch::aarch64::vqtbl1q_u8(t(table), t(lossy_encoded))) }; + let packed_bytes = !(chars.cmp_eq(lookup).move_mask() as u16); + + if i + 16 <= last { + self.seq[idx + 0] = packed_bytes as u8; + self.seq[idx + 1] = (packed_bytes >> 8) as u8; + idx += 2; + } else { + let mut b = 0; + while i + b < last { + self.seq[idx] = (packed_bytes >> b) as u8; + idx += 1; + b += 8; + } + } + } + } } let mut packed_byte = 0; From e92817f08c5da675d863258cbe4d38e409cb6475 Mon Sep 17 00:00:00 2001 From: Igor Martayan Date: Tue, 30 Sep 2025 12:12:01 +0200 Subject: [PATCH 2/2] Use `.reverse_bits` in `rev_u64` and `rev_u128` --- src/packed_seq.rs | 31 ++++--------------------------- 1 file changed, 4 insertions(+), 27 deletions(-) diff --git a/src/packed_seq.rs b/src/packed_seq.rs index b3ff586..40ab59b 100644 --- a/src/packed_seq.rs +++ b/src/packed_seq.rs @@ -196,39 +196,16 @@ pub fn char_is_ambiguous(base: u8) -> u8 { (table[pack_char_lossy(base) as usize] != (base & upper_mask)) as u8 } -/// Reverse the bits in the input. -#[inline(always)] -const fn rev_raw(word: u64) -> u64 { - #[cfg(any(target_arch = "arm", target_arch = "aarch64"))] - { - // ARM can reverse bits in a single instruction - word.reverse_bits() - } - - #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))] - { - let mut res = word.swap_bytes(); - res = ((res >> 4) & 0x0F0F_0F0F_0F0F_0F0F) | ((res & 0x0F0F_0F0F_0F0F_0F0F) << 4); - res = ((res >> 2) & 0x3333_3333_3333_3333) | ((res & 0x3333_3333_3333_3333) << 2); - res = ((res >> 1) & 0x5555_5555_5555_5555) | ((res & 0x5555_5555_5555_5555) << 1); - res ^ 0xAAAA_AAAA_AAAA_AAAA - } -} - -/// Compute the reverse complement of a short sequence packed in a `u64`. +/// Reverse `len` bits packed in a `u64`. #[inline(always)] pub const fn rev_u64(word: u64, len: usize) -> u64 { - rev_raw(word) >> (usize::BITS as usize - len) + word.reverse_bits() >> (usize::BITS as usize - len) } +/// Reverse `len` bits packed in a `u128`. #[inline(always)] pub const fn rev_u128(word: u128, len: usize) -> u128 { - let low = word as u64; - let high = (word >> 64) as u64; - let rlow = rev_raw(low); - let rhigh = rev_raw(high); - let out = ((rlow as u128) << 64) | rhigh as u128; - out >> (u128::BITS as usize - len) + word.reverse_bits() >> (u128::BITS as usize - len) } // ======================================================================