Add more tests

Author: Raghuveer Devulapalli

tests/test_argsort.cpp

@@ -30,11 +30,11 @@ std::vector<int64_t> std_argsort(const std::vector<T> &array)
     return indices;
 }
 
-#define EXPECT_UNIQUE(arr) \
-    std::sort(arr.begin(), arr.end()); \
-    std::vector<int64_t> vec(arr.size()); \
-    std::iota(vec.begin(), vec.end(), 0); \
-    EXPECT_EQ(arr, vec) << "Indices aren't unique. Array size = " << arr.size();
+#define EXPECT_UNIQUE(sorted_arg) \
+    std::sort(sorted_arg.begin(), sorted_arg.end()); \
+    std::vector<int64_t> expected_arg(sorted_arg.size()); \
+    std::iota(expected_arg.begin(), expected_arg.end(), 0); \
+    EXPECT_EQ(sorted_arg, expected_arg) << "Indices aren't unique. Array size = " << sorted_arg.size();
 
 TYPED_TEST_P(avx512argsort, test_random)
 {
@@ -225,7 +225,7 @@ TYPED_TEST_P(avx512argsort, test_max_value_at_end_of_array)
         GTEST_SKIP() << "Skipping this test, it requires avx512bw ISA";
     }
     std::vector<int64_t> arrsizes;
-    for (int64_t ii = 1; ii <= 64; ++ii) {
+    for (int64_t ii = 1; ii <= 256; ++ii) {
         arrsizes.push_back(ii);
     }
     std::vector<TypeParam> arr;
@@ -250,12 +250,52 @@ TYPED_TEST_P(avx512argsort, test_max_value_at_end_of_array)
     }
 }
 
+TYPED_TEST_P(avx512argsort, test_all_inf_array)
+{
+    if (!cpu_has_avx512bw()) {
+        GTEST_SKIP() << "Skipping this test, it requires avx512bw ISA";
+    }
+    std::vector<int64_t> arrsizes;
+    for (int64_t ii = 1; ii <= 256; ++ii) {
+        arrsizes.push_back(ii);
+    }
+    std::vector<TypeParam> arr;
+    for (auto &size : arrsizes) {
+        arr = get_uniform_rand_array<TypeParam>(size);
+        if (std::numeric_limits<TypeParam>::has_infinity) {
+            for (int64_t jj = 1; jj <= size; ++jj) {
+                if (rand() % 0x1) {
+                    arr.push_back(std::numeric_limits<TypeParam>::infinity());
+                }
+            }
+        }
+        else {
+            for (int64_t jj = 1; jj <= size; ++jj) {
+                if (rand() % 0x1) {
+                    arr.push_back(std::numeric_limits<TypeParam>::max());
+                }
+            }
+        }
+        std::vector<int64_t> inx = avx512_argsort(arr.data(), arr.size());
+        std::vector<TypeParam> sorted;
+        for (size_t jj = 0; jj < size; ++jj) {
+            sorted.push_back(arr[inx[jj]]);
+        }
+        if (!std::is_sorted(sorted.begin(), sorted.end())) {
+            EXPECT_TRUE(false) << "Array of size " << size << "is not sorted";
+        }
+        EXPECT_UNIQUE(inx)
+        arr.clear();
+    }
+}
+
 REGISTER_TYPED_TEST_SUITE_P(avx512argsort,
                             test_random,
                             test_reverse,
                             test_constant,
                             test_sorted,
                             test_small_range,
+                            test_all_inf_array,
                             test_array_with_nan,
                             test_max_value_at_end_of_array);
 

tests/test_qsort.hpp

@@ -128,8 +128,41 @@ TYPED_TEST_P(avx512_sort, test_small_range)
         GTEST_SKIP() << "Skipping this test, it requires avx512bw";
     }
 }
+
+TYPED_TEST_P(avx512_sort, test_max_value_at_end_of_array)
+{
+    if (!cpu_has_avx512bw()) {
+        GTEST_SKIP() << "Skipping this test, it requires avx512bw ISA";
+    }
+    if ((sizeof(TypeParam) == 2) && (!cpu_has_avx512_vbmi2())) {
+        GTEST_SKIP() << "Skipping this test, it requires avx512_vbmi2";
+    }
+    std::vector<int64_t> arrsizes;
+    for (int64_t ii = 1; ii <= 1024; ++ii) {
+        arrsizes.push_back(ii);
+    }
+    std::vector<TypeParam> arr;
+    std::vector<TypeParam> sortedarr;
+    for (auto &size : arrsizes) {
+        arr = get_uniform_rand_array<TypeParam>(size);
+        if (std::numeric_limits<TypeParam>::has_infinity) {
+            arr[size - 1] = std::numeric_limits<TypeParam>::infinity();
+        }
+        else {
+            arr[size - 1] = std::numeric_limits<TypeParam>::max();
+        }
+        sortedarr = arr;
+        avx512_qsort(arr.data(), arr.size());
+        std::sort(sortedarr.begin(), sortedarr.end());
+        EXPECT_EQ(sortedarr, arr) << "Array size = " << size;
+        arr.clear();
+        sortedarr.clear();
+    }
+}
+
 REGISTER_TYPED_TEST_SUITE_P(avx512_sort,
                             test_random,
                             test_reverse,
                             test_constant,
-                            test_small_range);
+                            test_small_range,
+                            test_max_value_at_end_of_array);