Optimise partial sort to use a single pass

The original partial sort function was a combination of a QuickSelect
pass followed by a QuickSort pass. This can lead to unnecessary
comparisons and movement of data. This update changes the internal
method used by `avx512_partial_qsort` so that it can more efficiently
sort or recurse into each partition as required and without needing to
take a second sweep over the array. This commit also tweaks the
conditional checks in `avx512_qselect` to avoid a redundant comparison.

Author: mosullivan93

src/avx512-16bit-common.h

@@ -316,10 +316,53 @@ static void qselect_16bit_(type_t *arr,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (pos < pivot_index))
-        qselect_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (pos >= pivot_index))
-        qselect_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselect_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != biggest)
+            qselect_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
+}
+
+template <typename vtype, typename type_t>
+static void qselsort_16bit_(type_t *arr,
+                            int64_t pos,
+                            int64_t left,
+                            int64_t right,
+                            int64_t max_iters)
+{
+    /*
+     * Resort to std::sort if quicksort isnt making any progress
+     */
+    if (max_iters <= 0) {
+        std::sort(arr + left, arr + right + 1, comparison_func<vtype>);
+        return;
+    }
+    /*
+     * Base case: use bitonic networks to sort arrays <= 128
+     */
+    if (right + 1 - left <= 128) {
+        sort_128_16bit<vtype>(arr + left, (int32_t)(right + 1 - left));
+        return;
+    }
+
+    type_t pivot = get_pivot_16bit<vtype>(arr, left, right);
+    type_t smallest = vtype::type_max();
+    type_t biggest = vtype::type_min();
+    int64_t pivot_index = partition_avx512<vtype>(
+            arr, left, right + 1, pivot, &smallest, &biggest);
+
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselsort_16bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != smallest)
+            qsort_16bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
+        if (pivot != biggest)
+            qselsort_16bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
 }
 
 #endif // AVX512_16BIT_COMMON

src/avx512-16bit-qsort.hpp

@@ -433,14 +433,44 @@ void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
 {
     int64_t indx_last_elem = arrsize - 1;
     if (UNLIKELY(hasnan)) {
-         indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
     }
     if (indx_last_elem >= k) {
         qselect_16bit_<zmm_vector<float16>, uint16_t>(
             arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
     }
 }
 
+template <>
+void avx512_partial_qsort(int16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_16bit_<zmm_vector<int16_t>, int16_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_partial_qsort(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_16bit_<zmm_vector<uint16_t>, uint16_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (LIKELY(k > 0)) {
+        int64_t indx_last_elem = arrsize - 1;
+        if (UNLIKELY(hasnan)) {
+            indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
+        qselsort_16bit_<zmm_vector<float16>, uint16_t>(
+            arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
+    }
+}
+
 template <>
 void avx512_qsort(int16_t *arr, int64_t arrsize)
 {

src/avx512-32bit-qsort.hpp

@@ -683,10 +683,52 @@ static void qselect_32bit_(type_t *arr,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512_unrolled<vtype, 2>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (pos < pivot_index))
-        qselect_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (pos >= pivot_index))
-        qselect_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselect_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != biggest)
+            qselect_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
+}
+
+template <typename vtype, typename type_t>
+static void qselsort_32bit_(type_t *arr,
+                            int64_t pos,
+                            int64_t left,
+                            int64_t right,
+                            int64_t max_iters)
+{
+    /*
+     * Resort to std::sort if quicksort isnt making any progress
+     */
+    if (max_iters <= 0) {
+        std::sort(arr + left, arr + right + 1);
+        return;
+    }
+    /*
+     * Base case: use bitonic networks to sort arrays <= 128
+     */
+    if (right + 1 - left <= 128) {
+        sort_128_32bit<vtype>(arr + left, (int32_t)(right + 1 - left));
+        return;
+    }
+
+    type_t pivot = get_pivot_32bit<vtype>(arr, left, right);
+    type_t smallest = vtype::type_max();
+    type_t biggest = vtype::type_min();
+    int64_t pivot_index = partition_avx512_unrolled<vtype, 2>(
+            arr, left, right + 1, pivot, &smallest, &biggest);
+
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselsort_32bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != smallest)
+            qsort_32bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
+        if (pivot != biggest)
+            qselsort_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
 }
 
 X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(float *arr, int64_t arrsize)
@@ -737,14 +779,45 @@ void avx512_qselect<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
 {
     int64_t indx_last_elem = arrsize - 1;
     if (UNLIKELY(hasnan)) {
-         indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
     }
     if (indx_last_elem >= k) {
         qselect_32bit_<zmm_vector<float>, float>(
             arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
     }
 }
 
+template <>
+void avx512_partial_qsort<int32_t>(int32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_32bit_<zmm_vector<int32_t>, int32_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_partial_qsort<uint32_t>(uint32_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_32bit_<zmm_vector<uint32_t>, uint32_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_partial_qsort<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (LIKELY(k > 0)) {
+        int64_t indx_last_elem = arrsize - 1;
+        if (UNLIKELY(hasnan)) {
+            indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
+        qselsort_32bit_<zmm_vector<float>, float>(
+            arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
+    }
+}
+
 template <>
 void avx512_qsort<int32_t>(int32_t *arr, int64_t arrsize)
 {

src/avx512-64bit-qsort.hpp

@@ -732,7 +732,7 @@ qsort_64bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
         return;
     }
     /*
-     * Base case: use bitonic networks to sort arrays <= 128
+     * Base case: use bitonic networks to sort arrays <= 256
      */
     if (right + 1 - left <= 256) {
         sort_256_64bit<vtype>(arr + left, (int32_t)(right + 1 - left));
@@ -777,10 +777,53 @@ static void qselect_64bit_(type_t *arr,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512_unrolled<vtype, 8>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (pos < pivot_index))
-        qselect_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (pos >= pivot_index))
-        qselect_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselect_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != biggest)
+            qselect_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
+}
+
+template <typename vtype, typename type_t>
+static void qselsort_64bit_(type_t *arr,
+                            int64_t pos,
+                            int64_t left,
+                            int64_t right,
+                            int64_t max_iters)
+{
+    /*
+     * Resort to std::sort if quicksort isnt making any progress
+     */
+    if (max_iters <= 0) {
+        std::sort(arr + left, arr + right + 1);
+        return;
+    }
+    /*
+     * Base case: use bitonic networks to sort arrays <= 128
+     */
+    if (right + 1 - left <= 128) {
+        sort_128_64bit<vtype>(arr + left, (int32_t)(right + 1 - left));
+        return;
+    }
+
+    type_t pivot = get_pivot_64bit<vtype>(arr, left, right);
+    type_t smallest = vtype::type_max();
+    type_t biggest = vtype::type_min();
+    int64_t pivot_index = partition_avx512_unrolled<vtype, 8>(
+            arr, left, right + 1, pivot, &smallest, &biggest);
+
+    if (pos < pivot_index) {
+        if (pivot != smallest)
+            qselsort_64bit_<vtype>(arr, pos, left, pivot_index - 1, max_iters - 1);
+    } else {
+        if (pivot != smallest)
+            qsort_64bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
+        if (pivot != biggest)
+            qselsort_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
+    }
 }
 
 template <>
@@ -806,14 +849,45 @@ void avx512_qselect<double>(double *arr, int64_t k, int64_t arrsize, bool hasnan
 {
     int64_t indx_last_elem = arrsize - 1;
     if (UNLIKELY(hasnan)) {
-         indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
     }
     if (indx_last_elem >= k) {
         qselect_64bit_<zmm_vector<double>, double>(
             arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
     }
 }
 
+template <>
+void avx512_partial_qsort<int64_t>(int64_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_64bit_<zmm_vector<int64_t>, int64_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_partial_qsort<uint64_t>(uint64_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (arrsize > 1) {
+        qselsort_64bit_<zmm_vector<uint64_t>, uint64_t>(
+                arr, k - 1, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_partial_qsort<double>(double *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (LIKELY(k > 0)) {
+        int64_t indx_last_elem = arrsize - 1;
+        if (UNLIKELY(hasnan)) {
+            indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
+        qselsort_64bit_<zmm_vector<double>, double>(
+            arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
+    }
+}
+
 template <>
 void avx512_qsort<int64_t>(int64_t *arr, int64_t arrsize)
 {

src/avx512-common-qsort.h

@@ -104,16 +104,8 @@ void avx512_qselect(T *arr, int64_t k, int64_t arrsize, bool hasnan = false);
 void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false);
 
 template <typename T>
-inline void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize, bool hasnan = false)
-{
-    avx512_qselect<T>(arr, k - 1, arrsize, hasnan);
-    avx512_qsort<T>(arr, k - 1);
-}
-inline void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false)
-{
-    avx512_qselect_fp16(arr, k - 1, arrsize, hasnan);
-    avx512_qsort_fp16(arr, k - 1);
-}
+void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize, bool hasnan = false);
+void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false);
 
 // key-value sort routines
 template <typename T>

src/avx512fp16-16bit-qsort.hpp

@@ -157,14 +157,27 @@ void avx512_qselect(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)
 {
     int64_t indx_last_elem = arrsize - 1;
     if (UNLIKELY(hasnan)) {
-         indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
     }
     if (indx_last_elem >= k) {
         qselect_16bit_<zmm_vector<_Float16>, _Float16>(
             arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
     }
 }
 
+template <>
+void avx512_partial_qsort(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)
+{
+    if (LIKELY(k > 0)) {
+        int64_t indx_last_elem = arrsize - 1;
+        if (UNLIKELY(hasnan)) {
+            indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
+        qselsort_16bit_<zmm_vector<_Float16>, _Float16>(
+            arr, k - 1, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
+    }
+}
+
 template <>
 void avx512_qsort(_Float16 *arr, int64_t arrsize)
 {