Detect NAN in float/double array and skip avx512_argsort

Author: Raghuveer Devulapalli

src/avx512-64bit-argsort.hpp

@@ -11,6 +11,19 @@
 #include "avx512-common-argsort.h"
 #include "avx512-64bit-keyvalue-networks.hpp"
 
+/* argsort using std::sort */
+template <typename T>
+void std_argsort_withnan(T *arr, int64_t *arg, int64_t left, int64_t right)
+{
+    std::sort(arg + left,
+              arg + right,
+              [arr](int64_t left, int64_t right) -> bool {
+              if ((!std::isnan(arr[left])) && (!std::isnan(arr[right]))) {return arr[left] < arr[right];}
+              else if (std::isnan(arr[left])) {return false;}
+              else {return true;}
+              });
+}
+
 /* argsort using std::sort */
 template <typename T>
 void std_argsort(T *arr, int64_t *arg, int64_t left, int64_t right)
@@ -270,6 +283,33 @@ inline void argsort_64bit_(type_t *arr,
         argsort_64bit_<vtype>(arr, arg, pivot_index, right, max_iters - 1);
 }
 
+template <typename vtype, typename type_t>
+bool has_nan(type_t* arr, int64_t arrsize)
+{
+    using opmask_t = typename vtype::opmask_t;
+    using zmm_t = typename vtype::zmm_t;
+    bool found_nan = false;
+    opmask_t loadmask = 0xFF;
+    zmm_t in;
+    while (arrsize > 0) {
+        if (arrsize < vtype::numlanes) {
+            loadmask = (0x01 << arrsize) - 0x01;
+            in = vtype::maskz_loadu(loadmask, arr);
+        }
+        else {
+            in = vtype::loadu(arr);
+        }
+        opmask_t nanmask = vtype::template fpclass<0x01|0x80>(in);
+        arr += vtype::numlanes;
+        arrsize -= vtype::numlanes;
+        if (nanmask != 0x00) {
+            found_nan = true;
+            break;
+        }
+    }
+    return found_nan;
+}
+
 template <typename T>
 void avx512_argsort(T* arr, int64_t *arg, int64_t arrsize)
 {
@@ -279,6 +319,21 @@ void avx512_argsort(T* arr, int64_t *arg, int64_t arrsize)
     }
 }
 
+template <>
+void avx512_argsort(double* arr, int64_t *arg, int64_t arrsize)
+{
+    if (arrsize > 1) {
+        if (has_nan<zmm_vector<double>>(arr, arrsize)) {
+            std_argsort_withnan(arr, arg, 0, arrsize);
+        }
+        else {
+            argsort_64bit_<zmm_vector<double>>(
+                    arr, arg, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        }
+    }
+}
+
+
 template <>
 void avx512_argsort(int32_t* arr, int64_t *arg, int64_t arrsize)
 {
@@ -301,8 +356,13 @@ template <>
 void avx512_argsort(float* arr, int64_t *arg, int64_t arrsize)
 {
     if (arrsize > 1) {
-        argsort_64bit_<ymm_vector<float>>(
-                arr, arg, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        if (has_nan<ymm_vector<float>>(arr, arrsize)) {
+            std_argsort_withnan(arr, arg, 0, arrsize);
+        }
+        else {
+            argsort_64bit_<ymm_vector<float>>(
+                    arr, arg, 0, arrsize - 1, 2 * (int64_t)log2(arrsize));
+        }
     }
 }
 

src/avx512-64bit-common.h

@@ -71,6 +71,11 @@ struct ymm_vector<float> {
     {
         return _mm256_cmp_ps_mask(x, y, _CMP_EQ_OQ);
     }
+    template <int type>
+    static opmask_t fpclass(zmm_t x)
+    {
+        return _mm256_fpclass_ps_mask(x, type);
+    }
     template <int scale>
     static zmm_t
     mask_i64gather(zmm_t src, opmask_t mask, __m512i index, void const *base)
@@ -682,6 +687,10 @@ struct zmm_vector<double> {
         return _mm512_set_epi64(v1, v2, v3, v4, v5, v6, v7, v8);
     }
 
+    static zmm_t maskz_loadu(opmask_t mask, void const *mem)
+    {
+        return _mm512_maskz_loadu_pd(mask, mem);
+    }
     static opmask_t knot_opmask(opmask_t x)
     {
         return _knot_mask8(x);
@@ -694,6 +703,11 @@ struct zmm_vector<double> {
     {
         return _mm512_cmp_pd_mask(x, y, _CMP_EQ_OQ);
     }
+    template <int type>
+    static opmask_t fpclass(zmm_t x)
+    {
+        return _mm512_fpclass_pd_mask(x, type);
+    }
     template <int scale>
     static zmm_t
     mask_i64gather(zmm_t src, opmask_t mask, __m512i index, void const *base)