diff --git a/Cargo.toml b/Cargo.toml index 87a3840..6fab290 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ authors = ["Ragnar Groot Koerkamp", "Igor Martayan"] description = "Constructing and iterating packed DNA sequences using SIMD" [dependencies] -wide = "1.5" +wide = "1.7" mem_dbg = "0.4" cfg-if = "1.0" ensure_simd = "0.1.0" diff --git a/src/intrinsics/transpose.rs b/src/intrinsics/transpose.rs index 5ba4e52..d111eed 100644 --- a/src/intrinsics/transpose.rs +++ b/src/intrinsics/transpose.rs @@ -62,31 +62,34 @@ fn _transpose(m: [S; 8]) -> [S; 8] { } #[inline(always)] -#[cfg(target_feature = "neon")] +#[cfg(any( + target_feature = "neon", + all(target_arch = "wasm32", target_feature = "simd128") +))] fn _transpose(m: [S; 8]) -> [S; 8] { unsafe { use core::mem::transmute; let m: [u32x4; 16] = transmute(m); - let t11 = transpose_4x4_neon( + let t11 = transpose_4x4( *m.get_unchecked(0), *m.get_unchecked(2), *m.get_unchecked(4), *m.get_unchecked(6), ); - let t21 = transpose_4x4_neon( + let t21 = transpose_4x4( *m.get_unchecked(1), *m.get_unchecked(3), *m.get_unchecked(5), *m.get_unchecked(7), ); - let t12 = transpose_4x4_neon( + let t12 = transpose_4x4( *m.get_unchecked(8), *m.get_unchecked(10), *m.get_unchecked(12), *m.get_unchecked(14), ); - let t22 = transpose_4x4_neon( + let t22 = transpose_4x4( *m.get_unchecked(9), *m.get_unchecked(11), *m.get_unchecked(13), @@ -116,7 +119,7 @@ fn _transpose(m: [S; 8]) -> [S; 8] { #[inline(always)] #[cfg(target_feature = "neon")] -fn transpose_4x4_neon(m0: u32x4, m1: u32x4, m2: u32x4, m3: u32x4) -> [u32x4; 4] { +fn transpose_4x4(m0: u32x4, m1: u32x4, m2: u32x4, m3: u32x4) -> [u32x4; 4] { unsafe { use core::arch::aarch64::vzipq_u32; use core::mem::transmute; @@ -128,7 +131,30 @@ fn transpose_4x4_neon(m0: u32x4, m1: u32x4, m2: u32x4, m3: u32x4) -> [u32x4; 4] } #[inline(always)] -#[cfg(not(any(target_feature = "avx", target_feature = "neon")))] +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +fn transpose_4x4(m0: u32x4, m1: u32x4, m2: u32x4, m3: u32x4) -> [u32x4; 4] { + unsafe { + use core::arch::wasm32::{i32x4_shuffle, v128}; + use core::mem::transmute; + + let zip = |a: v128, b: v128| { + ( + i32x4_shuffle::<0, 4, 1, 5>(a, b), + i32x4_shuffle::<2, 6, 3, 7>(a, b), + ) + }; + let x = zip(transmute(m0), transmute(m2)); + let y = zip(transmute(m1), transmute(m3)); + transmute((zip(x.0, y.0), zip(x.1, y.1))) + } +} + +#[inline(always)] +#[cfg(not(any( + target_feature = "avx", + target_feature = "neon", + all(target_arch = "wasm32", target_feature = "simd128") +)))] fn _transpose(m: [S; 8]) -> [S; 8] { unsafe { let m = m.map(|v| v.to_array()); diff --git a/src/packed_seq.rs b/src/packed_seq.rs index 050ea99..38d7fc6 100644 --- a/src/packed_seq.rs +++ b/src/packed_seq.rs @@ -1292,6 +1292,34 @@ where } } } + + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + use core::mem::transmute as t; + use wide::{u8x16, u32x4}; + + last = unaligned + (len - unaligned) / 16 * 16; + + let pick_low_bytes = u8x16::new([ + 0, 4, 8, 12, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, + ]); + + for i in (unaligned..last).step_by(16) { + unsafe { + let ascii = seq.as_ptr().add(i).cast::().read_unaligned(); + let bits: u32x4 = t((ascii >> 1) & u8x16::splat(3)); + let bits = bits | (bits >> 6); + let bits = bits | (bits >> 12); + let packed: [u32; 4] = + t(t::<_, u8x16>(bits).shuffle_zeroing(pick_low_bytes)); + self.seq + .get_unchecked_mut(idx..(idx + 4)) + .copy_from_slice(&packed[0].to_le_bytes()); + idx += 4; + self.len += 16; + } + } + } } if B == 1 { @@ -1384,6 +1412,33 @@ where } } } + + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + use wide::u8x16; + + last = unaligned + (len - unaligned) / 16 * 16; + + let upper_mask = u8x16::splat(!(b'a' - b'A')); + let two_bits_mask = u8x16::splat(6); + let table = u8x16::new(*b"AxCxTxGxxxxxxxxx"); + + for i in (unaligned..last).step_by(16) { + unsafe { + let chars = seq.as_ptr().add(i).cast::().read_unaligned(); + let chars = chars & upper_mask; + let lossy_encoded = chars & two_bits_mask; + let lookup = table.shuffle_zeroing(lossy_encoded); + let packed_bytes = !(chars.simd_eq(lookup).to_bitmask() as u16); + + self.seq + .get_unchecked_mut(idx..(idx + 2)) + .copy_from_slice(&packed_bytes.to_le_bytes()); + idx += 2; + self.len += 16; + } + } + } } let mut packed_byte = 0;