diff --git a/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh b/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh index 31292948a98..b62c48f58e4 100644 --- a/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh +++ b/cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh @@ -24,8 +24,11 @@ #include #include #include +#include #include +#include + #include #include #include @@ -35,6 +38,7 @@ #include #include #include +#include #include #include @@ -84,13 +88,72 @@ struct deterministic_sum_t template +struct DeterministicDeviceReduceKernelSource +{ + // PolicySelector must be stateless, so we can pass the type to the kernel + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER( + SingleTileKernel, + reduce::DeterministicDeviceReduceSingleTileKernel< + PolicySelector, + InputIteratorT, + OutputIteratorT, + ReductionOpT, + InitValueT, + DeterministicAccumT, + TransformOpT>) + + CUB_DEFINE_KERNEL_GETTER( + ReductionKernel, + reduce:: + DeterministicDeviceReduceKernel) + + // The second single-tile pass reduces the per-block partials and always uses an identity transform + CUB_DEFINE_KERNEL_GETTER( + SingleTileSecondKernel, + reduce::DeterministicDeviceReduceSingleTileKernel< + PolicySelector, + DeterministicAccumT*, + OutputIteratorT, + ReductionOpT, + InitValueT, + DeterministicAccumT>) +}; + +//! Kernel source instantiation matching the default used by `rfa::dispatch` below. Tests use this alias to +//! reference the exact kernel instantiations the dispatch launches. +template , + typename PolicySelector = + policy_selector_from_types, __determinism_t::__gpu_to_gpu>> +using default_kernel_source_t = DeterministicDeviceReduceKernelSource< + PolicySelector, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + OutputIteratorT, + deterministic_sum_t, + InitValueT, + typename deterministic_sum_t::DeterministicAcc, + TransformOpT>; + +template CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invoke_single_tile( + SingleTileKernelT single_tile_kernel, void* d_temp_storage, size_t& temp_storage_bytes, InputIteratorT d_in, @@ -122,20 +185,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok // Invoke single_reduce_sweep_kernel if (const auto error = CubDebug( launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) - .doit(detail::reduce::DeterministicDeviceReduceSingleTileKernel< - PolicySelector, - InputIteratorT, - OutputIteratorT, - ReductionOpT, - InitValueT, - DeterministicAccumT, - TransformOpT>, - d_in, - d_out, - static_cast(num_items), - reduction_op, - init, - transform_op))) + .doit(single_tile_kernel, d_in, d_out, static_cast(num_items), reduction_op, init, transform_op))) { return error; } @@ -150,16 +200,19 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok return CubDebug(detail::DebugSyncStream(stream)); } -template CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invoke_passes( + ReduceKernelT reduce_kernel, + SingleTileKernelT single_tile_second_kernel, void* d_temp_storage, size_t& temp_storage_bytes, InputIteratorT d_in, @@ -179,10 +232,7 @@ 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, - active_policy.multi_tile))) + if (const auto error = CubDebug(reduce_config.__init(reduce_kernel, active_policy.multi_tile))) { return error; } @@ -249,11 +299,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok if (const auto error = CubDebug( launcher_factory(current_grid_size, active_policy.multi_tile.threads_per_block, 0, stream) - .doit(detail::reduce::DeterministicDeviceReduceKernel, + .doit(reduce_kernel, d_in, d_chunk_block_reductions, num_current_items, @@ -295,13 +341,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok // 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>, + .doit(single_tile_second_kernel, d_block_reductions, d_out, reduce_grid_size, @@ -322,15 +362,18 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invok return CubDebug(detail::DebugSyncStream(stream)); } -template , - typename PolicySelector = - policy_selector_from_types, __determinism_t::__gpu_to_gpu>, - typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +template < + typename InputIteratorT, + typename OutputIteratorT, + typename OffsetT, + typename InitValueT, + typename TransformOpT = ::cuda::std::identity, + typename AccumT = accum_t, + typename PolicySelector = + policy_selector_from_types, __determinism_t::__gpu_to_gpu>, + typename KernelSource = + default_kernel_source_t, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -341,6 +384,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( cudaStream_t stream = {}, TransformOpT transform_op = {}, PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}) { // Get CC @@ -373,14 +417,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( if (num_items <= tile_items) { - return invoke_single_tile( + return invoke_single_tile( + kernel_source.SingleTileKernel(), d_temp_storage, temp_storage_bytes, d_in_unwrapped, @@ -394,14 +432,9 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( launcher_factory); } - return invoke_passes( + return invoke_passes( + kernel_source.ReductionKernel(), + kernel_source.SingleTileSecondKernel(), d_temp_storage, temp_storage_bytes, d_in_unwrapped, diff --git a/cub/test/catch2_test_device_reduce_env.cu b/cub/test/catch2_test_device_reduce_env.cu index 5df99e17abc..6cafdbc0a1a 100644 --- a/cub/test/catch2_test_device_reduce_env.cu +++ b/cub/test/catch2_test_device_reduce_env.cu @@ -10,6 +10,7 @@ struct stream_registry_factory_t; #include #include +#include #include #include @@ -332,6 +333,88 @@ using requirements = cuda::execution::determinism::run_to_run_t, cuda::execution::determinism::not_guaranteed_t>; +// Builds the kernel allowlist for the requirements-parametrized tests below and computes the expected temporary +// storage size. To check if a given algorithm implementation is used, the tests check if the associated kernels are +// invoked. The allowlists are built from the same kernel source types the dispatch uses, so the expected kernel +// instantiations cannot drift from the launched ones. This is deliberately a function template instead of a lambda +// in the tests: with nvcc 13.3 on MSVC, decltype() of a variable captured by reference inside a lambda is +// misevaluated as a reference type, which would make the allowlist reference (and instantiate) different kernels +// than the ones the dispatch launches. See: https://github.com/NVIDIA/cccl/issues/9643 +template +static auto expected_reduce_kernels( + InputItT d_in, + OutputItT d_out, + AccumulatorT* d_out_ptr, + NumItemsT num_items, + InitT init, + size_t& expected_bytes_allocated, + RunToRunSizeF run_to_run_size) +{ + if constexpr (std::is_same_v) + { + REQUIRE(cudaSuccess == run_to_run_size()); + + using policy_t = cub::detail::reduce::policy_selector_from_types; + using kernel_source_t = cub::detail::reduce:: + DeviceReduceKernelSource; + return cuda::std::array{ + reinterpret_cast(kernel_source_t::SingleTileKernel()), + reinterpret_cast(kernel_source_t::ReductionKernel()), + reinterpret_cast(kernel_source_t::SingleTileSecondKernel())}; + } + else if constexpr (cub::detail::is_non_deterministic_v) + { + REQUIRE( + cudaSuccess + == cub::detail::reduce::dispatch( + nullptr, + expected_bytes_allocated, + d_in, + d_out_ptr, + num_items, + OpT{}, + init, + /* stream */ nullptr, + TransformT{})); + + using policy_t = + cub::detail::reduce::policy_selector_from_types; + using kernel_source_t = cub::detail::reduce::DeviceReduceKernelSource< + policy_t, + InputItT, + AccumulatorT*, + OffsetT, + OpT, + InitT, + AccumulatorT, + TransformT, + /* StableReductionOrder */ false>; + return cuda::std::array{reinterpret_cast(kernel_source_t::ReductionKernel())}; + } + else + { + REQUIRE(cudaSuccess + == cub::detail::rfa::dispatch( + nullptr, expected_bytes_allocated, d_in, d_out, num_items, init)); + + using kernel_source_t = + cub::detail::rfa::default_kernel_source_t; + return cuda::std::array{ + reinterpret_cast(kernel_source_t::SingleTileKernel()), + reinterpret_cast(kernel_source_t::ReductionKernel()), + reinterpret_cast(kernel_source_t::SingleTileSecondKernel())}; + } +} + C2H_TEST("Device reduce uses environment", "[reduce][device]", requirements) { using determinism_t = c2h::get<0, TestType>; @@ -349,117 +432,10 @@ C2H_TEST("Device reduce uses environment", "[reduce][device]", requirements) init_value_t init = 0; size_t expected_bytes_allocated{}; - // To check if a given algorithm implementation is used, we check if associated kernels are invoked. - auto kernels = [&]() { - if constexpr (std::is_same_v) - { - REQUIRE( - cudaSuccess - == cub::DeviceReduce::Reduce(nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items, op_t{}, init)); - - using policy_t = cub::detail::reduce::policy_selector_from_types; - return cuda::std::array{ - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - decltype(d_in), - decltype(d_out.begin()), - offset_t, - op_t, - init_value_t, - accumulator_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ true, - decltype(d_in), - accumulator_t*, - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - accumulator_t*, - decltype(d_out.begin()), - int, // always used with int offset - op_t, - init_value_t, - accumulator_t>)}; - } - else if constexpr (cub::detail::is_non_deterministic_v) - { - using policy_t = - cub::detail::reduce::policy_selector_from_types; - auto* raw_ptr = thrust::raw_pointer_cast(d_out.data()); - - REQUIRE( - cudaSuccess - == cub::detail::reduce::dispatch( - nullptr, - expected_bytes_allocated, - d_in, - raw_ptr, - num_items, - op_t{}, - init, - /* stream */ nullptr, - transform_t{})); - - return cuda::std::array{reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ false, - decltype(d_in), - decltype(raw_ptr), - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>)}; - } - else - { - using deterministic_add_t = cub::detail::rfa::deterministic_sum_t; - using reduction_op_t = deterministic_add_t; - using policy_t = cub::detail::reduce:: - policy_selector_from_types; - using deterministic_accum_t = deterministic_add_t::DeterministicAcc; - using output_it_t = decltype(d_out.begin()); - - REQUIRE(cudaSuccess - == cub::detail::rfa:: - dispatch( - nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items, init)); - - auto k1 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< - policy_t, - decltype(d_in), - output_it_t, - reduction_op_t, - init_value_t, - deterministic_accum_t, - transform_t>; - auto k2 = cub::detail::reduce:: - DeterministicDeviceReduceKernel; - auto k3 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< - policy_t, - deterministic_accum_t*, - output_it_t, - reduction_op_t, - init_value_t, - deterministic_accum_t, - transform_t>; - // TODO(bgruber): enable this when we have Catch2 3.13+ - // UNSCOPED_CAPTURE(c2h::type_name(), c2h::type_name(), - // c2h::type_name()); - return cuda::std::array{ - reinterpret_cast(k1), reinterpret_cast(k2), reinterpret_cast(k3)}; - } - }(); + auto kernels = expected_reduce_kernels( + d_in, d_out.begin(), thrust::raw_pointer_cast(d_out.data()), num_items, init, expected_bytes_allocated, [&]() { + return cub::DeviceReduce::Reduce(nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items, op_t{}, init); + }); // Equivalent to `cuexec::require(cuexec::determinism::run_to_run)` and // `cuexec::require(cuexec::determinism::not_guaranteed)` @@ -486,118 +462,13 @@ C2H_TEST("Device sum uses environment", "[reduce][device]", requirements) auto d_in = cuda::constant_iterator(1.0f); auto d_out = thrust::device_vector(1); - [[maybe_unused]] init_value_t init = 0; + init_value_t init = 0; size_t expected_bytes_allocated{}; - // To check if a given algorithm implementation is used, we check if associated kernels are invoked. - auto kernels = [&]() { - if constexpr (std::is_same_v) - { - REQUIRE(cudaSuccess == cub::DeviceReduce::Sum(nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items)); - - using policy_t = cub::detail::reduce::policy_selector_from_types; - return cuda::std::array{ - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - decltype(d_in), - decltype(d_out.begin()), - offset_t, - op_t, - init_value_t, - accumulator_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ true, - decltype(d_in), - accumulator_t*, - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - accumulator_t*, - decltype(d_out.begin()), - int, // always used with int offset - op_t, - init_value_t, - accumulator_t>)}; - } - else if constexpr (cub::detail::is_non_deterministic_v) - { - using policy_t = - cub::detail::reduce::policy_selector_from_types; - auto* raw_ptr = thrust::raw_pointer_cast(d_out.data()); - - REQUIRE( - cudaSuccess - == cub::detail::reduce::dispatch( - nullptr, - expected_bytes_allocated, - d_in, - raw_ptr, - num_items, - op_t{}, - init, - /* stream */ nullptr, - transform_t{})); - - return cuda::std::array{reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ false, - decltype(d_in), - decltype(raw_ptr), - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>)}; - } - else - { - using deterministic_add_t = cub::detail::rfa::deterministic_sum_t; - using reduction_op_t = deterministic_add_t; - using policy_t = cub::detail::reduce:: - policy_selector_from_types; - using deterministic_accum_t = deterministic_add_t::DeterministicAcc; - using output_it_t = decltype(d_out.begin()); - - REQUIRE(cudaSuccess - == cub::detail::rfa:: - dispatch( - nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items, init)); - - auto k1 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< - policy_t, - decltype(d_in), - output_it_t, - reduction_op_t, - init_value_t, - deterministic_accum_t, - transform_t>; - auto k2 = cub::detail::reduce:: - DeterministicDeviceReduceKernel; - auto k3 = cub::detail::reduce::DeterministicDeviceReduceSingleTileKernel< - policy_t, - deterministic_accum_t*, - output_it_t, - reduction_op_t, - init_value_t, - deterministic_accum_t, - transform_t>; - // TODO(bgruber): enable this when we have Catch2 3.13+ - // UNSCOPED_CAPTURE(c2h::type_name(), c2h::type_name(), - // c2h::type_name()); - return cuda::std::array{ - reinterpret_cast(k1), reinterpret_cast(k2), reinterpret_cast(k3)}; - } - }(); + auto kernels = expected_reduce_kernels( + d_in, d_out.begin(), thrust::raw_pointer_cast(d_out.data()), num_items, init, expected_bytes_allocated, [&]() { + return cub::DeviceReduce::Sum(nullptr, expected_bytes_allocated, d_in, d_out.begin(), num_items); + }); // Equivalent to `cuexec::require(cuexec::determinism::run_to_run)` and // `cuexec::require(cuexec::determinism::not_guaranteed)` @@ -631,38 +502,22 @@ C2H_TEST("Device reduce not_guaranteed falls back when output type differs from == cub::DeviceReduce::Reduce( nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_items, op_t{}, init)); - using policy_t = cub::detail::reduce::policy_selector_from_types; - auto kernels = cuda::std::array{ - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - decltype(d_in.begin()), - decltype(d_out.begin()), - offset_t, - op_t, - init_value_t, - accumulator_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ true, - decltype(d_in.begin()), - accumulator_t*, - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - accumulator_t*, - decltype(d_out.begin()), - int, // always used with int offset - op_t, - init_value_t, - accumulator_t>)}; + // Build the allowlist from the same kernel source type the dispatch uses, so the expected kernel + // instantiations cannot drift from the launched ones. + using policy_t = cub::detail::reduce::policy_selector_from_types; + using kernel_source_t = cub::detail::reduce::DeviceReduceKernelSource< + policy_t, + decltype(d_in.begin()), + decltype(d_out.begin()), + offset_t, + op_t, + init_value_t, + accumulator_t, + transform_t>; + auto kernels = cuda::std::array{ + reinterpret_cast(kernel_source_t::SingleTileKernel()), + reinterpret_cast(kernel_source_t::ReductionKernel()), + reinterpret_cast(kernel_source_t::SingleTileSecondKernel())}; auto env = stdexec::env{cuda::execution::require(cuda::execution::determinism::not_guaranteed), allowed_kernels(kernels), @@ -692,38 +547,22 @@ C2H_TEST("Device sum not_guaranteed falls back when output type differs from acc REQUIRE( cudaSuccess == cub::DeviceReduce::Sum(nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_items)); - using policy_t = cub::detail::reduce::policy_selector_from_types; - auto kernels = cuda::std::array{ - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - decltype(d_in.begin()), - decltype(d_out.begin()), - offset_t, - op_t, - init_value_t, - accumulator_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceKernel< - policy_t, - /* StableReductionOrder */ true, - decltype(d_in.begin()), - accumulator_t*, - offset_t, - op_t, - accumulator_t, - init_value_t, - transform_t>), - reinterpret_cast( - cub::detail::reduce::DeviceReduceSingleTileKernel< - policy_t, - accumulator_t*, - decltype(d_out.begin()), - int, // always used with int offset - op_t, - init_value_t, - accumulator_t>)}; + // Build the allowlist from the same kernel source type the dispatch uses, so the expected kernel + // instantiations cannot drift from the launched ones. + using policy_t = cub::detail::reduce::policy_selector_from_types; + using kernel_source_t = cub::detail::reduce::DeviceReduceKernelSource< + policy_t, + decltype(d_in.begin()), + decltype(d_out.begin()), + offset_t, + op_t, + init_value_t, + accumulator_t, + transform_t>; + auto kernels = cuda::std::array{ + reinterpret_cast(kernel_source_t::SingleTileKernel()), + reinterpret_cast(kernel_source_t::ReductionKernel()), + reinterpret_cast(kernel_source_t::SingleTileSecondKernel())}; auto env = stdexec::env{cuda::execution::require(cuda::execution::determinism::not_guaranteed), allowed_kernels(kernels), diff --git a/cub/test/catch2_test_env_launch_helper.h b/cub/test/catch2_test_env_launch_helper.h index 950bd5654eb..ca078719273 100644 --- a/cub/test/catch2_test_env_launch_helper.h +++ b/cub/test/catch2_test_env_launch_helper.h @@ -55,6 +55,11 @@ expected_allocation_size(size_t expected) struct get_allowed_kernels_t {}; +//! The listed kernels must be the exact instantiations the dispatch launches, taken from the dispatch's kernel +//! source. Beware of computing the expected instantiations with decltype() of a variable captured by reference +//! inside a lambda: nvcc 13.3 with MSVC misevaluates such a decltype() as a reference type, silently yielding a +//! different kernel instantiation. Compute the types at function scope instead. +//! See: https://github.com/NVIDIA/cccl/issues/9643 __host__ __device__ static cuda::std::execution::prop> allowed_kernels(cuda::std::span allowed_kernels) {