Allow public tuning of cub::DeviceMergeSort (#8600)

bernhardmgruber · web-flow · commit 5a9ea633bfe6 · 2026-06-03T13:02:33.000+02:00
Fixes: #8574
diff --git a/cub/benchmarks/bench/merge_sort/keys.cu b/cub/benchmarks/bench/merge_sort/keys.cu
@@ -19,10 +19,9 @@
 template <typename KeyT>
 struct policy_selector
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-    -> cub::detail::merge_sort::merge_sort_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::MergeSortPolicy
   {
-    return cub::detail::merge_sort::merge_sort_policy{
+    return cub::MergeSortPolicy{
       TUNE_THREADS_PER_BLOCK,
       cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
       (TUNE_TRANSPOSE == 0 ? cub::BLOCK_LOAD_DIRECT : cub::BLOCK_LOAD_WARP_TRANSPOSE),
diff --git a/cub/benchmarks/bench/merge_sort/pairs.cu b/cub/benchmarks/bench/merge_sort/pairs.cu
@@ -19,10 +19,9 @@
 template <typename KeyT>
 struct policy_selector
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-    -> cub::detail::merge_sort::merge_sort_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::MergeSortPolicy
   {
-    return cub::detail::merge_sort::merge_sort_policy{
+    return cub::MergeSortPolicy{
       TUNE_THREADS_PER_BLOCK,
       cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
       (TUNE_TRANSPOSE == 0 ? cub::BLOCK_LOAD_DIRECT : cub::BLOCK_LOAD_WARP_TRANSPOSE),
diff --git a/cub/cub/agent/agent_merge_sort.cuh b/cub/cub/agent/agent_merge_sort.cuh
@@ -43,10 +43,10 @@ struct AgentBlockSort
 
   static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v<ValueT, NullType>;
 
-  static constexpr merge_sort_policy policy = PolicyGetter{}();
-  static constexpr int BLOCK_THREADS        = policy.threads_per_block;
-  static constexpr int ITEMS_PER_THREAD     = policy.items_per_thread;
-  static constexpr int ITEMS_PER_TILE       = policy.items_per_tile();
+  static constexpr MergeSortPolicy policy = PolicyGetter{}();
+  static constexpr int BLOCK_THREADS      = policy.threads_per_block;
+  static constexpr int ITEMS_PER_THREAD   = policy.items_per_thread;
+  static constexpr int ITEMS_PER_TILE     = BLOCK_THREADS * ITEMS_PER_THREAD;
 
   using BlockMergeSortT = BlockMergeSort<KeyT, BLOCK_THREADS, ITEMS_PER_THREAD, ValueT>;
 
@@ -378,10 +378,10 @@ struct AgentMerge
 
   static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v<ValueT, NullType>;
 
-  static constexpr merge_sort_policy policy = PolicyGetter{}();
-  static constexpr int BLOCK_THREADS        = policy.threads_per_block;
-  static constexpr int ITEMS_PER_THREAD     = policy.items_per_thread;
-  static constexpr int ITEMS_PER_TILE       = policy.items_per_tile();
+  static constexpr MergeSortPolicy policy = PolicyGetter{}();
+  static constexpr int BLOCK_THREADS      = policy.threads_per_block;
+  static constexpr int ITEMS_PER_THREAD   = policy.items_per_thread;
+  static constexpr int ITEMS_PER_TILE     = BLOCK_THREADS * ITEMS_PER_THREAD;
 
   using KeysLoadPingIt  = try_make_cache_modified_iterator_t<policy.load_modifier, KeyIteratorT>;
   using ItemsLoadPingIt = try_make_cache_modified_iterator_t<policy.load_modifier, ValueIteratorT>;
@@ -412,8 +412,8 @@ struct AgentMerge
     typename BlockStoreKeysPong::TempStorage store_keys_pong;
     typename BlockStoreItemsPong::TempStorage store_items_pong;
 
-    KeyT keys_shared[policy.items_per_tile() + 1];
-    ValueT items_shared[policy.items_per_tile() + 1];
+    KeyT keys_shared[ITEMS_PER_TILE + 1];
+    ValueT items_shared[ITEMS_PER_TILE + 1];
   };
 
   /// Alias wrapper allowing storage to be unioned
diff --git a/cub/cub/device/device_merge_sort.cuh b/cub/cub/device/device_merge_sort.cuh
@@ -86,6 +86,24 @@ CUB_NAMESPACE_BEGIN
  *   CustomLess());
  * @endcode
  *
+ * @par Tuning
+ * @rst
+ * All algorithms in DeviceMergeSort that accept an environment can be tuned by passing a custom :ref:`policy selector
+ * <cub-policy-selectors>` that returns a @ref MergeSortPolicy, as shown in the example below:
+ *
+ *  .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu
+ *      :language: c++
+ *      :dedent:
+ *      :start-after: example-begin sort-pairs-policy-selector
+ *      :end-before: example-end sort-pairs-policy-selector
+ *
+ *  .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu
+ *      :language: c++
+ *      :dedent:
+ *      :start-after: example-begin sort-pairs-tuning
+ *      :end-before: example-end sort-pairs-tuning
+ * @endrst
+ *
  * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable
  */
 struct DeviceMergeSort
diff --git a/cub/cub/device/dispatch/dispatch_merge_sort.cuh b/cub/cub/device/dispatch/dispatch_merge_sort.cuh
@@ -111,7 +111,7 @@ template <typename KeyInputIteratorT,
           typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY,
           typename KeyT                  = cub::detail::it_value_t<KeyIteratorT>,
           typename ValueT                = cub::detail::it_value_t<ValueIteratorT>>
-struct DispatchMergeSort
+struct CCCL_DEPRECATED_BECAUSE("Please use DeviceMergeSort") DispatchMergeSort
 {
   /// Whether or not there are values to be trucked along with keys
   static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v<ValueT, NullType>;
@@ -185,7 +185,7 @@ private:
   template <typename ActivePolicyT>
   struct policy_getter
   {
-    _CCCL_HOST_DEVICE_API constexpr auto operator()() -> detail::merge_sort::merge_sort_policy
+    _CCCL_HOST_DEVICE_API constexpr auto operator()() -> MergeSortPolicy
     {
       using mp = typename ActivePolicyT::MergeSortPolicy;
       return {mp::BLOCK_THREADS, mp::ITEMS_PER_THREAD, mp::LOAD_ALGORITHM, mp::LOAD_MODIFIER, mp::STORE_ALGORITHM};
@@ -195,7 +195,7 @@ private:
 public:
   // Invocation
   template <typename ActivePolicyT>
-  CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke([[maybe_unused]] ActivePolicyT policy = {})
+  CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke([[maybe_unused]] ActivePolicyT = {})
   {
     if (num_items == 0)
     {
@@ -206,18 +206,21 @@ public:
       return cudaSuccess;
     }
 
-    constexpr auto tile_size =
-      detail::merge_sort::merge_sort_vsmem_helper_t<
-        policy_getter<ActivePolicyT>,
-        KeyInputIteratorT,
-        ValueInputIteratorT,
-        KeyIteratorT,
-        ValueIteratorT,
-        OffsetT,
-        CompareOpT,
-        KeyT,
-        ValueT>::policy.items_per_tile();
-    const auto num_tiles = ::cuda::ceil_div(num_items, tile_size);
+    static constexpr auto policy = detail::merge_sort::merge_sort_vsmem_helper_t<
+      policy_getter<ActivePolicyT>,
+      KeyInputIteratorT,
+      ValueInputIteratorT,
+      KeyIteratorT,
+      ValueIteratorT,
+      OffsetT,
+      CompareOpT,
+      KeyT,
+      ValueT>::policy;
+    static_assert(1 <= policy.threads_per_block && policy.threads_per_block <= 1024,
+                  "Number of threads per block need to be inside [1;1024]");
+    static_assert(1 <= policy.items_per_thread, "Number of items per thread needs to be at least 1");
+    constexpr auto tile_size = policy.threads_per_block * policy.items_per_thread;
+    const auto num_tiles     = ::cuda::ceil_div(num_items, tile_size);
 
     const auto merge_partitions_size         = static_cast<size_t>(1 + num_tiles) * sizeof(OffsetT);
     const auto temporary_keys_storage_size   = static_cast<size_t>(num_items * kernel_source.KeySize());
@@ -279,17 +282,7 @@ public:
     auto keys_buffer      = static_cast<KeyT*>(allocations[1]);
     auto items_buffer     = static_cast<ValueT*>(allocations[2]);
 
-    const int threads_per_block =
-      detail::merge_sort::merge_sort_vsmem_helper_t<
-        policy_getter<ActivePolicyT>,
-        KeyInputIteratorT,
-        ValueInputIteratorT,
-        KeyIteratorT,
-        ValueIteratorT,
-        OffsetT,
-        CompareOpT,
-        KeyT,
-        ValueT>::policy.threads_per_block;
+    const int threads_per_block = policy.threads_per_block;
 
     // Invoke DeviceMergeSortBlockSortKernel
     launcher_factory(
@@ -495,7 +488,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
 
   return detail::dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) -> cudaError_t {
 #ifdef CUB_DEFINE_RUNTIME_POLICIES
-    const merge_sort_policy active_policy = policy_getter();
+    const MergeSortPolicy active_policy = policy_getter();
 #else // CUB_DEFINE_RUNTIME_POLICIES
     using vsmem_adapted_agents = merge_sort_vsmem_helper_t<
       decltype(policy_getter),
@@ -507,7 +500,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
       CompareOpT,
       KeyT,
       ValueT>;
-  constexpr merge_sort_policy active_policy = vsmem_adapted_agents::policy;
+  constexpr MergeSortPolicy active_policy = vsmem_adapted_agents::policy;
 #endif // CUB_DEFINE_RUNTIME_POLICIES
 
 #if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
@@ -521,7 +514,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
                  }))
 #endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
 
-    const auto tile_size = active_policy.items_per_tile();
+    _CCCL_ASSERT(1 <= active_policy.threads_per_block && active_policy.threads_per_block <= 1024,
+                 "Number of threads per block need to be inside [1;1024]");
+    _CCCL_ASSERT(1 <= active_policy.items_per_thread, "Number of items per thread needs to be at least 1");
+    const auto tile_size = active_policy.threads_per_block * active_policy.items_per_thread;
     const auto num_tiles = ::cuda::ceil_div(num_items, tile_size);
 
     const auto merge_partitions_size         = static_cast<size_t>(1 + num_tiles) * sizeof(OffsetT);
diff --git a/cub/cub/device/dispatch/kernels/kernel_merge_sort.cuh b/cub/cub/device/dispatch/kernels/kernel_merge_sort.cuh
@@ -30,7 +30,7 @@ struct fallback_policy_getter
   _CCCL_EXEC_CHECK_DISABLE
   _CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const
   {
-    merge_sort_policy policy = DefaultPolicyGetter{}();
+    MergeSortPolicy policy   = DefaultPolicyGetter{}();
     policy.threads_per_block = 64;
     policy.items_per_thread  = 1;
     return policy;
@@ -75,7 +75,7 @@ class merge_sort_vsmem_helper_impl
     (max_default_size > max_smem_per_block) && (max_fallback_size <= max_smem_per_block);
 
 public:
-  static constexpr merge_sort_policy policy = uses_fallback_policy ? FallbackPolicyGetter{}() : DefaultPolicyGetter{}();
+  static constexpr MergeSortPolicy policy = uses_fallback_policy ? FallbackPolicyGetter{}() : DefaultPolicyGetter{}();
   using block_sort_agent_t =
     ::cuda::std::_If<uses_fallback_policy, fallback_block_sort_agent_t, default_block_sort_agent_t>;
   using merge_agent_t = ::cuda::std::_If<uses_fallback_policy, fallback_merge_agent_t, default_merge_agent_t>;
@@ -166,9 +166,9 @@ __launch_bounds__(
     KeyT,
     ValueT>;
 
-  static constexpr merge_sort_policy active_policy = vsmem_adapted_agents::policy;
-  using agent_block_sort_t                         = typename vsmem_adapted_agents::block_sort_agent_t;
-  using vsmem_helper_t                             = vsmem_helper_impl<agent_block_sort_t>;
+  static constexpr MergeSortPolicy active_policy = vsmem_adapted_agents::policy;
+  using agent_block_sort_t                       = typename vsmem_adapted_agents::block_sort_agent_t;
+  using vsmem_helper_t                           = vsmem_helper_impl<agent_block_sort_t>;
 
   // Static shared memory allocation
   __shared__ typename vsmem_helper_t::static_temp_storage_t static_temp_storage;
@@ -267,9 +267,9 @@ __launch_bounds__(
     KeyT,
     ValueT>;
 
-  static constexpr merge_sort_policy active_policy = vsmem_adapted_agents::policy;
-  using agent_merge_t                              = typename vsmem_adapted_agents::merge_agent_t;
-  using vsmem_helper_t                             = vsmem_helper_impl<agent_merge_t>;
+  static constexpr MergeSortPolicy active_policy = vsmem_adapted_agents::policy;
+  using agent_merge_t                            = typename vsmem_adapted_agents::merge_agent_t;
+  using vsmem_helper_t                           = vsmem_helper_impl<agent_merge_t>;
 
   // Static shared memory allocation
   __shared__ typename vsmem_helper_t::static_temp_storage_t static_temp_storage;
diff --git a/cub/cub/device/dispatch/tuning/tuning_merge_sort.cuh b/cub/cub/device/dispatch/tuning/tuning_merge_sort.cuh
@@ -41,6 +41,39 @@ struct AgentMergeSortPolicy
   static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = StoreAlgorithm;
 };
 
+//! The tuning policy for all algorithms in @ref DeviceMergeSort.
+struct MergeSortPolicy
+{
+  int threads_per_block; //!< Number of threads in a CUDA block
+  int items_per_thread; //!< Number of items processed per thread
+  BlockLoadAlgorithm load_algorithm; //!< The @ref BlockLoadAlgorithm used for loading items from global memory
+  CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory
+  BlockStoreAlgorithm store_algorithm; //!< The @ref BlockStoreAlgorithm used for storing items to global memory
+
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
+  operator==(const MergeSortPolicy& lhs, const MergeSortPolicy& rhs)
+  {
+    return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
+        && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier
+        && lhs.store_algorithm == rhs.store_algorithm;
+  }
+
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
+  operator!=(const MergeSortPolicy& lhs, const MergeSortPolicy& rhs)
+  {
+    return !(lhs == rhs);
+  }
+
+#if _CCCL_HOSTED()
+  friend ::std::ostream& operator<<(::std::ostream& os, const MergeSortPolicy& p)
+  {
+    return os << "MergeSortPolicy { .threads_per_block = " << p.threads_per_block
+              << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
+              << ", .load_modifier = " << p.load_modifier << ", .store_algorithm = " << p.store_algorithm << " }";
+  }
+#endif // _CCCL_HOSTED()
+};
+
 namespace detail::merge_sort
 {
 // TODO(bgruber): drop in CCCL 4.0 when we remove all public CUB dispatchers
@@ -87,56 +120,19 @@ struct policy_hub
   using MaxPolicy = Policy600;
 };
 
-struct merge_sort_policy
-{
-  int threads_per_block;
-  int items_per_thread;
-  BlockLoadAlgorithm load_algorithm;
-  CacheLoadModifier load_modifier;
-  BlockStoreAlgorithm store_algorithm;
-
-  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int items_per_tile() const
-  {
-    return threads_per_block * items_per_thread;
-  }
-
-  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
-  operator==(const merge_sort_policy& lhs, const merge_sort_policy& rhs)
-  {
-    return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
-        && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier
-        && lhs.store_algorithm == rhs.store_algorithm;
-  }
-
-  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
-  operator!=(const merge_sort_policy& lhs, const merge_sort_policy& rhs)
-  {
-    return !(lhs == rhs);
-  }
-
-#if _CCCL_HOSTED()
-  friend ::std::ostream& operator<<(::std::ostream& os, const merge_sort_policy& p)
-  {
-    return os << "merge_sort_policy { .threads_per_block = " << p.threads_per_block
-              << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
-              << ", .load_modifier = " << p.load_modifier << ", .store_algorithm = " << p.store_algorithm << " }";
-  }
-#endif // _CCCL_HOSTED()
-};
-
 #if _CCCL_HAS_CONCEPTS()
 template <typename T>
-concept merge_sort_policy_selector = policy_selector<T, merge_sort_policy>;
+concept merge_sort_policy_selector = policy_selector<T, MergeSortPolicy>;
 #endif // _CCCL_HAS_CONCEPTS()
 
 struct policy_selector
 {
   int key_size;
 
-  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> merge_sort_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> MergeSortPolicy
   {
     // from SM60
-    return merge_sort_policy{
+    return MergeSortPolicy{
       256,
       detail::nominal_4B_items_to_items(17, key_size),
       BLOCK_LOAD_WARP_TRANSPOSE,
@@ -152,8 +148,7 @@ static_assert(merge_sort_policy_selector<policy_selector>);
 template <typename KeyIteratorT>
 struct policy_selector_from_types
 {
-  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const
-    -> merge_sort_policy
+  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> MergeSortPolicy
   {
     return policy_selector{int{sizeof(it_value_t<KeyIteratorT>)}}(cc);
   }
@@ -164,11 +159,11 @@ template <typename PolicyHub>
 struct policy_selector_from_hub
 {
   // this is only called in device code, so we can ignore the cc parameter
-  _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const -> merge_sort_policy
+  _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const -> MergeSortPolicy
   {
     using ap = typename PolicyHub::MaxPolicy::ActivePolicy;
     using mp = typename ap::MergeSortPolicy;
-    return merge_sort_policy{
+    return MergeSortPolicy{
       mp::BLOCK_THREADS, mp::ITEMS_PER_THREAD, mp::LOAD_ALGORITHM, mp::LOAD_MODIFIER, mp::STORE_ALGORITHM};
   }
 };
diff --git a/cub/test/catch2_test_device_merge_sort_custom_policy_hub.cu b/cub/test/catch2_test_device_merge_sort_custom_policy_hub.cu
@@ -1,6 +1,11 @@
 // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+// TODO(bgruber): drop this test with CCCL 4.0 when we drop the merge sort dispatcher
+
+// disable deprecation warnings for DispatchMergeSort
+#define CCCL_IGNORE_DEPRECATED_API
+
 #include "insert_nested_NVTX_range_guard.h"
 
 #include <cub/device/device_merge_sort.cuh>
@@ -14,8 +19,6 @@
 
 using namespace cub;
 
-// TODO(bgruber): drop this test with CCCL 4.0 when we drop the merge sort dispatcher after publishing the tuning API
-
 template <typename KeyIteratorT>
 struct my_policy_hub
 {
diff --git a/cub/test/catch2_test_device_merge_sort_env.cu b/cub/test/catch2_test_device_merge_sort_env.cu
diff --git a/cub/test/catch2_test_device_merge_sort_env_api.cu b/cub/test/catch2_test_device_merge_sort_env_api.cu

Original file line number	Diff line number	Diff line change
`@@ -19,10 +19,9 @@`
`19`	`19`	`template <typename KeyT>`
`20`	`20`	`struct policy_selector`
`21`	`21`	`{`
`22`		`- [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const`
`23`		`- -> cub::detail::merge_sort::merge_sort_policy`
	`22`	`+ [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::MergeSortPolicy`
`24`	`23`	`{`
`25`		`- return cub::detail::merge_sort::merge_sort_policy{`
	`24`	`+ return cub::MergeSortPolicy{`
`26`	`25`	`TUNE_THREADS_PER_BLOCK,`
`27`	`26`	`cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),`
`28`	`27`	`(TUNE_TRANSPOSE == 0 ? cub::BLOCK_LOAD_DIRECT : cub::BLOCK_LOAD_WARP_TRANSPOSE),`