From cb1a5fc5d1a712fabaab14803b91f6c95cf7bf0b Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Mon, 9 Mar 2026 22:24:18 +0100 Subject: [PATCH 01/17] Use the new tuning API for DeviceRadixSort Fixes: #7947 --- cub/benchmarks/bench/radix_sort/keys.cu | 79 ++----- cub/benchmarks/bench/radix_sort/pairs.cu | 88 +++----- .../bench/radix_sort/policy_selector.h | 35 ++- cub/cub/detail/temporary_storage.cuh | 2 + cub/cub/device/device_radix_sort.cuh | 210 ++++++++++++++---- .../dispatch_segmented_radix_sort.cuh | 22 +- cub/test/catch2_test_device_radix_sort_env.cu | 159 +++++++++++++ 7 files changed, 410 insertions(+), 185 deletions(-) diff --git a/cub/benchmarks/bench/radix_sort/keys.cu b/cub/benchmarks/bench/radix_sort/keys.cu index c05a72190382..7a5b78b72e16 100644 --- a/cub/benchmarks/bench/radix_sort/keys.cu +++ b/cub/benchmarks/bench/radix_sort/keys.cu @@ -14,82 +14,47 @@ template void radix_sort_keys(nvbench::state& state, nvbench::type_list) { - using offset_t = cub::detail::choose_offset_t; - - constexpr cub::SortOrder sort_order = cub::SortOrder::Ascending; - constexpr bool is_overwrite_ok = false; - using key_t = T; - using value_t = cub::NullType; - - if constexpr (!fits_in_default_shared_memory()) + using key_t = T; + using value_t = cub::NullType; + if constexpr (!fits_in_default_shared_memory()) { return; } - constexpr int begin_bit = 0; - constexpr int end_bit = sizeof(key_t) * 8; - // Retrieve axis parameters const auto elements = static_cast(state.get_int64("Elements{io}")); const bit_entropy entropy = str_to_entropy(state.get_string("Entropy")); thrust::device_vector buffer_1 = generate(elements, entropy); - thrust::device_vector buffer_2(elements); + thrust::device_vector buffer_2(elements, thrust::no_init); - key_t* d_buffer_1 = thrust::raw_pointer_cast(buffer_1.data()); - key_t* d_buffer_2 = thrust::raw_pointer_cast(buffer_2.data()); - - cub::DoubleBuffer d_keys(d_buffer_1, d_buffer_2); - cub::DoubleBuffer d_values; + const key_t* d_buffer_1 = thrust::raw_pointer_cast(buffer_1.data()); + key_t* d_buffer_2 = thrust::raw_pointer_cast(buffer_2.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(elements); - // Allocate temporary storage: - std::size_t temp_size{}; - - cub::detail::radix_sort::dispatch( - nullptr, - temp_size, - d_keys, - d_values, - static_cast(elements), - begin_bit, - end_bit, - is_overwrite_ok, - nullptr /* stream */ -#if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} -#endif // !TUNE_BASE - ); - - thrust::device_vector temp(temp_size, thrust::no_init); - auto* temp_storage = thrust::raw_pointer_cast(temp.data()); - + caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - cub::DoubleBuffer keys = d_keys; - cub::DoubleBuffer values = d_values; - - cub::detail::radix_sort::dispatch( - temp_storage, - temp_size, - keys, - values, - static_cast(elements), - begin_bit, - end_bit, - is_overwrite_ok, - launch.get_stream() + if (cub::DeviceRadixSort::SortKeys( + d_buffer_1, + d_buffer_2, + static_cast(elements), + 0, + sizeof(key_t) * 8, + cub_bench_env(alloc, + launch #if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} + , + cuda::execution::__tune(policy_selector{}) #endif // !TUNE_BASE - ); + )) + != cudaSuccess) + { + throw 42; // TODO(bgruber): what's the appropriate way to abort a benchmark? + } }); } diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index f293fffa7649..519196d701ee 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -16,37 +16,26 @@ template void radix_sort_values(nvbench::state& state, nvbench::type_list) { - using offset_t = cub::detail::choose_offset_t; - - constexpr cub::SortOrder sort_order = cub::SortOrder::Ascending; - constexpr bool is_overwrite_ok = false; - using key_t = KeyT; - using value_t = ValueT; - - if constexpr (!fits_in_default_shared_memory()) + using key_t = KeyT; + using value_t = ValueT; + if constexpr (!fits_in_default_shared_memory()) { return; } - constexpr int begin_bit = 0; - constexpr int end_bit = sizeof(key_t) * 8; - // Retrieve axis parameters const auto elements = static_cast(state.get_int64("Elements{io}")); const bit_entropy entropy = str_to_entropy(state.get_string("Entropy")); - thrust::device_vector keys_buffer_1 = generate(elements, entropy); - thrust::device_vector values_buffer_1 = generate(elements); - thrust::device_vector keys_buffer_2(elements); - thrust::device_vector values_buffer_2(elements); + thrust::device_vector keys_in = generate(elements, entropy); + thrust::device_vector keys_out(elements, thrust::no_init); + thrust::device_vector values_in = generate(elements); + thrust::device_vector values_out(elements, thrust::no_init); - key_t* d_keys_buffer_1 = thrust::raw_pointer_cast(keys_buffer_1.data()); - key_t* d_keys_buffer_2 = thrust::raw_pointer_cast(keys_buffer_2.data()); - value_t* d_values_buffer_1 = thrust::raw_pointer_cast(values_buffer_1.data()); - value_t* d_values_buffer_2 = thrust::raw_pointer_cast(values_buffer_2.data()); - - cub::DoubleBuffer d_keys(d_keys_buffer_1, d_keys_buffer_2); - cub::DoubleBuffer d_values(d_values_buffer_1, d_values_buffer_2); + const key_t* d_keys_in = thrust::raw_pointer_cast(keys_in.data()); + key_t* d_keys_out = thrust::raw_pointer_cast(keys_out.data()); + const value_t* d_values_in = thrust::raw_pointer_cast(values_in.data()); + value_t* d_values_out = thrust::raw_pointer_cast(values_out.data()); // Enable throughput calculations and add "Size" column to results. state.add_element_count(elements); @@ -55,48 +44,25 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list(elements); state.add_global_memory_writes(elements); - // Allocate temporary storage: - std::size_t temp_size{}; - cub::detail::radix_sort::dispatch( - nullptr, - temp_size, - d_keys, - d_values, - static_cast(elements), - begin_bit, - end_bit, - is_overwrite_ok, - nullptr /* stream */ -#if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} -#endif // !TUNE_BASE - ); - - thrust::device_vector temp(temp_size, thrust::no_init); - auto* temp_storage = thrust::raw_pointer_cast(temp.data()); - + caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - cub::DoubleBuffer keys = d_keys; - cub::DoubleBuffer values = d_values; - - cub::detail::radix_sort::dispatch( - temp_storage, - temp_size, - keys, - values, - static_cast(elements), - begin_bit, - end_bit, - is_overwrite_ok, - launch.get_stream() + if (cub::DeviceRadixSort::SortPairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + static_cast(elements), + cub_bench_env(alloc, + launch #if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} + , + cuda::execution::__tune(policy_selector{}) #endif // !TUNE_BASE - ); + )) + != cudaSuccess) + { + throw 42; // TODO(bgruber): what's the appropriate way to abort a benchmark? + } }); } diff --git a/cub/benchmarks/bench/radix_sort/policy_selector.h b/cub/benchmarks/bench/radix_sort/policy_selector.h index 051f29a900af..896311d10c59 100644 --- a/cub/benchmarks/bench/radix_sort/policy_selector.h +++ b/cub/benchmarks/bench/radix_sort/policy_selector.h @@ -77,13 +77,33 @@ struct policy_selector template constexpr std::size_t max_onesweep_temp_storage_size() { - using portion_offset = int; - using onesweep_policy = typename policy_hub_t::policy_t::OnesweepPolicy; + using portion_offset = int; + + constexpr auto active_policy = policy_selector{}(cuda::arch_id{}); + + constexpr auto onesweep = active_policy.onesweep; + using onesweep_policy_t = AgentRadixSortOnesweepPolicy< + 0, + 0, + void, + onesweep.rank_num_parts, + onesweep.rank_algorith, + onesweep.scan_algorithm, + onesweep.store_algorithm, + onesweep.radix_bits, + NoScaling>; + using agent_radix_sort_onesweep_t = - cub::AgentRadixSortOnesweep; + cub::AgentRadixSortOnesweep; - using hist_policy = typename policy_hub_t::policy_t::HistogramPolicy; - using hist_agent = cub::AgentRadixSortHistogram; + constexpr auto histogram = active_policy.histogram; + using histogram_policy_t = + AgentRadixSortHistogramPolicy; + using hist_agent = cub::AgentRadixSortHistogram; return cuda::std::max(sizeof(typename agent_radix_sort_onesweep_t::TempStorage), sizeof(typename hist_agent::TempStorage)); @@ -92,10 +112,11 @@ constexpr std::size_t max_onesweep_temp_storage_size() template constexpr std::size_t max_temp_storage_size() { - using policy_t = typename policy_hub_t::policy_t; + using offset_t = cub::detail::choose_offset_t; + using policy_t = typename policy_hub_t::policy_t; static_assert(policy_t::ONESWEEP); - return max_onesweep_temp_storage_size(); + return max_onesweep_temp_storage_size(); } template diff --git a/cub/cub/detail/temporary_storage.cuh b/cub/cub/detail/temporary_storage.cuh index 9863995fc8be..e6be1ecadf85 100644 --- a/cub/cub/detail/temporary_storage.cuh +++ b/cub/cub/detail/temporary_storage.cuh @@ -347,6 +347,7 @@ private: } }; +_CCCL_EXEC_CHECK_DISABLE template CUB_RUNTIME_FUNCTION cudaError_t allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_bytes, MRT& mr) @@ -362,6 +363,7 @@ allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_b } } +_CCCL_EXEC_CHECK_DISABLE template CUB_RUNTIME_FUNCTION cudaError_t deallocate(::cuda::stream_ref stream, void* d_temp_storage, size_t temp_storage_bytes, MRT& mr) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index d7f038daa48b..63a0b7882bc0 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -450,8 +450,8 @@ public: template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -466,14 +466,28 @@ public: using offset_t = detail::choose_offset_t; - // Dispatch with environment - handles all boilerplate DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, false, stream); - }); + // Dispatch with environment - handles all boilerplate + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -1140,8 +1154,8 @@ public: template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -1154,10 +1168,24 @@ public: using offset_t = detail::choose_offset_t; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, true, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -1798,8 +1826,8 @@ public: template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -1817,10 +1845,24 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, false, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -2297,8 +2339,8 @@ public: template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -2311,10 +2353,24 @@ public: using offset_t = detail::choose_offset_t; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, true, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -3032,8 +3088,8 @@ public: //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -3050,10 +3106,24 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, false, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -3635,8 +3705,8 @@ public: //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( DoubleBuffer& d_keys, NumItemsT num_items, int begin_bit = 0, int end_bit = sizeof(KeyT) * 8, EnvT env = {}) { @@ -3646,10 +3716,24 @@ public: DoubleBuffer d_values; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, true, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -4189,8 +4273,8 @@ public: //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -4207,10 +4291,24 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, false, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + policy_selector_t{}); + }); } //! @rst @@ -4645,8 +4743,8 @@ public: //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. template , - typename ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( DoubleBuffer& d_keys, NumItemsT num_items, int begin_bit = 0, int end_bit = sizeof(KeyT) * 8, EnvT env = {}) { @@ -4656,10 +4754,24 @@ public: DoubleBuffer d_values; - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { - return detail::radix_sort::dispatch( - storage, bytes, d_keys, d_values, static_cast(num_items), begin_bit, end_bit, true, stream); - }); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + policy_selector_t{}); + }); } //! @rst diff --git a/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh index 23bb954acb52..f39f5c7d4d76 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh @@ -883,17 +883,17 @@ template , - typename KernelSource = DeviceSegmentedRadixSortKernelSource< - PolicySelector, - Order, - KeyT, - ValueT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - SegmentSizeT, - DecomposerT>, + typename DecomposerT = identity_decomposer_t, + typename PolicySelector, + typename KernelSource = DeviceSegmentedRadixSortKernelSource< + PolicySelector, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_radix_sort_policy_selector diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index fc8605aea07f..4de139a928c0 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1106,3 +1106,162 @@ TEST_CASE("Device radix sort pairs descending DB decomposer+bits uses custom str REQUIRE(keys == expected_keys); REQUIRE(values == expected_values); } + +// using different block sizes yields to different temporary storage sizes, so use a custom policy to influence that +template +struct tiny_onesweep_policy_selector +{ + _CCCL_API constexpr auto operator()(cuda::arch_id arch) const -> cub::detail::radix_sort::radix_sort_policy + { + using default_selector_t = cub::detail::radix_sort::policy_selector_from_types; + auto policy = default_selector_t{}(arch); + policy.use_onesweep = true; + policy.onesweep.block_threads = BlockThreads; + policy.onesweep.items_per_thread = 1; + return policy; + } +}; + +template +std::size_t measure_allocated_bytes(CallableT&& run, PolicySelector policy_selector) +{ + cuda::stream_ref stream{cudaStream_t{}}; + size_t bytes_allocated = 0; + size_t bytes_deallocated = 0; + auto env = stdexec::env{ + cuda::std::execution::prop{cuda::mr::__get_memory_resource_t{}, + device_memory_resource{{}, stream.get(), &bytes_allocated, &bytes_deallocated}}, + stream, + cuda::execution::__tune(policy_selector)}; + REQUIRE(cudaSuccess == run(env)); + stream.sync(); + CHECK(bytes_allocated > 0); + CHECK(bytes_allocated == bytes_deallocated); + return bytes_allocated; +} + +TEST_CASE("DeviceRadixSort::SortPairs can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + return cub::DeviceRadixSort::SortPairs( + data.data().get(), + data.data().get(), + data.data().get(), + data.data().get(), + static_cast(data.size()), + 0, + 32, + env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortPairs DoubleBuffer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + cub::DoubleBuffer double_buf(data.data().get(), data.data().get()); + return cub::DeviceRadixSort::SortPairs(double_buf, double_buf, static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortPairsDescending can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + return cub::DeviceRadixSort::SortPairsDescending( + data.data().get(), + data.data().get(), + data.data().get(), + data.data().get(), + static_cast(data.size()), + 0, + 32, + env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortPairsDescending DoubleBuffer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + cub::DoubleBuffer double_buf(data.data().get(), data.data().get()); + return cub::DeviceRadixSort::SortPairsDescending(double_buf, double_buf, static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortKeys can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + return cub::DeviceRadixSort::SortKeys( + data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortKeys DoubleBuffer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + cub::DoubleBuffer double_buf(data.data().get(), data.data().get()); + return cub::DeviceRadixSort::SortKeys(double_buf, static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortKeysDescending can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + return cub::DeviceRadixSort::SortKeysDescending( + data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} + +TEST_CASE("DeviceRadixSort::SortKeysDescending DoubleBuffer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); // must be larger than the single tile path + cub::DoubleBuffer double_buf(data.data().get(), data.data().get()); + return cub::DeviceRadixSort::SortKeysDescending(double_buf, static_cast(data.size()), 0, 32, env); + }; + + auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + + CHECK(default_bytes != tuned_bytes); +} From ca6a1850e79003166e5c36ef8a931c8a4aaea71e Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 13:22:13 +0200 Subject: [PATCH 02/17] fixes --- cub/benchmarks/bench/radix_sort/keys.cu | 2 +- cub/benchmarks/bench/radix_sort/pairs.cu | 2 +- cub/test/catch2_test_device_radix_sort_env.cu | 67 +++++++++---------- 3 files changed, 35 insertions(+), 36 deletions(-) diff --git a/cub/benchmarks/bench/radix_sort/keys.cu b/cub/benchmarks/bench/radix_sort/keys.cu index 7a5b78b72e16..7b654fc3ab60 100644 --- a/cub/benchmarks/bench/radix_sort/keys.cu +++ b/cub/benchmarks/bench/radix_sort/keys.cu @@ -48,7 +48,7 @@ void radix_sort_keys(nvbench::state& state, nvbench::type_list) launch #if !TUNE_BASE , - cuda::execution::__tune(policy_selector{}) + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE )) != cudaSuccess) diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index 519196d701ee..49eeb7aeff42 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -56,7 +56,7 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list{}) + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE )) != cudaSuccess) diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index 4de139a928c0..4630b7805625 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1107,7 +1107,12 @@ TEST_CASE("Device radix sort pairs descending DB decomposer+bits uses custom str REQUIRE(values == expected_values); } -// using different block sizes yields to different temporary storage sizes, so use a custom policy to influence that +#if TEST_LAUNCH == 0 + +// Radix sort does not accept user-provided functors or iterators, so we cannot use the block_size_extracting_op or +// block_size_extracting_constant_iterator approach. Instead, we verify tuning by measuring allocation sizes: different +// block sizes produce different temporary storage requirements, which is an observable side effect of the tuning being +// applied. template struct tiny_onesweep_policy_selector { @@ -1132,7 +1137,7 @@ std::size_t measure_allocated_bytes(CallableT&& run, PolicySelector policy_selec cuda::std::execution::prop{cuda::mr::__get_memory_resource_t{}, device_memory_resource{{}, stream.get(), &bytes_allocated, &bytes_deallocated}}, stream, - cuda::execution::__tune(policy_selector)}; + cuda::execution::tune(policy_selector)}; REQUIRE(cudaSuccess == run(env)); stream.sync(); CHECK(bytes_allocated > 0); @@ -1155,10 +1160,9 @@ TEST_CASE("DeviceRadixSort::SortPairs can be tuned", "[radix_sort][device]") env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortPairs DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1169,10 +1173,9 @@ TEST_CASE("DeviceRadixSort::SortPairs DoubleBuffer can be tuned", "[radix_sort][ return cub::DeviceRadixSort::SortPairs(double_buf, double_buf, static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortPairsDescending can be tuned", "[radix_sort][device]") @@ -1190,10 +1193,9 @@ TEST_CASE("DeviceRadixSort::SortPairsDescending can be tuned", "[radix_sort][dev env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortPairsDescending DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1204,10 +1206,9 @@ TEST_CASE("DeviceRadixSort::SortPairsDescending DoubleBuffer can be tuned", "[ra return cub::DeviceRadixSort::SortPairsDescending(double_buf, double_buf, static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortKeys can be tuned", "[radix_sort][device]") @@ -1218,10 +1219,9 @@ TEST_CASE("DeviceRadixSort::SortKeys can be tuned", "[radix_sort][device]") data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortKeys DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1232,10 +1232,9 @@ TEST_CASE("DeviceRadixSort::SortKeys DoubleBuffer can be tuned", "[radix_sort][d return cub::DeviceRadixSort::SortKeys(double_buf, static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortKeysDescending can be tuned", "[radix_sort][device]") @@ -1246,10 +1245,9 @@ TEST_CASE("DeviceRadixSort::SortKeysDescending can be tuned", "[radix_sort][devi data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } TEST_CASE("DeviceRadixSort::SortKeysDescending DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1260,8 +1258,9 @@ TEST_CASE("DeviceRadixSort::SortKeysDescending DoubleBuffer can be tuned", "[rad return cub::DeviceRadixSort::SortKeysDescending(double_buf, static_cast(data.size()), 0, 32, env); }; - auto default_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - auto tuned_bytes = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - - CHECK(default_bytes != tuned_bytes); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes64); } + +#endif // TEST_LAUNCH != 1 From 64027db766afb75e680c5c81e59cd10a55dc5fe7 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:12:24 +0200 Subject: [PATCH 03/17] device layer --- cub/cub/device/device_radix_sort.cuh | 106 ++++++++++++++++----------- 1 file changed, 62 insertions(+), 44 deletions(-) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index 63a0b7882bc0..7ed47a9dd021 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -149,6 +149,45 @@ inline constexpr bool __can_use_radix_sort = struct DeviceRadixSort { private: + // TODO(bgruber): I would ideally like to have the logic of extracting the policy selector from the tuning environment + // inside the dispatch function, but this will not work with CCCL.C, which needs to pass a stateful policy selector. + // Refactor this once we have a host code JIT compiler. + template > + CUB_RUNTIME_FUNCTION static cudaError_t select_tuning_and_dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + OffsetT num_items, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + DecomposerT decomposer = {}, + TuningEnvT = {}) + { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + return detail::radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + begin_bit, + end_bit, + is_overwrite_okay, + stream, + decomposer, + policy_selector_t{}); + } + template CUB_RUNTIME_FUNCTION static cudaError_t radix_sort_with_decomposer( void* d_temp_storage, @@ -171,7 +210,7 @@ private: if constexpr (decomposer_check) { - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -359,7 +398,7 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -472,10 +511,7 @@ public: // Dispatch with environment - handles all boilerplate return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -486,7 +522,7 @@ public: false, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -1070,7 +1106,7 @@ public: constexpr bool is_overwrite_okay = true; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -1170,10 +1206,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -1184,7 +1217,7 @@ public: true, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -1735,7 +1768,7 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -1847,10 +1880,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -1861,7 +1891,7 @@ public: false, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -2255,7 +2285,7 @@ public: constexpr bool is_overwrite_okay = true; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -2355,10 +2385,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -2369,7 +2396,7 @@ public: true, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -3010,7 +3037,7 @@ public: // Null value type DoubleBuffer d_values; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -3108,10 +3135,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -3122,7 +3146,7 @@ public: false, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -3623,7 +3647,7 @@ public: // Null value type DoubleBuffer d_values; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -3718,10 +3742,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -3732,7 +3753,7 @@ public: true, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -4195,7 +4216,7 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -4293,10 +4314,7 @@ public: return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( storage, bytes, d_keys, @@ -4307,7 +4325,7 @@ public: false, stream, {}, - policy_selector_t{}); + tuning_env); }); } @@ -4661,7 +4679,7 @@ public: // Null value type DoubleBuffer d_values; - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, From 6f8e250bdef0fbe5b9b387165b078880986ab685 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:16:22 +0200 Subject: [PATCH 04/17] fix arch_id --- cub/benchmarks/bench/radix_sort/policy_selector.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/benchmarks/bench/radix_sort/policy_selector.h b/cub/benchmarks/bench/radix_sort/policy_selector.h index 896311d10c59..11ef201f5589 100644 --- a/cub/benchmarks/bench/radix_sort/policy_selector.h +++ b/cub/benchmarks/bench/radix_sort/policy_selector.h @@ -79,7 +79,7 @@ constexpr std::size_t max_onesweep_temp_storage_size() { using portion_offset = int; - constexpr auto active_policy = policy_selector{}(cuda::arch_id{}); + constexpr auto active_policy = policy_selector{}(cuda::compute_capability{}); constexpr auto onesweep = active_policy.onesweep; using onesweep_policy_t = AgentRadixSortOnesweepPolicy< From c6d54226388be4968622afa776431570518d3010 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:23:20 +0200 Subject: [PATCH 05/17] Fixes --- cub/cub/device/device_radix_sort.cuh | 2 +- cub/test/catch2_test_device_radix_sort_env.cu | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index 7ed47a9dd021..950ef604f367 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -156,7 +156,7 @@ private: typename KeyT, typename ValueT, typename OffsetT, - typename DecomposerT = identity_decomposer_t, + typename DecomposerT = detail::identity_decomposer_t, typename TuningEnvT = ::cuda::std::execution::env<>> CUB_RUNTIME_FUNCTION static cudaError_t select_tuning_and_dispatch( void* d_temp_storage, diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index 4630b7805625..5368b73d050d 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1116,13 +1116,13 @@ TEST_CASE("Device radix sort pairs descending DB decomposer+bits uses custom str template struct tiny_onesweep_policy_selector { - _CCCL_API constexpr auto operator()(cuda::arch_id arch) const -> cub::detail::radix_sort::radix_sort_policy + _CCCL_API constexpr auto operator()(cuda::compute_capability cc) const -> cub::detail::radix_sort::radix_sort_policy { - using default_selector_t = cub::detail::radix_sort::policy_selector_from_types; - auto policy = default_selector_t{}(arch); - policy.use_onesweep = true; - policy.onesweep.block_threads = BlockThreads; - policy.onesweep.items_per_thread = 1; + using default_selector_t = cub::detail::radix_sort::policy_selector_from_types; + auto policy = default_selector_t{}(cc); + policy.use_onesweep = true; + policy.onesweep.threads_per_block = BlockThreads; + policy.onesweep.items_per_thread = 1; return policy; } }; @@ -1134,8 +1134,8 @@ std::size_t measure_allocated_bytes(CallableT&& run, PolicySelector policy_selec size_t bytes_allocated = 0; size_t bytes_deallocated = 0; auto env = stdexec::env{ - cuda::std::execution::prop{cuda::mr::__get_memory_resource_t{}, - device_memory_resource{{}, stream.get(), &bytes_allocated, &bytes_deallocated}}, + cuda::std::execution::prop{ + cuda::mr::__get_memory_resource_t{}, device_memory_resource{stream.get(), &bytes_allocated, &bytes_deallocated}}, stream, cuda::execution::tune(policy_selector)}; REQUIRE(cudaSuccess == run(env)); From 37e8dca0f1e11c74e987a67ad4c7e8c16b7e588f Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:28:18 +0200 Subject: [PATCH 06/17] refactor --- cub/test/catch2_test_device_radix_sort_env.cu | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index 5368b73d050d..f9a1d16c0e89 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1133,11 +1133,9 @@ std::size_t measure_allocated_bytes(CallableT&& run, PolicySelector policy_selec cuda::stream_ref stream{cudaStream_t{}}; size_t bytes_allocated = 0; size_t bytes_deallocated = 0; - auto env = stdexec::env{ - cuda::std::execution::prop{ - cuda::mr::__get_memory_resource_t{}, device_memory_resource{stream.get(), &bytes_allocated, &bytes_deallocated}}, - stream, - cuda::execution::tune(policy_selector)}; + auto env = stdexec::env{device_memory_resource{stream.get(), &bytes_allocated, &bytes_deallocated}, + stream, + cuda::execution::tune(policy_selector)}; REQUIRE(cudaSuccess == run(env)); stream.sync(); CHECK(bytes_allocated > 0); From 59ac4845480c1986d6b444928c342726786e94d2 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:31:33 +0200 Subject: [PATCH 07/17] Fix benchmark --- cub/benchmarks/bench/radix_sort/keys.cu | 30 ++++++++++++------------ cub/benchmarks/bench/radix_sort/pairs.cu | 30 ++++++++++++------------ 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/cub/benchmarks/bench/radix_sort/keys.cu b/cub/benchmarks/bench/radix_sort/keys.cu index 7b654fc3ab60..682fe3c32e72 100644 --- a/cub/benchmarks/bench/radix_sort/keys.cu +++ b/cub/benchmarks/bench/radix_sort/keys.cu @@ -38,23 +38,23 @@ void radix_sort_keys(nvbench::state& state, nvbench::type_list) caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - if (cub::DeviceRadixSort::SortKeys( - d_buffer_1, - d_buffer_2, - static_cast(elements), - 0, - sizeof(key_t) * 8, - cub_bench_env(alloc, - launch + auto env = cub_bench_env( + alloc, + launch #if !TUNE_BASE - , - cuda::execution::tune(policy_selector{}) + , + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE - )) - != cudaSuccess) - { - throw 42; // TODO(bgruber): what's the appropriate way to abort a benchmark? - } + ); + _CCCL_TRY_CUDA_API( + cub::DeviceRadixSort::SortKeys, + "SortKeys failed", + d_buffer_1, + d_buffer_2, + static_cast(elements), + 0, + sizeof(key_t) * 8, + env); }); } diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index 49eeb7aeff42..8286d9f0aaf1 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -46,23 +46,23 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list(elements), - cub_bench_env(alloc, - launch + auto env = cub_bench_env( + alloc, + launch #if !TUNE_BASE - , - cuda::execution::tune(policy_selector{}) + , + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE - )) - != cudaSuccess) - { - throw 42; // TODO(bgruber): what's the appropriate way to abort a benchmark? - } + ); + _CCCL_TRY_CUDA_API( + cub::DeviceRadixSort::SortPairs, + "SortPairs failed", + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + static_cast(elements), + env); }); } From 71579d7f389184927fd08c0d5c59172209d9e8c0 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:36:34 +0200 Subject: [PATCH 08/17] fix --- cub/benchmarks/bench/radix_sort/pairs.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index 8286d9f0aaf1..16102716db22 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -62,6 +62,8 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list(elements), + 0, + sizeof(key_t) * 8, env); }); } From 0c76b96ba746f5f412ef3ce56f85c5e52950234a Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:42:23 +0200 Subject: [PATCH 09/17] refactor --- cub/benchmarks/bench/radix_sort/keys.cu | 9 ++++----- cub/benchmarks/bench/radix_sort/pairs.cu | 24 +++++++++++------------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/cub/benchmarks/bench/radix_sort/keys.cu b/cub/benchmarks/bench/radix_sort/keys.cu index 682fe3c32e72..4e64b5a212dd 100644 --- a/cub/benchmarks/bench/radix_sort/keys.cu +++ b/cub/benchmarks/bench/radix_sort/keys.cu @@ -14,7 +14,6 @@ template void radix_sort_keys(nvbench::state& state, nvbench::type_list) { - using key_t = T; using value_t = cub::NullType; if constexpr (!fits_in_default_shared_memory()) { @@ -28,8 +27,8 @@ void radix_sort_keys(nvbench::state& state, nvbench::type_list) thrust::device_vector buffer_1 = generate(elements, entropy); thrust::device_vector buffer_2(elements, thrust::no_init); - const key_t* d_buffer_1 = thrust::raw_pointer_cast(buffer_1.data()); - key_t* d_buffer_2 = thrust::raw_pointer_cast(buffer_2.data()); + const T* d_buffer_1 = thrust::raw_pointer_cast(buffer_1.data()); + T* d_buffer_2 = thrust::raw_pointer_cast(buffer_2.data()); // Enable throughput calculations and add "Size" column to results. state.add_element_count(elements); @@ -43,7 +42,7 @@ void radix_sort_keys(nvbench::state& state, nvbench::type_list) launch #if !TUNE_BASE , - cuda::execution::tune(policy_selector{}) + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE ); _CCCL_TRY_CUDA_API( @@ -53,7 +52,7 @@ void radix_sort_keys(nvbench::state& state, nvbench::type_list) d_buffer_2, static_cast(elements), 0, - sizeof(key_t) * 8, + sizeof(T) * 8, env); }); } diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index 16102716db22..5ae1f9491691 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -16,9 +16,7 @@ template void radix_sort_values(nvbench::state& state, nvbench::type_list) { - using key_t = KeyT; - using value_t = ValueT; - if constexpr (!fits_in_default_shared_memory()) + if constexpr (!fits_in_default_shared_memory()) { return; } @@ -27,15 +25,15 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list(state.get_int64("Elements{io}")); const bit_entropy entropy = str_to_entropy(state.get_string("Entropy")); - thrust::device_vector keys_in = generate(elements, entropy); - thrust::device_vector keys_out(elements, thrust::no_init); - thrust::device_vector values_in = generate(elements); - thrust::device_vector values_out(elements, thrust::no_init); + thrust::device_vector keys_in = generate(elements, entropy); + thrust::device_vector keys_out(elements, thrust::no_init); + thrust::device_vector values_in = generate(elements); + thrust::device_vector values_out(elements, thrust::no_init); - const key_t* d_keys_in = thrust::raw_pointer_cast(keys_in.data()); - key_t* d_keys_out = thrust::raw_pointer_cast(keys_out.data()); - const value_t* d_values_in = thrust::raw_pointer_cast(values_in.data()); - value_t* d_values_out = thrust::raw_pointer_cast(values_out.data()); + const KeyT* d_keys_in = thrust::raw_pointer_cast(keys_in.data()); + KeyT* d_keys_out = thrust::raw_pointer_cast(keys_out.data()); + const ValueT* d_values_in = thrust::raw_pointer_cast(values_in.data()); + ValueT* d_values_out = thrust::raw_pointer_cast(values_out.data()); // Enable throughput calculations and add "Size" column to results. state.add_element_count(elements); @@ -51,7 +49,7 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list{}) + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE ); _CCCL_TRY_CUDA_API( @@ -63,7 +61,7 @@ void radix_sort_values(nvbench::state& state, nvbench::type_list(elements), 0, - sizeof(key_t) * 8, + sizeof(KeyT) * 8, env); }); } From efd46684a85ec310937caa684dc05178c98cfb10 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:43:17 +0200 Subject: [PATCH 10/17] revert --- cub/cub/detail/temporary_storage.cuh | 2 -- 1 file changed, 2 deletions(-) diff --git a/cub/cub/detail/temporary_storage.cuh b/cub/cub/detail/temporary_storage.cuh index e6be1ecadf85..9863995fc8be 100644 --- a/cub/cub/detail/temporary_storage.cuh +++ b/cub/cub/detail/temporary_storage.cuh @@ -347,7 +347,6 @@ private: } }; -_CCCL_EXEC_CHECK_DISABLE template CUB_RUNTIME_FUNCTION cudaError_t allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_bytes, MRT& mr) @@ -363,7 +362,6 @@ allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_b } } -_CCCL_EXEC_CHECK_DISABLE template CUB_RUNTIME_FUNCTION cudaError_t deallocate(::cuda::stream_ref stream, void* d_temp_storage, size_t temp_storage_bytes, MRT& mr) From 6dc095e65b746c4325f776083bd3c59e88f6a81a Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 17:43:45 +0200 Subject: [PATCH 11/17] revert --- .../dispatch_segmented_radix_sort.cuh | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh b/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh index f39f5c7d4d76..23bb954acb52 100644 --- a/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh +++ b/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh @@ -883,17 +883,17 @@ template , + typename DecomposerT = identity_decomposer_t, + typename PolicySelector = policy_selector_from_types, + typename KernelSource = DeviceSegmentedRadixSortKernelSource< + PolicySelector, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_radix_sort_policy_selector From 7ba0ab8878b3e60a66d1aa9093b1e54eb6744d09 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 18:25:28 +0200 Subject: [PATCH 12/17] device layer --- cub/cub/device/device_radix_sort.cuh | 805 ++++++++++++++------------- 1 file changed, 409 insertions(+), 396 deletions(-) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index 950ef604f367..996ac1caad11 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -188,7 +188,12 @@ private: policy_selector_t{}); } - template + template > CUB_RUNTIME_FUNCTION static cudaError_t radix_sort_with_decomposer( void* d_temp_storage, size_t& temp_storage_bytes, @@ -197,6 +202,7 @@ private: NumItemsT num_items, DecomposerT decomposer, cudaStream_t stream, + TuningEnvT tuning_env = {}, int begin_bit = 0, int end_bit = detail::radix::traits_t::default_end_bit(DecomposerT{}), bool is_overwrite_okay = true) @@ -220,12 +226,18 @@ private: end_bit, is_overwrite_okay, stream, - decomposer); + decomposer, + tuning_env); } _CCCL_UNREACHABLE(); } - template + template > CUB_RUNTIME_FUNCTION static cudaError_t radix_sort_with_decomposer( void* d_temp_storage, size_t& temp_storage_bytes, @@ -236,8 +248,9 @@ private: NumItemsT num_items, DecomposerT decomposer, cudaStream_t stream, - int begin_bit = 0, - int end_bit = detail::radix::traits_t::default_end_bit(DecomposerT{})) + TuningEnvT tuning_env = {}, + int begin_bit = 0, + int end_bit = detail::radix::traits_t::default_end_bit(DecomposerT{})) { // We cast away const-ness, but will *not* write to these arrays. ``DispatchRadixSort::Dispatch`` will allocate // temporary storage and create a new double-buffer internally when the ``is_overwrite_ok`` flag is not set. @@ -251,6 +264,7 @@ private: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit, /* is_overwrite_okay */ false); @@ -509,21 +523,20 @@ public: DoubleBuffer d_values(const_cast(d_values_in), d_values_out); // Dispatch with environment - handles all boilerplate - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - false, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); } //! @rst @@ -629,22 +642,23 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairs(void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - const ValueT* d_values_in, - ValueT* d_values_out, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -657,6 +671,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -738,7 +753,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -751,7 +766,7 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -762,6 +777,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); }); @@ -862,20 +878,21 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairs(void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - const ValueT* d_values_in, - ValueT* d_values_out, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -959,7 +976,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -970,9 +987,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream); + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream, tuning_env); }); } @@ -1204,21 +1221,20 @@ public: using offset_t = detail::choose_offset_t; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - true, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); } //! @rst @@ -1319,18 +1335,19 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairs(void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - DoubleBuffer& d_values, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -1402,7 +1419,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -1411,9 +1428,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); }); } @@ -1526,24 +1543,25 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairs(void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - DoubleBuffer& d_values, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( - d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); } //! @rst @@ -1618,7 +1636,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -1629,9 +1647,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); }); } @@ -1878,21 +1896,20 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - false, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); } //! @rst @@ -1998,23 +2015,23 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairsDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - const ValueT* d_values_in, - ValueT* d_values_out, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -2027,6 +2044,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -2126,21 +2144,21 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairsDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - const ValueT* d_values_in, - ValueT* d_values_out, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -2383,21 +2401,20 @@ public: using offset_t = detail::choose_offset_t; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - true, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); } //! @rst @@ -2498,19 +2515,19 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairsDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - DoubleBuffer& d_values, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -2626,25 +2643,25 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortPairsDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - DoubleBuffer& d_values, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( - d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); } //! @rst @@ -2697,7 +2714,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -2710,7 +2727,7 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -2721,6 +2738,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -2773,7 +2791,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -2784,9 +2802,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream); + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream, tuning_env); }); } @@ -2835,7 +2853,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -2844,9 +2862,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); }); } @@ -2898,7 +2916,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( DoubleBuffer& d_keys, DoubleBuffer& d_values, @@ -2909,9 +2927,9 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); }); } @@ -3133,21 +3151,20 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - false, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); } //! @rst @@ -3244,20 +3261,20 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeys(void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -3270,6 +3287,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -3325,7 +3343,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -3336,7 +3354,7 @@ public: EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -3347,6 +3365,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -3438,18 +3457,18 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeys(void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -3512,12 +3531,12 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys(const KeyT* d_keys_in, KeyT* d_keys_out, NumItemsT num_items, DecomposerT decomposer, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -3527,7 +3546,8 @@ public: static_cast(nullptr), num_items, decomposer, - stream); + stream, + tuning_env); }); } @@ -3740,21 +3760,20 @@ public: DoubleBuffer d_values; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - true, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); } //! @rst @@ -3843,17 +3862,17 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeys(void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); DoubleBuffer d_values; @@ -3908,15 +3927,15 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys(DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { DoubleBuffer d_values; return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); }); } @@ -4017,24 +4036,24 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeys(void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); DoubleBuffer d_values; return radix_sort_with_decomposer( - d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); } //! @rst @@ -4087,16 +4106,16 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, int begin_bit, int end_bit, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { DoubleBuffer d_values; return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); }); } @@ -4312,21 +4331,20 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - return select_tuning_and_dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - false, - stream, - {}, - tuning_env); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); } //! @rst @@ -4423,21 +4441,20 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeysDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -4450,6 +4467,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -4537,19 +4555,18 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeysDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - const KeyT* d_keys_in, - KeyT* d_keys_out, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); return radix_sort_with_decomposer( @@ -4772,24 +4789,20 @@ public: DoubleBuffer d_values; - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, auto stream) { - using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; - return detail::radix_sort::dispatch( - storage, - bytes, - d_keys, - d_values, - static_cast(num_items), - begin_bit, - end_bit, - true, - stream, - {}, - policy_selector_t{}); - }); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); } //! @rst @@ -4878,18 +4891,17 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeysDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - NumItemsT num_items, - DecomposerT decomposer, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); DoubleBuffer d_values; @@ -4994,25 +5006,24 @@ public: //! @param[in] stream //! **[optional]** CUDA stream to launch kernels within. //! Default is stream0. - template - CUB_RUNTIME_FUNCTION static // - ::cuda::std::enable_if_t< // - !::cuda::std::is_convertible_v, // - cudaError_t> - SortKeysDescending( - void* d_temp_storage, - size_t& temp_storage_bytes, - DoubleBuffer& d_keys, - NumItemsT num_items, - DecomposerT decomposer, - int begin_bit, - int end_bit, - cudaStream_t stream = nullptr) + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); DoubleBuffer d_values; return radix_sort_with_decomposer( - d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); } //! @rst @@ -5060,7 +5071,7 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( const KeyT* d_keys_in, KeyT* d_keys_out, @@ -5072,7 +5083,7 @@ public: { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -5083,6 +5094,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -5130,12 +5142,12 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( const KeyT* d_keys_in, KeyT* d_keys_out, NumItemsT num_items, DecomposerT decomposer, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return radix_sort_with_decomposer( storage, bytes, @@ -5145,7 +5157,8 @@ public: static_cast(nullptr), num_items, decomposer, - stream); + stream, + tuning_env); }); } @@ -5190,15 +5203,15 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending(DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { DoubleBuffer d_values; return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); }); } @@ -5246,16 +5259,16 @@ public: typename NumItemsT, typename DecomposerT, typename EnvT = ::cuda::std::execution::env<>, - typename ::cuda::std::enable_if_t, int> = 0> + ::cuda::std::enable_if_t, int> = 0> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, int begin_bit, int end_bit, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE(GetName()); - return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { DoubleBuffer d_values; return radix_sort_with_decomposer( - storage, bytes, d_keys, d_values, num_items, decomposer, stream, begin_bit, end_bit); + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); }); } From 51960b22751cf63a4e4d3a7fdccd33a078974e6b Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 18:34:47 +0200 Subject: [PATCH 13/17] more tests --- cub/test/catch2_test_device_radix_sort_env.cu | 332 ++++++++++++++++-- 1 file changed, 303 insertions(+), 29 deletions(-) diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index f9a1d16c0e89..17992b9a4c60 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1118,11 +1118,17 @@ struct tiny_onesweep_policy_selector { _CCCL_API constexpr auto operator()(cuda::compute_capability cc) const -> cub::detail::radix_sort::radix_sort_policy { - using default_selector_t = cub::detail::radix_sort::policy_selector_from_types; - auto policy = default_selector_t{}(cc); - policy.use_onesweep = true; - policy.onesweep.threads_per_block = BlockThreads; - policy.onesweep.items_per_thread = 1; + using default_selector_t = cub::detail::radix_sort::policy_selector_from_types; + auto policy = default_selector_t{}(cc); + policy.use_onesweep = true; + policy.onesweep.threads_per_block = BlockThreads; + policy.onesweep.items_per_thread = 1; + policy.single_tile.threads_per_block = BlockThreads; + policy.single_tile.items_per_thread = 1; + policy.downsweep.threads_per_block = BlockThreads; + policy.downsweep.items_per_thread = 1; + policy.alt_downsweep.threads_per_block = BlockThreads; + policy.alt_downsweep.items_per_thread = 1; return policy; } }; @@ -1158,9 +1164,9 @@ TEST_CASE("DeviceRadixSort::SortPairs can be tuned", "[radix_sort][device]") env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortPairs DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1171,9 +1177,9 @@ TEST_CASE("DeviceRadixSort::SortPairs DoubleBuffer can be tuned", "[radix_sort][ return cub::DeviceRadixSort::SortPairs(double_buf, double_buf, static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortPairsDescending can be tuned", "[radix_sort][device]") @@ -1191,9 +1197,9 @@ TEST_CASE("DeviceRadixSort::SortPairsDescending can be tuned", "[radix_sort][dev env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortPairsDescending DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1204,9 +1210,9 @@ TEST_CASE("DeviceRadixSort::SortPairsDescending DoubleBuffer can be tuned", "[ra return cub::DeviceRadixSort::SortPairsDescending(double_buf, double_buf, static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortKeys can be tuned", "[radix_sort][device]") @@ -1217,9 +1223,9 @@ TEST_CASE("DeviceRadixSort::SortKeys can be tuned", "[radix_sort][device]") data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortKeys DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1230,9 +1236,9 @@ TEST_CASE("DeviceRadixSort::SortKeys DoubleBuffer can be tuned", "[radix_sort][d return cub::DeviceRadixSort::SortKeys(double_buf, static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortKeysDescending can be tuned", "[radix_sort][device]") @@ -1243,9 +1249,9 @@ TEST_CASE("DeviceRadixSort::SortKeysDescending can be tuned", "[radix_sort][devi data.data().get(), data.data().get(), static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } TEST_CASE("DeviceRadixSort::SortKeysDescending DoubleBuffer can be tuned", "[radix_sort][device]") @@ -1256,9 +1262,277 @@ TEST_CASE("DeviceRadixSort::SortKeysDescending DoubleBuffer can be tuned", "[rad return cub::DeviceRadixSort::SortKeysDescending(double_buf, static_cast(data.size()), 0, 32, env); }; - const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - const auto bytes64 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); - CHECK(bytes32 != bytes64); + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +// decomposer variants: keys + +TEST_CASE("DeviceRadixSort::SortKeys decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortKeys( + data.data().get(), data.data().get(), static_cast(data.size()), keys_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeys decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortKeys( + data.data().get(), data.data().get(), static_cast(data.size()), keys_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeys DB decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto buf0 = c2h::device_vector(10'000); + auto buf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(buf0.data().get(), buf1.data().get()); + return cub::DeviceRadixSort::SortKeys(d_keys, static_cast(buf0.size()), keys_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeys DB decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto buf0 = c2h::device_vector(10'000); + auto buf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(buf0.data().get(), buf1.data().get()); + return cub::DeviceRadixSort::SortKeys(d_keys, static_cast(buf0.size()), keys_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +// decomposer variants: keys descending + +TEST_CASE("DeviceRadixSort::SortKeysDescending decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortKeysDescending( + data.data().get(), data.data().get(), static_cast(data.size()), keys_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeysDescending decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto data = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortKeysDescending( + data.data().get(), data.data().get(), static_cast(data.size()), keys_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeysDescending DB decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto buf0 = c2h::device_vector(10'000); + auto buf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(buf0.data().get(), buf1.data().get()); + return cub::DeviceRadixSort::SortKeysDescending(d_keys, static_cast(buf0.size()), keys_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortKeysDescending DB decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto buf0 = c2h::device_vector(10'000); + auto buf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(buf0.data().get(), buf1.data().get()); + return cub::DeviceRadixSort::SortKeysDescending( + d_keys, static_cast(buf0.size()), keys_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +// decomposer variants: pairs + +TEST_CASE("DeviceRadixSort::SortPairs decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto kbuf0 = c2h::device_vector(10'000); + auto kbuf1 = c2h::device_vector(10'000); + auto vbuf0 = c2h::device_vector(10'000); + auto vbuf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(kbuf0.data().get(), kbuf1.data().get()); + cub::DoubleBuffer d_values(vbuf0.data().get(), vbuf1.data().get()); + return cub::DeviceRadixSort::SortPairs( + d_keys, d_values, static_cast(kbuf0.size()), pairs_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairs decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto keys = c2h::device_vector(10'000); + auto vals = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortPairs( + keys.data().get(), + keys.data().get(), + vals.data().get(), + vals.data().get(), + static_cast(keys.size()), + pairs_decomposer_t{}, + env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairs DB decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto kbuf0 = c2h::device_vector(10'000); + auto kbuf1 = c2h::device_vector(10'000); + auto vbuf0 = c2h::device_vector(10'000); + auto vbuf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(kbuf0.data().get(), kbuf1.data().get()); + cub::DoubleBuffer d_values(vbuf0.data().get(), vbuf1.data().get()); + return cub::DeviceRadixSort::SortPairs(d_keys, d_values, static_cast(kbuf0.size()), pairs_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairs DB decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto kbuf0 = c2h::device_vector(10'000); + auto kbuf1 = c2h::device_vector(10'000); + auto vbuf0 = c2h::device_vector(10'000); + auto vbuf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(kbuf0.data().get(), kbuf1.data().get()); + cub::DoubleBuffer d_values(vbuf0.data().get(), vbuf1.data().get()); + return cub::DeviceRadixSort::SortPairs( + d_keys, d_values, static_cast(kbuf0.size()), pairs_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +// decomposer variants: pairs descending + +TEST_CASE("DeviceRadixSort::SortPairsDescending decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto keys = c2h::device_vector(10'000); + auto vals = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortPairsDescending( + keys.data().get(), + keys.data().get(), + vals.data().get(), + vals.data().get(), + static_cast(keys.size()), + pairs_decomposer_t{}, + 0, + 32, + env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairsDescending decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto keys = c2h::device_vector(10'000); + auto vals = c2h::device_vector(10'000); + return cub::DeviceRadixSort::SortPairsDescending( + keys.data().get(), + keys.data().get(), + vals.data().get(), + vals.data().get(), + static_cast(keys.size()), + pairs_decomposer_t{}, + env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairsDescending DB decomposer can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto kbuf0 = c2h::device_vector(10'000); + auto kbuf1 = c2h::device_vector(10'000); + auto vbuf0 = c2h::device_vector(10'000); + auto vbuf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(kbuf0.data().get(), kbuf1.data().get()); + cub::DoubleBuffer d_values(vbuf0.data().get(), vbuf1.data().get()); + return cub::DeviceRadixSort::SortPairsDescending( + d_keys, d_values, static_cast(kbuf0.size()), pairs_decomposer_t{}, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); +} + +TEST_CASE("DeviceRadixSort::SortPairsDescending DB decomposer+bits can be tuned", "[radix_sort][device]") +{ + auto l = [&](auto env) { + auto kbuf0 = c2h::device_vector(10'000); + auto kbuf1 = c2h::device_vector(10'000); + auto vbuf0 = c2h::device_vector(10'000); + auto vbuf1 = c2h::device_vector(10'000); + cub::DoubleBuffer d_keys(kbuf0.data().get(), kbuf1.data().get()); + cub::DoubleBuffer d_values(vbuf0.data().get(), vbuf1.data().get()); + return cub::DeviceRadixSort::SortPairsDescending( + d_keys, d_values, static_cast(kbuf0.size()), pairs_decomposer_t{}, 0, 32, env); + }; + + const auto bytes32 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + const auto bytes128 = measure_allocated_bytes(l, tiny_onesweep_policy_selector{}); + CHECK(bytes32 != bytes128); } #endif // TEST_LAUNCH != 1 From f7b5459fbd03d62db5f7c72aa089c8c5fe930f2f Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 18:36:33 +0200 Subject: [PATCH 14/17] drop comments --- cub/cub/device/device_radix_sort.cuh | 3 --- 1 file changed, 3 deletions(-) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index 996ac1caad11..58509ba8fe1e 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -522,7 +522,6 @@ public: DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values(const_cast(d_values_in), d_values_out); - // Dispatch with environment - handles all boilerplate return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { return select_tuning_and_dispatch( storage, @@ -3147,7 +3146,6 @@ public: using offset_t = detail::choose_offset_t; - // Dispatch with environment - handles all boilerplate DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; @@ -4327,7 +4325,6 @@ public: using offset_t = detail::choose_offset_t; - // Dispatch with environment - handles all boilerplate DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); DoubleBuffer d_values; From 056327f1b69d363aa9e19ab4ec50ab61d0946d39 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 20:41:14 +0200 Subject: [PATCH 15/17] fix unused tuning_env --- cub/cub/device/device_radix_sort.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index 58509ba8fe1e..538fda775d98 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -776,7 +776,7 @@ public: num_items, decomposer, stream, - {}, + tuning_env, begin_bit, end_bit); }); From c01905482667d41a72db62a543fd0dfd9c92f408 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 20:42:44 +0200 Subject: [PATCH 16/17] Reduce SMEM --- cub/test/catch2_test_device_radix_sort_env.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index 17992b9a4c60..f95bd83340db 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1129,6 +1129,7 @@ struct tiny_onesweep_policy_selector policy.downsweep.items_per_thread = 1; policy.alt_downsweep.threads_per_block = BlockThreads; policy.alt_downsweep.items_per_thread = 1; + policy.histogram.num_parts = 1; return policy; } }; From fe72f1e955e7fcc1f9e680cc45a46640e83828bc Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 21:29:03 +0200 Subject: [PATCH 17/17] use selector over hub --- cub/benchmarks/bench/radix_sort/policy_selector.h | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/cub/benchmarks/bench/radix_sort/policy_selector.h b/cub/benchmarks/bench/radix_sort/policy_selector.h index 11ef201f5589..22bac8e49ca9 100644 --- a/cub/benchmarks/bench/radix_sort/policy_selector.h +++ b/cub/benchmarks/bench/radix_sort/policy_selector.h @@ -112,10 +112,9 @@ constexpr std::size_t max_onesweep_temp_storage_size() template constexpr std::size_t max_temp_storage_size() { - using offset_t = cub::detail::choose_offset_t; - using policy_t = typename policy_hub_t::policy_t; - - static_assert(policy_t::ONESWEEP); + using offset_t = cub::detail::choose_offset_t; + constexpr auto active_policy = policy_selector{}(cuda::compute_capability{}); + static_assert(active_policy.use_onesweep); return max_onesweep_temp_storage_size(); }