Merge pull request #179 from sterrettm2/qsort-openmp

r-devulap · web-flow · commit b12fe4d04475 · 2025-03-28T10:26:09.000-07:00
Adds OpenMP to qsort, should also improve test speed a bit
diff --git a/lib/meson.build b/lib/meson.build
@@ -34,7 +34,7 @@ if cpp.has_argument('-march=icelake-client')
       'x86simdsort-icl.cpp',
       ),
     include_directories : [src],
-    cpp_args : ['-march=icelake-client'],
+    cpp_args : ['-march=icelake-client', openmpflags],
     gnu_symbol_visibility : 'inlineshidden',
     )
 endif
@@ -45,7 +45,7 @@ if cancompilefp16
       'x86simdsort-spr.cpp',
       ),
     include_directories : [src],
-    cpp_args : ['-march=sapphirerapids'],
+    cpp_args : ['-march=sapphirerapids', openmpflags],
     gnu_symbol_visibility : 'inlineshidden',
     )
 endif
diff --git a/src/avx512-16bit-qsort.hpp b/src/avx512-16bit-qsort.hpp
@@ -556,6 +556,7 @@ avx512_qsort_fp16(uint16_t *arr,
 {
     using vtype = zmm_vector<float16>;
 
+    // TODO multithreading support here
     if (arrsize > 1) {
         arrsize_t nan_count = 0;
         if (UNLIKELY(hasnan)) {
@@ -564,11 +565,11 @@ avx512_qsort_fp16(uint16_t *arr,
         }
         if (descending) {
             qsort_<vtype, Comparator<vtype, true>, uint16_t>(
-                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
+                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
         }
         else {
             qsort_<vtype, Comparator<vtype, false>, uint16_t>(
-                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
+                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
         }
         replace_inf_with_nan(arr, arrsize, nan_count, descending);
     }
diff --git a/src/avx512fp16-16bit-qsort.hpp b/src/avx512fp16-16bit-qsort.hpp
@@ -22,7 +22,7 @@ struct zmm_vector<_Float16> {
     using opmask_t = __mmask32;
     static const uint8_t numlanes = 32;
     static constexpr int network_sort_threshold = 128;
-    static constexpr int partition_unroll_factor = 0;
+    static constexpr int partition_unroll_factor = 8;
     static constexpr simd_type vec_type = simd_type::AVX512;
 
     using swizzle_ops = avx512_16bit_swizzle_ops;
diff --git a/src/xss-common-includes.h b/src/xss-common-includes.h
@@ -82,6 +82,11 @@
     0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, \
             21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
 
+#if defined(XSS_USE_OPENMP) && defined(_OPENMP)
+#define XSS_COMPILE_OPENMP
+#include <omp.h>
+#endif
+
 template <class... T>
 constexpr bool always_false = false;
 
diff --git a/src/xss-common-keyvaluesort.hpp b/src/xss-common-keyvaluesort.hpp
@@ -11,11 +11,6 @@
 #include "xss-common-qsort.h"
 #include "xss-network-keyvaluesort.hpp"
 
-#if defined(XSS_USE_OPENMP) && defined(_OPENMP)
-#define XSS_COMPILE_OPENMP
-#include <omp.h>
-#endif
-
 /*
  * Sort all the NAN's to end of the array and return the index of the last elem
  * in the array which is not a nan
diff --git a/src/xss-common-qsort.h b/src/xss-common-qsort.h
@@ -521,8 +521,11 @@ template <typename vtype, int maxN>
 void sort_n(typename vtype::type_t *arr, int N);
 
 template <typename vtype, typename comparator, typename type_t>
-static void
-qsort_(type_t *arr, arrsize_t left, arrsize_t right, arrsize_t max_iters)
+static void qsort_(type_t *arr,
+                   arrsize_t left,
+                   arrsize_t right,
+                   arrsize_t max_iters,
+                   arrsize_t task_threshold)
 {
     /*
      * Resort to std::sort if quicksort isnt making any progress
@@ -559,10 +562,40 @@ qsort_(type_t *arr, arrsize_t left, arrsize_t right, arrsize_t max_iters)
     type_t leftmostValue = comparator::leftmost(smallest, biggest);
     type_t rightmostValue = comparator::rightmost(smallest, biggest);
 
+#ifdef XSS_COMPILE_OPENMP
+    if (pivot != leftmostValue) {
+        bool parallel_left = (pivot_index - left) > task_threshold;
+        if (parallel_left) {
+#pragma omp task
+            qsort_<vtype, comparator>(
+                    arr, left, pivot_index - 1, max_iters - 1, task_threshold);
+        }
+        else {
+            qsort_<vtype, comparator>(
+                    arr, left, pivot_index - 1, max_iters - 1, task_threshold);
+        }
+    }
+    if (pivot != rightmostValue) {
+        bool parallel_right = (right - pivot_index) > task_threshold;
+
+        if (parallel_right) {
+#pragma omp task
+            qsort_<vtype, comparator>(
+                    arr, pivot_index, right, max_iters - 1, task_threshold);
+        }
+        else {
+            qsort_<vtype, comparator>(
+                    arr, pivot_index, right, max_iters - 1, task_threshold);
+        }
+    }
+#else
+    UNUSED(task_threshold);
+
     if (pivot != leftmostValue)
-        qsort_<vtype, comparator>(arr, left, pivot_index - 1, max_iters - 1);
+        qsort_<vtype, comparator>(arr, left, pivot_index - 1, max_iters - 1, 0);
     if (pivot != rightmostValue)
-        qsort_<vtype, comparator>(arr, pivot_index, right, max_iters - 1);
+        qsort_<vtype, comparator>(arr, pivot_index, right, max_iters - 1, 0);
+#endif
 }
 
 template <typename vtype, typename comparator, typename type_t>
@@ -627,8 +660,41 @@ X86_SIMD_SORT_INLINE void xss_qsort(T *arr, arrsize_t arrsize, bool hasnan)
         }
 
         UNUSED(hasnan);
+
+#ifdef XSS_COMPILE_OPENMP
+
+        bool use_parallel = arrsize > 100000;
+
+        if (use_parallel) {
+            // This thread limit was determined experimentally; it may be better for it to be the number of physical cores on the system
+            constexpr int thread_limit = 8;
+            int thread_count = std::min(thread_limit, omp_get_max_threads());
+            arrsize_t task_threshold
+                    = std::max((arrsize_t)100000, arrsize / 100);
+
+            // We use omp parallel and then omp single to setup the threads that will run the omp task calls in qsort_
+            // The omp single prevents multiple threads from running the initial qsort_ simultaneously and causing problems
+            // Note that we do not use the if(...) clause built into OpenMP, because it causes a performance regression for small arrays
+#pragma omp parallel num_threads(thread_count)
+#pragma omp single
+            qsort_<vtype, comparator, T>(arr,
+                                         0,
+                                         arrsize - 1,
+                                         2 * (arrsize_t)log2(arrsize),
+                                         task_threshold);
+        }
+        else {
+            qsort_<vtype, comparator, T>(arr,
+                                         0,
+                                         arrsize - 1,
+                                         2 * (arrsize_t)log2(arrsize),
+                                         std::numeric_limits<arrsize_t>::max());
+        }
+#pragma omp taskwait
+#else
         qsort_<vtype, comparator, T>(
-                arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
+                arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
+#endif
 
         replace_inf_with_nan(arr, arrsize, nan_count, descending);
     }
diff --git a/tests/meson.build b/tests/meson.build
@@ -1,19 +1,26 @@
 libtests = []
 
+if get_option('use_openmp')
+  openmpflags = ['-DXSS_USE_OPENMP=true']
+endif
+
 libtests += static_library('tests_qsort',
   files('test-qsort.cpp', ),
   dependencies: gtest_dep,
   include_directories : [src, lib, utils],
+  cpp_args : [openmpflags],
   )
 
 libtests += static_library('tests_kvsort',
   files('test-keyvalue.cpp', ),
   dependencies: gtest_dep,
   include_directories : [src, lib, utils],
+  cpp_args : [openmpflags],
   )
 
 libtests += static_library('tests_objsort',
   files('test-objqsort.cpp', ),
   dependencies: gtest_dep,
   include_directories : [src, lib, utils],
+  cpp_args : [openmpflags],
   )
diff --git a/tests/test-keyvalue.cpp b/tests/test-keyvalue.cpp
@@ -15,9 +15,13 @@ class simdkvsort : public ::testing::Test {
     simdkvsort()
     {
         std::iota(arrsize.begin(), arrsize.end(), 1);
-        arrsize.push_back(10'000);
-        arrsize.push_back(100'000);
-        arrsize.push_back(1'000'000);
+        std::iota(arrsize_long.begin(), arrsize_long.end(), 1);
+#ifdef XSS_USE_OPENMP
+        // These extended tests are only needed for the OpenMP logic
+        arrsize_long.push_back(10'000);
+        arrsize_long.push_back(100'000);
+        arrsize_long.push_back(1'000'000);
+#endif
 
         arrtype = {"random",
                    "constant",
@@ -32,6 +36,7 @@ class simdkvsort : public ::testing::Test {
     }
     std::vector<std::string> arrtype;
     std::vector<size_t> arrsize = std::vector<size_t>(1024);
+    std::vector<size_t> arrsize_long = std::vector<size_t>(1024);
 };
 
 TYPED_TEST_SUITE_P(simdkvsort);
@@ -168,7 +173,7 @@ TYPED_TEST_P(simdkvsort, test_kvsort_ascending)
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
         bool hasnan = is_nan_test(type);
-        for (auto size : this->arrsize) {
+        for (auto size : this->arrsize_long) {
             std::vector<T1> key = get_array<T1>(type, size);
             std::vector<T2> val = get_array<T2>(type, size);
             std::vector<T1> key_bckp = key;
@@ -199,7 +204,7 @@ TYPED_TEST_P(simdkvsort, test_kvsort_descending)
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
         bool hasnan = is_nan_test(type);
-        for (auto size : this->arrsize) {
+        for (auto size : this->arrsize_long) {
             std::vector<T1> key = get_array<T1>(type, size);
             std::vector<T2> val = get_array<T2>(type, size);
             std::vector<T1> key_bckp = key;
diff --git a/tests/test-qsort.cpp b/tests/test-qsort.cpp
@@ -11,6 +11,14 @@ class simdsort : public ::testing::Test {
     simdsort()
     {
         std::iota(arrsize.begin(), arrsize.end(), 1);
+        std::iota(arrsize_long.begin(), arrsize_long.end(), 1);
+#ifdef XSS_USE_OPENMP
+        // These extended tests are only needed for the OpenMP logic
+        arrsize_long.push_back(10'000);
+        arrsize_long.push_back(100'000);
+        arrsize_long.push_back(1'000'000);
+#endif
+
         arrtype = {"random",
                    "constant",
                    "sorted",
@@ -24,6 +32,7 @@ class simdsort : public ::testing::Test {
     }
     std::vector<std::string> arrtype;
     std::vector<size_t> arrsize = std::vector<size_t>(1024);
+    std::vector<size_t> arrsize_long = std::vector<size_t>(1024);
 };
 
 TYPED_TEST_SUITE_P(simdsort);
@@ -32,7 +41,7 @@ TYPED_TEST_P(simdsort, test_qsort_ascending)
 {
     for (auto type : this->arrtype) {
         bool hasnan = is_nan_test(type);
-        for (auto size : this->arrsize) {
+        for (auto size : this->arrsize_long) {
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
 
             // Ascending order
@@ -54,7 +63,7 @@ TYPED_TEST_P(simdsort, test_qsort_descending)
 {
     for (auto type : this->arrtype) {
         bool hasnan = is_nan_test(type);
-        for (auto size : this->arrsize) {
+        for (auto size : this->arrsize_long) {
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
 
             // Descending order

Original file line number	Diff line number	Diff line change
`@@ -556,6 +556,7 @@ avx512_qsort_fp16(uint16_t *arr,`
`556`	`556`	`{`
`557`	`557`	`using vtype = zmm_vector<float16>;`
`558`	`558`
	`559`	`+ // TODO multithreading support here`
`559`	`560`	`if (arrsize > 1) {`
`560`	`561`	`arrsize_t nan_count = 0;`
`561`	`562`	`if (UNLIKELY(hasnan)) {`
`@@ -564,11 +565,11 @@ avx512_qsort_fp16(uint16_t *arr,`
`564`	`565`	`}`
`565`	`566`	`if (descending) {`
`566`	`567`	`qsort_<vtype, Comparator<vtype, true>, uint16_t>(`
`567`		`- arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));`
	`568`	`+ arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);`
`568`	`569`	`}`
`569`	`570`	`else {`
`570`	`571`	`qsort_<vtype, Comparator<vtype, false>, uint16_t>(`
`571`		`- arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));`
	`572`	`+ arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);`
`572`	`573`	`}`
`573`	`574`	`replace_inf_with_nan(arr, arrsize, nan_count, descending);`
`574`	`575`	`}`