From a03f9528d2f430d6b664e673d00accc93bcb5dc8 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 31 Jul 2026 08:04:54 +0000 Subject: [PATCH 01/21] [cub] Add cooperative global-memory histogram baseline --- cub/cub/detail/launcher/cuda_driver.cuh | 26 +++ cub/cub/detail/launcher/cuda_runtime.cuh | 38 +++ .../device/dispatch/dispatch_histogram.cuh | 221 ++++++++++++------ .../dispatch/kernels/kernel_histogram.cuh | 130 +++++++++++ cub/test/catch2_test_env_launch_helper.h | 42 +++- 5 files changed, 391 insertions(+), 66 deletions(-) diff --git a/cub/cub/detail/launcher/cuda_driver.cuh b/cub/cub/detail/launcher/cuda_driver.cuh index 23d631b37723..a84dafbbc5ec 100644 --- a/cub/cub/detail/launcher/cuda_driver.cuh +++ b/cub/cub/detail/launcher/cuda_driver.cuh @@ -125,6 +125,32 @@ struct CudaDriverLauncherFactory ::cuOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_fn, block_size, dynamic_smem_bytes)); } + _CCCL_HIDE_FROM_ABI ::cudaError_t CooperativeLaunchSupported(bool& supported) const + { + int attribute = 0; + const auto status = + static_cast<::cudaError_t>(::cuDeviceGetAttribute(&attribute, ::CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH, device_)); + supported = status == ::cudaSuccess && attribute != 0; + return status; + } + + template + _CCCL_HIDE_FROM_ABI ::cudaError_t LaunchCooperative( + dim3 grid, dim3 block, unsigned int shared_mem, ::CUstream stream, ::CUkernel kernel, Args const&... args) const + { + void* kernel_args[] = {const_cast(static_cast(&args))...}; + + ::CUfunction kernel_fn; + auto status = static_cast<::cudaError_t>(::cuKernelGetFunction(&kernel_fn, kernel)); + if (status != cudaSuccess) + { + return status; + } + + return static_cast<::cudaError_t>(::cuLaunchCooperativeKernel( + kernel_fn, grid.x, grid.y, grid.z, block.x, block.y, block.z, shared_mem, stream, kernel_args)); + } + _CCCL_HIDE_FROM_ABI ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const { return static_cast<::cudaError_t>( diff --git a/cub/cub/detail/launcher/cuda_runtime.cuh b/cub/cub/detail/launcher/cuda_runtime.cuh index 4d2580a88937..2c35f85ce64a 100644 --- a/cub/cub/detail/launcher/cuda_runtime.cuh +++ b/cub/cub/detail/launcher/cuda_runtime.cuh @@ -86,6 +86,44 @@ struct TripleChevronFactory return ::cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes); } + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t CooperativeLaunchSupported(bool& supported) const + { + NV_IF_ELSE_TARGET( + NV_IS_HOST, + ({ + int device_ordinal = 0; + if (const auto error = CubDebug(::cudaGetDevice(&device_ordinal))) + { + return error; + } + + int attribute = 0; + if (const auto error = + CubDebug(::cudaDeviceGetAttribute(&attribute, ::cudaDevAttrCooperativeLaunch, device_ordinal))) + { + return error; + } + + supported = attribute != 0; + return ::cudaSuccess; + }), + ({ + supported = false; + return ::cudaSuccess; + })) + } + + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t LaunchCooperative( + dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, Kernel kernel, Args const&... args) + const {NV_IF_ELSE_TARGET( + NV_IS_HOST, + ({ + void* kernel_args[] = {const_cast(static_cast(&args))...}; + return ::cudaLaunchCooperativeKernel(kernel, grid, block, kernel_args, shared_mem, stream); + }), + ({ return ::cudaErrorNotSupported; }))} + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const { int device_ordinal; diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 875f6744afef..07464f983319 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -129,6 +129,22 @@ struct DeviceHistogramKernelSource IsEven>; } + /// Returns the cooperative high-bin histogram sweep kernel. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramSweepKernelPersistent() + { + return &DeviceHistogramSweepPersistentKernel< + PolicyT, + PRIVATIZED_SMEM_BINS, + NUM_CHANNELS, + NUM_ACTIVE_CHANNELS, + SampleIteratorT, + CounterT, + PrivatizedDecodeOpT, + OutputDecodeOpT, + OffsetT>; + } + CUB_RUNTIME_FUNCTION static constexpr size_t CounterSize() { return sizeof(CounterT); @@ -254,6 +270,25 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( // Get device occupancy for sweep_kernel int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; + if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) + { + using output_decode_op_t = typename FirstLevelArrayT::value_type; + using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + const auto persistent_kernel = kernel_source.template HistogramSweepKernelPersistent< + PolicySelector, + PRIVATIZED_SMEM_BINS, + privatized_decode_op_t, + output_decode_op_t>(); + + int persistent_sm_occupancy = 0; + if (const auto error = + CubDebug(launcher_factory.MaxSmOccupancy(persistent_sm_occupancy, persistent_kernel, threads_per_block))) + { + return error; + } + histogram_sweep_occupancy = persistent_sm_occupancy * sm_count; + } + if (num_row_pixels * NUM_CHANNELS == row_stride_samples) { // Treat as a single linear array of samples @@ -322,82 +357,138 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( num_privatized_levels.begin(), num_privatized_levels.end(), num_privatized_bins_wrapper.begin(), minus_one); ::cuda::std::transform(num_output_levels.begin(), num_output_levels.end(), num_output_bins_wrapper.begin(), minus_one); - constexpr int histogram_init_threads_per_block = 256; - int histogram_init_grid_dims = - (max_num_output_bins + histogram_init_threads_per_block - 1) / histogram_init_threads_per_block; + bool launched_persistent = false; +#if _CCCL_HOSTED() + if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) + { + if (blocks_per_row > 0 && blocks_per_col > 0) + { + using output_decode_op_t = typename FirstLevelArrayT::value_type; + using privatized_decode_op_t = typename SecondLevelArrayT::value_type; - // Log DeviceHistogramInitKernel configuration + dim3 persistent_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; + const auto persistent_kernel = kernel_source.template HistogramSweepKernelPersistent< + PolicySelector, + PRIVATIZED_SMEM_BINS, + privatized_decode_op_t, + output_decode_op_t>(); + + bool cooperative_supported = false; + if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(cooperative_supported))) + { + return error; + } + if (cooperative_supported) + { + if (const auto error = CubDebug(launcher_factory.LaunchCooperative( + persistent_grid_dims, + dim3{static_cast(threads_per_block)}, + 0, + stream, + persistent_kernel, + d_samples, + num_output_bins_wrapper, + num_privatized_bins_wrapper, + d_output_histograms, + d_privatized_histograms_wrapper, + first_level_array, + second_level_array, + num_row_pixels, + num_rows, + row_stride_samples, + tiles_per_row, + tile_queue))) + { + return error; + } + launched_persistent = true; + } + } + } +#endif // _CCCL_HOSTED() + + if (!launched_persistent) + { + constexpr int histogram_init_threads_per_block = 256; + int histogram_init_grid_dims = + (max_num_output_bins + histogram_init_threads_per_block - 1) / histogram_init_threads_per_block; + + // Log DeviceHistogramInitKernel configuration #ifdef CUB_DEBUG_LOG - _CubLog("Invoking DeviceHistogramInitKernel<<<%d, %d, 0, %lld>>>()\n", - histogram_init_grid_dims, - histogram_init_threads_per_block, - (long long) stream); + _CubLog("Invoking DeviceHistogramInitKernel<<<%d, %d, 0, %lld>>>()\n", + histogram_init_grid_dims, + histogram_init_threads_per_block, + (long long) stream); #else // CUB_DEBUG_LOG - log("Invoking DeviceHistogramInitKernel<<<%d, %d, 0, %lld>>>()\n", - histogram_init_grid_dims, - histogram_init_threads_per_block, - (long long) stream); + log("Invoking DeviceHistogramInitKernel<<<%d, %d, 0, %lld>>>()\n", + histogram_init_grid_dims, + histogram_init_threads_per_block, + (long long) stream); #endif // CUB_DEBUG_LOG - // Invoke histogram_init_kernel - if (const auto error = CubDebug( - launcher_factory(histogram_init_grid_dims, - histogram_init_threads_per_block, - 0, - stream, - /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) - .doit(init_kernel, num_output_bins_wrapper, d_output_histograms, tile_queue))) - { - return error; - } + // Invoke histogram_init_kernel + if (const auto error = CubDebug( + launcher_factory(histogram_init_grid_dims, + histogram_init_threads_per_block, + 0, + stream, + /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(init_kernel, num_output_bins_wrapper, d_output_histograms, tile_queue))) + { + return error; + } - // Return if empty problem - if (blocks_per_row == 0 || blocks_per_col == 0) - { - return cudaSuccess; - } + // Return if empty problem + if (blocks_per_row == 0 || blocks_per_col == 0) + { + return cudaSuccess; + } - // Log histogram_sweep_kernel configuration + // Log histogram_sweep_kernel configuration #ifdef CUB_DEBUG_LOG - _CubLog("Invoking histogram_sweep_kernel<<<{%d, %d, %d}, %d, 0, %lld>>>(), %d pixels " - "per thread, %d SM occupancy\n", - sweep_grid_dims.x, - sweep_grid_dims.y, - sweep_grid_dims.z, - threads_per_block, - (long long) stream, - pixels_per_thread, - histogram_sweep_sm_occupancy); + _CubLog("Invoking histogram_sweep_kernel<<<{%d, %d, %d}, %d, 0, %lld>>>(), %d pixels " + "per thread, %d SM occupancy\n", + sweep_grid_dims.x, + sweep_grid_dims.y, + sweep_grid_dims.z, + threads_per_block, + (long long) stream, + pixels_per_thread, + histogram_sweep_sm_occupancy); #else // CUB_DEBUG_LOG - log("Invoking histogram_sweep_kernel<<<{%d, %d, %d}, %d, 0, %lld>>>(), %d pixels " - "per thread, %d SM occupancy\n", - sweep_grid_dims.x, - sweep_grid_dims.y, - sweep_grid_dims.z, - threads_per_block, - (long long) stream, - pixels_per_thread, - histogram_sweep_sm_occupancy); + log("Invoking histogram_sweep_kernel<<<{%d, %d, %d}, %d, 0, %lld>>>(), %d pixels " + "per thread, %d SM occupancy\n", + sweep_grid_dims.x, + sweep_grid_dims.y, + sweep_grid_dims.z, + threads_per_block, + (long long) stream, + pixels_per_thread, + histogram_sweep_sm_occupancy); #endif // CUB_DEBUG_LOG - if (const auto error = CubDebug( - launcher_factory( - sweep_grid_dims, threads_per_block, 0, stream, /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) - .doit(sweep_kernel, - d_samples, - num_output_bins_wrapper, - num_privatized_bins_wrapper, - d_output_histograms, - d_privatized_histograms_wrapper, - first_level_array, - second_level_array, - num_row_pixels, - num_rows, - row_stride_samples, - tiles_per_row, - tile_queue))) - { - return error; + if (const auto error = CubDebug( + launcher_factory(sweep_grid_dims, + threads_per_block, + 0, + stream, + /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(sweep_kernel, + d_samples, + num_output_bins_wrapper, + num_privatized_bins_wrapper, + d_output_histograms, + d_privatized_histograms_wrapper, + first_level_array, + second_level_array, + num_row_pixels, + num_rows, + row_stride_samples, + tiles_per_row, + tile_queue))) + { + return error; + } } // Check for failure to launch diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index cc64ddea1862..70d94a105c6a 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -21,6 +21,10 @@ #include #include +#if !_CCCL_COMPILER(NVRTC) +# include +#endif // !_CCCL_COMPILER(NVRTC) + CUB_NAMESPACE_BEGIN namespace detail::histogram { @@ -703,5 +707,131 @@ __launch_bounds__(int(current_policy().threads_per_block)) // Store output to global (if necessary) agent.StoreOutput(); } + +//! Cooperative high-bin histogram sweep with global-memory privatization. +//! +//! The kernel initializes output and private histograms, consumes the input, +//! and gathers the block-private counters after a grid-wide barrier. +template +#if _CCCL_HAS_CONCEPTS() + requires histogram_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramSweepPersistentKernel( + _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, + _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, + _CCCL_GRID_CONSTANT const ::cuda::std::array num_privatized_bins_wrapper, + ::cuda::std::array d_output_histograms_wrapper, + ::cuda::std::array d_privatized_histograms_wrapper, + _CCCL_GRID_CONSTANT const ::cuda::std::array output_decode_op_wrapper, + _CCCL_GRID_CONSTANT const ::cuda::std::array privatized_decode_op_wrapper, + _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, + _CCCL_GRID_CONSTANT const OffsetT num_rows, + _CCCL_GRID_CONSTANT const OffsetT row_stride_samples, + _CCCL_GRID_CONSTANT const int tiles_per_row, + GridQueue tile_queue) +{ + static constexpr HistogramPolicy hp = current_policy(); + + namespace cg = ::cooperative_groups; + + cg::grid_group grid = cg::this_grid(); + + const unsigned int blocks_per_grid = gridDim.x * gridDim.y * gridDim.z; + const unsigned int block_id = (blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x; + const unsigned int tid_global = block_id * blockDim.x + threadIdx.x; + const unsigned int total_threads = blocks_per_grid * blockDim.x; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const int channel_bins = num_output_bins_wrapper[ch]; + for (unsigned int bin = tid_global; bin < static_cast(channel_bins); bin += total_threads) + { + d_output_histograms_wrapper[ch][bin] = 0; + } + } + + if (tid_global == 0) + { + GridQueue queue = tile_queue; + queue.ResetDrain(); + } + + grid.sync(); + + using AgentHistogramPolicyT = agent_histogram_policy< + hp.threads_per_block, + hp.pixels_per_thread, + hp.load_algorithm, + hp.load_modifier, + hp.rle_compress, + hp.mem_preference, + hp.use_work_stealing, + hp.vec_size>; + using AgentHistogramT = + AgentHistogram; + + __shared__ typename AgentHistogramT::TempStorage temp_storage; + + CounterT* d_privatized_base[NumActiveChannels]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + d_privatized_base[ch] = d_privatized_histograms_wrapper[ch]; + } + + AgentHistogramT agent( + temp_storage, + d_samples, + num_output_bins_wrapper.data(), + num_privatized_bins_wrapper.data(), + d_output_histograms_wrapper.data(), + d_privatized_histograms_wrapper.data(), + output_decode_op_wrapper.data(), + privatized_decode_op_wrapper.data()); + + agent.InitBinCounters(); + agent.ConsumeTiles(num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue); + + static_assert(PrivatizedSmemBins == 0); + { + grid.sync(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const int num_bins = num_privatized_bins_wrapper[ch]; + const CounterT* base = d_privatized_base[ch]; + CounterT* d_out = d_output_histograms_wrapper[ch]; + const unsigned int num_bins_u = static_cast(num_bins); + for (unsigned int bin = tid_global; bin < num_bins_u; bin += total_threads) + { + CounterT total = 0; + for (unsigned int b = 0; b < blocks_per_grid; ++b) + { + total += base[b * num_bins_u + bin]; + } + d_out[bin] = total; + } + } + } +} } // namespace detail::histogram CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_env_launch_helper.h b/cub/test/catch2_test_env_launch_helper.h index 686179a6c1c3..02fe4e28b5ad 100644 --- a/cub/test/catch2_test_env_launch_helper.h +++ b/cub/test/catch2_test_env_launch_helper.h @@ -140,8 +140,48 @@ struct stream_registry_factory_t return cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes); } + CUB_RUNTIME_FUNCTION cudaError_t CooperativeLaunchSupported(bool& supported) const + { + NV_IF_ELSE_TARGET( + NV_IS_HOST, + ({ + int device_ordinal = 0; + if (const auto error = cudaGetDevice(&device_ordinal)) + { + return error; + } + + int attribute = 0; + if (const auto error = cudaDeviceGetAttribute(&attribute, cudaDevAttrCooperativeLaunch, device_ordinal)) + { + return error; + } + + supported = attribute != 0; + return cudaSuccess; + }), + ({ + supported = false; + return cudaSuccess; + })) + } + + template + CUB_RUNTIME_FUNCTION cudaError_t LaunchCooperative( + dim3 grid, dim3 block, size_t shared_mem, cudaStream_t stream, Kernel kernel, Args const&... args) const { + NV_IF_ELSE_TARGET(NV_IS_HOST, + ({ + if (get_stream_registry_factory_state()->m_stream) + { + REQUIRE(stream == get_stream_registry_factory_state()->m_stream); + } + void* kernel_args[] = {const_cast(static_cast(&args))...}; + return cudaLaunchCooperativeKernel(kernel, grid, block, kernel_args, shared_mem, stream); + }), + ({ return cudaErrorNotSupported; }))} + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t - MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const + MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const { NV_IF_TARGET(NV_IS_HOST, ({ if (get_stream_registry_factory_state()->m_stream) From 9a3fe3474a2bd83c4495e4b25fcf43c232b25f0b Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 31 Jul 2026 10:30:28 +0000 Subject: [PATCH 02/21] [cub] Add shared-memory cache for high-bin histograms --- .../device/dispatch/dispatch_histogram.cuh | 94 +++++--- .../dispatch/kernels/kernel_histogram.cuh | 202 ++++++++++++++++++ 2 files changed, 261 insertions(+), 35 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 07464f983319..473f064a5c6e 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -145,6 +145,21 @@ struct DeviceHistogramKernelSource OffsetT>; } + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramCacheKernel() + { + return &DeviceHistogramCacheKernel< + PolicyT, + NUM_CHANNELS, + NUM_ACTIVE_CHANNELS, + SampleIteratorT, + CounterT, + PrivatizedDecodeOpT, + OffsetT, + UseSecondProbe, + CountReplicas>; + } + CUB_RUNTIME_FUNCTION static constexpr size_t CounterSize() { return sizeof(CounterT); @@ -252,6 +267,8 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( const int threads_per_block = active_policy.threads_per_block; const int pixels_per_thread = active_policy.pixels_per_thread; + constexpr int cache_count_replicas = IsEven ? 1 : (NUM_ACTIVE_CHANNELS == 1 ? 2 : 4); + // Get SM count int sm_count; if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) @@ -272,21 +289,21 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) { - using output_decode_op_t = typename FirstLevelArrayT::value_type; using privatized_decode_op_t = typename SecondLevelArrayT::value_type; - const auto persistent_kernel = kernel_source.template HistogramSweepKernelPersistent< - PolicySelector, - PRIVATIZED_SMEM_BINS, - privatized_decode_op_t, - output_decode_op_t>(); + const auto cache_kernel = + kernel_source.template HistogramCacheKernel(); + + constexpr int cache_slots_per_channel = NUM_ACTIVE_CHANNELS == 1 ? (IsEven ? 4096 : 2048) : 512; + constexpr int cache_smem_bytes = + NUM_ACTIVE_CHANNELS * cache_slots_per_channel * (sizeof(int) + cache_count_replicas * sizeof(CounterT)); - int persistent_sm_occupancy = 0; - if (const auto error = - CubDebug(launcher_factory.MaxSmOccupancy(persistent_sm_occupancy, persistent_kernel, threads_per_block))) + int cache_sm_occupancy = 0; + if (const auto error = CubDebug( + launcher_factory.MaxSmOccupancy(cache_sm_occupancy, cache_kernel, threads_per_block, cache_smem_bytes))) { return error; } - histogram_sweep_occupancy = persistent_sm_occupancy * sm_count; + histogram_sweep_occupancy = cache_sm_occupancy * sm_count; } if (num_row_pixels * NUM_CHANNELS == row_stride_samples) @@ -367,11 +384,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( using privatized_decode_op_t = typename SecondLevelArrayT::value_type; dim3 persistent_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; - const auto persistent_kernel = kernel_source.template HistogramSweepKernelPersistent< - PolicySelector, - PRIVATIZED_SMEM_BINS, - privatized_decode_op_t, - output_decode_op_t>(); + constexpr int cache_slots_per_channel = NUM_ACTIVE_CHANNELS == 1 ? (IsEven ? 4096 : 2048) : 512; + constexpr int cache_smem_bytes = + NUM_ACTIVE_CHANNELS * cache_slots_per_channel * (sizeof(int) + cache_count_replicas * sizeof(CounterT)); + const bool use_second_probe = max_num_output_bins < 262144; bool cooperative_supported = false; if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(cooperative_supported))) @@ -380,24 +396,32 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } if (cooperative_supported) { - if (const auto error = CubDebug(launcher_factory.LaunchCooperative( - persistent_grid_dims, - dim3{static_cast(threads_per_block)}, - 0, - stream, - persistent_kernel, - d_samples, - num_output_bins_wrapper, - num_privatized_bins_wrapper, - d_output_histograms, - d_privatized_histograms_wrapper, - first_level_array, - second_level_array, - num_row_pixels, - num_rows, - row_stride_samples, - tiles_per_row, - tile_queue))) + const auto launch_cache = [&](auto cache_kernel) { + return launcher_factory.LaunchCooperative( + persistent_grid_dims, + dim3{static_cast(threads_per_block)}, + cache_smem_bytes, + stream, + cache_kernel, + d_samples, + num_output_bins_wrapper, + d_output_histograms, + second_level_array, + num_row_pixels, + num_rows, + row_stride_samples, + cache_slots_per_channel); + }; + + const auto error = + use_second_probe + ? launch_cache( + kernel_source + .template HistogramCacheKernel()) + : launch_cache( + kernel_source + .template HistogramCacheKernel()); + if (CubDebug(error) != cudaSuccess) { return error; } @@ -1213,7 +1237,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( NUM_ACTIVE_CHANNELS, PRIVATIZED_SMEM_BINS, /* IsDeviceInit = */ false, - /* IsEven = */ false, + /* IsEven = */ true, /* IsByteSample = */ false>( d_temp_storage, temp_storage_bytes, @@ -1244,7 +1268,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( NUM_ACTIVE_CHANNELS, PRIVATIZED_SMEM_BINS, /* IsDeviceInit = */ false, - /* IsEven = */ false, + /* IsEven = */ true, /* IsByteSample = */ false>( d_temp_storage, temp_storage_bytes, diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index 70d94a105c6a..936944d578dd 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -20,6 +20,7 @@ #include #include +#include #if !_CCCL_COMPILER(NVRTC) # include @@ -28,6 +29,19 @@ CUB_NAMESPACE_BEGIN namespace detail::histogram { +template +_CCCL_DEVICE _CCCL_FORCEINLINE void histogram_atomic_add(CounterT* address, CounterT value) +{ + if constexpr (::cuda::std::is_integral_v && sizeof(CounterT) == sizeof(unsigned long long)) + { + atomicAdd(reinterpret_cast(address), static_cast(value)); + } + else + { + atomicAdd(address, value); + } +} + template struct Transforms { @@ -833,5 +847,193 @@ __launch_bounds__(int(current_policy().threads_per_block)) } } } + +template +_CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( + int* keys, CounterT* counts, int count_replicas, int bin, CounterT contribution, int cache_mask, int cache_log2) +{ + const auto try_slot = [&](int slot) { + int key = keys[slot]; + if (key == -1) + { + key = atomicCAS_block(&keys[slot], -1, bin); + } + if (key == -1 || key == bin) + { + const int replica = static_cast((threadIdx.x >> 5) % count_replicas); + atomicAdd_block(&counts[replica * (cache_mask + 1) + slot], contribution); + return true; + } + return false; + }; + + const unsigned int hash = static_cast(bin) * 2654435761u; + const int primary = static_cast((hash >> (32 - cache_log2)) & static_cast(cache_mask)); + if (try_slot(primary)) + { + return true; + } + + if constexpr (UseSecondProbe) + { + const unsigned int hash2 = (static_cast(bin) ^ 0x9e3779b9u) * 2246822519u; + const int secondary = static_cast((hash2 >> (32 - cache_log2)) & static_cast(cache_mask)); + return try_slot(secondary); + } + return false; +} + +//! Cooperative high-bin histogram with a per-block shared-memory cache. +template +#if _CCCL_HAS_CONCEPTS() + requires histogram_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramCacheKernel( + _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, + _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, + ::cuda::std::array d_output_histograms_wrapper, + _CCCL_GRID_CONSTANT const ::cuda::std::array decode_op_wrapper, + _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, + _CCCL_GRID_CONSTANT const OffsetT num_rows, + _CCCL_GRID_CONSTANT const OffsetT row_stride_samples, + _CCCL_GRID_CONSTANT const int cache_slots_per_channel) +{ + namespace cg = ::cooperative_groups; + cg::grid_group grid = cg::this_grid(); + + const unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned int total_threads = gridDim.x * blockDim.x; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); + bin += total_threads) + { + d_output_histograms_wrapper[ch][bin] = CounterT{0}; + } + } + grid.sync(); + + extern __shared__ unsigned char dynamic_smem[]; + int* cache_keys = reinterpret_cast(dynamic_smem); + CounterT* cache_counts = + reinterpret_cast(cache_keys + static_cast(NumActiveChannels) * cache_slots_per_channel); + const int cache_mask = cache_slots_per_channel - 1; + const int cache_log2 = 31 - __clz(static_cast(cache_slots_per_channel)); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; + for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) + { + channel_keys[slot] = -1; + _CCCL_PRAGMA_UNROLL_FULL() + for (int replica = 0; replica < CountReplicas; ++replica) + { + channel_counts[static_cast(replica) * cache_slots_per_channel + slot] = CounterT{0}; + } + } + } + __syncthreads(); + + constexpr int unroll = NumActiveChannels == 1 ? 4 : 1; + const OffsetT total_pixels = num_rows * num_row_pixels; + const OffsetT step = static_cast(total_threads); + const OffsetT chunk = static_cast(unroll) * step; + const OffsetT chunk_count = ::cuda::ceil_div(total_pixels, chunk); + const unsigned int lane_id = threadIdx.x & 0x1f; + const bool contiguous_input = num_rows == 1; + + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + { + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + OffsetT pixel_offset = 0; + if (valid) + { + if (contiguous_input) + { + pixel_offset = pixel * NumChannels; + } + else + { + const OffsetT row = pixel / num_row_pixels; + pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + } + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int bin = -1; + if (valid) + { + const auto sample = d_samples[pixel_offset + ch]; + decode_op_wrapper[ch].template BinSelect(sample, bin, true); + if (bin < 0 || bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } + } + + const unsigned int active = __activemask(); + const unsigned int peers = __match_any_sync(active, static_cast(bin)); + const int leader = __ffs(static_cast(peers)) - 1; + if (bin >= 0 && static_cast(lane_id) == leader) + { + const CounterT contribution = static_cast(__popc(peers)); + int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; + if (!histogram_cache_probe( + channel_keys, channel_counts, CountReplicas, bin, contribution, cache_mask, cache_log2)) + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][bin], contribution); + } + } + } + } + } + + __syncthreads(); + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; + for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) + { + const int key = channel_keys[slot]; + if (key >= 0) + { + CounterT count = CounterT{0}; + _CCCL_PRAGMA_UNROLL_FULL() + for (int replica = 0; replica < CountReplicas; ++replica) + { + count += channel_counts[static_cast(replica) * cache_slots_per_channel + slot]; + } + if (count > CounterT{0}) + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][key], count); + } + } + } + } +} } // namespace detail::histogram CUB_NAMESPACE_END From 95ae80dc3409eec52aaad8d417bec7335c1a08e7 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Tue, 11 Aug 2026 16:14:19 +0000 Subject: [PATCH 03/21] [cub] Address cooperative histogram review --- .../device/dispatch/dispatch_histogram.cuh | 147 +++-- .../dispatch/kernels/kernel_histogram.cuh | 546 ++++++++++-------- .../dispatch/tuning/tuning_histogram.cuh | 200 ++++++- cub/test/catch2_test_device_histogram_env.cu | 78 ++- 4 files changed, 649 insertions(+), 322 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 473f064a5c6e..1cd47793a93d 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -129,37 +129,20 @@ struct DeviceHistogramKernelSource IsEven>; } - /// Returns the cooperative high-bin histogram sweep kernel. - template - _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramSweepKernelPersistent() + /// Returns the policy-configurable cooperative high-bin histogram kernel. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramCooperativeKernel() { - return &DeviceHistogramSweepPersistentKernel< + return &DeviceHistogramCooperativeKernel< PolicyT, - PRIVATIZED_SMEM_BINS, NUM_CHANNELS, NUM_ACTIVE_CHANNELS, SampleIteratorT, CounterT, PrivatizedDecodeOpT, - OutputDecodeOpT, OffsetT>; } - template - _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramCacheKernel() - { - return &DeviceHistogramCacheKernel< - PolicyT, - NUM_CHANNELS, - NUM_ACTIVE_CHANNELS, - SampleIteratorT, - CounterT, - PrivatizedDecodeOpT, - OffsetT, - UseSecondProbe, - CountReplicas>; - } - CUB_RUNTIME_FUNCTION static constexpr size_t CounterSize() { return sizeof(CounterT); @@ -267,8 +250,6 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( const int threads_per_block = active_policy.threads_per_block; const int pixels_per_thread = active_policy.pixels_per_thread; - constexpr int cache_count_replicas = IsEven ? 1 : (NUM_ACTIVE_CHANNELS == 1 ? 2 : 4); - // Get SM count int sm_count; if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) @@ -286,25 +267,49 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( // Get device occupancy for sweep_kernel int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; + bool use_cooperative = false; + int cooperative_smem_bytes = 0; +#if _CCCL_HOSTED() if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) { - using privatized_decode_op_t = typename SecondLevelArrayT::value_type; - const auto cache_kernel = - kernel_source.template HistogramCacheKernel(); - - constexpr int cache_slots_per_channel = NUM_ACTIVE_CHANNELS == 1 ? (IsEven ? 4096 : 2048) : 512; - constexpr int cache_smem_bytes = - NUM_ACTIVE_CHANNELS * cache_slots_per_channel * (sizeof(int) + cache_count_replicas * sizeof(CounterT)); - - int cache_sm_occupancy = 0; - if (const auto error = CubDebug( - launcher_factory.MaxSmOccupancy(cache_sm_occupancy, cache_kernel, threads_per_block, cache_smem_bytes))) + if (active_policy.high_bin_algorithm == HistogramHighBinAlgorithm::cooperative) { - return error; + if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(use_cooperative))) + { + return error; + } + if (use_cooperative) + { + using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + const auto cooperative_kernel = + kernel_source.template HistogramCooperativeKernel(); + const int cache_slots_per_channel = + active_policy.high_bin_cache == HistogramCacheAlgorithm::none + ? 0 + : active_policy.high_bin_cache_entries_per_channel; + cooperative_smem_bytes = + NUM_ACTIVE_CHANNELS * cache_slots_per_channel + * (int{sizeof(::cuda::std::uint32_t)} + active_policy.high_bin_cache_count_replicas * int{sizeof(CounterT)}); + + int cooperative_sm_occupancy = 0; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + cooperative_sm_occupancy, cooperative_kernel, threads_per_block, cooperative_smem_bytes))) + { + return error; + } + if (cooperative_sm_occupancy > 0) + { + histogram_sweep_occupancy = cooperative_sm_occupancy * sm_count; + } + else + { + use_cooperative = false; + } + } } - histogram_sweep_occupancy = cache_sm_occupancy * sm_count; } +#endif // _CCCL_HOSTED() if (num_row_pixels * NUM_CHANNELS == row_stride_samples) { @@ -315,9 +320,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy - int pixels_per_tile = threads_per_block * pixels_per_thread; - int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); - int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); + const int selected_pixels_per_thread = use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread; + int pixels_per_tile = threads_per_block * selected_pixels_per_thread; + int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); + int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); int blocks_per_col = (blocks_per_row > 0) ? int(::cuda::std::min(static_cast(histogram_sweep_occupancy / blocks_per_row), num_rows)) @@ -336,8 +342,12 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( for (int CHANNEL = 0; CHANNEL < NUM_ACTIVE_CHANNELS; ++CHANNEL) { + const bool needs_privatized_storage = + !use_cooperative || active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized; allocation_sizes[CHANNEL] = - size_t(num_thread_blocks) * (num_privatized_levels[CHANNEL] - 1) * kernel_source.CounterSize(); + needs_privatized_storage + ? size_t(num_thread_blocks) * (num_privatized_levels[CHANNEL] - 1) * kernel_source.CounterSize() + : 0; } allocation_sizes[NUM_ALLOCATIONS - 1] = GridQueue::AllocationSize(); @@ -374,64 +384,45 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( num_privatized_levels.begin(), num_privatized_levels.end(), num_privatized_bins_wrapper.begin(), minus_one); ::cuda::std::transform(num_output_levels.begin(), num_output_levels.end(), num_output_bins_wrapper.begin(), minus_one); - bool launched_persistent = false; + bool launched_cooperative = false; #if _CCCL_HOSTED() if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) { - if (blocks_per_row > 0 && blocks_per_col > 0) + if (use_cooperative && blocks_per_row > 0 && blocks_per_col > 0) { - using output_decode_op_t = typename FirstLevelArrayT::value_type; using privatized_decode_op_t = typename SecondLevelArrayT::value_type; - dim3 persistent_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; - constexpr int cache_slots_per_channel = NUM_ACTIVE_CHANNELS == 1 ? (IsEven ? 4096 : 2048) : 512; - constexpr int cache_smem_bytes = - NUM_ACTIVE_CHANNELS * cache_slots_per_channel * (sizeof(int) + cache_count_replicas * sizeof(CounterT)); - const bool use_second_probe = max_num_output_bins < 262144; - - bool cooperative_supported = false; - if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(cooperative_supported))) - { - return error; - } - if (cooperative_supported) - { - const auto launch_cache = [&](auto cache_kernel) { - return launcher_factory.LaunchCooperative( - persistent_grid_dims, + const dim3 cooperative_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; + const int cache_slots_per_channel = + active_policy.high_bin_cache == HistogramCacheAlgorithm::none + ? 0 + : active_policy.high_bin_cache_entries_per_channel; + const auto cooperative_kernel = + kernel_source.template HistogramCooperativeKernel(); + if (const auto error = CubDebug(launcher_factory.LaunchCooperative( + cooperative_grid_dims, dim3{static_cast(threads_per_block)}, - cache_smem_bytes, + cooperative_smem_bytes, stream, - cache_kernel, + cooperative_kernel, d_samples, num_output_bins_wrapper, d_output_histograms, + d_privatized_histograms_wrapper, second_level_array, num_row_pixels, num_rows, row_stride_samples, - cache_slots_per_channel); - }; - - const auto error = - use_second_probe - ? launch_cache( - kernel_source - .template HistogramCacheKernel()) - : launch_cache( - kernel_source - .template HistogramCacheKernel()); - if (CubDebug(error) != cudaSuccess) - { - return error; - } - launched_persistent = true; + cache_slots_per_channel))) + { + return error; } + launched_cooperative = true; } } #endif // _CCCL_HOSTED() - if (!launched_persistent) + if (!launched_cooperative) { constexpr int histogram_init_threads_per_block = 256; int histogram_init_grid_dims = diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index 936944d578dd..e493f09d04d2 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -20,6 +20,7 @@ #include #include +#include #include #if !_CCCL_COMPILER(NVRTC) @@ -32,8 +33,10 @@ namespace detail::histogram template _CCCL_DEVICE _CCCL_FORCEINLINE void histogram_atomic_add(CounterT* address, CounterT value) { - if constexpr (::cuda::std::is_integral_v && sizeof(CounterT) == sizeof(unsigned long long)) + if constexpr (::cuda::std::is_integral_v && sizeof(CounterT) == sizeof(::cuda::std::uint64_t)) { + // CUDA's 64-bit integer atomic overload is spelled in terms of unsigned long long. + // Keep the width decision explicit and use that spelling only at the API boundary. atomicAdd(reinterpret_cast(address), static_cast(value)); } else @@ -722,143 +725,32 @@ __launch_bounds__(int(current_policy().threads_per_block)) agent.StoreOutput(); } -//! Cooperative high-bin histogram sweep with global-memory privatization. -//! -//! The kernel initializes output and private histograms, consumes the input, -//! and gathers the block-private counters after a grid-wide barrier. -template -#if _CCCL_HAS_CONCEPTS() - requires histogram_policy_selector -#endif // _CCCL_HAS_CONCEPTS() -__launch_bounds__(int(current_policy().threads_per_block)) - _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramSweepPersistentKernel( - _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, - _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, - _CCCL_GRID_CONSTANT const ::cuda::std::array num_privatized_bins_wrapper, - ::cuda::std::array d_output_histograms_wrapper, - ::cuda::std::array d_privatized_histograms_wrapper, - _CCCL_GRID_CONSTANT const ::cuda::std::array output_decode_op_wrapper, - _CCCL_GRID_CONSTANT const ::cuda::std::array privatized_decode_op_wrapper, - _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, - _CCCL_GRID_CONSTANT const OffsetT num_rows, - _CCCL_GRID_CONSTANT const OffsetT row_stride_samples, - _CCCL_GRID_CONSTANT const int tiles_per_row, - GridQueue tile_queue) +template +_CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( + ::cuda::std::uint32_t* keys, + CounterT* counts, + int count_replicas, + int bin, + CounterT contribution, + int cache_mask, + int cache_log2, + HistogramCacheAlgorithm cache_algorithm, + bool use_second_probe) { - static constexpr HistogramPolicy hp = current_policy(); - - namespace cg = ::cooperative_groups; - - cg::grid_group grid = cg::this_grid(); - - const unsigned int blocks_per_grid = gridDim.x * gridDim.y * gridDim.z; - const unsigned int block_id = (blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x; - const unsigned int tid_global = block_id * blockDim.x + threadIdx.x; - const unsigned int total_threads = blocks_per_grid * blockDim.x; - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - const int channel_bins = num_output_bins_wrapper[ch]; - for (unsigned int bin = tid_global; bin < static_cast(channel_bins); bin += total_threads) - { - d_output_histograms_wrapper[ch][bin] = 0; - } - } - - if (tid_global == 0) - { - GridQueue queue = tile_queue; - queue.ResetDrain(); - } - - grid.sync(); - - using AgentHistogramPolicyT = agent_histogram_policy< - hp.threads_per_block, - hp.pixels_per_thread, - hp.load_algorithm, - hp.load_modifier, - hp.rle_compress, - hp.mem_preference, - hp.use_work_stealing, - hp.vec_size>; - using AgentHistogramT = - AgentHistogram; - - __shared__ typename AgentHistogramT::TempStorage temp_storage; - - CounterT* d_privatized_base[NumActiveChannels]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - d_privatized_base[ch] = d_privatized_histograms_wrapper[ch]; - } - - AgentHistogramT agent( - temp_storage, - d_samples, - num_output_bins_wrapper.data(), - num_privatized_bins_wrapper.data(), - d_output_histograms_wrapper.data(), - d_privatized_histograms_wrapper.data(), - output_decode_op_wrapper.data(), - privatized_decode_op_wrapper.data()); - - agent.InitBinCounters(); - agent.ConsumeTiles(num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue); - - static_assert(PrivatizedSmemBins == 0); + if (cache_algorithm == HistogramCacheAlgorithm::none) { - grid.sync(); - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - const int num_bins = num_privatized_bins_wrapper[ch]; - const CounterT* base = d_privatized_base[ch]; - CounterT* d_out = d_output_histograms_wrapper[ch]; - const unsigned int num_bins_u = static_cast(num_bins); - for (unsigned int bin = tid_global; bin < num_bins_u; bin += total_threads) - { - CounterT total = 0; - for (unsigned int b = 0; b < blocks_per_grid; ++b) - { - total += base[b * num_bins_u + bin]; - } - d_out[bin] = total; - } - } + return false; } -} -template -_CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( - int* keys, CounterT* counts, int count_replicas, int bin, CounterT contribution, int cache_mask, int cache_log2) -{ - const auto try_slot = [&](int slot) { - int key = keys[slot]; - if (key == -1) + constexpr ::cuda::std::uint32_t empty_key = UINT32_MAX; + const auto bin_key = static_cast<::cuda::std::uint32_t>(bin); + const auto try_slot = [&](int slot) { + ::cuda::std::uint32_t key = keys[slot]; + if (key == empty_key) { - key = atomicCAS_block(&keys[slot], -1, bin); + key = atomicCAS_block(&keys[slot], empty_key, bin_key); } - if (key == -1 || key == bin) + if (key == empty_key || key == bin_key) { const int replica = static_cast((threadIdx.x >> 5) % count_replicas); atomicAdd_block(&counts[replica * (cache_mask + 1) + slot], contribution); @@ -874,7 +766,7 @@ _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( return true; } - if constexpr (UseSecondProbe) + if (use_second_probe) { const unsigned int hash2 = (static_cast(bin) ^ 0x9e3779b9u) * 2246822519u; const int secondary = static_cast((hash2 >> (32 - cache_log2)) & static_cast(cache_mask)); @@ -883,157 +775,345 @@ _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( return false; } -//! Cooperative high-bin histogram with a per-block shared-memory cache. +//! Agent for the policy-configurable cooperative high-bin histogram kernel. template + typename OffsetT> #if _CCCL_HAS_CONCEPTS() requires histogram_policy_selector #endif // _CCCL_HAS_CONCEPTS() -__launch_bounds__(int(current_policy().threads_per_block)) - _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramCacheKernel( - _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, - _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, - ::cuda::std::array d_output_histograms_wrapper, - _CCCL_GRID_CONSTANT const ::cuda::std::array decode_op_wrapper, - _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, - _CCCL_GRID_CONSTANT const OffsetT num_rows, - _CCCL_GRID_CONSTANT const OffsetT row_stride_samples, - _CCCL_GRID_CONSTANT const int cache_slots_per_channel) +struct AgentHistogramCooperative { - namespace cg = ::cooperative_groups; - cg::grid_group grid = cg::this_grid(); - - const unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x; - const unsigned int total_threads = gridDim.x * blockDim.x; - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + _CCCL_DEVICE _CCCL_FORCEINLINE static void Consume( + const SampleIteratorT d_samples, + const ::cuda::std::array num_output_bins_wrapper, + ::cuda::std::array d_output_histograms_wrapper, + ::cuda::std::array d_privatized_histograms_wrapper, + const ::cuda::std::array decode_op_wrapper, + const OffsetT num_row_pixels, + const OffsetT num_rows, + const OffsetT row_stride_samples, + const int cache_slots_per_channel) { - for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); - bin += total_threads) - { - d_output_histograms_wrapper[ch][bin] = CounterT{0}; - } - } - grid.sync(); + static constexpr HistogramPolicy policy = current_policy(); + static constexpr int count_replicas = policy.high_bin_cache_count_replicas; + static_assert(policy.high_bin_pixels_per_thread > 0, "Histogram cooperative unroll must be positive"); + static_assert( + policy.high_bin_cache == HistogramCacheAlgorithm::none + || (policy.high_bin_cache_entries_per_channel > 0 + && (policy.high_bin_cache_entries_per_channel & (policy.high_bin_cache_entries_per_channel - 1)) == 0), + "Histogram cache entries per channel must be a power of two"); + namespace cg = ::cooperative_groups; + cg::grid_group grid = cg::this_grid(); + + const unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned int total_threads = gridDim.x * blockDim.x; - extern __shared__ unsigned char dynamic_smem[]; - int* cache_keys = reinterpret_cast(dynamic_smem); - CounterT* cache_counts = - reinterpret_cast(cache_keys + static_cast(NumActiveChannels) * cache_slots_per_channel); - const int cache_mask = cache_slots_per_channel - 1; - const int cache_log2 = 31 - __clz(static_cast(cache_slots_per_channel)); - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; - for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - channel_keys[slot] = -1; - _CCCL_PRAGMA_UNROLL_FULL() - for (int replica = 0; replica < CountReplicas; ++replica) + for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); + bin += total_threads) { - channel_counts[static_cast(replica) * cache_slots_per_channel + slot] = CounterT{0}; + d_output_histograms_wrapper[ch][bin] = CounterT{0}; } } - } - __syncthreads(); + grid.sync(); - constexpr int unroll = NumActiveChannels == 1 ? 4 : 1; - const OffsetT total_pixels = num_rows * num_row_pixels; - const OffsetT step = static_cast(total_threads); - const OffsetT chunk = static_cast(unroll) * step; - const OffsetT chunk_count = ::cuda::ceil_div(total_pixels, chunk); - const unsigned int lane_id = threadIdx.x & 0x1f; - const bool contiguous_input = num_rows == 1; + static_assert(count_replicas > 0, "Histogram cache replication must be positive"); + + extern __shared__ unsigned char dynamic_smem[]; + auto* cache_keys = reinterpret_cast<::cuda::std::uint32_t*>(dynamic_smem); + CounterT* cache_counts = + reinterpret_cast(cache_keys + static_cast(NumActiveChannels) * cache_slots_per_channel); + const int cache_mask = cache_slots_per_channel > 0 ? cache_slots_per_channel - 1 : 0; + const int cache_log2 = + cache_slots_per_channel > 0 ? 31 - __clz(static_cast(cache_slots_per_channel)) : 0; - for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) - { - const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + for (int ch = 0; ch < NumActiveChannels; ++ch) { - const OffsetT pixel = first_pixel + static_cast(item) * step; - const bool valid = pixel < total_pixels; - OffsetT pixel_offset = 0; - if (valid) + auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) { - if (contiguous_input) + channel_keys[slot] = ~::cuda::std::uint32_t{0}; + _CCCL_PRAGMA_UNROLL_FULL() + for (int replica = 0; replica < count_replicas; ++replica) { - pixel_offset = pixel * NumChannels; + channel_counts[static_cast(replica) * cache_slots_per_channel + slot] = CounterT{0}; } - else + } + + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + for (int bin = threadIdx.x; bin < num_output_bins_wrapper[ch]; bin += blockDim.x) { - const OffsetT row = pixel / num_row_pixels; - pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + block_histogram[bin] = CounterT{0}; } } + } + __syncthreads(); + + constexpr int unroll = policy.high_bin_pixels_per_thread; + const OffsetT total_pixels = num_rows * num_row_pixels; + const OffsetT step = static_cast(total_threads); + const OffsetT chunk = static_cast(unroll) * step; + const OffsetT chunk_count = ::cuda::ceil_div(total_pixels, chunk); + const unsigned int lane_id = threadIdx.x & 0x1f; + const bool contiguous_input = num_rows == 1; + int pending_bin[NumActiveChannels]; + CounterT pending_count[NumActiveChannels]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + pending_bin[ch] = -1; + pending_count[ch] = CounterT{0}; + } + + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + { + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + for (int item = 0; item < unroll; ++item) { - int bin = -1; + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + OffsetT pixel_offset = 0; if (valid) { - const auto sample = d_samples[pixel_offset + ch]; - decode_op_wrapper[ch].template BinSelect(sample, bin, true); - if (bin < 0 || bin >= num_output_bins_wrapper[ch]) + if (contiguous_input) + { + pixel_offset = pixel * NumChannels; + } + else { - bin = -1; + const OffsetT row = pixel / num_row_pixels; + pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; } } - const unsigned int active = __activemask(); - const unsigned int peers = __match_any_sync(active, static_cast(bin)); - const int leader = __ffs(static_cast(peers)) - 1; - if (bin >= 0 && static_cast(lane_id) == leader) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - const CounterT contribution = static_cast(__popc(peers)); - int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; - if (!histogram_cache_probe( - channel_keys, channel_counts, CountReplicas, bin, contribution, cache_mask, cache_log2)) + int bin = -1; + if (valid) + { + const auto sample = d_samples[pixel_offset + ch]; + decode_op_wrapper[ch].template BinSelect(sample, bin, true); + if (bin < 0 || bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } + } + + const auto update_bin = [&](int selected_bin, CounterT contribution) { + auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = + cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo + && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; + if (!histogram_cache_probe( + channel_keys, + channel_counts, + count_replicas, + selected_bin, + contribution, + cache_mask, + cache_log2, + policy.high_bin_cache, + use_second_probe)) + { + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + atomicAdd_block(&block_histogram[selected_bin], contribution); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][selected_bin], contribution); + } + } + }; + + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::warp_coalesced) { - histogram_atomic_add(&d_output_histograms_wrapper[ch][bin], contribution); + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (const unsigned int active = __activemask(); + const unsigned int peers = __match_any_sync(active, static_cast(bin)); + const int leader = __ffs(static_cast(peers)) - 1; + if (bin >= 0 && static_cast(lane_id) == leader) { + update_bin(bin, static_cast(__popc(peers))); + }), + (if (bin >= 0) { update_bin(bin, CounterT{1}); })); + } + else if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) + { + if (bin >= 0) + { + if (pending_bin[ch] == bin) + { + ++pending_count[ch]; + } + else + { + if (pending_bin[ch] >= 0) + { + update_bin(pending_bin[ch], pending_count[ch]); + } + pending_bin[ch] = bin; + pending_count[ch] = CounterT{1}; + } + } + } + else if (bin >= 0) + { + update_bin(bin, CounterT{1}); } } } } - } - __syncthreads(); - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - int* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = cache_counts + static_cast(ch) * CountReplicas * cache_slots_per_channel; - for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) { - const int key = channel_keys[slot]; - if (key >= 0) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - CounterT count = CounterT{0}; - _CCCL_PRAGMA_UNROLL_FULL() - for (int replica = 0; replica < CountReplicas; ++replica) + if (pending_bin[ch] >= 0) + { + auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo + && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; + if (!histogram_cache_probe( + channel_keys, + channel_counts, + count_replicas, + pending_bin[ch], + pending_count[ch], + cache_mask, + cache_log2, + policy.high_bin_cache, + use_second_probe)) + { + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][pending_bin[ch]], pending_count[ch]); + } + } + } + } + } + + __syncthreads(); + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) + { + const auto key = channel_keys[slot]; + if (key != ~::cuda::std::uint32_t{0}) { - count += channel_counts[static_cast(replica) * cache_slots_per_channel + slot]; + CounterT count = CounterT{0}; + _CCCL_PRAGMA_UNROLL_FULL() + for (int replica = 0; replica < count_replicas; ++replica) + { + count += channel_counts[static_cast(replica) * cache_slots_per_channel + slot]; + } + if (count > CounterT{0}) + { + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + block_histogram[key] += count; + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][key], count); + } + } } - if (count > CounterT{0}) + } + } + + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + grid.sync(); + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const unsigned int num_bins = static_cast(num_output_bins_wrapper[ch]); + for (unsigned int bin = tid_global; bin < num_bins; bin += total_threads) { - histogram_atomic_add(&d_output_histograms_wrapper[ch][key], count); + CounterT total = CounterT{0}; + for (unsigned int block = 0; block < gridDim.x; ++block) + { + total += d_privatized_histograms_wrapper[ch][static_cast(block) * num_bins + bin]; + } + d_output_histograms_wrapper[ch][bin] = total; } } } } +}; + +//! Policy-configurable cooperative high-bin histogram kernel. +template +#if _CCCL_HAS_CONCEPTS() + requires histogram_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramCooperativeKernel( + _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, + _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, + ::cuda::std::array d_output_histograms_wrapper, + ::cuda::std::array d_privatized_histograms_wrapper, + _CCCL_GRID_CONSTANT const ::cuda::std::array decode_op_wrapper, + _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, + _CCCL_GRID_CONSTANT const OffsetT num_rows, + _CCCL_GRID_CONSTANT const OffsetT row_stride_samples, + _CCCL_GRID_CONSTANT const int cache_slots_per_channel) +{ + AgentHistogramCooperative< + PolicySelector, + NumChannels, + NumActiveChannels, + SampleIteratorT, + CounterT, + PrivatizedDecodeOpT, + OffsetT>::Consume(d_samples, + num_output_bins_wrapper, + d_output_histograms_wrapper, + d_privatized_histograms_wrapper, + decode_op_wrapper, + num_row_pixels, + num_rows, + row_stride_samples, + cache_slots_per_channel); } } // namespace detail::histogram CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 2084666a5af3..37972e2334c4 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -25,6 +25,87 @@ CUB_NAMESPACE_BEGIN +enum class HistogramHighBinAlgorithm +{ + global_memory_privatized, + cooperative +}; + +enum class HistogramCacheAlgorithm +{ + none, + single_probe, + cuckoo +}; + +enum class HistogramSpillAlgorithm +{ + output, + global_memory_privatized +}; + +enum class HistogramAggregationAlgorithm +{ + direct, + warp_coalesced, + rle +}; + +namespace detail::histogram +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(HistogramHighBinAlgorithm value) noexcept +{ + return value == HistogramHighBinAlgorithm::cooperative + ? "HistogramHighBinAlgorithm::cooperative" + : "HistogramHighBinAlgorithm::global_memory_privatized"; +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(HistogramCacheAlgorithm value) noexcept +{ + return value == HistogramCacheAlgorithm::none ? "HistogramCacheAlgorithm::none" + : value == HistogramCacheAlgorithm::single_probe + ? "HistogramCacheAlgorithm::single_probe" + : "HistogramCacheAlgorithm::cuckoo"; +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(HistogramSpillAlgorithm value) noexcept +{ + return value == HistogramSpillAlgorithm::output + ? "HistogramSpillAlgorithm::output" + : "HistogramSpillAlgorithm::global_memory_privatized"; +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(HistogramAggregationAlgorithm value) noexcept +{ + return value == HistogramAggregationAlgorithm::direct ? "HistogramAggregationAlgorithm::direct" + : value == HistogramAggregationAlgorithm::warp_coalesced + ? "HistogramAggregationAlgorithm::warp_coalesced" + : "HistogramAggregationAlgorithm::rle"; +} +} // namespace detail::histogram + +#if _CCCL_HOSTED() +inline ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) +{ + return os << detail::histogram::to_string(value); +} + +inline ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) +{ + return os << detail::histogram::to_string(value); +} + +inline ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) +{ + return os << detail::histogram::to_string(value); +} + +inline ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) +{ + return os << detail::histogram::to_string(value); +} +#endif // _CCCL_HOSTED() + //! The tuning policy for all algorithms in @ref DeviceHistogram. struct HistogramPolicy { @@ -39,6 +120,14 @@ struct HistogramPolicy bool use_work_stealing; //!< Whether to dequeue tiles from a global work queue int init_kernel_pdl_trigger_max_bins; //!< Maximum number of bins for the init kernel to trigger the histogram kernel //!< early using PDL + HistogramHighBinAlgorithm high_bin_algorithm = HistogramHighBinAlgorithm::cooperative; + HistogramCacheAlgorithm high_bin_cache = HistogramCacheAlgorithm::cuckoo; + HistogramSpillAlgorithm high_bin_spill = HistogramSpillAlgorithm::output; + HistogramAggregationAlgorithm high_bin_aggregation = HistogramAggregationAlgorithm::warp_coalesced; + int high_bin_cache_entries_per_channel = 2048; + int high_bin_cache_count_replicas = 1; + int high_bin_cache_cuckoo_max_bins = 262144; + int high_bin_pixels_per_thread = 4; [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const HistogramPolicy& lhs, const HistogramPolicy& rhs) noexcept @@ -47,7 +136,13 @@ struct HistogramPolicy && lhs.vec_size == rhs.vec_size && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier && lhs.rle_compress == rhs.rle_compress && lhs.mem_preference == rhs.mem_preference && lhs.use_work_stealing == rhs.use_work_stealing - && lhs.init_kernel_pdl_trigger_max_bins == rhs.init_kernel_pdl_trigger_max_bins; + && lhs.init_kernel_pdl_trigger_max_bins == rhs.init_kernel_pdl_trigger_max_bins + && lhs.high_bin_algorithm == rhs.high_bin_algorithm && lhs.high_bin_cache == rhs.high_bin_cache + && lhs.high_bin_spill == rhs.high_bin_spill && lhs.high_bin_aggregation == rhs.high_bin_aggregation + && lhs.high_bin_cache_entries_per_channel == rhs.high_bin_cache_entries_per_channel + && lhs.high_bin_cache_count_replicas == rhs.high_bin_cache_count_replicas + && lhs.high_bin_cache_cuckoo_max_bins == rhs.high_bin_cache_cuckoo_max_bins + && lhs.high_bin_pixels_per_thread == rhs.high_bin_pixels_per_thread; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -64,7 +159,13 @@ struct HistogramPolicy << p.pixels_per_thread << ", .vec_size = " << p.vec_size << ", .load_algorithm = " << p.load_algorithm << ", .load_modifier = " << p.load_modifier << ", .rle_compress = " << p.rle_compress << ", .mem_preference = " << p.mem_preference << ", .use_work_stealing = " << p.use_work_stealing - << ", .init_kernel_pdl_trigger_max_bins = " << p.init_kernel_pdl_trigger_max_bins << " }"; + << ", .init_kernel_pdl_trigger_max_bins = " << p.init_kernel_pdl_trigger_max_bins + << ", .high_bin_algorithm = " << p.high_bin_algorithm << ", .high_bin_cache = " << p.high_bin_cache + << ", .high_bin_spill = " << p.high_bin_spill << ", .high_bin_aggregation = " << p.high_bin_aggregation + << ", .high_bin_cache_entries_per_channel = " << p.high_bin_cache_entries_per_channel + << ", .high_bin_cache_count_replicas = " << p.high_bin_cache_count_replicas + << ", .high_bin_cache_cuckoo_max_bins = " << p.high_bin_cache_cuckoo_max_bins + << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread << " }"; } #endif // _CCCL_HOSTED() }; @@ -305,12 +406,46 @@ public: if (is_even) { // ipt_12.tpb_928.rle_0.ws_0.mem_1.ld_2.laid_0.vec_2 1.033332 0.940517 1.031835 1.195876 - return HistogramPolicy{928, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_CA, false, SMEM, false, 2048}; + return HistogramPolicy{ + 928, + 12, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_CA, + false, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::cuckoo, + HistogramSpillAlgorithm::output, + HistogramAggregationAlgorithm::warp_coalesced, + 4096, + 1, + 262144, + 4}; } else { // ipt_12.tpb_448.rle_0.ws_0.mem_1.ld_1.laid_0.vec_2 1.078987 0.985542 1.085118 1.175637 - return HistogramPolicy{448, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_LDG, false, SMEM, false, 2048}; + return HistogramPolicy{ + 448, + 12, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + false, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::cuckoo, + HistogramSpillAlgorithm::output, + HistogramAggregationAlgorithm::warp_coalesced, + 2048, + 2, + 262144, + 4}; } } @@ -324,17 +459,68 @@ public: { if (sample_size == 1) { - return HistogramPolicy{768, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_LDG, false, SMEM, false, 2048}; + return HistogramPolicy{ + 768, + 12, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + false, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::cuckoo, + HistogramSpillAlgorithm::output, + HistogramAggregationAlgorithm::warp_coalesced, + is_even ? 4096 : 2048, + is_even ? 1 : 2, + 262144, + 4}; } else if (sample_size == 2) { - return HistogramPolicy{960, 10, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, true, SMEM, false, 2048}; + return HistogramPolicy{ + 960, + 10, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_DEFAULT, + true, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::cuckoo, + HistogramSpillAlgorithm::output, + HistogramAggregationAlgorithm::warp_coalesced, + is_even ? 4096 : 2048, + is_even ? 1 : 2, + 262144, + 4}; } } } // fallback from SM50 - return HistogramPolicy{384, t_scale(16), 4, BLOCK_LOAD_DIRECT, LOAD_LDG, true, SMEM, false, 0}; + return HistogramPolicy{ + 384, + t_scale(16), + 4, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + true, + SMEM, + false, + 0, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::cuckoo, + HistogramSpillAlgorithm::output, + HistogramAggregationAlgorithm::warp_coalesced, + num_active_channels == 1 ? (is_even ? 4096 : 2048) : 512, + is_even ? 1 : (num_active_channels == 1 ? 2 : 4), + 262144, + num_active_channels == 1 ? 4 : 1}; } }; diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index 6712bc678055..75a263e47a91 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1630,6 +1630,27 @@ struct histogram_tuning } }; +template +struct high_bin_histogram_tuning +{ + _CCCL_HOST_DEVICE_API constexpr auto operator()(cuda::compute_capability) const -> cub::HistogramPolicy + { + auto policy = histogram_tuning<128>{}(cuda::compute_capability{}); + policy.high_bin_algorithm = Algorithm; + policy.high_bin_cache = Cache; + policy.high_bin_spill = Spill; + policy.high_bin_aggregation = Aggregation; + policy.high_bin_cache_entries_per_channel = 512; + policy.high_bin_cache_count_replicas = 2; + policy.high_bin_cache_cuckoo_max_bins = 4096; + policy.high_bin_pixels_per_thread = 2; + return policy; + } +}; + using block_sizes = c2h::type_list, cuda::std::integral_constant>; @@ -1740,6 +1761,47 @@ CUB_TEST("DeviceHistogram::MultiHistogramRange can be tuned", "[histogram][devic REQUIRE(d_block_size[0] == target_block_size); } +CUB_TEST("DeviceHistogram high-bin cooperative strategies can be tuned", "[histogram][device]", CUB_SMALL) +{ + constexpr int num_levels = 1026; + c2h::device_vector d_samples{0, 1, 1, 1024}; + c2h::device_vector expected(num_levels - 1, 0); + expected[0] = 1; + expected[1] = 2; + expected[1024] = 1; + + const auto run = [&](auto tuning) { + c2h::device_vector d_histogram(num_levels - 1, 0); + auto env = cuda::execution::tune(tuning); + histogram_even( + thrust::raw_pointer_cast(d_samples.data()), + thrust::raw_pointer_cast(d_histogram.data()), + num_levels, + 0, + num_levels - 1, + static_cast(d_samples.size()), + env); + REQUIRE(d_histogram == expected); + }; + + run(high_bin_histogram_tuning{}); + run(high_bin_histogram_tuning{}); + run(high_bin_histogram_tuning{}); + run(high_bin_histogram_tuning{}); +} + #endif // TEST_LAUNCH != 1 #if _CCCL_COMPILER(GCC, >=, 8) // gcc 7 cannot preserve constexpr-ness from p1 to p2 @@ -1777,9 +1839,17 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) os << p; return os.str(); }; - REQUIRE(to_string(p1) - == "HistogramPolicy { .threads_per_block = 128, .pixels_per_thread = 7, .vec_size = 4" - ", .load_algorithm = BLOCK_LOAD_DIRECT, .load_modifier = LOAD_LDG, .rle_compress = 0" - ", .mem_preference = SMEM, .use_work_stealing = 0, .init_kernel_pdl_trigger_max_bins = 2048 }"); + REQUIRE( + to_string(p1) + == "HistogramPolicy { .threads_per_block = 128, .pixels_per_thread = 7, .vec_size = 4" + ", .load_algorithm = BLOCK_LOAD_DIRECT, .load_modifier = LOAD_LDG, .rle_compress = 0" + ", .mem_preference = SMEM, .use_work_stealing = 0, .init_kernel_pdl_trigger_max_bins = 2048" + ", .high_bin_algorithm = HistogramHighBinAlgorithm::cooperative" + ", .high_bin_cache = HistogramCacheAlgorithm::cuckoo" + ", .high_bin_spill = HistogramSpillAlgorithm::output" + ", .high_bin_aggregation = HistogramAggregationAlgorithm::warp_coalesced" + ", .high_bin_cache_entries_per_channel = 2048" + ", .high_bin_cache_count_replicas = 1, .high_bin_cache_cuckoo_max_bins = 262144" + ", .high_bin_pixels_per_thread = 4 }"); } #endif // _CCCL_COMPILER(GCC, >=, 8) From 9eb769249acacc1e9107cdd1929f4b0e2b842d64 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Sun, 30 Aug 2026 20:45:25 +0000 Subject: [PATCH 04/21] [cub] Harden cooperative histogram dispatch --- .../device/dispatch/dispatch_histogram.cuh | 91 ++++++++++++------- .../dispatch/kernels/kernel_histogram.cuh | 4 +- .../dispatch/tuning/tuning_histogram.cuh | 8 +- cub/test/catch2_test_device_histogram_env.cu | 91 +++++++++++++++++-- 4 files changed, 146 insertions(+), 48 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 1cd47793a93d..1fc03036d79b 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -271,44 +271,67 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( int cooperative_smem_bytes = 0; #if _CCCL_HOSTED() - if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) - { - if (active_policy.high_bin_algorithm == HistogramHighBinAlgorithm::cooperative) - { - if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(use_cooperative))) - { - return error; - } - if (use_cooperative) + NV_IF_TARGET( + NV_IS_HOST, ({ + if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) { - using privatized_decode_op_t = typename SecondLevelArrayT::value_type; - const auto cooperative_kernel = - kernel_source.template HistogramCooperativeKernel(); - const int cache_slots_per_channel = - active_policy.high_bin_cache == HistogramCacheAlgorithm::none - ? 0 - : active_policy.high_bin_cache_entries_per_channel; - cooperative_smem_bytes = - NUM_ACTIVE_CHANNELS * cache_slots_per_channel - * (int{sizeof(::cuda::std::uint32_t)} + active_policy.high_bin_cache_count_replicas * int{sizeof(CounterT)}); - - int cooperative_sm_occupancy = 0; - if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( - cooperative_sm_occupancy, cooperative_kernel, threads_per_block, cooperative_smem_bytes))) + if (active_policy.high_bin_algorithm == HistogramHighBinAlgorithm::cooperative) { - return error; - } - if (cooperative_sm_occupancy > 0) - { - histogram_sweep_occupancy = cooperative_sm_occupancy * sm_count; - } - else - { - use_cooperative = false; + if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(use_cooperative))) + { + return error; + } + if (use_cooperative) + { + using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + const auto cooperative_kernel = + kernel_source.template HistogramCooperativeKernel(); + const int cache_slots_per_channel = + active_policy.high_bin_cache == HistogramCacheAlgorithm::none + ? 0 + : active_policy.high_bin_cache_entries_per_channel; + cooperative_smem_bytes = + NUM_ACTIVE_CHANNELS * cache_slots_per_channel + * (int{sizeof(::cuda::std::uint32_t)} + + active_policy.high_bin_cache_count_replicas * int{sizeof(CounterT)}); + + int max_dynamic_smem_bytes = 0; + if (const auto error = + CubDebug(launcher_factory.max_dynamic_smem_size_for(max_dynamic_smem_bytes, cooperative_kernel))) + { + return error; + } + if (cooperative_smem_bytes > max_dynamic_smem_bytes) + { + use_cooperative = false; + } + else + { + if (const auto error = CubDebug( + launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, cooperative_smem_bytes))) + { + return error; + } + + int cooperative_sm_occupancy = 0; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + cooperative_sm_occupancy, cooperative_kernel, threads_per_block, cooperative_smem_bytes))) + { + return error; + } + if (cooperative_sm_occupancy > 0) + { + histogram_sweep_occupancy = cooperative_sm_occupancy * sm_count; + } + else + { + use_cooperative = false; + } + } + } } } - } - } + })) #endif // _CCCL_HOSTED() if (num_row_pixels * NUM_CHANNELS == row_stride_samples) diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index e493f09d04d2..0adf491843c7 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -804,9 +804,9 @@ struct AgentHistogramCooperative static_assert(policy.high_bin_pixels_per_thread > 0, "Histogram cooperative unroll must be positive"); static_assert( policy.high_bin_cache == HistogramCacheAlgorithm::none - || (policy.high_bin_cache_entries_per_channel > 0 + || (policy.high_bin_cache_entries_per_channel >= 32 && (policy.high_bin_cache_entries_per_channel & (policy.high_bin_cache_entries_per_channel - 1)) == 0), - "Histogram cache entries per channel must be a power of two"); + "Histogram cache entries per channel must be a power of two of at least 32"); namespace cg = ::cooperative_groups; cg::grid_group grid = cg::this_grid(); diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 37972e2334c4..230c1e956fe2 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -85,22 +85,22 @@ namespace detail::histogram } // namespace detail::histogram #if _CCCL_HOSTED() -inline ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) +_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) { return os << detail::histogram::to_string(value); } -inline ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) +_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) { return os << detail::histogram::to_string(value); } -inline ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) +_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) { return os << detail::histogram::to_string(value); } -inline ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) +_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) { return os << detail::histogram::to_string(value); } diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index 75a263e47a91..1694f5f77542 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1633,7 +1633,8 @@ struct histogram_tuning template + cub::HistogramAggregationAlgorithm Aggregation, + int CacheEntriesPerChannel = 512> struct high_bin_histogram_tuning { _CCCL_HOST_DEVICE_API constexpr auto operator()(cuda::compute_capability) const -> cub::HistogramPolicy @@ -1643,7 +1644,7 @@ struct high_bin_histogram_tuning policy.high_bin_cache = Cache; policy.high_bin_spill = Spill; policy.high_bin_aggregation = Aggregation; - policy.high_bin_cache_entries_per_channel = 512; + policy.high_bin_cache_entries_per_channel = CacheEntriesPerChannel; policy.high_bin_cache_count_replicas = 2; policy.high_bin_cache_cuckoo_max_bins = 4096; policy.high_bin_pixels_per_thread = 2; @@ -1763,12 +1764,18 @@ CUB_TEST("DeviceHistogram::MultiHistogramRange can be tuned", "[histogram][devic CUB_TEST("DeviceHistogram high-bin cooperative strategies can be tuned", "[histogram][device]", CUB_SMALL) { - constexpr int num_levels = 1026; - c2h::device_vector d_samples{0, 1, 1, 1024}; - c2h::device_vector expected(num_levels - 1, 0); - expected[0] = 1; - expected[1] = 2; - expected[1024] = 1; + constexpr int num_levels = 1026; + constexpr int num_samples = 32768; + c2h::host_vector h_samples(num_samples); + c2h::host_vector h_expected(num_levels - 1, 0); + for (int i = 0; i < num_samples; ++i) + { + const int sample = i % (num_levels - 1); + h_samples[i] = sample; + ++h_expected[sample]; + } + const c2h::device_vector d_samples = h_samples; + const c2h::device_vector expected = h_expected; const auto run = [&](auto tuning) { c2h::device_vector d_histogram(num_levels - 1, 0); @@ -1800,6 +1807,74 @@ CUB_TEST("DeviceHistogram high-bin cooperative strategies can be tuned", "[histo cub::HistogramCacheAlgorithm::cuckoo, cub::HistogramSpillAlgorithm::output, cub::HistogramAggregationAlgorithm::warp_coalesced>{}); + run(high_bin_histogram_tuning{}); +} + +CUB_TEST("DeviceHistogram high-bin cooperative strategy handles strided rows", "[histogram][device]", CUB_SMALL) +{ + constexpr int num_channels = 4; + constexpr int num_active_channels = 3; + constexpr int num_levels = 1026; + constexpr int num_row_pixels = 512; + constexpr int num_rows = 4; + constexpr int row_stride_pixels = num_row_pixels + 8; + constexpr int row_stride_samples = row_stride_pixels * num_channels; + + c2h::host_vector h_samples(row_stride_samples * num_rows, num_levels - 1); + cuda::std::array, num_active_channels> h_expected{ + c2h::host_vector(num_levels - 1, 0), + c2h::host_vector(num_levels - 1, 0), + c2h::host_vector(num_levels - 1, 0)}; + for (int row = 0; row < num_rows; ++row) + { + for (int pixel = 0; pixel < num_row_pixels; ++pixel) + { + for (int channel = 0; channel < num_active_channels; ++channel) + { + const int sample = (row * num_row_pixels + pixel + channel) % (num_levels - 1); + h_samples[row * row_stride_samples + pixel * num_channels + channel] = sample; + ++h_expected[channel][sample]; + } + } + } + + const c2h::device_vector d_samples = h_samples; + cuda::std::array, num_active_channels> d_histograms{ + c2h::device_vector(num_levels - 1, 0), + c2h::device_vector(num_levels - 1, 0), + c2h::device_vector(num_levels - 1, 0)}; + cuda::std::array histogram_ptrs{ + thrust::raw_pointer_cast(d_histograms[0].data()), + thrust::raw_pointer_cast(d_histograms[1].data()), + thrust::raw_pointer_cast(d_histograms[2].data())}; + constexpr cuda::std::array levels{num_levels, num_levels, num_levels}; + constexpr cuda::std::array lower_levels{0, 0, 0}; + constexpr cuda::std::array upper_levels{num_levels - 1, num_levels - 1, num_levels - 1}; + const auto env = cuda::execution::tune( + high_bin_histogram_tuning{}); + + multi_histogram_even( + thrust::raw_pointer_cast(d_samples.data()), + histogram_ptrs, + levels, + lower_levels, + upper_levels, + num_row_pixels, + num_rows, + row_stride_samples * sizeof(int), + env); + + for (int channel = 0; channel < num_active_channels; ++channel) + { + REQUIRE(d_histograms[channel] == h_expected[channel]); + } } #endif // TEST_LAUNCH != 1 From 975ff8d209cdc3688ed84f3a4cb0286cbbf6ad7d Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Sun, 30 Aug 2026 23:32:15 +0000 Subject: [PATCH 05/21] [cub] Fix cooperative launch compatibility --- cub/cub/detail/launcher/cuda_runtime.cuh | 14 +++++++------- .../device/dispatch/tuning/tuning_histogram.cuh | 8 ++++---- cub/test/catch2_test_env_launch_helper.h | 3 ++- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/cub/cub/detail/launcher/cuda_runtime.cuh b/cub/cub/detail/launcher/cuda_runtime.cuh index 2c35f85ce64a..ad5c8091f076 100644 --- a/cub/cub/detail/launcher/cuda_runtime.cuh +++ b/cub/cub/detail/launcher/cuda_runtime.cuh @@ -116,13 +116,13 @@ struct TripleChevronFactory template _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t LaunchCooperative( dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, Kernel kernel, Args const&... args) - const {NV_IF_ELSE_TARGET( - NV_IS_HOST, - ({ - void* kernel_args[] = {const_cast(static_cast(&args))...}; - return ::cudaLaunchCooperativeKernel(kernel, grid, block, kernel_args, shared_mem, stream); - }), - ({ return ::cudaErrorNotSupported; }))} + const {NV_IF_ELSE_TARGET(NV_IS_HOST, + ({ + void* kernel_args[] = {const_cast(static_cast(&args))...}; + return ::cudaLaunchCooperativeKernel( + reinterpret_cast(kernel), grid, block, kernel_args, shared_mem, stream); + }), + ({ return ::cudaErrorNotSupported; }))} _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const { diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 230c1e956fe2..4c7e745fcbbe 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -85,22 +85,22 @@ namespace detail::histogram } // namespace detail::histogram #if _CCCL_HOSTED() -_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) +_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) +_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) +_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API inline ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) +_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) { return os << detail::histogram::to_string(value); } diff --git a/cub/test/catch2_test_env_launch_helper.h b/cub/test/catch2_test_env_launch_helper.h index 02fe4e28b5ad..326bf4b87389 100644 --- a/cub/test/catch2_test_env_launch_helper.h +++ b/cub/test/catch2_test_env_launch_helper.h @@ -176,7 +176,8 @@ struct stream_registry_factory_t REQUIRE(stream == get_stream_registry_factory_state()->m_stream); } void* kernel_args[] = {const_cast(static_cast(&args))...}; - return cudaLaunchCooperativeKernel(kernel, grid, block, kernel_args, shared_mem, stream); + return cudaLaunchCooperativeKernel( + reinterpret_cast(kernel), grid, block, kernel_args, shared_mem, stream); }), ({ return cudaErrorNotSupported; }))} From 6ae0302f32bbc1ccb1d69dbbc399b8a48ff02cfc Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Mon, 31 Aug 2026 00:11:36 +0000 Subject: [PATCH 06/21] [cub] Keep histogram stream operators inline --- cub/cub/device/dispatch/tuning/tuning_histogram.cuh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 4c7e745fcbbe..37972e2334c4 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -85,22 +85,22 @@ namespace detail::histogram } // namespace detail::histogram #if _CCCL_HOSTED() -_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) +inline ::std::ostream& operator<<(::std::ostream& os, HistogramHighBinAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) +inline ::std::ostream& operator<<(::std::ostream& os, HistogramCacheAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) +inline ::std::ostream& operator<<(::std::ostream& os, HistogramSpillAlgorithm value) { return os << detail::histogram::to_string(value); } -_CCCL_HOST_API ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) +inline ::std::ostream& operator<<(::std::ostream& os, HistogramAggregationAlgorithm value) { return os << detail::histogram::to_string(value); } From ff77e2edd6b1d44ba2e7cf16bc9b96b099c26065 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Mon, 31 Aug 2026 01:02:19 +0000 Subject: [PATCH 07/21] [cub] Port production high-bin histogram optimizations --- .../device/dispatch/dispatch_histogram.cuh | 124 +++- .../dispatch/kernels/kernel_histogram.cuh | 678 +++++++++++++++--- .../dispatch/tuning/tuning_histogram.cuh | 125 +++- cub/test/catch2_test_device_histogram_env.cu | 70 +- 4 files changed, 820 insertions(+), 177 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 1fc03036d79b..c125dbd672c6 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -42,7 +42,9 @@ #include #include #include +#include #include +#include #include #include @@ -52,6 +54,35 @@ CUB_NAMESPACE_BEGIN namespace detail::histogram { +template +struct local_counter +{ + using type = OutputCounterT; +}; + +template +struct local_counter> +{ + using type = typename PolicySelector::local_counter_type; +}; + +template +struct local_counter, + OutputCounterT, + OffsetT> +{ + using type = ::cuda::std::conditional_t< + (sizeof(OutputCounterT) > sizeof(::cuda::std::uint32_t) && sizeof(OffsetT) <= sizeof(::cuda::std::uint32_t)), + ::cuda::std::uint32_t, + OutputCounterT>; +}; + +template +using local_counter_t = typename local_counter::type; + // Maximum number of bins per channel for which we will use a privatized smem strategy static constexpr int max_privatized_smem_bins = 256; @@ -133,11 +164,15 @@ struct DeviceHistogramKernelSource template _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramCooperativeKernel() { + using LocalCounterT = local_counter_t; + static_assert(sizeof(LocalCounterT) <= sizeof(CounterT), + "The output histogram counter must be at least as wide as the local counter"); return &DeviceHistogramCooperativeKernel< PolicyT, NUM_CHANNELS, NUM_ACTIVE_CHANNELS, SampleIteratorT, + LocalCounterT, CounterT, PrivatizedDecodeOpT, OffsetT>; @@ -205,6 +240,8 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}) { + using LocalCounterT = local_counter_t; + ::cuda::compute_capability cc{}; if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) { @@ -247,8 +284,9 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } }(); - const int threads_per_block = active_policy.threads_per_block; - const int pixels_per_thread = active_policy.pixels_per_thread; + const int threads_per_block = active_policy.threads_per_block; + const int high_bin_threads_per_block = active_policy.high_bin_threads(); + const int pixels_per_thread = active_policy.pixels_per_thread; // Get SM count int sm_count; @@ -266,9 +304,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } // Get device occupancy for sweep_kernel - int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; - bool use_cooperative = false; - int cooperative_smem_bytes = 0; + int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; + bool use_cooperative = false; + int cooperative_smem_bytes = 0; + int cooperative_cache_slots_per_channel = 0; #if _CCCL_HOSTED() NV_IF_TARGET( @@ -286,14 +325,14 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( using privatized_decode_op_t = typename SecondLevelArrayT::value_type; const auto cooperative_kernel = kernel_source.template HistogramCooperativeKernel(); - const int cache_slots_per_channel = + const int cache_slots_floor = active_policy.high_bin_cache == HistogramCacheAlgorithm::none ? 0 : active_policy.high_bin_cache_entries_per_channel; - cooperative_smem_bytes = - NUM_ACTIVE_CHANNELS * cache_slots_per_channel + const int cache_bytes_per_slot = + NUM_ACTIVE_CHANNELS * (int{sizeof(::cuda::std::uint32_t)} - + active_policy.high_bin_cache_count_replicas * int{sizeof(CounterT)}); + + active_policy.high_bin_cache_count_replicas * int{sizeof(LocalCounterT)}); int max_dynamic_smem_bytes = 0; if (const auto error = @@ -301,26 +340,55 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { return error; } - if (cooperative_smem_bytes > max_dynamic_smem_bytes) + const int max_slots_by_smem = cache_bytes_per_slot == 0 ? 0 : max_dynamic_smem_bytes / cache_bytes_per_slot; + if (cache_slots_floor > max_slots_by_smem) { use_cooperative = false; } else { - if (const auto error = CubDebug( - launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, cooperative_smem_bytes))) - { + const auto occupancy_for_slots = [&](int slots, int& occupancy) { + const int smem_bytes = slots * cache_bytes_per_slot; + cudaError_t error = launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, smem_bytes); + if (error == cudaSuccess) + { + error = launcher_factory.MaxSmOccupancy( + occupancy, cooperative_kernel, high_bin_threads_per_block, smem_bytes); + } return error; - } + }; int cooperative_sm_occupancy = 0; - if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( - cooperative_sm_occupancy, cooperative_kernel, threads_per_block, cooperative_smem_bytes))) + if (const auto error = CubDebug(occupancy_for_slots(cache_slots_floor, cooperative_sm_occupancy))) { return error; } if (cooperative_sm_occupancy > 0) { + cooperative_cache_slots_per_channel = cache_slots_floor; + const int floor_occupancy = cooperative_sm_occupancy; + for (int candidate = cache_slots_floor == 0 ? 0 : cache_slots_floor << 1; + candidate > 0 && candidate <= max_slots_by_smem; + candidate <<= 1) + { + int candidate_occupancy = 0; + if (const auto error = CubDebug(occupancy_for_slots(candidate, candidate_occupancy))) + { + return error; + } + if (candidate_occupancy < floor_occupancy) + { + break; + } + cooperative_cache_slots_per_channel = candidate; + cooperative_sm_occupancy = candidate_occupancy; + } + cooperative_smem_bytes = cooperative_cache_slots_per_channel * cache_bytes_per_slot; + if (const auto error = CubDebug( + launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, cooperative_smem_bytes))) + { + return error; + } histogram_sweep_occupancy = cooperative_sm_occupancy * sm_count; } else @@ -344,7 +412,8 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy const int selected_pixels_per_thread = use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread; - int pixels_per_tile = threads_per_block * selected_pixels_per_thread; + const int selected_threads_per_block = use_cooperative ? high_bin_threads_per_block : threads_per_block; + int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); int blocks_per_col = @@ -369,7 +438,11 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( !use_cooperative || active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized; allocation_sizes[CHANNEL] = needs_privatized_storage - ? size_t(num_thread_blocks) * (num_privatized_levels[CHANNEL] - 1) * kernel_source.CounterSize() + ? size_t(num_thread_blocks) + * ((use_cooperative && active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + ? (num_output_levels[CHANNEL] - 1) + : (num_privatized_levels[CHANNEL] - 1)) + * (use_cooperative ? sizeof(LocalCounterT) : kernel_source.CounterSize()) : 0; } @@ -394,11 +467,16 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( // Wrap arrays so we can pass them by-value to the kernel ::cuda::std::array d_privatized_histograms_wrapper; + ::cuda::std::array d_cooperative_privatized_histograms_wrapper; ::cuda::std::array num_privatized_bins_wrapper; ::cuda::std::array num_output_bins_wrapper; auto* typed_allocations = reinterpret_cast(allocations); ::cuda::std::copy(typed_allocations, typed_allocations + NUM_ACTIVE_CHANNELS, d_privatized_histograms_wrapper.begin()); + auto* local_typed_allocations = reinterpret_cast(allocations); + ::cuda::std::copy(local_typed_allocations, + local_typed_allocations + NUM_ACTIVE_CHANNELS, + d_cooperative_privatized_histograms_wrapper.begin()); auto minus_one = ::cuda::proclaim_return_type([](int levels) { return levels - 1; @@ -416,27 +494,23 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( using privatized_decode_op_t = typename SecondLevelArrayT::value_type; const dim3 cooperative_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; - const int cache_slots_per_channel = - active_policy.high_bin_cache == HistogramCacheAlgorithm::none - ? 0 - : active_policy.high_bin_cache_entries_per_channel; const auto cooperative_kernel = kernel_source.template HistogramCooperativeKernel(); if (const auto error = CubDebug(launcher_factory.LaunchCooperative( cooperative_grid_dims, - dim3{static_cast(threads_per_block)}, + dim3{static_cast(high_bin_threads_per_block)}, cooperative_smem_bytes, stream, cooperative_kernel, d_samples, num_output_bins_wrapper, d_output_histograms, - d_privatized_histograms_wrapper, + d_cooperative_privatized_histograms_wrapper, second_level_array, num_row_pixels, num_rows, row_stride_samples, - cache_slots_per_channel))) + cooperative_cache_slots_per_channel))) { return error; } diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index 0adf491843c7..b3251480bffa 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -17,19 +17,173 @@ #include #include #include +#include #include #include +#include #include #include -#if !_CCCL_COMPILER(NVRTC) -# include -#endif // !_CCCL_COMPILER(NVRTC) +#include CUB_NAMESPACE_BEGIN namespace detail::histogram { +template +struct fast_divide_by_constant +{ + static_assert(::cuda::std::is_unsigned_v, "fast_divide_by_constant requires an unsigned integer divisor type"); + static_assert(sizeof(UInt) == 4 || sizeof(UInt) == 8, "fast_divide_by_constant supports 32-bit or 64-bit divisors"); + + static constexpr int bits = static_cast(sizeof(UInt) * 8); + + UInt magic; + unsigned char shift; + unsigned char mode; + + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static int count_leading_zeros(::cuda::std::uint64_t value) + { + NV_IF_ELSE_TARGET(NV_IS_DEVICE, + (return value == 0 ? 64 : __clzll(static_cast(value));), + (return value == 0 ? 64 : __builtin_clzll(value);)); + } + + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static int count_leading_zeros(::cuda::std::uint32_t value) + { + NV_IF_ELSE_TARGET(NV_IS_DEVICE, + (return value == 0 ? 32 : __clz(static_cast(value));), + (return value == 0 ? 32 : __builtin_clz(value);)); + } + + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static int ceil_log2(UInt divisor) + { + if (divisor <= UInt{1}) + { + return 0; + } + if constexpr (sizeof(UInt) == 4) + { + return bits - count_leading_zeros(static_cast<::cuda::std::uint32_t>(divisor - UInt{1})); + } + else + { + return bits - count_leading_zeros(static_cast<::cuda::std::uint64_t>(divisor - UInt{1})); + } + } + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void Init(UInt divisor) + { + if (divisor <= UInt{1}) + { + magic = UInt{0}; + shift = 0; + mode = 0; + return; + } + if ((divisor & (divisor - UInt{1})) == UInt{0}) + { + magic = UInt{0}; + shift = static_cast(ceil_log2(divisor)); + mode = 1; + return; + } + + const int log2_divisor = ceil_log2(divisor); + if constexpr (sizeof(UInt) == 8) + { +#if _CCCL_HAS_INT128() + const __uint128_t numerator = static_cast<__uint128_t>(1) << (bits + log2_divisor); + const __uint128_t denominator = static_cast<__uint128_t>(divisor); + magic = static_cast((numerator + denominator - 1) / denominator); +#else + UInt quotient = 0; + UInt remainder = 0; + for (int bit = bits + log2_divisor; bit >= 0; --bit) + { + UInt next_remainder = (remainder << 1) | (bit == bits + log2_divisor ? UInt{1} : UInt{0}); + const bool carry = (remainder >> (bits - 1)) != 0; + const UInt quotient_bit = (carry || next_remainder >= divisor) ? UInt{1} : UInt{0}; + if (quotient_bit != 0) + { + next_remainder -= divisor; + } + remainder = next_remainder; + quotient = (quotient << 1) | quotient_bit; + } + magic = quotient + (remainder != 0 ? UInt{1} : UInt{0}); +#endif + } + else + { + const ::cuda::std::uint64_t numerator = ::cuda::std::uint64_t{1} << (bits + log2_divisor); + const ::cuda::std::uint64_t denominator = static_cast<::cuda::std::uint64_t>(divisor); + magic = static_cast((numerator + denominator - 1) / denominator); + } + shift = static_cast(log2_divisor); + mode = 2; + } + + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE UInt Divide(UInt numerator) const + { + if (mode == 0) + { + return numerator; + } + if (mode == 1) + { + return numerator >> shift; + } + + UInt high; + if constexpr (sizeof(UInt) == 8) + { + NV_IF_ELSE_TARGET( + NV_IS_DEVICE, + (high = static_cast( + __umul64hi(static_cast(magic), static_cast(numerator)));), + ({ +#if _CCCL_HAS_INT128() + high = static_cast((static_cast<__uint128_t>(magic) * static_cast<__uint128_t>(numerator)) >> bits); +#else + const ::cuda::std::uint64_t a_low = static_cast<::cuda::std::uint32_t>(magic); + const ::cuda::std::uint64_t a_high = magic >> 32; + const ::cuda::std::uint64_t b_low = static_cast<::cuda::std::uint32_t>(numerator); + const ::cuda::std::uint64_t b_high = numerator >> 32; + const ::cuda::std::uint64_t low_low = a_low * b_low; + const ::cuda::std::uint64_t low_high = a_low * b_high; + const ::cuda::std::uint64_t high_low = a_high * b_low; + const ::cuda::std::uint64_t high_high = a_high * b_high; + const ::cuda::std::uint64_t middle = + (low_low >> 32) + static_cast<::cuda::std::uint32_t>(low_high) + + static_cast<::cuda::std::uint32_t>(high_low); + high = high_high + (low_high >> 32) + (high_low >> 32) + (middle >> 32); +#endif + })); + } + else + { + high = static_cast( + (static_cast<::cuda::std::uint64_t>(magic) * static_cast<::cuda::std::uint64_t>(numerator)) >> bits); + } + return (((numerator - high) >> 1) + high) >> (shift - 1); + } +}; + +template +_CCCL_DEVICE _CCCL_FORCEINLINE const SampleValueT* sample_native_pointer(SampleIteratorT itr) +{ + if constexpr (::cuda::std::is_pointer_v) + { + return itr; + } + else + { + return NativePointer(itr); + } + _CCCL_UNREACHABLE(); +} + template _CCCL_DEVICE _CCCL_FORCEINLINE void histogram_atomic_add(CounterT* address, CounterT value) { @@ -56,8 +210,24 @@ struct Transforms template struct SearchTransform { + static constexpr bool is_range_transform = true; + + struct BracketCacheT + { + LevelT lo{}; + LevelT hi{}; + int bin = -1; + }; + LevelIteratorT d_levels; // Pointer to levels array int num_output_levels; // Number of levels in array + LevelT first{}; + LevelT middle{}; + LevelT last{}; + float inverse_scale_low = 0.0f; + float inverse_scale_high = 0.0f; + int middle_bin = 0; + bool has_precompute = false; //! @brief Initializer //! @@ -67,6 +237,58 @@ struct Transforms { this->d_levels = d_levels_; this->num_output_levels = num_output_levels_; + this->has_precompute = false; + } + + template + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static float level_distance(T upper, T lower) + { + if constexpr (::cuda::std::is_integral_v) + { + using UnsignedT = ::cuda::std::make_unsigned_t; + return static_cast(static_cast(upper) - static_cast(lower)); + } + else + { + return static_cast(upper - lower); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void PrecomputeOnDevice(int interpolation_min_bins) + { + const int num_bins = num_output_levels - 1; + if (num_bins < interpolation_min_bins) + { + return; + } + + using WrappedLevelIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + LevelIteratorT>; + WrappedLevelIteratorT wrapped_levels(d_levels); + const LevelT first_level = wrapped_levels[0]; + const LevelT last_level = wrapped_levels[num_bins]; + if (!(first_level < last_level)) + { + return; + } + + first = first_level; + last = last_level; + middle_bin = num_bins >> 1; + middle = wrapped_levels[middle_bin]; + if (!(first < middle) || !(middle < last)) + { + middle_bin = 0; + inverse_scale_low = static_cast(num_bins) / level_distance(last, first); + } + else + { + inverse_scale_low = static_cast(middle_bin) / level_distance(middle, first); + inverse_scale_high = static_cast(num_bins - middle_bin) / level_distance(last, middle); + } + has_precompute = true; } // Method for converting samples to bin-ids @@ -84,13 +306,60 @@ struct Transforms WrappedLevelIteratorT wrapped_levels(d_levels); const int num_bins = num_output_levels - 1; - if (valid) + if (!valid) { - bin = UpperBound(wrapped_levels, num_output_levels, static_cast(sample)) - 1; - if (bin >= num_bins) - { - bin = -1; - } + return; + } + + const LevelT value = static_cast(sample); + if (!has_precompute) + { + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + bin = bin < num_bins ? bin : -1; + return; + } + if (value < first || !(value < last)) + { + bin = -1; + return; + } + + int guess = value < middle || middle_bin == 0 + ? static_cast(level_distance(value, first) * inverse_scale_low) + : middle_bin + static_cast(level_distance(value, middle) * inverse_scale_high); + guess = guess < 0 ? 0 : (guess < num_bins ? guess : num_bins - 1); + const LevelT lo = wrapped_levels[guess]; + const LevelT hi = wrapped_levels[guess + 1]; + if (!(value < lo) && value < hi) + { + bin = guess; + return; + } + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + } + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void + BinSelect(_SampleT sample, int& bin, bool valid, BracketCacheT& bracket) const + { + const LevelT value = static_cast(sample); + if (valid && bracket.bin >= 0 && !(value < bracket.lo) && value < bracket.hi) + { + bin = bracket.bin; + return; + } + + BinSelect(sample, bin, valid); + if (bin >= 0) + { + using WrappedLevelIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + LevelIteratorT>; + WrappedLevelIteratorT wrapped_levels(d_levels); + bracket.lo = wrapped_levels[bin]; + bracket.hi = wrapped_levels[bin + 1]; + bracket.bin = bin; } } }; @@ -138,14 +407,18 @@ struct Transforms ::cuda::std::is_integral; #endif // !_CCCL_HAS_INT128() + using FractionStorageT = ::cuda::std::_If::value, IntArithmeticT, CommonT>; + union ScaleT { // Used when CommonT is not floating-point to avoid intermediate // rounding errors (see NVIDIA/cub#489). struct FractionT { - CommonT bins; - CommonT range; + FractionStorageT bins; + FractionStorageT range; + fast_divide_by_constant range_divider; + bool bins_equal_range; } fraction; // Used when CommonT is floating-point as an optimization. @@ -170,8 +443,20 @@ struct Transforms ComputeScale(int num_levels, T max_level, T min_level, ::cuda::std::false_type /* is_fp */) { ScaleT result; - result.fraction.bins = static_cast(num_levels - 1); - result.fraction.range = static_cast(max_level - min_level); + result.fraction.bins = static_cast(num_levels - 1); + if constexpr (::cuda::std::is_integral_v) + { + using UnsignedT = ::cuda::std::make_unsigned_t; + const UnsignedT distance = + static_cast(static_cast(max_level) - static_cast(min_level)); + result.fraction.range = static_cast(distance); + } + else + { + result.fraction.range = static_cast(max_level - min_level); + } + result.fraction.bins_equal_range = result.fraction.bins == result.fraction.range; + result.fraction.range_divider.Init(static_cast(result.fraction.range)); return result; } @@ -255,9 +540,15 @@ struct Transforms template ::value, int> = 0> _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int ComputeBin(T sample, T min_level, ScaleT scale) const { - return static_cast( - (static_cast(sample - min_level) * static_cast(scale.fraction.bins)) - / static_cast(scale.fraction.range)); + using UnsignedT = ::cuda::std::make_unsigned_t; + const IntArithmeticT distance = static_cast( + static_cast(static_cast(sample) - static_cast(min_level))); + if (scale.fraction.bins_equal_range) + { + return static_cast(distance); + } + const IntArithmeticT numerator = distance * static_cast(scale.fraction.bins); + return static_cast(scale.fraction.range_divider.Divide(numerator)); } template ::value, int> = 0> @@ -277,6 +568,12 @@ struct Transforms #endif // _CCCL_HAS_NVFP16() public: + static constexpr bool is_range_transform = false; + struct BracketCacheT + {}; + + _CCCL_DEVICE _CCCL_FORCEINLINE void PrecomputeOnDevice(int) {} + //! @brief Initializes the ScaleTransform for the given parameters _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void Init(int num_levels, LevelT max_level, LevelT min_level) { @@ -302,6 +599,11 @@ struct Transforms // Pass-through bin transform operator struct PassThruTransform { + static constexpr bool is_range_transform = false; + struct BracketCacheT + {}; + + _CCCL_DEVICE _CCCL_FORCEINLINE void PrecomputeOnDevice(int) {} // GCC 14 rightfully warns that when a value-initialized array of this struct is copied using memcpy, uninitialized // bytes may be accessed. To avoid this, we add a dummy member, so value initialization actually initializes the memory. #if _CCCL_COMPILER(GCC, >=, 13) @@ -781,6 +1083,7 @@ template #if _CCCL_HAS_CONCEPTS() @@ -791,7 +1094,7 @@ struct AgentHistogramCooperative _CCCL_DEVICE _CCCL_FORCEINLINE static void Consume( const SampleIteratorT d_samples, const ::cuda::std::array num_output_bins_wrapper, - ::cuda::std::array d_output_histograms_wrapper, + ::cuda::std::array d_output_histograms_wrapper, ::cuda::std::array d_privatized_histograms_wrapper, const ::cuda::std::array decode_op_wrapper, const OffsetT num_row_pixels, @@ -819,7 +1122,7 @@ struct AgentHistogramCooperative for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); bin += total_threads) { - d_output_histograms_wrapper[ch][bin] = CounterT{0}; + d_output_histograms_wrapper[ch][bin] = OutputCounterT{0}; } } grid.sync(); @@ -868,6 +1171,17 @@ struct AgentHistogramCooperative const OffsetT chunk_count = ::cuda::ceil_div(total_pixels, chunk); const unsigned int lane_id = threadIdx.x & 0x1f; const bool contiguous_input = num_rows == 1; + + PrivatizedDecodeOpT decode_op[NumActiveChannels]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + decode_op[ch] = decode_op_wrapper[ch]; + decode_op[ch].PrecomputeOnDevice(policy.high_bin_interpolation_min_bins); + } + + constexpr bool use_mru_cache = NumActiveChannels == 1 && PrivatizedDecodeOpT::is_range_transform; + typename PrivatizedDecodeOpT::BracketCacheT bracket_cache[NumActiveChannels]; int pending_bin[NumActiveChannels]; CounterT pending_count[NumActiveChannels]; @@ -878,108 +1192,250 @@ struct AgentHistogramCooperative pending_count[ch] = CounterT{0}; } - for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) - { - const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + const auto spill_bin = [&](int ch, int selected_bin, CounterT contribution) { + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) { - const OffsetT pixel = first_pixel + static_cast(item) * step; - const bool valid = pixel < total_pixels; - OffsetT pixel_offset = 0; - if (valid) + if (pending_bin[ch] == selected_bin) + { + pending_count[ch] += contribution; + return; + } + if (pending_bin[ch] >= 0) { - if (contiguous_input) + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - pixel_offset = pixel * NumChannels; + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); } else { - const OffsetT row = pixel / num_row_pixels; - pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + histogram_atomic_add(&d_output_histograms_wrapper[ch][pending_bin[ch]], + static_cast(pending_count[ch])); + } + } + pending_bin[ch] = selected_bin; + pending_count[ch] = contribution; + } + else if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + atomicAdd_block(&block_histogram[selected_bin], contribution); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][selected_bin], static_cast(contribution)); + } + }; + + const auto update_bin = [&](int ch, int selected_bin, CounterT contribution) { + auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo + && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; + if (!histogram_cache_probe( + channel_keys, + channel_counts, + count_replicas, + selected_bin, + contribution, + cache_mask, + cache_log2, + policy.high_bin_cache, + use_second_probe)) + { + spill_bin(ch, selected_bin, contribution); + } + }; + + const auto consume_bin = [&](int ch, int bin) { + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::warp_coalesced) + { + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (const unsigned int active = __activemask(); + const unsigned int peers = __match_any_sync(active, static_cast(bin)); + const int leader = __ffs(static_cast(peers)) - 1; + if (bin >= 0 && static_cast(lane_id) == leader) { + update_bin(ch, bin, static_cast(__popc(peers))); + }), + (if (bin >= 0) { update_bin(ch, bin, CounterT{1}); })); + } + else if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) + { + if (bin >= 0) + { + update_bin(ch, bin, CounterT{1}); + } + } + else if (bin >= 0) + { + update_bin(ch, bin, CounterT{1}); + } + }; + + using SampleValueT = it_value_t; + if constexpr (NumActiveChannels == 1) + { + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + { + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + SampleValueT staged_samples[unroll]; + bool valid_samples[unroll]; + int bins[unroll]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + const OffsetT pixel = first_pixel + static_cast(item) * step; + valid_samples[item] = pixel < total_pixels; + OffsetT pixel_offset{}; + if (valid_samples[item]) + { + if (contiguous_input) + { + pixel_offset = pixel * NumChannels; + } + else + { + const OffsetT row = pixel / num_row_pixels; + pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + } + staged_samples[item] = d_samples[pixel_offset]; } } _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + for (int item = 0; item < unroll; ++item) { int bin = -1; - if (valid) + if (valid_samples[item]) { - const auto sample = d_samples[pixel_offset + ch]; - decode_op_wrapper[ch].template BinSelect(sample, bin, true); - if (bin < 0 || bin >= num_output_bins_wrapper[ch]) + if constexpr (use_mru_cache) + { + decode_op[0].template BinSelect(staged_samples[item], bin, true, bracket_cache[0]); + } + else + { + decode_op[0].template BinSelect(staged_samples[item], bin, true); + } + if (bin >= num_output_bins_wrapper[0]) { bin = -1; } } + bins[item] = bin; + } - const auto update_bin = [&](int selected_bin, CounterT contribution) { - auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = - cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; - const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo - && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; - if (!histogram_cache_probe( - channel_keys, - channel_counts, - count_replicas, - selected_bin, - contribution, - cache_mask, - cache_log2, - policy.high_bin_cache, - use_second_probe)) + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + consume_bin(0, bins[item]); + } + } + } + else + { + const auto consume_pixels = [&](auto load_pixel) { + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + { + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + SampleValueT samples[NumActiveChannels]; + int bins[NumActiveChannels]; + load_pixel(pixel, valid, samples); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + int bin = -1; + if (valid) { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[selected_bin], contribution); - } - else - { - histogram_atomic_add(&d_output_histograms_wrapper[ch][selected_bin], contribution); + decode_op[ch].template BinSelect(samples[ch], bin, true); + if (bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } } + bins[ch] = bin; } - }; - if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::warp_coalesced) - { - NV_IF_ELSE_TARGET( - NV_PROVIDES_SM_70, - (const unsigned int active = __activemask(); - const unsigned int peers = __match_any_sync(active, static_cast(bin)); - const int leader = __ffs(static_cast(peers)) - 1; - if (bin >= 0 && static_cast(lane_id) == leader) { - update_bin(bin, static_cast(__popc(peers))); - }), - (if (bin >= 0) { update_bin(bin, CounterT{1}); })); + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + consume_bin(ch, bins[ch]); + } } - else if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) - { - if (bin >= 0) + } + }; + + if constexpr ((NumChannels == 2 || NumChannels == 4) && ::cuda::std::is_trivially_copyable_v) + { + using PixelT = typename CubVector::Type; + const auto* native_base = sample_native_pointer(d_samples); + const bool vectorizable = contiguous_input && native_base != nullptr + && (reinterpret_cast(native_base) & (alignof(PixelT) - 1)) == 0; + if (vectorizable) + { + const PixelT* pixels = reinterpret_cast(native_base); + consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { + if (valid) { - if (pending_bin[ch] == bin) + const PixelT packed_pixel = pixels[pixel]; + const SampleValueT* pixel_lanes = reinterpret_cast(&packed_pixel); + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - ++pending_count[ch]; + samples[ch] = pixel_lanes[ch]; } - else + } + }); + } + else + { + consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { + OffsetT pixel_offset{}; + if (valid) + { + const OffsetT row = contiguous_input ? OffsetT{0} : pixel / num_row_pixels; + const OffsetT col = contiguous_input ? pixel : pixel - row * num_row_pixels; + pixel_offset = row * row_stride_samples + col * NumChannels; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + if (valid) { - if (pending_bin[ch] >= 0) - { - update_bin(pending_bin[ch], pending_count[ch]); - } - pending_bin[ch] = bin; - pending_count[ch] = CounterT{1}; + samples[ch] = d_samples[pixel_offset + ch]; } } + }); + } + } + else + { + consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { + OffsetT pixel_offset{}; + if (valid) + { + const OffsetT row = contiguous_input ? OffsetT{0} : pixel / num_row_pixels; + const OffsetT col = contiguous_input ? pixel : pixel - row * num_row_pixels; + pixel_offset = row * row_stride_samples + col * NumChannels; } - else if (bin >= 0) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - update_bin(bin, CounterT{1}); + if (valid) + { + samples[ch] = d_samples[pixel_offset + ch]; + } } - } + }); } } @@ -990,31 +1446,16 @@ struct AgentHistogramCooperative { if (pending_bin[ch] >= 0) { - auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; - const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo - && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; - if (!histogram_cache_probe( - channel_keys, - channel_counts, - count_replicas, - pending_bin[ch], - pending_count[ch], - cache_mask, - cache_log2, - policy.high_bin_cache, - use_second_probe)) + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) - { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); - } - else - { - histogram_atomic_add(&d_output_histograms_wrapper[ch][pending_bin[ch]], pending_count[ch]); - } + CounterT* block_histogram = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][pending_bin[ch]], + static_cast(pending_count[ch])); } } } @@ -1043,11 +1484,11 @@ struct AgentHistogramCooperative { CounterT* block_histogram = d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - block_histogram[key] += count; + atomicAdd_block(&block_histogram[key], count); } else { - histogram_atomic_add(&d_output_histograms_wrapper[ch][key], count); + histogram_atomic_add(&d_output_histograms_wrapper[ch][key], static_cast(count)); } } } @@ -1063,10 +1504,11 @@ struct AgentHistogramCooperative const unsigned int num_bins = static_cast(num_output_bins_wrapper[ch]); for (unsigned int bin = tid_global; bin < num_bins; bin += total_threads) { - CounterT total = CounterT{0}; + OutputCounterT total = OutputCounterT{0}; for (unsigned int block = 0; block < gridDim.x; ++block) { - total += d_privatized_histograms_wrapper[ch][static_cast(block) * num_bins + bin]; + total += static_cast( + d_privatized_histograms_wrapper[ch][static_cast(block) * num_bins + bin]); } d_output_histograms_wrapper[ch][bin] = total; } @@ -1081,16 +1523,17 @@ template #if _CCCL_HAS_CONCEPTS() requires histogram_policy_selector #endif // _CCCL_HAS_CONCEPTS() -__launch_bounds__(int(current_policy().threads_per_block)) +__launch_bounds__(int(current_policy().high_bin_threads())) _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramCooperativeKernel( _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, - ::cuda::std::array d_output_histograms_wrapper, + ::cuda::std::array d_output_histograms_wrapper, ::cuda::std::array d_privatized_histograms_wrapper, _CCCL_GRID_CONSTANT const ::cuda::std::array decode_op_wrapper, _CCCL_GRID_CONSTANT const OffsetT num_row_pixels, @@ -1104,6 +1547,7 @@ __launch_bounds__(int(current_policy().threads_per_block)) NumActiveChannels, SampleIteratorT, CounterT, + OutputCounterT, PrivatizedDecodeOpT, OffsetT>::Consume(d_samples, num_output_bins_wrapper, diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 37972e2334c4..7c78a11ddd20 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -121,13 +121,20 @@ struct HistogramPolicy int init_kernel_pdl_trigger_max_bins; //!< Maximum number of bins for the init kernel to trigger the histogram kernel //!< early using PDL HistogramHighBinAlgorithm high_bin_algorithm = HistogramHighBinAlgorithm::cooperative; - HistogramCacheAlgorithm high_bin_cache = HistogramCacheAlgorithm::cuckoo; - HistogramSpillAlgorithm high_bin_spill = HistogramSpillAlgorithm::output; - HistogramAggregationAlgorithm high_bin_aggregation = HistogramAggregationAlgorithm::warp_coalesced; + HistogramCacheAlgorithm high_bin_cache = HistogramCacheAlgorithm::single_probe; + HistogramSpillAlgorithm high_bin_spill = HistogramSpillAlgorithm::global_memory_privatized; + HistogramAggregationAlgorithm high_bin_aggregation = HistogramAggregationAlgorithm::rle; int high_bin_cache_entries_per_channel = 2048; int high_bin_cache_count_replicas = 1; int high_bin_cache_cuckoo_max_bins = 262144; int high_bin_pixels_per_thread = 4; + int high_bin_threads_per_block = 0; //!< High-bin block size; 0 inherits threads_per_block + int high_bin_interpolation_min_bins = 512; + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_threads() const noexcept + { + return high_bin_threads_per_block != 0 ? high_bin_threads_per_block : threads_per_block; + } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const HistogramPolicy& lhs, const HistogramPolicy& rhs) noexcept @@ -142,7 +149,9 @@ struct HistogramPolicy && lhs.high_bin_cache_entries_per_channel == rhs.high_bin_cache_entries_per_channel && lhs.high_bin_cache_count_replicas == rhs.high_bin_cache_count_replicas && lhs.high_bin_cache_cuckoo_max_bins == rhs.high_bin_cache_cuckoo_max_bins - && lhs.high_bin_pixels_per_thread == rhs.high_bin_pixels_per_thread; + && lhs.high_bin_pixels_per_thread == rhs.high_bin_pixels_per_thread + && lhs.high_bin_threads_per_block == rhs.high_bin_threads_per_block + && lhs.high_bin_interpolation_min_bins == rhs.high_bin_interpolation_min_bins; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -165,7 +174,9 @@ struct HistogramPolicy << ", .high_bin_cache_entries_per_channel = " << p.high_bin_cache_entries_per_channel << ", .high_bin_cache_count_replicas = " << p.high_bin_cache_count_replicas << ", .high_bin_cache_cuckoo_max_bins = " << p.high_bin_cache_cuckoo_max_bins - << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread << " }"; + << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread + << ", .high_bin_threads_per_block = " << p.high_bin_threads_per_block + << ", .high_bin_interpolation_min_bins = " << p.high_bin_interpolation_min_bins << " }"; } #endif // _CCCL_HOSTED() }; @@ -417,13 +428,14 @@ public: false, 2048, HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::cuckoo, - HistogramSpillAlgorithm::output, - HistogramAggregationAlgorithm::warp_coalesced, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, 4096, 1, 262144, - 4}; + 4, + 0}; } else { @@ -439,16 +451,64 @@ public: false, 2048, HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::cuckoo, - HistogramSpillAlgorithm::output, - HistogramAggregationAlgorithm::warp_coalesced, - 2048, - 2, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + 4096, + 1, 262144, - 4}; + 4, + 0}; } } + if (counter_size == 4 && sample_is_primitive && num_channels == 1 && num_active_channels == 1 + && (sample_size == 4 || sample_size == 8)) + { + return HistogramPolicy{ + 384, + t_scale(16), + 4, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + true, + SMEM, + false, + 0, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + 4096, + 1, + 262144, + 4, + is_even ? 768 : 512}; + } + + if (counter_size == 4 && sample_is_primitive && num_channels >= 2) + { + return HistogramPolicy{ + 384, + t_scale(16), + 4, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + true, + SMEM, + false, + 0, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + 1024, + 4, + 262144, + 4, + 1024}; + } + // sample_size 2/4/8 showed no benefit over SM90 during verification benchmarks // multi.even and multi.range: none of the found tunings surpassed the SM90 tuning during verification benchmarks } @@ -470,13 +530,14 @@ public: false, 2048, HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::cuckoo, - HistogramSpillAlgorithm::output, - HistogramAggregationAlgorithm::warp_coalesced, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, is_even ? 4096 : 2048, - is_even ? 1 : 2, + 1, 262144, - 4}; + 4, + 0}; } else if (sample_size == 2) { @@ -491,13 +552,14 @@ public: false, 2048, HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::cuckoo, - HistogramSpillAlgorithm::output, - HistogramAggregationAlgorithm::warp_coalesced, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, is_even ? 4096 : 2048, - is_even ? 1 : 2, + 1, 262144, - 4}; + 4, + 0}; } } } @@ -514,13 +576,14 @@ public: false, 0, HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::cuckoo, - HistogramSpillAlgorithm::output, - HistogramAggregationAlgorithm::warp_coalesced, - num_active_channels == 1 ? (is_even ? 4096 : 2048) : 512, - is_even ? 1 : (num_active_channels == 1 ? 2 : 4), + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + num_active_channels == 1 ? 4096 : 1024, + num_active_channels == 1 ? 1 : 4, 262144, - num_active_channels == 1 ? 4 : 1}; + 4, + 0}; } }; diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index 1694f5f77542..7f820d13aa70 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -66,6 +66,33 @@ CUB_TEST_CASE("DeviceHistogram::HistogramEven works with default environment", " REQUIRE(d_histogram == expected); } +CUB_TEST_CASE("DeviceHistogram::HistogramEven supports wide output counters with default tuning", + "[histogram][device]", + CUB_SMALL) +{ + using counter_t = unsigned long long; + + const auto d_samples = c2h::device_vector{0, 2, 1, 0, 3, 4, 2, 1}; + const int num_samples = static_cast(d_samples.size()); + const int num_levels = 6; + const unsigned int lower_level = 0; + const unsigned int upper_level = 5; + auto d_histogram = c2h::device_vector(num_levels - 1, counter_t{0}); + + REQUIRE( + cudaSuccess + == cub::DeviceHistogram::HistogramEven( + thrust::raw_pointer_cast(d_samples.data()), + thrust::raw_pointer_cast(d_histogram.data()), + num_levels, + lower_level, + upper_level, + num_samples)); + + const c2h::device_vector expected{2, 2, 2, 1, 1}; + REQUIRE(d_histogram == expected); +} + CUB_TEST_CASE("DeviceHistogram::HistogramEven works with user provided memory and environment", "[histogram][device]", CUB_SMALL) @@ -1648,10 +1675,35 @@ struct high_bin_histogram_tuning policy.high_bin_cache_count_replicas = 2; policy.high_bin_cache_cuckoo_max_bins = 4096; policy.high_bin_pixels_per_thread = 2; + policy.high_bin_threads_per_block = 128; return policy; } }; +template +struct histogram_tuning_with_local_counter : histogram_tuning +{ + using local_counter_type = LocalCounterT; +}; + +using mixed_counter_tuning_t = histogram_tuning_with_local_counter<128, unsigned int>; +using legacy_counter_tuning_t = histogram_tuning<128>; + +static_assert( + cuda::std::is_same_v, + unsigned int>); +static_assert( + cuda::std::is_same_v, + unsigned long long>); +using production_counter_selector = + cub::detail::histogram::policy_selector_from_types; +static_assert( + cuda::std::is_same_v, + unsigned int>); +static_assert(cuda::std::is_same_v< + cub::detail::histogram::local_counter_t, + unsigned long long>); + using block_sizes = c2h::type_list, cuda::std::integral_constant>; @@ -1920,11 +1972,21 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) ", .load_algorithm = BLOCK_LOAD_DIRECT, .load_modifier = LOAD_LDG, .rle_compress = 0" ", .mem_preference = SMEM, .use_work_stealing = 0, .init_kernel_pdl_trigger_max_bins = 2048" ", .high_bin_algorithm = HistogramHighBinAlgorithm::cooperative" - ", .high_bin_cache = HistogramCacheAlgorithm::cuckoo" - ", .high_bin_spill = HistogramSpillAlgorithm::output" - ", .high_bin_aggregation = HistogramAggregationAlgorithm::warp_coalesced" + ", .high_bin_cache = HistogramCacheAlgorithm::single_probe" + ", .high_bin_spill = HistogramSpillAlgorithm::global_memory_privatized" + ", .high_bin_aggregation = HistogramAggregationAlgorithm::rle" ", .high_bin_cache_entries_per_channel = 2048" ", .high_bin_cache_count_replicas = 1, .high_bin_cache_cuckoo_max_bins = 262144" - ", .high_bin_pixels_per_thread = 4 }"); + ", .high_bin_pixels_per_thread = 4, .high_bin_threads_per_block = 0" + ", .high_bin_interpolation_min_bins = 512 }"); + + constexpr auto high_bin_policy = [] { + auto policy = + cub::HistogramPolicy{128, 1, 1, cub::BLOCK_LOAD_DIRECT, cub::LOAD_DEFAULT, false, cub::SMEM, false, 0}; + policy.high_bin_threads_per_block = 512; + return policy; + }(); + STATIC_REQUIRE(high_bin_policy.high_bin_threads() == 512); + STATIC_REQUIRE(p1.high_bin_threads() == p1.threads_per_block); } #endif // _CCCL_COMPILER(GCC, >=, 8) From 3b09ba181eaf4abca273c366bdc495a29dfd3d8d Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Mon, 31 Aug 2026 01:31:21 +0000 Subject: [PATCH 08/21] [cub] Handle full-width histogram ranges --- cub/cub/device/dispatch/dispatch_histogram.cuh | 6 +++--- cub/cub/device/dispatch/kernels/kernel_histogram.cuh | 11 +++++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index c125dbd672c6..2996c94ebdf6 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -413,7 +413,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy const int selected_pixels_per_thread = use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread; const int selected_threads_per_block = use_cooperative ? high_bin_threads_per_block : threads_per_block; - int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; + const int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); int blocks_per_col = @@ -471,9 +471,9 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( ::cuda::std::array num_privatized_bins_wrapper; ::cuda::std::array num_output_bins_wrapper; - auto* typed_allocations = reinterpret_cast(allocations); + auto* const typed_allocations = reinterpret_cast(allocations); ::cuda::std::copy(typed_allocations, typed_allocations + NUM_ACTIVE_CHANNELS, d_privatized_histograms_wrapper.begin()); - auto* local_typed_allocations = reinterpret_cast(allocations); + auto* const local_typed_allocations = reinterpret_cast(allocations); ::cuda::std::copy(local_typed_allocations, local_typed_allocations + NUM_ACTIVE_CHANNELS, d_cooperative_privatized_histograms_wrapper.begin()); diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index b3251480bffa..277ab99842d8 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -90,6 +90,13 @@ struct fast_divide_by_constant } const int log2_divisor = ceil_log2(divisor); + if (log2_divisor == bits) + { + magic = divisor; + shift = 0; + mode = 3; + return; + } if constexpr (sizeof(UInt) == 8) { #if _CCCL_HAS_INT128() @@ -134,6 +141,10 @@ struct fast_divide_by_constant { return numerator >> shift; } + if (mode == 3) + { + return numerator / magic; + } UInt high; if constexpr (sizeof(UInt) == 8) From 249253e270ff31aa30d5221f8cbf3c6592f701c9 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Thu, 3 Sep 2026 23:08:51 +0000 Subject: [PATCH 09/21] [cub] Fix full-width histogram range validation --- .../device/dispatch/dispatch_histogram.cuh | 7 ++++++- cub/test/catch2_test_device_histogram.cu | 21 ++++++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 2996c94ebdf6..73b44d30c794 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -195,7 +195,12 @@ struct DeviceHistogramKernelSource if constexpr (::cuda::std::is_integral_v) { using IntArithmeticT = typename TransformsT::ScaleTransform::IntArithmeticT; - return static_cast(upper_level[channel] - lower_level[channel]) + using ArrayLevelT = typename UpperLevelArrayT::value_type; + using ULevelT = ::cuda::std::make_unsigned_t; + + const ULevelT range = + static_cast(static_cast(upper_level[channel]) - static_cast(lower_level[channel])); + return static_cast(range) > (::cuda::std::numeric_limits::max() / static_cast(num_bins)); } else diff --git a/cub/test/catch2_test_device_histogram.cu b/cub/test/catch2_test_device_histogram.cu index 20e78b7079cc..ee6726d1ff10 100644 --- a/cub/test/catch2_test_device_histogram.cu +++ b/cub/test/catch2_test_device_histogram.cu @@ -720,17 +720,22 @@ catch (const std::exception& e) // Our bin computation for HistogramEven is guaranteed only for when (max_level - min_level) * num_bins does not // overflow using uint64_t arithmetic. In case of overflow, we expect cudaErrorInvalidValue to be returned. -CUB_TEST_LIST("DeviceHistogram::HistogramEven bin computation does not overflow", - "[histogram_even][device]", - CUB_SMALL, - uint8_t, - uint16_t, - uint32_t, - uint64_t) +CUB_TEST_LIST( + "DeviceHistogram::HistogramEven bin computation does not overflow", + "[histogram_even][device]", + CUB_SMALL, + int8_t, + int16_t, + int32_t, + int64_t, + uint8_t, + uint16_t, + uint32_t, + uint64_t) { using sample_t = TestType; using counter_t = uint32_t; - constexpr sample_t lower_level = 0; + constexpr sample_t lower_level = cs::numeric_limits::min(); constexpr sample_t upper_level = cs::numeric_limits::max(); constexpr auto num_samples = 1000; auto d_samples = cuda::counting_iterator{0UL}; From 86905611a46eaa0e679ac62a821be55b6e6e27ac Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Thu, 3 Sep 2026 23:32:58 +0000 Subject: [PATCH 10/21] [cub] Isolate high-bin histogram optimizations --- .../device/dispatch/dispatch_histogram.cuh | 60 ++- .../dispatch/kernels/kernel_histogram.cuh | 356 +++++++++++------- .../dispatch/tuning/tuning_histogram.cuh | 48 ++- cub/test/catch2_test_device_histogram_env.cu | 4 +- 4 files changed, 311 insertions(+), 157 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 73b44d30c794..62d31b94ee9c 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -309,17 +309,21 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } // Get device occupancy for sweep_kernel - int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; - bool use_cooperative = false; - int cooperative_smem_bytes = 0; - int cooperative_cache_slots_per_channel = 0; + int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; + [[maybe_unused]] const int privatized_storage_grid_limit = histogram_sweep_occupancy; + bool use_cooperative = false; + int cooperative_smem_bytes = 0; + int cooperative_cache_slots_per_channel = 0; #if _CCCL_HOSTED() NV_IF_TARGET( NV_IS_HOST, ({ if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) { - if (active_policy.high_bin_algorithm == HistogramHighBinAlgorithm::cooperative) + const size_t output_histogram_bytes = + static_cast(max_num_output_bins) * NUM_ACTIVE_CHANNELS * sizeof(LocalCounterT); + if (active_policy.high_bin_algorithm == HistogramHighBinAlgorithm::cooperative + && output_histogram_bytes > static_cast(active_policy.high_bin_min_histogram_bytes)) { if (const auto error = CubDebug(launcher_factory.CooperativeLaunchSupported(use_cooperative))) { @@ -394,7 +398,16 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { return error; } - histogram_sweep_occupancy = cooperative_sm_occupancy * sm_count; + const int cooperative_grid_capacity = cooperative_sm_occupancy * sm_count; + const int tuned_grid_capacity = + active_policy.high_bin_max_blocks_per_sm > 0 + ? active_policy.high_bin_max_blocks_per_sm * sm_count + : cooperative_grid_capacity; + histogram_sweep_occupancy = + active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized + ? ::cuda::std::min(privatized_storage_grid_limit, + ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity)) + : ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity); } else { @@ -1095,19 +1108,14 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( { using TransformsT = Transforms; - // Use the search transform op for converting samples to privatized bins - using PrivatizedDecodeOpT = typename TransformsT::template SearchTransform; - // Use the pass-thru transform op for converting privatized bins to output bins using OutputDecodeOpT = typename TransformsT::PassThruTransform; - ::cuda::std::array privatized_decode_op{}; ::cuda::std::array output_decode_op{}; int max_levels = num_output_levels[0]; for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) { - privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); if (num_output_levels[channel] > max_levels) { max_levels = num_output_levels[channel]; @@ -1120,6 +1128,12 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( { // Too many bins to keep in shared memory. constexpr int PRIVATIZED_SMEM_BINS = 0; + using PrivatizedDecodeOpT = typename TransformsT::template CachedSearchTransform; + ::cuda::std::array privatized_decode_op{}; + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); + } if (const auto error = CubDebug( (detail::histogram::dispatch; + ::cuda::std::array privatized_decode_op{}; + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); + } if (const auto error = CubDebug( (detail::histogram::dispatch; - // Use the scale transform op for converting samples to privatized bins - using PrivatizedDecodeOpT = typename TransformsT::ScaleTransform; - // Use the pass-thru transform op for converting privatized bins to output bins using OutputDecodeOpT = typename TransformsT::PassThruTransform; using CommonT = typename TransformsT::ScaleTransform::CommonT; - ::cuda::std::array privatized_decode_op{}; ::cuda::std::array output_decode_op{}; int max_levels = num_output_levels[0]; @@ -1312,8 +1328,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( return cudaErrorInvalidValue; } - privatized_decode_op[channel].Init(num_levels, upper_level[channel], lower_level[channel]); - if (num_levels > max_levels) { max_levels = num_levels; @@ -1324,6 +1338,12 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( if (max_num_output_bins > max_privatized_smem_bins) { constexpr int PRIVATIZED_SMEM_BINS = 0; + using PrivatizedDecodeOpT = typename TransformsT::FastScaleTransform; + ::cuda::std::array privatized_decode_op{}; + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + privatized_decode_op[channel].Init(num_output_levels[channel], upper_level[channel], lower_level[channel]); + } if (const auto error = CubDebug( (detail::histogram::dispatch privatized_decode_op{}; + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + privatized_decode_op[channel].Init(num_output_levels[channel], upper_level[channel], lower_level[channel]); + } if (const auto error = CubDebug( (detail::histogram::dispatch +struct scale_fraction; + +template +struct scale_fraction +{ + FractionStorageT bins; + FractionStorageT range; +}; + +template +struct scale_fraction +{ + FractionStorageT bins; + FractionStorageT range; + fast_divide_by_constant range_divider; + double reciprocal; + bool bins_equal_range; +}; + template _CCCL_DEVICE _CCCL_FORCEINLINE const SampleValueT* sample_native_pointer(SampleIteratorT itr) { @@ -221,24 +241,8 @@ struct Transforms template struct SearchTransform { - static constexpr bool is_range_transform = true; - - struct BracketCacheT - { - LevelT lo{}; - LevelT hi{}; - int bin = -1; - }; - LevelIteratorT d_levels; // Pointer to levels array int num_output_levels; // Number of levels in array - LevelT first{}; - LevelT middle{}; - LevelT last{}; - float inverse_scale_low = 0.0f; - float inverse_scale_high = 0.0f; - int middle_bin = 0; - bool has_precompute = false; //! @brief Initializer //! @@ -248,23 +252,79 @@ struct Transforms { this->d_levels = d_levels_; this->num_output_levels = num_output_levels_; - this->has_precompute = false; } + // Method for converting samples to bin-ids + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid) const + { + /// Level iterator wrapper type + // Wrap the native input pointer with CacheModifiedInputIterator + // or Directly use the supplied input iterator type + using WrappedLevelIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + LevelIteratorT>; + + WrappedLevelIteratorT wrapped_levels(d_levels); + + const int num_bins = num_output_levels - 1; + if (valid) + { + bin = UpperBound(wrapped_levels, num_output_levels, static_cast(sample)) - 1; + if (bin >= num_bins) + { + bin = -1; + } + } + } + }; + + //! @brief Finds a RANGE bin with piecewise-linear interpolation and a per-thread bracket cache. + template + struct CachedSearchTransform + { + static constexpr bool is_range_transform = true; + template - [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static float level_distance(T upper, T lower) + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static auto interpolation_difference(T lhs, T rhs) { if constexpr (::cuda::std::is_integral_v) { using UnsignedT = ::cuda::std::make_unsigned_t; - return static_cast(static_cast(upper) - static_cast(lower)); + return static_cast(lhs) - static_cast(rhs); } else { - return static_cast(upper - lower); + return lhs - rhs; } } + struct BracketCacheT + { + LevelT lo{}; + LevelT hi{}; + int bin = -1; + }; + + LevelIteratorT d_levels; + int num_output_levels; + LevelT first{}; + LevelT middle{}; + LevelT last{}; + float inverse_scale{}; + float inverse_scale_low{}; + float inverse_scale_high{}; + int middle_bin{}; + bool has_precompute{}; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void Init(LevelIteratorT d_levels_, int num_output_levels_) + { + d_levels = d_levels_; + num_output_levels = num_output_levels_; + has_precompute = false; + } + _CCCL_DEVICE _CCCL_FORCEINLINE void PrecomputeOnDevice(int interpolation_min_bins) { const int num_bins = num_output_levels - 1; @@ -285,98 +345,128 @@ struct Transforms return; } - first = first_level; - last = last_level; - middle_bin = num_bins >> 1; - middle = wrapped_levels[middle_bin]; - if (!(first < middle) || !(middle < last)) - { - middle_bin = 0; - inverse_scale_low = static_cast(num_bins) / level_distance(last, first); - } - else + first = first_level; + last = last_level; + inverse_scale = static_cast(num_bins) / static_cast(interpolation_difference(last, first)); + middle_bin = 0; + inverse_scale_low = inverse_scale; + inverse_scale_high = inverse_scale; + middle = first; + + const int split = num_bins >> 1; + if (split > 0 && split < num_bins) { - inverse_scale_low = static_cast(middle_bin) / level_distance(middle, first); - inverse_scale_high = static_cast(num_bins - middle_bin) / level_distance(last, middle); + const LevelT split_level = wrapped_levels[split]; + if (first < split_level && split_level < last) + { + middle = split_level; + middle_bin = split; + inverse_scale_low = + static_cast(split) / static_cast(interpolation_difference(split_level, first)); + inverse_scale_high = + static_cast(num_bins - split) / static_cast(interpolation_difference(last, split_level)); + } } has_precompute = true; } - // Method for converting samples to bin-ids template - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid) const + _CCCL_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid, BracketCacheT& bracket) const { - /// Level iterator wrapper type - // Wrap the native input pointer with CacheModifiedInputIterator - // or Directly use the supplied input iterator type + if (!valid) + { + return; + } + using WrappedLevelIteratorT = ::cuda::std::_If<::cuda::std::is_pointer_v, CacheModifiedInputIterator, LevelIteratorT>; - WrappedLevelIteratorT wrapped_levels(d_levels); - const int num_bins = num_output_levels - 1; - if (!valid) + const LevelT value = static_cast(sample); + + if (bracket.bin >= 0 && !(value < bracket.lo) && value < bracket.hi) { + bin = bracket.bin; return; } - const LevelT value = static_cast(sample); if (!has_precompute) { bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; - bin = bin < num_bins ? bin : -1; - return; + if (bin >= num_bins) + { + bin = -1; + } } - if (value < first || !(value < last)) + else if (value < first || !(value < last)) { bin = -1; - return; } - - int guess = value < middle || middle_bin == 0 - ? static_cast(level_distance(value, first) * inverse_scale_low) - : middle_bin + static_cast(level_distance(value, middle) * inverse_scale_high); - guess = guess < 0 ? 0 : (guess < num_bins ? guess : num_bins - 1); - const LevelT lo = wrapped_levels[guess]; - const LevelT hi = wrapped_levels[guess + 1]; - if (!(value < lo) && value < hi) + else { - bin = guess; - return; - } - bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; - } + int guess = + value < middle || middle_bin == 0 + ? static_cast(static_cast(interpolation_difference(value, first)) * inverse_scale_low) + : middle_bin + + static_cast(static_cast(interpolation_difference(value, middle)) * inverse_scale_high); + guess = guess < 0 ? 0 : (guess < num_bins ? guess : num_bins - 1); + const LevelT lo = wrapped_levels[guess]; + const LevelT hi = wrapped_levels[guess + 1]; + + if (!(value < lo) && value < hi) + { + bin = guess; + bracket = BracketCacheT{lo, hi, guess}; + return; + } - template - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void - BinSelect(_SampleT sample, int& bin, bool valid, BracketCacheT& bracket) const - { - const LevelT value = static_cast(sample); - if (valid && bracket.bin >= 0 && !(value < bracket.lo) && value < bracket.hi) - { - bin = bracket.bin; - return; + if (value < lo && guess > 0) + { + const LevelT adjacent_lo = wrapped_levels[guess - 1]; + if (!(value < adjacent_lo)) + { + bin = guess - 1; + bracket = BracketCacheT{adjacent_lo, lo, bin}; + return; + } + } + else if (!(value < hi) && guess + 1 < num_bins) + { + const LevelT adjacent_hi = wrapped_levels[guess + 2]; + if (value < adjacent_hi) + { + bin = guess + 1; + bracket = BracketCacheT{hi, adjacent_hi, bin}; + return; + } + } + + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + if (bin >= num_bins) + { + bin = -1; + } } - BinSelect(sample, bin, valid); if (bin >= 0) { - using WrappedLevelIteratorT = - ::cuda::std::_If<::cuda::std::is_pointer_v, - CacheModifiedInputIterator, - LevelIteratorT>; - WrappedLevelIteratorT wrapped_levels(d_levels); - bracket.lo = wrapped_levels[bin]; - bracket.hi = wrapped_levels[bin + 1]; - bracket.bin = bin; + bracket = BracketCacheT{wrapped_levels[bin], wrapped_levels[bin + 1], bin}; } } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid) const + { + BracketCacheT bracket{}; + BinSelect(sample, bin, valid, bracket); + } }; - // Scales samples to evenly-spaced bins - struct ScaleTransform + // Scales samples to evenly-spaced bins. + template + struct ScaleTransformImpl { using CommonT = ::cuda::std::common_type_t; static_assert(::cuda::std::is_convertible_v, @@ -424,13 +514,7 @@ struct Transforms { // Used when CommonT is not floating-point to avoid intermediate // rounding errors (see NVIDIA/cub#489). - struct FractionT - { - FractionStorageT bins; - FractionStorageT range; - fast_divide_by_constant range_divider; - bool bins_equal_range; - } fraction; + scale_fraction fraction; // Used when CommonT is floating-point as an optimization. CommonT reciprocal; @@ -466,8 +550,15 @@ struct Transforms { result.fraction.range = static_cast(max_level - min_level); } - result.fraction.bins_equal_range = result.fraction.bins == result.fraction.range; - result.fraction.range_divider.Init(static_cast(result.fraction.range)); + if constexpr (UseFastDivision) + { + result.fraction.bins_equal_range = result.fraction.bins == result.fraction.range; + result.fraction.range_divider.Init(static_cast(result.fraction.range)); + result.fraction.reciprocal = + result.fraction.range == FractionStorageT{0} + ? 0.0 + : static_cast(result.fraction.bins) / static_cast(result.fraction.range); + } return result; } @@ -554,12 +645,24 @@ struct Transforms using UnsignedT = ::cuda::std::make_unsigned_t; const IntArithmeticT distance = static_cast( static_cast(static_cast(sample) - static_cast(min_level))); - if (scale.fraction.bins_equal_range) + if constexpr (UseFastDivision) + { + if (scale.fraction.bins_equal_range) + { + return static_cast(distance); + } + if constexpr (sizeof(CommonT) <= 4) + { + return static_cast(static_cast(distance) * scale.fraction.reciprocal); + } + const IntArithmeticT numerator = distance * static_cast(scale.fraction.bins); + return static_cast(scale.fraction.range_divider.Divide(numerator)); + } + else { - return static_cast(distance); + return static_cast((distance * static_cast(scale.fraction.bins)) + / static_cast(scale.fraction.range)); } - const IntArithmeticT numerator = distance * static_cast(scale.fraction.bins); - return static_cast(scale.fraction.range_divider.Divide(numerator)); } template ::value, int> = 0> @@ -607,6 +710,9 @@ struct Transforms } }; + using ScaleTransform = ScaleTransformImpl; + using FastScaleTransform = ScaleTransformImpl; + // Pass-through bin transform operator struct PassThruTransform { @@ -1038,19 +1144,17 @@ __launch_bounds__(int(current_policy().threads_per_block)) agent.StoreOutput(); } -template +template _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( ::cuda::std::uint32_t* keys, CounterT* counts, - int count_replicas, int bin, CounterT contribution, int cache_mask, int cache_log2, - HistogramCacheAlgorithm cache_algorithm, bool use_second_probe) { - if (cache_algorithm == HistogramCacheAlgorithm::none) + if constexpr (CacheAlgorithm == HistogramCacheAlgorithm::none) { return false; } @@ -1065,8 +1169,7 @@ _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( } if (key == empty_key || key == bin_key) { - const int replica = static_cast((threadIdx.x >> 5) % count_replicas); - atomicAdd_block(&counts[replica * (cache_mask + 1) + slot], contribution); + atomicAdd_block(&counts[slot], contribution); return true; } return false; @@ -1079,11 +1182,14 @@ _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( return true; } - if (use_second_probe) + if constexpr (CacheAlgorithm == HistogramCacheAlgorithm::cuckoo) { - const unsigned int hash2 = (static_cast(bin) ^ 0x9e3779b9u) * 2246822519u; - const int secondary = static_cast((hash2 >> (32 - cache_log2)) & static_cast(cache_mask)); - return try_slot(secondary); + if (use_second_probe) + { + const unsigned int hash2 = (static_cast(bin) ^ 0x9e3779b9u) * 2246822519u; + const int secondary = static_cast((hash2 >> (32 - cache_log2)) & static_cast(cache_mask)); + return try_slot(secondary); + } } return false; } @@ -1127,16 +1233,19 @@ struct AgentHistogramCooperative const unsigned int tid_global = blockIdx.x * blockDim.x + threadIdx.x; const unsigned int total_threads = gridDim.x * blockDim.x; - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + if constexpr (policy.high_bin_spill != HistogramSpillAlgorithm::global_memory_privatized) { - for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); - bin += total_threads) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - d_output_histograms_wrapper[ch][bin] = OutputCounterT{0}; + for (unsigned int bin = tid_global; bin < static_cast(num_output_bins_wrapper[ch]); + bin += total_threads) + { + d_output_histograms_wrapper[ch][bin] = OutputCounterT{0}; + } } + grid.sync(); } - grid.sync(); static_assert(count_replicas > 0, "Histogram cache replication must be positive"); @@ -1156,11 +1265,10 @@ struct AgentHistogramCooperative for (int slot = threadIdx.x; slot < cache_slots_per_channel; slot += blockDim.x) { channel_keys[slot] = ~::cuda::std::uint32_t{0}; - _CCCL_PRAGMA_UNROLL_FULL() - for (int replica = 0; replica < count_replicas; ++replica) - { - channel_counts[static_cast(replica) * cache_slots_per_channel + slot] = CounterT{0}; - } + } + for (int count = threadIdx.x; count < count_replicas * cache_slots_per_channel; count += blockDim.x) + { + channel_counts[count] = CounterT{0}; } if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) @@ -1193,14 +1301,20 @@ struct AgentHistogramCooperative constexpr bool use_mru_cache = NumActiveChannels == 1 && PrivatizedDecodeOpT::is_range_transform; typename PrivatizedDecodeOpT::BracketCacheT bracket_cache[NumActiveChannels]; + ::cuda::std::uint32_t* channel_keys[NumActiveChannels]; + CounterT* thread_counts[NumActiveChannels]; int pending_bin[NumActiveChannels]; CounterT pending_count[NumActiveChannels]; _CCCL_PRAGMA_UNROLL_FULL() for (int ch = 0; ch < NumActiveChannels; ++ch) { - pending_bin[ch] = -1; - pending_count[ch] = CounterT{0}; + channel_keys[ch] = cache_keys + static_cast(ch) * cache_slots_per_channel; + CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; + const int replica = static_cast((threadIdx.x >> 5) % count_replicas); + thread_counts[ch] = channel_counts + static_cast(replica) * cache_slots_per_channel; + pending_bin[ch] = -1; + pending_count[ch] = CounterT{0}; } const auto spill_bin = [&](int ch, int selected_bin, CounterT contribution) { @@ -1241,20 +1355,10 @@ struct AgentHistogramCooperative }; const auto update_bin = [&](int ch, int selected_bin, CounterT contribution) { - auto* channel_keys = cache_keys + static_cast(ch) * cache_slots_per_channel; - CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; const bool use_second_probe = policy.high_bin_cache == HistogramCacheAlgorithm::cuckoo && num_output_bins_wrapper[ch] < policy.high_bin_cache_cuckoo_max_bins; - if (!histogram_cache_probe( - channel_keys, - channel_counts, - count_replicas, - selected_bin, - contribution, - cache_mask, - cache_log2, - policy.high_bin_cache, - use_second_probe)) + if (!histogram_cache_probe( + channel_keys[ch], thread_counts[ch], selected_bin, contribution, cache_mask, cache_log2, use_second_probe)) { spill_bin(ch, selected_bin, contribution); } diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 7c78a11ddd20..b5162629857b 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -130,6 +130,8 @@ struct HistogramPolicy int high_bin_pixels_per_thread = 4; int high_bin_threads_per_block = 0; //!< High-bin block size; 0 inherits threads_per_block int high_bin_interpolation_min_bins = 512; + int high_bin_min_histogram_bytes = 0; + int high_bin_max_blocks_per_sm = 0; //!< Zero uses the occupancy-derived capacity [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_threads() const noexcept { @@ -151,7 +153,9 @@ struct HistogramPolicy && lhs.high_bin_cache_cuckoo_max_bins == rhs.high_bin_cache_cuckoo_max_bins && lhs.high_bin_pixels_per_thread == rhs.high_bin_pixels_per_thread && lhs.high_bin_threads_per_block == rhs.high_bin_threads_per_block - && lhs.high_bin_interpolation_min_bins == rhs.high_bin_interpolation_min_bins; + && lhs.high_bin_interpolation_min_bins == rhs.high_bin_interpolation_min_bins + && lhs.high_bin_min_histogram_bytes == rhs.high_bin_min_histogram_bytes + && lhs.high_bin_max_blocks_per_sm == rhs.high_bin_max_blocks_per_sm; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -174,9 +178,10 @@ struct HistogramPolicy << ", .high_bin_cache_entries_per_channel = " << p.high_bin_cache_entries_per_channel << ", .high_bin_cache_count_replicas = " << p.high_bin_cache_count_replicas << ", .high_bin_cache_cuckoo_max_bins = " << p.high_bin_cache_cuckoo_max_bins - << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread - << ", .high_bin_threads_per_block = " << p.high_bin_threads_per_block - << ", .high_bin_interpolation_min_bins = " << p.high_bin_interpolation_min_bins << " }"; + << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread << ", .high_bin_threads_per_block = " + << p.high_bin_threads_per_block << ", .high_bin_interpolation_min_bins = " << p.high_bin_interpolation_min_bins + << ", .high_bin_min_histogram_bytes = " << p.high_bin_min_histogram_bytes + << ", .high_bin_max_blocks_per_sm = " << p.high_bin_max_blocks_per_sm << " }"; } #endif // _CCCL_HOSTED() }; @@ -412,12 +417,29 @@ public: { if (cc >= ::cuda::compute_capability{10, 0}) { + constexpr int sm100_smem_bytes = 228352; + constexpr int sm120_smem_bytes = 99 * 1024; + constexpr int range_smem_bytes_per_channel = 8192; + constexpr int even_smem_bytes_per_channel = 32768; + const int architecture_smem_bytes = cc >= ::cuda::compute_capability{12, 0} ? sm120_smem_bytes : sm100_smem_bytes; + const int candidate_smem_bytes = + num_active_channels == 1 ? architecture_smem_bytes + : is_even ? even_smem_bytes_per_channel * num_active_channels + : range_smem_bytes_per_channel * num_active_channels; + const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, architecture_smem_bytes); + const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : 1; + const auto with_high_bin_threshold = [=](HistogramPolicy policy) { + policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; + policy.high_bin_max_blocks_per_sm = high_bin_max_blocks_per_sm; + return policy; + }; + if (num_channels == 1 && num_active_channels == 1 && counter_size == 4 && sample_is_primitive && sample_size == 1) { if (is_even) { // ipt_12.tpb_928.rle_0.ws_0.mem_1.ld_2.laid_0.vec_2 1.033332 0.940517 1.031835 1.195876 - return HistogramPolicy{ + return with_high_bin_threshold(HistogramPolicy{ 928, 12, 1 << 2, @@ -435,12 +457,12 @@ public: 1, 262144, 4, - 0}; + 0}); } else { // ipt_12.tpb_448.rle_0.ws_0.mem_1.ld_1.laid_0.vec_2 1.078987 0.985542 1.085118 1.175637 - return HistogramPolicy{ + return with_high_bin_threshold(HistogramPolicy{ 448, 12, 1 << 2, @@ -458,14 +480,14 @@ public: 1, 262144, 4, - 0}; + 0}); } } if (counter_size == 4 && sample_is_primitive && num_channels == 1 && num_active_channels == 1 && (sample_size == 4 || sample_size == 8)) { - return HistogramPolicy{ + return with_high_bin_threshold(HistogramPolicy{ 384, t_scale(16), 4, @@ -483,12 +505,12 @@ public: 1, 262144, 4, - is_even ? 768 : 512}; + is_even ? 768 : 512}); } if (counter_size == 4 && sample_is_primitive && num_channels >= 2) { - return HistogramPolicy{ + return with_high_bin_threshold(HistogramPolicy{ 384, t_scale(16), 4, @@ -502,11 +524,11 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - 1024, + is_even ? 1024 : 2048, 4, 262144, 4, - 1024}; + 1024}); } // sample_size 2/4/8 showed no benefit over SM90 during verification benchmarks diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index 7f820d13aa70..b21afb4609bd 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1676,6 +1676,7 @@ struct high_bin_histogram_tuning policy.high_bin_cache_cuckoo_max_bins = 4096; policy.high_bin_pixels_per_thread = 2; policy.high_bin_threads_per_block = 128; + policy.high_bin_min_histogram_bytes = 0; return policy; } }; @@ -1978,7 +1979,8 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) ", .high_bin_cache_entries_per_channel = 2048" ", .high_bin_cache_count_replicas = 1, .high_bin_cache_cuckoo_max_bins = 262144" ", .high_bin_pixels_per_thread = 4, .high_bin_threads_per_block = 0" - ", .high_bin_interpolation_min_bins = 512 }"); + ", .high_bin_interpolation_min_bins = 512, .high_bin_min_histogram_bytes = 0" + ", .high_bin_max_blocks_per_sm = 0 }"); constexpr auto high_bin_policy = [] { auto policy = From 6b4487482e13019272910753b906a482d4d4ae50 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 00:14:33 +0000 Subject: [PATCH 11/21] [cub] Isolate high-bin histogram dispatch --- .../device/dispatch/dispatch_histogram.cuh | 39 ++++++++--- .../dispatch/kernels/kernel_histogram.cuh | 52 +++++++-------- .../dispatch/tuning/tuning_histogram.cuh | 64 ++----------------- cub/test/catch2_test_device_histogram_env.cu | 2 +- 4 files changed, 63 insertions(+), 94 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 62d31b94ee9c..9604f695b83b 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -220,6 +220,7 @@ template num_output_levels, FirstLevelArrayT first_level_array, SecondLevelArrayT second_level_array, + CooperativeSecondLevelArrayT cooperative_second_level_array, int max_num_output_bins, OffsetT num_row_pixels, OffsetT num_rows, @@ -331,7 +333,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } if (use_cooperative) { - using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + using privatized_decode_op_t = typename CooperativeSecondLevelArrayT::value_type; const auto cooperative_kernel = kernel_source.template HistogramCooperativeKernel(); const int cache_slots_floor = @@ -375,7 +377,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( if (cooperative_sm_occupancy > 0) { cooperative_cache_slots_per_channel = cache_slots_floor; - const int floor_occupancy = cooperative_sm_occupancy; + const int target_occupancy = + active_policy.high_bin_max_blocks_per_sm > 0 + ? ::cuda::std::min(cooperative_sm_occupancy, active_policy.high_bin_max_blocks_per_sm) + : cooperative_sm_occupancy; for (int candidate = cache_slots_floor == 0 ? 0 : cache_slots_floor << 1; candidate > 0 && candidate <= max_slots_by_smem; candidate <<= 1) @@ -385,7 +390,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { return error; } - if (candidate_occupancy < floor_occupancy) + if (candidate_occupancy < target_occupancy) { break; } @@ -509,7 +514,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { if (use_cooperative && blocks_per_row > 0 && blocks_per_col > 0) { - using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + using privatized_decode_op_t = typename CooperativeSecondLevelArrayT::value_type; const dim3 cooperative_grid_dims{static_cast(num_thread_blocks), 1u, 1u}; const auto cooperative_kernel = @@ -524,7 +529,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( num_output_bins_wrapper, d_output_histograms, d_cooperative_privatized_histograms_wrapper, - second_level_array, + cooperative_second_level_array, num_row_pixels, num_rows, row_stride_samples, @@ -769,6 +774,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __dispatch_even_device_init( num_output_levels, upper_level, lower_level, + lower_level, max_num_output_bins, num_row_pixels, num_rows, @@ -801,6 +807,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __dispatch_even_device_init( num_output_levels, upper_level, lower_level, + lower_level, max_num_output_bins, num_row_pixels, num_rows, @@ -941,6 +948,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __dispatch_even_device_init( num_output_levels, upper_level, lower_level, + lower_level, max_num_output_bins, num_row_pixels, num_rows, @@ -1092,6 +1100,7 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( num_output_levels, output_decode_op, privatized_decode_op, + privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, @@ -1127,12 +1136,15 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( if (max_num_output_bins > max_privatized_smem_bins) { // Too many bins to keep in shared memory. - constexpr int PRIVATIZED_SMEM_BINS = 0; - using PrivatizedDecodeOpT = typename TransformsT::template CachedSearchTransform; + constexpr int PRIVATIZED_SMEM_BINS = 0; + using PrivatizedDecodeOpT = typename TransformsT::template SearchTransform; + using CooperativePrivatizedDecodeOpT = typename TransformsT::template CachedSearchTransform; ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array cooperative_privatized_decode_op{}; for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) { privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); + cooperative_privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); } if (const auto error = CubDebug( @@ -1150,6 +1162,7 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( num_output_levels, output_decode_op, privatized_decode_op, + cooperative_privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, @@ -1188,6 +1201,7 @@ CUB_RUNTIME_FUNCTION cudaError_t dispatch_range( num_output_levels, output_decode_op, privatized_decode_op, + privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, @@ -1292,6 +1306,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( num_output_levels, output_decode_op, privatized_decode_op, + privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, @@ -1337,12 +1352,16 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( if (max_num_output_bins > max_privatized_smem_bins) { - constexpr int PRIVATIZED_SMEM_BINS = 0; - using PrivatizedDecodeOpT = typename TransformsT::FastScaleTransform; + constexpr int PRIVATIZED_SMEM_BINS = 0; + using PrivatizedDecodeOpT = typename TransformsT::ScaleTransform; + using CooperativePrivatizedDecodeOpT = typename TransformsT::FastScaleTransform; ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array cooperative_privatized_decode_op{}; for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) { privatized_decode_op[channel].Init(num_output_levels[channel], upper_level[channel], lower_level[channel]); + cooperative_privatized_decode_op[channel].Init( + num_output_levels[channel], upper_level[channel], lower_level[channel]); } if (const auto error = CubDebug( @@ -1360,6 +1379,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( num_output_levels, output_decode_op, privatized_decode_op, + cooperative_privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, @@ -1397,6 +1417,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( num_output_levels, output_decode_op, privatized_decode_op, + privatized_decode_op, max_num_output_bins, num_row_pixels, num_rows, diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index 879d788c2889..e1fac7269887 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -508,7 +508,8 @@ struct Transforms ::cuda::std::is_integral; #endif // !_CCCL_HAS_INT128() - using FractionStorageT = ::cuda::std::_If::value, IntArithmeticT, CommonT>; + using FractionStorageT = + ::cuda::std::_If::value, IntArithmeticT, CommonT>; union ScaleT { @@ -539,7 +540,7 @@ struct Transforms { ScaleT result; result.fraction.bins = static_cast(num_levels - 1); - if constexpr (::cuda::std::is_integral_v) + if constexpr (UseFastDivision && ::cuda::std::is_integral_v) { using UnsignedT = ::cuda::std::make_unsigned_t; const UnsignedT distance = @@ -642,11 +643,11 @@ struct Transforms template ::value, int> = 0> _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int ComputeBin(T sample, T min_level, ScaleT scale) const { - using UnsignedT = ::cuda::std::make_unsigned_t; - const IntArithmeticT distance = static_cast( - static_cast(static_cast(sample) - static_cast(min_level))); if constexpr (UseFastDivision) { + using UnsignedT = ::cuda::std::make_unsigned_t; + const IntArithmeticT distance = static_cast( + static_cast(static_cast(sample) - static_cast(min_level))); if (scale.fraction.bins_equal_range) { return static_cast(distance); @@ -660,8 +661,9 @@ struct Transforms } else { - return static_cast((distance * static_cast(scale.fraction.bins)) - / static_cast(scale.fraction.range)); + return static_cast( + (static_cast(sample - min_level) * static_cast(scale.fraction.bins)) + / static_cast(scale.fraction.range)); } } @@ -1303,6 +1305,7 @@ struct AgentHistogramCooperative typename PrivatizedDecodeOpT::BracketCacheT bracket_cache[NumActiveChannels]; ::cuda::std::uint32_t* channel_keys[NumActiveChannels]; CounterT* thread_counts[NumActiveChannels]; + CounterT* private_histograms[NumActiveChannels]; int pending_bin[NumActiveChannels]; CounterT pending_count[NumActiveChannels]; @@ -1313,8 +1316,13 @@ struct AgentHistogramCooperative CounterT* channel_counts = cache_counts + static_cast(ch) * count_replicas * cache_slots_per_channel; const int replica = static_cast((threadIdx.x >> 5) % count_replicas); thread_counts[ch] = channel_counts + static_cast(replica) * cache_slots_per_channel; - pending_bin[ch] = -1; - pending_count[ch] = CounterT{0}; + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + private_histograms[ch] = + d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; + } + pending_bin[ch] = -1; + pending_count[ch] = CounterT{0}; } const auto spill_bin = [&](int ch, int selected_bin, CounterT contribution) { @@ -1329,9 +1337,7 @@ struct AgentHistogramCooperative { if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); + atomicAdd_block(&private_histograms[ch][pending_bin[ch]], pending_count[ch]); } else { @@ -1344,9 +1350,7 @@ struct AgentHistogramCooperative } else if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[selected_bin], contribution); + atomicAdd_block(&private_histograms[ch][selected_bin], contribution); } else { @@ -1365,13 +1369,15 @@ struct AgentHistogramCooperative }; const auto consume_bin = [&](int ch, int bin) { - if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::warp_coalesced) + constexpr bool coalesce_before_probe = + policy.high_bin_cache != HistogramCacheAlgorithm::none && sizeof(CounterT) > sizeof(::cuda::std::uint32_t); + if constexpr (coalesce_before_probe + || policy.high_bin_aggregation == HistogramAggregationAlgorithm::warp_coalesced) { NV_IF_ELSE_TARGET( NV_PROVIDES_SM_70, - (const unsigned int active = __activemask(); - const unsigned int peers = __match_any_sync(active, static_cast(bin)); - const int leader = __ffs(static_cast(peers)) - 1; + (const unsigned int peers = __match_any_sync(0xffffffffu, static_cast(bin)); + const int leader = __ffs(static_cast(peers)) - 1; if (bin >= 0 && static_cast(lane_id) == leader) { update_bin(ch, bin, static_cast(__popc(peers))); }), @@ -1563,9 +1569,7 @@ struct AgentHistogramCooperative { if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[pending_bin[ch]], pending_count[ch]); + atomicAdd_block(&private_histograms[ch][pending_bin[ch]], pending_count[ch]); } else { @@ -1597,9 +1601,7 @@ struct AgentHistogramCooperative { if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) { - CounterT* block_histogram = - d_privatized_histograms_wrapper[ch] + static_cast(blockIdx.x) * num_output_bins_wrapper[ch]; - atomicAdd_block(&block_histogram[key], count); + atomicAdd_block(&private_histograms[ch][key], count); } else { diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index b5162629857b..17047008b94a 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -120,7 +120,7 @@ struct HistogramPolicy bool use_work_stealing; //!< Whether to dequeue tiles from a global work queue int init_kernel_pdl_trigger_max_bins; //!< Maximum number of bins for the init kernel to trigger the histogram kernel //!< early using PDL - HistogramHighBinAlgorithm high_bin_algorithm = HistogramHighBinAlgorithm::cooperative; + HistogramHighBinAlgorithm high_bin_algorithm = HistogramHighBinAlgorithm::global_memory_privatized; HistogramCacheAlgorithm high_bin_cache = HistogramCacheAlgorithm::single_probe; HistogramSpillAlgorithm high_bin_spill = HistogramSpillAlgorithm::global_memory_privatized; HistogramAggregationAlgorithm high_bin_aggregation = HistogramAggregationAlgorithm::rle; @@ -524,7 +524,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - is_even ? 1024 : 2048, + 1024, 4, 262144, 4, @@ -541,71 +541,17 @@ public: { if (sample_size == 1) { - return HistogramPolicy{ - 768, - 12, - 1 << 2, - BLOCK_LOAD_DIRECT, - LOAD_LDG, - false, - SMEM, - false, - 2048, - HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::single_probe, - HistogramSpillAlgorithm::global_memory_privatized, - HistogramAggregationAlgorithm::rle, - is_even ? 4096 : 2048, - 1, - 262144, - 4, - 0}; + return HistogramPolicy{768, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_LDG, false, SMEM, false, 2048}; } else if (sample_size == 2) { - return HistogramPolicy{ - 960, - 10, - 1 << 2, - BLOCK_LOAD_DIRECT, - LOAD_DEFAULT, - true, - SMEM, - false, - 2048, - HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::single_probe, - HistogramSpillAlgorithm::global_memory_privatized, - HistogramAggregationAlgorithm::rle, - is_even ? 4096 : 2048, - 1, - 262144, - 4, - 0}; + return HistogramPolicy{960, 10, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_DEFAULT, true, SMEM, false, 2048}; } } } // fallback from SM50 - return HistogramPolicy{ - 384, - t_scale(16), - 4, - BLOCK_LOAD_DIRECT, - LOAD_LDG, - true, - SMEM, - false, - 0, - HistogramHighBinAlgorithm::cooperative, - HistogramCacheAlgorithm::single_probe, - HistogramSpillAlgorithm::global_memory_privatized, - HistogramAggregationAlgorithm::rle, - num_active_channels == 1 ? 4096 : 1024, - num_active_channels == 1 ? 1 : 4, - 262144, - 4, - 0}; + return HistogramPolicy{384, t_scale(16), 4, BLOCK_LOAD_DIRECT, LOAD_LDG, true, SMEM, false, 0}; } }; diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index b21afb4609bd..c57bcb9165a7 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1972,7 +1972,7 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) == "HistogramPolicy { .threads_per_block = 128, .pixels_per_thread = 7, .vec_size = 4" ", .load_algorithm = BLOCK_LOAD_DIRECT, .load_modifier = LOAD_LDG, .rle_compress = 0" ", .mem_preference = SMEM, .use_work_stealing = 0, .init_kernel_pdl_trigger_max_bins = 2048" - ", .high_bin_algorithm = HistogramHighBinAlgorithm::cooperative" + ", .high_bin_algorithm = HistogramHighBinAlgorithm::global_memory_privatized" ", .high_bin_cache = HistogramCacheAlgorithm::single_probe" ", .high_bin_spill = HistogramSpillAlgorithm::global_memory_privatized" ", .high_bin_aggregation = HistogramAggregationAlgorithm::rle" From ae75ab7dc4e7fe3e3e71c7c5093c3fb88eef2cca Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 00:32:32 +0000 Subject: [PATCH 12/21] [cub] Match raw high-bin cache occupancy --- cub/cub/device/dispatch/dispatch_histogram.cuh | 16 ++++------------ .../device/dispatch/tuning/tuning_histogram.cuh | 2 -- 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 9604f695b83b..1b48f7ab77d8 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -377,10 +377,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( if (cooperative_sm_occupancy > 0) { cooperative_cache_slots_per_channel = cache_slots_floor; - const int target_occupancy = - active_policy.high_bin_max_blocks_per_sm > 0 - ? ::cuda::std::min(cooperative_sm_occupancy, active_policy.high_bin_max_blocks_per_sm) - : cooperative_sm_occupancy; + const int floor_occupancy = cooperative_sm_occupancy; for (int candidate = cache_slots_floor == 0 ? 0 : cache_slots_floor << 1; candidate > 0 && candidate <= max_slots_by_smem; candidate <<= 1) @@ -390,7 +387,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { return error; } - if (candidate_occupancy < target_occupancy) + if (candidate_occupancy < floor_occupancy) { break; } @@ -404,15 +401,10 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( return error; } const int cooperative_grid_capacity = cooperative_sm_occupancy * sm_count; - const int tuned_grid_capacity = - active_policy.high_bin_max_blocks_per_sm > 0 - ? active_policy.high_bin_max_blocks_per_sm * sm_count - : cooperative_grid_capacity; histogram_sweep_occupancy = active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized - ? ::cuda::std::min(privatized_storage_grid_limit, - ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity)) - : ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity); + ? ::cuda::std::min(privatized_storage_grid_limit, cooperative_grid_capacity) + : cooperative_grid_capacity; } else { diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 17047008b94a..3398196f9dce 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -427,10 +427,8 @@ public: : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, architecture_smem_bytes); - const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : 1; const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; - policy.high_bin_max_blocks_per_sm = high_bin_max_blocks_per_sm; return policy; }; From cae2d720e6c959025afbaa7ce69647a1cb82a6eb Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 00:46:13 +0000 Subject: [PATCH 13/21] [cub] Use explicit high-bin launch tuning --- .../device/dispatch/dispatch_histogram.cuh | 55 ++++++------------- .../dispatch/tuning/tuning_histogram.cuh | 10 ++-- 2 files changed, 21 insertions(+), 44 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 1b48f7ab77d8..35d175a635f0 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -336,7 +336,7 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( using privatized_decode_op_t = typename CooperativeSecondLevelArrayT::value_type; const auto cooperative_kernel = kernel_source.template HistogramCooperativeKernel(); - const int cache_slots_floor = + cooperative_cache_slots_per_channel = active_policy.high_bin_cache == HistogramCacheAlgorithm::none ? 0 : active_policy.high_bin_cache_entries_per_channel; @@ -352,59 +352,36 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( return error; } const int max_slots_by_smem = cache_bytes_per_slot == 0 ? 0 : max_dynamic_smem_bytes / cache_bytes_per_slot; - if (cache_slots_floor > max_slots_by_smem) + if (cooperative_cache_slots_per_channel > max_slots_by_smem) { use_cooperative = false; } else { - const auto occupancy_for_slots = [&](int slots, int& occupancy) { - const int smem_bytes = slots * cache_bytes_per_slot; - cudaError_t error = launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, smem_bytes); - if (error == cudaSuccess) - { - error = launcher_factory.MaxSmOccupancy( - occupancy, cooperative_kernel, high_bin_threads_per_block, smem_bytes); - } + cooperative_smem_bytes = cooperative_cache_slots_per_channel * cache_bytes_per_slot; + if (const auto error = CubDebug( + launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, cooperative_smem_bytes))) + { return error; - }; - + } int cooperative_sm_occupancy = 0; - if (const auto error = CubDebug(occupancy_for_slots(cache_slots_floor, cooperative_sm_occupancy))) + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + cooperative_sm_occupancy, cooperative_kernel, high_bin_threads_per_block, cooperative_smem_bytes))) { return error; } if (cooperative_sm_occupancy > 0) { - cooperative_cache_slots_per_channel = cache_slots_floor; - const int floor_occupancy = cooperative_sm_occupancy; - for (int candidate = cache_slots_floor == 0 ? 0 : cache_slots_floor << 1; - candidate > 0 && candidate <= max_slots_by_smem; - candidate <<= 1) - { - int candidate_occupancy = 0; - if (const auto error = CubDebug(occupancy_for_slots(candidate, candidate_occupancy))) - { - return error; - } - if (candidate_occupancy < floor_occupancy) - { - break; - } - cooperative_cache_slots_per_channel = candidate; - cooperative_sm_occupancy = candidate_occupancy; - } - cooperative_smem_bytes = cooperative_cache_slots_per_channel * cache_bytes_per_slot; - if (const auto error = CubDebug( - launcher_factory.set_max_dynamic_smem_size_for(cooperative_kernel, cooperative_smem_bytes))) - { - return error; - } const int cooperative_grid_capacity = cooperative_sm_occupancy * sm_count; + const int tuned_grid_capacity = + active_policy.high_bin_max_blocks_per_sm > 0 + ? active_policy.high_bin_max_blocks_per_sm * sm_count + : cooperative_grid_capacity; histogram_sweep_occupancy = active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized - ? ::cuda::std::min(privatized_storage_grid_limit, cooperative_grid_capacity) - : cooperative_grid_capacity; + ? ::cuda::std::min(privatized_storage_grid_limit, + ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity)) + : ::cuda::std::min(cooperative_grid_capacity, tuned_grid_capacity); } else { diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 3398196f9dce..d4b96a9efac3 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -415,20 +415,20 @@ private: public: [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> HistogramPolicy { - if (cc >= ::cuda::compute_capability{10, 0}) + if (cc == ::cuda::compute_capability{10, 0}) { constexpr int sm100_smem_bytes = 228352; - constexpr int sm120_smem_bytes = 99 * 1024; constexpr int range_smem_bytes_per_channel = 8192; constexpr int even_smem_bytes_per_channel = 32768; - const int architecture_smem_bytes = cc >= ::cuda::compute_capability{12, 0} ? sm120_smem_bytes : sm100_smem_bytes; const int candidate_smem_bytes = - num_active_channels == 1 ? architecture_smem_bytes + num_active_channels == 1 ? sm100_smem_bytes : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; - const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, architecture_smem_bytes); + const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); + const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : 0; const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; + policy.high_bin_max_blocks_per_sm = high_bin_max_blocks_per_sm; return policy; }; From a67568cad66775fca0b539c059c0c12a4fe2697c Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 01:50:20 +0000 Subject: [PATCH 14/21] [cub] Match raw high-bin launch and hot paths --- .../dispatch/kernels/kernel_histogram.cuh | 226 +++++++++++++----- .../dispatch/tuning/tuning_histogram.cuh | 13 +- 2 files changed, 178 insertions(+), 61 deletions(-) diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index e1fac7269887..36287e6ec6d8 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -459,8 +459,72 @@ struct Transforms template _CCCL_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid) const { - BracketCacheT bracket{}; - BinSelect(sample, bin, valid, bracket); + if (!valid) + { + return; + } + + using WrappedLevelIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + LevelIteratorT>; + WrappedLevelIteratorT wrapped_levels(d_levels); + const int num_bins = num_output_levels - 1; + const LevelT value = static_cast(sample); + + if (!has_precompute) + { + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + if (bin >= num_bins) + { + bin = -1; + } + return; + } + if (value < first || !(value < last)) + { + bin = -1; + return; + } + + int guess = + value < middle || middle_bin == 0 + ? static_cast(static_cast(interpolation_difference(value, first)) * inverse_scale_low) + : middle_bin + + static_cast(static_cast(interpolation_difference(value, middle)) * inverse_scale_high); + guess = guess < 0 ? 0 : (guess < num_bins ? guess : num_bins - 1); + const LevelT lo = wrapped_levels[guess]; + const LevelT hi = wrapped_levels[guess + 1]; + + if (!(value < lo) && value < hi) + { + bin = guess; + return; + } + if (value < lo && guess > 0) + { + const LevelT adjacent_lo = wrapped_levels[guess - 1]; + if (!(value < adjacent_lo)) + { + bin = guess - 1; + return; + } + } + else if (!(value < hi) && guess + 1 < num_bins) + { + const LevelT adjacent_hi = wrapped_levels[guess + 2]; + if (value < adjacent_hi) + { + bin = guess + 1; + return; + } + } + + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + if (bin >= num_bins) + { + bin = -1; + } } }; @@ -1165,15 +1229,20 @@ _CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( const auto bin_key = static_cast<::cuda::std::uint32_t>(bin); const auto try_slot = [&](int slot) { ::cuda::std::uint32_t key = keys[slot]; - if (key == empty_key) - { - key = atomicCAS_block(&keys[slot], empty_key, bin_key); - } - if (key == empty_key || key == bin_key) + if (key == bin_key) { atomicAdd_block(&counts[slot], contribution); return true; } + if (key == empty_key) + { + key = atomicCAS_block(&keys[slot], empty_key, bin_key); + if (key == empty_key || key == bin_key) + { + atomicAdd_block(&counts[slot], contribution); + return true; + } + } return false; }; @@ -1399,60 +1468,81 @@ struct AgentHistogramCooperative using SampleValueT = it_value_t; if constexpr (NumActiveChannels == 1) { - for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + if (contiguous_input) { - const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; - SampleValueT staged_samples[unroll]; - bool valid_samples[unroll]; - int bins[unroll]; - - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { - const OffsetT pixel = first_pixel + static_cast(item) * step; - valid_samples[item] = pixel < total_pixels; - OffsetT pixel_offset{}; - if (valid_samples[item]) + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + SampleValueT staged_samples[unroll]; + bool valid_samples[unroll]; + int bins[unroll]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) { - if (contiguous_input) - { - pixel_offset = pixel * NumChannels; - } - else + const OffsetT pixel = first_pixel + static_cast(item) * step; + valid_samples[item] = pixel < total_pixels; + const OffsetT safe_pixel = valid_samples[item] ? pixel : OffsetT{0}; + staged_samples[item] = d_samples[safe_pixel * NumChannels]; + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + int bin = -1; + if (valid_samples[item]) { - const OffsetT row = pixel / num_row_pixels; - pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + if constexpr (use_mru_cache) + { + decode_op[0].template BinSelect(staged_samples[item], bin, true, bracket_cache[0]); + } + else + { + decode_op[0].template BinSelect(staged_samples[item], bin, true); + } + if (bin >= num_output_bins_wrapper[0]) + { + bin = -1; + } } - staged_samples[item] = d_samples[pixel_offset]; + bins[item] = bin; } - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + consume_bin(0, bins[item]); + } + } + } + else + { + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { - int bin = -1; - if (valid_samples[item]) + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) { - if constexpr (use_mru_cache) - { - decode_op[0].template BinSelect(staged_samples[item], bin, true, bracket_cache[0]); - } - else - { - decode_op[0].template BinSelect(staged_samples[item], bin, true); - } - if (bin >= num_output_bins_wrapper[0]) + const OffsetT pixel = first_pixel + static_cast(item) * step; + if (pixel < total_pixels) { - bin = -1; + const OffsetT row = pixel / num_row_pixels; + const OffsetT pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + int bin = -1; + if constexpr (use_mru_cache) + { + decode_op[0].template BinSelect(d_samples[pixel_offset], bin, true, bracket_cache[0]); + } + else + { + decode_op[0].template BinSelect(d_samples[pixel_offset], bin, true); + } + if (bin < num_output_bins_wrapper[0]) + { + consume_bin(0, bin); + } } } - bins[item] = bin; - } - - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) - { - consume_bin(0, bins[item]); } } } @@ -1503,19 +1593,40 @@ struct AgentHistogramCooperative && (reinterpret_cast(native_base) & (alignof(PixelT) - 1)) == 0; if (vectorizable) { - const PixelT* pixels = reinterpret_cast(native_base); - consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { - if (valid) + const PixelT* const pixels = reinterpret_cast(native_base); + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + { + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) { - const PixelT packed_pixel = pixels[pixel]; - const SampleValueT* pixel_lanes = reinterpret_cast(&packed_pixel); + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + const PixelT packed = pixels[valid ? pixel : OffsetT{0}]; + const SampleValueT* lanes = reinterpret_cast(&packed); + int bins[NumActiveChannels]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int bin = -1; + if (valid) + { + decode_op[ch].template BinSelect(lanes[ch], bin, true); + if (bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } + } + bins[ch] = bin; + } _CCCL_PRAGMA_UNROLL_FULL() for (int ch = 0; ch < NumActiveChannels; ++ch) { - samples[ch] = pixel_lanes[ch]; + consume_bin(ch, bins[ch]); } } - }); + } } else { @@ -1646,7 +1757,8 @@ template #endif // _CCCL_HAS_CONCEPTS() -__launch_bounds__(int(current_policy().high_bin_threads())) +__launch_bounds__(int(current_policy().high_bin_threads()), + int(current_policy().high_bin_min_blocks())) _CCCL_KERNEL_ATTRIBUTES void DeviceHistogramCooperativeKernel( _CCCL_GRID_CONSTANT const SampleIteratorT d_samples, _CCCL_GRID_CONSTANT const ::cuda::std::array num_output_bins_wrapper, diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index d4b96a9efac3..087bd9fede88 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -138,6 +138,11 @@ struct HistogramPolicy return high_bin_threads_per_block != 0 ? high_bin_threads_per_block : threads_per_block; } + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_min_blocks() const noexcept + { + return high_bin_max_blocks_per_sm != 0 ? high_bin_max_blocks_per_sm : 1; + } + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const HistogramPolicy& lhs, const HistogramPolicy& rhs) noexcept { @@ -425,7 +430,7 @@ public: : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); - const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : 0; + const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : (is_even ? 1 : 2); const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; policy.high_bin_max_blocks_per_sm = high_bin_max_blocks_per_sm; @@ -451,7 +456,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - 4096, + 8192, 1, 262144, 4, @@ -499,7 +504,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - 4096, + is_even ? 8192 : 4096, 1, 262144, 4, @@ -522,7 +527,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - 1024, + is_even ? 2048 : 1024, 4, 262144, 4, From 4cdffc62ebced3f146b63ef6716c79fd01ede78a Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 03:20:06 +0000 Subject: [PATCH 15/21] [cub] Match raw histogram fallback and hot path --- .../device/dispatch/dispatch_histogram.cuh | 4 +- .../dispatch/kernels/kernel_histogram.cuh | 221 +++++++++--------- .../dispatch/tuning/tuning_histogram.cuh | 13 +- cub/test/catch2_test_device_histogram.cu | 7 +- cub/test/catch2_test_device_histogram_env.cu | 2 +- 5 files changed, 130 insertions(+), 117 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 35d175a635f0..fab557bd7c64 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -374,8 +374,8 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( { const int cooperative_grid_capacity = cooperative_sm_occupancy * sm_count; const int tuned_grid_capacity = - active_policy.high_bin_max_blocks_per_sm > 0 - ? active_policy.high_bin_max_blocks_per_sm * sm_count + active_policy.high_bin_blocks_per_sm > 0 + ? active_policy.high_bin_blocks_per_sm * sm_count : cooperative_grid_capacity; histogram_sweep_occupancy = active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized diff --git a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh index 36287e6ec6d8..b1cabe7c3fbc 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -201,6 +201,18 @@ struct scale_fraction bool bins_equal_range; }; +template +struct unsigned_if_integral +{ + using type = T; +}; + +template +struct unsigned_if_integral +{ + using type = ::cuda::std::make_unsigned_t; +}; + template _CCCL_DEVICE _CCCL_FORCEINLINE const SampleValueT* sample_native_pointer(SampleIteratorT itr) { @@ -572,8 +584,9 @@ struct Transforms ::cuda::std::is_integral; #endif // !_CCCL_HAS_INT128() + using UnsignedCommonT = typename unsigned_if_integral::value>::type; using FractionStorageT = - ::cuda::std::_If::value, IntArithmeticT, CommonT>; + ::cuda::std::_If::value, IntArithmeticT, UnsignedCommonT>; union ScaleT { @@ -604,7 +617,7 @@ struct Transforms { ScaleT result; result.fraction.bins = static_cast(num_levels - 1); - if constexpr (UseFastDivision && ::cuda::std::is_integral_v) + if constexpr (::cuda::std::is_integral_v) { using UnsignedT = ::cuda::std::make_unsigned_t; const UnsignedT distance = @@ -725,9 +738,11 @@ struct Transforms } else { - return static_cast( - (static_cast(sample - min_level) * static_cast(scale.fraction.bins)) - / static_cast(scale.fraction.range)); + using UnsignedT = ::cuda::std::make_unsigned_t; + const IntArithmeticT distance = static_cast( + static_cast(static_cast(sample) - static_cast(min_level))); + return static_cast((distance * static_cast(scale.fraction.bins)) + / static_cast(scale.fraction.range)); } } @@ -1293,6 +1308,7 @@ struct AgentHistogramCooperative static constexpr HistogramPolicy policy = current_policy(); static constexpr int count_replicas = policy.high_bin_cache_count_replicas; static_assert(policy.high_bin_pixels_per_thread > 0, "Histogram cooperative unroll must be positive"); + static_assert(policy.high_bin_blocks_per_sm >= 0, "Histogram cooperative blocks per SM must not be negative"); static_assert( policy.high_bin_cache == HistogramCacheAlgorithm::none || (policy.high_bin_cache_entries_per_channel >= 32 @@ -1466,9 +1482,9 @@ struct AgentHistogramCooperative }; using SampleValueT = it_value_t; - if constexpr (NumActiveChannels == 1) + if (contiguous_input) { - if (contiguous_input) + if constexpr (NumActiveChannels == 1) { for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { @@ -1517,102 +1533,107 @@ struct AgentHistogramCooperative } else { - for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) + if constexpr ((NumChannels == 2 || NumChannels == 4) && ::cuda::std::is_trivially_copyable_v) { - const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + using PixelT = typename CubVector::Type; + const auto* native_base = sample_native_pointer(d_samples); + const bool vectorizable = + native_base != nullptr && (reinterpret_cast(native_base) & (alignof(PixelT) - 1)) == 0; + if (vectorizable) { - const OffsetT pixel = first_pixel + static_cast(item) * step; - if (pixel < total_pixels) + const PixelT* const pixels = reinterpret_cast(native_base); + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { - const OffsetT row = pixel / num_row_pixels; - const OffsetT pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; - int bin = -1; - if constexpr (use_mru_cache) - { - decode_op[0].template BinSelect(d_samples[pixel_offset], bin, true, bracket_cache[0]); - } - else - { - decode_op[0].template BinSelect(d_samples[pixel_offset], bin, true); - } - if (bin < num_output_bins_wrapper[0]) + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) { - consume_bin(0, bin); + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + const PixelT packed = pixels[valid ? pixel : OffsetT{0}]; + const SampleValueT* lanes = reinterpret_cast(&packed); + int bins[NumActiveChannels]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int bin = -1; + if (valid) + { + decode_op[ch].template BinSelect(lanes[ch], bin, true); + if (bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } + } + bins[ch] = bin; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + consume_bin(ch, bins[ch]); + } } } } - } - } - } - else - { - const auto consume_pixels = [&](auto load_pixel) { - for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) - { - const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; - _CCCL_PRAGMA_UNROLL_FULL() - for (int item = 0; item < unroll; ++item) + else { - const OffsetT pixel = first_pixel + static_cast(item) * step; - const bool valid = pixel < total_pixels; - SampleValueT samples[NumActiveChannels]; - int bins[NumActiveChannels]; - load_pixel(pixel, valid, samples); - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { - int bin = -1; - if (valid) + const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) { - decode_op[ch].template BinSelect(samples[ch], bin, true); - if (bin >= num_output_bins_wrapper[ch]) + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + const OffsetT safe_pixel = valid ? pixel : OffsetT{0}; + int bins[NumActiveChannels]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const SampleValueT sample = d_samples[safe_pixel * NumChannels + ch]; + int bin = -1; + if (valid) + { + decode_op[ch].template BinSelect(sample, bin, true); + if (bin >= num_output_bins_wrapper[ch]) + { + bin = -1; + } + } + bins[ch] = bin; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) { - bin = -1; + consume_bin(ch, bins[ch]); } } - bins[ch] = bin; - } - - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - consume_bin(ch, bins[ch]); } } } - }; - - if constexpr ((NumChannels == 2 || NumChannels == 4) && ::cuda::std::is_trivially_copyable_v) - { - using PixelT = typename CubVector::Type; - const auto* native_base = sample_native_pointer(d_samples); - const bool vectorizable = contiguous_input && native_base != nullptr - && (reinterpret_cast(native_base) & (alignof(PixelT) - 1)) == 0; - if (vectorizable) + else { - const PixelT* const pixels = reinterpret_cast(native_base); for (OffsetT chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { const OffsetT first_pixel = static_cast(tid_global) + chunk_idx * chunk; _CCCL_PRAGMA_UNROLL_FULL() for (int item = 0; item < unroll; ++item) { - const OffsetT pixel = first_pixel + static_cast(item) * step; - const bool valid = pixel < total_pixels; - const PixelT packed = pixels[valid ? pixel : OffsetT{0}]; - const SampleValueT* lanes = reinterpret_cast(&packed); + const OffsetT pixel = first_pixel + static_cast(item) * step; + const bool valid = pixel < total_pixels; + const OffsetT safe_pixel = valid ? pixel : OffsetT{0}; int bins[NumActiveChannels]; _CCCL_PRAGMA_UNROLL_FULL() for (int ch = 0; ch < NumActiveChannels; ++ch) { - int bin = -1; + const SampleValueT sample = d_samples[safe_pixel * NumChannels + ch]; + int bin = -1; if (valid) { - decode_op[ch].template BinSelect(lanes[ch], bin, true); + decode_op[ch].template BinSelect(sample, bin, true); if (bin >= num_output_bins_wrapper[ch]) { bin = -1; @@ -1628,46 +1649,32 @@ struct AgentHistogramCooperative } } } - else - { - consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { - OffsetT pixel_offset{}; - if (valid) - { - const OffsetT row = contiguous_input ? OffsetT{0} : pixel / num_row_pixels; - const OffsetT col = contiguous_input ? pixel : pixel - row * num_row_pixels; - pixel_offset = row * row_stride_samples + col * NumChannels; - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) - { - if (valid) - { - samples[ch] = d_samples[pixel_offset + ch]; - } - } - }); - } } - else + } + else + { + for (OffsetT pixel = static_cast(tid_global); pixel < total_pixels; pixel += step) { - consume_pixels([&](OffsetT pixel, bool valid, SampleValueT(&samples)[NumActiveChannels]) { - OffsetT pixel_offset{}; - if (valid) + const OffsetT row = pixel / num_row_pixels; + const OffsetT pixel_offset = row * row_stride_samples + (pixel - row * num_row_pixels) * NumChannels; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int bin = -1; + if constexpr (use_mru_cache) { - const OffsetT row = contiguous_input ? OffsetT{0} : pixel / num_row_pixels; - const OffsetT col = contiguous_input ? pixel : pixel - row * num_row_pixels; - pixel_offset = row * row_stride_samples + col * NumChannels; + decode_op[ch].template BinSelect(d_samples[pixel_offset + ch], bin, true, bracket_cache[ch]); } - _CCCL_PRAGMA_UNROLL_FULL() - for (int ch = 0; ch < NumActiveChannels; ++ch) + else { - if (valid) - { - samples[ch] = d_samples[pixel_offset + ch]; - } + decode_op[ch].template BinSelect(d_samples[pixel_offset + ch], bin, true); + } + if (bin >= 0 && bin < num_output_bins_wrapper[ch]) + { + spill_bin(ch, bin, CounterT{1}); } - }); + } } } diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 087bd9fede88..619a6e55c27e 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -131,7 +131,8 @@ struct HistogramPolicy int high_bin_threads_per_block = 0; //!< High-bin block size; 0 inherits threads_per_block int high_bin_interpolation_min_bins = 512; int high_bin_min_histogram_bytes = 0; - int high_bin_max_blocks_per_sm = 0; //!< Zero uses the occupancy-derived capacity + //! Target resident cooperative blocks per SM. Zero keeps the occupancy-derived grid and a one-block launch bound. + int high_bin_blocks_per_sm = 0; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_threads() const noexcept { @@ -140,7 +141,7 @@ struct HistogramPolicy [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_min_blocks() const noexcept { - return high_bin_max_blocks_per_sm != 0 ? high_bin_max_blocks_per_sm : 1; + return high_bin_blocks_per_sm != 0 ? high_bin_blocks_per_sm : 1; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -160,7 +161,7 @@ struct HistogramPolicy && lhs.high_bin_threads_per_block == rhs.high_bin_threads_per_block && lhs.high_bin_interpolation_min_bins == rhs.high_bin_interpolation_min_bins && lhs.high_bin_min_histogram_bytes == rhs.high_bin_min_histogram_bytes - && lhs.high_bin_max_blocks_per_sm == rhs.high_bin_max_blocks_per_sm; + && lhs.high_bin_blocks_per_sm == rhs.high_bin_blocks_per_sm; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -186,7 +187,7 @@ struct HistogramPolicy << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread << ", .high_bin_threads_per_block = " << p.high_bin_threads_per_block << ", .high_bin_interpolation_min_bins = " << p.high_bin_interpolation_min_bins << ", .high_bin_min_histogram_bytes = " << p.high_bin_min_histogram_bytes - << ", .high_bin_max_blocks_per_sm = " << p.high_bin_max_blocks_per_sm << " }"; + << ", .high_bin_blocks_per_sm = " << p.high_bin_blocks_per_sm << " }"; } #endif // _CCCL_HOSTED() }; @@ -430,10 +431,10 @@ public: : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); - const int high_bin_max_blocks_per_sm = num_active_channels == 1 ? 2 : (is_even ? 1 : 2); + const int high_bin_blocks_per_sm = num_active_channels == 1 ? 2 : (is_even ? 1 : 2); const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; - policy.high_bin_max_blocks_per_sm = high_bin_max_blocks_per_sm; + policy.high_bin_blocks_per_sm = high_bin_blocks_per_sm; return policy; }; diff --git a/cub/test/catch2_test_device_histogram.cu b/cub/test/catch2_test_device_histogram.cu index ee6726d1ff10..2d8176dff245 100644 --- a/cub/test/catch2_test_device_histogram.cu +++ b/cub/test/catch2_test_device_histogram.cu @@ -738,7 +738,7 @@ CUB_TEST_LIST( constexpr sample_t lower_level = cs::numeric_limits::min(); constexpr sample_t upper_level = cs::numeric_limits::max(); constexpr auto num_samples = 1000; - auto d_samples = cuda::counting_iterator{0UL}; + auto d_samples = cuda::constant_iterator{lower_level}; auto d_histo_out = c2h::device_vector(1024); const auto num_bins = GENERATE(1, 2); @@ -775,6 +775,11 @@ CUB_TEST_LIST( // types, hence we expect cudaErrorInvalidValue to be returned to indicate of a potential overflow // Ensure we do not return an error on querying temporary storage requirements CHECK(error2 == (num_bins == 1 || sizeof(sample_t) <= 4UL ? cudaSuccess : cudaErrorInvalidValue)); + + if (error2 == cudaSuccess && sizeof(sample_t) > 1) + { + CHECK(c2h::host_vector(d_histo_out)[0] == num_samples); + } } // When the number of bins exceeds what LevelT can represent, the bin computation will overflow diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index c57bcb9165a7..419589de374e 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1980,7 +1980,7 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) ", .high_bin_cache_count_replicas = 1, .high_bin_cache_cuckoo_max_bins = 262144" ", .high_bin_pixels_per_thread = 4, .high_bin_threads_per_block = 0" ", .high_bin_interpolation_min_bins = 512, .high_bin_min_histogram_bytes = 0" - ", .high_bin_max_blocks_per_sm = 0 }"); + ", .high_bin_blocks_per_sm = 0 }"); constexpr auto high_bin_policy = [] { auto policy = From a4100ca85197b6a52764cec645a3cbd2a0a1b68d Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 06:55:41 +0000 Subject: [PATCH 16/21] [cub] Match raw RANGE cache capacity --- cub/cub/device/dispatch/tuning/tuning_histogram.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 619a6e55c27e..3699b522e364 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -505,7 +505,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - is_even ? 8192 : 4096, + 8192, 1, 262144, 4, From ef4d7b49d4d49f6a6036bd8255ad2aabff3807f8 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 07:12:45 +0000 Subject: [PATCH 17/21] [cub] Match raw cooperative grid sizing --- .../device/dispatch/dispatch_histogram.cuh | 20 ++++++++++++++----- .../dispatch/tuning/tuning_histogram.cuh | 4 ++-- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index fab557bd7c64..6d7a08e9b5e2 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -403,11 +403,21 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy - const int selected_pixels_per_thread = use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread; - const int selected_threads_per_block = use_cooperative ? high_bin_threads_per_block : threads_per_block; - const int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; - int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); - int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); + // A privatized high-bin block owns a full histogram slab. Preserve the ordinary + // sweep's useful-work grid limit so small inputs do not allocate, initialize, and + // gather more slabs than the corresponding sweep would launch. The cooperative + // kernel still launches with its independently tuned block size. + const bool use_privatized_work_limit = + use_cooperative && active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized; + const int selected_pixels_per_thread = + use_privatized_work_limit + ? pixels_per_thread + : (use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread); + const int selected_threads_per_block = + use_privatized_work_limit ? threads_per_block : (use_cooperative ? high_bin_threads_per_block : threads_per_block); + const int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; + int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); + int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); int blocks_per_col = (blocks_per_row > 0) ? int(::cuda::std::min(static_cast(histogram_sweep_occupancy / blocks_per_row), num_rows)) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 3699b522e364..8293edc434e1 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -431,7 +431,7 @@ public: : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); - const int high_bin_blocks_per_sm = num_active_channels == 1 ? 2 : (is_even ? 1 : 2); + const int high_bin_blocks_per_sm = num_active_channels == 1 ? 2 : 1; const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; policy.high_bin_blocks_per_sm = high_bin_blocks_per_sm; @@ -528,7 +528,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - is_even ? 2048 : 1024, + 2048, 4, 262144, 4, From fe59656837edad5a7c703d0c12bbef57d991899b Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 08:00:22 +0000 Subject: [PATCH 18/21] [cub] Tune high-bin work-grid size --- .../device/dispatch/dispatch_histogram.cuh | 23 +++++++------------ .../dispatch/tuning/tuning_histogram.cuh | 19 +++++++++++++-- 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 6d7a08e9b5e2..192b62eaa9a8 100644 --- a/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -403,21 +403,14 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( } // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy - // A privatized high-bin block owns a full histogram slab. Preserve the ordinary - // sweep's useful-work grid limit so small inputs do not allocate, initialize, and - // gather more slabs than the corresponding sweep would launch. The cooperative - // kernel still launches with its independently tuned block size. - const bool use_privatized_work_limit = - use_cooperative && active_policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized; - const int selected_pixels_per_thread = - use_privatized_work_limit - ? pixels_per_thread - : (use_cooperative ? active_policy.high_bin_pixels_per_thread : pixels_per_thread); - const int selected_threads_per_block = - use_privatized_work_limit ? threads_per_block : (use_cooperative ? high_bin_threads_per_block : threads_per_block); - const int pixels_per_tile = selected_threads_per_block * selected_pixels_per_thread; - int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); - int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); + // A privatized high-bin block owns a full histogram slab. Its policy therefore + // carries a separate useful-work tile size so small inputs do not initialize and + // gather more full slabs than needed. This is independent of the kernel's four-item + // processing unroll and its launch block size. + const int pixels_per_tile = + use_cooperative ? active_policy.high_bin_grid_pixels() : threads_per_block * pixels_per_thread; + int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); + int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); int blocks_per_col = (blocks_per_row > 0) ? int(::cuda::std::min(static_cast(histogram_sweep_occupancy / blocks_per_row), num_rows)) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 8293edc434e1..2510fef22936 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -133,6 +133,8 @@ struct HistogramPolicy int high_bin_min_histogram_bytes = 0; //! Target resident cooperative blocks per SM. Zero keeps the occupancy-derived grid and a one-block launch bound. int high_bin_blocks_per_sm = 0; + //! Input pixels represented by one block when limiting the cooperative grid. Zero uses the kernel tile size. + int high_bin_grid_pixels_per_block = 0; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_threads() const noexcept { @@ -144,6 +146,13 @@ struct HistogramPolicy return high_bin_blocks_per_sm != 0 ? high_bin_blocks_per_sm : 1; } + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int high_bin_grid_pixels() const noexcept + { + return high_bin_grid_pixels_per_block != 0 + ? high_bin_grid_pixels_per_block + : high_bin_threads() * high_bin_pixels_per_thread; + } + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const HistogramPolicy& lhs, const HistogramPolicy& rhs) noexcept { @@ -161,7 +170,8 @@ struct HistogramPolicy && lhs.high_bin_threads_per_block == rhs.high_bin_threads_per_block && lhs.high_bin_interpolation_min_bins == rhs.high_bin_interpolation_min_bins && lhs.high_bin_min_histogram_bytes == rhs.high_bin_min_histogram_bytes - && lhs.high_bin_blocks_per_sm == rhs.high_bin_blocks_per_sm; + && lhs.high_bin_blocks_per_sm == rhs.high_bin_blocks_per_sm + && lhs.high_bin_grid_pixels_per_block == rhs.high_bin_grid_pixels_per_block; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -187,7 +197,8 @@ struct HistogramPolicy << ", .high_bin_pixels_per_thread = " << p.high_bin_pixels_per_thread << ", .high_bin_threads_per_block = " << p.high_bin_threads_per_block << ", .high_bin_interpolation_min_bins = " << p.high_bin_interpolation_min_bins << ", .high_bin_min_histogram_bytes = " << p.high_bin_min_histogram_bytes - << ", .high_bin_blocks_per_sm = " << p.high_bin_blocks_per_sm << " }"; + << ", .high_bin_blocks_per_sm = " << p.high_bin_blocks_per_sm + << ", .high_bin_grid_pixels_per_block = " << p.high_bin_grid_pixels_per_block << " }"; } #endif // _CCCL_HOSTED() }; @@ -435,6 +446,10 @@ public: const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; policy.high_bin_blocks_per_sm = high_bin_blocks_per_sm; + policy.high_bin_grid_pixels_per_block = + num_active_channels == 1 + ? (sample_size == 1 ? policy.threads_per_block * policy.pixels_per_thread : 768 * t_scale(12)) + : 1024 * (is_even ? t_scale(8) : t_scale(16)); return policy; }; From 0db9dec1356c4f5f56bcea83228ebe8895a680be Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 08:30:33 +0000 Subject: [PATCH 19/21] [cub] Fix high-bin policy validation --- cub/cub/device/dispatch/tuning/tuning_histogram.cuh | 9 +++++---- cub/test/catch2_test_device_histogram_env.cu | 4 +++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 2510fef22936..ba883e86f92b 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -443,13 +443,14 @@ public: : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); const int high_bin_blocks_per_sm = num_active_channels == 1 ? 2 : 1; - const auto with_high_bin_threshold = [=](HistogramPolicy policy) { + const bool use_ordinary_grid_tile = num_active_channels == 1 && sample_size == 1; + const int high_bin_grid_pixels_per_block = + num_active_channels == 1 ? 768 * t_scale(12) : 1024 * (is_even ? t_scale(8) : t_scale(16)); + const auto with_high_bin_threshold = [=](HistogramPolicy policy) { policy.high_bin_min_histogram_bytes = high_bin_min_histogram_bytes; policy.high_bin_blocks_per_sm = high_bin_blocks_per_sm; policy.high_bin_grid_pixels_per_block = - num_active_channels == 1 - ? (sample_size == 1 ? policy.threads_per_block * policy.pixels_per_thread : 768 * t_scale(12)) - : 1024 * (is_even ? t_scale(8) : t_scale(16)); + use_ordinary_grid_tile ? policy.threads_per_block * policy.pixels_per_thread : high_bin_grid_pixels_per_block; return policy; }; diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index 419589de374e..f1ff3705461a 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1980,7 +1980,7 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) ", .high_bin_cache_count_replicas = 1, .high_bin_cache_cuckoo_max_bins = 262144" ", .high_bin_pixels_per_thread = 4, .high_bin_threads_per_block = 0" ", .high_bin_interpolation_min_bins = 512, .high_bin_min_histogram_bytes = 0" - ", .high_bin_blocks_per_sm = 0 }"); + ", .high_bin_blocks_per_sm = 0, .high_bin_grid_pixels_per_block = 0 }"); constexpr auto high_bin_policy = [] { auto policy = @@ -1990,5 +1990,7 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) }(); STATIC_REQUIRE(high_bin_policy.high_bin_threads() == 512); STATIC_REQUIRE(p1.high_bin_threads() == p1.threads_per_block); + STATIC_REQUIRE(high_bin_policy.high_bin_grid_pixels() == 512 * high_bin_policy.high_bin_pixels_per_thread); + STATIC_REQUIRE(p1.high_bin_grid_pixels() == p1.high_bin_threads() * p1.high_bin_pixels_per_thread); } #endif // _CCCL_COMPILER(GCC, >=, 8) From 73f31295a4227681d40bf25d912ed8e3258e7703 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Fri, 4 Sep 2026 08:37:08 +0000 Subject: [PATCH 20/21] [cub] Match raw multi-channel RANGE occupancy --- cub/cub/device/dispatch/tuning/tuning_histogram.cuh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index ba883e86f92b..036b7a51bff9 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -442,8 +442,9 @@ public: : is_even ? even_smem_bytes_per_channel * num_active_channels : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); - const int high_bin_blocks_per_sm = num_active_channels == 1 ? 2 : 1; - const bool use_ordinary_grid_tile = num_active_channels == 1 && sample_size == 1; + // The raw occupancy-sized cache resolves to two resident blocks for four-byte multi-channel RANGE samples. + const int high_bin_blocks_per_sm = num_active_channels == 1 || (!is_even && sample_size == 4) ? 2 : 1; + const bool use_ordinary_grid_tile = num_active_channels == 1 && sample_size == 1; const int high_bin_grid_pixels_per_block = num_active_channels == 1 ? 768 * t_scale(12) : 1024 * (is_even ? t_scale(8) : t_scale(16)); const auto with_high_bin_threshold = [=](HistogramPolicy policy) { @@ -544,7 +545,7 @@ public: HistogramCacheAlgorithm::single_probe, HistogramSpillAlgorithm::global_memory_privatized, HistogramAggregationAlgorithm::rle, - 2048, + is_even || sample_size == 8 ? 2048 : 1024, 4, 262144, 4, From 3fa42516e4ca1a5c1b0cba49f3d7766d314ce0e7 Mon Sep 17 00:00:00 2001 From: Bryce Adelstein Lelbach Date: Sat, 5 Sep 2026 00:30:40 +0000 Subject: [PATCH 21/21] [cub] Preserve multi-channel EVEN SMEM range --- .../dispatch/tuning/tuning_histogram.cuh | 10 +++++++--- cub/test/catch2_test_device_histogram_env.cu | 20 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh index 036b7a51bff9..f5bd0715e30b 100644 --- a/cub/cub/device/dispatch/tuning/tuning_histogram.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_histogram.cuh @@ -437,10 +437,14 @@ public: constexpr int sm100_smem_bytes = 228352; constexpr int range_smem_bytes_per_channel = 8192; constexpr int even_smem_bytes_per_channel = 32768; + // Preserve the raw selector's dynamic-SMEM region. Three-active-channel EVEN + // can use the full per-CTA capacity; four-channel EVEN and RANGE switch at + // their measured per-channel crossover budgets. + const bool use_full_smem_capacity = num_active_channels == 1 || (is_even && num_active_channels <= 3); const int candidate_smem_bytes = - num_active_channels == 1 ? sm100_smem_bytes - : is_even ? even_smem_bytes_per_channel * num_active_channels - : range_smem_bytes_per_channel * num_active_channels; + use_full_smem_capacity ? sm100_smem_bytes + : is_even ? even_smem_bytes_per_channel * num_active_channels + : range_smem_bytes_per_channel * num_active_channels; const int high_bin_min_histogram_bytes = (::cuda::std::min) (candidate_smem_bytes, sm100_smem_bytes); // The raw occupancy-sized cache resolves to two resident blocks for four-byte multi-channel RANGE samples. const int high_bin_blocks_per_sm = num_active_channels == 1 || (!is_even && sample_size == 4) ? 2 : 1; diff --git a/cub/test/catch2_test_device_histogram_env.cu b/cub/test/catch2_test_device_histogram_env.cu index f1ff3705461a..2dcf8b8cf52a 100644 --- a/cub/test/catch2_test_device_histogram_env.cu +++ b/cub/test/catch2_test_device_histogram_env.cu @@ -1992,5 +1992,25 @@ CUB_TEST("Test HistogramPolicy properties", "[histogram][device]", CUB_SMALL) STATIC_REQUIRE(p1.high_bin_threads() == p1.threads_per_block); STATIC_REQUIRE(high_bin_policy.high_bin_grid_pixels() == 512 * high_bin_policy.high_bin_pixels_per_thread); STATIC_REQUIRE(p1.high_bin_grid_pixels() == p1.high_bin_threads() * p1.high_bin_pixels_per_thread); + + constexpr auto sm100 = cuda::compute_capability{10, 0}; + constexpr auto single_channel_even_policy = + cub::detail::histogram::policy_selector_from_types{}(sm100); + constexpr auto three_channel_even_policy = + cub::detail::histogram::policy_selector_from_types{}(sm100); + constexpr auto four_channel_even_policy = + cub::detail::histogram::policy_selector_from_types{}(sm100); + constexpr auto three_channel_range_policy = + cub::detail::histogram::policy_selector_from_types{}(sm100); + + // Match the raw selector's dynamic-SMEM tiers: three-active-channel EVEN uses + // the full SM100 capacity, while four-channel EVEN and RANGE retain their + // measured per-channel crossover budgets. + STATIC_REQUIRE( + three_channel_even_policy.high_bin_min_histogram_bytes == single_channel_even_policy.high_bin_min_histogram_bytes); + STATIC_REQUIRE( + four_channel_even_policy.high_bin_min_histogram_bytes < three_channel_even_policy.high_bin_min_histogram_bytes); + STATIC_REQUIRE( + three_channel_range_policy.high_bin_min_histogram_bytes < three_channel_even_policy.high_bin_min_histogram_bytes); } #endif // _CCCL_COMPILER(GCC, >=, 8)