Fix overflow in reduce

NVIDIA · Nov 15, 2022 · f313fa0 · f313fa0
1 parent 5ae7439
commit f313fa0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/cub/agent/agent_reduce.cuh b/cub/agent/agent_reduce.cuh
@@ -355,7 +355,7 @@ struct AgentReduce
   {
     AccumT thread_aggregate{};
 
-    if (even_share.block_offset + TILE_ITEMS > even_share.block_end)
+    if (even_share.block_end - even_share.block_offset < TILE_ITEMS)
     {
       // First tile isn't full (not all threads have valid items)
       int valid_items = even_share.block_end - even_share.block_offset;
@@ -377,14 +377,17 @@ struct AgentReduce
     even_share.block_offset += even_share.block_stride;
 
     // Consume subsequent full tiles of input
-    while (even_share.block_offset + TILE_ITEMS <= even_share.block_end)
+    while (even_share.block_offset <= even_share.block_end - TILE_ITEMS)
     {
       ConsumeTile<false>(thread_aggregate,
                          even_share.block_offset,
                          TILE_ITEMS,
                          Int2Type<true>(),
                          can_vectorize);
-      even_share.block_offset += even_share.block_stride;
+      const OffsetT new_offset = even_share.block_offset + even_share.block_stride;
+      even_share.block_offset = (new_offset > even_share.block_offset) 
+                              ? new_offset
+                              : even_share.block_end;
     }
 
     // Consume a partially-full tile

diff --git a/test/test_device_reduce.cu b/test/test_device_reduce.cu
@@ -1333,10 +1333,10 @@ __global__ void InitializeTestAccumulatorTypes(int num_items,
   }
 }
 
-template <typename T>
-void TestBigIndicesHelper(int magnitude)
+template <typename T, 
+          typename OffsetT>
+void TestBigIndicesHelper(OffsetT num_items)
 {
-  const std::size_t num_items = 1ll << magnitude;
   thrust::constant_iterator<T> const_iter(T{1});
   thrust::device_vector<std::size_t> out(1);
   std::size_t* d_out = thrust::raw_pointer_cast(out.data());
@@ -1360,10 +1360,10 @@ void TestBigIndicesHelper(int magnitude)
 template <typename T>
 void TestBigIndices()
 {
-  TestBigIndicesHelper<T>(30);
-  TestBigIndicesHelper<T>(31);
-  TestBigIndicesHelper<T>(32);
-  TestBigIndicesHelper<T>(33);
+  TestBigIndicesHelper<T, std::uint32_t>(1ull << 30);
+  TestBigIndicesHelper<T, std::uint32_t>(1ull << 31);
+  TestBigIndicesHelper<T, std::uint32_t>((1ull << 32) - 1);
+  TestBigIndicesHelper<T, std::uint64_t>(1ull << 33);
 }
 
 void TestAccumulatorTypes()