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..ad5c8091f076 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( + 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 { int device_ordinal; diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh index 875f6744afef..192b62eaa9a8 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; @@ -129,6 +160,24 @@ struct DeviceHistogramKernelSource IsEven>; } + /// Returns the policy-configurable cooperative high-bin histogram kernel. + 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>; + } + CUB_RUNTIME_FUNCTION static constexpr size_t CounterSize() { return sizeof(CounterT); @@ -146,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 @@ -166,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, @@ -191,6 +247,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))) { @@ -233,8 +291,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; @@ -252,7 +311,88 @@ 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; + 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) + { + 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))) + { + return error; + } + if (use_cooperative) + { + using privatized_decode_op_t = typename CooperativeSecondLevelArrayT::value_type; + const auto cooperative_kernel = + kernel_source.template HistogramCooperativeKernel(); + cooperative_cache_slots_per_channel = + active_policy.high_bin_cache == HistogramCacheAlgorithm::none + ? 0 + : active_policy.high_bin_cache_entries_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(LocalCounterT)}); + + 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; + } + const int max_slots_by_smem = cache_bytes_per_slot == 0 ? 0 : max_dynamic_smem_bytes / cache_bytes_per_slot; + if (cooperative_cache_slots_per_channel > max_slots_by_smem) + { + use_cooperative = false; + } + else + { + 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(launcher_factory.MaxSmOccupancy( + cooperative_sm_occupancy, cooperative_kernel, high_bin_threads_per_block, cooperative_smem_bytes))) + { + return error; + } + if (cooperative_sm_occupancy > 0) + { + const int cooperative_grid_capacity = cooperative_sm_occupancy * sm_count; + const int tuned_grid_capacity = + 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 + ? ::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 + { + use_cooperative = false; + } + } + } + } + } + })) +#endif // _CCCL_HOSTED() if (num_row_pixels * NUM_CHANNELS == row_stride_samples) { @@ -263,9 +403,14 @@ 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); + // 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)) @@ -284,8 +429,16 @@ 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) + * ((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; } allocation_sizes[NUM_ALLOCATIONS - 1] = GridQueue::AllocationSize(); @@ -309,11 +462,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); + auto* const typed_allocations = reinterpret_cast(allocations); ::cuda::std::copy(typed_allocations, typed_allocations + NUM_ACTIVE_CHANNELS, d_privatized_histograms_wrapper.begin()); + 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()); auto minus_one = ::cuda::proclaim_return_type([](int levels) { return levels - 1; @@ -322,82 +480,122 @@ 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_cooperative = false; +#if _CCCL_HOSTED() + if constexpr (!IsDeviceInit && PRIVATIZED_SMEM_BINS == 0) + { + if (use_cooperative && blocks_per_row > 0 && blocks_per_col > 0) + { + 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 = + kernel_source.template HistogramCooperativeKernel(); + if (const auto error = CubDebug(launcher_factory.LaunchCooperative( + cooperative_grid_dims, + dim3{static_cast(high_bin_threads_per_block)}, + cooperative_smem_bytes, + stream, + cooperative_kernel, + d_samples, + num_output_bins_wrapper, + d_output_histograms, + d_cooperative_privatized_histograms_wrapper, + cooperative_second_level_array, + num_row_pixels, + num_rows, + row_stride_samples, + cooperative_cache_slots_per_channel))) + { + return error; + } + launched_cooperative = true; + } + } +#endif // _CCCL_HOSTED() + + if (!launched_cooperative) + { + 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 + // 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 @@ -548,6 +746,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, @@ -580,6 +779,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, @@ -720,6 +920,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, @@ -871,6 +1072,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, @@ -887,19 +1089,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]; @@ -911,7 +1108,16 @@ 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; + 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( (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]; @@ -1104,8 +1315,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; @@ -1115,14 +1324,24 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( if (max_num_output_bins > max_privatized_smem_bins) { - constexpr int PRIVATIZED_SMEM_BINS = 0; + 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( (detail::histogram::dispatch( d_temp_storage, temp_storage_bytes, @@ -1132,6 +1351,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, @@ -1147,13 +1367,19 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_even( else { constexpr int PRIVATIZED_SMEM_BINS = max_privatized_smem_bins; + using PrivatizedDecodeOpT = typename TransformsT::ScaleTransform; + ::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( d_temp_storage, temp_storage_bytes, @@ -1163,6 +1389,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 cc64ddea1862..b1cabe7c3fbc 100644 --- a/cub/cub/device/dispatch/kernels/kernel_histogram.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_histogram.cuh @@ -17,13 +17,231 @@ #include #include #include +#include #include #include +#include +#include +#include + +#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 (log2_divisor == bits) + { + magic = divisor; + shift = 0; + mode = 3; + return; + } + 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; + } + if (mode == 3) + { + return numerator / magic; + } + + 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 +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 +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) +{ + 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) +{ + 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 + { + atomicAdd(address, value); + } +} + template struct Transforms { @@ -74,8 +292,257 @@ struct Transforms } }; - // Scales samples to evenly-spaced bins - struct ScaleTransform + //! @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 auto interpolation_difference(T lhs, T rhs) + { + if constexpr (::cuda::std::is_integral_v) + { + using UnsignedT = ::cuda::std::make_unsigned_t; + return static_cast(lhs) - static_cast(rhs); + } + else + { + 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; + 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; + 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) + { + 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; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid, BracketCacheT& bracket) const + { + 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 (bracket.bin >= 0 && !(value < bracket.lo) && value < bracket.hi) + { + bin = bracket.bin; + return; + } + + if (!has_precompute) + { + bin = UpperBound(wrapped_levels, num_output_levels, value) - 1; + if (bin >= num_bins) + { + bin = -1; + } + } + else if (value < first || !(value < last)) + { + bin = -1; + } + else + { + 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; + } + + 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; + } + } + + if (bin >= 0) + { + bracket = BracketCacheT{wrapped_levels[bin], wrapped_levels[bin + 1], bin}; + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BinSelect(_SampleT sample, int& bin, bool valid) const + { + 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; + } + } + }; + + // Scales samples to evenly-spaced bins. + template + struct ScaleTransformImpl { using CommonT = ::cuda::std::common_type_t; static_assert(::cuda::std::is_convertible_v, @@ -117,15 +584,15 @@ 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, UnsignedCommonT>; + union ScaleT { // Used when CommonT is not floating-point to avoid intermediate // rounding errors (see NVIDIA/cub#489). - struct FractionT - { - CommonT bins; - CommonT range; - } fraction; + scale_fraction fraction; // Used when CommonT is floating-point as an optimization. CommonT reciprocal; @@ -149,8 +616,27 @@ 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); + } + 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; } @@ -234,9 +720,30 @@ 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)); + 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); + } + 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 + { + 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)); + } } template ::value, int> = 0> @@ -256,6 +763,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) { @@ -278,9 +791,17 @@ struct Transforms } }; + using ScaleTransform = ScaleTransformImpl; + using FastScaleTransform = ScaleTransformImpl; + // 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) @@ -703,5 +1224,576 @@ __launch_bounds__(int(current_policy().threads_per_block)) // Store output to global (if necessary) agent.StoreOutput(); } + +template +_CCCL_DEVICE _CCCL_FORCEINLINE bool histogram_cache_probe( + ::cuda::std::uint32_t* keys, + CounterT* counts, + int bin, + CounterT contribution, + int cache_mask, + int cache_log2, + bool use_second_probe) +{ + if constexpr (CacheAlgorithm == HistogramCacheAlgorithm::none) + { + return false; + } + + 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 == 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; + }; + + 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 (CacheAlgorithm == HistogramCacheAlgorithm::cuckoo) + { + 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; +} + +//! Agent for the policy-configurable cooperative high-bin histogram kernel. +template +#if _CCCL_HAS_CONCEPTS() + requires histogram_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +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_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) + { + 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 + && (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 of at least 32"); + 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; + + if constexpr (policy.high_bin_spill != HistogramSpillAlgorithm::global_memory_privatized) + { + _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] = OutputCounterT{0}; + } + } + grid.sync(); + } + + 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; + + _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) + { + channel_keys[slot] = ~::cuda::std::uint32_t{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) + { + 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) + { + 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; + + 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]; + ::cuda::std::uint32_t* channel_keys[NumActiveChannels]; + CounterT* thread_counts[NumActiveChannels]; + CounterT* private_histograms[NumActiveChannels]; + int pending_bin[NumActiveChannels]; + CounterT pending_count[NumActiveChannels]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + 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; + 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) { + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) + { + if (pending_bin[ch] == selected_bin) + { + pending_count[ch] += contribution; + return; + } + if (pending_bin[ch] >= 0) + { + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + atomicAdd_block(&private_histograms[ch][pending_bin[ch]], pending_count[ch]); + } + else + { + 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) + { + atomicAdd_block(&private_histograms[ch][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) { + 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[ch], thread_counts[ch], selected_bin, contribution, cache_mask, cache_log2, use_second_probe)) + { + spill_bin(ch, selected_bin, contribution); + } + }; + + const auto consume_bin = [&](int ch, int bin) { + 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 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))); + }), + (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 (contiguous_input) + { + 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; + 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]) + { + 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; + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < unroll; ++item) + { + consume_bin(0, bins[item]); + } + } + } + else + { + 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 = + native_base != nullptr && (reinterpret_cast(native_base) & (alignof(PixelT) - 1)) == 0; + if (vectorizable) + { + 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); + 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 + { + 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 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) + { + consume_bin(ch, bins[ch]); + } + } + } + } + } + else + { + 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 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) + { + consume_bin(ch, bins[ch]); + } + } + } + } + } + } + else + { + for (OffsetT pixel = static_cast(tid_global); pixel < total_pixels; pixel += step) + { + 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) + { + decode_op[ch].template BinSelect(d_samples[pixel_offset + ch], bin, true, bracket_cache[ch]); + } + else + { + 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}); + } + } + } + } + + if constexpr (policy.high_bin_aggregation == HistogramAggregationAlgorithm::rle) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + if (pending_bin[ch] >= 0) + { + if constexpr (policy.high_bin_spill == HistogramSpillAlgorithm::global_memory_privatized) + { + atomicAdd_block(&private_histograms[ch][pending_bin[ch]], pending_count[ch]); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][pending_bin[ch]], + static_cast(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}) + { + 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) + { + atomicAdd_block(&private_histograms[ch][key], count); + } + else + { + histogram_atomic_add(&d_output_histograms_wrapper[ch][key], static_cast(count)); + } + } + } + } + } + + 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) + { + OutputCounterT total = OutputCounterT{0}; + for (unsigned int block = 0; block < gridDim.x; ++block) + { + total += static_cast( + 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().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, + ::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, + OutputCounterT, + 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..f5bd0715e30b 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,38 @@ 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::global_memory_privatized; + 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; + 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 + { + 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_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 @@ -47,7 +160,18 @@ 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 + && 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_grid_pixels_per_block == rhs.high_bin_grid_pixels_per_block; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -64,7 +188,17 @@ 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 << ", .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_grid_pixels_per_block = " << p.high_bin_grid_pixels_per_block << " }"; } #endif // _CCCL_HOSTED() }; @@ -298,22 +432,130 @@ 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 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 = + 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; + 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 = + use_ordinary_grid_tile ? policy.threads_per_block * policy.pixels_per_thread : high_bin_grid_pixels_per_block; + 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{928, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_CA, false, SMEM, false, 2048}; + return with_high_bin_threshold(HistogramPolicy{ + 928, + 12, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_CA, + false, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + 8192, + 1, + 262144, + 4, + 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{448, 12, 1 << 2, BLOCK_LOAD_DIRECT, LOAD_LDG, false, SMEM, false, 2048}; + return with_high_bin_threshold(HistogramPolicy{ + 448, + 12, + 1 << 2, + BLOCK_LOAD_DIRECT, + LOAD_LDG, + false, + SMEM, + false, + 2048, + HistogramHighBinAlgorithm::cooperative, + HistogramCacheAlgorithm::single_probe, + HistogramSpillAlgorithm::global_memory_privatized, + HistogramAggregationAlgorithm::rle, + 4096, + 1, + 262144, + 4, + 0}); } } + if (counter_size == 4 && sample_is_primitive && num_channels == 1 && num_active_channels == 1 + && (sample_size == 4 || sample_size == 8)) + { + return with_high_bin_threshold(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, + 8192, + 1, + 262144, + 4, + is_even ? 768 : 512}); + } + + if (counter_size == 4 && sample_is_primitive && num_channels >= 2) + { + return with_high_bin_threshold(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, + is_even || sample_size == 8 ? 2048 : 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 } diff --git a/cub/test/catch2_test_device_histogram.cu b/cub/test/catch2_test_device_histogram.cu index 20e78b7079cc..2d8176dff245 100644 --- a/cub/test/catch2_test_device_histogram.cu +++ b/cub/test/catch2_test_device_histogram.cu @@ -720,20 +720,25 @@ 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}; + auto d_samples = cuda::constant_iterator{lower_level}; auto d_histo_out = c2h::device_vector(1024); const auto num_bins = GENERATE(1, 2); @@ -770,6 +775,11 @@ CUB_TEST_LIST("DeviceHistogram::HistogramEven bin computation does not overflow" // 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 6712bc678055..2dcf8b8cf52a 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) @@ -1630,6 +1657,54 @@ 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 = CacheEntriesPerChannel; + 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; + policy.high_bin_min_histogram_bytes = 0; + 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>; @@ -1740,6 +1815,121 @@ 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; + 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); + 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{}); + 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 #if _CCCL_COMPILER(GCC, >=, 8) // gcc 7 cannot preserve constexpr-ness from p1 to p2 @@ -1777,9 +1967,50 @@ 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::global_memory_privatized" + ", .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_threads_per_block = 0" + ", .high_bin_interpolation_min_bins = 512, .high_bin_min_histogram_bytes = 0" + ", .high_bin_blocks_per_sm = 0, .high_bin_grid_pixels_per_block = 0 }"); + + 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); + 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) diff --git a/cub/test/catch2_test_env_launch_helper.h b/cub/test/catch2_test_env_launch_helper.h index 686179a6c1c3..326bf4b87389 100644 --- a/cub/test/catch2_test_env_launch_helper.h +++ b/cub/test/catch2_test_env_launch_helper.h @@ -140,8 +140,49 @@ 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( + reinterpret_cast(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)