minimize template specialization for avx512_qselect

Raghuveer Devulapalli · Raghuveer Devulapalli · commit 472c7d00fd44 · 2023-08-08T09:59:33.000-07:00
diff --git a/src/avx512-16bit-qsort.hpp b/src/avx512-16bit-qsort.hpp
@@ -403,56 +403,51 @@ bool is_a_nan<uint16_t>(uint16_t elem)
     return (elem & 0x7c00) == 0x7c00;
 }
 
+/* Specialized template function for 16-bit qsort_ funcs*/
 template <>
-void avx512_qselect(int16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+void qsort_<zmm_vector<int16_t>>(int16_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_16bit_<zmm_vector<int16_t>, int16_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_16bit_<zmm_vector<int16_t>>(arr, left, right, maxiters);
 }
 
 template <>
-void avx512_qselect(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+void qsort_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_16bit_<zmm_vector<uint16_t>, uint16_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qsort_16bit_<zmm_vector<uint16_t>>(arr, left, right, maxiters);
 }
 
-void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
+void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)
 {
-    int64_t indx_last_elem = arrsize - 1;
-    if (UNLIKELY(hasnan)) {
-        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));
+    if (arrsize > 1) {
+        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);
     }
 }
 
+/* Specialized template function for 16-bit qselect_ funcs*/
 template <>
-void qsort_<zmm_vector<int16_t>>(int16_t* arr, int64_t left, int64_t right, int64_t maxiters)
+void qselect_<zmm_vector<int16_t>>(int16_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    qsort_16bit_<zmm_vector<int16_t>>(arr, left, right, maxiters);
+    qselect_16bit_<zmm_vector<int16_t>>(arr, k, left, right, maxiters);
 }
 
 template <>
-void qsort_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t left, int64_t right, int64_t maxiters)
+void qselect_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    qsort_16bit_<zmm_vector<uint16_t>>(arr, left, right, maxiters);
+    qselect_16bit_<zmm_vector<uint16_t>>(arr, k, left, right, maxiters);
 }
 
-void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)
+void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)
 {
-    if (arrsize > 1) {
-        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);
+    int64_t indx_last_elem = arrsize - 1;
+    if (UNLIKELY(hasnan)) {
+        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));
     }
 }
-
 #endif // AVX512_QSORT_16BIT
diff --git a/src/avx512-32bit-qsort.hpp b/src/avx512-32bit-qsort.hpp
@@ -702,43 +702,26 @@ static void qselect_32bit_(type_t *arr,
         qselect_32bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
 }
 
+/* Specialized template function for 32-bit qselect_ funcs*/
 template <>
-void avx512_qselect<int32_t>(int32_t *arr,
-                             int64_t k,
-                             int64_t arrsize,
-                             bool hasnan)
+void qselect_<zmm_vector<int32_t>>(int32_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_32bit_<zmm_vector<int32_t>, int32_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qselect_32bit_<zmm_vector<int32_t>>(arr, k, left, right, maxiters);
 }
 
 template <>
-void avx512_qselect<uint32_t>(uint32_t *arr,
-                              int64_t k,
-                              int64_t arrsize,
-                              bool hasnan)
+void qselect_<zmm_vector<uint32_t>>(uint32_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_32bit_<zmm_vector<uint32_t>, uint32_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qselect_32bit_<zmm_vector<uint32_t>>(arr, k, left, right, maxiters);
 }
 
 template <>
-void avx512_qselect<float>(float *arr, int64_t k, int64_t arrsize, bool hasnan)
+void qselect_<zmm_vector<float>>(float* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    int64_t indx_last_elem = arrsize - 1;
-    if (UNLIKELY(hasnan)) {
-        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));
-    }
+    qselect_32bit_<zmm_vector<float>>(arr, k, left, right, maxiters);
 }
 
+/* Specialized template function for 32-bit qsort_ funcs*/
 template <>
 void qsort_<zmm_vector<int32_t>>(int32_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
diff --git a/src/avx512-64bit-qsort.hpp b/src/avx512-64bit-qsort.hpp
@@ -783,46 +783,26 @@ static void qselect_64bit_(type_t *arr,
         qselect_64bit_<vtype>(arr, pos, pivot_index, right, max_iters - 1);
 }
 
+/* Specialized template function for 64-bit qselect_ funcs*/
 template <>
-void avx512_qselect<int64_t>(int64_t *arr,
-                             int64_t k,
-                             int64_t arrsize,
-                             bool hasnan)
+void qselect_<zmm_vector<int64_t>>(int64_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_64bit_<zmm_vector<int64_t>, int64_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qselect_64bit_<zmm_vector<int64_t>>(arr, k, left, right, maxiters);
 }
 
 template <>
-void avx512_qselect<uint64_t>(uint64_t *arr,
-                              int64_t k,
-                              int64_t arrsize,
-                              bool hasnan)
+void qselect_<zmm_vector<uint64_t>>(uint64_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    if (arrsize > 1) {
-        qselect_64bit_<zmm_vector<uint64_t>, uint64_t>(
-                arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
-    }
+    qselect_64bit_<zmm_vector<uint64_t>>(arr, k, left, right, maxiters);
 }
 
 template <>
-void avx512_qselect<double>(double *arr,
-                            int64_t k,
-                            int64_t arrsize,
-                            bool hasnan)
+void qselect_<zmm_vector<double>>(double* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    int64_t indx_last_elem = arrsize - 1;
-    if (UNLIKELY(hasnan)) {
-        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));
-    }
+    qselect_64bit_<zmm_vector<double>>(arr, k, left, right, maxiters);
 }
 
+/* Specialized template function for 64-bit qsort_ funcs*/
 template <>
 void qsort_<zmm_vector<int64_t>>(int64_t* arr, int64_t left, int64_t right, int64_t maxiters)
 {
diff --git a/src/avx512-common-qsort.h b/src/avx512-common-qsort.h
@@ -673,11 +673,15 @@ static inline int64_t partition_avx512(type_t1 *keys,
 template <typename vtype, typename type_t>
 void qsort_(type_t* arr, int64_t left, int64_t right, int64_t maxiters);
 
+template <typename vtype, typename type_t>
+void qselect_(type_t* arr, int64_t pos, 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) {
+        /* std::is_floating_point_v<_Float16> == False, unless c++-23*/
         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>(
@@ -694,7 +698,21 @@ 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(T *arr, int64_t k, int64_t arrsize, bool hasnan = false)
+{
+    int64_t indx_last_elem = arrsize - 1;
+    /* std::is_floating_point_v<_Float16> == False, unless c++-23*/
+    if constexpr (std::is_floating_point_v<T>) {
+        if (UNLIKELY(hasnan)) {
+            indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
+    }
+    if (indx_last_elem >= k) {
+        qselect_<zmm_vector<T>, T>(
+                arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));
+    }
+}
+
 void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan = false);
 
 template <typename T>
diff --git a/src/avx512fp16-16bit-qsort.hpp b/src/avx512fp16-16bit-qsort.hpp
@@ -135,31 +135,26 @@ bool is_a_nan<_Float16>(_Float16 elem)
     return (temp.i_ & 0x7c00) == 0x7c00;
 }
 
-template <>
-void avx512_qselect(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)
+template<>
+void replace_inf_with_nan(_Float16 *arr, int64_t arrsize, int64_t nan_count)
 {
-    int64_t indx_last_elem = arrsize - 1;
-    if (UNLIKELY(hasnan)) {
-        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));
-    }
+    memset(arr + arrsize - nan_count, 0xFF, nan_count * 2);
 }
 
 template <>
-void qsort_<zmm_vector<_Float16>>(_Float16* arr, int64_t left, int64_t right, int64_t maxiters)
+void qselect_<zmm_vector<_Float16>>(_Float16* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)
 {
-    qsort_16bit_<zmm_vector<_Float16>>(arr, left, right, maxiters);
+    qselect_16bit_<zmm_vector<_Float16>>(arr, k, left, right, maxiters);
 }
 
-template<>
-void replace_inf_with_nan(_Float16 *arr, int64_t arrsize, int64_t nan_count)
+
+template <>
+void qsort_<zmm_vector<_Float16>>(_Float16* arr, int64_t left, int64_t right, int64_t maxiters)
 {
-    memset(arr + arrsize - nan_count, 0xFF, nan_count * 2);
+    qsort_16bit_<zmm_vector<_Float16>>(arr, left, right, maxiters);
 }
 
+/* Specialized template function for _Float16 qsort_*/
 template<>
 void avx512_qsort(_Float16 *arr, int64_t arrsize)
 {

Original file line number	Diff line number	Diff line change
`@@ -403,56 +403,51 @@ bool is_a_nan<uint16_t>(uint16_t elem)`
`403`	`403`	`return (elem & 0x7c00) == 0x7c00;`
`404`	`404`	`}`
`405`	`405`
	`406`	`+/* Specialized template function for 16-bit qsort_ funcs*/`
`406`	`407`	`template <>`
`407`		`-void avx512_qselect(int16_t *arr, int64_t k, int64_t arrsize, bool hasnan)`
	`408`	`+void qsort_<zmm_vector<int16_t>>(int16_t* arr, int64_t left, int64_t right, int64_t maxiters)`
`408`	`409`	`{`
`409`		`- if (arrsize > 1) {`
`410`		`- qselect_16bit_<zmm_vector<int16_t>, int16_t>(`
`411`		`- arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));`
`412`		`- }`
	`410`	`+ qsort_16bit_<zmm_vector<int16_t>>(arr, left, right, maxiters);`
`413`	`411`	`}`
`414`	`412`
`415`	`413`	`template <>`
`416`		`-void avx512_qselect(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)`
	`414`	`+void qsort_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t left, int64_t right, int64_t maxiters)`
`417`	`415`	`{`
`418`		`- if (arrsize > 1) {`
`419`		`- qselect_16bit_<zmm_vector<uint16_t>, uint16_t>(`
`420`		`- arr, k, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));`
`421`		`- }`
	`416`	`+ qsort_16bit_<zmm_vector<uint16_t>>(arr, left, right, maxiters);`
`422`	`417`	`}`
`423`	`418`
`424`		`-void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)`
	`419`	`+void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)`
`425`	`420`	`{`
`426`		`- int64_t indx_last_elem = arrsize - 1;`
`427`		`- if (UNLIKELY(hasnan)) {`
`428`		`- indx_last_elem = move_nans_to_end_of_array(arr, arrsize);`
`429`		`- }`
`430`		`- if (indx_last_elem >= k) {`
`431`		`- qselect_16bit_<zmm_vector<float16>, uint16_t>(`
`432`		`- arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));`
	`421`	`+ if (arrsize > 1) {`
	`422`	`+ int64_t nan_count = replace_nan_with_inf<zmm_vector<float16>, uint16_t>(arr, arrsize);`
	`423`	`+ qsort_16bit_<zmm_vector<float16>, uint16_t>(`
	`424`	`+ arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));`
	`425`	`+ replace_inf_with_nan(arr, arrsize, nan_count);`
`433`	`426`	`}`
`434`	`427`	`}`
`435`	`428`
	`429`	`+/* Specialized template function for 16-bit qselect_ funcs*/`
`436`	`430`	`template <>`
`437`		`-void qsort_<zmm_vector<int16_t>>(int16_t* arr, int64_t left, int64_t right, int64_t maxiters)`
	`431`	`+void qselect_<zmm_vector<int16_t>>(int16_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)`
`438`	`432`	`{`
`439`		`- qsort_16bit_<zmm_vector<int16_t>>(arr, left, right, maxiters);`
	`433`	`+ qselect_16bit_<zmm_vector<int16_t>>(arr, k, left, right, maxiters);`
`440`	`434`	`}`
`441`	`435`
`442`	`436`	`template <>`
`443`		`-void qsort_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t left, int64_t right, int64_t maxiters)`
	`437`	`+void qselect_<zmm_vector<uint16_t>>(uint16_t* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)`
`444`	`438`	`{`
`445`		`- qsort_16bit_<zmm_vector<uint16_t>>(arr, left, right, maxiters);`
	`439`	`+ qselect_16bit_<zmm_vector<uint16_t>>(arr, k, left, right, maxiters);`
`446`	`440`	`}`
`447`	`441`
`448`		`-void avx512_qsort_fp16(uint16_t *arr, int64_t arrsize)`
	`442`	`+void avx512_qselect_fp16(uint16_t *arr, int64_t k, int64_t arrsize, bool hasnan)`
`449`	`443`	`{`
`450`		`- if (arrsize > 1) {`
`451`		`- int64_t nan_count = replace_nan_with_inf<zmm_vector<float16>, uint16_t>(arr, arrsize);`
`452`		`- qsort_16bit_<zmm_vector<float16>, uint16_t>(`
`453`		`- arr, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));`
`454`		`- replace_inf_with_nan(arr, arrsize, nan_count);`
	`444`	`+ int64_t indx_last_elem = arrsize - 1;`
	`445`	`+ if (UNLIKELY(hasnan)) {`
	`446`	`+ indx_last_elem = move_nans_to_end_of_array(arr, arrsize);`
	`447`	`+ }`
	`448`	`+ if (indx_last_elem >= k) {`
	`449`	`+ qselect_16bit_<zmm_vector<float16>, uint16_t>(`
	`450`	`+ arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));`
`455`	`451`	`}`
`456`	`452`	`}`
`457`		`-`
`458`	`453`	`#endif // AVX512_QSORT_16BIT`
Original file line number	Diff line number	Diff line change
`@@ -135,31 +135,26 @@ bool is_a_nan<_Float16>(_Float16 elem)`
`135`	`135`	`return (temp.i_ & 0x7c00) == 0x7c00;`
`136`	`136`	`}`
`137`	`137`
`138`		`-template <>`
`139`		`-void avx512_qselect(_Float16 *arr, int64_t k, int64_t arrsize, bool hasnan)`
	`138`	`+template<>`
	`139`	`+void replace_inf_with_nan(_Float16 *arr, int64_t arrsize, int64_t nan_count)`
`140`	`140`	`{`
`141`		`- int64_t indx_last_elem = arrsize - 1;`
`142`		`- if (UNLIKELY(hasnan)) {`
`143`		`- indx_last_elem = move_nans_to_end_of_array(arr, arrsize);`
`144`		`- }`
`145`		`- if (indx_last_elem >= k) {`
`146`		`- qselect_16bit_<zmm_vector<_Float16>, _Float16>(`
`147`		`- arr, k, 0, indx_last_elem, 2 * (int64_t)log2(indx_last_elem));`
`148`		`- }`
	`141`	`+ memset(arr + arrsize - nan_count, 0xFF, nan_count * 2);`
`149`	`142`	`}`
`150`	`143`
`151`	`144`	`template <>`
`152`		`-void qsort_<zmm_vector<_Float16>>(_Float16* arr, int64_t left, int64_t right, int64_t maxiters)`
	`145`	`+void qselect_<zmm_vector<_Float16>>(_Float16* arr, int64_t k, int64_t left, int64_t right, int64_t maxiters)`
`153`	`146`	`{`
`154`		`- qsort_16bit_<zmm_vector<_Float16>>(arr, left, right, maxiters);`
	`147`	`+ qselect_16bit_<zmm_vector<_Float16>>(arr, k, left, right, maxiters);`
`155`	`148`	`}`
`156`	`149`
`157`		`-template<>`
`158`		`-void replace_inf_with_nan(_Float16 *arr, int64_t arrsize, int64_t nan_count)`
	`150`	`+`
	`151`	`+template <>`
	`152`	`+void qsort_<zmm_vector<_Float16>>(_Float16* arr, int64_t left, int64_t right, int64_t maxiters)`
`159`	`153`	`{`
`160`		`- memset(arr + arrsize - nan_count, 0xFF, nan_count * 2);`
	`154`	`+ qsort_16bit_<zmm_vector<_Float16>>(arr, left, right, maxiters);`
`161`	`155`	`}`
`162`	`156`
	`157`	`+/* Specialized template function for _Float16 qsort_*/`
`163`	`158`	`template<>`
`164`	`159`	`void avx512_qsort(_Float16 *arr, int64_t arrsize)`
`165`	`160`	`{`