diff --git a/c/parallel/src/reduce.cu b/c/parallel/src/reduce.cu index 5beabcfc319..ba870128cd3 100644 --- a/c/parallel/src/reduce.cu +++ b/c/parallel/src/reduce.cu @@ -126,12 +126,13 @@ std::string get_device_reduce_kernel_name( check(cccl_type_name_from_nvrtc(&transform_op_t)); return std::format( - "cub::detail::reduce::DeviceReduceKernel<{0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}>", + "cub::detail::reduce::DeviceReduceKernel<{0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}>", policy_selector_t, stable_reduction_order ? "true" : "false", input_iterator_t, stable_reduction_order ? std::string(accum_t) + "*" : output_iterator_t, offset_t, + offset_t, reduction_op_t, accum_t, init_t, diff --git a/cub/benchmarks/bench/reduce/deferred_nondeterministic.cu b/cub/benchmarks/bench/reduce/deferred_nondeterministic.cu new file mode 100644 index 00000000000..d0e7e504256 --- /dev/null +++ b/cub/benchmarks/bench/reduce/deferred_nondeterministic.cu @@ -0,0 +1,98 @@ +// 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 +// %RANGE% TUNE_ITEMS_PER_VEC_LOAD_POW2 ipv 1:2:1 + +#if !TUNE_BASE +template +struct policy_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::ReducePolicy + { + const auto [items, threads] = + cub::detail::scale_mem_bound(TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, int{sizeof(AccumT)}); + const auto policy = cub::ReducePassPolicy{ + threads, + items, + 1 << TUNE_ITEMS_PER_VEC_LOAD_POW2, + cub::BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC, + cub::LOAD_DEFAULT}; + return {policy, {}}; + } +}; +#endif // !TUNE_BASE + +template +void nondeterministic_sum(nvbench::state& state, nvbench::type_list) +{ + using op_t = cuda::std::plus<>; + using init_value_t = T; + + // Retrieve axis parameters + 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, static_cast(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::not_guaranteed) +#if !TUNE_BASE + , + cuda::execution::tune(policy_selector>{}) +#endif // !TUNE_BASE + ); + _CCCL_TRY_CUDA_API( + cub::DeviceReduce::Reduce, + "Reduce failed", + d_in, + d_out, + cuda::args::deferred{d_num_items}, + op_t{}, + init_value_t{}, + env); + }); +} + +#ifdef TUNE_T +using value_types = nvbench::type_list; +#else +using value_types = nvbench::type_list; +#endif + +NVBENCH_BENCH_TYPES(nondeterministic_sum, NVBENCH_TYPE_AXES(value_types, offset_types)) + .set_name("base") + .set_type_axes_names({"T{ct}", "OffsetT{ct}"}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4)); diff --git a/cub/benchmarks/bench/reduce/deferred_sum.cu b/cub/benchmarks/bench/reduce/deferred_sum.cu new file mode 100644 index 00000000000..974fb43700f --- /dev/null +++ b/cub/benchmarks/bench/reduce/deferred_sum.cu @@ -0,0 +1,86 @@ +// 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 + +// %RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1 +// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32 +// %RANGE% TUNE_ITEMS_PER_VEC_LOAD_POW2 ipv 1:2:1 + +#if !TUNE_BASE +template +struct policy_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::ReducePolicy + { + const auto [items, threads] = + cub::detail::scale_mem_bound(TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, int{sizeof(AccumT)}); + const auto policy = cub::ReducePassPolicy{ + threads, items, 1 << TUNE_ITEMS_PER_VEC_LOAD_POW2, cub::BLOCK_REDUCE_WARP_REDUCTIONS, cub::LOAD_DEFAULT}; + return {policy, policy}; + } +}; +#endif // !TUNE_BASE + +using op_t = cuda::std::plus<>; + +template +void reduce(nvbench::state& state, nvbench::type_list) +{ + using init_value_t = T; + + // Retrieve axis parameters + const auto elements = state.get_int64("Elements{io}"); + + thrust::device_vector in = generate(elements); + thrust::device_vector out(1, thrust::default_init); + thrust::device_vector device_num_items(1, static_cast(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 +#if !TUNE_BASE + , + cuda::execution::tune(policy_selector>{}) +#endif // !TUNE_BASE + ); + _CCCL_TRY_CUDA_API( + cub::DeviceReduce::Reduce, + "Reduce failed", + d_in, + d_out, + cuda::args::deferred{d_num_items}, + op_t{}, + init_value_t{}, + env); + }); +} + +using value_types = all_types; + +NVBENCH_BENCH_TYPES(reduce, NVBENCH_TYPE_AXES(value_types, offset_types)) + .set_name("base") + .set_type_axes_names({"T{ct}", "OffsetT{ct}"}) + .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4)); diff --git a/cub/cub/detail/deferred_parameter.cuh b/cub/cub/detail/deferred_parameter.cuh new file mode 100644 index 00000000000..e808cbd2407 --- /dev/null +++ b/cub/cub/detail/deferred_parameter.cuh @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +#if !_CCCL_COMPILER(NVRTC) +// Preserve deferred problem sizes for dispatch and canonicalize immediate values to CUB's offset type. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto make_num_items_dispatch_arg(NumItemsT num_items) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + + if constexpr (args_traits_t::is_deferred) + { + return num_items; + } + else + { + using offset_t = choose_offset_t; + return static_cast(::cuda::args::__unwrap(num_items)); + } +} + +// Forms a kernel parameter from a single-value argument without reading a deferred source. +// Immediate values are converted to TargetT. Deferred arguments are unwrapped to their source, erasing bounds from +// the kernel type and payload. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto parameter_from_host(ParameterT parameter) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + static_assert(args_traits_t::is_single_value, "parameter must contain a single value"); + + if constexpr (args_traits_t::is_deferred) + { + return ::cuda::args::__unwrap(parameter); + } + else + { + return static_cast(::cuda::args::__unwrap(parameter)); + } +} + +template +using parameter_from_host_t = decltype(parameter_from_host(::cuda::std::declval())); +#endif // !_CCCL_COMPILER(NVRTC) + +// Forms a value from a kernel parameter, reading element zero when the parameter is a deferred source. +template +[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE TargetT parameter_from_device(ParameterT parameter) noexcept +{ + if constexpr (::cuda::std::is_same_v) + { + return parameter; + } + else + { + return static_cast(parameter[0]); + } +} +} // namespace detail + +CUB_NAMESPACE_END diff --git a/cub/cub/device/device_reduce.cuh b/cub/cub/device/device_reduce.cuh index 863c07520cd..7d1c80734cb 100644 --- a/cub/cub/device/device_reduce.cuh +++ b/cub/cub/device/device_reduce.cuh @@ -19,6 +19,7 @@ #endif // no system header #include +#include #include #include #include @@ -39,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -121,6 +123,27 @@ inline constexpr bool is_non_deterministic_v = //! :start-after: example-begin reduce-by-key-tuning //! :end-before: example-end reduce-by-key-tuning //! +//! Deferred problem sizes +//! ==================================== +//! +//! ``Reduce``, ``Sum``, ``Min``, ``Max``, and ``TransformReduce`` allow specifying the problem size from a value that +//! resides in device memory through a single-value ``cuda::args::deferred`` argument. The deferred source may be a +//! device pointer, a one-element span, or a random-access fancy iterator whose element is a non-``bool`` 32- or 64-bit +//! integer. +//! +//! The problem size is read in stream order by the reduction kernels. Work that produces the count in the same stream +//! is ordered automatically; a producer in another stream requires an event or an equivalent dependency. The source +//! and its value must remain accessible and unchanged until all reduction kernels complete. The value must be +//! nonnegative and ``[d_in, d_in + num_items)`` must be accessible. A temporary-storage query does not dereference the +//! deferred source. +//! +//! Deferred reductions are CUDA Graph capturable. The pointed-to count may change between graph replays without +//! 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. +//! //! Determinism //! ==================================== //! @@ -175,21 +198,39 @@ private: ::cuda::execution::determinism::__determinism_holder_t, EnvT env) { - using offset_t = detail::choose_offset_t; - using accum_t = decltype(detail::reduce::template select_accum_t( + using args_traits_t = ::cuda::args::__traits; + using offset_t = detail::choose_offset_t; + using accum_t = decltype(detail::reduce::select_accum_t( static_cast(nullptr))); if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__gpu_to_gpu) { - // 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, static_cast(num_items), init, stream, transform_op, policy_selector); - }); + 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); + }); + } } else if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__not_guaranteed) { @@ -202,7 +243,7 @@ private: bytes, d_in, THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(d_out), - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), reduction_op, init, stream, @@ -221,7 +262,7 @@ private: bytes, d_in, d_out, - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), reduction_op, init, stream, @@ -232,16 +273,15 @@ private: } //! @brief Internal implementation shared by Reduce and TransformReduce env overloads - template < - typename InputIteratorT, - typename OutputIteratorT, - typename ReductionOpT, - typename TransformOpT, - typename T, - typename NumItemsT, - typename EnvT, - typename AccumT = decltype(detail::reduce::template select_accum_t( - static_cast(nullptr)))> + template ( + static_cast(nullptr)))> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __transform_reduce( InputIteratorT d_in, OutputIteratorT d_out, @@ -470,11 +510,15 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Reduce"); - // Signed integer type for global offsets - using OffsetT = detail::choose_offset_t; - return detail::reduce::dispatch( - d_temp_storage, temp_storage_bytes, d_in, d_out, static_cast(num_items), reduction_op, init, stream); + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + reduction_op, + init, + stream); } //! @rst @@ -728,9 +772,6 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Sum"); - // Signed integer type for global offsets - using OffsetT = detail::choose_offset_t; - // The output value type using OutputT = cub::detail::non_void_value_t>; @@ -741,7 +782,7 @@ public: temp_storage_bytes, d_in, d_out, - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), ::cuda::std::plus<>{}, init_value_t{}, // zero-initialize stream); @@ -836,7 +877,6 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Min"); - using OffsetT = detail::choose_offset_t; // Signed integer type for global offsets using InputT = detail::it_value_t; using init_value_t = InputT; using limits_t = ::cuda::std::numeric_limits; @@ -853,7 +893,7 @@ public: temp_storage_bytes, d_in, d_out, - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), ::cuda::minimum<>{}, limits_t::max(), stream); @@ -1484,8 +1524,6 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Max"); - // Signed integer type for global offsets - using OffsetT = detail::choose_offset_t; using InputT = detail::it_value_t; using init_value_t = InputT; using limits_t = ::cuda::std::numeric_limits; @@ -1502,7 +1540,7 @@ public: temp_storage_bytes, d_in, d_out, - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), ::cuda::maximum<>{}, limits_t::lowest(), stream); @@ -2093,14 +2131,12 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::TransformReduce"); - using OffsetT = detail::choose_offset_t; - return detail::reduce::dispatch( d_temp_storage, temp_storage_bytes, d_in, d_out, - static_cast(num_items), + detail::make_num_items_dispatch_arg(num_items), reduction_op, init, stream, diff --git a/cub/cub/device/dispatch/dispatch_reduce.cuh b/cub/cub/device/dispatch/dispatch_reduce.cuh index f9a5352e638..09f33b161cc 100644 --- a/cub/cub/device/dispatch/dispatch_reduce.cuh +++ b/cub/cub/device/dispatch/dispatch_reduce.cuh @@ -1,5 +1,5 @@ // SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. -// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: BSD-3 /** @@ -22,6 +22,8 @@ #endif // no system header #include +#include +#include #include #include #include @@ -32,9 +34,14 @@ #include #include // for cub::detail::non_void_value_t, cub::detail::it_value_t +#include #include #include #include +#include +#include +#include +#include #include // TODO(bgruber): included to not break users when moving DeviceSegmentedReduce to its own file. Remove in CCCL 4.0. @@ -48,6 +55,7 @@ template ) + CUB_DEFINE_KERNEL_GETTER( + DeferredSingleTileSecondKernel, + DeviceReduceDeferredSingleTileKernel< + PolicySelector, + AccumT*, + OutputIteratorT, + OffsetT, + KernelNumItemsT, + ReductionOpT, + InitValueT, + AccumT>) + CUB_RUNTIME_FUNCTION static constexpr size_t AccumSize() { return sizeof(AccumT); @@ -173,6 +194,7 @@ template < InputIteratorT, OutputIteratorT, OffsetT, + OffsetT, ReductionOpT, InitValueT, AccumT, @@ -572,6 +594,7 @@ template < InputIteratorT, OutputIteratorT, OffsetT, + OffsetT, ReductionOpT, InitValueT, AccumT, @@ -602,6 +625,32 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE void* get_device_ptr(void* ptr) return *reinterpret_cast(ptr); } +//! Preserve caller-selected immediate offset types; select a concrete offset type for deferred arguments. +template +using num_items_offset_t = + ::cuda::std::conditional_t<::cuda::args::__traits::is_deferred, + detail::choose_offset_t::element_type>, + typename ::cuda::args::__traits::element_type>; + +//! Creates the kernel argument for an immediate or deferred problem size without reading a deferred source. +//! Immediate values are cast to the selected offset type; deferred arguments are stripped to their source. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto make_num_items_kernel_arg(OffsetT num_items) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + using element_t = typename args_traits_t::element_type; + + if constexpr (args_traits_t::is_deferred) + { + static_assert(args_traits_t::is_single_value, "num_items must be a single value wrapped in cuda::args::deferred"); + static_assert(::cuda::std::__cccl_is_integer_v, "the num_items element type must be an integer"); + static_assert( + sizeof(element_t) == sizeof(::cuda::std::int32_t) || sizeof(element_t) == sizeof(::cuda::std::int64_t)); + } + + return CUB_NS_QUALIFIER::detail::parameter_from_host>(num_items); +} + template ; + + const auto kernel_num_items = make_num_items_kernel_arg(num_items); + // Get SM count int sm_count = 0; if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) @@ -634,8 +687,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( } // Init regular kernel configuration - const auto tile_size = active_policy.multi_tile.threads_per_block * active_policy.multi_tile.items_per_thread; - int sm_occupancy = 0; + int sm_occupancy = 0; if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( sm_occupancy, kernel_source.ReductionKernel(), active_policy.multi_tile.threads_per_block))) { @@ -643,11 +695,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( } const int reduce_device_occupancy = sm_occupancy * sm_count; - - // Even-share work distribution - const int max_blocks = reduce_device_occupancy * detail::subscription_factor; - GridEvenShare even_share; - even_share.DispatchInit(num_items, max_blocks, tile_size); + const int max_blocks = reduce_device_occupancy * detail::subscription_factor; [[maybe_unused]] AccumT* d_block_reductions = nullptr; // buffer for per-block aggregates for the two-phase code path if constexpr (!StableReductionOrder) @@ -684,14 +732,30 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( d_block_reductions = static_cast(allocations[0]); } - // The grid size for DeviceReduceKernel can be zero if the input size is zero. Since the atomic code path does not run - // a second kernel, we need to handle the empty grid in first kernel already - int reduce_grid_size = even_share.grid_size; - if constexpr (!StableReductionOrder) + GridEvenShare even_share; + if constexpr (!::cuda::args::__traits::is_deferred) { - reduce_grid_size = ::cuda::std::max(1, reduce_grid_size); + const auto tile_size = active_policy.multi_tile.threads_per_block * active_policy.multi_tile.items_per_thread; + even_share.DispatchInit(kernel_num_items, max_blocks, tile_size); } + const int reduce_grid_size = [&] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return max_blocks; + } + else if constexpr (!StableReductionOrder) + { + // The grid size for DeviceReduceKernel can be zero if the input size is zero. + // The atomic code path does not run a second kernel, so block zero handles an empty input. + return ::cuda::std::max(1, even_share.grid_size); + } + else + { + return even_share.grid_size; + } + }(); + // Log device_reduce_sweep_kernel configuration #ifdef CUB_DEBUG_LOG _CubLog("Invoking DeviceReduceKernel<<<%lu, %d, 0, %lld>>>(), %d items " @@ -719,7 +783,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( .doit(kernel_source.ReductionKernel(), d_in, reduce_kernel_output, - num_items, + kernel_num_items, even_share, reduction_op, init, @@ -751,18 +815,37 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( active_policy.single_tile.items_per_thread); #endif // CUB_DEBUG_LOG - // Invoke DeviceReduceSingleTileKernel - if (const auto error = CubDebug( - launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) - .doit(kernel_source.SingleTileSecondKernel(), - d_block_reductions, - d_out, - reduce_grid_size, - reduction_op, - init, - ::cuda::std::identity{}))) + // Invoke DeviceReduceSingleTileKernel/DeviceReduceDeferredSingleTileKernel + if constexpr (::cuda::args::__traits::is_deferred) { - return error; + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(kernel_source.DeferredSingleTileSecondKernel(), + d_block_reductions, + d_out, + kernel_num_items, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}))) + { + return error; + } + } + else + { + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(kernel_source.SingleTileSecondKernel(), + d_block_reductions, + d_out, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}))) + { + return error; + } } // Check for failure to launch @@ -814,14 +897,15 @@ template < static_cast(nullptr))), typename PolicySelector = policy_selector_from_types< AccumT, - OffsetT, + num_items_offset_t, ReductionOpT, StableReductionOrder ? __determinism_t::__run_to_run : __determinism_t::__not_guaranteed>, typename KernelSource = DeviceReduceKernelSource< PolicySelector, InputIteratorT, OutputIteratorT, - OffsetT, + num_items_offset_t, + CUB_NS_QUALIFIER::detail::parameter_from_host_t, OffsetT>, ReductionOpT, InitValueT, AccumT, @@ -845,12 +929,25 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}) { + using offset_t = num_items_offset_t; + ::cuda::compute_capability cc{}; if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) { return error; } + // TODO: Remove this workaround once nvcc versions older than 12.4 are no longer supported. + // Older nvcc versions eagerly instantiate discarded statements in generic lambdas, so perform this conversion here. + // Both suppressions are needed for "never referenced" and "set but never used" diagnostics across supported nvcc + // and MSVC combinations. + [[maybe_unused]] offset_t offset_num_items{}; + if constexpr (StableReductionOrder && !::cuda::args::__traits::is_deferred) + { + offset_num_items = static_cast(num_items); + } + (void) offset_num_items; + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { CUB_DETAIL_CONSTEXPR_ISH const ReducePolicy active_policy = policy_getter(); @@ -877,10 +974,11 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( })) #endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) - if constexpr (StableReductionOrder) + if constexpr (StableReductionOrder && !::cuda::args::__traits::is_deferred) { - const bool single_tile_problem = num_items <= (static_cast(active_policy.single_tile.threads_per_block) - * active_policy.single_tile.items_per_thread); + const bool single_tile_problem = + offset_num_items <= (static_cast(active_policy.single_tile.threads_per_block) + * active_policy.single_tile.items_per_thread); // if the problem is small enough to fit into a single tile, just handle it and return early if (single_tile_problem) @@ -903,7 +1001,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( // Invoke single_reduce_sweep_kernel if (const auto error = CubDebug( launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) - .doit(kernel_source.SingleTileKernel(), d_in, d_out, num_items, reduction_op, init, transform_op))) + .doit( + kernel_source.SingleTileKernel(), d_in, d_out, offset_num_items, reduction_op, init, transform_op))) { return error; } @@ -923,7 +1022,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( return cudaSuccess; } } - else + else if constexpr (!StableReductionOrder) { // Return if the caller is simply requesting the size of the storage allocation if (d_temp_storage == nullptr) diff --git a/cub/cub/device/dispatch/kernels/kernel_reduce.cuh b/cub/cub/device/dispatch/kernels/kernel_reduce.cuh index 3b7155532ec..f8c988d0d1e 100644 --- a/cub/cub/device/dispatch/kernels/kernel_reduce.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_reduce.cuh @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: BSD-3 #pragma once @@ -14,12 +14,14 @@ #endif // no system header #include +#include #include #include #include #include #include +#include CUB_NAMESPACE_BEGIN @@ -112,6 +114,9 @@ finalize_and_store_aggregate(OutputIteratorT d_out, ReductionOpT, no_init_t, Acc * @tparam OffsetT * Signed integer type for global offsets * + * @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)` @@ -128,8 +133,8 @@ finalize_and_store_aggregate(OutputIteratorT d_out, ReductionOpT, no_init_t, Acc * @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] even_share * Even-share descriptor for mapping an equal number of tiles onto each @@ -143,6 +148,7 @@ template ().multi_tile.threads_per_block)) void DeviceReduceKernel( _CCCL_GRID_CONSTANT const InputIteratorT d_in, _CCCL_GRID_CONSTANT const OutputIteratorT d_out, - _CCCL_GRID_CONSTANT const OffsetT num_items, + _CCCL_GRID_CONSTANT const KernelNumItemsT kernel_num_items, GridEvenShare even_share, ReductionOpT reduction_op, [[maybe_unused]] _CCCL_GRID_CONSTANT const InitValueT init, TransformOpT transform_op) { static constexpr ReducePassPolicy policy = current_policy().multi_tile; + const OffsetT num_items = CUB_NS_QUALIFIER::detail::parameter_from_device(kernel_num_items); + + // Early return from surplus blocks for deferred num_items + if constexpr (!::cuda::std::is_integral_v) + { + constexpr int tile_size = policy.threads_per_block * policy.items_per_thread; + even_share.DispatchInit(num_items, static_cast(gridDim.x), tile_size); + + if constexpr (StableReductionOrder) + { + if (static_cast(blockIdx.x) >= even_share.grid_size) + { + return; + } + } + else + { + // Only block zero handles an empty atomic reduction. For non-empty problems, all surplus blocks return. + if ((num_items == 0 && blockIdx.x != 0) + || (num_items != 0 && static_cast(blockIdx.x) >= even_share.grid_size)) + { + return; + } + } + } + if constexpr (!StableReductionOrder) { static_assert(detail::is_cuda_std_plus_v, @@ -288,8 +320,8 @@ _CCCL_KERNEL_ATTRIBUTES __launch_bounds__( static constexpr ReducePassPolicy policy = current_policy().single_tile; // TODO(bgruber): pass policy directly as template argument to AgentReduce in C++20 using agent_policy_t = detail::agent_reduce_policy< - 0, - 0, + /* NominalThreadsPerBlock4B = */ 0, + /* NominalItemsPerThread4B = */ 0, AccumT, policy.vec_size, policy.reduce_algorithm, @@ -306,6 +338,10 @@ _CCCL_KERNEL_ATTRIBUTES __launch_bounds__( // Shared memory storage __shared__ typename AgentReduceT::TempStorage temp_storage; + // TODO(NaderAlAwar): This code is intentionally duplicated in DeviceReduceDeferredSingleTileKernel 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) { @@ -327,6 +363,82 @@ _CCCL_KERNEL_ATTRIBUTES __launch_bounds__( detail::reduce::finalize_and_store_aggregate(d_out, reduction_op, init, block_aggregate); } } + +//! Single-tile entry point that derives the number of first-pass partials from a deferred problem size. +template +#if _CCCL_HAS_CONCEPTS() + requires reduce_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +_CCCL_KERNEL_ATTRIBUTES __launch_bounds__( + int{current_policy().single_tile.threads_per_block}, + 1) void DeviceReduceDeferredSingleTileKernel(_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) +{ + const OffsetT actual_num_items = CUB_NS_QUALIFIER::detail::parameter_from_device(kernel_num_items); + static constexpr ReducePassPolicy first_pass_policy = current_policy().multi_tile; + constexpr int first_pass_tile_size = first_pass_policy.threads_per_block * first_pass_policy.items_per_thread; + GridEvenShare even_share; + even_share.DispatchInit(actual_num_items, first_pass_grid_size, first_pass_tile_size); + const int num_items = even_share.grid_size; + + static constexpr ReducePassPolicy policy = current_policy().single_tile; + + // TODO(bgruber): pass policy directly as template argument to AgentReduce in C++20 + using agent_policy_t = detail::agent_reduce_policy< + /* NominalThreadsPerBlock4B = */ 0, + /* NominalItemsPerThread4B = */ 0, + AccumT, + policy.vec_size, + policy.reduce_algorithm, + policy.load_modifier, + NoScaling>; + + // Thread block type for reducing input tiles + using AgentReduceT = AgentReduce; + + static_assert(sizeof(typename AgentReduceT::TempStorage) <= max_smem_per_block, + "cub::DeviceReduce ran out of CUDA shared memory, which we judged to be extremely unlikely. Please " + "file an issue at: https://github.com/NVIDIA/cccl/issues"); + + // Shared memory storage + __shared__ typename AgentReduceT::TempStorage temp_storage; + + // TODO(NaderAlAwar): This code is intentionally duplicated in DeviceReduceSingleTileKernel 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 (num_items == 0) + { + if (threadIdx.x == 0) + { + detail::reduce::handle_empty_problem(d_out, init); + } + + return; + } + + // Consume input tiles + AccumT block_aggregate = AgentReduceT(temp_storage, d_in, reduction_op, transform_op).ConsumeRange(int(0), num_items); + + // Output result + if (threadIdx.x == 0) + { + detail::reduce::finalize_and_store_aggregate(d_out, reduction_op, init, block_aggregate); + } +} } // 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 new file mode 100644 index 00000000000..192e91e4446 --- /dev/null +++ b/cub/test/catch2_test_device_reduce_deferred.cu @@ -0,0 +1,699 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "insert_nested_NVTX_range_guard.h" + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include "catch2_test_device_reduce.cuh" +#include "catch2_test_launch_helper.h" +#include +#include + +DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Reduce, device_reduce); +DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Sum, device_sum); +DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Min, device_min); +DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::Max, device_max); +DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::TransformReduce, device_transform_reduce); + +static_assert(cuda::std::is_same_v, int32_t>); +using deferred_count_t = decltype(cuda::args::deferred{static_cast(nullptr)}); +static_assert(cuda::std::is_same_v, uint32_t>); + +// %PARAM% TEST_LAUNCH lid 0:1:2 + +template +struct select_less_than_device_value_t +{ + const T* bound; + + _CCCL_DEVICE_API bool operator()(const T& value) const + { + return value < *bound; + } +}; + +struct fixed_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_WARP_REDUCTIONS, cub::LOAD_DEFAULT}; + return {policy, policy}; + } +}; + +struct fixed_grid_factory_t : cub::detail::TripleChevronFactory +{ + CUB_RUNTIME_FUNCTION cudaError_t MultiProcessorCount(int& sm_count) const + { + sm_count = 1; + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION cudaError_t + MaxSmOccupancy(int& sm_occupancy, KernelT, int, [[maybe_unused]] int dynamic_smem_bytes = 0) + { + sm_occupancy = 1; + return cudaSuccess; + } +}; + +template +struct fixed_grid_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, + ReductionOpT reduction_op, + InitT init, + cudaStream_t stream = nullptr) const + { + return cub::detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_items, + reduction_op, + init, + stream, + cuda::std::identity{}, + fixed_reduce_policy_selector_t{}, + {}, + fixed_grid_factory_t{}); + } +}; + +struct select_then_reduce_t +{ + template + CUB_RUNTIME_FUNCTION cudaError_t operator()( + uint8_t* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + SelectedOutputIteratorT d_selected_out, + NumSelectedIteratorT d_num_selected_out, + ReduceOutputIteratorT d_reduce_out, + NumItemsT num_items, + SelectOpT select_op, + ReductionOpT reduction_op, + InitT init, + cudaStream_t stream = nullptr) const + { + size_t select_temp_storage_bytes{}; + if (const cudaError_t error = cub::DeviceSelect::If( + nullptr, select_temp_storage_bytes, d_in, d_selected_out, d_num_selected_out, num_items, select_op, stream); + error != cudaSuccess) + { + return error; + } + + const auto deferred_num_items = cuda::args::deferred{d_num_selected_out}; + size_t reduce_temp_storage_bytes{}; + if (const cudaError_t error = cub::DeviceReduce::Reduce( + nullptr, + reduce_temp_storage_bytes, + d_selected_out, + d_reduce_out, + deferred_num_items, + reduction_op, + init, + stream); + error != cudaSuccess) + { + return error; + } + + temp_storage_bytes = cuda::std::max(select_temp_storage_bytes, reduce_temp_storage_bytes); + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + if (const cudaError_t error = cub::DeviceSelect::If( + d_temp_storage, + select_temp_storage_bytes, + d_in, + d_selected_out, + d_num_selected_out, + num_items, + select_op, + stream); + error != cudaSuccess) + { + return error; + } + + return cub::DeviceReduce::Reduce( + d_temp_storage, + reduce_temp_storage_bytes, + d_selected_out, + d_reduce_out, + deferred_num_items, + reduction_op, + init, + stream); + } +}; + +using count_types = c2h::type_list; + +C2H_TEST("DeviceReduce::Reduce accepts deferred num_items", "[device][reduce][deferred]", count_types) +{ + using value_t = int64_t; + using count_t = typename c2h::get<0, TestType>; + + constexpr count_t max_num_items = 100'000; + const count_t num_items = GENERATE_COPY(count_t{0}, count_t{1}, count_t{1'000}, max_num_items); + constexpr value_t init = 42; + + c2h::device_vector input(static_cast(num_items), thrust::no_init); + c2h::gen(C2H_SEED(2), input); + + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto d_input = thrust::raw_pointer_cast(input.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto d_num_items = thrust::raw_pointer_cast(device_num_items.data()); + + device_reduce(d_input, d_output, cuda::args::deferred{d_num_items}, cuda::std::plus<>{}, init); + + const value_t expected = compute_single_problem_reference(input, cuda::std::plus<>{}, init); + REQUIRE(output[0] == expected); +} + +C2H_TEST("DeviceReduce::Reduce with deferred num_items is run-to-run reproducible", + "[device][reduce][deferred][run_to_run]") +{ + using value_t = float; + using count_t = int32_t; + + constexpr count_t num_items = 100'000; + constexpr int num_runs = 3; + + c2h::device_vector input(num_items, 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 d_input = thrust::raw_pointer_cast(input.data()); + const auto d_reference = thrust::raw_pointer_cast(reference.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto count = cuda::args::deferred{device_num_items.begin()}; + + device_reduce(d_input, d_reference, count, cuda::std::plus<>{}, value_t{}); + for (int run = 1; run < num_runs; ++run) + { + device_reduce(d_input, d_output, count, cuda::std::plus<>{}, value_t{}); + REQUIRE_THAT(detail::to_vec(reference), detail::BitwiseEqualsRange(detail::to_vec(output))); + } +} + +C2H_TEST("DeviceReduce::Reduce with a deferred size handles grid boundaries", "[device][reduce][deferred]") +{ + using value_t = int64_t; + using count_t = int32_t; + + constexpr count_t tile_size = 32; + constexpr count_t max_blocks = cub::detail::subscription_factor; + const count_t num_items = GENERATE_COPY( + tile_size - 1, + tile_size, + tile_size + 1, + max_blocks * tile_size - 1, + max_blocks * tile_size, + max_blocks * tile_size + 1); + + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto count = cuda::args::deferred{device_num_items.begin()}; + launch(fixed_grid_reduce_t<>{}, + cuda::constant_iterator{value_t{1}}, + d_output, + count, + cuda::std::plus<>{}, + value_t{7}); + REQUIRE(output[0] == value_t{7} + num_items); +} + +C2H_TEST("DeviceReduce atomic dispatch with a deferred size handles grid boundaries", + "[device][reduce][deferred][not_guaranteed]") +{ + using value_t = float; + using count_t = int32_t; + + constexpr count_t tile_size = 32; + constexpr count_t max_blocks = cub::detail::subscription_factor; + const count_t num_items = GENERATE_COPY( + count_t{0}, + count_t{1}, + tile_size - 1, + tile_size, + tile_size + 1, + max_blocks * tile_size - 1, + max_blocks * tile_size, + max_blocks * tile_size + 1); + 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()); + + // Select the atomic implementation directly so this test cannot pass through the run-to-run fallback. + launch(fixed_grid_reduce_t{}, + cuda::constant_iterator{value_t{1}}, + d_output, + cuda::args::deferred{d_num_items}, + cuda::std::plus<>{}, + 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; + using count_t = int32_t; + + constexpr count_t capacity = 100'000; + constexpr count_t num_items = 1'000; + + const c2h::device_vector input(capacity, value_t{1}); + c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto d_input = thrust::raw_pointer_cast(input.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto d_num_items = thrust::raw_pointer_cast(device_num_items.data()); + + const auto const_count_span = cuda::std::span{d_num_items, 1}; + device_reduce(d_input, d_output, cuda::args::deferred{const_count_span}, cuda::std::plus<>{}, value_t{}); + REQUIRE(output[0] == num_items); + + output[0] = value_t{-2}; + const auto mutable_count_span = cuda::std::span{d_num_items, 1}; + device_reduce(d_input, d_output, cuda::args::deferred{mutable_count_span}, cuda::std::plus<>{}, value_t{}); + REQUIRE(output[0] == num_items); + + output[0] = value_t{-3}; + const auto count_transform = cuda::transform_iterator(d_num_items, cuda::std::identity{}); + device_reduce(d_input, d_output, cuda::args::deferred{count_transform}, cuda::std::plus<>{}, value_t{}); + REQUIRE(output[0] == num_items); + + output[0] = value_t{-4}; + const auto bounded_count = cuda::args::deferred{ + d_num_items, cuda::args::bounds(), cuda::args::bounds(count_t{0}, capacity)}; + device_reduce(d_input, d_output, bounded_count, cuda::std::plus<>{}, value_t{}); + REQUIRE(output[0] == num_items); +} + +C2H_TEST("DeviceReduce entry points sharing reduce dispatch accept deferred num_items", "[device][reduce][deferred]") +{ + using value_t = int64_t; + using count_t = int32_t; + + constexpr count_t capacity = 100'000; + constexpr count_t num_items = 1'000; + + c2h::device_vector input(capacity, value_t{2}); + input[num_items] = value_t{1}; + input[num_items + 1] = value_t{3}; + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto d_input = thrust::raw_pointer_cast(input.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto d_num_items = thrust::raw_pointer_cast(device_num_items.data()); + const auto count = cuda::args::deferred{d_num_items}; + + output[0] = value_t{-2}; + device_sum(d_input, d_output, count); + REQUIRE(output[0] == value_t{2} * num_items); + + output[0] = value_t{-3}; + device_min(d_input, d_output, count); + REQUIRE(output[0] == value_t{2}); + + output[0] = value_t{-4}; + device_max(d_input, d_output, count); + REQUIRE(output[0] == value_t{2}); + + output[0] = value_t{-5}; + device_transform_reduce(d_input, d_output, count, cuda::std::plus<>{}, thrust::square{}, value_t{}); + REQUIRE(output[0] == value_t{4} * num_items); +} + +C2H_TEST("DeviceReduce::Reduce consumes a count produced by DeviceSelect::If", "[device][reduce][deferred]") +{ + using value_t = int64_t; + using count_t = int32_t; + + constexpr count_t num_items = 100'000; + const count_t selected_count = GENERATE_COPY(count_t{0}, count_t{1}, count_t{1'000}, num_items); + CAPTURE(selected_count); + + const auto d_input = cuda::counting_iterator{value_t{0}}; + c2h::device_vector selected_items(num_items, value_t{num_items + 1}); + c2h::device_vector device_num_selected(1, count_t{-1}); + c2h::device_vector output(1, value_t{-1}); + + const auto d_selected_items = thrust::raw_pointer_cast(selected_items.data()); + const auto d_num_selected = thrust::raw_pointer_cast(device_num_selected.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto select_op = thrust::placeholders::_1 < value_t{selected_count}; + + launch(select_then_reduce_t{}, + d_input, + d_selected_items, + d_num_selected, + d_output, + num_items, + select_op, + cuda::std::plus<>{}, + value_t{}); + + const auto expected_sum = + static_cast(selected_count) * static_cast(selected_count - 1) / value_t{2}; + REQUIRE(device_num_selected[0] == selected_count); + REQUIRE(output[0] == expected_sum); +} + +#if TEST_LAUNCH == 0 +C2H_TEST("DeviceReduce::Reduce consumes a deferred count produced in another stream after an event", + "[device][reduce][deferred]") +{ + using value_t = int64_t; + using count_t = int32_t; + + constexpr count_t capacity = 100'000; + constexpr count_t selected_count = 1'000; + + int current_device{}; + REQUIRE(cudaSuccess == cudaGetDevice(¤t_device)); + cuda::stream producer{cuda::devices[current_device]}; + cuda::stream consumer{cuda::devices[current_device]}; + + const auto d_input = cuda::counting_iterator{value_t{0}}; + c2h::device_vector selected_items(capacity, value_t{capacity + 1}); + c2h::device_vector device_num_selected(1, count_t{-1}); + c2h::device_vector output(1, value_t{-1}); + + const auto d_selected_items = thrust::raw_pointer_cast(selected_items.data()); + const auto d_num_selected = thrust::raw_pointer_cast(device_num_selected.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto num_selected = cuda::args::deferred{d_num_selected}; + const auto select_op = thrust::placeholders::_1 < value_t{selected_count}; + + size_t select_temp_storage_bytes{}; + const auto select_query_error = cub::DeviceSelect::If( + nullptr, select_temp_storage_bytes, d_input, d_selected_items, d_num_selected, capacity, select_op, producer.get()); + REQUIRE(cudaSuccess == select_query_error); + + size_t reduce_temp_storage_bytes{}; + const auto reduce_query_error = cub::DeviceReduce::Reduce( + nullptr, + reduce_temp_storage_bytes, + d_selected_items, + d_output, + num_selected, + cuda::std::plus<>{}, + value_t{}, + consumer.get()); + REQUIRE(cudaSuccess == reduce_query_error); + + const size_t temp_storage_bytes = cuda::std::max(select_temp_storage_bytes, reduce_temp_storage_bytes); + c2h::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + const auto d_temp_storage = thrust::raw_pointer_cast(temp_storage.data()); + + const auto select_error = cub::DeviceSelect::If( + d_temp_storage, + select_temp_storage_bytes, + d_input, + d_selected_items, + d_num_selected, + capacity, + select_op, + producer.get()); + REQUIRE(cudaSuccess == select_error); + + const auto count_ready = producer.record_event(); + consumer.wait(count_ready); + + const auto reduce_error = cub::DeviceReduce::Reduce( + d_temp_storage, + reduce_temp_storage_bytes, + d_selected_items, + d_output, + num_selected, + cuda::std::plus<>{}, + value_t{}, + consumer.get()); + REQUIRE(cudaSuccess == reduce_error); + consumer.sync(); + + REQUIRE(device_num_selected[0] == selected_count); + const value_t expected_sum = + static_cast(selected_count) * static_cast(selected_count - 1) / value_t{2}; + REQUIRE(output[0] == expected_sum); +} + +C2H_TEST("DeviceReduce environment entry points accept deferred num_items", "[device][reduce][deferred]") +{ + using value_t = int64_t; + using count_t = int32_t; + + constexpr count_t num_items = 1'000; + c2h::device_vector input(num_items + 2, value_t{2}); + input[num_items] = value_t{1}; + input[num_items + 1] = value_t{3}; + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto count = cuda::args::deferred{device_num_items.begin()}; + + const auto reduce_error = + cub::DeviceReduce::Reduce(input.begin(), output.begin(), count, cuda::std::plus<>{}, value_t{7}); + REQUIRE(cudaSuccess == reduce_error); + REQUIRE(output[0] == value_t{7} + value_t{2} * num_items); + + output[0] = value_t{-2}; + const auto sum_error = cub::DeviceReduce::Sum(input.begin(), output.begin(), count); + REQUIRE(cudaSuccess == sum_error); + REQUIRE(output[0] == value_t{2} * num_items); + + output[0] = value_t{-3}; + const auto min_error = cub::DeviceReduce::Min(input.begin(), output.begin(), count); + REQUIRE(cudaSuccess == min_error); + REQUIRE(output[0] == value_t{2}); + + output[0] = value_t{-4}; + const auto max_error = cub::DeviceReduce::Max(input.begin(), output.begin(), count); + REQUIRE(cudaSuccess == max_error); + REQUIRE(output[0] == value_t{2}); + + output[0] = value_t{-5}; + const auto transform_reduce_error = cub::DeviceReduce::TransformReduce( + input.begin(), output.begin(), count, cuda::std::plus<>{}, thrust::square{}, value_t{}); + REQUIRE(cudaSuccess == transform_reduce_error); + REQUIRE(output[0] == value_t{4} * num_items); +} + +C2H_TEST("DeviceReduce::Reduce supports deferred num_items with not_guaranteed determinism", + "[device][reduce][deferred]") +{ + using value_t = float; + using count_t = int32_t; + + 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; + + const c2h::device_vector input(capacity, value_t{1}); + const c2h::device_vector device_num_items(1, num_items); + c2h::device_vector output(1, value_t{-1}); + + const auto env = cuda::execution::require(cuda::execution::determinism::not_guaranteed); + const auto reduce_error = cub::DeviceReduce::Reduce( + input.begin(), output.begin(), cuda::args::deferred{device_num_items.begin()}, cuda::std::plus<>{}, init, env); + REQUIRE(cudaSuccess == reduce_error); + REQUIRE(output[0] == init + static_cast(num_items)); +} +#endif // TEST_LAUNCH == 0 + +#if TEST_LAUNCH == 2 +C2H_TEST("captured DeviceReduce atomic dispatch replays with zero and nonzero deferred counts", + "[device][reduce][deferred][not_guaranteed]") +{ + 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_reduce_t reduce; + size_t temp_storage_bytes{}; + const auto query_error = + reduce(nullptr, temp_storage_bytes, input, d_output, count, cuda::std::plus<>{}, init, stream); + REQUIRE(cudaSuccess == query_error); + 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)); + const auto reduce_error = + reduce(d_temp_storage, temp_storage_bytes, input, d_output, count, cuda::std::plus<>{}, init, stream); + REQUIRE(cudaSuccess == reduce_error); + 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]") +{ + using value_t = int64_t; + using count_t = int32_t; + using offset_t = cub::detail::choose_offset_t; + + constexpr count_t capacity = 100'000; + cuda::compute_capability cc{}; + REQUIRE(cudaSuccess == cub::detail::ptx_compute_cap(cc)); + const auto policy = cub::detail::reduce::policy_selector_from_types>{}(cc); + const count_t tile_size = + static_cast(policy.multi_tile.threads_per_block * policy.multi_tile.items_per_thread); + + const auto d_input = cuda::counting_iterator{value_t{0}}; + c2h::device_vector selected_items(capacity, value_t{capacity + 1}); + c2h::device_vector device_bound(1, value_t{-1}); + c2h::device_vector device_num_selected(1, count_t{-1}); + c2h::device_vector output(1, value_t{-1}); + + const auto d_selected_items = thrust::raw_pointer_cast(selected_items.data()); + const auto d_bound = thrust::raw_pointer_cast(device_bound.data()); + const auto d_num_selected = thrust::raw_pointer_cast(device_num_selected.data()); + const auto d_output = thrust::raw_pointer_cast(output.data()); + const auto select_op = select_less_than_device_value_t{d_bound}; + const auto num_selected = cuda::args::deferred{d_num_selected}; + + cudaStream_t stream{}; + REQUIRE(cudaSuccess == cudaStreamCreate(&stream)); + + size_t select_temp_storage_bytes{}; + const auto select_query_error = cub::DeviceSelect::If( + nullptr, select_temp_storage_bytes, d_input, d_selected_items, d_num_selected, capacity, select_op, stream); + REQUIRE(cudaSuccess == select_query_error); + + size_t reduce_temp_storage_bytes{}; + const auto reduce_query_error = cub::DeviceReduce::Reduce( + nullptr, reduce_temp_storage_bytes, d_selected_items, d_output, num_selected, cuda::std::plus<>{}, value_t{}, stream); + REQUIRE(cudaSuccess == reduce_query_error); + + const size_t temp_storage_bytes = cuda::std::max(select_temp_storage_bytes, reduce_temp_storage_bytes); + 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)); + const auto select_error = cub::DeviceSelect::If( + d_temp_storage, select_temp_storage_bytes, d_input, d_selected_items, d_num_selected, capacity, select_op, stream); + REQUIRE(cudaSuccess == select_error); + const auto reduce_error = cub::DeviceReduce::Reduce( + d_temp_storage, + reduce_temp_storage_bytes, + d_selected_items, + d_output, + num_selected, + cuda::std::plus<>{}, + value_t{}, + stream); + REQUIRE(cudaSuccess == reduce_error); + REQUIRE(cudaSuccess == cudaStreamEndCapture(stream, &graph)); + + cudaGraphExec_t executable{}; + REQUIRE(cudaSuccess == cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); + + for (const count_t selected_count : + {count_t{0}, count_t{1}, tile_size - 1, tile_size, tile_size + 1, capacity, count_t{0}}) + { + const value_t bound = selected_count; + REQUIRE(cudaSuccess == cudaMemcpyAsync(d_bound, &bound, sizeof(value_t), cudaMemcpyHostToDevice, stream)); + REQUIRE(cudaSuccess == cudaGraphLaunch(executable, stream)); + REQUIRE(cudaSuccess == cudaStreamSynchronize(stream)); + REQUIRE(device_num_selected[0] == selected_count); + const value_t expected_sum = + static_cast(selected_count) * static_cast(selected_count - 1) / value_t{2}; + REQUIRE(output[0] == expected_sum); + } + + REQUIRE(cudaSuccess == cudaGraphExecDestroy(executable)); + REQUIRE(cudaSuccess == cudaGraphDestroy(graph)); + REQUIRE(cudaSuccess == cudaStreamDestroy(stream)); +} +#endif // TEST_LAUNCH == 2 diff --git a/cub/test/catch2_test_device_reduce_env.cu b/cub/test/catch2_test_device_reduce_env.cu index 5df99e17abc..f520211afd6 100644 --- a/cub/test/catch2_test_device_reduce_env.cu +++ b/cub/test/catch2_test_device_reduce_env.cu @@ -376,6 +376,7 @@ C2H_TEST("Device reduce uses environment", "[reduce][device]", requirements) decltype(d_in), accumulator_t*, offset_t, + offset_t, op_t, accumulator_t, init_value_t, @@ -416,6 +417,7 @@ C2H_TEST("Device reduce uses environment", "[reduce][device]", requirements) decltype(d_in), decltype(raw_ptr), offset_t, + offset_t, op_t, accumulator_t, init_value_t, @@ -514,6 +516,7 @@ C2H_TEST("Device sum uses environment", "[reduce][device]", requirements) decltype(d_in), accumulator_t*, offset_t, + offset_t, op_t, accumulator_t, init_value_t, @@ -554,6 +557,7 @@ C2H_TEST("Device sum uses environment", "[reduce][device]", requirements) decltype(d_in), decltype(raw_ptr), offset_t, + offset_t, op_t, accumulator_t, init_value_t, @@ -650,6 +654,7 @@ C2H_TEST("Device reduce not_guaranteed falls back when output type differs from decltype(d_in.begin()), accumulator_t*, offset_t, + offset_t, op_t, accumulator_t, init_value_t, @@ -711,6 +716,7 @@ C2H_TEST("Device sum not_guaranteed falls back when output type differs from acc decltype(d_in.begin()), accumulator_t*, offset_t, + offset_t, op_t, accumulator_t, init_value_t,