Fixed fp16 code issues

Author: sterrettm2

src/avx512fp16-16bit-qsort.hpp

@@ -179,51 +179,101 @@ X86_SIMD_SORT_INLINE_ONLY bool is_a_nan<_Float16>(_Float16 elem)
 }
 
 template <>
-X86_SIMD_SORT_INLINE_ONLY void
-replace_inf_with_nan(_Float16 *arr, arrsize_t size, arrsize_t nan_count)
+X86_SIMD_SORT_INLINE_ONLY void replace_inf_with_nan(_Float16 *arr,
+                                                    arrsize_t size,
+                                                    arrsize_t nan_count,
+                                                    bool descending)
 {
     Fp16Bits val;
     val.i_ = 0x7c01;
-    for (arrsize_t ii = size - 1; nan_count > 0; --ii) {
-        arr[ii] = val.f_;
-        nan_count -= 1;
+
+    if (descending) {
+        for (arrsize_t ii = 0; nan_count > 0; ++ii) {
+            arr[ii] = val.f_;
+            nan_count -= 1;
+        }
+    }
+    else {
+        for (arrsize_t ii = size - 1; nan_count > 0; --ii) {
+            arr[ii] = val.f_;
+            nan_count -= 1;
+        }
     }
 }
 /* Specialized template function for _Float16 qsort_*/
 template <>
 X86_SIMD_SORT_INLINE_ONLY void
-avx512_qsort(_Float16 *arr, arrsize_t arrsize, bool hasnan)
+avx512_qsort(_Float16 *arr, arrsize_t arrsize, bool hasnan, bool descending)
 {
+    using vtype = zmm_vector<_Float16>;
+
     if (arrsize > 1) {
         arrsize_t nan_count = 0;
         if (UNLIKELY(hasnan)) {
-            nan_count = replace_nan_with_inf<zmm_vector<_Float16>, _Float16>(
-                    arr, arrsize);
+            nan_count = replace_nan_with_inf<vtype, _Float16>(arr, arrsize);
+        }
+        if (descending) {
+            qsort_<vtype, DescendingComparator<vtype>, _Float16>(
+                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
+        }
+        else {
+            qsort_<vtype, AscendingComparator<vtype>, _Float16>(
+                    arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
         }
-        qsort_<zmm_vector<_Float16>, _Float16>(
-                arr, 0, arrsize - 1, 2 * (arrsize_t)log2(arrsize));
-        replace_inf_with_nan(arr, arrsize, nan_count);
+        replace_inf_with_nan(arr, arrsize, nan_count, descending);
     }
 }
 
 template <>
-X86_SIMD_SORT_INLINE_ONLY void
-avx512_qselect(_Float16 *arr, arrsize_t k, arrsize_t arrsize, bool hasnan)
+X86_SIMD_SORT_INLINE_ONLY void avx512_qselect(_Float16 *arr,
+                                              arrsize_t k,
+                                              arrsize_t arrsize,
+                                              bool hasnan,
+                                              bool descending)
 {
-    arrsize_t indx_last_elem = arrsize - 1;
-    if (UNLIKELY(hasnan)) {
-        indx_last_elem = move_nans_to_end_of_array(arr, arrsize);
+    using vtype = zmm_vector<_Float16>;
+
+    if (descending) {
+        arrsize_t index_first_elem = 0;
+        if (UNLIKELY(hasnan)) {
+            index_first_elem = move_nans_to_start_of_array(arr, arrsize);
+        }
+
+        arrsize_t size_without_nans = arrsize - index_first_elem;
+
+        if (index_first_elem <= k) {
+            qselect_<vtype, DescendingComparator<vtype>, _Float16>(
+                    arr,
+                    k,
+                    index_first_elem,
+                    arrsize - 1,
+                    2 * (arrsize_t)log2(size_without_nans));
+        }
     }
-    if (indx_last_elem >= k) {
-        qselect_<zmm_vector<_Float16>, _Float16>(
-                arr, k, 0, indx_last_elem, 2 * (arrsize_t)log2(indx_last_elem));
+    else {
+        arrsize_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_<vtype, AscendingComparator<vtype>, _Float16>(
+                    arr,
+                    k,
+                    0,
+                    indx_last_elem,
+                    2 * (arrsize_t)log2(indx_last_elem));
+        }
     }
 }
 template <>
-X86_SIMD_SORT_INLINE_ONLY void
-avx512_partial_qsort(_Float16 *arr, arrsize_t k, arrsize_t arrsize, bool hasnan)
+X86_SIMD_SORT_INLINE_ONLY void avx512_partial_qsort(_Float16 *arr,
+                                                    arrsize_t k,
+                                                    arrsize_t arrsize,
+                                                    bool hasnan,
+                                                    bool descending)
 {
-    avx512_qselect(arr, k - 1, arrsize, hasnan);
-    avx512_qsort(arr, k - 1, hasnan);
+    avx512_qselect(arr, k - 1, arrsize, hasnan, descending);
+    avx512_qsort(arr, k - 1, hasnan, descending);
 }
 #endif // AVX512FP16_QSORT_16BIT