Align qselect parameter meaning with nth_element

The QuickSelect internal method is now phrased such that the position to
be sorted is given as an offset (in the same way that left points to the
first element and right points to the last element). Similarly, the
avx512_qselect method also now uses this interpretation.

Author: mosullivan93

benchmarks/bench_qselect.hpp

@@ -18,7 +18,7 @@ static void avx512_qselect(benchmark::State& state) {
     arr_bkp = arr;
 
     /* Choose random index to make sorted */
-    int k = get_uniform_rand_array<int64_t>(1, ARRSIZE, 1).front();
+    int k = get_uniform_rand_array<int64_t>(1, ARRSIZE - 1, 0).front();
 
     /* call avx512 quickselect */
     for (auto _ : state) {
@@ -42,7 +42,7 @@ static void stdnthelement(benchmark::State& state) {
     arr_bkp = arr;
 
     /* Choose random index to make sorted */
-    int k = get_uniform_rand_array<int64_t>(1, ARRSIZE, 1).front();
+    int k = get_uniform_rand_array<int64_t>(1, ARRSIZE - 1, 0).front();
 
     /* call std::nth_element */
     for (auto _ : state) {

benchmarks/bench_qsortfp16.cpp

@@ -78,7 +78,7 @@ static void avx512_qselect(benchmark::State& state) {
         arr_bkp = arr;
 
         /* Choose random index to make sorted */
-        int k = get_uniform_rand_array<int64_t>(1, ARRSIZE, 1).front();
+        int k = get_uniform_rand_array<int64_t>(1, ARRSIZE - 1, 0).front();
 
         /* call avx512 quickselect */
         for (auto _ : state) {
@@ -110,7 +110,7 @@ static void stdnthelement(benchmark::State& state) {
         arr_bkp = arr;
 
         /* Choose random index to sort until */
-        int k = get_uniform_rand_array<int64_t>(1, ARRSIZE, 1).front();
+        int k = get_uniform_rand_array<int64_t>(1, ARRSIZE - 1, 0).front();
 
         /* call std::nth_element */
         for (auto _ : state) {

src/avx512-16bit-common.h

@@ -261,9 +261,7 @@ X86_SIMD_SORT_INLINE type_t get_pivot_16bit(type_t *arr,
 
 template <typename vtype, typename type_t>
 static void
-qselect_16bit_(type_t *arr, int64_t k,
-               int64_t left, int64_t right, 
-               int64_t max_iters)
+qsort_16bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
 {
     /*
      * Resort to std::sort if quicksort isnt making any progress
@@ -285,15 +283,17 @@ qselect_16bit_(type_t *arr, int64_t k,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (k <= pivot_index))
-        qselect_16bit_<vtype>(arr, k, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (k > pivot_index))
-        qselect_16bit_<vtype>(arr, k, pivot_index, right, max_iters - 1);
+    if (pivot != smallest)
+        qsort_16bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
+    if (pivot != biggest)
+        qsort_16bit_<vtype>(arr, pivot_index, right, max_iters - 1);
 }
 
 template <typename vtype, typename type_t>
 static void
-qsort_16bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
+qselect_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
@@ -315,10 +315,10 @@ qsort_16bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if (pivot != smallest)
-        qsort_16bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
-    if (pivot != biggest)
-        qsort_16bit_<vtype>(arr, pivot_index, right, max_iters - 1);
+    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);
 }
 
 #endif // AVX512_16BIT_COMMON

src/avx512-32bit-qsort.hpp

@@ -628,9 +628,7 @@ X86_SIMD_SORT_INLINE type_t get_pivot_32bit(type_t *arr,
 
 template <typename vtype, typename type_t>
 static void
-qselect_32bit_(type_t *arr, int64_t k,
-               int64_t left, int64_t right,
-               int64_t max_iters)
+qsort_32bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
 {
     /*
      * Resort to std::sort if quicksort isnt making any progress
@@ -652,15 +650,17 @@ qselect_32bit_(type_t *arr, int64_t k,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (k <= pivot_index))
-        qselect_32bit_<vtype>(arr, k, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (k > pivot_index))
-        qselect_32bit_<vtype>(arr, k, pivot_index, right, max_iters - 1);
+    if (pivot != smallest)
+        qsort_32bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
+    if (pivot != biggest)
+        qsort_32bit_<vtype>(arr, pivot_index, right, max_iters - 1);
 }
 
 template <typename vtype, typename type_t>
 static void
-qsort_32bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
+qselect_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
@@ -682,10 +682,10 @@ qsort_32bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if (pivot != smallest)
-        qsort_32bit_<vtype>(arr, left, pivot_index - 1, max_iters - 1);
-    if (pivot != biggest)
-        qsort_32bit_<vtype>(arr, pivot_index, right, max_iters - 1);
+    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);
 }
 
 X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(float *arr, int64_t arrsize)

src/avx512-64bit-qsort.hpp

@@ -403,7 +403,7 @@ qsort_64bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
 
 template <typename vtype, typename type_t>
 static void
-qselect_64bit_(type_t *arr, int64_t k,
+qselect_64bit_(type_t *arr, int64_t pos,
                int64_t left, int64_t right,
                int64_t max_iters)
 {
@@ -427,18 +427,18 @@ qselect_64bit_(type_t *arr, int64_t k,
     type_t biggest = vtype::type_min();
     int64_t pivot_index = partition_avx512<vtype>(
             arr, left, right + 1, pivot, &smallest, &biggest);
-    if ((pivot != smallest) && (k <= pivot_index))
-        qselect_64bit_<vtype>(arr, k, left, pivot_index - 1, max_iters - 1);
-    else if ((pivot != biggest) && (k > pivot_index))
-        qselect_64bit_<vtype>(arr, k, pivot_index, right, max_iters - 1);
+    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);
 }
 
 template <>
 void avx512_qselect<int64_t>(int64_t *arr, int64_t k, int64_t arrsize)
 {
     if (arrsize > 1) {
         qselect_64bit_<zmm_vector<int64_t>, int64_t>(
-            arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
     }
 }
 

src/avx512-common-qsort.h

@@ -95,14 +95,14 @@ void avx512_qselect(T *arr, int64_t k, int64_t arrsize);
 void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize);
 
 template <typename T>
-void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize) {
-    avx512_qselect<T>(arr, k, arrsize);
-    avx512_qsort<T>(arr, k);
+inline void avx512_partial_qsort(T *arr, int64_t k, int64_t arrsize) {
+    avx512_qselect<T>(arr, k - 1, arrsize);
+    avx512_qsort<T>(arr, k - 1);
 }
 inline void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize)
 {
-    avx512_qselect_fp16(arr, k, arrsize);
-    avx512_qsort_fp16(arr, k);
+    avx512_qselect_fp16(arr, k - 1, arrsize);
+    avx512_qsort_fp16(arr, k - 1);
 }
 
 template <typename vtype, typename T = typename vtype::type_t>

tests/test_qselect.hpp

@@ -26,8 +26,17 @@ TYPED_TEST_P(avx512_select, test_arrsizes)
             std::sort(sortedarr.begin(), sortedarr.end());
             for (size_t k = 0; k < arr.size(); ++k) {
                 psortedarr = arr;
-                avx512_qselect<TypeParam>(psortedarr.data(), k+1, psortedarr.size());
+                avx512_qselect<TypeParam>(psortedarr.data(), k, psortedarr.size());
+                /* index k is correct */
                 ASSERT_EQ(sortedarr[k], psortedarr[k]);
+                /* Check left partition */
+                for (size_t jj = 0; jj < k; jj++) {
+                    ASSERT_LE(psortedarr[jj], psortedarr[k]);
+                }
+                /* Check right partition */
+                for (size_t jj = k+1; jj < arr.size(); jj++) {
+                    ASSERT_GE(psortedarr[jj], psortedarr[k]);
+                }
                 psortedarr.clear();
             }
             arr.clear();

tests/test_qsortfp16.cpp

@@ -95,8 +95,17 @@ TEST(avx512_qselect_float16, test_arrsizes)
             std::sort(sortedarr.begin(), sortedarr.end());
             for (size_t k = 0; k < arr.size(); ++k) {
                 psortedarr = arr;
-                avx512_qselect<_Float16>(psortedarr.data(), k+1, psortedarr.size());
+                avx512_qselect<_Float16>(psortedarr.data(), k, psortedarr.size());
+                /* index k is correct */
                 ASSERT_EQ(sortedarr[k], psortedarr[k]);
+                /* Check left partition */
+                for (size_t jj = 0; jj < k; jj++) {
+                    ASSERT_LE(psortedarr[jj], psortedarr[k]);
+                }
+                /* Check right partition */
+                for (size_t jj = k+1; jj < arr.size(); jj++) {
+                    ASSERT_GE(psortedarr[jj], psortedarr[k]);
+                }
                 psortedarr.clear();
             }
             arr.clear();