diff --git a/cub/benchmarks/bench/reduce/deferred_deterministic.cu b/cub/benchmarks/bench/reduce/deferred_deterministic.cu new file mode 100644 index 00000000000..dd03e3eb745 --- /dev/null +++ b/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 + +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include + +// %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 +void deterministic_sum(nvbench::state& state, nvbench::type_list) +try +{ + using init_value_t = T; + + if (!cuda::std::in_range(state.get_int64("Elements{io}"))) + { + state.skip("Skipping: Elements{io} is not representable by OffsetT."); + return; + } + const auto elements = static_cast(state.get_int64("Elements{io}")); + + thrust::device_vector in = generate(elements); + thrust::device_vector out(1, thrust::no_init); + thrust::device_vector 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(elements, "Size"); + state.add_global_memory_writes(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; +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}); diff --git a/cub/benchmarks/bench/reduce/deterministic.cu b/cub/benchmarks/bench/reduce/deterministic.cu index 31d1f074c0d..02981685778 100644 --- a/cub/benchmarks/bench/reduce/deterministic.cu +++ b/cub/benchmarks/bench/reduce/deterministic.cu @@ -5,6 +5,7 @@ #include #include +#include #include @@ -26,12 +27,18 @@ struct policy_selector_t }; #endif // !TUNE_BASE -template -void deterministic_sum(nvbench::state& state, nvbench::type_list) +template +void deterministic_sum(nvbench::state& state, nvbench::type_list) +try { using init_value_t = T; - const auto elements = static_cast(state.get_int64("Elements{io}")); + if (!cuda::std::in_range(state.get_int64("Elements{io}"))) + { + state.skip("Skipping: Elements{io} is not representable by OffsetT."); + return; + } + const auto elements = static_cast(state.get_int64("Elements{io}")); thrust::device_vector in = generate(elements); thrust::device_vector out(1); @@ -43,7 +50,7 @@ void deterministic_sum(nvbench::state& state, nvbench::type_list) state.add_global_memory_writes(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) 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; -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}); diff --git a/cub/cub/device/device_reduce.cuh b/cub/cub/device/device_reduce.cuh index 7d1c80734cb..483d71cc1d3 100644 --- a/cub/cub/device/device_reduce.cuh +++ b/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`; RFA hardcodes `deterministic_sum_t`. - (void) reduction_op; - using default_policy_selector = detail::reduce:: - policy_selector_from_types, Determinism>; - return detail::dispatch_with_env_and_tuning( - env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { - return detail::rfa::dispatch( - storage, - bytes, - d_in, - d_out, - detail::make_num_items_dispatch_arg(num_items), - init, - stream, - transform_op, - policy_selector); - }); - } + // Only instantiated with `plus`; RFA hardcodes `deterministic_sum_t`. + (void) reduction_op; + using default_policy_selector = detail::reduce:: + policy_selector_from_types, Determinism>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::rfa::dispatch( + 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) { diff --git a/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh b/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh index e0751ceb97c..e5fdea03ead 100644 --- a/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh +++ b/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: BSD-3 //! @file @@ -18,6 +18,7 @@ #endif // no system header #include +#include #include #include #include @@ -26,6 +27,7 @@ #include #include +#include #include #include #include @@ -172,6 +174,9 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok ReducePolicy active_policy, KernelLauncherFactory launcher_factory) { + // Immediate chunk sizes are passed to the kernels as-is; deferred problem sizes are read on device. + using num_items_kernel_t = detail::parameter_from_host_t; + int sm_count; if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) { @@ -180,8 +185,13 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok KernelConfig reduce_config; if (const auto error = CubDebug(reduce_config.__init( - detail::reduce:: - DeterministicDeviceReduceKernel, + detail::reduce::DeterministicDeviceReduceKernel< + PolicySelector, + InputIteratorT, + num_items_kernel_t, + ReductionOpT, + DeterministicAccumT, + TransformOpT>, active_policy.multi_tile))) { return error; @@ -191,16 +201,27 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok const int max_blocks = reduce_device_occupancy * detail::subscription_factor; const int num_items_per_chunk = ::cuda::std::numeric_limits<::cuda::std::int32_t>::max(); - const int num_chunks = static_cast(::cuda::ceil_div(num_items, num_items_per_chunk)); - const int chunk_tile_grid_size = ::cuda::ceil_div(num_items_per_chunk, reduce_config.tile_size); - const int chunk_grid_size = ::cuda::std::min(max_blocks, chunk_tile_grid_size); + // A deferred problem size cannot be read on the host, so these defaults stand: the kernel consumes the whole + // problem in a single launch with the worst-case grid, whose surplus blocks exit early. + int num_chunks = 1; + int chunk_grid_size = max_blocks; + int partial_chunk_size = 0; + bool has_partial_chunk = false; + int last_chunk_grid_size = max_blocks; + if constexpr (!::cuda::args::__traits::is_deferred) + { + num_chunks = static_cast(::cuda::ceil_div(num_items, num_items_per_chunk)); + + const int chunk_tile_grid_size = ::cuda::ceil_div(num_items_per_chunk, reduce_config.tile_size); + chunk_grid_size = ::cuda::std::min(max_blocks, chunk_tile_grid_size); - const int partial_chunk_size = num_items % num_items_per_chunk; - const bool has_partial_chunk = partial_chunk_size != 0; - const int last_chunk_tile_grid_size = ::cuda::ceil_div(partial_chunk_size, reduce_config.tile_size); + partial_chunk_size = num_items % num_items_per_chunk; + has_partial_chunk = partial_chunk_size != 0; + const int last_chunk_tile_grid_size = ::cuda::ceil_div(partial_chunk_size, reduce_config.tile_size); - const int last_chunk_grid_size = ::cuda::std::min(max_blocks, last_chunk_tile_grid_size); + last_chunk_grid_size = ::cuda::std::min(max_blocks, last_chunk_tile_grid_size); + } const int reduce_grid_size = chunk_grid_size * (num_chunks - 1) + last_chunk_grid_size; @@ -236,6 +257,19 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok const auto current_grid_size = static_cast(num_current_items == num_items_per_chunk ? chunk_grid_size : last_chunk_grid_size); + // An immediate problem size is passed as the current chunk size; a deferred problem size is read on device and + // `num_current_items` holds the worst-case chunk size, which must not be passed to the kernel. + const auto kernel_num_items = [&] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return detail::reduce::make_num_items_kernel_arg(num_items); + } + else + { + return num_current_items; + } + }(); + // Log device_reduce_sweep_kernel configuration #ifdef CUB_DEBUG_LOG _CubLog("Invoking DeterministicDeviceReduceKernel<<<%d, %d, 0, %lld>>>(), %d items " @@ -251,12 +285,13 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok launcher_factory(current_grid_size, active_policy.multi_tile.threads_per_block, 0, stream) .doit(detail::reduce::DeterministicDeviceReduceKernel, d_in, d_chunk_block_reductions, - num_current_items, + kernel_num_items, reduction_op, transform_op, current_grid_size))) @@ -292,22 +327,46 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok active_policy.single_tile.items_per_thread); #endif // CUB_DEBUG_LOG - // Invoke DeterministicDeviceReduceSingleTileKernel - if (const auto error = CubDebug( - launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) - .doit(detail::reduce::DeterministicDeviceReduceSingleTileKernel< - PolicySelector, - DeterministicAccumT*, - OutputIteratorT, - ReductionOpT, - InitValueT, - DeterministicAccumT>, - d_block_reductions, - d_out, - reduce_grid_size, - reduction_op, - init, - ::cuda::std::identity{}))) + // Invoke DeterministicDeviceReduceSingleTileKernel/DeterministicDeviceReduceDeferredSingleTileKernel + const auto second_pass_error = [&] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceDeferredSingleTileKernel< + PolicySelector, + DeterministicAccumT*, + OutputIteratorT, + num_items_kernel_t, + ReductionOpT, + InitValueT, + DeterministicAccumT>, + d_block_reductions, + d_out, + detail::reduce::make_num_items_kernel_arg(num_items), + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}); + } + else + { + return launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceSingleTileKernel< + PolicySelector, + DeterministicAccumT*, + OutputIteratorT, + ReductionOpT, + InitValueT, + DeterministicAccumT>, + d_block_reductions, + d_out, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}); + } + }(); + if (const auto error = CubDebug(second_pass_error)) { return error; } @@ -326,10 +385,12 @@ template , - typename PolicySelector = - policy_selector_from_types, __determinism_t::__gpu_to_gpu>, + typename TransformOpT = ::cuda::std::identity, + typename AccumT = accum_t, + typename PolicySelector = policy_selector_from_types, + deterministic_sum_t, + __determinism_t::__gpu_to_gpu>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( void* d_temp_storage, @@ -363,35 +424,40 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( })) #endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) - const auto tile_items = static_cast(active_policy.single_tile.threads_per_block) - * static_cast(active_policy.single_tile.items_per_thread); - using deterministic_add_t = deterministic_sum_t; using input_unwrapped_it_t = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t; input_unwrapped_it_t d_in_unwrapped = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in); - if (num_items <= tile_items) + // A deferred problem size cannot be compared against the single-tile capacity on the host, so a deferred + // reduction always takes the two-pass path. + if constexpr (!::cuda::args::__traits::is_deferred) { - return invoke_single_tile( - d_temp_storage, - temp_storage_bytes, - d_in_unwrapped, - d_out, - num_items, - deterministic_add_t{}, - init, - stream, - transform_op, - active_policy, - launcher_factory); + const auto tile_items = static_cast(active_policy.single_tile.threads_per_block) + * static_cast(active_policy.single_tile.items_per_thread); + + if (num_items <= tile_items) + { + return invoke_single_tile( + d_temp_storage, + temp_storage_bytes, + d_in_unwrapped, + d_out, + num_items, + deterministic_add_t{}, + init, + stream, + transform_op, + active_policy, + launcher_factory); + } } return invoke_passes #include #include #include #include +#include #include +#include #include +#include +#include +#include +#include +#include +#include CUB_NAMESPACE_BEGIN namespace detail::reduce { +// Type a kernel problem-size argument resolves to on device: an immediate problem size keeps its integral type, a +// deferred source resolves to an integer wide enough for its element. Only a signed 32-bit element is guaranteed +// to fit a single 32-bit chunk; every other element type is widened to 64 bits and consumed in chunks, with an +// unsigned element resolving to an unsigned type so that its full value range is representable. +template > +struct deterministic_num_items +{ + using type = KernelNumItemsT; +}; + +template +struct deterministic_num_items +{ + using element_t = it_value_t; + using type = ::cuda::std::conditional_t< + ::cuda::std::is_signed_v && sizeof(element_t) == 4, + int, + ::cuda::std::conditional_t<::cuda::std::is_signed_v, ::cuda::std::int64_t, ::cuda::std::uint64_t>>; +}; + +template +using deterministic_num_items_t = typename deterministic_num_items::type; + +// Number of first-pass blocks that produce a partial for a deferred problem size: the worst-case launch is trimmed +// to the grid the host would have computed had it been able to read the problem size. Must be used consistently by +// both reduction passes. +template +[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE int +deferred_reduce_grid_size(NumItemsT num_items, int launched_grid_size) noexcept +{ + constexpr ReducePassPolicy policy = detail::current_policy().multi_tile; + constexpr int tile_size = policy.threads_per_block * policy.items_per_thread; + const NumItemsT num_tiles = ::cuda::ceil_div(num_items, NumItemsT{tile_size}); + return static_cast(::cuda::std::min(static_cast(launched_grid_size), num_tiles)); +} + /** * @brief Deterministically Reduce region kernel entry point (multi-block). Computes privatized * reductions, one per thread block in deterministic fashion @@ -36,6 +81,9 @@ namespace detail::reduce * @tparam InputIteratorT * Random-access input iterator type for reading input items @iterator * + * @tparam KernelNumItemsT + * Type of integral problem size or a deferred problem-size descriptor + * * @tparam ReductionOpT * Binary reduction functor type having member * `auto operator()(const T &a, const U &b)` @@ -49,18 +97,23 @@ namespace detail::reduce * @param[out] d_out * Pointer to the output aggregate * - * @param[in] num_items - * Total number of input data items + * @param[in] kernel_num_items + * Immediate problem size or a deferred problem-size descriptor * * @param[in] reduction_op * Binary reduction functor */ -template +template _CCCL_KERNEL_ATTRIBUTES __launch_bounds__(int(current_policy().multi_tile.threads_per_block)) void DeterministicDeviceReduceKernel( InputIteratorT d_in, AccumT* d_out, - int num_items, + const KernelNumItemsT kernel_num_items, ReductionOpT reduction_op, TransformOpT transform_op, const int reduce_grid_size) @@ -69,6 +122,32 @@ __launch_bounds__(int(current_policy().multi_tile.threads_per_bl constexpr int items_per_thread = policy.items_per_thread; constexpr int threads_per_block = policy.threads_per_block; + // A 64-bit deferred problem size is consumed in a single launch that loops over 32-bit chunks in the kernel. + using num_items_t = deterministic_num_items_t; + + const num_items_t num_items = detail::parameter_from_device(kernel_num_items); + + // The worst-case grid of a deferred problem size is trimmed to the blocks that receive at least one tile. Both the + // early exit and the loop stride must use the trimmed grid so that the remaining blocks cover the whole input. + const int active_grid_size = [&] { + if constexpr (::cuda::std::is_integral_v) + { + return reduce_grid_size; + } + else + { + return detail::reduce::deferred_reduce_grid_size(num_items, reduce_grid_size); + } + }(); + + if constexpr (!::cuda::std::is_integral_v) + { + if (static_cast(blockIdx.x) >= active_grid_size) + { + return; + } + } + using block_reduce_t = BlockReduce; // Shared memory storage @@ -91,40 +170,100 @@ __launch_bounds__(int(current_policy().multi_tile.threads_per_bl AccumT thread_aggregate{}; int count = 0; - int n_threads = reduce_grid_size * threads_per_block; + int n_threads = active_grid_size * threads_per_block; - _CCCL_PRAGMA_UNROLL_FULL() - for (unsigned i = tid; i < static_cast(num_items); i += (n_threads * items_per_thread)) + if constexpr (sizeof(num_items_t) == 8) { - ftype items[items_per_thread] = {}; - for (int j = 0; j < items_per_thread; j++) + // Loop over 32-bit chunks so that the hot loop keeps the index arithmetic and register footprint of the 32-bit + // path below; only the count of remaining items is 64-bit and the iterator advances once per chunk. Binned + // accumulation is exact, so accumulating across chunk boundaries produces the same bits as the host-side chunking + // of an immediate problem size. + // TODO(NaderAlAwar): The chunk body is intentionally duplicated from the 32-bit loop below because extracting it + // into a device function changes the SASS of the immediate instantiations. Changes here must also be applied to + // the copy below. + constexpr auto num_items_per_chunk = num_items_t{::cuda::std::numeric_limits::max()}; + for (num_items_t remaining = num_items; remaining != 0;) { - const unsigned idx = i + j * n_threads; - if (idx < static_cast(num_items)) + const auto chunk_num_items = static_cast(::cuda::std::min(remaining, num_items_per_chunk)); + + _CCCL_PRAGMA_UNROLL_FULL() + for (unsigned i = tid; i < static_cast(chunk_num_items); i += (n_threads * items_per_thread)) { - items[j] = transform_op(d_in[idx]); + ftype items[items_per_thread] = {}; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; j++) + { + const unsigned idx = i + j * n_threads; + if (idx < static_cast(chunk_num_items)) + { + items[j] = transform_op(d_in[idx]); + } + } + + ftype abs_max_val = ::cuda::std::fabs(items[0]); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 1; j < items_per_thread; j++) + { + abs_max_val = ::cuda::std::fmax(::cuda::std::fabs(items[j]), abs_max_val); + } + + thread_aggregate.set_max_val(abs_max_val); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; j++) + { + thread_aggregate.unsafe_add(items[j]); + count++; + if (count >= thread_aggregate.endurance()) + { + thread_aggregate.renorm(); + count = 0; + } + } } - } - - ftype abs_max_val = ::cuda::std::fabs(items[0]); + d_in += chunk_num_items; + // chunk_num_items == min(remaining, num_items_per_chunk) <= remaining, so the countdown lands exactly on zero + // and cannot wrap when num_items_t is unsigned. + remaining -= chunk_num_items; + } + } + else + { _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 1; j < items_per_thread; j++) + for (unsigned i = tid; i < static_cast(num_items); i += (n_threads * items_per_thread)) { - abs_max_val = ::cuda::std::fmax(::cuda::std::fabs(items[j]), abs_max_val); - } + ftype items[items_per_thread] = {}; + for (int j = 0; j < items_per_thread; j++) + { + const unsigned idx = i + j * n_threads; + if (idx < static_cast(num_items)) + { + items[j] = transform_op(d_in[idx]); + } + } - thread_aggregate.set_max_val(abs_max_val); + ftype abs_max_val = ::cuda::std::fabs(items[0]); - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; j++) - { - thread_aggregate.unsafe_add(items[j]); - count++; - if (count >= thread_aggregate.endurance()) + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 1; j < items_per_thread; j++) { - thread_aggregate.renorm(); - count = 0; + abs_max_val = ::cuda::std::fmax(::cuda::std::fabs(items[j]), abs_max_val); + } + + thread_aggregate.set_max_val(abs_max_val); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; j++) + { + thread_aggregate.unsafe_add(items[j]); + count++; + if (count >= thread_aggregate.endurance()) + { + thread_aggregate.renorm(); + count = 0; + } } } } @@ -205,6 +344,10 @@ _CCCL_KERNEL_ATTRIBUTES __launch_bounds__( // Shared memory storage __shared__ typename block_reduce_t::TempStorage temp_storage; + // TODO(NaderAlAwar): This code is intentionally duplicated in DeterministicDeviceReduceDeferredSingleTileKernel + // because extracting it into a device function changes the SASS of this kernel. Changes here must also be applied to + // the copy below. + // Check if empty problem if (num_items == 0) { @@ -245,6 +388,85 @@ _CCCL_KERNEL_ATTRIBUTES __launch_bounds__( detail::reduce::finalize_and_store_aggregate(d_out, reduction_op, init, block_aggregate.conv_to_fp()); } } + +//! Single-tile entry point for a reduction with a deferred problem size. Derives on device how many blocks of the +//! worst-case first-pass grid produced a partial (the surplus blocks exited without writing one) and consumes exactly +//! those partials. An empty problem stores `init` unmodified. +template +_CCCL_KERNEL_ATTRIBUTES __launch_bounds__( + int(current_policy().single_tile.threads_per_block), + 1) void DeterministicDeviceReduceDeferredSingleTileKernel(_CCCL_GRID_CONSTANT const InputIteratorT d_in, + _CCCL_GRID_CONSTANT const OutputIteratorT d_out, + _CCCL_GRID_CONSTANT const KernelNumItemsT kernel_num_items, + _CCCL_GRID_CONSTANT const int first_pass_grid_size, + ReductionOpT reduction_op, + _CCCL_GRID_CONSTANT const InitValueT init, + TransformOpT transform_op) +{ + using actual_num_items_t = deterministic_num_items_t; + const actual_num_items_t actual_num_items = detail::parameter_from_device(kernel_num_items); + const int num_items = + detail::reduce::deferred_reduce_grid_size(actual_num_items, first_pass_grid_size); + + constexpr ReducePassPolicy policy = current_policy().single_tile; + constexpr int threads_per_block = policy.threads_per_block; + + using block_reduce_t = BlockReduce; + + // Shared memory storage + __shared__ typename block_reduce_t::TempStorage temp_storage; + + // TODO(NaderAlAwar): This code is intentionally duplicated in DeterministicDeviceReduceSingleTileKernel because + // extracting it into a device function changes the SASS of that kernel. Changes here must also be applied to the + // copy above. + + // Check if empty problem + if (actual_num_items == 0) + { + if (threadIdx.x == 0) + { + *d_out = init; + } + return; + } + + using float_type = typename AccumT::ftype; + constexpr int bin_length = AccumT::max_index + AccumT::max_fold; + + float_type* shared_bins = detail::rfa::get_shared_bin_array(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int index = static_cast(threadIdx.x); index < bin_length; index += threads_per_block) + { + shared_bins[index] = AccumT::initialize_bin(index); + } + + __syncthreads(); + + AccumT thread_aggregate{}; + + // Consume block aggregates of previous kernel + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = static_cast(threadIdx.x); i < num_items; i += threads_per_block) + { + thread_aggregate += transform_op(d_in[i]); + } + + AccumT block_aggregate = block_reduce_t(temp_storage).Reduce(thread_aggregate, reduction_op, num_items); + + // Output result + if (threadIdx.x == 0) + { + detail::reduce::finalize_and_store_aggregate(d_out, reduction_op, init, block_aggregate.conv_to_fp()); + } +} } // namespace detail::reduce CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_device_reduce_deferred.cu b/cub/test/catch2_test_device_reduce_deferred.cu index 192e91e4446..90c64fd125d 100644 --- a/cub/test/catch2_test_device_reduce_deferred.cu +++ b/cub/test/catch2_test_device_reduce_deferred.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -112,6 +113,41 @@ struct fixed_grid_reduce_t } }; +struct fixed_deterministic_reduce_policy_selector_t +{ + _CCCL_HOST_DEVICE_API constexpr auto operator()(cuda::compute_capability) const -> cub::ReducePolicy + { + const auto policy = cub::ReducePassPolicy{32, 1, 1, cub::BLOCK_REDUCE_RAKING, cub::LOAD_DEFAULT}; + return {policy, policy}; + } +}; + +struct fixed_grid_deterministic_reduce_t +{ + template + CUB_RUNTIME_FUNCTION cudaError_t operator()( + uint8_t* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + InitT init, + cudaStream_t stream = nullptr) const + { + return cub::detail::rfa::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_items, + init, + stream, + cuda::std::identity{}, + fixed_deterministic_reduce_policy_selector_t{}, + fixed_grid_factory_t{}); + } +}; + struct select_then_reduce_t { template (num_items)); } +C2H_TEST("DeviceReduce deterministic dispatch with a deferred size handles grid boundaries", + "[device][reduce][deferred][gpu_to_gpu]") +{ + using value_t = float; + using count_t = int32_t; + + // The deterministic dispatch derives its occupancy without consulting the launcher factory, so the worst-case grid + // size is not fixed here; large_num_items exceeds the capacity of any single-pass grid of 32-thread tiles. + constexpr count_t tile_size = 32; + constexpr count_t large_num_items = 100'000; + const count_t num_items = + GENERATE_COPY(count_t{0}, count_t{1}, tile_size - 1, tile_size, tile_size + 1, large_num_items); + constexpr value_t init = 7.0f; + + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto d_num_items = thrust::raw_pointer_cast(device_num_items.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + + launch(fixed_grid_deterministic_reduce_t{}, + cuda::constant_iterator{value_t{1}}, + d_output, + cuda::args::deferred{d_num_items}, + init); + + REQUIRE(output[0] == init + static_cast(num_items)); +} + C2H_TEST("DeviceReduce::Reduce accepts deferred span and fancy-iterator sources", "[device][reduce][deferred]") { using value_t = int64_t; @@ -559,6 +624,68 @@ C2H_TEST("DeviceReduce::Reduce supports deferred num_items with not_guaranteed d REQUIRE(cudaSuccess == reduce_error); REQUIRE(output[0] == init + static_cast(num_items)); } + +using gpu_to_gpu_count_types = c2h::type_list; + +C2H_TEST("DeviceReduce::Reduce with deferred num_items matches the immediate result bitwise with gpu_to_gpu " + "determinism", + "[device][reduce][deferred][gpu_to_gpu]", + gpu_to_gpu_count_types) +{ + using value_t = float; + using count_t = typename c2h::get<0, TestType>; + + constexpr count_t capacity = 100'000; + const count_t num_items = GENERATE_COPY(count_t{0}, count_t{1}, count_t{1'000}, capacity); + constexpr value_t init = 3.0f; + + c2h::device_vector input(capacity, thrust::no_init); + c2h::gen(C2H_SEED(1), input, value_t{-100}, value_t{100}); + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector reference(1, thrust::no_init); + c2h::device_vector output(1, thrust::no_init); + + const auto env = cuda::execution::require(cuda::execution::determinism::gpu_to_gpu); + + REQUIRE(cudaSuccess + == cub::DeviceReduce::Reduce(input.begin(), reference.begin(), num_items, cuda::std::plus<>{}, init, env)); + REQUIRE( + cudaSuccess + == cub::DeviceReduce::Reduce( + input.begin(), output.begin(), cuda::args::deferred{device_num_items.begin()}, cuda::std::plus<>{}, init, env)); + REQUIRE_THAT(detail::to_vec(reference), detail::BitwiseEqualsRange(detail::to_vec(output))); +} + +using gpu_to_gpu_large_count_types = c2h::type_list; + +C2H_TEST("DeviceReduce::Reduce with a large deferred num_items matches the immediate result bitwise with gpu_to_gpu " + "determinism", + "[device][reduce][deferred][gpu_to_gpu]", + gpu_to_gpu_large_count_types) +{ + using value_t = float; + using count_t = typename c2h::get<0, TestType>; + + // Exceeds INT32_MAX, so the immediate reference reduces two host-side chunks while the deferred reduction consumes + // the whole problem in a single launch with 64-bit indexing; RFA results are partition independent, so the results + // must match bitwise. + const count_t num_items = (count_t{1} << 31) + GENERATE_COPY(count_t{0}, count_t{12'345}); + constexpr value_t init = 3.0f; + + const auto input = cuda::constant_iterator{value_t{1}}; + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector reference(1, thrust::no_init); + c2h::device_vector output(1, thrust::no_init); + + const auto env = cuda::execution::require(cuda::execution::determinism::gpu_to_gpu); + + REQUIRE( + cudaSuccess == cub::DeviceReduce::Reduce(input, reference.begin(), num_items, cuda::std::plus<>{}, init, env)); + REQUIRE(cudaSuccess + == cub::DeviceReduce::Reduce( + input, output.begin(), cuda::args::deferred{device_num_items.begin()}, cuda::std::plus<>{}, init, env)); + REQUIRE_THAT(detail::to_vec(reference), detail::BitwiseEqualsRange(detail::to_vec(output))); +} #endif // TEST_LAUNCH == 0 #if TEST_LAUNCH == 2 @@ -615,6 +742,55 @@ C2H_TEST("captured DeviceReduce atomic dispatch replays with zero and nonzero de REQUIRE(cudaSuccess == cudaStreamDestroy(stream)); } +C2H_TEST("captured DeviceReduce deterministic dispatch replays with zero and nonzero deferred counts", + "[device][reduce][deferred][gpu_to_gpu]") +{ + using value_t = float; + using count_t = int32_t; + + constexpr count_t capacity = 100'000; + constexpr value_t init = 3.0f; + constexpr value_t poison = -1234.0f; + + c2h::device_vector device_num_items(1, count_t{-1}); + c2h::device_vector output(1, poison); + + const auto d_num_items = thrust::raw_pointer_cast(device_num_items.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto count = cuda::args::deferred{d_num_items}; + const auto input = cuda::constant_iterator{value_t{1}}; + + cudaStream_t stream{}; + REQUIRE(cudaSuccess == cudaStreamCreate(&stream)); + + const fixed_grid_deterministic_reduce_t reduce; + size_t temp_storage_bytes{}; + REQUIRE(cudaSuccess == reduce(nullptr, temp_storage_bytes, input, d_output, count, init, stream)); + c2h::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + const auto d_temp_storage = thrust::raw_pointer_cast(temp_storage.data()); + + cudaGraph_t graph{}; + REQUIRE(cudaSuccess == cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal)); + REQUIRE(cudaSuccess == reduce(d_temp_storage, temp_storage_bytes, input, d_output, count, init, stream)); + REQUIRE(cudaSuccess == cudaStreamEndCapture(stream, &graph)); + + cudaGraphExec_t executable{}; + REQUIRE(cudaSuccess == cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); + + for (const count_t num_items : {capacity, count_t{0}, capacity}) + { + REQUIRE(cudaSuccess == cudaMemcpyAsync(d_num_items, &num_items, sizeof(num_items), cudaMemcpyHostToDevice, stream)); + REQUIRE(cudaSuccess == cudaMemcpyAsync(d_output, &poison, sizeof(poison), cudaMemcpyHostToDevice, stream)); + REQUIRE(cudaSuccess == cudaGraphLaunch(executable, stream)); + REQUIRE(cudaSuccess == cudaStreamSynchronize(stream)); + REQUIRE(output[0] == init + static_cast(num_items)); + } + + REQUIRE(cudaSuccess == cudaGraphExecDestroy(executable)); + REQUIRE(cudaSuccess == cudaGraphDestroy(graph)); + REQUIRE(cudaSuccess == cudaStreamDestroy(stream)); +} + C2H_TEST("captured DeviceSelect::If to DeviceReduce::Reduce pipeline replays with changing counts", "[device][reduce][deferred]") { diff --git a/cub/test/catch2_test_device_reduce_env.cu b/cub/test/catch2_test_device_reduce_env.cu index f520211afd6..9e7310750c7 100644 --- a/cub/test/catch2_test_device_reduce_env.cu +++ b/cub/test/catch2_test_device_reduce_env.cu @@ -446,7 +446,7 @@ C2H_TEST("Device reduce uses environment", "[reduce][device]", requirements) deterministic_accum_t, transform_t>; auto k2 = cub::detail::reduce:: - DeterministicDeviceReduceKernel; + DeterministicDeviceReduceKernel; auto k3 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< policy_t, deterministic_accum_t*, @@ -586,7 +586,7 @@ C2H_TEST("Device sum uses environment", "[reduce][device]", requirements) deterministic_accum_t, transform_t>; auto k2 = cub::detail::reduce:: - DeterministicDeviceReduceKernel; + DeterministicDeviceReduceKernel; auto k3 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< policy_t, deterministic_accum_t*,