Workaround memory allocation issues on Windows for large size test cases (#690)

Radix sort and merge sort test cases that use the largest data sizes
hang or cause an SEH exception on Windows when using the latest version
of the HIP SDK.

Unfortunately, the HIP SDK issue interferes with our current strategies
for skipping tests that require too much memory for a device (querying
available device memory or examining the return value from hipMalloc),
making them unreliable.

While we wait for a fix, this change introduces a workaround by
reducing the maximum size these test cases can use on Windows.

Author: umfranzw

test/rocprim/test_device_merge_sort.cpp

@@ -423,7 +423,15 @@ void testLargeIndices()
     // at least some sizes that fit into device memory.
     using config = rocprim::merge_sort_config<256, 256, 1, 128, 128, 1, (1 << 17)>;
 
-    for(size_t size : test_utils::get_large_sizes(seeds[0]))
+    // On Windows, sizes above 2^34 cause issues that we can't currently catch by examining
+    // the hipMalloc return value or querying available memory. Workaround this for now
+    // by setting a different maximum size for that platform.
+#if defined(_WIN32)
+    const size_t max_pow2 = 34;
+#else
+    const size_t max_pow2 = 37;
+#endif
+    for(size_t size : test_utils::get_large_sizes<max_pow2>(seeds[0]))
     {
         SCOPED_TRACE(testing::Message() << "with size = " << size);
 

test/rocprim/test_device_radix_sort.hpp

@@ -1409,7 +1409,16 @@ inline void sort_keys_large_sizes()
 
     // Currently, CI enforces a hard limit of 96 GB on memory allocations.
     // Temporarily use sizes that will require less space than the limit.
-    const std::vector<size_t> sizes = test_utils::get_large_sizes<35>(seeds[0]);
+    // On Windows, sizes above 2^34 (that are still under the 96 GB limit)
+    // can hang due to issues that we can't currently catch by examining
+    // the hipMalloc return value or querying available memory. Workaround
+    // this for now by setting a different maximum size for that platform.
+#if defined(_WIN32)
+    const size_t max_pow2 = 34;
+#else
+    const size_t max_pow2 = 35;
+#endif
+    const std::vector<size_t> sizes = test_utils::get_large_sizes<max_pow2>(seeds[0]);
     for(const size_t size : sizes)
     {
         SCOPED_TRACE(testing::Message() << "with size = " << size);

test/rocprim/test_utils_data_generation.hpp

@@ -35,6 +35,7 @@
 #include <iterator>
 #include <random>
 #include <vector>
+#include <algorithm>
 
 namespace test_utils {
 
@@ -505,7 +506,7 @@ std::vector<size_t> get_sizes(T seed_value)
 template<unsigned int MaxPow2 = 37, class T>
 std::vector<size_t> get_large_sizes(T seed_value)
 {
-    std::vector<size_t> sizes = {
+    std::vector<size_t> test_sizes = {
         (size_t{1} << 30) - 1,
         size_t{1} << 31,
         (size_t{1} << 32) - 15,
@@ -519,7 +520,18 @@ std::vector<size_t> get_large_sizes(T seed_value)
                                               (size_t{1} << 30) + 1,
                                               (size_t{1} << MaxPow2) - 2,
                                               seed_value);
-    sizes.insert(sizes.end(), random_sizes.begin(), random_sizes.end());
+
+    std::vector<size_t> sizes(test_sizes.size() + random_sizes.size());
+    int count = 0;
+    auto predicate = [&count](const size_t& val) {
+        const bool result = (val <= (size_t{1} << MaxPow2));
+        count += (result ? 1 : 0);
+        return result;
+    };
+    std::copy_if(test_sizes.begin(), test_sizes.end(), sizes.begin(), predicate);
+    std::copy_if(random_sizes.begin(), random_sizes.end(), sizes.begin() + count, predicate);
+    sizes.resize(count);
+
     std::sort(sizes.begin(), sizes.end());
     return sizes;
 }