Remove template specializations for quicksort methods

Author: Raghuveer Devulapalli

src/avx512-16bit-qsort.hpp

@@ -377,8 +377,9 @@ bool comparison_func<zmm_vector<float16>>(const uint16_t &a, const uint16_t &b)
     //return npy_half_to_float(a) < npy_half_to_float(b);
 }
 
-X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(uint16_t *arr,
-                                                  int64_t arrsize)
+template<>
+int64_t
+replace_nan_with_inf<zmm_vector<float16>>(uint16_t *arr, int64_t arrsize)
 {
     int64_t nan_count = 0;
     __mmask16 loadmask = 0xFFFF;
@@ -396,15 +397,6 @@ X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(uint16_t *arr,
     return nan_count;
 }
 
-X86_SIMD_SORT_INLINE void
-replace_inf_with_nan(uint16_t *arr, int64_t arrsize, int64_t nan_count)
-{
-    for (int64_t ii = arrsize - 1; nan_count > 0; --ii) {
-        arr[ii] = 0xFFFF;
-        nan_count -= 1;
-    }
-}
-
 template <>
 bool is_a_nan<uint16_t>(uint16_t elem)
 {
@@ -442,27 +434,21 @@ void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
 }
 
 template <>
-void avx512_qsort(int16_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<int16_t>>(int16_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_16bit_<zmm_vector<int16_t>, int16_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_16bit_<zmm_vector<int16_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qsort(uint16_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_16bit_<zmm_vector<uint16_t>, uint16_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_16bit_<zmm_vector<uint16_t>>(arr, left, right, maxiters);
 }
 
 void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)
 {
     if (arrsize > 1) {
-        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
+        int64_t nan_count = replace_nan_with_inf<zmm_vector<float16>, uint16_t>(arr, arrsize);
         qsort_16bit_<zmm_vector<float16>, uint16_t>(
                 arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
         replace_inf_with_nan(arr, arrsize, nan_count);

src/avx512-32bit-qsort.hpp

@@ -256,6 +256,11 @@ struct zmm_vector<float> {
     {
         return _mm512_cmp_ps_mask(x, y, _CMP_GE_OQ);
     }
+    template <int type>
+    static opmask_t fpclass(zmm_t x)
+    {
+        return _mm512_fpclass_ps_mask(x, type);
+    }
     template <int scale>
     static ymm_t i64gather(__m512i index, void const *base)
     {
@@ -279,6 +284,10 @@ struct zmm_vector<float> {
     {
         return _mm512_mask_compressstoreu_ps(mem, mask, x);
     }
+    static zmm_t maskz_loadu(opmask_t mask, void const *mem)
+    {
+        return _mm512_maskz_loadu_ps(mask, mem);
+    }
     static zmm_t mask_loadu(zmm_t x, opmask_t mask, void const *mem)
     {
         return _mm512_mask_loadu_ps(x, mask, mem);
@@ -689,31 +698,6 @@ static void qselect_32bit_(type_t *arr,
         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)
-{
-    int64_t nan_count = 0;
-    __mmask16 loadmask = 0xFFFF;
-    while (arrsize > 0) {
-        if (arrsize < 16) { loadmask = (0x0001 << arrsize) - 0x0001; }
-        __m512 in_zmm = _mm512_maskz_loadu_ps(loadmask, arr);
-        __mmask16 nanmask = _mm512_cmp_ps_mask(in_zmm, in_zmm, _CMP_NEQ_UQ);
-        nan_count += _mm_popcnt_u32((int32_t)nanmask);
-        _mm512_mask_storeu_ps(arr, nanmask, ZMM_MAX_FLOAT);
-        arr += 16;
-        arrsize -= 16;
-    }
-    return nan_count;
-}
-
-X86_SIMD_SORT_INLINE void
-replace_inf_with_nan(float *arr, int64_t arrsize, int64_t nan_count)
-{
-    for (int64_t ii = arrsize - 1; nan_count > 0; --ii) {
-        arr[ii] = std::nanf("1");
-        nan_count -= 1;
-    }
-}
-
 template <>
 void avx512_qselect<int32_t>(int32_t *arr,
                              int64_t k,
@@ -752,32 +736,20 @@ void avx512_qselect<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
 }
 
 template <>
-void avx512_qsort<int32_t>(int32_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<int32_t>>(int32_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_32bit_<zmm_vector<int32_t>, int32_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_32bit_<zmm_vector<int32_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qsort<uint32_t>(uint32_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<uint32_t>>(uint32_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_32bit_<zmm_vector<uint32_t>, uint32_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_32bit_<zmm_vector<uint32_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qsort<float>(float *arr, int64_t arrsize)
+void qsort_<zmm_vector<float>>(float* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
-        qsort_32bit_<zmm_vector<float>, float>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-        replace_inf_with_nan(arr, arrsize, nan_count);
-    }
+    qsort_32bit_<zmm_vector<float>>(arr, left, right, maxiters);
 }
-
 #endif //AVX512_QSORT_32BIT

src/avx512-64bit-common.h

@@ -773,30 +773,7 @@ struct zmm_vector<double> {
         _mm512_storeu_pd(mem, x);
     }
 };
-X86_SIMD_SORT_INLINE int64_t replace_nan_with_inf(double *arr, int64_t arrsize)
-{
-    int64_t nan_count = 0;
-    __mmask8 loadmask = 0xFF;
-    while (arrsize > 0) {
-        if (arrsize < 8) { loadmask = (0x01 << arrsize) - 0x01; }
-        __m512d in_zmm = _mm512_maskz_loadu_pd(loadmask, arr);
-        __mmask8 nanmask = _mm512_cmp_pd_mask(in_zmm, in_zmm, _CMP_NEQ_UQ);
-        nan_count += _mm_popcnt_u32((int32_t)nanmask);
-        _mm512_mask_storeu_pd(arr, nanmask, ZMM_MAX_DOUBLE);
-        arr += 8;
-        arrsize -= 8;
-    }
-    return nan_count;
-}
 
-X86_SIMD_SORT_INLINE void
-replace_inf_with_nan(double *arr, int64_t arrsize, int64_t nan_count)
-{
-    for (int64_t ii = arrsize - 1; nan_count > 0; --ii) {
-        arr[ii] = std::nan("1");
-        nan_count -= 1;
-    }
-}
 /*
  * Assumes zmm is random and performs a full sorting network defined in
  * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg

src/avx512-64bit-keyvaluesort.hpp

@@ -463,7 +463,7 @@ template <>
 void avx512_qsort_kv<double>(double *keys, uint64_t *indexes, int64_t arrsize)
 {
     if (arrsize > 1) {
-        int64_t nan_count = replace_nan_with_inf(keys, arrsize);
+        int64_t nan_count = replace_nan_with_inf<zmm_vector<double>>(keys, arrsize);
         qsort_64bit_<zmm_vector<double>, zmm_vector<uint64_t>>(
                 keys, indexes, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
         replace_inf_with_nan(keys, arrsize, nan_count);

src/avx512-64bit-qsort.hpp

@@ -824,31 +824,20 @@ void avx512_qselect<double>(double *arr,
 }
 
 template <>
-void avx512_qsort<int64_t>(int64_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<int64_t>>(int64_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_64bit_<zmm_vector<int64_t>, int64_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_64bit_<zmm_vector<int64_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qsort<uint64_t>(uint64_t *arr, int64_t arrsize)
+void qsort_<zmm_vector<uint64_t>>(uint64_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qsort_64bit_<zmm_vector<uint64_t>, uint64_t>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_64bit_<zmm_vector<uint64_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qsort<double>(double *arr, int64_t arrsize)
+void qsort_<zmm_vector<double>>(double* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        int64_t nan_count = replace_nan_with_inf(arr, arrsize);
-        qsort_64bit_<zmm_vector<double>, double>(
-                arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-        replace_inf_with_nan(arr, arrsize, nan_count);
-    }
+    qsort_64bit_<zmm_vector<double>>(arr, left, right, maxiters);
 }
 #endif // AVX512_QSORT_64BIT

src/avx512-common-qsort.h

@@ -94,35 +94,50 @@ struct zmm_vector;
 template <typename type>
 struct ymm_vector;
 
-// Regular quicksort routines:
 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_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)
+bool is_a_nan(T elem)
 {
-    avx512_qselect<T>(arr, k - 1, arrsize, hasnan);
-    avx512_qsort<T>(arr, k - 1);
+    return std::isnan(elem);
 }
-inline void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false)
+
+template <typename vtype, typename type_t>
+int64_t replace_nan_with_inf(type_t *arr, int64_t arrsize)
 {
-    avx512_qselect_fp16(arr, k - 1, arrsize, hasnan);
-    avx512_qsort_fp16(arr, k - 1);
+    int64_t nan_count = 0;
+    using opmask_t = typename vtype::opmask_t;
+    using zmm_t = typename vtype::zmm_t;
+    bool found_nan = false;
+    opmask_t loadmask = 0xFF;
+    zmm_t in;
+    while (arrsize > 0) {
+        if (arrsize < vtype::numlanes) {
+            loadmask = (0x01 << arrsize) - 0x01;
+            in = vtype::maskz_loadu(loadmask, arr);
+        }
+        else {
+            in = vtype::loadu(arr);
+        }
+        opmask_t nanmask = vtype::template fpclass<0x01 | 0x80>(in);
+        nan_count += _mm_popcnt_u32((int32_t)nanmask);
+        vtype::mask_storeu(arr, nanmask, vtype::zmm_max());
+        arr += vtype::numlanes;
+        arrsize -= vtype::numlanes;
+    }
+    return nan_count;
 }
 
-// key-value sort routines
-template <typename T>
-void avx512_qsort_kv(T *keys, uint64_t *indexes, int64_t arrsize);
-
-template <typename T>
-bool is_a_nan(T elem)
+template<typename type_t>
+void replace_inf_with_nan(type_t *arr, int64_t arrsize, int64_t nan_count)
 {
-    return std::isnan(elem);
+    for (int64_t ii = arrsize - 1; nan_count > 0; --ii) {
+        if constexpr (std::is_floating_point_v<type_t>) {
+            arr[ii] = std::numeric_limits<type_t>::quiet_NaN();
+        }
+        else {
+            arr[ii] = 0xFFFF;
+        }
+        nan_count -= 1;
+    }
 }
 
 /*
@@ -628,4 +643,48 @@ static inline int64_t partition_avx512(type_t1 *keys,
     *biggest = vtype1::reducemax(max_vec);
     return l_store;
 }
+
+template <typename vtype, typename type_t>
+void qsort_(type_t* arr, int64_t left, int64_t right, int64_t maxiters);
+
+// Regular quicksort routines:
+template <typename T>
+void avx512_qsort(T *arr, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        if constexpr (std::is_floating_point_v<T>) {
+            int64_t nan_count = replace_nan_with_inf<zmm_vector<T>>(arr, arrsize);
+            qsort_<zmm_vector<T>, T>(
+                    arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+            replace_inf_with_nan(arr, arrsize, nan_count);
+        }
+        else {
+            qsort_<zmm_vector<T>, T>(
+                    arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        }
+    }
+}
+
+void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize);
+
+template <typename T>
+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);
+}
+
+// key-value sort routines
+template <typename T>
+void avx512_qsort_kv(T *keys, uint64_t *indexes, int64_t arrsize);
+
 #endif // AVX512_QSORT_COMMON