Implement the new tuning API for DeviceAdjacentDifference part 2 (#7918)

The work in #7524 was incomplete. Here is the rest.

Fixes: #7522

Co-authored-by: Nader Al Awar <naderalawar@gmail.com>

Author: bernhardmgruber

cub/benchmarks/bench/adjacent_difference/subtract_left.cu

@@ -9,19 +9,17 @@
 // %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
 
 #if !TUNE_BASE
-struct policy_hub_t
+struct policy_selector_t
 {
-  struct Policy500 : cub::ChainedPolicy<500, Policy500, Policy500>
+  _CCCL_API constexpr auto operator()(::cuda::arch_id) const
+    -> cub::detail::adjacent_difference::adjacent_difference_policy
   {
-    using AdjacentDifferencePolicy =
-      cub::AgentAdjacentDifferencePolicy<TUNE_THREADS_PER_BLOCK,
-                                         TUNE_ITEMS_PER_THREAD,
-                                         cub::BLOCK_LOAD_WARP_TRANSPOSE,
-                                         cub::LOAD_CA,
-                                         cub::BLOCK_STORE_WARP_TRANSPOSE>;
-  };
-
-  using MaxPolicy = Policy500;
+    return {TUNE_THREADS_PER_BLOCK,
+            TUNE_ITEMS_PER_THREAD,
+            cub::BLOCK_LOAD_WARP_TRANSPOSE,
+            cub::LOAD_CA,
+            cub::BLOCK_STORE_WARP_TRANSPOSE};
+  }
 };
 #endif // !TUNE_BASE
 
@@ -33,25 +31,6 @@ void left(nvbench::state& state, nvbench::type_list<T, OffsetT>)
   using difference_op_t = ::cuda::std::minus<>;
   using offset_t        = cub::detail::choose_offset_t<OffsetT>;
 
-#if !TUNE_BASE
-  using dispatch_t = cub::DispatchAdjacentDifference<
-    input_it_t,
-    output_it_t,
-    difference_op_t,
-    offset_t,
-    cub::MayAlias::No,
-    cub::ReadOption::Left,
-    policy_hub_t>;
-#else
-  using dispatch_t = cub::DispatchAdjacentDifference<
-    input_it_t,
-    output_it_t,
-    difference_op_t,
-    offset_t,
-    cub::MayAlias::No,
-    cub::ReadOption::Left>;
-#endif // TUNE_BASE
-
   const auto elements         = static_cast<std::size_t>(state.get_int64("Elements{io}"));
   thrust::device_vector<T> in = generate(elements);
   thrust::device_vector<T> out(elements);
@@ -64,20 +43,40 @@ void left(nvbench::state& state, nvbench::type_list<T, OffsetT>)
   state.add_global_memory_writes<T>(elements);
 
   std::size_t temp_storage_bytes{};
-  dispatch_t::Dispatch(nullptr, temp_storage_bytes, d_in, d_out, static_cast<offset_t>(elements), difference_op_t{}, 0);
 
-  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes);
-  std::uint8_t* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
-
-  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
-    dispatch_t::Dispatch(
-      d_temp_storage,
+  cub::detail::adjacent_difference::
+    dispatch<cub::MayAlias::No, cub::ReadOption::Left, input_it_t, output_it_t, offset_t, difference_op_t>(
+      nullptr,
       temp_storage_bytes,
       d_in,
       d_out,
       static_cast<offset_t>(elements),
       difference_op_t{},
-      launch.get_stream());
+      0
+#if !TUNE_BASE
+      ,
+      policy_selector_t{}
+#endif // TUNE_BASE
+    );
+
+  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes, thrust::no_init);
+  std::uint8_t* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
+
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    cub::detail::adjacent_difference::
+      dispatch<cub::MayAlias::No, cub::ReadOption::Left, input_it_t, output_it_t, offset_t, difference_op_t>(
+        d_temp_storage,
+        temp_storage_bytes,
+        d_in,
+        d_out,
+        static_cast<offset_t>(elements),
+        difference_op_t{},
+        launch.get_stream()
+#if !TUNE_BASE
+          ,
+        policy_selector_t{}
+#endif // TUNE_BASE
+      );
   });
 }
 

cub/cub/device/device_adjacent_difference.cuh

@@ -84,32 +84,6 @@ CUB_NAMESPACE_BEGIN
 //! @endrst
 struct DeviceAdjacentDifference
 {
-private:
-  template <MayAlias AliasOpt,
-            ReadOption ReadOpt,
-            typename NumItemsT,
-            typename InputIteratorT,
-            typename OutputIteratorT,
-            typename DifferenceOpT>
-  static CUB_RUNTIME_FUNCTION cudaError_t AdjacentDifference(
-    void* d_temp_storage,
-    size_t& temp_storage_bytes,
-    InputIteratorT d_input,
-    OutputIteratorT d_output,
-    NumItemsT num_items,
-    DifferenceOpT difference_op,
-    cudaStream_t stream)
-  {
-    using OffsetT = detail::choose_offset_t<NumItemsT>;
-
-    using DispatchT =
-      DispatchAdjacentDifference<InputIteratorT, OutputIteratorT, DifferenceOpT, OffsetT, AliasOpt, ReadOpt>;
-
-    return DispatchT::Dispatch(
-      d_temp_storage, temp_storage_bytes, d_input, d_output, static_cast<OffsetT>(num_items), difference_op, stream);
-  }
-
-public:
   //! @rst
   //! Subtracts the left element of each adjacent pair of elements residing within device-accessible memory
   //!
@@ -234,9 +208,9 @@ public:
     cudaStream_t stream         = 0)
   {
     _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractLeftCopy");
-
-    return AdjacentDifference<MayAlias::No, ReadOption::Left>(
-      d_temp_storage, temp_storage_bytes, d_input, d_output, num_items, difference_op, stream);
+    using OffsetT = detail::choose_offset_t<NumItemsT>;
+    return detail::adjacent_difference::dispatch<MayAlias::No, ReadOption::Left>(
+      d_temp_storage, temp_storage_bytes, d_input, d_output, static_cast<OffsetT>(num_items), difference_op, stream);
   }
 
   //! @rst
@@ -345,9 +319,9 @@ public:
     cudaStream_t stream         = 0)
   {
     _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractLeft");
-
-    return AdjacentDifference<MayAlias::Yes, ReadOption::Left>(
-      d_temp_storage, temp_storage_bytes, d_input, d_input, num_items, difference_op, stream);
+    using OffsetT = detail::choose_offset_t<NumItemsT>;
+    return detail::adjacent_difference::dispatch<MayAlias::Yes, ReadOption::Left>(
+      d_temp_storage, temp_storage_bytes, d_input, d_input, static_cast<OffsetT>(num_items), difference_op, stream);
   }
 
   //! @rst
@@ -475,9 +449,9 @@ public:
     cudaStream_t stream         = 0)
   {
     _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractRightCopy");
-
-    return AdjacentDifference<MayAlias::No, ReadOption::Right>(
-      d_temp_storage, temp_storage_bytes, d_input, d_output, num_items, difference_op, stream);
+    using OffsetT = detail::choose_offset_t<NumItemsT>;
+    return detail::adjacent_difference::dispatch<MayAlias::No, ReadOption::Right>(
+      d_temp_storage, temp_storage_bytes, d_input, d_output, static_cast<OffsetT>(num_items), difference_op, stream);
   }
 
   //! @rst
@@ -575,9 +549,9 @@ public:
     cudaStream_t stream         = 0)
   {
     _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractRight");
-
-    return AdjacentDifference<MayAlias::Yes, ReadOption::Right>(
-      d_temp_storage, temp_storage_bytes, d_input, d_input, num_items, difference_op, stream);
+    using OffsetT = detail::choose_offset_t<NumItemsT>;
+    return detail::adjacent_difference::dispatch<MayAlias::Yes, ReadOption::Right>(
+      d_temp_storage, temp_storage_bytes, d_input, d_input, static_cast<OffsetT>(num_items), difference_op, stream);
   }
 };
 

cub/cub/device/dispatch/dispatch_adjacent_difference.cuh

@@ -313,12 +313,12 @@ struct DispatchAdjacentDifference
 
 namespace detail::adjacent_difference
 {
-template <typename InputIteratorT,
+template <MayAlias AliasOpt,
+          ReadOption ReadOpt,
+          typename InputIteratorT,
           typename OutputIteratorT,
-          typename DifferenceOpT,
           typename OffsetT,
-          MayAlias AliasOpt,
-          ReadOption ReadOpt,
+          typename DifferenceOpT,
           typename PolicySelector        = policy_selector_from_types<InputIteratorT, AliasOpt == MayAlias::Yes>,
           typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
 #if _CCCL_HAS_CONCEPTS()

thrust/thrust/system/cuda/detail/adjacent_difference.h

@@ -65,16 +65,10 @@ cudaError_t THRUST_RUNTIME_FUNCTION doit_step(
     return cudaSuccess;
   }
 
-  constexpr cub::ReadOption read_left = cub::ReadOption::Left;
-
-  using Dispatch32 = cub::DispatchAdjacentDifference<InputIt, OutputIt, BinaryOp, std::int32_t, AliasOpt, read_left>;
-  using Dispatch64 = cub::DispatchAdjacentDifference<InputIt, OutputIt, BinaryOp, std::int64_t, AliasOpt, read_left>;
-
   cudaError_t status;
-  THRUST_INDEX_TYPE_DISPATCH2(
+  THRUST_INDEX_TYPE_DISPATCH(
     status,
-    Dispatch32::Dispatch,
-    Dispatch64::Dispatch,
+    (cub::detail::adjacent_difference::dispatch<AliasOpt, cub::ReadOption::Left>),
     num_items,
     (d_temp_storage, temp_storage_bytes, first, result, num_items_fixed, binary_op, stream));
   return status;