[cub] Add static and runtime-sized shared-memory histogram privatization

c/parallel/src/histogram.cu

@@ -56,7 +56,7 @@ struct histogram_kernel_source
   }
 
   template <typename PolicyT,
-            int PRIVATIZED_SMEM_BINS,
+            typename PrivatizationMode,
             typename FirstLevelArrayT,
             typename SecondLevelArrayT,
             bool IsEven,
@@ -96,7 +96,7 @@ std::string get_init_kernel_name(int num_active_channels, std::string_view count
 }
 
 std::string get_sweep_kernel_name(
-  int privatized_smem_bins,
+  std::string_view privatization_mode_t,
   int num_channels,
   int num_active_channels,
   cccl_iterator_t d_samples,
@@ -147,7 +147,7 @@ std::string get_sweep_kernel_name(
     "cub::detail::histogram::DeviceHistogramSweepDeviceInitKernel<{0}, {1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}, {9}, "
     "{10}, {11}>",
     chained_policy_t,
-    privatized_smem_bins,
+    privatization_mode_t,
     num_channels,
     num_active_channels,
     samples_iterator_t,
@@ -324,17 +324,26 @@ static_assert(device_histogram_policy()(detail::current_tuning_cc()) == {4}, "Ho
   fflush(stdout);
 #endif
 
-  // TODO: This is tricky because we need to know the input to set this to a
-  // value greater than 0 (see dispatch_histogram.cuh), but we don't have this
-  // information here.
-  const int privatized_smem_bins =
-    num_output_levels_val - 1 > cub::detail::histogram::max_privatized_smem_bins ? 0 : 256;
-
   const bool is_byte_sample = d_samples.value_type.size == 1;
+  const int num_privatized_bins =
+    is_byte_sample ? cub::detail::histogram::byte_sample_privatized_levels - 1 : num_output_levels_val - 1;
+  const int counter_size_bytes = static_cast<int>(d_output_histograms.value_type.size);
+  const auto privatization =
+    is_evenly_segmented
+      ? cub::detail::histogram::select_privatization_mode_for_counter_size<true, 1>(
+          active_policy, num_privatized_bins, counter_size_bytes)
+      : cub::detail::histogram::select_privatization_mode_for_counter_size<false, 1>(
+          active_policy, num_privatized_bins, counter_size_bytes);
+  const std::string_view privatization_mode_t =
+    privatization == cub::detail::histogram::privatization_mode::static_smem
+      ? "cub::detail::histogram::HistogramPrivatizedStaticSmem"
+    : privatization == cub::detail::histogram::privatization_mode::dynamic_smem
+      ? "cub::detail::histogram::HistogramPrivatizedDynamicSmem"
+      : "cub::detail::histogram::HistogramPrivatizedGmem";
 
   std::string init_kernel_name  = histogram::get_init_kernel_name(num_active_channels, counter_cpp, offset_cpp);
   std::string sweep_kernel_name = histogram::get_sweep_kernel_name(
-    privatized_smem_bins,
+    privatization_mode_t,
     num_channels,
     num_active_channels,
     d_samples,
@@ -574,6 +583,7 @@ CUresult cccl_device_histogram_even_impl(
       indirect_arg_t, // LevelT
       OffsetT, // OffsetT
       cub::detail::histogram::policy_selector, // PolicySelector
+      void, // PrivatizedCounterT: C Parallel preserves the counter width after type erasure
       indirect_arg_t, // SampleT
       histogram::histogram_kernel_source, // KernelSource
       cub::detail::CudaDriverLauncherFactory // KernelLauncherFactory

cub/benchmarks/bench/histogram/even.cu

@@ -9,7 +9,6 @@
 // %RANGE% TUNE_THREADS tpb 128:1024:32
 // %RANGE% TUNE_RLE_COMPRESS rle 0:1:1
 // %RANGE% TUNE_WORK_STEALING ws 0:1:1
-// %RANGE% TUNE_MEM_PREFERENCE mem 0:2:1
 // %RANGE% TUNE_LOAD ld 0:2:1
 // %RANGE% TUNE_LOAD_ALGORITHM_ID laid 0:2:1
 // %RANGE% TUNE_VEC_SIZE_POW vec 0:2:1
@@ -52,7 +51,7 @@ static void even(nvbench::state& state, nvbench::type_list<SampleT, CounterT, Of
       launch
 #if !TUNE_BASE
       ,
-      cuda::execution::tune(bench_policy_selector<key_t, 1, 1>{})
+      cuda::execution::tune(histogram_tuning_policy_selector<SampleT, CounterT, 1, 1, true>{})
 #endif // !TUNE_BASE
     );
     _CCCL_TRY_CUDA_API(

cub/benchmarks/bench/histogram/histogram_common.cuh

@@ -19,14 +19,6 @@
 
 #  define TUNE_VEC_SIZE (1 << TUNE_VEC_SIZE_POW)
 
-#  if TUNE_MEM_PREFERENCE == 0
-constexpr cub::BlockHistogramMemoryPreference MEM_PREFERENCE = cub::GMEM;
-#  elif TUNE_MEM_PREFERENCE == 1
-constexpr cub::BlockHistogramMemoryPreference MEM_PREFERENCE = cub::SMEM;
-#  else // TUNE_MEM_PREFERENCE == 2
-constexpr cub::BlockHistogramMemoryPreference MEM_PREFERENCE = cub::BLEND;
-#  endif // TUNE_MEM_PREFERENCE
-
 #  if TUNE_LOAD_ALGORITHM_ID == 0
 #    define TUNE_LOAD_ALGORITHM cub::BLOCK_LOAD_DIRECT
 #  elif TUNE_LOAD_ALGORITHM_ID == 1
@@ -35,25 +27,38 @@ constexpr cub::BlockHistogramMemoryPreference MEM_PREFERENCE = cub::BLEND;
 #    define TUNE_LOAD_ALGORITHM cub::BLOCK_LOAD_STRIPED
 #  endif // TUNE_LOAD_ALGORITHM_ID
 
-template <typename SampleT, int NUM_CHANNELS, int NUM_ACTIVE_CHANNELS>
-struct bench_policy_selector
+// Only generated tuning variants instantiate this selector. The `.base` target used for
+// production-policy comparisons defines TUNE_BASE=1 and calls DeviceHistogram without a
+// tuning environment, so it exercises the shipping selector unchanged.
+//
+// A generated tuning point supplies one candidate kernel configuration. Apply that same
+// candidate to every privatization mode so the tuner measures the candidate independently
+// of the runtime bin count selected by the production storage thresholds.
+template <typename SampleT, typename CounterT, int NUM_CHANNELS, int NUM_ACTIVE_CHANNELS, bool IS_EVEN>
+struct histogram_tuning_policy_selector
 {
-  _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cub::HistogramPolicy
+  _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> cub::HistogramPolicy
   {
     constexpr cub::BlockLoadAlgorithm load_algorithm =
       (TUNE_LOAD_ALGORITHM == cub::BLOCK_LOAD_STRIPED)
         ? (NUM_CHANNELS == 1 ? cub::BLOCK_LOAD_STRIPED : cub::BLOCK_LOAD_DIRECT)
         : TUNE_LOAD_ALGORITHM;
 
-    return {TUNE_THREADS,
-            TUNE_ITEMS,
-            TUNE_VEC_SIZE,
-            load_algorithm,
-            TUNE_LOAD_MODIFIER,
-            TUNE_RLE_COMPRESS,
-            MEM_PREFERENCE,
-            TUNE_WORK_STEALING,
-            2048}; // TODO(bgruber): make tunable
+    constexpr auto sweep = cub::HistogramPrivatizationPolicy{
+      TUNE_THREADS,
+      TUNE_ITEMS,
+      TUNE_VEC_SIZE,
+      load_algorithm,
+      TUNE_LOAD_MODIFIER,
+      TUNE_RLE_COMPRESS,
+      TUNE_WORK_STEALING};
+    auto policy =
+      cub::detail::histogram::policy_selector_from_types<SampleT, CounterT, NUM_CHANNELS, NUM_ACTIVE_CHANNELS, IS_EVEN>{}(
+        cc);
+    policy.gmem         = sweep;
+    policy.static_smem  = sweep;
+    policy.dynamic_smem = sweep;
+    return policy;
   }
 };
 #endif // !TUNE_BASE

cub/benchmarks/bench/histogram/multi/even.cu

@@ -9,7 +9,6 @@
 // %RANGE% TUNE_THREADS tpb 128:1024:32
 // %RANGE% TUNE_RLE_COMPRESS rle 0:1:1
 // %RANGE% TUNE_WORK_STEALING ws 0:1:1
-// %RANGE% TUNE_MEM_PREFERENCE mem 0:2:1
 // %RANGE% TUNE_LOAD ld 0:2:1
 // %RANGE% TUNE_LOAD_ALGORITHM_ID laid 0:2:1
 // %RANGE% TUNE_VEC_SIZE_POW vec 0:2:1
@@ -65,7 +64,8 @@ static void even(nvbench::state& state, nvbench::type_list<SampleT, CounterT, Of
       launch
 #if !TUNE_BASE
       ,
-      cuda::execution::tune(bench_policy_selector<key_t, num_channels, num_active_channels>{})
+      cuda::execution::tune(
+        histogram_tuning_policy_selector<SampleT, CounterT, num_channels, num_active_channels, true>{})
 #endif // !TUNE_BASE
     );
     _CCCL_TRY_CUDA_API(

cub/benchmarks/bench/histogram/multi/range.cu

@@ -11,7 +11,6 @@
 // %RANGE% TUNE_THREADS tpb 128:1024:32
 // %RANGE% TUNE_RLE_COMPRESS rle 0:1:1
 // %RANGE% TUNE_WORK_STEALING ws 0:1:1
-// %RANGE% TUNE_MEM_PREFERENCE mem 0:2:1
 // %RANGE% TUNE_LOAD ld 0:2:1
 // %RANGE% TUNE_LOAD_ALGORITHM_ID laid 0:2:1
 // %RANGE% TUNE_VEC_SIZE_POW vec 0:2:1
@@ -65,7 +64,8 @@ static void range(nvbench::state& state, nvbench::type_list<SampleT, CounterT, O
       launch
 #if !TUNE_BASE
       ,
-      cuda::execution::tune(bench_policy_selector<key_t, num_channels, num_active_channels>{})
+      cuda::execution::tune(
+        histogram_tuning_policy_selector<SampleT, CounterT, num_channels, num_active_channels, false>{})
 #endif // !TUNE_BASE
     );
     _CCCL_TRY_CUDA_API(

cub/benchmarks/bench/histogram/range.cu

@@ -11,7 +11,6 @@
 // %RANGE% TUNE_THREADS tpb 128:1024:32
 // %RANGE% TUNE_RLE_COMPRESS rle 0:1:1
 // %RANGE% TUNE_WORK_STEALING ws 0:1:1
-// %RANGE% TUNE_MEM_PREFERENCE mem 0:2:1
 // %RANGE% TUNE_LOAD ld 0:2:1
 // %RANGE% TUNE_LOAD_ALGORITHM_ID laid 0:2:1
 // %RANGE% TUNE_VEC_SIZE_POW vec 0:2:1
@@ -51,7 +50,7 @@ static void range(nvbench::state& state, nvbench::type_list<SampleT, CounterT, O
       launch
 #if !TUNE_BASE
       ,
-      cuda::execution::tune(bench_policy_selector<key_t, 1, 1>{})
+      cuda::execution::tune(histogram_tuning_policy_selector<SampleT, CounterT, 1, 1, false>{})
 #endif // !TUNE_BASE
     );
     _CCCL_TRY_CUDA_API(