diff --git a/cub/benchmarks/bench/radix_sort/keys.cu b/cub/benchmarks/bench/radix_sort/keys.cu index c05a72190382..4e64b5a212dd 100644 --- a/cub/benchmarks/bench/radix_sort/keys.cu +++ b/cub/benchmarks/bench/radix_sort/keys.cu @@ -14,82 +14,46 @@ 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 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 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); 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() + auto env = cub_bench_env( + alloc, + launch #if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} + , + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE ); + _CCCL_TRY_CUDA_API( + cub::DeviceRadixSort::SortKeys, + "SortKeys failed", + d_buffer_1, + d_buffer_2, + static_cast(elements), + 0, + sizeof(T) * 8, + env); }); } diff --git a/cub/benchmarks/bench/radix_sort/pairs.cu b/cub/benchmarks/bench/radix_sort/pairs.cu index f293fffa7649..5ae1f9491691 100644 --- a/cub/benchmarks/bench/radix_sort/pairs.cu +++ b/cub/benchmarks/bench/radix_sort/pairs.cu @@ -16,37 +16,24 @@ 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()) + 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 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); @@ -55,48 +42,27 @@ 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() + auto env = cub_bench_env( + alloc, + launch #if !TUNE_BASE - , - cub::detail::identity_decomposer_t{}, - policy_selector{} + , + cuda::execution::tune(policy_selector{}) #endif // !TUNE_BASE ); + _CCCL_TRY_CUDA_API( + cub::DeviceRadixSort::SortPairs, + "SortPairs failed", + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + static_cast(elements), + 0, + sizeof(KeyT) * 8, + env); }); } diff --git a/cub/benchmarks/bench/radix_sort/policy_selector.h b/cub/benchmarks/bench/radix_sort/policy_selector.h index 051f29a900af..22bac8e49ca9 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::compute_capability{}); + + 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,10 @@ 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; - - static_assert(policy_t::ONESWEEP); - return max_onesweep_temp_storage_size(); + 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(); } template diff --git a/cub/cub/device/device_radix_sort.cuh b/cub/cub/device/device_radix_sort.cuh index d7f038daa48b..538fda775d98 100644 --- a/cub/cub/device/device_radix_sort.cuh +++ b/cub/cub/device/device_radix_sort.cuh @@ -149,7 +149,51 @@ inline constexpr bool __can_use_radix_sort = struct DeviceRadixSort { private: - template + // 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, size_t& temp_storage_bytes, @@ -158,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) @@ -171,7 +216,7 @@ private: if constexpr (decomposer_check) { - return detail::radix_sort::dispatch( + return select_tuning_and_dispatch( d_temp_storage, temp_storage_bytes, d_keys, @@ -181,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, @@ -197,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. @@ -212,6 +264,7 @@ private: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit, /* is_overwrite_okay */ false); @@ -359,7 +412,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, @@ -450,8 +503,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,13 +519,22 @@ 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); + 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); }); } @@ -579,22 +641,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( @@ -607,6 +670,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -688,7 +752,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, @@ -701,7 +765,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, @@ -712,6 +776,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -812,20 +877,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( @@ -909,7 +975,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, @@ -920,9 +986,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); }); } @@ -1056,7 +1122,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, @@ -1140,8 +1206,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,9 +1220,19 @@ 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, [&](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); }); } @@ -1258,18 +1334,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( @@ -1341,7 +1418,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, @@ -1350,9 +1427,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); }); } @@ -1465,24 +1542,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 @@ -1557,7 +1635,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, @@ -1568,9 +1646,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); }); } @@ -1707,7 +1785,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, @@ -1798,8 +1876,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,9 +1895,19 @@ 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, [&](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); }); } @@ -1926,23 +2014,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( @@ -1955,6 +2043,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -2054,21 +2143,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( @@ -2213,7 +2302,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, @@ -2297,8 +2386,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,9 +2400,19 @@ 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, [&](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); }); } @@ -2415,19 +2514,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( @@ -2543,25 +2642,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 @@ -2614,7 +2713,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, @@ -2627,7 +2726,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, @@ -2638,6 +2737,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -2690,7 +2790,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, @@ -2701,9 +2801,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); }); } @@ -2752,7 +2852,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, @@ -2761,9 +2861,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); }); } @@ -2815,7 +2915,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, @@ -2826,9 +2926,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); }); } @@ -2954,7 +3054,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, @@ -3032,8 +3132,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, @@ -3046,13 +3146,22 @@ 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; - 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, [&](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); }); } @@ -3150,20 +3259,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( @@ -3176,6 +3285,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -3231,7 +3341,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, @@ -3242,7 +3352,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, @@ -3253,6 +3363,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -3344,18 +3455,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( @@ -3418,12 +3529,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, @@ -3433,7 +3544,8 @@ public: static_cast(nullptr), num_items, decomposer, - stream); + stream, + tuning_env); }); } @@ -3553,7 +3665,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, @@ -3635,8 +3747,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,9 +3758,19 @@ 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, [&](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); }); } @@ -3738,17 +3860,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; @@ -3803,15 +3925,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); }); } @@ -3912,24 +4034,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 @@ -3982,16 +4104,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); }); } @@ -4111,7 +4233,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, @@ -4189,8 +4311,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, @@ -4203,13 +4325,22 @@ 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; - 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, [&](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); }); } @@ -4307,21 +4438,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( @@ -4334,6 +4464,7 @@ public: num_items, decomposer, stream, + {}, begin_bit, end_bit); } @@ -4421,19 +4552,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( @@ -4563,7 +4693,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, @@ -4645,8 +4775,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,9 +4786,19 @@ 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, [&](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); }); } @@ -4748,18 +4888,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; @@ -4864,25 +5003,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 @@ -4930,7 +5068,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, @@ -4942,7 +5080,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, @@ -4953,6 +5091,7 @@ public: num_items, decomposer, stream, + tuning_env, begin_bit, end_bit); }); @@ -5000,12 +5139,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, @@ -5015,7 +5154,8 @@ public: static_cast(nullptr), num_items, decomposer, - stream); + stream, + tuning_env); }); } @@ -5060,15 +5200,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); }); } @@ -5116,16 +5256,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); }); } diff --git a/cub/test/catch2_test_device_radix_sort_env.cu b/cub/test/catch2_test_device_radix_sort_env.cu index fc8605aea07f..f95bd83340db 100644 --- a/cub/test/catch2_test_device_radix_sort_env.cu +++ b/cub/test/catch2_test_device_radix_sort_env.cu @@ -1106,3 +1106,434 @@ TEST_CASE("Device radix sort pairs descending DB decomposer+bits uses custom str REQUIRE(keys == expected_keys); REQUIRE(values == expected_values); } + +#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 +{ + _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; + 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; + policy.histogram.num_parts = 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{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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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]") +{ + 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); + }; + + 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