Guard against an overflow in sort tests (NVIDIA#1980)

Fixes: NVIDIA#1979
pciolkosz · Jul 17, 2024 · 347c051 · 347c051
1 parent be3a654
commit 347c051
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/cub/test/catch2_test_device_merge_sort.cu b/cub/test/catch2_test_device_merge_sort.cu
@@ -140,6 +140,8 @@ struct index_to_key_value_op
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::max());
   static constexpr std::size_t lowest_key_value =
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::lowest());
+  static_assert(sizeof(UnsignedIntegralKeyT) < sizeof(std::size_t),
+                "Calculation of num_distinct_key_values would overflow");
   static constexpr std::size_t num_distinct_key_values = (max_key_value - lowest_key_value + std::size_t{1ULL});
 
   __device__ __host__ UnsignedIntegralKeyT operator()(std::size_t index)
@@ -166,6 +168,8 @@ private:
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::max());
   static constexpr std::size_t lowest_key_value =
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::lowest());
+  static_assert(sizeof(UnsignedIntegralKeyT) < sizeof(std::size_t),
+                "Calculation of num_distinct_key_values would overflow");
   static constexpr std::size_t num_distinct_key_values = (max_key_value - lowest_key_value + std::size_t{1ULL});
 
   // item_count / num_distinct_key_values
@@ -406,7 +410,7 @@ CUB_TEST("DeviceMergeSort::StableSortPairs works for large inputs", "[merge][sor
 {
   using testing_types_tuple = c2h::get<0, TestType>;
   using key_t               = typename testing_types_tuple::key_t;
-  using offset_t            = typename c2h::get<0, TestType>::offset_t;
+  using offset_t            = typename testing_types_tuple::offset_t;
 
   // Clamp 64-bit offset type problem sizes to just slightly larger than 2^32 items
   auto num_items_ull =

diff --git a/thrust/testing/cuda/sort.cu b/thrust/testing/cuda/sort.cu
@@ -178,6 +178,8 @@ struct index_to_key_value_op
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::max());
   static constexpr std::size_t lowest_key_value =
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::lowest());
+  static_assert(sizeof(UnsignedIntegralKeyT) < sizeof(std::size_t),
+                "Calculation of num_distinct_key_values would overflow");
   static constexpr std::size_t num_distinct_key_values = (max_key_value - lowest_key_value + std::size_t{1ULL});
 
   __device__ __host__ UnsignedIntegralKeyT operator()(std::size_t index)
@@ -197,6 +199,8 @@ private:
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::max());
   static constexpr std::size_t lowest_key_value =
     static_cast<std::size_t>(::cuda::std::numeric_limits<UnsignedIntegralKeyT>::lowest());
+  static_assert(sizeof(UnsignedIntegralKeyT) < sizeof(std::size_t),
+                "Calculation of num_distinct_key_values would overflow");
   static constexpr std::size_t num_distinct_key_values = (max_key_value - lowest_key_value + std::size_t{1ULL});
 
   // item_count / num_distinct_key_values