Improve: hybrid bench sort performance

ashvardanian · Jan 5, 2025 · 9880f26 · 9880f26
1 parent 455508f
commit 9880f26
Showing 1 changed file with 62 additions and 6 deletions.
diff --git a/scripts/bench_sort.cpp b/scripts/bench_sort.cpp
@@ -85,9 +85,37 @@ static idx_t hybrid_sort_cpp(strings_t const &strings, sz_u64_t *order) {
         return *(uint32_t *)(i_bytes + offset_in_word) < *(uint32_t *)(j_bytes + offset_in_word);
     });
 
-    for (size_t i = 0; i != strings.size(); ++i) std::memset((char *)&order[i] + offset_in_word, 0, 4ul);
+    const auto extract_bytes = [](sz_u64_t v) -> uint32_t {
+        char *bytes = (char *)&v;
+        return *(uint32_t *)(bytes + offset_in_word);
+    };
+
+    if (strings.size() >= 2) {
+        size_t prev_index = 0;
+        uint64_t prev_bytes = extract_bytes(order[0]);
+
+        for (size_t i = 1; i < strings.size(); ++i) {
+            uint32_t bytes = extract_bytes(order[i]);
+            if (bytes != prev_bytes) {
+                std::sort(order + prev_index, order + i, [&](sz_u64_t i, sz_u64_t j) {
+                    // Assumes: offset_in_word==4
+                    sz_size_t i_index = i & 0xFFFF'FFFF;
+                    sz_size_t j_index = j & 0xFFFF'FFFF;
+                    return strings[i_index] < strings[j_index];
+                });
+                prev_index = i;
+                prev_bytes = bytes;
+            }
+        }
 
-    std::sort(order, order + strings.size(), [&](sz_u64_t i, sz_u64_t j) { return strings[i] < strings[j]; });
+        std::sort(order + prev_index, order + strings.size(), [&](sz_u64_t i, sz_u64_t j) {
+            sz_size_t i_index = i & 0xFFFF'FFFF;
+            sz_size_t j_index = j & 0xFFFF'FFFF;
+            return strings[i_index] < strings[j_index];
+        });
+    }
+
+    for (size_t i = 0; i != strings.size(); ++i) std::memset((char *)&order[i] + offset_in_word, 0, 4ul);
 
     return strings.size();
 }
@@ -109,9 +137,37 @@ static idx_t hybrid_stable_sort_cpp(strings_t const &strings, sz_u64_t *order) {
         return *(uint32_t *)(i_bytes + offset_in_word) < *(uint32_t *)(j_bytes + offset_in_word);
     });
 
-    for (size_t i = 0; i != strings.size(); ++i) std::memset((char *)&order[i] + offset_in_word, 0, 4ul);
+    const auto extract_bytes = [](sz_u64_t v) -> uint32_t {
+        char *bytes = (char *)&v;
+        return *(uint32_t *)(bytes + offset_in_word);
+    };
+
+    if (strings.size() >= 2) {
+        size_t prev_index = 0;
+        uint64_t prev_bytes = extract_bytes(order[0]);
+
+        for (size_t i = 1; i < strings.size(); ++i) {
+            uint32_t bytes = extract_bytes(order[i]);
+            if (bytes != prev_bytes) {
+                std::stable_sort(order + prev_index, order + i, [&](sz_u64_t i, sz_u64_t j) {
+                    // Assumes: offset_in_word==4
+                    sz_size_t i_index = i & 0xFFFF'FFFF;
+                    sz_size_t j_index = j & 0xFFFF'FFFF;
+                    return strings[i_index] < strings[j_index];
+                });
+                prev_index = i;
+                prev_bytes = bytes;
+            }
+        }
+
+        std::stable_sort(order + prev_index, order + strings.size(), [&](sz_u64_t i, sz_u64_t j) {
+            sz_size_t i_index = i & 0xFFFF'FFFF;
+            sz_size_t j_index = j & 0xFFFF'FFFF;
+            return strings[i_index] < strings[j_index];
+        });
+    }
 
-    std::stable_sort(order, order + strings.size(), [&](sz_u64_t i, sz_u64_t j) { return strings[i] < strings[j]; });
+    for (size_t i = 0; i != strings.size(); ++i) std::memset((char *)&order[i] + offset_in_word, 0, 4ul);
 
     return strings.size();
 }
@@ -204,7 +260,7 @@ int main(int argc, char const **argv) {
         });
         expect_sorted(strings, permute_new);
 
-#if __linux__ && defined(_GNU_SOURCE) & !defined(__BIONIC__)
+#if __linux__ && defined(_GNU_SOURCE) && !defined(__BIONIC__)
         bench_permute("qsort_r", strings, permute_new, [](strings_t const &strings, permute_t &permute) {
             sz_sequence_t array;
             array.order = permute.data();
@@ -248,4 +304,4 @@ int main(int argc, char const **argv) {
     }
 
     return 0;
-}
+}