Fix NAN check for _Float16

Author: Raghuveer Devulapalli

src/avx512-16bit-qsort.hpp

@@ -445,7 +445,8 @@ arrsize_t replace_nan_with_inf<zmm_vector<float16>>(uint16_t *arr,
 template <>
 bool is_a_nan<uint16_t>(uint16_t elem)
 {
-    return (elem & 0x7c00) == 0x7c00;
+    return ((elem & 0x7c00u) == 0x7c00u) &&
+           ((elem & 0x03ffu) != 0);
 }
 
 X86_SIMD_SORT_INLINE

src/avx512-common-qsort.h

@@ -191,7 +191,7 @@ X86_SIMD_SORT_INLINE arrsize_t move_nans_to_end_of_array(T *arr, arrsize_t size)
     arrsize_t jj = size - 1;
     arrsize_t ii = 0;
     arrsize_t count = 0;
-    while (ii <= jj) {
+    while (ii < jj) {
         if (is_a_nan(arr[ii])) {
             std::swap(arr[ii], arr[jj]);
             jj -= 1;
@@ -201,6 +201,10 @@ X86_SIMD_SORT_INLINE arrsize_t move_nans_to_end_of_array(T *arr, arrsize_t size)
             ii += 1;
         }
     }
+    /* Haven't checked for nan when ii == jj */
+    if (is_a_nan(arr[ii])) {
+        count++;
+    }
     return size - count - 1;
 }
 

src/avx512fp16-16bit-qsort.hpp

@@ -145,17 +145,19 @@ struct zmm_vector<_Float16> {
 template <>
 bool is_a_nan<_Float16>(_Float16 elem)
 {
-    Fp16Bits temp;
-    temp.f_ = elem;
-    return (temp.i_ & 0x7c00) == 0x7c00;
+    return elem != elem;
 }
 
 template <>
-void replace_inf_with_nan(_Float16 *arr, arrsize_t arrsize, arrsize_t nan_count)
+void replace_inf_with_nan(_Float16 *arr, arrsize_t size, arrsize_t nan_count)
 {
-    memset(arr + arrsize - nan_count, 0xFF, nan_count * 2);
+    Fp16Bits val;
+    val.i_ = 0x7c01;
+    for (arrsize_t ii = size - 1; nan_count > 0; --ii) {
+        arr[ii] = val.f_;
+        nan_count -= 1;
+    }
 }
-
 /* Specialized template function for _Float16 qsort_*/
 template <>
 void avx512_qsort(_Float16 *arr, arrsize_t arrsize)
@@ -169,4 +171,17 @@ void avx512_qsort(_Float16 *arr, arrsize_t arrsize)
         replace_inf_with_nan(arr, arrsize, nan_count);
     }
 }
+
+template <>
+void avx512_qselect(_Float16 *arr, arrsize_t k, arrsize_t arrsize, bool hasnan)
+{
+    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_<zmm_vector<_Float16>, _Float16>(
+                arr, k, 0, indx_last_elem, 2 * (arrsize_t)log2(indx_last_elem));
+    }
+}
 #endif // AVX512FP16_QSORT_16BIT