Use PDL in DeviceHistogram

c/parallel/src/histogram.cu

@@ -75,6 +75,7 @@ struct histogram_kernel_source
 {
   cccl_device_histogram_build_result_t& build;
 
+  template <typename PolicyT>
   CUkernel HistogramInitKernel() const
   {
     return build.init_kernel;
@@ -102,10 +103,15 @@ histogram_runtime_tuning_policy get_policy(int /*cc*/, cccl_type_info sample_t,
   return {384, pixels_per_thread};
 }
 
-std::string get_init_kernel_name(int num_active_channels, std::string_view counter_t, std::string_view offset_t)
+std::string get_init_kernel_name(
+  std::string_view chained_policy_t, int num_active_channels, std::string_view counter_t, std::string_view offset_t)
 {
   return std::format(
-    "cub::detail::histogram::DeviceHistogramInitKernel<{0}, {1}, {2}>", num_active_channels, counter_t, offset_t);
+    "cub::detail::histogram::DeviceHistogramInitKernel<{0}, {1}, {2}, {3}>",
+    chained_policy_t,
+    num_active_channels,
+    counter_t,
+    offset_t);
 }
 
 std::string get_sweep_kernel_name(
@@ -225,6 +231,7 @@ struct agent_policy_t {{
 struct {5} {{
   struct ActivePolicy {{
     using AgentHistogramPolicyT = agent_policy_t;
+    static constexpr int pdl_trigger_next_launch_in_init_kernel_max_bin_count = 2048;
   }};
 }};
 )XXX";
@@ -253,7 +260,8 @@ struct {5} {{
 
     const bool is_byte_sample = d_samples.value_type.size == 1;
 
-    std::string init_kernel_name  = histogram::get_init_kernel_name(num_active_channels, counter_cpp, offset_cpp);
+    std::string init_kernel_name =
+      histogram::get_init_kernel_name(chained_policy_t, num_active_channels, counter_cpp, offset_cpp);
     std::string sweep_kernel_name = histogram::get_sweep_kernel_name(
       chained_policy_t,
       privatized_smem_bins,

cub/cub/agent/agent_histogram.cuh

@@ -23,6 +23,7 @@
 #include <cub/iterator/cache_modified_input_iterator.cuh>
 #include <cub/util_type.cuh>
 
+#include <cuda/std/__cccl/cuda_capabilities.h>
 #include <cuda/std/__type_traits/conditional.h>
 #include <cuda/std/__type_traits/integral_constant.h>
 #include <cuda/std/__type_traits/is_pointer.h>
@@ -628,6 +629,8 @@ struct AgentHistogram
       ((size_t(d_native_samples) & pixel_mask) == 0) && // ptr is pixel-aligned
       ((row_bytes & pixel_mask) == 0); // number of row-samples is a multiple of the alignment of the pixel
 
+    _CCCL_PDL_GRID_DEPENDENCY_SYNC();
+
     // Whether rows are aligned and can be vectorized
     if ((d_native_samples != nullptr) && (vec_aligned_rows || pixel_aligned_rows))
     {
@@ -639,6 +642,8 @@ struct AgentHistogram
       ConsumeTiles<false>(
         num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v<is_work_stealing>);
     }
+
+    _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); // omitting makes no difference in cub.bench.histogram.even.base
   }
 
   //! Initialize privatized bin counters.  Specialized for privatized shared-memory counters

cub/cub/device/dispatch/dispatch_histogram.cuh

@@ -54,10 +54,15 @@ namespace detail::histogram
 template <int NUM_CHANNELS, int NUM_ACTIVE_CHANNELS, typename SampleIteratorT, typename CounterT, typename OffsetT>
 struct DeviceHistogramKernelSource
 {
-  CUB_DEFINE_KERNEL_GETTER(HistogramInitKernel, DeviceHistogramInitKernel<NUM_ACTIVE_CHANNELS, CounterT, OffsetT>);
-
   // We define this differently than the other kernel getters because there are
   // different dispatch paths which affect which policy is used and the decode operators.
+
+  template <typename PolicyT>
+  _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramInitKernel()
+  {
+    return &DeviceHistogramInitKernel<PolicyT, NUM_ACTIVE_CHANNELS, CounterT, OffsetT>;
+  }
+
   template <typename PolicyT, int PRIVATIZED_SMEM_BINS, typename PrivatizedDecodeOpT, typename OutputDecodeOpT>
   _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramSweepKernel()
   {
@@ -227,7 +232,7 @@ struct dispatch_histogram
 #endif // CUB_DEBUG_LOG
 
       // Invoke histogram_init_kernel
-      launcher_factory(histogram_init_grid_dims, histogram_init_block_threads, 0, stream)
+      launcher_factory(histogram_init_grid_dims, histogram_init_block_threads, 0, stream, true)
         .doit(histogram_init_kernel, num_output_bins_wrapper, d_output_histograms, tile_queue);
 
       // Return if empty problem
@@ -250,7 +255,7 @@ struct dispatch_histogram
 #endif // CUB_DEBUG_LOG
 
       // Invoke histogram_sweep_kernel
-      launcher_factory(sweep_grid_dims, block_threads, 0, stream)
+      launcher_factory(sweep_grid_dims, block_threads, 0, stream, true)
         .doit(histogram_sweep_kernel,
               d_samples,
               num_output_bins_wrapper,
@@ -287,7 +292,7 @@ struct dispatch_histogram
   CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {})
   {
     return Invoke<ActivePolicyT>(
-      kernel_source.HistogramInitKernel(),
+      kernel_source.template HistogramInitKernel<MaxPolicyT>(),
       kernel_source
         .template HistogramSweepKernel<MaxPolicyT, PRIVATIZED_SMEM_BINS, PrivatizedDecodeOpT, OutputDecodeOpT>(),
       active_policy);

cub/cub/device/dispatch/kernels/histogram.cuh

@@ -16,6 +16,8 @@
 #include <cub/agent/agent_histogram.cuh>
 #include <cub/grid/grid_queue.cuh>
 
+#include <cuda/std/__numeric/reduce.h>
+
 CUB_NAMESPACE_BEGIN
 namespace detail::histogram
 {
@@ -328,12 +330,22 @@ struct Transforms
 //!
 //! @param tile_queue
 //!   Drain queue descriptor for dynamically mapping tile data onto thread blocks
-template <int NumActiveChannels, typename CounterT, typename OffsetT>
+template <typename ChainedPolicyT, int NumActiveChannels, typename CounterT, typename OffsetT>
 CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceHistogramInitKernel(
   ::cuda::std::array<int, NumActiveChannels> num_output_bins_wrapper,
   ::cuda::std::array<CounterT*, NumActiveChannels> d_output_histograms_wrapper,
   GridQueue<int> tile_queue)
 {
+  _CCCL_PDL_GRID_DEPENDENCY_SYNC(); // TODO(bgruber): if we had the guarantee that there would be no pending
+                                    // writes/reads to the temp storage, we could omit the sync here
+
+  // we trigger the sweep kernel only if we have a small number of remaining writes in this kernel
+  NV_IF_TARGET(NV_PROVIDES_SM_90,
+               (if (::cuda::std::reduce(num_output_bins_wrapper.begin(), num_output_bins_wrapper.end())
+                    <= ChainedPolicyT::ActivePolicy::pdl_trigger_next_launch_in_init_kernel_max_bin_count) {
+                 _CCCL_PDL_TRIGGER_NEXT_LAUNCH();
+               }));
+
   if ((threadIdx.x == 0) && (blockIdx.x == 0))
   {
     tile_queue.ResetDrain();

cub/cub/device/dispatch/tuning/tuning_histogram.cuh

@@ -219,6 +219,8 @@ struct policy_hub
     using AgentHistogramPolicyT =
       decltype(select_agent_policy<
                sm90_tuning<SampleT, NumChannels, NumActiveChannels, histogram::classify_counter_size<CounterT>()>>(0));
+
+    static constexpr int pdl_trigger_next_launch_in_init_kernel_max_bin_count = 2048;
   };
 
   struct Policy1000 : ChainedPolicy<1000, Policy1000, Policy900>
@@ -242,6 +244,8 @@ struct policy_hub
       decltype(select_agent_policy<
                sm100_tuning<IsEven, SampleT, NumChannels, NumActiveChannels, histogram::classify_counter_size<CounterT>()>>(
         0));
+
+    static constexpr int pdl_trigger_next_launch_in_init_kernel_max_bin_count = 2048;
   };
 
   using MaxPolicy = Policy1000;