generalize key-value sort for all 64-bit combinations

Author: Raghuveer Devulapalli

src/avx512-64bit-keyvaluesort.hpp

@@ -439,34 +439,20 @@ void qsort_64bit_(type1_t *keys,
     }
 }
 
-template <>
-void avx512_qsort_kv<int64_t>(int64_t *keys, uint64_t *indexes, int64_t arrsize)
+template <typename T1, typename T2>
+void avx512_qsort_kv(T1 *keys, T2 *indexes, int64_t arrsize)
 {
     if (arrsize > 1) {
-        qsort_64bit_<zmm_vector<int64_t>, zmm_vector<uint64_t>>(
-                keys, indexes, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
-}
-
-template <>
-void avx512_qsort_kv<uint64_t>(uint64_t *keys,
-                               uint64_t *indexes,
-                               int64_t arrsize)
-{
-    if (arrsize > 1) {
-        qsort_64bit_<zmm_vector<uint64_t>, zmm_vector<uint64_t>>(
-                keys, indexes, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
-}
-
-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<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);
+        if constexpr (std::is_floating_point_v<T1>) {
+            int64_t nan_count = replace_nan_with_inf<zmm_vector<double>>(keys, arrsize);
+            qsort_64bit_<zmm_vector<T1>, zmm_vector<T2>>(
+                    keys, indexes, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+            replace_inf_with_nan(keys, arrsize, nan_count);
+        }
+        else {
+            qsort_64bit_<zmm_vector<T1>, zmm_vector<T2>>(
+                    keys, indexes, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        }
     }
 }
 #endif // AVX512_QSORT_64BIT_KV

src/avx512-common-qsort.h

@@ -726,9 +726,4 @@ inline void avx512_partial_qsort_fp16(uint16_t *arr, int64_t k, int64_t arrsize,
     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

tests/test-keyvalue.cpp

@@ -54,7 +54,7 @@ TYPED_TEST_P(KeyValueSort, test_64bit_random_data)
             std::sort(sortedarr.begin(),
                       sortedarr.end(),
                       compare<TypeParam, uint64_t>);
-            avx512_qsort_kv<TypeParam>(keys.data(), values.data(), keys.size());
+            avx512_qsort_kv(keys.data(), values.data(), keys.size());
             for (size_t i = 0; i < keys.size(); i++) {
                 ASSERT_EQ(keys[i], sortedarr[i].key);
                 ASSERT_EQ(values[i], sortedarr[i].value);