Enable fp16 sorting without AVX512FP16

sterrettm2 · sterrettm2 · commit e0103be46398 · 2025-04-24T14:28:40.000-07:00
diff --git a/lib/x86simdsort-icl.cpp b/lib/x86simdsort-icl.cpp
@@ -50,5 +50,30 @@ namespace avx512 {
     {
         x86simdsortStatic::partial_qsort(arr, k, arrsize, hasnan, descending);
     }
+#ifdef __FLT16_MAX__
+    template <>
+    void qsort(_Float16 *arr, size_t size, bool hasnan, bool descending)
+    {
+        x86simdsortStatic::qsort(arr, size, hasnan, descending);
+    }
+    template <>
+    void qselect(_Float16 *arr,
+                 size_t k,
+                 size_t arrsize,
+                 bool hasnan,
+                 bool descending)
+    {
+        x86simdsortStatic::qselect(arr, k, arrsize, hasnan, descending);
+    }
+    template <>
+    void partial_qsort(_Float16 *arr,
+                       size_t k,
+                       size_t arrsize,
+                       bool hasnan,
+                       bool descending)
+    {
+        x86simdsortStatic::partial_qsort(arr, k, arrsize, hasnan, descending);
+    }
+#endif
 } // namespace avx512
 } // namespace xss
diff --git a/lib/x86simdsort.cpp b/lib/x86simdsort.cpp
@@ -137,9 +137,9 @@ namespace x86simdsort {
     }
 
 #ifdef __FLT16_MAX__
-DISPATCH(qsort, _Float16, ISA_LIST("avx512_spr"))
-DISPATCH(qselect, _Float16, ISA_LIST("avx512_spr"))
-DISPATCH(partial_qsort, _Float16, ISA_LIST("avx512_spr"))
+DISPATCH(qsort, _Float16, ISA_LIST("avx512_spr", "avx512_icl"))
+DISPATCH(qselect, _Float16, ISA_LIST("avx512_spr", "avx512_icl"))
+DISPATCH(partial_qsort, _Float16, ISA_LIST("avx512_spr", "avx512_icl"))
 DISPATCH(argsort, _Float16, ISA_LIST("none"))
 DISPATCH(argselect, _Float16, ISA_LIST("none"))
 #endif
diff --git a/src/avx512-16bit-qsort.hpp b/src/avx512-16bit-qsort.hpp
@@ -9,10 +9,6 @@
 
 #include "avx512-16bit-common.h"
 
-struct float16 {
-    uint16_t val;
-};
-
 template <>
 struct zmm_vector<float16> {
     using type_t = uint16_t;
@@ -545,10 +541,65 @@ replace_nan_with_inf<zmm_vector<float16>>(uint16_t *arr, arrsize_t arrsize)
     return nan_count;
 }
 
-template <>
-X86_SIMD_SORT_INLINE_ONLY bool is_a_nan<uint16_t>(uint16_t elem)
+X86_SIMD_SORT_INLINE_ONLY void replace_inf_with_nan_fp16(_Float16 *arr,
+                                                         arrsize_t size,
+                                                         arrsize_t nan_count,
+                                                         bool descending
+                                                         = false)
+{
+    if (descending) {
+        for (arrsize_t ii = 0; nan_count > 0; ++ii) {
+            arr[ii] = xss::fp::quiet_NaN<_Float16>();
+            nan_count -= 1;
+        }
+    }
+    else {
+        for (arrsize_t ii = size - 1; nan_count > 0; --ii) {
+            arr[ii] = xss::fp::quiet_NaN<_Float16>();
+            nan_count -= 1;
+        }
+    }
+}
+
+template <typename comparator>
+[[maybe_unused]] X86_SIMD_SORT_INLINE void
+avx512_qsort_fp16_helper(uint16_t *arr, arrsize_t arrsize)
 {
-    return ((elem & 0x7c00u) == 0x7c00u) && ((elem & 0x03ffu) != 0);
+    using T = uint16_t;
+    using vtype = zmm_vector<float16>;
+
+#ifdef XSS_COMPILE_OPENMP
+    bool use_parallel = arrsize > 100000;
+
+    if (use_parallel) {
+        // This thread limit was determined experimentally; it may be better for it to be the number of physical cores on the system
+        constexpr int thread_limit = 8;
+        int thread_count = std::min(thread_limit, omp_get_max_threads());
+        arrsize_t task_threshold = std::max((arrsize_t)100000, arrsize / 100);
+
+        // We use omp parallel and then omp single to setup the threads that will run the omp task calls in qsort_
+        // The omp single prevents multiple threads from running the initial qsort_ simultaneously and causing problems
+        // Note that we do not use the if(...) clause built into OpenMP, because it causes a performance regression for small arrays
+#pragma omp parallel num_threads(thread_count)
+#pragma omp single
+        qsort_<vtype, comparator, T>(arr,
+                                     0,
+                                     arrsize - 1,
+                                     2 * (arrsize_t)log2(arrsize),
+                                     task_threshold);
+    }
+    else {
+        qsort_<vtype, comparator, T>(arr,
+                                     0,
+                                     arrsize - 1,
+                                     2 * (arrsize_t)log2(arrsize),
+                                     std::numeric_limits<arrsize_t>::max());
+    }
+#pragma omp taskwait
+#else
+    qsort_<vtype, comparator, T>(
+            arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
+#endif
 }
 
 [[maybe_unused]] X86_SIMD_SORT_INLINE void
@@ -559,22 +610,19 @@ avx512_qsort_fp16(uint16_t *arr,
 {
     using vtype = zmm_vector<float16>;
 
-    // TODO multithreading support here
     if (arrsize > 1) {
         arrsize_t nan_count = 0;
         if (UNLIKELY(hasnan)) {
-            nan_count = replace_nan_with_inf<zmm_vector<float16>, uint16_t>(
-                    arr, arrsize);
+            nan_count = replace_nan_with_inf<vtype, uint16_t>(arr, arrsize);
         }
         if (descending) {
-            qsort_<vtype, Comparator<vtype, true>, uint16_t>(
-                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
+            avx512_qsort_fp16_helper<Comparator<vtype, true>>(arr, arrsize);
         }
         else {
-            qsort_<vtype, Comparator<vtype, false>, uint16_t>(
-                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize), 0);
+            avx512_qsort_fp16_helper<Comparator<vtype, false>>(arr, arrsize);
         }
-        replace_inf_with_nan(arr, arrsize, nan_count, descending);
+        replace_inf_with_nan_fp16(
+                (_Float16 *)arr, arrsize, nan_count, descending);
     }
 
 #ifdef __MMX__
@@ -592,26 +640,37 @@ avx512_qselect_fp16(uint16_t *arr,
 {
     using vtype = zmm_vector<float16>;
 
-    arrsize_t indx_last_elem = arrsize - 1;
+    // Exit early if no work would be done
+    if (arrsize <= 1) return;
+
+    arrsize_t index_first_elem = 0;
+    arrsize_t index_last_elem = arrsize - 1;
+
     if (UNLIKELY(hasnan)) {
-        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        if (descending) {
+            index_first_elem = move_nans_to_start_of_array(arr, arrsize);
+        }
+        else {
+            index_last_elem = move_nans_to_end_of_array(arr, arrsize);
+        }
     }
-    if (indx_last_elem >= k) {
+
+    if (index_first_elem <= k && index_last_elem >= k) {
         if (descending) {
             qselect_<vtype, Comparator<vtype, true>, uint16_t>(
                     arr,
                     k,
-                    0,
-                    indx_last_elem,
-                    2 * (arrsize_t)log2(indx_last_elem));
+                    index_first_elem,
+                    index_last_elem,
+                    2 * (arrsize_t)log2(arrsize));
         }
         else {
             qselect_<vtype, Comparator<vtype, false>, uint16_t>(
                     arr,
                     k,
-                    0,
-                    indx_last_elem,
-                    2 * (arrsize_t)log2(indx_last_elem));
+                    index_first_elem,
+                    index_last_elem,
+                    2 * (arrsize_t)log2(arrsize));
         }
     }
 
@@ -628,7 +687,8 @@ avx512_partial_qsort_fp16(uint16_t *arr,
                           bool hasnan = false,
                           bool descending = false)
 {
+    if (k == 0) return;
     avx512_qselect_fp16(arr, k - 1, arrsize, hasnan, descending);
-    avx512_qsort_fp16(arr, k - 1, descending);
+    avx512_qsort_fp16(arr, k - 1, hasnan, descending);
 }
 #endif // AVX512_QSORT_16BIT
diff --git a/src/x86simdsort-static-incl.h b/src/x86simdsort-static-incl.h
@@ -173,6 +173,27 @@ X86_SIMD_SORT_FINLINE void keyvalue_partial_sort(T1 *key,
 
 XSS_METHODS(avx512)
 
+#if defined(__FLT16_MAX__) && defined(__AVX512BW__) && defined(__AVX512VBMI2__) && !defined(__AVX512FP16__)
+template <>
+void x86simdsortStatic::qsort<_Float16>(
+		_Float16 *arr, size_t size, bool hasnan, bool descending)
+{
+	avx512_qsort_fp16((uint16_t *)arr, size, hasnan, descending);
+}
+template <>
+void x86simdsortStatic::qselect<_Float16>(
+		_Float16 *arr, size_t k, size_t size, bool hasnan, bool descending)
+{
+	avx512_qselect_fp16((uint16_t *)arr, k, size, hasnan, descending);
+}
+template <>
+void x86simdsortStatic::partial_qsort<_Float16>(
+		_Float16 *arr, size_t k, size_t size, bool hasnan, bool descending)
+{
+	avx512_partial_qsort_fp16((uint16_t *)arr, k, size, hasnan, descending);
+}
+#endif
+
 #elif defined(__AVX512F__)
 #error "x86simdsort requires AVX512DQ and AVX512VL to be enabled in addition to AVX512F to use AVX512"
 
diff --git a/src/xss-common-includes.h b/src/xss-common-includes.h
@@ -109,4 +109,8 @@ enum class simd_type : int { AVX2, AVX512 };
 template <typename vtype, typename T = typename vtype::type_t>
 X86_SIMD_SORT_INLINE bool comparison_func(const T &a, const T &b);
 
+struct float16 {
+    uint16_t val;
+};
+
 #endif // XSS_COMMON_INCLUDES
diff --git a/src/xss-common-qsort.h b/src/xss-common-qsort.h
@@ -45,6 +45,12 @@ bool is_a_nan(T elem)
     return std::isnan(elem);
 }
 
+template <>
+X86_SIMD_SORT_INLINE_ONLY bool is_a_nan<uint16_t>(uint16_t elem)
+{
+    return ((elem & 0x7c00u) == 0x7c00u) && ((elem & 0x03ffu) != 0);
+}
+
 template <typename vtype, typename T>
 X86_SIMD_SORT_INLINE arrsize_t replace_nan_with_inf(T *arr, arrsize_t size)
 {

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,12 @@ bool is_a_nan(T elem)`
`45`	`45`	`return std::isnan(elem);`
`46`	`46`	`}`
`47`	`47`
	`48`	`+template <>`
	`49`	`+X86_SIMD_SORT_INLINE_ONLY bool is_a_nan<uint16_t>(uint16_t elem)`
	`50`	`+{`
	`51`	`+ return ((elem & 0x7c00u) == 0x7c00u) && ((elem & 0x03ffu) != 0);`
	`52`	`+}`
	`53`	`+`
`48`	`54`	`template <typename vtype, typename T>`
`49`	`55`	`X86_SIMD_SORT_INLINE arrsize_t replace_nan_with_inf(T *arr, arrsize_t size)`
`50`	`56`	`{`