Fix kv-sort and kv-select NAN handling

sterrettm2 · sterrettm2 · commit a55a65557a98 · 2025-01-15T12:55:41.000-08:00
diff --git a/src/xss-common-keyvaluesort.hpp b/src/xss-common-keyvaluesort.hpp
@@ -16,6 +16,34 @@
 #include <omp.h>
 #endif
 
+/*
+ * Sort all the NAN's to end of the array and return the index of the last elem
+ * in the array which is not a nan
+ */
+template <typename T1, typename T2>
+X86_SIMD_SORT_INLINE arrsize_t move_nans_to_end_of_array(T1 *keys,
+                                                         T2 *vals,
+                                                         arrsize_t size)
+{
+    arrsize_t jj = size - 1;
+    arrsize_t ii = 0;
+    arrsize_t count = 0;
+    while (ii < jj) {
+        if (is_a_nan(keys[ii])) {
+            std::swap(keys[ii], keys[jj]);
+            std::swap(vals[ii], vals[jj]);
+            jj -= 1;
+            count++;
+        }
+        else {
+            ii += 1;
+        }
+    }
+    /* Haven't checked for nan when ii == jj */
+    if (is_a_nan(keys[ii])) { count++; }
+    return size - count - 1;
+}
+
 /*
  * Parition one ZMM register based on the pivot and returns the index of the
  * last element that is less than equal to the pivot.
@@ -538,11 +566,11 @@ X86_SIMD_SORT_INLINE void xss_qsort_kv(
 #endif // XSS_TEST_KEYVALUE_BASE_CASE
 
     if (minarrsize) {
-        arrsize_t nan_count = 0;
+        arrsize_t index_last_elem = arrsize - 1;
         if constexpr (xss::fp::is_floating_point_v<T1>) {
             if (UNLIKELY(hasnan)) {
-                nan_count
-                        = replace_nan_with_inf<full_vector<T1>>(keys, arrsize);
+                index_last_elem
+                        = move_nans_to_end_of_array(keys, indexes, arrsize);
             }
         }
         else {
@@ -565,24 +593,27 @@ X86_SIMD_SORT_INLINE void xss_qsort_kv(
             // Note that we do not use the if(...) clause built into OpenMP, because it causes a performance regression for small arrays
 #pragma omp parallel num_threads(thread_count)
 #pragma omp single
-            kvsort_<keytype, valtype>(
-                    keys, indexes, 0, arrsize - 1, maxiters, task_threshold);
+            kvsort_<keytype, valtype>(keys,
+                                      indexes,
+                                      0,
+                                      index_last_elem,
+                                      maxiters,
+                                      task_threshold);
         }
         else {
             kvsort_<keytype, valtype>(keys,
                                       indexes,
                                       0,
-                                      arrsize - 1,
+                                      index_last_elem,
                                       maxiters,
                                       std::numeric_limits<arrsize_t>::max());
         }
 #pragma omp taskwait
 #else
-        kvsort_<keytype, valtype>(keys, indexes, 0, arrsize - 1, maxiters, 0);
+        kvsort_<keytype, valtype>(
+                keys, indexes, 0, index_last_elem, maxiters, 0);
 #endif
 
-        replace_inf_with_nan(keys, arrsize, nan_count);
-
         if (descending) {
             std::reverse(keys, keys + arrsize);
             std::reverse(indexes, indexes + arrsize);
@@ -625,20 +656,18 @@ X86_SIMD_SORT_INLINE void xss_select_kv(T1 *keys,
     if (minarrsize) {
         if (descending) { k = arrsize - 1 - k; }
 
-        if constexpr (std::is_floating_point_v<T1>) {
-            arrsize_t nan_count = 0;
+        arrsize_t index_last_elem = arrsize - 1;
+        if constexpr (xss::fp::is_floating_point_v<T1>) {
             if (UNLIKELY(hasnan)) {
-                nan_count
-                        = replace_nan_with_inf<full_vector<T1>>(keys, arrsize);
+                index_last_elem
+                        = move_nans_to_end_of_array(keys, indexes, arrsize);
             }
-            kvselect_<keytype, valtype>(
-                    keys, indexes, k, 0, arrsize - 1, maxiters);
-            replace_inf_with_nan(keys, arrsize, nan_count);
         }
-        else {
-            UNUSED(hasnan);
+
+        UNUSED(hasnan);
+        if (index_last_elem >= k) {
             kvselect_<keytype, valtype>(
-                    keys, indexes, k, 0, arrsize - 1, maxiters);
+                    keys, indexes, k, 0, index_last_elem, maxiters);
         }
 
         if (descending) {
diff --git a/tests/test-keyvalue.cpp b/tests/test-keyvalue.cpp
@@ -26,7 +26,9 @@ class simdkvsort : public ::testing::Test {
                    "smallrange",
                    "max_at_the_end",
                    "random_5d",
-                   "rand_max"};
+                   "rand_max",
+                   "rand_with_nan",
+                   "rand_with_max_and_nan"};
     }
     std::vector<std::string> arrtype;
     std::vector<size_t> arrsize = std::vector<size_t>(1024);
@@ -123,27 +125,36 @@ bool is_kv_partialsorted(T1 *keys_comp,
     }
 
     // Now, we need to do some more work to handle keys exactly equal to the true kth
+    // There may be more values after the kth element with the same key,
+    // and thus we can find that the values of the kth elements do not match,
+    // even though the sort is correct.
+
     // First, fully kvsort both arrays
     xss::scalar::keyvalue_qsort<T1, T2>(keys_ref, vals_ref, size, true, false);
     xss::scalar::keyvalue_qsort<T1, T2>(
             keys_comp, vals_comp, size, true, false);
 
-    auto trueKth = keys_ref[k];
-    bool notFoundFirst = true;
+    auto trueKthKey = keys_ref[k];
+    bool foundFirstKthKey = false;
     size_t i = 0;
 
+    // Search forwards until we find the block of keys that match the kth key,
+    // then find where it ends
     for (; i < size; i++) {
-        if (notFoundFirst && cmp_eq(keys_ref[i], trueKth)) {
-            notFoundFirst = false;
+        if (!foundFirstKthKey && cmp_eq(keys_ref[i], trueKthKey)) {
+            foundFirstKthKey = true;
             i_start = i;
         }
-        else if (!notFoundFirst && !cmp_eq(keys_ref[i], trueKth)) {
+        else if (foundFirstKthKey && !cmp_eq(keys_ref[i], trueKthKey)) {
             break;
         }
     }
 
-    if (notFoundFirst) return false;
+    // kth key is somehow missing? Since we got that value from keys_ref, should be impossible
+    if (!foundFirstKthKey) { return false; }
 
+    // Check that the values in the kth key block match, so they are equivalent
+    // up to permutation, which is allowed since the sort is not stable
     if (!same_values(vals_ref + i_start, vals_comp + i_start, i - i_start)) {
         return false;
     }
@@ -156,7 +167,7 @@ TYPED_TEST_P(simdkvsort, test_kvsort_ascending)
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<T1> key = get_array<T1>(type, size);
             std::vector<T2> val = get_array<T2>(type, size);
@@ -187,7 +198,7 @@ TYPED_TEST_P(simdkvsort, test_kvsort_descending)
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<T1> key = get_array<T1>(type, size);
             std::vector<T2> val = get_array<T2>(type, size);
@@ -217,8 +228,9 @@ TYPED_TEST_P(simdkvsort, test_kvselect_ascending)
 {
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
+    auto cmp_eq = compare<T1, std::equal_to<T1>>();
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
 
@@ -237,7 +249,7 @@ TYPED_TEST_P(simdkvsort, test_kvselect_ascending)
             xss::scalar::keyvalue_qsort(
                     key.data(), val.data(), k, hasnan, false);
 
-            ASSERT_EQ(key[k], key_bckp[k]);
+            ASSERT_EQ(cmp_eq(key[k], key_bckp[k]), true);
 
             bool is_kv_partialsorted_
                     = is_kv_partialsorted<T1, T2>(key.data(),
@@ -260,8 +272,9 @@ TYPED_TEST_P(simdkvsort, test_kvselect_descending)
 {
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
+    auto cmp_eq = compare<T1, std::equal_to<T1>>();
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
 
@@ -280,7 +293,7 @@ TYPED_TEST_P(simdkvsort, test_kvselect_descending)
             xss::scalar::keyvalue_qsort(
                     key.data(), val.data(), k, hasnan, true);
 
-            ASSERT_EQ(key[k], key_bckp[k]);
+            ASSERT_EQ(cmp_eq(key[k], key_bckp[k]), true);
 
             bool is_kv_partialsorted_
                     = is_kv_partialsorted<T1, T2>(key.data(),
@@ -304,7 +317,7 @@ TYPED_TEST_P(simdkvsort, test_kvpartial_sort_ascending)
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
 
@@ -341,7 +354,7 @@ TYPED_TEST_P(simdkvsort, test_kvpartial_sort_descending)
     using T1 = typename std::tuple_element<0, decltype(TypeParam())>::type;
     using T2 = typename std::tuple_element<1, decltype(TypeParam())>::type;
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
 
diff --git a/tests/test-qsort-common.h b/tests/test-qsort-common.h
@@ -20,6 +20,12 @@
     ASSERT_TRUE(false) << msg << ". arr size = " << size \
                        << ", type = " << type << ", k = " << k;
 
+inline bool is_nan_test(std::string type)
+{
+    // Currently, determine whether the test uses nan just be checking if nan is in its name
+    return type.find("nan") != std::string::npos;
+}
+
 template <typename T>
 void IS_SORTED(std::vector<T> sorted, std::vector<T> arr, std::string type)
 {
diff --git a/tests/test-qsort.cpp b/tests/test-qsort.cpp
@@ -19,7 +19,8 @@ class simdsort : public ::testing::Test {
                    "max_at_the_end",
                    "random_5d",
                    "rand_max",
-                   "rand_with_nan"};
+                   "rand_with_nan",
+                   "rand_with_max_and_nan"};
     }
     std::vector<std::string> arrtype;
     std::vector<size_t> arrsize = std::vector<size_t>(1024);
@@ -30,7 +31,7 @@ TYPED_TEST_SUITE_P(simdsort);
 TYPED_TEST_P(simdsort, test_qsort_ascending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
 
@@ -52,7 +53,7 @@ TYPED_TEST_P(simdsort, test_qsort_ascending)
 TYPED_TEST_P(simdsort, test_qsort_descending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
 
@@ -74,7 +75,7 @@ TYPED_TEST_P(simdsort, test_qsort_descending)
 TYPED_TEST_P(simdsort, test_argsort_ascending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<TypeParam> arr = get_array<TypeParam>(type, size);
             std::vector<TypeParam> sortedarr = arr;
@@ -92,7 +93,7 @@ TYPED_TEST_P(simdsort, test_argsort_ascending)
 TYPED_TEST_P(simdsort, test_argsort_descending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             std::vector<TypeParam> arr = get_array<TypeParam>(type, size);
             std::vector<TypeParam> sortedarr = arr;
@@ -111,7 +112,7 @@ TYPED_TEST_P(simdsort, test_argsort_descending)
 TYPED_TEST_P(simdsort, test_qselect_ascending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
@@ -135,7 +136,7 @@ TYPED_TEST_P(simdsort, test_qselect_ascending)
 TYPED_TEST_P(simdsort, test_qselect_descending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
@@ -159,7 +160,7 @@ TYPED_TEST_P(simdsort, test_qselect_descending)
 TYPED_TEST_P(simdsort, test_argselect)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
             std::vector<TypeParam> arr = get_array<TypeParam>(type, size);
@@ -179,7 +180,7 @@ TYPED_TEST_P(simdsort, test_argselect)
 TYPED_TEST_P(simdsort, test_partial_qsort_ascending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             size_t k = rand() % size;
             std::vector<TypeParam> basearr = get_array<TypeParam>(type, size);
@@ -202,7 +203,7 @@ TYPED_TEST_P(simdsort, test_partial_qsort_ascending)
 TYPED_TEST_P(simdsort, test_partial_qsort_descending)
 {
     for (auto type : this->arrtype) {
-        bool hasnan = (type == "rand_with_nan") ? true : false;
+        bool hasnan = is_nan_test(type);
         for (auto size : this->arrsize) {
             // k should be at least 1
             size_t k = std::max((size_t)1, rand() % size);
diff --git a/utils/rand_array.h b/utils/rand_array.h
@@ -140,6 +140,26 @@ static std::vector<T> get_array(std::string arrtype,
             if (rand() & 0x1) { arr[ii] = val; }
         }
     }
+    else if (arrtype == "rand_with_max_and_nan") {
+        arr = get_uniform_rand_array<T>(arrsize, max, min);
+        T max_val;
+        T nan_val;
+        if constexpr (xss::fp::is_floating_point_v<T>) {
+            max_val = xss::fp::infinity<T>();
+            nan_val = xss::fp::quiet_NaN<T>();
+        }
+        else {
+            max_val = std::numeric_limits<T>::max();
+            nan_val = std::numeric_limits<T>::max();
+        }
+        for (size_t ii = 0; ii < arrsize; ++ii) {
+            int res = rand() % 4;
+            if (res == 2) { arr[ii] = max_val; }
+            else if (res == 3) {
+                arr[ii] = nan_val;
+            }
+        }
+    }
     else {
         std::cout << "Warning: unrecognized array type " << arrtype
                   << std::endl;
