Allow public tuning of the cub::DeviceReduce::ReduceByKey and cub::DeviceRunLengthEncode::Encode

cub/benchmarks/bench/reduce/by_key.cu

@@ -17,8 +17,7 @@
 #if !TUNE_BASE
 struct bench_reduce_by_key_policy_selector
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-    -> cub::detail::reduce_by_key::reduce_by_key_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::ReduceByKeyPolicy
   {
     return {
       TUNE_THREADS,

cub/benchmarks/bench/run_length_encode/encode.cu

@@ -17,8 +17,7 @@
 #if !TUNE_BASE
 struct bench_encode_policy_selector
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-    -> cub::detail::reduce_by_key::reduce_by_key_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::RleEncodePolicy
   {
     return {
       TUNE_THREADS,

cub/cub/agent/agent_reduce_by_key.cuh

@@ -36,56 +36,38 @@ CUB_NAMESPACE_BEGIN
  * Tuning policy types
  ******************************************************************************/
 
-/**
- * @brief Parameterizable tuning policy type for AgentReduceByKey
- *
- * @tparam ThreadsPerBlock
- *   Threads per thread block
- *
- * @tparam ItemsPerThread
- *   Items per thread (per tile of input)
- *
- * @tparam LoadAlgorithm
- *   The BlockLoad algorithm to use
- *
- * @tparam LoadModifier
- *   Cache load modifier for reading input elements
- *
- * @tparam ScanAlgorithm
- *   The BlockScan algorithm to use
- *
- * @tparam DelayConstructorT
- *   Implementation detail, do not specify directly, requirements on the
- *   content of this type are subject to breaking change.
- */
+namespace detail
+{
 template <int ThreadsPerBlock,
           int ItemsPerThread,
           BlockLoadAlgorithm LoadAlgorithm,
           CacheLoadModifier LoadModifier,
           BlockScanAlgorithm ScanAlgorithm,
           typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>>
-struct AgentReduceByKeyPolicy
+struct agent_reduce_by_key_policy
 {
-  ///< Threads per thread block
-  static constexpr int BLOCK_THREADS = ThreadsPerBlock;
-
-  ///< Items per thread (per tile of input)
-  static constexpr int ITEMS_PER_THREAD = ItemsPerThread;
-
-  ///< The BlockLoad algorithm to use
+  static constexpr int BLOCK_THREADS                 = ThreadsPerBlock;
+  static constexpr int ITEMS_PER_THREAD              = ItemsPerThread;
   static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm;
-
-  ///< Cache load modifier for reading input elements
-  static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier;
-
-  ///< The BlockScan algorithm to use
+  static constexpr CacheLoadModifier LOAD_MODIFIER   = LoadModifier;
   static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;
 
   struct detail
   {
     using delay_constructor_t = DelayConstructorT;
   };
 };
+} // namespace detail
+
+//! Deprecated [Since 3.5]
+template <int ThreadsPerBlock,
+          int ItemsPerThread,
+          BlockLoadAlgorithm LoadAlgorithm,
+          CacheLoadModifier LoadModifier,
+          BlockScanAlgorithm ScanAlgorithm,
+          typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>>
+using AgentReduceByKeyPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey") = detail::
+  agent_reduce_by_key_policy<ThreadsPerBlock, ItemsPerThread, LoadAlgorithm, LoadModifier, ScanAlgorithm, DelayConstructorT>;
 
 /******************************************************************************
  * Thread block abstractions

cub/cub/device/device_reduce.cuh

@@ -85,6 +85,25 @@ inline constexpr bool is_non_deterministic_v =
 //!
 //! @linear_performance{reduction, reduce-by-key, and run-length encode}
 //!
+//! Tuning
+//! +++++++++++++++++++++++++++++++++++++++++++++
+//!
+//! ``DeviceReduce::ReduceByKey`` can be tuned by passing a custom
+//! :ref:`policy selector <cub-policy-selectors>` that returns a @ref ReduceByKeyPolicy, as shown in the
+//! example below:
+//!
+//!  .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_by_key_env_api.cu
+//!      :language: c++
+//!      :dedent:
+//!      :start-after: example-begin reduce-by-key-policy-selector
+//!      :end-before: example-end reduce-by-key-policy-selector
+//!
+//!  .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_by_key_env_api.cu
+//!      :language: c++
+//!      :dedent:
+//!      :start-after: example-begin reduce-by-key-tuning
+//!      :end-before: example-end reduce-by-key-tuning
+//!
 //! @endrst
 struct DeviceReduce
 {

cub/cub/device/device_run_length_encode.cuh

@@ -59,6 +59,30 @@ CUB_NAMESPACE_BEGIN
 //! @endrst
 struct DeviceRunLengthEncode
 {
+#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
+  // DeviceRunLengthEncode::Encode dispatches to ReduceByKey, but we want to have a dedicated tuning policy, so we need
+  // to adapt the policy selector to convert the tuning policy
+  template <typename PolicySelector>
+#  if _CCCL_HAS_CONCEPTS()
+    requires detail::rle::encode::rle_encode_policy_selector<PolicySelector>
+#  endif // _CCCL_HAS_CONCEPTS()
+  struct __policy_selector_adapter
+  {
+    [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const
+      -> ReduceByKeyPolicy
+    {
+      const RleEncodePolicy policy = PolicySelector{}(cc);
+      return ReduceByKeyPolicy{
+        policy.threads_per_block,
+        policy.items_per_thread,
+        policy.load_algorithm,
+        policy.load_modifier,
+        policy.scan_algorithm,
+        policy.lookback_delay};
+    }
+  };
+#endif // _CCCL_DOXYGEN_INVOKED
+
   //! @rst
   //! Computes a run-length encoding of the sequence ``d_in``.
   //!
@@ -203,7 +227,7 @@ struct DeviceRunLengthEncode
       reduction_op{},
       static_cast<offset_t>(num_items),
       stream,
-      policy_selector_t{});
+      __policy_selector_adapter<policy_selector_t>{});
   }
 
   //! @rst
@@ -303,7 +327,7 @@ struct DeviceRunLengthEncode
     using default_policy_selector = detail::rle::encode::policy_selector_from_types<accum_t, key_t>;
 
     return detail::dispatch_with_env_and_tuning<default_policy_selector>(
-      env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) {
+      env, [&]([[maybe_unused]] auto policy_selector, void* storage, size_t& bytes, auto stream) {
         return detail::reduce_by_key::dispatch_streaming(
           storage,
           bytes,
@@ -316,7 +340,7 @@ struct DeviceRunLengthEncode
           reduction_op{},
           static_cast<offset_t>(num_items),
           stream,
-          policy_selector);
+          __policy_selector_adapter<decltype(policy_selector)>{});
       });
   }
 

cub/cub/device/dispatch/dispatch_reduce_by_key.cuh

@@ -203,16 +203,14 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
     _CCCL_GRID_CONSTANT const StreamingContextT streaming_context,
     vsmem_t vsmem)
 {
-  static constexpr reduce_by_key_policy policy = current_policy<PolicySelector>();
-  using AgentReduceByKeyPolicyT                = AgentReduceByKeyPolicy<
-                   policy.threads_per_block,
-                   policy.items_per_thread,
-                   policy.load_algorithm,
-                   policy.load_modifier,
-                   policy.scan_algorithm,
-                   delay_constructor_t<policy.delay_constructor.kind,
-                                       policy.delay_constructor.delay,
-                                       policy.delay_constructor.l2_write_latency>>;
+  static constexpr ReduceByKeyPolicy policy = current_policy<PolicySelector>();
+  using AgentReduceByKeyPolicyT             = agent_reduce_by_key_policy<
+                policy.threads_per_block,
+                policy.items_per_thread,
+                policy.load_algorithm,
+                policy.load_modifier,
+                policy.scan_algorithm,
+                delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
 
   using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<
     AgentReduceByKeyPolicyT,
@@ -292,7 +290,6 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
  *   Implementation detail, do not specify directly, requirements on the
  *   content of this type are subject to breaking change.
  */
-// TODO(bgruber): deprecate when we make the tuning API public and remove in CCCL 4.0
 template <typename KeysInputIteratorT,
           typename UniqueOutputIteratorT,
           typename ValuesInputIteratorT,
@@ -308,7 +305,7 @@ template <typename KeysInputIteratorT,
             ReductionOpT,
             AccumT,
             cub::detail::non_void_value_t<UniqueOutputIteratorT, cub::detail::it_value_t<KeysInputIteratorT>>>>
-struct DispatchReduceByKey
+struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey") DispatchReduceByKey
 {
   //-------------------------------------------------------------------------
   // Types and constants
@@ -650,16 +647,14 @@ template <typename PolicyGetter, typename... Args>
 _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_getter)
 {
   // TODO(bgruber): refactor this in the future
-  constexpr reduce_by_key_policy policy = policy_getter();
-  using Policy                          = AgentReduceByKeyPolicy<
-                             policy.threads_per_block,
-                             policy.items_per_thread,
-                             policy.load_algorithm,
-                             policy.load_modifier,
-                             policy.scan_algorithm,
-                             delay_constructor_t<policy.delay_constructor.kind,
-                                                 policy.delay_constructor.delay,
-                                                 policy.delay_constructor.l2_write_latency>>;
+  constexpr ReduceByKeyPolicy policy = policy_getter();
+  using Policy                       = agent_reduce_by_key_policy<
+                          policy.threads_per_block,
+                          policy.items_per_thread,
+                          policy.load_algorithm,
+                          policy.load_modifier,
+                          policy.scan_algorithm,
+                          delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
   using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<Policy, AgentReduceByKey, Args...>;
   return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS,
                             vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD,

cub/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh

@@ -68,7 +68,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_streaming(
     return error;
   }
 
-  const reduce_by_key_policy policy = policy_selector(cc);
+  const ReduceByKeyPolicy policy = policy_selector(cc);
 
 #if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
   NV_IF_TARGET(NV_IS_HOST, ({