Implement partial sorting algorithms

Each datatype now supports two partial sorting algorithms:
1) Sort such that a particular index is valid, and
2) Sort such that a range of indices is valid,
where 'valid' means that the kth smallest element is in position k.

Additionally transferred a few lingering comments from a refactor
earlier in the project.

Author: mosullivan93

src/avx512-16bit-common.h

@@ -259,6 +259,38 @@ X86_SIMD_SORT_INLINE type_t get_pivot_16bit(type_t *arr,
     return ((type_t *)&sort)[16];
 }
 
+template <typename vtype, typename type_t>
+static void
+qsort_partial_16bit_(int64_t k, type_t *arr,
+                     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 ((pivot != smallest) && (k <= pivot_index))
+        qsort_partial_16bit_<vtype>(k, arr, left, pivot_index - 1, max_iters - 1);
+    else if ((pivot != biggest) && (k > pivot_index))
+        qsort_partial_16bit_<vtype>(k, 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)

src/avx512-16bit-qsort.hpp

@@ -405,6 +405,34 @@ replace_inf_with_nan(uint16_t *arr, int64_t arrsize, int64_t nan_count)
     }
 }
 
+template <>
+void avx512_qsort_partial(int64_t k, int16_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_16bit_<zmm_vector<int16_t>, int16_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_qsort_partial(int64_t k, uint16_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_16bit_<zmm_vector<uint16_t>, uint16_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+void avx512_qsort_fp16_partial(int64_t k, uint16_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
+        qsort_partial_16bit_<zmm_vector<float16>, uint16_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        replace_inf_with_nan(arr, arrsize, nan_count);
+    }
+}
+
 template <>
 void avx512_qsort(int16_t *arr, int64_t arrsize)
 {
@@ -432,4 +460,5 @@ void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)
         replace_inf_with_nan(arr, arrsize, nan_count);
     }
 }
+
 #endif // AVX512_QSORT_16BIT

src/avx512-32bit-qsort.hpp

@@ -626,6 +626,38 @@ X86_SIMD_SORT_INLINE type_t get_pivot_32bit(type_t *arr,
     return ((type_t *)&sort)[8];
 }
 
+template <typename vtype, typename type_t>
+static void
+qsort_partial_32bit_(int64_t k, type_t *arr,
+                     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<vtype>(
+            arr, left, right + 1, pivot, &smallest, &biggest);
+    if ((pivot != smallest) && (k <= pivot_index))
+        qsort_partial_32bit_<vtype>(k, arr, left, pivot_index - 1, max_iters - 1);
+    else if ((pivot != biggest) && (k > pivot_index))
+        qsort_partial_32bit_<vtype>(k, 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)
@@ -681,6 +713,35 @@ replace_inf_with_nan(float *arr, int64_t arrsize, int64_t nan_count)
     }
 }
 
+template <>
+void avx512_qsort_partial<int32_t>(int64_t k, int32_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_32bit_<zmm_vector<int32_t>, int32_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_qsort_partial<uint32_t>(int64_t k, uint32_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_32bit_<zmm_vector<uint32_t>, uint32_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_qsort_partial<float>(int64_t k, float *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
+        qsort_partial_32bit_<zmm_vector<float>, float>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        replace_inf_with_nan(arr, arrsize, nan_count);
+    }
+}
+
 template <>
 void avx512_qsort<int32_t>(int32_t *arr, int64_t arrsize)
 {

src/avx512-64bit-common.h

@@ -4,10 +4,16 @@
  * Authors: Raghuveer Devulapalli <raghuveer.devulapalli@intel.com>
  * ****************************************************************/
 
-#ifndef AVX512_64BIT_COMMOM
-#define AVX512_64BIT_COMMOM
+#ifndef AVX512_64BIT_COMMON
+#define AVX512_64BIT_COMMON
 #include "avx512-common-qsort.h"
 
+/*
+ * Constants used in sorting 8 elements in a ZMM registers. Based on Bitonic
+ * sorting network (see
+ * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg)
+ */
+// ZMM                  7, 6, 5, 4, 3, 2, 1, 0
 #define NETWORK_64BIT_1 4, 5, 6, 7, 0, 1, 2, 3
 #define NETWORK_64BIT_2 0, 1, 2, 3, 4, 5, 6, 7
 #define NETWORK_64BIT_3 5, 4, 7, 6, 1, 0, 3, 2

src/avx512-64bit-qsort.hpp

@@ -9,13 +9,6 @@
 
 #include "avx512-64bit-common.h"
 
-/*
- * Constants used in sorting 8 elements in a ZMM registers. Based on Bitonic
- * sorting network (see
- * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg)
- */
-// ZMM                  7, 6, 5, 4, 3, 2, 1, 0
-
 // Assumes zmm is bitonic and performs a recursive half cleaner
 template <typename vtype, typename zmm_t = typename vtype::zmm_t>
 X86_SIMD_SORT_INLINE zmm_t bitonic_merge_zmm_64bit(zmm_t zmm)
@@ -408,6 +401,67 @@ qsort_64bit_(type_t *arr, int64_t left, int64_t right, int64_t max_iters)
         qsort_64bit_<vtype>(arr, pivot_index, right, max_iters - 1);
 }
 
+template <typename vtype, typename type_t>
+static void
+qsort_partial_64bit_(int64_t k, type_t *arr,
+                     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<vtype>(
+            arr, left, right + 1, pivot, &smallest, &biggest);
+    if ((pivot != smallest) && (k <= pivot_index))
+        qsort_partial_64bit_<vtype>(k, arr, left, pivot_index - 1, max_iters - 1);
+    else if ((pivot != biggest) && (k > pivot_index))
+        qsort_partial_64bit_<vtype>(k, arr, pivot_index, right, max_iters - 1);
+}
+
+template <>
+void avx512_qsort_partial<int64_t>(int64_t k, int64_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_64bit_<zmm_vector<int64_t>, int64_t>(
+            k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_qsort_partial<uint64_t>(int64_t k, uint64_t *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        qsort_partial_64bit_<zmm_vector<uint64_t>, uint64_t>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+    }
+}
+
+template <>
+void avx512_qsort_partial<double>(int64_t k, double *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
+        qsort_partial_64bit_<zmm_vector<double>, double>(
+                k, arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        replace_inf_with_nan(arr, arrsize, nan_count);
+    }
+}
+
 template <>
 void avx512_qsort<int64_t>(int64_t *arr, int64_t arrsize)
 {

src/avx512-common-qsort.h

@@ -88,9 +88,25 @@ struct zmm_vector;
 
 template <typename T>
 void avx512_qsort(T *arr, int64_t arrsize);
-
 void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize);
 
+template <typename T>
+void avx512_qsort_partial(int64_t k, T *arr, int64_t arrsize);
+void avx512_qsort_fp16_partial(int64_t k, uint16_t *arr, int64_t arrsize);
+
+template <typename T>
+void avx512_qsort_partialrange(int64_t kfrom, int64_t kto, T *arr, int64_t arrsize) {
+    avx512_qsort_partial<T>(kto, arr, arrsize);
+    avx512_qsort_partial<T>(kfrom, arr, kto);
+    avx512_qsort<T>(arr + kfrom, kto - kfrom);
+}
+inline void avx512_qsort_fp16_partialrange(int64_t kfrom, int64_t kto, uint16_t *arr, int64_t arrsize)
+{
+    avx512_qsort_fp16_partial(kto, arr, arrsize);
+    avx512_qsort_fp16_partial(kfrom, arr, kto);
+    avx512_qsort_fp16(arr + kfrom, kto - kfrom);
+}
+
 template <typename vtype, typename T = typename vtype::type_t>
 bool comparison_func(const T &a, const T &b)
 {