Add CCCL_SEED_COUNT_OVERRIDE env var.

NVIDIA · Jun 28, 2024 · fa4857b · fa4857b
1 parent 9596cd0
commit fa4857b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/ci/test_cub.sh b/ci/test_cub.sh
@@ -77,6 +77,7 @@ fi
 if $COMPUTE_SANITIZER; then
   echo "Setting CCCL_TEST_MODE=compute-sanitizer-${TOOL}"
   export CCCL_TEST_MODE=compute-sanitizer-${TOOL}
+  export CCCL_SEED_COUNT_OVERRIDE=1
 fi
 
 for PRESET in ${PRESETS[@]}; do

diff --git a/cub/test/catch2_test_helper.h b/cub/test/catch2_test_helper.h
@@ -30,6 +30,7 @@
 #include <cub/util_compiler.cuh>
 
 #include <cstdint>
+#include <cstdlib>
 #include <tuple>
 #include <type_traits>
 
@@ -250,10 +251,22 @@ struct Catch::StringMaker<cudaError>
 
 #define CUB_TEST_STR(a) #a
 
+namespace detail
+{
+std::size_t adjust_seed_count(std::size_t requested)
+{
+  // Setting this environment variable forces a fixed number of seeds to be generated, regardless of the requested
+  // count. Set to 1 to reduce redundant, expensive testing when using sanitizers, etc.
+  static const char* override_str = std::getenv("CCCL_SEED_COUNT_OVERRIDE");
+  static int override             = override_str ? std::atoi(override_str) : 0;
+  return override_str ? override : requested;
+}
+} // namespace detail
+
 #define CUB_SEED(N)                                                                                                    \
   c2h::seed_t                                                                                                          \
   {                                                                                                                    \
     GENERATE_COPY(take(                                                                                                \
-      N,                                                                                                               \
+      detail::adjust_seed_count(N),                                                                                    \
       random(std::numeric_limits<unsigned long long int>::min(), std::numeric_limits<unsigned long long int>::max()))) \
   }