Replace detail::merge::dispatch by CUB's public API

This includes using the new tuning API internally where possible.

Fixes: #7955

Author: bernhardmgruber

cub/benchmarks/bench/merge/keys.cu

@@ -18,24 +18,10 @@
 // %RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1
 // %RANGE% TUNE_THREADS_PER_BLOCK_POW2 tpb 6:10:1
 
-#if !TUNE_BASE
-struct bench_policy_selector
-{
-  _CCCL_API constexpr auto operator()(::cuda::arch_id /*arch*/) const -> cub::detail::merge::merge_policy
-  {
-    return cub::detail::merge::merge_policy{
-      TUNE_THREADS_PER_BLOCK,
-      cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
-      TUNE_LOAD_MODIFIER,
-      TUNE_STORE_ALGORITHM,
-      TUNE_USE_BL2SH};
-  }
-};
-#endif // !TUNE_BASE
-
-template <typename KeyT, typename OffsetT>
-void keys(nvbench::state& state, nvbench::type_list<KeyT, OffsetT>)
+template <typename KeyT>
+void keys(nvbench::state& state, nvbench::type_list<KeyT>)
 {
+  using offset_t     = int64_t;
   using compare_op_t = less_t;
 
   // Retrieve axis parameters
@@ -46,7 +32,7 @@ void keys(nvbench::state& state, nvbench::type_list<KeyT, OffsetT>)
   const auto num_items_rhs  = elements - num_items_lhs;
   auto [keys_lhs, keys_rhs] = generate_lhs_rhs<KeyT>(num_items_lhs, num_items_rhs, entropy);
 
-  thrust::device_vector<KeyT> keys_out(elements);
+  thrust::device_vector<KeyT> keys_out(elements, thrust::no_init);
   KeyT* d_keys_lhs = thrust::raw_pointer_cast(keys_lhs.data());
   KeyT* d_keys_rhs = thrust::raw_pointer_cast(keys_rhs.data());
   KeyT* d_keys_out = thrust::raw_pointer_cast(keys_out.data());
@@ -56,51 +42,26 @@ void keys(nvbench::state& state, nvbench::type_list<KeyT, OffsetT>)
   state.add_global_memory_reads<KeyT>(elements);
   state.add_global_memory_writes<KeyT>(elements);
 
-  auto value_nullptr = static_cast<cub::NullType*>(nullptr);
-
-  // Allocate temporary storage:
-  std::size_t temp_size{};
-  cub::detail::merge::dispatch(
-    nullptr,
-    temp_size,
-    d_keys_lhs,
-    value_nullptr,
-    static_cast<OffsetT>(num_items_lhs),
-    d_keys_rhs,
-    value_nullptr,
-    static_cast<OffsetT>(num_items_rhs),
-    d_keys_out,
-    value_nullptr,
-    compare_op_t{},
-    cudaStream_t{}
+  caching_allocator_t alloc;
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cub_bench_env(
+      alloc,
+      launch
 #if !TUNE_BASE
-    ,
-    bench_policy_selector{}
+      ,
+      cuda::execution::__tune(policy_selector<key_t, value_t, offset_t>{})
 #endif // !TUNE_BASE
-  );
-
-  thrust::device_vector<nvbench::uint8_t> temp(temp_size);
-  auto* temp_storage = thrust::raw_pointer_cast(temp.data());
-
-  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
-    cub::detail::merge::dispatch(
-      temp_storage,
-      temp_size,
+    );
+    _CCCL_TRY_CUDA_API(
+      cub::DeviceMerge::MergeKeys,
+      "MergePairs failed",
       d_keys_lhs,
-      value_nullptr,
-      static_cast<OffsetT>(num_items_lhs),
+      static_cast<offset_t>(num_items_lhs),
       d_keys_rhs,
-      value_nullptr,
-      static_cast<OffsetT>(num_items_rhs),
+      static_cast<offset_t>(num_items_rhs),
       d_keys_out,
-      value_nullptr,
       compare_op_t{},
-      launch.get_stream()
-#if !TUNE_BASE
-        ,
-      bench_policy_selector{}
-#endif // !TUNE_BASE
-    );
+      env);
   });
 }
 
@@ -110,8 +71,8 @@ using key_types = nvbench::type_list<TUNE_KeyT>;
 using key_types = fundamental_types;
 #endif // TUNE_KeyT
 
-NVBENCH_BENCH_TYPES(keys, NVBENCH_TYPE_AXES(key_types, offset_types))
+NVBENCH_BENCH_TYPES(keys, NVBENCH_TYPE_AXES(key_types))
   .set_name("base")
-  .set_type_axes_names({"KeyT{ct}", "OffsetT{ct}"})
+  .set_type_axes_names({"KeyT{ct}"})
   .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
   .add_string_axis("Entropy", {"1.000", "0.201"});

cub/benchmarks/bench/merge/merge_common.cuh

@@ -13,13 +13,6 @@
 #include <nvbench_helper.cuh>
 
 #if !TUNE_BASE
-#  define TUNE_THREADS_PER_BLOCK (1 << TUNE_THREADS_PER_BLOCK_POW2)
-#  if TUNE_TRANSPOSE == 0
-#    define TUNE_STORE_ALGORITHM cub::BLOCK_STORE_DIRECT
-#  else // TUNE_TRANSPOSE == 1
-#    define TUNE_STORE_ALGORITHM cub::BLOCK_STORE_WARP_TRANSPOSE
-#  endif // TUNE_TRANSPOSE
-
 #  if TUNE_LOAD == 0
 #    define TUNE_LOAD_MODIFIER cub::LOAD_DEFAULT
 #    define TUNE_USE_BL2SH     false
@@ -35,19 +28,17 @@
 #  endif // TUNE_LOAD
 
 template <typename KeyT>
-struct policy_hub_t
+struct bench_policy_selector
 {
-  struct policy_t : cub::ChainedPolicy<500, policy_t, policy_t>
+  _CCCL_API constexpr auto operator()(::cuda::arch_id /*arch*/) const -> cub::detail::merge::merge_policy
   {
-    using merge_policy =
-      cub::agent_policy_t<TUNE_THREADS_PER_BLOCK,
-                          cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
-                          TUNE_LOAD_MODIFIER,
-                          TUNE_STORE_ALGORITHM,
-                          TUNE_USE_BL2SH>;
-  };
-
-  using MaxPolicy = policy_t;
+    return cub::detail::merge::merge_policy{
+      (1 << TUNE_THREADS_PER_BLOCK_POW2),
+      cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
+      TUNE_LOAD_MODIFIER,
+      TUNE_TRANSPOSE == 0 ? cub::BLOCK_STORE_DIRECT : cub::BLOCK_STORE_WARP_TRANSPOSE,
+      TUNE_USE_BL2SH};
+  }
 };
 #endif // TUNE_BASE
 

cub/benchmarks/bench/merge/pairs.cu

@@ -18,24 +18,10 @@
 // %RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1
 // %RANGE% TUNE_THREADS_PER_BLOCK_POW2 tpb 6:10:1
 
-#if !TUNE_BASE
-struct bench_policy_selector
-{
-  _CCCL_API constexpr auto operator()(::cuda::arch_id /*arch*/) const -> cub::detail::merge::merge_policy
-  {
-    return cub::detail::merge::merge_policy{
-      TUNE_THREADS_PER_BLOCK,
-      cub::Nominal4BItemsToItems<KeyT>(TUNE_ITEMS_PER_THREAD),
-      TUNE_LOAD_MODIFIER,
-      TUNE_STORE_ALGORITHM,
-      TUNE_USE_BL2SH};
-  }
-};
-#endif // !TUNE_BASE
-
-template <typename KeyT, typename ValueT, typename OffsetT>
-void pairs(nvbench::state& state, nvbench::type_list<KeyT, ValueT, OffsetT>)
+template <typename KeyT, typename ValueT>
+void pairs(nvbench::state& state, nvbench::type_list<KeyT, ValueT>)
 {
+  using offset_t     = int64_t;
   using compare_op_t = less_t;
 
   // Retrieve axis parameters
@@ -45,10 +31,10 @@ void pairs(nvbench::state& state, nvbench::type_list<KeyT, ValueT, OffsetT>)
   const auto num_items_lhs = elements / 2;
   const auto num_items_rhs = elements - num_items_lhs;
 
-  thrust::device_vector<KeyT> keys_out(elements);
-  thrust::device_vector<ValueT> values_lhs(num_items_lhs);
-  thrust::device_vector<ValueT> values_rhs(num_items_rhs);
-  thrust::device_vector<ValueT> values_out(elements);
+  thrust::device_vector<KeyT> keys_out(elements, thrust::no_init);
+  thrust::device_vector<ValueT> values_lhs(num_items_lhs, thrust::no_init);
+  thrust::device_vector<ValueT> values_rhs(num_items_rhs, thrust::no_init);
+  thrust::device_vector<ValueT> values_out(elements, thrust::no_init);
 
   auto [keys_lhs, keys_rhs] = generate_lhs_rhs<KeyT>(num_items_lhs, num_items_rhs, entropy);
 
@@ -66,49 +52,29 @@ void pairs(nvbench::state& state, nvbench::type_list<KeyT, ValueT, OffsetT>)
   state.add_global_memory_writes<KeyT>(elements);
   state.add_global_memory_writes<ValueT>(elements);
 
-  // Allocate temporary storage:
-  std::size_t temp_size{};
-  cub::detail::merge::dispatch(
-    nullptr,
-    temp_size,
-    d_keys_lhs,
-    d_values_lhs,
-    static_cast<OffsetT>(num_items_lhs),
-    d_keys_rhs,
-    d_values_rhs,
-    static_cast<OffsetT>(num_items_rhs),
-    d_keys_out,
-    d_values_out,
-    compare_op_t{},
-    cudaStream_t{}
+  caching_allocator_t alloc;
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cub_bench_env(
+      alloc,
+      launch
 #if !TUNE_BASE
-    ,
-    bench_policy_selector{}
+      ,
+      cuda::execution::__tune(policy_selector<key_t, value_t, offset_t>{})
 #endif // !TUNE_BASE
-  );
-
-  thrust::device_vector<nvbench::uint8_t> temp(temp_size);
-  auto* temp_storage = thrust::raw_pointer_cast(temp.data());
-
-  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
-    cub::detail::merge::dispatch(
-      temp_storage,
-      temp_size,
+    );
+    _CCCL_TRY_CUDA_API(
+      cub::DeviceMerge::MergePairs,
+      "MergePairs failed",
       d_keys_lhs,
       d_values_lhs,
-      static_cast<OffsetT>(num_items_lhs),
+      static_cast<offset_t>(num_items_lhs),
       d_keys_rhs,
       d_values_rhs,
-      static_cast<OffsetT>(num_items_rhs),
+      static_cast<offset_t>(num_items_rhs),
       d_keys_out,
       d_values_out,
       compare_op_t{},
-      launch.get_stream()
-#if !TUNE_BASE
-        ,
-      bench_policy_selector{}
-#endif // !TUNE_BASE
-    );
+      env);
   });
 }
 
@@ -130,8 +96,8 @@ using value_types = nvbench::type_list<int8_t, int16_t, int32_t, int64_t
                                        >;
 #endif // TUNE_ValueT
 
-NVBENCH_BENCH_TYPES(pairs, NVBENCH_TYPE_AXES(key_types, value_types, offset_types))
+NVBENCH_BENCH_TYPES(pairs, NVBENCH_TYPE_AXES(key_types, value_types))
   .set_name("base")
-  .set_type_axes_names({"KeyT{ct}", "ValueT{ct}", "OffsetT{ct}"})
+  .set_type_axes_names({"KeyT{ct}", "ValueT{ct}"})
   .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
   .add_string_axis("Entropy", {"1.000", "0.201"});

cub/cub/detail/env_dispatch.cuh

@@ -74,6 +74,18 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(EnvT env, AlgorithmCal
   return (error != cudaSuccess) ? error : deallocate_error;
 }
 //! @endcond
+
+template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
+CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env_and_tuning(EnvT env, AlgorithmCallable&& algorithm_callable)
+{
+  return detail::dispatch_with_env(
+    env, [&]([[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
+      using policy_t = decltype(DefaultPolicySelector{}(::cuda::arch_id{}));
+      using policy_selector =
+        ::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
+      return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
+    });
+}
 } // namespace detail
 
 CUB_NAMESPACE_END

cub/cub/device/device_merge.cuh

@@ -192,8 +192,10 @@ struct DeviceMerge
   {
     _CCCL_NVTX_RANGE_SCOPE("cub::DeviceMerge::MergeKeys");
 
-    return detail::dispatch_with_env(
-      env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
+    using default_policy_selector =
+      detail::merge::policy_selector_from_types<detail::it_value_t<KeyIteratorIn1>, NullType, int64_t>;
+    return detail::dispatch_with_env_and_tuning<default_policy_selector>(
+      env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
         return detail::merge::dispatch(
           d_temp_storage,
           temp_storage_bytes,
@@ -206,7 +208,8 @@ struct DeviceMerge
           keys_out,
           static_cast<NullType*>(nullptr),
           compare_op,
-          stream);
+          stream,
+          policy_selector);
       });
   }
 
@@ -413,9 +416,10 @@ struct DeviceMerge
     EnvT env             = {})
   {
     _CCCL_NVTX_RANGE_SCOPE("cub::DeviceMerge::MergePairs");
-
-    return detail::dispatch_with_env(
-      env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
+    using default_policy_selector = detail::merge::
+      policy_selector_from_types<detail::it_value_t<KeyIteratorIn1>, detail::it_value_t<ValueIteratorIn1>, int64_t>;
+    return detail::dispatch_with_env_and_tuning<default_policy_selector>(
+      env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
         return detail::merge::dispatch(
           d_temp_storage,
           temp_storage_bytes,
@@ -428,7 +432,8 @@ struct DeviceMerge
           keys_out,
           values_out,
           compare_op,
-          stream);
+          stream,
+          policy_selector);
       });
   }
 };