diff --git a/cub/cub/agent/agent_reduce_by_key.cuh b/cub/cub/agent/agent_reduce_by_key.cuh index e44d6c665e42..b9051a652089 100644 --- a/cub/cub/agent/agent_reduce_by_key.cuh +++ b/cub/cub/agent/agent_reduce_by_key.cuh @@ -108,8 +108,12 @@ namespace detail::reduce_by_key * * @tparam AccumT * The type of intermediate accumulator (according to P2322R6) + * + * @tparam StableReductionOrder + * Whether to use a stable reduction order across tiles */ template ; + TilePrefixCallbackOp; // Key and value exchange types using KeyExchangeT = KeyOutputT[TILE_ITEMS + 1]; diff --git a/cub/cub/device/device_reduce.cuh b/cub/cub/device/device_reduce.cuh index b515cffe1f17..d46a76f82e36 100644 --- a/cub/cub/device/device_reduce.cuh +++ b/cub/cub/device/device_reduce.cuh @@ -46,6 +46,7 @@ #include #include #include +#include #include #include #include @@ -153,7 +154,9 @@ inline constexpr bool is_non_deterministic_v = //! ==================================== //! //! ``cub::DeviceReduce`` supports all three :ref:`determinism guarantees `; the -//! default is ``run_to_run``. +//! default is ``run_to_run``. ``ReduceByKey`` has separate type/operator constraints, documented on its +//! environment overload below. The implementation details below apply to the other reductions; +//! ``ReduceByKey`` does not use RFA or atomic accumulation. //! //! - ``run_to_run`` (the default) is reproducible because, for a given GPU, every launch with the same input, //! build, and launch configuration selects the *same* tuning policy and therefore performs the *same* fixed @@ -2734,8 +2737,13 @@ public: //! respectively. The total number of runs encountered is written to ``d_num_runs_out``. //! //! - The ``==`` equality operator is used to determine whether keys are equivalent - //! - Provides "run-to-run" determinism for pseudo-associative reduction - //! (e.g., addition of floating point types) on the same GPU device. + //! - Requests "run-to-run" determinism by default. It is supported for integral types with known operators, + //! primitive types with min/max, and floating-point types with ``cuda::std::plus``. + //! - Pass ``cuda::execution::require(cuda::execution::determinism::not_guaranteed)`` in the environment to opt out. + //! This removes the determinism requirement without changing the algorithm's other requirements. + //! - "gpu-to-gpu" determinism is supported for integral types with known operators and primitive types with min/max. + //! Floating-point addition is unsupported; ``ReduceByKey`` does not use RFA. + //! Other combinations under "run-to-run" or "gpu-to-gpu" are rejected at compile time. //! - Let ``out`` be any of //! ``[d_unique_out, d_unique_out + *d_num_runs_out)`` //! ``[d_aggregates_out, d_aggregates_out + *d_num_runs_out)`` @@ -2743,6 +2751,12 @@ public: //! ``[d_keys_in, d_keys_in + num_items)``, //! ``[d_values_in, d_values_in + num_items)`` nor ``out`` in any way. //! + //! These determinism constraints apply to this environment overload and its accumulator type. Known CUDA binary + //! operators are ``cuda::minimum``, ``cuda::maximum``, and ``cuda::std::plus``, ``multiplies``, ``bit_and``, + //! ``bit_or``, ``bit_xor``, ``logical_and``, and ``logical_or``. Primitive accumulators are the types recognized by + //! ``cub::Traits``; floating-point accumulators (including CUDA extended types) are recognized by + //! ``cuda::is_floating_point_v``. The operator must support the accumulator type. + //! //! Snippet //! +++++++++++++++++++++++++++++++++++++++++++++ //! @@ -2780,7 +2794,8 @@ public: //! //! @tparam EnvT //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. - //! Supports customization of stream via ``cuda::get_stream``. + //! Supports customization of stream via ``cuda::get_stream`` and determinism through + //! ``cuda::execution::require``. //! //! @param[in] d_keys_in //! Pointer to the input sequence of keys @@ -2830,15 +2845,43 @@ public: { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ReduceByKey"); - using OffsetT = detail::choose_offset_t; - using EqualityOp = ::cuda::std::equal_to<>; - using default_policy_selector = detail::reduce_by_key::policy_selector_from_types< - ReductionOpT, - ::cuda::std::__accumulator_t>, - detail::non_void_value_t>>; + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + using accum_t = ::cuda::std::__accumulator_t>; + using key_t = detail::non_void_value_t>; + using OffsetT = detail::choose_offset_t; + using EqualityOp = ::cuda::std::equal_to<>; + + constexpr bool is_run_to_run_required = + ::cuda::std::is_same_v; + constexpr bool is_gpu_to_gpu_required = + ::cuda::std::is_same_v; + constexpr bool is_safe_integral_op = + ::cuda::std::is_integral_v && detail::is_cuda_binary_operator; + constexpr bool is_primitive_min_max_op = + detail::is_primitive_v && detail::is_cuda_minimum_maximum_v; + constexpr bool is_fp_plus_op = + ::cuda::is_floating_point_v && detail::is_cuda_std_plus_v; + + static_assert(!is_run_to_run_required || is_safe_integral_op || is_primitive_min_max_op || is_fp_plus_op, + "run_to_run deterministic reduce-by-key requires integral types with known operators, " + "primitive types with min/max, or floating-point types with plus operator"); + static_assert(!is_gpu_to_gpu_required || is_safe_integral_op || is_primitive_min_max_op, + "gpu_to_gpu deterministic reduce-by-key requires integral types with known operators or " + "primitive types with min/max"); + + static constexpr bool stable_reduction_order = is_run_to_run_required && is_fp_plus_op; + using default_policy_selector = detail::reduce_by_key::policy_selector_from_types; return detail::dispatch_with_env_and_tuning( env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { - return detail::reduce_by_key::dispatch( + return detail::reduce_by_key::dispatch( storage, bytes, d_keys_in, @@ -2867,8 +2910,7 @@ public: //! respectively. The total number of runs encountered is written to ``d_num_runs_out``. //! //! - The ``==`` equality operator is used to determine whether keys are equivalent - //! - Provides "run-to-run" determinism for pseudo-associative reduction - //! (e.g., addition of floating point types) on the same GPU device. + //! - Provides "run-to-run" determinism for floating-point reduction with ``cuda::std::plus`` on the same GPU device. //! However, results for pseudo-associative reduction may be inconsistent //! from one device to a another device of a different compute-capability //! because CUB can employ different tile-sizing for different architectures. @@ -3011,10 +3053,14 @@ public: { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ReduceByKey"); + using accum_t = ::cuda::std::__accumulator_t>; using OffsetT = detail::choose_offset_t; using EqualityOp = ::cuda::std::equal_to<>; - return detail::reduce_by_key::dispatch( + static constexpr bool stable_reduction_order = + ::cuda::is_floating_point_v && detail::is_cuda_std_plus_v; + + return detail::reduce_by_key::dispatch( d_temp_storage, temp_storage_bytes, d_keys_in, diff --git a/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh b/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh index 926f40ef2032..4056d696eb48 100644 --- a/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh @@ -112,6 +112,14 @@ struct streaming_context } }; +// Bind the stability flag while selecting the default and fallback agent policies. +template +using vsmem_helper_t = + vsmem_helper_with_fallback_impl, + policy_wrapper_t, + AgentReduceByKey, StableReductionOrder, Args...>>; + /** * @brief Multi-block reduce-by-key sweep kernel entry point * @@ -145,6 +153,9 @@ struct streaming_context * @tparam OffsetT * Signed integer type for global offsets * + * @tparam StableReductionOrder + * Whether to use a stable reduction order across tiles + * * @param d_keys_in * Pointer to the input sequence of keys * @@ -187,7 +198,8 @@ template + typename StreamingContextT, + bool StableReductionOrder> #if _CCCL_HAS_CONCEPTS() requires reduce_by_key_policy_selector #endif @@ -217,9 +229,9 @@ __launch_bounds__(int(current_policy().lookback.threads_per_bloc policy.lookback.lookback_delay.delay, policy.lookback.lookback_delay.l2_write_latency>>; - using vsmem_helper_t = vsmem_helper_default_fallback_policy_t< + using vsmem_helper_t = reduce_by_key::vsmem_helper_t< + StableReductionOrder, AgentReduceByKeyPolicyT, - AgentReduceByKey, KeysInputIteratorT, UniqueOutputIteratorT, ValuesInputIteratorT, @@ -374,9 +386,9 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t Invoke(ScanInitKernelT init_kernel, ReduceByKeyKernelT reduce_by_key_kernel) { - using vsmem_helper_t = detail::vsmem_helper_default_fallback_policy_t< + using vsmem_helper_t = detail::reduce_by_key::vsmem_helper_t< + false, typename ActivePolicyT::ReduceByKeyPolicyT, - detail::reduce_by_key::AgentReduceByKey, KeysInputIteratorT, UniqueOutputIteratorT, ValuesInputIteratorT, @@ -555,7 +567,8 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey ReductionOpT, OffsetT, AccumT, - streaming_context_t>); + streaming_context_t, + false>); } /** @@ -650,7 +663,7 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey namespace detail::reduce_by_key { // we move the conversion of the policy to the agent policy and its use out of the lambda below, so MSVC does not ICE -template +template _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_getter) { // TODO(bgruber): refactor this in the future @@ -664,13 +677,14 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get delay_constructor_t>; - using vsmem_helper_t = vsmem_helper_default_fallback_policy_t; + using vsmem_helper_t = reduce_by_key::vsmem_helper_t; return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS, vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD, vsmem_helper_t::vsmem_per_block}; } -template ; + streaming_context_t, + StableReductionOrder>; int reduce_by_key_sm_occupancy{}; if (const auto error = diff --git a/cub/cub/device/dispatch/dispatch_rle_encode.cuh b/cub/cub/device/dispatch/dispatch_rle_encode.cuh index f67632dc00f2..a07285a88190 100644 --- a/cub/cub/device/dispatch/dispatch_rle_encode.cuh +++ b/cub/cub/device/dispatch/dispatch_rle_encode.cuh @@ -111,9 +111,9 @@ __launch_bounds__(int(current_policy().lookback.threads_per_bloc policy.lookback.lookback_delay.delay, policy.lookback.lookback_delay.l2_write_latency>>; - using vsmem_helper_t = vsmem_helper_default_fallback_policy_t< + using vsmem_helper_t = reduce_by_key::vsmem_helper_t< + false, AgentReduceByKeyPolicyT, - reduce_by_key::AgentReduceByKey, KeysInputIteratorT, UniqueOutputIteratorT, ValuesInputIteratorT, diff --git a/cub/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh b/cub/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh index 7dde6ba68ba0..78d3a088bdbd 100644 --- a/cub/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh @@ -36,7 +36,7 @@ struct reduce_by_key_kernel_source template [[nodiscard]] CUB_RUNTIME_FUNCTION static constexpr auto reduce_by_key_kernel() { - return &DeviceReduceByKeyKernel; + return &DeviceReduceByKeyKernel; } }; diff --git a/cub/test/catch2_test_device_reduce_by_key.cu b/cub/test/catch2_test_device_reduce_by_key.cu index eff1ad0449d2..1259e3ee0da3 100644 --- a/cub/test/catch2_test_device_reduce_by_key.cu +++ b/cub/test/catch2_test_device_reduce_by_key.cu @@ -5,6 +5,9 @@ #include +#include +#include + #include "catch2_test_device_reduce.cuh" #include "catch2_test_launch_helper.h" #include "cub_test_macros.h" @@ -156,3 +159,89 @@ CUB_TEST("Device reduce-by-key works", "[by_key][reduce][device]", CUB_SMALL, fu REQUIRE(expected_keys == out_unique_keys); } } + +#if TEST_LAUNCH == 0 && TEST_TYPES == 0 +CUB_TEST_CASE("Device reduce-by-key is run-to-run deterministic for fp64 sums", "[by_key][reduce][device]", CUB_SMALL) +{ + // Span several waves of tiles so that the unfixed lookback grouping can vary on large GPUs. + constexpr int num_items = 262'144; + constexpr int long_run = 200'000; + constexpr int short_run = 64; + constexpr int num_repetitions = 20; + constexpr int num_runs = 1 + (num_items - long_run + short_run - 1) / short_run; + + c2h::host_vector keys(num_items); + + for (int i = 0; i < num_items; ++i) + { + keys[i] = i < long_run ? 0 : 1 + (i - long_run) / short_run; + } + + c2h::device_vector keys_in = keys; + c2h::device_vector values_in(num_items); + c2h::gen(C2H_SEED(2), values_in, -1e-2, 1e-2); + c2h::device_vector unique_out(num_runs); + c2h::device_vector aggregates_out(num_runs); + c2h::device_vector num_runs_out(1); + + const auto d_keys_in = thrust::raw_pointer_cast(keys_in.data()); + const auto d_values_in = thrust::raw_pointer_cast(values_in.data()); + const auto d_unique_out = thrust::raw_pointer_cast(unique_out.data()); + const auto d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); + const auto d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); + + std::size_t temp_storage_bytes{}; + REQUIRE( + cudaSuccess + == cub::DeviceReduce::ReduceByKey( + nullptr, + temp_storage_bytes, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + cuda::std::plus{}, + num_items)); + c2h::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + const auto d_temp_storage = thrust::raw_pointer_cast(temp_storage.data()); + + c2h::host_vector expected_keys; + c2h::host_vector expected_aggregates; + for (int repetition = 0; repetition < num_repetitions; ++repetition) + { + REQUIRE(cudaSuccess == cudaMemset(d_aggregates_out, 0xee, num_runs * sizeof(double))); + REQUIRE( + cudaSuccess + == cub::DeviceReduce::ReduceByKey( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + cuda::std::plus{}, + num_items)); + + REQUIRE(num_runs_out[0] == num_runs); + const c2h::host_vector result_keys = unique_out; + const c2h::host_vector result_aggregates = aggregates_out; + c2h::host_vector result_bits(num_runs); + for (int run = 0; run < num_runs; ++run) + { + result_bits[run] = cuda::std::bit_cast(result_aggregates[run]); + } + if (repetition == 0) + { + expected_keys = result_keys; + expected_aggregates = result_bits; + } + else + { + REQUIRE(result_keys == expected_keys); + REQUIRE(result_bits == expected_aggregates); + } + } +} +#endif // TEST_LAUNCH == 0 && TEST_TYPES == 0 diff --git a/cub/test/catch2_test_device_reduce_env_api.cu b/cub/test/catch2_test_device_reduce_env_api.cu index d71617440b5e..8dc27bf76c29 100644 --- a/cub/test/catch2_test_device_reduce_env_api.cu +++ b/cub/test/catch2_test_device_reduce_env_api.cu @@ -12,6 +12,7 @@ #include #include #include +#include #include #include "catch2_test_memory_resources.h" @@ -599,6 +600,78 @@ CUB_TEST("cub::DeviceReduce::ReduceByKey accepts not_guaranteed determinism requ REQUIRE(aggregates_out == expected_aggregates); } +CUB_TEST("cub::DeviceReduce::ReduceByKey accepts run_to_run for floating-point addition", "[reduce][env]", CUB_SMALL) +{ + auto keys_in = thrust::device_vector{0, 2, 2, 9, 5, 5, 5, 8}; + auto values_in = thrust::device_vector{0, 7, 1, 6, 2, 5, 3, 4}; + auto unique_keys_out = thrust::device_vector(5); + auto aggregates_out = thrust::device_vector(5); + auto num_runs_out = thrust::device_vector(1); + + auto env = cuda::execution::require(cuda::execution::determinism::run_to_run); + + auto error = cub::DeviceReduce::ReduceByKey( + keys_in.begin(), + unique_keys_out.begin(), + values_in.begin(), + aggregates_out.begin(), + num_runs_out.begin(), + cuda::std::plus<>{}, + static_cast(keys_in.size()), + env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceReduce::ReduceByKey failed with status: " << error << '\n'; + } + + thrust::device_vector expected_keys{0, 2, 9, 5, 8}; + thrust::device_vector expected_aggregates{0, 8, 6, 10, 4}; + + REQUIRE(error == cudaSuccess); + REQUIRE(num_runs_out[0] == 5); + unique_keys_out.resize(5); + aggregates_out.resize(5); + REQUIRE(unique_keys_out == expected_keys); + REQUIRE(aggregates_out == expected_aggregates); +} + +CUB_TEST("cub::DeviceReduce::ReduceByKey accepts not_guaranteed for floating-point multiplication", + "[reduce][env]", + CUB_SMALL) +{ + auto keys_in = thrust::device_vector{0, 2, 2, 9, 5, 5, 5, 8}; + auto values_in = thrust::device_vector{0, 7, 1, 6, 2, 5, 3, 4}; + auto unique_keys_out = thrust::device_vector(5); + auto aggregates_out = thrust::device_vector(5); + auto num_runs_out = thrust::device_vector(1); + + auto env = cuda::execution::require(cuda::execution::determinism::not_guaranteed); + + auto error = cub::DeviceReduce::ReduceByKey( + keys_in.begin(), + unique_keys_out.begin(), + values_in.begin(), + aggregates_out.begin(), + num_runs_out.begin(), + cuda::std::multiplies<>{}, + static_cast(keys_in.size()), + env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceReduce::ReduceByKey failed with status: " << error << '\n'; + } + + thrust::device_vector expected_keys{0, 2, 9, 5, 8}; + thrust::device_vector expected_aggregates{0, 7, 6, 30, 4}; + + REQUIRE(error == cudaSuccess); + REQUIRE(num_runs_out[0] == 5); + unique_keys_out.resize(5); + aggregates_out.resize(5); + REQUIRE(unique_keys_out == expected_keys); + REQUIRE(aggregates_out == expected_aggregates); +} + CUB_TEST("cub::DeviceReduce::ReduceByKey accepts stream", "[reduce][env]", CUB_SMALL) { // example-begin reduce-by-key-env-stream diff --git a/cub/test/test_device_reduce_by_key_determinism_fail.cu b/cub/test/test_device_reduce_by_key_determinism_fail.cu new file mode 100644 index 000000000000..ddc498492b07 --- /dev/null +++ b/cub/test/test_device_reduce_by_key_determinism_fail.cu @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// %PARAM% TEST_ERR err 0:1 + +#include + +#include +#include +#include + +int main() +{ + int* const keys{}; + double* const values{}; + int* const num_runs{}; + +#if TEST_ERR == 0 + // clang-format off + // expected-error-0 {{"run_to_run deterministic reduce-by-key requires integral types with known operators, primitive types with min/max, or floating-point types with plus operator"}} + // clang-format on + const auto error = cub::DeviceReduce::ReduceByKey( + keys, + keys, + values, + values, + num_runs, + cuda::std::multiplies<>{}, + 0, + cuda::execution::require(cuda::execution::determinism::run_to_run)); +#else + // clang-format off + // expected-error-1 {{"gpu_to_gpu deterministic reduce-by-key requires integral types with known operators or primitive types with min/max"}} + // clang-format on + const auto error = cub::DeviceReduce::ReduceByKey( + keys, + keys, + values, + values, + num_runs, + cuda::std::plus<>{}, + 0, + cuda::execution::require(cuda::execution::determinism::gpu_to_gpu)); +#endif + + return error == cudaSuccess ? 0 : 1; +} diff --git a/docs/cub/determinism.rst b/docs/cub/determinism.rst index 301e11653924..312ac32836e4 100644 --- a/docs/cub/determinism.rst +++ b/docs/cub/determinism.rst @@ -41,11 +41,16 @@ Support matrix - ``run_to_run`` - ``gpu_to_gpu`` - Default - * - ``cub::DeviceReduce`` (``Reduce``, ``Sum``, ``Min``, ``Max``, ``TransformReduce``, ...) + * - ``cub::DeviceReduce`` (``Reduce``, ``Sum``, ``Min``, ``Max``, ``TransformReduce``, ...; excluding ``ReduceByKey``) - Yes - Yes - Yes (partial) - ``run_to_run`` + * - ``cub::DeviceReduce::ReduceByKey`` + - Yes + - Yes (partial) + - Yes (partial) + - ``run_to_run`` * - ``cub::DeviceScan`` (``ExclusiveSum``, ``ExclusiveScan``, ``InclusiveSum``, ``InclusiveScan``, ...) - Yes - Yes (partial)