C2H/CUB: Add checked buffer utilities, stream-aware test support, and migrate segmented-scan tests

c2h/generators.cu

@@ -3,20 +3,28 @@
 
 #include <cub/device/device_copy.cuh>
 
+#include <thrust/detail/config/device_system.h>
+#include <thrust/detail/raw_pointer_cast.h>
 #include <thrust/for_each.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
 #include <thrust/tabulate.h>
+#include <thrust/version.h>
 
 #include <cuda/iterator>
 #include <cuda/std/optional>
+#include <cuda/std/span>
+#include <cuda/stream>
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
 
 #include <c2h/bfloat16.cuh>
 #include <c2h/custom_type.h>
 #include <c2h/detail/generators.cuh>
 #include <c2h/device_policy.h>
 #include <c2h/extended_types.h>
-#include <c2h/generators.h>
 #include <c2h/half.cuh>
 #include <c2h/vector.h>
 
@@ -65,31 +73,46 @@ public:
 
   float* prepare_random_generator(seed_t seed, std::size_t num_items)
   {
-    m_distribution.resize(num_items);
+    return prepare_random_generator(::cuda::stream_ref{::cudaStream_t{}}, seed, num_items);
+  }
+
+  float* prepare_random_generator(::cuda::stream_ref stream, seed_t seed, std::size_t num_items)
+  {
+    resize_distribution(num_items);
 
 #if C2H_HAS_CURAND
     curandSetPseudoRandomGeneratorSeed(m_gen, seed.get());
-#else
+#else // C2H_HAS_CURAND
     m_gen.seed(seed.get());
-#endif
+#endif // C2H_HAS_CURAND
 
-    generate();
+    generate(stream);
 
     return thrust::raw_pointer_cast(m_distribution.data());
   }
 
   // re-fills the currently held distribution vector with new random values
-  void generate()
+  void generate(::cuda::stream_ref stream)
   {
 #if C2H_HAS_CURAND
+    curandSetStream(m_gen, stream.get());
     curandGenerateUniform(m_gen, thrust::raw_pointer_cast(m_distribution.data()), m_distribution.size());
 #else
-    thrust::tabulate(device_policy, m_distribution.begin(), m_distribution.end(), i_to_rnd_t{m_gen});
+    thrust::tabulate(device_policy.on(stream.get()), m_distribution.begin(), m_distribution.end(), i_to_rnd_t{m_gen});
     m_gen.discard(m_distribution.size());
 #endif
   }
 
 private:
+  void resize_distribution(std::size_t num_items)
+  {
+#if THRUST_VERSION >= 300100
+    m_distribution.resize(num_items, thrust::no_init);
+#else // THRUST_VERSION >= 300100
+    m_distribution.resize(num_items);
+#endif // THRUST_VERSION >= 300100
+  }
+
 #if C2H_HAS_CURAND
   curandGenerator_t
 #else
@@ -113,6 +136,11 @@ float* prepare_random_data(seed_t seed, std::size_t num_items)
   return generator.value().prepare_random_generator(seed, num_items);
 }
 
+float* prepare_random_data(::cuda::stream_ref stream, seed_t seed, std::size_t num_items)
+{
+  return generator.value().prepare_random_generator(stream, seed, num_items);
+}
+
 void cleanup_generator()
 {
   _CCCL_VERIFY(generator.has_value(), "");
@@ -136,6 +164,18 @@ struct random_to_custom_t
 };
 
 void gen_custom_type_state(
+  seed_t seed,
+  char* d_out,
+  custom_type_state_t min,
+  custom_type_state_t max,
+  std::size_t elements,
+  std::size_t element_size)
+{
+  gen_custom_type_state(::cuda::stream_ref{::cudaStream_t{}}, seed, d_out, min, max, elements, element_size);
+}
+
+void gen_custom_type_state(
+  ::cuda::stream_ref stream,
   seed_t seed,
   char* d_out,
   custom_type_state_t /* min */,
@@ -144,8 +184,8 @@ void gen_custom_type_state(
   std::size_t element_size)
 {
   // FIXME(bgruber): implement min/max handling for custom_type_state_t
-  float* d_in = prepare_random_data(seed, elements * 2);
-  thrust::for_each(device_policy,
+  float* d_in = prepare_random_data(stream, seed, elements * 2);
+  thrust::for_each(device_policy.on(stream.get()),
                    thrust::counting_iterator<std::size_t>{0},
                    thrust::counting_iterator<std::size_t>{elements},
                    random_to_custom_t{d_in, d_out, element_size});

c2h/generators_gen_values.cu

@@ -4,20 +4,28 @@
 #include <thrust/tabulate.h>
 #include <thrust/transform.h>
 
+#include <cuda/stream>
+
 #include <c2h/bfloat16.cuh>
 #include <c2h/detail/generators.cuh>
 #include <c2h/device_policy.h>
 #include <c2h/extended_types.h>
-#include <c2h/generators.h>
 #include <c2h/half.cuh>
 
 namespace c2h::detail
 {
 template <typename T>
 void gen_values_between(seed_t seed, ::cuda::std::span<T> data, T min, T max)
 {
-  const auto* dist = prepare_random_data(seed, data.size());
-  thrust::transform(device_policy, dist, dist + data.size(), data.begin(), random_to_item_t<T>(min, max));
+  gen_values_between(::cuda::stream_ref{::cudaStream_t{}}, seed, data, min, max);
+}
+
+template <typename T>
+void gen_values_between(::cuda::stream_ref stream, seed_t seed, ::cuda::std::span<T> data, T min, T max)
+{
+  const auto* dist = prepare_random_data(stream, seed, data.size());
+  thrust::transform(
+    device_policy.on(stream.get()), dist, dist + data.size(), data.begin(), random_to_item_t<T>(min, max));
 }
 
 template <typename T>
@@ -39,8 +47,9 @@ void gen_values_cyclic(modulo_t mod, ::cuda::std::span<T> data)
   thrust::tabulate(device_policy, data.begin(), data.end(), counter_to_cyclic_item_t<T>{mod.get()});
 }
 
-#define INSTANTIATE_RND(TYPE) \
-  template void gen_values_between<TYPE>(seed_t, ::cuda::std::span<TYPE> data, TYPE min, TYPE max)
+#define INSTANTIATE_RND(TYPE)                                                                       \
+  template void gen_values_between<TYPE>(seed_t, ::cuda::std::span<TYPE> data, TYPE min, TYPE max); \
+  template void gen_values_between<TYPE>(::cuda::stream_ref, seed_t, ::cuda::std::span<TYPE> data, TYPE min, TYPE max)
 #define INSTANTIATE_MOD(TYPE) template void gen_values_cyclic<TYPE>(modulo_t, ::cuda::std::span<TYPE> data)
 
 #define INSTANTIATE(TYPE) \

c2h/generators_uniform_offsets.cu

@@ -1,16 +1,17 @@
 // SPDX-FileCopyrightText: Copyright (c) 2011-2025, NVIDIA CORPORATION. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-#include <c2h/generators.h>
-// #include <c2h/detail/generators.cuh>
-
-#include <thrust/device_ptr.h>
+#include <thrust/fill.h>
 #include <thrust/find.h>
 #include <thrust/scan.h>
 
 #include <cuda/std/cstdint>
 #include <cuda/std/span>
+#include <cuda/stream>
+
+#include <stdexcept>
 
+#include <c2h/detail/generators.cuh>
 #include <c2h/device_policy.h>
 
 namespace c2h::detail
@@ -30,12 +31,34 @@ template <typename T>
 std::size_t gen_uniform_offsets(
   seed_t seed, cuda::std::span<T> segment_offsets, T total_elements, T min_segment_size, T max_segment_size)
 {
-  gen_values_between(seed, segment_offsets, min_segment_size, max_segment_size);
-  *thrust::device_ptr<T>(&segment_offsets[total_elements]) = total_elements + 1;
-  thrust::exclusive_scan(device_policy, segment_offsets.begin(), segment_offsets.end(), segment_offsets.begin());
+  return gen_uniform_offsets(
+    ::cuda::stream_ref{::cudaStream_t{}}, seed, segment_offsets, total_elements, min_segment_size, max_segment_size);
+}
+
+template <typename T>
+std::size_t gen_uniform_offsets(
+  ::cuda::stream_ref stream,
+  seed_t seed,
+  cuda::std::span<T> segment_offsets,
+  T total_elements,
+  T min_segment_size,
+  T max_segment_size)
+{
+  const auto offsets_size = checked_uniform_offsets_size(total_elements);
+  if (segment_offsets.size() < offsets_size)
+  {
+    throw std::invalid_argument{"segment_offsets is too small for total_elements"};
+  }
+
+  const auto policy              = device_policy.on(stream.get());
+  const auto total_elements_size = offsets_size - 2;
+
+  gen_values_between(stream, seed, segment_offsets, min_segment_size, max_segment_size);
+  thrust::fill_n(policy, segment_offsets.begin() + total_elements_size, 1, total_elements + 1);
+  thrust::exclusive_scan(policy, segment_offsets.begin(), segment_offsets.end(), segment_offsets.begin());
   const auto iter =
-    thrust::find_if(device_policy, segment_offsets.begin(), segment_offsets.end(), greater_equal_op<T>{total_elements});
-  *thrust::device_ptr<T>(&*iter) = total_elements;
+    thrust::find_if(policy, segment_offsets.begin(), segment_offsets.end(), greater_equal_op<T>{total_elements});
+  thrust::fill_n(policy, iter, 1, total_elements);
   return iter - segment_offsets.begin() + 1;
 }
 
@@ -63,4 +86,33 @@ template std::size_t gen_uniform_offsets(
   uint64_t total_elements,
   uint64_t min_segment_size,
   uint64_t max_segment_size);
+
+template std::size_t gen_uniform_offsets(
+  ::cuda::stream_ref stream,
+  seed_t seed,
+  cuda::std::span<int32_t> segment_offsets,
+  int32_t total_elements,
+  int32_t min_segment_size,
+  int32_t max_segment_size);
+template std::size_t gen_uniform_offsets(
+  ::cuda::stream_ref stream,
+  seed_t seed,
+  cuda::std::span<uint32_t> segment_offsets,
+  uint32_t total_elements,
+  uint32_t min_segment_size,
+  uint32_t max_segment_size);
+template std::size_t gen_uniform_offsets(
+  ::cuda::stream_ref stream,
+  seed_t seed,
+  cuda::std::span<int64_t> segment_offsets,
+  int64_t total_elements,
+  int64_t min_segment_size,
+  int64_t max_segment_size);
+template std::size_t gen_uniform_offsets(
+  ::cuda::stream_ref stream,
+  seed_t seed,
+  cuda::std::span<uint64_t> segment_offsets,
+  uint64_t total_elements,
+  uint64_t min_segment_size,
+  uint64_t max_segment_size);
 } // namespace c2h::detail

c2h/generators_vector.cu

@@ -1,16 +1,20 @@
 // SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
+#include <thrust/detail/config/device_system.h>
+#include <thrust/for_each.h>
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/tabulate.h>
 
 #include <cuda/std/cstddef>
 #include <cuda/std/cstdint>
+#include <cuda/std/span>
+#include <cuda/stream>
 
 #include <c2h/detail/generators.cuh>
 #include <c2h/device_policy.h>
 #include <c2h/extended_types.h>
 #include <c2h/fill_striped.h>
-#include <c2h/generators.h>
 #include <c2h/vector.h>
 
 namespace c2h::detail
@@ -49,14 +53,21 @@ struct random_to_vec_item_t
 };
 
 #if THRUST_DEVICE_SYSTEM == THRUST_DEVICE_SYSTEM_CUDA
-#  define VEC_SPECIALIZATION(T)                                                                                   \
-    template <>                                                                                                   \
-    void gen_values_between(seed_t seed, ::cuda::std::span<T> data, T min, T max)                                 \
-    {                                                                                                             \
-      const auto* dist = prepare_random_data(seed, data.size());                                                  \
-      auto op          = random_to_vec_item_t<T, ::cuda::std::tuple_size_v<T>>{min, max, dist, data.data()};      \
-      thrust::for_each(                                                                                           \
-        device_policy, thrust::counting_iterator<size_t>{0}, thrust::counting_iterator<size_t>{data.size()}, op); \
+#  define VEC_SPECIALIZATION(T)                                                                              \
+    template <>                                                                                              \
+    void gen_values_between(::cuda::stream_ref stream, seed_t seed, ::cuda::std::span<T> data, T min, T max) \
+    {                                                                                                        \
+      const auto* dist = prepare_random_data(stream, seed, data.size());                                     \
+      auto op          = random_to_vec_item_t<T, ::cuda::std::tuple_size_v<T>>{min, max, dist, data.data()}; \
+      thrust::for_each(device_policy.on(stream.get()),                                                       \
+                       thrust::counting_iterator<size_t>{0},                                                 \
+                       thrust::counting_iterator<size_t>{data.size()},                                       \
+                       op);                                                                                  \
+    }                                                                                                        \
+    template <>                                                                                              \
+    void gen_values_between(seed_t seed, ::cuda::std::span<T> data, T min, T max)                            \
+    {                                                                                                        \
+      gen_values_between(::cuda::stream_ref{::cudaStream_t{}}, seed, data, min, max);                        \
     }
 
 VEC_SPECIALIZATION(char2);