Merge pull request #37 from UM-Applied-Algorithms-Lab/version-0.2.1

Sawwave · web-flow · commit 0752671052ff · 2025-01-10T16:38:55.000-07:00
Version 0.2.1
diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@
 
 Generates an Fm-Index of a given biological sequence text (Fasta or Fastq file), and implements Locate() and Search() functionalities.
 
-AWRY is a port of a state-of-the-art, fastest in its class FM-index implementation (https://doi.org/10.1186/s13015-021-00204-6). AWRY supports parallelized searching, with parallel_count() and parallel_locate() functions.
+AWRY is a port of a state-of-the-art, fastest in its class FM-index implementation (https://doi.org/10.1186/s13015-021-00204-6). AWRY supports parallelized searching, with parallel_count() and parallel_locate() functions. As AWRY is powered by custom SIMD operations, the library currently supports x86_64 and aarch64 machines with AVX2 and ARM Neon.
 
 AWRY supports DNA, RNA, and protein alphabets, and is able to search at lightning speed by leveraging SIMD vectorization and multithreading over collections of queries. These alphabets are case-insensitive- `A` and `a` are considered the same symbol under the hood.
 
diff --git a/src/bwt.rs b/src/bwt.rs
@@ -1,7 +1,9 @@
+
+
 use crate::{
     alphabet::{Symbol, SymbolAlphabet},
     search::SearchPtr,
-    simd_instructions::{Vec256, SimdVec256},
+    simd_instructions::{masked_popcount, simd_and, simd_andnot, simd_or, Vec256},
 };
 use mem_dbg::MemSize;
 use serde::{Deserialize, Serialize};
@@ -112,21 +114,21 @@ impl NucleotideBwtBlock {
     pub  (crate) fn global_occurrence(&self, local_query_position: u64, symbol: &Symbol) -> u64 {
         unsafe{
         let milestone_count = self.milestone(&symbol);
-        let vec0 = SimdVec256::from(*self.bit_vectors.get_unchecked(0));
-        let vec1 = SimdVec256::from(*self.bit_vectors.get_unchecked(1));
-        let vec2 = SimdVec256::from(*self.bit_vectors.get_unchecked(2));
+        let vec0 = self.bit_vectors.get_unchecked(0).to_simd();
+        let vec1 = self.bit_vectors.get_unchecked(1).to_simd();
+        let vec2 = self.bit_vectors.get_unchecked(2).to_simd();
         let occurrence_vector = match &symbol.index() {
-            1 => vec2.and(&vec1), //A:    0b110
-            2 => vec2.and(&vec0), //C:    0b101
-            3 => vec1.and(&vec0), //G:    0b011
-            4 => vec2.andnot(&vec0.andnot(&vec1)), //N:    0b010
-            5 => vec2.andnot(&vec1.andnot(&vec0)), //T:    0b001
+            1 => simd_and(vec1, vec2),  //A:    0b110
+            2 => simd_and(vec0, vec2), //C:    0b101
+            3 => simd_and(vec0, vec1), //G:    0b011
+            4 => simd_andnot(vec2, simd_andnot(vec0, vec1)), //N:    0b010
+            5 => simd_andnot(vec2, simd_andnot(vec1, vec0)), //T:    0b001
             _ => {
                 panic!("illegal letter index given in global occurrence function symbol idx given: {}", symbol.index());
             } //assume every other character is an N, since it's illegal to search for a sentinel
         };
 
-        let popcount = occurrence_vector.masked_popcount(local_query_position);
+        let popcount = masked_popcount(occurrence_vector, local_query_position);
 
         return milestone_count + popcount as u64;
     }
@@ -225,44 +227,44 @@ impl AminoBwtBlock {
     /// The occurrence function uses the milestone value and the masked occurrenc vector to
     /// determine how many instances of the given character were before this position.
     #[inline]
-    pub  (crate) fn global_occurrence(&self, local_query_position: SearchPtr, symbol: &Symbol) -> SearchPtr {
+    pub (crate) fn global_occurrence(&self, local_query_position: SearchPtr, symbol: &Symbol) -> SearchPtr {
         let milestone_count = self.milestone(symbol);
 
         unsafe{
-            let vec0 = SimdVec256::from(*self.bit_vectors.get_unchecked(0));
-            let vec1 = SimdVec256::from(*self.bit_vectors.get_unchecked(1));
-            let vec2 = SimdVec256::from(*self.bit_vectors.get_unchecked(2));
-            let vec3 = SimdVec256::from(*self.bit_vectors.get_unchecked(3));   
-            let vec4 = SimdVec256::from(*self.bit_vectors.get_unchecked(4));
+            let vec0 = self.bit_vectors.get_unchecked(0).to_simd();
+            let vec1 = self.bit_vectors.get_unchecked(1).to_simd();
+            let vec2 = self.bit_vectors.get_unchecked(2).to_simd();
+            let vec3 = self.bit_vectors.get_unchecked(3).to_simd();   
+            let vec4 = self.bit_vectors.get_unchecked(4).to_simd();
             let occurrence_vector = match symbol.index() {
-                1 => vec3.and(&vec4.andnot(&vec2)), //A:    0b01100
-                2 => vec3.andnot(&vec2).and(&vec1.and(&vec0)), //C:    0b10111
-                3 => vec1.and(&vec4.andnot(&vec0)), //D:    0b00011
-                4 => vec4.andnot(&vec2.and(&vec1)), //E: 0b00110
-                5 => vec0.andnot(&vec3).and(&vec2.and(&vec1)), //F:    0b11110
-                6 => vec2.andnot(&vec0.andnot(&vec4)), //G:    0b11010
-                7 => vec2.andnot(&vec3).and(&vec1.and(&vec0)), //H: 0b11011
-                8 => vec2.andnot(&vec1.andnot(&vec4)), //I:    0b11001
-                9 => vec3.andnot(&vec1.andnot(&vec4)), //K:    0b10101
-                10 => vec1.andnot(&vec0.andnot(&vec4)), //L:    0b11100
-                11 => vec1.andnot(&vec3).and(&vec2.and(&vec0)), //M:    0b11101
-                12 => vec0.or(&vec1).andnot(&vec2.andnot(&vec3)), //N:    0b01000
-                13 => vec3.and(&vec4.andnot(&vec0)), //P:    0b01001,
-                14 => vec3.or(&vec1).andnot(&vec0.andnot(&vec2)), //Q:    0b00100
-                15 => vec3.andnot(&vec2.andnot(&vec4)), //R:    0b10011
-                16 => vec3.and(&vec4.andnot(&vec1)), //S:    0b01010
-                17 => vec2.and(&vec4.andnot(&vec0)), //T:    0b00101
-                18 => vec3.andnot(&vec0.andnot(&vec4)), //V:    0b10110
-                19 => vec3.or(&vec2).andnot(&vec1.andnot(&vec0)), //W:    0b00001
-                20 => vec3.and(&vec2).and(&vec1.and(&vec0)), //Ambiguity character X:  0b11111
-                21 => vec0.or(&vec2).andnot(&vec3.andnot(&vec1)), //Y:    0b00010
+                1 => simd_and(vec2,simd_andnot(vec4, vec3)), //A:    0b01100
+                2 => simd_andnot(vec3, simd_and(simd_and(vec0, vec1), vec2)), //C:    0b10111
+                3 => simd_andnot(vec4, simd_and(vec0, vec1)), //D:    0b00011
+                4 => simd_andnot(vec4, simd_and(vec1, vec2)), //E: 0b00110
+                5 => simd_andnot(vec0, simd_and(simd_and(vec1, vec2), vec3)), //F:    0b11110
+                6 => simd_andnot(vec2, simd_andnot(vec0, vec4)), //G:    0b11010
+                7 => simd_andnot(vec2, simd_and(vec0, simd_and(vec1, vec3))), //H: 0b11011
+                8 => simd_andnot(vec2, simd_andnot(vec1, vec4)), //I:    0b11001
+                9 => simd_andnot(vec1, simd_andnot(vec3, vec4)), //K:    0b10101
+                10 => simd_andnot(vec1, simd_andnot(vec0, vec4)), //L:    0b11100
+                11 => simd_andnot(vec1, simd_and(vec3, simd_and(vec2, vec0))), //M:    0b11101
+                12 => simd_andnot(simd_or(vec0, vec1), simd_andnot(vec2, vec3)), //N:    0b01000
+                13 => simd_and(vec3, simd_andnot(vec4, vec0)), //P:    0b01001,
+                14 => simd_andnot(simd_or(vec0, vec1), simd_andnot(vec3, vec2)),//Q:    0b00100
+                15 => simd_andnot(vec2, simd_andnot(vec3, vec4)), //R:    0b10011
+                16 => simd_and(vec1, simd_andnot(vec4, vec3)), //S:    0b01010
+                17 => simd_and(vec0, simd_andnot(vec4, vec2)), //T:    0b00101
+                18 => simd_andnot(vec3, simd_andnot(vec0, vec4)), //V:    0b10110
+                19 => simd_andnot(simd_or(vec1, vec2), simd_andnot(vec3, vec0)), //W:    0b00001
+                20 => simd_and(simd_and(vec0, vec1), simd_and(vec2, vec3)), //Ambiguity character X:  0b11111
+                21 => simd_andnot(simd_or(vec0, vec2), simd_andnot(vec3, vec1)), //Y:    0b00010
                 // 0b00000 is sentinel, but since you can't search for sentinels, it is not included here.
                 _ => {
                     panic!("illegal letter index given in global occurrence function");
                 }
         };
         
-        let popcount = occurrence_vector.masked_popcount(local_query_position);
+        let popcount = masked_popcount(occurrence_vector, local_query_position);
         
         return milestone_count + popcount as u64;
         }
diff --git a/src/fm_index.rs b/src/fm_index.rs
@@ -390,7 +390,9 @@ impl FmIndex {
             let mut search_range =
                 SearchRange::new(self, Symbol::new_index(alphabet, final_query_index));
             for query_char in query.chars().rev().skip(1) {
-                if !search_range.is_empty() {
+                if search_range.is_empty() {
+                    break;
+                } else {
                     let query_symbol = Symbol::new_ascii(alphabet, query_char);
                     search_range = self.update_range_with_symbol(search_range, query_symbol);
                 }
@@ -404,10 +406,14 @@ impl FmIndex {
                 .rev()
                 .skip(self.kmer_lookup_table.kmer_len() as usize)
             {
-                search_range = self.update_range_with_symbol(
-                    search_range,
-                    Symbol::new_ascii(self.alphabet(), query_char),
-                );
+                if search_range.is_empty() {
+                    break;
+                } else {
+                    search_range = self.update_range_with_symbol(
+                        search_range,
+                        Symbol::new_ascii(self.alphabet(), query_char),
+                    );
+                }
             }
 
             return search_range;
@@ -429,9 +435,7 @@ impl FmIndex {
     ///     println!("count: {}", count);
     /// }
     /// ```
-    pub fn parallel_count<'a>(&self, queries: impl ParallelIterator<Item= &'a str>) -> Vec<u64> 
-    
-    {
+    pub fn parallel_count<'a>(&self, queries: impl ParallelIterator<Item = &'a str>) -> Vec<u64> {
         queries
             .into_par_iter()
             .map(|query| self.count_string(&query))
@@ -457,9 +461,8 @@ impl FmIndex {
     /// ```
     pub fn parallel_locate<'a>(
         &self,
-        queries: impl ParallelIterator<Item= &'a str>,
-    ) -> Vec<Vec<LocalizedSequencePosition>>
-    {
+        queries: impl ParallelIterator<Item = &'a str>,
+    ) -> Vec<Vec<LocalizedSequencePosition>> {
         queries
             .into_par_iter()
             .map(|query| self.locate_string(query))
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,3 +1,4 @@
+
 pub mod alphabet;
 pub mod bwt;
 pub mod compressed_suffix_array;
diff --git a/src/search.rs b/src/search.rs
@@ -60,6 +60,7 @@ impl SearchRange {
     }
 
     ///gets the number of elements represented by the search range
+    #[inline]
     pub  (crate) fn len(&self) -> SearchPtr {
         match self.is_empty() {
             true => 0,
@@ -68,6 +69,7 @@ impl SearchRange {
     }
 
     /// Returns an interator over the BWT positions corresponding to this search range
+    #[inline]
     pub  (crate) fn range_iter(&self) -> core::ops::Range<SearchPtr> {
         match self.is_empty() {
             true => 0..0,
diff --git a/src/simd_instructions.rs b/src/simd_instructions.rs

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+`
`1`	`2`	`pub mod alphabet;`
`2`	`3`	`pub mod bwt;`
`3`	`4`	`pub mod compressed_suffix_array;`
Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,7 @@ impl SearchRange {`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`///gets the number of elements represented by the search range`
	`63`	`+ #[inline]`
`63`	`64`	`pub (crate) fn len(&self) -> SearchPtr {`
`64`	`65`	`match self.is_empty() {`
`65`	`66`	`true => 0,`
`@@ -68,6 +69,7 @@ impl SearchRange {`
`68`	`69`	`}`
`69`	`70`
`70`	`71`	`/// Returns an interator over the BWT positions corresponding to this search range`
	`72`	`+ #[inline]`
`71`	`73`	`pub (crate) fn range_iter(&self) -> core::ops::Range<SearchPtr> {`
`72`	`74`	`match self.is_empty() {`
`73`	`75`	`true => 0..0,`