Simd128 paths for transpose and push_ascii

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"

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());

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::<u8x16>().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::<u8x16>().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;