[CUB] GPU-to-GPU DeviceReduce with device resident problem size

cub/benchmarks/bench/reduce/deferred_deterministic.cu

@@ -0,0 +1,93 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cub/device/device_reduce.cuh>
+
+#include <thrust/detail/raw_pointer_cast.h>
+#include <thrust/device_vector.h>
+
+#include <cuda/argument>
+#include <cuda/execution.determinism.h>
+#include <cuda/execution.require.h>
+#include <cuda/std/functional>
+#include <cuda/std/utility>
+
+#include <nvbench_helper.cuh>
+
+#include <nvbench/range.cuh>
+#include <nvbench/types.cuh>
+
+// %RANGE% TUNE_ITEMS_PER_THREAD ipt 3:24:1
+// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
+
+#if !TUNE_BASE
+struct policy_selector_t
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::ReducePolicy
+  {
+    const auto p = cub::ReducePassPolicy{
+      TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, 1, cub::BLOCK_REDUCE_RAKING, cub::LOAD_DEFAULT};
+    return {p, p};
+  }
+};
+#endif // !TUNE_BASE
+
+template <class T, class OffsetT>
+void deterministic_sum(nvbench::state& state, nvbench::type_list<T, OffsetT>)
+try
+{
+  using init_value_t = T;
+
+  if (!cuda::std::in_range<OffsetT>(state.get_int64("Elements{io}")))
+  {
+    state.skip("Skipping: Elements{io} is not representable by OffsetT.");
+    return;
+  }
+  const auto elements = static_cast<OffsetT>(state.get_int64("Elements{io}"));
+
+  thrust::device_vector<T> in = generate(elements);
+  thrust::device_vector<T> out(1, thrust::no_init);
+  thrust::device_vector<OffsetT> device_num_items(1, elements);
+
+  auto d_in        = thrust::raw_pointer_cast(in.data());
+  auto d_out       = thrust::raw_pointer_cast(out.data());
+  auto d_num_items = thrust::raw_pointer_cast(device_num_items.data());
+
+  // Enable throughput calculations and add "Size" column to results.
+  state.add_element_count(elements);
+  state.add_global_memory_reads<T>(elements, "Size");
+  state.add_global_memory_writes<T>(1);
+
+  caching_allocator_t alloc;
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cub_bench_env(
+      alloc,
+      launch,
+      cuda::execution::require(cuda::execution::determinism::gpu_to_gpu)
+#if !TUNE_BASE
+        ,
+      cuda::execution::tune(policy_selector_t{})
+#endif // !TUNE_BASE
+    );
+    _CCCL_TRY_CUDA_API(
+      cub::DeviceReduce::Reduce,
+      "Reduce failed",
+      d_in,
+      d_out,
+      cuda::args::deferred{d_num_items},
+      cuda::std::plus<>{},
+      init_value_t{},
+      env);
+  });
+}
+catch (const std::bad_alloc&)
+{
+  state.skip("Skipping: out of memory.");
+}
+
+using types = nvbench::type_list<float, double>;
+NVBENCH_BENCH_TYPES(deterministic_sum, NVBENCH_TYPE_AXES(types, offset_types))
+  .set_name("base")
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  // 2^31 and 2^32 straddle INT32_MAX to cover the code paths for problem sizes that exceed a single 32-bit chunk
+  .add_int64_power_of_two_axis("Elements{io}", {16, 20, 24, 28, 31, 32});

cub/benchmarks/bench/reduce/deterministic.cu

@@ -5,6 +5,7 @@
 
 #include <cuda/execution.determinism.h>
 #include <cuda/execution.require.h>
+#include <cuda/std/utility>
 
 #include <nvbench_helper.cuh>
 
@@ -26,12 +27,18 @@ struct policy_selector_t
 };
 #endif // !TUNE_BASE
 
-template <class T>
-void deterministic_sum(nvbench::state& state, nvbench::type_list<T>)
+template <class T, class OffsetT>
+void deterministic_sum(nvbench::state& state, nvbench::type_list<T, OffsetT>)
+try
 {
   using init_value_t = T;
 
-  const auto elements = static_cast<int>(state.get_int64("Elements{io}"));
+  if (!cuda::std::in_range<OffsetT>(state.get_int64("Elements{io}")))
+  {
+    state.skip("Skipping: Elements{io} is not representable by OffsetT.");
+    return;
+  }
+  const auto elements = static_cast<OffsetT>(state.get_int64("Elements{io}"));
 
   thrust::device_vector<T> in = generate(elements);
   thrust::device_vector<T> out(1);
@@ -43,7 +50,7 @@ void deterministic_sum(nvbench::state& state, nvbench::type_list<T>)
   state.add_global_memory_writes<T>(out.size());
 
   caching_allocator_t alloc;
-  state.exec(nvbench::exec_tag::no_batch | nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
     auto env = cub_bench_env(
       alloc,
       launch,
@@ -57,9 +64,14 @@ void deterministic_sum(nvbench::state& state, nvbench::type_list<T>)
       cub::DeviceReduce::Reduce, "Reduce failed", d_in, d_out, elements, cuda::std::plus<>{}, init_value_t{}, env);
   });
 }
+catch (const std::bad_alloc&)
+{
+  state.skip("Skipping: out of memory.");
+}
 
 using types = nvbench::type_list<float, double>;
-NVBENCH_BENCH_TYPES(deterministic_sum, NVBENCH_TYPE_AXES(types))
+NVBENCH_BENCH_TYPES(deterministic_sum, NVBENCH_TYPE_AXES(types, offset_types))
   .set_name("base")
-  .set_type_axes_names({"T{ct}"})
-  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4));
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  // 2^31 and 2^32 straddle INT32_MAX to cover the code paths for problem sizes that exceed a single 32-bit chunk
+  .add_int64_power_of_two_axis("Elements{io}", {16, 20, 24, 28, 31, 32});

cub/cub/device/device_reduce.cuh

@@ -141,8 +141,7 @@ inline constexpr bool is_non_deterministic_v =
 //! updating or recapturing the graph. Compile-time and runtime bounds are accepted as caller preconditions, but do not
 //! currently change temporary storage, grid dimensions, or pass selection. Since the problem size is not available to
 //! the host, the first reduction pass launches CUB's maximum grid and may launch more blocks than the actual problem
-//! size needs. Extra blocks return without reducing input. Deferred problem sizes are not currently supported with
-//! ``gpu_to_gpu`` determinism.
+//! size needs. Extra blocks return without reducing input.
 //!
 //! Determinism
 //! ====================================
@@ -205,32 +204,28 @@ private:
 
     if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__gpu_to_gpu)
     {
-      if constexpr (args_traits_t::is_deferred)
-      {
-        static_assert(!args_traits_t::is_deferred, "cuda::args::deferred is not supported with gpu_to_gpu determinism");
-        // NVHPC requires a return statement even though the static assertion makes this branch ill-formed.
-        return cudaErrorNotSupported;
-      }
-      else
-      {
-        // Only instantiated with `plus<float|double>`; RFA hardcodes `deterministic_sum_t<accum_t>`.
-        (void) reduction_op;
-        using default_policy_selector = detail::reduce::
-          policy_selector_from_types<accum_t, offset_t, detail::rfa::deterministic_sum_t<accum_t>, Determinism>;
-        return detail::dispatch_with_env_and_tuning<default_policy_selector>(
-          env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) {
-            return detail::rfa::dispatch<InputIteratorT, OutputIteratorT, offset_t, T, TransformOpT, accum_t>(
-              storage,
-              bytes,
-              d_in,
-              d_out,
-              detail::make_num_items_dispatch_arg(num_items),
-              init,
-              stream,
-              transform_op,
-              policy_selector);
-          });
-      }
+      // Only instantiated with `plus<float|double>`; RFA hardcodes `deterministic_sum_t<accum_t>`.
+      (void) reduction_op;
+      using default_policy_selector = detail::reduce::
+        policy_selector_from_types<accum_t, offset_t, detail::rfa::deterministic_sum_t<accum_t>, Determinism>;
+      return detail::dispatch_with_env_and_tuning<default_policy_selector>(
+        env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) {
+          return detail::rfa::dispatch<InputIteratorT,
+                                       OutputIteratorT,
+                                       decltype(detail::make_num_items_dispatch_arg(num_items)),
+                                       T,
+                                       TransformOpT,
+                                       accum_t>(
+            storage,
+            bytes,
+            d_in,
+            d_out,
+            detail::make_num_items_dispatch_arg(num_items),
+            init,
+            stream,
+            transform_op,
+            policy_selector);
+        });
     }
     else if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__not_guaranteed)
     {