Merge pull request #38 from r-devulap/handle-nan

Skip avx-512 argsort for arrays with NAN

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)

tests/test_argsort.cpp

@@ -174,12 +174,46 @@ TYPED_TEST_P(avx512argsort, test_reverse)
     }
 }
 
+TYPED_TEST_P(avx512argsort, test_array_with_nan)
+{
+    if (!cpu_has_avx512bw()) {
+        GTEST_SKIP() << "Skipping this test, it requires avx512bw ISA";
+    }
+    if (!std::is_floating_point<TypeParam>::value) {
+        GTEST_SKIP() << "Skipping this test, it is meant for float/double";
+    }
+    std::vector<int64_t> arrsizes;
+    for (int64_t ii = 2; ii <= 1024; ++ii) {
+        arrsizes.push_back(ii);
+    }
+    std::vector<TypeParam> arr;
+    for (auto &size : arrsizes) {
+        arr = get_uniform_rand_array<TypeParam>(size);
+        arr[0] = std::numeric_limits<TypeParam>::quiet_NaN();
+        arr[1] = std::numeric_limits<TypeParam>::quiet_NaN();
+        std::vector<int64_t> inx
+                = avx512_argsort<TypeParam>(arr.data(), arr.size());
+        std::vector<TypeParam> sort1;
+        for (size_t jj = 0; jj < size; ++jj) {
+            sort1.push_back(arr[inx[jj]]);
+        }
+        if ((!std::isnan(sort1[size-1])) || (!std::isnan(sort1[size-2]))) {
+            FAIL() << "NAN's aren't sorted to the end";
+        }
+        if (!std::is_sorted(sort1.begin(), sort1.end() - 2)) {
+            FAIL() << "Array isn't sorted";
+        }
+        arr.clear();
+    }
+}
+
 REGISTER_TYPED_TEST_SUITE_P(avx512argsort,
                             test_random,
                             test_reverse,
                             test_constant,
                             test_sorted,
-                            test_small_range);
+                            test_small_range,
+                            test_array_with_nan);
 
 using ArgSortTestTypes = testing::Types<int32_t,
                                         uint32_t,