Merge pull request #1832 from senior-zero/fix-main/github/reduce_offsets

Let CUB select reduce offsets

testing/cuda/reduce.cu

@@ -1,6 +1,7 @@
 #include <unittest/unittest.h>
 #include <thrust/reduce.h>
 #include <thrust/execution_policy.h>
+#include <thrust/iterator/constant_iterator.h>
 
 
 template<typename ExecutionPolicy, typename Iterator, typename T, typename Iterator2>
@@ -99,3 +100,22 @@ void TestReduceCudaStreamsNoSync()
 }
 DECLARE_UNITTEST(TestReduceCudaStreamsNoSync);
 
+#if defined(THRUST_RDC_ENABLED)
+void TestReduceLargeInput()
+{
+  using T = unsigned long long;
+  using OffsetT = std::size_t;
+  const OffsetT num_items = 1ull << 32;
+
+  thrust::constant_iterator<T> d_data(T{1});
+  thrust::device_vector<T> d_result(1);
+
+  reduce_kernel<<<1,1>>>(thrust::device, d_data, d_data + num_items, T{}, d_result.begin());
+  cudaError_t const err = cudaDeviceSynchronize();
+  ASSERT_EQUAL(cudaSuccess, err);
+
+  ASSERT_EQUAL(num_items, d_result[0]);
+}
+DECLARE_UNITTEST(TestReduceLargeInput);
+#endif
+

thrust/system/cuda/detail/reduce.h

@@ -943,11 +943,8 @@ T reduce_n_impl(execution_policy<Derived>& policy,
 
   size_t tmp_size = 0;
 
-  THRUST_INDEX_TYPE_DISPATCH2(status,
+  THRUST_INDEX_TYPE_DISPATCH(status,
     cub::DeviceReduce::Reduce,
-    (cub::DispatchReduce<
-        InputIt, T*, Size, BinaryOp, T
-    >::Dispatch),
     num_items,
     (NULL, tmp_size, first, reinterpret_cast<T*>(NULL),
         num_items_fixed, binary_op, init, stream));
@@ -970,11 +967,8 @@ T reduce_n_impl(execution_policy<Derived>& policy,
   // make this guarantee.
   T* ret_ptr = thrust::detail::aligned_reinterpret_cast<T*>(tmp.data().get());
   void* tmp_ptr = static_cast<void*>((tmp.data() + sizeof(T)).get());
-  THRUST_INDEX_TYPE_DISPATCH2(status,
+  THRUST_INDEX_TYPE_DISPATCH(status,
     cub::DeviceReduce::Reduce,
-    (cub::DispatchReduce<
-        InputIt, T*, Size, BinaryOp, T
-    >::Dispatch),
     num_items,
     (tmp_ptr, tmp_size, first, ret_ptr,
         num_items_fixed, binary_op, init, stream));