From 493ad08834ad9cc8fd0020ce00c94c5aaef64ce5 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Mon, 10 Nov 2025 12:35:11 +0100 Subject: [PATCH 01/10] [CUB][device_segmented_reduce] add a env-based overload of device_segmented_reduce --- cub/cub/device/device_segmented_reduce.cuh | 210 +++++++++++++++++++++ 1 file changed, 210 insertions(+) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index b82205edec40..44f6091c47e3 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -407,6 +407,216 @@ struct DeviceSegmentedReduce _CCCL_UNREACHABLE(); } + /* + Remi Bourgeois: + env-based overload of the device segmented reduce taking offset iterators. + Heavily inspired from the env-based device reduce. + */ + + //! @rst + //! Computes a device-wide segmented sum using the addition (``+``) operator. + //! + //! - Uses ``0`` as the initial value of the reduction for each segment. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``+`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the `env` parameter + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! TODO (a new snippet) + //! The code snippet below illustrates the sum reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-sum + //! :end-before: example-end segmented-reduce-sum + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is `cuda::std::execution::env<>`. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments`, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>, + typename EnvT = ::cuda::std::execution::env<>> + CUB_RUNTIME_FUNCTION static cudaError_t + Sum(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Sum"); + + using OffsetT = detail::common_iterator_value_t; + using OutputT = detail::non_void_value_t>; + using init_t = OutputT; + + // Remi Bourgeois: copy-pasted from device_reduce.cuh but device segmented reduce has no determinism option ATM, + // left as comment + /* + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using default_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + constexpr auto no_determinism = detail::is_non_deterministic_v; + + // The output iterator must be a contiguous iterator or we fall back to + // run-to-run determinism. + constexpr auto is_contiguous_fallback = + !no_determinism || THRUST_NS_QUALIFIER::is_contiguous_iterator_v; + + using OutputT = cub::detail::non_void_value_t>; + + // Since atomics for types of size < 4B are emulated, they perform poorly, so we fall back to the run-to-run + // determinism. + constexpr auto is_4b_or_greater = !no_determinism || sizeof(OutputT) >= 4; + + using determinism_t = + ::cuda::std::conditional_t; + */ + + // Remi Bourgeois copy-pasted from device_reduce.cuh but device segmented reduce has no tuning option ATM, left as + // comment + /* + using tuning_t = + ::cuda::std::execution::__query_result_or_t>; + */ + + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + auto stream = ::cuda::std::execution::__query_or(env, ::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}); + auto mr = + ::cuda::std::execution::__query_or(env, ::cuda::mr::__get_memory_resource, detail::device_memory_resource{}); + + void* d_temp_storage = nullptr; + size_t temp_storage_bytes = 0; + + // Query the required temporary storage size + cudaError_t error = DispatchSegmentedReduce< + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ::cuda::std::plus<>, + init_t>::Dispatch(d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::std::plus<>{}, + init_t{}, // zero-initialize + stream.get()); + if (error != cudaSuccess) + { + return error; + } + + // TODO(gevtushenko): use uninitialized buffer whenit's available + error = CubDebug(detail::temporary_storage::allocate(stream, d_temp_storage, temp_storage_bytes, mr)); + if (error != cudaSuccess) + { + return error; + } + + // Run the algorithm + error = DispatchSegmentedReduce< + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ::cuda::std::plus<>, + init_t>::Dispatch(d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::std::plus<>{}, + init_t{}, // zero-initialize + stream.get()); + + // Try to deallocate regardless of the error to avoid memory leaks + cudaError_t deallocate_error = + CubDebug(detail::temporary_storage::deallocate(stream, d_temp_storage, temp_storage_bytes, mr)); + + if (error != cudaSuccess) + { + // Reduction error takes precedence over deallocation error since it happens first + return error; + } + return deallocate_error; + } + _CCCL_UNREACHABLE(); + } + //! @rst //! Computes a device-wide segmented sum using the addition (``+``) operator. //! From 41998d83ab1c8a3a5790e0269209106d84a85361 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Mon, 10 Nov 2025 16:12:24 +0100 Subject: [PATCH 02/10] [DO NOT MERGE] sample code for env-based device reduce --- sample_device_segmented_reduce_env.cu | 53 +++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 sample_device_segmented_reduce_env.cu diff --git a/sample_device_segmented_reduce_env.cu b/sample_device_segmented_reduce_env.cu new file mode 100644 index 000000000000..588a9d34a39f --- /dev/null +++ b/sample_device_segmented_reduce_env.cu @@ -0,0 +1,53 @@ +#include + +#include +#include + +#include +#include + +#include +#include + +namespace cudax = cuda::experimental; + +// Compile with nvcc -Icub -Ilibcudacxx/include -Icudax/include -Ithrust/ -Icudax/include/ sample_device_segmented_reduce_env.cu -o sample_device_segmented_reduce_env + +// In cccl/ repo + +int main() +{ + + int num_segments = 3; + + + thrust::device_vector d_offsets = {0, 3, 3, 7}; + int* d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + int* d_in_it = thrust::raw_pointer_cast(d_in.data()); + + thrust::device_vector d_out(num_segments); + int* d_out_it = thrust::raw_pointer_cast(d_out.data()); + + // Build an env (from https://github.com/NVIDIA/cccl/blob/main/cudax/examples/cub_reduce.cu) + // A CUDA stream on which to execute the reduction + cuda::stream stream{cuda::devices[0]}; + cuda::device_memory_pool_ref mr = cuda::device_default_memory_pool(cuda::devices[0]); + + // An environment we use to pass all necessary information to CUB + cudax::env_t env{mr, stream}; + + cub::DeviceSegmentedReduce::Sum(d_in_it, d_out_it, num_segments, d_offsets_it, d_offsets_it + 1, env); + + thrust::host_vector h_out = d_out; + thrust::host_vector expected{21, 0, 17}; + + std::cout << "Segmented reduce results:\n"; + for (int i = 0; i < num_segments; ++i) + { + std::cout << "Segment " << i << ": " << h_out[i] << " (expected " << expected[i] << ")\n"; + } + + return 0; +} From 83f1b72a2c0dd131934d20bb25353187f18474de Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Tue, 18 Nov 2025 15:50:27 +0100 Subject: [PATCH 03/10] static assert determinism --- cub/cub/device/device_segmented_reduce.cuh | 40 +++------------------- sample_device_segmented_reduce_env.cu | 11 ++++-- 2 files changed, 12 insertions(+), 39 deletions(-) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index 44f6091c47e3..7f3fe907072c 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -407,12 +407,6 @@ struct DeviceSegmentedReduce _CCCL_UNREACHABLE(); } - /* - Remi Bourgeois: - env-based overload of the device segmented reduce taking offset iterators. - Heavily inspired from the env-based device reduce. - */ - //! @rst //! Computes a device-wide segmented sum using the addition (``+``) operator. //! @@ -507,42 +501,16 @@ struct DeviceSegmentedReduce using OutputT = detail::non_void_value_t>; using init_t = OutputT; - // Remi Bourgeois: copy-pasted from device_reduce.cuh but device segmented reduce has no determinism option ATM, - // left as comment - /* using requirements_t = ::cuda::std::execution:: __query_result_or_t>; - using default_determinism_t = + using requested_determinism_t = ::cuda::std::execution::__query_result_or_t; - constexpr auto no_determinism = detail::is_non_deterministic_v; - - // The output iterator must be a contiguous iterator or we fall back to - // run-to-run determinism. - constexpr auto is_contiguous_fallback = - !no_determinism || THRUST_NS_QUALIFIER::is_contiguous_iterator_v; - - using OutputT = cub::detail::non_void_value_t>; - - // Since atomics for types of size < 4B are emulated, they perform poorly, so we fall back to the run-to-run - // determinism. - constexpr auto is_4b_or_greater = !no_determinism || sizeof(OutputT) >= 4; - - using determinism_t = - ::cuda::std::conditional_t; - */ - - // Remi Bourgeois copy-pasted from device_reduce.cuh but device segmented reduce has no tuning option ATM, left as - // comment - /* - using tuning_t = - ::cuda::std::execution::__query_result_or_t>; - */ + // Static assert to reject gpu_to_gpu determinism since it's not properly implemented + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); if constexpr (::cuda::std::is_integral_v) diff --git a/sample_device_segmented_reduce_env.cu b/sample_device_segmented_reduce_env.cu index 588a9d34a39f..36df8e4e4ec3 100644 --- a/sample_device_segmented_reduce_env.cu +++ b/sample_device_segmented_reduce_env.cu @@ -11,16 +11,15 @@ namespace cudax = cuda::experimental; -// Compile with nvcc -Icub -Ilibcudacxx/include -Icudax/include -Ithrust/ -Icudax/include/ sample_device_segmented_reduce_env.cu -o sample_device_segmented_reduce_env +// Compile with nvcc -Icub -Ilibcudacxx/include -Icudax/include -Ithrust/ -Icudax/include/ +// sample_device_segmented_reduce_env.cu -o sample_device_segmented_reduce_env // In cccl/ repo int main() { - int num_segments = 3; - thrust::device_vector d_offsets = {0, 3, 3, 7}; int* d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); @@ -38,7 +37,13 @@ int main() // An environment we use to pass all necessary information to CUB cudax::env_t env{mr, stream}; + const auto envdet = cuda::execution::require(cuda::execution::determinism::run_to_run); + // const auto envdet = cuda::execution::require(cuda::execution::determinism::not_guaranteed); + // const auto envdet = cuda::execution::require(cuda::execution::determinism::gpu_to_gpu); + cub::DeviceSegmentedReduce::Sum(d_in_it, d_out_it, num_segments, d_offsets_it, d_offsets_it + 1, env); + //test specified determinism + cub::DeviceSegmentedReduce::Sum(d_in_it, d_out_it, num_segments, d_offsets_it, d_offsets_it + 1, envdet); thrust::host_vector h_out = d_out; thrust::host_vector expected{21, 0, 17}; From e7db3397770f7fde33ba7f511552618381657fce Mon Sep 17 00:00:00 2001 From: moletteremi <98650511+rbourgeois33@users.noreply.github.com> Date: Fri, 21 Nov 2025 20:10:54 +0100 Subject: [PATCH 04/10] Update cub/cub/device/device_segmented_reduce.cuh Co-authored-by: Srinivas Yadav <43375352+srinivasyadav18@users.noreply.github.com> --- cub/cub/device/device_segmented_reduce.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index 7f3fe907072c..68707324e5ec 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -432,8 +432,8 @@ struct DeviceSegmentedReduce //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu //! :language: c++ //! :dedent: - //! :start-after: example-begin segmented-reduce-sum - //! :end-before: example-end segmented-reduce-sum + //! :start-after: example-begin segmented-reduce-sum-env + //! :end-before: example-end segmented-reduce-sum-env //! //! @endrst //! From 2094c015eedc7b3656668ca73f0771b94dda43b1 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Fri, 21 Nov 2025 21:14:21 +0100 Subject: [PATCH 05/10] Add tests for the env-based segmented sum --- cub/cub/device/device_segmented_reduce.cuh | 4 +- ...h2_test_device_segmented_reduce_env_api.cu | 77 +++++++++++++++++++ sample_device_segmented_reduce_env.cu | 58 -------------- 3 files changed, 79 insertions(+), 60 deletions(-) create mode 100644 cub/test/catch2_test_device_segmented_reduce_env_api.cu delete mode 100644 sample_device_segmented_reduce_env.cu diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index 68707324e5ec..fc5afd6e5e29 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -432,8 +432,8 @@ struct DeviceSegmentedReduce //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu //! :language: c++ //! :dedent: - //! :start-after: example-begin segmented-reduce-sum-env - //! :end-before: example-end segmented-reduce-sum-env + //! :start-after: example-begin segmented-reduce-sum-env-stream + //! :end-before: example-end segmented-reduce-sum-env-stream //! //! @endrst //! diff --git a/cub/test/catch2_test_device_segmented_reduce_env_api.cu b/cub/test/catch2_test_device_segmented_reduce_env_api.cu new file mode 100644 index 000000000000..05296a48b2d9 --- /dev/null +++ b/cub/test/catch2_test_device_segmented_reduce_env_api.cu @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#include "insert_nested_NVTX_range_guard.h" +// above header needs to be included first + +#include + +#include + +#include +#include + +#include + +#include + +C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts stream", "[segmented_reduce][env]") +{ + // example-begin segmented-reduce-sum-env-stream + int num_segments = 3; + c2h::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + c2h::device_vector d_out(3); + + cudaStream_t legacy_stream = 0; + cuda::stream_ref stream_ref{legacy_stream}; + + auto error = cub::DeviceSegmentedReduce::Sum( + d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, stream_ref); + c2h::device_vector expected{21, 0, 17}; + + // example-end segmented-reduce-reduce-sum-env-stream + REQUIRE(d_out == expected); + REQUIRE(error == cudaSuccess); +} + +C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism requirements", "[segmented_reduce][env]") +{ + // example-begin segmented-reduce-sum-env-determinism + int num_segments = 3; + c2h::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + c2h::device_vector d_out(3); + + auto env = cuda::execution::require(cuda::execution::determinism::run_to_run); + + auto error = + cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); + c2h::device_vector expected{21, 0, 17}; + + // example-end segmented-reduce-reduce-sum-env-determinism + REQUIRE(d_out == expected); + REQUIRE(error == cudaSuccess); +} + +C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism requirements", "[segmented_reduce][env]") +{ + // example-begin segmented-reduce-sum-env-non-determinism + int num_segments = 3; + c2h::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + c2h::device_vector d_out(3); + + auto env = cuda::execution::require(cuda::execution::determinism::not_guaranteed); + + auto error = + cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); + c2h::device_vector expected{21, 0, 17}; + + // example-end segmented-reduce-reduce-sum-env-non-determinism + REQUIRE(d_out == expected); + REQUIRE(error == cudaSuccess); +} diff --git a/sample_device_segmented_reduce_env.cu b/sample_device_segmented_reduce_env.cu deleted file mode 100644 index 36df8e4e4ec3..000000000000 --- a/sample_device_segmented_reduce_env.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include - -#include -#include - -#include -#include - -#include -#include - -namespace cudax = cuda::experimental; - -// Compile with nvcc -Icub -Ilibcudacxx/include -Icudax/include -Ithrust/ -Icudax/include/ -// sample_device_segmented_reduce_env.cu -o sample_device_segmented_reduce_env - -// In cccl/ repo - -int main() -{ - int num_segments = 3; - - thrust::device_vector d_offsets = {0, 3, 3, 7}; - int* d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); - - thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; - int* d_in_it = thrust::raw_pointer_cast(d_in.data()); - - thrust::device_vector d_out(num_segments); - int* d_out_it = thrust::raw_pointer_cast(d_out.data()); - - // Build an env (from https://github.com/NVIDIA/cccl/blob/main/cudax/examples/cub_reduce.cu) - // A CUDA stream on which to execute the reduction - cuda::stream stream{cuda::devices[0]}; - cuda::device_memory_pool_ref mr = cuda::device_default_memory_pool(cuda::devices[0]); - - // An environment we use to pass all necessary information to CUB - cudax::env_t env{mr, stream}; - - const auto envdet = cuda::execution::require(cuda::execution::determinism::run_to_run); - // const auto envdet = cuda::execution::require(cuda::execution::determinism::not_guaranteed); - // const auto envdet = cuda::execution::require(cuda::execution::determinism::gpu_to_gpu); - - cub::DeviceSegmentedReduce::Sum(d_in_it, d_out_it, num_segments, d_offsets_it, d_offsets_it + 1, env); - //test specified determinism - cub::DeviceSegmentedReduce::Sum(d_in_it, d_out_it, num_segments, d_offsets_it, d_offsets_it + 1, envdet); - - thrust::host_vector h_out = d_out; - thrust::host_vector expected{21, 0, 17}; - - std::cout << "Segmented reduce results:\n"; - for (int i = 0; i < num_segments; ++i) - { - std::cout << "Segment " << i << ": " << h_out[i] << " (expected " << expected[i] << ")\n"; - } - - return 0; -} From 4aad8339d26d735533f1ba8a7c59f29307ddfa53 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Tue, 25 Nov 2025 13:56:08 +0100 Subject: [PATCH 06/10] Adress suggestions --- cub/cub/device/device_segmented_reduce.cuh | 14 +++- ...h2_test_device_segmented_reduce_env_api.cu | 64 +++++++++++++------ 2 files changed, 55 insertions(+), 23 deletions(-) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index fc5afd6e5e29..6a8791f49211 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -19,6 +19,8 @@ #endif // no system header #include +#include +#include #include #include #include @@ -26,12 +28,20 @@ #include #include +#include +#include #include #include +#include +#include +#include +#include #include #include +#include #include #include +#include #include #include #include @@ -432,8 +442,8 @@ struct DeviceSegmentedReduce //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu //! :language: c++ //! :dedent: - //! :start-after: example-begin segmented-reduce-sum-env-stream - //! :end-before: example-end segmented-reduce-sum-env-stream + //! :start-after: example-begin segmented-reduce-sum-env + //! :end-before: example-end segmented-reduce-sum-env //! //! @endrst //! diff --git a/cub/test/catch2_test_device_segmented_reduce_env_api.cu b/cub/test/catch2_test_device_segmented_reduce_env_api.cu index 05296a48b2d9..a87b6e5705a3 100644 --- a/cub/test/catch2_test_device_segmented_reduce_env_api.cu +++ b/cub/test/catch2_test_device_segmented_reduce_env_api.cu @@ -11,45 +11,67 @@ #include #include -#include - #include +C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts env with stream and determinism requirements", + "[segmented_reduce][env]") +{ + // example-begin segmented-reduce-sum-env + int num_segments = 3; + thrust::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + thrust::device_vector d_out(3); + + auto req_env = cuda::execution::require(cuda::execution::determinism::not_guaranteed); + cudaStream_t legacy_stream = 0; + cuda::stream_ref stream_ref{legacy_stream}; + auto env = ::cuda::std::execution::env{req_env, stream_ref}; + + auto error = + cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); + thrust::device_vector expected{21, 0, 17}; + + // example-end segmented-reduce-reduce-sum-env + REQUIRE(d_out == expected); + REQUIRE(error == cudaSuccess); +} + C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts stream", "[segmented_reduce][env]") { // example-begin segmented-reduce-sum-env-stream - int num_segments = 3; - c2h::device_vector d_offsets = {0, 3, 3, 7}; - auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); - c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; - c2h::device_vector d_out(3); + int num_segments = 3; + thrust::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + thrust::device_vector d_out(3); cudaStream_t legacy_stream = 0; cuda::stream_ref stream_ref{legacy_stream}; auto error = cub::DeviceSegmentedReduce::Sum( d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, stream_ref); - c2h::device_vector expected{21, 0, 17}; + thrust::device_vector expected{21, 0, 17}; // example-end segmented-reduce-reduce-sum-env-stream REQUIRE(d_out == expected); REQUIRE(error == cudaSuccess); } -C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism requirements", "[segmented_reduce][env]") +C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts run_to_run determinism requirements", "[segmented_reduce][env]") { // example-begin segmented-reduce-sum-env-determinism - int num_segments = 3; - c2h::device_vector d_offsets = {0, 3, 3, 7}; - auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); - c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; - c2h::device_vector d_out(3); + int num_segments = 3; + thrust::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + thrust::device_vector d_out(3); auto env = cuda::execution::require(cuda::execution::determinism::run_to_run); auto error = cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); - c2h::device_vector expected{21, 0, 17}; + thrust::device_vector expected{21, 0, 17}; // example-end segmented-reduce-reduce-sum-env-determinism REQUIRE(d_out == expected); @@ -59,17 +81,17 @@ C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism req C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism requirements", "[segmented_reduce][env]") { // example-begin segmented-reduce-sum-env-non-determinism - int num_segments = 3; - c2h::device_vector d_offsets = {0, 3, 3, 7}; - auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); - c2h::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; - c2h::device_vector d_out(3); + int num_segments = 3; + thrust::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + thrust::device_vector d_out(3); auto env = cuda::execution::require(cuda::execution::determinism::not_guaranteed); auto error = cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); - c2h::device_vector expected{21, 0, 17}; + thrust::device_vector expected{21, 0, 17}; // example-end segmented-reduce-reduce-sum-env-non-determinism REQUIRE(d_out == expected); From 0d70afca6b76e31d6f4f653b71b3229ab22103a4 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Wed, 26 Nov 2025 09:29:43 +0100 Subject: [PATCH 07/10] Applied pre-commit --- cub/cub/device/device_segmented_reduce.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index 6a8791f49211..a73dfaf47a93 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -436,7 +436,7 @@ struct DeviceSegmentedReduce //! //! Snippet //! +++++++++++++++++++++++++++++++++++++++++++++ - //! TODO (a new snippet) + //! //! The code snippet below illustrates the sum reduction of a device vector of ``int`` data elements. //! //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu @@ -518,7 +518,7 @@ struct DeviceSegmentedReduce ::cuda::execution::determinism::__get_determinism_t, ::cuda::execution::determinism::run_to_run_t>; - // Static assert to reject gpu_to_gpu determinism since it's not properly implemented + // Static assert to reject gpu_to_gpu determinism since it's not properly implemented atm static_assert(!::cuda::std::is_same_v, "gpu_to_gpu determinism is not supported for device segmented reductions "); From 05194d7eb3d8d243631c0bef007187ddd9cbf35e Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Wed, 26 Nov 2025 11:45:38 +0100 Subject: [PATCH 08/10] 1st attempt at extracting tuning from env --- cub/cub/device/device_segmented_reduce.cuh | 98 ++++++++++++++-------- 1 file changed, 64 insertions(+), 34 deletions(-) diff --git a/cub/cub/device/device_segmented_reduce.cuh b/cub/cub/device/device_segmented_reduce.cuh index 37fe67e73757..72f34e45ab06 100644 --- a/cub/cub/device/device_segmented_reduce.cuh +++ b/cub/cub/device/device_segmented_reduce.cuh @@ -50,6 +50,30 @@ CUB_NAMESPACE_BEGIN +namespace detail +{ +namespace segmented_reduce +{ +struct get_tuning_query_t +{}; + +template +struct tuning +{ + [[nodiscard]] _CCCL_NODEBUG_API constexpr auto query(const get_tuning_query_t&) const noexcept -> Derived + { + return static_cast(*this); + } +}; + +struct default_tuning : tuning +{ + template + using fn = detail::reduce::policy_hub; +}; +} // namespace segmented_reduce +} // namespace detail + //! @rst //! DeviceSegmentedReduce provides device-wide, parallel operations for //! computing a reduction across multiple sequences of data items @@ -511,14 +535,32 @@ struct DeviceSegmentedReduce using OffsetT = detail::common_iterator_value_t; using OutputT = detail::non_void_value_t>; using init_t = OutputT; + using AccumT = ::cuda::std::__accumulator_t<::cuda::std::plus<>, cub::detail::it_value_t, init_t>; + + using segmented_reduce_tuning_t = ::cuda::std::execution:: + __query_result_or_t; + + using policy_t = typename segmented_reduce_tuning_t::template fn>; using requirements_t = ::cuda::std::execution:: __query_result_or_t>; + using requested_determinism_t = ::cuda::std::execution::__query_result_or_t; + using dispatch_t = DispatchSegmentedReduce< + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ::cuda::std::plus<>, + init_t, + AccumT, + policy_t>; + // Static assert to reject gpu_to_gpu determinism since it's not properly implemented atm static_assert(!::cuda::std::is_same_v, "gpu_to_gpu determinism is not supported for device segmented reductions "); @@ -534,23 +576,17 @@ struct DeviceSegmentedReduce size_t temp_storage_bytes = 0; // Query the required temporary storage size - cudaError_t error = DispatchSegmentedReduce< - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT, - ::cuda::std::plus<>, - init_t>::Dispatch(d_temp_storage, - temp_storage_bytes, - d_in, - d_out, - num_segments, - d_begin_offsets, - d_end_offsets, - ::cuda::std::plus<>{}, - init_t{}, // zero-initialize - stream.get()); + cudaError_t error = dispatch_t::Dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::std::plus<>{}, + init_t{}, // zero-initialize + stream.get()); if (error != cudaSuccess) { return error; @@ -564,23 +600,17 @@ struct DeviceSegmentedReduce } // Run the algorithm - error = DispatchSegmentedReduce< - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT, - ::cuda::std::plus<>, - init_t>::Dispatch(d_temp_storage, - temp_storage_bytes, - d_in, - d_out, - num_segments, - d_begin_offsets, - d_end_offsets, - ::cuda::std::plus<>{}, - init_t{}, // zero-initialize - stream.get()); + error = dispatch_t::Dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::std::plus<>{}, + init_t{}, // zero-initialize + stream.get()); // Try to deallocate regardless of the error to avoid memory leaks cudaError_t deallocate_error = From d352b3300475b8ac36cfcf349e9ed3483329f589 Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Sun, 30 Nov 2025 20:36:14 +0100 Subject: [PATCH 09/10] add unit test for tuning --- ...catch2_test_device_segmented_reduce_env.cu | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 cub/test/catch2_test_device_segmented_reduce_env.cu diff --git a/cub/test/catch2_test_device_segmented_reduce_env.cu b/cub/test/catch2_test_device_segmented_reduce_env.cu new file mode 100644 index 000000000000..2f9f7ee30f33 --- /dev/null +++ b/cub/test/catch2_test_device_segmented_reduce_env.cu @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#include "insert_nested_NVTX_range_guard.h" + +#include + +#include + +#include + +template +struct reduce_tuning : cub::detail::reduce::tuning> +{ + template + struct fn + { + struct Policy500 : cub::ChainedPolicy<500, Policy500, Policy500> + { + struct ReducePolicy + { + static constexpr int VECTOR_LOAD_LENGTH = 1; + + static constexpr cub::BlockReduceAlgorithm BLOCK_ALGORITHM = cub::BLOCK_REDUCE_WARP_REDUCTIONS; + + static constexpr cub::CacheLoadModifier LOAD_MODIFIER = cub::LOAD_DEFAULT; + + static constexpr int ITEMS_PER_THREAD = 1; + static constexpr int BLOCK_THREADS = BlockThreads; + }; + + using SingleTilePolicy = ReducePolicy; + using SegmentedReducePolicy = ReducePolicy; + }; + + using MaxPolicy = Policy500; + }; +}; + +struct get_scan_tuning_query_t +{}; + +struct scan_tuning +{ + [[nodiscard]] _CCCL_NODEBUG_API constexpr auto query(const get_scan_tuning_query_t&) const noexcept + { + return *this; + } + + // Make sure this is not used + template + struct fn + {}; +}; + +using block_sizes = c2h::type_list, cuda::std::integral_constant>; + +C2H_TEST("Device segmented sum can be tuned", "[reduce][device]", block_sizes) +{ + constexpr int target_block_size = c2h::get<0, TestType>::value; + + int num_segments = 3; + thrust::device_vector d_offsets = {0, 3, 3, 7}; + auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data()); + thrust::device_vector d_in{8, 6, 7, 5, 3, 0, 9}; + thrust::device_vector d_out(3); + + // We are expecting that `scan_tuning` is ignored + auto env = cuda::execution::__tune(reduce_tuning{}, scan_tuning{}); + + auto error = + cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); + thrust::device_vector expected{21, 0, 17}; + + REQUIRE(d_out == expected); + REQUIRE(error == cudaSuccess); +} From a20271bf9074130dd3695d04db67d19c41913efc Mon Sep 17 00:00:00 2001 From: Remi Bourgeois Date: Tue, 2 Dec 2025 19:42:27 +0100 Subject: [PATCH 10/10] Add proper error handling to the new tests --- ...h2_test_device_segmented_reduce_env_api.cu | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/cub/test/catch2_test_device_segmented_reduce_env_api.cu b/cub/test/catch2_test_device_segmented_reduce_env_api.cu index a87b6e5705a3..38d92aa6d659 100644 --- a/cub/test/catch2_test_device_segmented_reduce_env_api.cu +++ b/cub/test/catch2_test_device_segmented_reduce_env_api.cu @@ -32,7 +32,12 @@ C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts env with stream and determinis cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); thrust::device_vector expected{21, 0, 17}; + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceSegmentedReduce::Sum failed with status: " << error << std::endl; + } // example-end segmented-reduce-reduce-sum-env + REQUIRE(d_out == expected); REQUIRE(error == cudaSuccess); } @@ -53,7 +58,12 @@ C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts stream", "[segmented_reduce][e d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, stream_ref); thrust::device_vector expected{21, 0, 17}; + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceSegmentedReduce::Sum failed with status: " << error << std::endl; + } // example-end segmented-reduce-reduce-sum-env-stream + REQUIRE(d_out == expected); REQUIRE(error == cudaSuccess); } @@ -73,7 +83,12 @@ C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts run_to_run determinism require cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); thrust::device_vector expected{21, 0, 17}; + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceSegmentedReduce::Sum failed with status: " << error << std::endl; + } // example-end segmented-reduce-reduce-sum-env-determinism + REQUIRE(d_out == expected); REQUIRE(error == cudaSuccess); } @@ -93,7 +108,12 @@ C2H_TEST("cub::DeviceSegmentedReduce::Sum accepts not_guaranteed determinism req cub::DeviceSegmentedReduce::Sum(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env); thrust::device_vector expected{21, 0, 17}; + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceSegmentedReduce::Sum failed with status: " << error << std::endl; + } // example-end segmented-reduce-reduce-sum-env-non-determinism + REQUIRE(d_out == expected); REQUIRE(error == cudaSuccess); }