Merge pull request #103 from smarco/development

Merge AVX2/512 extend from development to main
smarco · Jan 16, 2025 · 787a051 · 787a051
2 parents 42f8bf3 + baf62cf
commit 787a051
Show file tree

Hide file tree

Showing 4 changed files with 644 additions and 64 deletions.
diff --git a/wavefront/wavefront_extend.c b/wavefront/wavefront_extend.c
@@ -54,11 +54,7 @@ void wavefront_extend_end2end_dispatcher_seq(
   wavefront_sequences_t* const seqs = &wf_aligner->sequences;
   // Check the sequence mode
   if (seqs->mode == wf_sequences_ascii) {
-//#if __AVX2__ // TODO
-//    wavefront_extend_matches_packed_end2end_avx2(wf_aligner,mwavefront,lo,hi);
-//#else
     wavefront_extend_matches_packed_end2end(wf_aligner,mwavefront,lo,hi);
-//#endif
   } else {
     wf_offset_t dummy;
     wavefront_extend_matches_custom(wf_aligner,mwavefront,score,lo,hi,false,&dummy);

diff --git a/wavefront/wavefront_extend_kernels.c b/wavefront/wavefront_extend_kernels.c
@@ -35,6 +35,7 @@
 
 #include "wavefront_extend_kernels.h"
 #include "wavefront_termination.h"
+#include "wavefront_extend_kernels_avx.h"
 
 #if __BYTE_ORDER == __LITTLE_ENDIAN
 #define wavefront_extend_matches_kernel wavefront_extend_matches_kernel_blockwise
@@ -63,6 +64,7 @@ FORCE_INLINE wf_offset_t wavefront_extend_matches_kernel_charwise(
   // Return extended offset
   return offset;
 }
+
 FORCE_INLINE wf_offset_t wavefront_extend_matches_kernel_blockwise(
     wavefront_aligner_t* const wf_aligner,
     const int k,
@@ -88,6 +90,7 @@ FORCE_INLINE wf_offset_t wavefront_extend_matches_kernel_blockwise(
   // Return extended offset
   return offset;
 }
+
 /*
  * Wavefront-Extend Inner Kernels
  *   Wavefront offset extension comparing characters
@@ -100,42 +103,67 @@ FORCE_NO_INLINE void wavefront_extend_matches_packed_end2end(
     wavefront_t* const mwavefront,
     const int lo,
     const int hi) {
-  wf_offset_t* const offsets = mwavefront->offsets;
-  int k;
-  for (k=lo;k<=hi;++k) {
-    // Fetch offset
-    const wf_offset_t offset = offsets[k];
-    if (offset == WAVEFRONT_OFFSET_NULL) continue;
-    // Extend offset
-    offsets[k] = wavefront_extend_matches_kernel(wf_aligner,k,offset);
-  }
+  #if __AVX2__ &&  __BYTE_ORDER == __LITTLE_ENDIAN
+    #if __AVX512CD__ && __AVX512VL__
+      wavefront_extend_matches_packed_end2end_avx512(wf_aligner, mwavefront, lo, hi);
+    #else
+      wavefront_extend_matches_packed_end2end_avx2(wf_aligner, mwavefront, lo, hi);
+    #endif
+  #else
+    wf_offset_t* const offsets = mwavefront->offsets;
+    int k;
+    for (k=lo;k<=hi;++k) {
+      // Fetch offset
+      const wf_offset_t offset = offsets[k];
+      if (offset == WAVEFRONT_OFFSET_NULL) continue;
+      // Extend offset
+      offsets[k] = wavefront_extend_matches_kernel(wf_aligner,k,offset);
+    }
+  #endif
 }
+
 FORCE_NO_INLINE wf_offset_t wavefront_extend_matches_packed_end2end_max(
     wavefront_aligner_t* const wf_aligner,
     wavefront_t* const mwavefront,
     const int lo,
     const int hi) {
-  wf_offset_t* const offsets = mwavefront->offsets;
-  wf_offset_t max_antidiag = 0;
-  int k;
-  for (k=lo;k<=hi;++k) {
-    // Fetch offset
-    const wf_offset_t offset = offsets[k];
-    if (offset == WAVEFRONT_OFFSET_NULL) continue;
-    // Extend offset
-    offsets[k] = wavefront_extend_matches_kernel(wf_aligner,k,offset);
-    // Compute max
-    const wf_offset_t antidiag = WAVEFRONT_ANTIDIAGONAL(k,offsets[k]);
-    if (max_antidiag < antidiag) max_antidiag = antidiag;
-  }
-  return max_antidiag;
+  #if __AVX2__ &&  __BYTE_ORDER == __LITTLE_ENDIAN
+    #if __AVX512CD__ && __AVX512VL__
+      return wavefront_extend_matches_packed_end2end_max_avx512(wf_aligner, mwavefront, lo, hi);
+    #else
+      return wavefront_extend_matches_packed_end2end_max_avx2(wf_aligner, mwavefront, lo, hi);
+    #endif
+  #else
+    wf_offset_t* const offsets = mwavefront->offsets;
+    wf_offset_t max_antidiag = 0;
+    int k;
+    for (k=lo;k<=hi;++k) {
+      // Fetch offset
+      const wf_offset_t offset = offsets[k];
+      if (offset == WAVEFRONT_OFFSET_NULL) continue;
+      // Extend offset
+      offsets[k] = wavefront_extend_matches_kernel(wf_aligner,k,offset);
+      // Compute max
+      const wf_offset_t antidiag = WAVEFRONT_ANTIDIAGONAL(k,offsets[k]);
+      if (max_antidiag < antidiag) max_antidiag = antidiag;
+    }
+    return max_antidiag;
+  #endif 
 }
+
 FORCE_NO_INLINE bool wavefront_extend_matches_packed_endsfree(
     wavefront_aligner_t* const wf_aligner,
     wavefront_t* const mwavefront,
     const int score,
     const int lo,
     const int hi) {
+  #if __AVX2__ &&  __BYTE_ORDER == __LITTLE_ENDIAN
+    #if __AVX512CD__ && __AVX512VL__
+      return wavefront_extend_matches_packed_endsfree_avx512(wf_aligner, mwavefront, score, lo, hi);
+    #else
+      return wavefront_extend_matches_packed_endsfree_avx2(wf_aligner, mwavefront, score, lo, hi);
+    #endif
+  #else
   // Parameters
   wf_offset_t* const offsets = mwavefront->offsets;
   int k;
@@ -162,6 +190,7 @@ FORCE_NO_INLINE bool wavefront_extend_matches_packed_endsfree(
   }
   // Alignment not finished
   return false;
+  #endif
 }
 /*
  * Wavefront-Extend Inner Kernel (Custom match function)