Add: OpenMP score_diagonally

Author: ashvardanian

include/stringzilla/similarity.hpp

@@ -8,8 +8,7 @@
  *  - `sz::openmp::levenshtein_distances` & `sz::openmp::levenshtein_distances_utf8` for Levenshtein edit-distances.
  *  - `sz::openmp::needleman_wunsch_score` for weighted Needleman-Wunsch global alignment.
  *
- *  Those are mostly providing specialized overloads of the @b `sz::openmp::global_score_kernel` template.
- *
+ *  Those are mostly providing specialized overloads of the @b `sz::openmp::score_diagonally` template.
  */
 #ifndef STRINGZILLA_SIMILARITY_HPP_
 #define STRINGZILLA_SIMILARITY_HPP_
@@ -25,23 +24,26 @@ struct uniform_substitution_cost_t {
 };
 
 /**
- *  @brief  Levenshtein edit-distance computation using skewed diagonal order evaluation on CPU.
- *          Type-agnostic, unlike `sz_levenshtein_distance`, and annotated for OpenMP parallelization.
+ *  @brief  Alignment Score and Edit Distance algorithm evaluating the Dynamic Programming matrix
+ *          @b three skewed (reverse) diagonals at a time on a CPU, leveraging OpenMP for parallelization.
+ *  @sa     sz_levenshtein_distance, sz_levenshtein_distance_utf8, sz_needleman_wunsch_score
  *
  *  @param[in] first The first string.
  *  @param[in] first_length The length of the first string.
  *  @param[in] second The second string.
  *  @param[in] second_length The length of the second string.
  *
- *  There are smarter algorithms for computing the Levenshtein distance, but this one is extremely flexible.
+ *  There are smarter algorithms for computing the Levenshtein distance, mostly based on bit-level operations.
+ *  Those, however, don't generalize well to arbitrary length inputs or non-uniform substitution costs.
+ *  This algorithm provides a more flexible baseline implementation for future SIMD and GPGPU optimizations.
  */
 template <                                                         //
     typename char_type_ = char,                                    //
     typename distance_type_ = sz_size_t,                           //
     typename get_substitution_cost_ = uniform_substitution_cost_t, //
     sz_error_cost_t gap_cost_ = 1                                  //
     >
-sz_status_t global_score_kernel(                         //
+sz_status_t score_diagonally(                            //
     char_type_ const *shorter, sz_size_t shorter_length, //
     char_type_ const *longer, sz_size_t longer_length,   //
     get_substitution_cost_ &&get_substitution_cost,      //
@@ -95,13 +97,16 @@ sz_status_t global_score_kernel(                         //
         for (sz_size_t offset_in_diagonal = 1; offset_in_diagonal + 1 < next_diagonal_length; ++offset_in_diagonal) {
             // ? Note that here we are still traversing both buffers in the same order,
             // ? because the shorter string has been reversed into `shorter_reversed`.
-            char_type_ shorter_char = shorter_reversed[shorter_length - next_diagonal_index + offset_in_diagonal - 1];
-            char_type_ longer_char = longer[offset_in_diagonal];
+            char_type_ shorter_char = shorter_reversed[shorter_length - next_diagonal_index + offset_within_diagonal];
+            char_type_ longer_char = longer[offset_in_diagonal - 1];
             sz_error_cost_t cost_of_substitution = get_substitution_cost(shorter_char, longer_char);
             distance_type_ cost_if_substitution = previous_distances[offset_in_diagonal] + cost_of_substitution;
-            distance_type_ cost_if_deletion_or_insertion =
-                gap_cost_ +
-                sz_min_of_two(current_distances[offset_in_diagonal - 1], current_distances[offset_in_diagonal]);
+            distance_type_ cost_if_deletion_or_insertion =     //
+                sz_min_of_two(                                 //
+                    current_distances[offset_in_diagonal - 1], //
+                    current_distances[offset_in_diagonal]      //
+                    ) +
+                gap_cost_;
             next_distances[offset_in_diagonal] = sz_min_of_two(cost_if_deletion_or_insertion, cost_if_substitution);
         }
         // Don't forget to populate the first row and the first column of the Levenshtein matrix.
@@ -116,19 +121,50 @@ sz_status_t global_score_kernel(                         //
     // Now let's handle the anti-diagonal band of the matrix, between the top and bottom triangles.
     for (; next_diagonal_index != longer_dim; ++next_diagonal_index) {
         sz_size_t const next_diagonal_length = shorter_dim;
-        for (sz_size_t offset_in_diagonal = 1; offset_in_diagonal + 1 < next_diagonal_length; ++offset_in_diagonal) {
-            char_type_ shorter_char = shorter[shorter_length - 1 - i]; // ! Walks in reverse order.
-            char_type_ longer_char = longer[next_diagonal_index - n + i];
+#pragma omp simd
+        for (sz_size_t offset_in_diagonal = 0; offset_in_diagonal + 1 < next_diagonal_length; ++offset_in_diagonal) {
+            char_type_ shorter_char = shorter_reversed[shorter_length - shorter_dim + offset_in_diagonal + 1];
+            char_type_ longer_char = longer[next_diagonal_index - shorter_dim + offset_in_diagonal];
             sz_error_cost_t cost_of_substitution = get_substitution_cost(shorter_char, longer_char);
-            distance_type_ cost_if_substitution = previous_distances[i] + cost_of_substitution;
-            distance_type_ cost_if_deletion_or_insertion =
-                sz_min_of_two(current_distances[i], current_distances[i + 1]) + gap_cost_;
+            distance_type_ cost_if_substitution = previous_distances[offset_in_diagonal] + cost_of_substitution;
+            distance_type_ cost_if_deletion_or_insertion =    //
+                sz_min_of_two(                                //
+                    current_distances[offset_in_diagonal],    //
+                    current_distances[offset_in_diagonal + 1] //
+                    ) +
+                gap_cost_;
+            next_distances[i] = sz_min_of_two(cost_if_deletion_or_insertion, cost_if_substitution);
+        }
+        next_distances[next_diagonal_length - 1] = next_diagonal_index;
+        // Perform a circular rotation of those buffers, to reuse the memory, this time, with a shift,
+        // dropping the first element in the current array.
+        distance_type_ *temporary = previous_distances;
+        previous_distances = current_distances + 1; // ! Note how we shift forward here
+        current_distances = next_distances;
+        next_distances = temporary;
+    }
+
+    // Now let's handle the bottom-right triangle of the matrix.
+    for (; next_diagonal_index != longer_dim; ++next_diagonal_index) {
+        sz_size_t const next_diagonal_length = diagonals_count - next_diagonal_index;
+#pragma omp simd
+        for (sz_size_t offset_in_diagonal = 0; offset_in_diagonal < next_diagonal_length; ++offset_in_diagonal) {
+            char_type_ shorter_char = shorter_reversed[shorter_length - shorter_dim + offset_in_diagonal + 1];
+            char_type_ longer_char = longer[next_diagonal_index - shorter_dim + offset_in_diagonal];
+            sz_error_cost_t cost_of_substitution = get_substitution_cost(shorter_char, longer_char);
+            distance_type_ cost_if_substitution = previous_distances[offset_in_diagonal] + cost_of_substitution;
+            distance_type_ cost_if_deletion_or_insertion =    //
+                sz_min_of_two(                                //
+                    current_distances[offset_in_diagonal],    //
+                    current_distances[offset_in_diagonal + 1] //
+                    ) +
+                gap_cost_;
             next_distances[i] = sz_min_of_two(cost_if_deletion_or_insertion, cost_if_substitution);
         }
         // Perform a circular rotation of those buffers, to reuse the memory, this time, with a shift,
         // dropping the first element in the current array.
         distance_type_ *temporary = previous_distances;
-        previous_distances = current_distances + 1;
+        previous_distances = current_distances + 1; // ! Note how we shift forward here
         current_distances = next_distances;
         next_distances = temporary;
     }