Fix scalar methods to treat NAN correctlt

Author: Raghuveer Devulapalli

lib/x86simdsort-scalar.h

@@ -1,36 +1,40 @@
 #include <algorithm>
 #include <numeric>
-#define UNUSED(x) (void)(x)
+#include "custom-compare.h"
+
 namespace xss {
 namespace scalar {
-    /* TODO: handle NAN */
     template <typename T>
     void qsort(T *arr, int64_t arrsize)
     {
-        std::sort(arr, arr + arrsize);
+        std::sort(arr, arr + arrsize, compare<T, std::less<T>>());
     }
     template <typename T>
     void qselect(T *arr, int64_t k, int64_t arrsize, bool hasnan)
     {
-        UNUSED(hasnan);
-        std::nth_element(arr, arr + k, arr + arrsize);
+        if (hasnan) {
+            std::nth_element(arr, arr + k, arr + arrsize, compare<T, std::less<T>>());
+        }
+        else {
+            std::nth_element(arr, arr + k, arr + arrsize);
+        }
     }
     template <typename T>
     void partial_qsort(T *arr, int64_t k, int64_t arrsize, bool hasnan)
     {
-        UNUSED(hasnan);
-        std::partial_sort(arr, arr + k, arr + arrsize);
+        if (hasnan) {
+            std::partial_sort(arr, arr + k, arr + arrsize, compare<T, std::less<T>>());
+        }
+        else {
+            std::partial_sort(arr, arr + k, arr + arrsize);
+        }
     }
     template <typename T>
     std::vector<int64_t> argsort(T *arr, int64_t arrsize)
     {
         std::vector<int64_t> arg(arrsize);
         std::iota(arg.begin(), arg.end(), 0);
-        std::sort(arg.begin(),
-                  arg.end(),
-                  [arr](int64_t left, int64_t right) -> bool {
-                      return arr[left] < arr[right];
-                  });
+        std::sort(arg.begin(), arg.end(), compare_arg<T, std::less<T>>(arr));
         return arg;
     }
     template <typename T>
@@ -41,9 +45,7 @@ namespace scalar {
         std::nth_element(arg.begin(),
                          arg.begin() + k,
                          arg.end(),
-                         [arr](int64_t left, int64_t right) -> bool {
-                             return arr[left] < arr[right];
-                         });
+                         compare_arg<T, std::less<T>>(arr));
         return arg;
     }
 

meson.build

@@ -26,7 +26,7 @@ subdir('benchmarks')
 
 libsimdsort = shared_library('x86simdsort',
                              'lib/x86simdsort.cpp',
-                             include_directories : [lib],
+                             include_directories : [utils, lib],
                              link_whole : [libtargets],
                              cpp_args : ['-O3'],
                             )

tests/test-qsort-common.h

@@ -2,6 +2,7 @@
 #define AVX512_TEST_COMMON
 
 #include "rand_array.h"
+#include "custom-compare.h"
 #include "x86simdsort.h"
 #include <gtest/gtest.h>
 
@@ -17,46 +18,6 @@
     ASSERT_TRUE(false) << msg << ". arr size = " << size \
                        << ", type = " << type << ", k = " << k;
 
-/*
- * Custom comparator class to handle NAN's: treats NAN  > INF
- */
-template <typename T, typename Comparator>
-struct compare {
-    static constexpr auto op = Comparator {};
-    bool operator()(const T a, const T b)
-    {
-        if constexpr (std::is_floating_point_v<T>) {
-            T inf = std::numeric_limits<T>::infinity();
-            if (!std::isunordered(a, b)) { return op(a, b); }
-            else if ((std::isnan(a)) && (!std::isnan(b))) {
-                return b == inf ? op(inf, 1.) : op(inf, b);
-            }
-            else if ((!std::isnan(a)) && (std::isnan(b))) {
-                return a == inf ? op(1., inf) : op(a, inf);
-            }
-            else {
-                return op(1., 1.);
-            }
-        }
-        else {
-            return op(a, b);
-        }
-    }
-};
-
-//template <typename T, typename Comparator>
-//struct compare_arg {
-//    compare_arg(std::vector<T> arr)
-//    {
-//        this->arr = arr;
-//    }
-//    bool operator()(const int64_t a, const int64_t b)
-//    {
-//        return compare<T, Comparator>()(arr[a], arr[b]);
-//    }
-//    std::vector<T> arr;
-//};
-
 template <typename T>
 void IS_SORTED(std::vector<T> sorted, std::vector<T> arr, std::string type)
 {

utils/custom-compare.h

@@ -0,0 +1,41 @@
+#include <limits>
+#include <cmath>
+/*
+ * Custom comparator class to handle NAN's: treats NAN  > INF
+ */
+template <typename T, typename Comparator>
+struct compare {
+    static constexpr auto op = Comparator {};
+    bool operator()(const T a, const T b)
+    {
+        if constexpr (std::is_floating_point_v<T>) {
+            T inf = std::numeric_limits<T>::infinity();
+            if (!std::isunordered(a, b)) { return op(a, b); }
+            else if ((std::isnan(a)) && (!std::isnan(b))) {
+                return b == inf ? op(inf, 1.) : op(inf, b);
+            }
+            else if ((!std::isnan(a)) && (std::isnan(b))) {
+                return a == inf ? op(1., inf) : op(a, inf);
+            }
+            else {
+                return op(1., 1.);
+            }
+        }
+        else {
+            return op(a, b);
+        }
+    }
+};
+
+template <typename T, typename Comparator>
+struct compare_arg {
+    compare_arg(const T* arr)
+    {
+        this->arr = arr;
+    }
+    bool operator()(const int64_t a, const int64_t b)
+    {
+        return compare<T, Comparator>()(arr[a], arr[b]);
+    }
+    const T* arr;
+};