From 5c45cdd610a79beb1787c5b0ca0242dba912d33d Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Tue, 19 May 2026 18:18:02 -0700 Subject: [PATCH 001/126] Simple cluster --- .../bench/segmented_topk/variable/keys.cu | 36 ++ cub/cub/agent/agent_batched_topk_cluster.cuh | 441 ++++++++++++++++++ .../dispatch_batched_topk_cluster.cuh | 220 +++++++++ .../catch2_test_device_segmented_topk_keys.cu | 136 ++++++ 4 files changed, 833 insertions(+) create mode 100644 cub/cub/agent/agent_batched_topk_cluster.cuh create mode 100644 cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index e0f2396e012..602f4b8006e 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -3,6 +3,14 @@ #include #include +#include + +// Compile-time switch: define BENCH_CLUSTER_TOPK=1 to drive the prototype +// cluster-based dispatch instead of the small-segment baseline. We default to +// the baseline so existing benchmark scripts keep their behavior. +#ifndef BENCH_CLUSTER_TOPK +# define BENCH_CLUSTER_TOPK 0 +#endif #include #include @@ -198,6 +206,19 @@ void variable_seg_size_topk_keys(nvbench::state& state, state.add_global_memory_writes(output_elements, "OutputKeys"); size_t temp_size{}; +#if BENCH_CLUSTER_TOPK + cub::detail::batched_topk_cluster::dispatch( + nullptr, + temp_size, + d_keys_in, + d_keys_out, + segment_sizes_param, + k_param, + select_directions, + num_segments_uniform_param, + total_num_items, + nullptr); +#else cub::detail::batched_topk::dispatch( nullptr, temp_size, @@ -211,11 +232,25 @@ void variable_seg_size_topk_keys(nvbench::state& state, num_segments_uniform_param, total_num_items, nullptr); +#endif thrust::device_vector temp(temp_size, thrust::no_init); auto* temp_storage = thrust::raw_pointer_cast(temp.data()); state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { +#if BENCH_CLUSTER_TOPK + cub::detail::batched_topk_cluster::dispatch( + temp_storage, + temp_size, + d_keys_in, + d_keys_out, + segment_sizes_param, + k_param, + select_directions, + num_segments_uniform_param, + total_num_items, + launch.get_stream()); +#else cub::detail::batched_topk::dispatch( temp_storage, temp_size, @@ -229,6 +264,7 @@ void variable_seg_size_topk_keys(nvbench::state& state, num_segments_uniform_param, total_num_items, launch.get_stream()); +#endif }); } diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh new file mode 100644 index 00000000000..2b4c7e199dd --- /dev/null +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -0,0 +1,441 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! Cluster-based per-segment top-k agent. +//! +//! Prototype that exercises CUDA thread block clusters as a replacement for +//! the multi-kernel + global histogram pipeline used by cub::DeviceTopK. Each +//! cluster processes exactly one segment. +//! +//! Histogram strategy (Pattern C): +//! 1. Every block lays out `hist[num_buckets]` at the same offset in its own +//! shared memory. Each block accumulates a block-private histogram using +//! block-scope atomicAdd_block (cheap, SMEM-local). +//! 2. After a cluster-wide barrier, every non-leader block walks its +//! histogram and folds its bucket counts into the leader block's `hist` +//! via cluster-scope DSMEM atomics. The leader's `hist` therefore plays +//! a dual role: its own block-private histogram in step 1, then the +//! cluster-merged histogram after the second cluster sync. +//! 3. The leader's thread 0 prefix-scans `hist`, identifies the bucket of +//! the k-th key, and updates the cluster-shared `state`. Every block +//! reads `state.kth_key_bits` from the leader via DSMEM at the start of +//! the next pass. +//! +//! Output cursors live in the same cluster-shared `state` and are reached the +//! same way (cluster-scope DSMEM atomics). + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk_cluster +{ +// ----------------------------------------------------------------------------- +// Tuning policy +// ----------------------------------------------------------------------------- +struct cluster_topk_policy +{ + int cluster_size; + int threads_per_block; + int items_per_thread; + int bits_per_pass; +}; + +// ----------------------------------------------------------------------------- +// Cluster-shared state. Lives in the leader block's shared memory and is +// reached from every block of the cluster through DSMEM. +// ----------------------------------------------------------------------------- +template +struct alignas(16) cluster_topk_state +{ + using key_prefix_t = detail::topk::key_prefix_storage_t; + + OffsetT len; + OutOffsetT k; + key_prefix_t kth_key_bits; + OutOffsetT out_cnt; + OutOffsetT out_back_cnt; +}; + +// ----------------------------------------------------------------------------- +// Cluster top-k agent +// ----------------------------------------------------------------------------- +template +struct agent_batched_topk_cluster +{ + // --------------------------------------------------------------------------- + // Types / constants + // --------------------------------------------------------------------------- + using key_it_t = it_value_t; + using key_t = it_value_t; + + using segment_size_val_t = typename SegmentSizeParameterT::value_type; + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + + using offset_t = ::cuda::std::uint32_t; + using out_offset_t = ::cuda::std::uint32_t; + using state_t = cluster_topk_state; + using key_prefix_t = typename state_t::key_prefix_t; + + static constexpr int cluster_size = ClusterSize; + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int items_per_thread = ItemsPerThread; + static constexpr int bits_per_pass = BitsPerPass; + static constexpr int tile_items = threads_per_block * items_per_thread; + static constexpr int cluster_tile = cluster_size * tile_items; + static constexpr int num_buckets = 1 << bits_per_pass; + + using decomposer_t = detail::identity_decomposer_t; + + // --------------------------------------------------------------------------- + // Shared memory storage + // --------------------------------------------------------------------------- + // The same layout is allocated by every block of the cluster so that any + // block can reach the leader's fields at a known offset via DSMEM. Each + // block populates its own `hist` block-locally; after the first cluster + // sync the non-leader blocks fold their bucket counts into the leader's + // `hist` through DSMEM atomics. `state` is meaningful only in the leader + // block; the other blocks reach it exclusively through the DSMEM mapping. + // `broadcast_kth` is a per-block fan-out slot for `kth_key_bits`. + struct _TempStorage + { + offset_t hist[num_buckets]; + state_t state; + key_prefix_t broadcast_kth; + }; + + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // --------------------------------------------------------------------------- + // Per-thread members + // --------------------------------------------------------------------------- + _TempStorage& temp_storage; + KeyInputItItT d_key_segments_it; + KeyOutputItItT d_key_segments_out_it; + SegmentSizeParameterT segment_sizes; + KParameterT k_param; + SelectDirectionParameterT select_directions; + NumSegmentsParameterT num_segments; + + _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( + TempStorage& temp_storage_, + KeyInputItItT d_key_segments_it_, + KeyOutputItItT d_key_segments_out_it_, + SegmentSizeParameterT segment_sizes_, + KParameterT k_param_, + SelectDirectionParameterT select_directions_, + NumSegmentsParameterT num_segments_) + : temp_storage(temp_storage_.Alias()) + , d_key_segments_it(d_key_segments_it_) + , d_key_segments_out_it(d_key_segments_out_it_) + , segment_sizes(segment_sizes_) + , k_param(k_param_) + , select_directions(select_directions_) + , num_segments(num_segments_) + {} + + // --------------------------------------------------------------------------- + // Main entry point + // --------------------------------------------------------------------------- + // Prototype targets SM 9.0+ only; older architectures are unsupported by the + // dispatch and never reach this kernel. + _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() + { + process_impl(); + } + +private: + _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() + { + for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + { + temp_storage.hist[i] = 0; + } + } + + // Sequential prefix sum on the leader's histogram + locate the bucket + // containing the k-th item. Runs on thread 0 of the leader block. + // num_buckets is small (256 for 8 bits/pass) so this serial pass is fine + // for the prototype. + _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) + { + if (threadIdx.x == 0) + { + const out_offset_t target_k = temp_storage.state.k; + offset_t cumulative = 0; + out_offset_t selected_prev = 0; + offset_t selected_cur = 0; + int selected_bucket = num_buckets - 1; + bool found = false; + for (int b = 0; b < num_buckets; ++b) + { + const offset_t cur = temp_storage.hist[b]; + if (!found && cumulative + cur >= target_k) + { + selected_bucket = b; + selected_prev = static_cast(cumulative); + selected_cur = cur; + found = true; + } + cumulative += cur; + } + temp_storage.state.len = selected_cur; + temp_storage.state.k = target_k - selected_prev; + detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, selected_bucket); + } + } + + // ------------------------------------------------------------------------- + // Per-direction implementation + // ------------------------------------------------------------------------- + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run( + ::cooperative_groups::cluster_group& cluster, + num_segments_val_t segment_id, + unsigned int cluster_rank, + segment_size_val_t segment_size, + out_offset_t k) + { + using extract_bin_op_t = detail::topk::extract_bin_op_t; + using identify_candidates_op_t = detail::topk::identify_candidates_op_t; + + constexpr int total_bits = int{sizeof(key_t)} * 8; + constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + + // Sentinel key used to pad partial tiles. Worst-possible value for the + // requested direction so the sentinel can't land in the selected set. + const key_t pad_key = (SelectDirection == detail::topk::select::max) + ? ::cuda::std::numeric_limits::lowest() + : ::cuda::std::numeric_limits::max(); + + auto block_keys_in = d_key_segments_it[segment_id]; + const auto block_offset_in_cluster = + static_cast(static_cast(cluster_rank) * tile_items); + + // Striped load into per-thread registers. Each thread holds + // `items_per_thread` keys spanning the block's tile within the cluster + // tile. Out-of-segment slots receive sentinel padding. + key_t thread_keys[items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; + } + + // DSMEM pointers into the leader block's shared memory. + offset_t* leader_hist = cluster.map_shared_rank(temp_storage.hist, 0); + state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, 0); + + // Per-block local copy of `kth_key_bits` so each key check hits the + // block's own SMEM rather than DSMEM during the histogram loop. + key_prefix_t kth_key_bits_local = {}; + + for (int pass = 0; pass < num_passes; ++pass) + { + const bool is_first_pass = (pass == 0); + + // Refresh per-block local kth_key_bits from the leader's state. For + // pass 0 the bits are all zero (no filtering yet) so we skip the read. + if (!is_first_pass) + { + if (threadIdx.x == 0) + { + temp_storage.broadcast_kth = leader_state->kth_key_bits; + } + __syncthreads(); + kth_key_bits_local = temp_storage.broadcast_kth; + } + + // Every block (including the leader) starts each pass with a fresh, + // empty `hist`. For pass 0 the leader's initial reset in process_impl() + // covered the leader's slot, but every non-leader block must also + // reset its own. We just always reset here for symmetry. + reset_hist(); + __syncthreads(); + + identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); + extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); + + // Step 1: block-private histogram, cheap block-scope atomic adds. + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const key_t key = thread_keys[j]; + const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); + if (keep) + { + const int bucket = extract_op(key); + atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + } + } + // Cluster-wide barrier: every block must finish its block-scope atomic + // adds before any non-leader block starts DSMEM atomic adds into the + // leader's `hist`. A plain __syncthreads() would only order the local + // block's writes; the leader's `atomicAdd_block`s would still race with + // remote blocks' cluster-scope `atomicAdd`s targeting the same memory. + cluster.sync(); + + // Step 2: non-leader blocks fold their bucket counts into the leader's + // `hist` via cluster-scope DSMEM atomics. The leader's data is already + // present in `hist`; if the leader also merged from itself, its + // contribution would be counted twice. + if (cluster_rank != 0) + { + for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + { + const offset_t v = temp_storage.hist[i]; + if (v != 0) + { + atomicAdd(leader_hist + i, v); + } + } + } + cluster.sync(); + + // Step 3: the leader prefix-scans its (now merged) `hist` and updates + // the cluster-shared `state`. Subsequent reads (next-pass refresh, last + // filter) all observe these writes after the next cluster sync. + if (cluster_rank == 0) + { + leader_identify_kth_bucket(pass); + } + cluster.sync(); + } + + // ----------------------------------------------------------------------- + // Final filter pass: write strictly-selected items to the front of the + // output; back-fill candidates that share the k-th key's prefix bits. + // ----------------------------------------------------------------------- + if (threadIdx.x == 0) + { + temp_storage.broadcast_kth = leader_state->kth_key_bits; + } + __syncthreads(); + kth_key_bits_local = temp_storage.broadcast_kth; + + auto block_keys_out = d_key_segments_out_it[segment_id]; + const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes + + identify_candidates_op_t identify_op(&kth_key_bits_local, num_passes, total_bits, decomposer_t{}); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + if (global_idx >= segment_size) + { + continue; + } + const key_t key = thread_keys[j]; + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[pos] = key; + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + } + } + } + + // Final cluster barrier: hold every block in the cluster until all DSMEM + // atomics into the leader's state are complete. Without this, a fast + // block (e.g. one whose tile is entirely padding) can return while another + // block is still writing to leader-resident memory through DSMEM, which + // surfaces as a "cluster target block not present" exception. + cluster.sync(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() + { + ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); + const unsigned int cluster_rank = cluster.block_rank(); + const auto segment_id = static_cast(blockIdx.x / cluster_size); + + if (segment_id >= num_segments.get_param(0)) + { + return; + } + + const auto segment_size = static_cast(segment_sizes.get_param(segment_id)); + const auto k_requested = static_cast(k_param.get_param(segment_id)); + const auto k = + static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); + + if (k == 0) + { + return; + } + + // The leader block initializes the cluster-shared state. Every block + // (including the leader) will reset its own `hist` at the top of the + // per-pass loop. + if (cluster_rank == 0 && threadIdx.x == 0) + { + temp_storage.state.len = static_cast(segment_size); + temp_storage.state.k = k; + temp_storage.state.kth_key_bits = {}; + temp_storage.state.out_cnt = 0; + temp_storage.state.out_back_cnt = 0; + } + cluster.sync(); + + const bool ok = detail::params::dispatch_discrete( + select_directions, + segment_id, + [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { + constexpr detail::topk::select Direction = decltype(direction_tag)::value; + this->template run(cluster, segment_id, cluster_rank, segment_size, k); + }); + _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); + (void) ok; + } +}; + +} // namespace detail::batched_topk_cluster + +CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh new file mode 100644 index 00000000000..c4c71ac22c0 --- /dev/null +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -0,0 +1,220 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! Single-kernel cluster-based batched top-k dispatch. +//! +//! Prototype that launches a single grid of thread block clusters to compute a +//! segmented top-k. Each cluster processes one segment end-to-end: private +//! histograms are reduced into the leader block via DSMEM atomics, then every +//! block reads the merged histogram back through DSMEM, locally identifies the +//! k-th bucket, and refines its in-register key set across radix passes. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk_cluster +{ +// ----------------------------------------------------------------------------- +// Default tuning +// ----------------------------------------------------------------------------- +// Defaults chosen so a single cluster tile holds the largest segment we expect +// in the existing test and benchmark coverage (8 KiB items @ 4 bytes each). +inline constexpr int default_cluster_size = 4; +inline constexpr int default_threads_per_block = 128; +inline constexpr int default_items_per_thread = 16; +inline constexpr int default_bits_per_pass = 8; + +// ----------------------------------------------------------------------------- +// Kernel entry point +// ----------------------------------------------------------------------------- +template +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KERNEL_ATTRIBUTES void + device_segmented_topk_cluster_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments) +{ + using agent_t = agent_batched_topk_cluster< + ClusterSize, + ThreadsPerBlock, + ItemsPerThread, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + + agent_t agent(temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + + agent.Process(); +} + +// ----------------------------------------------------------------------------- +// Dispatch +// ----------------------------------------------------------------------------- +// The dispatch is intentionally narrow: it currently only supports the +// keys-only path and assumes the maximum segment size fits in a single cluster +// tile (`ClusterSize * ThreadsPerBlock * ItemsPerThread`). Both restrictions +// match the small/decode-style segmented top-k workloads we are prototyping for. +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream = nullptr) +{ + static_assert(ClusterSize >= 1 && ClusterSize <= 8, "ClusterSize must be in [1, 8]"); + + constexpr int cluster_tile = ClusterSize * ThreadsPerBlock * ItemsPerThread; + + static_assert(detail::params::static_max_value_v <= cluster_tile, + "Cluster top-k prototype only supports segments that fit in a single cluster tile."); + + // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. + size_t allocation_sizes[1] = {1}; + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + static_assert(!detail::params::is_per_segment_param_v, + "Number of segments must be resolved on the host."); + + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); + const auto num_seg_unsigned = static_cast(num_seg_val); + if (num_seg_unsigned == 0) + { + return cudaSuccess; + } + + // Cluster launches require compute capability 9.0+. + int sm_version = 0; + if (const auto error = CubDebug(SmVersionUncached(sm_version))) + { + return error; + } + if (sm_version < 900) + { + return cudaErrorNotSupported; + } + + const auto grid_blocks = static_cast(num_seg_unsigned) * static_cast(ClusterSize); + if (grid_blocks > static_cast(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + auto kernel = device_segmented_topk_cluster_kernel< + ClusterSize, + ThreadsPerBlock, + ItemsPerThread, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + // The kernel declares `__cluster_dims__(ClusterSize, 1, 1)`, so the + // cluster dimension is fixed at compile time. Launch via the standard + // triple-chevron syntax; the cluster setup is applied automatically. + if (const auto error = + CubDebug(THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + dim3(static_cast(grid_blocks), 1, 1), + dim3(static_cast(ThreadsPerBlock), 1, 1), + 0, + stream) + .doit( + kernel, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments))) + { + return error; + } + + // Synchronously surface any launch-time errors back to the caller. The + // prototype is keyed on `cudaLaunchKernelEx`, which may return success even + // when the cluster launch fails on the device (e.g. insufficient resources), + // so we explicitly check `cudaDeviceSynchronize` here while bringing up the + // implementation. + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + return CubDebug(detail::DebugSyncStream(stream)); +} + +} // namespace detail::batched_topk_cluster + +CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 3465cebd023..0cd57660c10 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -4,6 +4,7 @@ #include "insert_nested_NVTX_range_guard.h" #include +#include #include #include @@ -65,6 +66,43 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( // %PARAM% TEST_LAUNCH lid 0:1:2 DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_keys, batched_topk_keys); +// ---------------------------------------------------------------------------- +// Cluster-based prototype +// ---------------------------------------------------------------------------- +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys_cluster( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParamT segment_sizes, + KParamT k, + SelectDirectionParamT select_directions, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream = nullptr) +{ + return cub::detail::batched_topk_cluster::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items_guarantee, + stream); +} + +DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_keys_cluster, batched_topk_keys_cluster); + // Total segment size using max_segment_size_list = c2h::enum_type_list; @@ -266,6 +304,104 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } +// ---------------------------------------------------------------------------- +// Cluster-based prototype tests +// ---------------------------------------------------------------------------- +// The cluster prototype currently restricts the maximum segment size to a +// single cluster tile (ClusterSize * ThreadsPerBlock * ItemsPerThread). The +// default policy gives a cluster tile of 4 * 128 * 16 = 8192 items, so we use +// 8 KiB as the upper bound on segment size. Sizes that exceed 3/4 of the +// cluster tile (>= 6 KiB) exercise the path where every block in the cluster +// holds real (non-padding) items; smaller sizes exercise the partially-padded +// path where some blocks are entirely sentinel keys. +using cluster_max_segment_size_list = c2h::enum_type_list; +using cluster_max_num_k_list = c2h::enum_type_list; + +using cluster_key_types = + c2h::type_list; +// clang-format on + +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys cluster prototype works with small fixed-size segments", + "[keys][segmented][topk][device][cluster]", + cluster_key_types, + cluster_max_segment_size_list, + cluster_max_num_k_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using key_t = c2h::get<0, TestType>; + + constexpr segment_size_t static_max_segment_size = c2h::get<1, TestType>::value; + constexpr segment_size_t static_max_k = c2h::get<2, TestType>::value; + + const auto direction = GENERATE_COPY(cub::detail::topk::select::min, cub::detail::topk::select::max); + + constexpr segment_size_t min_segment_size = 1; + constexpr auto max_segment_size = static_max_segment_size; + // Cover both partially-padded clusters (small sizes; some cluster blocks + // are entirely sentinel keys) and the fully-utilized path (sizes large + // enough to put real keys on every block of the cluster). With the default + // cluster tile of 8192 and a 4-block cluster, 7000 reaches into block 3's + // tile so all four blocks see real work. + const segment_size_t segment_size = GENERATE_COPY( + values({min_segment_size, segment_size_t{3}, segment_size_t{7000}, max_segment_size}), + take(2, random(min_segment_size, max_segment_size))); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k}), take(2, random(segment_size_t{1}, max_k))); + + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{17}}), take(2, random(segment_index_t{1}, segment_index_t{128}))); + + CAPTURE(c2h::type_name(), + c2h::type_name(), + c2h::type_name(), + static_max_segment_size, + static_max_k, + segment_size, + k, + num_segments, + direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + const int num_key_seeds = 1; + c2h::gen(C2H_SEED(num_key_seeds), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys_cluster( + d_keys_in, + d_keys_out, + cub::detail::batched_topk::segment_size_uniform<1, max_segment_size>{segment_size}, + cub::detail::batched_topk::k_uniform<1, static_max_k>{k}, + cub::detail::batched_topk::select_direction_uniform{direction}, + cub::detail::batched_topk::num_segments_uniform<>{num_segments}, + cub::detail::batched_topk::total_num_items_guarantee{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segmented][topk][device][float]") { From ef0c358435cb3a56de7bc1b0b98294891708c52d Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Tue, 19 May 2026 18:26:26 -0700 Subject: [PATCH 002/126] Early --- cub/cub/agent/agent_batched_topk_cluster.cuh | 49 ++++++++++++++++++-- 1 file changed, 45 insertions(+), 4 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2b4c7e199dd..44ad1502df6 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -80,6 +80,12 @@ struct alignas(16) cluster_topk_state key_prefix_t kth_key_bits; OutOffsetT out_cnt; OutOffsetT out_back_cnt; + // Set by the leader after `leader_identify_kth_bucket` whenever the + // identified bucket holds exactly `k` items (every candidate is part of + // the top-k). Read by every block of the cluster at the top of the next + // radix pass through DSMEM. Carried in the cluster-shared state so the + // value survives the cluster sync that ends the current pass. + ::cuda::std::uint32_t early_stop; }; // ----------------------------------------------------------------------------- @@ -130,12 +136,14 @@ struct agent_batched_topk_cluster // sync the non-leader blocks fold their bucket counts into the leader's // `hist` through DSMEM atomics. `state` is meaningful only in the leader // block; the other blocks reach it exclusively through the DSMEM mapping. - // `broadcast_kth` is a per-block fan-out slot for `kth_key_bits`. + // `broadcast_kth` / `broadcast_early_stop` are per-block fan-out slots used + // to share the leader-computed values across threads of each block. struct _TempStorage { offset_t hist[num_buckets]; state_t state; key_prefix_t broadcast_kth; + ::cuda::std::uint32_t broadcast_early_stop; }; struct TempStorage : Uninitialized<_TempStorage> @@ -214,8 +222,16 @@ private: } cumulative += cur; } - temp_storage.state.len = selected_cur; - temp_storage.state.k = target_k - selected_prev; + const out_offset_t new_k = target_k - selected_prev; + const offset_t new_len = selected_cur; + temp_storage.state.len = new_len; + temp_storage.state.k = new_k; + // Early-stop opportunity: the bucket holding the k-th key contains + // exactly the remaining `k` items. Every candidate is therefore part of + // the top-k, so subsequent radix passes only redistribute the same + // items across finer buckets without changing the final result. + temp_storage.state.early_stop = + (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, selected_bucket); } } @@ -267,20 +283,38 @@ private: // block's own SMEM rather than DSMEM during the histogram loop. key_prefix_t kth_key_bits_local = {}; + // Tracks the highest pass count that actually executed. Without early + // stop this stays at `num_passes`; with early stop it captures the pass + // at which we broke out so the final filter can construct its identify + // operator at the matching radix level. + int last_pass = num_passes; + for (int pass = 0; pass < num_passes; ++pass) { const bool is_first_pass = (pass == 0); // Refresh per-block local kth_key_bits from the leader's state. For // pass 0 the bits are all zero (no filtering yet) so we skip the read. + // We also pull the leader's `early_stop` flag here so every block has + // an opportunity to bail out before doing any more histogram work. if (!is_first_pass) { if (threadIdx.x == 0) { temp_storage.broadcast_kth = leader_state->kth_key_bits; +#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP + temp_storage.broadcast_early_stop = leader_state->early_stop; +#endif } __syncthreads(); kth_key_bits_local = temp_storage.broadcast_kth; +#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP + if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) + { + last_pass = pass; + break; + } +#endif } // Every block (including the leader) starts each pass with a fresh, @@ -353,7 +387,13 @@ private: auto block_keys_out = d_key_segments_out_it[segment_id]; const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes - identify_candidates_op_t identify_op(&kth_key_bits_local, num_passes, total_bits, decomposer_t{}); + // The pass argument controls how many radix levels of `kth_key_bits` are + // considered significant. After an early-stop break at the start of pass + // `last_pass`, only the first `last_pass` digits of the splitter have + // been set; comparing all bits would treat the (still-zero) trailing + // digits as smaller and erroneously reject candidates that share the + // identified prefix. + identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < items_per_thread; ++j) @@ -421,6 +461,7 @@ private: temp_storage.state.kth_key_bits = {}; temp_storage.state.out_cnt = 0; temp_storage.state.out_back_cnt = 0; + temp_storage.state.early_stop = 0; } cluster.sync(); From f7b835943e2e337b1f2dc17adb293ae6b417bfd0 Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Tue, 19 May 2026 18:39:53 -0700 Subject: [PATCH 003/126] Coop --- cub/cub/agent/agent_batched_topk_cluster.cuh | 69 +++++++++++-------- .../dispatch_batched_topk_cluster.cuh | 12 ++-- .../catch2_test_device_segmented_topk_keys.cu | 10 ++- 3 files changed, 54 insertions(+), 37 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 44ad1502df6..f9df6d2158c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -127,6 +127,16 @@ struct agent_batched_topk_cluster using decomposer_t = detail::identity_decomposer_t; + // --------------------------------------------------------------------------- + // Block-scan used by the leader block to prefix-sum its merged histogram + // --------------------------------------------------------------------------- + // Constraint: every leader thread owns at most one bucket. Threads beyond + // num_buckets contribute zero to the scan, which keeps the prefix-sum loop + // body trivial (1 item per thread, no inner loops). + static_assert(num_buckets <= threads_per_block, + "Cluster top-k requires num_buckets <= threads_per_block (1 bin per leader thread or fewer)"); + using block_scan_t = BlockScan; + // --------------------------------------------------------------------------- // Shared memory storage // --------------------------------------------------------------------------- @@ -144,6 +154,7 @@ struct agent_batched_topk_cluster state_t state; key_prefix_t broadcast_kth; ::cuda::std::uint32_t broadcast_early_stop; + typename block_scan_t::TempStorage scan_storage; }; struct TempStorage : Uninitialized<_TempStorage> @@ -196,43 +207,41 @@ private: } } - // Sequential prefix sum on the leader's histogram + locate the bucket - // containing the k-th item. Runs on thread 0 of the leader block. - // num_buckets is small (256 for 8 bits/pass) so this serial pass is fine - // for the prototype. + // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram + // plus identification of the bucket holding the k-th item. Every thread of + // the leader block participates with one bucket; threads with rank beyond + // num_buckets contribute zero. The single thread that owns the k-th bucket + // writes the per-pass state. The caller must guarantee the leader block + // has finished its DSMEM merge before invoking this. _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) { - if (threadIdx.x == 0) + // Capture `state.k` before the scan: this is the only legal window where + // every thread is guaranteed to read the previous pass's value. The + // owning thread overwrites `state.k` in the if-block below, so any read + // after that point would race with that write. + const out_offset_t target_k = temp_storage.state.k; + + const int bucket = static_cast(threadIdx.x); + const bool owns_bucket = bucket < num_buckets; + const offset_t hist_val = owns_bucket ? temp_storage.hist[bucket] : offset_t{0}; + offset_t prefix = 0; + + block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_val, prefix); + + // Exactly one thread satisfies `prefix < target_k <= prefix + hist_val`. + if (owns_bucket && prefix < target_k && prefix + hist_val >= target_k) { - const out_offset_t target_k = temp_storage.state.k; - offset_t cumulative = 0; - out_offset_t selected_prev = 0; - offset_t selected_cur = 0; - int selected_bucket = num_buckets - 1; - bool found = false; - for (int b = 0; b < num_buckets; ++b) - { - const offset_t cur = temp_storage.hist[b]; - if (!found && cumulative + cur >= target_k) - { - selected_bucket = b; - selected_prev = static_cast(cumulative); - selected_cur = cur; - found = true; - } - cumulative += cur; - } - const out_offset_t new_k = target_k - selected_prev; - const offset_t new_len = selected_cur; - temp_storage.state.len = new_len; - temp_storage.state.k = new_k; + const out_offset_t new_k = target_k - static_cast(prefix); + const offset_t new_len = hist_val; + temp_storage.state.len = new_len; + temp_storage.state.k = new_k; // Early-stop opportunity: the bucket holding the k-th key contains - // exactly the remaining `k` items. Every candidate is therefore part of - // the top-k, so subsequent radix passes only redistribute the same + // exactly the remaining `k` items. Every candidate is therefore part + // of the top-k, so subsequent radix passes only redistribute the same // items across finer buckets without changing the final result. temp_storage.state.early_stop = (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, selected_bucket); + detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); } } diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index c4c71ac22c0..91d5d748315 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -43,10 +43,14 @@ namespace detail::batched_topk_cluster // Default tuning // ----------------------------------------------------------------------------- // Defaults chosen so a single cluster tile holds the largest segment we expect -// in the existing test and benchmark coverage (8 KiB items @ 4 bytes each). -inline constexpr int default_cluster_size = 4; -inline constexpr int default_threads_per_block = 128; -inline constexpr int default_items_per_thread = 16; +// in the existing test and benchmark coverage (8 KiB items). The leader +// block's BlockScan uses 1 bin per thread, so threads_per_block (256) must be +// >= num_buckets (1 << bits_per_pass = 256). Cluster size is 8 — the largest +// portable cluster size — so all eight blocks of the cluster participate in +// the histogram merge. +inline constexpr int default_cluster_size = 8; +inline constexpr int default_threads_per_block = 256; +inline constexpr int default_items_per_thread = 4; inline constexpr int default_bits_per_pass = 8; // ----------------------------------------------------------------------------- diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 0cd57660c10..ba26314dfa7 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -352,10 +352,14 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys cluster prototype works with small fi // Cover both partially-padded clusters (small sizes; some cluster blocks // are entirely sentinel keys) and the fully-utilized path (sizes large // enough to put real keys on every block of the cluster). With the default - // cluster tile of 8192 and a 4-block cluster, 7000 reaches into block 3's - // tile so all four blocks see real work. + // 8-block cluster (per-block tile = 1024), sizes >= 7169 reach block 7; + // 8192 fills every block exactly. const segment_size_t segment_size = GENERATE_COPY( - values({min_segment_size, segment_size_t{3}, segment_size_t{7000}, max_segment_size}), + values({min_segment_size, + segment_size_t{3}, + segment_size_t{4097}, // straddles the block-2 / block-3 boundary in 4-block cfgs + segment_size_t{7500}, // exercises every block in the default 8-block cluster + max_segment_size}), take(2, random(min_segment_size, max_segment_size))); const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); From 55bdf24f2ac04844acdb5a6e9b0fc7a3ecae65d4 Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Tue, 19 May 2026 19:17:46 -0700 Subject: [PATCH 004/126] Larger tile --- .../bench/segmented_topk/variable/keys.cu | 6 +++--- .../dispatch/dispatch_batched_topk_cluster.cuh | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index 602f4b8006e..5e5136d007c 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -275,11 +275,11 @@ using max_segment_size_list = nvbench::enum_type_list< // 1024, 2048, 4096, - 8192 + 8192, + 16384, + 32768 #if 0 // need these, waiting for implementation to catch up , - 16384, - 32768, 65536, 131072, 262144, diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 91d5d748315..db6f9ce7d1f 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -42,15 +42,15 @@ namespace detail::batched_topk_cluster // ----------------------------------------------------------------------------- // Default tuning // ----------------------------------------------------------------------------- -// Defaults chosen so a single cluster tile holds the largest segment we expect -// in the existing test and benchmark coverage (8 KiB items). The leader -// block's BlockScan uses 1 bin per thread, so threads_per_block (256) must be -// >= num_buckets (1 << bits_per_pass = 256). Cluster size is 8 — the largest -// portable cluster size — so all eight blocks of the cluster participate in -// the histogram merge. +// Defaults give a 32 Ki-item cluster tile (cluster_size * threads_per_block * +// items_per_thread = 8 * 256 * 16). The leader block's BlockScan uses 1 bin +// per thread, so threads_per_block (256) must be >= num_buckets +// (1 << bits_per_pass = 256). Cluster size is 8 — the largest portable +// cluster size — so all eight blocks of the cluster participate in the +// histogram merge. inline constexpr int default_cluster_size = 8; inline constexpr int default_threads_per_block = 256; -inline constexpr int default_items_per_thread = 4; +inline constexpr int default_items_per_thread = 16; inline constexpr int default_bits_per_pass = 8; // ----------------------------------------------------------------------------- From ba545e33913129abf11ca8f1d5429977b1be05ea Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 08:23:53 -0700 Subject: [PATCH 005/126] Refactor test --- cub/cub/agent/agent_batched_topk_atomic.cuh | 355 ++++++++++++++++++ .../dispatch/dispatch_batched_topk_atomic.cuh | 186 +++++++++ .../catch2_test_device_segmented_topk_keys.cu | 208 +++------- 3 files changed, 596 insertions(+), 153 deletions(-) create mode 100644 cub/cub/agent/agent_batched_topk_atomic.cuh create mode 100644 cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh new file mode 100644 index 00000000000..248ed10ec0e --- /dev/null +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -0,0 +1,355 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk_atomic +{ +struct atomic_topk_policy +{ + int cluster_size; + int threads_per_block; + int items_per_thread; + int bits_per_pass; +}; + +template +struct alignas(16) atomic_topk_state +{ + using key_prefix_t = detail::topk::key_prefix_storage_t; + + OffsetT len; + OutOffsetT k; + key_prefix_t kth_key_bits; + OutOffsetT out_cnt; + OutOffsetT out_back_cnt; + ::cuda::std::uint32_t early_stop; +}; + +template +struct agent_batched_topk_atomic +{ + using key_it_t = it_value_t; + using key_t = it_value_t; + + using segment_size_val_t = typename SegmentSizeParameterT::value_type; + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + + using offset_t = ::cuda::std::uint32_t; + using out_offset_t = ::cuda::std::uint32_t; + using state_t = atomic_topk_state; + using key_prefix_t = typename state_t::key_prefix_t; + + static constexpr int cluster_size = ClusterSize; + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int items_per_thread = ItemsPerThread; + static constexpr int bits_per_pass = BitsPerPass; + static constexpr int tile_items = threads_per_block * items_per_thread; + static constexpr int cluster_tile = cluster_size * tile_items; + static constexpr int num_buckets = 1 << bits_per_pass; + + using decomposer_t = detail::identity_decomposer_t; + + static_assert(num_buckets <= threads_per_block, + "Atomic top-k requires num_buckets <= threads_per_block (1 bin per leader thread or fewer)"); + using block_scan_t = BlockScan; + + struct _TempStorage + { + offset_t hist[num_buckets]; + state_t state; + key_prefix_t broadcast_kth; + ::cuda::std::uint32_t broadcast_early_stop; + typename block_scan_t::TempStorage scan_storage; + }; + + struct TempStorage : Uninitialized<_TempStorage> + {}; + + _TempStorage& temp_storage; + KeyInputItItT d_key_segments_it; + KeyOutputItItT d_key_segments_out_it; + SegmentSizeParameterT segment_sizes; + KParameterT k_param; + SelectDirectionParameterT select_directions; + NumSegmentsParameterT num_segments; + + _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_atomic( + TempStorage& temp_storage_, + KeyInputItItT d_key_segments_it_, + KeyOutputItItT d_key_segments_out_it_, + SegmentSizeParameterT segment_sizes_, + KParameterT k_param_, + SelectDirectionParameterT select_directions_, + NumSegmentsParameterT num_segments_) + : temp_storage(temp_storage_.Alias()) + , d_key_segments_it(d_key_segments_it_) + , d_key_segments_out_it(d_key_segments_out_it_) + , segment_sizes(segment_sizes_) + , k_param(k_param_) + , select_directions(select_directions_) + , num_segments(num_segments_) + {} + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() + { + process_impl(); + } + +private: + _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() + { + for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + { + temp_storage.hist[i] = 0; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) + { + const out_offset_t target_k = temp_storage.state.k; + + const int bucket = static_cast(threadIdx.x); + const bool owns_bucket = bucket < num_buckets; + const offset_t hist_val = owns_bucket ? temp_storage.hist[bucket] : offset_t{0}; + offset_t prefix = 0; + + block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_val, prefix); + + if (owns_bucket && prefix < target_k && prefix + hist_val >= target_k) + { + const out_offset_t new_k = target_k - static_cast(prefix); + const offset_t new_len = hist_val; + temp_storage.state.len = new_len; + temp_storage.state.k = new_k; + temp_storage.state.early_stop = + (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; + detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run( + ::cooperative_groups::cluster_group& cluster, + num_segments_val_t segment_id, + unsigned int cluster_rank, + segment_size_val_t segment_size, + out_offset_t k) + { + using extract_bin_op_t = detail::topk::extract_bin_op_t; + using identify_candidates_op_t = detail::topk::identify_candidates_op_t; + + constexpr int total_bits = int{sizeof(key_t)} * 8; + constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + + const key_t pad_key = (SelectDirection == detail::topk::select::max) + ? ::cuda::std::numeric_limits::lowest() + : ::cuda::std::numeric_limits::max(); + + auto block_keys_in = d_key_segments_it[segment_id]; + const auto block_offset_in_cluster = + static_cast(static_cast(cluster_rank) * tile_items); + + key_t thread_keys[items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; + } + + offset_t* leader_hist = cluster.map_shared_rank(temp_storage.hist, 0); + state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, 0); + + key_prefix_t kth_key_bits_local = {}; + + int last_pass = num_passes; + + for (int pass = 0; pass < num_passes; ++pass) + { + const bool is_first_pass = (pass == 0); + + if (!is_first_pass) + { + if (threadIdx.x == 0) + { + temp_storage.broadcast_kth = leader_state->kth_key_bits; +#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP + temp_storage.broadcast_early_stop = leader_state->early_stop; +#endif + } + __syncthreads(); + kth_key_bits_local = temp_storage.broadcast_kth; +#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP + if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) + { + last_pass = pass; + break; + } +#endif + } + + reset_hist(); + __syncthreads(); + + identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); + extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const key_t key = thread_keys[j]; + const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); + if (keep) + { + const int bucket = extract_op(key); + atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + } + } + cluster.sync(); + + if (cluster_rank != 0) + { + for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + { + const offset_t v = temp_storage.hist[i]; + if (v != 0) + { + atomicAdd(leader_hist + i, v); + } + } + } + cluster.sync(); + + if (cluster_rank == 0) + { + leader_identify_kth_bucket(pass); + } + cluster.sync(); + } + + if (threadIdx.x == 0) + { + temp_storage.broadcast_kth = leader_state->kth_key_bits; + } + __syncthreads(); + kth_key_bits_local = temp_storage.broadcast_kth; + + auto block_keys_out = d_key_segments_out_it[segment_id]; + const out_offset_t num_kth = leader_state->k; + + identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + if (global_idx >= segment_size) + { + continue; + } + const key_t key = thread_keys[j]; + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[pos] = key; + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + } + } + } + + cluster.sync(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() + { + ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); + const unsigned int cluster_rank = cluster.block_rank(); + const auto segment_id = static_cast(blockIdx.x / cluster_size); + + if (segment_id >= num_segments.get_param(0)) + { + return; + } + + const auto segment_size = static_cast(segment_sizes.get_param(segment_id)); + const auto k_requested = static_cast(k_param.get_param(segment_id)); + const auto k = + static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); + + if (k == 0) + { + return; + } + + if (cluster_rank == 0 && threadIdx.x == 0) + { + temp_storage.state.len = static_cast(segment_size); + temp_storage.state.k = k; + temp_storage.state.kth_key_bits = {}; + temp_storage.state.out_cnt = 0; + temp_storage.state.out_back_cnt = 0; + temp_storage.state.early_stop = 0; + } + cluster.sync(); + + const bool ok = detail::params::dispatch_discrete( + select_directions, + segment_id, + [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { + constexpr detail::topk::select Direction = decltype(direction_tag)::value; + this->template run(cluster, segment_id, cluster_rank, segment_size, k); + }); + _CCCL_ASSERT(ok, "Unsupported select direction for atomic top-k"); + (void) ok; + } +}; + +} // namespace detail::batched_topk_atomic + +CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh new file mode 100644 index 00000000000..ecd46ac8105 --- /dev/null +++ b/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh @@ -0,0 +1,186 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk_atomic +{ +inline constexpr int default_cluster_size = 8; +inline constexpr int default_threads_per_block = 256; +inline constexpr int default_items_per_thread = 16; +inline constexpr int default_bits_per_pass = 8; + +template +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KERNEL_ATTRIBUTES void + device_segmented_topk_atomic_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments) +{ + using agent_t = agent_batched_topk_atomic< + ClusterSize, + ThreadsPerBlock, + ItemsPerThread, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + + agent_t agent(temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + + agent.Process(); +} + +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream = nullptr) +{ + static_assert(ClusterSize >= 1 && ClusterSize <= 8, "ClusterSize must be in [1, 8]"); + + constexpr int cluster_tile = ClusterSize * ThreadsPerBlock * ItemsPerThread; + + static_assert(detail::params::static_max_value_v <= cluster_tile, + "Atomic top-k prototype only supports segments that fit in a single cluster tile."); + + size_t allocation_sizes[1] = {1}; + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + static_assert(!detail::params::is_per_segment_param_v, + "Number of segments must be resolved on the host."); + + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); + const auto num_seg_unsigned = static_cast(num_seg_val); + if (num_seg_unsigned == 0) + { + return cudaSuccess; + } + + int sm_version = 0; + if (const auto error = CubDebug(SmVersionUncached(sm_version))) + { + return error; + } + if (sm_version < 900) + { + return cudaErrorNotSupported; + } + + const auto grid_blocks = static_cast(num_seg_unsigned) * static_cast(ClusterSize); + if (grid_blocks > static_cast(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + auto kernel = device_segmented_topk_atomic_kernel< + ClusterSize, + ThreadsPerBlock, + ItemsPerThread, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + if (const auto error = + CubDebug(THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + dim3(static_cast(grid_blocks), 1, 1), + dim3(static_cast(ThreadsPerBlock), 1, 1), + 0, + stream) + .doit( + kernel, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + return CubDebug(detail::DebugSyncStream(stream)); +} + +} // namespace detail::batched_topk_atomic + +CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index ba26314dfa7..8cbae999638 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -4,6 +4,7 @@ #include "insert_nested_NVTX_range_guard.h" #include +#include #include #include @@ -28,47 +29,15 @@ struct is_minus_zero } }; -template -CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( - void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - SegmentSizeParamT segment_sizes, - KParamT k, - SelectDirectionParamT select_directions, - NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr) +enum class topk_backend { - auto values_it = static_cast(nullptr); - return cub::detail::batched_topk::dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - values_it, - values_it, - segment_sizes, - k, - select_directions, - num_segments, - total_num_items_guarantee, - stream); -} + baseline, + cluster, + atomic, +}; -// %PARAM% TEST_LAUNCH lid 0:1:2 -DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_keys, batched_topk_keys); +inline constexpr topk_backend selected_backend = topk_backend::atomic; -// ---------------------------------------------------------------------------- -// Cluster-based prototype -// ---------------------------------------------------------------------------- template -CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys_cluster( +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, @@ -88,20 +57,55 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys_cluster( TotalNumItemsGuaranteeT total_num_items_guarantee, cudaStream_t stream = nullptr) { - return cub::detail::batched_topk_cluster::dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - total_num_items_guarantee, - stream); + if constexpr (selected_backend == topk_backend::cluster) + { + return cub::detail::batched_topk_cluster::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items_guarantee, + stream); + } + else if constexpr (selected_backend == topk_backend::atomic) + { + return cub::detail::batched_topk_atomic::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items_guarantee, + stream); + } + else + { + auto values_it = static_cast(nullptr); + return cub::detail::batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + values_it, + values_it, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items_guarantee, + stream); + } } -DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_keys_cluster, batched_topk_keys_cluster); +// %PARAM% TEST_LAUNCH lid 0:1:2 +DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_keys, batched_topk_keys); // Total segment size using max_segment_size_list = c2h::enum_type_list; @@ -304,108 +308,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } -// ---------------------------------------------------------------------------- -// Cluster-based prototype tests -// ---------------------------------------------------------------------------- -// The cluster prototype currently restricts the maximum segment size to a -// single cluster tile (ClusterSize * ThreadsPerBlock * ItemsPerThread). The -// default policy gives a cluster tile of 4 * 128 * 16 = 8192 items, so we use -// 8 KiB as the upper bound on segment size. Sizes that exceed 3/4 of the -// cluster tile (>= 6 KiB) exercise the path where every block in the cluster -// holds real (non-padding) items; smaller sizes exercise the partially-padded -// path where some blocks are entirely sentinel keys. -using cluster_max_segment_size_list = c2h::enum_type_list; -using cluster_max_num_k_list = c2h::enum_type_list; - -using cluster_key_types = - c2h::type_list; -// clang-format on - -C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys cluster prototype works with small fixed-size segments", - "[keys][segmented][topk][device][cluster]", - cluster_key_types, - cluster_max_segment_size_list, - cluster_max_num_k_list) -{ - using segment_size_t = cuda::std::int64_t; - using segment_index_t = cuda::std::int64_t; - - using key_t = c2h::get<0, TestType>; - - constexpr segment_size_t static_max_segment_size = c2h::get<1, TestType>::value; - constexpr segment_size_t static_max_k = c2h::get<2, TestType>::value; - - const auto direction = GENERATE_COPY(cub::detail::topk::select::min, cub::detail::topk::select::max); - - constexpr segment_size_t min_segment_size = 1; - constexpr auto max_segment_size = static_max_segment_size; - // Cover both partially-padded clusters (small sizes; some cluster blocks - // are entirely sentinel keys) and the fully-utilized path (sizes large - // enough to put real keys on every block of the cluster). With the default - // 8-block cluster (per-block tile = 1024), sizes >= 7169 reach block 7; - // 8192 fills every block exactly. - const segment_size_t segment_size = GENERATE_COPY( - values({min_segment_size, - segment_size_t{3}, - segment_size_t{4097}, // straddles the block-2 / block-3 boundary in 4-block cfgs - segment_size_t{7500}, // exercises every block in the default 8-block cluster - max_segment_size}), - take(2, random(min_segment_size, max_segment_size))); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k}), take(2, random(segment_size_t{1}, max_k))); - - const segment_index_t num_segments = GENERATE_COPY( - values({segment_index_t{1}, segment_index_t{17}}), take(2, random(segment_index_t{1}, segment_index_t{128}))); - - CAPTURE(c2h::type_name(), - c2h::type_name(), - c2h::type_name(), - static_max_segment_size, - static_max_k, - segment_size, - k, - num_segments, - direction); - - c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); - c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); - const int num_key_seeds = 1; - c2h::gen(C2H_SEED(num_key_seeds), keys_in_buffer); - auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); - auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); - auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); - auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); - - c2h::device_vector expected_keys(keys_in_buffer); - - batched_topk_keys_cluster( - d_keys_in, - d_keys_out, - cub::detail::batched_topk::segment_size_uniform<1, max_segment_size>{segment_size}, - cub::detail::batched_topk::k_uniform<1, static_max_k>{k}, - cub::detail::batched_topk::select_direction_uniform{direction}, - cub::detail::batched_topk::num_segments_uniform<>{num_segments}, - cub::detail::batched_topk::total_num_items_guarantee{num_segments * segment_size}); - - fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); - compact_sorted_keys_to_topk(expected_keys, segment_size, k); - - fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); - - REQUIRE(expected_keys == keys_out_buffer); -} - // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segmented][topk][device][float]") { From fd2f4f32848947d3ecff926bc503c000a5b2b5ac Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 10:40:25 -0700 Subject: [PATCH 006/126] Working on atomic --- .clangd | 1 + cub/cub/agent/agent_batched_topk_atomic.cuh | 288 ++++-------------- .../dispatch/dispatch_batched_topk_atomic.cuh | 66 +--- .../catch2_test_device_segmented_topk_keys.cu | 14 +- 4 files changed, 83 insertions(+), 286 deletions(-) diff --git a/.clangd b/.clangd index 2200fd689ba..079f1b7b8b8 100644 --- a/.clangd +++ b/.clangd @@ -73,6 +73,7 @@ CompileFlags: Add: - -x - cuda + - --offload-arch=sm_90 - -Wno-unknown-cuda-version - -Wno-pragma-system-header-outside-header - --no-cuda-version-check diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh index 248ed10ec0e..8d8b2beadbd 100644 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -35,27 +36,9 @@ struct atomic_topk_policy { int cluster_size; int threads_per_block; - int items_per_thread; - int bits_per_pass; }; -template -struct alignas(16) atomic_topk_state -{ - using key_prefix_t = detail::topk::key_prefix_storage_t; - - OffsetT len; - OutOffsetT k; - key_prefix_t kth_key_bits; - OutOffsetT out_cnt; - OutOffsetT out_back_cnt; - ::cuda::std::uint32_t early_stop; -}; - -template ; - using key_prefix_t = typename state_t::key_prefix_t; - static constexpr int cluster_size = ClusterSize; - static constexpr int threads_per_block = ThreadsPerBlock; - static constexpr int items_per_thread = ItemsPerThread; - static constexpr int bits_per_pass = BitsPerPass; - static constexpr int tile_items = threads_per_block * items_per_thread; - static constexpr int cluster_tile = cluster_size * tile_items; - static constexpr int num_buckets = 1 << bits_per_pass; + // TODO(gevtushenko): direction + using twiddle_t = cub::RadixSortTwiddle; + using bit_ordered_t = typename cub::Traits::UnsignedBits; - using decomposer_t = detail::identity_decomposer_t; + static constexpr auto threads_per_block = ThreadsPerBlock; + static constexpr auto max_k = params::static_max_value_v; - static_assert(num_buckets <= threads_per_block, - "Atomic top-k requires num_buckets <= threads_per_block (1 bin per leader thread or fewer)"); + using decomposer_t = detail::identity_decomposer_t; using block_scan_t = BlockScan; struct _TempStorage { - offset_t hist[num_buckets]; - state_t state; - key_prefix_t broadcast_kth; - ::cuda::std::uint32_t broadcast_early_stop; - typename block_scan_t::TempStorage scan_storage; + bit_ordered_t top_k[max_k]; }; struct TempStorage : Uninitialized<_TempStorage> @@ -128,228 +101,81 @@ struct agent_batched_topk_atomic _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() { - process_impl(); - } + const auto segment_id = static_cast(blockIdx.x); + const auto direction = select_directions.get_param(segment_id); + const auto segment_size = static_cast(segment_sizes.get_param(segment_id)); + const auto k_requested = static_cast(k_param.get_param(segment_id)); + const auto k = + static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); -private: - _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() - { - for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + if (k == 0) { - temp_storage.hist[i] = 0; + return; } - } - - _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) - { - const out_offset_t target_k = temp_storage.state.k; - const int bucket = static_cast(threadIdx.x); - const bool owns_bucket = bucket < num_buckets; - const offset_t hist_val = owns_bucket ? temp_storage.hist[bucket] : offset_t{0}; - offset_t prefix = 0; + // TODO(gevtushenk); direction + (void) direction; + bit_ordered_t sentinel{0}; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_val, prefix); + init_smem(k, sentinel); + deposit(segment_id, segment_size, k, sentinel); + merge(segment_id, k); + store(segment_id, k); + } - if (owns_bucket && prefix < target_k && prefix + hist_val >= target_k) +private: + __device__ void init_smem(unsigned k, bit_ordered_t sentinel) + { + for (unsigned i = threadIdx.x; i < k; i += threads_per_block) { - const out_offset_t new_k = target_k - static_cast(prefix); - const offset_t new_len = hist_val; - temp_storage.state.len = new_len; - temp_storage.state.k = new_k; - temp_storage.state.early_stop = - (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + temp_storage.top_k[i] = sentinel; } + __syncthreads(); } - template - _CCCL_DEVICE _CCCL_FORCEINLINE void run( - ::cooperative_groups::cluster_group& cluster, - num_segments_val_t segment_id, - unsigned int cluster_rank, - segment_size_val_t segment_size, - out_offset_t k) + __device__ void deposit(unsigned segment_id, unsigned segment_size, unsigned k, bit_ordered_t sentinel) { - using extract_bin_op_t = detail::topk::extract_bin_op_t; - using identify_candidates_op_t = detail::topk::identify_candidates_op_t; - - constexpr int total_bits = int{sizeof(key_t)} * 8; - constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); - - const key_t pad_key = (SelectDirection == detail::topk::select::max) - ? ::cuda::std::numeric_limits::lowest() - : ::cuda::std::numeric_limits::max(); - + auto min = sentinel; auto block_keys_in = d_key_segments_it[segment_id]; - const auto block_offset_in_cluster = - static_cast(static_cast(cluster_rank) * tile_items); - - key_t thread_keys[items_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) + for (unsigned i = threadIdx.x; i < segment_size; i += threads_per_block) { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; - } - - offset_t* leader_hist = cluster.map_shared_rank(temp_storage.hist, 0); - state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, 0); - - key_prefix_t kth_key_bits_local = {}; - - int last_pass = num_passes; - - for (int pass = 0; pass < num_passes; ++pass) - { - const bool is_first_pass = (pass == 0); - - if (!is_first_pass) - { - if (threadIdx.x == 0) - { - temp_storage.broadcast_kth = leader_state->kth_key_bits; -#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - temp_storage.broadcast_early_stop = leader_state->early_stop; -#endif - } - __syncthreads(); - kth_key_bits_local = temp_storage.broadcast_kth; -#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) - { - last_pass = pass; - break; - } -#endif - } - - reset_hist(); - __syncthreads(); - - identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); - extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); - - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) - { - const key_t key = thread_keys[j]; - const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); - if (keep) - { - const int bucket = extract_op(key); - atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); - } - } - cluster.sync(); - - if (cluster_rank != 0) - { - for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) - { - const offset_t v = temp_storage.hist[i]; - if (v != 0) - { - atomicAdd(leader_hist + i, v); - } - } - } - cluster.sync(); - - if (cluster_rank == 0) - { - leader_identify_kth_bucket(pass); - } - cluster.sync(); - } - - if (threadIdx.x == 0) - { - temp_storage.broadcast_kth = leader_state->kth_key_bits; + auto bit_ordered_key = twiddle_t::In(reinterpret_cast(block_keys_in[i])); + min = top_k(bit_ordered_key, temp_storage.top_k, k); } __syncthreads(); - kth_key_bits_local = temp_storage.broadcast_kth; - - auto block_keys_out = d_key_segments_out_it[segment_id]; - const out_offset_t num_kth = leader_state->k; - - identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); - - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) - { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - if (global_idx >= segment_size) - { - continue; - } - const key_t key = thread_keys[j]; - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[pos] = key; - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); - if (back_pos < num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; - } - } - } - - cluster.sync(); } - _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() + __device__ void merge(unsigned segment_id, unsigned k) { - ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); - const unsigned int cluster_rank = cluster.block_rank(); - const auto segment_id = static_cast(blockIdx.x / cluster_size); - - if (segment_id >= num_segments.get_param(0)) - { - return; - } + // TODO(gevtushenko): cluster + } - const auto segment_size = static_cast(segment_sizes.get_param(segment_id)); - const auto k_requested = static_cast(k_param.get_param(segment_id)); - const auto k = - static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); + __device__ void store(unsigned segment_id, unsigned k) + { + auto block_keys_out = d_key_segments_out_it[segment_id]; - if (k == 0) + for (unsigned i = threadIdx.x; i < k; i += threads_per_block) { - return; + auto key = twiddle_t::Out(temp_storage.top_k[i]); + block_keys_out[i] = reinterpret_cast(key); } + } - if (cluster_rank == 0 && threadIdx.x == 0) + template + __device__ bit_ordered_t top_k(bit_ordered_t key, bit_ordered_t* top, unsigned k) + { + for (unsigned j = 0; j < k; ++j) { - temp_storage.state.len = static_cast(segment_size); - temp_storage.state.k = k; - temp_storage.state.kth_key_bits = {}; - temp_storage.state.out_cnt = 0; - temp_storage.state.out_back_cnt = 0; - temp_storage.state.early_stop = 0; + cuda::atomic_ref shared_top_ref(top[j]); + const bit_ordered_t old_max_j = shared_top_ref.fetch_max(key, cuda::memory_order_relaxed); + if (old_max_j < key) + { + key = old_max_j; + } } - cluster.sync(); - - const bool ok = detail::params::dispatch_discrete( - select_directions, - segment_id, - [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { - constexpr detail::topk::select Direction = decltype(direction_tag)::value; - this->template run(cluster, segment_id, cluster_rank, segment_size, k); - }); - _CCCL_ASSERT(ok, "Unsupported select direction for atomic top-k"); - (void) ok; + return key; } }; - } // namespace detail::batched_topk_atomic CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh index ecd46ac8105..93d67f360a5 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh @@ -19,34 +19,27 @@ #include #include +#include + #include #include #include -#include - CUB_NAMESPACE_BEGIN namespace detail::batched_topk_atomic { -inline constexpr int default_cluster_size = 8; inline constexpr int default_threads_per_block = 256; -inline constexpr int default_items_per_thread = 16; -inline constexpr int default_bits_per_pass = 8; -template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KERNEL_ATTRIBUTES void - device_segmented_topk_atomic_kernel( +__launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_atomic_kernel( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, SegmentSizeParameterT segment_sizes, @@ -55,10 +48,7 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KER NumSegmentsParameterT num_segments) { using agent_t = agent_batched_topk_atomic< - ClusterSize, ThreadsPerBlock, - ItemsPerThread, - BitsPerPass, KeyInputItItT, KeyOutputItItT, SegmentSizeParameterT, @@ -68,15 +58,13 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KER __shared__ typename agent_t::TempStorage temp_storage; - agent_t agent(temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + agent_t agent( + temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); agent.Process(); } -template = 1 && ClusterSize <= 8, "ClusterSize must be in [1, 8]"); - - constexpr int cluster_tile = ClusterSize * ThreadsPerBlock * ItemsPerThread; - - static_assert(detail::params::static_max_value_v <= cluster_tile, - "Atomic top-k prototype only supports segments that fit in a single cluster tile."); - size_t allocation_sizes[1] = {1}; void* allocations[1] = {}; if (const auto error = @@ -119,8 +100,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( static_assert(!detail::params::is_per_segment_param_v, "Number of segments must be resolved on the host."); - using num_segments_val_t = typename NumSegmentsParameterT::value_type; - const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); const auto num_seg_unsigned = static_cast(num_seg_val); if (num_seg_unsigned == 0) { @@ -137,17 +118,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaErrorNotSupported; } - const auto grid_blocks = static_cast(num_seg_unsigned) * static_cast(ClusterSize); - if (grid_blocks > static_cast(::cuda::std::numeric_limits::max())) - { - return cudaErrorInvalidValue; - } + const auto grid_blocks = static_cast(num_seg_unsigned); auto kernel = device_segmented_topk_atomic_kernel< - ClusterSize, ThreadsPerBlock, - ItemsPerThread, - BitsPerPass, KeyInputItItT, KeyOutputItItT, SegmentSizeParameterT, @@ -155,20 +129,11 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionParameterT, NumSegmentsParameterT>; - if (const auto error = - CubDebug(THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - dim3(static_cast(grid_blocks), 1, 1), - dim3(static_cast(ThreadsPerBlock), 1, 1), - 0, - stream) - .doit( - kernel, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments))) + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + static_cast(grid_blocks), static_cast(ThreadsPerBlock), 0, stream) + .doit( + kernel, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments))) { return error; } @@ -180,7 +145,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return CubDebug(detail::DebugSyncStream(stream)); } - } // namespace detail::batched_topk_atomic CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 8cbae999638..c075b6da03e 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -113,6 +113,7 @@ using max_segment_size_list = c2h::enum_type_list; // Segment size: static, uniform using max_num_k_list = c2h::enum_type_list; +#if 0 using key_types = c2h::type_list; +#else +using key_types = c2h::type_list; +#endif // clang-format on C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", @@ -143,7 +147,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", constexpr segment_size_t static_max_k = c2h::get<2, TestType>::value; // Test both directions (as runtime value) - const auto direction = GENERATE_COPY(cub::detail::topk::select::min, cub::detail::topk::select::max); + const auto direction = cub::detail::topk::select::max; // GENERATE_COPY(cub::detail::topk::select::min, + // cub::detail::topk::select::max); // Generate segment size constexpr segment_size_t min_segment_size = 1; @@ -224,7 +229,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment constexpr segment_size_t static_max_k = c2h::get<2, TestType>::value; // Test both directions (as runtime value) - const auto direction = GENERATE_COPY(cub::detail::topk::select::min, cub::detail::topk::select::max); + const auto direction = cub::detail::topk::select::max; // GENERATE_COPY(cub::detail::topk::select::min, + // cub::detail::topk::select::max); constexpr segment_size_t min_items = 1; constexpr segment_size_t max_items = 1'000'000; @@ -317,7 +323,7 @@ C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segment constexpr cuda::std::size_t max_segment_size = 64; // Input: one segment containing -0.0f and +0.0f; top-5 min should include both zeros. - c2h::device_vector d_keys_in{3.0f, -0.0f, 1.0f, 2.0f, 0.0f, -1.0f, 4.0f, 5.0f}; + c2h::device_vector d_keys_in{-3.0f, -0.0f, -1.0f, -2.0f, 0.0f, -1.0f, -4.0f, -5.0f}; c2h::device_vector d_keys_out(k, thrust::no_init); auto d_keys_in_it = @@ -330,7 +336,7 @@ C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segment d_keys_out_it, cub::detail::batched_topk::segment_size_uniform<1, max_segment_size>{segment_size}, cub::detail::batched_topk::k_uniform<1, static_cast(k)>{k}, - cub::detail::batched_topk::select_direction_uniform{cub::detail::topk::select::min}, + cub::detail::batched_topk::select_direction_uniform{cub::detail::topk::select::max}, cub::detail::batched_topk::num_segments_uniform<>{num_segments}, cub::detail::batched_topk::total_num_items_guarantee{num_segments * segment_size}); From 01c68385106b10aa50897b299ab144dac296686c Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 11:30:08 -0700 Subject: [PATCH 007/126] Bench --- .../bench/segmented_topk/fixed/keys.cu | 166 ++++++++++++++-- .../bench/segmented_topk/variable/keys.cu | 186 ++++++++++++++---- 2 files changed, 291 insertions(+), 61 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index 1465c9b8c44..78dec4bd1d2 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -2,9 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include +#include #include +#include +#include +#include +#include +#include #include +#include +#include + +#include +#include #include @@ -12,6 +23,16 @@ // %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32 // %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1 +enum class topk_backend +{ + baseline, + cluster, + atomic, + device, +}; + +inline constexpr topk_backend selected_backend = topk_backend::baseline; + #if !TUNE_BASE struct tuned_policy_selector { @@ -39,6 +60,124 @@ struct tuned_policy_selector }; #endif // !TUNE_BASE +template +CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + SegmentSizeParamT segment_sizes, + KParamT k, + SelectDirectionParamT select_directions, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + const HostSegSizeT* h_segment_sizes, + cudaStream_t stream = nullptr) +{ + if constexpr (selected_backend == topk_backend::cluster) + { + (void) h_segment_sizes; + return cub::detail::batched_topk_cluster::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream); + } + else if constexpr (selected_backend == topk_backend::atomic) + { + (void) h_segment_sizes; + return cub::detail::batched_topk_atomic::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream); + } + else if constexpr (selected_backend == topk_backend::device) + { + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_segs = num_segments.get_param(num_segments_val_t{0}); + + auto requirements = cuda::execution::require( + cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); + + if (d_temp_storage == nullptr) + { + const auto max_size = *std::max_element(h_segment_sizes, h_segment_sizes + num_segs); + const auto k_value = k.get_param(num_segments_val_t{0}); + return cub::DeviceTopK::MaxKeys( + nullptr, + temp_storage_bytes, + d_keys_in[num_segments_val_t{0}], + d_keys_out[num_segments_val_t{0}], + static_cast(max_size), + static_cast(k_value), + cuda::std::execution::env{requirements}); + } + + cuda::stream_ref stream_ref{stream}; + auto env = cuda::std::execution::env{stream_ref, requirements}; + for (num_segments_val_t i = 0; i < num_segs; ++i) + { + const auto k_value = k.get_param(i); + const auto seg_size = h_segment_sizes[i]; + if (const auto err = cub::DeviceTopK::MaxKeys( + d_temp_storage, + temp_storage_bytes, + d_keys_in[i], + d_keys_out[i], + static_cast(seg_size), + static_cast(k_value), + env); + err != cudaSuccess) + { + return err; + } + } + return cudaSuccess; + } + else + { + (void) h_segment_sizes; + return cub::detail::batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream +#if !TUNE_BASE + , + tuned_policy_selector{} +#endif // !TUNE_BASE + ); + } +} + template void fixed_seg_size_topk_keys( nvbench::state& state, @@ -97,50 +236,41 @@ void fixed_seg_size_topk_keys( state.add_global_memory_reads(elements, "InputKeys"); state.add_global_memory_writes(selected_elements * num_segments, "OutputKeys"); + // Host copy of segment sizes — all entries equal MaxSegmentSize for fixed-size segments. + std::vector h_segment_sizes(num_segments, static_cast(MaxSegmentSize)); + // allocate temporary storage size_t temp_size; - cub::detail::batched_topk::dispatch( + batched_topk_keys( nullptr, temp_size, d_keys_in, d_keys_out, - static_cast(nullptr), - static_cast(nullptr), segment_sizes, k, select_directions, num_segments_uniform_t{static_cast<::cuda::std::int64_t>(num_segments)}, total_num_items, - nullptr -#if !TUNE_BASE - , - tuned_policy_selector{} -#endif // !TUNE_BASE - ); + h_segment_sizes.data(), + nullptr); thrust::device_vector temp(temp_size, thrust::no_init); auto* temp_storage = thrust::raw_pointer_cast(temp.data()); // run the algorithm state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - cub::detail::batched_topk::dispatch( + batched_topk_keys( temp_storage, temp_size, d_keys_in, d_keys_out, - static_cast(nullptr), - static_cast(nullptr), segment_sizes, k, select_directions, num_segments_uniform_t{static_cast<::cuda::std::int64_t>(num_segments)}, total_num_items, - launch.get_stream() -#if !TUNE_BASE - , - tuned_policy_selector{} -#endif // !TUNE_BASE - ); + h_segment_sizes.data(), + launch.get_stream()); }); } diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index 5e5136d007c..5450fbcb9a0 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -2,15 +2,141 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include +#include #include +#include #include -// Compile-time switch: define BENCH_CLUSTER_TOPK=1 to drive the prototype -// cluster-based dispatch instead of the small-segment baseline. We default to -// the baseline so existing benchmark scripts keep their behavior. -#ifndef BENCH_CLUSTER_TOPK -# define BENCH_CLUSTER_TOPK 0 -#endif +#include +#include +#include +#include +#include + +#include + +enum class topk_backend +{ + baseline, + cluster, + atomic, + device, +}; + +inline constexpr topk_backend selected_backend = topk_backend::baseline; + +template +CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + SegmentSizeParamT segment_sizes, + KParamT k, + SelectDirectionParamT select_directions, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + const HostSegSizeT* h_segment_sizes, + cudaStream_t stream = nullptr) +{ + if constexpr (selected_backend == topk_backend::cluster) + { + (void) h_segment_sizes; + return cub::detail::batched_topk_cluster::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream); + } + else if constexpr (selected_backend == topk_backend::atomic) + { + (void) h_segment_sizes; + return cub::detail::batched_topk_atomic::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream); + } + else if constexpr (selected_backend == topk_backend::device) + { + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_segs = num_segments.get_param(num_segments_val_t{0}); + + auto requirements = cuda::execution::require( + cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); + + if (d_temp_storage == nullptr) + { + const auto max_size = *std::max_element(h_segment_sizes, h_segment_sizes + num_segs); + const auto k_value = k.get_param(num_segments_val_t{0}); + return cub::DeviceTopK::MaxKeys( + nullptr, + temp_storage_bytes, + d_keys_in[num_segments_val_t{0}], + d_keys_out[num_segments_val_t{0}], + static_cast(max_size), + static_cast(k_value), + cuda::std::execution::env{requirements}); + } + + cuda::stream_ref stream_ref{stream}; + auto env = cuda::std::execution::env{stream_ref, requirements}; + for (num_segments_val_t i = 0; i < num_segs; ++i) + { + const auto k_value = k.get_param(i); + const auto seg_size = h_segment_sizes[i]; + if (const auto err = cub::DeviceTopK::MaxKeys( + d_temp_storage, + temp_storage_bytes, + d_keys_in[i], + d_keys_out[i], + static_cast(seg_size), + static_cast(k_value), + env); + err != cudaSuccess) + { + return err; + } + } + return cudaSuccess; + } + else + { + (void) h_segment_sizes; + return cub::detail::batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + select_directions, + num_segments, + total_num_items, + stream); + } +} #include #include @@ -205,66 +331,40 @@ void variable_seg_size_topk_keys(nvbench::state& state, state.add_global_memory_reads(input_elements, "InputKeys"); state.add_global_memory_writes(output_elements, "OutputKeys"); + // Host copy of segment sizes — consumed by the per-segment device backend. + std::vector h_segment_sizes(static_cast(num_segments)); + thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), h_segment_sizes.begin()); + size_t temp_size{}; -#if BENCH_CLUSTER_TOPK - cub::detail::batched_topk_cluster::dispatch( - nullptr, - temp_size, - d_keys_in, - d_keys_out, - segment_sizes_param, - k_param, - select_directions, - num_segments_uniform_param, - total_num_items, - nullptr); -#else - cub::detail::batched_topk::dispatch( + batched_topk_keys( nullptr, temp_size, d_keys_in, d_keys_out, - static_cast(nullptr), - static_cast(nullptr), segment_sizes_param, k_param, select_directions, num_segments_uniform_param, total_num_items, + h_segment_sizes.data(), nullptr); -#endif thrust::device_vector temp(temp_size, thrust::no_init); auto* temp_storage = thrust::raw_pointer_cast(temp.data()); state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { -#if BENCH_CLUSTER_TOPK - cub::detail::batched_topk_cluster::dispatch( - temp_storage, - temp_size, - d_keys_in, - d_keys_out, - segment_sizes_param, - k_param, - select_directions, - num_segments_uniform_param, - total_num_items, - launch.get_stream()); -#else - cub::detail::batched_topk::dispatch( + batched_topk_keys( temp_storage, temp_size, d_keys_in, d_keys_out, - static_cast(nullptr), - static_cast(nullptr), segment_sizes_param, k_param, select_directions, num_segments_uniform_param, total_num_items, + h_segment_sizes.data(), launch.get_stream()); -#endif }); } @@ -275,11 +375,11 @@ using max_segment_size_list = nvbench::enum_type_list< // 1024, 2048, 4096, - 8192, - 16384, - 32768 + 8192 #if 0 // need these, waiting for implementation to catch up , + 16384, + 32768, 65536, 131072, 262144, From 82eae9a9054cff5329adbb4585861f37002a9bde Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 11:47:32 -0700 Subject: [PATCH 008/126] Force atoms --- cub/cub/agent/agent_batched_topk_atomic.cuh | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh index 8d8b2beadbd..2306149db9c 100644 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -140,7 +140,7 @@ private: for (unsigned i = threadIdx.x; i < segment_size; i += threads_per_block) { auto bit_ordered_key = twiddle_t::In(reinterpret_cast(block_keys_in[i])); - min = top_k(bit_ordered_key, temp_storage.top_k, k); + min = block_top_k(bit_ordered_key, temp_storage.top_k, k); } __syncthreads(); } @@ -161,13 +161,11 @@ private: } } - template - __device__ bit_ordered_t top_k(bit_ordered_t key, bit_ordered_t* top, unsigned k) + __device__ bit_ordered_t block_top_k(bit_ordered_t key, bit_ordered_t* top, unsigned k) { for (unsigned j = 0; j < k; ++j) { - cuda::atomic_ref shared_top_ref(top[j]); - const bit_ordered_t old_max_j = shared_top_ref.fetch_max(key, cuda::memory_order_relaxed); + const bit_ordered_t old_max_j = atomicMax_block(top + j, key); if (old_max_j < key) { key = old_max_j; From 34e3ab9c18f5405b47666a729e4e22919284122c Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 11:57:44 -0700 Subject: [PATCH 009/126] Early exit --- cub/cub/agent/agent_batched_topk_atomic.cuh | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh index 2306149db9c..d208f6b1df1 100644 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -139,8 +139,12 @@ private: auto block_keys_in = d_key_segments_it[segment_id]; for (unsigned i = threadIdx.x; i < segment_size; i += threads_per_block) { - auto bit_ordered_key = twiddle_t::In(reinterpret_cast(block_keys_in[i])); - min = block_top_k(bit_ordered_key, temp_storage.top_k, k); + auto key = twiddle_t::In(reinterpret_cast(block_keys_in[i])); + // early exit check: -65% + if (key > min) + { + min = block_top_k(key, temp_storage.top_k, k); + } } __syncthreads(); } From 00bc1e2b5ee1959c2048f4996f026a7a291d057c Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 12:30:55 -0700 Subject: [PATCH 010/126] Late entry --- cub/cub/agent/agent_batched_topk_atomic.cuh | 23 ++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh index d208f6b1df1..fbca00b2eb2 100644 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -167,7 +168,27 @@ private: __device__ bit_ordered_t block_top_k(bit_ordered_t key, bit_ordered_t* top, unsigned k) { - for (unsigned j = 0; j < k; ++j) + // late entry: binary-search the descending-sorted prefix to skip slots already >= key. + // Static iteration count (compile-time unrolled via max_k) mirrors cub::StaticUpperBound. + unsigned lo = 0; + unsigned hi = k; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i <= Log2::VALUE; i++) + { + unsigned mid = cub::MidPoint(lo, hi); + mid = (::cuda::std::min) (mid, k - 1); + const bit_ordered_t v_at = *static_cast(top + mid); + if (v_at >= key) + { + lo = mid + 1; + } + else + { + hi = mid; + } + } + + for (unsigned j = lo; j < k; ++j) { const bit_ordered_t old_max_j = atomicMax_block(top + j, key); if (old_max_j < key) From b36aa62e6d371df19fb81cf5a440a9acd912a432 Mon Sep 17 00:00:00 2001 From: Georgii Evtushenko Date: Sat, 23 May 2026 12:35:43 -0700 Subject: [PATCH 011/126] Bound search --- cub/cub/agent/agent_batched_topk_atomic.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh index fbca00b2eb2..66e318754d0 100644 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ b/cub/cub/agent/agent_batched_topk_atomic.cuh @@ -173,7 +173,7 @@ private: unsigned lo = 0; unsigned hi = k; _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i <= Log2::VALUE; i++) + for (int i = 0; i <= (cuda::std::min) (4, Log2::VALUE); i++) { unsigned mid = cub::MidPoint(lo, hi); mid = (::cuda::std::min) (mid, k - 1); From 2477d87ecaff6a7d70944cbd80068a27c4da49b3 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 26 May 2026 02:37:51 +0200 Subject: [PATCH 012/126] Remove restriction of one bucket per thread --- cub/cub/agent/agent_batched_topk_cluster.cuh | 106 ++++++++++-------- .../dispatch_batched_topk_cluster.cuh | 73 ++++++------ 2 files changed, 93 insertions(+), 86 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f9df6d2158c..40a5ba037c3 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -130,12 +130,12 @@ struct agent_batched_topk_cluster // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram // --------------------------------------------------------------------------- - // Constraint: every leader thread owns at most one bucket. Threads beyond - // num_buckets contribute zero to the scan, which keeps the prefix-sum loop - // body trivial (1 item per thread, no inner loops). - static_assert(num_buckets <= threads_per_block, - "Cluster top-k requires num_buckets <= threads_per_block (1 bin per leader thread or fewer)"); - using block_scan_t = BlockScan; + // The leader-block prefix sum spans `num_buckets` entries. Each thread owns + // `buckets_per_thread` consecutive buckets in a blocked arrangement; entries + // past the last valid bucket contribute zero so the histogram size is no + // longer constrained to be `<= threads_per_block`. + static constexpr int buckets_per_thread = ::cuda::ceil_div(num_buckets, threads_per_block); + using block_scan_t = BlockScan; // --------------------------------------------------------------------------- // Shared memory storage @@ -208,11 +208,12 @@ private: } // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram - // plus identification of the bucket holding the k-th item. Every thread of - // the leader block participates with one bucket; threads with rank beyond - // num_buckets contribute zero. The single thread that owns the k-th bucket - // writes the per-pass state. The caller must guarantee the leader block - // has finished its DSMEM merge before invoking this. + // plus identification of the bucket holding the k-th item. Each thread of + // the leader block contributes `buckets_per_thread` consecutive buckets in + // a blocked arrangement; entries past `num_buckets` contribute zero. The + // single (thread, slot) pair that owns the k-th bucket writes the per-pass + // state. The caller must guarantee the leader block has finished its DSMEM + // merge before invoking this. _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) { // Capture `state.k` before the scan: this is the only legal window where @@ -221,27 +222,38 @@ private: // after that point would race with that write. const out_offset_t target_k = temp_storage.state.k; - const int bucket = static_cast(threadIdx.x); - const bool owns_bucket = bucket < num_buckets; - const offset_t hist_val = owns_bucket ? temp_storage.hist[bucket] : offset_t{0}; - offset_t prefix = 0; + offset_t hist_vals[buckets_per_thread]; + offset_t prefixes[buckets_per_thread]; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_val, prefix); + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) + { + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + hist_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; + } + + block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_vals, prefixes); - // Exactly one thread satisfies `prefix < target_k <= prefix + hist_val`. - if (owns_bucket && prefix < target_k && prefix + hist_val >= target_k) + // Exactly one (thread, slot) pair satisfies + // `prefix < target_k <= prefix + hist_val`. + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) { - const out_offset_t new_k = target_k - static_cast(prefix); - const offset_t new_len = hist_val; - temp_storage.state.len = new_len; - temp_storage.state.k = new_k; - // Early-stop opportunity: the bucket holding the k-th key contains - // exactly the remaining `k` items. Every candidate is therefore part - // of the top-k, so subsequent radix passes only redistribute the same - // items across finer buckets without changing the final result. - temp_storage.state.early_stop = - (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + if (bucket < num_buckets && prefixes[j] < target_k && prefixes[j] + hist_vals[j] >= target_k) + { + const out_offset_t new_k = target_k - static_cast(prefixes[j]); + const offset_t new_len = hist_vals[j]; + temp_storage.state.len = new_len; + temp_storage.state.k = new_k; + // Early-stop opportunity: the bucket holding the k-th key contains + // exactly the remaining `k` items. Every candidate is therefore part + // of the top-k, so subsequent radix passes only redistribute the same + // items across finer buckets without changing the final result. + temp_storage.state.early_stop = + (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; + detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + } } } @@ -249,28 +261,29 @@ private: // Per-direction implementation // ------------------------------------------------------------------------- template - _CCCL_DEVICE _CCCL_FORCEINLINE void run( - ::cooperative_groups::cluster_group& cluster, - num_segments_val_t segment_id, - unsigned int cluster_rank, - segment_size_val_t segment_size, - out_offset_t k) + _CCCL_DEVICE _CCCL_FORCEINLINE void + run(::cooperative_groups::cluster_group& cluster, + num_segments_val_t segment_id, + unsigned int cluster_rank, + segment_size_val_t segment_size, + out_offset_t k) { - using extract_bin_op_t = detail::topk::extract_bin_op_t; - using identify_candidates_op_t = detail::topk::identify_candidates_op_t; + using extract_bin_op_t = detail::topk::extract_bin_op_t; + using identify_candidates_op_t = + detail::topk::identify_candidates_op_t; constexpr int total_bits = int{sizeof(key_t)} * 8; constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); // Sentinel key used to pad partial tiles. Worst-possible value for the // requested direction so the sentinel can't land in the selected set. - const key_t pad_key = (SelectDirection == detail::topk::select::max) - ? ::cuda::std::numeric_limits::lowest() - : ::cuda::std::numeric_limits::max(); + const key_t pad_key = + (SelectDirection == detail::topk::select::max) + ? ::cuda::std::numeric_limits::lowest() + : ::cuda::std::numeric_limits::max(); - auto block_keys_in = d_key_segments_it[segment_id]; - const auto block_offset_in_cluster = - static_cast(static_cast(cluster_rank) * tile_items); + auto block_keys_in = d_key_segments_it[segment_id]; + const auto block_offset_in_cluster = static_cast(static_cast(cluster_rank) * tile_items); // Striped load into per-thread registers. Each thread holds // `items_per_thread` keys spanning the block's tile within the cluster @@ -281,7 +294,7 @@ private: { const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; + thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; } // DSMEM pointers into the leader block's shared memory. @@ -475,9 +488,7 @@ private: cluster.sync(); const bool ok = detail::params::dispatch_discrete( - select_directions, - segment_id, - [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { + select_directions, segment_id, [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; this->template run(cluster, segment_id, cluster_rank, segment_size, k); }); @@ -485,7 +496,6 @@ private: (void) ok; } }; - } // namespace detail::batched_topk_cluster CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index db6f9ce7d1f..e25e4f27971 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -28,13 +28,13 @@ #include #include +#include + #include #include #include -#include - CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster @@ -43,11 +43,13 @@ namespace detail::batched_topk_cluster // Default tuning // ----------------------------------------------------------------------------- // Defaults give a 32 Ki-item cluster tile (cluster_size * threads_per_block * -// items_per_thread = 8 * 256 * 16). The leader block's BlockScan uses 1 bin -// per thread, so threads_per_block (256) must be >= num_buckets -// (1 << bits_per_pass = 256). Cluster size is 8 — the largest portable -// cluster size — so all eight blocks of the cluster participate in the -// histogram merge. +// items_per_thread = 8 * 256 * 16). The leader block's BlockScan partitions +// `num_buckets` (1 << bits_per_pass = 256) entries across threads_per_block +// (256) threads, so each thread happens to own a single bin in the default +// configuration; the agent itself supports `num_buckets > threads_per_block` +// via `ceil_div(num_buckets, threads_per_block)` items per thread. Cluster +// size is 8 — the largest portable cluster size — so all eight blocks of the +// cluster participate in the histogram merge. inline constexpr int default_cluster_size = 8; inline constexpr int default_threads_per_block = 256; inline constexpr int default_items_per_thread = 16; @@ -66,14 +68,14 @@ template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KERNEL_ATTRIBUTES void - device_segmented_topk_cluster_kernel( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments) +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) + _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments) { using agent_t = agent_batched_topk_cluster< ClusterSize, @@ -89,7 +91,8 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KER __shared__ typename agent_t::TempStorage temp_storage; - agent_t agent(temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + agent_t agent( + temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); agent.Process(); } @@ -101,10 +104,10 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) _CCCL_KER // keys-only path and assumes the maximum segment size fits in a single cluster // tile (`ClusterSize * ThreadsPerBlock * ItemsPerThread`). Both restrictions // match the small/decode-style segmented top-k workloads we are prototyping for. -template , "Number of segments must be resolved on the host."); - using num_segments_val_t = typename NumSegmentsParameterT::value_type; - const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); const auto num_seg_unsigned = static_cast(num_seg_val); if (num_seg_unsigned == 0) { @@ -167,7 +170,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaErrorNotSupported; } - const auto grid_blocks = static_cast(num_seg_unsigned) * static_cast(ClusterSize); + const auto grid_blocks = + static_cast(num_seg_unsigned) * static_cast(ClusterSize); if (grid_blocks > static_cast(::cuda::std::numeric_limits::max())) { return cudaErrorInvalidValue; @@ -188,20 +192,14 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // The kernel declares `__cluster_dims__(ClusterSize, 1, 1)`, so the // cluster dimension is fixed at compile time. Launch via the standard // triple-chevron syntax; the cluster setup is applied automatically. - if (const auto error = - CubDebug(THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - dim3(static_cast(grid_blocks), 1, 1), - dim3(static_cast(ThreadsPerBlock), 1, 1), - 0, - stream) - .doit( - kernel, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments))) + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + dim3(static_cast(grid_blocks), 1, 1), + dim3(static_cast(ThreadsPerBlock), 1, 1), + 0, + stream) + .doit( + kernel, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments))) { return error; } @@ -218,7 +216,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return CubDebug(detail::DebugSyncStream(stream)); } - } // namespace detail::batched_topk_cluster CUB_NAMESPACE_END From dc1c00d844d84f99889b2db511df59ccc2b0e8ef Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 26 May 2026 14:36:15 +0200 Subject: [PATCH 013/126] Remove experimental atomic backend --- .../bench/segmented_topk/fixed/keys.cu | 17 -- .../bench/segmented_topk/variable/keys.cu | 17 -- cub/cub/agent/agent_batched_topk_atomic.cuh | 204 ------------------ .../dispatch/dispatch_batched_topk_atomic.cuh | 150 ------------- .../catch2_test_device_segmented_topk_keys.cu | 18 +- 5 files changed, 1 insertion(+), 405 deletions(-) delete mode 100644 cub/cub/agent/agent_batched_topk_atomic.cuh delete mode 100644 cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index 78dec4bd1d2..2fe2ce50d62 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -4,7 +4,6 @@ #include #include #include -#include #include #include @@ -27,7 +26,6 @@ enum class topk_backend { baseline, cluster, - atomic, device, }; @@ -96,21 +94,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( total_num_items, stream); } - else if constexpr (selected_backend == topk_backend::atomic) - { - (void) h_segment_sizes; - return cub::detail::batched_topk_atomic::dispatch( - d_temp_storage, - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k, - select_directions, - num_segments, - total_num_items, - stream); - } else if constexpr (selected_backend == topk_backend::device) { using num_segments_val_t = typename NumSegmentsParameterT::value_type; diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index 5450fbcb9a0..486f6f55ff2 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -4,7 +4,6 @@ #include #include #include -#include #include #include @@ -19,7 +18,6 @@ enum class topk_backend { baseline, cluster, - atomic, device, }; @@ -61,21 +59,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( total_num_items, stream); } - else if constexpr (selected_backend == topk_backend::atomic) - { - (void) h_segment_sizes; - return cub::detail::batched_topk_atomic::dispatch( - d_temp_storage, - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k, - select_directions, - num_segments, - total_num_items, - stream); - } else if constexpr (selected_backend == topk_backend::device) { using num_segments_val_t = typename NumSegmentsParameterT::value_type; diff --git a/cub/cub/agent/agent_batched_topk_atomic.cuh b/cub/cub/agent/agent_batched_topk_atomic.cuh deleted file mode 100644 index 66e318754d0..00000000000 --- a/cub/cub/agent/agent_batched_topk_atomic.cuh +++ /dev/null @@ -1,204 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -#pragma once - -#include - -#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) -# pragma GCC system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) -# pragma clang system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) -# pragma system_header -#endif // no system header - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -#include - -CUB_NAMESPACE_BEGIN - -namespace detail::batched_topk_atomic -{ -struct atomic_topk_policy -{ - int cluster_size; - int threads_per_block; -}; - -template -struct agent_batched_topk_atomic -{ - using key_it_t = it_value_t; - using key_t = it_value_t; - - using segment_size_val_t = typename SegmentSizeParameterT::value_type; - using num_segments_val_t = typename NumSegmentsParameterT::value_type; - - using offset_t = ::cuda::std::uint32_t; - using out_offset_t = ::cuda::std::uint32_t; - - // TODO(gevtushenko): direction - using twiddle_t = cub::RadixSortTwiddle; - using bit_ordered_t = typename cub::Traits::UnsignedBits; - - static constexpr auto threads_per_block = ThreadsPerBlock; - static constexpr auto max_k = params::static_max_value_v; - - using decomposer_t = detail::identity_decomposer_t; - using block_scan_t = BlockScan; - - struct _TempStorage - { - bit_ordered_t top_k[max_k]; - }; - - struct TempStorage : Uninitialized<_TempStorage> - {}; - - _TempStorage& temp_storage; - KeyInputItItT d_key_segments_it; - KeyOutputItItT d_key_segments_out_it; - SegmentSizeParameterT segment_sizes; - KParameterT k_param; - SelectDirectionParameterT select_directions; - NumSegmentsParameterT num_segments; - - _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_atomic( - TempStorage& temp_storage_, - KeyInputItItT d_key_segments_it_, - KeyOutputItItT d_key_segments_out_it_, - SegmentSizeParameterT segment_sizes_, - KParameterT k_param_, - SelectDirectionParameterT select_directions_, - NumSegmentsParameterT num_segments_) - : temp_storage(temp_storage_.Alias()) - , d_key_segments_it(d_key_segments_it_) - , d_key_segments_out_it(d_key_segments_out_it_) - , segment_sizes(segment_sizes_) - , k_param(k_param_) - , select_directions(select_directions_) - , num_segments(num_segments_) - {} - - _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() - { - const auto segment_id = static_cast(blockIdx.x); - const auto direction = select_directions.get_param(segment_id); - const auto segment_size = static_cast(segment_sizes.get_param(segment_id)); - const auto k_requested = static_cast(k_param.get_param(segment_id)); - const auto k = - static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); - - if (k == 0) - { - return; - } - - // TODO(gevtushenk); direction - (void) direction; - bit_ordered_t sentinel{0}; - - init_smem(k, sentinel); - deposit(segment_id, segment_size, k, sentinel); - merge(segment_id, k); - store(segment_id, k); - } - -private: - __device__ void init_smem(unsigned k, bit_ordered_t sentinel) - { - for (unsigned i = threadIdx.x; i < k; i += threads_per_block) - { - temp_storage.top_k[i] = sentinel; - } - __syncthreads(); - } - - __device__ void deposit(unsigned segment_id, unsigned segment_size, unsigned k, bit_ordered_t sentinel) - { - auto min = sentinel; - auto block_keys_in = d_key_segments_it[segment_id]; - for (unsigned i = threadIdx.x; i < segment_size; i += threads_per_block) - { - auto key = twiddle_t::In(reinterpret_cast(block_keys_in[i])); - // early exit check: -65% - if (key > min) - { - min = block_top_k(key, temp_storage.top_k, k); - } - } - __syncthreads(); - } - - __device__ void merge(unsigned segment_id, unsigned k) - { - // TODO(gevtushenko): cluster - } - - __device__ void store(unsigned segment_id, unsigned k) - { - auto block_keys_out = d_key_segments_out_it[segment_id]; - - for (unsigned i = threadIdx.x; i < k; i += threads_per_block) - { - auto key = twiddle_t::Out(temp_storage.top_k[i]); - block_keys_out[i] = reinterpret_cast(key); - } - } - - __device__ bit_ordered_t block_top_k(bit_ordered_t key, bit_ordered_t* top, unsigned k) - { - // late entry: binary-search the descending-sorted prefix to skip slots already >= key. - // Static iteration count (compile-time unrolled via max_k) mirrors cub::StaticUpperBound. - unsigned lo = 0; - unsigned hi = k; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i <= (cuda::std::min) (4, Log2::VALUE); i++) - { - unsigned mid = cub::MidPoint(lo, hi); - mid = (::cuda::std::min) (mid, k - 1); - const bit_ordered_t v_at = *static_cast(top + mid); - if (v_at >= key) - { - lo = mid + 1; - } - else - { - hi = mid; - } - } - - for (unsigned j = lo; j < k; ++j) - { - const bit_ordered_t old_max_j = atomicMax_block(top + j, key); - if (old_max_j < key) - { - key = old_max_j; - } - } - return key; - } -}; -} // namespace detail::batched_topk_atomic - -CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh deleted file mode 100644 index 93d67f360a5..00000000000 --- a/cub/cub/device/dispatch/dispatch_batched_topk_atomic.cuh +++ /dev/null @@ -1,150 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -#pragma once - -#include - -#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) -# pragma GCC system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) -# pragma clang system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) -# pragma system_header -#endif // no system header - -#include -#include -#include -#include -#include - -#include - -#include -#include - -#include - -CUB_NAMESPACE_BEGIN - -namespace detail::batched_topk_atomic -{ -inline constexpr int default_threads_per_block = 256; - -template -__launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_atomic_kernel( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments) -{ - using agent_t = agent_batched_topk_atomic< - ThreadsPerBlock, - KeyInputItItT, - KeyOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - __shared__ typename agent_t::TempStorage temp_storage; - - agent_t agent( - temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); - - agent.Process(); -} - -template -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( - void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr) -{ - size_t allocation_sizes[1] = {1}; - void* allocations[1] = {}; - if (const auto error = - CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) - { - return error; - } - - if (d_temp_storage == nullptr) - { - return cudaSuccess; - } - - static_assert(!detail::params::is_per_segment_param_v, - "Number of segments must be resolved on the host."); - - using num_segments_val_t = typename NumSegmentsParameterT::value_type; - const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); - const auto num_seg_unsigned = static_cast(num_seg_val); - if (num_seg_unsigned == 0) - { - return cudaSuccess; - } - - int sm_version = 0; - if (const auto error = CubDebug(SmVersionUncached(sm_version))) - { - return error; - } - if (sm_version < 900) - { - return cudaErrorNotSupported; - } - - const auto grid_blocks = static_cast(num_seg_unsigned); - - auto kernel = device_segmented_topk_atomic_kernel< - ThreadsPerBlock, - KeyInputItItT, - KeyOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - if (const auto error = CubDebug( - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - static_cast(grid_blocks), static_cast(ThreadsPerBlock), 0, stream) - .doit( - kernel, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments))) - { - return error; - } - - if (const auto error = CubDebug(cudaPeekAtLastError())) - { - return error; - } - - return CubDebug(detail::DebugSyncStream(stream)); -} -} // namespace detail::batched_topk_atomic - -CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index c075b6da03e..3560b133d50 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -4,7 +4,6 @@ #include "insert_nested_NVTX_range_guard.h" #include -#include #include #include @@ -33,10 +32,9 @@ enum class topk_backend { baseline, cluster, - atomic, }; -inline constexpr topk_backend selected_backend = topk_backend::atomic; +inline constexpr topk_backend selected_backend = topk_backend::cluster; template (nullptr); From 784416a5fe9bf7eef6e0e674b484e4c0f4121f4f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 26 May 2026 14:56:22 +0200 Subject: [PATCH 014/126] Local atomic output index accumulation --- cub/cub/agent/agent_batched_topk_cluster.cuh | 109 ++++++++++++++++++- 1 file changed, 108 insertions(+), 1 deletion(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 40a5ba037c3..1ee2820f12b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -23,7 +23,12 @@ //! the next pass. //! //! Output cursors live in the same cluster-shared `state` and are reached the -//! same way (cluster-scope DSMEM atomics). +//! same way (cluster-scope DSMEM atomics). Defining +//! `CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT` switches the final filter pass +//! to a block-aggregated variant that first counts selected / candidate keys +//! in block-local shared-memory counters and then issues a single DSMEM +//! atomic per block per counter to reserve a contiguous output range. This +//! lets us A/B the two strategies without touching the rest of the kernel. #pragma once @@ -155,6 +160,15 @@ struct agent_batched_topk_cluster key_prefix_t broadcast_kth; ::cuda::std::uint32_t broadcast_early_stop; typename block_scan_t::TempStorage scan_storage; +#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT + // Final-filter block-aggregation slots. First populated by per-thread + // `atomicAdd_block` calls that count the block's selected / candidate + // keys, then overwritten by thread 0 with the leader-side base position + // returned from a single DSMEM atomic per counter, so the same slots + // double as the broadcast channel for the base to the rest of the block. + out_offset_t block_out_cnt; + out_offset_t block_out_back_cnt; +#endif }; struct TempStorage : Uninitialized<_TempStorage> @@ -402,6 +416,14 @@ private: if (threadIdx.x == 0) { temp_storage.broadcast_kth = leader_state->kth_key_bits; +#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT + // Piggyback the block-local output-counter init on the broadcast_kth + // publish: the `__syncthreads()` below already orders these writes + // against the per-thread `atomicAdd_block` calls in the first phase + // of the block-aggregated path, so no dedicated init+sync is needed. + temp_storage.block_out_cnt = 0; + temp_storage.block_out_back_cnt = 0; +#endif } __syncthreads(); kth_key_bits_local = temp_storage.broadcast_kth; @@ -417,6 +439,90 @@ private: // identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); +#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT + // Block-aggregated variant: each thread first reserves its slot in the + // block's local counters via cheap `atomicAdd_block`, then thread 0 of + // every block claims the block's contiguous range in the leader's + // counters with a single DSMEM `atomicAdd` per counter. This collapses + // up to `threads_per_block * items_per_thread` DSMEM atomics per block + // down to two, at the cost of one extra block-wide barrier (the local + // counters are zeroed inside the pre-existing broadcast_kth publish) + // and one additional pass over the per-thread keys to materialize the + // writes. + + // Per-thread, per-item local slot index within the block. Only the + // entries corresponding to `selected` / `candidate` keys are written and + // later read; the rest stay undefined because the write loop below + // re-runs `identify_op` and reproduces the same control flow. + out_offset_t local_pos[items_per_thread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + if (global_idx >= segment_size) + { + continue; + } + const key_t key = thread_keys[j]; + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + local_pos[j] = atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); + } + else if (res == detail::topk::candidate_class::candidate) + { + local_pos[j] = atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); + } + } + __syncthreads(); + + // Thread 0 of every block (leader block included) reserves the block's + // contiguous range in the leader's counters. The returned old value is + // the block's base; we stash it back into the same shared-memory slot + // so the rest of the block can pick it up after the barrier without a + // separate broadcast field. For the leader block these atomics target + // its own shared memory; for non-leaders they go through DSMEM. + if (threadIdx.x == 0) + { + const out_offset_t fwd_count = temp_storage.block_out_cnt; + const out_offset_t back_count = temp_storage.block_out_back_cnt; + temp_storage.block_out_cnt = atomicAdd(&leader_state->out_cnt, fwd_count); + temp_storage.block_out_back_cnt = atomicAdd(&leader_state->out_back_cnt, back_count); + } + __syncthreads(); + + const out_offset_t fwd_base = temp_storage.block_out_cnt; + const out_offset_t back_base = temp_storage.block_out_back_cnt; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); + const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; + if (global_idx >= segment_size) + { + continue; + } + const key_t key = thread_keys[j]; + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = fwd_base + local_pos[j]; + block_keys_out[pos] = key; + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = back_base + local_pos[j]; + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + } + } + } +#else _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < items_per_thread; ++j) { @@ -443,6 +549,7 @@ private: } } } +#endif // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast From 5ea4399284f51882e72f6ef1a603ade46e749b86 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 27 May 2026 15:55:48 +0200 Subject: [PATCH 015/126] Fix race introduced with early exit --- cub/cub/agent/agent_batched_topk_cluster.cuh | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 1ee2820f12b..2f09bceb634 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -348,6 +348,11 @@ private: if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) { last_pass = pass; + // Order the `broadcast_kth` load above against thread 0's + // overwrite of `broadcast_kth` in the final filter pass below. + // The normal loop exit gets this from the last iteration's + // `cluster.sync()`; the early-stop path skips that sync. + __syncthreads(); break; } #endif From 8657a21fa8d10fb862b163a945b3df6d70418cf0 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 27 May 2026 16:02:39 +0200 Subject: [PATCH 016/126] Add scan-then-reduce path Also avoid unnecessary grid sync in reduce-then-scan path --- cub/cub/agent/agent_batched_topk_cluster.cuh | 170 ++++++++++++++++--- 1 file changed, 146 insertions(+), 24 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2f09bceb634..f1019c9f01b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -29,9 +29,17 @@ //! in block-local shared-memory counters and then issues a single DSMEM //! atomic per block per counter to reserve a contiguous output range. This //! lets us A/B the two strategies without touching the rest of the kernel. +//! +//! Defining `CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE` flips the histogram +//! pipeline from reduce-then-scan to scan-then-reduce: every block does a +//! block-local `InclusiveSum` over its own histogram, then folds the scans +//! (not the raw counts) into the leader. The leader then identifies the +//! k-th bucket directly from `hist[bucket]` / `hist[bucket - 1]`. #pragma once +#define CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE + #include #if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) @@ -271,6 +279,41 @@ private: } } +#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE + // Scan-then-reduce counterpart of `leader_identify_kth_bucket`. Assumes + // `hist[i]` already holds the cluster-wide inclusive prefix sum. The + // per-bucket count is recovered as `inclusive - exclusive`, where + // `exclusive` is `hist[bucket - 1]` (or `0` for `bucket == 0`). + // + // `target_k` is `state.k` read by the caller before the `cluster.sync()` + // that precedes this call, so that barrier orders the read against the + // write this function may issue to `state.k`. + _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket_from_inclusive(int pass, out_offset_t target_k) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) + { + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + if (bucket >= num_buckets) + { + continue; + } + const offset_t inclusive = temp_storage.hist[bucket]; + const offset_t exclusive = (bucket == 0) ? offset_t{0} : temp_storage.hist[bucket - 1]; + if (exclusive < target_k && inclusive >= target_k) + { + const out_offset_t new_k = target_k - static_cast(exclusive); + const offset_t new_len = inclusive - exclusive; + temp_storage.state.len = new_len; + temp_storage.state.k = new_k; + temp_storage.state.early_stop = + (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; + detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + } + } + } +#endif + // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- @@ -368,29 +411,89 @@ private: identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); - // Step 1: block-private histogram, cheap block-scope atomic adds. - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) + // Step 1: block-private histogram. The leader uses device-scope + // `atomicAdd` so it is mutually atomic per the PTX ISA with the + // remote device-scope `atomicAdd`s that non-leaders issue against + // the same SMEM through DSMEM in Step 2. Non-leaders only write to + // their own `hist[]` and keep the cheaper `atomicAdd_block`. + // TODO(https://github.com/NVIDIA/cccl/issues/73): collapse both + // branches onto cluster-scope atomics once + // `cuda::thread_scope_cluster` is exposed in libcudacxx. + if (cluster_rank == 0) { - const key_t key = thread_keys[j]; - const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); - if (keep) + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) { - const int bucket = extract_op(key); - atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + const key_t key = thread_keys[j]; + const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); + if (keep) + { + const int bucket = extract_op(key); + atomicAdd(&temp_storage.hist[bucket], offset_t{1}); + } } } - // Cluster-wide barrier: every block must finish its block-scope atomic - // adds before any non-leader block starts DSMEM atomic adds into the - // leader's `hist`. A plain __syncthreads() would only order the local - // block's writes; the leader's `atomicAdd_block`s would still race with - // remote blocks' cluster-scope `atomicAdd`s targeting the same memory. + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < items_per_thread; ++j) + { + const key_t key = thread_keys[j]; + const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); + if (keep) + { + const int bucket = extract_op(key); + atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + } + } + } + +#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE + // Step 1b (scan-then-reduce): in-place block-local `InclusiveSum` + // over `hist[]`. The barrier publishes Step 1's atomic writes to + // the scan's reads. + __syncthreads(); + { + offset_t bucket_vals[buckets_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) + { + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + bucket_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; + } + block_scan_t(temp_storage.scan_storage).InclusiveSum(bucket_vals, bucket_vals); + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) + { + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + if (bucket < num_buckets) + { + temp_storage.hist[bucket] = bucket_vals[j]; + } + } + } + // Cluster-wide barrier: the leader's scan write-back to `hist[]` + // is non-atomic and would be lost if a remote Step 2 RMW + // interleaved with it, so every block must finish its write-back + // before anyone starts the fold. cluster.sync(); +#else + // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` + // are atomic at compatible scopes (see Step 1 dispatch). The + // cluster-wide ordering before Step 3's leader read of `hist[]` + // is supplied by the `cluster.sync()` further below. + __syncthreads(); +#endif - // Step 2: non-leader blocks fold their bucket counts into the leader's - // `hist` via cluster-scope DSMEM atomics. The leader's data is already - // present in `hist`; if the leader also merged from itself, its - // contribution would be counted twice. + // Step 2: non-leader blocks fold their per-bucket values + // (raw counts in the reduce-then-scan path, block-local inclusive + // scans in the scan-then-reduce path) into the leader's `hist` + // via DSMEM atomics. The leader skips this to avoid double-counting + // its own contribution. `atomicAdd` matches the leader's + // device-scope Step 1 atomic (see comment there). + // TODO(https://github.com/NVIDIA/cccl/issues/73): use a + // cluster-scope atomic once `cuda::thread_scope_cluster` is + // exposed in libcudacxx. if (cluster_rank != 0) { for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) @@ -402,14 +505,30 @@ private: } } } + +#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE + // Hoist the leader's read of `state.k` above the upcoming + // `cluster.sync()` so that barrier separates it from the write + // `leader_identify_kth_bucket_from_inclusive` may issue to + // `state.k`. Unconditional in every block to avoid predication; + // safe because `process_impl` initializes `state` everywhere. + const out_offset_t leader_target_k = temp_storage.state.k; +#endif + cluster.sync(); - // Step 3: the leader prefix-scans its (now merged) `hist` and updates - // the cluster-shared `state`. Subsequent reads (next-pass refresh, last - // filter) all observe these writes after the next cluster sync. + // Step 3: the leader walks the merged `hist` (raw counts in the + // reduce-then-scan path, cluster-wide inclusive scan in the + // scan-then-reduce path) and updates the cluster-shared `state`. + // Subsequent reads (next-pass refresh, last filter) all observe + // these writes after the next cluster sync. if (cluster_rank == 0) { +#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE + leader_identify_kth_bucket_from_inclusive(pass, leader_target_k); +#else leader_identify_kth_bucket(pass); +#endif } cluster.sync(); } @@ -585,10 +704,13 @@ private: return; } - // The leader block initializes the cluster-shared state. Every block - // (including the leader) will reset its own `hist` at the top of the - // per-pass loop. - if (cluster_rank == 0 && threadIdx.x == 0) + // Every block's thread 0 initializes its local `state`. Only the + // leader's copy is semantically read (non-leaders reach the cluster + // state through `leader_state`), but mirroring the writes everywhere + // keeps the scan-then-reduce path's unconditional `state.k` load + // safe under compute-sanitizer. Every block will reset its own + // `hist` at the top of the per-pass loop. + if (threadIdx.x == 0) { temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; From f8442af71215f811810f6b6a55ce1fa5bec01e5f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 29 May 2026 15:30:08 +0200 Subject: [PATCH 017/126] Make cluster size dynamic --- cub/cub/agent/agent_batched_topk_cluster.cuh | 73 ++-- .../dispatch_batched_topk_cluster.cuh | 365 ++++++++++++++---- .../tuning/tuning_batched_topk_cluster.cuh | 47 +++ 3 files changed, 381 insertions(+), 104 deletions(-) create mode 100644 cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f1019c9f01b..f456a8e5033 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -62,23 +62,14 @@ #include #include +#include + #include CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { -// ----------------------------------------------------------------------------- -// Tuning policy -// ----------------------------------------------------------------------------- -struct cluster_topk_policy -{ - int cluster_size; - int threads_per_block; - int items_per_thread; - int bits_per_pass; -}; - // ----------------------------------------------------------------------------- // Cluster-shared state. Lives in the leader block's shared memory and is // reached from every block of the cluster through DSMEM. @@ -104,8 +95,10 @@ struct alignas(16) cluster_topk_state // ----------------------------------------------------------------------------- // Cluster top-k agent // ----------------------------------------------------------------------------- -template ; using key_prefix_t = typename state_t::key_prefix_t; - static constexpr int cluster_size = ClusterSize; static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int items_per_thread = ItemsPerThread; static constexpr int bits_per_pass = BitsPerPass; static constexpr int tile_items = threads_per_block * items_per_thread; - static constexpr int cluster_tile = cluster_size * tile_items; static constexpr int num_buckets = 1 << bits_per_pass; using decomposer_t = detail::identity_decomposer_t; @@ -213,11 +204,14 @@ struct agent_batched_topk_cluster // --------------------------------------------------------------------------- // Main entry point // --------------------------------------------------------------------------- - // Prototype targets SM 9.0+ only; older architectures are unsupported by the - // dispatch and never reach this kernel. + // SM 9.0+ only. `_CG_HAS_CLUSTER_GROUP` keeps the body and the + // `process_impl` definition consistent across NVCC and clang-cuda/clangd; + // `NV_IF_TARGET` strips the call from NVCC's sub-SM-9.0 device passes. _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() { - process_impl(); +#if defined(_CG_HAS_CLUSTER_GROUP) + NV_IF_TARGET(NV_PROVIDES_SM_90, (process_impl();)); +#endif } private: @@ -317,6 +311,9 @@ private: // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- + // Stripped on sub-SM-9.0 device passes; uses `cluster_group`, which is only + // declared when `_CG_HAS_CLUSTER_GROUP` is set. +#if defined(_CG_HAS_CLUSTER_GROUP) template _CCCL_DEVICE _CCCL_FORCEINLINE void run(::cooperative_groups::cluster_group& cluster, @@ -381,13 +378,13 @@ private: if (threadIdx.x == 0) { temp_storage.broadcast_kth = leader_state->kth_key_bits; -#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP +# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP temp_storage.broadcast_early_stop = leader_state->early_stop; -#endif +# endif } __syncthreads(); kth_key_bits_local = temp_storage.broadcast_kth; -#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP +# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) { last_pass = pass; @@ -398,7 +395,7 @@ private: __syncthreads(); break; } -#endif +# endif } // Every block (including the leader) starts each pass with a fresh, @@ -448,7 +445,7 @@ private: } } -#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE +# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE // Step 1b (scan-then-reduce): in-place block-local `InclusiveSum` // over `hist[]`. The barrier publishes Step 1's atomic writes to // the scan's reads. @@ -477,13 +474,13 @@ private: // interleaved with it, so every block must finish its write-back // before anyone starts the fold. cluster.sync(); -#else +# else // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` // are atomic at compatible scopes (see Step 1 dispatch). The // cluster-wide ordering before Step 3's leader read of `hist[]` // is supplied by the `cluster.sync()` further below. __syncthreads(); -#endif +# endif // Step 2: non-leader blocks fold their per-bucket values // (raw counts in the reduce-then-scan path, block-local inclusive @@ -506,14 +503,14 @@ private: } } -#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE +# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE // Hoist the leader's read of `state.k` above the upcoming // `cluster.sync()` so that barrier separates it from the write // `leader_identify_kth_bucket_from_inclusive` may issue to // `state.k`. Unconditional in every block to avoid predication; // safe because `process_impl` initializes `state` everywhere. const out_offset_t leader_target_k = temp_storage.state.k; -#endif +# endif cluster.sync(); @@ -524,11 +521,11 @@ private: // these writes after the next cluster sync. if (cluster_rank == 0) { -#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE +# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE leader_identify_kth_bucket_from_inclusive(pass, leader_target_k); -#else +# else leader_identify_kth_bucket(pass); -#endif +# endif } cluster.sync(); } @@ -540,14 +537,14 @@ private: if (threadIdx.x == 0) { temp_storage.broadcast_kth = leader_state->kth_key_bits; -#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT +# ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT // Piggyback the block-local output-counter init on the broadcast_kth // publish: the `__syncthreads()` below already orders these writes // against the per-thread `atomicAdd_block` calls in the first phase // of the block-aggregated path, so no dedicated init+sync is needed. temp_storage.block_out_cnt = 0; temp_storage.block_out_back_cnt = 0; -#endif +# endif } __syncthreads(); kth_key_bits_local = temp_storage.broadcast_kth; @@ -563,7 +560,7 @@ private: // identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); -#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT +# ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT // Block-aggregated variant: each thread first reserves its slot in the // block's local counters via cheap `atomicAdd_block`, then thread 0 of // every block claims the block's contiguous range in the leader's @@ -646,7 +643,7 @@ private: } } } -#else +# else _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < items_per_thread; ++j) { @@ -673,7 +670,7 @@ private: } } } -#endif +# endif // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast @@ -687,7 +684,10 @@ private: { ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); const unsigned int cluster_rank = cluster.block_rank(); - const auto segment_id = static_cast(blockIdx.x / cluster_size); + // Runtime cluster width matches the launch attribute the dispatch passed + // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). + const unsigned int cluster_blocks = cluster.num_blocks(); + const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= num_segments.get_param(0)) { @@ -729,6 +729,7 @@ private: _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); (void) ok; } +#endif // _CG_HAS_CLUSTER_GROUP }; } // namespace detail::batched_topk_cluster diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index e25e4f27971..88725618ebc 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -2,13 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception //! @file -//! Single-kernel cluster-based batched top-k dispatch. +//! Cluster-based batched top-k dispatch. //! -//! Prototype that launches a single grid of thread block clusters to compute a +//! Prototype that launches a grid of thread block clusters to compute a //! segmented top-k. Each cluster processes one segment end-to-end: private //! histograms are reduced into the leader block via DSMEM atomics, then every //! block reads the merged histogram back through DSMEM, locally identifies the //! k-th bucket, and refines its in-register key set across radix passes. +//! +//! Two kernels share the agent body: +//! * Host: no `__cluster_dims__`, launched via `cudaLaunchKernelExC` with the +//! cluster width chosen at runtime (up to 16 on Hopper). +//! * CDP: static `__cluster_dims__(max_portable_cluster_blocks, 1, 1)`, +//! gated on `CUB_RDC_ENABLED` so CDP-disabled builds skip emitting it +//! (same pattern as `dispatch_segmented_sort.cuh`). #pragma once @@ -26,13 +33,19 @@ #include #include #include +#include #include #include +#include +#include +#include #include #include +#include + #include CUB_NAMESPACE_BEGIN @@ -40,26 +53,97 @@ CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { // ----------------------------------------------------------------------------- -// Default tuning +// Hardware constants // ----------------------------------------------------------------------------- -// Defaults give a 32 Ki-item cluster tile (cluster_size * threads_per_block * -// items_per_thread = 8 * 256 * 16). The leader block's BlockScan partitions -// `num_buckets` (1 << bits_per_pass = 256) entries across threads_per_block -// (256) threads, so each thread happens to own a single bin in the default -// configuration; the agent itself supports `num_buckets > threads_per_block` -// via `ceil_div(num_buckets, threads_per_block)` items per thread. Cluster -// size is 8 — the largest portable cluster size — so all eight blocks of the -// cluster participate in the histogram merge. -inline constexpr int default_cluster_size = 8; -inline constexpr int default_threads_per_block = 256; -inline constexpr int default_items_per_thread = 16; -inline constexpr int default_bits_per_pass = 8; +// Largest cluster width guaranteed on every SM 9.0+ device. +inline constexpr int max_portable_cluster_blocks = 8; + +// CUDA's hardware ceiling on cluster width (Hopper supports up to 16). +inline constexpr int max_supported_cluster_blocks = 16; // ----------------------------------------------------------------------------- -// Kernel entry point +// Cluster-size selection // ----------------------------------------------------------------------------- -template ::max()`. +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto +runtime_max_segment_size([[maybe_unused]] const SegmentSizeParameterT& segment_sizes) noexcept +{ + if constexpr (detail::params::is_static_param_v) + { + return detail::params::static_max_value_v; + } + else if constexpr (detail::params::is_per_segment_param_v) + { + return segment_sizes.max_value; + } + else + { + return segment_sizes.value; + } +} + +// Smallest cluster width covering `max_seg_size`, capped at +// `hw_max_cluster_blocks`. Clamps first so a loose `numeric_limits::max()` +// returns the cap instead of overflowing; whether the chosen tile actually +// covers the original bound is the dispatch's job to verify. +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr int +select_cluster_blocks(SegmentSizeValueT max_seg_size, int items_per_block, int hw_max_cluster_blocks) noexcept +{ + using value_t = SegmentSizeValueT; + + const auto max_tile = static_cast(hw_max_cluster_blocks * items_per_block); + const auto bounded = ::cuda::std::min(max_seg_size, max_tile); + return ::cuda::ceil_div(::cuda::narrow(bounded), items_per_block); +} + +// ----------------------------------------------------------------------------- +// Kernel entry points +// ----------------------------------------------------------------------------- +// Dynamic-cluster kernel for host launches; the agent reads the active cluster +// width via cooperative groups. +template +__launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments) +{ + using agent_t = agent_batched_topk_cluster< + ThreadsPerBlock, + ItemsPerThread, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + + agent_t agent( + temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + + agent.Process(); +} + +#ifdef CUB_RDC_ENABLED +// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so the +// triple-chevron launch from device code needs no `cudaFuncSetAttribute`. +template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) - _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) + _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, SegmentSizeParameterT segment_sizes, @@ -78,7 +162,6 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) NumSegmentsParameterT num_segments) { using agent_t = agent_batched_topk_cluster< - ClusterSize, ThreadsPerBlock, ItemsPerThread, BitsPerPass, @@ -96,25 +179,85 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterSize, 1, 1) agent.Process(); } +#endif // CUB_RDC_ENABLED + +// ----------------------------------------------------------------------------- +// NVCC kernel-emission workaround +// ----------------------------------------------------------------------------- +// NVCC `-rdc=false` only emits the device side of a templated `__global__` +// when it sees a triple-chevron in the same TU; CUDA 13's +// `-static-global-template-stub=true` makes the host stub internally linked +// on top. Address-of and `cudaLaunchKernelEx` are not enough. The host path +// must use `cudaLaunchKernelExC` for the cluster attribute, so instantiating +// `force_emit_kernel::emit` parses a (dead) chevron in its place. +// `Args` are deduced from the function-pointer type to avoid repeating the +// dispatch's template parameter list. +// +// See https://developer.nvidia.com/blog/cuda-c-compiler-updates-impacting-elf-visibility-and-linkage/. +template +struct force_emit_kernel; + +template +struct force_emit_kernel +{ + [[noreturn]] _CCCL_HOST static void emit(Args... args) + { + _CCCL_ASSERT(false, "force_emit_kernel::emit must never be called"); + // Unreachable; present only so NVCC emits `Kernel` for this TU. + if (false) + { + Kernel<<<1, 1>>>(args...); + } + _CCCL_UNREACHABLE(); + } +}; // ----------------------------------------------------------------------------- // Dispatch // ----------------------------------------------------------------------------- -// The dispatch is intentionally narrow: it currently only supports the -// keys-only path and assumes the maximum segment size fits in a single cluster -// tile (`ClusterSize * ThreadsPerBlock * ItemsPerThread`). Both restrictions -// match the small/decode-style segmented top-k workloads we are prototyping for. -template ; \ + if (const auto error = CubDebug( \ + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ + static_cast(grid_blocks), ThreadsPerBlock, 0, stream) \ + .doit(static_kernel, \ + d_key_segments_it, \ + d_key_segments_out_it, \ + segment_sizes, \ + k_param, \ + select_directions, \ + num_segments))) \ + { \ + return error; \ + } +#endif // CUB_RDC_ENABLED + +template + typename TotalNumItemsGuaranteeT, + typename PolicySelector = policy_selector> CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -125,14 +268,26 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionParameterT select_directions, NumSegmentsParameterT num_segments, [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr) + cudaStream_t stream = nullptr, + [[maybe_unused]] PolicySelector policy_selector = {}) { - static_assert(ClusterSize >= 1 && ClusterSize <= 8, "ClusterSize must be in [1, 8]"); + // Clusters are SM 9.0+ only and the tuning is currently identical across + // CCs, so pin the selector query to the minimum supported CC. + constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); + constexpr int ThreadsPerBlock = policy.threads_per_block; + constexpr int ItemsPerThread = policy.items_per_thread; + constexpr int BitsPerPass = policy.bits_per_pass; - constexpr int cluster_tile = ClusterSize * ThreadsPerBlock * ItemsPerThread; + constexpr int items_per_block = ThreadsPerBlock * ItemsPerThread; + constexpr int max_cluster_tile = max_supported_cluster_blocks * items_per_block; - static_assert(detail::params::static_max_value_v <= cluster_tile, - "Cluster top-k prototype only supports segments that fit in a single cluster tile."); + // Only static-size shapes can be rejected at compile time; runtime shapes + // are re-checked below against the chosen tile. + static_assert(!detail::params::is_static_param_v + || detail::params::static_max_value_v <= max_cluster_tile, + "Static segment size exceeds max_supported_cluster_blocks * ThreadsPerBlock * ItemsPerThread."); + + const auto max_seg_size = runtime_max_segment_size(segment_sizes); // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. size_t allocation_sizes[1] = {1}; @@ -151,10 +306,15 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( static_assert(!detail::params::is_per_segment_param_v, "Number of segments must be resolved on the host."); - using num_segments_val_t = typename NumSegmentsParameterT::value_type; - const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); - const auto num_seg_unsigned = static_cast(num_seg_val); - if (num_seg_unsigned == 0) + using num_segments_val_t = typename NumSegmentsParameterT::value_type; + const auto num_seg_val = num_segments.get_param(num_segments_val_t{0}); + if (num_seg_val == 0) + { + return cudaSuccess; + } + + // A zero bound would drive `clusterDim.x = 0`, which the runtime rejects. + if (max_seg_size == 0) { return cudaSuccess; } @@ -170,15 +330,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaErrorNotSupported; } - const auto grid_blocks = - static_cast(num_seg_unsigned) * static_cast(ClusterSize); - if (grid_blocks > static_cast(::cuda::std::numeric_limits::max())) - { - return cudaErrorInvalidValue; - } - - auto kernel = device_segmented_topk_cluster_kernel< - ClusterSize, + constexpr auto dynamic_kernel = &device_segmented_topk_cluster_kernel< ThreadsPerBlock, ItemsPerThread, BitsPerPass, @@ -189,26 +341,101 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionParameterT, NumSegmentsParameterT>; - // The kernel declares `__cluster_dims__(ClusterSize, 1, 1)`, so the - // cluster dimension is fixed at compile time. Launch via the standard - // triple-chevron syntax; the cluster setup is applied automatically. - if (const auto error = CubDebug( - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( - dim3(static_cast(grid_blocks), 1, 1), - dim3(static_cast(ThreadsPerBlock), 1, 1), - 0, - stream) - .doit( - kernel, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments))) - { - return error; - } + // Force NVCC to emit the device side of `dynamic_kernel` for this TU; see + // `force_emit_kernel` above. + [[maybe_unused]] constexpr auto force_emit = &force_emit_kernel::emit; + + NV_IF_TARGET( + NV_IS_HOST, + ({ + // Opt in to non-portable cluster widths (>8 on Hopper). + if (const auto error = CubDebug(cudaFuncSetAttribute( + reinterpret_cast(dynamic_kernel), cudaFuncAttributeNonPortableClusterSizeAllowed, 1))) + { + return error; + } + + // Reused across the probe and the launch; `clusterDim.x` is a placeholder + // until after `cudaOccupancyMaxPotentialClusterSize` (which ignores it). + ::cudaLaunchAttribute cluster_attr{}; + cluster_attr.id = ::cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim.x = 1; + cluster_attr.val.clusterDim.y = 1; + cluster_attr.val.clusterDim.z = 1; - // Synchronously surface any launch-time errors back to the caller. The - // prototype is keyed on `cudaLaunchKernelEx`, which may return success even - // when the cluster launch fails on the device (e.g. insufficient resources), - // so we explicitly check `cudaDeviceSynchronize` here while bringing up the - // implementation. + ::cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(1, 1, 1); + cfg.blockDim = dim3(static_cast(ThreadsPerBlock), 1, 1); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cfg.attrs = &cluster_attr; + cfg.numAttrs = 1; + + int hw_max_cluster_blocks = 0; + if (const auto error = CubDebug(cudaOccupancyMaxPotentialClusterSize( + &hw_max_cluster_blocks, reinterpret_cast(dynamic_kernel), &cfg))) + { + return error; + } + hw_max_cluster_blocks = ::cuda::std::min(hw_max_cluster_blocks, max_supported_cluster_blocks); + + const int cluster_blocks = select_cluster_blocks(max_seg_size, items_per_block, hw_max_cluster_blocks); + + // Exact bounds must fit the chosen tile; `per_segment_param` may carry + // a loose upper bound and is left to the agent's contract. + if constexpr (!detail::params::is_per_segment_param_v) + { + if (max_seg_size > static_cast(cluster_blocks) * items_per_block) + { + return cudaErrorInvalidValue; + } + } + + const auto grid_blocks = + static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); + if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + cfg.gridDim = dim3(static_cast(grid_blocks), 1, 1); + cluster_attr.val.clusterDim.x = static_cast(cluster_blocks); + + if (const auto error = CubDebug(::cudaLaunchKernelEx( + &cfg, + dynamic_kernel, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments))) + { + return error; + } + }), + ({ + // CDP path: same exact-vs-loose rule against the fixed static tile. + constexpr int static_cluster_tile = max_portable_cluster_blocks * items_per_block; + if constexpr (!detail::params::is_per_segment_param_v) + { + if (max_seg_size > static_cast(static_cluster_tile)) + { + return cudaErrorInvalidValue; + } + } + + const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) + * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); + if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + CUB_TOPK_CLUSTER_DEVICE_LAUNCH + })); + + // Cluster launches can fail on the device while reporting success; sync. if (const auto error = CubDebug(cudaPeekAtLastError())) { return error; @@ -216,6 +443,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return CubDebug(detail::DebugSyncStream(stream)); } + +#undef CUB_TOPK_CLUSTER_DEVICE_LAUNCH } // namespace detail::batched_topk_cluster CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh new file mode 100644 index 00000000000..423d40c5598 --- /dev/null +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +CUB_NAMESPACE_BEGIN +namespace detail::batched_topk_cluster +{ +// Per-block tile shape; the dispatch picks `cluster_blocks` at runtime. +struct cluster_topk_policy +{ + int threads_per_block; + int items_per_thread; + int bits_per_pass; +}; + +inline constexpr cluster_topk_policy default_policy{ + /*threads_per_block=*/256, + /*items_per_thread=*/16, + /*bits_per_pass=*/8, +}; + +// Default selector: one policy for every cluster-capable architecture +// (SM 9.0+). New tunings can be wired in by passing a custom selector to +// `dispatch` without changing the kernel signatures. +struct policy_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy + { + return default_policy; + } +}; +} // namespace detail::batched_topk_cluster + +CUB_NAMESPACE_END From c03daaf5adba6a475245f1551e5a2ba53149cffc Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 1 Jun 2026 14:34:00 +0200 Subject: [PATCH 018/126] Increase tile size using smem --- cub/cub/agent/agent_batched_topk_cluster.cuh | 542 ++++++++++++++---- .../dispatch_batched_topk_cluster.cuh | 332 ++++++++--- .../tuning/tuning_batched_topk_cluster.cuh | 19 +- .../catch2_test_device_segmented_topk_keys.cu | 125 ++++ ...tch2_test_segmented_topk_cluster_layout.cu | 98 ++++ 5 files changed, 937 insertions(+), 179 deletions(-) create mode 100644 cub/test/catch2_test_segmented_topk_cluster_layout.cu diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f456a8e5033..1a9c7279e22 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -51,16 +51,27 @@ #endif // no system header #include +#include #include #include #include #include +#include #include +#include +#include +#include + #include +#include +#include #include #include +#include #include +#include +#include #include @@ -92,6 +103,39 @@ struct alignas(16) cluster_topk_state ::cuda::std::uint32_t early_stop; }; +// Dynamic-SMEM layout shared by dispatch and the agent. `tile_capacity` is the physical per-CTA resident capacity +// passed to the kernel; `cluster_coverage` reserves one chunk of logical coverage for a possible unaligned head chunk. +template +struct smem_tile_layout +{ + static constexpr int chunk_items = ChunkBytes / int{sizeof(KeyT)}; + static constexpr int slot_alignment = + (::cuda::std::max) (LoadAlignBytes, detail::LoadToSharedBufferAlignBytes()); + static constexpr int slot_stride_bytes = + ::cuda::round_up(detail::LoadToSharedBufferSizeBytes(chunk_items) + slot_alignment, slot_alignment); + static constexpr int base_padding_bytes = (alignof(KeyT) > 16) ? slot_alignment : 0; + + [[nodiscard]] _CCCL_HOST_DEVICE static constexpr ::cuda::std::uint32_t tile_capacity(int dynamic_smem_bytes) noexcept + { + const int usable_bytes = dynamic_smem_bytes - base_padding_bytes; + if (usable_bytes <= 0) + { + return 0; + } + const int slots = usable_bytes / slot_stride_bytes; + return static_cast<::cuda::std::uint32_t>(slots * chunk_items); + } + + template + [[nodiscard]] _CCCL_HOST_DEVICE static constexpr SizeT + cluster_coverage(int cluster_blocks, ::cuda::std::uint32_t physical_tile_capacity) noexcept + { + const auto physical_coverage = static_cast(cluster_blocks) * static_cast(physical_tile_capacity); + const auto head_chunk_reserve = static_cast(chunk_items); + return (physical_coverage > head_chunk_reserve) ? physical_coverage - head_chunk_reserve : SizeT{0}; + } +}; + // ----------------------------------------------------------------------------- // Cluster top-k agent // ----------------------------------------------------------------------------- @@ -99,7 +143,10 @@ struct alignas(16) cluster_topk_state // it is not a template parameter; per-block tile layout is still controlled // by the template parameters below. template ; + static constexpr int chunk_items = smem_layout_t::chunk_items; + static constexpr int slot_alignment = smem_layout_t::slot_alignment; + static constexpr int slot_stride_bytes = smem_layout_t::slot_stride_bytes; + + static_assert(PipelineStages > 0); + static_assert(ChunkBytes > 0); + static_assert(LoadAlignBytes > 0); + static_assert(ChunkBytes % LoadAlignBytes == 0); + static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + static_assert(chunk_items > 0); using decomposer_t = detail::identity_decomposer_t; + using block_load_t = BlockLoadToShared; + + static constexpr bool use_block_load_to_shared = + THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v; // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram @@ -159,6 +220,7 @@ struct agent_batched_topk_cluster key_prefix_t broadcast_kth; ::cuda::std::uint32_t broadcast_early_stop; typename block_scan_t::TempStorage scan_storage; + typename block_load_t::TempStorage load_storage[PipelineStages]; #ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT // Final-filter block-aggregation slots. First populated by per-thread // `atomicAdd_block` calls that count the block's selected / candidate @@ -170,6 +232,137 @@ struct agent_batched_topk_cluster #endif }; + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span key_tile_buffer() const + { + const int slots = static_cast(tile_capacity / chunk_items); + return {key_slots, static_cast<::cuda::std::size_t>(slots * slot_stride_bytes)}; + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE key_t* slot_keys_unpadded(int slot) const + { + return reinterpret_cast(key_slots + slot * slot_stride_bytes); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span + available_key_tile_buffer(char* buffer_begin) const + { + const auto buffer = key_tile_buffer(); + char* const end = ::cuda::std::data(buffer) + static_cast(::cuda::std::size(buffer)); + _CCCL_ASSERT(buffer_begin >= ::cuda::std::data(buffer) && buffer_begin <= end, "Invalid key tile buffer cursor"); + _CCCL_ASSERT(::cuda::is_aligned(buffer_begin, detail::LoadToSharedBufferAlignBytes()), + "Key tile buffer cursor must satisfy BlockLoadToShared's shared-memory alignment"); + return {buffer_begin, static_cast<::cuda::std::size_t>(end - buffer_begin)}; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + append_contiguous_span(::cuda::std::span& merged, ::cuda::std::span next) const + { + const int next_count = static_cast(::cuda::std::size(next)); + _CCCL_ASSERT(static_cast<::cuda::std::size_t>(next_count) == ::cuda::std::size(next), + "Resident key span length must fit in int"); + if (next_count == 0) + { + return; + } + + const int merged_count = static_cast(::cuda::std::size(merged)); + _CCCL_ASSERT(static_cast<::cuda::std::size_t>(merged_count) == ::cuda::std::size(merged), + "Resident key span length must fit in int"); + if (merged_count == 0) + { + merged = next; + return; + } + + _CCCL_ASSERT(::cuda::std::data(merged) + merged_count == ::cuda::std::data(next), + "BlockLoadToShared returned non-contiguous resident key spans"); + merged = {::cuda::std::data(merged), static_cast<::cuda::std::size_t>(merged_count + next_count)}; + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int span_size(::cuda::std::span keys) const + { + const int count = static_cast(::cuda::std::size(keys)); + _CCCL_ASSERT(static_cast<::cuda::std::size_t>(count) == ::cuda::std::size(keys), + "Resident key span length must fit in int"); + return count; + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE char* span_end(::cuda::std::span keys) const + { + return reinterpret_cast(::cuda::std::data(keys) + span_size(keys)); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t + aligned_head_items(const key_t* base, offset_t segment_size) const + { + const auto base_addr = reinterpret_cast<::cuda::std::uintptr_t>(base); + const auto rem = base_addr % static_cast<::cuda::std::uintptr_t>(load_align_bytes); + const auto bytes = + (rem == 0) ? ::cuda::std::uintptr_t{0} : static_cast<::cuda::std::uintptr_t>(load_align_bytes) - rem; + const auto items = static_cast(bytes / static_cast<::cuda::std::uintptr_t>(sizeof(key_t))); + return (::cuda::std::min) (items, segment_size); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t num_chunks(offset_t segment_size, offset_t head_items) const + { + const offset_t remaining = segment_size - head_items; + return ((head_items != 0) ? offset_t{1} : offset_t{0}) + + static_cast(::cuda::ceil_div(remaining, offset_t{chunk_items})); + } + + struct chunk_desc + { + offset_t offset; + int count; + }; + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_desc + get_chunk(offset_t chunk_idx, offset_t segment_size, offset_t head_items) const + { + offset_t offset = chunk_idx * offset_t{chunk_items}; + if (head_items != 0) + { + if (chunk_idx == 0) + { + return {offset_t{0}, static_cast(head_items)}; + } + offset = head_items + (chunk_idx - 1) * offset_t{chunk_items}; + } + const offset_t remaining = segment_size - offset; + return {offset, static_cast((::cuda::std::min) (remaining, offset_t{chunk_items}))}; + } + + template + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_aligned_chunk(PtrT base, const chunk_desc chunk) const + { + const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); + const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); + return (begin % static_cast<::cuda::std::uintptr_t>(load_align_bytes) == 0) + && (end % static_cast<::cuda::std::uintptr_t>(load_align_bytes) == 0); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t + num_rank_chunks(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const + { + return (cluster_rank < chunks) + ? static_cast((chunks - 1 - cluster_rank) / cluster_blocks + 1) + : offset_t{0}; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const + { + const int chunk_count = span_size(chunk_keys); + const int iterations = ::cuda::ceil_div(chunk_count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk_count) + { + f(::cuda::std::data(chunk_keys)[local]); + } + }); + } + struct TempStorage : Uninitialized<_TempStorage> {}; @@ -183,6 +376,8 @@ struct agent_batched_topk_cluster KParameterT k_param; SelectDirectionParameterT select_directions; NumSegmentsParameterT num_segments; + char* key_slots; + offset_t tile_capacity; _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -191,7 +386,9 @@ struct agent_batched_topk_cluster SegmentSizeParameterT segment_sizes_, KParameterT k_param_, SelectDirectionParameterT select_directions_, - NumSegmentsParameterT num_segments_) + NumSegmentsParameterT num_segments_, + char* key_slots_, + offset_t tile_capacity_) : temp_storage(temp_storage_.Alias()) , d_key_segments_it(d_key_segments_it_) , d_key_segments_out_it(d_key_segments_out_it_) @@ -199,6 +396,8 @@ struct agent_batched_topk_cluster , k_param(k_param_) , select_directions(select_directions_) , num_segments(num_segments_) + , key_slots(key_slots_) + , tile_capacity(tile_capacity_) {} // --------------------------------------------------------------------------- @@ -329,27 +528,9 @@ private: constexpr int total_bits = int{sizeof(key_t)} * 8; constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); - // Sentinel key used to pad partial tiles. Worst-possible value for the - // requested direction so the sentinel can't land in the selected set. - const key_t pad_key = - (SelectDirection == detail::topk::select::max) - ? ::cuda::std::numeric_limits::lowest() - : ::cuda::std::numeric_limits::max(); - - auto block_keys_in = d_key_segments_it[segment_id]; - const auto block_offset_in_cluster = static_cast(static_cast(cluster_rank) * tile_items); - - // Striped load into per-thread registers. Each thread holds - // `items_per_thread` keys spanning the block's tile within the cluster - // tile. Out-of-segment slots receive sentinel padding. - key_t thread_keys[items_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) - { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - thread_keys[j] = (global_idx < segment_size) ? block_keys_in[global_idx] : pad_key; - } + auto block_keys_in = d_key_segments_it[segment_id]; + const auto segment_size_u32 = static_cast(segment_size); + const unsigned int cluster_size = cluster.num_blocks(); // DSMEM pointers into the leader block's shared memory. offset_t* leader_hist = cluster.map_shared_rank(temp_storage.hist, 0); @@ -365,6 +546,128 @@ private: // operator at the matching radix level. int last_pass = num_passes; + offset_t head_items = 0; + if constexpr (use_block_load_to_shared) + { + auto* block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); + head_items = aligned_head_items(block_keys_base, segment_size_u32); + } + // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple + // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only + // depends on the multiset of keys covered by the cluster. + const offset_t chunks = num_chunks(segment_size_u32, head_items); + const offset_t my_chunks = num_rank_chunks(chunks, cluster_rank, cluster_size); + // The launch coverage check reserves one extra chunk for the possible unaligned head, so every local chunk remains + // resident for later radix passes while only the BlockLoadToShared instances are reused round-robin. + _CCCL_ASSERT(my_chunks * offset_t{chunk_items} <= tile_capacity, "Dynamic shared memory tile is too small"); + + ::cuda::std::span resident_keys; + + reset_hist(); + __syncthreads(); + + { + extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); + auto add_first_pass = [&](const key_t& key) { + const int bucket = extract_op(key); + if (cluster_rank == 0) + { + atomicAdd(&temp_storage.hist[bucket], offset_t{1}); + } + else + { + atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + } + }; + + if constexpr (use_block_load_to_shared) + { + auto* block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); + // BlockLoadToShared is non-copyable and non-movable; keep the active pipeline local and emplace-only. + ::cuda::std::inplace_vector loaders; + ::cuda::std::inplace_vector tokens; + ::cuda::std::inplace_vector<::cuda::std::span, PipelineStages> pending_spans; + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_chunks)); + char* next_dst = key_slots; + + for (int stage = 0; stage < prologue; ++stage) + { + const offset_t chunk_idx = static_cast(cluster_rank) + static_cast(stage * cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const bool aligned_chunk = is_aligned_chunk(block_keys_base, chunk); + auto& loader = loaders.emplace_back(temp_storage.load_storage[stage]); + const ::cuda::std::span src{ + block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; + if (aligned_chunk) + { + pending_spans.emplace_back( + loader.template CopyAsync(available_key_tile_buffer(next_dst), src)); + } + else + { + pending_spans.emplace_back(loader.template CopyAsync(available_key_tile_buffer(next_dst), src)); + } + next_dst = span_end(pending_spans[stage]); + tokens.emplace_back(loader.Commit()); + } + + for (offset_t p = 0; p < my_chunks; ++p) + { + const int stage = static_cast(p % static_cast(prologue)); + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + loaders[stage].Wait(::cuda::std::move(tokens[stage])); + append_contiguous_span(resident_keys, pending_spans[stage]); + for_each_chunk_key(pending_spans[stage], add_first_pass); + + if (p + static_cast(prologue) < my_chunks) + { + const offset_t next_chunk_idx = static_cast(cluster_rank) + + (p + static_cast(prologue)) * static_cast(cluster_size); + const auto next_chunk = get_chunk(next_chunk_idx, segment_size_u32, head_items); + const bool next_aligned_chunk = is_aligned_chunk(block_keys_base, next_chunk); + const ::cuda::std::span src{ + block_keys_base + next_chunk.offset, static_cast<::cuda::std::size_t>(next_chunk.count)}; + if (next_aligned_chunk) + { + pending_spans[stage] = + loaders[stage].template CopyAsync(available_key_tile_buffer(next_dst), src); + } + else + { + pending_spans[stage] = loaders[stage].template CopyAsync(available_key_tile_buffer(next_dst), src); + } + next_dst = span_end(pending_spans[stage]); + tokens[stage] = loaders[stage].Commit(); + } + } + } + else + { + for (offset_t p = 0; p < my_chunks; ++p) + { + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk.count) + { + const key_t key = + block_keys_in[static_cast(chunk.offset + static_cast(local))]; + chunk_keys[local] = key; + add_first_pass(key); + } + }); + } + } + const int resident_count = span_size(resident_keys); + _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= tile_capacity, + "Dynamic shared memory tile is too small"); + __syncthreads(); + } + for (int pass = 0; pass < num_passes; ++pass) { const bool is_first_pass = (pass == 0); @@ -398,49 +701,51 @@ private: # endif } - // Every block (including the leader) starts each pass with a fresh, - // empty `hist`. For pass 0 the leader's initial reset in process_impl() - // covered the leader's slot, but every non-leader block must also - // reset its own. We just always reset here for symmetry. - reset_hist(); - __syncthreads(); - - identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); - extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); - - // Step 1: block-private histogram. The leader uses device-scope - // `atomicAdd` so it is mutually atomic per the PTX ISA with the - // remote device-scope `atomicAdd`s that non-leaders issue against - // the same SMEM through DSMEM in Step 2. Non-leaders only write to - // their own `hist[]` and keep the cheaper `atomicAdd_block`. - // TODO(https://github.com/NVIDIA/cccl/issues/73): collapse both - // branches onto cluster-scope atomics once - // `cuda::thread_scope_cluster` is exposed in libcudacxx. - if (cluster_rank == 0) + if (!is_first_pass) { - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) - { - const key_t key = thread_keys[j]; - const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); - if (keep) + // Every block (including the leader) starts each non-first pass with + // a fresh, empty `hist`. Pass 0 was fused with the load pipeline above. + reset_hist(); + __syncthreads(); + + identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); + extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); + + // Step 1: block-private histogram. The leader uses device-scope + // `atomicAdd` so it is mutually atomic per the PTX ISA with the + // remote device-scope `atomicAdd`s that non-leaders issue against + // the same SMEM through DSMEM in Step 2. Non-leaders only write to + // their own `hist[]` and keep the cheaper `atomicAdd_block`. + // TODO(https://github.com/NVIDIA/cccl/issues/73): collapse both + // branches onto cluster-scope atomics once + // `cuda::thread_scope_cluster` is exposed in libcudacxx. + auto add_hist = [&](const key_t& key) { + if (identify_op(key) == detail::topk::candidate_class::candidate) { const int bucket = extract_op(key); - atomicAdd(&temp_storage.hist[bucket], offset_t{1}); + if (cluster_rank == 0) + { + atomicAdd(&temp_storage.hist[bucket], offset_t{1}); + } + else + { + atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + } } + }; + + if constexpr (use_block_load_to_shared) + { + for_each_chunk_key(resident_keys, add_hist); } - } - else - { - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) + else { - const key_t key = thread_keys[j]; - const bool keep = is_first_pass || (identify_op(key) == detail::topk::candidate_class::candidate); - if (keep) + for (offset_t p = 0; p < my_chunks; ++p) { - const int bucket = extract_op(key); - atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } } } @@ -565,36 +870,35 @@ private: // block's local counters via cheap `atomicAdd_block`, then thread 0 of // every block claims the block's contiguous range in the leader's // counters with a single DSMEM `atomicAdd` per counter. This collapses - // up to `threads_per_block * items_per_thread` DSMEM atomics per block + // up to the block's resident key count DSMEM atomics per block // down to two, at the cost of one extra block-wide barrier (the local // counters are zeroed inside the pre-existing broadcast_kth publish) // and one additional pass over the per-thread keys to materialize the // writes. - // Per-thread, per-item local slot index within the block. Only the - // entries corresponding to `selected` / `candidate` keys are written and - // later read; the rest stay undefined because the write loop below - // re-runs `identify_op` and reproduces the same control flow. - out_offset_t local_pos[items_per_thread]; - - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) - { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - if (global_idx >= segment_size) - { - continue; - } - const key_t key = thread_keys[j]; - const auto res = identify_op(key); + auto count_selected = [&](const key_t& key) { + const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { - local_pos[j] = atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); + atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); } else if (res == detail::topk::candidate_class::candidate) { - local_pos[j] = atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); + atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); + } + }; + if constexpr (use_block_load_to_shared) + { + for_each_chunk_key(resident_keys, count_selected); + } + else + { + for (offset_t p = 0; p < my_chunks; ++p) + { + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, count_selected); } } __syncthreads(); @@ -617,44 +921,47 @@ private: const out_offset_t fwd_base = temp_storage.block_out_cnt; const out_offset_t back_base = temp_storage.block_out_back_cnt; - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) + if (threadIdx.x == 0) { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - if (global_idx >= segment_size) - { - continue; - } - const key_t key = thread_keys[j]; - const auto res = identify_op(key); + temp_storage.block_out_cnt = 0; + temp_storage.block_out_back_cnt = 0; + } + __syncthreads(); + + auto write_selected = [&](const key_t& key) { + const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = fwd_base + local_pos[j]; + const out_offset_t pos = fwd_base + atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); block_keys_out[pos] = key; } else if (res == detail::topk::candidate_class::candidate) { - const out_offset_t back_pos = back_base + local_pos[j]; + const out_offset_t back_pos = back_base + atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); if (back_pos < num_kth) { const out_offset_t pos = k - 1 - back_pos; block_keys_out[pos] = key; } } + }; + if constexpr (use_block_load_to_shared) + { + for_each_chunk_key(resident_keys, write_selected); } -# else - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < items_per_thread; ++j) + else { - const segment_size_val_t local_idx = static_cast(j * threads_per_block + threadIdx.x); - const segment_size_val_t global_idx = block_offset_in_cluster + local_idx; - if (global_idx >= segment_size) + for (offset_t p = 0; p < my_chunks; ++p) { - continue; + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } - const key_t key = thread_keys[j]; - const auto res = identify_op(key); + } +# else + auto write_selected = [&](const key_t& key) { + const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); @@ -669,6 +976,20 @@ private: block_keys_out[pos] = key; } } + }; + if constexpr (use_block_load_to_shared) + { + for_each_chunk_key(resident_keys, write_selected); + } + else + { + for (offset_t p = 0; p < my_chunks; ++p) + { + const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + } } # endif @@ -704,6 +1025,20 @@ private: return; } + bool segment_fits_offset = true; + if constexpr (sizeof(segment_size_val_t) > sizeof(offset_t)) + { + segment_fits_offset = + segment_size <= static_cast(::cuda::std::numeric_limits::max()); + } + const auto max_cluster_coverage = + smem_layout_t::template cluster_coverage(static_cast(cluster_blocks), tile_capacity); + if (!segment_fits_offset || segment_size > max_cluster_coverage) + { + _CCCL_ASSERT(false, "Segment exceeds the selected cluster top-k tile capacity"); + return; + } + // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster // state through `leader_state`), but mirroring the writes everywhere @@ -721,13 +1056,12 @@ private: } cluster.sync(); - const bool ok = detail::params::dispatch_discrete( + [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( select_directions, segment_id, [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; this->template run(cluster, segment_id, cluster_rank, segment_size, k); }); _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); - (void) ok; } #endif // _CG_HAS_CLUSTER_GROUP }; diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 88725618ebc..2165cea4cb2 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -8,7 +8,7 @@ //! segmented top-k. Each cluster processes one segment end-to-end: private //! histograms are reduced into the leader block via DSMEM atomics, then every //! block reads the merged histogram back through DSMEM, locally identifies the -//! k-th bucket, and refines its in-register key set across radix passes. +//! k-th bucket, and refines its shared-memory key tile across radix passes. //! //! Two kernels share the agent body: //! * Host: no `__cluster_dims__`, launched via `cudaLaunchKernelExC` with the @@ -39,6 +39,7 @@ #include #include +#include #include #include #include @@ -62,7 +63,7 @@ inline constexpr int max_portable_cluster_blocks = 8; inline constexpr int max_supported_cluster_blocks = 16; // ----------------------------------------------------------------------------- -// Cluster-size selection +// Cluster-size / dynamic-SMEM selection // ----------------------------------------------------------------------------- // Tightest upper bound carried by `SegmentSizeParameterT`. For // `per_segment_param` this can be the loose `numeric_limits::max()`. @@ -84,20 +85,12 @@ runtime_max_segment_size([[maybe_unused]] const SegmentSizeParameterT& segment_s } } -// Smallest cluster width covering `max_seg_size`, capped at -// `hw_max_cluster_blocks`. Clamps first so a loose `numeric_limits::max()` -// returns the cap instead of overflowing; whether the chosen tile actually -// covers the original bound is the dispatch's job to verify. -template -[[nodiscard]] _CCCL_HOST_DEVICE constexpr int -select_cluster_blocks(SegmentSizeValueT max_seg_size, int items_per_block, int hw_max_cluster_blocks) noexcept +struct launch_config { - using value_t = SegmentSizeValueT; - - const auto max_tile = static_cast(hw_max_cluster_blocks * items_per_block); - const auto bounded = ::cuda::std::min(max_seg_size, max_tile); - return ::cuda::ceil_div(::cuda::narrow(bounded), items_per_block); -} + int cluster_blocks; + int dynamic_smem_bytes; + ::cuda::std::uint32_t tile_capacity; +}; // ----------------------------------------------------------------------------- // Kernel entry points @@ -105,7 +98,10 @@ select_cluster_blocks(SegmentSizeValueT max_seg_size, int items_per_block, int h // Dynamic-cluster kernel for host launches; the agent reads the active cluster // width via cooperative groups. template ; __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Without manual realignment (`alignof(key_t) <= 16`), `slot_alignment` is a stride quantum; the base is only + // required to satisfy BlockLoadToShared's 16-byte shared-memory destination alignment. + if constexpr (alignof(typename agent_t::key_t) > 16) + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } agent_t agent( - temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + key_slots, + tile_capacity); agent.Process(); } @@ -144,7 +163,10 @@ __launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented // CDP-only static-cluster kernel: compile-time `__cluster_dims__` so the // triple-chevron launch from device code needs no `cudaFuncSetAttribute`. template ; __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Without manual realignment (`alignof(key_t) <= 16`), `slot_alignment` is a stride quantum; the base is only + // required to satisfy BlockLoadToShared's 16-byte shared-memory destination alignment. + if constexpr (alignof(typename agent_t::key_t) > 16) + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } agent_t agent( - temp_storage, d_key_segments_it, d_key_segments_out_it, segment_sizes, k_param, select_directions, num_segments); + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + key_slots, + tile_capacity); agent.Process(); } @@ -216,37 +261,41 @@ struct force_emit_kernel // Dispatch // ----------------------------------------------------------------------------- // Keys-only; every segment must fit in one cluster tile. Host picks -// `cluster_blocks` at runtime (capped by `cudaOccupancyMaxPotentialClusterSize`); -// CDP uses the static kernel at `max_portable_cluster_blocks`. +// `(cluster_blocks, dynamic_smem_bytes)` at runtime from a finite table; CDP +// uses the static kernel at `max_portable_cluster_blocks` and portable SMEM. // CDP launch body, empty when CDP is disabled. Wrapped in a macro because // `#ifdef` can't sit inside `NV_IF_TARGET`. #ifndef CUB_RDC_ENABLED # define CUB_TOPK_CLUSTER_DEVICE_LAUNCH #else // CUB_RDC_ENABLED -# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ - auto static_kernel = device_segmented_topk_cluster_kernel_static< \ - ThreadsPerBlock, \ - ItemsPerThread, \ - BitsPerPass, \ - KeyInputItItT, \ - KeyOutputItItT, \ - SegmentSizeParameterT, \ - KParameterT, \ - SelectDirectionParameterT, \ - NumSegmentsParameterT>; \ - if (const auto error = CubDebug( \ - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ - static_cast(grid_blocks), ThreadsPerBlock, 0, stream) \ - .doit(static_kernel, \ - d_key_segments_it, \ - d_key_segments_out_it, \ - segment_sizes, \ - k_param, \ - select_directions, \ - num_segments))) \ - { \ - return error; \ +# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ + auto static_kernel = device_segmented_topk_cluster_kernel_static< \ + ThreadsPerBlock, \ + UnrollFactor, \ + PipelineStages, \ + ChunkBytes, \ + LoadAlignBytes, \ + BitsPerPass, \ + KeyInputItItT, \ + KeyOutputItItT, \ + SegmentSizeParameterT, \ + KParameterT, \ + SelectDirectionParameterT, \ + NumSegmentsParameterT>; \ + if (const auto error = CubDebug( \ + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ + static_cast(grid_blocks), ThreadsPerBlock, dynamic_smem_bytes, stream) \ + .doit(static_kernel, \ + d_key_segments_it, \ + d_key_segments_out_it, \ + segment_sizes, \ + k_param, \ + select_directions, \ + num_segments, \ + tile_capacity))) \ + { \ + return error; \ } #endif // CUB_RDC_ENABLED @@ -275,19 +324,21 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // CCs, so pin the selector query to the minimum supported CC. constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); constexpr int ThreadsPerBlock = policy.threads_per_block; - constexpr int ItemsPerThread = policy.items_per_thread; + constexpr int UnrollFactor = policy.unroll_factor; + constexpr int PipelineStages = policy.pipeline_stages; + constexpr int ChunkBytes = policy.chunk_bytes; + constexpr int LoadAlignBytes = policy.load_align_bytes; constexpr int BitsPerPass = policy.bits_per_pass; - constexpr int items_per_block = ThreadsPerBlock * ItemsPerThread; - constexpr int max_cluster_tile = max_supported_cluster_blocks * items_per_block; + using key_it_t = it_value_t; + using key_t = it_value_t; + using layout_t = smem_tile_layout; - // Only static-size shapes can be rejected at compile time; runtime shapes - // are re-checked below against the chosen tile. - static_assert(!detail::params::is_static_param_v - || detail::params::static_max_value_v <= max_cluster_tile, - "Static segment size exceeds max_supported_cluster_blocks * ThreadsPerBlock * ItemsPerThread."); + static_assert(ChunkBytes % LoadAlignBytes == 0); + static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); const auto max_seg_size = runtime_max_segment_size(segment_sizes); + using max_seg_size_t = decltype(+max_seg_size); // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. size_t allocation_sizes[1] = {1}; @@ -332,7 +383,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr auto dynamic_kernel = &device_segmented_topk_cluster_kernel< ThreadsPerBlock, - ItemsPerThread, + UnrollFactor, + PipelineStages, + ChunkBytes, + LoadAlignBytes, BitsPerPass, KeyInputItItT, KeyOutputItItT, @@ -371,25 +425,158 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( cfg.attrs = &cluster_attr; cfg.numAttrs = 1; - int hw_max_cluster_blocks = 0; - if (const auto error = CubDebug(cudaOccupancyMaxPotentialClusterSize( - &hw_max_cluster_blocks, reinterpret_cast(dynamic_kernel), &cfg))) + int max_dynamic_smem_bytes = 0; + if (const auto error = CubDebug(MaxPotentialDynamicSmemBytes(max_dynamic_smem_bytes, dynamic_kernel))) { return error; } - hw_max_cluster_blocks = ::cuda::std::min(hw_max_cluster_blocks, max_supported_cluster_blocks); - const int cluster_blocks = select_cluster_blocks(max_seg_size, items_per_block, hw_max_cluster_blocks); + constexpr int portable_dynamic_smem_bytes = 48 * 1024; + int configured_dynamic_smem_limit = portable_dynamic_smem_bytes; + const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { + if (dynamic_smem_bytes <= configured_dynamic_smem_limit) + { + return cudaSuccess; + } - // Exact bounds must fit the chosen tile; `per_segment_param` may carry - // a loose upper bound and is left to the agent's contract. - if constexpr (!detail::params::is_per_segment_param_v) + if (const auto error = CubDebug(cudaFuncSetAttribute( + reinterpret_cast(dynamic_kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + dynamic_smem_bytes))) + { + return error; + } + configured_dynamic_smem_limit = dynamic_smem_bytes; + return cudaSuccess; + }; + + const int smem_96k = (::cuda::std::min) (96 * 1024, max_dynamic_smem_bytes); + const int smem_160k = (::cuda::std::min) (160 * 1024, max_dynamic_smem_bytes); + const int optin_dynamic_smem_bytes = max_dynamic_smem_bytes; + + constexpr int num_candidate_configs = 8; + const int candidate_clusters[num_candidate_configs] = { + max_portable_cluster_blocks, + max_supported_cluster_blocks, + max_portable_cluster_blocks, + max_supported_cluster_blocks, + max_portable_cluster_blocks, + max_supported_cluster_blocks, + max_portable_cluster_blocks, + max_supported_cluster_blocks}; + const int candidate_smem[num_candidate_configs] = { + portable_dynamic_smem_bytes, + portable_dynamic_smem_bytes, + smem_96k, + smem_96k, + smem_160k, + smem_160k, + optin_dynamic_smem_bytes, + optin_dynamic_smem_bytes}; + + launch_config selected_config{0, 0, 0}; + launch_config largest_supported_config{0, 0, 0}; + // The table is intentionally tiny; scan all entries and select the lowest coverage that satisfies the bound + // rather than relying on the declaration order to be sorted by coverage. + for (int i = 0; i < num_candidate_configs; ++i) + { + const int candidate_cluster_blocks = candidate_clusters[i]; + const int candidate_dynamic_smem = candidate_smem[i]; + if (candidate_dynamic_smem > max_dynamic_smem_bytes) + { + continue; + } + + const auto candidate_tile_capacity = layout_t::tile_capacity(candidate_dynamic_smem); + if (candidate_tile_capacity == 0) + { + continue; + } + + launch_config candidate{candidate_cluster_blocks, candidate_dynamic_smem, candidate_tile_capacity}; + const auto candidate_coverage = + layout_t::template cluster_coverage(candidate.cluster_blocks, candidate.tile_capacity); + + bool candidate_can_improve_selection = false; + if constexpr (detail::params::is_per_segment_param_v) + { + // Runtime-sized segments may exceed every finite candidate, so find the largest supported fallback. + candidate_can_improve_selection = true; + } + else if (candidate_coverage >= max_seg_size) + { + const auto selected_coverage = layout_t::template cluster_coverage( + selected_config.cluster_blocks, selected_config.tile_capacity); + candidate_can_improve_selection = + selected_config.cluster_blocks == 0 || candidate_coverage < selected_coverage; + } + if (!candidate_can_improve_selection) + { + continue; + } + + if (const auto error = ensure_dynamic_smem_limit(candidate_dynamic_smem)) + { + return error; + } + + cfg.dynamicSmemBytes = static_cast(candidate_dynamic_smem); + int candidate_hw_max_cluster_blocks = 0; + if (const auto error = CubDebug(cudaOccupancyMaxPotentialClusterSize( + &candidate_hw_max_cluster_blocks, reinterpret_cast(dynamic_kernel), &cfg))) + { + return error; + } + candidate_hw_max_cluster_blocks = + (::cuda::std::min) (candidate_hw_max_cluster_blocks, max_supported_cluster_blocks); + if (candidate_cluster_blocks > candidate_hw_max_cluster_blocks) + { + continue; + } + + if constexpr (detail::params::is_per_segment_param_v) + { + const auto largest_supported_coverage = layout_t::template cluster_coverage( + largest_supported_config.cluster_blocks, largest_supported_config.tile_capacity); + if (largest_supported_config.cluster_blocks == 0 || candidate_coverage > largest_supported_coverage) + { + largest_supported_config = candidate; + } + } + + if (candidate_coverage >= max_seg_size) + { + selected_config = candidate; + } + } + + if (selected_config.cluster_blocks == 0) { - if (max_seg_size > static_cast(cluster_blocks) * items_per_block) + if constexpr (detail::params::is_per_segment_param_v) + { + selected_config = largest_supported_config; + } + else { return cudaErrorInvalidValue; } } + if (selected_config.cluster_blocks == 0) + { + return cudaErrorInvalidValue; + } + + if (const auto error = CubDebug(cudaFuncSetAttribute( + reinterpret_cast(dynamic_kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + selected_config.dynamic_smem_bytes))) + { + return error; + } + + const int cluster_blocks = selected_config.cluster_blocks; + const auto tile_capacity = selected_config.tile_capacity; + cfg.dynamicSmemBytes = static_cast(selected_config.dynamic_smem_bytes); const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); @@ -409,17 +596,22 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( segment_sizes, k_param, select_directions, - num_segments))) + num_segments, + tile_capacity))) { return error; } }), ({ - // CDP path: same exact-vs-loose rule against the fixed static tile. - constexpr int static_cluster_tile = max_portable_cluster_blocks * items_per_block; + // CDP path: device-side launches cannot opt in to more than portable + // dynamic SMEM or non-portable cluster widths. + constexpr int dynamic_smem_bytes = 48 * 1024; + constexpr auto tile_capacity = layout_t::tile_capacity(dynamic_smem_bytes); + constexpr auto static_cluster_tile = + layout_t::template cluster_coverage(max_portable_cluster_blocks, tile_capacity); if constexpr (!detail::params::is_per_segment_param_v) { - if (max_seg_size > static_cast(static_cluster_tile)) + if (max_seg_size > static_cast(static_cluster_tile)) { return cudaErrorInvalidValue; } diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 423d40c5598..6916e40925f 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -18,23 +18,32 @@ CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { -// Per-block tile shape; the dispatch picks `cluster_blocks` at runtime. +// Per-block execution shape; the dispatch picks `cluster_blocks` and the +// dynamic shared-memory tile capacity at runtime. struct cluster_topk_policy { int threads_per_block; - int items_per_thread; + int unroll_factor; + int pipeline_stages; + int chunk_bytes; + int load_align_bytes; int bits_per_pass; }; inline constexpr cluster_topk_policy default_policy{ /*threads_per_block=*/256, - /*items_per_thread=*/16, + /*unroll_factor=*/8, + /*pipeline_stages=*/3, + /*chunk_bytes=*/8 * 1024, + /*load_align_bytes=*/128, /*bits_per_pass=*/8, }; +static_assert(default_policy.chunk_bytes % default_policy.load_align_bytes == 0); // Default selector: one policy for every cluster-capable architecture -// (SM 9.0+). New tunings can be wired in by passing a custom selector to -// `dispatch` without changing the kernel signatures. +// (SM 9.0+). `unroll_factor` only controls ILP in runtime-sized chunk loops; +// it no longer defines the tile size. New tunings can be wired in by passing a +// custom selector to `dispatch` without changing the kernel signatures. struct policy_selector { [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 3560b133d50..95218c5c851 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -197,6 +198,57 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", REQUIRE(expected_keys == keys_out_buffer); } +#if TEST_LAUNCH != 1 +TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr segment_size_t static_max_segment_size = 128 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 3; + + const auto direction = cub::detail::topk::select::max; + const int pad = GENERATE(0, 1, 3, 7); + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 1, static_max_segment_size - 31})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + + CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(pad + num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(num_segments * segment_size, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cub::detail::batched_topk::segment_size_uniform<1, static_max_segment_size>{segment_size}, + cub::detail::batched_topk::k_uniform<1, static_max_k>{k}, + cub::detail::batched_topk::select_direction_uniform{direction}, + cub::detail::batched_topk::num_segments_uniform<>{num_segments}, + cub::detail::batched_topk::total_num_items_guarantee{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_LAUNCH != 1 + C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segments", "[keys][segmented][topk][device]", key_types, @@ -298,6 +350,79 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } +#if TEST_LAUNCH != 1 +TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr segment_size_t static_max_segment_size = 128 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + + const auto direction = cub::detail::topk::select::max; + const int pad = GENERATE(1, 3, 7); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); + + c2h::host_vector h_segment_offsets{ + 0, + static_max_segment_size, + static_max_segment_size + (static_max_segment_size - 31), + static_max_segment_size + (static_max_segment_size - 31) + (96 * 1024 + 17), + static_max_segment_size + (static_max_segment_size - 31) + (96 * 1024 + 17) + 257}; + c2h::device_vector segment_offsets = h_segment_offsets; + const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); + const segment_size_t num_items = h_segment_offsets.back(); + + auto segment_offsets_it = thrust::raw_pointer_cast(segment_offsets.data()); + auto segment_size_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), segment_size_op{segment_offsets_it}); + + CAPTURE(pad, static_max_segment_size, static_max_k, k, num_segments, num_items, direction); + + auto compacted_output_sizes_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), + get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k)}); + c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); + thrust::exclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + segment_size_t total_output_size = compacted_offsets.back(); + + c2h::device_vector keys_in_buffer(pad + num_items, thrust::no_init); + c2h::device_vector keys_out_buffer(total_output_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_offsets.cbegin()); + auto d_keys_out = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_out_ptr), compacted_offsets.cbegin()); + + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cub::detail::batched_topk::segment_size_per_segment{ + segment_size_it}, + cub::detail::batched_topk::k_uniform<1, static_max_k>{k}, + cub::detail::batched_topk::select_direction_uniform{direction}, + cub::detail::batched_topk::num_segments_uniform<>{num_segments}, + cub::detail::batched_topk::total_num_items_guarantee{num_items}); + + segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); + expected_keys = compact_to_topk_batched(expected_keys, segment_offsets, k); + + segmented_sort_keys( + keys_out_buffer, num_segments, compacted_offsets.cbegin(), compacted_offsets.cbegin() + 1, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_LAUNCH != 1 + // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segmented][topk][device][float]") { diff --git a/cub/test/catch2_test_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_segmented_topk_cluster_layout.cu new file mode 100644 index 00000000000..b53c3fa91fe --- /dev/null +++ b/cub/test/catch2_test_segmented_topk_cluster_layout.cu @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include + +#include + +#include + +namespace +{ +template +[[nodiscard]] constexpr SizeT host_ceil_div(SizeT numerator, SizeT denominator) +{ + return (numerator + denominator - 1) / denominator; +} + +template +void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) +{ + using layout_t = cub::detail::batched_topk_cluster::smem_tile_layout; + + const int usable_bytes = dynamic_smem_bytes - layout_t::base_padding_bytes; + REQUIRE(usable_bytes > 0); + + const int slots = usable_bytes / layout_t::slot_stride_bytes; + REQUIRE(slots > 0); + + const auto tile_capacity = layout_t::tile_capacity(dynamic_smem_bytes); + REQUIRE(tile_capacity == static_cast(slots * layout_t::chunk_items)); + + const auto coverage = layout_t::template cluster_coverage(cluster_blocks, tile_capacity); + REQUIRE(coverage > 0); + + const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { + using size_t = cuda::std::int64_t; + const size_t head_chunks = head_items == 0 ? size_t{0} : size_t{1}; + const size_t tail_items = segment_size - head_items; + const size_t chunks = head_chunks + host_ceil_div(tail_items, size_t{layout_t::chunk_items}); + return host_ceil_div(chunks, static_cast(blocks)); + }; + + const int heads[] = {0, 1, layout_t::chunk_items / 2, layout_t::chunk_items - 1}; + for (const int head_items : heads) + { + CAPTURE(c2h::type_name(), + ChunkBytes, + LoadAlignBytes, + dynamic_smem_bytes, + cluster_blocks, + slots, + tile_capacity, + coverage, + head_items); + REQUIRE(max_rank_chunks(coverage, head_items, cluster_blocks) <= slots); + } + + const auto unreserved_coverage = static_cast(cluster_blocks) * tile_capacity; + CAPTURE(c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, tile_capacity); + REQUIRE(max_rank_chunks(unreserved_coverage, 1, cluster_blocks) == slots + 1); +} + +template +void check_layout_matrix() +{ + constexpr int dynamic_smem_cases[] = {48 * 1024, 96 * 1024, 160 * 1024, 228 * 1024}; + constexpr int cluster_block_cases[] = {8, 16}; + + for (const int dynamic_smem_bytes : dynamic_smem_cases) + { + for (const int cluster_blocks : cluster_block_cases) + { + check_layout_case(dynamic_smem_bytes, cluster_blocks); + } + } +} +} // namespace + +TEST_CASE("Segmented TopK cluster SMEM layout reserves the unaligned head chunk", + "[keys][segmented][topk][cluster][layout]") +{ + using default_policy = cub::detail::batched_topk_cluster::policy_selector; + constexpr auto policy = default_policy{}(cuda::compute_capability{9, 0}); + + using default_float_layout = + cub::detail::batched_topk_cluster::smem_tile_layout; + static_assert(default_float_layout::tile_capacity(0) == 0); + static_assert(default_float_layout::template cluster_coverage(8, 0) == 0); + static_assert(default_float_layout::template cluster_coverage(1, default_float_layout::chunk_items) == 0); + + check_layout_matrix(); + check_layout_matrix(); + check_layout_matrix(); + + check_layout_matrix(); + check_layout_matrix(); + check_layout_matrix(); +} From 082c3dfddc17b4390465c14c2c66ce2a2a76287d Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 3 Jun 2026 00:18:19 +0200 Subject: [PATCH 019/126] Improve dispatch and naming --- cub/cub/agent/agent_batched_topk_cluster.cuh | 75 ++++---- .../dispatch_batched_topk_cluster.cuh | 166 +++++++++++------- .../tuning/tuning_batched_topk_cluster.cuh | 116 ++++++++++-- cub/cub/util_device.cuh | 2 + ...tch2_test_segmented_topk_cluster_layout.cu | 32 ++-- 5 files changed, 262 insertions(+), 129 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 1a9c7279e22..7ab40a6ffcf 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -103,10 +103,11 @@ struct alignas(16) cluster_topk_state ::cuda::std::uint32_t early_stop; }; -// Dynamic-SMEM layout shared by dispatch and the agent. `tile_capacity` is the physical per-CTA resident capacity -// passed to the kernel; `cluster_coverage` reserves one chunk of logical coverage for a possible unaligned head chunk. +// Dynamic-SMEM layout shared by dispatch and the agent. `block_tile_capacity` is the physical per-CTA +// resident capacity passed to the kernel; `cluster_tile_capacity` reserves one chunk of logical coverage +// for a possible unaligned head chunk. template -struct smem_tile_layout +struct smem_block_tile_layout { static constexpr int chunk_items = ChunkBytes / int{sizeof(KeyT)}; static constexpr int slot_alignment = @@ -115,24 +116,24 @@ struct smem_tile_layout ::cuda::round_up(detail::LoadToSharedBufferSizeBytes(chunk_items) + slot_alignment, slot_alignment); static constexpr int base_padding_bytes = (alignof(KeyT) > 16) ? slot_alignment : 0; - [[nodiscard]] _CCCL_HOST_DEVICE static constexpr ::cuda::std::uint32_t tile_capacity(int dynamic_smem_bytes) noexcept + [[nodiscard]] _CCCL_HOST_DEVICE static constexpr ::cuda::std::uint32_t + block_tile_capacity(int dynamic_smem_bytes) noexcept { - const int usable_bytes = dynamic_smem_bytes - base_padding_bytes; - if (usable_bytes <= 0) - { - return 0; - } - const int slots = usable_bytes / slot_stride_bytes; + const int usable_bytes = (::cuda::std::max) (0, dynamic_smem_bytes - base_padding_bytes); + const int slots = usable_bytes / slot_stride_bytes; return static_cast<::cuda::std::uint32_t>(slots * chunk_items); } template [[nodiscard]] _CCCL_HOST_DEVICE static constexpr SizeT - cluster_coverage(int cluster_blocks, ::cuda::std::uint32_t physical_tile_capacity) noexcept + cluster_tile_capacity(int cluster_blocks, ::cuda::std::uint32_t physical_block_tile_capacity) noexcept { - const auto physical_coverage = static_cast(cluster_blocks) * static_cast(physical_tile_capacity); + const auto physical_cluster_tile_items = + static_cast(cluster_blocks) * static_cast(physical_block_tile_capacity); const auto head_chunk_reserve = static_cast(chunk_items); - return (physical_coverage > head_chunk_reserve) ? physical_coverage - head_chunk_reserve : SizeT{0}; + return (physical_cluster_tile_items > head_chunk_reserve) + ? physical_cluster_tile_items - head_chunk_reserve + : SizeT{0}; } }; @@ -140,7 +141,7 @@ struct smem_tile_layout // Cluster top-k agent // ----------------------------------------------------------------------------- // Cluster width is a runtime value (see `process_impl` for the readback), so -// it is not a template parameter; per-block tile layout is still controlled +// it is not a template parameter; per-block block_tile layout is still controlled // by the template parameters below. template ; + using smem_layout_t = smem_block_tile_layout; static constexpr int chunk_items = smem_layout_t::chunk_items; static constexpr int slot_alignment = smem_layout_t::slot_alignment; static constexpr int slot_stride_bytes = smem_layout_t::slot_stride_bytes; @@ -232,9 +233,9 @@ struct agent_batched_topk_cluster #endif }; - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span key_tile_buffer() const + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const { - const int slots = static_cast(tile_capacity / chunk_items); + const int slots = static_cast(block_tile_capacity / chunk_items); return {key_slots, static_cast<::cuda::std::size_t>(slots * slot_stride_bytes)}; } @@ -244,13 +245,13 @@ struct agent_batched_topk_cluster } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span - available_key_tile_buffer(char* buffer_begin) const + available_block_tile_buffer(char* buffer_begin) const { - const auto buffer = key_tile_buffer(); + const auto buffer = block_tile_buffer(); char* const end = ::cuda::std::data(buffer) + static_cast(::cuda::std::size(buffer)); - _CCCL_ASSERT(buffer_begin >= ::cuda::std::data(buffer) && buffer_begin <= end, "Invalid key tile buffer cursor"); + _CCCL_ASSERT(buffer_begin >= ::cuda::std::data(buffer) && buffer_begin <= end, "Invalid block_tile buffer cursor"); _CCCL_ASSERT(::cuda::is_aligned(buffer_begin, detail::LoadToSharedBufferAlignBytes()), - "Key tile buffer cursor must satisfy BlockLoadToShared's shared-memory alignment"); + "block_tile buffer cursor must satisfy BlockLoadToShared's shared-memory alignment"); return {buffer_begin, static_cast<::cuda::std::size_t>(end - buffer_begin)}; } @@ -377,7 +378,7 @@ struct agent_batched_topk_cluster SelectDirectionParameterT select_directions; NumSegmentsParameterT num_segments; char* key_slots; - offset_t tile_capacity; + offset_t block_tile_capacity; _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -388,7 +389,7 @@ struct agent_batched_topk_cluster SelectDirectionParameterT select_directions_, NumSegmentsParameterT num_segments_, char* key_slots_, - offset_t tile_capacity_) + offset_t block_tile_capacity_) : temp_storage(temp_storage_.Alias()) , d_key_segments_it(d_key_segments_it_) , d_key_segments_out_it(d_key_segments_out_it_) @@ -397,7 +398,7 @@ struct agent_batched_topk_cluster , select_directions(select_directions_) , num_segments(num_segments_) , key_slots(key_slots_) - , tile_capacity(tile_capacity_) + , block_tile_capacity(block_tile_capacity_) {} // --------------------------------------------------------------------------- @@ -559,7 +560,8 @@ private: const offset_t my_chunks = num_rank_chunks(chunks, cluster_rank, cluster_size); // The launch coverage check reserves one extra chunk for the possible unaligned head, so every local chunk remains // resident for later radix passes while only the BlockLoadToShared instances are reused round-robin. - _CCCL_ASSERT(my_chunks * offset_t{chunk_items} <= tile_capacity, "Dynamic shared memory tile is too small"); + _CCCL_ASSERT(my_chunks * offset_t{chunk_items} <= block_tile_capacity, + "Dynamic shared memory block_tile is too small"); ::cuda::std::span resident_keys; @@ -601,11 +603,11 @@ private: if (aligned_chunk) { pending_spans.emplace_back( - loader.template CopyAsync(available_key_tile_buffer(next_dst), src)); + loader.template CopyAsync(available_block_tile_buffer(next_dst), src)); } else { - pending_spans.emplace_back(loader.template CopyAsync(available_key_tile_buffer(next_dst), src)); + pending_spans.emplace_back(loader.template CopyAsync(available_block_tile_buffer(next_dst), src)); } next_dst = span_end(pending_spans[stage]); tokens.emplace_back(loader.Commit()); @@ -631,11 +633,12 @@ private: if (next_aligned_chunk) { pending_spans[stage] = - loaders[stage].template CopyAsync(available_key_tile_buffer(next_dst), src); + loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); } else { - pending_spans[stage] = loaders[stage].template CopyAsync(available_key_tile_buffer(next_dst), src); + pending_spans[stage] = + loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); } next_dst = span_end(pending_spans[stage]); tokens[stage] = loaders[stage].Commit(); @@ -663,8 +666,8 @@ private: } } const int resident_count = span_size(resident_keys); - _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= tile_capacity, - "Dynamic shared memory tile is too small"); + _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, + "Dynamic shared memory block_tile is too small"); __syncthreads(); } @@ -995,7 +998,7 @@ private: // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast - // block (e.g. one whose tile is entirely padding) can return while another + // block (e.g. one whose block_tile is entirely padding) can return while another // block is still writing to leader-resident memory through DSMEM, which // surfaces as a "cluster target block not present" exception. cluster.sync(); @@ -1031,11 +1034,11 @@ private: segment_fits_offset = segment_size <= static_cast(::cuda::std::numeric_limits::max()); } - const auto max_cluster_coverage = - smem_layout_t::template cluster_coverage(static_cast(cluster_blocks), tile_capacity); - if (!segment_fits_offset || segment_size > max_cluster_coverage) + const auto max_cluster_tile_capacity = smem_layout_t::template cluster_tile_capacity( + static_cast(cluster_blocks), block_tile_capacity); + if (!segment_fits_offset || segment_size > max_cluster_tile_capacity) { - _CCCL_ASSERT(false, "Segment exceeds the selected cluster top-k tile capacity"); + _CCCL_ASSERT(false, "Segment exceeds the selected cluster top-k cluster_tile capacity"); return; } diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 2165cea4cb2..28b3bb01c3b 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -8,7 +8,7 @@ //! segmented top-k. Each cluster processes one segment end-to-end: private //! histograms are reduced into the leader block via DSMEM atomics, then every //! block reads the merged histogram back through DSMEM, locally identifies the -//! k-th bucket, and refines its shared-memory key tile across radix passes. +//! k-th bucket, and refines its shared-memory cluster_tile across radix passes. //! //! Two kernels share the agent body: //! * Host: no `__cluster_dims__`, launched via `cudaLaunchKernelExC` with the @@ -89,7 +89,7 @@ struct launch_config { int cluster_blocks; int dynamic_smem_bytes; - ::cuda::std::uint32_t tile_capacity; + ::cuda::std::uint32_t block_tile_capacity; }; // ----------------------------------------------------------------------------- @@ -116,7 +116,7 @@ __launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented KParameterT k_param, SelectDirectionParameterT select_directions, NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t tile_capacity) + ::cuda::std::uint32_t block_tile_capacity) { using agent_t = agent_batched_topk_cluster< ThreadsPerBlock, @@ -154,7 +154,7 @@ __launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented select_directions, num_segments, key_slots, - tile_capacity); + block_tile_capacity); agent.Process(); } @@ -182,7 +182,7 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, KParameterT k_param, SelectDirectionParameterT select_directions, NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t tile_capacity) + ::cuda::std::uint32_t block_tile_capacity) { using agent_t = agent_batched_topk_cluster< ThreadsPerBlock, @@ -220,7 +220,7 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, select_directions, num_segments, key_slots, - tile_capacity); + block_tile_capacity); agent.Process(); } @@ -260,7 +260,7 @@ struct force_emit_kernel // ----------------------------------------------------------------------------- // Dispatch // ----------------------------------------------------------------------------- -// Keys-only; every segment must fit in one cluster tile. Host picks +// Keys-only; every segment must fit in one cluster_tile. Host picks // `(cluster_blocks, dynamic_smem_bytes)` at runtime from a finite table; CDP // uses the static kernel at `max_portable_cluster_blocks` and portable SMEM. @@ -293,7 +293,7 @@ struct force_emit_kernel k_param, \ select_directions, \ num_segments, \ - tile_capacity))) \ + block_tile_capacity))) \ { \ return error; \ } @@ -332,10 +332,24 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( using key_it_t = it_value_t; using key_t = it_value_t; - using layout_t = smem_tile_layout; + using layout_t = smem_block_tile_layout; + using agent_t = agent_batched_topk_cluster< + ThreadsPerBlock, + UnrollFactor, + PipelineStages, + ChunkBytes, + LoadAlignBytes, + BitsPerPass, + KeyInputItItT, + KeyOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; static_assert(ChunkBytes % LoadAlignBytes == 0); static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + constexpr int static_smem_bytes = static_cast(sizeof(typename agent_t::TempStorage)); const auto max_seg_size = runtime_max_segment_size(segment_sizes); using max_seg_size_t = decltype(+max_seg_size); @@ -431,9 +445,39 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return error; } - constexpr int portable_dynamic_smem_bytes = 48 * 1024; - int configured_dynamic_smem_limit = portable_dynamic_smem_bytes; - const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { + // The table entries express the documented total per-block shared-memory budget + // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`). The usable dynamic portion is that budget minus the + // kernel's static shared memory and the CUDA driver's per-block reserved shared memory, matching + // `MaxPotentialDynamicSmemBytes`. Subtracting the reserved bytes here keeps a documented entry (e.g. 99 KiB) + // from deriving a dynamic size that overshoots `max_dynamic_smem_bytes` and being discarded below. + int device_id = 0; + if (const auto error = CubDebug(cudaGetDevice(&device_id))) + { + return error; + } + int reserved_smem_bytes = 0; + if (const auto error = + CubDebug(cudaDeviceGetAttribute(&reserved_smem_bytes, cudaDevAttrReservedSharedMemoryPerBlock, device_id))) + { + return error; + } + // TODO: subtracting `reserved` mirrors `MaxPotentialDynamicSmemBytes`, but that helper double-counts it + // (opt-in already excludes reserved), so this is ~`reserved` bytes too conservative. Revisit once the helper is + // fixed. + const int nondynamic_smem_bytes = static_smem_bytes + reserved_smem_bytes; + + const auto runtime_policy = PolicySelector{}(::cuda::compute_capability{sm_version / 10}); + _CCCL_ASSERT(runtime_policy.launch_configs.size() <= max_launch_configs, + "Cluster TopK launch config table exceeds policy capacity"); + + const auto total_to_dynamic_smem = [&](int total_smem_bytes) { + return (total_smem_bytes > nondynamic_smem_bytes) ? total_smem_bytes - nondynamic_smem_bytes : 0; + }; + + constexpr int portable_total_smem_bytes = 48 * 1024; + const int portable_dynamic_smem_bytes = total_to_dynamic_smem(portable_total_smem_bytes); + int configured_dynamic_smem_limit = portable_dynamic_smem_bytes; + const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { if (dynamic_smem_bytes <= configured_dynamic_smem_limit) { return cudaSuccess; @@ -450,52 +494,28 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaSuccess; }; - const int smem_96k = (::cuda::std::min) (96 * 1024, max_dynamic_smem_bytes); - const int smem_160k = (::cuda::std::min) (160 * 1024, max_dynamic_smem_bytes); - const int optin_dynamic_smem_bytes = max_dynamic_smem_bytes; - - constexpr int num_candidate_configs = 8; - const int candidate_clusters[num_candidate_configs] = { - max_portable_cluster_blocks, - max_supported_cluster_blocks, - max_portable_cluster_blocks, - max_supported_cluster_blocks, - max_portable_cluster_blocks, - max_supported_cluster_blocks, - max_portable_cluster_blocks, - max_supported_cluster_blocks}; - const int candidate_smem[num_candidate_configs] = { - portable_dynamic_smem_bytes, - portable_dynamic_smem_bytes, - smem_96k, - smem_96k, - smem_160k, - smem_160k, - optin_dynamic_smem_bytes, - optin_dynamic_smem_bytes}; - launch_config selected_config{0, 0, 0}; launch_config largest_supported_config{0, 0, 0}; // The table is intentionally tiny; scan all entries and select the lowest coverage that satisfies the bound // rather than relying on the declaration order to be sorted by coverage. - for (int i = 0; i < num_candidate_configs; ++i) + for (const auto policy_config : runtime_policy.launch_configs) { - const int candidate_cluster_blocks = candidate_clusters[i]; - const int candidate_dynamic_smem = candidate_smem[i]; + const int candidate_cluster_blocks = policy_config.cluster_blocks; + const int candidate_dynamic_smem = total_to_dynamic_smem(policy_config.total_smem_bytes); if (candidate_dynamic_smem > max_dynamic_smem_bytes) { continue; } - const auto candidate_tile_capacity = layout_t::tile_capacity(candidate_dynamic_smem); - if (candidate_tile_capacity == 0) + const auto candidate_block_tile_capacity = layout_t::block_tile_capacity(candidate_dynamic_smem); + if (candidate_block_tile_capacity == 0) { continue; } - launch_config candidate{candidate_cluster_blocks, candidate_dynamic_smem, candidate_tile_capacity}; - const auto candidate_coverage = - layout_t::template cluster_coverage(candidate.cluster_blocks, candidate.tile_capacity); + launch_config candidate{candidate_cluster_blocks, candidate_dynamic_smem, candidate_block_tile_capacity}; + const auto candidate_cluster_tile_capacity = layout_t::template cluster_tile_capacity( + candidate.cluster_blocks, candidate.block_tile_capacity); bool candidate_can_improve_selection = false; if constexpr (detail::params::is_per_segment_param_v) @@ -503,12 +523,12 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // Runtime-sized segments may exceed every finite candidate, so find the largest supported fallback. candidate_can_improve_selection = true; } - else if (candidate_coverage >= max_seg_size) + else if (candidate_cluster_tile_capacity >= max_seg_size) { - const auto selected_coverage = layout_t::template cluster_coverage( - selected_config.cluster_blocks, selected_config.tile_capacity); + const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( + selected_config.cluster_blocks, selected_config.block_tile_capacity); candidate_can_improve_selection = - selected_config.cluster_blocks == 0 || candidate_coverage < selected_coverage; + selected_config.cluster_blocks == 0 || candidate_cluster_tile_capacity < selected_cluster_tile_capacity; } if (!candidate_can_improve_selection) { @@ -536,15 +556,16 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( if constexpr (detail::params::is_per_segment_param_v) { - const auto largest_supported_coverage = layout_t::template cluster_coverage( - largest_supported_config.cluster_blocks, largest_supported_config.tile_capacity); - if (largest_supported_config.cluster_blocks == 0 || candidate_coverage > largest_supported_coverage) + const auto largest_supported_cluster_tile_capacity = layout_t::template cluster_tile_capacity( + largest_supported_config.cluster_blocks, largest_supported_config.block_tile_capacity); + if (largest_supported_config.cluster_blocks == 0 + || candidate_cluster_tile_capacity > largest_supported_cluster_tile_capacity) { largest_supported_config = candidate; } } - if (candidate_coverage >= max_seg_size) + if (candidate_cluster_tile_capacity >= max_seg_size) { selected_config = candidate; } @@ -566,6 +587,25 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaErrorInvalidValue; } + const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( + selected_config.cluster_blocks, selected_config.block_tile_capacity); + if (selected_cluster_tile_capacity >= max_seg_size) + { + const auto required_physical_cluster_tile_items = + static_cast<::cuda::std::uint64_t>(max_seg_size) + static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); + selected_config.cluster_blocks = static_cast( + ::cuda::ceil_div(required_physical_cluster_tile_items, + static_cast<::cuda::std::uint64_t>(selected_config.block_tile_capacity))); + + const auto required_block_tile_capacity = ::cuda::ceil_div( + required_physical_cluster_tile_items, static_cast<::cuda::std::uint64_t>(selected_config.cluster_blocks)); + const auto required_slots = + ::cuda::ceil_div(required_block_tile_capacity, static_cast<::cuda::std::uint64_t>(layout_t::chunk_items)); + selected_config.dynamic_smem_bytes = + layout_t::base_padding_bytes + static_cast(required_slots) * layout_t::slot_stride_bytes; + selected_config.block_tile_capacity = layout_t::block_tile_capacity(selected_config.dynamic_smem_bytes); + } + if (const auto error = CubDebug(cudaFuncSetAttribute( reinterpret_cast(dynamic_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -574,9 +614,9 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return error; } - const int cluster_blocks = selected_config.cluster_blocks; - const auto tile_capacity = selected_config.tile_capacity; - cfg.dynamicSmemBytes = static_cast(selected_config.dynamic_smem_bytes); + const int cluster_blocks = selected_config.cluster_blocks; + const auto block_tile_capacity = selected_config.block_tile_capacity; + cfg.dynamicSmemBytes = static_cast(selected_config.dynamic_smem_bytes); const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); @@ -597,21 +637,23 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( k_param, select_directions, num_segments, - tile_capacity))) + block_tile_capacity))) { return error; } }), ({ // CDP path: device-side launches cannot opt in to more than portable - // dynamic SMEM or non-portable cluster widths. - constexpr int dynamic_smem_bytes = 48 * 1024; - constexpr auto tile_capacity = layout_t::tile_capacity(dynamic_smem_bytes); - constexpr auto static_cluster_tile = - layout_t::template cluster_coverage(max_portable_cluster_blocks, tile_capacity); + // total SMEM or non-portable cluster widths. + constexpr int portable_total_smem_bytes = 48 * 1024; + constexpr int dynamic_smem_bytes = + (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; + constexpr auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + constexpr auto portable_cluster_tile_capacity = + layout_t::template cluster_tile_capacity(max_portable_cluster_blocks, block_tile_capacity); if constexpr (!detail::params::is_per_segment_param_v) { - if (max_seg_size > static_cast(static_cluster_tile)) + if (max_seg_size > static_cast(portable_cluster_tile_capacity)) { return cudaErrorInvalidValue; } diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 6916e40925f..b80c7642c2c 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -14,12 +14,22 @@ #endif // no system header #include +#include CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { +inline constexpr int max_launch_configs = 19; + +struct cluster_topk_launch_config +{ + int cluster_blocks; + // Total per-block shared memory budget, including the agent's static shared memory. + int total_smem_bytes; +}; + // Per-block execution shape; the dispatch picks `cluster_blocks` and the -// dynamic shared-memory tile capacity at runtime. +// dynamic shared-memory block_tile capacity at runtime. struct cluster_topk_policy { int threads_per_block; @@ -28,27 +38,101 @@ struct cluster_topk_policy int chunk_bytes; int load_align_bytes; int bits_per_pass; + ::cuda::std::inplace_vector launch_configs; }; -inline constexpr cluster_topk_policy default_policy{ - /*threads_per_block=*/256, - /*unroll_factor=*/8, - /*pipeline_stages=*/3, - /*chunk_bytes=*/8 * 1024, - /*load_align_bytes=*/128, - /*bits_per_pass=*/8, -}; -static_assert(default_policy.chunk_bytes % default_policy.load_align_bytes == 0); +template +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_launch_configs(Configs... configs) + -> ::cuda::std::inplace_vector +{ + ::cuda::std::inplace_vector result; + (result.emplace_back(configs), ...); + return result; +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto +make_policy(::cuda::std::inplace_vector launch_configs) + -> cluster_topk_policy +{ + return cluster_topk_policy{ + /*threads_per_block=*/256, + /*unroll_factor=*/8, + /*pipeline_stages=*/3, + /*chunk_bytes=*/8 * 1024, + /*load_align_bytes=*/128, + /*bits_per_pass=*/8, + launch_configs}; +} -// Default selector: one policy for every cluster-capable architecture -// (SM 9.0+). `unroll_factor` only controls ILP in runtime-sized chunk loops; -// it no longer defines the tile size. New tunings can be wired in by passing a -// custom selector to `dispatch` without changing the kernel signatures. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto sm90_100_110_policy() -> cluster_topk_policy +{ + return make_policy(make_launch_configs( + cluster_topk_launch_config{2, 31 * 1024}, // 62 KiB + // cluster_topk_launch_config{4, 31 * 1024}, // 124 KiB + cluster_topk_launch_config{2, 63 * 1024}, // 126 KiB + cluster_topk_launch_config{2, 99 * 1024}, // 198 KiB + // cluster_topk_launch_config{8, 31 * 1024}, // 248 KiB + // cluster_topk_launch_config{4, 63 * 1024}, // 252 KiB + cluster_topk_launch_config{2, 131 * 1024}, // 262 KiB + cluster_topk_launch_config{2, 163 * 1024}, // 326 KiB + cluster_topk_launch_config{2, 195 * 1024}, // 390 KiB + // cluster_topk_launch_config{4, 99 * 1024}, // 396 KiB + cluster_topk_launch_config{2, 227 * 1024}, // 454 KiB + // cluster_topk_launch_config{8, 63 * 1024}, // 504 KiB + cluster_topk_launch_config{4, 131 * 1024}, // 524 KiB + cluster_topk_launch_config{4, 163 * 1024}, // 652 KiB + cluster_topk_launch_config{4, 195 * 1024}, // 780 KiB + // cluster_topk_launch_config{8, 99 * 1024}, // 792 KiB + cluster_topk_launch_config{4, 227 * 1024}, // 908 KiB + cluster_topk_launch_config{8, 131 * 1024}, // 1048 KiB + cluster_topk_launch_config{8, 163 * 1024}, // 1304 KiB + cluster_topk_launch_config{8, 195 * 1024}, // 1560 KiB + cluster_topk_launch_config{8, 227 * 1024}, // 1816 KiB + cluster_topk_launch_config{16, 131 * 1024}, // 2096 KiB + cluster_topk_launch_config{16, 163 * 1024}, // 2608 KiB + cluster_topk_launch_config{16, 195 * 1024}, // 3120 KiB + cluster_topk_launch_config{16, 227 * 1024} // 3632 KiB + )); +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto sm120_policy() -> cluster_topk_policy +{ + return make_policy(make_launch_configs( + cluster_topk_launch_config{2, 31 * 1024}, // 62 KiB + // cluster_topk_launch_config{4, 31 * 1024}, // 124 KiB + cluster_topk_launch_config{2, 63 * 1024}, // 126 KiB + cluster_topk_launch_config{2, 99 * 1024}, // 198 KiB + // cluster_topk_launch_config{8, 31 * 1024}, // 248 KiB + cluster_topk_launch_config{4, 63 * 1024}, // 252 KiB + cluster_topk_launch_config{4, 99 * 1024}, // 396 KiB + cluster_topk_launch_config{8, 63 * 1024}, // 504 KiB + cluster_topk_launch_config{8, 99 * 1024} // 792 KiB + )); +} + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool +{ + return policy.chunk_bytes % policy.load_align_bytes == 0 && policy.launch_configs.size() <= max_launch_configs; +} + +static_assert(is_valid_policy(sm90_100_110_policy())); +static_assert(is_valid_policy(sm120_policy())); +// Default selector for cluster-capable architectures (SM 9.0+). +// `unroll_factor` only controls ILP in runtime-sized chunk loops; it no longer +// defines the block_tile size. The launch config table is consumed at runtime, so +// selectors can vary it by architecture without forcing static dispatch over +// all cluster/SMEM combinations. struct policy_selector { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> cluster_topk_policy { - return default_policy; + if (cc >= ::cuda::compute_capability{12, 0}) + { + return sm120_policy(); + } + + return sm90_100_110_policy(); } }; } // namespace detail::batched_topk_cluster diff --git a/cub/cub/util_device.cuh b/cub/cub/util_device.cuh index 09caa064b2a..a629f77531b 100644 --- a/cub/cub/util_device.cuh +++ b/cub/cub/util_device.cuh @@ -496,6 +496,8 @@ MaxPotentialDynamicSmemBytes(int& max_dyn_smem_bytes, KernelPtr kernel_ptr) noex return error; } + // TODO: over-conservative by ~`reserved_smem_size`. `MaxSharedMemoryPerBlockOptin` already excludes reserved + // (opt_in == perSM - reserved), so subtracting reserved here double-counts it. max_dyn_smem_bytes = max_smem_size_optin - reserved_smem_size - static_cast(kernel_attrs.sharedSizeBytes); return cudaSuccess; } diff --git a/cub/test/catch2_test_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_segmented_topk_cluster_layout.cu index b53c3fa91fe..cb6da825ccc 100644 --- a/cub/test/catch2_test_segmented_topk_cluster_layout.cu +++ b/cub/test/catch2_test_segmented_topk_cluster_layout.cu @@ -18,7 +18,7 @@ template template void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) { - using layout_t = cub::detail::batched_topk_cluster::smem_tile_layout; + using layout_t = cub::detail::batched_topk_cluster::smem_block_tile_layout; const int usable_bytes = dynamic_smem_bytes - layout_t::base_padding_bytes; REQUIRE(usable_bytes > 0); @@ -26,11 +26,12 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) const int slots = usable_bytes / layout_t::slot_stride_bytes; REQUIRE(slots > 0); - const auto tile_capacity = layout_t::tile_capacity(dynamic_smem_bytes); - REQUIRE(tile_capacity == static_cast(slots * layout_t::chunk_items)); + const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + REQUIRE(block_tile_capacity == static_cast(slots * layout_t::chunk_items)); - const auto coverage = layout_t::template cluster_coverage(cluster_blocks, tile_capacity); - REQUIRE(coverage > 0); + const auto cluster_tile_capacity = + layout_t::template cluster_tile_capacity(cluster_blocks, block_tile_capacity); + REQUIRE(cluster_tile_capacity > 0); const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { using size_t = cuda::std::int64_t; @@ -49,15 +50,16 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) dynamic_smem_bytes, cluster_blocks, slots, - tile_capacity, - coverage, + block_tile_capacity, + cluster_tile_capacity, head_items); - REQUIRE(max_rank_chunks(coverage, head_items, cluster_blocks) <= slots); + REQUIRE(max_rank_chunks(cluster_tile_capacity, head_items, cluster_blocks) <= slots); } - const auto unreserved_coverage = static_cast(cluster_blocks) * tile_capacity; - CAPTURE(c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, tile_capacity); - REQUIRE(max_rank_chunks(unreserved_coverage, 1, cluster_blocks) == slots + 1); + const auto unreserved_cluster_tile_capacity = static_cast(cluster_blocks) * block_tile_capacity; + CAPTURE( + c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, block_tile_capacity); + REQUIRE(max_rank_chunks(unreserved_cluster_tile_capacity, 1, cluster_blocks) == slots + 1); } template @@ -83,10 +85,10 @@ TEST_CASE("Segmented TopK cluster SMEM layout reserves the unaligned head chunk" constexpr auto policy = default_policy{}(cuda::compute_capability{9, 0}); using default_float_layout = - cub::detail::batched_topk_cluster::smem_tile_layout; - static_assert(default_float_layout::tile_capacity(0) == 0); - static_assert(default_float_layout::template cluster_coverage(8, 0) == 0); - static_assert(default_float_layout::template cluster_coverage(1, default_float_layout::chunk_items) == 0); + cub::detail::batched_topk_cluster::smem_block_tile_layout; + static_assert(default_float_layout::block_tile_capacity(0) == 0); + static_assert(default_float_layout::template cluster_tile_capacity(8, 0) == 0); + static_assert(default_float_layout::template cluster_tile_capacity(1, default_float_layout::chunk_items) == 0); check_layout_matrix(); check_layout_matrix(); From 0b8111ddc2efb28fb7ba3cf1d34ec2347d5497d2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 3 Jun 2026 06:05:23 +0200 Subject: [PATCH 020/126] Handle very big segments by repeated streaming from gmem (L2) --- cub/cub/agent/agent_batched_topk_cluster.cuh | 293 ++++++++++++++++-- .../dispatch_batched_topk_cluster.cuh | 96 +++--- .../catch2_test_device_segmented_topk_keys.cu | 123 +++++++- ...catch2_test_device_segmented_topk_pairs.cu | 8 +- cub/test/catch2_test_device_topk_common.cuh | 7 +- 5 files changed, 426 insertions(+), 101 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 7ab40a6ffcf..31149e704ff 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -70,6 +70,7 @@ #include #include #include +#include #include #include @@ -508,6 +509,196 @@ private: } #endif + // --------------------------------------------------------------------------- + // Overflow streamer + // --------------------------------------------------------------------------- + // Re-streams the per-rank "overflow" chunks (those that do not fit in the + // resident SMEM region) from gmem through a small, fixed, round-robin set of + // `PipelineStages` streaming slots. The same object is reused for every radix + // pass and the final filter. It ping-pongs the iteration order across calls so + // the `PipelineStages` boundary chunks that one pass leaves resident in the + // streaming slots are reused by the next pass with no reload; in the limit + // where the overflow fits entirely in the streaming slots (`overflow <= + // PipelineStages`), the chunks are loaded once and never reloaded. The + // resident region is unaffected: it lives in the slots `[0, resident_slots)`, + // the streaming region in `[stream_slot_base, stream_slot_base + + // PipelineStages)`. + struct overflow_streamer + { + agent_batched_topk_cluster& agent; + key_it_t block_keys_in; + const key_t* block_keys_base; // unwrapped contiguous base (pipeline path only; null otherwise) + offset_t segment_size; + offset_t head_items; + unsigned int cluster_rank; + unsigned int cluster_size; + offset_t resident_chunks; // rank-local chunk index at which the overflow begins + int stream_slot_base; // SMEM slot index at which the streaming region begins + offset_t overflow_chunks; // number of overflow chunks for this rank (M) + int p_eff; // active streaming depth = min(PipelineStages, M) (>= 1) + bool forward = true; + bool primed = false; + + // Loaders are shared with the first-pass resident load: that load constructs + // them here (via `ensure_loaders`) and the streamer keeps reusing the very + // same objects for the overflow passes, continuing their Commit/Wait pipeline + // rather than tearing them down and re-initializing the shared `load_storage` + // mbarriers. During streaming, `loaders[stage]` targets the streaming slot + // `stream_slot_base + stage`. A `tokens` entry is engaged only while a copy is + // in flight (committed, not yet waited); after a wait the slot's data settles + // in SMEM and `pending`/`slot_chunk` keep describing it so the next ping-pong + // pass can reuse it without a reload. + ::cuda::std::inplace_vector loaders; + ::cuda::std::optional tokens[PipelineStages]; + ::cuda::std::span pending[PipelineStages]; + offset_t slot_chunk[PipelineStages]; + + _CCCL_DEVICE _CCCL_FORCEINLINE overflow_streamer( + agent_batched_topk_cluster& agent_, + key_it_t block_keys_in_, + const key_t* block_keys_base_, + offset_t segment_size_, + offset_t head_items_, + unsigned int cluster_rank_, + unsigned int cluster_size_, + offset_t resident_chunks_, + int stream_slot_base_, + offset_t my_chunks_) + : agent(agent_) + , block_keys_in(block_keys_in_) + , block_keys_base(block_keys_base_) + , segment_size(segment_size_) + , head_items(head_items_) + , cluster_rank(cluster_rank_) + , cluster_size(cluster_size_) + , resident_chunks(resident_chunks_) + , stream_slot_base(stream_slot_base_) + , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) + { + const int m = static_cast((::cuda::std::min) (overflow_chunks, static_cast(PipelineStages))); + p_eff = (m > 0) ? m : 1; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE offset_t chunk_index_of(offset_t overflow_idx) const + { + return cluster_rank + (resident_chunks + overflow_idx) * static_cast(cluster_size); + } + + // Grow the loader set to `count` entries (idempotent). The first-pass resident + // load calls this to construct the loaders, then keeps using them; the + // streamer reuses those same objects for the overflow passes. Each new loader + // binds to a distinct, previously-unused `load_storage` slot, so a live + // mbarrier is never re-initialized (which is why no `Invalidate()` is needed). + _CCCL_DEVICE _CCCL_FORCEINLINE void ensure_loaders(int count) + { + while (static_cast(loaders.size()) < count) + { + loaders.emplace_back(agent.temp_storage.load_storage[loaders.size()]); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) + { + const offset_t chunk_idx = chunk_index_of(overflow_idx); + const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); + const bool aligned_chunk = agent.is_aligned_chunk(block_keys_base, chunk); + char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; + const ::cuda::std::span src{ + block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; + if (aligned_chunk) + { + pending[stage] = + loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); + } + else + { + pending[stage] = loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); + } + tokens[stage].emplace(loaders[stage].Commit()); + slot_chunk[stage] = overflow_idx; + } + + // Apply `f` to every overflow key once, in the current ping-pong direction. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) + { + if (overflow_chunks == 0) + { + return; + } + + const offset_t m = overflow_chunks; + + if constexpr (use_block_load_to_shared) + { + // Normally a no-op: the resident load already built `PipelineStages >= p_eff` + // loaders. Only the (asserted-against) tiny-config edge would grow here. + ensure_loaders(p_eff); + const offset_t pe = static_cast(p_eff); + + // First ever call: prime the streaming slots. Subsequent calls inherit + // the previous pass's resident tail, which (because the order + // ping-pongs) is exactly the first `p_eff` chunks of this direction. + if (!primed) + { + for (int i = 0; i < p_eff; ++i) + { + const offset_t o = forward ? static_cast(i) : (m - 1 - static_cast(i)); + issue_load(static_cast(o % pe), o); + } + primed = true; + } + + for (offset_t i = 0; i < m; ++i) + { + const offset_t o = forward ? i : (m - 1 - i); + const int stage = static_cast(o % pe); + if (tokens[stage].has_value()) + { + loaders[stage].Wait(::cuda::std::move(*tokens[stage])); + tokens[stage].reset(); + } + _CCCL_ASSERT(slot_chunk[stage] == o, "overflow streamer slot/chunk mapping diverged"); + agent.for_each_chunk_key(pending[stage], f); + + // Prefetch the chunk `p_eff` visits ahead in this direction. It maps + // to the slot we just finished, so a barrier is required before the + // async copy can overwrite the data the block was just reading. + const offset_t ni = i + pe; + if (ni < m) + { + const offset_t no = forward ? ni : (m - 1 - ni); + __syncthreads(); + issue_load(stage, no); + } + } + forward = !forward; + } + else + { + // Generic fallback: overflow keys are read straight from gmem each pass + // (no SMEM reuse), but the walk still snakes for L2 locality. + for (offset_t i = 0; i < m; ++i) + { + const offset_t o = forward ? i : (m - 1 - i); + const offset_t chunk_idx = chunk_index_of(o); + const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk.count) + { + const key_t key = + block_keys_in[static_cast(chunk.offset + static_cast(local))]; + f(key); + } + }); + } + forward = !forward; + } + } + }; + // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- @@ -547,22 +738,55 @@ private: // operator at the matching radix level. int last_pass = num_passes; - offset_t head_items = 0; + const key_t* block_keys_base = nullptr; + offset_t head_items = 0; if constexpr (use_block_load_to_shared) { - auto* block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); - head_items = aligned_head_items(block_keys_base, segment_size_u32); + block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); + head_items = aligned_head_items(block_keys_base, segment_size_u32); } // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only // depends on the multiset of keys covered by the cluster. const offset_t chunks = num_chunks(segment_size_u32, head_items); const offset_t my_chunks = num_rank_chunks(chunks, cluster_rank, cluster_size); - // The launch coverage check reserves one extra chunk for the possible unaligned head, so every local chunk remains - // resident for later radix passes while only the BlockLoadToShared instances are reused round-robin. - _CCCL_ASSERT(my_chunks * offset_t{chunk_items} <= block_tile_capacity, + + // Resident vs. streaming split. Segments that fit the all-resident coverage behave exactly as before + // (`resident_slots_cap == full_slots`, no streaming). Larger segments reserve the last `PipelineStages` slots of + // the block_tile as a round-robin streaming region and keep `full_slots - PipelineStages` slots resident; the + // overflow chunks are re-streamed from gmem on every pass by `streamer`. The launch coverage check still reserves + // one extra chunk for the possible unaligned head. + const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); + const offset_t all_resident_capacity = + smem_layout_t::template cluster_tile_capacity(static_cast(cluster_size), block_tile_capacity); + const bool needs_streaming = segment_size_u32 > all_resident_capacity; + _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), + "block_tile too small to reserve a streaming region"); + const offset_t resident_slots_cap = + needs_streaming + ? ((full_slots > static_cast(PipelineStages)) + ? full_slots - static_cast(PipelineStages) + : offset_t{1}) + : full_slots; + const offset_t my_resident_chunks = (::cuda::std::min) (my_chunks, resident_slots_cap); + // Resident chunks stay within the first `resident_slots_cap` slots; the streaming region occupies the slots + // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. + _CCCL_ASSERT(my_resident_chunks * offset_t{chunk_items} <= resident_slots_cap * offset_t{chunk_items}, "Dynamic shared memory block_tile is too small"); + // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. + overflow_streamer streamer( + *this, + block_keys_in, + block_keys_base, + segment_size_u32, + head_items, + cluster_rank, + cluster_size, + my_resident_chunks, + static_cast(resident_slots_cap), + my_chunks); + ::cuda::std::span resident_keys; reset_hist(); @@ -584,20 +808,22 @@ private: if constexpr (use_block_load_to_shared) { - auto* block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); - // BlockLoadToShared is non-copyable and non-movable; keep the active pipeline local and emplace-only. - ::cuda::std::inplace_vector loaders; + // BlockLoadToShared is non-copyable and non-movable. The loaders are owned by `streamer` and constructed + // here; the streamer then reuses these very objects for the overflow passes (continuing their Commit/Wait + // pipeline) instead of tearing them down and re-initializing the shared `load_storage`. The per-chunk tokens + // and pending spans, however, are only needed for this resident load, so they stay local. ::cuda::std::inplace_vector tokens; ::cuda::std::inplace_vector<::cuda::std::span, PipelineStages> pending_spans; - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_chunks)); + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); char* next_dst = key_slots; + streamer.ensure_loaders(prologue); for (int stage = 0; stage < prologue; ++stage) { const offset_t chunk_idx = static_cast(cluster_rank) + static_cast(stage * cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); const bool aligned_chunk = is_aligned_chunk(block_keys_base, chunk); - auto& loader = loaders.emplace_back(temp_storage.load_storage[stage]); + auto& loader = streamer.loaders[stage]; const ::cuda::std::span src{ block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; if (aligned_chunk) @@ -613,16 +839,16 @@ private: tokens.emplace_back(loader.Commit()); } - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const int stage = static_cast(p % static_cast(prologue)); const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - loaders[stage].Wait(::cuda::std::move(tokens[stage])); + streamer.loaders[stage].Wait(::cuda::std::move(tokens[stage])); append_contiguous_span(resident_keys, pending_spans[stage]); for_each_chunk_key(pending_spans[stage], add_first_pass); - if (p + static_cast(prologue) < my_chunks) + if (p + static_cast(prologue) < my_resident_chunks) { const offset_t next_chunk_idx = static_cast(cluster_rank) + (p + static_cast(prologue)) * static_cast(cluster_size); @@ -632,22 +858,22 @@ private: block_keys_base + next_chunk.offset, static_cast<::cuda::std::size_t>(next_chunk.count)}; if (next_aligned_chunk) { - pending_spans[stage] = - loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); + pending_spans[stage] = streamer.loaders[stage].template CopyAsync( + available_block_tile_buffer(next_dst), src); } else { pending_spans[stage] = - loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); + streamer.loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); } next_dst = span_end(pending_spans[stage]); - tokens[stage] = loaders[stage].Commit(); + tokens[stage] = streamer.loaders[stage].Commit(); } } } else { - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -665,6 +891,12 @@ private: }); } } + + // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. + // The streamer reuses the resident loaders (no re-init), so this just continues their Commit/Wait pipeline; + // `Commit()`'s own block barrier provides the necessary synchronization. + streamer.process_pass(add_first_pass); + const int resident_count = span_size(resident_keys); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, "Dynamic shared memory block_tile is too small"); @@ -743,7 +975,7 @@ private: } else { - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -751,6 +983,10 @@ private: for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } } + + // Re-stream the overflow chunks into this pass's histogram. Ping-pongs direction and reuses the boundary + // chunks left resident by the previous pass. + streamer.process_pass(add_hist); } # ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE @@ -896,7 +1132,7 @@ private: } else { - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -904,6 +1140,7 @@ private: for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, count_selected); } } + streamer.process_pass(count_selected); __syncthreads(); // Thread 0 of every block (leader block included) reserves the block's @@ -954,7 +1191,7 @@ private: } else { - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -962,6 +1199,7 @@ private: for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } + streamer.process_pass(write_selected); # else auto write_selected = [&](const key_t& key) { const auto res = identify_op(key); @@ -986,7 +1224,7 @@ private: } else { - for (offset_t p = 0; p < my_chunks; ++p) + for (offset_t p = 0; p < my_resident_chunks; ++p) { const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -994,6 +1232,7 @@ private: for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } + streamer.process_pass(write_selected); # endif // Final cluster barrier: hold every block in the cluster until all DSMEM @@ -1028,17 +1267,17 @@ private: return; } + // Segments larger than the resident cluster_tile capacity are handled by re-streaming the overflow chunks from + // gmem (see `overflow_streamer`), so the only hard limit left is the 32-bit offset range used internally. bool segment_fits_offset = true; if constexpr (sizeof(segment_size_val_t) > sizeof(offset_t)) { segment_fits_offset = segment_size <= static_cast(::cuda::std::numeric_limits::max()); } - const auto max_cluster_tile_capacity = smem_layout_t::template cluster_tile_capacity( - static_cast(cluster_blocks), block_tile_capacity); - if (!segment_fits_offset || segment_size > max_cluster_tile_capacity) + if (!segment_fits_offset) { - _CCCL_ASSERT(false, "Segment exceeds the selected cluster top-k cluster_tile capacity"); + _CCCL_ASSERT(false, "Segment exceeds the 32-bit offset range supported by cluster top-k"); return; } diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 28b3bb01c3b..682e0056fe7 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -349,6 +349,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( static_assert(ChunkBytes % LoadAlignBytes == 0); static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + // Static-footprint estimate for the device-side CDP fallback, which cannot query `cudaFuncGetAttributes`. + // The host path instead uses the driver-reported `sharedSizeBytes` (see below), which is padding-aware. constexpr int static_smem_bytes = static_cast(sizeof(typename agent_t::TempStorage)); const auto max_seg_size = runtime_max_segment_size(segment_sizes); @@ -439,32 +441,35 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( cfg.attrs = &cluster_attr; cfg.numAttrs = 1; - int max_dynamic_smem_bytes = 0; - if (const auto error = CubDebug(MaxPotentialDynamicSmemBytes(max_dynamic_smem_bytes, dynamic_kernel))) - { - return error; - } - - // The table entries express the documented total per-block shared-memory budget - // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`). The usable dynamic portion is that budget minus the - // kernel's static shared memory and the CUDA driver's per-block reserved shared memory, matching - // `MaxPotentialDynamicSmemBytes`. Subtracting the reserved bytes here keeps a documented entry (e.g. 99 KiB) - // from deriving a dynamic size that overshoots `max_dynamic_smem_bytes` and being discarded below. + // Resolve the per-block opt-in shared-memory budget and the kernel's static footprint from the driver so + // the dynamic-SMEM math below matches exactly what the launch permits. The table entries express the + // documented total per-block budget (`cudaDevAttrMaxSharedMemoryPerBlockOptin`); the usable dynamic portion + // is that budget minus the static footprint. int device_id = 0; if (const auto error = CubDebug(cudaGetDevice(&device_id))) { return error; } - int reserved_smem_bytes = 0; + int max_smem_optin_bytes = 0; if (const auto error = - CubDebug(cudaDeviceGetAttribute(&reserved_smem_bytes, cudaDevAttrReservedSharedMemoryPerBlock, device_id))) + CubDebug(cudaDeviceGetAttribute(&max_smem_optin_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id))) { return error; } - // TODO: subtracting `reserved` mirrors `MaxPotentialDynamicSmemBytes`, but that helper double-counts it - // (opt-in already excludes reserved), so this is ~`reserved` bytes too conservative. Revisit once the helper is - // fixed. - const int nondynamic_smem_bytes = static_smem_bytes + reserved_smem_bytes; + // Use the driver-reported static footprint (`sharedSizeBytes`) rather than `sizeof(TempStorage)`: it reflects + // any padding the toolchain inserts to align the dynamic shared-memory section after the static one, so the + // derived dynamic sizes neither overshoot the budget nor conservatively drop the top table tier. + cudaFuncAttributes kernel_attrs{}; + if (const auto error = CubDebug(cudaFuncGetAttributes(&kernel_attrs, dynamic_kernel))) + { + return error; + } + // `cudaDevAttrMaxSharedMemoryPerBlockOptin` already excludes the driver's per-block reserved shared memory + // (opt-in == per-SM - reserved), so the dynamic budget is just the opt-in budget minus the static footprint; + // reserved must not be subtracted a second time. + const int nondynamic_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); + const int max_dynamic_smem_bytes = + (max_smem_optin_bytes > nondynamic_smem_bytes) ? max_smem_optin_bytes - nondynamic_smem_bytes : 0; const auto runtime_policy = PolicySelector{}(::cuda::compute_capability{sm_version / 10}); _CCCL_ASSERT(runtime_policy.launch_configs.size() <= max_launch_configs, @@ -517,24 +522,9 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const auto candidate_cluster_tile_capacity = layout_t::template cluster_tile_capacity( candidate.cluster_blocks, candidate.block_tile_capacity); - bool candidate_can_improve_selection = false; - if constexpr (detail::params::is_per_segment_param_v) - { - // Runtime-sized segments may exceed every finite candidate, so find the largest supported fallback. - candidate_can_improve_selection = true; - } - else if (candidate_cluster_tile_capacity >= max_seg_size) - { - const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( - selected_config.cluster_blocks, selected_config.block_tile_capacity); - candidate_can_improve_selection = - selected_config.cluster_blocks == 0 || candidate_cluster_tile_capacity < selected_cluster_tile_capacity; - } - if (!candidate_can_improve_selection) - { - continue; - } - + // Every smem-fitting candidate is considered: it may either tighten the covering selection or improve the + // largest-coverage fallback used for oversize segments. The table is tiny, so the extra driver queries below + // are negligible. if (const auto error = ensure_dynamic_smem_limit(candidate_dynamic_smem)) { return error; @@ -554,7 +544,9 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( continue; } - if constexpr (detail::params::is_per_segment_param_v) + // Track the largest-coverage hardware-supported config for all parameter kinds. Segments that exceed every + // finite candidate (always possible for per-segment sizes, and now also for static/uniform bounds) fall back + // to this config and re-stream the overflow from gmem in the agent. { const auto largest_supported_cluster_tile_capacity = layout_t::template cluster_tile_capacity( largest_supported_config.cluster_blocks, largest_supported_config.block_tile_capacity); @@ -565,22 +557,23 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( } } + // Among the candidates that fully cover the segment(s), keep the one with the lowest coverage. if (candidate_cluster_tile_capacity >= max_seg_size) { - selected_config = candidate; + const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( + selected_config.cluster_blocks, selected_config.block_tile_capacity); + if (selected_config.cluster_blocks == 0 || candidate_cluster_tile_capacity < selected_cluster_tile_capacity) + { + selected_config = candidate; + } } } if (selected_config.cluster_blocks == 0) { - if constexpr (detail::params::is_per_segment_param_v) - { - selected_config = largest_supported_config; - } - else - { - return cudaErrorInvalidValue; - } + // No finite candidate covers the segment(s); fall back to the largest hardware-supported config and let the + // agent stream the overflow. Applies to per-segment sizes and static/uniform bounds alike. + selected_config = largest_supported_config; } if (selected_config.cluster_blocks == 0) { @@ -644,20 +637,13 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( }), ({ // CDP path: device-side launches cannot opt in to more than portable - // total SMEM or non-portable cluster widths. + // total SMEM or non-portable cluster widths. Segments that exceed the + // portable resident coverage are still handled: the agent re-streams the + // overflow chunks from gmem. constexpr int portable_total_smem_bytes = 48 * 1024; constexpr int dynamic_smem_bytes = (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; constexpr auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); - constexpr auto portable_cluster_tile_capacity = - layout_t::template cluster_tile_capacity(max_portable_cluster_blocks, block_tile_capacity); - if constexpr (!detail::params::is_per_segment_param_v) - { - if (max_seg_size > static_cast(portable_cluster_tile_capacity)) - { - return cudaErrorInvalidValue; - } - } const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 95218c5c851..0bf2312e21f 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -206,14 +207,18 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - constexpr segment_size_t static_max_segment_size = 128 * 1024; + // `static_max_segment_size` is chosen to exceed the largest all-resident cluster coverage (~16 blocks worth of + // resident SMEM), so the 1 Mi-element segments force the agent's gmem-streaming overflow path (including an + // unaligned overflow tail via `- 31`), while the 128 Ki-element segment still runs fully resident under the same + // streaming-capable launch configuration. + constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; const auto direction = cub::detail::topk::select::max; const int pad = GENERATE(0, 1, 3, 7); const segment_size_t segment_size = - GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 1, static_max_segment_size - 31})); + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); @@ -247,6 +252,90 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned REQUIRE(expected_keys == keys_out_buffer); } + +template +struct cast_to_key_op +{ + template + __host__ __device__ KeyT operator()(T x) const + { + return static_cast(x); + } +}; + +// Yields, for each segment, a *non-contiguous* iterator over that segment's keys (an integral counting iterator +// cast to the key type). Feeding the cluster top-k a non-contiguous key iterator makes `use_block_load_to_shared` +// false, so the agent takes its generic (non-BlockLoadToShared) overflow-streaming path. Segment `seg` produces +// keys [seg * segment_size, (seg + 1) * segment_size), so the flattened input equals the identity sequence and the +// expected top-k is exact. +template +struct counting_segment_keys_op +{ + SegmentSizeT segment_size; + + template + __host__ __device__ auto operator()(IndexT seg) const + { + return cuda::make_transform_iterator( + cuda::make_counting_iterator(static_cast(seg) * segment_size), cast_to_key_op{}); + } +}; + +TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-contiguous key iterator", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + // The counting-iterator key source is non-contiguous, so the agent uses its generic overflow-streaming path rather + // than BlockLoadToShared. `static_max_segment_size` exceeds the largest all-resident cluster coverage, so the 1 Mi + // -element segments stream (incl. an unaligned `- 31` tail), while the 128 Ki-element segment validates the generic + // resident path (no streaming) through the same code. Keeping the largest total below 2^24 makes every key an exact + // float. + constexpr segment_size_t static_max_segment_size = 1024 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 3; + + const auto direction = cub::detail::topk::select::max; + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + // Non-contiguous input: segment `seg` is the counting iterator [seg * segment_size, (seg + 1) * segment_size). + auto d_keys_in = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), counting_segment_keys_op{segment_size}); + + // Output is a real buffer (the output iterator stays contiguous; only the input drives the streaming path). + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cub::detail::batched_topk::segment_size_uniform<1, static_max_segment_size>{segment_size}, + cub::detail::batched_topk::k_uniform<1, static_max_k>{k}, + cub::detail::batched_topk::select_direction_uniform{direction}, + cub::detail::batched_topk::num_segments_uniform<>{num_segments}, + cub::detail::batched_topk::total_num_items_guarantee{num_items}); + + // The flattened input is the identity sequence, so build the expected keys directly and reuse the standard + // sort + compact verification. + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::sequence(expected_keys.begin(), expected_keys.end()); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} #endif // TEST_LAUNCH != 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segments", @@ -308,9 +397,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment auto compacted_output_sizes_it = cuda::make_transform_iterator( cuda::make_counting_iterator(segment_index_t{0}), get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k)}); - c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); - thrust::exclusive_scan( - compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + // Exclusive prefix sum of the per-segment output sizes. Scan only the `num_segments` valid sizes (each reads + // `offset[seg]`/`offset[seg + 1]`, staying in bounds) into indices [1, num_segments]; index 0 stays 0. + c2h::device_vector compacted_offsets(num_segments + 1); + thrust::inclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments, compacted_offsets.begin() + 1); segment_size_t total_output_size = compacted_offsets.back(); // Prepare keys input & output @@ -358,19 +449,23 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalig using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - constexpr segment_size_t static_max_segment_size = 128 * 1024; + // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch + // mixes streaming segments (the 1 Mi-element ones, one with an unaligned `- 31` overflow tail) with fully-resident + // segments (96 Ki + 17 and 257 elements). + constexpr segment_size_t static_max_segment_size = 1100 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; const auto direction = cub::detail::topk::select::max; const int pad = GENERATE(1, 3, 7); const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); + constexpr segment_size_t big_segment_size = 1024 * 1024; c2h::host_vector h_segment_offsets{ 0, - static_max_segment_size, - static_max_segment_size + (static_max_segment_size - 31), - static_max_segment_size + (static_max_segment_size - 31) + (96 * 1024 + 17), - static_max_segment_size + (static_max_segment_size - 31) + (96 * 1024 + 17) + 257}; + big_segment_size, + big_segment_size + (big_segment_size - 31), + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17), + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257}; c2h::device_vector segment_offsets = h_segment_offsets; const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); const segment_size_t num_items = h_segment_offsets.back(); @@ -384,9 +479,11 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalig auto compacted_output_sizes_it = cuda::make_transform_iterator( cuda::make_counting_iterator(segment_index_t{0}), get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k)}); - c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); - thrust::exclusive_scan( - compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + // Exclusive prefix sum of the per-segment output sizes. Scan only the `num_segments` valid sizes (each reads + // `offset[seg]`/`offset[seg + 1]`, staying in bounds) into indices [1, num_segments]; index 0 stays 0. + c2h::device_vector compacted_offsets(num_segments + 1); + thrust::inclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments, compacted_offsets.begin() + 1); segment_size_t total_output_size = compacted_offsets.back(); c2h::device_vector keys_in_buffer(pad + num_items, thrust::no_init); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index a519cff9959..b9823eb4f67 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -305,9 +305,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segmen auto compacted_output_sizes_it = cuda::make_transform_iterator( cuda::make_counting_iterator(segment_index_t{0}), get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k)}); - c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); - thrust::exclusive_scan( - compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + // Exclusive prefix sum of the per-segment output sizes. Scan only the `num_segments` valid sizes (each reads + // `offset[seg]`/`offset[seg + 1]`, staying in bounds) into indices [1, num_segments]; index 0 stays 0. + c2h::device_vector compacted_offsets(num_segments + 1); + thrust::inclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments, compacted_offsets.begin() + 1); segment_size_t total_output_size = compacted_offsets.back(); // Prepare keys input & output diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index aebf2ce7fdd..3585074b08d 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -299,9 +299,10 @@ c2h::device_vector compact_to_topk_batched( auto copy_sizes_it = cuda::make_transform_iterator( cuda::make_counting_iterator(0), get_output_size_op{d_offsets.cbegin(), cuda::constant_iterator(k)}); - // Calculate destination offsets via prefix sum - c2h::device_vector d_output_offsets(num_segments + 1, thrust::no_init); - thrust::exclusive_scan(copy_sizes_it, copy_sizes_it + num_segments + 1, d_output_offsets.begin()); + // Calculate destination offsets via prefix sum. Scan only the `num_segments` valid sizes (each reads + // `offset[seg]`/`offset[seg + 1]`, which stay in bounds) into indices [1, num_segments]; index 0 stays 0. + c2h::device_vector d_output_offsets(num_segments + 1); + thrust::inclusive_scan(copy_sizes_it, copy_sizes_it + num_segments, d_output_offsets.begin() + 1); OffsetT total_compacted_size = d_output_offsets.back(); c2h::device_vector d_keys_out(total_compacted_size, thrust::no_init); From 8e29b292484614c7fe8d9ac0d6311a29fb32e8b7 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 8 Jun 2026 05:31:28 +0200 Subject: [PATCH 021/126] Fix variable benchmark --- cub/benchmarks/bench/segmented_topk/variable/keys.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index 486f6f55ff2..d5b4673b565 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -299,7 +299,8 @@ void variable_seg_size_topk_keys(nvbench::state& state, cub::detail::batched_topk::segment_size_per_segment segment_sizes_param{ thrust::raw_pointer_cast(d_segment_sizes.data())}; cub::detail::batched_topk::k_static k_param{}; - cub::detail::batched_topk::select_direction_static select_directions{}; + cub::detail::batched_topk::select_direction_static select_directions{ + cub::detail::topk::select::max}; cub::detail::batched_topk::num_segments_uniform<> num_segments_uniform_param{ static_cast(num_segments)}; From 69e7932b3ef6b79b4fa1546279acb606756b4e9c Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 8 Jun 2026 05:32:19 +0200 Subject: [PATCH 022/126] Manually tune for B200 Also remove ifdefs since perf is clearly worse with the alternative code paths. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 209 +----------------- .../dispatch_batched_topk_cluster.cuh | 5 +- .../tuning/tuning_batched_topk_cluster.cuh | 19 +- 3 files changed, 17 insertions(+), 216 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 31149e704ff..04a76479546 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -23,23 +23,10 @@ //! the next pass. //! //! Output cursors live in the same cluster-shared `state` and are reached the -//! same way (cluster-scope DSMEM atomics). Defining -//! `CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT` switches the final filter pass -//! to a block-aggregated variant that first counts selected / candidate keys -//! in block-local shared-memory counters and then issues a single DSMEM -//! atomic per block per counter to reserve a contiguous output range. This -//! lets us A/B the two strategies without touching the rest of the kernel. -//! -//! Defining `CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE` flips the histogram -//! pipeline from reduce-then-scan to scan-then-reduce: every block does a -//! block-local `InclusiveSum` over its own histogram, then folds the scans -//! (not the raw counts) into the leader. The leader then identifies the -//! k-th bucket directly from `hist[bucket]` / `hist[bucket - 1]`. +//! same way (cluster-scope DSMEM atomics). #pragma once -#define CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE - #include #if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) @@ -223,15 +210,6 @@ struct agent_batched_topk_cluster ::cuda::std::uint32_t broadcast_early_stop; typename block_scan_t::TempStorage scan_storage; typename block_load_t::TempStorage load_storage[PipelineStages]; -#ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT - // Final-filter block-aggregation slots. First populated by per-thread - // `atomicAdd_block` calls that count the block's selected / candidate - // keys, then overwritten by thread 0 with the leader-side base position - // returned from a single DSMEM atomic per counter, so the same slots - // double as the broadcast channel for the base to the rest of the block. - out_offset_t block_out_cnt; - out_offset_t block_out_back_cnt; -#endif }; [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const @@ -474,41 +452,6 @@ private: } } -#ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE - // Scan-then-reduce counterpart of `leader_identify_kth_bucket`. Assumes - // `hist[i]` already holds the cluster-wide inclusive prefix sum. The - // per-bucket count is recovered as `inclusive - exclusive`, where - // `exclusive` is `hist[bucket - 1]` (or `0` for `bucket == 0`). - // - // `target_k` is `state.k` read by the caller before the `cluster.sync()` - // that precedes this call, so that barrier orders the read against the - // write this function may issue to `state.k`. - _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket_from_inclusive(int pass, out_offset_t target_k) - { - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < buckets_per_thread; ++j) - { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; - if (bucket >= num_buckets) - { - continue; - } - const offset_t inclusive = temp_storage.hist[bucket]; - const offset_t exclusive = (bucket == 0) ? offset_t{0} : temp_storage.hist[bucket - 1]; - if (exclusive < target_k && inclusive >= target_k) - { - const out_offset_t new_k = target_k - static_cast(exclusive); - const offset_t new_len = inclusive - exclusive; - temp_storage.state.len = new_len; - temp_storage.state.k = new_k; - temp_storage.state.early_stop = - (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); - } - } - } -#endif - // --------------------------------------------------------------------------- // Overflow streamer // --------------------------------------------------------------------------- @@ -989,42 +932,11 @@ private: streamer.process_pass(add_hist); } -# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE - // Step 1b (scan-then-reduce): in-place block-local `InclusiveSum` - // over `hist[]`. The barrier publishes Step 1's atomic writes to - // the scan's reads. - __syncthreads(); - { - offset_t bucket_vals[buckets_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < buckets_per_thread; ++j) - { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; - bucket_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; - } - block_scan_t(temp_storage.scan_storage).InclusiveSum(bucket_vals, bucket_vals); - _CCCL_PRAGMA_UNROLL_FULL() - for (int j = 0; j < buckets_per_thread; ++j) - { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; - if (bucket < num_buckets) - { - temp_storage.hist[bucket] = bucket_vals[j]; - } - } - } - // Cluster-wide barrier: the leader's scan write-back to `hist[]` - // is non-atomic and would be lost if a remote Step 2 RMW - // interleaved with it, so every block must finish its write-back - // before anyone starts the fold. - cluster.sync(); -# else // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` // are atomic at compatible scopes (see Step 1 dispatch). The // cluster-wide ordering before Step 3's leader read of `hist[]` // is supplied by the `cluster.sync()` further below. __syncthreads(); -# endif // Step 2: non-leader blocks fold their per-bucket values // (raw counts in the reduce-then-scan path, block-local inclusive @@ -1047,15 +959,6 @@ private: } } -# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE - // Hoist the leader's read of `state.k` above the upcoming - // `cluster.sync()` so that barrier separates it from the write - // `leader_identify_kth_bucket_from_inclusive` may issue to - // `state.k`. Unconditional in every block to avoid predication; - // safe because `process_impl` initializes `state` everywhere. - const out_offset_t leader_target_k = temp_storage.state.k; -# endif - cluster.sync(); // Step 3: the leader walks the merged `hist` (raw counts in the @@ -1065,11 +968,7 @@ private: // these writes after the next cluster sync. if (cluster_rank == 0) { -# ifdef CUB_ENABLE_CLUSTER_TOPK_SCAN_THEN_REDUCE - leader_identify_kth_bucket_from_inclusive(pass, leader_target_k); -# else leader_identify_kth_bucket(pass); -# endif } cluster.sync(); } @@ -1081,14 +980,6 @@ private: if (threadIdx.x == 0) { temp_storage.broadcast_kth = leader_state->kth_key_bits; -# ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT - // Piggyback the block-local output-counter init on the broadcast_kth - // publish: the `__syncthreads()` below already orders these writes - // against the per-thread `atomicAdd_block` calls in the first phase - // of the block-aggregated path, so no dedicated init+sync is needed. - temp_storage.block_out_cnt = 0; - temp_storage.block_out_back_cnt = 0; -# endif } __syncthreads(); kth_key_bits_local = temp_storage.broadcast_kth; @@ -1104,103 +995,6 @@ private: // identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); -# ifdef CUB_ENABLE_CLUSTER_TOPK_BLOCK_AGG_OUTPUT - // Block-aggregated variant: each thread first reserves its slot in the - // block's local counters via cheap `atomicAdd_block`, then thread 0 of - // every block claims the block's contiguous range in the leader's - // counters with a single DSMEM `atomicAdd` per counter. This collapses - // up to the block's resident key count DSMEM atomics per block - // down to two, at the cost of one extra block-wide barrier (the local - // counters are zeroed inside the pre-existing broadcast_kth publish) - // and one additional pass over the per-thread keys to materialize the - // writes. - - auto count_selected = [&](const key_t& key) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); - } - else if (res == detail::topk::candidate_class::candidate) - { - atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); - } - }; - if constexpr (use_block_load_to_shared) - { - for_each_chunk_key(resident_keys, count_selected); - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) - { - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, count_selected); - } - } - streamer.process_pass(count_selected); - __syncthreads(); - - // Thread 0 of every block (leader block included) reserves the block's - // contiguous range in the leader's counters. The returned old value is - // the block's base; we stash it back into the same shared-memory slot - // so the rest of the block can pick it up after the barrier without a - // separate broadcast field. For the leader block these atomics target - // its own shared memory; for non-leaders they go through DSMEM. - if (threadIdx.x == 0) - { - const out_offset_t fwd_count = temp_storage.block_out_cnt; - const out_offset_t back_count = temp_storage.block_out_back_cnt; - temp_storage.block_out_cnt = atomicAdd(&leader_state->out_cnt, fwd_count); - temp_storage.block_out_back_cnt = atomicAdd(&leader_state->out_back_cnt, back_count); - } - __syncthreads(); - - const out_offset_t fwd_base = temp_storage.block_out_cnt; - const out_offset_t back_base = temp_storage.block_out_back_cnt; - - if (threadIdx.x == 0) - { - temp_storage.block_out_cnt = 0; - temp_storage.block_out_back_cnt = 0; - } - __syncthreads(); - - auto write_selected = [&](const key_t& key) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = fwd_base + atomicAdd_block(&temp_storage.block_out_cnt, out_offset_t{1}); - block_keys_out[pos] = key; - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = back_base + atomicAdd_block(&temp_storage.block_out_back_cnt, out_offset_t{1}); - if (back_pos < num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; - } - } - }; - if constexpr (use_block_load_to_shared) - { - for_each_chunk_key(resident_keys, write_selected); - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) - { - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); - } - } - streamer.process_pass(write_selected); -# else auto write_selected = [&](const key_t& key) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) @@ -1233,7 +1027,6 @@ private: } } streamer.process_pass(write_selected); -# endif // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 682e0056fe7..a86e64d8c5f 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -98,6 +98,7 @@ struct launch_config // Dynamic-cluster kernel for host launches; the agent reads the active cluster // width via cooperative groups. template -__launch_bounds__(ThreadsPerBlock) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( +__launch_bounds__(ThreadsPerBlock, MinBlocksPerSm) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, SegmentSizeParameterT segment_sizes, @@ -324,6 +325,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // CCs, so pin the selector query to the minimum supported CC. constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); constexpr int ThreadsPerBlock = policy.threads_per_block; + constexpr int MinBlocksPerSm = policy.min_blocks_per_sm; constexpr int UnrollFactor = policy.unroll_factor; constexpr int PipelineStages = policy.pipeline_stages; constexpr int ChunkBytes = policy.chunk_bytes; @@ -399,6 +401,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr auto dynamic_kernel = &device_segmented_topk_cluster_kernel< ThreadsPerBlock, + MinBlocksPerSm, UnrollFactor, PipelineStages, ChunkBytes, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index b80c7642c2c..8562bf0b7e0 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -38,6 +38,10 @@ struct cluster_topk_policy int chunk_bytes; int load_align_bytes; int bits_per_pass; + // Minimum number of CTAs per SM passed to the dynamic-cluster kernel's launch bounds. Defaults to 1: most cluster + // configurations already have an occupancy of 1 due to their shared-memory usage, so allowing the compiler to + // optimize for higher occupancy (fewer registers per thread) provides no benefit. + int min_blocks_per_sm; ::cuda::std::inplace_vector launch_configs; }; @@ -55,12 +59,13 @@ make_policy(::cuda::std::inplace_vector cluster_topk_policy { return cluster_topk_policy{ - /*threads_per_block=*/256, - /*unroll_factor=*/8, + /*threads_per_block=*/512, + /*unroll_factor=*/0, /*pipeline_stages=*/3, - /*chunk_bytes=*/8 * 1024, + /*chunk_bytes=*/16 * 1024, /*load_align_bytes=*/128, - /*bits_per_pass=*/8, + /*bits_per_pass=*/11, + /*min_blocks_per_sm=*/1, launch_configs}; } @@ -88,9 +93,9 @@ make_policy(::cuda::std::inplace_vector Date: Tue, 9 Jun 2026 05:12:14 +0200 Subject: [PATCH 023/126] Clean up chunking Also fixes use of generic addressing due to spilling of smem addresses to local memory. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 356 ++++++++++++++----- 1 file changed, 269 insertions(+), 87 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 04a76479546..c8327cb13bd 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -97,12 +97,19 @@ struct alignas(16) cluster_topk_state template struct smem_block_tile_layout { - static constexpr int chunk_items = ChunkBytes / int{sizeof(KeyT)}; + static constexpr int chunk_items = ChunkBytes / int{sizeof(KeyT)}; + static constexpr int load_align_items = LoadAlignBytes / int{sizeof(KeyT)}; static constexpr int slot_alignment = (::cuda::std::max) (LoadAlignBytes, detail::LoadToSharedBufferAlignBytes()); - static constexpr int slot_stride_bytes = - ::cuda::round_up(detail::LoadToSharedBufferSizeBytes(chunk_items) + slot_alignment, slot_alignment); + // Each chunk maps to exactly one slot of ChunkBytes: boundary chunks are loaded as an aligned bulk plus a small + // hand-rolled edge, so no chunk ever needs the scalar-load guard headroom. + static constexpr int slot_stride_bytes = ::cuda::round_up(ChunkBytes, slot_alignment); // == ChunkBytes normally static constexpr int base_padding_bytes = (alignof(KeyT) > 16) ? slot_alignment : 0; + static_assert(chunk_items >= load_align_items, "ChunkBytes must hold at least one load_align unit"); + // The aligned full-chunk buffer (no guard, since LoadAlignBytes >= bulk_copy_min_align) is exactly its byte count + // and must fit one slot. + static_assert(detail::LoadToSharedBufferSizeBytes(chunk_items) <= slot_stride_bytes, + "An aligned full chunk must fit one slot"); [[nodiscard]] _CCCL_HOST_DEVICE static constexpr ::cuda::std::uint32_t block_tile_capacity(int dynamic_smem_bytes) noexcept @@ -165,6 +172,7 @@ struct agent_batched_topk_cluster static constexpr int num_buckets = 1 << bits_per_pass; using smem_layout_t = smem_block_tile_layout; static constexpr int chunk_items = smem_layout_t::chunk_items; + static constexpr int load_align_items = smem_layout_t::load_align_items; static constexpr int slot_alignment = smem_layout_t::slot_alignment; static constexpr int slot_stride_bytes = smem_layout_t::slot_stride_bytes; @@ -173,6 +181,9 @@ struct agent_batched_topk_cluster static_assert(LoadAlignBytes > 0); static_assert(ChunkBytes % LoadAlignBytes == 0); static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + // The hybrid load relies on the aligned CopyAsync path being exact (no scalar guard), which requires the load + // alignment to be at least the bulk-copy minimum alignment. + static_assert(LoadAlignBytes >= detail::bulk_copy_min_align, "LoadAlignBytes must be >= bulk_copy_min_align"); static_assert(chunk_items > 0); using decomposer_t = detail::identity_decomposer_t; @@ -267,11 +278,6 @@ struct agent_batched_topk_cluster return count; } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE char* span_end(::cuda::std::span keys) const - { - return reinterpret_cast(::cuda::std::data(keys) + span_size(keys)); - } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t aligned_head_items(const key_t* base, offset_t segment_size) const { @@ -283,11 +289,25 @@ struct agent_batched_topk_cluster return (::cuda::std::min) (items, segment_size); } + // Item count of the near-full head chunk (chunk 0) for an unaligned base. The unaligned prefix (`head_items`) plus + // an aligned bulk filling the rest of the slot up to a `load_align` boundary, so chunk 0 ends aligned and every + // later chunk begins aligned. Clamped to the segment for tiny segments. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t + head_chunk_items(offset_t segment_size, offset_t head_items) const + { + const offset_t bulk0 = offset_t{chunk_items} - offset_t{load_align_items}; + return (::cuda::std::min) (segment_size, head_items + bulk0); + } + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t num_chunks(offset_t segment_size, offset_t head_items) const { - const offset_t remaining = segment_size - head_items; - return ((head_items != 0) ? offset_t{1} : offset_t{0}) - + static_cast(::cuda::ceil_div(remaining, offset_t{chunk_items})); + if (head_items == 0) + { + return static_cast(::cuda::ceil_div(segment_size, offset_t{chunk_items})); + } + const offset_t head0 = head_chunk_items(segment_size, head_items); + const offset_t remaining = segment_size - head0; + return offset_t{1} + static_cast(::cuda::ceil_div(remaining, offset_t{chunk_items})); } struct chunk_desc @@ -299,26 +319,61 @@ struct agent_batched_topk_cluster [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_desc get_chunk(offset_t chunk_idx, offset_t segment_size, offset_t head_items) const { - offset_t offset = chunk_idx * offset_t{chunk_items}; + offset_t offset; if (head_items != 0) { + const offset_t head0 = head_chunk_items(segment_size, head_items); if (chunk_idx == 0) { - return {offset_t{0}, static_cast(head_items)}; + return {offset_t{0}, static_cast(head0)}; } - offset = head_items + (chunk_idx - 1) * offset_t{chunk_items}; + offset = head0 + (chunk_idx - 1) * offset_t{chunk_items}; + } + else + { + offset = chunk_idx * offset_t{chunk_items}; } const offset_t remaining = segment_size - offset; return {offset, static_cast((::cuda::std::min) (remaining, offset_t{chunk_items}))}; } + // Splits a chunk into its unaligned front edge, aligned interior (bulk), and unaligned back edge, relative to the + // gmem base. The interior begins and ends on a `load_align` boundary so it can be loaded with the aligned, + // guard-free BlockLoadToShared path; the edges (each `< load_align_items` items) are loaded with `copy_edge`. Only + // chunk 0 (head) can have a nonzero prefix and only the last chunk (tail) a nonzero suffix. + struct chunk_split + { + offset_t prefix; + offset_t bulk; + offset_t suffix; + }; + + template + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_split split_chunk(PtrT base, const chunk_desc chunk) const + { + const auto la = static_cast<::cuda::std::uintptr_t>(load_align_bytes); + const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); + const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); + const auto aligned_begin = ::cuda::round_up(begin, la); + const auto aligned_end = (end / la) * la; + if (aligned_begin > aligned_end) + { + // The chunk lies strictly between two load_align boundaries (no aligned point inside): the whole chunk is an + // unaligned edge. This only happens for a tiny chunk with an unaligned begin (the head or a single-chunk + // segment), so attributing it entirely to the front edge is correct. A tail always begins on a boundary, so it + // takes the `aligned_begin <= aligned_end` path below and its unaligned remainder becomes the suffix. + return {static_cast(chunk.count), offset_t{0}, offset_t{0}}; + } + const offset_t prefix = static_cast((aligned_begin - begin) / sizeof(key_t)); + const offset_t bulk = static_cast((aligned_end - aligned_begin) / sizeof(key_t)); + return {prefix, bulk, static_cast(chunk.count) - prefix - bulk}; + } + template [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_aligned_chunk(PtrT base, const chunk_desc chunk) const { - const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); - const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); - return (begin % static_cast<::cuda::std::uintptr_t>(load_align_bytes) == 0) - && (end % static_cast<::cuda::std::uintptr_t>(load_align_bytes) == 0); + const auto s = split_chunk(base, chunk); + return s.prefix == 0 && s.suffix == 0; } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t @@ -329,6 +384,22 @@ struct agent_batched_topk_cluster : offset_t{0}; } + // Hand-rolled per-thread copy of a small (`< load_align_items` items) unaligned edge from gmem to smem. Used for + // the head prefix and tail suffix, which cannot go through the aligned (16-byte-aligned dst, guard-free) + // BlockLoadToShared path. Each thread copies the same indices it later reads in the first-pass histogram, so no + // extra synchronization is needed for the edge's own pass. + // + // TODO(cccl): an asymmetric-alignment BlockLoadToShared API (independent begin/end alignment, e.g. an aligned begin + // with an arbitrary end) would let a boundary chunk be loaded with a single aligned-bulk + in-place edge call and + // remove these hand-rolled copies entirely. + _CCCL_DEVICE _CCCL_FORCEINLINE void copy_edge(key_t* dst, const key_t* src, int count) const + { + for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) + { + dst[local] = src[local]; + } + } + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const { @@ -343,6 +414,22 @@ struct agent_batched_topk_cluster }); } + // A resident/streamed bulk in the block_tile, described by its 32-bit shared-window address plus length instead of + // a 64-bit generic pointer. Pipeline arrays of these are read across long loops and spill to local memory; a + // spilled generic pointer reloads as generic (`LDL.64`) and demotes the key reads from `LDS` to a generic `LD`, + // whereas a spilled 32-bit shared address rebuilt with `__cvta_shared_to_generic` keeps the reads `LDS`. + struct shared_bulk + { + ::cuda::std::uint32_t smem32; + int len; + }; + + // Rebuild a bulk's span from its 32-bit shared address at the point of use (spill-proof `LDS`). + [[nodiscard]] static _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span bulk_span(shared_bulk b) + { + return {reinterpret_cast(__cvta_shared_to_generic(b.smem32)), static_cast<::cuda::std::size_t>(b.len)}; + } + struct TempStorage : Uninitialized<_TempStorage> {}; @@ -475,7 +562,8 @@ private: offset_t head_items; unsigned int cluster_rank; unsigned int cluster_size; - offset_t resident_chunks; // rank-local chunk index at which the overflow begins + offset_t resident_chunks; // number of rank-local chunks kept resident + offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk int stream_slot_base; // SMEM slot index at which the streaming region begins offset_t overflow_chunks; // number of overflow chunks for this rank (M) int p_eff; // active streaming depth = min(PipelineStages, M) (>= 1) @@ -493,7 +581,7 @@ private: // pass can reuse it without a reload. ::cuda::std::inplace_vector loaders; ::cuda::std::optional tokens[PipelineStages]; - ::cuda::std::span pending[PipelineStages]; + shared_bulk pending[PipelineStages]; offset_t slot_chunk[PipelineStages]; _CCCL_DEVICE _CCCL_FORCEINLINE overflow_streamer( @@ -505,6 +593,7 @@ private: unsigned int cluster_rank_, unsigned int cluster_size_, offset_t resident_chunks_, + offset_t overflow_base_, int stream_slot_base_, offset_t my_chunks_) : agent(agent_) @@ -515,6 +604,7 @@ private: , cluster_rank(cluster_rank_) , cluster_size(cluster_size_) , resident_chunks(resident_chunks_) + , overflow_base(overflow_base_) , stream_slot_base(stream_slot_base_) , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) { @@ -524,7 +614,7 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE offset_t chunk_index_of(offset_t overflow_idx) const { - return cluster_rank + (resident_chunks + overflow_idx) * static_cast(cluster_size); + return cluster_rank + (overflow_base + overflow_idx) * static_cast(cluster_size); } // Grow the loader set to `count` entries (idempotent). The first-pass resident @@ -544,19 +634,18 @@ private: { const offset_t chunk_idx = chunk_index_of(overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - const bool aligned_chunk = agent.is_aligned_chunk(block_keys_base, chunk); - char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; + // The boundary chunks (unaligned head and tail) are kept resident, so every streamed chunk is fully aligned + // and uses the guard-free aligned path. + _CCCL_ASSERT(agent.is_aligned_chunk(block_keys_base, chunk), "overflow streamer received an unaligned chunk"); + char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; const ::cuda::std::span src{ block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; - if (aligned_chunk) - { - pending[stage] = - loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); - } - else - { - pending[stage] = loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); - } + // The returned span aliases `dst`, but its shared address-space provenance is not preserved across the spill of + // `pending[]`; record the rooted 32-bit shared address and rebuild the span via `bulk_span` at the read. + [[maybe_unused]] const auto loaded = + loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); + pending[stage] = {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), + static_cast(chunk.count)}; tokens[stage].emplace(loaders[stage].Commit()); slot_chunk[stage] = overflow_idx; } @@ -602,7 +691,7 @@ private: tokens[stage].reset(); } _CCCL_ASSERT(slot_chunk[stage] == o, "overflow streamer slot/chunk mapping diverged"); - agent.for_each_chunk_key(pending[stage], f); + agent.for_each_chunk_key(agent.bulk_span(pending[stage]), f); // Prefetch the chunk `p_eff` visits ahead in this direction. It maps // to the slot we just finished, so a barrier is required before the @@ -717,6 +806,31 @@ private: _CCCL_ASSERT(my_resident_chunks * offset_t{chunk_items} <= resident_slots_cap * offset_t{chunk_items}, "Dynamic shared memory block_tile is too small"); + // Boundary chunks (the unaligned head = chunk 0 and the unaligned tail = chunk chunks-1) carry hand-rolled edges + // and are kept resident so the overflow streamer only ever sees fully-aligned middle chunks. The head is already + // resident (local index 0 of its owning rank). The tail is forced into its owner's resident set - occupying the + // last resident slot - when this rank owns it, it has an unaligned end, and it would otherwise stream. The + // overflow then begins at `overflow_base`, shifted back by one in the forced case so the tail is skipped while the + // middle chunk it displaces is streamed instead. + const offset_t overflow_count = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; + offset_t tail_local = offset_t{0}; + bool force_tail_resident = false; + if constexpr (use_block_load_to_shared) + { + if (overflow_count > 0 + && (chunks - offset_t{1}) % static_cast(cluster_size) == static_cast(cluster_rank)) + { + tail_local = (chunks - offset_t{1}) / static_cast(cluster_size); + const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); + const bool tail_has_suffix = split_chunk(block_keys_base, tail_chunk).suffix != offset_t{0}; + force_tail_resident = tail_has_suffix && (tail_local >= my_resident_chunks); + } + } + const offset_t overflow_base = force_tail_resident ? (my_resident_chunks - offset_t{1}) : my_resident_chunks; + // A rank that owns both boundary chunks (the head on rank 0 and a forced tail) needs two distinct resident slots. + _CCCL_ASSERT(!(force_tail_resident && head_items != 0 && cluster_rank == 0) || my_resident_chunks >= offset_t{2}, + "streaming needs >= 2 resident slots to keep both head and tail resident"); + // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. overflow_streamer streamer( *this, @@ -727,10 +841,18 @@ private: cluster_rank, cluster_size, my_resident_chunks, + overflow_base, static_cast(resident_slots_cap), my_chunks); ::cuda::std::span resident_keys; + // 32-bit shared-window address of `resident_keys.data()`. The resident span is read once per radix pass and in + // the final filter; a 64-bit generic pointer kept live across that loop spills and reloads as a *generic* + // pointer (`LDL.64`), which demotes every key read from `LDS` to a generic `LD`. Carrying the base as a 32-bit + // shared address and rebuilding the pointer with `__cvta_shared_to_generic` at each use keeps the reads `LDS` + // even when the value spills (the cvta intrinsic re-confers shared provenance; a spilled 64-bit generic pointer + // cannot be re-anchored after the fact). + ::cuda::std::uint32_t resident_smem32 = 0; reset_hist(); __syncthreads(); @@ -751,67 +873,123 @@ private: if constexpr (use_block_load_to_shared) { - // BlockLoadToShared is non-copyable and non-movable. The loaders are owned by `streamer` and constructed - // here; the streamer then reuses these very objects for the overflow passes (continuing their Commit/Wait - // pipeline) instead of tearing them down and re-initializing the shared `load_storage`. The per-chunk tokens - // and pending spans, however, are only needed for this resident load, so they stay local. - ::cuda::std::inplace_vector tokens; - ::cuda::std::inplace_vector<::cuda::std::span, PipelineStages> pending_spans; - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - char* next_dst = key_slots; - streamer.ensure_loaders(prologue); - - for (int stage = 0; stage < prologue; ++stage) + if (my_resident_chunks > 0) { - const offset_t chunk_idx = static_cast(cluster_rank) + static_cast(stage * cluster_size); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const bool aligned_chunk = is_aligned_chunk(block_keys_base, chunk); - auto& loader = streamer.loaders[stage]; - const ::cuda::std::span src{ - block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; - if (aligned_chunk) + // BlockLoadToShared is non-copyable and non-movable. The loaders are owned by `streamer` and constructed + // here; the streamer then reuses these very objects for the overflow passes (continuing their Commit/Wait + // pipeline) instead of tearing them down and re-initializing the shared `load_storage`. The per-chunk tokens + // and pending (bulk) spans, however, are only needed for this resident load, so they stay local. Each + // resident chunk is loaded as its aligned bulk only; the two boundary edges are filled afterwards. + ::cuda::std::inplace_vector tokens; + // Per-stage resident bulk, carried as a 32-bit shared address + length rather than a 64-bit pointer (see + // `shared_bulk`): the pipeline array spills to local memory, and a spilled generic pointer would demote the + // first-pass key reads from `LDS` to a generic `LD`. + ::cuda::std::inplace_vector pending_spans; + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); + streamer.ensure_loaders(prologue); + + // Resident slot -> rank-local chunk index. Identity, except the last slot holds the forced-resident tail. + const auto resident_local = [&](offset_t slot) -> offset_t { + return (force_tail_resident && slot == my_resident_chunks - offset_t{1}) ? tail_local : slot; + }; + // Aligned bulk of the resident chunk in `slot`; empty when the chunk has no aligned interior. + const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { + const offset_t chunk_idx = + static_cast(cluster_rank) + resident_local(slot) * static_cast(cluster_size); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); + if (split.bulk == 0) + { + return {}; + } + return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; + }; + + // First resident slot's unaligned front edge (the head prefix on rank 0). Reserve an aligned-up gap in + // front of its bulk so the bulk stays 16-aligned and the edge sits contiguously right before it. + const auto first_chunk = get_chunk( + static_cast(cluster_rank) + resident_local(offset_t{0}) * static_cast(cluster_size), + segment_size_u32, + head_items); + const auto first_split = split_chunk(block_keys_base, first_chunk); + const int front_edge = static_cast(first_split.prefix); + const int sba = detail::LoadToSharedBufferAlignBytes(); + const int front_bytes = front_edge * int{sizeof(key_t)}; + const int head_bulk_off = ::cuda::round_up(front_bytes, sba); + // Write cursor as a byte offset from `key_slots`. The destination is always formed as `key_slots + offset` + // (rooted at the extern-shared base) rather than by chaining the pointers returned by `CopyAsync`: those + // returned pointers do not preserve their shared address-space provenance, so a chained cursor would be + // rejected by BlockLoadToShared's "destination must be shared memory" check even though it is numerically + // inside the block_tile. + const int resident_begin_off = head_bulk_off - front_bytes; + int next_off = head_bulk_off; + + // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail + // bulk in the forced case) ends the packed region and its suffix can be appended right after it. + for (int stage = 0; stage < prologue; ++stage) { - pending_spans.emplace_back( - loader.template CopyAsync(available_block_tile_buffer(next_dst), src)); + auto& loader = streamer.loaders[stage]; + const auto src = bulk_src(static_cast(stage)); + // The returned span aliases `key_slots + next_off`, but its shared address-space provenance is not + // preserved; we record the rooted shared address and rebuild the read span via `bulk_span` below. + [[maybe_unused]] const auto loaded = loader.template CopyAsync( + available_block_tile_buffer(key_slots + next_off), src); + pending_spans.push_back({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + next_off)), + static_cast(::cuda::std::size(src))}); + next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + tokens.emplace_back(loader.Commit()); } - else + + for (offset_t p = 0; p < my_resident_chunks; ++p) { - pending_spans.emplace_back(loader.template CopyAsync(available_block_tile_buffer(next_dst), src)); + const int stage = static_cast(p % static_cast(prologue)); + streamer.loaders[stage].Wait(::cuda::std::move(tokens[stage])); + for_each_chunk_key(bulk_span(pending_spans[stage]), add_first_pass); + + const offset_t next_p = p + static_cast(prologue); + if (next_p < my_resident_chunks) + { + const auto src = bulk_src(next_p); + [[maybe_unused]] const auto loaded = streamer.loaders[stage].template CopyAsync( + available_block_tile_buffer(key_slots + next_off), src); + pending_spans[stage] = { + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + next_off)), + static_cast(::cuda::std::size(src))}; + next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + tokens[stage] = streamer.loaders[stage].Commit(); + } } - next_dst = span_end(pending_spans[stage]); - tokens.emplace_back(loader.Commit()); - } - for (offset_t p = 0; p < my_resident_chunks; ++p) - { - const int stage = static_cast(p % static_cast(prologue)); - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - streamer.loaders[stage].Wait(::cuda::std::move(tokens[stage])); - append_contiguous_span(resident_keys, pending_spans[stage]); - for_each_chunk_key(pending_spans[stage], add_first_pass); + // Head prefix: copy into the reserved front gap, contiguous right before the first bulk. + if (front_edge > 0) + { + key_t* const edge_dst = reinterpret_cast(key_slots + resident_begin_off); + copy_edge(edge_dst, block_keys_base + first_chunk.offset, front_edge); + for_each_chunk_key({edge_dst, static_cast<::cuda::std::size_t>(front_edge)}, add_first_pass); + } - if (p + static_cast(prologue) < my_resident_chunks) + // Tail suffix: append right after the last (tail) bulk. Nothing follows, so its non-16 length is harmless. + const offset_t last_slot = my_resident_chunks - offset_t{1}; + const auto last_chunk = get_chunk( + static_cast(cluster_rank) + resident_local(last_slot) * static_cast(cluster_size), + segment_size_u32, + head_items); + const auto last_split = split_chunk(block_keys_base, last_chunk); + if (last_split.suffix > 0) { - const offset_t next_chunk_idx = static_cast(cluster_rank) - + (p + static_cast(prologue)) * static_cast(cluster_size); - const auto next_chunk = get_chunk(next_chunk_idx, segment_size_u32, head_items); - const bool next_aligned_chunk = is_aligned_chunk(block_keys_base, next_chunk); - const ::cuda::std::span src{ - block_keys_base + next_chunk.offset, static_cast<::cuda::std::size_t>(next_chunk.count)}; - if (next_aligned_chunk) - { - pending_spans[stage] = streamer.loaders[stage].template CopyAsync( - available_block_tile_buffer(next_dst), src); - } - else - { - pending_spans[stage] = - streamer.loaders[stage].template CopyAsync(available_block_tile_buffer(next_dst), src); - } - next_dst = span_end(pending_spans[stage]); - tokens[stage] = streamer.loaders[stage].Commit(); + key_t* const edge_dst = reinterpret_cast(key_slots + next_off); + copy_edge(edge_dst, + block_keys_base + last_chunk.offset + last_split.prefix + last_split.bulk, + static_cast(last_split.suffix)); + for_each_chunk_key({edge_dst, static_cast<::cuda::std::size_t>(last_split.suffix)}, add_first_pass); + next_off += static_cast(last_split.suffix) * int{sizeof(key_t)}; } + + // The resident region is one contiguous span [head edge | bulks | tail edge] for the later passes. + resident_keys = {reinterpret_cast(key_slots + resident_begin_off), + static_cast<::cuda::std::size_t>((next_off - resident_begin_off) / int{sizeof(key_t)})}; + resident_smem32 = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + resident_begin_off)); } } else @@ -914,7 +1092,10 @@ private: if constexpr (use_block_load_to_shared) { - for_each_chunk_key(resident_keys, add_hist); + // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value + // spilled across the pass loop (see `resident_smem32`). + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } else { @@ -1014,7 +1195,8 @@ private: }; if constexpr (use_block_load_to_shared) { - for_each_chunk_key(resident_keys, write_selected); + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } else { From c24d75e17b78ed42ec6ae4d7914f0da5a37e39f6 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 9 Jun 2026 19:17:08 +0200 Subject: [PATCH 024/126] Get rid of local memory and generic atomics BlockLoadToShared seems to have some problems when used for pipelining. So for now it was replaced with direct PTX. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 323 ++++++++++++------- 1 file changed, 199 insertions(+), 124 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index c8327cb13bd..2649858cd83 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -11,7 +11,8 @@ //! Histogram strategy (Pattern C): //! 1. Every block lays out `hist[num_buckets]` at the same offset in its own //! shared memory. Each block accumulates a block-private histogram using -//! block-scope atomicAdd_block (cheap, SMEM-local). +//! shared-space `red` reductions (cheap, SMEM-local): cta scope for +//! non-leaders, cluster scope for the leader (see `hist_inc`). //! 2. After a cluster-wide barrier, every non-leader block walks its //! histogram and folds its bucket counts into the leader block's `hist` //! via cluster-scope DSMEM atomics. The leader's `hist` therefore plays @@ -52,12 +53,14 @@ #include #include +#include +#include +#include +#include #include #include #include -#include #include -#include #include #include @@ -181,13 +184,14 @@ struct agent_batched_topk_cluster static_assert(LoadAlignBytes > 0); static_assert(ChunkBytes % LoadAlignBytes == 0); static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); - // The hybrid load relies on the aligned CopyAsync path being exact (no scalar guard), which requires the load + // The hybrid load relies on the aligned bulk-copy path being exact (no scalar guard), which requires the load // alignment to be at least the bulk-copy minimum alignment. static_assert(LoadAlignBytes >= detail::bulk_copy_min_align, "LoadAlignBytes must be >= bulk_copy_min_align"); static_assert(chunk_items > 0); using decomposer_t = detail::identity_decomposer_t; - using block_load_t = BlockLoadToShared; + // Resident/streamed chunks are pulled into the block_tile with elected-thread `cp.async.bulk`/TMA copies against + // raw per-stage mbarriers; see the async bulk-copy pipeline helpers below. static constexpr bool use_block_load_to_shared = THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v; @@ -220,7 +224,9 @@ struct agent_batched_topk_cluster key_prefix_t broadcast_kth; ::cuda::std::uint32_t broadcast_early_stop; typename block_scan_t::TempStorage scan_storage; - typename block_load_t::TempStorage load_storage[PipelineStages]; + // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused + // (ping-ponged) across radix passes; initialized once by `init_load_barriers`. + ::cuda::std::uint64_t load_mbar[PipelineStages]; }; [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const @@ -414,10 +420,8 @@ struct agent_batched_topk_cluster }); } - // A resident/streamed bulk in the block_tile, described by its 32-bit shared-window address plus length instead of - // a 64-bit generic pointer. Pipeline arrays of these are read across long loops and spill to local memory; a - // spilled generic pointer reloads as generic (`LDL.64`) and demotes the key reads from `LDS` to a generic `LD`, - // whereas a spilled 32-bit shared address rebuilt with `__cvta_shared_to_generic` keeps the reads `LDS`. + // A bulk in the block_tile as a 32-bit shared address + length. A spilled 32-bit shared address (rebuilt with + // `__cvta_shared_to_generic`) keeps the key reads `LDS`; a spilled 64-bit generic pointer would demote them to `LD`. struct shared_bulk { ::cuda::std::uint32_t smem32; @@ -430,6 +434,79 @@ struct agent_batched_topk_cluster return {reinterpret_cast(__cvta_shared_to_generic(b.smem32)), static_cast<::cuda::std::size_t>(b.len)}; } + // --------------------------------------------------------------------------- + // Async bulk-copy pipeline helpers (raw mbarrier + cp.async.bulk via cuda::ptx) + // --------------------------------------------------------------------------- + // Inlines BlockLoadToShared's internals (one mbarrier per stage, a single elected thread issuing the TMA copy + + // transaction arrival) without its reference member or per-call CommitToken: the only per-stage wait state is one + // bit of the per-thread `load_phase` mask, so the pipeline loops spill nothing per stage. SM 9.0+ only (the agent is + // gated behind `NV_PROVIDES_SM_90` in `Process`). + + // Init each stage mbarrier with arrival count 1: only the elected thread arrives (registering the tx byte count) and + // the `cp.async.bulk` delivers the matching count, so the phase completes. Call once, followed by a block sync. + _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barriers() + { + if (threadIdx.x == 0) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int s = 0; s < PipelineStages; ++s) + { + ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[s], 1u); + } + } + } + + // Issue one aligned global->shared (TMA) bulk copy into `dst` on stage `stage`'s mbarrier from the elected thread, + // which also arrives with the transaction byte count (an empty copy arrives with zero so the phase still completes). + // Each call must be paired, in issue order per stage, with a matching `wait_stage(stage)`. + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) + { + if (threadIdx.x != 0) + { + return; + } + const int num_bytes = static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + if (num_bytes > 0) + { + // The TMA path requires `bulk_copy_min_align` (>= 16); the cursor carries the stronger `load_align_bytes`. + _CCCL_ASSERT(::cuda::is_aligned(dst, detail::bulk_copy_min_align), + "block_tile destination must satisfy the shared-memory bulk-copy alignment"); +#if __cccl_ptx_isa >= 860 + ::cuda::ptx::cp_async_bulk( + ::cuda::ptx::space_shared, + ::cuda::ptx::space_global, + dst, + ::cuda::std::data(src), + num_bytes, + &temp_storage.load_mbar[stage]); +#else // __cccl_ptx_isa < 860 + ::cuda::ptx::cp_async_bulk( + ::cuda::ptx::space_cluster, + ::cuda::ptx::space_global, + dst, + ::cuda::std::data(src), + num_bytes, + &temp_storage.load_mbar[stage]); +#endif // __cccl_ptx_isa >= 860 + } + ::cuda::ptx::mbarrier_arrive_expect_tx( + ::cuda::ptx::sem_release, + ::cuda::ptx::scope_cta, + ::cuda::ptx::space_shared, + &temp_storage.load_mbar[stage], + static_cast<::cuda::std::uint32_t>(num_bytes)); + } + + // Wait for stage `stage`'s copy to land (all threads spin on its current parity), then flip that stage's parity bit. + _CCCL_DEVICE _CCCL_FORCEINLINE void wait_stage(int stage) + { + const ::cuda::std::uint32_t parity = (load_phase >> stage) & 1u; + while (!::cuda::ptx::mbarrier_try_wait_parity(&temp_storage.load_mbar[stage], parity)) + { + } + load_phase ^= (::cuda::std::uint32_t{1} << stage); + } + struct TempStorage : Uninitialized<_TempStorage> {}; @@ -445,6 +522,9 @@ struct agent_batched_topk_cluster NumSegmentsParameterT num_segments; char* key_slots; offset_t block_tile_capacity; + // Per-thread mbarrier phase parity, one bit per pipeline stage (see `wait_stage`); the resident load and the + // overflow streamer keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. + ::cuda::std::uint32_t load_phase{}; _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -489,6 +569,48 @@ private: } } + // --------------------------------------------------------------------------- + // Block-private histogram atomics (shared-space `red` via cuda::ptx-style inline PTX) + // --------------------------------------------------------------------------- + // A builtin `atomicAdd(&temp_storage.hist[bucket], 1)` compiles to a generic atomic (`ATOM.E`) whose 64-bit base is + // spilled and reloaded (`LDL.64`) at every update across this huge agent. A shared-space `red` instead addresses with + // the 32-bit shared address (no base to spill). Shared atomics only allow cta/cluster scope, which is exactly right: + // every writer of a given `hist` is in the same cluster. `red` (no return) matches the discarded `atomicAdd` result. + + // The 32-bit shared address of `hist[0]`. Hoisted once per histogram region so the per-key address math is a pure + // 32-bit add; recomputing it per key would reload the 64-bit generic base of `temp_storage` from the stack. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t hist_base32() const + { + return static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(temp_storage.hist)); + } + + // Increment this block's own histogram bucket by one via its 32-bit shared address. The leader (rank 0) also receives + // remote folds from the other cluster blocks (Step 2, `hist_fold_remote`), so its add must be cluster-scoped to be + // mutually atomic with them; non-leaders only touch their own `hist` before the fold, so cta scope suffices. + _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket, bool leader) + { + const ::cuda::std::uint32_t addr = base32 + static_cast<::cuda::std::uint32_t>(bucket) * sizeof(offset_t); + if (leader) + { + asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); + } + else + { + asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); + } + } + + // Step 2: a non-leader folds one bucket into the leader's histogram through DSMEM. `own_bucket_addr32` is the + // bucket's address in this block's own window; since every CTA's shared window is laid out identically, `mapa` remaps + // it to rank 0 to form the leader's `shared::cluster` address (no 64-bit pointer, no memory descriptor). Cluster + // scope makes it mutually atomic with the leader's `hist_inc` adds. + _CCCL_DEVICE _CCCL_FORCEINLINE void hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v) + { + ::cuda::std::uint32_t remote; + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "n"(0)); + asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); + } + // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram // plus identification of the bucket holding the k-th item. Each thread of // the leader block contributes `buckets_per_thread` consecutive buckets in @@ -570,19 +692,12 @@ private: bool forward = true; bool primed = false; - // Loaders are shared with the first-pass resident load: that load constructs - // them here (via `ensure_loaders`) and the streamer keeps reusing the very - // same objects for the overflow passes, continuing their Commit/Wait pipeline - // rather than tearing them down and re-initializing the shared `load_storage` - // mbarriers. During streaming, `loaders[stage]` targets the streaming slot - // `stream_slot_base + stage`. A `tokens` entry is engaged only while a copy is - // in flight (committed, not yet waited); after a wait the slot's data settles - // in SMEM and `pending`/`slot_chunk` keep describing it so the next ping-pong - // pass can reuse it without a reload. - ::cuda::std::inplace_vector loaders; - ::cuda::std::optional tokens[PipelineStages]; - shared_bulk pending[PipelineStages]; - offset_t slot_chunk[PipelineStages]; + // Stage mbarriers are shared with the resident load (`agent.temp_storage.load_mbar`); stage `stage` targets slot + // `stream_slot_base + stage`. `inflight_mask` bit `stage` is set only while a copy is in flight (issued, not yet + // waited). The slot/stage mapping is fixed, so the read span is recomputed on demand by `stage_span` rather than + // held in a spillable per-stage array; the only per-stage state is one bit each of `inflight_mask` and + // `load_phase`. + ::cuda::std::uint32_t inflight_mask = 0; _CCCL_DEVICE _CCCL_FORCEINLINE overflow_streamer( agent_batched_topk_cluster& agent_, @@ -617,37 +732,28 @@ private: return cluster_rank + (overflow_base + overflow_idx) * static_cast(cluster_size); } - // Grow the loader set to `count` entries (idempotent). The first-pass resident - // load calls this to construct the loaders, then keeps using them; the - // streamer reuses those same objects for the overflow passes. Each new loader - // binds to a distinct, previously-unused `load_storage` slot, so a live - // mbarrier is never re-initialized (which is why no `Invalidate()` is needed). - _CCCL_DEVICE _CCCL_FORCEINLINE void ensure_loaders(int count) - { - while (static_cast(loaders.size()) < count) - { - loaders.emplace_back(agent.temp_storage.load_storage[loaders.size()]); - } - } - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) { const offset_t chunk_idx = chunk_index_of(overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); // The boundary chunks (unaligned head and tail) are kept resident, so every streamed chunk is fully aligned - // and uses the guard-free aligned path. + // and uses the guard-free aligned (TMA bulk) path. _CCCL_ASSERT(agent.is_aligned_chunk(block_keys_base, chunk), "overflow streamer received an unaligned chunk"); char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; const ::cuda::std::span src{ block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; - // The returned span aliases `dst`, but its shared address-space provenance is not preserved across the spill of - // `pending[]`; record the rooted 32-bit shared address and rebuild the span via `bulk_span` at the read. - [[maybe_unused]] const auto loaded = - loaders[stage].template CopyAsync(agent.available_block_tile_buffer(dst), src); - pending[stage] = {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), - static_cast(chunk.count)}; - tokens[stage].emplace(loaders[stage].Commit()); - slot_chunk[stage] = overflow_idx; + agent.issue_bulk_copy(stage, dst, src); + inflight_mask |= (::cuda::std::uint32_t{1} << stage); + } + + // Rebuild the shared span for the chunk currently resident in `stage`'s slot without storing per-stage state: the + // slot address is a pure function of `stage` and the length is recomputed from chunk index `o`, so there is no + // spillable `pending[]` array (see `bulk_span`). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span stage_span(int stage, offset_t o) const + { + char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; + const auto chunk = agent.get_chunk(chunk_index_of(o), segment_size, head_items); + return agent.bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), chunk.count}); } // Apply `f` to every overflow key once, in the current ping-pong direction. @@ -663,9 +769,6 @@ private: if constexpr (use_block_load_to_shared) { - // Normally a no-op: the resident load already built `PipelineStages >= p_eff` - // loaders. Only the (asserted-against) tiny-config edge would grow here. - ensure_loaders(p_eff); const offset_t pe = static_cast(p_eff); // First ever call: prime the streaming slots. Subsequent calls inherit @@ -685,13 +788,12 @@ private: { const offset_t o = forward ? i : (m - 1 - i); const int stage = static_cast(o % pe); - if (tokens[stage].has_value()) + if (inflight_mask & (::cuda::std::uint32_t{1} << stage)) { - loaders[stage].Wait(::cuda::std::move(*tokens[stage])); - tokens[stage].reset(); + agent.wait_stage(stage); + inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); } - _CCCL_ASSERT(slot_chunk[stage] == o, "overflow streamer slot/chunk mapping diverged"); - agent.for_each_chunk_key(agent.bulk_span(pending[stage]), f); + agent.for_each_chunk_key(stage_span(stage, o), f); // Prefetch the chunk `p_eff` visits ahead in this direction. It maps // to the slot we just finished, so a barrier is required before the @@ -756,8 +858,8 @@ private: const auto segment_size_u32 = static_cast(segment_size); const unsigned int cluster_size = cluster.num_blocks(); - // DSMEM pointers into the leader block's shared memory. - offset_t* leader_hist = cluster.map_shared_rank(temp_storage.hist, 0); + // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` + // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, 0); // Per-block local copy of `kth_key_bits` so each key check hits the @@ -855,38 +957,33 @@ private: ::cuda::std::uint32_t resident_smem32 = 0; reset_hist(); + if constexpr (use_block_load_to_shared) + { + // Arm the stage barriers once; reused (ping-ponged) by the resident load and the overflow streamer across passes. + init_load_barriers(); + } __syncthreads(); { extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); - auto add_first_pass = [&](const key_t& key) { + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); + const bool is_leader_block = cluster_rank == 0; + auto add_first_pass = [&](const key_t& key) { const int bucket = extract_op(key); - if (cluster_rank == 0) - { - atomicAdd(&temp_storage.hist[bucket], offset_t{1}); - } - else - { - atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); - } + hist_inc(hist_smem32, bucket, is_leader_block); }; if constexpr (use_block_load_to_shared) { if (my_resident_chunks > 0) { - // BlockLoadToShared is non-copyable and non-movable. The loaders are owned by `streamer` and constructed - // here; the streamer then reuses these very objects for the overflow passes (continuing their Commit/Wait - // pipeline) instead of tearing them down and re-initializing the shared `load_storage`. The per-chunk tokens - // and pending (bulk) spans, however, are only needed for this resident load, so they stay local. Each - // resident chunk is loaded as its aligned bulk only; the two boundary edges are filled afterwards. - ::cuda::std::inplace_vector tokens; - // Per-stage resident bulk, carried as a 32-bit shared address + length rather than a 64-bit pointer (see - // `shared_bulk`): the pipeline array spills to local memory, and a spilled generic pointer would demote the - // first-pass key reads from `LDS` to a generic `LD`. - ::cuda::std::inplace_vector pending_spans; + // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). + // Chunks are written densely in slot order and read back in the same order, so the read cursor (`read_off`) + // mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed + // `pending_spans` array that would anchor surrounding state to local memory. Each chunk loads its aligned + // bulk only (boundary edges are filled afterwards); the read span is rebuilt from a rooted 32-bit shared + // address (see `bulk_span`) so a spilled cursor cannot demote the first-pass reads from `LDS` to `LD`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - streamer.ensure_loaders(prologue); // Resident slot -> rank-local chunk index. Identity, except the last slot holds the forced-resident tail. const auto resident_local = [&](offset_t slot) -> offset_t { @@ -916,11 +1013,8 @@ private: const int sba = detail::LoadToSharedBufferAlignBytes(); const int front_bytes = front_edge * int{sizeof(key_t)}; const int head_bulk_off = ::cuda::round_up(front_bytes, sba); - // Write cursor as a byte offset from `key_slots`. The destination is always formed as `key_slots + offset` - // (rooted at the extern-shared base) rather than by chaining the pointers returned by `CopyAsync`: those - // returned pointers do not preserve their shared address-space provenance, so a chained cursor would be - // rejected by BlockLoadToShared's "destination must be shared memory" check even though it is numerically - // inside the block_tile. + // Write cursor as a byte offset from `key_slots`; the destination is always `key_slots + offset` (rooted at + // the extern-shared base) so it keeps shared address-space provenance. const int resident_begin_off = head_bulk_off - front_bytes; int next_off = head_bulk_off; @@ -928,35 +1022,30 @@ private: // bulk in the forced case) ends the packed region and its suffix can be appended right after it. for (int stage = 0; stage < prologue; ++stage) { - auto& loader = streamer.loaders[stage]; const auto src = bulk_src(static_cast(stage)); - // The returned span aliases `key_slots + next_off`, but its shared address-space provenance is not - // preserved; we record the rooted shared address and rebuild the read span via `bulk_span` below. - [[maybe_unused]] const auto loaded = loader.template CopyAsync( - available_block_tile_buffer(key_slots + next_off), src); - pending_spans.push_back({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + next_off)), - static_cast(::cuda::std::size(src))}); + issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - tokens.emplace_back(loader.Commit()); } + // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in + // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. + int read_off = head_bulk_off; for (offset_t p = 0; p < my_resident_chunks; ++p) { const int stage = static_cast(p % static_cast(prologue)); - streamer.loaders[stage].Wait(::cuda::std::move(tokens[stage])); - for_each_chunk_key(bulk_span(pending_spans[stage]), add_first_pass); + wait_stage(stage); + const int read_len = static_cast(::cuda::std::size(bulk_src(p))); + for_each_chunk_key( + bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), + add_first_pass); + read_off += read_len * int{sizeof(key_t)}; const offset_t next_p = p + static_cast(prologue); if (next_p < my_resident_chunks) { - const auto src = bulk_src(next_p); - [[maybe_unused]] const auto loaded = streamer.loaders[stage].template CopyAsync( - available_block_tile_buffer(key_slots + next_off), src); - pending_spans[stage] = { - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + next_off)), - static_cast(::cuda::std::size(src))}; + const auto src = bulk_src(next_p); + issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - tokens[stage] = streamer.loaders[stage].Commit(); } } @@ -1013,9 +1102,9 @@ private: } } - // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. - // The streamer reuses the resident loaders (no re-init), so this just continues their Commit/Wait pipeline; - // `Commit()`'s own block barrier provides the necessary synchronization. + // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The + // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer + // sync. streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); @@ -1067,26 +1156,15 @@ private: identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); - // Step 1: block-private histogram. The leader uses device-scope - // `atomicAdd` so it is mutually atomic per the PTX ISA with the - // remote device-scope `atomicAdd`s that non-leaders issue against - // the same SMEM through DSMEM in Step 2. Non-leaders only write to - // their own `hist[]` and keep the cheaper `atomicAdd_block`. - // TODO(https://github.com/NVIDIA/cccl/issues/73): collapse both - // branches onto cluster-scope atomics once - // `cuda::thread_scope_cluster` is exposed in libcudacxx. - auto add_hist = [&](const key_t& key) { + // Step 1: block-private histogram via shared-space `red` (see `hist_inc`): leader uses cluster scope to be + // mutually atomic with the non-leaders' Step 2 DSMEM folds, non-leaders use the cheaper cta scope. + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); + const bool is_leader_block = cluster_rank == 0; + auto add_hist = [&](const key_t& key) { if (identify_op(key) == detail::topk::candidate_class::candidate) { const int bucket = extract_op(key); - if (cluster_rank == 0) - { - atomicAdd(&temp_storage.hist[bucket], offset_t{1}); - } - else - { - atomicAdd_block(&temp_storage.hist[bucket], offset_t{1}); - } + hist_inc(hist_smem32, bucket, is_leader_block); } }; @@ -1122,20 +1200,17 @@ private: // Step 2: non-leader blocks fold their per-bucket values // (raw counts in the reduce-then-scan path, block-local inclusive // scans in the scan-then-reduce path) into the leader's `hist` - // via DSMEM atomics. The leader skips this to avoid double-counting - // its own contribution. `atomicAdd` matches the leader's - // device-scope Step 1 atomic (see comment there). - // TODO(https://github.com/NVIDIA/cccl/issues/73): use a - // cluster-scope atomic once `cuda::thread_scope_cluster` is - // exposed in libcudacxx. + // via cluster-scope DSMEM atomics (see `hist_fold_remote`). The + // leader skips this to avoid double-counting its own contribution. if (cluster_rank != 0) { + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) { const offset_t v = temp_storage.hist[i]; if (v != 0) { - atomicAdd(leader_hist + i, v); + hist_fold_remote(hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), v); } } } From b2b4494c168de6db6d2697786386a1d6ec9da36a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 9 Jun 2026 23:57:45 +0200 Subject: [PATCH 025/126] Fix deadlocks --- cub/cub/agent/agent_batched_topk_cluster.cuh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2649858cd83..c26f15a521b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -776,6 +776,9 @@ private: // ping-pongs) is exactly the first `p_eff` chunks of this direction. if (!primed) { + // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else + // the phase advances twice and a lagging thread misses the flip and spins forever. + __syncthreads(); for (int i = 0; i < p_eff; ++i) { const offset_t o = forward ? static_cast(i) : (m - 1 - static_cast(i)); @@ -1044,6 +1047,9 @@ private: if (next_p < my_resident_chunks) { const auto src = bulk_src(next_p); + // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads + // leave the wait above would advance the phase twice, stranding a lagging waiter forever. + __syncthreads(); issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } From e56abb0969b4a82f7b61c7ea3db2fe7a2ebef292 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 10 Jun 2026 01:55:12 +0200 Subject: [PATCH 026/126] Fix smem alignment to match gmem alignment This might help on Hopper. Need to check that it does not pessimize Blackwell. Also some cleanup. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 62 ++++++++++--------- .../dispatch_batched_topk_cluster.cuh | 12 ++-- 2 files changed, 39 insertions(+), 35 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index c26f15a521b..6fbbed54eb2 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -100,25 +100,23 @@ struct alignas(16) cluster_topk_state template struct smem_block_tile_layout { + static constexpr int chunk_bytes = ChunkBytes; // one chunk == one slot; this is the slot stride static constexpr int chunk_items = ChunkBytes / int{sizeof(KeyT)}; static constexpr int load_align_items = LoadAlignBytes / int{sizeof(KeyT)}; - static constexpr int slot_alignment = - (::cuda::std::max) (LoadAlignBytes, detail::LoadToSharedBufferAlignBytes()); - // Each chunk maps to exactly one slot of ChunkBytes: boundary chunks are loaded as an aligned bulk plus a small - // hand-rolled edge, so no chunk ever needs the scalar-load guard headroom. - static constexpr int slot_stride_bytes = ::cuda::round_up(ChunkBytes, slot_alignment); // == ChunkBytes normally - static constexpr int base_padding_bytes = (alignof(KeyT) > 16) ? slot_alignment : 0; - static_assert(chunk_items >= load_align_items, "ChunkBytes must hold at least one load_align unit"); - // The aligned full-chunk buffer (no guard, since LoadAlignBytes >= bulk_copy_min_align) is exactly its byte count - // and must fit one slot. - static_assert(detail::LoadToSharedBufferSizeBytes(chunk_items) <= slot_stride_bytes, - "An aligned full chunk must fit one slot"); + // Base/slot alignment: at least the load alignment, bumped up for over-aligned key types. The slot stride is exactly + // one chunk (`ChunkBytes`), so `ChunkBytes` must be a multiple of this for consecutive slots to stay aligned, which + // also rejects a key alignment larger than a whole chunk at compile time (it makes no sense). + static constexpr int slot_alignment = (::cuda::std::max) (LoadAlignBytes, int{alignof(KeyT)}); + // Reserve one alignment quantum so the agent can round the dynamic-SMEM base up to `slot_alignment` (>= LoadAlign), + // giving every bulk-copy destination the same `load_align` alignment the gmem sources already have. + static constexpr int base_padding_bytes = slot_alignment; + static_assert(ChunkBytes % slot_alignment == 0, "ChunkBytes must be a multiple of the load and key alignment"); [[nodiscard]] _CCCL_HOST_DEVICE static constexpr ::cuda::std::uint32_t block_tile_capacity(int dynamic_smem_bytes) noexcept { const int usable_bytes = (::cuda::std::max) (0, dynamic_smem_bytes - base_padding_bytes); - const int slots = usable_bytes / slot_stride_bytes; + const int slots = usable_bytes / ChunkBytes; return static_cast<::cuda::std::uint32_t>(slots * chunk_items); } @@ -177,13 +175,11 @@ struct agent_batched_topk_cluster static constexpr int chunk_items = smem_layout_t::chunk_items; static constexpr int load_align_items = smem_layout_t::load_align_items; static constexpr int slot_alignment = smem_layout_t::slot_alignment; - static constexpr int slot_stride_bytes = smem_layout_t::slot_stride_bytes; static_assert(PipelineStages > 0); static_assert(ChunkBytes > 0); static_assert(LoadAlignBytes > 0); - static_assert(ChunkBytes % LoadAlignBytes == 0); - static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + static_assert(ChunkBytes % LoadAlignBytes == 0, "ChunkBytes must be a multiple of LoadAlignBytes"); // The hybrid load relies on the aligned bulk-copy path being exact (no scalar guard), which requires the load // alignment to be at least the bulk-copy minimum alignment. static_assert(LoadAlignBytes >= detail::bulk_copy_min_align, "LoadAlignBytes must be >= bulk_copy_min_align"); @@ -193,8 +189,15 @@ struct agent_batched_topk_cluster // Resident/streamed chunks are pulled into the block_tile with elected-thread `cp.async.bulk`/TMA copies against // raw per-stage mbarriers; see the async bulk-copy pipeline helpers below. + // The aligned bulk (TMA) path tiles gmem into `load_align`-aligned units and dense-packs them by bytes into smem. + // That is only sound when a `load_align` unit is a whole number of keys and a key has no internal padding; otherwise + // (e.g. `float3`: 12 bytes, 4-byte aligned) a `load_align` boundary would slice a key or leave bubbles in the + // resident range, so such types fall back to plain per-element loads/stores. + static constexpr bool key_is_bulk_tileable = + int{sizeof(key_t)} == int{alignof(key_t)} && LoadAlignBytes % int{sizeof(key_t)} == 0; static constexpr bool use_block_load_to_shared = - THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v; + THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v + && key_is_bulk_tileable; // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram @@ -232,12 +235,12 @@ struct agent_batched_topk_cluster [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const { const int slots = static_cast(block_tile_capacity / chunk_items); - return {key_slots, static_cast<::cuda::std::size_t>(slots * slot_stride_bytes)}; + return {key_slots, static_cast<::cuda::std::size_t>(slots * ChunkBytes)}; } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE key_t* slot_keys_unpadded(int slot) const { - return reinterpret_cast(key_slots + slot * slot_stride_bytes); + return reinterpret_cast(key_slots + slot * ChunkBytes); } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span @@ -468,8 +471,9 @@ struct agent_batched_topk_cluster const int num_bytes = static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; if (num_bytes > 0) { - // The TMA path requires `bulk_copy_min_align` (>= 16); the cursor carries the stronger `load_align_bytes`. - _CCCL_ASSERT(::cuda::is_aligned(dst, detail::bulk_copy_min_align), + // The TMA path only requires `bulk_copy_min_align` (>= 16), but the aligned base + `load_align`-multiple bulk + // packing give every destination the stronger `load_align_bytes` alignment (matching the gmem sources). + _CCCL_ASSERT(::cuda::is_aligned(dst, load_align_bytes), "block_tile destination must satisfy the shared-memory bulk-copy alignment"); #if __cccl_ptx_isa >= 860 ::cuda::ptx::cp_async_bulk( @@ -739,7 +743,7 @@ private: // The boundary chunks (unaligned head and tail) are kept resident, so every streamed chunk is fully aligned // and uses the guard-free aligned (TMA bulk) path. _CCCL_ASSERT(agent.is_aligned_chunk(block_keys_base, chunk), "overflow streamer received an unaligned chunk"); - char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; + char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; const ::cuda::std::span src{ block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; agent.issue_bulk_copy(stage, dst, src); @@ -751,7 +755,7 @@ private: // spillable `pending[]` array (see `bulk_span`). [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span stage_span(int stage, offset_t o) const { - char* const dst = agent.key_slots + (stream_slot_base + stage) * slot_stride_bytes; + char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; const auto chunk = agent.get_chunk(chunk_index_of(o), segment_size, head_items); return agent.bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), chunk.count}); } @@ -1006,16 +1010,18 @@ private: }; // First resident slot's unaligned front edge (the head prefix on rank 0). Reserve an aligned-up gap in - // front of its bulk so the bulk stays 16-aligned and the edge sits contiguously right before it. + // front of its bulk so the bulk stays `load_align`-aligned and the edge sits contiguously right before it. const auto first_chunk = get_chunk( static_cast(cluster_rank) + resident_local(offset_t{0}) * static_cast(cluster_size), segment_size_u32, head_items); - const auto first_split = split_chunk(block_keys_base, first_chunk); - const int front_edge = static_cast(first_split.prefix); - const int sba = detail::LoadToSharedBufferAlignBytes(); - const int front_bytes = front_edge * int{sizeof(key_t)}; - const int head_bulk_off = ::cuda::round_up(front_bytes, sba); + const auto first_split = split_chunk(block_keys_base, first_chunk); + const int front_edge = static_cast(first_split.prefix); + const int front_bytes = front_edge * int{sizeof(key_t)}; + // Round the first bulk up to `load_align` (not just BlockLoadToShared's 16B): with a `load_align`-aligned + // base and `load_align`-multiple bulk sizes, every densely packed resident bulk lands on a `load_align` + // boundary. + const int head_bulk_off = ::cuda::round_up(front_bytes, load_align_bytes); // Write cursor as a byte offset from `key_slots`; the destination is always `key_slots + offset` (rooted at // the extern-shared base) so it keeps shared address-space provenance. const int resident_begin_off = head_bulk_off - front_bytes; diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index a86e64d8c5f..f39415a1f86 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -136,9 +136,8 @@ __launch_bounds__(ThreadsPerBlock, MinBlocksPerSm) _CCCL_KERNEL_ATTRIBUTES void __shared__ typename agent_t::TempStorage temp_storage; extern __shared__ char topk_cluster_smem[]; char* key_slots = topk_cluster_smem; - // Without manual realignment (`alignof(key_t) <= 16`), `slot_alignment` is a stride quantum; the base is only - // required to satisfy BlockLoadToShared's 16-byte shared-memory destination alignment. - if constexpr (alignof(typename agent_t::key_t) > 16) + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` + // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. { ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); @@ -202,9 +201,8 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, __shared__ typename agent_t::TempStorage temp_storage; extern __shared__ char topk_cluster_smem[]; char* key_slots = topk_cluster_smem; - // Without manual realignment (`alignof(key_t) <= 16`), `slot_alignment` is a stride quantum; the base is only - // required to satisfy BlockLoadToShared's 16-byte shared-memory destination alignment. - if constexpr (alignof(typename agent_t::key_t) > 16) + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` + // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. { ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); @@ -598,7 +596,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const auto required_slots = ::cuda::ceil_div(required_block_tile_capacity, static_cast<::cuda::std::uint64_t>(layout_t::chunk_items)); selected_config.dynamic_smem_bytes = - layout_t::base_padding_bytes + static_cast(required_slots) * layout_t::slot_stride_bytes; + layout_t::base_padding_bytes + static_cast(required_slots) * layout_t::chunk_bytes; selected_config.block_tile_capacity = layout_t::block_tile_capacity(selected_config.dynamic_smem_bytes); } From 68eca9a142d7debf7b8bacd172f27f03bd098959 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 10 Jun 2026 02:45:05 +0200 Subject: [PATCH 027/126] Use blocked chunks This is in preparation of implementing deterministic TopK (via counting scan) --- cub/cub/agent/agent_batched_topk_cluster.cuh | 94 +++++++++++++------ ...tch2_test_segmented_topk_cluster_layout.cu | 2 +- 2 files changed, 66 insertions(+), 30 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 6fbbed54eb2..0afca927e8c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -68,6 +68,8 @@ #include +#define CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS + CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster @@ -393,6 +395,48 @@ struct agent_batched_topk_cluster : offset_t{0}; } + // Assignment of the cluster's global chunk indices `[0, chunks)` to its CTAs. A rank owns `count` chunks; its i-th + // owned chunk has global index `global_index(i) = first + i * stride`. The single mapping point lets the rest of the + // agent (resident load, streamer, per-pass scans) stay agnostic to the layout chosen by `make_chunk_partition`. + struct chunk_partition + { + offset_t first; // global index of this rank's first owned chunk + offset_t stride; // distance between consecutive owned chunks (`cluster_size` strided, `1` blocked) + offset_t count; // number of chunks owned by this rank + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t global_index(offset_t local) const + { + return first + local * stride; + } + }; + + // Decides which global chunks a cluster rank owns. Both layouts keep the unaligned head (chunk 0) on rank 0 and the + // unaligned tail (chunk `chunks-1`) on a single rank, and leave the per-chunk alignment, the resident/streaming + // split, and the streamer ping-pong untouched, because all of those depend only on the global chunk index, not on + // which rank owns it. + // + // * Strided (default): chunk `i` goes to rank `i % cluster_size`, so each CTA walks `first, first+S, first+2S, + // ...`. + // * Blocked (`CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS`): each CTA owns a contiguous run of `ceil_div(chunks, S)` + // chunks (the last non-empty rank gets the short remainder). Chunks are large enough that the per-CTA contiguous + // gmem footprint does not change L2/cache locality versus the strided walk. + // + // The macro is an off-by-default opt-in so the two layouts can be A/B benchmarked without touching call sites. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition + make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_size) const + { +#if defined(CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS) + const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_size)); + const offset_t first = static_cast(cluster_rank) * chunks_per_cta; + const offset_t count = (first < chunks) ? (::cuda::std::min) (chunks_per_cta, chunks - first) : offset_t{0}; + return {first, offset_t{1}, count}; +#else + return {static_cast(cluster_rank), + static_cast(cluster_size), + num_rank_chunks(chunks, cluster_rank, cluster_size)}; +#endif + } + // Hand-rolled per-thread copy of a small (`< load_align_items` items) unaligned edge from gmem to smem. Used for // the head prefix and tail suffix, which cannot go through the aligned (16-byte-aligned dst, guard-free) // BlockLoadToShared path. Each thread copies the same indices it later reads in the first-pass histogram, so no @@ -686,8 +730,7 @@ private: const key_t* block_keys_base; // unwrapped contiguous base (pipeline path only; null otherwise) offset_t segment_size; offset_t head_items; - unsigned int cluster_rank; - unsigned int cluster_size; + chunk_partition part; // rank -> global chunk index mapping (strided or blocked) offset_t resident_chunks; // number of rank-local chunks kept resident offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk int stream_slot_base; // SMEM slot index at which the streaming region begins @@ -709,8 +752,7 @@ private: const key_t* block_keys_base_, offset_t segment_size_, offset_t head_items_, - unsigned int cluster_rank_, - unsigned int cluster_size_, + chunk_partition part_, offset_t resident_chunks_, offset_t overflow_base_, int stream_slot_base_, @@ -720,8 +762,7 @@ private: , block_keys_base(block_keys_base_) , segment_size(segment_size_) , head_items(head_items_) - , cluster_rank(cluster_rank_) - , cluster_size(cluster_size_) + , part(part_) , resident_chunks(resident_chunks_) , overflow_base(overflow_base_) , stream_slot_base(stream_slot_base_) @@ -733,7 +774,7 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE offset_t chunk_index_of(offset_t overflow_idx) const { - return cluster_rank + (overflow_base + overflow_idx) * static_cast(cluster_size); + return part.global_index(overflow_base + overflow_idx); } _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) @@ -889,8 +930,9 @@ private: // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only // depends on the multiset of keys covered by the cluster. - const offset_t chunks = num_chunks(segment_size_u32, head_items); - const offset_t my_chunks = num_rank_chunks(chunks, cluster_rank, cluster_size); + const offset_t chunks = num_chunks(segment_size_u32, head_items); + const chunk_partition part = make_chunk_partition(chunks, cluster_rank, cluster_size); + const offset_t my_chunks = part.count; // Resident vs. streaming split. Segments that fit the all-resident coverage behave exactly as before // (`resident_slots_cap == full_slots`, no streaming). Larger segments reserve the last `PipelineStages` slots of @@ -926,10 +968,11 @@ private: bool force_tail_resident = false; if constexpr (use_block_load_to_shared) { - if (overflow_count > 0 - && (chunks - offset_t{1}) % static_cast(cluster_size) == static_cast(cluster_rank)) + // This rank owns the global tail iff its last owned chunk is chunk `chunks-1`; that chunk is then this rank's + // local index `my_chunks-1` (true for both the strided and blocked partitions). + if (overflow_count > 0 && my_chunks > 0 && part.global_index(my_chunks - offset_t{1}) == chunks - offset_t{1}) { - tail_local = (chunks - offset_t{1}) / static_cast(cluster_size); + tail_local = my_chunks - offset_t{1}; const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); const bool tail_has_suffix = split_chunk(block_keys_base, tail_chunk).suffix != offset_t{0}; force_tail_resident = tail_has_suffix && (tail_local >= my_resident_chunks); @@ -947,8 +990,7 @@ private: block_keys_base, segment_size_u32, head_items, - cluster_rank, - cluster_size, + part, my_resident_chunks, overflow_base, static_cast(resident_slots_cap), @@ -998,10 +1040,9 @@ private: }; // Aligned bulk of the resident chunk in `slot`; empty when the chunk has no aligned interior. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { - const offset_t chunk_idx = - static_cast(cluster_rank) + resident_local(slot) * static_cast(cluster_size); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const auto split = split_chunk(block_keys_base, chunk); + const offset_t chunk_idx = part.global_index(resident_local(slot)); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); if (split.bulk == 0) { return {}; @@ -1011,10 +1052,8 @@ private: // First resident slot's unaligned front edge (the head prefix on rank 0). Reserve an aligned-up gap in // front of its bulk so the bulk stays `load_align`-aligned and the edge sits contiguously right before it. - const auto first_chunk = get_chunk( - static_cast(cluster_rank) + resident_local(offset_t{0}) * static_cast(cluster_size), - segment_size_u32, - head_items); + const auto first_chunk = + get_chunk(part.global_index(resident_local(offset_t{0})), segment_size_u32, head_items); const auto first_split = split_chunk(block_keys_base, first_chunk); const int front_edge = static_cast(first_split.prefix); const int front_bytes = front_edge * int{sizeof(key_t)}; @@ -1071,10 +1110,7 @@ private: // Tail suffix: append right after the last (tail) bulk. Nothing follows, so its non-16 length is harmless. const offset_t last_slot = my_resident_chunks - offset_t{1}; - const auto last_chunk = get_chunk( - static_cast(cluster_rank) + resident_local(last_slot) * static_cast(cluster_size), - segment_size_u32, - head_items); + const auto last_chunk = get_chunk(part.global_index(resident_local(last_slot)), segment_size_u32, head_items); const auto last_split = split_chunk(block_keys_base, last_chunk); if (last_split.suffix > 0) { @@ -1097,7 +1133,7 @@ private: { for (offset_t p = 0; p < my_resident_chunks; ++p) { - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const offset_t chunk_idx = part.global_index(p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1191,7 +1227,7 @@ private: { for (offset_t p = 0; p < my_resident_chunks; ++p) { - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const offset_t chunk_idx = part.global_index(p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); @@ -1289,7 +1325,7 @@ private: { for (offset_t p = 0; p < my_resident_chunks; ++p) { - const offset_t chunk_idx = static_cast(cluster_rank) + p * static_cast(cluster_size); + const offset_t chunk_idx = part.global_index(p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); diff --git a/cub/test/catch2_test_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_segmented_topk_cluster_layout.cu index cb6da825ccc..dd1c537126b 100644 --- a/cub/test/catch2_test_segmented_topk_cluster_layout.cu +++ b/cub/test/catch2_test_segmented_topk_cluster_layout.cu @@ -23,7 +23,7 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) const int usable_bytes = dynamic_smem_bytes - layout_t::base_padding_bytes; REQUIRE(usable_bytes > 0); - const int slots = usable_bytes / layout_t::slot_stride_bytes; + const int slots = usable_bytes / layout_t::chunk_bytes; REQUIRE(slots > 0); const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); From 50abf5955fdeb69b92a7f203b2bb51678ec9915e Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 06:59:19 +0200 Subject: [PATCH 028/126] Implement determinism --- cub/cub/agent/agent_batched_topk_cluster.cuh | 446 ++++++++++++++---- .../dispatch_batched_topk_cluster.cuh | 8 + .../tuning/tuning_batched_topk_cluster.cuh | 5 + .../catch2_test_device_segmented_topk_keys.cu | 76 +++ 4 files changed, 456 insertions(+), 79 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f13b8a1f438..0f2c6bbe837 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -20,8 +20,8 @@ //! cluster-merged histogram after the second cluster sync. //! 3. The leader's thread 0 prefix-scans `hist`, identifies the bucket of //! the k-th key, and updates the cluster-shared `state`. Every block -//! reads `state.kth_key_bits` from the leader via DSMEM at the start of -//! the next pass. +//! reads `state.kth_bucket` from the leader via DSMEM at the end of each +//! pass and folds it into its own local splitter key. //! //! Output cursors live in the same cluster-shared `state` and are reached the //! same way (cluster-scope DSMEM atomics). @@ -68,7 +68,11 @@ #include -#define CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS +// Opt-in deterministic tie-break for the cluster top-k final filter. When enabled, candidates tied at the k-th key's +// prefix are selected by a cluster-wide, index-ordered scan (smallest global indices by default, largest when +// `CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST` is also defined) instead of the nondeterministic racing atomics. +// Enabling it also switches `make_chunk_partition` to the blocked layout, on which the deterministic scan depends +// (CTA-rank order == ascending contiguous global-index ranges). CUB_NAMESPACE_BEGIN @@ -85,14 +89,17 @@ struct alignas(16) cluster_topk_state OffsetT len; OutOffsetT k; - key_prefix_t kth_key_bits; + // Splitter bucket identified by the leader for the current pass. Replaces broadcasting the full multi-word + // `kth_key_bits`: every block reads this single digit through DSMEM at the end of each pass and folds it into its own + // `kth_key_bits_local` via `set_kth_key_bits`, so the full splitter key is reconstructed locally and never broadcast. + ::cuda::std::uint32_t kth_bucket; OutOffsetT out_cnt; OutOffsetT out_back_cnt; // Set by the leader after `leader_identify_kth_bucket` whenever the // identified bucket holds exactly `k` items (every candidate is part of - // the top-k). Read by every block of the cluster at the top of the next - // radix pass through DSMEM. Carried in the cluster-shared state so the - // value survives the cluster sync that ends the current pass. + // the top-k). Read by every block of the cluster at the end of each radix + // pass through DSMEM. Carried in the cluster-shared state so the value + // survives the cluster sync that ends the current pass. ::cuda::std::uint32_t early_stop; }; @@ -147,6 +154,7 @@ template ; using key_prefix_t = typename state_t::key_prefix_t; - static constexpr int threads_per_block = ThreadsPerBlock; - static constexpr int load_align_bytes = LoadAlignBytes; - static constexpr int bits_per_pass = BitsPerPass; - static constexpr int num_buckets = 1 << bits_per_pass; - using smem_layout_t = smem_block_tile_layout; - static constexpr int chunk_items = smem_layout_t::chunk_items; - static constexpr int load_align_items = smem_layout_t::load_align_items; - static constexpr int slot_alignment = smem_layout_t::slot_alignment; + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int load_align_bytes = LoadAlignBytes; + static constexpr int bits_per_pass = BitsPerPass; + static constexpr int tie_break_items_per_thread = TieBreakItemsPerThread; + static constexpr int num_buckets = 1 << bits_per_pass; + using smem_layout_t = smem_block_tile_layout; + static constexpr int chunk_items = smem_layout_t::chunk_items; + static constexpr int load_align_items = smem_layout_t::load_align_items; + static constexpr int slot_alignment = smem_layout_t::slot_alignment; static_assert(PipelineStages > 0); static_assert(ChunkBytes > 0); @@ -220,14 +229,13 @@ struct agent_batched_topk_cluster // sync the non-leader blocks fold their bucket counts into the leader's // `hist` through DSMEM atomics. `state` is meaningful only in the leader // block; the other blocks reach it exclusively through the DSMEM mapping. - // `broadcast_kth` / `broadcast_early_stop` are per-block fan-out slots used - // to share the leader-computed values across threads of each block. + // `cand_prefix` is each block's own exclusive cross-CTA candidate count for the deterministic tie-break; other blocks + // add into it through DSMEM (`add_remote_prefix`), so it must sit at an identical offset in every block's storage. struct _TempStorage { offset_t hist[num_buckets]; state_t state; - key_prefix_t broadcast_kth; - ::cuda::std::uint32_t broadcast_early_stop; + offset_t cand_prefix; typename block_scan_t::TempStorage scan_storage; // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused // (ping-ponged) across radix passes; initialized once by `init_load_barriers`. @@ -417,15 +425,16 @@ struct agent_batched_topk_cluster // // * Strided (default): chunk `i` goes to rank `i % cluster_size`, so each CTA walks `first, first+S, first+2S, // ...`. - // * Blocked (`CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS`): each CTA owns a contiguous run of `ceil_div(chunks, S)` + // * Blocked (`CUB_ENABLE_CLUSTER_TOPK_DETERMINISM`): each CTA owns a contiguous run of `ceil_div(chunks, S)` // chunks (the last non-empty rank gets the short remainder). Chunks are large enough that the per-CTA contiguous // gmem footprint does not change L2/cache locality versus the strided walk. // - // The macro is an off-by-default opt-in so the two layouts can be A/B benchmarked without touching call sites. + // The blocked layout is a hard requirement of the deterministic tie-break (its cross-CTA scan assumes CTA-rank order + // matches ascending contiguous global-index ranges), so it is selected exactly when the determinism path is enabled. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_size) const { -#if defined(CUB_ENABLE_CLUSTER_TOPK_BLOCKED_CHUNKS) +#if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_size)); const offset_t first = static_cast(cluster_rank) * chunks_per_cta; const offset_t count = (first < chunks) ? (::cuda::std::min) (chunks_per_cta, chunks - first) : offset_t{0}; @@ -652,10 +661,23 @@ private: // bucket's address in this block's own window; since every CTA's shared window is laid out identically, `mapa` remaps // it to rank 0 to form the leader's `shared::cluster` address (no 64-bit pointer, no memory descriptor). Cluster // scope makes it mutually atomic with the leader's `hist_inc` adds. - _CCCL_DEVICE _CCCL_FORCEINLINE void hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v) + _CCCL_DEVICE _CCCL_FORCEINLINE void + hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, unsigned int leader_rank) { ::cuda::std::uint32_t remote; - asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "n"(0)); + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "r"(leader_rank)); + asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); + } + + // Adds `v` to the `cand_prefix` of the CTA at cluster rank `target_rank` through DSMEM (mirrors `hist_fold_remote`: + // `mapa` this block's own `cand_prefix` address to `target_rank`, then a cluster-scope `red.add`). Used by the + // deterministic tie-break's exclusive cross-CTA candidate-count scan. + _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, offset_t v) + { + const ::cuda::std::uint32_t own = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.cand_prefix)); + ::cuda::std::uint32_t remote; + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own), "r"(target_rank)); asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); } @@ -666,7 +688,7 @@ private: // single (thread, slot) pair that owns the k-th bucket writes the per-pass // state. The caller must guarantee the leader block has finished its DSMEM // merge before invoking this. - _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket(int pass) + _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket() { // Capture `state.k` before the scan: this is the only legal window where // every thread is guaranteed to read the previous pass's value. The @@ -704,7 +726,9 @@ private: // items across finer buckets without changing the final result. temp_storage.state.early_stop = (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - detail::topk::set_kth_key_bits(temp_storage.state.kth_key_bits, pass, bucket); + // Publish only the splitter bucket; every block folds it into its own `kth_key_bits_local` at the end of the + // pass (see the pass loop), so the full splitter key is never broadcast. + temp_storage.state.kth_bucket = static_cast<::cuda::std::uint32_t>(bucket); } } } @@ -906,14 +930,29 @@ private: const auto segment_size_u32 = static_cast(segment_size); const unsigned int cluster_size = cluster.num_blocks(); + // Leader rank. The leader owns the cluster-merged histogram and the shared `state`. The deterministic tie-break + // makes the leader the *last* CTA in scan order so it never needs its own (merged-away) local candidate count: + // prefer-smallest scans ascending by rank (leader = last rank), prefer-largest scans descending (leader = rank 0). + // The nondeterministic path keeps rank 0 (unchanged codegen). +# if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) && !defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) + const unsigned int leader_rank = cluster_size - 1u; +# else + const unsigned int leader_rank = 0u; +# endif + // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, 0); + state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, leader_rank); // Per-block local copy of `kth_key_bits` so each key check hits the - // block's own SMEM rather than DSMEM during the histogram loop. + // block's own SMEM rather than DSMEM during the histogram loop. Built up one splitter digit per pass from the + // leader's published `kth_bucket` (see the pass loop), so the full key is never broadcast. key_prefix_t kth_key_bits_local = {}; + // Last splitter bucket folded into `kth_key_bits_local` (the last executed pass's bucket). Used by the + // deterministic tie-break to read this CTA's local candidate count from `hist[last_bucket]`. + [[maybe_unused]] int last_bucket = 0; + // Tracks the highest pass count that actually executed. Without early // stop this stays at `num_passes`; with early stop it captures the pass // at which we broke out so the final filter can construct its identify @@ -1016,7 +1055,7 @@ private: { extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool is_leader_block = cluster_rank == 0; + const bool is_leader_block = cluster_rank == leader_rank; auto add_first_pass = [&](const key_t& key) { const int bucket = extract_op(key); hist_inc(hist_smem32, bucket, is_leader_block); @@ -1165,35 +1204,9 @@ private: { const bool is_first_pass = (pass == 0); - // Refresh per-block local kth_key_bits from the leader's state. For - // pass 0 the bits are all zero (no filtering yet) so we skip the read. - // We also pull the leader's `early_stop` flag here so every block has - // an opportunity to bail out before doing any more histogram work. - if (!is_first_pass) - { - if (threadIdx.x == 0) - { - temp_storage.broadcast_kth = leader_state->kth_key_bits; -# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - temp_storage.broadcast_early_stop = leader_state->early_stop; -# endif - } - __syncthreads(); - kth_key_bits_local = temp_storage.broadcast_kth; -# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - if (temp_storage.broadcast_early_stop != ::cuda::std::uint32_t{0}) - { - last_pass = pass; - // Order the `broadcast_kth` load above against thread 0's - // overwrite of `broadcast_kth` in the final filter pass below. - // The normal loop exit gets this from the last iteration's - // `cluster.sync()`; the early-stop path skips that sync. - __syncthreads(); - break; - } -# endif - } - + // `kth_key_bits_local` already holds the splitter digits for passes `0..pass-1`: each block folded the leader's + // published `kth_bucket` into it at the end of the previous pass (see the end of this loop). Pass 0 needs no + // filtering (all keys are candidates), so it is handled by the fused first-pass load above. if (!is_first_pass) { // Every block (including the leader) starts each non-first pass with @@ -1207,7 +1220,7 @@ private: // Step 1: block-private histogram via shared-space `red` (see `hist_inc`): leader uses cluster scope to be // mutually atomic with the non-leaders' Step 2 DSMEM folds, non-leaders use the cheaper cta scope. const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool is_leader_block = cluster_rank == 0; + const bool is_leader_block = cluster_rank == leader_rank; auto add_hist = [&](const key_t& key) { if (identify_op(key) == detail::topk::candidate_class::candidate) { @@ -1250,7 +1263,7 @@ private: // scans in the scan-then-reduce path) into the leader's `hist` // via cluster-scope DSMEM atomics (see `hist_fold_remote`). The // leader skips this to avoid double-counting its own contribution. - if (cluster_rank != 0) + if (cluster_rank != leader_rank) { const ::cuda::std::uint32_t hist_smem32 = hist_base32(); for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) @@ -1258,7 +1271,7 @@ private: const offset_t v = temp_storage.hist[i]; if (v != 0) { - hist_fold_remote(hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), v); + hist_fold_remote(hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), v, leader_rank); } } } @@ -1268,37 +1281,311 @@ private: // Step 3: the leader walks the merged `hist` (raw counts in the // reduce-then-scan path, cluster-wide inclusive scan in the // scan-then-reduce path) and updates the cluster-shared `state`. - // Subsequent reads (next-pass refresh, last filter) all observe + // Subsequent reads (end-of-pass fold, last filter) all observe // these writes after the next cluster sync. - if (cluster_rank == 0) + if (cluster_rank == leader_rank) { - leader_identify_kth_bucket(pass); + leader_identify_kth_bucket(); } cluster.sync(); + + // End-of-pass splitter fold. Every thread reads the leader's just-published `kth_bucket` directly through DSMEM + // (a single `u32`, ordered by the `cluster.sync()` above) and folds it into its own `kth_key_bits_local`, so the + // full splitter key is reconstructed locally without any block-wide broadcast or `__syncthreads`. Runs uniformly + // for every pass (including the last), leaving `kth_key_bits_local` complete for the final filter. + { + const int bucket = static_cast(leader_state->kth_bucket); + detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); + last_bucket = bucket; + last_pass = pass + 1; + } +# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP + // Early stop: the leader sets `early_stop` when the splitter bucket holds exactly the remaining `k` candidates, + // so no further radix refinement can change the result. Every block reads the same flag from the leader and + // breaks together; `last_pass`/`kth_key_bits_local` then match what the original top-of-next-pass break produced. + if (leader_state->early_stop != ::cuda::std::uint32_t{0}) + { + break; + } +# endif } // ----------------------------------------------------------------------- - // Final filter pass: write strictly-selected items to the front of the - // output; back-fill candidates that share the k-th key's prefix bits. + // Final filter pass: write the top-k keys for this segment. Strictly- + // selected keys go to the front; the `num_kth` tied candidates fill the + // back. `kth_key_bits_local` already holds the full splitter key (folded + // from each pass's bucket above), so no broadcast is needed here. // ----------------------------------------------------------------------- + auto block_keys_out = d_key_segments_out_it[segment_id]; + const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes + + // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, + // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) + // trailing digits as smaller and erroneously reject candidates that share the identified prefix. + identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); + +# if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) + // Deterministic tie-break: select the `num_kth` tied candidates with the smallest (default) or largest + // (`CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST`) global indices via a cluster-wide, index-ordered scan instead of + // the nondeterministic racing atomics. Strictly-selected keys still go to the front; ties fill the back by rank. +# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) + constexpr bool tie_reversed = true; +# else + constexpr bool tie_reversed = false; +# endif + const out_offset_t num_selected = k - num_kth; + + // Exclusive cross-CTA candidate-count scan. Each non-leader adds its local candidate count (its surviving + // `hist[last_bucket]`) into the `cand_prefix` of every CTA that follows it in scan order; the leader is last in + // scan order (see `leader_rank`) so it adds nothing and receives the sum of all preceding CTAs. + const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; if (threadIdx.x == 0) { - temp_storage.broadcast_kth = leader_state->kth_key_bits; + temp_storage.cand_prefix = 0; } - __syncthreads(); - kth_key_bits_local = temp_storage.broadcast_kth; + cluster.sync(); + if (threadIdx.x == 0 && local_count != offset_t{0}) + { +# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) + for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow + { + add_remote_prefix(r, local_count); + } +# else + for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + { + add_remote_prefix(r, local_count); + } +# endif + } + cluster.sync(); - auto block_keys_out = d_key_segments_out_it[segment_id]; - const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes + offset_t running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) + bool tie_active = running < static_cast(num_kth); + + // Process a flat span of `count` keys already in scan order (`get_key(pos)`), tiled by + // `threads_per_block * tie_break_items_per_thread`. Selected keys go to the front via `out_cnt`; candidates get a + // BlockScan-exclusive index rank (seeded by `running`) and, if `rank < num_kth`, are written in reverse at + // `block_keys_out[k - 1 - rank]`. The running aggregate carries across tiles and across regions. + auto process_flat = [&](auto get_key, int count) { + constexpr int items = tie_break_items_per_thread; + constexpr int tile = threads_per_block * items; + for (int tile_base = 0; tile_base < count; tile_base += tile) + { + key_t keys[items]; + offset_t flags[items]; + detail::topk::candidate_class cls[items]; + bool valid[items]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + valid[i] = pos < count; + flags[i] = offset_t{0}; + if (valid[i]) + { + keys[i] = get_key(pos); + cls[i] = identify_op(keys[i]); + flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + } + } - // The pass argument controls how many radix levels of `kth_key_bits` are - // considered significant. After an early-stop break at the start of pass - // `last_pass`, only the first `last_pass` digits of the splitter have - // been set; comparing all bits would treat the (still-zero) trailing - // digits as smaller and erroneously reject candidates that share the - // identified prefix. - identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); + // Strictly-selected keys: atomic into the front. Exactly `num_selected` keys are selected cluster-wide, so + // `out_cnt` never overruns the `[0, num_selected)` region. + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && cls[i] == detail::topk::candidate_class::selected) + { + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[pos] = keys[i]; + } + } + + if (tie_active) + { + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_kth)) + { + block_keys_out[static_cast(k - 1) - static_cast(global_rank)] = keys[i]; + } + } + } + running += tile_total; + if (running >= static_cast(num_kth)) + { + tie_active = false; + } + // The next tile reuses `scan_storage`; order this tile's reads against the next ExclusiveSum's writes. + __syncthreads(); + } + } + }; + + // Uniform early-exit between regions: stop once this CTA's ties are placed and all selected are placed + // cluster-wide. Lets the common prefer-smallest case skip the (expensive) overflow re-stream entirely. + auto should_stop = [&]() -> bool { + if (tie_active) + { + return false; + } + const out_offset_t out_now = *static_cast(&leader_state->out_cnt); + return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; + }; + + // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, which is + // visited separately after the overflow because it is the globally-last chunk. + int tail_count = 0; + int front_count = span_size(resident_keys); + if constexpr (use_block_load_to_shared) + { + if (force_tail_resident) + { + tail_count = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).count; + front_count = front_count - tail_count; + } + } + + auto process_resident = [&](bool reversed) { + if constexpr (use_block_load_to_shared) + { + key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + const int fc = front_count; + if (reversed) + { + process_flat( + [&](int pos) { + return rfront[fc - 1 - pos]; + }, + fc); + } + else + { + process_flat( + [&](int pos) { + return rfront[pos]; + }, + fc); + } + } + else + { + const int rc = static_cast(my_resident_chunks); + for (int s = 0; s < rc; ++s) + { + const int local_slot = reversed ? (rc - 1 - s) : s; + const offset_t chunk_idx = part.global_index(static_cast(local_slot)); + const int cc = get_chunk(chunk_idx, segment_size_u32, head_items).count; + key_t* const ck = slot_keys_unpadded(local_slot); + if (reversed) + { + process_flat( + [&](int pos) { + return ck[cc - 1 - pos]; + }, + cc); + } + else + { + process_flat( + [&](int pos) { + return ck[pos]; + }, + cc); + } + } + } + }; + + auto process_overflow = [&](bool reversed) { + for (offset_t oo = 0; oo < overflow_count; ++oo) + { + const offset_t o = reversed ? (overflow_count - 1 - oo) : oo; + const offset_t chunk_idx = part.global_index(overflow_base + o); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const int cc = chunk.count; + const offset_t base_off = chunk.offset; + if (reversed) + { + process_flat( + [&](int pos) { + return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; + }, + cc); + } + else + { + process_flat( + [&](int pos) { + return block_keys_in[static_cast(base_off + static_cast(pos))]; + }, + cc); + } + } + }; + auto process_tail = [&](bool reversed) { + if constexpr (use_block_load_to_shared) + { + if (!force_tail_resident) + { + return; + } + key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + key_t* const tptr = rfront + front_count; + const int tc = tail_count; + if (reversed) + { + process_flat( + [&](int pos) { + return tptr[tc - 1 - pos]; + }, + tc); + } + else + { + process_flat( + [&](int pos) { + return tptr[pos]; + }, + tc); + } + } + }; + + if constexpr (tie_reversed) + { + process_tail(true); + if (!should_stop()) + { + process_overflow(true); + } + if (!should_stop()) + { + process_resident(true); + } + } + else + { + process_resident(false); + if (!should_stop()) + { + process_overflow(false); + } + if (!should_stop()) + { + process_tail(false); + } + } +# else // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM auto write_selected = [&](const key_t& key) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) @@ -1332,6 +1619,7 @@ private: } } streamer.process_pass(write_selected); +# endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast @@ -1389,7 +1677,7 @@ private: { temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; - temp_storage.state.kth_key_bits = {}; + temp_storage.state.kth_bucket = 0; temp_storage.state.out_cnt = 0; temp_storage.state.out_back_cnt = 0; temp_storage.state.early_stop = 0; diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 9adcedc25aa..67611f82c5d 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -96,6 +96,7 @@ template ChunkBytes, \ LoadAlignBytes, \ BitsPerPass, \ + TieBreakItemsPerThread, \ KeyInputItItT, \ KeyOutputItItT, \ SegmentSizeParameterT, \ @@ -325,6 +330,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr int ChunkBytes = policy.chunk_bytes; constexpr int LoadAlignBytes = policy.load_align_bytes; constexpr int BitsPerPass = policy.bits_per_pass; + constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; using key_it_t = it_value_t; using key_t = it_value_t; @@ -336,6 +342,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( ChunkBytes, LoadAlignBytes, BitsPerPass, + TieBreakItemsPerThread, KeyInputItItT, KeyOutputItItT, SegmentSizeParameterT, @@ -401,6 +408,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( ChunkBytes, LoadAlignBytes, BitsPerPass, + TieBreakItemsPerThread, KeyInputItItT, KeyOutputItItT, SegmentSizeParameterT, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 8562bf0b7e0..60674305612 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -42,6 +42,10 @@ struct cluster_topk_policy // configurations already have an occupancy of 1 due to their shared-memory usage, so allowing the compiler to // optimize for higher occupancy (fewer registers per thread) provides no benefit. int min_blocks_per_sm; + // Items-per-thread for the deterministic tie-break BlockScan sweep (final filter). Independent of `chunk_bytes`; the + // sweep tiles each region by `threads_per_block * tie_break_items_per_thread`. Only used when the deterministic + // tie-break path is enabled (`CUB_ENABLE_CLUSTER_TOPK_DETERMINISM`). + int tie_break_items_per_thread; ::cuda::std::inplace_vector launch_configs; }; @@ -66,6 +70,7 @@ make_policy(::cuda::std::inplace_vector #include #include +#include #include #include @@ -30,6 +31,20 @@ struct is_minus_zero } }; +// Maps a flat element index to one of only 8 distinct key values, so a large segment contains many duplicates and the +// k-th key's bucket holds a large tied-candidate set. Used to stress the cluster agent's candidate/tie-break path (and +// the deterministic tie-break scan when CUB_ENABLE_CLUSTER_TOPK_DETERMINISM is defined). +template +struct heavy_tie_key_op +{ + template + __host__ __device__ KeyT operator()(IndexT i) const + { + const unsigned h = static_cast(static_cast(i) * 2654435761u); + return static_cast((h >> 13) & 7u); + } +}; + enum class topk_backend { baseline, @@ -526,6 +541,67 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalig } #endif // TEST_LAUNCH != 1 +#if TEST_LAUNCH != 1 +// Heavy-tie stress/regression: collapse the keys to only a handful of distinct values so the k-th key's bucket holds a +// large set of tied candidates. This exercises the cluster agent's candidate path and, when built with +// CUB_ENABLE_CLUSTER_TOPK_DETERMINISM, the deterministic cross-CTA tie-break scan (cand_prefix + BlockScan ranks). The +// returned top-k *value* set must be correct in either mode (tied keys are equal, so which tied index is chosen is not +// observable in keys-only output, but the count of each value at the boundary must be exact). +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle heavy ties at the k-th boundary", + "[keys][segmented][topk][device][cluster]", + key_types, + select_direction_list) +{ + using key_t = c2h::get<0, TestType>; + constexpr auto direction = c2h::get<1, TestType>::value; + + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr segment_size_t static_max_segment_size = 64 * 1024; + constexpr segment_size_t static_max_k = 64 * 1024; + constexpr segment_index_t num_segments = 3; + + const segment_size_t segment_size = + GENERATE_COPY(values({segment_size_t{257}, segment_size_t{4096}, segment_size_t{64 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + // Deterministic, duplicate-heavy input (8 distinct values), materialized into a contiguous buffer so the cluster + // agent takes its resident BlockLoadToShared path. + c2h::device_vector keys_in_buffer(num_items, thrust::no_init); + thrust::tabulate(keys_in_buffer.begin(), keys_in_buffer.end(), heavy_tie_key_op{}); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys( + d_keys_in, + d_keys_out, + ::cuda::__argument::__immediate{ + segment_size, ::cuda::__argument::__bounds()}, + ::cuda::__argument::__immediate{k, ::cuda::__argument::__bounds()}, + ::cuda::__argument::__constant{}, + ::cuda::__argument::__immediate{num_segments}, + ::cuda::__argument::__immediate{num_items}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_LAUNCH != 1 + // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segmented][topk][device][float]") { From 6313b0283fb67a0564b7a0cb103267c084834134 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 16:11:12 +0200 Subject: [PATCH 029/126] Add value handling --- .../bench/segmented_topk/fixed/keys.cu | 11 +- .../bench/segmented_topk/variable/indexed.cu | 65 +++++- .../bench/segmented_topk/variable/keys.cu | 11 +- cub/cub/agent/agent_batched_topk_cluster.cuh | 219 +++++++++++++++--- .../dispatch_batched_topk_cluster.cuh | 36 +++ .../catch2_test_device_segmented_topk_keys.cu | 2 + ...catch2_test_device_segmented_topk_pairs.cu | 71 +++++- 7 files changed, 383 insertions(+), 32 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index f0a9b48050d..f0974d2bd9c 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -84,7 +84,16 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( if constexpr (selected_backend == topk_backend::cluster) { return cub::detail::batched_topk_cluster::dispatch_with_env( - d_keys_in, d_keys_out, segment_sizes, k, select_direction, num_segments, total_num_items, env); + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + select_direction, + num_segments, + total_num_items, + env); } else if constexpr (selected_backend == topk_backend::device) { diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu index 4f416d46ab6..f1bf0ec4002 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -15,6 +16,68 @@ #include "common.cuh" +enum class topk_backend +{ + baseline, + cluster, +}; + +// Which backend computes the indexed (key + segment-local-index value) top-k. The cluster backend exercises the new +// key-value-pair path through `agent_batched_topk_cluster`; the baseline backend is kept for A/B comparison. +inline constexpr topk_backend selected_backend = topk_backend::cluster; + +template +CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_indexed( + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + ValueInputItItT d_values_in, + ValueOutputItItT d_values_out, + SegmentSizeParamT segment_sizes, + KParamT k, + SelectDirectionParamT select_direction, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env) +{ + if constexpr (selected_backend == topk_backend::cluster) + { + return cub::detail::batched_topk_cluster::dispatch_with_env( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + select_direction, + num_segments, + total_num_items, + env); + } + else + { + return cub::detail::batched_topk::dispatch_with_env( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + select_direction, + num_segments, + total_num_items, + env); + } +} + // Indexed (arg-top-k) variant: each key carries a segment-local index as its value payload. The input values are // produced by a counting iterator that restarts at 0 for every segment, so indices are not (pre-)materialized in global // memory @@ -73,7 +136,7 @@ void decode_style_variable_topk_indexed( auto env = cub_bench_env(alloc, launch); // TODO(bgruber): call the public API once available _CCCL_TRY_CUDA_API( - cub::detail::batched_topk::dispatch_with_env, + batched_topk_indexed, "batched topk failed", d_keys_in, d_keys_out, diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index f792de307b4..0c9af8b7eaf 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -59,7 +59,16 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( if constexpr (selected_backend == topk_backend::cluster) { return cub::detail::batched_topk_cluster::dispatch_with_env( - d_keys_in, d_keys_out, segment_sizes, k, select_direction, num_segments, total_num_items, env); + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + select_direction, + num_segments, + total_num_items, + env); } else if constexpr (selected_backend == topk_backend::device) { diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 0f2c6bbe837..65a955c4715 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -59,6 +59,7 @@ #include #include #include +#include #include #include #include @@ -157,6 +158,8 @@ template ; - using key_t = it_value_t; + using key_it_t = it_value_t; + using key_t = it_value_t; + using value_it_t = it_value_t; + using value_t = it_value_t; + + // Keys-only when the value payload type is `cub::NullType` (mirrors the baseline batched top-k agent). The value + // iterators are then never dereferenced and the final filter's value writes are compiled out. + static constexpr bool is_keys_only = ::cuda::std::is_same_v; using segment_size_val_t = typename ::cuda::__argument::__traits::element_type; using num_segments_val_t = typename ::cuda::__argument::__traits::element_type; @@ -573,6 +582,8 @@ struct agent_batched_topk_cluster _TempStorage& temp_storage; KeyInputItItT d_key_segments_it; KeyOutputItItT d_key_segments_out_it; + ValueInputItItT d_value_segments_it; + ValueOutputItItT d_value_segments_out_it; SegmentSizeParameterT segment_sizes; KParameterT k_param; SelectDirectionParameterT select_directions; @@ -587,6 +598,8 @@ struct agent_batched_topk_cluster TempStorage& temp_storage_, KeyInputItItT d_key_segments_it_, KeyOutputItItT d_key_segments_out_it_, + ValueInputItItT d_value_segments_it_, + ValueOutputItItT d_value_segments_out_it_, SegmentSizeParameterT segment_sizes_, KParameterT k_param_, SelectDirectionParameterT select_directions_, @@ -596,6 +609,8 @@ struct agent_batched_topk_cluster : temp_storage(temp_storage_.Alias()) , d_key_segments_it(d_key_segments_it_) , d_key_segments_out_it(d_key_segments_out_it_) + , d_value_segments_it(d_value_segments_it_) + , d_value_segments_out_it(d_value_segments_out_it_) , segment_sizes(segment_sizes_) , k_param(k_param_) , select_directions(select_directions_) @@ -1319,6 +1334,20 @@ private: auto block_keys_out = d_key_segments_out_it[segment_id]; const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes + // For each key written to `block_keys_out[pos]`, the associated input value at the key's segment-local index + // `seg_idx` is loaded naively from gmem and written to `block_vals_out[pos]`. `seg_idx` is recomputed per region in + // the sweeps below. The per-segment value iterators are derived *inside* the `is_keys_only` guard: in keys-only + // builds the value iterators-of-iterators are `cub::NullType**` (null), so indexing them with `segment_id` here + // would dereference a null pointer; `segment_id` is loop-invariant, so the compiler hoists these out of the writes. + const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { + if constexpr (!is_keys_only) + { + auto block_vals_in = d_value_segments_it[segment_id]; + auto block_vals_out = d_value_segments_out_it[segment_id]; + block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; + } + }; + // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) // trailing digits as smaller and erroneously reject candidates that share the identified prefix. @@ -1367,7 +1396,9 @@ private: // `threads_per_block * tie_break_items_per_thread`. Selected keys go to the front via `out_cnt`; candidates get a // BlockScan-exclusive index rank (seeded by `running`) and, if `rank < num_kth`, are written in reverse at // `block_keys_out[k - 1 - rank]`. The running aggregate carries across tiles and across regions. - auto process_flat = [&](auto get_key, int count) { + // `get_idx(pos)` returns the segment-local index of the key `get_key(pos)` returns, used only to load the value + // payload for written keys (compiled out in keys-only builds). + auto process_flat = [&](auto get_key, auto get_idx, int count) { constexpr int items = tie_break_items_per_thread; constexpr int tile = threads_per_block * items; for (int tile_base = 0; tile_base < count; tile_base += tile) @@ -1397,8 +1428,10 @@ private: { if (valid[i] && cls[i] == detail::topk::candidate_class::selected) { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[pos] = keys[i]; + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); } } @@ -1415,7 +1448,10 @@ private: const offset_t global_rank = running + excl[i]; if (global_rank < static_cast(num_kth)) { - block_keys_out[static_cast(k - 1) - static_cast(global_rank)] = keys[i]; + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); } } } @@ -1454,6 +1490,11 @@ private: } } + // Segment-local base of the resident-front span. The deterministic path always uses the blocked partition, so the + // front chunks `[part.first, part.first + front_chunks)` are contiguous in segment order and pack densely into the + // resident SMEM region; element `pos` of the front therefore maps to `front_seg_base + pos`. + const offset_t front_seg_base = get_chunk(part.first, segment_size_u32, head_items).offset; + auto process_resident = [&](bool reversed) { if constexpr (use_block_load_to_shared) { @@ -1465,6 +1506,9 @@ private: [&](int pos) { return rfront[fc - 1 - pos]; }, + [&](int pos) { + return front_seg_base + static_cast(fc - 1 - pos); + }, fc); } else @@ -1473,6 +1517,9 @@ private: [&](int pos) { return rfront[pos]; }, + [&](int pos) { + return front_seg_base + static_cast(pos); + }, fc); } } @@ -1483,7 +1530,9 @@ private: { const int local_slot = reversed ? (rc - 1 - s) : s; const offset_t chunk_idx = part.global_index(static_cast(local_slot)); - const int cc = get_chunk(chunk_idx, segment_size_u32, head_items).count; + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const int cc = chunk.count; + const offset_t base_off = chunk.offset; key_t* const ck = slot_keys_unpadded(local_slot); if (reversed) { @@ -1491,6 +1540,9 @@ private: [&](int pos) { return ck[cc - 1 - pos]; }, + [&](int pos) { + return base_off + static_cast(cc - 1 - pos); + }, cc); } else @@ -1499,6 +1551,9 @@ private: [&](int pos) { return ck[pos]; }, + [&](int pos) { + return base_off + static_cast(pos); + }, cc); } } @@ -1519,6 +1574,9 @@ private: [&](int pos) { return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; }, + [&](int pos) { + return base_off + static_cast(cc - 1 - pos); + }, cc); } else @@ -1527,6 +1585,9 @@ private: [&](int pos) { return block_keys_in[static_cast(base_off + static_cast(pos))]; }, + [&](int pos) { + return base_off + static_cast(pos); + }, cc); } } @@ -1542,12 +1603,18 @@ private: key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); key_t* const tptr = rfront + front_count; const int tc = tail_count; + // The forced-resident tail is the globally-last chunk `chunks-1`; its segment-local base is that chunk's + // offset. + const offset_t tail_seg_base = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).offset; if (reversed) { process_flat( [&](int pos) { return tptr[tc - 1 - pos]; }, + [&](int pos) { + return tail_seg_base + static_cast(tc - 1 - pos); + }, tc); } else @@ -1556,6 +1623,9 @@ private: [&](int pos) { return tptr[pos]; }, + [&](int pos) { + return tail_seg_base + static_cast(pos); + }, tc); } } @@ -1586,39 +1656,132 @@ private: } } # else // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM - auto write_selected = [&](const key_t& key) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) + if constexpr (is_keys_only) + { + auto write_selected = [&](const key_t& key) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[pos] = key; + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + } + } + }; + if constexpr (use_block_load_to_shared) { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[pos] = key; + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } - else if (res == detail::topk::candidate_class::candidate) + else { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); - if (back_pos < num_kth) + for (offset_t p = 0; p < my_resident_chunks; ++p) { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; + const offset_t chunk_idx = part.global_index(p); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } - }; - if constexpr (use_block_load_to_shared) - { - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + streamer.process_pass(write_selected); } else { - for (offset_t p = 0; p < my_resident_chunks; ++p) + // Pair (key + value) path. The key write order and `out_cnt`/`out_back_cnt` atomics are unchanged; for each + // written key we additionally load its value payload from gmem at the key's segment-local index `seg_idx` and + // store it at the same output slot. Resident keys are still reused from SMEM (only the value is fetched from + // gmem); overflow keys *and* values are re-read straight from gmem in chunk order, so no value streaming is + // needed for this first cut. + auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[pos] = key; + write_value(pos, seg_idx); + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + write_value(pos, seg_idx); + } + } + }; + + // Iterate a contiguous run of `count` keys whose element `local` has segment-local index `base_off + local`. + // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). + auto write_run = [&](auto get_key, offset_t base_off, int count) { + const int iterations = ::cuda::ceil_div(count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < count) + { + write_selected_idx(get_key(local), base_off + static_cast(local)); + } + }); + }; + + if constexpr (use_block_load_to_shared) { - const offset_t chunk_idx = part.global_index(p); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + // Resident keys are densely packed in slot order; each chunk's keys are contiguous, so a running cursor over + // the packed region recovers per-chunk spans. The last slot holds the forced-resident tail when applicable. + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + int cursor = 0; + for (offset_t s = 0; s < my_resident_chunks; ++s) + { + const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; + const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + const int cc = chunk.count; + write_run( + [&](int local) { + return rk[cursor + local]; + }, + base_off, + cc); + cursor += cc; + } + } + else + { + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + write_run( + [&](int local) { + return chunk_keys[local]; + }, + base_off, + chunk.count); + } + } + + // Overflow chunks: re-read keys and values straight from gmem in chunk order. + for (offset_t o = 0; o < overflow_count; ++o) + { + const auto chunk = get_chunk(part.global_index(overflow_base + o), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + write_run( + [&](int local) { + return block_keys_in[static_cast(base_off + static_cast(local))]; + }, + base_off, + chunk.count); } } - streamer.process_pass(write_selected); # endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM // Final cluster barrier: hold every block in the cluster until all DSMEM diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 67611f82c5d..47e66df2dad 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -99,6 +99,8 @@ template TieBreakItemsPerThread, \ KeyInputItItT, \ KeyOutputItItT, \ + ValueInputItItT, \ + ValueOutputItItT, \ SegmentSizeParameterT, \ KParameterT, \ SelectDirectionParameterT, \ @@ -285,6 +303,8 @@ struct force_emit_kernel .doit(static_kernel, \ d_key_segments_it, \ d_key_segments_out_it, \ + d_value_segments_it, \ + d_value_segments_out_it, \ segment_sizes, \ k_param, \ select_directions, \ @@ -297,6 +317,8 @@ struct force_emit_kernel template (nullptr), + static_cast(nullptr), segment_sizes, k, select_direction, diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 5564a71110b..d678e8fa241 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -4,6 +4,7 @@ #include "insert_nested_NVTX_range_guard.h" #include +#include #include #include @@ -47,8 +48,76 @@ struct flag_intra_segment_duplicates template flag_intra_segment_duplicates(ItemItT, SegIdItT) -> flag_intra_segment_duplicates; +enum class topk_backend +{ + baseline, + cluster, +}; + +inline constexpr topk_backend selected_backend = topk_backend::cluster; + +// Routes the key-value (pairs) top-k to either the baseline or the cluster backend, threading both the key and value +// iterators-of-iterators through. The cluster backend exercises the key-value-pair path in +// `agent_batched_topk_cluster`. +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParamT segment_sizes, + KParamT k, + SelectDirectionT select_direction, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream = nullptr) +{ + if constexpr (selected_backend == topk_backend::cluster) + { + return cub::detail::batched_topk_cluster::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_direction, + num_segments, + total_num_items_guarantee, + stream); + } + else + { + return cub::detail::batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_direction, + num_segments, + total_num_items_guarantee, + stream); + } +} + // %PARAM% TEST_LAUNCH lid 0:1:2 -DECLARE_LAUNCH_WRAPPER(cub::detail::batched_topk::dispatch, batched_topk_pairs); +DECLARE_LAUNCH_WRAPPER(dispatch_batched_topk_pairs, batched_topk_pairs); // Total segment size using max_segment_size_list = c2h::enum_type_list; From aa96a9aeb8e4640f69ea1ab025a4b64444f8a5ad Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 17:39:27 +0200 Subject: [PATCH 030/126] Make streaming decision more flexible Now each CTA can decide on their own at runtime if/how many chunks need to be streamed. This should help us keeping more chunks resident in smem especially in edge cases where the segment does almost fit into dsmem. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 30 ++++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 65a955c4715..5fbfab936aa 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -988,15 +988,27 @@ private: const chunk_partition part = make_chunk_partition(chunks, cluster_rank, cluster_size); const offset_t my_chunks = part.count; - // Resident vs. streaming split. Segments that fit the all-resident coverage behave exactly as before - // (`resident_slots_cap == full_slots`, no streaming). Larger segments reserve the last `PipelineStages` slots of - // the block_tile as a round-robin streaming region and keep `full_slots - PipelineStages` slots resident; the - // overflow chunks are re-streamed from gmem on every pass by `streamer`. The launch coverage check still reserves - // one extra chunk for the possible unaligned head. - const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); - const offset_t all_resident_capacity = - smem_layout_t::template cluster_tile_capacity(static_cast(cluster_size), block_tile_capacity); - const bool needs_streaming = segment_size_u32 > all_resident_capacity; + // Resident vs. streaming split, decided independently per CTA. A CTA whose owned-chunk count fits its resident + // slots (`my_chunks <= full_slots`) keeps every chunk resident and streams nothing; only a CTA that actually + // overflows reserves the last `PipelineStages` slots of its block_tile as a round-robin streaming region (keeping + // `full_slots - PipelineStages` resident) and re-streams its overflow chunks from gmem on every pass via + // `streamer`. + // + // This is purely a local decision: each CTA only ever loads/scans its own chunks (resident SMEM or its own gmem + // overflow), and the cross-CTA traffic (histogram fold, leader `state`, the deterministic `cand_prefix` scan) plus + // every `cluster.sync()` are reached uniformly regardless of how many chunks any CTA streams - so CTAs need not + // agree on the split. `my_chunks` already folds in this segment's actual base alignment (an unaligned head costs + // exactly one extra chunk via `num_chunks`/`head_chunk_items`), so no head reserve is needed here; that differs + // from the host-side launch selection, which must provision a one-chunk margin (`cluster_tile_capacity`) because it + // picks a single launch-wide `(cluster_size, smem)` from only the segment-size *upper bound*. + // + // Versus a cluster-uniform split (`chunks > full_slots * cluster_size`, which forces every CTA to stream): the + // busiest rank streams the same amount either way, but every other rank now stays fully resident, cutting cluster + // gmem traffic and SMEM pressure. Both schemes stream under the exact same global condition (some rank overflows + // iff `ceil_div(chunks, cluster_size) > full_slots`), so the `full_slots > PipelineStages` reservation guarantee + // the dispatch already provisions for is unchanged. + const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); + const bool needs_streaming = my_chunks > full_slots; _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), "block_tile too small to reserve a streaming region"); const offset_t resident_slots_cap = From 406f4c8185a76d697c34ac0a565953e66fe2558b Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 18:10:22 +0200 Subject: [PATCH 031/126] Hide streaming latency behind resident work This applies to the histogram passes. Naturally it only applies to the first wave of streamed data. Any more waves still rely on pipeline depth. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 93 ++++++++++++++------ 1 file changed, 65 insertions(+), 28 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5fbfab936aa..49b061b1de7 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -840,12 +840,17 @@ private: return agent.bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), chunk.count}); } - // Apply `f` to every overflow key once, in the current ping-pong direction. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) + // Apply `f` to every overflow key once, in the current ping-pong direction. `mid()` is invoked exactly once, + // after the prefetch loads for this pass's first reload wave (the first `p_eff` visits) have been issued but before + // they are waited on, so the caller's resident-chunk work overlaps those in-flight bulk copies. The histogram is + // order-independent, so folding resident keys between the streamer's load issue and its wait is safe. `mid` must be + // uniform across the block and contain no unmatched block barrier (`for_each_chunk_key` has none). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { if (overflow_chunks == 0) { + mid(); return; } @@ -871,8 +876,10 @@ private: primed = true; } - for (offset_t i = 0; i < m; ++i) - { + // Consume overflow visit `i`: wait for its slot, fold its keys via `f`, then prefetch the chunk `p_eff` visits + // ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite the data the + // block was just reading). + const auto consume = [&](offset_t i) { const offset_t o = forward ? i : (m - 1 - i); const int stage = static_cast(o % pe); if (inflight_mask & (::cuda::std::uint32_t{1} << stage)) @@ -882,9 +889,6 @@ private: } agent.for_each_chunk_key(stage_span(stage, o), f); - // Prefetch the chunk `p_eff` visits ahead in this direction. It maps - // to the slot we just finished, so a barrier is required before the - // async copy can overwrite the data the block was just reading. const offset_t ni = i + pe; if (ni < m) { @@ -892,13 +896,32 @@ private: __syncthreads(); issue_load(stage, no); } + }; + + // Phase 1: consume the first `p_eff` visits (the chunks reused from the previous pass, already resident in the + // streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. + const offset_t split = (::cuda::std::min) (pe, m); + for (offset_t i = 0; i < split; ++i) + { + consume(i); + } + + // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. + mid(); + + // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). + for (offset_t i = split; i < m; ++i) + { + consume(i); } forward = !forward; } else { - // Generic fallback: overflow keys are read straight from gmem each pass - // (no SMEM reuse), but the walk still snakes for L2 locality. + // Generic fallback: no async SMEM pipeline, so resident work cannot hide load latency here. Fold the resident + // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no + // SMEM reuse), with the walk still snaking for L2 locality. + mid(); for (offset_t i = 0; i < m; ++i) { const offset_t o = forward ? i : (m - 1 - i); @@ -918,6 +941,14 @@ private: forward = !forward; } } + + // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in + // by the BlockLoadToShared pipeline (rather than already resident in SMEM). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) + { + process_pass(static_cast(f), [] {}); + } }; // ------------------------------------------------------------------------- @@ -1256,27 +1287,33 @@ private: } }; - if constexpr (use_block_load_to_shared) - { - // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value - // spilled across the pass loop (see `resident_smem32`). - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) + // Resident-chunk histogram, deferred into the streamer so it overlaps the streamer's in-flight first reload + // wave (see `process_pass`). The histogram is order-independent, so folding resident keys between the + // streamer's load issue and its wait does not change the result. + const auto fold_resident_hist = [&] { + if constexpr (use_block_load_to_shared) { - const offset_t chunk_idx = part.global_index(p); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); + // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value + // spilled across the pass loop (see `resident_smem32`). + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } - } + else + { + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const offset_t chunk_idx = part.global_index(p); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); + } + } + }; - // Re-stream the overflow chunks into this pass's histogram. Ping-pongs direction and reuses the boundary - // chunks left resident by the previous pass. - streamer.process_pass(add_hist); + // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the + // first wave of reload bulk copies. Ping-pongs direction and reuses the boundary chunks left resident by the + // previous pass. + streamer.process_pass(add_hist, fold_resident_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` From 34e22a7f29ab006013452f6a1a7ffd71218d25a5 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 19:57:15 +0200 Subject: [PATCH 032/126] Fix non-deterministic key-value case It was reading streaming chunks from gmem naively instead of using the existing staging/pipeline. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 123 ++++++++++++++----- 1 file changed, 90 insertions(+), 33 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 49b061b1de7..d53d7847ad3 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -485,6 +485,24 @@ struct agent_batched_topk_cluster }); } + // Like `for_each_chunk_key`, but also hands `f` each key's segment-local index `base_off + local`, where `base_off` + // is the segment-local offset of the chunk's first element. The pair path uses that index to fetch the key's value + // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const + { + const int chunk_count = span_size(chunk_keys); + const int iterations = ::cuda::ceil_div(chunk_count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk_count) + { + f(::cuda::std::data(chunk_keys)[local], base_off + static_cast(local)); + } + }); + } + // A bulk in the block_tile as a 32-bit shared address + length. A spilled 32-bit shared address (rebuilt with // `__cvta_shared_to_generic`) keeps the key reads `LDS`; a spilled 64-bit generic pointer would demote them to `LD`. struct shared_bulk @@ -840,13 +858,15 @@ private: return agent.bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), chunk.count}); } - // Apply `f` to every overflow key once, in the current ping-pong direction. `mid()` is invoked exactly once, - // after the prefetch loads for this pass's first reload wave (the first `p_eff` visits) have been issued but before - // they are waited on, so the caller's resident-chunk work overlaps those in-flight bulk copies. The histogram is - // order-independent, so folding resident keys between the streamer's load issue and its wait is safe. `mid` must be - // uniform across the block and contain no unmatched block barrier (`for_each_chunk_key` has none). - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` currently resident in + // the streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from + // gmem (generic fallback). `mid()` is invoked exactly once, after the prefetch loads for this pass's first reload + // wave (the first `p_eff` visits) have been issued but before they are waited on, so the caller's resident-chunk + // work overlaps those in-flight bulk copies. The two public entry points (`process_pass` / `process_pass_indexed`) + // only differ in whether the applied callable receives the key alone or the key plus its segment-local index, so + // they share this loop verbatim. `mid` must be uniform across the block and contain no unmatched block barrier. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid) { if (overflow_chunks == 0) { @@ -876,9 +896,9 @@ private: primed = true; } - // Consume overflow visit `i`: wait for its slot, fold its keys via `f`, then prefetch the chunk `p_eff` visits - // ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite the data the - // block was just reading). + // Consume overflow visit `i`: wait for its slot, fold its keys via `block_apply`, then prefetch the chunk + // `p_eff` visits ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite + // the data the block was just reading). const auto consume = [&](offset_t i) { const offset_t o = forward ? i : (m - 1 - i); const int stage = static_cast(o % pe); @@ -887,7 +907,7 @@ private: agent.wait_stage(stage); inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); } - agent.for_each_chunk_key(stage_span(stage, o), f); + block_apply(stage, o); const offset_t ni = i + pe; if (ni < m) @@ -927,19 +947,32 @@ private: const offset_t o = forward ? i : (m - 1 - i); const offset_t chunk_idx = chunk_index_of(o); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + generic_apply(chunk); + } + forward = !forward; + } + } + + // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap + // semantics of `mid`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + { + run_pass( + [&](int stage, offset_t o) { + agent.for_each_chunk_key(stage_span(stage, o), f); + }, + [&](const auto& chunk) { + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); detail::transform::unrolled_for(iterations, [&](int j) { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { - const key_t key = - block_keys_in[static_cast(chunk.offset + static_cast(local))]; - f(key); + f(block_keys_in[static_cast(chunk.offset + static_cast(local))]); } }); - } - forward = !forward; - } + }, + static_cast(mid)); } // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in @@ -949,6 +982,38 @@ private: { process_pass(static_cast(f), [] {}); } + + // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. The pair + // final filter needs that index to fetch each selected key's value payload from gmem, while still reusing the + // overflow keys from the streaming SMEM pipeline (block-load path) instead of re-reading them from gmem. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) + { + run_pass( + [&](int stage, offset_t o) { + const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; + agent.for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + }, + [&](const auto& chunk) { + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + detail::transform::unrolled_for(iterations, [&](int j) { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk.count) + { + const offset_t seg_idx = chunk.offset + static_cast(local); + f(block_keys_in[static_cast(seg_idx)], seg_idx); + } + }); + }, + static_cast(mid)); + } + + // Overload with no interleaved work (resident keys are processed before the streamer in the pair final filter). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f) + { + process_pass_indexed(static_cast(f), [] {}); + } }; // ------------------------------------------------------------------------- @@ -1745,9 +1810,9 @@ private: { // Pair (key + value) path. The key write order and `out_cnt`/`out_back_cnt` atomics are unchanged; for each // written key we additionally load its value payload from gmem at the key's segment-local index `seg_idx` and - // store it at the same output slot. Resident keys are still reused from SMEM (only the value is fetched from - // gmem); overflow keys *and* values are re-read straight from gmem in chunk order, so no value streaming is - // needed for this first cut. + // store it at the same output slot. Keys are reused exactly as in the keys-only path - resident keys from SMEM, + // overflow keys from the streaming SMEM pipeline via `process_pass_indexed` (the generic fallback re-reads them + // from gmem); only the values are fetched from gmem (no value streaming). auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) @@ -1818,18 +1883,10 @@ private: } } - // Overflow chunks: re-read keys and values straight from gmem in chunk order. - for (offset_t o = 0; o < overflow_count; ++o) - { - const auto chunk = get_chunk(part.global_index(overflow_base + o), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - write_run( - [&](int local) { - return block_keys_in[static_cast(base_off + static_cast(local))]; - }, - base_off, - chunk.count); - } + // Overflow chunks: reuse the keys from the streaming SMEM pipeline (block-load path; only the generic fallback + // re-reads them from gmem), and fetch each selected key's value at its segment-local index `seg_idx`. `write_run` + // already wrote out resident chunks above, so no interleaved `mid` work is left to overlap here. + streamer.process_pass_indexed(write_selected_idx); } # endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM From 0f605246237e540f97f704d9e656e3b37477b6cc Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 20:15:39 +0200 Subject: [PATCH 033/126] Apply latency hiding to non-deterministic filter Same as previous commit for the histogram passes. Hide first wave of streaming behind work on resident items. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 127 ++++++++++--------- 1 file changed, 66 insertions(+), 61 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index d53d7847ad3..2300b8a4f4c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1007,13 +1007,6 @@ private: }, static_cast(mid)); } - - // Overload with no interleaved work (resident keys are processed before the streamer in the pair final filter). - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f) - { - process_pass_indexed(static_cast(f), [] {}); - } }; // ------------------------------------------------------------------------- @@ -1789,30 +1782,37 @@ private: } } }; - if constexpr (use_block_load_to_shared) - { - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) + // Fold the resident keys as the streamer's `mid` work so they overlap the first wave of overflow reloads. The + // writes are order-independent atomics, so interleaving resident and overflow output is safe, and the resident + // SMEM slots are disjoint from the streaming slots, so `mid` reads never race the in-flight loads. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) { - const offset_t chunk_idx = part.global_index(p); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } - } - streamer.process_pass(write_selected); + else + { + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const offset_t chunk_idx = part.global_index(p); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + } + } + }; + streamer.process_pass(write_selected, fold_resident); } else { - // Pair (key + value) path. The key write order and `out_cnt`/`out_back_cnt` atomics are unchanged; for each - // written key we additionally load its value payload from gmem at the key's segment-local index `seg_idx` and - // store it at the same output slot. Keys are reused exactly as in the keys-only path - resident keys from SMEM, - // overflow keys from the streaming SMEM pipeline via `process_pass_indexed` (the generic fallback re-reads them - // from gmem); only the values are fetched from gmem (no value streaming). + // Pair (key + value) path. The `out_cnt`/`out_back_cnt` atomics are unchanged; for each written key we + // additionally load its value payload from gmem at the key's segment-local index `seg_idx` and store it at the + // same output slot. Keys are reused exactly as in the keys-only path - resident keys from SMEM, overflow keys + // from the streaming SMEM pipeline via `process_pass_indexed` (the generic fallback re-reads them from gmem); + // only the values are fetched from gmem (no value streaming). As in the keys-only path the output order is not + // preserved (the resident keys are folded in as the streamer's `mid` work), which the non-deterministic path + // permits. auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) @@ -1846,47 +1846,52 @@ private: }); }; - if constexpr (use_block_load_to_shared) - { - // Resident keys are densely packed in slot order; each chunk's keys are contiguous, so a running cursor over - // the packed region recovers per-chunk spans. The last slot holds the forced-resident tail when applicable. - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - int cursor = 0; - for (offset_t s = 0; s < my_resident_chunks; ++s) + // Fold the resident keys (and their values) as the streamer's `mid` work so they overlap the first wave of + // overflow reloads, exactly as in the keys-only path: order-independent atomic writes into a disjoint output, and + // resident SMEM slots disjoint from the streaming slots, so `mid` never races the in-flight loads. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) { - const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; - const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - const int cc = chunk.count; - write_run( - [&](int local) { - return rk[cursor + local]; - }, - base_off, - cc); - cursor += cc; + // Resident keys are densely packed in slot order; each chunk's keys are contiguous, so a running cursor over + // the packed region recovers per-chunk spans. The last slot holds the forced-resident tail when applicable. + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + int cursor = 0; + for (offset_t s = 0; s < my_resident_chunks; ++s) + { + const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; + const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + const int cc = chunk.count; + write_run( + [&](int local) { + return rk[cursor + local]; + }, + base_off, + cc); + cursor += cc; + } } - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) + else { - const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - write_run( - [&](int local) { - return chunk_keys[local]; - }, - base_off, - chunk.count); + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + write_run( + [&](int local) { + return chunk_keys[local]; + }, + base_off, + chunk.count); + } } - } + }; // Overflow chunks: reuse the keys from the streaming SMEM pipeline (block-load path; only the generic fallback - // re-reads them from gmem), and fetch each selected key's value at its segment-local index `seg_idx`. `write_run` - // already wrote out resident chunks above, so no interleaved `mid` work is left to overlap here. - streamer.process_pass_indexed(write_selected_idx); + // re-reads them from gmem), and fetch each selected key's value at its segment-local index `seg_idx`. The + // resident keys above are folded in as the streamer's `mid` work to hide the first reload wave's latency. + streamer.process_pass_indexed(write_selected_idx, fold_resident); } # endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM From f830031f92d4908e0f321388708797f5a9378cf0 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 11 Jun 2026 20:48:06 +0200 Subject: [PATCH 034/126] Avoid inflating the number of streamed chunks When the segment size is only slightly bigger than what can be held resident in dsmem, avoid sizing the number of streaming slots to the full pipeline depth. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 59 +++++++++++++------- 1 file changed, 38 insertions(+), 21 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2300b8a4f4c..474fea1f9ef 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -771,15 +771,18 @@ private: // --------------------------------------------------------------------------- // Re-streams the per-rank "overflow" chunks (those that do not fit in the // resident SMEM region) from gmem through a small, fixed, round-robin set of - // `PipelineStages` streaming slots. The same object is reused for every radix - // pass and the final filter. It ping-pongs the iteration order across calls so - // the `PipelineStages` boundary chunks that one pass leaves resident in the - // streaming slots are reused by the next pass with no reload; in the limit - // where the overflow fits entirely in the streaming slots (`overflow <= - // PipelineStages`), the chunks are loaded once and never reloaded. The - // resident region is unaffected: it lives in the slots `[0, resident_slots)`, - // the streaming region in `[stream_slot_base, stream_slot_base + - // PipelineStages)`. + // `p_eff` (<= `PipelineStages`) streaming slots. The same object is reused for + // every radix pass and the final filter. It ping-pongs the iteration order + // across calls so the `p_eff` boundary chunks that one pass leaves resident in + // the streaming slots are reused by the next pass with no reload; the remaining + // `overflow_chunks - p_eff` chunks are reloaded from gmem on each pass. The + // caller right-sizes the reservation to `p_eff = min(PipelineStages, excess)` + // (where `excess = my_chunks - full_slots`), so a streaming rank always has + // `overflow_chunks = excess + p_eff` and reloads exactly `excess` chunks per + // pass - the reserved slots only ever buy reuse of the `p_eff` boundary chunks, + // never a reload-free pass. The resident region is unaffected: it lives in the + // slots `[0, resident_slots)`, the streaming region in `[stream_slot_base, + // stream_slot_base + p_eff)`. struct overflow_streamer { agent_batched_topk_cluster& agent; @@ -792,7 +795,7 @@ private: offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk int stream_slot_base; // SMEM slot index at which the streaming region begins offset_t overflow_chunks; // number of overflow chunks for this rank (M) - int p_eff; // active streaming depth = min(PipelineStages, M) (>= 1) + int p_eff; // streaming region size = reserved streaming slots (<= PipelineStages, <= M, >= 1) bool forward = true; bool primed = false; @@ -813,6 +816,7 @@ private: offset_t resident_chunks_, offset_t overflow_base_, int stream_slot_base_, + int stream_slots_, offset_t my_chunks_) : agent(agent_) , block_keys_in(block_keys_in_) @@ -825,8 +829,13 @@ private: , stream_slot_base(stream_slot_base_) , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) { - const int m = static_cast((::cuda::std::min) (overflow_chunks, static_cast(PipelineStages))); - p_eff = (m > 0) ? m : 1; + // `p_eff` is the streaming region size the caller reserved at `[stream_slot_base, stream_slot_base + + // stream_slots)`; using all of it as pipeline stages keeps the ping-pong reuse maximal. It is `<= M` whenever + // there is overflow (`M = excess + stream_slots >= stream_slots`); the `>= 1` floor only matters for the no-op + // (`M == 0`) case, where the streamer never touches a slot. + p_eff = (stream_slots_ > 0) ? stream_slots_ : 1; + _CCCL_ASSERT(overflow_chunks == 0 || p_eff <= static_cast(overflow_chunks), + "streaming depth exceeds the overflow chunk count"); } _CCCL_DEVICE _CCCL_FORCEINLINE offset_t chunk_index_of(offset_t overflow_idx) const @@ -1079,9 +1088,8 @@ private: // Resident vs. streaming split, decided independently per CTA. A CTA whose owned-chunk count fits its resident // slots (`my_chunks <= full_slots`) keeps every chunk resident and streams nothing; only a CTA that actually - // overflows reserves the last `PipelineStages` slots of its block_tile as a round-robin streaming region (keeping - // `full_slots - PipelineStages` resident) and re-streams its overflow chunks from gmem on every pass via - // `streamer`. + // overflows reserves a round-robin streaming region at the tail of its block_tile and re-streams its overflow + // chunks from gmem on every pass via `streamer`. // // This is purely a local decision: each CTA only ever loads/scans its own chunks (resident SMEM or its own gmem // overflow), and the cross-CTA traffic (histogram fold, leader `state`, the deterministic `cand_prefix` scan) plus @@ -1096,16 +1104,24 @@ private: // gmem traffic and SMEM pressure. Both schemes stream under the exact same global condition (some rank overflows // iff `ceil_div(chunks, cluster_size) > full_slots`), so the `full_slots > PipelineStages` reservation guarantee // the dispatch already provisions for is unchanged. + // + // Right-size the streaming region. A CTA that overflows its resident slots by only `excess = my_chunks - + // full_slots` chunks needs at most `excess` streaming slots to cycle that overflow through gmem; reserving the full + // `PipelineStages` would needlessly route up-to-`PipelineStages` extra chunks through the streaming machinery (and + // its per-prefetch `__syncthreads()`) that could instead stay resident and be read once. So reserve + // `stream_slots = min(PipelineStages, excess)` slots: deep overflows (`excess >= PipelineStages`) behave exactly as + // before, while barely-overflowing segments keep the rest resident. The per-pass gmem reload is `excess` either way + // (the streamer reuses its `p_eff = stream_slots` boundary chunks across passes); this only shrinks the streaming + // region - and grows the resident region - which can only relax the `>= 2 resident slots` head+tail guarantee + // below. const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); const bool needs_streaming = my_chunks > full_slots; _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), "block_tile too small to reserve a streaming region"); - const offset_t resident_slots_cap = - needs_streaming - ? ((full_slots > static_cast(PipelineStages)) - ? full_slots - static_cast(PipelineStages) - : offset_t{1}) - : full_slots; + const offset_t stream_slots = + needs_streaming ? (::cuda::std::min) (static_cast(PipelineStages), my_chunks - full_slots) + : offset_t{0}; + const offset_t resident_slots_cap = full_slots - stream_slots; const offset_t my_resident_chunks = (::cuda::std::min) (my_chunks, resident_slots_cap); // Resident chunks stay within the first `resident_slots_cap` slots; the streaming region occupies the slots // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. @@ -1149,6 +1165,7 @@ private: my_resident_chunks, overflow_base, static_cast(resident_slots_cap), + static_cast(stream_slots), my_chunks); ::cuda::std::span resident_keys; From cb52dac5bea7162e670bbf50ed7582e529ad5c80 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 12 Jun 2026 20:04:54 +0200 Subject: [PATCH 035/126] Optimize histogram loop --- cub/cub/agent/agent_batched_topk_cluster.cuh | 88 ++++++++++++++----- .../dispatch_batched_topk_cluster.cuh | 32 +++---- .../tuning/tuning_batched_topk_cluster.cuh | 14 +-- 3 files changed, 84 insertions(+), 50 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 474fea1f9ef..ea71d456082 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -52,6 +52,7 @@ #include #include +#include #include #include #include @@ -150,7 +151,7 @@ struct smem_block_tile_layout // it is not a template parameter; per-block block_tile layout is still controlled // by the template parameters below. template 0); + static_assert(HistogramItemsPerThread > 0, "histogram_items_per_thread must be positive"); static_assert(ChunkBytes > 0); static_assert(LoadAlignBytes > 0); static_assert(ChunkBytes % LoadAlignBytes == 0, "ChunkBytes must be a multiple of LoadAlignBytes"); @@ -383,7 +386,7 @@ struct agent_batched_topk_cluster const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); const auto aligned_begin = ::cuda::round_up(begin, la); - const auto aligned_end = (end / la) * la; + const auto aligned_end = ::cuda::round_down(end, la); if (aligned_begin > aligned_end) { // The chunk lies strictly between two load_align boundaries (no aligned point inside): the whole chunk is an @@ -471,17 +474,64 @@ struct agent_batched_topk_cluster } } - template - _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const + // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in + // `[0, chunk_count)`), processing tiles of `histogram_items_per_thread * threads_per_block` keys. Each tile is loaded + // into registers by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the + // histogram passes: `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would + // interleave each load with its atomic instead of hoisting the whole load wave ahead. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - const int chunk_count = span_size(chunk_keys); - const int iterations = ::cuda::ceil_div(chunk_count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < chunk_count) + constexpr int tile = histogram_items_per_thread * threads_per_block; + const int tid = static_cast(threadIdx.x); + const int full_tiles = ::cuda::round_down(chunk_count, tile); + + _CCCL_PRAGMA_NOUNROLL() + for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) + { + key_t regs[histogram_items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + regs[t] = data[tile_base + t * threads_per_block + tid]; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) { - f(::cuda::std::data(chunk_keys)[local]); + const int local = tile_base + t * threads_per_block + tid; + apply(regs[t], local); } + } + + if (full_tiles < chunk_count) + { + key_t regs[histogram_items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + regs[t] = data[local]; + } + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + apply(regs[t], local); + } + } + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const + { + for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { + f(key); }); } @@ -492,14 +542,8 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const { - const int chunk_count = span_size(chunk_keys); - const int iterations = ::cuda::ceil_div(chunk_count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < chunk_count) - { - f(::cuda::std::data(chunk_keys)[local], base_off + static_cast(local)); - } + for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { + f(key, base_off + static_cast(local)); }); } @@ -973,7 +1017,7 @@ private: }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + detail::transform::unrolled_for(iterations, [&](int j) { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { @@ -1005,7 +1049,7 @@ private: }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + detail::transform::unrolled_for(iterations, [&](int j) { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { @@ -1309,7 +1353,7 @@ private: const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + detail::transform::unrolled_for(iterations, [&](int j) { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { @@ -1854,7 +1898,7 @@ private: // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). auto write_run = [&](auto get_key, offset_t base_off, int count) { const int iterations = ::cuda::ceil_div(count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + detail::transform::unrolled_for(iterations, [&](int j) { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < count) { diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 47e66df2dad..4f2f346d1c5 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -91,7 +91,7 @@ struct launch_config // width via cooperative groups. template # define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ auto static_kernel = device_segmented_topk_cluster_kernel_static< \ ThreadsPerBlock, \ - UnrollFactor, \ + HistogramItemsPerThread, \ PipelineStages, \ ChunkBytes, \ LoadAlignBytes, \ @@ -346,22 +346,22 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( using SelectDirectionParameterT = decltype(select_directions); // Clusters are SM 9.0+ only and the tuning is currently identical across // CCs, so pin the selector query to the minimum supported CC. - constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); - constexpr int ThreadsPerBlock = policy.threads_per_block; - constexpr int MinBlocksPerSm = policy.min_blocks_per_sm; - constexpr int UnrollFactor = policy.unroll_factor; - constexpr int PipelineStages = policy.pipeline_stages; - constexpr int ChunkBytes = policy.chunk_bytes; - constexpr int LoadAlignBytes = policy.load_align_bytes; - constexpr int BitsPerPass = policy.bits_per_pass; - constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; + constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); + constexpr int ThreadsPerBlock = policy.threads_per_block; + constexpr int MinBlocksPerSm = policy.min_blocks_per_sm; + constexpr int HistogramItemsPerThread = policy.histogram_items_per_thread; + constexpr int PipelineStages = policy.pipeline_stages; + constexpr int ChunkBytes = policy.chunk_bytes; + constexpr int LoadAlignBytes = policy.load_align_bytes; + constexpr int BitsPerPass = policy.bits_per_pass; + constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; using key_it_t = it_value_t; using key_t = it_value_t; using layout_t = smem_block_tile_layout; using agent_t = agent_batched_topk_cluster< ThreadsPerBlock, - UnrollFactor, + HistogramItemsPerThread, PipelineStages, ChunkBytes, LoadAlignBytes, @@ -429,7 +429,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr auto dynamic_kernel = &device_segmented_topk_cluster_kernel< ThreadsPerBlock, MinBlocksPerSm, - UnrollFactor, + HistogramItemsPerThread, PipelineStages, ChunkBytes, LoadAlignBytes, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 60674305612..23a65f39b57 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -33,18 +33,12 @@ struct cluster_topk_launch_config struct cluster_topk_policy { int threads_per_block; - int unroll_factor; + int histogram_items_per_thread; int pipeline_stages; int chunk_bytes; int load_align_bytes; int bits_per_pass; - // Minimum number of CTAs per SM passed to the dynamic-cluster kernel's launch bounds. Defaults to 1: most cluster - // configurations already have an occupancy of 1 due to their shared-memory usage, so allowing the compiler to - // optimize for higher occupancy (fewer registers per thread) provides no benefit. int min_blocks_per_sm; - // Items-per-thread for the deterministic tie-break BlockScan sweep (final filter). Independent of `chunk_bytes`; the - // sweep tiles each region by `threads_per_block * tie_break_items_per_thread`. Only used when the deterministic - // tie-break path is enabled (`CUB_ENABLE_CLUSTER_TOPK_DETERMINISM`). int tie_break_items_per_thread; ::cuda::std::inplace_vector launch_configs; }; @@ -64,7 +58,7 @@ make_policy(::cuda::std::inplace_vector Date: Sat, 13 Jun 2026 01:39:44 +0200 Subject: [PATCH 036/126] Improve non-TMA fallback path and its testing Fallback path streams directly through registers, so use all smem slots for resident chunks. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 24 ++-- ...catch2_test_device_segmented_topk_pairs.cu | 108 ++++++++++++++++++ 2 files changed, 125 insertions(+), 7 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index ea71d456082..8ffc7d4a311 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1158,13 +1158,23 @@ private: // (the streamer reuses its `p_eff = stream_slots` boundary chunks across passes); this only shrinks the streaming // region - and grows the resident region - which can only relax the `>= 2 resident slots` head+tail guarantee // below. - const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); - const bool needs_streaming = my_chunks > full_slots; - _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), - "block_tile too small to reserve a streaming region"); - const offset_t stream_slots = - needs_streaming ? (::cuda::std::min) (static_cast(PipelineStages), my_chunks - full_slots) - : offset_t{0}; + // + // The streaming region is the async SMEM pipeline, which only exists on the TMA path. The generic fallback has no + // pipeline: it still keeps its resident chunks in SMEM (read once, reused across passes), but re-reads its overflow + // chunks straight from gmem every pass without ever staging them in a slot. Reserving streaming slots there would + // just leave SMEM idle, so the fallback reserves none and devotes the whole block_tile to resident chunks - the + // more chunks it keeps resident, the fewer it re-reads from gmem each pass. + const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); + [[maybe_unused]] const bool needs_streaming = my_chunks > full_slots; + offset_t stream_slots = offset_t{0}; + if constexpr (use_block_load_to_shared) + { + _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), + "block_tile too small to reserve a streaming region"); + stream_slots = needs_streaming + ? (::cuda::std::min) (static_cast(PipelineStages), my_chunks - full_slots) + : offset_t{0}; + } const offset_t resident_slots_cap = full_slots - stream_slots; const offset_t my_resident_chunks = (::cuda::std::min) (my_chunks, resident_slots_cap); // Resident chunks stay within the first `resident_slots_cap` slots; the streaming region occupies the slots diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index d678e8fa241..9e0d527c724 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -10,6 +10,7 @@ #include #include #include +#include #include @@ -320,6 +321,113 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small fixed-size segments" REQUIRE(expected_keys == keys_out_buffer); } +// Exercise the pair (key+value) generic fallback only in the float build (`TEST_TYPES == 1`); the test fixes its own +// key/value types, so repeating it for every key-type axis would only waste time on the expensive 1 Mi-element runs. +#if TEST_LAUNCH != 1 && TEST_TYPES == 1 +template +struct cast_to_key_op +{ + template + __host__ __device__ KeyT operator()(T x) const + { + return static_cast(x); + } +}; + +// Yields, for each segment, a *non-contiguous* iterator over that segment's keys (an integral counting iterator cast +// to the key type). Feeding the cluster top-k a non-contiguous key iterator makes `use_block_load_to_shared` false, so +// the agent takes its generic (non-BlockLoadToShared) path. Segment `seg` produces keys +// [seg * segment_size, (seg + 1) * segment_size), so the flattened input equals the identity sequence. +template +struct counting_segment_keys_op +{ + SegmentSizeT segment_size; + + template + __host__ __device__ auto operator()(IndexT seg) const + { + return cuda::make_transform_iterator( + cuda::make_counting_iterator(static_cast(seg) * segment_size), cast_to_key_op{}); + } +}; + +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments through a non-contiguous key iterator", + "[pairs][segmented][topk][device][cluster]", + select_direction_list) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + // Selection direction comes from the compile-time test axis. + constexpr auto direction = c2h::get<0, TestType>::value; + + // The counting-iterator key source is non-contiguous, so the agent uses its generic path rather than + // BlockLoadToShared - the only place the pair value writes flow through the generic resident/overflow code. + // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so the 1 Mi-element segments stream + // (incl. an unaligned `- 31` tail), while the 128 Ki-element segment validates the generic resident path (no + // streaming) through the same code. Keeping the largest total below 2^24 makes every key an exact float, and every + // value index an exact `int32`. + constexpr segment_size_t static_max_segment_size = 1024 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 3; + + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + // Non-contiguous key input: segment `seg` is the counting iterator [seg * segment_size, (seg + 1) * segment_size). + auto d_keys_in = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), counting_segment_keys_op{segment_size}); + + // Value payload = global flattened index, so each output value indexes back into the identity `expected_keys`. Only + // the key iterator drives the resident/streaming path; the value iterator type is immaterial (values are fetched + // lazily per selected key), so a strided counting iterator suffices. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + + // Outputs are real buffers (the output iterators stay contiguous; only the key input drives the generic path). + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + ::cuda::__argument::__immediate{ + segment_size, ::cuda::__argument::__bounds()}, + ::cuda::__argument::__immediate{k, ::cuda::__argument::__bounds()}, + ::cuda::__argument::__constant{}, + ::cuda::__argument::__immediate{num_segments}, + ::cuda::__argument::__immediate{num_items}); + + // The flattened input is the identity sequence, so materialize it for the standard pair verification. + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::sequence(expected_keys.begin(), expected_keys.end()); + + // Verify (before sorting) that values stayed associated with their keys and that no index repeats within a segment. + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // Verify the selected keys are the correct top-k. + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_LAUNCH != 1 && TEST_TYPES == 1 + C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segments", "[pairs][segmented][topk][device]", key_types, From 9a23a333326c63b3104e1d460a0ab2c54f1b0bde Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 13 Jun 2026 02:54:57 +0200 Subject: [PATCH 037/126] Clean up - Extent triple-chevron launch facilities to allow for dynamic cluster launch. - Remove references to transform's loop abstraction. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 47 ++++++---- cub/cub/detail/launcher/cuda_runtime.cuh | 28 +++++- .../dispatch_batched_topk_cluster.cuh | 86 +++++++------------ .../cuda/detail/core/triple_chevron_launch.h | 48 ++++++++--- 4 files changed, 121 insertions(+), 88 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 8ffc7d4a311..7f29cba1c3d 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -39,12 +39,11 @@ #endif // no system header #include -#include #include #include #include #include -#include +#include #include #include @@ -1007,8 +1006,10 @@ private: } // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap - // semantics of `mid`. - template + // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass the + // tuning parameter matching their context (`histogram_items_per_thread` for the histogram passes, + // `tie_break_items_per_thread` for the non-deterministic final filter). + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { run_pass( @@ -1017,29 +1018,31 @@ private: }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + _CCCL_PRAGMA_UNROLL(UnrollFactor) + for (int j = 0; j < iterations; ++j) + { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { f(block_keys_in[static_cast(chunk.offset + static_cast(local))]); } - }); + } }, static_cast(mid)); } // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in // by the BlockLoadToShared pipeline (rather than already resident in SMEM). - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) { - process_pass(static_cast(f), [] {}); + process_pass(static_cast(f), [] {}); } // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. The pair // final filter needs that index to fetch each selected key's value payload from gmem, while still reusing the // overflow keys from the streaming SMEM pipeline (block-load path) instead of re-reading them from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) { run_pass( @@ -1049,14 +1052,16 @@ private: }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + _CCCL_PRAGMA_UNROLL(UnrollFactor) + for (int j = 0; j < iterations; ++j) + { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { const offset_t seg_idx = chunk.offset + static_cast(local); f(block_keys_in[static_cast(seg_idx)], seg_idx); } - }); + } }, static_cast(mid)); } @@ -1363,7 +1368,9 @@ private: const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + _CCCL_PRAGMA_UNROLL(histogram_items_per_thread) + for (int j = 0; j < iterations; ++j) + { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { @@ -1372,14 +1379,14 @@ private: chunk_keys[local] = key; add_first_pass(key); } - }); + } } } // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer // sync. - streamer.process_pass(add_first_pass); + streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -1442,7 +1449,7 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the boundary chunks left resident by the // previous pass. - streamer.process_pass(add_hist, fold_resident_hist); + streamer.process_pass(add_hist, fold_resident_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -1873,7 +1880,7 @@ private: } } }; - streamer.process_pass(write_selected, fold_resident); + streamer.process_pass(write_selected, fold_resident); } else { @@ -1908,13 +1915,15 @@ private: // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). auto write_run = [&](auto get_key, offset_t base_off, int count) { const int iterations = ::cuda::ceil_div(count, threads_per_block); - detail::transform::unrolled_for(iterations, [&](int j) { + _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread) + for (int j = 0; j < iterations; ++j) + { const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < count) { write_selected_idx(get_key(local), base_off + static_cast(local)); } - }); + } }; // Fold the resident keys (and their values) as the streamer's `mid` work so they overlap the first wave of @@ -1962,7 +1971,7 @@ private: // Overflow chunks: reuse the keys from the streaming SMEM pipeline (block-load path; only the generic fallback // re-reads them from gmem), and fetch each selected key's value at its segment-local index `seg_idx`. The // resident keys above are folded in as the streamer's `mid` work to hide the first reload wave's latency. - streamer.process_pass_indexed(write_selected_idx, fold_resident); + streamer.process_pass_indexed(write_selected_idx, fold_resident); } # endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM diff --git a/cub/cub/detail/launcher/cuda_runtime.cuh b/cub/cub/detail/launcher/cuda_runtime.cuh index ab8e3a95cdb..c67d3f9c79d 100644 --- a/cub/cub/detail/launcher/cuda_runtime.cuh +++ b/cub/cub/detail/launcher/cuda_runtime.cuh @@ -42,10 +42,16 @@ struct TripleChevronFactory } CUB_RUNTIME_FUNCTION THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron operator()( - dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, bool dependent_launch = false) const + dim3 grid, + dim3 block, + ::cuda::std::size_t shared_mem, + ::cudaStream_t stream, + bool dependent_launch = false, + dim3 cluster_dim = dim3{0, 0, 0}) const { __assert_pdl_allowed(dependent_launch); - return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid, block, shared_mem, stream, dependent_launch); + return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + grid, block, shared_mem, stream, dependent_launch, cluster_dim); } template @@ -141,6 +147,24 @@ struct TripleChevronFactory { return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); } + + // Opt the kernel in to non-portable cluster widths (i.e. >8). + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_non_portable_cluster_allowed(Kernel kernel_ptr) + { + return CubDebug(::cudaFuncSetAttribute( + reinterpret_cast(kernel_ptr), cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + } + + // Query the largest thread-block cluster the hardware can launch for `kernel_ptr` under `config`. The cluster + // dimension in `config` is ignored by the query. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + max_potential_cluster_size(int& cluster_size, Kernel kernel_ptr, const ::cudaLaunchConfig_t* config) + { + return CubDebug( + ::cudaOccupancyMaxPotentialClusterSize(&cluster_size, reinterpret_cast(kernel_ptr), config)); + } }; } // namespace detail diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 4f2f346d1c5..26608d7e6ae 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -31,6 +31,7 @@ #include #include +#include #include #include #include @@ -237,37 +238,6 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, } #endif // CUB_RDC_ENABLED -// ----------------------------------------------------------------------------- -// NVCC kernel-emission workaround -// ----------------------------------------------------------------------------- -// NVCC `-rdc=false` only emits the device side of a templated `__global__` -// when it sees a triple-chevron in the same TU; CUDA 13's -// `-static-global-template-stub=true` makes the host stub internally linked -// on top. Address-of and `cudaLaunchKernelEx` are not enough. The host path -// must use `cudaLaunchKernelExC` for the cluster attribute, so instantiating -// `force_emit_kernel::emit` parses a (dead) chevron in its place. -// `Args` are deduced from the function-pointer type to avoid repeating the -// dispatch's template parameter list. -// -// See https://developer.nvidia.com/blog/cuda-c-compiler-updates-impacting-elf-visibility-and-linkage/. -template -struct force_emit_kernel; - -template -struct force_emit_kernel -{ - [[noreturn]] _CCCL_HOST static void emit(Args... args) - { - _CCCL_ASSERT(false, "force_emit_kernel::emit must never be called"); - // Unreachable; present only so NVCC emits `Kernel` for this TU. - if (false) - { - Kernel<<<1, 1>>>(args...); - } - _CCCL_UNREACHABLE(); - } -}; - // ----------------------------------------------------------------------------- // Dispatch // ----------------------------------------------------------------------------- @@ -444,22 +414,21 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionParameterT, NumSegmentsParameterT>; - // Force NVCC to emit the device side of `dynamic_kernel` for this TU; see - // `force_emit_kernel` above. - [[maybe_unused]] constexpr auto force_emit = &force_emit_kernel::emit; - NV_IF_TARGET( NV_IS_HOST, ({ + // The launcher's `doit` carries the triple-chevron that makes NVCC emit `dynamic_kernel` for this TU, and + // performs the cluster launch via `cudaLaunchKernelEx`. The factory also wraps the pre-launch driver queries. + detail::TripleChevronFactory launcher_factory{}; + // Opt in to non-portable cluster widths (>8 on Hopper). - if (const auto error = CubDebug(cudaFuncSetAttribute( - reinterpret_cast(dynamic_kernel), cudaFuncAttributeNonPortableClusterSizeAllowed, 1))) + if (const auto error = launcher_factory.set_non_portable_cluster_allowed(dynamic_kernel)) { return error; } - // Reused across the probe and the launch; `clusterDim.x` is a placeholder - // until after `cudaOccupancyMaxPotentialClusterSize` (which ignores it). + // Config used only for the occupancy probe below; the final launch goes through `launcher_factory`. + // `clusterDim.x` is a placeholder since `cudaOccupancyMaxPotentialClusterSize` ignores it. ::cudaLaunchAttribute cluster_attr{}; cluster_attr.id = ::cudaLaunchAttributeClusterDimension; cluster_attr.val.clusterDim.x = 1; @@ -565,8 +534,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( cfg.dynamicSmemBytes = static_cast(candidate_dynamic_smem); int candidate_hw_max_cluster_blocks = 0; - if (const auto error = CubDebug(cudaOccupancyMaxPotentialClusterSize( - &candidate_hw_max_cluster_blocks, reinterpret_cast(dynamic_kernel), &cfg))) + if (const auto error = + launcher_factory.max_potential_cluster_size(candidate_hw_max_cluster_blocks, dynamic_kernel, &cfg)) { return error; } @@ -642,7 +611,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const int cluster_blocks = selected_config.cluster_blocks; const auto block_tile_capacity = selected_config.block_tile_capacity; - cfg.dynamicSmemBytes = static_cast(selected_config.dynamic_smem_bytes); const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); @@ -651,21 +619,25 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaErrorInvalidValue; } - cfg.gridDim = dim3(static_cast(grid_blocks), 1, 1); - cluster_attr.val.clusterDim.x = static_cast(cluster_blocks); - - if (const auto error = CubDebug(::cudaLaunchKernelEx( - &cfg, - dynamic_kernel, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - block_tile_capacity))) + // The cluster dimension routes the launch through `cudaLaunchKernelEx`; the sibling triple-chevron in + // `doit_host` forces NVCC to emit `dynamic_kernel` for this TU. + if (const auto error = CubDebug( + launcher_factory(dim3(static_cast(grid_blocks), 1, 1), + dim3(static_cast(ThreadsPerBlock), 1, 1), + static_cast<::cuda::std::size_t>(selected_config.dynamic_smem_bytes), + stream, + /*dependent_launch=*/false, + dim3(static_cast(cluster_blocks), 1, 1)) + .doit(dynamic_kernel, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + block_tile_capacity))) { return error; } diff --git a/thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h b/thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h index ccea3deea0d..93466b2d8cb 100644 --- a/thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h +++ b/thrust/thrust/system/cuda/detail/core/triple_chevron_launch.h @@ -28,15 +28,23 @@ struct _CCCL_VISIBILITY_HIDDEN triple_chevron Size const shared_mem; bool const dependent_launch; cudaStream_t const stream; + dim3 const cluster_dim; /// @param dependent_launch Launches the kernel using programmatic dependent launch if available. + /// @param cluster_dim Launches the kernel with the given thread-block cluster dimension. A zero `x` means no cluster. THRUST_RUNTIME_FUNCTION triple_chevron( - dim3 grid_, dim3 block_, Size shared_mem_ = 0, cudaStream_t stream_ = nullptr, bool dependent_launch = false) + dim3 grid_, + dim3 block_, + Size shared_mem_ = 0, + cudaStream_t stream_ = nullptr, + bool dependent_launch = false, + dim3 cluster_dim_ = dim3{0, 0, 0}) : grid(grid_) , block(block_) , shared_mem(shared_mem_) , dependent_launch(dependent_launch) , stream(stream_) + , cluster_dim(cluster_dim_) {} // cudaLaunchKernelEx requires C++11, but unfortunately checks this using the __cplusplus macro, @@ -59,12 +67,33 @@ struct _CCCL_VISIBILITY_HIDDEN triple_chevron template cudaError_t _CCCL_HOST doit_host(K k, Args const&... args) const { + const bool has_cluster = cluster_dim.x != 0; # if _CCCL_HAS_PDL() - if (dependent_launch) + const bool needs_launch_ex = dependent_launch || has_cluster; +# else // _CCCL_HAS_PDL() + const bool needs_launch_ex = has_cluster; +# endif // _CCCL_HAS_PDL() + if (needs_launch_ex) { - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute[0].val.programmaticStreamSerializationAllowed = 1; + // Up to two attributes: programmatic dependent launch and/or the cluster dimension. + cudaLaunchAttribute attribute[2]; + int num_attrs = 0; +# if _CCCL_HAS_PDL() + if (dependent_launch) + { + attribute[num_attrs].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[num_attrs].val.programmaticStreamSerializationAllowed = 1; + ++num_attrs; + } +# endif // _CCCL_HAS_PDL() + if (has_cluster) + { + attribute[num_attrs].id = cudaLaunchAttributeClusterDimension; + attribute[num_attrs].val.clusterDim.x = cluster_dim.x; + attribute[num_attrs].val.clusterDim.y = cluster_dim.y; + attribute[num_attrs].val.clusterDim.z = cluster_dim.z; + ++num_attrs; + } cudaLaunchConfig_t config{}; config.gridDim = grid; @@ -72,15 +101,14 @@ struct _CCCL_VISIBILITY_HIDDEN triple_chevron config.dynamicSmemBytes = shared_mem; config.stream = stream; config.attrs = attribute; - config.numAttrs = 1; -# if _CCCL_COMPILER(MSVC) && _CCCL_CUDACC_BELOW(12, 3) + config.numAttrs = num_attrs; +# if _CCCL_COMPILER(MSVC) && _CCCL_CUDACC_BELOW(12, 3) cudaLaunchKernelEx_MSVC_workaround(&config, k, args...); -# else +# else cudaLaunchKernelEx(&config, k, args...); -# endif +# endif } else -# endif // _CCCL_HAS_PDL() { k<<>>(args...); } From 488322144285d9bf6ab0834695f543f385c85b41 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 13 Jun 2026 03:30:48 +0200 Subject: [PATCH 038/126] Move hardware constants to appropriate header --- .../device/dispatch/dispatch_batched_topk_cluster.cuh | 9 --------- cub/cub/util_device.cuh | 6 ++++++ 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 26608d7e6ae..9a55eb76196 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -57,15 +57,6 @@ CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { -// ----------------------------------------------------------------------------- -// Hardware constants -// ----------------------------------------------------------------------------- -// Largest cluster width guaranteed on every SM 9.0+ device. -inline constexpr int max_portable_cluster_blocks = 8; - -// CUDA's hardware ceiling on cluster width (Hopper supports up to 16). -inline constexpr int max_supported_cluster_blocks = 16; - // ----------------------------------------------------------------------------- // Cluster-size / dynamic-SMEM selection // ----------------------------------------------------------------------------- diff --git a/cub/cub/util_device.cuh b/cub/cub/util_device.cuh index a629f77531b..7b5eca2e7b3 100644 --- a/cub/cub/util_device.cuh +++ b/cub/cub/util_device.cuh @@ -600,6 +600,12 @@ namespace detail // This should stay an implementation detail even when below functions become public. inline constexpr int bulk_copy_min_align = 16; +// Largest thread-block cluster width guaranteed on every SM 9.0+ device. +inline constexpr int max_portable_cluster_blocks = 8; + +// CUDA's hardware ceiling on thread-block cluster width (Hopper supports up to 16). +inline constexpr int max_supported_cluster_blocks = 16; + //! @brief Returns the alignment needed for the shared memory destination buffer of BlockLoadToShared. //! @tparam T //! Value type to be loaded. From 34b613a3c27e995952835e237056fd47cbe5a949 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 18 Jun 2026 04:52:35 +0200 Subject: [PATCH 039/126] Fix races due to missplaced histogram resets --- cub/cub/agent/agent_batched_topk_cluster.cuh | 95 ++++++++++++-------- 1 file changed, 57 insertions(+), 38 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 4c08e97ab62..65b87e62d5c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1236,7 +1236,6 @@ private: // cannot be re-anchored after the fact). ::cuda::std::uint32_t resident_smem32 = 0; - reset_hist(); if constexpr (use_block_load_to_shared) { // Arm the stage barriers once; reused (ping-ponged) by the resident load and the overflow streamer across passes. @@ -1403,11 +1402,6 @@ private: // filtering (all keys are candidates), so it is handled by the fused first-pass load above. if (!is_first_pass) { - // Every block (including the leader) starts each non-first pass with - // a fresh, empty `hist`. Pass 0 was fused with the load pipeline above. - reset_hist(); - __syncthreads(); - identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); @@ -1487,6 +1481,15 @@ private: { leader_identify_kth_bucket(); } + + if (pass + 1 < num_passes) + { + if (cluster_rank == leader_rank) + { + __syncthreads(); // order the leader's `hist` reads (BlockScan in `leader_identify_kth_bucket`) before resets + } + reset_hist(); + } cluster.sync(); // End-of-pass splitter fold. Every thread reads the leader's just-published `kth_bucket` directly through DSMEM @@ -1547,39 +1550,49 @@ private: # else constexpr bool tie_reversed = false; # endif - const out_offset_t num_selected = k - num_kth; - - // Exclusive cross-CTA candidate-count scan. Each non-leader adds its local candidate count (its surviving - // `hist[last_bucket]`) into the `cand_prefix` of every CTA that follows it in scan order; the leader is last in - // scan order (see `leader_rank`) so it adds nothing and receives the sum of all preceding CTAs. - const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; - if (threadIdx.x == 0) + // Early stop means the splitter bucket holds *exactly* the remaining `k` candidates cluster-wide, so all are + // winners with no surplus ties: fold them into the front (`num_back == 0`, `num_selected == k`) and skip the + // index-ordered cross-CTA scan, including its `hist[last_bucket]` read (the end-of-pass reset may already have + // cleared `hist` on an early break). Valid even when the early-stop *break* is compiled out. + const out_offset_t num_back = (leader_state->early_stop != ::cuda::std::uint32_t{0}) ? out_offset_t{0} : num_kth; + const out_offset_t num_selected = k - num_back; // front region; == k on early stop + + offset_t running = 0; + bool tie_active = false; + if (num_back != out_offset_t{0}) { - temp_storage.cand_prefix = 0; - } - cluster.sync(); - if (threadIdx.x == 0 && local_count != offset_t{0}) - { -# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) - for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow + // Exclusive cross-CTA candidate-count scan. Each non-leader adds its local candidate count (its surviving + // `hist[last_bucket]`) into the `cand_prefix` of every CTA that follows it in scan order; the leader is last in + // scan order (see `leader_rank`) so it adds nothing and receives the sum of all preceding CTAs. + const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; + if (threadIdx.x == 0) { - add_remote_prefix(r, local_count); + temp_storage.cand_prefix = 0; } -# else - for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + cluster.sync(); + if (threadIdx.x == 0 && local_count != offset_t{0}) { - add_remote_prefix(r, local_count); - } +# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) + for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow + { + add_remote_prefix(r, local_count); + } +# else + for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + { + add_remote_prefix(r, local_count); + } # endif - } - cluster.sync(); + } + cluster.sync(); - offset_t running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) - bool tie_active = running < static_cast(num_kth); + running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) + tie_active = running < static_cast(num_back); + } // Process a flat span of `count` keys already in scan order (`get_key(pos)`), tiled by - // `threads_per_block * tie_break_items_per_thread`. Selected keys go to the front via `out_cnt`; candidates get a - // BlockScan-exclusive index rank (seeded by `running`) and, if `rank < num_kth`, are written in reverse at + // `threads_per_block * tie_break_items_per_thread`. Front keys go via `out_cnt`; surplus ties get a + // BlockScan-exclusive index rank (seeded by `running`) and, if `rank < num_back`, are written in reverse at // `block_keys_out[k - 1 - rank]`. The running aggregate carries across tiles and across regions. // `get_idx(pos)` returns the segment-local index of the key `get_key(pos)` returns, used only to load the value // payload for written keys (compiled out in keys-only builds). @@ -1606,12 +1619,16 @@ private: } } - // Strictly-selected keys: atomic into the front. Exactly `num_selected` keys are selected cluster-wide, so - // `out_cnt` never overruns the `[0, num_selected)` region. + // Strictly-selected keys go to the front via `out_cnt`; on early stop (`num_back == 0`) the splitter bucket's + // candidates are winners too and fold in here. Exactly `num_selected` are placed in `[0, num_selected)`. _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && cls[i] == detail::topk::candidate_class::selected) + const bool to_front = + valid[i] + && (cls[i] == detail::topk::candidate_class::selected + || (num_back == out_offset_t{0} && cls[i] == detail::topk::candidate_class::candidate)); + if (to_front) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; const out_offset_t out = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); @@ -1620,6 +1637,8 @@ private: } } + // Surplus ties (`num_back > 0`): each candidate gets a BlockScan rank (seeded by `running`), written reversed + // at `block_keys_out[k - 1 - rank]` when `rank < num_back`. Skipped on early stop (`tie_active == false`). if (tie_active) { offset_t excl[items]; @@ -1631,7 +1650,7 @@ private: if (valid[i] && flags[i] != offset_t{0}) { const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_kth)) + if (global_rank < static_cast(num_back)) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); @@ -1641,7 +1660,7 @@ private: } } running += tile_total; - if (running >= static_cast(num_kth)) + if (running >= static_cast(num_back)) { tie_active = false; } @@ -2025,8 +2044,7 @@ private: // leader's copy is semantically read (non-leaders reach the cluster // state through `leader_state`), but mirroring the writes everywhere // keeps the scan-then-reduce path's unconditional `state.k` load - // safe under compute-sanitizer. Every block will reset its own - // `hist` at the top of the per-pass loop. + // safe under compute-sanitizer. if (threadIdx.x == 0) { temp_storage.state.len = static_cast(segment_size); @@ -2036,6 +2054,7 @@ private: temp_storage.state.out_back_cnt = 0; temp_storage.state.early_stop = 0; } + reset_hist(); cluster.sync(); [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( From 9fc8e46fbcaac2ef10a5a7398671bf971c530f3b Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 22 Jun 2026 15:55:43 +0200 Subject: [PATCH 040/126] Add special handling for head/tail storage By keeping the unaligned head (and if needed tail) in small "edge" areas we can remove the requirement to have resident chunks for them and can use the full chunk budget on streaming. This lets us enable large-segment tests in CDP mode but more importantly means that we can portably run on potential future GPUs with only 48 KiB smem. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 479 +++++++++++------- .../dispatch_batched_topk_cluster.cuh | 14 +- .../catch2_test_device_segmented_topk_keys.cu | 12 +- ...catch2_test_device_segmented_topk_pairs.cu | 85 +++- ...tch2_test_segmented_topk_cluster_layout.cu | 23 +- 5 files changed, 421 insertions(+), 192 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 65b87e62d5c..66deaa620da 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -104,9 +104,9 @@ struct alignas(16) cluster_topk_state ::cuda::std::uint32_t early_stop; }; -// Dynamic-SMEM layout shared by dispatch and the agent. `block_tile_capacity` is the physical per-CTA -// resident capacity passed to the kernel; `cluster_tile_capacity` reserves one chunk of logical coverage -// for a possible unaligned head chunk. +// Dynamic-SMEM layout shared by dispatch and the agent. `block_tile_capacity` is the physical per-CTA resident +// capacity; `cluster_tile_capacity` is the cluster's total resident coverage. The unaligned head is staged as an edge +// in static SMEM (`edge_keys`), not a chunk slot, so the full physical capacity is usable (no head reservation). template struct smem_block_tile_layout { @@ -134,12 +134,7 @@ struct smem_block_tile_layout [[nodiscard]] _CCCL_HOST_DEVICE static constexpr SizeT cluster_tile_capacity(int cluster_blocks, ::cuda::std::uint32_t physical_block_tile_capacity) noexcept { - const auto physical_cluster_tile_items = - static_cast(cluster_blocks) * static_cast(physical_block_tile_capacity); - const auto head_chunk_reserve = static_cast(chunk_items); - return (physical_cluster_tile_items > head_chunk_reserve) - ? physical_cluster_tile_items - head_chunk_reserve - : SizeT{0}; + return static_cast(cluster_blocks) * static_cast(physical_block_tile_capacity); } }; @@ -251,7 +246,14 @@ struct agent_batched_topk_cluster // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused // (ping-ponged) across radix passes; initialized once by `init_load_barriers`. ::cuda::std::uint64_t load_mbar[PipelineStages]; + // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap)`, on rank 0) and + // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner when `full_slots == 1`), each + // strictly `< load_align_items` keys. Loaded once in the first pass and folded into every pass + the final filter. + // Block-local (never reached through DSMEM). + key_t edge_keys[2 * load_align_items]; }; + // Split point of `edge_keys`: head edge in `[0, head_edge_cap)`, tail edge in `[head_edge_cap, 2 * head_edge_cap)`. + static constexpr int head_edge_cap = load_align_items; [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const { @@ -319,25 +321,14 @@ struct agent_batched_topk_cluster return (::cuda::std::min) (items, segment_size); } - // Item count of the near-full head chunk (chunk 0) for an unaligned base. The unaligned prefix (`head_items`) plus - // an aligned bulk filling the rest of the slot up to a `load_align` boundary, so chunk 0 ends aligned and every - // later chunk begins aligned. Clamped to the segment for tiny segments. - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t - head_chunk_items(offset_t segment_size, offset_t head_items) const - { - const offset_t bulk0 = offset_t{chunk_items} - offset_t{load_align_items}; - return (::cuda::std::min) (segment_size, head_items + bulk0); - } - + // Number of aligned chunks covering the segment. The unaligned head prefix (`head_items`) is handled as a separate + // edge, not a chunk, so the aligned region `[head_items, segment_size)` chunks uniformly at `chunk_items` (a segment + // lying entirely before the first boundary has no chunks). `head_items == 0` (aligned base or generic fallback) is + // then plain uniform chunking from offset 0. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t num_chunks(offset_t segment_size, offset_t head_items) const { - if (head_items == 0) - { - return static_cast(::cuda::ceil_div(segment_size, offset_t{chunk_items})); - } - const offset_t head0 = head_chunk_items(segment_size, head_items); - const offset_t remaining = segment_size - head0; - return offset_t{1} + static_cast(::cuda::ceil_div(remaining, offset_t{chunk_items})); + const offset_t aligned_items = (segment_size > head_items) ? (segment_size - head_items) : offset_t{0}; + return static_cast(::cuda::ceil_div(aligned_items, offset_t{chunk_items})); } struct chunk_desc @@ -346,31 +337,22 @@ struct agent_batched_topk_cluster int count; }; + // Chunk `chunk_idx` of the aligned region: it begins on a `load_align` boundary (`head_items + chunk_idx * + // chunk_items`, with `head_items` itself aligning the base), so every chunk has a zero prefix and only the last chunk + // can carry an unaligned suffix (the segment's trailing `< load_align_items` items). [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_desc get_chunk(offset_t chunk_idx, offset_t segment_size, offset_t head_items) const { - offset_t offset; - if (head_items != 0) - { - const offset_t head0 = head_chunk_items(segment_size, head_items); - if (chunk_idx == 0) - { - return {offset_t{0}, static_cast(head0)}; - } - offset = head0 + (chunk_idx - 1) * offset_t{chunk_items}; - } - else - { - offset = chunk_idx * offset_t{chunk_items}; - } + const offset_t offset = head_items + chunk_idx * offset_t{chunk_items}; const offset_t remaining = segment_size - offset; return {offset, static_cast((::cuda::std::min) (remaining, offset_t{chunk_items}))}; } // Splits a chunk into its unaligned front edge, aligned interior (bulk), and unaligned back edge, relative to the // gmem base. The interior begins and ends on a `load_align` boundary so it can be loaded with the aligned, - // guard-free BlockLoadToShared path; the edges (each `< load_align_items` items) are loaded with `copy_edge`. Only - // chunk 0 (head) can have a nonzero prefix and only the last chunk (tail) a nonzero suffix. + // guard-free BlockLoadToShared path; the edges (each `< load_align_items` items) are staged via + // `stage_and_fold_edge`. The head prefix is peeled as a separate edge before chunking, so a `get_chunk` chunk begins + // on a boundary (zero prefix) and only the last chunk (tail) carries a nonzero suffix. struct chunk_split { offset_t prefix; @@ -389,9 +371,9 @@ struct agent_batched_topk_cluster if (aligned_begin > aligned_end) { // The chunk lies strictly between two load_align boundaries (no aligned point inside): the whole chunk is an - // unaligned edge. This only happens for a tiny chunk with an unaligned begin (the head or a single-chunk - // segment), so attributing it entirely to the front edge is correct. A tail always begins on a boundary, so it - // takes the `aligned_begin <= aligned_end` path below and its unaligned remainder becomes the suffix. + // unaligned edge, attributed entirely to the front edge. `get_chunk` chunks begin on a boundary (the head is + // peeled separately), so this only guards a degenerate sub-`load_align` chunk. A tail always begins on a + // boundary, so it takes the `aligned_begin <= aligned_end` path below and its unaligned remainder is the suffix. return {static_cast(chunk.count), offset_t{0}, offset_t{0}}; } const offset_t prefix = static_cast((aligned_begin - begin) / sizeof(key_t)); @@ -429,8 +411,8 @@ struct agent_batched_topk_cluster } }; - // Decides which global chunks a cluster rank owns. Both layouts keep the unaligned head (chunk 0) on rank 0 and the - // unaligned tail (chunk `chunks-1`) on a single rank, and leave the per-chunk alignment, the resident/streaming + // Decides which global chunks a cluster rank owns. Both layouts keep chunk 0 on rank 0 (which also stages the head + // edge) and the tail (chunk `chunks-1`) on a single rank, and leave the per-chunk alignment, the resident/streaming // split, and the streamer ping-pong untouched, because all of those depend only on the global chunk index, not on // which rank owns it. // @@ -457,19 +439,26 @@ struct agent_batched_topk_cluster #endif } - // Hand-rolled per-thread copy of a small (`< load_align_items` items) unaligned edge from gmem to smem. Used for - // the head prefix and tail suffix, which cannot go through the aligned (16-byte-aligned dst, guard-free) - // BlockLoadToShared path. Each thread copies the same indices it later reads in the first-pass histogram, so no - // extra synchronization is needed for the edge's own pass. + // Stage a small (`< load_align_items` items) unaligned run -- a boundary edge (head prefix / tail suffix) or a + // resident chunk's suffix -- from gmem `src` into SMEM `dst` and fold it into the first pass in one strided sweep: + // each thread copies *and* folds the same indices it owns (`local % threads_per_block == threadIdx.x`), folding from + // the just-loaded register rather than from SMEM. These runs cannot go through the aligned (16-byte-aligned dst, + // guard-free) BlockLoadToShared path. Fusing the copy and the fold means no thread ever reads a key another thread + // wrote, so the first-pass fold needs no barrier after the staging *by construction* -- not by a coincidental match + // between two separate traversals. Later passes and the final filter re-read `dst` (from `edge_keys` / the resident + // span) only after a pass-boundary barrier, so their reads are independently safe. // // TODO(cccl): an asymmetric-alignment BlockLoadToShared API (independent begin/end alignment, e.g. an aligned begin // with an arbitrary end) would let a boundary chunk be loaded with a single aligned-bulk + in-place edge call and // remove these hand-rolled copies entirely. - _CCCL_DEVICE _CCCL_FORCEINLINE void copy_edge(key_t* dst, const key_t* src, int count) const + template + _CCCL_DEVICE _CCCL_FORCEINLINE void stage_and_fold_edge(key_t* dst, const key_t* src, int count, Apply&& apply) const { for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) { - dst[local] = src[local]; + const key_t key = src[local]; + dst[local] = key; + apply(key); } } @@ -816,14 +805,14 @@ private: // resident SMEM region) from gmem through a small, fixed, round-robin set of // `p_eff` (<= `PipelineStages`) streaming slots. The same object is reused for // every radix pass and the final filter. It ping-pongs the iteration order - // across calls so the `p_eff` boundary chunks that one pass leaves resident in - // the streaming slots are reused by the next pass with no reload; the remaining - // `overflow_chunks - p_eff` chunks are reloaded from gmem on each pass. The - // caller right-sizes the reservation to `p_eff = min(PipelineStages, excess)` + // across calls so the `p_eff` turn-around chunks that one pass leaves resident + // in the streaming slots are reused by the next pass with no reload; the + // remaining `overflow_chunks - p_eff` chunks are reloaded from gmem on each pass. + // The caller right-sizes the reservation to `p_eff = min(PipelineStages, excess)` // (where `excess = my_chunks - full_slots`), so a streaming rank always has // `overflow_chunks = excess + p_eff` and reloads exactly `excess` chunks per - // pass - the reserved slots only ever buy reuse of the `p_eff` boundary chunks, - // never a reload-free pass. The resident region is unaffected: it lives in the + // pass - the reserved slots only ever buy reuse of those `p_eff` turn-around + // chunks, never a reload-free pass. The resident region is unaffected: it lives in the // slots `[0, resident_slots)`, the streaming region in `[stream_slot_base, // stream_slot_base + p_eff)`. struct overflow_streamer @@ -890,24 +879,30 @@ private: { const offset_t chunk_idx = chunk_index_of(overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - // The boundary chunks (unaligned head and tail) are kept resident, so every streamed chunk is fully aligned - // and uses the guard-free aligned (TMA bulk) path. - _CCCL_ASSERT(agent.is_aligned_chunk(block_keys_base, chunk), "overflow streamer received an unaligned chunk"); + // Every chunk begins on a `load_align` boundary (zero prefix), so the guard-free aligned (TMA bulk) path applies. + // Only the global-last chunk can carry an unaligned suffix, and it only reaches the streamer when that tail was + // peeled (its suffix lives in `edge_keys`); streaming just the aligned bulk excludes it. For every interior chunk + // `bulk == count`. + const auto split = agent.split_chunk(block_keys_base, chunk); + _CCCL_ASSERT(split.prefix == offset_t{0}, "overflow streamer received a chunk with an unaligned start"); char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; const ::cuda::std::span src{ - block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(chunk.count)}; + block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(split.bulk)}; agent.issue_bulk_copy(stage, dst, src); inflight_mask |= (::cuda::std::uint32_t{1} << stage); } // Rebuild the shared span for the chunk currently resident in `stage`'s slot without storing per-stage state: the // slot address is a pure function of `stage` and the length is recomputed from chunk index `o`, so there is no - // spillable `pending[]` array (see `bulk_span`). + // spillable `pending[]` array (see `bulk_span`). Returns the aligned bulk only (a peeled tail's suffix is + // excluded). [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span stage_span(int stage, offset_t o) const { char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; const auto chunk = agent.get_chunk(chunk_index_of(o), segment_size, head_items); - return agent.bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), chunk.count}); + const auto split = agent.split_chunk(block_keys_base, chunk); + return agent.bulk_span( + {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), static_cast(split.bulk)}); } // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` currently resident in @@ -1135,83 +1130,80 @@ private: const chunk_partition part = make_chunk_partition(chunks, cluster_rank, cluster_size); const offset_t my_chunks = part.count; - // Resident vs. streaming split, decided independently per CTA. A CTA whose owned-chunk count fits its resident - // slots (`my_chunks <= full_slots`) keeps every chunk resident and streams nothing; only a CTA that actually - // overflows reserves a round-robin streaming region at the tail of its block_tile and re-streams its overflow - // chunks from gmem on every pass via `streamer`. - // - // This is purely a local decision: each CTA only ever loads/scans its own chunks (resident SMEM or its own gmem - // overflow), and the cross-CTA traffic (histogram fold, leader `state`, the deterministic `cand_prefix` scan) plus - // every `cluster.sync()` are reached uniformly regardless of how many chunks any CTA streams - so CTAs need not - // agree on the split. `my_chunks` already folds in this segment's actual base alignment (an unaligned head costs - // exactly one extra chunk via `num_chunks`/`head_chunk_items`), so no head reserve is needed here; that differs - // from the host-side launch selection, which must provision a one-chunk margin (`cluster_tile_capacity`) because it - // picks a single launch-wide `(cluster_size, smem)` from only the segment-size *upper bound*. + // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every + // `cluster.sync()` are reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) + // keeps them all resident and streams nothing; an overflowing CTA reserves a round-robin streaming region at the + // tail of its block_tile and re-streams its overflow chunks from gmem each pass via `streamer`. // - // Versus a cluster-uniform split (`chunks > full_slots * cluster_size`, which forces every CTA to stream): the - // busiest rank streams the same amount either way, but every other rank now stays fully resident, cutting cluster - // gmem traffic and SMEM pressure. Both schemes stream under the exact same global condition (some rank overflows - // iff `ceil_div(chunks, cluster_size) > full_slots`), so the `full_slots > PipelineStages` reservation guarantee - // the dispatch already provisions for is unchanged. + // Boundary edges (the unaligned head prefix on rank 0 and the unaligned tail suffix on the tail owner) cannot use + // the aligned TMA streamer. The head is always peeled into the persistent `edge_keys` buffer; the tail suffix is + // kept resident in its (partial, single-slot) tail chunk, peeled into `edge_keys` only when a single slot fits + // (`full_slots == 1`). Peeling the head removes the dual-boundary pressure, so streaming needs only `full_slots >= + // 1`. // - // Right-size the streaming region. A CTA that overflows its resident slots by only `excess = my_chunks - - // full_slots` chunks needs at most `excess` streaming slots to cycle that overflow through gmem; reserving the full - // `PipelineStages` would needlessly route up-to-`PipelineStages` extra chunks through the streaming machinery (and - // its per-prefetch `__syncthreads()`) that could instead stay resident and be read once. So reserve - // `stream_slots = min(PipelineStages, excess)` slots: deep overflows (`excess >= PipelineStages`) behave exactly as - // before, while barely-overflowing segments keep the rest resident. The per-pass gmem reload is `excess` either way - // (the streamer reuses its `p_eff = stream_slots` boundary chunks across passes); this only shrinks the streaming - // region - and grows the resident region - which can only relax the `>= 2 resident slots` head+tail guarantee - // below. - // - // The streaming region is the async SMEM pipeline, which only exists on the TMA path. The generic fallback has no - // pipeline: it still keeps its resident chunks in SMEM (read once, reused across passes), but re-reads its overflow - // chunks straight from gmem every pass without ever staging them in a slot. Reserving streaming slots there would - // just leave SMEM idle, so the fallback reserves none and devotes the whole block_tile to resident chunks - the - // more chunks it keeps resident, the fewer it re-reads from gmem each pass. + // `stream_slots` is right-sized: clamped to leave a resident slot for a non-peeled tail without exceeding the + // available slots; deep overflows still get the full `PipelineStages` depth. The generic fallback has no async + // pipeline (it re-reads overflow from gmem each pass and peels nothing), so it reserves no streaming slots. const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); [[maybe_unused]] const bool needs_streaming = my_chunks > full_slots; - offset_t stream_slots = offset_t{0}; + + // Does this rank own the global tail, and does that tail carry an unaligned suffix? (block-load path only; the + // generic fallback reads any trailing items straight from gmem and never peels.) + [[maybe_unused]] offset_t tail_local = offset_t{0}; + [[maybe_unused]] offset_t tail_suffix_items = offset_t{0}; + bool owns_suffix_tail = false; + if constexpr (use_block_load_to_shared) + { + // This rank owns the global tail iff its last owned chunk is chunk `chunks-1` (its local index `my_chunks-1`, + // true for both the strided and blocked partitions). + if (my_chunks > 0 && part.global_index(my_chunks - offset_t{1}) == chunks - offset_t{1}) + { + tail_local = my_chunks - offset_t{1}; + const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); + tail_suffix_items = split_chunk(block_keys_base, tail_chunk).suffix; + owns_suffix_tail = tail_suffix_items != offset_t{0}; + } + } + + // Peel the suffix tail only when streaming and no resident slot can be spared for it (`full_slots == 1`); otherwise + // keep it resident. `reserved_resident` holds back one slot for a non-peeled tail; `stream_slots` is then clamped + // into `[1, full_slots - reserved_resident]`. + offset_t stream_slots = offset_t{0}; + bool peel_tail = false; + bool force_tail_resident = false; if constexpr (use_block_load_to_shared) { - _CCCL_ASSERT(!needs_streaming || full_slots > static_cast(PipelineStages), - "block_tile too small to reserve a streaming region"); - stream_slots = needs_streaming - ? (::cuda::std::min) (static_cast(PipelineStages), my_chunks - full_slots) - : offset_t{0}; + if (needs_streaming) + { + peel_tail = owns_suffix_tail && full_slots < offset_t{2}; + const offset_t reserved_resident = (owns_suffix_tail && !peel_tail) ? offset_t{1} : offset_t{0}; + const offset_t excess = my_chunks - full_slots; + const offset_t want_stream = (::cuda::std::min) (static_cast(PipelineStages), excess); + const offset_t max_stream = full_slots - reserved_resident; // >= 1 (full_slots >= 1; reserve <= f-1) + stream_slots = (::cuda::std::max) (offset_t{1}, (::cuda::std::min) (want_stream, max_stream)); + } } const offset_t resident_slots_cap = full_slots - stream_slots; const offset_t my_resident_chunks = (::cuda::std::min) (my_chunks, resident_slots_cap); // Resident chunks stay within the first `resident_slots_cap` slots; the streaming region occupies the slots // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. - _CCCL_ASSERT(my_resident_chunks * offset_t{chunk_items} <= resident_slots_cap * offset_t{chunk_items}, - "Dynamic shared memory block_tile is too small"); - - // Boundary chunks (the unaligned head = chunk 0 and the unaligned tail = chunk chunks-1) carry hand-rolled edges - // and are kept resident so the overflow streamer only ever sees fully-aligned middle chunks. The head is already - // resident (local index 0 of its owning rank). The tail is forced into its owner's resident set - occupying the - // last resident slot - when this rank owns it, it has an unaligned end, and it would otherwise stream. The - // overflow then begins at `overflow_base`, shifted back by one in the forced case so the tail is skipped while the - // middle chunk it displaces is streamed instead. + _CCCL_ASSERT(my_resident_chunks <= resident_slots_cap, "Dynamic shared memory block_tile is too small"); + + // A non-peeled suffix tail is forced into the last resident slot iff it would otherwise stream (it falls in the + // overflow region). The overflow then begins one chunk earlier so the displaced middle chunk streams in its place. const offset_t overflow_count = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; - offset_t tail_local = offset_t{0}; - bool force_tail_resident = false; if constexpr (use_block_load_to_shared) { - // This rank owns the global tail iff its last owned chunk is chunk `chunks-1`; that chunk is then this rank's - // local index `my_chunks-1` (true for both the strided and blocked partitions). - if (overflow_count > 0 && my_chunks > 0 && part.global_index(my_chunks - offset_t{1}) == chunks - offset_t{1}) - { - tail_local = my_chunks - offset_t{1}; - const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); - const bool tail_has_suffix = split_chunk(block_keys_base, tail_chunk).suffix != offset_t{0}; - force_tail_resident = tail_has_suffix && (tail_local >= my_resident_chunks); - } + force_tail_resident = owns_suffix_tail && !peel_tail && overflow_count > 0 && (tail_local >= my_resident_chunks); } const offset_t overflow_base = force_tail_resident ? (my_resident_chunks - offset_t{1}) : my_resident_chunks; - // A rank that owns both boundary chunks (the head on rank 0 and a forced tail) needs two distinct resident slots. - _CCCL_ASSERT(!(force_tail_resident && head_items != 0 && cluster_rank == 0) || my_resident_chunks >= offset_t{2}, - "streaming needs >= 2 resident slots to keep both head and tail resident"); + _CCCL_ASSERT(!force_tail_resident || my_resident_chunks >= offset_t{1}, + "a forced-resident tail needs at least one resident slot"); + + // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and + // for an aligned base); the peeled tail suffix lives on the tail owner only when `peel_tail`. + [[maybe_unused]] const int head_edge_len = (cluster_rank == 0u) ? static_cast(head_items) : 0; + [[maybe_unused]] const int tail_edge_len = peel_tail ? static_cast(tail_suffix_items) : 0; // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. overflow_streamer streamer( @@ -1236,6 +1228,24 @@ private: // cannot be re-anchored after the fact). ::cuda::std::uint32_t resident_smem32 = 0; + // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), + // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final + // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). + const auto fold_edges = [&](auto&& apply) { + if constexpr (use_block_load_to_shared) + { + if (head_edge_len > 0) + { + for_each_chunk_key({temp_storage.edge_keys, static_cast<::cuda::std::size_t>(head_edge_len)}, apply); + } + if (tail_edge_len > 0) + { + for_each_chunk_key({temp_storage.edge_keys + head_edge_cap, static_cast<::cuda::std::size_t>(tail_edge_len)}, + apply); + } + } + }; + if constexpr (use_block_load_to_shared) { // Arm the stage barriers once; reused (ping-ponged) by the resident load and the overflow streamer across passes. @@ -1257,18 +1267,19 @@ private: if (my_resident_chunks > 0) { // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). - // Chunks are written densely in slot order and read back in the same order, so the read cursor (`read_off`) - // mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed - // `pending_spans` array that would anchor surrounding state to local memory. Each chunk loads its aligned - // bulk only (boundary edges are filled afterwards); the read span is rebuilt from a rooted 32-bit shared - // address (see `bulk_span`) so a spilled cursor cannot demote the first-pass reads from `LDS` to `LD`. + // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor + // (`read_off`) mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed + // `pending_spans` array that would anchor surrounding state to local memory. Every chunk begins on a + // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except a resident suffix + // tail, whose bulk is loaded here and whose suffix is appended afterwards. The read span is rebuilt from a + // rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot demote the reads to `LD`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); // Resident slot -> rank-local chunk index. Identity, except the last slot holds the forced-resident tail. const auto resident_local = [&](offset_t slot) -> offset_t { return (force_tail_resident && slot == my_resident_chunks - offset_t{1}) ? tail_local : slot; }; - // Aligned bulk of the resident chunk in `slot`; empty when the chunk has no aligned interior. + // Aligned bulk of the resident chunk in `slot` (its `count` minus any tail suffix); empty when it has none. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { const offset_t chunk_idx = part.global_index(resident_local(slot)); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -1280,21 +1291,9 @@ private: return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; }; - // First resident slot's unaligned front edge (the head prefix on rank 0). Reserve an aligned-up gap in - // front of its bulk so the bulk stays `load_align`-aligned and the edge sits contiguously right before it. - const auto first_chunk = - get_chunk(part.global_index(resident_local(offset_t{0})), segment_size_u32, head_items); - const auto first_split = split_chunk(block_keys_base, first_chunk); - const int front_edge = static_cast(first_split.prefix); - const int front_bytes = front_edge * int{sizeof(key_t)}; - // Round the first bulk up to `load_align` (not just BlockLoadToShared's 16B): with a `load_align`-aligned - // base and `load_align`-multiple bulk sizes, every densely packed resident bulk lands on a `load_align` - // boundary. - const int head_bulk_off = ::cuda::round_up(front_bytes, load_align_bytes); - // Write cursor as a byte offset from `key_slots`; the destination is always `key_slots + offset` (rooted at - // the extern-shared base) so it keeps shared address-space provenance. - const int resident_begin_off = head_bulk_off - front_bytes; - int next_off = head_bulk_off; + // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives + // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. + int next_off = 0; // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail // bulk in the forced case) ends the packed region and its suffix can be appended right after it. @@ -1307,7 +1306,7 @@ private: // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. - int read_off = head_bulk_off; + int read_off = 0; for (offset_t p = 0; p < my_resident_chunks; ++p) { const int stage = static_cast(p % static_cast(prologue)); @@ -1330,33 +1329,26 @@ private: } } - // Head prefix: copy into the reserved front gap, contiguous right before the first bulk. - if (front_edge > 0) - { - key_t* const edge_dst = reinterpret_cast(key_slots + resident_begin_off); - copy_edge(edge_dst, block_keys_base + first_chunk.offset, front_edge); - for_each_chunk_key({edge_dst, static_cast<::cuda::std::size_t>(front_edge)}, add_first_pass); - } - - // Tail suffix: append right after the last (tail) bulk. Nothing follows, so its non-16 length is harmless. + // Tail suffix of a resident (non-peeled) suffix tail: append right after the last (tail) bulk. Nothing + // follows, so its non-16 length is harmless. A peeled tail's suffix lives in `edge_keys` instead. const offset_t last_slot = my_resident_chunks - offset_t{1}; const auto last_chunk = get_chunk(part.global_index(resident_local(last_slot)), segment_size_u32, head_items); const auto last_split = split_chunk(block_keys_base, last_chunk); if (last_split.suffix > 0) { key_t* const edge_dst = reinterpret_cast(key_slots + next_off); - copy_edge(edge_dst, - block_keys_base + last_chunk.offset + last_split.prefix + last_split.bulk, - static_cast(last_split.suffix)); - for_each_chunk_key({edge_dst, static_cast<::cuda::std::size_t>(last_split.suffix)}, add_first_pass); + stage_and_fold_edge(edge_dst, + block_keys_base + last_chunk.offset + last_split.prefix + last_split.bulk, + static_cast(last_split.suffix), + add_first_pass); next_off += static_cast(last_split.suffix) * int{sizeof(key_t)}; } - // The resident region is one contiguous span [head edge | bulks | tail edge] for the later passes. - resident_keys = {reinterpret_cast(key_slots + resident_begin_off), - static_cast<::cuda::std::size_t>((next_off - resident_begin_off) / int{sizeof(key_t)})}; - resident_smem32 = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + resident_begin_off)); + // The resident region is one contiguous span [bulks | tail suffix] for the later passes; the head prefix is + // folded separately from `edge_keys`. + resident_keys = {reinterpret_cast(key_slots), + static_cast<::cuda::std::size_t>(next_off / int{sizeof(key_t)})}; + resident_smem32 = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots)); } } else @@ -1382,6 +1374,27 @@ private: } } + // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see + // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix + // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, `full_slots == 1`) trails the last chunk. The + // `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via + // `fold_edges` after a barrier). + if constexpr (use_block_load_to_shared) + { + if (head_edge_len > 0) + { + stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len, add_first_pass); + } + if (tail_edge_len > 0) + { + const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); + stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap, + block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len), + tail_edge_len, + add_first_pass); + } + } + // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer // sync. @@ -1441,9 +1454,14 @@ private: }; // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the - // first wave of reload bulk copies. Ping-pongs direction and reuses the boundary chunks left resident by the + // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. streamer.process_pass(add_hist, fold_resident_hist); + + // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the + // resident and overflow keys. Keeps every owner's `hist[last_bucket]` (the deterministic cross-CTA prefix + // input) inclusive of its edge candidates. + fold_edges(add_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -1770,8 +1788,19 @@ private: const offset_t o = reversed ? (overflow_count - 1 - oo) : oo; const offset_t chunk_idx = part.global_index(overflow_base + o); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const int cc = chunk.count; - const offset_t base_off = chunk.offset; + // Visit only the aligned bulk: a peeled suffix tail's suffix is handled by `process_tail_edge`, so it is + // excluded here; every other overflow chunk has `bulk == count`. (The generic fallback never peels and reads + // the trailing items straight from gmem, so it uses the full count.) + int cc; + if constexpr (use_block_load_to_shared) + { + cc = static_cast(split_chunk(block_keys_base, chunk).bulk); + } + else + { + cc = chunk.count; + } + const offset_t base_off = chunk.offset; if (reversed) { process_flat( @@ -1835,9 +1864,80 @@ private: } }; + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. Processed + // before every chunk (ascending) / after every chunk (descending) so the global index order holds across regions. + // `process_flat` with `count == 0` is a barrier-free no-op, so non-head ranks (`head_edge_len == 0`) skip it. + auto process_head_edge = [&](bool reversed) { + if constexpr (use_block_load_to_shared) + { + key_t* const he = temp_storage.edge_keys; + const int hc = head_edge_len; + if (reversed) + { + process_flat( + [&](int pos) { + return he[hc - 1 - pos]; + }, + [&](int pos) { + return static_cast(hc - 1 - pos); + }, + hc); + } + else + { + process_flat( + [&](int pos) { + return he[pos]; + }, + [&](int pos) { + return static_cast(pos); + }, + hc); + } + } + }; + + // Peeled tail suffix edge (tail owner, `full_slots == 1`): the segment's highest indices, staged in `edge_keys`. + // Processed after every chunk (ascending) / before every chunk (descending). `tail_edge_len == 0` (the tail was + // kept resident, or this rank does not own it) makes this a barrier-free no-op. + auto process_tail_edge = [&](bool reversed) { + if constexpr (use_block_load_to_shared) + { + key_t* const te = temp_storage.edge_keys + head_edge_cap; + const int tc = tail_edge_len; + const offset_t tail_base = segment_size_u32 - static_cast(tc); + if (reversed) + { + process_flat( + [&](int pos) { + return te[tc - 1 - pos]; + }, + [&](int pos) { + return tail_base + static_cast(tc - 1 - pos); + }, + tc); + } + else + { + process_flat( + [&](int pos) { + return te[pos]; + }, + [&](int pos) { + return tail_base + static_cast(pos); + }, + tc); + } + } + }; + if constexpr (tie_reversed) { - process_tail(true); + process_tail_edge(true); + if (!should_stop()) + { + process_tail(true); + } if (!should_stop()) { process_overflow(true); @@ -1846,10 +1946,18 @@ private: { process_resident(true); } + if (!should_stop()) + { + process_head_edge(true); + } } else { - process_resident(false); + process_head_edge(false); + if (!should_stop()) + { + process_resident(false); + } if (!should_stop()) { process_overflow(false); @@ -1858,6 +1966,10 @@ private: { process_tail(false); } + if (!should_stop()) + { + process_tail_edge(false); + } } # else // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM if constexpr (is_keys_only) @@ -1898,6 +2010,8 @@ private: for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } + // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). + fold_edges(write_selected); }; streamer.process_pass(write_selected, fold_resident); } @@ -1985,6 +2099,29 @@ private: chunk.count); } } + // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the + // peeled tail suffix at `segment_size - tail_edge_len`. Order-independent atomic writes, value fetched per key. + if constexpr (use_block_load_to_shared) + { + if (head_edge_len > 0) + { + write_run( + [&](int local) { + return temp_storage.edge_keys[local]; + }, + offset_t{0}, + head_edge_len); + } + if (tail_edge_len > 0) + { + write_run( + [&](int local) { + return temp_storage.edge_keys[head_edge_cap + local]; + }, + segment_size_u32 - static_cast(tail_edge_len), + tail_edge_len); + } + } }; // Overflow chunks: reuse the keys from the streaming SMEM pipeline (block-load path; only the generic fallback diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 8e4d5305dc4..efbd9b2e44c 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -577,9 +577,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( selected_config.cluster_blocks, selected_config.block_tile_capacity); if (selected_cluster_tile_capacity >= max_seg_size) { - const auto required_physical_cluster_tile_items = - static_cast<::cuda::std::uint64_t>(max_seg_size) + static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); - selected_config.cluster_blocks = static_cast( + // `cluster_tile_capacity == physical` (the head is an edge, not a reserved chunk), so the physical capacity + // need only reach `max_seg_size`; the chunk-granular round-up below grows it as needed. + const auto required_physical_cluster_tile_items = static_cast<::cuda::std::uint64_t>(max_seg_size); + selected_config.cluster_blocks = static_cast( ::cuda::ceil_div(required_physical_cluster_tile_items, static_cast<::cuda::std::uint64_t>(selected_config.block_tile_capacity))); @@ -641,7 +642,14 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr int portable_total_smem_bytes = 48 * 1024; constexpr int dynamic_smem_bytes = (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; + + // The compile-time `ChunkBytes` is reused verbatim for the device launch: the agent peels the unaligned boundary + // edges into a tiny per-block buffer, so streaming needs only a single resident-or-streaming slot, and segments + // exceeding the small portable-SMEM block tile are handled by re-streaming overflow from gmem. The only hard + // requirement is that at least one load-aligned chunk fits the worst-case portable SMEM block tile. constexpr auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + static_assert(block_tile_capacity >= static_cast<::cuda::std::uint32_t>(layout_t::chunk_items), + "Portable SMEM is too small to fit even one load-aligned chunk for the device-launch (CDP) path"); const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 77f45ed7f8d..3565307af14 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -225,7 +225,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", REQUIRE(expected_keys == keys_out_buffer); } -#if TEST_LAUNCH != 1 && TEST_TYPES == 1 +#if TEST_TYPES == 1 TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", "[keys][segmented][topk][device][cluster]") { @@ -360,7 +360,7 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non- REQUIRE(expected_keys == keys_out_buffer); } -#endif // TEST_LAUNCH != 1 && TEST_TYPES == 1 +#endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segments", "[keys][segmented][topk][device]", @@ -461,7 +461,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } -#if TEST_LAUNCH != 1 && TEST_TYPES == 1 +#if TEST_TYPES == 1 TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", "[keys][segmented][topk][device][cluster]") { @@ -535,7 +535,7 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalig REQUIRE(expected_keys == keys_out_buffer); } -#endif // TEST_LAUNCH != 1 && TEST_TYPES == 1 +#endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per-segment k", "[keys][segmented][topk][device]", @@ -733,7 +733,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with variable-size segments and REQUIRE(expected_keys == keys_out_buffer); } -#if TEST_LAUNCH != 1 && TEST_TYPES == 2 +#if TEST_TYPES == 2 // Heavy-tie stress/regression: collapse the keys to only a handful of distinct values so the k-th key's bucket holds a // large set of tied candidates. This exercises the cluster agent's candidate path and, when built with // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM, the deterministic cross-CTA tie-break scan (cand_prefix + BlockScan ranks). The @@ -795,7 +795,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle heavy ties at the k-th boundar REQUIRE(expected_keys == keys_out_buffer); } -#endif // TEST_LAUNCH != 1 && TEST_TYPES == 2 +#endif // TEST_TYPES == 2 // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::MinKeys preserves -0.0f in output", "[keys][segmented][topk][device][float]") diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index ce01d8b2718..1bd0b123d12 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -324,9 +324,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small fixed-size segments" REQUIRE(expected_keys == keys_out_buffer); } -// Exercise the pair (key+value) generic fallback only in the float build (`TEST_TYPES == 1`); the test fixes its own -// key/value types, so repeating it for every key-type axis would only waste time on the expensive 1 Mi-element runs. -#if TEST_LAUNCH != 1 && TEST_TYPES == 1 +// Large-segment cluster pair tests, built only in the float build (`TEST_TYPES == 1`): they fix their own key/value +// types, so repeating them per key-type axis would only waste time on the expensive 1 Mi-element runs. They run on +// every launch id, including device launch (`lid_1`): the CDP static config's small resident capacity is what streams +// big segments and peels the unaligned tail edge, so it is the path that must cover them. +#if TEST_TYPES == 1 template struct cast_to_key_op { @@ -427,7 +429,82 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments through a non- REQUIRE(expected_keys == keys_out_buffer); } -#endif // TEST_LAUNCH != 1 && TEST_TYPES == 1 + +// Pair analog of the keys "large fixed-size unaligned segments" test: contiguous keys offset by `pad` take the +// block-load path with an unaligned head edge, and the 1 Mi-element segments stream, so the value payloads exercise +// the boundary-edge value writes that the small and non-contiguous pair tests above do not. Which boundary the launch +// stresses depends on its resident capacity: the host/graph configs keep the tail suffix resident (`full_slots > 1`), +// covering the head edge and the resident-suffix write; the device-launch (`lid_1`) static config has a small +// `full_slots`, so the tail suffix is peeled into `edge_keys` and the persistent `tail_edge_len`/`process_tail_edge` +// path is covered there. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned segments", + "[pairs][segmented][topk][device][cluster]", + select_direction_list) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = c2h::get<0, TestType>::value; + + constexpr segment_size_t static_max_segment_size = 1024 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 3; + + const int pad = GENERATE(0, 1, 3, 7); + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + // Contiguous key storage offset by `pad` so each segment base is unaligned (forces the head boundary edge). + c2h::device_vector keys_in_buffer(pad + num_items, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + // Value payload = global flattened index into the (pad-excluded) logical input, so each output value indexes back + // into `expected_keys`. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + // Logical input (pad excluded), flattened, for verification. + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_items}); + + // Verify (before sorting) that values stayed associated with their keys and that no index repeats within a segment. + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // Verify the selected keys are the correct top-k. + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segments", "[pairs][segmented][topk][device]", diff --git a/cub/test/catch2_test_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_segmented_topk_cluster_layout.cu index dd1c537126b..a943e172c15 100644 --- a/cub/test/catch2_test_segmented_topk_cluster_layout.cu +++ b/cub/test/catch2_test_segmented_topk_cluster_layout.cu @@ -32,12 +32,16 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) const auto cluster_tile_capacity = layout_t::template cluster_tile_capacity(cluster_blocks, block_tile_capacity); REQUIRE(cluster_tile_capacity > 0); + // The head is an edge (static SMEM), not a reserved chunk, so coverage is the full physical capacity. + REQUIRE(cluster_tile_capacity == static_cast(cluster_blocks) * block_tile_capacity); + // Worst-case resident chunks on any single rank, mirroring the agent: only the aligned region `[head_items, + // segment_size)` is chunked (`num_chunks`), then spread across `blocks` ranks (the `ceil_div` max for both the + // strided and blocked partitions). const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { - using size_t = cuda::std::int64_t; - const size_t head_chunks = head_items == 0 ? size_t{0} : size_t{1}; - const size_t tail_items = segment_size - head_items; - const size_t chunks = head_chunks + host_ceil_div(tail_items, size_t{layout_t::chunk_items}); + using size_t = cuda::std::int64_t; + const size_t tail_items = segment_size - head_items; + const size_t chunks = host_ceil_div(tail_items, size_t{layout_t::chunk_items}); return host_ceil_div(chunks, static_cast(blocks)); }; @@ -56,10 +60,11 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) REQUIRE(max_rank_chunks(cluster_tile_capacity, head_items, cluster_blocks) <= slots); } - const auto unreserved_cluster_tile_capacity = static_cast(cluster_blocks) * block_tile_capacity; + // Tightness: the full physical capacity fits resident (`slots` chunks per rank), one item beyond overflows. CAPTURE( c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, block_tile_capacity); - REQUIRE(max_rank_chunks(unreserved_cluster_tile_capacity, 1, cluster_blocks) == slots + 1); + REQUIRE(max_rank_chunks(cluster_tile_capacity, 0, cluster_blocks) == slots); + REQUIRE(max_rank_chunks(cluster_tile_capacity + 1, 0, cluster_blocks) == slots + 1); } template @@ -78,7 +83,7 @@ void check_layout_matrix() } } // namespace -TEST_CASE("Segmented TopK cluster SMEM layout reserves the unaligned head chunk", +TEST_CASE("Segmented TopK cluster SMEM layout exposes the full physical capacity (head is an edge, not a chunk)", "[keys][segmented][topk][cluster][layout]") { using default_policy = cub::detail::batched_topk_cluster::policy_selector; @@ -88,7 +93,9 @@ TEST_CASE("Segmented TopK cluster SMEM layout reserves the unaligned head chunk" cub::detail::batched_topk_cluster::smem_block_tile_layout; static_assert(default_float_layout::block_tile_capacity(0) == 0); static_assert(default_float_layout::template cluster_tile_capacity(8, 0) == 0); - static_assert(default_float_layout::template cluster_tile_capacity(1, default_float_layout::chunk_items) == 0); + // No head reservation: a one-chunk cluster reports its full chunk as coverage. + static_assert(default_float_layout::template cluster_tile_capacity(1, default_float_layout::chunk_items) + == default_float_layout::chunk_items); check_layout_matrix(); check_layout_matrix(); From 4bedd272214b8dbb0dfa82ff9c97c32e202b28e7 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 24 Jun 2026 04:05:26 +0200 Subject: [PATCH 041/126] Wire up determinism/tie-breaking requirements Add testing and switches for the benchmarks. --- .../bench/segmented_topk/fixed/keys.cu | 17 +- .../bench/segmented_topk/variable/indexed.cu | 16 +- .../bench/segmented_topk/variable/keys.cu | 17 +- cub/cub/agent/agent_batched_topk_cluster.cuh | 968 +++++++++--------- .../dispatch_batched_topk_cluster.cuh | 77 +- .../catch2_test_device_segmented_topk_keys.cu | 141 ++- ...catch2_test_device_segmented_topk_pairs.cu | 300 +++++- 7 files changed, 991 insertions(+), 545 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index d22900aa9b1..ad7dabd74e0 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -6,6 +6,8 @@ #include #include +#include +#include #include #include #include @@ -28,6 +30,17 @@ enum class topk_backend inline constexpr topk_backend selected_backend = topk_backend::baseline; +// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). +inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; +inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + +// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. +static_assert(selected_backend == topk_backend::cluster + || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed + && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), + "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " + "selected_tie_break at their defaults for the baseline/device backends."); + #if !TUNE_BASE struct tuned_policy_selector { @@ -80,7 +93,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { if constexpr (selected_backend == topk_backend::cluster) { - return cub::detail::batched_topk_cluster::dispatch_with_env( + return cub::detail::batched_topk_cluster::dispatch_with_env( d_keys_in, d_keys_out, static_cast(nullptr), @@ -94,7 +107,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( } else if constexpr (selected_backend == topk_backend::device) { - using num_segments_val_t = typename ::cuda::__argument::__traits::element_type; + using num_segments_val_t = typename ::cuda::args::__traits::element_type; const auto num_segs = cub::detail::params::get_param(num_segments, num_segments_val_t{0}); // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu index 1db819b8a9f..b6fa25d8b33 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include #include @@ -26,6 +28,18 @@ enum class topk_backend // key-value-pair path through `agent_batched_topk_cluster`; the baseline backend is kept for A/B comparison. inline constexpr topk_backend selected_backend = topk_backend::cluster; +// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). Only forwarded +// to the cluster backend -- the baseline backend does not implement these requirements. +inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; +inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + +// The baseline backend ignores these requirements, so require the defaults there to avoid a silently ignored selection. +static_assert(selected_backend == topk_backend::cluster + || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed + && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), + "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " + "selected_tie_break at their defaults for the baseline backend."); + template ( d_keys_in, d_keys_out, d_values_in, diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index 39679f5d630..0562f43ab74 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -10,6 +10,8 @@ #include #include +#include +#include #include #include #include @@ -30,6 +32,17 @@ enum class topk_backend inline constexpr topk_backend selected_backend = topk_backend::baseline; +// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). +inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; +inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + +// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. +static_assert(selected_backend == topk_backend::cluster + || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed + && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), + "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " + "selected_tie_break at their defaults for the baseline/device backends."); + // Env-based dispatch over the selected backend. The cluster and baseline backends route through their respective // `dispatch_with_env` entry points (temporary storage is allocated from the memory resource carried by `env`); the // device backend issues one `cub::DeviceTopK::MaxKeys` per segment, reading the host-side segment sizes. @@ -55,7 +68,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { if constexpr (selected_backend == topk_backend::cluster) { - return cub::detail::batched_topk_cluster::dispatch_with_env( + return cub::detail::batched_topk_cluster::dispatch_with_env( d_keys_in, d_keys_out, static_cast(nullptr), @@ -69,7 +82,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( } else if constexpr (selected_backend == topk_backend::device) { - using num_segments_val_t = typename ::cuda::__argument::__traits::element_type; + using num_segments_val_t = typename ::cuda::args::__traits::element_type; const auto num_segs = cub::detail::params::get_param(num_segments, num_segments_val_t{0}); // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 66deaa620da..1781bc1fee9 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -53,6 +53,8 @@ #include #include #include +#include +#include #include #include #include @@ -69,12 +71,6 @@ #include -// Opt-in deterministic tie-break for the cluster top-k final filter. When enabled, candidates tied at the k-th key's -// prefix are selected by a cluster-wide, index-ordered scan (smallest global indices by default, largest when -// `CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST` is also defined) instead of the nondeterministic racing atomics. -// Enabling it also switches `make_chunk_partition` to the blocked layout, on which the deterministic scan depends -// (CTA-rank order == ascending contiguous global-index ranges). - CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster @@ -144,6 +140,11 @@ struct smem_block_tile_layout // Cluster width is a runtime value (see `process_impl` for the readback), so // it is not a template parameter; per-block block_tile layout is still controlled // by the template parameters below. +// +// `Determinism` / `TieBreak` carry the requested `cuda::execution` requirements down from the dispatch layer. Any +// deterministic guarantee selects the cluster-wide, index-ordered tie-break scan (and the blocked chunk partition it +// depends on) over the nondeterministic racing atomics; `not_guaranteed` keeps the atomics. On the deterministic path +// `prefer_larger_index` reverses the scan order so the largest-index ties win (else the smallest). template ; using key_prefix_t = typename state_t::key_prefix_t; + // See the class comment: deterministic guarantees select the index-ordered scan; `tie_reversed` flips its order. + static constexpr bool deterministic = + Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed; + static constexpr bool tie_reversed = TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int histogram_items_per_thread = HistogramItemsPerThread; static constexpr int load_align_bytes = LoadAlignBytes; @@ -418,25 +426,27 @@ struct agent_batched_topk_cluster // // * Strided (default): chunk `i` goes to rank `i % cluster_size`, so each CTA walks `first, first+S, first+2S, // ...`. - // * Blocked (`CUB_ENABLE_CLUSTER_TOPK_DETERMINISM`): each CTA owns a contiguous run of `ceil_div(chunks, S)` - // chunks (the last non-empty rank gets the short remainder). Chunks are large enough that the per-CTA contiguous - // gmem footprint does not change L2/cache locality versus the strided walk. + // * Blocked (deterministic path): each CTA owns a contiguous run of `ceil_div(chunks, S)` chunks (the last + // non-empty rank gets the short remainder). // - // The blocked layout is a hard requirement of the deterministic tie-break (its cross-CTA scan assumes CTA-rank order - // matches ascending contiguous global-index ranges), so it is selected exactly when the determinism path is enabled. + // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches + // ascending contiguous global-index ranges), so it is selected exactly when `deterministic` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_size) const { -#if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) - const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_size)); - const offset_t first = static_cast(cluster_rank) * chunks_per_cta; - const offset_t count = (first < chunks) ? (::cuda::std::min) (chunks_per_cta, chunks - first) : offset_t{0}; - return {first, offset_t{1}, count}; -#else - return {static_cast(cluster_rank), - static_cast(cluster_size), - num_rank_chunks(chunks, cluster_rank, cluster_size)}; -#endif + if constexpr (deterministic) + { + const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_size)); + const offset_t first = static_cast(cluster_rank) * chunks_per_cta; + const offset_t count = (first < chunks) ? (::cuda::std::min) (chunks_per_cta, chunks - first) : offset_t{0}; + return {first, offset_t{1}, count}; + } + else + { + return {static_cast(cluster_rank), + static_cast(cluster_size), + num_rank_chunks(chunks, cluster_rank, cluster_size)}; + } } // Stage a small (`< load_align_items` items) unaligned run -- a boundary edge (head prefix / tail suffix) or a @@ -1091,11 +1101,7 @@ private: // makes the leader the *last* CTA in scan order so it never needs its own (merged-away) local candidate count: // prefer-smallest scans ascending by rank (leader = last rank), prefer-largest scans descending (leader = rank 0). // The nondeterministic path keeps rank 0 (unchanged codegen). -# if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) && !defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) - const unsigned int leader_rank = cluster_size - 1u; -# else - const unsigned int leader_rank = 0u; -# endif + const unsigned int leader_rank = (deterministic && !tie_reversed) ? (cluster_size - 1u) : 0u; // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). @@ -1559,208 +1565,245 @@ private: // trailing digits as smaller and erroneously reject candidates that share the identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); -# if defined(CUB_ENABLE_CLUSTER_TOPK_DETERMINISM) - // Deterministic tie-break: select the `num_kth` tied candidates with the smallest (default) or largest - // (`CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST`) global indices via a cluster-wide, index-ordered scan instead of - // the nondeterministic racing atomics. Strictly-selected keys still go to the front; ties fill the back by rank. -# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) - constexpr bool tie_reversed = true; -# else - constexpr bool tie_reversed = false; -# endif - // Early stop means the splitter bucket holds *exactly* the remaining `k` candidates cluster-wide, so all are - // winners with no surplus ties: fold them into the front (`num_back == 0`, `num_selected == k`) and skip the - // index-ordered cross-CTA scan, including its `hist[last_bucket]` read (the end-of-pass reset may already have - // cleared `hist` on an early break). Valid even when the early-stop *break* is compiled out. - const out_offset_t num_back = (leader_state->early_stop != ::cuda::std::uint32_t{0}) ? out_offset_t{0} : num_kth; - const out_offset_t num_selected = k - num_back; // front region; == k on early stop - - offset_t running = 0; - bool tie_active = false; - if (num_back != out_offset_t{0}) + if constexpr (deterministic) { - // Exclusive cross-CTA candidate-count scan. Each non-leader adds its local candidate count (its surviving - // `hist[last_bucket]`) into the `cand_prefix` of every CTA that follows it in scan order; the leader is last in - // scan order (see `leader_rank`) so it adds nothing and receives the sum of all preceding CTAs. - const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; - if (threadIdx.x == 0) + // Deterministic tie-break: select the `num_kth` tied candidates with the smallest (or largest, if `tie_reversed`) + // global indices via a cluster-wide index-ordered scan. Strictly-selected keys go to the front; ties fill the + // back by rank. Early stop means the splitter bucket holds *exactly* the remaining `k` candidates, so all are + // winners (`num_back == 0`, `num_selected == k`) and we skip the scan and its `hist[last_bucket]` read (the + // end-of-pass reset may already have cleared `hist`). Valid even when the early-stop *break* is compiled out. + const out_offset_t num_back = (leader_state->early_stop != ::cuda::std::uint32_t{0}) ? out_offset_t{0} : num_kth; + const out_offset_t num_selected = k - num_back; // front region; == k on early stop + + offset_t running = 0; + bool tie_active = false; + if (num_back != out_offset_t{0}) { - temp_storage.cand_prefix = 0; - } - cluster.sync(); - if (threadIdx.x == 0 && local_count != offset_t{0}) - { -# if defined(CUB_CLUSTER_TOPK_DETERMINISM_PREFER_LARGEST) - for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow - { - add_remote_prefix(r, local_count); - } -# else - for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + // Exclusive cross-CTA candidate-count scan: each non-leader adds its candidate count (`hist[last_bucket]`) into + // every later CTA's `cand_prefix` in scan order; the leader is last (see `leader_rank`) and receives the sum. + const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; + if (threadIdx.x == 0) { - add_remote_prefix(r, local_count); + temp_storage.cand_prefix = 0; } -# endif - } - cluster.sync(); - - running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) - tie_active = running < static_cast(num_back); - } - - // Process a flat span of `count` keys already in scan order (`get_key(pos)`), tiled by - // `threads_per_block * tie_break_items_per_thread`. Front keys go via `out_cnt`; surplus ties get a - // BlockScan-exclusive index rank (seeded by `running`) and, if `rank < num_back`, are written in reverse at - // `block_keys_out[k - 1 - rank]`. The running aggregate carries across tiles and across regions. - // `get_idx(pos)` returns the segment-local index of the key `get_key(pos)` returns, used only to load the value - // payload for written keys (compiled out in keys-only builds). - auto process_flat = [&](auto get_key, auto get_idx, int count) { - constexpr int items = tie_break_items_per_thread; - constexpr int tile = threads_per_block * items; - for (int tile_base = 0; tile_base < count; tile_base += tile) - { - key_t keys[items]; - offset_t flags[items]; - detail::topk::candidate_class cls[items]; - bool valid[items]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + cluster.sync(); + if (threadIdx.x == 0 && local_count != offset_t{0}) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - valid[i] = pos < count; - flags[i] = offset_t{0}; - if (valid[i]) + if constexpr (tie_reversed) { - keys[i] = get_key(pos); - cls[i] = identify_op(keys[i]); - flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow + { + add_remote_prefix(r, local_count); + } + } + else + { + for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + { + add_remote_prefix(r, local_count); + } } } + cluster.sync(); - // Strictly-selected keys go to the front via `out_cnt`; on early stop (`num_back == 0`) the splitter bucket's - // candidates are winners too and fold in here. Exactly `num_selected` are placed in `[0, num_selected)`. - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) + tie_active = running < static_cast(num_back); + } + + // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * items`. Front keys go + // via `out_cnt`; surplus ties get a BlockScan-exclusive rank (seeded by `running`) and, if `rank < num_back`, are + // written reversed at `block_keys_out[k - 1 - rank]`. `running` carries across tiles and regions. `get_idx(pos)` + // gives the segment-local index, used only to load the value payload (compiled out in keys-only builds). + auto process_flat = [&](auto get_key, auto get_idx, int count) { + constexpr int items = tie_break_items_per_thread; + constexpr int tile = threads_per_block * items; + for (int tile_base = 0; tile_base < count; tile_base += tile) { - const bool to_front = - valid[i] - && (cls[i] == detail::topk::candidate_class::selected - || (num_back == out_offset_t{0} && cls[i] == detail::topk::candidate_class::candidate)); - if (to_front) + key_t keys[items]; + offset_t flags[items]; + detail::topk::candidate_class cls[items]; + bool valid[items]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + valid[i] = pos < count; + flags[i] = offset_t{0}; + if (valid[i]) + { + keys[i] = get_key(pos); + cls[i] = identify_op(keys[i]); + flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + } } - } - // Surplus ties (`num_back > 0`): each candidate gets a BlockScan rank (seeded by `running`), written reversed - // at `block_keys_out[k - 1 - rank]` when `rank < num_back`. Skipped on early stop (`tie_active == false`). - if (tie_active) - { - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + // Strictly-selected keys go to the front via `out_cnt`; on early stop (`num_back == 0`) the splitter bucket's + // candidates are winners too and fold in here. Exactly `num_selected` are placed in `[0, num_selected)`. _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && flags[i] != offset_t{0}) + const bool to_front = + valid[i] + && (cls[i] == detail::topk::candidate_class::selected + || (num_back == out_offset_t{0} && cls[i] == detail::topk::candidate_class::candidate)); + if (to_front) { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); } } - running += tile_total; - if (running >= static_cast(num_back)) + + // Surplus ties (`num_back > 0`): each candidate gets a BlockScan rank (seeded by `running`), written reversed + // at `block_keys_out[k - 1 - rank]` when `rank < num_back`. Skipped on early stop (`tie_active == false`). + if (tie_active) { - tie_active = false; + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); + } + } + } + running += tile_total; + if (running >= static_cast(num_back)) + { + tie_active = false; + } + // The next tile reuses `scan_storage`; order this tile's reads against the next ExclusiveSum's writes. + __syncthreads(); } - // The next tile reuses `scan_storage`; order this tile's reads against the next ExclusiveSum's writes. - __syncthreads(); } - } - }; + }; - // Uniform early-exit between regions: stop once this CTA's ties are placed and all selected are placed - // cluster-wide. Lets the common prefer-smallest case skip the (expensive) overflow re-stream entirely. - auto should_stop = [&]() -> bool { - if (tie_active) - { - return false; - } - const out_offset_t out_now = *static_cast(&leader_state->out_cnt); - return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; - }; + // Uniform early-exit between regions: stop once this CTA's ties are placed and all selected are placed + // cluster-wide. Lets the common prefer-smallest case skip the (expensive) overflow re-stream entirely. + auto should_stop = [&]() -> bool { + if (tie_active) + { + return false; + } + const out_offset_t out_now = *static_cast(&leader_state->out_cnt); + return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; + }; - // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, which is - // visited separately after the overflow because it is the globally-last chunk. - int tail_count = 0; - int front_count = span_size(resident_keys); - if constexpr (use_block_load_to_shared) - { - if (force_tail_resident) + // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, which is + // visited separately after the overflow because it is the globally-last chunk. + int tail_count = 0; + int front_count = span_size(resident_keys); + if constexpr (use_block_load_to_shared) { - tail_count = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).count; - front_count = front_count - tail_count; + if (force_tail_resident) + { + tail_count = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).count; + front_count = front_count - tail_count; + } } - } - // Segment-local base of the resident-front span. The deterministic path always uses the blocked partition, so the - // front chunks `[part.first, part.first + front_chunks)` are contiguous in segment order and pack densely into the - // resident SMEM region; element `pos` of the front therefore maps to `front_seg_base + pos`. - const offset_t front_seg_base = get_chunk(part.first, segment_size_u32, head_items).offset; + // Segment-local base of the resident-front span. The deterministic path's blocked partition keeps the front + // chunks contiguous and densely packed, so element `pos` of the front maps to `front_seg_base + pos`. + const offset_t front_seg_base = get_chunk(part.first, segment_size_u32, head_items).offset; - auto process_resident = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - const int fc = front_count; - if (reversed) + auto process_resident = [&](bool reversed) { + if constexpr (use_block_load_to_shared) { - process_flat( - [&](int pos) { - return rfront[fc - 1 - pos]; - }, - [&](int pos) { - return front_seg_base + static_cast(fc - 1 - pos); - }, - fc); + key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + const int fc = front_count; + if (reversed) + { + process_flat( + [&](int pos) { + return rfront[fc - 1 - pos]; + }, + [&](int pos) { + return front_seg_base + static_cast(fc - 1 - pos); + }, + fc); + } + else + { + process_flat( + [&](int pos) { + return rfront[pos]; + }, + [&](int pos) { + return front_seg_base + static_cast(pos); + }, + fc); + } } else { - process_flat( - [&](int pos) { - return rfront[pos]; - }, - [&](int pos) { - return front_seg_base + static_cast(pos); - }, - fc); + const int rc = static_cast(my_resident_chunks); + for (int s = 0; s < rc; ++s) + { + const int local_slot = reversed ? (rc - 1 - s) : s; + const offset_t chunk_idx = part.global_index(static_cast(local_slot)); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const int cc = chunk.count; + const offset_t base_off = chunk.offset; + key_t* const ck = slot_keys_unpadded(local_slot); + if (reversed) + { + process_flat( + [&](int pos) { + return ck[cc - 1 - pos]; + }, + [&](int pos) { + return base_off + static_cast(cc - 1 - pos); + }, + cc); + } + else + { + process_flat( + [&](int pos) { + return ck[pos]; + }, + [&](int pos) { + return base_off + static_cast(pos); + }, + cc); + } + } } - } - else - { - const int rc = static_cast(my_resident_chunks); - for (int s = 0; s < rc; ++s) + }; + + auto process_overflow = [&](bool reversed) { + for (offset_t oo = 0; oo < overflow_count; ++oo) { - const int local_slot = reversed ? (rc - 1 - s) : s; - const offset_t chunk_idx = part.global_index(static_cast(local_slot)); + const offset_t o = reversed ? (overflow_count - 1 - oo) : oo; + const offset_t chunk_idx = part.global_index(overflow_base + o); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const int cc = chunk.count; - const offset_t base_off = chunk.offset; - key_t* const ck = slot_keys_unpadded(local_slot); + // Visit only the aligned bulk: a peeled tail suffix is handled by `process_tail_edge`; every other overflow + // chunk has `bulk == count`. (The generic fallback never peels, so it uses the full count.) + int cc; + if constexpr (use_block_load_to_shared) + { + cc = static_cast(split_chunk(block_keys_base, chunk).bulk); + } + else + { + cc = chunk.count; + } + const offset_t base_off = chunk.offset; if (reversed) { process_flat( [&](int pos) { - return ck[cc - 1 - pos]; + return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; }, [&](int pos) { return base_off + static_cast(cc - 1 - pos); @@ -1771,7 +1814,7 @@ private: { process_flat( [&](int pos) { - return ck[pos]; + return block_keys_in[static_cast(base_off + static_cast(pos))]; }, [&](int pos) { return base_off + static_cast(pos); @@ -1779,357 +1822,306 @@ private: cc); } } - } - }; + }; - auto process_overflow = [&](bool reversed) { - for (offset_t oo = 0; oo < overflow_count; ++oo) - { - const offset_t o = reversed ? (overflow_count - 1 - oo) : oo; - const offset_t chunk_idx = part.global_index(overflow_base + o); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - // Visit only the aligned bulk: a peeled suffix tail's suffix is handled by `process_tail_edge`, so it is - // excluded here; every other overflow chunk has `bulk == count`. (The generic fallback never peels and reads - // the trailing items straight from gmem, so it uses the full count.) - int cc; + auto process_tail = [&](bool reversed) { if constexpr (use_block_load_to_shared) { - cc = static_cast(split_chunk(block_keys_base, chunk).bulk); - } - else - { - cc = chunk.count; + if (!force_tail_resident) + { + return; + } + key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + key_t* const tptr = rfront + front_count; + const int tc = tail_count; + // The forced-resident tail is the globally-last chunk `chunks-1`; its segment-local base is that chunk's + // offset. + const offset_t tail_seg_base = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).offset; + if (reversed) + { + process_flat( + [&](int pos) { + return tptr[tc - 1 - pos]; + }, + [&](int pos) { + return tail_seg_base + static_cast(tc - 1 - pos); + }, + tc); + } + else + { + process_flat( + [&](int pos) { + return tptr[pos]; + }, + [&](int pos) { + return tail_seg_base + static_cast(pos); + }, + tc); + } } - const offset_t base_off = chunk.offset; - if (reversed) + }; + + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. Processed + // before every chunk (ascending) / after (descending) to keep global index order. `count == 0` is a no-op, so + // non-head ranks skip it. + auto process_head_edge = [&](bool reversed) { + if constexpr (use_block_load_to_shared) { - process_flat( - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; - }, - [&](int pos) { - return base_off + static_cast(cc - 1 - pos); - }, - cc); + key_t* const he = temp_storage.edge_keys; + const int hc = head_edge_len; + if (reversed) + { + process_flat( + [&](int pos) { + return he[hc - 1 - pos]; + }, + [&](int pos) { + return static_cast(hc - 1 - pos); + }, + hc); + } + else + { + process_flat( + [&](int pos) { + return he[pos]; + }, + [&](int pos) { + return static_cast(pos); + }, + hc); + } } - else + }; + + // Peeled tail suffix edge (tail owner, `full_slots == 1`): the segment's highest indices, staged in `edge_keys`. + // Processed after every chunk (ascending) / before (descending). `tail_edge_len == 0` (tail kept resident or not + // owned here) makes it a no-op. + auto process_tail_edge = [&](bool reversed) { + if constexpr (use_block_load_to_shared) { - process_flat( - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(pos))]; - }, - [&](int pos) { - return base_off + static_cast(pos); - }, - cc); + key_t* const te = temp_storage.edge_keys + head_edge_cap; + const int tc = tail_edge_len; + const offset_t tail_base = segment_size_u32 - static_cast(tc); + if (reversed) + { + process_flat( + [&](int pos) { + return te[tc - 1 - pos]; + }, + [&](int pos) { + return tail_base + static_cast(tc - 1 - pos); + }, + tc); + } + else + { + process_flat( + [&](int pos) { + return te[pos]; + }, + [&](int pos) { + return tail_base + static_cast(pos); + }, + tc); + } } - } - }; + }; - auto process_tail = [&](bool reversed) { - if constexpr (use_block_load_to_shared) + if constexpr (tie_reversed) { - if (!force_tail_resident) + process_tail_edge(true); + if (!should_stop()) { - return; + process_tail(true); } - key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - key_t* const tptr = rfront + front_count; - const int tc = tail_count; - // The forced-resident tail is the globally-last chunk `chunks-1`; its segment-local base is that chunk's - // offset. - const offset_t tail_seg_base = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).offset; - if (reversed) + if (!should_stop()) { - process_flat( - [&](int pos) { - return tptr[tc - 1 - pos]; - }, - [&](int pos) { - return tail_seg_base + static_cast(tc - 1 - pos); - }, - tc); + process_overflow(true); } - else + if (!should_stop()) + { + process_resident(true); + } + if (!should_stop()) { - process_flat( - [&](int pos) { - return tptr[pos]; - }, - [&](int pos) { - return tail_seg_base + static_cast(pos); - }, - tc); + process_head_edge(true); } } - }; - - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. Processed - // before every chunk (ascending) / after every chunk (descending) so the global index order holds across regions. - // `process_flat` with `count == 0` is a barrier-free no-op, so non-head ranks (`head_edge_len == 0`) skip it. - auto process_head_edge = [&](bool reversed) { - if constexpr (use_block_load_to_shared) + else { - key_t* const he = temp_storage.edge_keys; - const int hc = head_edge_len; - if (reversed) + process_head_edge(false); + if (!should_stop()) { - process_flat( - [&](int pos) { - return he[hc - 1 - pos]; - }, - [&](int pos) { - return static_cast(hc - 1 - pos); - }, - hc); + process_resident(false); } - else + if (!should_stop()) { - process_flat( - [&](int pos) { - return he[pos]; - }, - [&](int pos) { - return static_cast(pos); - }, - hc); + process_overflow(false); } - } - }; - - // Peeled tail suffix edge (tail owner, `full_slots == 1`): the segment's highest indices, staged in `edge_keys`. - // Processed after every chunk (ascending) / before every chunk (descending). `tail_edge_len == 0` (the tail was - // kept resident, or this rank does not own it) makes this a barrier-free no-op. - auto process_tail_edge = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - key_t* const te = temp_storage.edge_keys + head_edge_cap; - const int tc = tail_edge_len; - const offset_t tail_base = segment_size_u32 - static_cast(tc); - if (reversed) + if (!should_stop()) { - process_flat( - [&](int pos) { - return te[tc - 1 - pos]; - }, - [&](int pos) { - return tail_base + static_cast(tc - 1 - pos); - }, - tc); + process_tail(false); } - else + if (!should_stop()) { - process_flat( - [&](int pos) { - return te[pos]; - }, - [&](int pos) { - return tail_base + static_cast(pos); - }, - tc); + process_tail_edge(false); } } - }; - - if constexpr (tie_reversed) - { - process_tail_edge(true); - if (!should_stop()) - { - process_tail(true); - } - if (!should_stop()) - { - process_overflow(true); - } - if (!should_stop()) - { - process_resident(true); - } - if (!should_stop()) - { - process_head_edge(true); - } } else { - process_head_edge(false); - if (!should_stop()) - { - process_resident(false); - } - if (!should_stop()) - { - process_overflow(false); - } - if (!should_stop()) - { - process_tail(false); - } - if (!should_stop()) + if constexpr (is_keys_only) { - process_tail_edge(false); - } - } -# else // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM - if constexpr (is_keys_only) - { - auto write_selected = [&](const key_t& key) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[pos] = key; - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); - if (back_pos < num_kth) + auto write_selected = [&](const key_t& key) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = k - 1 - back_pos; + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); block_keys_out[pos] = key; } - } - }; - // Fold the resident keys as the streamer's `mid` work so they overlap the first wave of overflow reloads. The - // writes are order-independent atomics, so interleaving resident and overflow output is safe, and the resident - // SMEM slots are disjoint from the streaming slots, so `mid` reads never race the in-flight loads. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) - { - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) + else if (res == detail::topk::candidate_class::candidate) { - const offset_t chunk_idx = part.global_index(p); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + } } - } - // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_edges(write_selected); - }; - streamer.process_pass(write_selected, fold_resident); - } - else - { - // Pair (key + value) path. The `out_cnt`/`out_back_cnt` atomics are unchanged; for each written key we - // additionally load its value payload from gmem at the key's segment-local index `seg_idx` and store it at the - // same output slot. Keys are reused exactly as in the keys-only path - resident keys from SMEM, overflow keys - // from the streaming SMEM pipeline via `process_pass_indexed` (the generic fallback re-reads them from gmem); - // only the values are fetched from gmem (no value streaming). As in the keys-only path the output order is not - // preserved (the resident keys are folded in as the streamer's `mid` work), which the non-deterministic path - // permits. - auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); - block_keys_out[pos] = key; - write_value(pos, seg_idx); - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); - if (back_pos < num_kth) + }; + // Fold the resident keys as the streamer's `mid` work so they overlap the first overflow reloads. Writes are + // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) + { + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + } + else + { + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const offset_t chunk_idx = part.global_index(p); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + } + } + // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). + fold_edges(write_selected); + }; + streamer.process_pass(write_selected, fold_resident); + } + else + { + // Pair (key + value) path. Same `out_cnt`/`out_back_cnt` atomics and key reuse as the keys-only path; for each + // written key we additionally load its value from gmem at the segment-local index `seg_idx` and store it at the + // same slot (values are never streamed). Output order is not preserved, as the non-deterministic path permits. + auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = k - 1 - back_pos; + const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); block_keys_out[pos] = key; write_value(pos, seg_idx); } - } - }; - - // Iterate a contiguous run of `count` keys whose element `local` has segment-local index `base_off + local`. - // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). - auto write_run = [&](auto get_key, offset_t base_off, int count) { - const int iterations = ::cuda::ceil_div(count, threads_per_block); - _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread) - for (int j = 0; j < iterations; ++j) - { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < count) + else if (res == detail::topk::candidate_class::candidate) { - write_selected_idx(get_key(local), base_off + static_cast(local)); + const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + write_value(pos, seg_idx); + } } - } - }; + }; - // Fold the resident keys (and their values) as the streamer's `mid` work so they overlap the first wave of - // overflow reloads, exactly as in the keys-only path: order-independent atomic writes into a disjoint output, and - // resident SMEM slots disjoint from the streaming slots, so `mid` never races the in-flight loads. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) - { - // Resident keys are densely packed in slot order; each chunk's keys are contiguous, so a running cursor over - // the packed region recovers per-chunk spans. The last slot holds the forced-resident tail when applicable. - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - int cursor = 0; - for (offset_t s = 0; s < my_resident_chunks; ++s) + // Iterate a contiguous run of `count` keys whose element `local` has segment-local index `base_off + local`. + // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). + auto write_run = [&](auto get_key, offset_t base_off, int count) { + const int iterations = ::cuda::ceil_div(count, threads_per_block); + _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread) + for (int j = 0; j < iterations; ++j) { - const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; - const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - const int cc = chunk.count; - write_run( - [&](int local) { - return rk[cursor + local]; - }, - base_off, - cc); - cursor += cc; + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < count) + { + write_selected_idx(get_key(local), base_off + static_cast(local)); + } } - } - else - { - for (offset_t p = 0; p < my_resident_chunks; ++p) + }; + + // Fold the resident keys (and their values) as the streamer's `mid` work, exactly as in the keys-only path. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) { - const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - write_run( - [&](int local) { - return chunk_keys[local]; - }, - base_off, - chunk.count); + // Resident keys are densely packed in slot order, so a running cursor recovers per-chunk spans. The last + // slot holds the forced-resident tail when applicable. + key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + int cursor = 0; + for (offset_t s = 0; s < my_resident_chunks; ++s) + { + const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; + const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + const int cc = chunk.count; + write_run( + [&](int local) { + return rk[cursor + local]; + }, + base_off, + cc); + cursor += cc; + } } - } - // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the - // peeled tail suffix at `segment_size - tail_edge_len`. Order-independent atomic writes, value fetched per key. - if constexpr (use_block_load_to_shared) - { - if (head_edge_len > 0) + else { - write_run( - [&](int local) { - return temp_storage.edge_keys[local]; - }, - offset_t{0}, - head_edge_len); + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + write_run( + [&](int local) { + return chunk_keys[local]; + }, + base_off, + chunk.count); + } } - if (tail_edge_len > 0) + // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the + // peeled tail suffix at `segment_size - tail_edge_len`. Value fetched per key. + if constexpr (use_block_load_to_shared) { - write_run( - [&](int local) { - return temp_storage.edge_keys[head_edge_cap + local]; - }, - segment_size_u32 - static_cast(tail_edge_len), - tail_edge_len); + if (head_edge_len > 0) + { + write_run( + [&](int local) { + return temp_storage.edge_keys[local]; + }, + offset_t{0}, + head_edge_len); + } + if (tail_edge_len > 0) + { + write_run( + [&](int local) { + return temp_storage.edge_keys[head_edge_cap + local]; + }, + segment_size_u32 - static_cast(tail_edge_len), + tail_edge_len); + } } - } - }; + }; - // Overflow chunks: reuse the keys from the streaming SMEM pipeline (block-load path; only the generic fallback - // re-reads them from gmem), and fetch each selected key's value at its segment-local index `seg_idx`. The - // resident keys above are folded in as the streamer's `mid` work to hide the first reload wave's latency. - streamer.process_pass_indexed(write_selected_idx, fold_resident); + // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected + // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. + streamer.process_pass_indexed(write_selected_idx, fold_resident); + } } -# endif // CUB_ENABLE_CLUSTER_TOPK_DETERMINISM // Final cluster barrier: hold every block in the cluster until all DSMEM // atomics into the leader's state are complete. Without this, a fast diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index efbd9b2e44c..b9740eaf42d 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -42,6 +42,8 @@ #include #include +#include +#include #include #include #include @@ -89,6 +91,8 @@ template +// `Determinism`/`TieBreak` carry the requested `cuda::execution` requirements down to the kernel and agent (the public +// env-based interface is handled separately). They default to the nondeterministic, racing-atomics behavior. +template < + ::cuda::execution::determinism::__determinism_t Determinism = + ::cuda::execution::determinism::__determinism_t::__not_guaranteed, + ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, + typename KeyInputItItT, + typename KeyOutputItItT, + typename ValueInputItItT, + typename ValueOutputItItT, + typename SegmentSizeParameterT, + typename KParameterT, + typename SelectDirectionT, + typename NumSegmentsParameterT, + typename TotalNumItemsGuaranteeT, + typename PolicySelector = policy_selector> CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -328,6 +346,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( LoadAlignBytes, BitsPerPass, TieBreakItemsPerThread, + Determinism, + TieBreak, KeyInputItItT, KeyOutputItItT, ValueInputItItT, @@ -337,6 +357,12 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SelectDirectionParameterT, NumSegmentsParameterT>; + // TODO: This should be taken care of in the public env-based interface. + // A tie-break preference is only meaningful once the result set itself is deterministic. + static_assert(Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed + || TieBreak == ::cuda::execution::tie_break::__tie_break_t::__unspecified, + "A tie-break preference requires a deterministic execution requirement"); + static_assert(ChunkBytes % LoadAlignBytes == 0); static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); // Static-footprint estimate for the device-side CDP fallback, which cannot query `cudaFuncGetAttributes`. @@ -396,6 +422,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( LoadAlignBytes, BitsPerPass, TieBreakItemsPerThread, + Determinism, + TieBreak, KeyInputItItT, KeyOutputItItT, ValueInputItItT, @@ -677,16 +705,23 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // algorithm is single-phase (one kernel launch over a placeholder allocation), but routing it through the shared // env-based machinery keeps the call shape identical to the baseline backend and lets it pick up the stream, memory // resource, and tuning carried by the environment. -template > +// +// `Determinism`/`TieBreak` are forwarded verbatim to `dispatch` (the public env-based interface is handled separately); +// they default to the nondeterministic, unspecified-tie-break behavior. +template < + ::cuda::execution::determinism::__determinism_t Determinism = + ::cuda::execution::determinism::__determinism_t::__not_guaranteed, + ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, + typename KeyInputItItT, + typename KeyOutputItItT, + typename ValueInputItItT, + typename ValueOutputItItT, + typename SegmentSizeParameterT, + typename KParameterT, + typename SelectDirectionT, + typename NumSegmentsParameterT, + typename TotalNumItemsGuaranteeT, + typename EnvT = ::cuda::std::execution::env<>> [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, @@ -701,7 +736,7 @@ template ( env, [&](auto policy_sel, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { - return dispatch( + return dispatch( d_temp_storage, temp_storage_bytes, d_key_segments_it, diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 3565307af14..627626b257f 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include @@ -31,9 +33,8 @@ struct is_minus_zero } }; -// Maps a flat element index to one of only 8 distinct key values, so a large segment contains many duplicates and the -// k-th key's bucket holds a large tied-candidate set. Used to stress the cluster agent's candidate/tie-break path (and -// the deterministic tie-break scan when CUB_ENABLE_CLUSTER_TOPK_DETERMINISM is defined). +// Maps a flat element index to one of only 8 distinct key values, so a large segment has many duplicates and the k-th +// key's bucket holds a large tied-candidate set. Stresses the cluster agent's candidate/tie-break path. template struct heavy_tie_key_op { @@ -54,6 +55,9 @@ enum class topk_backend inline constexpr topk_backend selected_backend = topk_backend::cluster; template ( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -108,7 +112,14 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( // %PARAM% TEST_LAUNCH lid 0:1:2 DECLARE_TMPL_LAUNCH_WRAPPER( - dispatch_batched_topk_keys, batched_topk_keys, cub::detail::topk::select Direction, Direction); + dispatch_batched_topk_keys, + batched_topk_keys, + ESCAPE_LIST( + cub::detail::topk::select Direction, + cuda::execution::determinism::__determinism_t Determinism = + cuda::execution::determinism::__determinism_t::__not_guaranteed, + cuda::execution::tie_break::__tie_break_t TieBreak = cuda::execution::tie_break::__tie_break_t::__unspecified), + ESCAPE_LIST(Direction, Determinism, TieBreak)); // Total segment size using max_segment_size_list = c2h::enum_type_list; @@ -144,6 +155,30 @@ using key_types = using select_direction_list = c2h::enum_type_list; +// A (determinism, tie-break) requirement pair, used as a single compile-time test axis. Not a full cross product: a +// tie-break preference is only meaningful with a deterministic requirement and is `gpu_to_gpu` by definition. +template +struct det_tie +{ + static constexpr auto determinism = Determinism; + static constexpr auto tie_break = TieBreak; +}; + +// The 5 valid determinism/tie-break combinations. The selected key multiset is identical for all of them (which tied +// index wins is not observable in keys-only output), so each is verified against the same reference -- confirming every +// code path computes the correct top-k, including the boundary tie count. +using det_tie_combos = + c2h::type_list, + det_tie, + det_tie, + det_tie, + det_tie>; + C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", "[keys][segmented][topk][device]", key_types, @@ -226,13 +261,29 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", } #if TEST_TYPES == 1 -TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", - "[keys][segmented][topk][device][cluster]") +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", + "[keys][segmented][topk][device][cluster][determinism]", + det_tie_combos) { using key_t = float; using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; + // Requirement under test. Random keys rarely tie at the boundary, so this mainly exercises the blocked-partition + + // streaming/edge interaction; the key result is invariant to the requirement, so all combinations match the + // reference. + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + // `static_max_segment_size` is chosen to exceed the largest all-resident cluster coverage (~16 blocks worth of // resident SMEM), so the 1 Mi-element segments force the agent's gmem-streaming overflow path (including an // unaligned overflow tail via `- 31`), while the 128 Ki-element segment still runs fully resident under the same @@ -262,7 +313,7 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned c2h::device_vector expected_keys(num_segments * segment_size, thrust::no_init); thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); - batched_topk_keys( + batched_topk_keys( d_keys_in, d_keys_out, cuda::args::immediate{segment_size, cuda::args::bounds()}, @@ -306,13 +357,27 @@ struct counting_segment_keys_op } }; -TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-contiguous key iterator", - "[keys][segmented][topk][device][cluster]") +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-contiguous key iterator", + "[keys][segmented][topk][device][cluster][determinism]", + det_tie_combos) { using key_t = float; using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; + // Requirement under test (the key result is invariant to it; exercises the generic streaming path). + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + // The counting-iterator key source is non-contiguous, so the agent uses its generic overflow-streaming path rather // than BlockLoadToShared. `static_max_segment_size` exceeds the largest all-resident cluster coverage, so the 1 Mi // -element segments stream (incl. an unaligned `- 31` tail), while the 128 Ki-element segment validates the generic @@ -340,7 +405,7 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non- auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); - batched_topk_keys( + batched_topk_keys( d_keys_in, d_keys_out, cuda::args::immediate{segment_size, cuda::args::bounds()}, @@ -462,13 +527,27 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment } #if TEST_TYPES == 1 -TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", - "[keys][segmented][topk][device][cluster]") +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", + "[keys][segmented][topk][device][cluster][determinism]", + det_tie_combos) { using key_t = float; using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; + // Requirement under test (the key result is invariant to it; exercises mixed resident/streaming segments). + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch // mixes streaming segments (the 1 Mi-element ones, one with an unaligned `- 31` overflow tail) with fully-resident // segments (96 Ki + 17 and 257 elements). @@ -519,7 +598,7 @@ TEST_CASE("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalig c2h::device_vector expected_keys(num_items, thrust::no_init); thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); - batched_topk_keys( + batched_topk_keys( d_keys_in, d_keys_out, cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, @@ -735,23 +814,33 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with variable-size segments and } #if TEST_TYPES == 2 // Heavy-tie stress/regression: collapse the keys to only a handful of distinct values so the k-th key's bucket holds a -// large set of tied candidates. This exercises the cluster agent's candidate path and, when built with -// CUB_ENABLE_CLUSTER_TOPK_DETERMINISM, the deterministic cross-CTA tie-break scan (cand_prefix + BlockScan ranks). The -// returned top-k *value* set must be correct in either mode (tied keys are equal, so which tied index is chosen is not -// observable in keys-only output, but the count of each value at the boundary must be exact). +// large set of tied candidates. Exercises the cluster agent's candidate path and, on a deterministic requirement, the +// cross-CTA tie-break scan (cand_prefix + BlockScan ranks); the boundary value counts must be exact in either mode. // -// The tie-break / candidate-count path is independent of the key type: after radix conversion the 8 small, non-negative -// values twiddle trivially, so a single representative key exercises the same code. We therefore fix the key type and -// build this test only in the matching `TEST_TYPES` variant (the one whose axis already targets this type) instead of -// redundantly across every type axis. The float key path is additionally covered by the float-keyed cluster tests -// above, and per-type behavior by the generic tests. +// The tie-break path is key-type-independent (the 8 small values twiddle trivially), so we fix the key type and build +// this only in the matching `TEST_TYPES` variant. Float keys and per-type behavior are covered by the tests above. C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle heavy ties at the k-th boundary", - "[keys][segmented][topk][device][cluster]", - select_direction_list) + "[keys][segmented][topk][device][cluster][determinism]", + select_direction_list, + det_tie_combos) { using key_t = cuda::std::uint64_t; constexpr auto direction = c2h::get<0, TestType>::value; + // Requirement under test. The key multiset is invariant to the preference, so every combination matches the same + // reference; tie-rich data makes the deterministic scan do real work, so a boundary miscount would surface here. + using combo = c2h::get<1, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; @@ -780,7 +869,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle heavy ties at the k-th boundar c2h::device_vector expected_keys(keys_in_buffer); - batched_topk_keys( + batched_topk_keys( d_keys_in, d_keys_out, cuda::args::immediate{segment_size, cuda::args::bounds()}, diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 1bd0b123d12..8dc576fa26d 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -12,8 +12,15 @@ #include #include +#include +#include #include +#include +#include +#include +#include + #include "catch2_test_device_topk_common.cuh" #include "catch2_test_launch_helper.h" #include @@ -57,10 +64,12 @@ enum class topk_backend inline constexpr topk_backend selected_backend = topk_backend::cluster; -// Routes the key-value (pairs) top-k to either the baseline or the cluster backend, threading both the key and value -// iterators-of-iterators through. The cluster backend exercises the key-value-pair path in -// `agent_batched_topk_cluster`. +// Routes the key-value (pairs) top-k to the baseline or cluster backend. `Determinism`/`TieBreak` are forwarded to the +// cluster dispatch to request the deterministic tie-break path; they default to the nondeterministic behavior. template ( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -118,7 +127,14 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( // %PARAM% TEST_LAUNCH lid 0:1:2 DECLARE_TMPL_LAUNCH_WRAPPER( - dispatch_batched_topk_pairs, batched_topk_pairs, cub::detail::topk::select Direction, Direction); + dispatch_batched_topk_pairs, + batched_topk_pairs, + ESCAPE_LIST( + cub::detail::topk::select Direction, + cuda::execution::determinism::__determinism_t Determinism = + cuda::execution::determinism::__determinism_t::__not_guaranteed, + cuda::execution::tie_break::__tie_break_t TieBreak = cuda::execution::tie_break::__tie_break_t::__unspecified), + ESCAPE_LIST(Direction, Determinism, TieBreak)); // Total segment size using max_segment_size_list = c2h::enum_type_list; @@ -504,6 +520,280 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned REQUIRE(expected_keys == keys_out_buffer); } + +// Tie-break preferences exercised by the deterministic tests below (only meaningful with a deterministic requirement). +using tie_break_pref_list = + c2h::enum_type_list; + +// Deterministic requirements exercised with an *unspecified* tie-break: both must yield a reproducible result. +using determinism_list = + c2h::enum_type_list; + +// Host reference for the deterministic top-k: per segment, stably picks the k best (key, global-index) pairs via a +// plain lexicographic sort -- max sorts descending (larger index wins ties), min ascending (smaller wins). When the +// requested preference is the opposite, the index is reversed (last_index - idx) so the natural order breaks ties the +// desired way without a custom comparator (and, unlike negation, stays non-negative for unsigned indices). Returns the +// selected indices sorted ascending per segment, to compare against the unordered device output as a set. +template +c2h::host_vector reference_deterministic_topk_indices( + const c2h::host_vector& h_keys, + SegSizeT num_segments, + SegSizeT segment_size, + SegSizeT k, + cub::detail::topk::select direction, + bool prefer_larger_index) +{ + const bool want_max = direction == cub::detail::topk::select::max; + const bool reverse_index = want_max != prefer_larger_index; + const IndexT last_index = static_cast(num_segments * segment_size - 1); + // Tie-break sort key (an involution, so it also decodes): reverses index order when requested, else identity. + const auto encode = [&](IndexT idx) { + return reverse_index ? static_cast(last_index - idx) : idx; + }; + + c2h::host_vector selected(static_cast(num_segments * k)); + std::vector> pairs(static_cast(segment_size)); + for (SegSizeT seg = 0; seg < num_segments; ++seg) + { + const SegSizeT base = seg * segment_size; + for (SegSizeT i = 0; i < segment_size; ++i) + { + const IndexT idx = static_cast(base + i); + pairs[static_cast(i)] = {h_keys[static_cast(base + i)], encode(idx)}; + } + if (want_max) + { + std::partial_sort(pairs.begin(), pairs.begin() + k, pairs.end(), std::greater>{}); + } + else + { + std::partial_sort(pairs.begin(), pairs.begin() + k, pairs.end()); + } + const auto seg_begin = selected.begin() + static_cast(seg * k); + for (SegSizeT i = 0; i < k; ++i) + { + seg_begin[i] = encode(pairs[static_cast(i)].second); + } + std::sort(seg_begin, seg_begin + k); + } + return selected; +} + +// Deterministic tie-break: a specified preference is `gpu_to_gpu` deterministic by definition, so the cluster path +// returns a uniquely defined top-k. Few distinct key values pack many ties into the k-th bucket so the preference (not +// the key comparison) drives the result; the value payload is the global index, so we compare per-segment index sets +// against the host reference (within-top-k order is unspecified). `lid_1` streams the 64 Ki segments. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the index-ordered top-k", + "[pairs][segmented][topk][device][cluster][determinism]", + select_direction_list, + tie_break_pref_list) +{ + // Only the cluster backend provides the deterministic tie-break; the strict index-set check is meaningless elsewhere. + if constexpr (selected_backend != topk_backend::cluster) + { + SKIP("Deterministic tie-break is only provided by the cluster backend"); + } + + using key_t = cuda::std::uint32_t; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = c2h::get<0, TestType>::value; + constexpr auto tie_break = c2h::get<1, TestType>::value; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr bool prefer_larger = tie_break == cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + + constexpr segment_size_t static_max_segment_size = 64 * 1024; + constexpr segment_size_t static_max_k = 64 * 1024; + constexpr segment_index_t num_segments = 2; + + const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{4096}, segment_size_t{64 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(segment_size, k, num_segments, direction, prefer_larger); + + // Few distinct key values -> many tied candidates in the k-th bucket. Contiguous -> resident BlockLoadToShared path. + c2h::device_vector keys_in_buffer(num_items, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer, key_t{0}, key_t{7}); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + // Values = global flattened index, so each selected value points back into the flattened input. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_items}); + + // Values still belong to their keys, and no source index is selected twice. + c2h::device_vector expected_keys(keys_in_buffer); + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // The deterministic path must return *exactly* the index-ordered top-k. Compare per-segment selected index sets. + c2h::host_vector h_keys = keys_in_buffer; + const c2h::host_vector ref = + reference_deterministic_topk_indices(h_keys, num_segments, segment_size, k, direction, prefer_larger); + + c2h::host_vector h_values_out = values_out_buffer; + for (segment_index_t seg = 0; seg < num_segments; ++seg) + { + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + std::sort(seg_begin, seg_begin + k); + } + + REQUIRE(ref == h_values_out); +} + +// A second valid cluster tuning for the gpu_to_gpu determinism test: starts from the default policy (for its valid +// launch configs) and overrides the shape knobs (block size, items-per-thread, pipeline depth, tie-break granularity). +struct alt_cluster_tuning +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> cub::detail::batched_topk_cluster::cluster_topk_policy + { + auto policy = cub::detail::batched_topk_cluster::policy_selector{}(cc); + policy.threads_per_block = 256; + policy.histogram_items_per_thread = 2; + policy.pipeline_stages = 2; + policy.tie_break_items_per_thread = 2; + return policy; + } +}; + +// Reproducibility with an *unspecified* tie-break (which tied candidate wins is an implementation detail). We run twice +// and require the same selected index set, per each requirement's contract: `run_to_run` only promises repeated runs of +// the *same* config agree (so both runs share a tuning); `gpu_to_gpu` must be config-independent (so the second run +// uses a different valid tuning), mirroring the reduce/scan deterministic tests. Within-top-k order is unspecified, so +// we compare sorted sets. `gpu_to_gpu` cannot be checked more strictly without a second device. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break is reproducible", + "[pairs][segmented][topk][device][cluster][determinism]", + select_direction_list, + determinism_list) +{ + // Only the cluster backend honors determinism; the cross-tuning check is meaningless elsewhere. + if constexpr (selected_backend != topk_backend::cluster) + { + SKIP("Determinism is only provided by the cluster backend"); + } + + using key_t = cuda::std::uint32_t; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = c2h::get<0, TestType>::value; + constexpr auto determinism = c2h::get<1, TestType>::value; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + + constexpr segment_size_t static_max_segment_size = 64 * 1024; + constexpr segment_size_t static_max_k = 64 * 1024; + constexpr segment_index_t num_segments = 2; + + const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{4096}, segment_size_t{64 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(segment_size, k, num_segments, direction); + + // Few distinct key values -> many tied candidates at the k-th bucket, so the deterministic scan actually does work. + c2h::device_vector keys_in_buffer(num_items, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer, key_t{0}, key_t{7}); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + + // One output buffer per tuning. + c2h::device_vector keys_out_a(num_segments * k, thrust::no_init); + c2h::device_vector keys_out_b(num_segments * k, thrust::no_init); + c2h::device_vector values_out_a(num_segments * k, thrust::no_init); + c2h::device_vector values_out_b(num_segments * k, thrust::no_init); + + // Calls the dispatch directly (the launch wrapper does not thread a policy selector) so each run picks a tuning. + const auto run_with_tuning = + [&](auto policy_selector, c2h::device_vector& keys_out, c2h::device_vector& values_out) { + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::detail::batched_topk_cluster::dispatch( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::constant{}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_items}, + cudaStream_t{0}, + policy_selector); + }; + REQUIRE(invoke(nullptr) == cudaSuccess); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(invoke(thrust::raw_pointer_cast(temp_storage.data())) == cudaSuccess); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + }; + + run_with_tuning(cub::detail::batched_topk_cluster::policy_selector{}, keys_out_a, values_out_a); + // run_to_run: same tuning twice. gpu_to_gpu: a different tuning for the stronger config-independent check. + if constexpr (determinism == cuda::execution::determinism::__determinism_t::__gpu_to_gpu) + { + run_with_tuning(alt_cluster_tuning{}, keys_out_b, values_out_b); + } + else + { + run_with_tuning(cub::detail::batched_topk_cluster::policy_selector{}, keys_out_b, values_out_b); + } + + // Sanity: values still belong to their keys and no source index repeats within a segment. + c2h::device_vector expected_keys(keys_in_buffer); + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_a, values_out_a) == true); + REQUIRE(verify_unique_indices(values_out_a, num_segments, k) == true); + + // Determinism: both runs must select the same per-segment index set. Sort each segment (within-top-k order is free). + c2h::host_vector h_values_a = values_out_a; + c2h::host_vector h_values_b = values_out_b; + for (segment_index_t seg = 0; seg < num_segments; ++seg) + { + const auto a_begin = h_values_a.begin() + static_cast(seg * k); + const auto b_begin = h_values_b.begin() + static_cast(seg * k); + std::sort(a_begin, a_begin + k); + std::sort(b_begin, b_begin + k); + } + REQUIRE(h_values_a == h_values_b); +} #endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segments", From 57a5385798e9f8cfd5baec6720b8efb59adf9618 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 24 Jun 2026 04:08:07 +0200 Subject: [PATCH 042/126] Improve tuning parameters (unrolling) --- .../device/dispatch/tuning/tuning_batched_topk_cluster.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 23a65f39b57..632d3aa58d1 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -58,13 +58,13 @@ make_policy(::cuda::std::inplace_vector Date: Wed, 24 Jun 2026 18:15:17 +0200 Subject: [PATCH 043/126] Improve dispatch by minimizing waves Also gets rid of launch config table in tuning params and improves testing. --- cub/cub/detail/launcher/cuda_runtime.cuh | 10 + .../dispatch_batched_topk_cluster.cuh | 223 +++++++++--------- .../tuning/tuning_batched_topk_cluster.cuh | 94 +------- .../catch2_test_device_segmented_topk_keys.cu | 62 +++++ 4 files changed, 190 insertions(+), 199 deletions(-) diff --git a/cub/cub/detail/launcher/cuda_runtime.cuh b/cub/cub/detail/launcher/cuda_runtime.cuh index c67d3f9c79d..2cf8dbb374b 100644 --- a/cub/cub/detail/launcher/cuda_runtime.cuh +++ b/cub/cub/detail/launcher/cuda_runtime.cuh @@ -165,6 +165,16 @@ struct TripleChevronFactory return CubDebug( ::cudaOccupancyMaxPotentialClusterSize(&cluster_size, reinterpret_cast(kernel_ptr), config)); } + + // Query how many thread-block clusters can be co-resident on the whole device for `kernel_ptr` under `config` + // (device-wide clusters per wave). `config` must carry the cluster dimension; the query accounts for the kernel's + // static shared memory, the requested dynamic shared memory, register pressure, and GPC co-scheduling. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + max_active_clusters(int& num_clusters, Kernel kernel_ptr, const ::cudaLaunchConfig_t* config) + { + return CubDebug(::cudaOccupancyMaxActiveClusters(&num_clusters, reinterpret_cast(kernel_ptr), config)); + } }; } // namespace detail diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index b9740eaf42d..535979dc067 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -46,6 +46,7 @@ #include #include #include +#include #include #include #include @@ -71,13 +72,6 @@ template return ::cuda::args::__highest_(segment_sizes); } -struct launch_config -{ - int cluster_blocks; - int dynamic_smem_bytes; - ::cuda::std::uint32_t block_tile_capacity; -}; - // ----------------------------------------------------------------------------- // Kernel entry points // ----------------------------------------------------------------------------- @@ -242,8 +236,8 @@ __launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, // ----------------------------------------------------------------------------- // Dispatch // ----------------------------------------------------------------------------- -// Keys-only; every segment must fit in one cluster_tile. Host picks -// `(cluster_blocks, dynamic_smem_bytes)` at runtime from a finite table; CDP +// Keys and key/value pairs; segments larger than the resident block tile are streamed from gmem by the agent. The host +// path picks `(cluster_blocks, dynamic_smem_bytes)` at runtime via a wave-aware cluster-size search (see below); CDP // uses the static kernel at `max_portable_cluster_blocks` and portable SMEM. // CDP launch body, empty when CDP is disabled. Wrapped in a macro because @@ -370,7 +364,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( constexpr int static_smem_bytes = static_cast(sizeof(typename agent_t::TempStorage)); const auto max_seg_size = runtime_max_segment_size(segment_sizes); - using max_seg_size_t = decltype(+max_seg_size); // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. size_t allocation_sizes[1] = {1}; @@ -463,9 +456,9 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( cfg.numAttrs = 1; // Resolve the per-block opt-in shared-memory budget and the kernel's static footprint from the driver so - // the dynamic-SMEM math below matches exactly what the launch permits. The table entries express the - // documented total per-block budget (`cudaDevAttrMaxSharedMemoryPerBlockOptin`); the usable dynamic portion - // is that budget minus the static footprint. + // the dynamic-SMEM math below matches exactly what the launch permits. The opt-in budget + // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`) is the documented total per-block budget; the usable dynamic + // portion (`max_dynamic_smem_bytes`) is that budget minus the static footprint. int device_id = 0; if (const auto error = CubDebug(cudaGetDevice(&device_id))) { @@ -492,18 +485,13 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const int max_dynamic_smem_bytes = (max_smem_optin_bytes > nondynamic_smem_bytes) ? max_smem_optin_bytes - nondynamic_smem_bytes : 0; - const auto runtime_policy = PolicySelector{}(::cuda::compute_capability{sm_version / 10}); - _CCCL_ASSERT(runtime_policy.launch_configs.size() <= max_launch_configs, - "Cluster TopK launch config table exceeds policy capacity"); - - const auto total_to_dynamic_smem = [&](int total_smem_bytes) { - return (total_smem_bytes > nondynamic_smem_bytes) ? total_smem_bytes - nondynamic_smem_bytes : 0; - }; - + // Raise the kernel's dynamic-SMEM opt-in lazily: occupancy queries and the launch must not request more than the + // currently configured `cudaFuncAttributeMaxDynamicSharedMemorySize`. The kernel's compile-time default already + // permits the portable 48 KiB total, i.e. that budget minus the static footprint. constexpr int portable_total_smem_bytes = 48 * 1024; - const int portable_dynamic_smem_bytes = total_to_dynamic_smem(portable_total_smem_bytes); - int configured_dynamic_smem_limit = portable_dynamic_smem_bytes; - const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { + int configured_dynamic_smem_limit = + (portable_total_smem_bytes > nondynamic_smem_bytes) ? portable_total_smem_bytes - nondynamic_smem_bytes : 0; + const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { if (dynamic_smem_bytes <= configured_dynamic_smem_limit) { return cudaSuccess; @@ -520,117 +508,124 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaSuccess; }; - launch_config selected_config{0, 0, 0}; - launch_config largest_supported_config{0, 0, 0}; - // The table is intentionally tiny; scan all entries and select the lowest coverage that satisfies the bound - // rather than relying on the declaration order to be sorted by coverage. - for (const auto policy_config : runtime_policy.launch_configs) + // Wave-aware cluster-size selection. The free variable is the cluster size `C` (one cluster per segment); each + // `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more + // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We + // pick the `C` that minimizes the number of waves, breaking ties toward the largest `C` (= smallest SMEM = most + // L1), which matches the profiled fast configs. We enumerate `C` analytically rather than discovering SMEM tiers + // via occupancy, so a register-limited occupancy (e.g. 1 CTA/SM) cannot collapse the candidate set. + const auto seg = static_cast<::cuda::std::uint64_t>(max_seg_size); + const auto chunk_items_u64 = static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); + const int max_block_tile_capacity = static_cast(layout_t::block_tile_capacity(max_dynamic_smem_bytes)); + if (max_block_tile_capacity <= 0) { - const int candidate_cluster_blocks = policy_config.cluster_blocks; - const int candidate_dynamic_smem = total_to_dynamic_smem(policy_config.total_smem_bytes); - if (candidate_dynamic_smem > max_dynamic_smem_bytes) - { - continue; - } + // Not even one load-aligned chunk fits in the opt-in budget; the kernel cannot run. + return cudaErrorInvalidValue; + } - const auto candidate_block_tile_capacity = layout_t::block_tile_capacity(candidate_dynamic_smem); - if (candidate_block_tile_capacity == 0) - { - continue; - } + // `S_res(items)`: smallest chunk-granular dynamic SMEM whose per-CTA capacity reaches `items`. + const auto smem_for_block_capacity = [&](::cuda::std::uint64_t items) { + const auto slots = ::cuda::ceil_div(items, chunk_items_u64); + return layout_t::base_padding_bytes + static_cast(slots) * layout_t::chunk_bytes; + }; - launch_config candidate{candidate_cluster_blocks, candidate_dynamic_smem, candidate_block_tile_capacity}; - const auto candidate_cluster_tile_capacity = layout_t::template cluster_tile_capacity( - candidate.cluster_blocks, candidate.block_tile_capacity); + // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (cap HW + // max). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest fully-resident `C`. + // Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. + // `numeric_limits::max()` for an unbounded deferred sequence); narrowing such a `C_lo` to `int` could wrap to + // a small (or negative) value and wrongly enter the resident branch instead of the oversize/streaming fallback. + const int c_full = static_cast( + (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), + ::cuda::ceil_div(seg, chunk_items_u64))); + const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); - // Every smem-fitting candidate is considered: it may either tighten the covering selection or improve the - // largest-coverage fallback used for oversize segments. The table is tiny, so the extra driver queries below - // are negligible. - if (const auto error = ensure_dynamic_smem_limit(candidate_dynamic_smem)) - { - return error; - } + int cluster_blocks = 0; + int dynamic_smem_sel = 0; - cfg.dynamicSmemBytes = static_cast(candidate_dynamic_smem); - int candidate_hw_max_cluster_blocks = 0; - if (const auto error = - launcher_factory.max_potential_cluster_size(candidate_hw_max_cluster_blocks, dynamic_kernel, &cfg)) - { - return error; - } - candidate_hw_max_cluster_blocks = - (::cuda::std::min) (candidate_hw_max_cluster_blocks, max_supported_cluster_blocks); - if (candidate_cluster_blocks > candidate_hw_max_cluster_blocks) + if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) + { + // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every + // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. + // Scan `C` in `[max(1, C_lo), C_full]`, minimize waves, tie-break largest `C`. + const int c_begin = (::cuda::std::max) (1, static_cast(c_lo)); + const int c_end = (::cuda::std::max) (c_begin, c_full); + ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); + for (int c = c_begin; c <= c_end; ++c) { - continue; - } + const auto per_block_items = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(c)); + const int s_res = smem_for_block_capacity(per_block_items); + if (s_res > max_dynamic_smem_bytes) + { + continue; // unreachable for c >= C_lo, but guards the SMEM budget regardless. + } - // Track the largest-coverage hardware-supported config for all parameter kinds. Segments that exceed every - // finite candidate (always possible for per-segment sizes, and now also for static/uniform bounds) fall back - // to this config and re-stream the overflow from gmem in the agent. - { - const auto largest_supported_cluster_tile_capacity = layout_t::template cluster_tile_capacity( - largest_supported_config.cluster_blocks, largest_supported_config.block_tile_capacity); - if (largest_supported_config.cluster_blocks == 0 - || candidate_cluster_tile_capacity > largest_supported_cluster_tile_capacity) + if (const auto error = ensure_dynamic_smem_limit(s_res)) { - largest_supported_config = candidate; + return error; } - } - // Among the candidates that fully cover the segment(s), keep the one with the lowest coverage. - if (candidate_cluster_tile_capacity >= max_seg_size) - { - const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( - selected_config.cluster_blocks, selected_config.block_tile_capacity); - if (selected_config.cluster_blocks == 0 || candidate_cluster_tile_capacity < selected_cluster_tile_capacity) + // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must + // be a multiple of the cluster size. The returned value is the device-wide clusters-per-wave (capacity), + // independent of grid size, and accounts for the static footprint and register pressure internally. + cluster_attr.val.clusterDim.x = static_cast(c); + cfg.gridDim = dim3(static_cast(c), 1, 1); + cfg.dynamicSmemBytes = static_cast(s_res); + int clusters_per_wave = 0; + if (const auto error = launcher_factory.max_active_clusters(clusters_per_wave, dynamic_kernel, &cfg)) { - selected_config = candidate; + return error; + } + if (clusters_per_wave <= 0) + { + continue; // cluster size not launchable at this SMEM. } - } - } - if (selected_config.cluster_blocks == 0) - { - // No finite candidate covers the segment(s); fall back to the largest hardware-supported config and let the - // agent stream the overflow. Applies to per-segment sizes and static/uniform bounds alike. - selected_config = largest_supported_config; - } - if (selected_config.cluster_blocks == 0) - { - return cudaErrorInvalidValue; + const auto waves = ::cuda::ceil_div( + static_cast<::cuda::std::uint64_t>(num_seg_val), static_cast<::cuda::std::uint64_t>(clusters_per_wave)); + // Min waves; tie-break largest `C`. The loop ascends in `C`, so `<=` keeps the largest at equal waves. + if (cluster_blocks == 0 || waves <= best_waves) + { + best_waves = waves; + cluster_blocks = c; + dynamic_smem_sel = s_res; + } + } } - const auto selected_cluster_tile_capacity = layout_t::template cluster_tile_capacity( - selected_config.cluster_blocks, selected_config.block_tile_capacity); - if (selected_cluster_tile_capacity >= max_seg_size) + if (cluster_blocks == 0) { - // `cluster_tile_capacity == physical` (the head is an edge, not a reserved chunk), so the physical capacity - // need only reach `max_seg_size`; the chunk-granular round-up below grows it as needed. - const auto required_physical_cluster_tile_items = static_cast<::cuda::std::uint64_t>(max_seg_size); - selected_config.cluster_blocks = static_cast( - ::cuda::ceil_div(required_physical_cluster_tile_items, - static_cast<::cuda::std::uint64_t>(selected_config.block_tile_capacity))); - - const auto required_block_tile_capacity = ::cuda::ceil_div( - required_physical_cluster_tile_items, static_cast<::cuda::std::uint64_t>(selected_config.cluster_blocks)); - const auto required_slots = - ::cuda::ceil_div(required_block_tile_capacity, static_cast<::cuda::std::uint64_t>(layout_t::chunk_items)); - selected_config.dynamic_smem_bytes = - layout_t::base_padding_bytes + static_cast(required_slots) * layout_t::chunk_bytes; - selected_config.block_tile_capacity = layout_t::block_tile_capacity(selected_config.dynamic_smem_bytes); + // Oversize (`C_lo > HW max`) or nothing launchable in range: full residency is impossible, so maximize + // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. + if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) + { + return error; + } + cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` + cfg.gridDim = dim3(1, 1, 1); + cfg.dynamicSmemBytes = static_cast(max_dynamic_smem_bytes); + int hw_max_cluster_blocks = 0; + if (const auto error = launcher_factory.max_potential_cluster_size(hw_max_cluster_blocks, dynamic_kernel, &cfg)) + { + return error; + } + hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, max_supported_cluster_blocks); + if (hw_max_cluster_blocks <= 0) + { + return cudaErrorInvalidValue; + } + cluster_blocks = hw_max_cluster_blocks; + dynamic_smem_sel = max_dynamic_smem_bytes; } - if (const auto error = CubDebug(cudaFuncSetAttribute( - reinterpret_cast(dynamic_kernel), - cudaFuncAttributeMaxDynamicSharedMemorySize, - selected_config.dynamic_smem_bytes))) + // The launch needs `MaxDynamicSharedMemorySize >= dynamic_smem_sel`; the scan already raised the limit past the + // largest probed SMEM, so this is a no-op unless the selected config skipped the scan. + if (const auto error = ensure_dynamic_smem_limit(dynamic_smem_sel)) { return error; } - const int cluster_blocks = selected_config.cluster_blocks; - const auto block_tile_capacity = selected_config.block_tile_capacity; + const int dynamic_smem_bytes = dynamic_smem_sel; + const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); @@ -644,7 +639,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( if (const auto error = CubDebug( launcher_factory(dim3(static_cast(grid_blocks), 1, 1), dim3(static_cast(ThreadsPerBlock), 1, 1), - static_cast<::cuda::std::size_t>(selected_config.dynamic_smem_bytes), + static_cast<::cuda::std::size_t>(dynamic_smem_bytes), stream, /*dependent_launch=*/false, dim3(static_cast(cluster_blocks), 1, 1)) diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 632d3aa58d1..117ad681cfc 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -14,22 +14,12 @@ #endif // no system header #include -#include CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { -inline constexpr int max_launch_configs = 19; - -struct cluster_topk_launch_config -{ - int cluster_blocks; - // Total per-block shared memory budget, including the agent's static shared memory. - int total_smem_bytes; -}; - -// Per-block execution shape; the dispatch picks `cluster_blocks` and the -// dynamic shared-memory block_tile capacity at runtime. +// Per-block execution shape. The dispatch picks the cluster size and the dynamic shared-memory block_tile capacity at +// runtime (occupancy/wave-aware), so the policy carries only the per-block tuning knobs. struct cluster_topk_policy { int threads_per_block; @@ -40,21 +30,9 @@ struct cluster_topk_policy int bits_per_pass; int min_blocks_per_sm; int tie_break_items_per_thread; - ::cuda::std::inplace_vector launch_configs; }; -template -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_launch_configs(Configs... configs) - -> ::cuda::std::inplace_vector -{ - ::cuda::std::inplace_vector result; - (result.emplace_back(configs), ...); - return result; -} - -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto -make_policy(::cuda::std::inplace_vector launch_configs) - -> cluster_topk_policy +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy { return cluster_topk_policy{ /*threads_per_block=*/512, @@ -64,75 +42,21 @@ make_policy(::cuda::std::inplace_vector cluster_topk_policy -{ - return make_policy(make_launch_configs( - cluster_topk_launch_config{2, 31 * 1024}, // 62 KiB - // cluster_topk_launch_config{4, 31 * 1024}, // 124 KiB - cluster_topk_launch_config{2, 63 * 1024}, // 126 KiB - cluster_topk_launch_config{2, 99 * 1024}, // 198 KiB - // cluster_topk_launch_config{8, 31 * 1024}, // 248 KiB - // cluster_topk_launch_config{4, 63 * 1024}, // 252 KiB - cluster_topk_launch_config{2, 131 * 1024}, // 262 KiB - cluster_topk_launch_config{2, 163 * 1024}, // 326 KiB - cluster_topk_launch_config{2, 195 * 1024}, // 390 KiB - // cluster_topk_launch_config{4, 99 * 1024}, // 396 KiB - cluster_topk_launch_config{2, 227 * 1024}, // 454 KiB - // cluster_topk_launch_config{8, 63 * 1024}, // 504 KiB - cluster_topk_launch_config{4, 131 * 1024}, // 524 KiB - cluster_topk_launch_config{4, 163 * 1024}, // 652 KiB - cluster_topk_launch_config{4, 195 * 1024}, // 780 KiB - // cluster_topk_launch_config{8, 99 * 1024}, // 792 KiB - cluster_topk_launch_config{4, 227 * 1024}, // 908 KiB - cluster_topk_launch_config{8, 131 * 1024}, // 1048 KiB - cluster_topk_launch_config{8, 163 * 1024}, // 1304 KiB - cluster_topk_launch_config{8, 195 * 1024}, // 1560 KiB - cluster_topk_launch_config{8, 227 * 1024}, // 1816 KiB - // cluster_topk_launch_config{16, 131 * 1024}, // 2096 KiB - // cluster_topk_launch_config{16, 163 * 1024}, // 2608 KiB - // cluster_topk_launch_config{16, 195 * 1024}, // 3120 KiB - cluster_topk_launch_config{16, 227 * 1024} // 3632 KiB - )); -} - -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto sm120_policy() -> cluster_topk_policy -{ - return make_policy(make_launch_configs( - cluster_topk_launch_config{2, 31 * 1024}, // 62 KiB - // cluster_topk_launch_config{4, 31 * 1024}, // 124 KiB - cluster_topk_launch_config{2, 63 * 1024}, // 126 KiB - cluster_topk_launch_config{2, 99 * 1024}, // 198 KiB - // cluster_topk_launch_config{8, 31 * 1024}, // 248 KiB - cluster_topk_launch_config{4, 63 * 1024}, // 252 KiB - cluster_topk_launch_config{4, 99 * 1024}, // 396 KiB - cluster_topk_launch_config{8, 63 * 1024}, // 504 KiB - cluster_topk_launch_config{8, 99 * 1024} // 792 KiB - )); + /*tie_break_items_per_thread=*/8}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool { - return policy.chunk_bytes % policy.load_align_bytes == 0 && policy.launch_configs.size() <= max_launch_configs; + return policy.chunk_bytes % policy.load_align_bytes == 0; } -static_assert(is_valid_policy(sm90_100_110_policy())); -static_assert(is_valid_policy(sm120_policy())); -// Default selector for cluster-capable architectures (SM 9.0+). +static_assert(is_valid_policy(make_policy())); +// Default selector for cluster-capable architectures (SM 9.0+). The tuning is currently identical across CCs. struct policy_selector { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const - -> cluster_topk_policy + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy { - if (cc >= ::cuda::compute_capability{12, 0}) - { - return sm120_policy(); - } - - return sm90_100_110_policy(); + return make_policy(); } }; } // namespace detail::batched_topk_cluster diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 627626b257f..4df5ac73e27 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -614,6 +614,68 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign REQUIRE(expected_keys == keys_out_buffer); } + +// An un-annotated argument is auto-wrapped as an immediate with *loose* bounds, so its reported highest segment size is +// `numeric_limits::max()`. No fully-resident cluster config can cover that, so the dispatch must take the +// oversize/streaming fallback (max cluster + max SMEM) without the analytic cluster-size math overflowing while +// narrowing the derived sizes to `int`. The real (uniform) segment size stays small, so the streamed result must still +// be exact. Exercises both the auto-wrapping of un-annotated args and the loose-bound fallback path. +C2H_TEST("DeviceBatchedTopK::MaxKeys work with an un-annotated (loose-bound) segment size", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + // The oversize/streaming fallback for a loose bound is a cluster-dispatch concern; the baseline backend selects its + // config differently, so only assert it where the new analytic selector runs. + if constexpr (selected_backend != topk_backend::cluster) + { + SKIP("The oversize/streaming fallback for loose segment-size bounds is exercised by the cluster backend"); + } + + constexpr auto direction = cub::detail::topk::select::max; + constexpr segment_size_t static_max_k = 1024; + + // Real, uniform segment size kept small; it is passed un-annotated below so the dispatch instead sees a loose + // `numeric_limits::max()` upper bound. + const segment_size_t segment_size = + GENERATE_COPY(values({segment_size_t{1}, segment_size_t{4096}, segment_size_t{8192}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + const segment_index_t num_segments = + GENERATE_COPY(values({segment_index_t{1}, segment_index_t{8}, segment_index_t{64}})); + + CAPTURE(segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // `segment_size` is passed un-annotated (a raw scalar): the args framework auto-wraps it as an immediate whose + // static upper bound is the loose `numeric_limits::max()`. + batched_topk_keys( + d_keys_in, + d_keys_out, + segment_size, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + // Top-k output is unordered; sort each output segment before comparison. + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} #endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per-segment k", From f98b50c13c9d788ecde788dabe126469227fc2a2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 25 Jun 2026 17:29:53 +0200 Subject: [PATCH 044/126] Optimize single CTA case Use CTA barrier instead of cluster barrier and also lets the "leader" downgrade atomics scope and avoid MAPA. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 51 +++++++++++++++----- 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 1781bc1fee9..f27ba1b8998 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1078,6 +1078,23 @@ private: // Stripped on sub-SM-9.0 device passes; uses `cluster_group`, which is only // declared when `_CG_HAS_CLUSTER_GROUP` is set. #if defined(_CG_HAS_CLUSTER_GROUP) + // Synchronize the segment's cluster. A width-1 "cluster" is a single CTA whose state is entirely block-local, so a + // cluster-scoped barrier is unnecessary and `__syncthreads()` provides sufficient ordering. `single_cta` is derived + // from `cluster.num_blocks()` and is therefore uniform across the cluster, so the branch is reached uniformly and + // barrier reachability is preserved. + _CCCL_DEVICE _CCCL_FORCEINLINE static void + cluster_or_block_sync(::cooperative_groups::cluster_group& cluster, bool single_cta) + { + if (single_cta) + { + __syncthreads(); + } + else + { + cluster.sync(); + } + } + template _CCCL_DEVICE _CCCL_FORCEINLINE void run(::cooperative_groups::cluster_group& cluster, @@ -1096,6 +1113,9 @@ private: auto block_keys_in = d_key_segments_it[segment_id]; const auto segment_size_u32 = static_cast(segment_size); const unsigned int cluster_size = cluster.num_blocks(); + // A width-1 cluster is a lone CTA: route barriers to `__syncthreads()`, keep `state`/atomics block-local, and use + // CTA-scope histogram atomics (no cross-rank DSMEM folds exist to be mutually atomic with). + const bool single_cta = (cluster_size == 1u); // Leader rank. The leader owns the cluster-merged histogram and the shared `state`. The deterministic tie-break // makes the leader the *last* CTA in scan order so it never needs its own (merged-away) local candidate count: @@ -1105,7 +1125,8 @@ private: // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = cluster.map_shared_rank(&temp_storage.state, leader_rank); + state_t* leader_state = + single_cta ? &temp_storage.state : cluster.map_shared_rank(&temp_storage.state, leader_rank); // Per-block local copy of `kth_key_bits` so each key check hits the // block's own SMEM rather than DSMEM during the histogram loop. Built up one splitter digit per pass from the @@ -1262,10 +1283,12 @@ private: { extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool is_leader_block = cluster_rank == leader_rank; - auto add_first_pass = [&](const key_t& key) { + // Cluster-scope leader atomics are only needed to stay mutually atomic with the non-leaders' DSMEM folds; a lone + // CTA has none, so it uses the cheaper CTA scope. + const bool leader_cluster_scope = (cluster_rank == leader_rank) && !single_cta; + auto add_first_pass = [&](const key_t& key) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, is_leader_block); + hist_inc(hist_smem32, bucket, leader_cluster_scope); }; if constexpr (use_block_load_to_shared) @@ -1427,12 +1450,12 @@ private: // Step 1: block-private histogram via shared-space `red` (see `hist_inc`): leader uses cluster scope to be // mutually atomic with the non-leaders' Step 2 DSMEM folds, non-leaders use the cheaper cta scope. const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool is_leader_block = cluster_rank == leader_rank; + const bool leader_cluster_scope = (cluster_rank == leader_rank) && !single_cta; auto add_hist = [&](const key_t& key) { if (identify_op(key) == detail::topk::candidate_class::candidate) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, is_leader_block); + hist_inc(hist_smem32, bucket, leader_cluster_scope); } }; @@ -1494,7 +1517,7 @@ private: } } - cluster.sync(); + cluster_or_block_sync(cluster, single_cta); // Step 3: the leader walks the merged `hist` (raw counts in the // reduce-then-scan path, cluster-wide inclusive scan in the @@ -1514,7 +1537,7 @@ private: } reset_hist(); } - cluster.sync(); + cluster_or_block_sync(cluster, single_cta); // End-of-pass splitter fold. Every thread reads the leader's just-published `kth_bucket` directly through DSMEM // (a single `u32`, ordered by the `cluster.sync()` above) and folds it into its own `kth_key_bits_local`, so the @@ -1586,7 +1609,7 @@ private: { temp_storage.cand_prefix = 0; } - cluster.sync(); + cluster_or_block_sync(cluster, single_cta); if (threadIdx.x == 0 && local_count != offset_t{0}) { if constexpr (tie_reversed) @@ -1604,7 +1627,7 @@ private: } } } - cluster.sync(); + cluster_or_block_sync(cluster, single_cta); running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) tie_active = running < static_cast(num_back); @@ -2127,8 +2150,9 @@ private: // atomics into the leader's state are complete. Without this, a fast // block (e.g. one whose block_tile is entirely padding) can return while another // block is still writing to leader-resident memory through DSMEM, which - // surfaces as a "cluster target block not present" exception. - cluster.sync(); + // surfaces as a "cluster target block not present" exception. A lone CTA has no remote writers, so a block barrier + // suffices (and orders its own final writes). + cluster_or_block_sync(cluster, single_cta); } _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() @@ -2138,6 +2162,7 @@ private: // Runtime cluster width matches the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). const unsigned int cluster_blocks = cluster.num_blocks(); + const bool single_cta = (cluster_blocks == 1u); const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) @@ -2184,7 +2209,7 @@ private: temp_storage.state.early_stop = 0; } reset_hist(); - cluster.sync(); + cluster_or_block_sync(cluster, single_cta); [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( select_directions, segment_id, [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { From ca10336b8fc15b9ebdb3f5c1e9778afc11319bb0 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 25 Jun 2026 18:18:05 +0200 Subject: [PATCH 045/126] Prefer single-CTA when possible --- .../device/dispatch/dispatch_batched_topk_cluster.cuh | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 535979dc067..fc30719be85 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -542,7 +542,16 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( int cluster_blocks = 0; int dynamic_smem_sel = 0; - if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) + if (c_lo == 1) + { + // Single-CTA fast path: the segment fits resident in one CTA, so launch a width-1 "cluster" (the agent's + // cluster-barrier-free path) instead of spreading it across more CTAs for parallelism. `S_res(seg)` is within + // budget and one CTA is always launchable, so the occupancy probe is skipped (the shared + // `ensure_dynamic_smem_limit` below raises the opt-in for the selected SMEM). + cluster_blocks = 1; + dynamic_smem_sel = smem_for_block_capacity(seg); + } + else if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) { // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. From 347fd7f5b138d7a354c8de51202621d78a664ccd Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 25 Jun 2026 19:01:31 +0200 Subject: [PATCH 046/126] Prefer single-CTA only up to a tunable threshold --- .../dispatch_batched_topk_cluster.cuh | 25 +++++++++++++------ .../tuning/tuning_batched_topk_cluster.cuh | 4 ++- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index fc30719be85..02ce6f4054d 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -537,17 +537,18 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const int c_full = static_cast( (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), ::cuda::ceil_div(seg, chunk_items_u64))); - const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); + const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); + const bool prefer_single_cta = seg <= static_cast<::cuda::std::uint64_t>(policy.single_cta_max_segment_size); int cluster_blocks = 0; int dynamic_smem_sel = 0; - if (c_lo == 1) + if (c_lo == 1 && prefer_single_cta) { - // Single-CTA fast path: the segment fits resident in one CTA, so launch a width-1 "cluster" (the agent's - // cluster-barrier-free path) instead of spreading it across more CTAs for parallelism. `S_res(seg)` is within - // budget and one CTA is always launchable, so the occupancy probe is skipped (the shared - // `ensure_dynamic_smem_limit` below raises the opt-in for the selected SMEM). + // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's + // cluster-barrier-free path beats spreading it across more CTAs. `S_res(seg)` is within budget and one CTA is + // always launchable, so the occupancy probe is skipped (the shared `ensure_dynamic_smem_limit` below raises the + // opt-in for the selected SMEM). Larger fully-resident segments fall through to the wave-aware search below. cluster_blocks = 1; dynamic_smem_sel = smem_for_block_capacity(seg); } @@ -555,8 +556,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( { // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. - // Scan `C` in `[max(1, C_lo), C_full]`, minimize waves, tie-break largest `C`. - const int c_begin = (::cuda::std::max) (1, static_cast(c_lo)); + // Scan `C` in `[max(C_lo, 2), C_full]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. + const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); const int c_end = (::cuda::std::max) (c_begin, c_full); ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); for (int c = c_begin; c <= c_end; ++c) @@ -599,6 +600,14 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( dynamic_smem_sel = s_res; } } + + if (cluster_blocks == 0 && c_lo == 1) + { + // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, + // but `C_lo == 1` guarantees `S_res(seg)` fits the budget and one CTA is always launchable. + cluster_blocks = 1; + dynamic_smem_sel = smem_for_block_capacity(seg); + } } if (cluster_blocks == 0) diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 117ad681cfc..3e2bc339b4f 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -30,6 +30,7 @@ struct cluster_topk_policy int bits_per_pass; int min_blocks_per_sm; int tie_break_items_per_thread; + int single_cta_max_segment_size; }; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy @@ -42,7 +43,8 @@ struct cluster_topk_policy /*load_align_bytes=*/128, /*bits_per_pass=*/11, /*min_blocks_per_sm=*/1, - /*tie_break_items_per_thread=*/8}; + /*tie_break_items_per_thread=*/8, + /*single_cta_max_segment_size=*/8 * 1024}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool From 488f0a38e0b32b0ec1250224445c27012be11fe5 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 26 Jun 2026 03:52:10 +0200 Subject: [PATCH 047/126] Dynamically shrink clusters to single CTA --- cub/cub/agent/agent_batched_topk_cluster.cuh | 77 +++++-- cub/cub/detail/segmented_params.cuh | 4 +- .../dispatch_batched_topk_cluster.cuh | 8 + .../catch2_test_device_segmented_topk_keys.cu | 199 ++++++++++++------ 4 files changed, 212 insertions(+), 76 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f27ba1b8998..5e95a5a4522 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -59,6 +59,7 @@ #include #include #include +#include #include #include #include @@ -134,6 +135,17 @@ struct smem_block_tile_layout } }; +// True for `cuda::args` segment-size arguments whose exact per-segment value the host cannot know at dispatch time, so +// the launch is sized for a static upper bound: the `deferred` forms. Combined with `!is_single_value` (any +// per-segment sequence) this gates the runtime effective-single-CTA path; host-exact `immediate`/`constant` singles, +// which the dispatch already sizes exactly, are excluded. +template +inline constexpr bool __is_deferred_arg_v = false; +template +inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred<_Arg, _Bounds>> = true; +template +inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, _Bounds>> = true; + // ----------------------------------------------------------------------------- // Cluster top-k agent // ----------------------------------------------------------------------------- @@ -152,6 +164,7 @@ template ::is_single_value || __is_deferred_arg_v; + static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int histogram_items_per_thread = HistogramItemsPerThread; static constexpr int load_align_bytes = LoadAlignBytes; @@ -1078,10 +1101,11 @@ private: // Stripped on sub-SM-9.0 device passes; uses `cluster_group`, which is only // declared when `_CG_HAS_CLUSTER_GROUP` is set. #if defined(_CG_HAS_CLUSTER_GROUP) - // Synchronize the segment's cluster. A width-1 "cluster" is a single CTA whose state is entirely block-local, so a - // cluster-scoped barrier is unnecessary and `__syncthreads()` provides sufficient ordering. `single_cta` is derived - // from `cluster.num_blocks()` and is therefore uniform across the cluster, so the branch is reached uniformly and - // barrier reachability is preserved. + // Synchronize the segment's cluster. An effective width-1 "cluster" is a single CTA whose state is entirely + // block-local, so a cluster-scoped barrier is unnecessary and `__syncthreads()` provides sufficient ordering. + // `single_cta` derives from the effective cluster width chosen in `process_impl` (which may collapse a multi-CTA + // launch onto rank 0); it is per-segment uniform across the surviving block(s), so the branch is reached uniformly + // and barrier reachability is preserved. _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(::cooperative_groups::cluster_group& cluster, bool single_cta) { @@ -1100,6 +1124,7 @@ private: run(::cooperative_groups::cluster_group& cluster, num_segments_val_t segment_id, unsigned int cluster_rank, + unsigned int cluster_size, segment_size_val_t segment_size, out_offset_t k) { @@ -1110,11 +1135,12 @@ private: constexpr int total_bits = int{sizeof(key_t)} * 8; constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); - auto block_keys_in = d_key_segments_it[segment_id]; - const auto segment_size_u32 = static_cast(segment_size); - const unsigned int cluster_size = cluster.num_blocks(); - // A width-1 cluster is a lone CTA: route barriers to `__syncthreads()`, keep `state`/atomics block-local, and use - // CTA-scope histogram atomics (no cross-rank DSMEM folds exist to be mutually atomic with). + auto block_keys_in = d_key_segments_it[segment_id]; + const auto segment_size_u32 = static_cast(segment_size); + // `cluster_size`/`cluster_rank` are the *effective* width/rank chosen by `process_impl` (see the + // effective-single-CTA path), which may be smaller than the launched `cluster.num_blocks()`. A width-1 cluster is a + // lone CTA: route barriers to `__syncthreads()`, keep `state`/atomics block-local, and use CTA-scope histogram + // atomics (no cross-rank DSMEM folds exist to be mutually atomic with). const bool single_cta = (cluster_size == 1u); // Leader rank. The leader owns the cluster-merged histogram and the shared `state`. The deterministic tie-break @@ -2162,7 +2188,6 @@ private: // Runtime cluster width matches the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). const unsigned int cluster_blocks = cluster.num_blocks(); - const bool single_cta = (cluster_blocks == 1u); const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) @@ -2194,6 +2219,32 @@ private: return; } + // Effective cluster width/rank. For a per-segment (deferred) size argument the launch is sized for the maximum + // segment, so an individual small segment can be served by a single CTA: when it fits resident in one CTA and is at + // or below the single-CTA tuning threshold, rank 0 processes it alone via the cluster-barrier-free path and the + // cluster's other CTAs exit immediately, freeing their SM slots for other clusters/waves. The decision is uniform + // across the block (its inputs are per-segment uniform), so a redundant CTA returns whole. Compiled out for + // host-exact sizes, which the dispatch already sized to an exact cluster width. + unsigned int eff_cluster_blocks = cluster_blocks; + unsigned int eff_cluster_rank = cluster_rank; + if constexpr (enable_runtime_single_cta) + { + const bool fits_single_cta = + static_cast<::cuda::std::uint64_t>(segment_size) <= static_cast<::cuda::std::uint64_t>(block_tile_capacity) + && static_cast<::cuda::std::uint64_t>(segment_size) + <= static_cast<::cuda::std::uint64_t>(single_cta_max_segment_size); + if (fits_single_cta) + { + if (cluster_rank != 0u) + { + return; + } + eff_cluster_blocks = 1u; + eff_cluster_rank = 0u; + } + } + const bool single_cta = (eff_cluster_blocks == 1u); + // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster // state through `leader_state`), but mirroring the writes everywhere @@ -2212,9 +2263,11 @@ private: cluster_or_block_sync(cluster, single_cta); [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( - select_directions, segment_id, [this, &cluster, segment_id, cluster_rank, segment_size, k](auto direction_tag) { + select_directions, + segment_id, + [this, &cluster, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; - this->template run(cluster, segment_id, cluster_rank, segment_size, k); + this->template run(cluster, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k); }); _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); } diff --git a/cub/cub/detail/segmented_params.cuh b/cub/cub/detail/segmented_params.cuh index 0ece8092d0c..18be7f651c7 100644 --- a/cub/cub/detail/segmented_params.cuh +++ b/cub/cub/detail/segmented_params.cuh @@ -64,7 +64,9 @@ template [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(const ::cuda::args::deferred<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept { - return ::cuda::args::__unwrap(__arg); + // A single `deferred` wraps a device-accessible 1-element handle (pointer / iterator / 1-element span), not a value; + // the value is read on the device via element 0 (mirroring the `deferred_sequence` per-index read below). + return ::cuda::args::__unwrap(__arg)[0]; } template diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 02ce6f4054d..c23a1377ccf 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -85,6 +85,7 @@ template ; using key_t = it_value_t; @@ -340,6 +346,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( LoadAlignBytes, BitsPerPass, TieBreakItemsPerThread, + SingleCtaMaxSegmentSize, Determinism, TieBreak, KeyInputItItT, @@ -415,6 +422,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( LoadAlignBytes, BitsPerPass, TieBreakItemsPerThread, + SingleCtaMaxSegmentSize, Determinism, TieBreak, KeyInputItItT, diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 4df5ac73e27..5bee83560c5 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -155,6 +155,38 @@ using key_types = using select_direction_list = c2h::enum_type_list; +// Segment-size argument form for the single-value (uniform) fixed-size tests. The host-value forms (`immediate` and an +// un-annotated raw scalar, which auto-wraps as an `immediate` with loose bounds) are distributed across the TEST_TYPES +// variants so each variant compiles only one extra dispatch instantiation (balancing build time). The un-annotated form +// reports a `numeric_limits::max()` upper bound, so the types_2 variant also exercises the loose-bound +// oversize/streaming fallback (and its int-narrowing overflow guard) on the cluster backend. The remaining public forms +// get their own dedicated tests below (`constant` needs a compile-time size; `deferred` needs a device-accessible +// handle), and `deferred_sequence` is covered by the variable-size tests. +enum class seg_size_arg +{ + immediate_form, + unannotated_form, +}; + +template +auto make_segment_size_arg(SizeT segment_size) +{ + if constexpr (Form == seg_size_arg::immediate_form) + { + return cuda::args::immediate{segment_size, cuda::args::bounds()}; + } + else + { + return segment_size; // un-annotated: auto-wrapped as an immediate with loose bounds + } +} + +#if TEST_TYPES == 2 +inline constexpr seg_size_arg fixed_seg_size_arg = seg_size_arg::unannotated_form; +#else +inline constexpr seg_size_arg fixed_seg_size_arg = seg_size_arg::immediate_form; +#endif + // A (determinism, tie-break) requirement pair, used as a single compile-time test axis. Not a full cross product: a // tie-break preference is only meaningful with a deterministic requirement and is `gpu_to_gpu` by definition. template @@ -242,11 +274,12 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", // Copy input for verification c2h::device_vector expected_keys(keys_in_buffer); - // Run the top-k algorithm + // Run the top-k algorithm. The segment-size argument form (immediate / un-annotated) is selected per TEST_TYPES so + // the suite covers each form without any single variant compiling all of them. batched_topk_keys( d_keys_in, d_keys_out, - cuda::args::immediate{segment_size, cuda::args::bounds()}, + make_segment_size_arg(segment_size), cuda::args::immediate{k, cuda::args::bounds()}, cuda::args::immediate{num_segments}, cuda::args::immediate{num_segments * segment_size}); @@ -260,6 +293,107 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", REQUIRE(expected_keys == keys_out_buffer); } +#if TEST_TYPES == 0 +// `constant` is the one public arg form whose segment size must be a compile-time constant expression, so it gets its +// own fixed-size test. Gated to a single TEST_TYPES variant to balance build time across the suite. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with a compile-time-constant segment size", + "[keys][segmented][topk][device]", + key_types, + select_direction_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + using key_t = c2h::get<0, TestType>; + + constexpr auto direction = c2h::get<1, TestType>::value; + + constexpr segment_size_t segment_size = 384; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, segment_size_t{17}, segment_size})); + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::constant{}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_TYPES == 0 + +#if TEST_TYPES == 2 +// `deferred` wraps a device-accessible handle whose value is read in stream order on the device, so unlike the +// host-value forms it cannot be produced from a plain scalar; the uniform segment size lives in a 1-element device +// buffer. Gated to a single TEST_TYPES variant to balance build time across the suite. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with a deferred (device-resident) segment size", + "[keys][segmented][topk][device]", + key_types, + select_direction_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + using key_t = c2h::get<0, TestType>; + + constexpr auto direction = c2h::get<1, TestType>::value; + + constexpr segment_size_t segment_size = 384; + constexpr segment_size_t max_segment_size = 512; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, segment_size_t{17}, segment_size})); + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + // The uniform segment size is read by the device through a device-resident pointer; the static upper bound drives the + // host-side launch sizing. + c2h::device_vector d_segment_size(1, segment_size); + auto d_segment_size_ptr = thrust::raw_pointer_cast(d_segment_size.data()); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::deferred{d_segment_size_ptr, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_TYPES == 2 + #if TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", "[keys][segmented][topk][device][cluster][determinism]", @@ -615,67 +749,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign REQUIRE(expected_keys == keys_out_buffer); } -// An un-annotated argument is auto-wrapped as an immediate with *loose* bounds, so its reported highest segment size is -// `numeric_limits::max()`. No fully-resident cluster config can cover that, so the dispatch must take the -// oversize/streaming fallback (max cluster + max SMEM) without the analytic cluster-size math overflowing while -// narrowing the derived sizes to `int`. The real (uniform) segment size stays small, so the streamed result must still -// be exact. Exercises both the auto-wrapping of un-annotated args and the loose-bound fallback path. -C2H_TEST("DeviceBatchedTopK::MaxKeys work with an un-annotated (loose-bound) segment size", - "[keys][segmented][topk][device][cluster]") -{ - using key_t = float; - using segment_size_t = cuda::std::int64_t; - using segment_index_t = cuda::std::int64_t; - - // The oversize/streaming fallback for a loose bound is a cluster-dispatch concern; the baseline backend selects its - // config differently, so only assert it where the new analytic selector runs. - if constexpr (selected_backend != topk_backend::cluster) - { - SKIP("The oversize/streaming fallback for loose segment-size bounds is exercised by the cluster backend"); - } - - constexpr auto direction = cub::detail::topk::select::max; - constexpr segment_size_t static_max_k = 1024; - - // Real, uniform segment size kept small; it is passed un-annotated below so the dispatch instead sees a loose - // `numeric_limits::max()` upper bound. - const segment_size_t segment_size = - GENERATE_COPY(values({segment_size_t{1}, segment_size_t{4096}, segment_size_t{8192}})); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); - const segment_index_t num_segments = - GENERATE_COPY(values({segment_index_t{1}, segment_index_t{8}, segment_index_t{64}})); - - CAPTURE(segment_size, k, num_segments, direction); - - c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); - c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); - c2h::gen(C2H_SEED(1), keys_in_buffer); - auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); - auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); - auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); - auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); - - c2h::device_vector expected_keys(keys_in_buffer); - - // `segment_size` is passed un-annotated (a raw scalar): the args framework auto-wraps it as an immediate whose - // static upper bound is the loose `numeric_limits::max()`. - batched_topk_keys( - d_keys_in, - d_keys_out, - segment_size, - cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); - - fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); - compact_sorted_keys_to_topk(expected_keys, segment_size, k); - - // Top-k output is unordered; sort each output segment before comparison. - fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); - - REQUIRE(expected_keys == keys_out_buffer); -} #endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per-segment k", From c8f3e0547083927fff59cb9108def4fe399ac53d Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 26 Jun 2026 05:44:19 +0200 Subject: [PATCH 048/126] Disregard CTAs with less chunks than threshold Threshold is a tuning paramter and 1 for now. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 162 ++++++++++++------ .../dispatch_batched_topk_cluster.cuh | 41 +++-- .../tuning/tuning_batched_topk_cluster.cuh | 9 +- .../catch2_test_device_segmented_topk_keys.cu | 14 +- ...catch2_test_device_segmented_topk_pairs.cu | 119 +++++++++++++ 5 files changed, 276 insertions(+), 69 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5e95a5a4522..2362e738730 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -60,6 +60,7 @@ #include #include #include +#include #include #include #include @@ -146,10 +147,37 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred<_Arg, _Bounds>> template inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, _Bounds>> = true; +// ----------------------------------------------------------------------------- +// Occupancy-free cluster-blocks arithmetic (shared by host dispatch and device agent) +// ----------------------------------------------------------------------------- +// Whether a segment takes the barrier-free single-CTA path: resident in one CTA (`<= block_tile_capacity`) and at/below +// the single-CTA tuning threshold. Occupancy- and head-alignment-independent, so the host fast path and the device +// collapse decision agree exactly. 64-bit math for wide segment-size types / loose bounds. +[[nodiscard]] _CCCL_HOST_DEVICE constexpr bool single_cta_eligible( + ::cuda::std::uint64_t segment_size, + ::cuda::std::uint64_t block_tile_capacity, + int single_cta_max_segment_size) noexcept +{ + return segment_size <= block_tile_capacity + && segment_size <= static_cast<::cuda::std::uint64_t>(single_cta_max_segment_size); +} + +// Effective cluster blocks implied by a chunk count: a CTA joins the effective cluster iff it would own at least +// `min_chunks_per_cta` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap +// differs (live `cluster.num_blocks()` on the device, max launchable blocks on the host). `min_chunks_per_cta` is +// `static_assert`ed positive, so the divide is well-defined. 64-bit math. +[[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( + ::cuda::std::uint64_t chunks, int min_chunks_per_cta, unsigned int cluster_blocks_cap) noexcept +{ + const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_cta); + return static_cast( + (::cuda::std::clamp) (blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); +} + // ----------------------------------------------------------------------------- // Cluster top-k agent // ----------------------------------------------------------------------------- -// Cluster width is a runtime value (see `process_impl` for the readback), so +// Cluster blocks is a runtime value (see `process_impl` for the readback), so // it is not a template parameter; per-block block_tile layout is still controlled // by the template parameters below. // @@ -165,6 +193,7 @@ template = 1, "min_chunks_per_cta must be positive"); + // Enable the per-segment effective-single-CTA runtime path only when the host could not size the launch to each // segment's exact size: any per-segment sequence (`!is_single_value`) or a `deferred` single value. For host-exact - // `immediate`/`constant` singles the dispatch already picks the matching cluster size, so the logic is compiled out. + // `immediate`/`constant` singles the dispatch already picks the matching cluster blocks, so the logic is compiled + // out. static constexpr bool enable_runtime_single_cta = !::cuda::args::__traits::is_single_value || __is_deferred_arg_v; @@ -433,7 +469,7 @@ struct agent_batched_topk_cluster struct chunk_partition { offset_t first; // global index of this rank's first owned chunk - offset_t stride; // distance between consecutive owned chunks (`cluster_size` strided, `1` blocked) + offset_t stride; // distance between consecutive owned chunks (`cluster_blocks` strided, `1` blocked) offset_t count; // number of chunks owned by this rank [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t global_index(offset_t local) const @@ -447,7 +483,7 @@ struct agent_batched_topk_cluster // split, and the streamer ping-pong untouched, because all of those depend only on the global chunk index, not on // which rank owns it. // - // * Strided (default): chunk `i` goes to rank `i % cluster_size`, so each CTA walks `first, first+S, first+2S, + // * Strided (default): chunk `i` goes to rank `i % cluster_blocks`, so each CTA walks `first, first+S, first+2S, // ...`. // * Blocked (deterministic path): each CTA owns a contiguous run of `ceil_div(chunks, S)` chunks (the last // non-empty rank gets the short remainder). @@ -455,11 +491,11 @@ struct agent_batched_topk_cluster // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches // ascending contiguous global-index ranges), so it is selected exactly when `deterministic` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition - make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_size) const + make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { if constexpr (deterministic) { - const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_size)); + const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_blocks)); const offset_t first = static_cast(cluster_rank) * chunks_per_cta; const offset_t count = (first < chunks) ? (::cuda::std::min) (chunks_per_cta, chunks - first) : offset_t{0}; return {first, offset_t{1}, count}; @@ -467,8 +503,8 @@ struct agent_batched_topk_cluster else { return {static_cast(cluster_rank), - static_cast(cluster_size), - num_rank_chunks(chunks, cluster_rank, cluster_size)}; + static_cast(cluster_blocks), + num_rank_chunks(chunks, cluster_rank, cluster_blocks)}; } } @@ -1101,11 +1137,10 @@ private: // Stripped on sub-SM-9.0 device passes; uses `cluster_group`, which is only // declared when `_CG_HAS_CLUSTER_GROUP` is set. #if defined(_CG_HAS_CLUSTER_GROUP) - // Synchronize the segment's cluster. An effective width-1 "cluster" is a single CTA whose state is entirely - // block-local, so a cluster-scoped barrier is unnecessary and `__syncthreads()` provides sufficient ordering. - // `single_cta` derives from the effective cluster width chosen in `process_impl` (which may collapse a multi-CTA - // launch onto rank 0); it is per-segment uniform across the surviving block(s), so the branch is reached uniformly - // and barrier reachability is preserved. + // Synchronize the segment's cluster. A single-CTA "cluster" keeps all state block-local, so `__syncthreads()` orders + // it and the cluster-scoped barrier is unnecessary. `single_cta` is computed in `run()` from the collapsed cluster + // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw `cluster.num_blocks()`. + // It is per-segment uniform across the surviving block(s), so the branch is reached uniformly. _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(::cooperative_groups::cluster_group& cluster, bool single_cta) { @@ -1124,7 +1159,7 @@ private: run(::cooperative_groups::cluster_group& cluster, num_segments_val_t segment_id, unsigned int cluster_rank, - unsigned int cluster_size, + unsigned int cluster_blocks, segment_size_val_t segment_size, out_offset_t k) { @@ -1137,22 +1172,11 @@ private: auto block_keys_in = d_key_segments_it[segment_id]; const auto segment_size_u32 = static_cast(segment_size); - // `cluster_size`/`cluster_rank` are the *effective* width/rank chosen by `process_impl` (see the - // effective-single-CTA path), which may be smaller than the launched `cluster.num_blocks()`. A width-1 cluster is a - // lone CTA: route barriers to `__syncthreads()`, keep `state`/atomics block-local, and use CTA-scope histogram - // atomics (no cross-rank DSMEM folds exist to be mutually atomic with). - const bool single_cta = (cluster_size == 1u); - - // Leader rank. The leader owns the cluster-merged histogram and the shared `state`. The deterministic tie-break - // makes the leader the *last* CTA in scan order so it never needs its own (merged-away) local candidate count: - // prefer-smallest scans ascending by rank (leader = last rank), prefer-largest scans descending (leader = rank 0). - // The nondeterministic path keeps rank 0 (unchanged codegen). - const unsigned int leader_rank = (deterministic && !tie_reversed) ? (cluster_size - 1u) : 0u; - - // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` - // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = - single_cta ? &temp_storage.state : cluster.map_shared_rank(&temp_storage.state, leader_rank); + // `cluster_blocks` is what `process_impl` runs at: the launched `cluster.num_blocks()`, or `1` when it collapsed a + // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, + // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, + // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. + const bool single_cta = (cluster_blocks == 1u); // Per-block local copy of `kth_key_bits` so each key check hits the // block's own SMEM rather than DSMEM during the histogram loop. Built up one splitter digit per pass from the @@ -1179,10 +1203,40 @@ private: // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only // depends on the multiset of keys covered by the cluster. - const offset_t chunks = num_chunks(segment_size_u32, head_items); - const chunk_partition part = make_chunk_partition(chunks, cluster_rank, cluster_size); + const offset_t chunks = num_chunks(segment_size_u32, head_items); + + // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_cta` each), <= the + // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- + // but stay resident and still arrive at every `cluster.sync()` (a returned CTA would hang the barrier; see the + // TODOs at the barrier sites). Derived from this CTA's head-aligned `chunks` so it matches the partition exactly. + // Stays at `cluster_blocks` for host-exact sizes (the dispatch already matched it) and on the single-CTA path. + unsigned int eff_cluster_blocks = cluster_blocks; + if constexpr (enable_runtime_single_cta) + { + if (!single_cta) + { + eff_cluster_blocks = effective_cluster_blocks_from_chunks( + static_cast<::cuda::std::uint64_t>(chunks), min_chunks_per_cta, cluster_blocks); + } + } + const bool is_idle_rank = cluster_rank >= eff_cluster_blocks; + + // Idle ranks own no chunks; `make_chunk_partition` assumes `rank < size`, so hand them an explicit empty partition. + const chunk_partition part = is_idle_rank ? chunk_partition{offset_t{0}, offset_t{1}, offset_t{0}} + : make_chunk_partition(chunks, cluster_rank, eff_cluster_blocks); const offset_t my_chunks = part.count; + // Leader rank. The leader owns the cluster-merged histogram and the shared `state`, and is always a working rank + // (`< eff_cluster_blocks`). The deterministic tie-break makes the leader the *last* CTA in scan order so it never + // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last + // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. + const unsigned int leader_rank = (deterministic && !tie_reversed) ? (eff_cluster_blocks - 1u) : 0u; + + // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` + // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). + state_t* leader_state = + single_cta ? &temp_storage.state : cluster.map_shared_rank(&temp_storage.state, leader_rank); + // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every // `cluster.sync()` are reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) // keeps them all resident and streams nothing; an overflowing CTA reserves a round-robin streaming region at the @@ -1529,8 +1583,10 @@ private: // (raw counts in the reduce-then-scan path, block-local inclusive // scans in the scan-then-reduce path) into the leader's `hist` // via cluster-scope DSMEM atomics (see `hist_fold_remote`). The - // leader skips this to avoid double-counting its own contribution. - if (cluster_rank != leader_rank) + // leader skips this to avoid double-counting its own contribution; + // idle ranks (`>= eff_cluster_blocks`) have an all-zero histogram, so + // they skip the fold entirely (the loop would only read zeros). + if (cluster_rank != leader_rank && !is_idle_rank) { const ::cuda::std::uint32_t hist_smem32 = hist_base32(); for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) @@ -1543,6 +1599,8 @@ private: } } + // TODO(cccl): idle ranks arrive here only because `cluster.sync()` spans the whole launched cluster. An mbarrier + // over just the active ranks would let them exit and free their SM slots instead of spinning on this barrier. cluster_or_block_sync(cluster, single_cta); // Step 3: the leader walks the merged `hist` (raw counts in the @@ -1563,12 +1621,14 @@ private: } reset_hist(); } + // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. cluster_or_block_sync(cluster, single_cta); - // End-of-pass splitter fold. Every thread reads the leader's just-published `kth_bucket` directly through DSMEM - // (a single `u32`, ordered by the `cluster.sync()` above) and folds it into its own `kth_key_bits_local`, so the - // full splitter key is reconstructed locally without any block-wide broadcast or `__syncthreads`. Runs uniformly - // for every pass (including the last), leaving `kth_key_bits_local` complete for the final filter. + // End-of-pass splitter fold. Every working thread reads the leader's just-published `kth_bucket` through DSMEM (a + // single `u32`, ordered by the `cluster.sync()` above) and folds it into its own `kth_key_bits_local`, so the + // full splitter key is reconstructed locally without a broadcast. Idle ranks produce no output, so they skip the + // read (trimming the leader's fan-out); the early-stop check below stays uniform so every block breaks together. + if (!is_idle_rank) { const int bucket = static_cast(leader_state->kth_bucket); detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); @@ -1635,6 +1695,8 @@ private: { temp_storage.cand_prefix = 0; } + // TODO(cccl): idle ranks (`local_count == 0`) arrive only to keep this barrier reachable; a sub-cluster + // mbarrier over the working ranks would let them exit (see the pass loop). cluster_or_block_sync(cluster, single_cta); if (threadIdx.x == 0 && local_count != offset_t{0}) { @@ -1647,7 +1709,8 @@ private: } else { - for (unsigned int r = cluster_rank + 1u; r < cluster_size; ++r) // ascending scan order: higher ranks follow + // Ascending scan order: higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing. + for (unsigned int r = cluster_rank + 1u; r < eff_cluster_blocks; ++r) { add_remote_prefix(r, local_count); } @@ -2185,7 +2248,7 @@ private: { ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); const unsigned int cluster_rank = cluster.block_rank(); - // Runtime cluster width matches the launch attribute the dispatch passed + // Runtime cluster blocks match the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). const unsigned int cluster_blocks = cluster.num_blocks(); const auto segment_id = static_cast(blockIdx.x / cluster_blocks); @@ -2219,20 +2282,19 @@ private: return; } - // Effective cluster width/rank. For a per-segment (deferred) size argument the launch is sized for the maximum - // segment, so an individual small segment can be served by a single CTA: when it fits resident in one CTA and is at - // or below the single-CTA tuning threshold, rank 0 processes it alone via the cluster-barrier-free path and the - // cluster's other CTAs exit immediately, freeing their SM slots for other clusters/waves. The decision is uniform - // across the block (its inputs are per-segment uniform), so a redundant CTA returns whole. Compiled out for - // host-exact sizes, which the dispatch already sized to an exact cluster width. + // Effective cluster blocks/rank. For a per-segment (deferred) size argument the launch is sized for the maximum + // segment, so a small segment that fits resident in one CTA and is at/below the single-CTA tuning threshold is + // served by rank 0 alone via the barrier-free path; the cluster's other CTAs exit immediately, freeing their SM + // slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. Compiled out for + // host-exact sizes, which the dispatch already sized to exact cluster blocks. unsigned int eff_cluster_blocks = cluster_blocks; unsigned int eff_cluster_rank = cluster_rank; if constexpr (enable_runtime_single_cta) { - const bool fits_single_cta = - static_cast<::cuda::std::uint64_t>(segment_size) <= static_cast<::cuda::std::uint64_t>(block_tile_capacity) - && static_cast<::cuda::std::uint64_t>(segment_size) - <= static_cast<::cuda::std::uint64_t>(single_cta_max_segment_size); + const bool fits_single_cta = single_cta_eligible( + static_cast<::cuda::std::uint64_t>(segment_size), + static_cast<::cuda::std::uint64_t>(block_tile_capacity), + single_cta_max_segment_size); if (fits_single_cta) { if (cluster_rank != 0u) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index c23a1377ccf..258e188f01f 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -12,7 +12,7 @@ //! //! Two kernels share the agent body: //! * Host: no `__cluster_dims__`, launched via `cudaLaunchKernelExC` with the -//! cluster width chosen at runtime (up to 16 on Hopper). +//! cluster blocks chosen at runtime (up to 16 on Hopper). //! * CDP: static `__cluster_dims__(max_portable_cluster_blocks, 1, 1)`, //! gated on `CUB_RDC_ENABLED` so CDP-disabled builds skip emitting it //! (same pattern as `dispatch_segmented_sort.cuh`). @@ -86,6 +86,7 @@ template ; using key_t = it_value_t; @@ -347,6 +353,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( BitsPerPass, TieBreakItemsPerThread, SingleCtaMaxSegmentSize, + MinChunksPerCta, Determinism, TieBreak, KeyInputItItT, @@ -423,6 +430,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( BitsPerPass, TieBreakItemsPerThread, SingleCtaMaxSegmentSize, + MinChunksPerCta, Determinism, TieBreak, KeyInputItItT, @@ -441,7 +449,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // performs the cluster launch via `cudaLaunchKernelEx`. The factory also wraps the pre-launch driver queries. detail::TripleChevronFactory launcher_factory{}; - // Opt in to non-portable cluster widths (>8 on Hopper). + // Opt in to non-portable cluster blocks (>8 on Hopper). if (const auto error = launcher_factory.set_non_portable_cluster_allowed(dynamic_kernel)) { return error; @@ -516,8 +524,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( return cudaSuccess; }; - // Wave-aware cluster-size selection. The free variable is the cluster size `C` (one cluster per segment); each - // `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more + // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); + // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We // pick the `C` that minimizes the number of waves, breaking ties toward the largest `C` (= smallest SMEM = most // L1), which matches the profiled fast configs. We enumerate `C` analytically rather than discovering SMEM tiers @@ -545,13 +553,18 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( const int c_full = static_cast( (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), ::cuda::ceil_div(seg, chunk_items_u64))); - const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); - const bool prefer_single_cta = seg <= static_cast<::cuda::std::uint64_t>(policy.single_cta_max_segment_size); + const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); + // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than + // necessary). At `min_chunks_per_cta == 1` this equals `c_full`; a larger knob shrinks it. + const int desired_cluster_blocks = static_cast(effective_cluster_blocks_from_chunks( + ::cuda::ceil_div(seg, chunk_items_u64), + MinChunksPerCta, + static_cast(max_supported_cluster_blocks))); int cluster_blocks = 0; int dynamic_smem_sel = 0; - if (c_lo == 1 && prefer_single_cta) + if (single_cta_eligible(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity), SingleCtaMaxSegmentSize)) { // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's // cluster-barrier-free path beats spreading it across more CTAs. `S_res(seg)` is within budget and one CTA is @@ -564,9 +577,11 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( { // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. - // Scan `C` in `[max(C_lo, 2), C_full]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. - const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); - const int c_end = (::cuda::std::max) (c_begin, c_full); + // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. The + // upper bound is capped at the cluster blocks the max segment needs (`desired_cluster_blocks`), so the host + // never launches a wider cluster than necessary; at `min_chunks_per_cta == 1` the cap equals `c_full`. + const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); + const int c_end = (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks)); ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); for (int c = c_begin; c <= c_end; ++c) { @@ -583,7 +598,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( } // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must - // be a multiple of the cluster size. The returned value is the device-wide clusters-per-wave (capacity), + // be a multiple of the cluster blocks. The returned value is the device-wide clusters-per-wave (capacity), // independent of grid size, and accounts for the static footprint and register pressure internally. cluster_attr.val.clusterDim.x = static_cast(c); cfg.gridDim = dim3(static_cast(c), 1, 1); @@ -595,7 +610,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( } if (clusters_per_wave <= 0) { - continue; // cluster size not launchable at this SMEM. + continue; // cluster blocks not launchable at this SMEM. } const auto waves = ::cuda::ceil_div( @@ -685,7 +700,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( }), ({ // CDP path: device-side launches cannot opt in to more than portable - // total SMEM or non-portable cluster widths. Segments that exceed the + // total SMEM or non-portable cluster blocks. Segments that exceed the // portable resident coverage are still handled: the agent re-streams the // overflow chunks from gmem. constexpr int portable_total_smem_bytes = 48 * 1024; diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 3e2bc339b4f..d3a3ef6c00a 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -18,7 +18,7 @@ CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster { -// Per-block execution shape. The dispatch picks the cluster size and the dynamic shared-memory block_tile capacity at +// Per-block execution shape. The dispatch picks the cluster blocks and the dynamic shared-memory block_tile capacity at // runtime (occupancy/wave-aware), so the policy carries only the per-block tuning knobs. struct cluster_topk_policy { @@ -31,6 +31,10 @@ struct cluster_topk_policy int min_blocks_per_sm; int tie_break_items_per_thread; int single_cta_max_segment_size; + // Minimum chunks a CTA must own to belong to a segment's effective cluster. At 1 the effective cluster is just the + // CTAs that receive any chunk; a larger value trades parallelism for fewer, busier CTAs (and shrinks the + // host-launched cluster). Used by both the host cluster-blocks cap and the device per-segment effective cluster. + int min_chunks_per_cta; }; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy @@ -44,7 +48,8 @@ struct cluster_topk_policy /*bits_per_pass=*/11, /*min_blocks_per_sm=*/1, /*tie_break_items_per_thread=*/8, - /*single_cta_max_segment_size=*/8 * 1024}; + /*single_cta_max_segment_size=*/8 * 1024, + /*min_chunks_per_cta=*/1}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 5bee83560c5..85d1d07a2e6 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -682,9 +682,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign SKIP("Determinism requirements are only provided by the cluster backend"); } - // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch - // mixes streaming segments (the 1 Mi-element ones, one with an unaligned `- 31` overflow tail) with fully-resident - // segments (96 Ki + 17 and 257 elements). + // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch sizes a + // wide (16-CTA) cluster and mixes every effective-width regime: streaming segments (the 1 Mi-element ones, one with + // an unaligned `- 31` overflow tail), a fully-resident multi-CTA segment (96 Ki + 17), two *medium* segments that + // exceed the single-CTA threshold but need only a few chunks (so surplus cluster CTAs go idle -- the runtime + // effective-cluster-width path), and a tiny segment that collapses onto a single CTA (257 elements). constexpr segment_size_t static_max_segment_size = 1100 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; @@ -693,12 +695,16 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); constexpr segment_size_t big_segment_size = 1024 * 1024; + constexpr segment_size_t med_segment_a = 12 * 1024 + 1; // a few chunks, unaligned tail -> most cluster CTAs idle + constexpr segment_size_t med_segment_b = 40 * 1024; // more chunks, still well below the 16-CTA launch width c2h::host_vector h_segment_offsets{ 0, big_segment_size, big_segment_size + (big_segment_size - 31), big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17), - big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257}; + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257, + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257 + med_segment_a, + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257 + med_segment_a + med_segment_b}; c2h::device_vector segment_offsets = h_segment_offsets; const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); const segment_size_t num_items = h_segment_offsets.back(); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 8dc576fa26d..a75ac8c4258 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -794,6 +794,125 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break } REQUIRE(h_values_a == h_values_b); } + +// Determinism/tie-break combinations for the mixed-width pairs test below. The selected key/value multiset is invariant +// to the tie-break preference, so every combo is verified the same way (pairs consistency + per-segment uniqueness + +// sorted-key equality); tie-break preferences only pair with a deterministic requirement. +template +struct det_tie_pair +{ + static constexpr auto determinism = Determinism; + static constexpr auto tie_break = TieBreak; +}; +using det_tie_pair_combos = + c2h::type_list, + det_tie_pair, + det_tie_pair, + det_tie_pair>; + +// Effective-cluster-width coverage for pairs: the large `deferred_sequence` bound sizes a wide (16-CTA) launch, but the +// actual segments are small, so a single launch mixes a fully-resident multi-CTA segment (96 Ki + 17, every rank +// works), two medium segments that leave surplus cluster CTAs idle (the runtime effective-width path, one unaligned), +// and a tiny single-CTA-collapsed segment (257). Confirms each value payload stays attached to its key through that +// path, for every determinism requirement. Streaming (overflowing) pairs are covered by the large-segment tests above, +// so this one stays small enough for `compute-sanitizer --tool racecheck`. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with mixed effective-width variable-size segments", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + + // The bound far exceeds the actual segments, so the launch picks a wide cluster while the segments stay small. + constexpr segment_size_t static_max_segment_size = 1100 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + constexpr auto direction = cub::detail::topk::select::max; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); + + constexpr segment_size_t full_segment_size = 96 * 1024 + 17; // enough chunks to fill the launched cluster (no idle) + constexpr segment_size_t med_segment_a = 12 * 1024 + 1; // a few chunks -> most cluster CTAs idle (unaligned) + constexpr segment_size_t med_segment_b = 40 * 1024; // more chunks, still below the launched width + c2h::host_vector h_segment_offsets{ + 0, + full_segment_size, + full_segment_size + med_segment_a, + full_segment_size + med_segment_a + med_segment_b, + full_segment_size + med_segment_a + med_segment_b + 257}; + c2h::device_vector segment_offsets = h_segment_offsets; + const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); + const segment_size_t num_items = h_segment_offsets.back(); + + auto segment_offsets_it = thrust::raw_pointer_cast(segment_offsets.data()); + auto segment_size_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), segment_size_op{segment_offsets_it}); + + CAPTURE(static_max_segment_size, static_max_k, k, num_segments, num_items); + + // Each output segment holds exactly min(k, segment_size[i]) items, tightly packed. + auto compacted_output_sizes_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), + get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k), num_segments}); + c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); + thrust::exclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + const segment_size_t total_output_size = compacted_offsets.back(); + + c2h::device_vector keys_in_buffer(num_items, thrust::no_init); + c2h::device_vector keys_out_buffer(total_output_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_offsets.cbegin()); + auto d_keys_out = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_out_ptr), compacted_offsets.cbegin()); + + // Values = global flattened index, so each selected value points back into the flattened input. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + c2h::device_vector values_out_buffer(total_output_size, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_in = + cuda::make_permutation_iterator(cuda::make_counting_iterator(values_in_it), segment_offsets.cbegin()); + auto d_values_out = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_values_out_ptr), compacted_offsets.cbegin()); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_items}); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, compacted_offsets, num_segments) == true); + + segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); + expected_keys = compact_to_topk_batched(expected_keys, segment_offsets, k); + segmented_sort_keys( + keys_out_buffer, num_segments, compacted_offsets.cbegin(), compacted_offsets.cbegin() + 1, direction); + REQUIRE(expected_keys == keys_out_buffer); +} #endif // TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segments", From 8acc49ce7be9858bb8c9d50b1e6b5156b52cb296 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 26 Jun 2026 21:52:08 +0200 Subject: [PATCH 049/126] Add fast path for k > segment size Also adds more testing specifically for small segment size types. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 82 ++++++ .../catch2_test_device_segmented_topk_keys.cu | 187 ++++++++++++ ...catch2_test_device_segmented_topk_pairs.cu | 275 ++++++++++++++++-- 3 files changed, 525 insertions(+), 19 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2362e738730..1cd82e9806c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -2244,6 +2244,79 @@ private: cluster_or_block_sync(cluster, single_cta); } + // Copies an entire segment `input[i] -> output[i]` (keys, plus values for pairs) for the select-all fast path + // (`k == segment_size`); top-k output order is unspecified, so a straight copy is valid. The whole launched cluster + // grid-strides the segment with no chunk partition or cluster barriers -- each CTA copies its part and exits. Each + // thread issues all `copy_items` loads of a step before its stores to keep many requests in flight (ILP/MLP). Input + // and output bases are mutually unaligned in general, so we use per-element (not vectorized/TMA) coalesced accesses. + _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all( + num_segments_val_t segment_id, + segment_size_val_t segment_size, + unsigned int cluster_rank, + unsigned int cluster_blocks) + { + constexpr int copy_items = 8; + const offset_t n = static_cast(segment_size); + const offset_t cluster_tid = cluster_rank * static_cast(blockDim.x) + threadIdx.x; + const offset_t cluster_tids = cluster_blocks * static_cast(blockDim.x); + const offset_t step = cluster_tids * static_cast(copy_items); + auto keys_in_it = d_key_segments_it[segment_id]; + auto keys_out_it = d_key_segments_out_it[segment_id]; + for (offset_t tile = 0; tile < n; tile += step) + { + // `idx[]` stays in `offset_t` (not the possibly-narrower `segment_size_val_t`) so the `idx[j] < n` bound check + // happens before any narrowing; the cast at each indexing site is then safe. + offset_t idx[copy_items]; + key_t keys[copy_items]; + [[maybe_unused]] value_t vals[copy_items]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + idx[j] = tile + static_cast(j) * cluster_tids + cluster_tid; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + if (idx[j] < n) + { + keys[j] = keys_in_it[static_cast(idx[j])]; + } + } + if constexpr (!is_keys_only) + { + auto vals_in_it = d_value_segments_it[segment_id]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + if (idx[j] < n) + { + vals[j] = vals_in_it[static_cast(idx[j])]; + } + } + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + if (idx[j] < n) + { + keys_out_it[static_cast(idx[j])] = keys[j]; + } + } + if constexpr (!is_keys_only) + { + auto vals_out_it = d_value_segments_out_it[segment_id]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + if (idx[j] < n) + { + vals_out_it[static_cast(idx[j])] = vals[j]; + } + } + } + } + } + _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() { ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); @@ -2282,6 +2355,15 @@ private: return; } + // Select-all fast path: when `k` reaches the full segment, every element wins, so we skip the radix passes, + // histogram, and output-ordering and just copy. Runs on the full launched cluster (before the effective-cluster + // collapse); the decision is per-segment uniform, so the branch is cluster-uniform. + if (static_cast(k) == segment_size) + { + copy_segment_select_all(segment_id, segment_size, cluster_rank, cluster_blocks); + return; + } + // Effective cluster blocks/rank. For a per-segment (deferred) size argument the launch is sized for the maximum // segment, so a small segment that fits resident in one CTA and is at/below the single-CTA tuning threshold is // served by rank 0 alone via the barrier-free path; the cluster's other CTAs exit immediately, freeing their SM diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 85d1d07a2e6..8fd8e77e501 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -338,6 +338,133 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with a compile-time-constant seg REQUIRE(expected_keys == keys_out_buffer); } + +// A `k` larger than the segment is clamped to `segment_size` (every element selected), routing through the select-all +// fast path. The output then holds exactly `segment_size` items per segment, so we verify against the full sorted +// segment. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys clamp k larger than the segment size", + "[keys][segmented][topk][device]", + key_types, + select_direction_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + using key_t = c2h::get<0, TestType>; + + constexpr auto direction = c2h::get<1, TestType>::value; + + constexpr segment_size_t segment_size = 384; + // Deliberately request more than the segment holds; the algorithm must clamp to `segment_size`. + const segment_size_t k_requested = + GENERATE_COPY(values({segment_size + 1, segment_size + 100, 2 * segment_size, 10 * segment_size})); + const segment_size_t effective_k = (cuda::std::min) (k_requested, segment_size); // == segment_size + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), segment_size, k_requested, effective_k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * effective_k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), effective_k); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::constant{}, + cuda::args::immediate{k_requested, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, effective_k); // clamped k == segment_size keeps everything + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, effective_k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Segment-size types narrower than the internal `offset_t`: a signed type narrows a too-large index to a negative one +// (indexing before the segment); an unsigned one wraps to a small in-range index (duplicate racing stores). +using narrow_seg_size_list = c2h::type_list; + +// Regression for a narrow segment-size type. A block launches 512 threads -- past an 8-bit type's range -- so any path +// that narrows an index to `segment_size_val_t` before its bound check would go out of bounds. We sweep `segment_size` +// and `k` to hit every path (`k == segment_size` -> select-all copy; `k < segment_size` -> radix / histogram / +// output-ordering) and run over `det_tie_combos` to also cover the deterministic scan/atomics under the narrow type. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle a segment-size type narrower than the internal offset", + "[keys][segmented][topk][device][cluster][determinism]", + narrow_seg_size_list, + det_tie_combos) +{ + using seg_size_t = c2h::get<0, TestType>; + using segment_index_t = cuda::std::int64_t; + using key_t = cuda::std::uint8_t; // key type is immaterial to this index-arithmetic regression + + using combo = c2h::get<1, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + + // Sizes fit both 8-bit types but sit far below the 512 threads a block launches; `127` probes the signed type's max. + const seg_size_t segment_size = static_cast(GENERATE(values({3, 100, 127}))); + const seg_size_t k = static_cast( + GENERATE_COPY(values({1, int{segment_size} / 2, int{segment_size}}), take(2, random(1, int{segment_size})))); + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), int{segment_size}, int{k}, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), int{segment_size}); + + // Output buffer with front/back canary padding: a signed-narrowing underrun lands in the front guard, an overrun in + // the back guard. The in-range comparison below misses both (the underrun also stays inside the allocation, hiding it + // from memcheck), so we assert the guards stay untouched. `guard` covers the worst-case 8-bit offset (<= 256). + constexpr segment_index_t guard = 256; + const key_t key_canary = static_cast(0x5A); + c2h::device_vector keys_out_storage(guard + num_segments * k + guard, key_canary); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_storage.data()) + guard; + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), int{k}); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + // Guards must be untouched by the algorithm. + const c2h::device_vector expected_guard(guard, key_canary); + CHECK(c2h::device_vector(keys_out_storage.begin(), keys_out_storage.begin() + guard) == expected_guard); + CHECK(c2h::device_vector(keys_out_storage.end() - guard, keys_out_storage.end()) == expected_guard); + + // Extract the in-range output for the standard top-k comparison. + c2h::device_vector keys_out_buffer(keys_out_storage.begin() + guard, keys_out_storage.end() - guard); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} #endif // TEST_TYPES == 0 #if TEST_TYPES == 2 @@ -463,6 +590,66 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned REQUIRE(expected_keys == keys_out_buffer); } +// Streaming counterpart of the narrow-segment-size regression. Streaming needs segments larger than the resident +// cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- same width as the +// (unsigned) internal `offset_t` but signed -- to exercise the streaming path's index arithmetic across det/non-det. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 32-bit segment-size type", + "[keys][segmented][topk][device][cluster][determinism]", + det_tie_combos) +{ + using key_t = float; + using seg_size_t = int; // signed 32-bit: same width as offset_t, but signed + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + + constexpr seg_size_t static_max_segment_size = 1024 * 1024; + constexpr seg_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 2; + + const int pad = GENERATE(0, 7); + const seg_size_t segment_size = static_max_segment_size - 31; // unaligned -> forces streaming + unaligned tail edge + const seg_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const seg_size_t k = GENERATE_COPY(values({seg_size_t{1}, max_k})); + + CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(pad + num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(num_segments * segment_size, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + template struct cast_to_key_op { diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index a75ac8c4258..de87a2d199d 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -173,6 +173,25 @@ using uint_key_types = c2h::type_list; +// Determinism/tie-break combinations used as a single compile-time axis by the determinism-aware pairs tests. The +// selected multiset is invariant to the tie-break preference, so every combo is verified the same way; tie-break +// preferences only pair with a deterministic requirement. +template +struct det_tie_pair +{ + static constexpr auto determinism = Determinism; + static constexpr auto tie_break = TieBreak; +}; +using det_tie_pair_combos = + c2h::type_list, + det_tie_pair, + det_tie_pair, + det_tie_pair>; + // Consistency check: ensures values remain associated with their corresponding keys template bool verify_pairs_consistency(const c2h::device_vector& keys_in, @@ -340,6 +359,169 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small fixed-size segments" REQUIRE(expected_keys == keys_out_buffer); } +#if TEST_TYPES == 0 +// A `k` larger than the segment is clamped to `segment_size` (every pair is selected), routing through the select-all +// fast path. The output then holds exactly `segment_size` pairs per segment -- identical to a `k == segment_size` +// request -- so we verify with a `segment_size`-wide output and a key/value-consistency check on the full segment. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs clamp k larger than the segment size", + "[pairs][segmented][topk][device]", + key_types, + select_direction_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + using key_t = c2h::get<0, TestType>; + using val_t = cuda::std::int32_t; + + constexpr auto direction = c2h::get<1, TestType>::value; + + constexpr segment_size_t segment_size = 384; + // Deliberately request more than the segment holds; the algorithm must clamp to `segment_size`. + const segment_size_t k_requested = + GENERATE_COPY(values({segment_size + 1, segment_size + 100, 2 * segment_size, 10 * segment_size})); + const segment_size_t effective_k = (cuda::std::min) (k_requested, segment_size); // == segment_size + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), segment_size, k_requested, effective_k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * effective_k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), effective_k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + c2h::device_vector values_out_buffer(num_segments * effective_k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), effective_k); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::constant{}, + cuda::args::immediate{k_requested, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, effective_k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, effective_k); // clamped k == segment_size keeps everything + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, effective_k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Segment-size types narrower than the internal `offset_t`, both signed and unsigned: a signed type that narrows a +// too-large index goes negative (indexing before the segment); an unsigned one wraps to a small in-range index +// (duplicate racing stores). +using narrow_seg_size_list = c2h::type_list; + +// Regression for a segment-size type narrower than the internal `offset_t`. The select-all copy must bound-check in +// `offset_t` *before* narrowing: a block launches 512 threads -- far past an 8-bit type's range -- so any path that +// narrows a thread/element index to `segment_size_val_t` before its bound check (or wraps an unsigned one) would +// read/write out of bounds (a signed type indexes before the segment; an unsigned one wraps into duplicate racing +// stores -- caught by racecheck even when the value is coincidentally right). We sweep `segment_size` and `k` with +// each narrow size type to exercise every path under it: `k == segment_size` hits the select-all copy, while +// `k < segment_size` drives the radix-select / cluster-histogram / output-ordering paths (and, for pairs, the value +// gather). Running over `det_tie_pair_combos` additionally covers the deterministic output-ordering scan/atomics. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs handle a segment-size type narrower than the internal offset", + "[pairs][segmented][topk][device][cluster][determinism]", + narrow_seg_size_list, + det_tie_pair_combos) +{ + using seg_size_t = c2h::get<0, TestType>; + using segment_index_t = cuda::std::int64_t; + using key_t = cuda::std::uint8_t; // key type is immaterial to this index-arithmetic regression + using val_t = cuda::std::int32_t; + + using combo = c2h::get<1, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the + // nondeterministic path, so skip them (the base nondeterministic combo still runs). + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + + // Sizes fit both 8-bit types but sit far below the 512 threads a block launches; `127` probes the signed type's max. + const seg_size_t segment_size = static_cast(GENERATE(values({3, 100, 127}))); + const seg_size_t k = static_cast( + GENERATE_COPY(values({1, int{segment_size} / 2, int{segment_size}}), take(2, random(1, int{segment_size})))); + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{37}}), take(2, random(segment_index_t{1}, segment_index_t{500}))); + + CAPTURE(c2h::type_name(), int{segment_size}, int{k}, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), int{segment_size}); + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), int{segment_size}); + + // Key/value output buffers with front/back canary padding: a signed-narrowing underrun lands in the front guard, an + // overrun in the back guard. The in-range checks below miss both (the underrun also stays inside the allocation, + // hiding it from memcheck), so we assert the guards stay untouched. `guard` covers the worst-case 8-bit offset + // (<= 256); the value canary `-1` is impossible for a real (non-negative) payload, so the value guard is exact. + constexpr segment_index_t guard = 256; + const key_t key_canary = static_cast(0x5A); + constexpr val_t val_canary = -1; + c2h::device_vector keys_out_storage(guard + num_segments * k + guard, key_canary); + c2h::device_vector values_out_storage(guard + num_segments * k + guard, val_canary); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_storage.data()) + guard; + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_storage.data()) + guard; + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), int{k}); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), int{k}); + + c2h::device_vector expected_keys(keys_in_buffer); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_segments * segment_size}); + + // Guards must be untouched by the algorithm. + const c2h::device_vector expected_kguard(guard, key_canary); + const c2h::device_vector expected_vguard(guard, val_canary); + CHECK(c2h::device_vector(keys_out_storage.begin(), keys_out_storage.begin() + guard) == expected_kguard); + CHECK(c2h::device_vector(keys_out_storage.end() - guard, keys_out_storage.end()) == expected_kguard); + CHECK(c2h::device_vector(values_out_storage.begin(), values_out_storage.begin() + guard) == expected_vguard); + CHECK(c2h::device_vector(values_out_storage.end() - guard, values_out_storage.end()) == expected_vguard); + + // Extract the in-range outputs for the standard pair verification. + c2h::device_vector keys_out_buffer(keys_out_storage.begin() + guard, keys_out_storage.end() - guard); + c2h::device_vector values_out_buffer(values_out_storage.begin() + guard, values_out_storage.end() - guard); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +#endif // TEST_TYPES == 0 + // Large-segment cluster pair tests, built only in the float build (`TEST_TYPES == 1`): they fix their own key/value // types, so repeating them per key-type axis would only waste time on the expensive 1 Mi-element runs. They run on // every launch id, including device launch (`lid_1`): the CDP static config's small resident capacity is what streams @@ -521,6 +703,80 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned REQUIRE(expected_keys == keys_out_buffer); } +// Streaming counterpart of the narrow-segment-size regression (pairs). Streaming needs segments larger than the +// resident cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- same width as +// the (unsigned) internal `offset_t` but signed -- to exercise the streaming path's index arithmetic (including the +// value gather) across det/non-det. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments with a signed 32-bit segment-size type", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using seg_size_t = int; // signed 32-bit: same width as offset_t, but signed + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + if constexpr (selected_backend != topk_backend::cluster + && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + SKIP("Determinism requirements are only provided by the cluster backend"); + } + + constexpr seg_size_t static_max_segment_size = 1024 * 1024; + constexpr seg_size_t static_max_k = 4 * 1024; + constexpr segment_index_t num_segments = 2; + + const int pad = GENERATE(0, 7); + const seg_size_t segment_size = static_max_segment_size - 31; // unaligned -> forces streaming + unaligned tail edge + const seg_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const seg_size_t k = GENERATE_COPY(values({seg_size_t{1}, max_k})); + const segment_index_t num_items = num_segments * segment_size; + + CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(pad + num_items, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + cuda::args::immediate{num_items}); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + // Tie-break preferences exercised by the deterministic tests below (only meaningful with a deterministic requirement). using tie_break_pref_list = c2h::enum_type_list -struct det_tie_pair -{ - static constexpr auto determinism = Determinism; - static constexpr auto tie_break = TieBreak; -}; -using det_tie_pair_combos = - c2h::type_list, - det_tie_pair, - det_tie_pair, - det_tie_pair>; - // Effective-cluster-width coverage for pairs: the large `deferred_sequence` bound sizes a wide (16-CTA) launch, but the // actual segments are small, so a single launch mixes a fully-resident multi-CTA segment (96 Ki + 17, every rank // works), two medium segments that leave surplus cluster CTAs idle (the runtime effective-width path, one unaligned), From 9268c5bb36a4026e87aecda8f3672f928aa584c6 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 03:45:29 +0200 Subject: [PATCH 050/126] Symmetrize prefer-large-index perf --- cub/cub/agent/agent_batched_topk_cluster.cuh | 51 ++++++++++++++------ 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 1cd82e9806c..373901587ae 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -231,6 +231,12 @@ struct agent_batched_topk_cluster Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed; static constexpr bool tie_reversed = TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + // The deterministic final scan visits chunks in global-index order and bails early (`should_stop`), so keeping the + // *first-visited* chunks resident lets it skip re-reading the streamed overflow. Ascending visits low indices first + // (the default low-resident split); descending (prefer-larger) visits high indices first, so we flip the split to + // keep the high-index chunks resident (see `run`), restoring symmetry. Never set on the non-deterministic path. + static constexpr bool reverse_residency = deterministic && tie_reversed; + // Segments at or below this size that also fit resident in one CTA take the single-CTA fast path (see // `process_impl`). static constexpr int single_cta_max_segment_size = SingleCtaMaxSegmentSize; @@ -1301,9 +1307,18 @@ private: const offset_t overflow_count = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; if constexpr (use_block_load_to_shared) { - force_tail_resident = owns_suffix_tail && !peel_tail && overflow_count > 0 && (tail_local >= my_resident_chunks); + // Under `reverse_residency` the suffix tail (global-last chunk) is the top resident chunk by construction, so no + // forced-tail exception is needed (its suffix is appended to the resident span like any other resident tail). + force_tail_resident = !reverse_residency && owns_suffix_tail && !peel_tail && overflow_count > 0 + && (tail_local >= my_resident_chunks); } - const offset_t overflow_base = force_tail_resident ? (my_resident_chunks - offset_t{1}) : my_resident_chunks; + // Rank-local base of the resident and overflow (streamed) chunk windows. Default: resident `[0, + // my_resident_chunks)`, overflow the high-index rest. `reverse_residency` swaps them: resident `[overflow_count, + // my_chunks)`, overflow + // `[0, overflow_count)`. + const offset_t resident_base = reverse_residency ? overflow_count : offset_t{0}; + const offset_t overflow_base = + reverse_residency ? offset_t{0} : (force_tail_resident ? (my_resident_chunks - offset_t{1}) : my_resident_chunks); _CCCL_ASSERT(!force_tail_resident || my_resident_chunks >= offset_t{1}, "a forced-resident tail needs at least one resident slot"); @@ -1384,9 +1399,12 @@ private: // rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot demote the reads to `LD`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - // Resident slot -> rank-local chunk index. Identity, except the last slot holds the forced-resident tail. + // Resident slot -> rank-local chunk index. `resident_base + slot` (identity unless `reverse_residency` shifts + // the resident window to the high-index chunks), except the last slot holds a forced-resident tail. const auto resident_local = [&](offset_t slot) -> offset_t { - return (force_tail_resident && slot == my_resident_chunks - offset_t{1}) ? tail_local : slot; + return (force_tail_resident && slot == my_resident_chunks - offset_t{1}) + ? tail_local + : (resident_base + slot); }; // Aligned bulk of the resident chunk in `slot` (its `count` minus any tail suffix); empty when it has none. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { @@ -1464,7 +1482,7 @@ private: { for (offset_t p = 0; p < my_resident_chunks; ++p) { - const offset_t chunk_idx = part.global_index(p); + const offset_t chunk_idx = part.global_index(resident_base + p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1554,7 +1572,7 @@ private: { for (offset_t p = 0; p < my_resident_chunks; ++p) { - const offset_t chunk_idx = part.global_index(p); + const offset_t chunk_idx = part.global_index(resident_base + p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); @@ -1811,8 +1829,9 @@ private: return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; }; - // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, which is - // visited separately after the overflow because it is the globally-last chunk. + // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, visited + // separately after the overflow because it is the globally-last chunk. Only the ascending `force_tail_resident` + // case splits off a tail; under `reverse_residency` the tail stays folded into the span (`tail_count == 0`). int tail_count = 0; int front_count = span_size(resident_keys); if constexpr (use_block_load_to_shared) @@ -1824,9 +1843,10 @@ private: } } - // Segment-local base of the resident-front span. The deterministic path's blocked partition keeps the front - // chunks contiguous and densely packed, so element `pos` of the front maps to `front_seg_base + pos`. - const offset_t front_seg_base = get_chunk(part.first, segment_size_u32, head_items).offset; + // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to + // the high-index window under `reverse_residency`). The blocked partition packs the front contiguously, so + // element `pos` maps to `front_seg_base + pos`. + const offset_t front_seg_base = get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset; auto process_resident = [&](bool reversed) { if constexpr (use_block_load_to_shared) @@ -1862,7 +1882,7 @@ private: for (int s = 0; s < rc; ++s) { const int local_slot = reversed ? (rc - 1 - s) : s; - const offset_t chunk_idx = part.global_index(static_cast(local_slot)); + const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); const int cc = chunk.count; const offset_t base_off = chunk.offset; @@ -2043,6 +2063,9 @@ private: if constexpr (tie_reversed) { + // Descending global-index order: peeled suffix tail, resident high-index chunks, streamed low-index overflow, + // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the + // gmem-re-read overflow -- mirroring the ascending path. process_tail_edge(true); if (!should_stop()) { @@ -2050,11 +2073,11 @@ private: } if (!should_stop()) { - process_overflow(true); + process_resident(true); } if (!should_stop()) { - process_resident(true); + process_overflow(true); } if (!should_stop()) { From d28d69e2c6662d8d7a1b8053d48ee62751b88df1 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 04:44:01 +0200 Subject: [PATCH 051/126] Improve histogram unrolling Should especially help small segments that did not profit from the previous unrolling implementation. Instead of a singular unrolling followed by a leftover loop, we now have a cascade of halving unroll factors. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 69 +++++++++----------- 1 file changed, 31 insertions(+), 38 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 373901587ae..f4eed81e869 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -63,6 +63,7 @@ #include #include #include +#include #include #include #include @@ -538,56 +539,48 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`), processing tiles of `histogram_items_per_thread * threads_per_block` keys. Each tile is loaded - // into registers by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the - // histogram passes: `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would - // interleave each load with its atomic instead of hoisting the whole load wave ahead. + // `[0, chunk_count)`). Keys are visited in `threads_per_block`-wide rounds via an unroll cascade: phase `p` drains + // blocks of `U = histogram_items_per_thread >> p` rounds while at least `U` remain, so the bulk runs at `U == IPT` + // and smaller `U` mop up the leftover (the last phase is always `U == 1`); this avoids the `IPT`-wide per-item + // predication a single drain loop would emit on sub-tile segments. The fully unrolled phase loop makes `U` a per-copy + // constant so the inner block loops unroll. Loads are hoisted ahead of the applies: `apply`'s SMEM atomics can't be + // proven disjoint from the SMEM key reads, so a fused loop would serialize each load with its atomic. template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - constexpr int tile = histogram_items_per_thread * threads_per_block; - const int tid = static_cast(threadIdx.x); - const int full_tiles = ::cuda::round_down(chunk_count, tile); + const int tid = static_cast(threadIdx.x); + int rounds = chunk_count / threads_per_block; // full `threads_per_block`-wide rounds + const int partial = chunk_count - rounds * threads_per_block; + int base = 0; - _CCCL_PRAGMA_NOUNROLL() - for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) + constexpr int phases = ::cuda::std::bit_width(static_cast(histogram_items_per_thread)); + _CCCL_PRAGMA_UNROLL_FULL() + for (int phase = 0; phase < phases; ++phase) { - key_t regs[histogram_items_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) - { - regs[t] = data[tile_base + t * threads_per_block + tid]; - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) + const int unroll = histogram_items_per_thread >> phase; + _CCCL_PRAGMA_NOUNROLL() + while (rounds >= unroll) { - const int local = tile_base + t * threads_per_block + tid; - apply(regs[t], local); - } - } - - if (full_tiles < chunk_count) - { - key_t regs[histogram_items_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) - { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) + key_t regs[histogram_items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < unroll; ++t) { - regs[t] = data[local]; + regs[t] = data[base + t * threads_per_block + tid]; } - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) - { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < unroll; ++t) { - apply(regs[t], local); + apply(regs[t], base + t * threads_per_block + tid); } + base += unroll * threads_per_block; + rounds -= unroll; } } + + if (tid < partial) + { + apply(data[base + tid], base + tid); + } } template From 8e7e0b6ddf472ebcb1db3a2da152f3844e2a5221 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 06:33:57 +0200 Subject: [PATCH 052/126] Revert "Improve histogram unrolling" This reverts commit d28d69e2c6662d8d7a1b8053d48ee62751b88df1. It regressed performance for large and medium segments too much. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 69 +++++++++++--------- 1 file changed, 38 insertions(+), 31 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index f4eed81e869..373901587ae 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -63,7 +63,6 @@ #include #include #include -#include #include #include #include @@ -539,48 +538,56 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`). Keys are visited in `threads_per_block`-wide rounds via an unroll cascade: phase `p` drains - // blocks of `U = histogram_items_per_thread >> p` rounds while at least `U` remain, so the bulk runs at `U == IPT` - // and smaller `U` mop up the leftover (the last phase is always `U == 1`); this avoids the `IPT`-wide per-item - // predication a single drain loop would emit on sub-tile segments. The fully unrolled phase loop makes `U` a per-copy - // constant so the inner block loops unroll. Loads are hoisted ahead of the applies: `apply`'s SMEM atomics can't be - // proven disjoint from the SMEM key reads, so a fused loop would serialize each load with its atomic. + // `[0, chunk_count)`), processing tiles of `histogram_items_per_thread * threads_per_block` keys. Each tile is loaded + // into registers by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the + // histogram passes: `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would + // interleave each load with its atomic instead of hoisting the whole load wave ahead. template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - const int tid = static_cast(threadIdx.x); - int rounds = chunk_count / threads_per_block; // full `threads_per_block`-wide rounds - const int partial = chunk_count - rounds * threads_per_block; - int base = 0; + constexpr int tile = histogram_items_per_thread * threads_per_block; + const int tid = static_cast(threadIdx.x); + const int full_tiles = ::cuda::round_down(chunk_count, tile); - constexpr int phases = ::cuda::std::bit_width(static_cast(histogram_items_per_thread)); - _CCCL_PRAGMA_UNROLL_FULL() - for (int phase = 0; phase < phases; ++phase) + _CCCL_PRAGMA_NOUNROLL() + for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) { - const int unroll = histogram_items_per_thread >> phase; - _CCCL_PRAGMA_NOUNROLL() - while (rounds >= unroll) + key_t regs[histogram_items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) { - key_t regs[histogram_items_per_thread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < unroll; ++t) + regs[t] = data[tile_base + t * threads_per_block + tid]; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + const int local = tile_base + t * threads_per_block + tid; + apply(regs[t], local); + } + } + + if (full_tiles < chunk_count) + { + key_t regs[histogram_items_per_thread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) { - regs[t] = data[base + t * threads_per_block + tid]; + regs[t] = data[local]; } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < unroll; ++t) + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < histogram_items_per_thread; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) { - apply(regs[t], base + t * threads_per_block + tid); + apply(regs[t], local); } - base += unroll * threads_per_block; - rounds -= unroll; } } - - if (tid < partial) - { - apply(data[base + tid], base + tid); - } } template From 13fb7e336b21ef16ba5b65472ea9cdca5e8580c2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 20:08:56 +0200 Subject: [PATCH 053/126] Clamp unrolling for small segments Only possible when their size is know statically. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 130 ++++++++++++------- 1 file changed, 84 insertions(+), 46 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 373901587ae..53cdd34e7fc 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -256,17 +256,48 @@ struct agent_batched_topk_cluster static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int histogram_items_per_thread = HistogramItemsPerThread; + + // Static upper bound on segment size: exact for constant/immediate sizes, the type maximum for runtime sizes. + static constexpr ::cuda::std::int64_t static_max_segment_size = + ::cuda::args::__traits::highest; + + // Segments small enough to always be single-CTA resident (one contiguous SMEM span, see `run`) never need more than + // `ceil(static_max_segment_size / threads_per_block)` sweep rounds, so we clamp the per-thread unroll to that bound + // to trim predication/registers on sub-tile segments. Larger/unbounded types keep the full unroll, so their codegen + // is unchanged; the guard also keeps the rounds arithmetic in `int` range. + static constexpr bool clamp_items_to_segment = + static_max_segment_size > 0 && static_max_segment_size < single_cta_max_segment_size; + + // Histogram sweeps clamp to `floor`: the bulk runs unpredicated, leaving a sub-tile predicated tail. + static constexpr int histogram_rounds_floor = + clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : histogram_items_per_thread; + static constexpr int histogram_items_per_thread_clamped = + ::cuda::std::clamp(histogram_rounds_floor, 1, histogram_items_per_thread); + static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); + static constexpr int load_align_bytes = LoadAlignBytes; static constexpr int bits_per_pass = BitsPerPass; static constexpr int tie_break_items_per_thread = TieBreakItemsPerThread; - static constexpr int num_buckets = 1 << bits_per_pass; - using smem_layout_t = smem_block_tile_layout; - static constexpr int chunk_items = smem_layout_t::chunk_items; - static constexpr int load_align_items = smem_layout_t::load_align_items; - static constexpr int slot_alignment = smem_layout_t::slot_alignment; + + // The final filter clamps to `ceil` so one predicated tile spans the whole resident segment (uniform `pos < count` + // guard, no full-vs-partial-tile split for the reversed-tie BlockScan or the atomic writes). + static constexpr int tie_break_rounds_ceil = + clamp_items_to_segment + ? static_cast( + ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) + : tie_break_items_per_thread; + static constexpr int tie_break_items_per_thread_clamped = + ::cuda::std::clamp(tie_break_rounds_ceil, 1, tie_break_items_per_thread); + static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); + static constexpr int num_buckets = 1 << bits_per_pass; + using smem_layout_t = smem_block_tile_layout; + static constexpr int chunk_items = smem_layout_t::chunk_items; + static constexpr int load_align_items = smem_layout_t::load_align_items; + static constexpr int slot_alignment = smem_layout_t::slot_alignment; static_assert(PipelineStages > 0); static_assert(HistogramItemsPerThread > 0, "histogram_items_per_thread must be positive"); + static_assert(TieBreakItemsPerThread > 0, "tie_break_items_per_thread must be positive"); static_assert(ChunkBytes > 0); static_assert(LoadAlignBytes > 0); static_assert(ChunkBytes % LoadAlignBytes == 0, "ChunkBytes must be a multiple of LoadAlignBytes"); @@ -538,28 +569,28 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`), processing tiles of `histogram_items_per_thread * threads_per_block` keys. Each tile is loaded - // into registers by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the - // histogram passes: `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would - // interleave each load with its atomic instead of hoisting the whole load wave ahead. - template + // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each tile is loaded into registers by + // one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: + // `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would interleave each load + // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped items-per-thread. + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - constexpr int tile = histogram_items_per_thread * threads_per_block; + constexpr int tile = Unroll * threads_per_block; const int tid = static_cast(threadIdx.x); const int full_tiles = ::cuda::round_down(chunk_count, tile); _CCCL_PRAGMA_NOUNROLL() for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) { - key_t regs[histogram_items_per_thread]; + key_t regs[Unroll]; _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) + for (int t = 0; t < Unroll; ++t) { regs[t] = data[tile_base + t * threads_per_block + tid]; } _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) + for (int t = 0; t < Unroll; ++t) { const int local = tile_base + t * threads_per_block + tid; apply(regs[t], local); @@ -568,9 +599,9 @@ struct agent_batched_topk_cluster if (full_tiles < chunk_count) { - key_t regs[histogram_items_per_thread]; + key_t regs[Unroll]; _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) + for (int t = 0; t < Unroll; ++t) { const int local = full_tiles + t * threads_per_block + tid; if (local < chunk_count) @@ -579,7 +610,7 @@ struct agent_batched_topk_cluster } } _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < histogram_items_per_thread; ++t) + for (int t = 0; t < Unroll; ++t) { const int local = full_tiles + t * threads_per_block + tid; if (local < chunk_count) @@ -590,10 +621,10 @@ struct agent_batched_topk_cluster } } - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const { - for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { + for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { f(key); }); } @@ -601,13 +632,14 @@ struct agent_batched_topk_cluster // Like `for_each_chunk_key`, but also hands `f` each key's segment-local index `base_off + local`, where `base_off` // is the segment-local offset of the chunk's first element. The pair path uses that index to fetch the key's value // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const { - for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { - f(key, base_off + static_cast(local)); - }); + for_each_chunk_key_impl( + ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { + f(key, base_off + static_cast(local)); + }); } // A bulk in the block_tile as a 32-bit shared address + length. A spilled 32-bit shared address (rebuilt with @@ -1076,15 +1108,14 @@ private: } // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap - // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass the - // tuning parameter matching their context (`histogram_items_per_thread` for the histogram passes, - // `tie_break_items_per_thread` for the non-deterministic final filter). + // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass their + // clamped items-per-thread. template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { - agent.for_each_chunk_key(stage_span(stage, o), f); + agent.template for_each_chunk_key(stage_span(stage, o), f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1118,7 +1149,7 @@ private: run_pass( [&](int stage, offset_t o) { const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; - agent.for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1353,17 +1384,20 @@ private: // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). - const auto fold_edges = [&](auto&& apply) { + // `unroll_ic` is an `integral_constant` carrying the caller's clamped items-per-thread. + const auto fold_edges = [&](auto unroll_ic, auto&& apply) { if constexpr (use_block_load_to_shared) { + constexpr int unroll = unroll_ic(); if (head_edge_len > 0) { - for_each_chunk_key({temp_storage.edge_keys, static_cast<::cuda::std::size_t>(head_edge_len)}, apply); + this->template for_each_chunk_key( + {temp_storage.edge_keys, static_cast<::cuda::std::size_t>(head_edge_len)}, apply); } if (tail_edge_len > 0) { - for_each_chunk_key({temp_storage.edge_keys + head_edge_cap, static_cast<::cuda::std::size_t>(tail_edge_len)}, - apply); + this->template for_each_chunk_key( + {temp_storage.edge_keys + head_edge_cap, static_cast<::cuda::std::size_t>(tail_edge_len)}, apply); } } }; @@ -1439,7 +1473,7 @@ private: const int stage = static_cast(p % static_cast(prologue)); wait_stage(stage); const int read_len = static_cast(::cuda::std::size(bulk_src(p))); - for_each_chunk_key( + for_each_chunk_key( bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), add_first_pass); read_off += read_len * int{sizeof(key_t)}; @@ -1486,7 +1520,7 @@ private: const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - _CCCL_PRAGMA_UNROLL(histogram_items_per_thread) + _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) { const int local = j * threads_per_block + static_cast(threadIdx.x); @@ -1525,7 +1559,7 @@ private: // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer // sync. - streamer.process_pass(add_first_pass); + streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -1566,7 +1600,8 @@ private: // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value // spilled across the pass loop (see `resident_smem32`). key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); + for_each_chunk_key( + {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } else { @@ -1575,7 +1610,8 @@ private: const offset_t chunk_idx = part.global_index(resident_base + p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); + for_each_chunk_key( + {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } } }; @@ -1583,12 +1619,12 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. - streamer.process_pass(add_hist, fold_resident_hist); + streamer.process_pass(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's `hist[last_bucket]` (the deterministic cross-CTA prefix // input) inclusive of its edge candidates. - fold_edges(add_hist); + fold_edges(::cuda::std::integral_constant{}, add_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -1745,7 +1781,7 @@ private: // written reversed at `block_keys_out[k - 1 - rank]`. `running` carries across tiles and regions. `get_idx(pos)` // gives the segment-local index, used only to load the value payload (compiled out in keys-only builds). auto process_flat = [&](auto get_key, auto get_idx, int count) { - constexpr int items = tie_break_items_per_thread; + constexpr int items = tie_break_items_per_thread_clamped; constexpr int tile = threads_per_block * items; for (int tile_base = 0; tile_base < count; tile_base += tile) { @@ -2132,7 +2168,8 @@ private: if constexpr (use_block_load_to_shared) { key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key({rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + for_each_chunk_key( + {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } else { @@ -2141,13 +2178,14 @@ private: const offset_t chunk_idx = part.global_index(p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key({chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + for_each_chunk_key( + {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_edges(write_selected); + fold_edges(::cuda::std::integral_constant{}, write_selected); }; - streamer.process_pass(write_selected, fold_resident); + streamer.process_pass(write_selected, fold_resident); } else { @@ -2178,7 +2216,7 @@ private: // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). auto write_run = [&](auto get_key, offset_t base_off, int count) { const int iterations = ::cuda::ceil_div(count, threads_per_block); - _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread) + _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) { const int local = j * threads_per_block + static_cast(threadIdx.x); @@ -2254,7 +2292,7 @@ private: // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. - streamer.process_pass_indexed(write_selected_idx, fold_resident); + streamer.process_pass_indexed(write_selected_idx, fold_resident); } } From e44878797a602fc2ecf0b52ab9bead96362ab4f5 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 21:00:41 +0200 Subject: [PATCH 054/126] Also clamp unrolling of the k>=segment_size path Also add a tuning parameter for the unrolling on that path instead of relying on an un-tunable magic constant. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 16 +++++++++++++++- .../dispatch/dispatch_batched_topk_cluster.cuh | 8 ++++++++ .../tuning/tuning_batched_topk_cluster.cuh | 7 +++---- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 53cdd34e7fc..4a5331192e4 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -194,6 +194,7 @@ template = 1, "tie_break_items_per_thread_clamped must be positive"); + + // Per-thread unroll for the select-all copy fast path (`copy_segment_select_all`). Clamps to `ceil` like the final + // filter: the loops already guard each item with `idx[j] < n`, so a ceil-rounded count stays correct. + static constexpr int copy_items_per_thread = CopyItemsPerThread; + static constexpr int copy_rounds_ceil = + clamp_items_to_segment + ? static_cast( + ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) + : copy_items_per_thread; + static constexpr int copy_items_per_thread_clamped = ::cuda::std::clamp(copy_rounds_ceil, 1, copy_items_per_thread); + static_assert(copy_items_per_thread_clamped >= 1, "copy_items_per_thread_clamped must be positive"); + static constexpr int num_buckets = 1 << bits_per_pass; using smem_layout_t = smem_block_tile_layout; static constexpr int chunk_items = smem_layout_t::chunk_items; @@ -296,6 +309,7 @@ struct agent_batched_topk_cluster static constexpr int slot_alignment = smem_layout_t::slot_alignment; static_assert(PipelineStages > 0); + static_assert(CopyItemsPerThread > 0, "copy_items_per_thread must be positive"); static_assert(HistogramItemsPerThread > 0, "histogram_items_per_thread must be positive"); static_assert(TieBreakItemsPerThread > 0, "tie_break_items_per_thread must be positive"); static_assert(ChunkBytes > 0); @@ -2316,7 +2330,7 @@ private: unsigned int cluster_rank, unsigned int cluster_blocks) { - constexpr int copy_items = 8; + constexpr int copy_items = copy_items_per_thread_clamped; const offset_t n = static_cast(segment_size); const offset_t cluster_tid = cluster_rank * static_cast(blockDim.x) + threadIdx.x; const offset_t cluster_tids = cluster_blocks * static_cast(blockDim.x); diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 258e188f01f..d234e3f65f2 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -87,6 +87,7 @@ template ; using key_t = it_value_t; @@ -354,6 +360,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( TieBreakItemsPerThread, SingleCtaMaxSegmentSize, MinChunksPerCta, + CopyItemsPerThread, Determinism, TieBreak, KeyInputItItT, @@ -431,6 +438,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( TieBreakItemsPerThread, SingleCtaMaxSegmentSize, MinChunksPerCta, + CopyItemsPerThread, Determinism, TieBreak, KeyInputItItT, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index d3a3ef6c00a..5bf9c7fc127 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -31,10 +31,8 @@ struct cluster_topk_policy int min_blocks_per_sm; int tie_break_items_per_thread; int single_cta_max_segment_size; - // Minimum chunks a CTA must own to belong to a segment's effective cluster. At 1 the effective cluster is just the - // CTAs that receive any chunk; a larger value trades parallelism for fewer, busier CTAs (and shrinks the - // host-launched cluster). Used by both the host cluster-blocks cap and the device per-segment effective cluster. int min_chunks_per_cta; + int copy_items_per_thread; }; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy @@ -49,7 +47,8 @@ struct cluster_topk_policy /*min_blocks_per_sm=*/1, /*tie_break_items_per_thread=*/8, /*single_cta_max_segment_size=*/8 * 1024, - /*min_chunks_per_cta=*/1}; + /*min_chunks_per_cta=*/1, + /*copy_items_per_thread=*/8}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool From 0bc7b85a280a963260f864117316a42364a48b03 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 27 Jun 2026 21:18:32 +0200 Subject: [PATCH 055/126] Add histogram unrolling experiment --- cub/cub/agent/agent_batched_topk_cluster.cuh | 209 +++++++++++++------ 1 file changed, 140 insertions(+), 69 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 4a5331192e4..6dac60dd56b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -175,6 +175,29 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, } // ----------------------------------------------------------------------------- +// A/B/C benchmarking selector for the histogram sweep's small-segment unroll strategy. Only the histogram path is +// affected; the final filter and copy fast path always use `split_ceil` (mode 1). +// 1 (default): `ceil` clamp; the predicated remainder is split into a load wave and an apply wave, both unrolled. +// 2 : `floor` clamp; the remainder is a single non-unrolled fused loop (the pre-optimization baseline). +// 3 : `ceil` clamp; the remainder is a single fully-unrolled fused loop with the bound check inside each +// step. +#ifndef CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE +# define CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE 1 +#endif +#if CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE < 1 || CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE > 3 +# error "CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE must be 1, 2, or 3" +#endif + +// How `for_each_chunk_key_impl` handles the sub-tile remainder (the main full-tile loop is always the unrolled +// load-wave/apply-wave form). `split_ceil` is the default for every caller; the histogram path may override it via the +// macro above. +enum class chunk_remainder_mode +{ + split_ceil, // load wave then apply wave, both unrolled (mode 1) + fused_floor, // single non-unrolled fused loop (mode 2) + fused_ceil, // single fully-unrolled fused loop, bound check per step (mode 3) +}; + // Cluster top-k agent // ----------------------------------------------------------------------------- // Cluster blocks is a runtime value (see `process_impl` for the readback), so @@ -257,49 +280,56 @@ struct agent_batched_topk_cluster static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int histogram_items_per_thread = HistogramItemsPerThread; + static constexpr int load_align_bytes = LoadAlignBytes; + static constexpr int bits_per_pass = BitsPerPass; + static constexpr int tie_break_items_per_thread = TieBreakItemsPerThread; + static constexpr int copy_items_per_thread = CopyItemsPerThread; // select-all copy fast path // Static upper bound on segment size: exact for constant/immediate sizes, the type maximum for runtime sizes. static constexpr ::cuda::std::int64_t static_max_segment_size = ::cuda::args::__traits::highest; - // Segments small enough to always be single-CTA resident (one contiguous SMEM span, see `run`) never need more than - // `ceil(static_max_segment_size / threads_per_block)` sweep rounds, so we clamp the per-thread unroll to that bound - // to trim predication/registers on sub-tile segments. Larger/unbounded types keep the full unroll, so their codegen - // is unchanged; the guard also keeps the rounds arithmetic in `int` range. + // Segments small enough to always be single-CTA resident (one contiguous SMEM span, see `run`) need at most + // `ceil(static_max_segment_size / threads_per_block)` sweep rounds. We clamp each per-thread unroll down to that + // bound to trim predication/registers on sub-tile segments. Larger/unbounded types keep the full unroll (codegen + // unchanged); the guard also keeps the rounds arithmetic in `int` range. static constexpr bool clamp_items_to_segment = static_max_segment_size > 0 && static_max_segment_size < single_cta_max_segment_size; - - // Histogram sweeps clamp to `floor`: the bulk runs unpredicated, leaving a sub-tile predicated tail. - static constexpr int histogram_rounds_floor = - clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : histogram_items_per_thread; - static constexpr int histogram_items_per_thread_clamped = - ::cuda::std::clamp(histogram_rounds_floor, 1, histogram_items_per_thread); - static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); - - static constexpr int load_align_bytes = LoadAlignBytes; - static constexpr int bits_per_pass = BitsPerPass; - static constexpr int tie_break_items_per_thread = TieBreakItemsPerThread; - - // The final filter clamps to `ceil` so one predicated tile spans the whole resident segment (uniform `pos < count` - // guard, no full-vs-partial-tile split for the reversed-tie BlockScan or the atomic writes). - static constexpr int tie_break_rounds_ceil = + static constexpr int segment_rounds_ceil = clamp_items_to_segment ? static_cast( ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) - : tie_break_items_per_thread; - static constexpr int tie_break_items_per_thread_clamped = - ::cuda::std::clamp(tie_break_rounds_ceil, 1, tie_break_items_per_thread); - static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); + : 0; - // Per-thread unroll for the select-all copy fast path (`copy_segment_select_all`). Clamps to `ceil` like the final - // filter: the loops already guard each item with `idx[j] < n`, so a ceil-rounded count stays correct. - static constexpr int copy_items_per_thread = CopyItemsPerThread; - static constexpr int copy_rounds_ceil = + // `ceil` (not `floor`) so the whole resident segment fits in a single tile. In `for_each_chunk_key_impl` (histogram + // and non-deterministic filter) the unpredicated main loop then runs zero times and the unrolled, predicated + // remainder covers every key -- giving small/non-tight segments the same ILP as the hot loop without adding + // conditionals to it. Each unroll is capped at its tuning value; non-clamped (large/unbounded) segments keep the full + // width. + // + // The histogram path is the A/B/C subject (see `CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE`): mode 2 clamps to `floor` and + // uses a non-unrolled remainder; modes 1/3 clamp to `ceil` and differ only in the remainder shape (split vs. fused); + // the final filter and copy always use mode 1's `ceil` + split remainder. + static constexpr chunk_remainder_mode histogram_rem_mode = + CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 2 ? chunk_remainder_mode::fused_floor + : CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 3 + ? chunk_remainder_mode::fused_ceil + : chunk_remainder_mode::split_ceil; + static constexpr bool histogram_use_floor = histogram_rem_mode == chunk_remainder_mode::fused_floor; + static constexpr int histogram_rounds_clamped = clamp_items_to_segment - ? static_cast( - ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) - : copy_items_per_thread; - static constexpr int copy_items_per_thread_clamped = ::cuda::std::clamp(copy_rounds_ceil, 1, copy_items_per_thread); + ? (histogram_use_floor ? static_cast(static_max_segment_size / threads_per_block) : segment_rounds_ceil) + : histogram_items_per_thread; + static constexpr int histogram_items_per_thread_clamped = + ::cuda::std::clamp(histogram_rounds_clamped, 1, histogram_items_per_thread); + static constexpr int tie_break_items_per_thread_clamped = + clamp_items_to_segment + ? ::cuda::std::clamp(segment_rounds_ceil, 1, tie_break_items_per_thread) + : tie_break_items_per_thread; + static constexpr int copy_items_per_thread_clamped = + clamp_items_to_segment ? ::cuda::std::clamp(segment_rounds_ceil, 1, copy_items_per_thread) : copy_items_per_thread; + static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); + static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); static_assert(copy_items_per_thread_clamped >= 1, "copy_items_per_thread_clamped must be positive"); static constexpr int num_buckets = 1 << bits_per_pass; @@ -583,11 +613,14 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each tile is loaded into registers by - // one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: + // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each full tile is loaded into registers + // by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: // `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would interleave each load // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped items-per-thread. - template + // `RemMode` selects the sub-tile remainder shape (the main full-tile loop is always the unrolled split form): + // `split_ceil` mirrors the main loop (load wave then apply wave, default), `fused_ceil` is one fully-unrolled fused + // loop with the bound check per step, `fused_floor` the same but non-unrolled. See `chunk_remainder_mode`. + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { constexpr int tile = Unroll * threads_per_block; @@ -613,44 +646,76 @@ struct agent_batched_topk_cluster if (full_tiles < chunk_count) { - key_t regs[Unroll]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) + if constexpr (RemMode == chunk_remainder_mode::split_ceil) { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) + key_t regs[Unroll]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + regs[t] = data[local]; + } + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) { - regs[t] = data[local]; + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + apply(regs[t], local); + } } } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) + else { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) + // Single fused loop, bound check per step; `fused_ceil` unrolls it, `fused_floor` does not. + if constexpr (RemMode == chunk_remainder_mode::fused_ceil) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + apply(data[local], local); + } + } + } + else { - apply(regs[t], local); + _CCCL_PRAGMA_NOUNROLL() + for (int t = 0; t < Unroll; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + apply(data[local], local); + } + } } } } } - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const { - for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { - f(key); - }); + for_each_chunk_key_impl( + ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { + f(key); + }); } // Like `for_each_chunk_key`, but also hands `f` each key's segment-local index `base_off + local`, where `base_off` // is the segment-local offset of the chunk's first element. The pair path uses that index to fetch the key's value // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const { - for_each_chunk_key_impl( + for_each_chunk_key_impl( ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { f(key, base_off + static_cast(local)); }); @@ -1124,12 +1189,12 @@ private: // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass their // clamped items-per-thread. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { - agent.template for_each_chunk_key(stage_span(stage, o), f); + agent.template for_each_chunk_key(stage_span(stage, o), f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1148,22 +1213,22 @@ private: // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in // by the BlockLoadToShared pipeline (rather than already resident in SMEM). - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) { - process_pass(static_cast(f), [] {}); + process_pass(static_cast(f), [] {}); } // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. The pair // final filter needs that index to fetch each selected key's value payload from gmem, while still reusing the // overflow keys from the streaming SMEM pipeline (block-load path) instead of re-reading them from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; - agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1398,19 +1463,21 @@ private: // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). - // `unroll_ic` is an `integral_constant` carrying the caller's clamped items-per-thread. - const auto fold_edges = [&](auto unroll_ic, auto&& apply) { + // `unroll_ic` / `rem_mode_ic` are `integral_constant`s carrying the caller's clamped items-per-thread and its + // remainder mode (the histogram may pass a `floor`/fused mode in the A/B/C experiment). + const auto fold_edges = [&](auto unroll_ic, auto rem_mode_ic, auto&& apply) { if constexpr (use_block_load_to_shared) { - constexpr int unroll = unroll_ic(); + constexpr int unroll = unroll_ic(); + constexpr chunk_remainder_mode remode = rem_mode_ic(); if (head_edge_len > 0) { - this->template for_each_chunk_key( + this->template for_each_chunk_key( {temp_storage.edge_keys, static_cast<::cuda::std::size_t>(head_edge_len)}, apply); } if (tail_edge_len > 0) { - this->template for_each_chunk_key( + this->template for_each_chunk_key( {temp_storage.edge_keys + head_edge_cap, static_cast<::cuda::std::size_t>(tail_edge_len)}, apply); } } @@ -1487,7 +1554,7 @@ private: const int stage = static_cast(p % static_cast(prologue)); wait_stage(stage); const int read_len = static_cast(::cuda::std::size(bulk_src(p))); - for_each_chunk_key( + for_each_chunk_key( bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), add_first_pass); read_off += read_len * int{sizeof(key_t)}; @@ -1573,7 +1640,7 @@ private: // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer // sync. - streamer.process_pass(add_first_pass); + streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -1614,7 +1681,7 @@ private: // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value // spilled across the pass loop (see `resident_smem32`). key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key( + for_each_chunk_key( {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } else @@ -1624,7 +1691,7 @@ private: const offset_t chunk_idx = part.global_index(resident_base + p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key( + for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } } @@ -1633,12 +1700,14 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. - streamer.process_pass(add_hist, fold_resident_hist); + streamer.process_pass(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's `hist[last_bucket]` (the deterministic cross-CTA prefix // input) inclusive of its edge candidates. - fold_edges(::cuda::std::integral_constant{}, add_hist); + fold_edges(::cuda::std::integral_constant{}, + ::cuda::std::integral_constant{}, + add_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -2197,7 +2266,9 @@ private: } } // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_edges(::cuda::std::integral_constant{}, write_selected); + fold_edges(::cuda::std::integral_constant{}, + ::cuda::std::integral_constant{}, + write_selected); }; streamer.process_pass(write_selected, fold_resident); } From 04e7f65496734752845811629ba7f1a6025d836d Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 28 Jun 2026 00:16:49 +0200 Subject: [PATCH 056/126] Do another histogram unrolling experiment --- cub/cub/agent/agent_batched_topk_cluster.cuh | 135 ++++++++++--------- 1 file changed, 74 insertions(+), 61 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 6dac60dd56b..5ceeaf30499 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -175,27 +175,28 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, } // ----------------------------------------------------------------------------- -// A/B/C benchmarking selector for the histogram sweep's small-segment unroll strategy. Only the histogram path is -// affected; the final filter and copy fast path always use `split_ceil` (mode 1). -// 1 (default): `ceil` clamp; the predicated remainder is split into a load wave and an apply wave, both unrolled. -// 2 : `floor` clamp; the remainder is a single non-unrolled fused loop (the pre-optimization baseline). -// 3 : `ceil` clamp; the remainder is a single fully-unrolled fused loop with the bound check inside each -// step. +// A/B benchmarking selector for the histogram sweep's small-segment unroll strategy. Only the histogram path is +// affected; the final filter and copy fast path always use `split_ceil`. +// 1 (default): `fused_floor` -- `floor` clamp; a separate, non-unrolled fused remainder loop after the unpredicated +// main tiles. +// 2 : `unified_ceil` -- `ceil` clamp; no separate remainder. A single outer tile loop whose unrolled, split +// load/apply waves are predicated and cover the tail too (pays a per-step bound check on every tile). #ifndef CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE # define CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE 1 #endif -#if CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE < 1 || CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE > 3 -# error "CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE must be 1, 2, or 3" +#if CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE < 1 || CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE > 2 +# error "CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE must be 1 or 2" #endif -// How `for_each_chunk_key_impl` handles the sub-tile remainder (the main full-tile loop is always the unrolled -// load-wave/apply-wave form). `split_ceil` is the default for every caller; the histogram path may override it via the -// macro above. +// How `for_each_chunk_key_impl` covers a chunk. `split_ceil` (the default for every caller, including the final filter) +// and `fused_floor` keep the unpredicated main tile loop plus a separate remainder (split vs. single fused). +// `unified_ceil` has no separate remainder: one outer tile loop whose predicated, unrolled load/apply waves also cover +// the tail. enum class chunk_remainder_mode { - split_ceil, // load wave then apply wave, both unrolled (mode 1) - fused_floor, // single non-unrolled fused loop (mode 2) - fused_ceil, // single fully-unrolled fused loop, bound check per step (mode 3) + split_ceil, // unpredicated main tiles + a split (load wave/apply wave) unrolled, predicated remainder + fused_floor, // unpredicated main tiles + a single non-unrolled fused remainder loop + unified_ceil, // no separate remainder: every tile (incl. the partial last) uses predicated split unrolled waves }; // Cluster top-k agent @@ -301,20 +302,16 @@ struct agent_batched_topk_cluster ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) : 0; - // `ceil` (not `floor`) so the whole resident segment fits in a single tile. In `for_each_chunk_key_impl` (histogram - // and non-deterministic filter) the unpredicated main loop then runs zero times and the unrolled, predicated - // remainder covers every key -- giving small/non-tight segments the same ILP as the hot loop without adding - // conditionals to it. Each unroll is capped at its tuning value; non-clamped (large/unbounded) segments keep the full - // width. + // A `ceil` clamp keeps the whole resident segment inside a single tile (the predicated path in + // `for_each_chunk_key_impl` then covers every key in one unrolled pass with the same ILP as the hot loop); a `floor` + // clamp leaves a sub-tile remainder for the separate non-unrolled loop. Each unroll is capped at its tuning value; + // non-clamped (large/unbounded) segments keep the full width, and the guard keeps the rounds arithmetic in `int`. // - // The histogram path is the A/B/C subject (see `CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE`): mode 2 clamps to `floor` and - // uses a non-unrolled remainder; modes 1/3 clamp to `ceil` and differ only in the remainder shape (split vs. fused); - // the final filter and copy always use mode 1's `ceil` + split remainder. + // The histogram path is the A/B subject (see `CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE`): mode 1 (`fused_floor`) clamps + // to `floor` with a separate non-unrolled remainder; mode 2 (`unified_ceil`) clamps to `ceil` and drops the separate + // remainder. The final filter and copy keep `split_ceil` regardless. static constexpr chunk_remainder_mode histogram_rem_mode = - CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 2 ? chunk_remainder_mode::fused_floor - : CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 3 - ? chunk_remainder_mode::fused_ceil - : chunk_remainder_mode::split_ceil; + CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 2 ? chunk_remainder_mode::unified_ceil : chunk_remainder_mode::fused_floor; static constexpr bool histogram_use_floor = histogram_rem_mode == chunk_remainder_mode::fused_floor; static constexpr int histogram_rounds_clamped = clamp_items_to_segment @@ -613,46 +610,31 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each full tile is loaded into registers - // by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: + // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each tile is loaded into registers by + // one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: // `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would interleave each load // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped items-per-thread. - // `RemMode` selects the sub-tile remainder shape (the main full-tile loop is always the unrolled split form): - // `split_ceil` mirrors the main loop (load wave then apply wave, default), `fused_ceil` is one fully-unrolled fused - // loop with the bound check per step, `fused_floor` the same but non-unrolled. See `chunk_remainder_mode`. + // `RemMode` selects how the chunk is covered (see `chunk_remainder_mode`): + // `split_ceil` (default): unpredicated main tiles + a split (load wave/apply wave) unrolled, predicated remainder. + // `fused_floor` : unpredicated main tiles + a single non-unrolled fused remainder loop. + // `unified_ceil` : no separate remainder -- one outer tile loop whose split, unrolled load/apply waves are + // predicated and cover the partial last tile too (a per-step bound check on every tile). template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - constexpr int tile = Unroll * threads_per_block; - const int tid = static_cast(threadIdx.x); - const int full_tiles = ::cuda::round_down(chunk_count, tile); + constexpr int tile = Unroll * threads_per_block; + const int tid = static_cast(threadIdx.x); - _CCCL_PRAGMA_NOUNROLL() - for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) + if constexpr (RemMode == chunk_remainder_mode::unified_ceil) { - key_t regs[Unroll]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - regs[t] = data[tile_base + t * threads_per_block + tid]; - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = tile_base + t * threads_per_block + tid; - apply(regs[t], local); - } - } - - if (full_tiles < chunk_count) - { - if constexpr (RemMode == chunk_remainder_mode::split_ceil) + _CCCL_PRAGMA_NOUNROLL() + for (int tile_base = 0; tile_base < chunk_count; tile_base += tile) { key_t regs[Unroll]; _CCCL_PRAGMA_UNROLL_FULL() for (int t = 0; t < Unroll; ++t) { - const int local = full_tiles + t * threads_per_block + tid; + const int local = tile_base + t * threads_per_block + tid; if (local < chunk_count) { regs[t] = data[local]; @@ -661,29 +643,60 @@ struct agent_batched_topk_cluster _CCCL_PRAGMA_UNROLL_FULL() for (int t = 0; t < Unroll; ++t) { - const int local = full_tiles + t * threads_per_block + tid; + const int local = tile_base + t * threads_per_block + tid; if (local < chunk_count) { apply(regs[t], local); } } } - else + } + else + { + const int full_tiles = ::cuda::round_down(chunk_count, tile); + + _CCCL_PRAGMA_NOUNROLL() + for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) + { + key_t regs[Unroll]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) + { + regs[t] = data[tile_base + t * threads_per_block + tid]; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) + { + const int local = tile_base + t * threads_per_block + tid; + apply(regs[t], local); + } + } + + if (full_tiles < chunk_count) { - // Single fused loop, bound check per step; `fused_ceil` unrolls it, `fused_floor` does not. - if constexpr (RemMode == chunk_remainder_mode::fused_ceil) + if constexpr (RemMode == chunk_remainder_mode::split_ceil) { + key_t regs[Unroll]; _CCCL_PRAGMA_UNROLL_FULL() for (int t = 0; t < Unroll; ++t) { const int local = full_tiles + t * threads_per_block + tid; if (local < chunk_count) { - apply(data[local], local); + regs[t] = data[local]; + } + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) + { + const int local = full_tiles + t * threads_per_block + tid; + if (local < chunk_count) + { + apply(regs[t], local); } } } - else + else // fused_floor: single non-unrolled fused loop, bound check per step { _CCCL_PRAGMA_NOUNROLL() for (int t = 0; t < Unroll; ++t) @@ -1464,7 +1477,7 @@ private: // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). // `unroll_ic` / `rem_mode_ic` are `integral_constant`s carrying the caller's clamped items-per-thread and its - // remainder mode (the histogram may pass a `floor`/fused mode in the A/B/C experiment). + // remainder mode (the histogram may pass a `floor`/fused mode in the A/B experiment). const auto fold_edges = [&](auto unroll_ic, auto rem_mode_ic, auto&& apply) { if constexpr (use_block_load_to_shared) { From b269e7ac3bdc65c710db452661ce4830b67b679f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 28 Jun 2026 03:44:46 +0200 Subject: [PATCH 057/126] Clean up unrolling experiments --- cub/cub/agent/agent_batched_topk_cluster.cuh | 319 +++++++------------ 1 file changed, 120 insertions(+), 199 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5ceeaf30499..8938ed21415 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -174,31 +174,6 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, (::cuda::std::clamp) (blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } -// ----------------------------------------------------------------------------- -// A/B benchmarking selector for the histogram sweep's small-segment unroll strategy. Only the histogram path is -// affected; the final filter and copy fast path always use `split_ceil`. -// 1 (default): `fused_floor` -- `floor` clamp; a separate, non-unrolled fused remainder loop after the unpredicated -// main tiles. -// 2 : `unified_ceil` -- `ceil` clamp; no separate remainder. A single outer tile loop whose unrolled, split -// load/apply waves are predicated and cover the tail too (pays a per-step bound check on every tile). -#ifndef CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE -# define CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE 1 -#endif -#if CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE < 1 || CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE > 2 -# error "CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE must be 1 or 2" -#endif - -// How `for_each_chunk_key_impl` covers a chunk. `split_ceil` (the default for every caller, including the final filter) -// and `fused_floor` keep the unpredicated main tile loop plus a separate remainder (split vs. single fused). -// `unified_ceil` has no separate remainder: one outer tile loop whose predicated, unrolled load/apply waves also cover -// the tail. -enum class chunk_remainder_mode -{ - split_ceil, // unpredicated main tiles + a split (load wave/apply wave) unrolled, predicated remainder - fused_floor, // unpredicated main tiles + a single non-unrolled fused remainder loop - unified_ceil, // no separate remainder: every tile (incl. the partial last) uses predicated split unrolled waves -}; - // Cluster top-k agent // ----------------------------------------------------------------------------- // Cluster blocks is a runtime value (see `process_impl` for the readback), so @@ -301,32 +276,34 @@ struct agent_batched_topk_cluster ? static_cast( ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) : 0; - - // A `ceil` clamp keeps the whole resident segment inside a single tile (the predicated path in - // `for_each_chunk_key_impl` then covers every key in one unrolled pass with the same ILP as the hot loop); a `floor` - // clamp leaves a sub-tile remainder for the separate non-unrolled loop. Each unroll is capped at its tuning value; - // non-clamped (large/unbounded) segments keep the full width, and the guard keeps the rounds arithmetic in `int`. - // - // The histogram path is the A/B subject (see `CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE`): mode 1 (`fused_floor`) clamps - // to `floor` with a separate non-unrolled remainder; mode 2 (`unified_ceil`) clamps to `ceil` and drops the separate - // remainder. The final filter and copy keep `split_ceil` regardless. - static constexpr chunk_remainder_mode histogram_rem_mode = - CUB_DETAIL_TOPK_HISTOGRAM_UNROLL_MODE == 2 ? chunk_remainder_mode::unified_ceil : chunk_remainder_mode::fused_floor; - static constexpr bool histogram_use_floor = histogram_rem_mode == chunk_remainder_mode::fused_floor; - static constexpr int histogram_rounds_clamped = + static constexpr int segment_rounds_floor = + clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : 0; + + // Two clamp flavors. The `floor` clamp pairs with an unpredicated main loop over full tiles plus a single + // non-unrolled remainder loop (the chunk helper `for_each_chunk_key_impl` and the copy fast path); the `ceil` clamp + // keeps the whole resident segment inside one tile for fully-predicated loops (the deterministic filter). Each unroll + // is capped at its tuning value; non-clamped (large/unbounded) segments keep the full width, and the guard keeps the + // rounds arithmetic in `int`. + static constexpr int histogram_items_per_thread_clamped = clamp_items_to_segment - ? (histogram_use_floor ? static_cast(static_max_segment_size / threads_per_block) : segment_rounds_ceil) + ? ::cuda::std::clamp(segment_rounds_floor, 1, histogram_items_per_thread) : histogram_items_per_thread; - static constexpr int histogram_items_per_thread_clamped = - ::cuda::std::clamp(histogram_rounds_clamped, 1, histogram_items_per_thread); + // `ceil`-clamped for the deterministic filter's fully-predicated loops; `floor`-clamped for the non-deterministic + // path, which reuses the chunk helper. static constexpr int tie_break_items_per_thread_clamped = clamp_items_to_segment ? ::cuda::std::clamp(segment_rounds_ceil, 1, tie_break_items_per_thread) : tie_break_items_per_thread; + static constexpr int tie_break_items_per_thread_floor_clamped = + clamp_items_to_segment + ? ::cuda::std::clamp(segment_rounds_floor, 1, tie_break_items_per_thread) + : tie_break_items_per_thread; static constexpr int copy_items_per_thread_clamped = - clamp_items_to_segment ? ::cuda::std::clamp(segment_rounds_ceil, 1, copy_items_per_thread) : copy_items_per_thread; + clamp_items_to_segment ? ::cuda::std::clamp(segment_rounds_floor, 1, copy_items_per_thread) : copy_items_per_thread; static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); + static_assert(tie_break_items_per_thread_floor_clamped >= 1, + "tie_break_items_per_thread_floor_clamped must be positive"); static_assert(copy_items_per_thread_clamped >= 1, "copy_items_per_thread_clamped must be positive"); static constexpr int num_buckets = 1 << bits_per_pass; @@ -610,125 +587,59 @@ struct agent_batched_topk_cluster } // Apply `apply(key, local)` to each key of a contiguous chunk (`local` is the key's strided lane index in - // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each tile is loaded into registers by - // one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: + // `[0, chunk_count)`), processing tiles of `Unroll * threads_per_block` keys. Each full tile is loaded into registers + // by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: // `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would interleave each load - // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped items-per-thread. - // `RemMode` selects how the chunk is covered (see `chunk_remainder_mode`): - // `split_ceil` (default): unpredicated main tiles + a split (load wave/apply wave) unrolled, predicated remainder. - // `fused_floor` : unpredicated main tiles + a single non-unrolled fused remainder loop. - // `unified_ceil` : no separate remainder -- one outer tile loop whose split, unrolled load/apply waves are - // predicated and cover the partial last tile too (a per-step bound check on every tile). - template + // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped (floor) items per + // thread, so the sub-tile remainder is handled by a single non-unrolled fused block-stride loop bounded by the count. + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const { - constexpr int tile = Unroll * threads_per_block; - const int tid = static_cast(threadIdx.x); + constexpr int tile = Unroll * threads_per_block; + const int tid = static_cast(threadIdx.x); + const int full_tiles = ::cuda::round_down(chunk_count, tile); - if constexpr (RemMode == chunk_remainder_mode::unified_ceil) + _CCCL_PRAGMA_NOUNROLL() + for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) { - _CCCL_PRAGMA_NOUNROLL() - for (int tile_base = 0; tile_base < chunk_count; tile_base += tile) + key_t regs[Unroll]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) { - key_t regs[Unroll]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = tile_base + t * threads_per_block + tid; - if (local < chunk_count) - { - regs[t] = data[local]; - } - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = tile_base + t * threads_per_block + tid; - if (local < chunk_count) - { - apply(regs[t], local); - } - } + regs[t] = data[tile_base + t * threads_per_block + tid]; } - } - else - { - const int full_tiles = ::cuda::round_down(chunk_count, tile); - - _CCCL_PRAGMA_NOUNROLL() - for (int tile_base = 0; tile_base < full_tiles; tile_base += tile) + _CCCL_PRAGMA_UNROLL_FULL() + for (int t = 0; t < Unroll; ++t) { - key_t regs[Unroll]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - regs[t] = data[tile_base + t * threads_per_block + tid]; - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = tile_base + t * threads_per_block + tid; - apply(regs[t], local); - } + const int local = tile_base + t * threads_per_block + tid; + apply(regs[t], local); } + } - if (full_tiles < chunk_count) - { - if constexpr (RemMode == chunk_remainder_mode::split_ceil) - { - key_t regs[Unroll]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) - { - regs[t] = data[local]; - } - } - _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) - { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) - { - apply(regs[t], local); - } - } - } - else // fused_floor: single non-unrolled fused loop, bound check per step - { - _CCCL_PRAGMA_NOUNROLL() - for (int t = 0; t < Unroll; ++t) - { - const int local = full_tiles + t * threads_per_block + tid; - if (local < chunk_count) - { - apply(data[local], local); - } - } - } - } + // Sub-tile remainder + _CCCL_PRAGMA_NOUNROLL() + for (int local = full_tiles + tid; local < chunk_count; local += threads_per_block) + { + apply(data[local], local); } } - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const { - for_each_chunk_key_impl( - ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { - f(key); - }); + for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { + f(key); + }); } // Like `for_each_chunk_key`, but also hands `f` each key's segment-local index `base_off + local`, where `base_off` // is the segment-local offset of the chunk's first element. The pair path uses that index to fetch the key's value // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const { - for_each_chunk_key_impl( + for_each_chunk_key_impl( ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { f(key, base_off + static_cast(local)); }); @@ -1202,12 +1113,12 @@ private: // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass their // clamped items-per-thread. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { - agent.template for_each_chunk_key(stage_span(stage, o), f); + agent.template for_each_chunk_key(stage_span(stage, o), f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1226,22 +1137,22 @@ private: // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in // by the BlockLoadToShared pipeline (rather than already resident in SMEM). - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) { - process_pass(static_cast(f), [] {}); + process_pass(static_cast(f), [] {}); } // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. The pair // final filter needs that index to fetch each selected key's value payload from gmem, while still reusing the // overflow keys from the streaming SMEM pipeline (block-load path) instead of re-reading them from gmem. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; - agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1476,22 +1387,19 @@ private: // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). - // `unroll_ic` / `rem_mode_ic` are `integral_constant`s carrying the caller's clamped items-per-thread and its - // remainder mode (the histogram may pass a `floor`/fused mode in the A/B experiment). - const auto fold_edges = [&](auto unroll_ic, auto rem_mode_ic, auto&& apply) { + const auto fold_edges = [&](auto&& apply) { if constexpr (use_block_load_to_shared) { - constexpr int unroll = unroll_ic(); - constexpr chunk_remainder_mode remode = rem_mode_ic(); - if (head_edge_len > 0) + const int tid = static_cast(threadIdx.x); + _CCCL_PRAGMA_NOUNROLL() + for (int local = tid; local < head_edge_len; local += threads_per_block) { - this->template for_each_chunk_key( - {temp_storage.edge_keys, static_cast<::cuda::std::size_t>(head_edge_len)}, apply); + apply(temp_storage.edge_keys[local]); } - if (tail_edge_len > 0) + _CCCL_PRAGMA_NOUNROLL() + for (int local = tid; local < tail_edge_len; local += threads_per_block) { - this->template for_each_chunk_key( - {temp_storage.edge_keys + head_edge_cap, static_cast<::cuda::std::size_t>(tail_edge_len)}, apply); + apply(temp_storage.edge_keys[head_edge_cap + local]); } } }; @@ -1567,7 +1475,7 @@ private: const int stage = static_cast(p % static_cast(prologue)); wait_stage(stage); const int read_len = static_cast(::cuda::std::size(bulk_src(p))); - for_each_chunk_key( + for_each_chunk_key( bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), add_first_pass); read_off += read_len * int{sizeof(key_t)}; @@ -1653,7 +1561,7 @@ private: // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer // sync. - streamer.process_pass(add_first_pass); + streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -1694,7 +1602,7 @@ private: // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value // spilled across the pass loop (see `resident_smem32`). key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key( + for_each_chunk_key( {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } else @@ -1704,7 +1612,7 @@ private: const offset_t chunk_idx = part.global_index(resident_base + p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key( + for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } } @@ -1713,14 +1621,12 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. - streamer.process_pass(add_hist, fold_resident_hist); + streamer.process_pass(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's `hist[last_bucket]` (the deterministic cross-CTA prefix // input) inclusive of its edge candidates. - fold_edges(::cuda::std::integral_constant{}, - ::cuda::std::integral_constant{}, - add_hist); + fold_edges(add_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -2264,7 +2170,7 @@ private: if constexpr (use_block_load_to_shared) { key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key( + for_each_chunk_key( {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } else @@ -2274,16 +2180,14 @@ private: const offset_t chunk_idx = part.global_index(p); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - for_each_chunk_key( + for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } } // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_edges(::cuda::std::integral_constant{}, - ::cuda::std::integral_constant{}, - write_selected); + fold_edges(write_selected); }; - streamer.process_pass(write_selected, fold_resident); + streamer.process_pass(write_selected, fold_resident); } else { @@ -2390,7 +2294,7 @@ private: // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. - streamer.process_pass_indexed(write_selected_idx, fold_resident); + streamer.process_pass_indexed(write_selected_idx, fold_resident); } } @@ -2403,11 +2307,7 @@ private: cluster_or_block_sync(cluster, single_cta); } - // Copies an entire segment `input[i] -> output[i]` (keys, plus values for pairs) for the select-all fast path - // (`k == segment_size`); top-k output order is unspecified, so a straight copy is valid. The whole launched cluster - // grid-strides the segment with no chunk partition or cluster barriers -- each CTA copies its part and exits. Each - // thread issues all `copy_items` loads of a step before its stores to keep many requests in flight (ILP/MLP). Input - // and output bases are mutually unaligned in general, so we use per-element (not vectorized/TMA) coalesced accesses. + // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all( num_segments_val_t segment_id, segment_size_val_t segment_size, @@ -2416,64 +2316,85 @@ private: { constexpr int copy_items = copy_items_per_thread_clamped; const offset_t n = static_cast(segment_size); - const offset_t cluster_tid = cluster_rank * static_cast(blockDim.x) + threadIdx.x; - const offset_t cluster_tids = cluster_blocks * static_cast(blockDim.x); + const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; + const offset_t cluster_tids = cluster_blocks * static_cast(threads_per_block); const offset_t step = cluster_tids * static_cast(copy_items); + const offset_t full_tiles = ::cuda::round_down(n, step); auto keys_in_it = d_key_segments_it[segment_id]; auto keys_out_it = d_key_segments_out_it[segment_id]; - for (offset_t tile = 0; tile < n; tile += step) + // Per-segment value iterators (pairs only), hoisted once like the key iterators and reused by both loops. For + // keys-only the value iterators-of-iterators are null, so the discarded `if constexpr` branch keeps the indexing + // out of those builds (the variables stay unused there). + [[maybe_unused]] auto vals_in_it = [&]() -> value_it_t { + if constexpr (!is_keys_only) + { + return d_value_segments_it[segment_id]; + } + else + { + return value_it_t{}; + } + }(); + [[maybe_unused]] auto vals_out_it = [&]() -> it_value_t { + if constexpr (!is_keys_only) + { + return d_value_segments_out_it[segment_id]; + } + else + { + return it_value_t{}; + } + }(); + + for (offset_t base = 0; base < full_tiles; base += step) { - // `idx[]` stays in `offset_t` (not the possibly-narrower `segment_size_val_t`) so the `idx[j] < n` bound check - // happens before any narrowing; the cast at each indexing site is then safe. offset_t idx[copy_items]; key_t keys[copy_items]; [[maybe_unused]] value_t vals[copy_items]; _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - idx[j] = tile + static_cast(j) * cluster_tids + cluster_tid; + idx[j] = base + static_cast(j) * cluster_tids + cluster_tid; } _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - if (idx[j] < n) - { - keys[j] = keys_in_it[static_cast(idx[j])]; - } + keys[j] = keys_in_it[static_cast(idx[j])]; } if constexpr (!is_keys_only) { - auto vals_in_it = d_value_segments_it[segment_id]; _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - if (idx[j] < n) - { - vals[j] = vals_in_it[static_cast(idx[j])]; - } + vals[j] = vals_in_it[static_cast(idx[j])]; } } _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - if (idx[j] < n) - { - keys_out_it[static_cast(idx[j])] = keys[j]; - } + keys_out_it[static_cast(idx[j])] = keys[j]; } if constexpr (!is_keys_only) { - auto vals_out_it = d_value_segments_out_it[segment_id]; _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - if (idx[j] < n) - { - vals_out_it[static_cast(idx[j])] = vals[j]; - } + vals_out_it[static_cast(idx[j])] = vals[j]; } } } + + // Sub-tile remainder + _CCCL_PRAGMA_NOUNROLL() + for (offset_t idx = full_tiles + cluster_tid; idx < n; idx += cluster_tids) + { + const auto si = static_cast(idx); + keys_out_it[si] = keys_in_it[si]; + if constexpr (!is_keys_only) + { + vals_out_it[si] = vals_in_it[si]; + } + } } _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() From bc4bfd327729ea8ef0ac046ef3b8d2d5d3470cb2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 28 Jun 2026 05:56:35 +0200 Subject: [PATCH 058/126] Improve mbarrier/bulk copy codegen --- cub/cub/agent/agent_batched_topk_cluster.cuh | 69 +++++++++++++++----- 1 file changed, 53 insertions(+), 16 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 8938ed21415..11bcda7dec2 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -55,7 +55,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -313,6 +315,10 @@ struct agent_batched_topk_cluster static constexpr int slot_alignment = smem_layout_t::slot_alignment; static_assert(PipelineStages > 0); + // `load_phase` and `load_mbar_inited` track one bit per stage in a 32-bit mask. + static_assert(PipelineStages <= 32, "PipelineStages must fit in the 32-bit per-stage phase/init masks"); + // `__block_elect_one`'s `elect.sync` + full-mask `__shfl_sync` over warp 0 require whole warps. + static_assert(ThreadsPerBlock % detail::warp_threads == 0, "ThreadsPerBlock must be a multiple of the warp size"); static_assert(CopyItemsPerThread > 0, "copy_items_per_thread must be positive"); static_assert(HistogramItemsPerThread > 0, "histogram_items_per_thread must be positive"); static_assert(TieBreakItemsPerThread > 0, "tie_break_items_per_thread must be positive"); @@ -366,7 +372,7 @@ struct agent_batched_topk_cluster offset_t cand_prefix; typename block_scan_t::TempStorage scan_storage; // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused - // (ping-ponged) across radix passes; initialized once by `init_load_barriers`. + // (ping-ponged) across radix passes; each is initialized once on its first use (see `issue_bulk_copy`). ::cuda::std::uint64_t load_mbar[PipelineStages]; // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap)`, on rank 0) and // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner when `full_slots == 1`), each @@ -667,26 +673,44 @@ struct agent_batched_topk_cluster // bit of the per-thread `load_phase` mask, so the pipeline loops spill nothing per stage. SM 9.0+ only (the agent is // gated behind `NV_PROVIDES_SM_90` in `Process`). - // Init each stage mbarrier with arrival count 1: only the elected thread arrives (registering the tx byte count) and - // the `cp.async.bulk` delivers the matching count, so the phase completes. Call once, followed by a block sync. - _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barriers() + // Initialize stage `s`'s mbarrier on first touch, so each init overlaps the previous stage's in-flight bulk copy + // instead of front-loading them. Idempotent per stage (resident load and streamer share the stages; double-init would + // reset the phase and deadlock). The leader fences the init (generic proxy) ahead of its own `cp.async.bulk` (async + // proxy); `maybe_publish_load_barriers` then publishes the inits to the other threads. + _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barrier(int s) { - if (threadIdx.x == 0) + const ::cuda::std::uint32_t bit = ::cuda::std::uint32_t{1} << s; + if (!(load_mbar_inited & bit)) { - _CCCL_PRAGMA_UNROLL_FULL() - for (int s = 0; s < PipelineStages; ++s) + if (load_leader) { ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[s], 1u); + ::cuda::ptx::fence_proxy_async(); // order the init ahead of this thread's async bulk copy on the same mbarrier } + load_mbar_inited |= bit; // all threads flip the bit so the mask stays uniform across the block + } + } + + // Publish stage mbarriers armed since the `inited_before` snapshot to all threads, before their first `wait_stage`. + // Skips the (uniform) barrier when the priming wave armed no fresh stage. + _CCCL_DEVICE _CCCL_FORCEINLINE void maybe_publish_load_barriers(::cuda::std::uint32_t inited_before) + { + if (load_mbar_inited != inited_before) + { + __syncthreads(); } } // Issue one aligned global->shared (TMA) bulk copy into `dst` on stage `stage`'s mbarrier from the elected thread, // which also arrives with the transaction byte count (an empty copy arrives with zero so the phase still completes). - // Each call must be paired, in issue order per stage, with a matching `wait_stage(stage)`. + // First touch of a stage initializes its mbarrier (see `init_load_barrier`). Each call must be paired, in issue order + // per stage, with a matching `wait_stage(stage)`. _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) { - if (threadIdx.x != 0) + init_load_barrier(stage); + // Only the elected leader (see `load_leader`) drives the mbarrier, for a uniform branch and better mbarrier + // codegen. + if (!load_leader) { return; } @@ -753,6 +777,12 @@ struct agent_batched_topk_cluster // Per-thread mbarrier phase parity, one bit per pipeline stage (see `wait_stage`); the resident load and the // overflow streamer keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. ::cuda::std::uint32_t load_phase{}; + // One bit per stage, set the first time that stage's mbarrier is initialized (see `issue_bulk_copy`). Kept uniform + // across the block (all threads flip the bit) so a priming wave can publish exactly the inits it added with one sync. + ::cuda::std::uint32_t load_mbar_inited{}; + // The single thread elected (once at construction, see `cuda::device::__block_elect_one`) to drive the bulk-copy + // mbarriers, cached so the pipeline reuses it instead of re-electing per copy. The block constructs convergently. + const bool load_leader = ::cuda::device::__block_elect_one(); _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -1045,11 +1075,14 @@ private: // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else // the phase advances twice and a lagging thread misses the flip and spins forever. __syncthreads(); + // Stages not already armed by the resident load are armed on first touch here, then published before consume. + const ::cuda::std::uint32_t inited_before = agent.load_mbar_inited; for (int i = 0; i < p_eff; ++i) { const offset_t o = forward ? static_cast(i) : (m - 1 - static_cast(i)); issue_load(static_cast(o % pe), o); } + agent.maybe_publish_load_barriers(inited_before); primed = true; } @@ -1404,12 +1437,13 @@ private: } }; - if constexpr (use_block_load_to_shared) + // The stage mbarriers are now armed lazily on first use (see `issue_bulk_copy`). The histogram reset and `state` + // were already published by `process_impl`'s sync before `run`, so the block-load path needs no barrier here; the + // generic fallback keeps its pass-start barrier. + if constexpr (!use_block_load_to_shared) { - // Arm the stage barriers once; reused (ping-ponged) by the resident load and the overflow streamer across passes. - init_load_barriers(); + __syncthreads(); } - __syncthreads(); { extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); @@ -1459,13 +1493,16 @@ private: int next_off = 0; // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail - // bulk in the forced case) ends the packed region and its suffix can be appended right after it. + // bulk in the forced case) ends the packed region and its suffix can be appended right after it. First touch + // arms each stage's mbarrier; `maybe_publish_load_barriers` makes the inits visible before the read loop. + const ::cuda::std::uint32_t inited_before = load_mbar_inited; for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } + maybe_publish_load_barriers(inited_before); // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. @@ -1559,8 +1596,8 @@ private: } // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The - // streamer reuses the resident load's stage barriers (no re-init); `wait_stage` provides the producer/consumer - // sync. + // streamer shares the resident load's stage mbarriers: already-armed stages are reused, others are armed on first + // touch during priming. streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); From 2f398aabf549236085b99f51d6997d56546df56c Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 28 Jun 2026 17:30:52 +0200 Subject: [PATCH 059/126] Parallelize mbarrier init Undo the lazy init from the previous commit, use elected thread only for bulk copy where it is needed for good code-gen. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 84 +++++++------------- 1 file changed, 28 insertions(+), 56 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 11bcda7dec2..e262c3093f8 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -57,7 +57,6 @@ #include #include #include -#include #include #include #include @@ -315,9 +314,9 @@ struct agent_batched_topk_cluster static constexpr int slot_alignment = smem_layout_t::slot_alignment; static_assert(PipelineStages > 0); - // `load_phase` and `load_mbar_inited` track one bit per stage in a 32-bit mask. - static_assert(PipelineStages <= 32, "PipelineStages must fit in the 32-bit per-stage phase/init masks"); - // `__block_elect_one`'s `elect.sync` + full-mask `__shfl_sync` over warp 0 require whole warps. + // `load_phase` tracks one parity bit per stage in a 32-bit mask. + static_assert(PipelineStages <= 32, "PipelineStages must fit in the 32-bit per-stage phase mask"); + // `__block_elect_one` elects via a full-warp `__shfl_sync` + `elect.sync`, so the block must be whole warps. static_assert(ThreadsPerBlock % detail::warp_threads == 0, "ThreadsPerBlock must be a multiple of the warp size"); static_assert(CopyItemsPerThread > 0, "copy_items_per_thread must be positive"); static_assert(HistogramItemsPerThread > 0, "histogram_items_per_thread must be positive"); @@ -372,7 +371,7 @@ struct agent_batched_topk_cluster offset_t cand_prefix; typename block_scan_t::TempStorage scan_storage; // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused - // (ping-ponged) across radix passes; each is initialized once on its first use (see `issue_bulk_copy`). + // (ping-ponged) across radix passes; all are initialized once up front by `init_load_barriers`. ::cuda::std::uint64_t load_mbar[PipelineStages]; // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap)`, on rank 0) and // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner when `full_slots == 1`), each @@ -669,47 +668,29 @@ struct agent_batched_topk_cluster // Async bulk-copy pipeline helpers (raw mbarrier + cp.async.bulk via cuda::ptx) // --------------------------------------------------------------------------- // Inlines BlockLoadToShared's internals (one mbarrier per stage, a single elected thread issuing the TMA copy + - // transaction arrival) without its reference member or per-call CommitToken: the only per-stage wait state is one - // bit of the per-thread `load_phase` mask, so the pipeline loops spill nothing per stage. SM 9.0+ only (the agent is - // gated behind `NV_PROVIDES_SM_90` in `Process`). - - // Initialize stage `s`'s mbarrier on first touch, so each init overlaps the previous stage's in-flight bulk copy - // instead of front-loading them. Idempotent per stage (resident load and streamer share the stages; double-init would - // reset the phase and deadlock). The leader fences the init (generic proxy) ahead of its own `cp.async.bulk` (async - // proxy); `maybe_publish_load_barriers` then publishes the inits to the other threads. - _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barrier(int s) + // transaction arrival) without its reference member or per-call CommitToken: per-stage wait state is one bit of the + // per-thread `load_phase` mask, so the pipeline loops spill nothing. SM 9.0+ only (gated by `NV_PROVIDES_SM_90`). + + // Front-load every stage mbarrier before any bulk copy uses it. Init needs no leader or uniform path, so all threads + // init distinct stages in parallel via a block-stride loop. The caller's block `__syncthreads()` orders these + // generic-proxy inits before the first `cp.async.bulk` issue, so no `fence.proxy.async` is needed (cf. + // `BlockLoadToShared`, which front-loads init + `__syncthreads()` likewise). + _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barriers() { - const ::cuda::std::uint32_t bit = ::cuda::std::uint32_t{1} << s; - if (!(load_mbar_inited & bit)) - { - if (load_leader) - { - ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[s], 1u); - ::cuda::ptx::fence_proxy_async(); // order the init ahead of this thread's async bulk copy on the same mbarrier - } - load_mbar_inited |= bit; // all threads flip the bit so the mask stays uniform across the block - } - } - - // Publish stage mbarriers armed since the `inited_before` snapshot to all threads, before their first `wait_stage`. - // Skips the (uniform) barrier when the priming wave armed no fresh stage. - _CCCL_DEVICE _CCCL_FORCEINLINE void maybe_publish_load_barriers(::cuda::std::uint32_t inited_before) - { - if (load_mbar_inited != inited_before) + _CCCL_PRAGMA_NOUNROLL() + for (int s = static_cast(threadIdx.x); s < PipelineStages; s += threads_per_block) { - __syncthreads(); + ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[s], 1u); } } - // Issue one aligned global->shared (TMA) bulk copy into `dst` on stage `stage`'s mbarrier from the elected thread, + // Issue one aligned global->shared (TMA) bulk copy into `dst` on stage `stage`'s mbarrier from the block leader, // which also arrives with the transaction byte count (an empty copy arrives with zero so the phase still completes). - // First touch of a stage initializes its mbarrier (see `init_load_barrier`). Each call must be paired, in issue order + // The stage mbarrier must already be initialized (see `init_load_barriers`). Each call must be paired, in issue order // per stage, with a matching `wait_stage(stage)`. _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) { - init_load_barrier(stage); - // Only the elected leader (see `load_leader`) drives the mbarrier, for a uniform branch and better mbarrier - // codegen. + // Only the block leader (see `load_leader`) drives the mbarrier, for a uniform branch and better mbarrier codegen. if (!load_leader) { return; @@ -777,11 +758,8 @@ struct agent_batched_topk_cluster // Per-thread mbarrier phase parity, one bit per pipeline stage (see `wait_stage`); the resident load and the // overflow streamer keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. ::cuda::std::uint32_t load_phase{}; - // One bit per stage, set the first time that stage's mbarrier is initialized (see `issue_bulk_copy`). Kept uniform - // across the block (all threads flip the bit) so a priming wave can publish exactly the inits it added with one sync. - ::cuda::std::uint32_t load_mbar_inited{}; - // The single thread elected (once at construction, see `cuda::device::__block_elect_one`) to drive the bulk-copy - // mbarriers, cached so the pipeline reuses it instead of re-electing per copy. The block constructs convergently. + // The single block leader (warp 0's elected lane) that drives the bulk copies. Elected once at construction (the + // block constructs convergently) and cached so the pipeline reuses it instead of re-electing per copy. const bool load_leader = ::cuda::device::__block_elect_one(); _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( @@ -1075,14 +1053,11 @@ private: // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else // the phase advances twice and a lagging thread misses the flip and spins forever. __syncthreads(); - // Stages not already armed by the resident load are armed on first touch here, then published before consume. - const ::cuda::std::uint32_t inited_before = agent.load_mbar_inited; for (int i = 0; i < p_eff; ++i) { const offset_t o = forward ? static_cast(i) : (m - 1 - static_cast(i)); issue_load(static_cast(o % pe), o); } - agent.maybe_publish_load_barriers(inited_before); primed = true; } @@ -1437,13 +1412,13 @@ private: } }; - // The stage mbarriers are now armed lazily on first use (see `issue_bulk_copy`). The histogram reset and `state` - // were already published by `process_impl`'s sync before `run`, so the block-load path needs no barrier here; the - // generic fallback keeps its pass-start barrier. - if constexpr (!use_block_load_to_shared) + // Front-load all stage mbarrier inits before any bulk copy issues; the barrier below orders them ahead of the first + // issue (see `init_load_barriers`). The generic fallback has no mbarriers but keeps the same pass-start barrier. + if constexpr (use_block_load_to_shared) { - __syncthreads(); + init_load_barriers(); } + __syncthreads(); { extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); @@ -1493,16 +1468,13 @@ private: int next_off = 0; // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail - // bulk in the forced case) ends the packed region and its suffix can be appended right after it. First touch - // arms each stage's mbarrier; `maybe_publish_load_barriers` makes the inits visible before the read loop. - const ::cuda::std::uint32_t inited_before = load_mbar_inited; + // bulk in the forced case) ends the packed region and its suffix can be appended right after it. for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } - maybe_publish_load_barriers(inited_before); // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. @@ -1596,8 +1568,8 @@ private: } // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The - // streamer shares the resident load's stage mbarriers: already-armed stages are reused, others are armed on first - // touch during priming. + // streamer reuses the resident load's stage mbarriers (all front-loaded at `run` entry); `wait_stage` provides + // the producer/consumer sync. streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); From aa333410103563f178445c748f944190e27ab159 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 29 Jun 2026 19:18:45 +0200 Subject: [PATCH 060/126] Fix deterministic filter not using pipeline --- cub/cub/agent/agent_batched_topk_cluster.cuh | 297 +++++++++++++------ 1 file changed, 212 insertions(+), 85 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index e262c3093f8..fec442e2f3b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -64,6 +64,7 @@ #include #include #include +#include #include #include #include @@ -313,6 +314,13 @@ struct agent_batched_topk_cluster static constexpr int load_align_items = smem_layout_t::load_align_items; static constexpr int slot_alignment = smem_layout_t::slot_alignment; + // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a + // time: clamp items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit + // to <= chunk granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. + static constexpr int tie_break_items_streamed = + ::cuda::std::clamp(chunk_items / threads_per_block, 1, tie_break_items_per_thread_clamped); + static_assert(tie_break_items_streamed >= 1, "tie_break_items_streamed must be positive"); + static_assert(PipelineStages > 0); // `load_phase` tracks one parity bit per stage in a 32-bit mask. static_assert(PipelineStages <= 32, "PipelineStages must fit in the 32-bit per-stage phase mask"); @@ -952,8 +960,8 @@ private: // Stage mbarriers are shared with the resident load (`agent.temp_storage.load_mbar`); stage `stage` targets slot // `stream_slot_base + stage`. `inflight_mask` bit `stage` is set only while a copy is in flight (issued, not yet // waited). The slot/stage mapping is fixed, so the read span is recomputed on demand by `stage_span` rather than - // held in a spillable per-stage array; the only per-stage state is one bit each of `inflight_mask` and - // `load_phase`. + // held in a spillable per-stage array; the only per-stage state is one bit each of `inflight_mask` (here) and the + // agent's `load_phase` parity. ::cuda::std::uint32_t inflight_mask = 0; _CCCL_DEVICE _CCCL_FORCEINLINE overflow_streamer( @@ -1023,15 +1031,20 @@ private: {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), static_cast(split.bulk)}); } - // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` currently resident in - // the streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from - // gmem (generic fallback). `mid()` is invoked exactly once, after the prefetch loads for this pass's first reload - // wave (the first `p_eff` visits) have been issued but before they are waited on, so the caller's resident-chunk - // work overlaps those in-flight bulk copies. The two public entry points (`process_pass` / `process_pass_indexed`) - // only differ in whether the applied callable receives the key alone or the key plus its segment-local index, so - // they share this loop verbatim. `mid` must be uniform across the block and contain no unmatched block barrier. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid) + // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` resident in streaming + // slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from gmem + // (fallback). The two entry points (`process_pass` / `process_pass_indexed`) differ only in whether the callable + // also gets the segment-local index, so share this loop. `mid()` runs once on a full pass -- after the first reload + // wave (`p_eff` visits) is issued but before it is waited on, overlapping the caller's resident-chunk work with + // those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with no unmatched barrier. + // `should_continue()` is polled after each consumed chunk; returning false breaks the stream so the final filter + // bails once the top-k is placed. It must be block-uniform and barrier-free (evaluated after `block_apply`'s + // syncs). The histogram and non-deterministic filter pass an always-true predicate (folding back to the + // unconditional loop). An early break can leave prefetches in flight, so the pass drains the remaining stages + // before returning (a full pass ends with an empty `inflight_mask`, so the drain is a no-op). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) { if (overflow_chunks == 0) { @@ -1063,8 +1076,10 @@ private: // Consume overflow visit `i`: wait for its slot, fold its keys via `block_apply`, then prefetch the chunk // `p_eff` visits ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite - // the data the block was just reading). - const auto consume = [&](offset_t i) { + // the data the block was just reading). Returns false once `block_apply` has set `stop_now` -- polled before + // the prefetch so we never launch a copy we would only drain again; the up-to-`p_eff - 1` prefetches already in + // flight (from earlier visits or priming) are drained after the loop. + const auto consume = [&](offset_t i) -> bool { const offset_t o = forward ? i : (m - 1 - i); const int stage = static_cast(o % pe); if (inflight_mask & (::cuda::std::uint32_t{1} << stage)) @@ -1074,6 +1089,11 @@ private: } block_apply(stage, o); + if (!should_continue()) + { + return false; + } + const offset_t ni = i + pe; if (ni < m) { @@ -1081,23 +1101,46 @@ private: __syncthreads(); issue_load(stage, no); } + return true; }; // Phase 1: consume the first `p_eff` visits (the chunks reused from the previous pass, already resident in the // streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. + bool stop = false; const offset_t split = (::cuda::std::min) (pe, m); for (offset_t i = 0; i < split; ++i) { - consume(i); + if (!consume(i)) + { + stop = true; + break; + } } - // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. - mid(); + if (!stop) + { + // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. + mid(); - // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). - for (offset_t i = split; i < m; ++i) + // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). + for (offset_t i = split; i < m; ++i) + { + if (!consume(i)) + { + break; + } + } + } + + // Drain prefetches still in flight before returning: an early break leaves outstanding bulk copies whose + // mbarriers were never waited, and they must complete before the block can exit (their slots are never read). + // `inflight_mask` is block-uniform (set/cleared under uniform control flow), so the trip count and each + // collective `wait_stage` are uniform across the block. + while (inflight_mask != ::cuda::std::uint32_t{0}) { - consume(i); + const int drain_stage = __ffs(static_cast(inflight_mask)) - 1; + agent.wait_stage(drain_stage); + inflight_mask &= ~(::cuda::std::uint32_t{1} << drain_stage); } forward = !forward; } @@ -1113,6 +1156,10 @@ private: const offset_t chunk_idx = chunk_index_of(o); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); generic_apply(chunk); + if (!should_continue()) + { + break; + } } forward = !forward; } @@ -1140,7 +1187,10 @@ private: } } }, - static_cast(mid)); + static_cast(mid), + [] { + return true; + }); } // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in @@ -1175,7 +1225,10 @@ private: } } }, - static_cast(mid)); + static_cast(mid), + [] { + return true; + }); } }; @@ -1787,13 +1840,33 @@ private: tie_active = running < static_cast(num_back); } - // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * items`. Front keys go - // via `out_cnt`; surplus ties get a BlockScan-exclusive rank (seeded by `running`) and, if `rank < num_back`, are - // written reversed at `block_keys_out[k - 1 - rank]`. `running` carries across tiles and regions. `get_idx(pos)` - // gives the segment-local index, used only to load the value payload (compiled out in keys-only builds). - auto process_flat = [&](auto get_key, auto get_idx, int count) { - constexpr int items = tie_break_items_per_thread_clamped; + // Uniform early-exit predicate: true once this CTA's ties are placed and all selected are placed cluster-wide + // (the segment's whole top-k is done). Used between regions (skip the overflow stream) and per tile in + // `process_tiles` (bail the instant a committed stream has placed everything). + auto should_stop = [&]() -> bool { + if (tie_active) + { + return false; + } + const out_offset_t out_now = *static_cast(&leader_state->out_cnt); + return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; + }; + // Sticky, block-uniform flag set by `process_tiles` when `should_stop()` first holds. It short-circuits the + // remaining tiles, the overflow stream (via `run_pass`'s `should_continue`), and every later region. + bool stop_now = false; + + // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * Items` (`Items` passed + // as an `integral_constant`). Front keys go via `out_cnt`; surplus ties get a BlockScan-exclusive rank (seeded by + // `running`) and, if `rank < num_back`, are written reversed at `block_keys_out[k - 1 - rank]`. `running` carries + // across tiles and regions. `get_idx(pos)` gives the segment-local index, used only to load the value payload + // (compiled out in keys-only builds). After each tile we poll `should_stop()` and bail once the top-k is placed. + auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count) { + constexpr int items = items_ic(); constexpr int tile = threads_per_block * items; + if (stop_now) + { + return; + } for (int tile_base = 0; tile_base < count; tile_base += tile) { key_t keys[items]; @@ -1862,18 +1935,21 @@ private: // The next tile reuses `scan_storage`; order this tile's reads against the next ExclusiveSum's writes. __syncthreads(); } + + // Bail the instant the whole top-k is placed (this CTA's ties and all cluster-wide selected). Uniform. + if (should_stop()) + { + stop_now = true; + break; + } } }; - // Uniform early-exit between regions: stop once this CTA's ties are placed and all selected are placed - // cluster-wide. Lets the common prefer-smallest case skip the (expensive) overflow re-stream entirely. - auto should_stop = [&]() -> bool { - if (tie_active) - { - return false; - } - const out_offset_t out_now = *static_cast(&leader_state->out_cnt); - return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; + // Full-unroll flat scan for the resident and boundary-edge regions (one contiguous span; tile may exceed a + // chunk). + auto process_flat = [&](auto get_key, auto get_idx, int count) { + process_tiles( + ::cuda::std::integral_constant{}, get_key, get_idx, count); }; // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, visited @@ -1960,47 +2036,98 @@ private: } }; + // Overflow chunks, in strict segment-index order (the tie-break scan demands it). The histogram leaves the + // streamer ping-ponging for L2 locality, but the final filter is a single ordered pass, so we drive the bulk-copy + // pipeline here in one fixed direction (`forward == !tie_reversed`). On the block-load path this folds each + // landed slot through `process_tiles` (reusing the TMA pipeline) instead of re-reading gmem; the generic fallback + // reads gmem chunk by chunk. The streamed unroll (`tie_break_items_streamed`) keeps a tile within a chunk, and + // `run_pass`'s `should_continue` breaks the stream once `process_tiles` reports the top-k fully placed. + // + // Slot reuse: the histogram leaves its last pass's `p_eff` turn-around chunks resident in the streaming slots, + // which (ping-pong) are exactly the first `p_eff` chunks of the *next* direction. Absent an early stop, exactly + // `num_passes` passes ran, so that direction is `num_passes % 2 == 0` (compile time); when it matches our scan + // direction we keep those chunks (`primed`) and only stream the rest, otherwise we re-prime (overwrite) the slots + // in our direction. An early stop makes the pass count -- hence the resident direction -- a runtime value, but + // then `num_back == 0` and this walk does only the order-independent front writes, so re-priming is always safe. auto process_overflow = [&](bool reversed) { - for (offset_t oo = 0; oo < overflow_count; ++oo) - { - const offset_t o = reversed ? (overflow_count - 1 - oo) : oo; - const offset_t chunk_idx = part.global_index(overflow_base + o); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - // Visit only the aligned bulk: a peeled tail suffix is handled by `process_tail_edge`; every other overflow - // chunk has `bulk == count`. (The generic fallback never peels, so it uses the full count.) - int cc; - if constexpr (use_block_load_to_shared) - { - cc = static_cast(split_chunk(block_keys_base, chunk).bulk); - } - else - { - cc = chunk.count; - } - const offset_t base_off = chunk.offset; - if (reversed) - { - process_flat( - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; - }, - [&](int pos) { - return base_off + static_cast(cc - 1 - pos); - }, - cc); - } - else - { - process_flat( - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(pos))]; - }, - [&](int pos) { - return base_off + static_cast(pos); - }, - cc); - } - } + _CCCL_ASSERT(reversed == tie_reversed, "deterministic overflow walk must run in the tie-break scan direction"); + streamer.forward = !tie_reversed; + constexpr bool resident_dir_forward = (num_passes % 2) == 0; // streamer `forward` at filter entry (no early + // stop) + constexpr bool reuse_resident = (!tie_reversed) == resident_dir_forward; + streamer.primed = reuse_resident && (leader_state->early_stop == ::cuda::std::uint32_t{0}); + + streamer.run_pass( + // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. + // `stage_span` rebuilds the slot pointer from its 32-bit shared address (spill-proof `LDS`) and returns only + // the aligned bulk (a peeled tail suffix is handled by `process_tail_edge`). + [&](int stage, offset_t o) { + const auto span = streamer.stage_span(stage, o); + const key_t* const sm = span.data(); + const int cc = static_cast(span.size()); + const offset_t base_off = get_chunk(streamer.chunk_index_of(o), segment_size_u32, head_items).offset; + constexpr auto items = ::cuda::std::integral_constant{}; + if (reversed) + { + process_tiles( + items, + [&](int pos) { + return sm[cc - 1 - pos]; + }, + [&](int pos) { + return base_off + static_cast(cc - 1 - pos); + }, + cc); + } + else + { + process_tiles( + items, + [&](int pos) { + return sm[pos]; + }, + [&](int pos) { + return base_off + static_cast(pos); + }, + cc); + } + }, + // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). + [&](const auto& chunk) { + const offset_t base_off = chunk.offset; + const int cc = chunk.count; + constexpr auto items = ::cuda::std::integral_constant{}; + if (reversed) + { + process_tiles( + items, + [&](int pos) { + return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; + }, + [&](int pos) { + return base_off + static_cast(cc - 1 - pos); + }, + cc); + } + else + { + process_tiles( + items, + [&](int pos) { + return block_keys_in[static_cast(base_off + static_cast(pos))]; + }, + [&](int pos) { + return base_off + static_cast(pos); + }, + cc); + } + }, + // No interleaved resident work: the deterministic filter folds its resident span separately. + [] {}, + // Break the stream the moment the whole top-k is placed (set by `process_tiles` via `should_stop`). + [&] { + return !stop_now; + }); }; auto process_tail = [&](bool reversed) { @@ -2112,21 +2239,21 @@ private: { // Descending global-index order: peeled suffix tail, resident high-index chunks, streamed low-index overflow, // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the - // gmem-re-read overflow -- mirroring the ascending path. + // overflow re-stream -- mirroring the ascending path. process_tail_edge(true); - if (!should_stop()) + if (!stop_now && !should_stop()) { process_tail(true); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_resident(true); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_overflow(true); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_head_edge(true); } @@ -2134,19 +2261,19 @@ private: else { process_head_edge(false); - if (!should_stop()) + if (!stop_now && !should_stop()) { process_resident(false); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_overflow(false); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_tail(false); } - if (!should_stop()) + if (!stop_now && !should_stop()) { process_tail_edge(false); } From dcc39e420dee726dba865e8383446e71c46c354f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 01:35:59 +0200 Subject: [PATCH 061/126] Improve early exit - Let each CTA accumulate its number of strictly selected keys (let non-leaders scan their histograms to do this without significant additions to the critical path) - Do the inter-CTA scan using both counts (single 64b atomic) - After the scan every CTA knows exactly what it needs for early exit and more importantly does not need to do DSMEM atomics to get output positions. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 367 +++++++++++++------ 1 file changed, 250 insertions(+), 117 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index fec442e2f3b..5cac22d5d3a 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -23,8 +23,10 @@ //! reads `state.kth_bucket` from the leader via DSMEM at the end of each //! pass and folds it into its own local splitter key. //! -//! Output cursors live in the same cluster-shared `state` and are reached the -//! same way (cluster-scope DSMEM atomics). +//! The final filter places each block's output through per-CTA shared-memory +//! atomics, seeded by a single combined 64-bit cross-CTA prefix scan (each +//! block's selected-front and candidate-back base offsets); no cluster-wide +//! output cursor is kept in `state`. #pragma once @@ -94,8 +96,6 @@ struct alignas(16) cluster_topk_state // `kth_key_bits`: every block reads this single digit through DSMEM at the end of each pass and folds it into its own // `kth_key_bits_local` via `set_kth_key_bits`, so the full splitter key is reconstructed locally and never broadcast. ::cuda::std::uint32_t kth_bucket; - OutOffsetT out_cnt; - OutOffsetT out_back_cnt; // Set by the leader after `leader_identify_kth_bucket` whenever the // identified bucket holds exactly `k` items (every candidate is part of // the top-k). Read by every block of the cluster at the end of each radix @@ -233,6 +233,13 @@ struct agent_batched_topk_cluster Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed; static constexpr bool tie_reversed = TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + // Push direction of the combined cross-CTA prefix scan (`combined_prefix_scan`). The leader must be *last* in scan + // order so it derives its own (merged-away) counts from the predecessor sum: deterministic-prefer-smaller puts the + // leader at the last effective rank and scans ascending; every other config (deterministic-prefer-larger and the + // whole non-deterministic path) keeps the leader at rank 0 and scans descending, which makes rank 0 last. Matches the + // `leader_rank` computed in `run`. + static constexpr bool scan_descending = !(deterministic && !tie_reversed); + // The deterministic final scan visits chunks in global-index order and bails early (`should_stop`), so keeping the // *first-visited* chunks resident lets it skip re-reading the streamed overflow. Ascending visits low indices first // (the default low-resident split); descending (prefer-larger) visits high indices first, so we flip the split to @@ -370,13 +377,20 @@ struct agent_batched_topk_cluster // sync the non-leader blocks fold their bucket counts into the leader's // `hist` through DSMEM atomics. `state` is meaningful only in the leader // block; the other blocks reach it exclusively through the DSMEM mapping. - // `cand_prefix` is each block's own exclusive cross-CTA candidate count for the deterministic tie-break; other blocks + // `prefix_pair` packs this block's exclusive cross-CTA scan result (high = `sel_prefix`, low = `cand_prefix`); peers // add into it through DSMEM (`add_remote_prefix`), so it must sit at an identical offset in every block's storage. + // The remaining scalars are block-local: `front_local_cnt`/`back_local_cnt` hand out output slots via SMEM atomics, + // `num_strictly_selected` accumulates this block's strictly-selected count across passes, and `my_candidates` holds + // the last pass's splitter-bucket count. struct _TempStorage { offset_t hist[num_buckets]; state_t state; - offset_t cand_prefix; + ::cuda::std::uint64_t prefix_pair; + offset_t front_local_cnt; + offset_t back_local_cnt; + offset_t num_strictly_selected; + offset_t my_candidates; typename block_scan_t::TempStorage scan_storage; // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused // (ping-ponged) across radix passes; all are initialized once up front by `init_load_barriers`. @@ -387,6 +401,9 @@ struct agent_batched_topk_cluster // Block-local (never reached through DSMEM). key_t edge_keys[2 * load_align_items]; }; + // The `red.add.u64` on `prefix_pair`'s `.shared::cluster` address needs 8-byte alignment; the `uint64_t` member is + // already at an 8-aligned struct offset, so guarding the struct's alignment covers the absolute (and peer) address. + static_assert(alignof(_TempStorage) >= 8, "prefix_pair must be 8-byte aligned for the u64 DSMEM atomic"); // Split point of `edge_keys`: head edge in `[0, head_edge_cap)`, tail edge in `[head_edge_cap, 2 * head_edge_cap)`. static constexpr int head_edge_cap = load_align_items; @@ -860,16 +877,45 @@ private: asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); } - // Adds `v` to the `cand_prefix` of the CTA at cluster rank `target_rank` through DSMEM (mirrors `hist_fold_remote`: - // `mapa` this block's own `cand_prefix` address to `target_rank`, then a cluster-scope `red.add`). Used by the - // deterministic tie-break's exclusive cross-CTA candidate-count scan. - _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, offset_t v) + // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors + // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes + // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. + _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, ::cuda::std::uint64_t v) { const ::cuda::std::uint32_t own = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.cand_prefix)); + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.prefix_pair)); ::cuda::std::uint32_t remote; asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own), "r"(target_rank)); - asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); + asm volatile("red.relaxed.cluster.shared::cluster.add.u64 [%0], %1;" : : "r"(remote), "l"(v) : "memory"); + } + + // Block-local SMEM atomics for the final filter: `front_local_inc`/`back_local_inc` hand out this block's next + // front/back output slot (pre-increment value), `add_local_selected` accumulates its strictly-selected count across + // passes. Each addresses the 32-bit shared slot directly (like `hist_inc`) to dodge the generic-atomic base spill of + // `atomicAdd(&shared)`, at cta scope since every writer of a block's counter lives in that block. + _CCCL_DEVICE _CCCL_FORCEINLINE offset_t front_local_inc() + { + const ::cuda::std::uint32_t addr = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.front_local_cnt)); + offset_t old; + asm volatile("atom.relaxed.cta.shared::cta.add.u32 %0, [%1], 1;" : "=r"(old) : "r"(addr) : "memory"); + return old; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE offset_t back_local_inc() + { + const ::cuda::std::uint32_t addr = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)); + offset_t old; + asm volatile("atom.relaxed.cta.shared::cta.add.u32 %0, [%1], 1;" : "=r"(old) : "r"(addr) : "memory"); + return old; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t v) + { + const ::cuda::std::uint32_t addr = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.num_strictly_selected)); + asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], %1;" : : "r"(addr), "r"(v) : "memory"); } // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram @@ -1255,6 +1301,47 @@ private: } } + // Combined 64-bit exclusive cross-CTA prefix scan over each working CTA's packed `(front_count << 32) | cand_count`. + // Each CTA pushes its counts into every successor's `prefix_pair` in `scan_descending` order; the leader is last and + // pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` + // and untouched until here, so a single post-push barrier suffices. Idle ranks and the leader pass `packed == 0`. + // Returns this CTA's exclusive prefix packed the same way; `single_cta` returns 0. + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( + ::cooperative_groups::cluster_group& cluster, + bool single_cta, + unsigned int cluster_rank, + unsigned int eff_cluster_blocks, + ::cuda::std::uint64_t packed) + { + if (single_cta) + { + return ::cuda::std::uint64_t{0}; + } + if (threadIdx.x == 0 && packed != ::cuda::std::uint64_t{0}) + { + if constexpr (scan_descending) + { + for (unsigned int r = 0; r < cluster_rank; ++r) // lower ranks follow; the leader at rank 0 is last + { + add_remote_prefix(r, packed); + } + } + else + { + // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last + // effective rank is last. + for (unsigned int r = cluster_rank + 1u; r < eff_cluster_blocks; ++r) + { + add_remote_prefix(r, packed); + } + } + } + // TODO(cccl): idle ranks arrive here only to keep this barrier reachable; a sub-cluster mbarrier over the working + // ranks would let them exit (see the pass loop). + cluster_or_block_sync(cluster, single_cta); + return temp_storage.prefix_pair; + } + template _CCCL_DEVICE _CCCL_FORCEINLINE void run(::cooperative_groups::cluster_group& cluster, @@ -1284,10 +1371,6 @@ private: // leader's published `kth_bucket` (see the pass loop), so the full key is never broadcast. key_prefix_t kth_key_bits_local = {}; - // Last splitter bucket folded into `kth_key_bits_local` (the last executed pass's bucket). Used by the - // deterministic tie-break to read this CTA's local candidate count from `hist[last_bucket]`. - [[maybe_unused]] int last_bucket = 0; - // Tracks the highest pass count that actually executed. Without early // stop this stays at `num_passes`; with early stop it captures the pass // at which we broke out so the final filter can construct its identify @@ -1686,8 +1769,8 @@ private: streamer.process_pass(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the - // resident and overflow keys. Keeps every owner's `hist[last_bucket]` (the deterministic cross-CTA prefix - // input) inclusive of its edge candidates. + // resident and overflow keys. Keeps every owner's per-bucket counts (the source of its `num_strictly_selected` + // and `my_candidates`, the cross-CTA scan inputs) inclusive of its edge candidates. fold_edges(add_hist); } @@ -1697,10 +1780,9 @@ private: // is supplied by the `cluster.sync()` further below. __syncthreads(); - // Step 2: non-leader blocks fold their per-bucket values - // (raw counts in the reduce-then-scan path, block-local inclusive - // scans in the scan-then-reduce path) into the leader's `hist` - // via cluster-scope DSMEM atomics (see `hist_fold_remote`). The + // Step 2: non-leader blocks fold their per-bucket raw counts into + // the leader's `hist` via cluster-scope DSMEM atomics (see + // `hist_fold_remote`). The // leader skips this to avoid double-counting its own contribution; // idle ranks (`>= eff_cluster_blocks`) have an all-zero histogram, so // they skip the fold entirely (the loop would only read zeros). @@ -1721,22 +1803,37 @@ private: // over just the active ranks would let them exit and free their SM slots instead of spinning on this barrier. cluster_or_block_sync(cluster, single_cta); - // Step 3: the leader walks the merged `hist` (raw counts in the - // reduce-then-scan path, cluster-wide inclusive scan in the - // scan-then-reduce path) and updates the cluster-shared `state`. - // Subsequent reads (end-of-pass fold, last filter) all observe - // these writes after the next cluster sync. + // Step 3: the leader prefix-scans the merged `hist` (raw counts) and + // updates the cluster-shared `state`. Subsequent reads (end-of-pass + // fold, last filter) all observe these writes after the next cluster sync. + // + // In parallel, each non-leader exclusive-scans its *own* (un-merged) histogram into registers (the leader was + // otherwise the only block doing useful work here). Once the leader publishes `kth_bucket` below, the lane that + // owns it reads its exclusive prefix (= this block's keys strictly above the splitter this pass -> accumulated + // into `num_strictly_selected`) and its raw bucket count (-> `my_candidates`). Keeping the scan in registers lets + // `hist` reset on the normal schedule (the regs survive the reset and the next cluster sync). + offset_t local_prefixes[buckets_per_thread]{}; + offset_t local_hist_vals[buckets_per_thread]{}; if (cluster_rank == leader_rank) { leader_identify_kth_bucket(); } - - if (pass + 1 < num_passes) + else if (!is_idle_rank) { - if (cluster_rank == leader_rank) + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < buckets_per_thread; ++j) { - __syncthreads(); // order the leader's `hist` reads (BlockScan in `leader_identify_kth_bucket`) before resets + const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + local_hist_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; } + block_scan_t(temp_storage.scan_storage).ExclusiveSum(local_hist_vals, local_prefixes); + } + + if (pass + 1 < num_passes) + { + // Unconditional barrier: every working rank just read `hist` (the leader via its BlockScan, non-leaders via the + // scan above), so order all of those reads before any rank resets `hist`. + __syncthreads(); reset_hist(); } // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. @@ -1750,8 +1847,22 @@ private: { const int bucket = static_cast(leader_state->kth_bucket); detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); - last_bucket = bucket; - last_pass = pass + 1; + last_pass = pass + 1; + + // Non-leader: the lane owning the splitter bucket holds its exclusive prefix and raw count in registers from + // the scan above. Accumulate the strictly-selected count and overwrite `my_candidates` with this pass's + // splitter-bucket count (the last pass's value is what the filter reads, however the loop exits). The leader's + // `hist` is merged, so it derives its own counts from the scan total instead (see the filter). + if (cluster_rank != leader_rank) + { + const int owner = bucket / buckets_per_thread; + if (static_cast(threadIdx.x) == owner) + { + const int slot = bucket - owner * buckets_per_thread; + add_local_selected(local_prefixes[slot]); + temp_storage.my_candidates = local_hist_vals[slot]; + } + } } # ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP // Early stop: the leader sets `early_stop` when the splitter bucket holds exactly the remaining `k` candidates, @@ -1794,72 +1905,70 @@ private: if constexpr (deterministic) { - // Deterministic tie-break: select the `num_kth` tied candidates with the smallest (or largest, if `tie_reversed`) - // global indices via a cluster-wide index-ordered scan. Strictly-selected keys go to the front; ties fill the - // back by rank. Early stop means the splitter bucket holds *exactly* the remaining `k` candidates, so all are - // winners (`num_back == 0`, `num_selected == k`) and we skip the scan and its `hist[last_bucket]` read (the - // end-of-pass reset may already have cleared `hist`). Valid even when the early-stop *break* is compiled out. - const out_offset_t num_back = (leader_state->early_stop != ::cuda::std::uint32_t{0}) ? out_offset_t{0} : num_kth; + // Deterministic tie-break: strictly-selected keys fill the front `[0, num_selected)`; the `num_kth` tied + // candidates fill the back by global index rank. The combined scan below gives this block its front base + // (`sel_prefix`) and candidate base (`cand_prefix`); front uses a SMEM atomic offset by `sel_prefix`, back keeps + // its index-ordered BlockScan seeded by `cand_prefix`. Early stop means the splitter bucket holds *exactly* the + // remaining `k`, so all are winners (`num_back == 0`, candidates fold into the front, no back scan). + // + // Cache `early_stop` now, while every block is still tightly coupled to the pass loop's final `cluster.sync` (no + // block has reached `process_tiles`, so the leader's DSMEM is present). The overflow walk reuses this local -- a + // post-scan re-read of `leader_state` could touch an already-returned leader now that the final barrier is gone. + const bool early_stopped = (leader_state->early_stop != ::cuda::std::uint32_t{0}); + const out_offset_t num_back = early_stopped ? out_offset_t{0} : num_kth; const out_offset_t num_selected = k - num_back; // front region; == k on early stop - offset_t running = 0; - bool tie_active = false; - if (num_back != out_offset_t{0}) - { - // Exclusive cross-CTA candidate-count scan: each non-leader adds its candidate count (`hist[last_bucket]`) into - // every later CTA's `cand_prefix` in scan order; the leader is last (see `leader_rank`) and receives the sum. - const offset_t local_count = (cluster_rank == leader_rank) ? offset_t{0} : temp_storage.hist[last_bucket]; - if (threadIdx.x == 0) - { - temp_storage.cand_prefix = 0; - } - // TODO(cccl): idle ranks (`local_count == 0`) arrive only to keep this barrier reachable; a sub-cluster - // mbarrier over the working ranks would let them exit (see the pass loop). - cluster_or_block_sync(cluster, single_cta); - if (threadIdx.x == 0 && local_count != offset_t{0}) - { - if constexpr (tie_reversed) - { - for (unsigned int r = 0; r < cluster_rank; ++r) // descending scan order: lower ranks follow - { - add_remote_prefix(r, local_count); - } - } - else - { - // Ascending scan order: higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing. - for (unsigned int r = cluster_rank + 1u; r < eff_cluster_blocks; ++r) - { - add_remote_prefix(r, local_count); - } - } - } - cluster_or_block_sync(cluster, single_cta); - - running = temp_storage.cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive prefix) - tie_active = running < static_cast(num_back); - } - - // Uniform early-exit predicate: true once this CTA's ties are placed and all selected are placed cluster-wide - // (the segment's whole top-k is done). Used between regions (skip the overflow stream) and per tile in - // `process_tiles` (bail the instant a committed stream has placed everything). + // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final + // `cluster.sync`) block-wide before they feed the scan and `front_count`. + __syncthreads(); + const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; + const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; + // Front count pushed by this block: its strictly-selected count, plus (on early stop) its folded splitter + // candidates. The leader and idle ranks push 0 -- the leader because its merged histogram cannot self-count (it + // derives its own front from the total below), idle ranks because they own nothing. + const offset_t push_front = (num_back == out_offset_t{0}) ? static_cast(my_sel + my_cand) : my_sel; + const ::cuda::std::uint64_t packed = + (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); + const ::cuda::std::uint64_t pp = + combined_prefix_scan(cluster, single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(pp >> 32); + const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so + // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. + const out_offset_t my_front = + (cluster_rank == leader_rank) + ? static_cast(num_selected - static_cast(sel_prefix)) + : static_cast(push_front); + + offset_t running = cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive back prefix) + bool tie_active = (num_back != out_offset_t{0}) && (cand_prefix < static_cast(num_back)); + + // Exact, block-local front-complete flag: `front_local_inc` counts only this block's front placements, so a lane + // seeing `local + 1 >= my_front` means all `my_front` keys are emitted. Seeded true when this block owns no + // front. No remote read, no cross-CTA dependency. + bool front_done_seen = (my_front == out_offset_t{0}); + + // Uniform early-exit predicate: true once this block's ties are placed and all of its strictly-selected keys are + // emitted (it has no further output to contribute). Used between regions (skip the overflow stream); + // `process_tiles` folds the equivalent OR into its per-tile scan-reuse barrier instead of calling this. auto should_stop = [&]() -> bool { if (tie_active) { return false; } - const out_offset_t out_now = *static_cast(&leader_state->out_cnt); - return __syncthreads_or(static_cast(out_now >= num_selected)) != 0; + return __syncthreads_or(static_cast(front_done_seen)) != 0; }; // Sticky, block-uniform flag set by `process_tiles` when `should_stop()` first holds. It short-circuits the // remaining tiles, the overflow stream (via `run_pass`'s `should_continue`), and every later region. bool stop_now = false; - // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * Items` (`Items` passed - // as an `integral_constant`). Front keys go via `out_cnt`; surplus ties get a BlockScan-exclusive rank (seeded by - // `running`) and, if `rank < num_back`, are written reversed at `block_keys_out[k - 1 - rank]`. `running` carries - // across tiles and regions. `get_idx(pos)` gives the segment-local index, used only to load the value payload - // (compiled out in keys-only builds). After each tile we poll `should_stop()` and bail once the top-k is placed. + // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * Items` (`Items` as an + // `integral_constant`). Front keys go via a block-local SMEM atomic offset by `sel_prefix`; surplus ties get a + // BlockScan-exclusive rank (seeded by `running`) and, if `rank < num_back`, are written reversed at + // `block_keys_out[k - 1 - rank]`. `running` carries across tiles and regions. `get_idx(pos)` gives the + // segment-local index for the value payload (compiled out in keys-only builds). Each tile ends with a single + // barrier that both orders `scan_storage` reuse and reduces the done flag, bailing once the top-k is placed. auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count) { constexpr int items = items_ic(); constexpr int tile = threads_per_block * items; @@ -1887,8 +1996,9 @@ private: } } - // Strictly-selected keys go to the front via `out_cnt`; on early stop (`num_back == 0`) the splitter bucket's - // candidates are winners too and fold in here. Exactly `num_selected` are placed in `[0, num_selected)`. + // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`; + // on early stop (`num_back == 0`) the splitter bucket's candidates are winners too and fold in here. The + // per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { @@ -1899,9 +2009,11 @@ private: if (to_front) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + const offset_t local = front_local_inc(); + const out_offset_t out = static_cast(sel_prefix + local); block_keys_out[out] = keys[i]; write_value(out, get_idx(pos)); + front_done_seen = front_done_seen || (local + offset_t{1} >= static_cast(my_front)); } } @@ -1932,12 +2044,13 @@ private: { tie_active = false; } - // The next tile reuses `scan_storage`; order this tile's reads against the next ExclusiveSum's writes. - __syncthreads(); } - // Bail the instant the whole top-k is placed (this CTA's ties and all cluster-wide selected). Uniform. - if (should_stop()) + // One barrier per tile doing double duty: it orders this tile's `scan_storage` reads against the next tile's + // ExclusiveSum writes (needed whenever `tie_active`), and it reduces the block-local done flag so the bail is + // uniform. Bail the instant the whole top-k is placed: this CTA's ties (`!tie_active`) plus a front-complete + // sighting from some lane's atomic (`front_done_seen`). The OR is over registers -- no remote read. + if (__syncthreads_or((!tie_active && front_done_seen) ? 1 : 0) != 0) { stop_now = true; break; @@ -2055,7 +2168,7 @@ private: constexpr bool resident_dir_forward = (num_passes % 2) == 0; // streamer `forward` at filter entry (no early // stop) constexpr bool reuse_resident = (!tie_reversed) == resident_dir_forward; - streamer.primed = reuse_resident && (leader_state->early_stop == ::cuda::std::uint32_t{0}); + streamer.primed = reuse_resident && !early_stopped; streamer.run_pass( // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. @@ -2281,18 +2394,34 @@ private: } else { + // Non-deterministic tie-break: strictly-selected keys fill the front `[0, num_selected)`, the first `num_kth` + // candidates (arrival order) fill the back. The same combined scan as the deterministic path gives this block + // disjoint front/back bases (`sel_prefix`/`cand_prefix`); placement then uses block-local SMEM atomics since + // output order is not preserved. A perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, + // not a correctness one. + __syncthreads(); + const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; + const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; + const ::cuda::std::uint64_t packed = + (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); + const ::cuda::std::uint64_t pp = + combined_prefix_scan(cluster, single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(pp >> 32); + const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + if constexpr (is_keys_only) { auto write_selected = [&](const key_t& key) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); block_keys_out[pos] = key; } else if (res == detail::topk::candidate_class::candidate) { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); if (back_pos < num_kth) { const out_offset_t pos = k - 1 - back_pos; @@ -2327,20 +2456,21 @@ private: } else { - // Pair (key + value) path. Same `out_cnt`/`out_back_cnt` atomics and key reuse as the keys-only path; for each - // written key we additionally load its value from gmem at the segment-local index `seg_idx` and store it at the - // same slot (values are never streamed). Output order is not preserved, as the non-deterministic path permits. + // Pair (key + value) path. Same block-local SMEM-atomic placement (`sel_prefix`/`cand_prefix`) and key reuse as + // the keys-only path; for each written key we additionally load its value from gmem at the segment-local index + // `seg_idx` and store it at the same slot (values are never streamed). Output order is not preserved, as the + // non-deterministic path permits. auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = atomicAdd(&leader_state->out_cnt, out_offset_t{1}); + const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); block_keys_out[pos] = key; write_value(pos, seg_idx); } else if (res == detail::topk::candidate_class::candidate) { - const out_offset_t back_pos = atomicAdd(&leader_state->out_back_cnt, out_offset_t{1}); + const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); if (back_pos < num_kth) { const out_offset_t pos = k - 1 - back_pos; @@ -2434,13 +2564,10 @@ private: } } - // Final cluster barrier: hold every block in the cluster until all DSMEM - // atomics into the leader's state are complete. Without this, a fast - // block (e.g. one whose block_tile is entirely padding) can return while another - // block is still writing to leader-resident memory through DSMEM, which - // surfaces as a "cluster target block not present" exception. A lone CTA has no remote writers, so a block barrier - // suffices (and orders its own final writes). - cluster_or_block_sync(cluster, single_cta); + // No final cluster barrier: both filter paths place output via block-local SMEM atomics into gmem, so the last + // cross-CTA DSMEM access is the combined scan's `prefix_pair` push, already fenced by its post-push + // `cluster.sync()` (and `early_stop` is cached pre-scan). With no shared-memory access to another block after the + // scan, a block can return without risking a "cluster target block not present" fault from a straggler. } // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). @@ -2608,16 +2735,22 @@ private: // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster // state through `leader_state`), but mirroring the writes everywhere - // keeps the scan-then-reduce path's unconditional `state.k` load - // safe under compute-sanitizer. + // keeps every block's unconditional `state.k` load safe under + // compute-sanitizer. if (threadIdx.x == 0) { - temp_storage.state.len = static_cast(segment_size); - temp_storage.state.k = k; - temp_storage.state.kth_bucket = 0; - temp_storage.state.out_cnt = 0; - temp_storage.state.out_back_cnt = 0; - temp_storage.state.early_stop = 0; + temp_storage.state.len = static_cast(segment_size); + temp_storage.state.k = k; + temp_storage.state.kth_bucket = 0; + temp_storage.state.early_stop = 0; + // Front-load every counter the final filter relies on so the first cluster barrier below publishes the zeros to + // all ranks (the combined scan's DSMEM pushes then add into an already-zeroed `prefix_pair`, needing only a + // post-push barrier). `my_candidates` is zeroed too so idle/leader ranks (which never write it) read 0. + temp_storage.prefix_pair = 0; + temp_storage.front_local_cnt = 0; + temp_storage.back_local_cnt = 0; + temp_storage.num_strictly_selected = 0; + temp_storage.my_candidates = 0; } reset_hist(); cluster_or_block_sync(cluster, single_cta); From 4069a2d5cc4a98d43701ad54f426d45371049d38 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 04:24:07 +0200 Subject: [PATCH 062/126] Reduce filter-scans on deterministic path Only one tile needs to actually do a BlockScan, the rest is fine with atomics. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 273 ++++++++++++++----- 1 file changed, 205 insertions(+), 68 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5cac22d5d3a..dc7a3eaa1f7 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -911,6 +911,18 @@ private: return old; } + // Relaxed atomic load of `back_local_cnt`. Used by the `all_below_cta` bail to snapshot how many of this block's + // candidates have been placed, concurrently with the `back_local_inc` atomics of faster lanes: an atomic-vs-atomic + // pair, so it carries no read/write hazard (unlike a plain SMEM read) and needs no extra barrier. + _CCCL_DEVICE _CCCL_FORCEINLINE offset_t back_local_load() + { + const ::cuda::std::uint32_t addr = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)); + offset_t v; + asm volatile("ld.relaxed.cta.shared::cta.u32 %0, [%1];" : "=r"(v) : "r"(addr) : "memory"); + return v; + } + _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t v) { const ::cuda::std::uint32_t addr = @@ -1905,18 +1917,17 @@ private: if constexpr (deterministic) { - // Deterministic tie-break: strictly-selected keys fill the front `[0, num_selected)`; the `num_kth` tied - // candidates fill the back by global index rank. The combined scan below gives this block its front base - // (`sel_prefix`) and candidate base (`cand_prefix`); front uses a SMEM atomic offset by `sel_prefix`, back keeps - // its index-ordered BlockScan seeded by `cand_prefix`. Early stop means the splitter bucket holds *exactly* the - // remaining `k`, so all are winners (`num_back == 0`, candidates fold into the front, no back scan). + // Deterministic tie-break: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); + // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on + // the boundary-crossing tile (see `all_below_cta` and the per-tile back logic). Early stop is not special-cased: + // `total_candidates == num_kth` then makes every CTA `all_below_cta`. // - // Cache `early_stop` now, while every block is still tightly coupled to the pass loop's final `cluster.sync` (no - // block has reached `process_tiles`, so the leader's DSMEM is present). The overflow walk reuses this local -- a - // post-scan re-read of `leader_state` could touch an already-returned leader now that the final barrier is gone. + // Cache `early_stop`/`total_candidates` now, while every block is still tightly coupled to the pass loop's final + // `cluster.sync` -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). const bool early_stopped = (leader_state->early_stop != ::cuda::std::uint32_t{0}); - const out_offset_t num_back = early_stopped ? out_offset_t{0} : num_kth; - const out_offset_t num_selected = k - num_back; // front region; == k on early stop + const offset_t total_candidates = leader_state->len; + const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys + const out_offset_t num_selected = k - num_back; // front region // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final // `cluster.sync`) block-wide before they feed the scan and `front_count`. @@ -1924,16 +1935,24 @@ private: const bool participates = !is_idle_rank && (cluster_rank != leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; - // Front count pushed by this block: its strictly-selected count, plus (on early stop) its folded splitter - // candidates. The leader and idle ranks push 0 -- the leader because its merged histogram cannot self-count (it - // derives its own front from the total below), idle ranks because they own nothing. - const offset_t push_front = (num_back == out_offset_t{0}) ? static_cast(my_sel + my_cand) : my_sel; + // Front count pushed by this block: its strictly-selected count. Candidates always route through the back, so + // nothing folds into the front here. The leader and idle ranks push 0 -- the leader because its merged histogram + // cannot self-count (it derives its own front from the total below), idle ranks because they own nothing. + const offset_t push_front = my_sel; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); const ::cuda::std::uint64_t pp = combined_prefix_scan(cluster, single_cta, cluster_rank, eff_cluster_blocks, packed); const offset_t sel_prefix = static_cast(pp >> 32); const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so + // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A + // CTA is `all_below_cta` when all of its candidates sit at or below the K-boundary (`cand_prefix + my_cand_count + // <= num_back`): every one wins, so the back places them with arrival-order SMEM atomics and skips the + // index-ordered scan. While `tie_active`, a non-`all_below_cta` CTA is the single boundary-crossing (straddling) + // CTA cluster-wide. + const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; + const bool all_below_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. const out_offset_t my_front = @@ -1963,13 +1982,13 @@ private: // remaining tiles, the overflow stream (via `run_pass`'s `should_continue`), and every later region. bool stop_now = false; - // Process a flat span of `count` keys already in scan order, tiled by `threads_per_block * Items` (`Items` as an - // `integral_constant`). Front keys go via a block-local SMEM atomic offset by `sel_prefix`; surplus ties get a - // BlockScan-exclusive rank (seeded by `running`) and, if `rank < num_back`, are written reversed at - // `block_keys_out[k - 1 - rank]`. `running` carries across tiles and regions. `get_idx(pos)` gives the - // segment-local index for the value payload (compiled out in keys-only builds). Each tile ends with a single - // barrier that both orders `scan_storage` reuse and reduces the done flag, bailing once the top-k is placed. - auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count) { + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. Strictly-selected keys + // go to the front via a SMEM atomic (offset by `sel_prefix`); candidates go to the back (see the per-tile logic + // below). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are unresolved -- + // holds the boundary and scans directly. `running` carries across tiles/regions; `get_idx(pos)` is the + // segment-local index for the value payload. Each tile ends with one barrier that reduces the done flag and + // bails. + auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count, bool region_is_terminal) { constexpr int items = items_ic(); constexpr int tile = threads_per_block * items; if (stop_now) @@ -1996,16 +2015,13 @@ private: } } - // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`; - // on early stop (`num_back == 0`) the splitter bucket's candidates are winners too and fold in here. The - // per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. + // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`. + // The per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in + // here -- they always route through the back below, even on early stop. _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - const bool to_front = - valid[i] - && (cls[i] == detail::topk::candidate_class::selected - || (num_back == out_offset_t{0} && cls[i] == detail::topk::candidate_class::candidate)); + const bool to_front = valid[i] && (cls[i] == detail::topk::candidate_class::selected); if (to_front) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; @@ -2017,40 +2033,119 @@ private: } } - // Surplus ties (`num_back > 0`): each candidate gets a BlockScan rank (seeded by `running`), written reversed - // at `block_keys_out[k - 1 - rank]` when `rank < num_back`. Skipped on early stop (`tie_active == false`). + // Snapshot for the `all_below_cta` bail (relaxed atomic load -> no hazard vs the back atomics, no barrier). + // Monotonic and capped at `my_cand_count`, so observing prior/partial increments can only delay the bail, + // never trip it early. Unused by straddling/above CTAs (they bail on `!tie_active`). + const offset_t back_seen = all_below_cta ? back_local_load() : offset_t{0}; + + // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: + // * all_below_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. + // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. + // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the + // counter and, on the crossing tile only, overwrite the arrival slots in index order. if (tie_active) { - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + const bool terminal_tile = region_is_terminal && (tile_base + tile >= count); + if (all_below_cta) { - if (valid[i] && flags[i] != offset_t{0}) + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) + if (valid[i] && flags[i] != offset_t{0}) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); + const offset_t global_rank = cand_prefix + back_local_inc(); + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); + } } } } - running += tile_total; - if (running >= static_cast(num_back)) + else if (terminal_tile) { + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); + } + } + } + running += tile_total; tie_active = false; } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = cand_prefix + back_local_inc(); + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); + } + } + } + // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make + // the counter read race-free and block-uniform. + __syncthreads(); + const offset_t placed = temp_storage.back_local_cnt; + if ((cand_prefix + placed) > static_cast(num_back)) + { + // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered + // winners (identical slot set, different candidate->slot mapping). + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = + static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, get_idx(pos)); + } + } + } + } + running = cand_prefix + placed; + if (running >= static_cast(num_back)) + { + tie_active = false; + } + } } - // One barrier per tile doing double duty: it orders this tile's `scan_storage` reads against the next tile's - // ExclusiveSum writes (needed whenever `tie_active`), and it reduces the block-local done flag so the bail is - // uniform. Bail the instant the whole top-k is placed: this CTA's ties (`!tie_active`) plus a front-complete - // sighting from some lane's atomic (`front_done_seen`). The OR is over registers -- no remote read. - if (__syncthreads_or((!tie_active && front_done_seen) ? 1 : 0) != 0) + // One barrier per tile: reduce the done flag (and order `scan_storage` reuse against the next tile's scan). + // Back-complete is `!tie_active` (straddling/above) or `back_seen >= my_cand_count` (an `all_below_cta` that + // never clears `tie_active`), AND a front-complete sighting. The OR only fires once the whole output is + // placed (no false positive), so all lanes break together. + const bool back_complete = all_below_cta ? (back_seen >= my_cand_count) : !tie_active; + if (__syncthreads_or((back_complete && front_done_seen) ? 1 : 0) != 0) { stop_now = true; break; @@ -2060,9 +2155,12 @@ private: // Full-unroll flat scan for the resident and boundary-edge regions (one contiguous span; tile may exceed a // chunk). - auto process_flat = [&](auto get_key, auto get_idx, int count) { - process_tiles( - ::cuda::std::integral_constant{}, get_key, get_idx, count); + auto process_flat = [&](auto get_key, auto get_idx, int count, bool region_is_terminal) { + process_tiles(::cuda::std::integral_constant{}, + get_key, + get_idx, + count, + region_is_terminal); }; // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, visited @@ -2079,6 +2177,25 @@ private: } } + // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop + // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later + // region in the sweep (head/resident/overflow/tail/tail-edge ascending, reversed descending) carries work. A + // forced tail/edge after the last overflow chunk is not special-cased -- a stray direct scan there is still + // correct, just rare. + const bool has_head = head_edge_len > 0; + const bool has_resident = front_count > 0; + const bool has_overflow = overflow_count > offset_t{0}; + const bool has_tail = tail_count > 0; + const bool has_tail_edge = tail_edge_len > 0; + [[maybe_unused]] const bool resident_terminal = + tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail && !has_tail_edge); + [[maybe_unused]] const bool tail_terminal = + tie_reversed ? (!has_resident && !has_overflow && !has_head) : (!has_tail_edge); + [[maybe_unused]] const bool head_edge_terminal = + tie_reversed ? true : (!has_resident && !has_overflow && !has_tail && !has_tail_edge); + [[maybe_unused]] const bool tail_edge_terminal = + tie_reversed ? (!has_tail && !has_resident && !has_overflow && !has_head) : true; + // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to // the high-index window under `reverse_residency`). The blocked partition packs the front contiguously, so // element `pos` maps to `front_seg_base + pos`. @@ -2098,7 +2215,8 @@ private: [&](int pos) { return front_seg_base + static_cast(fc - 1 - pos); }, - fc); + fc, + resident_terminal); } else { @@ -2109,7 +2227,8 @@ private: [&](int pos) { return front_seg_base + static_cast(pos); }, - fc); + fc, + resident_terminal); } } else @@ -2123,6 +2242,8 @@ private: const int cc = chunk.count; const offset_t base_off = chunk.offset; key_t* const ck = slot_keys_unpadded(local_slot); + // Generic multi-chunk resident loop: never the terminal-tile fast path (the lazy per-tile boundary + // detection handles any boundary), so pass `false`. if (reversed) { process_flat( @@ -2132,7 +2253,8 @@ private: [&](int pos) { return base_off + static_cast(cc - 1 - pos); }, - cc); + cc, + false); } else { @@ -2143,7 +2265,8 @@ private: [&](int pos) { return base_off + static_cast(pos); }, - cc); + cc, + false); } } } @@ -2160,8 +2283,9 @@ private: // which (ping-pong) are exactly the first `p_eff` chunks of the *next* direction. Absent an early stop, exactly // `num_passes` passes ran, so that direction is `num_passes % 2 == 0` (compile time); when it matches our scan // direction we keep those chunks (`primed`) and only stream the rest, otherwise we re-prime (overwrite) the slots - // in our direction. An early stop makes the pass count -- hence the resident direction -- a runtime value, but - // then `num_back == 0` and this walk does only the order-independent front writes, so re-priming is always safe. + // in our direction. An early stop makes the pass count -- hence the resident direction -- a runtime value, so we + // always re-prime then: it is still safe because on early stop every candidate wins (all CTAs are + // `all_below_cta`), so both the front and back writes here are order-independent. auto process_overflow = [&](bool reversed) { _CCCL_ASSERT(reversed == tie_reversed, "deterministic overflow walk must run in the tie-break scan direction"); streamer.forward = !tie_reversed; @@ -2180,6 +2304,9 @@ private: const int cc = static_cast(span.size()); const offset_t base_off = get_chunk(streamer.chunk_index_of(o), segment_size_u32, head_items).offset; constexpr auto items = ::cuda::std::integral_constant{}; + // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == + // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving + // is one barrier on one tile. if (reversed) { process_tiles( @@ -2190,7 +2317,8 @@ private: [&](int pos) { return base_off + static_cast(cc - 1 - pos); }, - cc); + cc, + false); } else { @@ -2202,7 +2330,8 @@ private: [&](int pos) { return base_off + static_cast(pos); }, - cc); + cc, + false); } }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). @@ -2220,7 +2349,8 @@ private: [&](int pos) { return base_off + static_cast(cc - 1 - pos); }, - cc); + cc, + false); } else { @@ -2232,7 +2362,8 @@ private: [&](int pos) { return base_off + static_cast(pos); }, - cc); + cc, + false); } }, // No interleaved resident work: the deterministic filter folds its resident span separately. @@ -2265,7 +2396,8 @@ private: [&](int pos) { return tail_seg_base + static_cast(tc - 1 - pos); }, - tc); + tc, + tail_terminal); } else { @@ -2276,7 +2408,8 @@ private: [&](int pos) { return tail_seg_base + static_cast(pos); }, - tc); + tc, + tail_terminal); } } }; @@ -2298,7 +2431,8 @@ private: [&](int pos) { return static_cast(hc - 1 - pos); }, - hc); + hc, + head_edge_terminal); } else { @@ -2309,7 +2443,8 @@ private: [&](int pos) { return static_cast(pos); }, - hc); + hc, + head_edge_terminal); } } }; @@ -2332,7 +2467,8 @@ private: [&](int pos) { return tail_base + static_cast(tc - 1 - pos); }, - tc); + tc, + tail_edge_terminal); } else { @@ -2343,7 +2479,8 @@ private: [&](int pos) { return tail_base + static_cast(pos); }, - tc); + tc, + tail_edge_terminal); } } }; From 9fd05c3361012587ee27cff31b7a71234caa7022 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 04:46:45 +0200 Subject: [PATCH 063/126] Parallelize inter-CTA scan atomics --- cub/cub/agent/agent_batched_topk_cluster.cuh | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index dc7a3eaa1f7..29a0fef0abd 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1318,6 +1318,10 @@ private: // pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` // and untouched until here, so a single post-push barrier suffices. Idle ranks and the leader pass `packed == 0`. // Returns this CTA's exclusive prefix packed the same way; `single_cta` returns 0. + // + // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, + // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All + // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( ::cooperative_groups::cluster_group& cluster, bool single_cta, @@ -1329,11 +1333,12 @@ private: { return ::cuda::std::uint64_t{0}; } - if (threadIdx.x == 0 && packed != ::cuda::std::uint64_t{0}) + if (packed != ::cuda::std::uint64_t{0}) { if constexpr (scan_descending) { - for (unsigned int r = 0; r < cluster_rank; ++r) // lower ranks follow; the leader at rank 0 is last + _CCCL_PRAGMA_NOUNROLL() + for (unsigned int r = threadIdx.x; r < cluster_rank; r += threads_per_block) // lower ranks follow; leader last { add_remote_prefix(r, packed); } @@ -1342,7 +1347,8 @@ private: { // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. - for (unsigned int r = cluster_rank + 1u; r < eff_cluster_blocks; ++r) + _CCCL_PRAGMA_NOUNROLL() + for (unsigned int r = cluster_rank + 1u + threadIdx.x; r < eff_cluster_blocks; r += threads_per_block) { add_remote_prefix(r, packed); } From 431a5f3e8a1d9e5bb98c64759c5c72a2c7898a08 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 13:35:57 +0200 Subject: [PATCH 064/126] Move from pull-based broadcast to push-based --- cub/cub/agent/agent_batched_topk_cluster.cuh | 110 ++++++++++++------- 1 file changed, 71 insertions(+), 39 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 29a0fef0abd..d7800615eba 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -19,8 +19,9 @@ //! a dual role: its own block-private histogram in step 1, then the //! cluster-merged histogram after the second cluster sync. //! 3. The leader's thread 0 prefix-scans `hist`, identifies the bucket of -//! the k-th key, and updates the cluster-shared `state`. Every block -//! reads `state.kth_bucket` from the leader via DSMEM at the end of each +//! the k-th key, and updates the cluster-shared `state`. The leader then +//! broadcasts the packed `state` result into every block's SMEM via DSMEM +//! stores; each block reads its splitter bucket locally at the end of the //! pass and folds it into its own local splitter key. //! //! The final filter places each block's output through per-CTA shared-memory @@ -92,16 +93,24 @@ struct alignas(16) cluster_topk_state OffsetT len; OutOffsetT k; - // Splitter bucket identified by the leader for the current pass. Replaces broadcasting the full multi-word - // `kth_key_bits`: every block reads this single digit through DSMEM at the end of each pass and folds it into its own - // `kth_key_bits_local` via `set_kth_key_bits`, so the full splitter key is reconstructed locally and never broadcast. - ::cuda::std::uint32_t kth_bucket; - // Set by the leader after `leader_identify_kth_bucket` whenever the - // identified bucket holds exactly `k` items (every candidate is part of - // the top-k). Read by every block of the cluster at the end of each radix - // pass through DSMEM. Carried in the cluster-shared state so the value - // survives the cluster sync that ends the current pass. - ::cuda::std::uint32_t early_stop; + // Per-pass leader result, packed so the leader broadcasts it cluster-wide in one naturally-aligned DSMEM store (see + // `broadcast_state_remote`). Low 32 bits: the splitter `kth_bucket`, which every block folds into its own + // `kth_key_bits_local` (so the full splitter key is never broadcast). High 32 bits: the `early_stop` flag, set when + // the splitter bucket holds exactly the remaining `k` (every candidate wins, so further passes cannot change it). + ::cuda::std::uint64_t result_pair; + + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t kth_bucket() const + { + return static_cast<::cuda::std::uint32_t>(result_pair); + } + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool early_stop() const + { + return (result_pair >> 32) != ::cuda::std::uint64_t{0}; + } + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void set_result(::cuda::std::uint32_t bucket, bool stop) + { + result_pair = static_cast<::cuda::std::uint64_t>(bucket) | (static_cast<::cuda::std::uint64_t>(stop) << 32); + } }; // Dynamic-SMEM layout shared by dispatch and the agent. `block_tile_capacity` is the physical per-CTA resident @@ -877,6 +886,18 @@ private: asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); } + // Stores the packed `state.result_pair` into the same field of the CTA at rank `target_rank` through DSMEM (mirrors + // `hist_fold_remote`'s addressing), broadcasting the leader's per-pass result. Single writer per target, so a plain + // store -- ordered against the reads by the next `cluster_or_block_sync` -- suffices (no atomic). + _CCCL_DEVICE _CCCL_FORCEINLINE void broadcast_state_remote(unsigned int target_rank, ::cuda::std::uint64_t packed) + { + const ::cuda::std::uint32_t own = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.state.result_pair)); + ::cuda::std::uint32_t remote; + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own), "r"(target_rank)); + asm volatile("st.relaxed.cluster.shared::cluster.u64 [%0], %1;" : : "r"(remote), "l"(packed) : "memory"); + } + // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. @@ -969,15 +990,10 @@ private: const offset_t new_len = hist_vals[j]; temp_storage.state.len = new_len; temp_storage.state.k = new_k; - // Early-stop opportunity: the bucket holding the k-th key contains - // exactly the remaining `k` items. Every candidate is therefore part - // of the top-k, so subsequent radix passes only redistribute the same - // items across finer buckets without changing the final result. - temp_storage.state.early_stop = - (static_cast(new_len) == new_k) ? ::cuda::std::uint32_t{1} : ::cuda::std::uint32_t{0}; - // Publish only the splitter bucket; every block folds it into its own `kth_key_bits_local` at the end of the - // pass (see the pass loop), so the full splitter key is never broadcast. - temp_storage.state.kth_bucket = static_cast<::cuda::std::uint32_t>(bucket); + // Publish the splitter bucket and the early-stop flag (set when the bucket holds exactly the remaining `k`); + // see `result_pair`. The owning lane writes; the leader then broadcasts it to every rank in the pass loop. + temp_storage.state.set_result( + static_cast<::cuda::std::uint32_t>(bucket), static_cast(new_len) == new_k); } } } @@ -1847,6 +1863,26 @@ private: block_scan_t(temp_storage.scan_storage).ExclusiveSum(local_hist_vals, local_prefixes); } + // Broadcast the leader's result into every other co-resident rank's local `state` (lane `r` -> rank `r`), so the + // consumers below read their own SMEM instead of pulling the leader's word across DSMEM (no remote-read + // round-trip on the critical path, no read hotspot). Push to all `cluster_blocks`, not just working ranks: idle + // ranks skip the fold but still read `early_stop()` for the uniform break. Sync first so the owning lane's + // `set_result` is visible to the pushing threads; `cluster_or_block_sync` below orders the stores against the + // reads. `single_cta` has no peers and skips the push. + if (!single_cta && cluster_rank == leader_rank) + { + __syncthreads(); + const ::cuda::std::uint64_t packed = temp_storage.state.result_pair; + _CCCL_PRAGMA_NOUNROLL() + for (unsigned int r = threadIdx.x; r < cluster_blocks; r += threads_per_block) + { + if (r != leader_rank) + { + broadcast_state_remote(r, packed); + } + } + } + if (pass + 1 < num_passes) { // Unconditional barrier: every working rank just read `hist` (the leader via its BlockScan, non-leaders via the @@ -1857,13 +1893,13 @@ private: // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. cluster_or_block_sync(cluster, single_cta); - // End-of-pass splitter fold. Every working thread reads the leader's just-published `kth_bucket` through DSMEM (a - // single `u32`, ordered by the `cluster.sync()` above) and folds it into its own `kth_key_bits_local`, so the - // full splitter key is reconstructed locally without a broadcast. Idle ranks produce no output, so they skip the - // read (trimming the leader's fan-out); the early-stop check below stays uniform so every block breaks together. + // End-of-pass splitter fold. Every working thread reads its own (leader-broadcast) `kth_bucket` from local SMEM, + // ordered by the `cluster_or_block_sync()` above, and folds it into its own `kth_key_bits_local`, so the full + // splitter key is reconstructed locally without a broadcast. Idle ranks produce no output, so they skip the fold + // (they were still pushed to above, so the early-stop check below stays uniform and every block breaks together). if (!is_idle_rank) { - const int bucket = static_cast(leader_state->kth_bucket); + const int bucket = static_cast(temp_storage.state.kth_bucket()); detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); last_pass = pass + 1; @@ -1883,10 +1919,9 @@ private: } } # ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - // Early stop: the leader sets `early_stop` when the splitter bucket holds exactly the remaining `k` candidates, - // so no further radix refinement can change the result. Every block reads the same flag from the leader and + // Early stop (see `result_pair`): every block reads the same leader-broadcast flag from its local `state` and // breaks together; `last_pass`/`kth_key_bits_local` then match what the original top-of-next-pass break produced. - if (leader_state->early_stop != ::cuda::std::uint32_t{0}) + if (temp_storage.state.early_stop()) { break; } @@ -1930,7 +1965,7 @@ private: // // Cache `early_stop`/`total_candidates` now, while every block is still tightly coupled to the pass loop's final // `cluster.sync` -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). - const bool early_stopped = (leader_state->early_stop != ::cuda::std::uint32_t{0}); + const bool early_stopped = leader_state->early_stop(); const offset_t total_candidates = leader_state->len; const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region @@ -2875,17 +2910,14 @@ private: } const bool single_cta = (eff_cluster_blocks == 1u); - // Every block's thread 0 initializes its local `state`. Only the - // leader's copy is semantically read (non-leaders reach the cluster - // state through `leader_state`), but mirroring the writes everywhere - // keeps every block's unconditional `state.k` load safe under - // compute-sanitizer. + // Every block's thread 0 initializes its local `state`. The leader's copy is canonical (`len`/`k`, reached through + // `leader_state`); each block's own `result_pair` doubles as the mailbox the leader broadcasts into each pass. + // Mirroring the writes everywhere also keeps every block's unconditional `state.k` load sanitizer-safe. if (threadIdx.x == 0) { - temp_storage.state.len = static_cast(segment_size); - temp_storage.state.k = k; - temp_storage.state.kth_bucket = 0; - temp_storage.state.early_stop = 0; + temp_storage.state.len = static_cast(segment_size); + temp_storage.state.k = k; + temp_storage.state.result_pair = 0; // Front-load every counter the final filter relies on so the first cluster barrier below publishes the zeros to // all ranks (the combined scan's DSMEM pushes then add into an already-zeroed `prefix_pair`, needing only a // post-push barrier). `my_candidates` is zeroed too so idle/leader ranks (which never write it) read 0. From 982d48c88094dbbabc6bea00c83d59cde2e5ca06 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 17:31:58 +0200 Subject: [PATCH 065/126] Revert "Move from pull-based broadcast to push-based" This reverts commit 431a5f3e8a1d9e5bb98c64759c5c72a2c7898a08. But keeps the change that packs the broadcasted values into one 64b value. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 100 +++++++++---------- 1 file changed, 45 insertions(+), 55 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index d7800615eba..0f530df595c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -19,10 +19,10 @@ //! a dual role: its own block-private histogram in step 1, then the //! cluster-merged histogram after the second cluster sync. //! 3. The leader's thread 0 prefix-scans `hist`, identifies the bucket of -//! the k-th key, and updates the cluster-shared `state`. The leader then -//! broadcasts the packed `state` result into every block's SMEM via DSMEM -//! stores; each block reads its splitter bucket locally at the end of the -//! pass and folds it into its own local splitter key. +//! the k-th key, and updates the cluster-shared `state`. Every block +//! reads the leader's packed `state.result_pair` (splitter bucket plus +//! early-stop flag) from the leader via DSMEM at the end of each pass and +//! folds the bucket into its own local splitter key. //! //! The final filter places each block's output through per-CTA shared-memory //! atomics, seeded by a single combined 64-bit cross-CTA prefix scan (each @@ -93,19 +93,32 @@ struct alignas(16) cluster_topk_state OffsetT len; OutOffsetT k; - // Per-pass leader result, packed so the leader broadcasts it cluster-wide in one naturally-aligned DSMEM store (see - // `broadcast_state_remote`). Low 32 bits: the splitter `kth_bucket`, which every block folds into its own - // `kth_key_bits_local` (so the full splitter key is never broadcast). High 32 bits: the `early_stop` flag, set when - // the splitter bucket holds exactly the remaining `k` (every candidate wins, so further passes cannot change it). + // Per-pass leader result, packed so every block pulls it from the leader cluster-wide in one naturally-aligned + // 64-bit DSMEM load (instead of two separate reads). Low 32 bits: the splitter `kth_bucket`, which every block folds + // into its own `kth_key_bits_local` (so the full splitter key is never broadcast). High 32 bits: the `early_stop` + // flag, set when the splitter bucket holds exactly the remaining `k` (every candidate wins, so further passes cannot + // change it). ::cuda::std::uint64_t result_pair; + // Decoders that operate on an already-loaded `result_pair`. The hot pass loop reads the leader's word once through + // DSMEM and decodes both halves from that single value, so the splitter fold and the early-stop check never issue a + // second remote load. The instance accessors below delegate here, keeping the bit layout defined in one place. + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static ::cuda::std::uint32_t + kth_bucket_of(::cuda::std::uint64_t packed) + { + return static_cast<::cuda::std::uint32_t>(packed); + } + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static bool early_stop_of(::cuda::std::uint64_t packed) + { + return (packed >> 32) != ::cuda::std::uint64_t{0}; + } [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t kth_bucket() const { - return static_cast<::cuda::std::uint32_t>(result_pair); + return kth_bucket_of(result_pair); } [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool early_stop() const { - return (result_pair >> 32) != ::cuda::std::uint64_t{0}; + return early_stop_of(result_pair); } _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void set_result(::cuda::std::uint32_t bucket, bool stop) { @@ -886,18 +899,6 @@ private: asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); } - // Stores the packed `state.result_pair` into the same field of the CTA at rank `target_rank` through DSMEM (mirrors - // `hist_fold_remote`'s addressing), broadcasting the leader's per-pass result. Single writer per target, so a plain - // store -- ordered against the reads by the next `cluster_or_block_sync` -- suffices (no atomic). - _CCCL_DEVICE _CCCL_FORCEINLINE void broadcast_state_remote(unsigned int target_rank, ::cuda::std::uint64_t packed) - { - const ::cuda::std::uint32_t own = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.state.result_pair)); - ::cuda::std::uint32_t remote; - asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own), "r"(target_rank)); - asm volatile("st.relaxed.cluster.shared::cluster.u64 [%0], %1;" : : "r"(remote), "l"(packed) : "memory"); - } - // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. @@ -990,8 +991,11 @@ private: const offset_t new_len = hist_vals[j]; temp_storage.state.len = new_len; temp_storage.state.k = new_k; - // Publish the splitter bucket and the early-stop flag (set when the bucket holds exactly the remaining `k`); - // see `result_pair`. The owning lane writes; the leader then broadcasts it to every rank in the pass loop. + // Publish the splitter bucket and the early-stop flag (set when the bucket holds exactly the remaining `k`, so + // every candidate is part of the top-k and subsequent radix passes only redistribute the same items across + // finer buckets without changing the result); see `result_pair`. Every block pulls this packed word from the + // leader at the end of the pass and folds the bucket into its own `kth_key_bits_local`, so the full splitter + // key is never broadcast. temp_storage.state.set_result( static_cast<::cuda::std::uint32_t>(bucket), static_cast(new_len) == new_k); } @@ -1863,26 +1867,6 @@ private: block_scan_t(temp_storage.scan_storage).ExclusiveSum(local_hist_vals, local_prefixes); } - // Broadcast the leader's result into every other co-resident rank's local `state` (lane `r` -> rank `r`), so the - // consumers below read their own SMEM instead of pulling the leader's word across DSMEM (no remote-read - // round-trip on the critical path, no read hotspot). Push to all `cluster_blocks`, not just working ranks: idle - // ranks skip the fold but still read `early_stop()` for the uniform break. Sync first so the owning lane's - // `set_result` is visible to the pushing threads; `cluster_or_block_sync` below orders the stores against the - // reads. `single_cta` has no peers and skips the push. - if (!single_cta && cluster_rank == leader_rank) - { - __syncthreads(); - const ::cuda::std::uint64_t packed = temp_storage.state.result_pair; - _CCCL_PRAGMA_NOUNROLL() - for (unsigned int r = threadIdx.x; r < cluster_blocks; r += threads_per_block) - { - if (r != leader_rank) - { - broadcast_state_remote(r, packed); - } - } - } - if (pass + 1 < num_passes) { // Unconditional barrier: every working rank just read `hist` (the leader via its BlockScan, non-leaders via the @@ -1893,13 +1877,15 @@ private: // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. cluster_or_block_sync(cluster, single_cta); - // End-of-pass splitter fold. Every working thread reads its own (leader-broadcast) `kth_bucket` from local SMEM, - // ordered by the `cluster_or_block_sync()` above, and folds it into its own `kth_key_bits_local`, so the full - // splitter key is reconstructed locally without a broadcast. Idle ranks produce no output, so they skip the fold - // (they were still pushed to above, so the early-stop check below stays uniform and every block breaks together). + // End-of-pass splitter fold. Every block pulls the leader's just-published `result_pair` once through DSMEM (a + // single naturally-aligned `u64`, ordered by the `cluster_or_block_sync()` above) and decodes both halves from + // that one load: working threads fold the `kth_bucket` digit into their own `kth_key_bits_local` (so the full + // splitter key is reconstructed locally without a broadcast), and every block -- including idle ranks -- reuses + // the same value for the uniform early-stop check below, so the whole cluster breaks together. + const ::cuda::std::uint64_t pass_result = leader_state->result_pair; if (!is_idle_rank) { - const int bucket = static_cast(temp_storage.state.kth_bucket()); + const int bucket = static_cast(state_t::kth_bucket_of(pass_result)); detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); last_pass = pass + 1; @@ -1919,9 +1905,11 @@ private: } } # ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP - // Early stop (see `result_pair`): every block reads the same leader-broadcast flag from its local `state` and - // breaks together; `last_pass`/`kth_key_bits_local` then match what the original top-of-next-pass break produced. - if (temp_storage.state.early_stop()) + // Early stop: the leader sets `early_stop` when the splitter bucket holds exactly the remaining `k` candidates, + // so no further radix refinement can change the result. Every block decodes the same flag from the `pass_result` + // it just loaded and breaks together; `last_pass`/`kth_key_bits_local` then match what the original + // top-of-next-pass break produced. + if (state_t::early_stop_of(pass_result)) { break; } @@ -2910,9 +2898,11 @@ private: } const bool single_cta = (eff_cluster_blocks == 1u); - // Every block's thread 0 initializes its local `state`. The leader's copy is canonical (`len`/`k`, reached through - // `leader_state`); each block's own `result_pair` doubles as the mailbox the leader broadcasts into each pass. - // Mirroring the writes everywhere also keeps every block's unconditional `state.k` load sanitizer-safe. + // Every block's thread 0 initializes its local `state`. Only the + // leader's copy is semantically read (non-leaders reach the cluster + // state through `leader_state`), but mirroring the writes everywhere + // keeps every block's unconditional `state.k` load safe under + // compute-sanitizer. if (threadIdx.x == 0) { temp_storage.state.len = static_cast(segment_size); From ac0c6a16e35abd7736b81d0a63d08c8d28068bdd Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 18:31:10 +0200 Subject: [PATCH 066/126] Subdivide first resident chunk for better ramp-up Big chunk sizes cause less latency hiding otherwise. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 114 +++++++++++++----- .../dispatch_batched_topk_cluster.cuh | 32 +++-- .../tuning/tuning_batched_topk_cluster.cuh | 11 +- 3 files changed, 110 insertions(+), 47 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 0f530df595c..2531efd9d6c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -218,6 +218,7 @@ template = 1, "FirstResidentSlotSubdivision must be positive"); + static constexpr bool first_slot_subdivisible = + chunk_items % FirstResidentSlotSubdivision == 0 + && (chunk_items / FirstResidentSlotSubdivision) * int{sizeof(key_t)} % slot_alignment == 0; + static constexpr int first_resident_slot_subdivision = first_slot_subdivisible ? FirstResidentSlotSubdivision : 1; + static constexpr int first_resident_slot_sub_items = chunk_items / first_resident_slot_subdivision; + // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a // time: clamp items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit // to <= chunk granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. @@ -1610,13 +1623,12 @@ private: if (my_resident_chunks > 0) { // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). - // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor + // Load units are written densely from offset 0 and read back in the same order, so the read cursor // (`read_off`) mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed // `pending_spans` array that would anchor surrounding state to local memory. Every chunk begins on a // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except a resident suffix // tail, whose bulk is loaded here and whose suffix is appended afterwards. The read span is rebuilt from a // rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot demote the reads to `LD`. - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); // Resident slot -> rank-local chunk index. `resident_base + slot` (identity unless `reverse_residency` shifts // the resident window to the high-index chunks), except the last slot holds a forced-resident tail. @@ -1637,42 +1649,80 @@ private: return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; }; - // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives - // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. + // Bulks are densely packed from the start of the block_tile (no reserved front gap: the head prefix lives in + // `edge_keys`), so the write cursor starts at offset 0 and the last unit's bulk ends the packed region with + // its suffix appended right after. int next_off = 0; - // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail - // bulk in the forced case) ends the packed region and its suffix can be appended right after it. - for (int stage = 0; stage < prologue; ++stage) - { - const auto src = bulk_src(static_cast(stage)); - issue_bulk_copy(stage, key_slots + next_off, src); - next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - } - - // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in - // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. - int read_off = 0; - for (offset_t p = 0; p < my_resident_chunks; ++p) - { - const int stage = static_cast(p % static_cast(prologue)); - wait_stage(stage); - const int read_len = static_cast(::cuda::std::size(bulk_src(p))); - for_each_chunk_key( - bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), - add_first_pass); - read_off += read_len * int{sizeof(key_t)}; - - const offset_t next_p = p + static_cast(prologue); - if (next_p < my_resident_chunks) + // Pipeline `n_units` densely-packed load units through the shared `prologue` mbarriers: issue the prologue, + // then consume each unit (folding its keys into the first-pass histogram) and prefetch the unit `prologue` + // ahead into the slot just freed. `unit_span(u)` returns unit `u`'s gmem span, written/consumed in index + // order. + const auto run_load = [&](int n_units, auto&& unit_span) { + const int prologue = (::cuda::std::min) (PipelineStages, n_units); + for (int stage = 0; stage < prologue; ++stage) { - const auto src = bulk_src(next_p); - // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads - // leave the wait above would advance the phase twice, stranding a lagging waiter forever. - __syncthreads(); + const auto src = unit_span(stage); issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } + int read_off = 0; + for (int u = 0; u < n_units; ++u) + { + const int stage = u % prologue; + wait_stage(stage); + const int read_len = static_cast(::cuda::std::size(unit_span(u))); + for_each_chunk_key( + bulk_span( + {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), + add_first_pass); + read_off += read_len * int{sizeof(key_t)}; + + const int next_u = u + prologue; + if (next_u < n_units) + { + const auto src = unit_span(next_u); + // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads + // leave the wait above would advance the phase twice, stranding a lagging waiter forever. + __syncthreads(); + issue_bulk_copy(stage, key_slots + next_off, src); + next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + } + } + }; + + const int resident_units = static_cast(my_resident_chunks); + if constexpr (first_resident_slot_subdivision == 1) + { + // One load unit per resident chunk (the un-split load). + run_load(resident_units, [&](int u) { + return bulk_src(static_cast(u)); + }); + } + else + { + // Split only the first chunk into `<= first_resident_slot_subdivision` sub-chunks; chunks 1+ stay whole. A + // short first bulk (also the global tail) needs fewer sub-chunks, its final one carrying the aligned + // remainder, so the count is clamped to what the bulk covers. + const int sub_items = first_resident_slot_sub_items; // local copy: avoids ODR-using the static member by + // ref + const auto first_bulk = bulk_src(offset_t{0}); + const int first_len = static_cast(::cuda::std::size(first_bulk)); + const int first_subs = (::cuda::std::max) (1, ::cuda::ceil_div(first_len, sub_items)); + const int n_units = (first_subs - 1) + resident_units; + run_load(n_units, [&](int u) -> ::cuda::std::span { + if (u < first_subs) + { + const int off = u * sub_items; + const int len = (::cuda::std::min) (sub_items, first_len - off); + if (len == 0) + { + return {}; // empty first bulk (sole resident chunk is a short global tail): avoid `nullptr + 0` UB + } + return {first_bulk.data() + off, static_cast<::cuda::std::size_t>(len)}; + } + return bulk_src(static_cast(u - first_subs + 1)); + }); } // Tail suffix of a resident (non-peeled) suffix tail: append right after the last (tail) bulk. Nothing diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index d234e3f65f2..6829ffc1810 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -88,6 +88,7 @@ template ; using key_t = it_value_t; @@ -361,6 +367,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SingleCtaMaxSegmentSize, MinChunksPerCta, CopyItemsPerThread, + FirstResidentSlotSubdivision, Determinism, TieBreak, KeyInputItItT, @@ -439,6 +446,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SingleCtaMaxSegmentSize, MinChunksPerCta, CopyItemsPerThread, + FirstResidentSlotSubdivision, Determinism, TieBreak, KeyInputItItT, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 5bf9c7fc127..94c043d05e5 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -33,6 +33,7 @@ struct cluster_topk_policy int single_cta_max_segment_size; int min_chunks_per_cta; int copy_items_per_thread; + int first_resident_slot_subdivision; }; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy @@ -41,19 +42,23 @@ struct cluster_topk_policy /*threads_per_block=*/512, /*histogram_items_per_thread=*/8, /*pipeline_stages=*/8, - /*chunk_bytes=*/16 * 1024, + /*chunk_bytes=*/32 * 1024, /*load_align_bytes=*/128, /*bits_per_pass=*/11, /*min_blocks_per_sm=*/1, /*tie_break_items_per_thread=*/8, /*single_cta_max_segment_size=*/8 * 1024, /*min_chunks_per_cta=*/1, - /*copy_items_per_thread=*/8}; + /*copy_items_per_thread=*/8, + /*first_resident_slot_subdivision=*/2}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool { - return policy.chunk_bytes % policy.load_align_bytes == 0; + // Whether `first_resident_slot_subdivision` actually splits the chunk depends on the runtime key type (sub-chunks + // must be whole and `slot_alignment`-aligned), so the agent decides per instantiation and falls back to the un-split + // load when it cannot. Here only its positivity is a policy-level invariant. + return policy.chunk_bytes % policy.load_align_bytes == 0 && policy.first_resident_slot_subdivision >= 1; } static_assert(is_valid_policy(make_policy())); From 94205d161b131346fe1b579f2ee4b56a038a4fda Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 30 Jun 2026 21:53:38 +0200 Subject: [PATCH 067/126] Revert "Subdivide first resident chunk for better ramp-up" This reverts commit ac0c6a16e35abd7736b81d0a63d08c8d28068bdd. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 114 +++++------------- .../dispatch_batched_topk_cluster.cuh | 32 ++--- .../tuning/tuning_batched_topk_cluster.cuh | 11 +- 3 files changed, 47 insertions(+), 110 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2531efd9d6c..0f530df595c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -218,7 +218,6 @@ template = 1, "FirstResidentSlotSubdivision must be positive"); - static constexpr bool first_slot_subdivisible = - chunk_items % FirstResidentSlotSubdivision == 0 - && (chunk_items / FirstResidentSlotSubdivision) * int{sizeof(key_t)} % slot_alignment == 0; - static constexpr int first_resident_slot_subdivision = first_slot_subdivisible ? FirstResidentSlotSubdivision : 1; - static constexpr int first_resident_slot_sub_items = chunk_items / first_resident_slot_subdivision; - // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a // time: clamp items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit // to <= chunk granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. @@ -1623,12 +1610,13 @@ private: if (my_resident_chunks > 0) { // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). - // Load units are written densely from offset 0 and read back in the same order, so the read cursor + // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor // (`read_off`) mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed // `pending_spans` array that would anchor surrounding state to local memory. Every chunk begins on a // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except a resident suffix // tail, whose bulk is loaded here and whose suffix is appended afterwards. The read span is rebuilt from a // rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot demote the reads to `LD`. + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); // Resident slot -> rank-local chunk index. `resident_base + slot` (identity unless `reverse_residency` shifts // the resident window to the high-index chunks), except the last slot holds a forced-resident tail. @@ -1649,80 +1637,42 @@ private: return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; }; - // Bulks are densely packed from the start of the block_tile (no reserved front gap: the head prefix lives in - // `edge_keys`), so the write cursor starts at offset 0 and the last unit's bulk ends the packed region with - // its suffix appended right after. + // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives + // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. int next_off = 0; - // Pipeline `n_units` densely-packed load units through the shared `prologue` mbarriers: issue the prologue, - // then consume each unit (folding its keys into the first-pass histogram) and prefetch the unit `prologue` - // ahead into the slot just freed. `unit_span(u)` returns unit `u`'s gmem span, written/consumed in index - // order. - const auto run_load = [&](int n_units, auto&& unit_span) { - const int prologue = (::cuda::std::min) (PipelineStages, n_units); - for (int stage = 0; stage < prologue; ++stage) + // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail + // bulk in the forced case) ends the packed region and its suffix can be appended right after it. + for (int stage = 0; stage < prologue; ++stage) + { + const auto src = bulk_src(static_cast(stage)); + issue_bulk_copy(stage, key_slots + next_off, src); + next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + } + + // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in + // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. + int read_off = 0; + for (offset_t p = 0; p < my_resident_chunks; ++p) + { + const int stage = static_cast(p % static_cast(prologue)); + wait_stage(stage); + const int read_len = static_cast(::cuda::std::size(bulk_src(p))); + for_each_chunk_key( + bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), + add_first_pass); + read_off += read_len * int{sizeof(key_t)}; + + const offset_t next_p = p + static_cast(prologue); + if (next_p < my_resident_chunks) { - const auto src = unit_span(stage); + const auto src = bulk_src(next_p); + // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads + // leave the wait above would advance the phase twice, stranding a lagging waiter forever. + __syncthreads(); issue_bulk_copy(stage, key_slots + next_off, src); next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } - int read_off = 0; - for (int u = 0; u < n_units; ++u) - { - const int stage = u % prologue; - wait_stage(stage); - const int read_len = static_cast(::cuda::std::size(unit_span(u))); - for_each_chunk_key( - bulk_span( - {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), - add_first_pass); - read_off += read_len * int{sizeof(key_t)}; - - const int next_u = u + prologue; - if (next_u < n_units) - { - const auto src = unit_span(next_u); - // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads - // leave the wait above would advance the phase twice, stranding a lagging waiter forever. - __syncthreads(); - issue_bulk_copy(stage, key_slots + next_off, src); - next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - } - } - }; - - const int resident_units = static_cast(my_resident_chunks); - if constexpr (first_resident_slot_subdivision == 1) - { - // One load unit per resident chunk (the un-split load). - run_load(resident_units, [&](int u) { - return bulk_src(static_cast(u)); - }); - } - else - { - // Split only the first chunk into `<= first_resident_slot_subdivision` sub-chunks; chunks 1+ stay whole. A - // short first bulk (also the global tail) needs fewer sub-chunks, its final one carrying the aligned - // remainder, so the count is clamped to what the bulk covers. - const int sub_items = first_resident_slot_sub_items; // local copy: avoids ODR-using the static member by - // ref - const auto first_bulk = bulk_src(offset_t{0}); - const int first_len = static_cast(::cuda::std::size(first_bulk)); - const int first_subs = (::cuda::std::max) (1, ::cuda::ceil_div(first_len, sub_items)); - const int n_units = (first_subs - 1) + resident_units; - run_load(n_units, [&](int u) -> ::cuda::std::span { - if (u < first_subs) - { - const int off = u * sub_items; - const int len = (::cuda::std::min) (sub_items, first_len - off); - if (len == 0) - { - return {}; // empty first bulk (sole resident chunk is a short global tail): avoid `nullptr + 0` UB - } - return {first_bulk.data() + off, static_cast<::cuda::std::size_t>(len)}; - } - return bulk_src(static_cast(u - first_subs + 1)); - }); } // Tail suffix of a resident (non-peeled) suffix tail: append right after the last (tail) bulk. Nothing diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index 6829ffc1810..d234e3f65f2 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -88,7 +88,6 @@ template ; using key_t = it_value_t; @@ -367,7 +361,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SingleCtaMaxSegmentSize, MinChunksPerCta, CopyItemsPerThread, - FirstResidentSlotSubdivision, Determinism, TieBreak, KeyInputItItT, @@ -446,7 +439,6 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( SingleCtaMaxSegmentSize, MinChunksPerCta, CopyItemsPerThread, - FirstResidentSlotSubdivision, Determinism, TieBreak, KeyInputItItT, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh index 94c043d05e5..5bf9c7fc127 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh @@ -33,7 +33,6 @@ struct cluster_topk_policy int single_cta_max_segment_size; int min_chunks_per_cta; int copy_items_per_thread; - int first_resident_slot_subdivision; }; [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy @@ -42,23 +41,19 @@ struct cluster_topk_policy /*threads_per_block=*/512, /*histogram_items_per_thread=*/8, /*pipeline_stages=*/8, - /*chunk_bytes=*/32 * 1024, + /*chunk_bytes=*/16 * 1024, /*load_align_bytes=*/128, /*bits_per_pass=*/11, /*min_blocks_per_sm=*/1, /*tie_break_items_per_thread=*/8, /*single_cta_max_segment_size=*/8 * 1024, /*min_chunks_per_cta=*/1, - /*copy_items_per_thread=*/8, - /*first_resident_slot_subdivision=*/2}; + /*copy_items_per_thread=*/8}; } [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool { - // Whether `first_resident_slot_subdivision` actually splits the chunk depends on the runtime key type (sub-chunks - // must be whole and `slot_alignment`-aligned), so the agent decides per instantiation and falls back to the un-split - // load when it cannot. Here only its positivity is a policy-level invariant. - return policy.chunk_bytes % policy.load_align_bytes == 0 && policy.first_resident_slot_subdivision >= 1; + return policy.chunk_bytes % policy.load_align_bytes == 0; } static_assert(is_valid_policy(make_policy())); From 204cfc594218319fb61a75d8f5377615039641dc Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 1 Jul 2026 04:07:39 +0200 Subject: [PATCH 068/126] Reduce early exit checking It is too expensive to do it after every tile. Just try to avoid unnecessary global loads. Also avoid unnecessary global loads in the CTAs that don't straddle the candidate boundary. They do not have to care about iterating through keys in any special order. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 148 ++++++++----------- 1 file changed, 59 insertions(+), 89 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 0f530df595c..36b0abca988 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -933,18 +933,6 @@ private: return old; } - // Relaxed atomic load of `back_local_cnt`. Used by the `all_below_cta` bail to snapshot how many of this block's - // candidates have been placed, concurrently with the `back_local_inc` atomics of faster lanes: an atomic-vs-atomic - // pair, so it carries no read/write hazard (unlike a plain SMEM read) and needs no extra barrier. - _CCCL_DEVICE _CCCL_FORCEINLINE offset_t back_local_load() - { - const ::cuda::std::uint32_t addr = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)); - offset_t v; - asm volatile("ld.relaxed.cta.shared::cta.u32 %0, [%1];" : "=r"(v) : "r"(addr) : "memory"); - return v; - } - _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t v) { const ::cuda::std::uint32_t addr = @@ -1115,10 +1103,11 @@ private: // also gets the segment-local index, so share this loop. `mid()` runs once on a full pass -- after the first reload // wave (`p_eff` visits) is issued but before it is waited on, overlapping the caller's resident-chunk work with // those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with no unmatched barrier. - // `should_continue()` is polled after each consumed chunk; returning false breaks the stream so the final filter - // bails once the top-k is placed. It must be block-uniform and barrier-free (evaluated after `block_apply`'s - // syncs). The histogram and non-deterministic filter pass an always-true predicate (folding back to the - // unconditional loop). An early break can leave prefetches in flight, so the pass drains the remaining stages + // `should_continue()` is polled once after each consumed chunk (before its refill copy is issued); returning false + // breaks the stream so the final filter bails once the top-k is placed. It must be block-uniform and is evaluated + // uniformly by every lane, so it may contain a barrier (the deterministic filter's does, to read its placement + // counters block-wide). The histogram and non-deterministic filter pass an always-true predicate (folding back to + // the unconditional loop). An early break can leave prefetches in flight, so the pass drains the remaining stages // before returning (a full pass ends with an empty `inflight_mask`, so the drain is a no-op). template _CCCL_DEVICE _CCCL_FORCEINLINE void @@ -1154,9 +1143,9 @@ private: // Consume overflow visit `i`: wait for its slot, fold its keys via `block_apply`, then prefetch the chunk // `p_eff` visits ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite - // the data the block was just reading). Returns false once `block_apply` has set `stop_now` -- polled before - // the prefetch so we never launch a copy we would only drain again; the up-to-`p_eff - 1` prefetches already in - // flight (from earlier visits or priming) are drained after the loop. + // the data the block was just reading). Returns false once `should_continue()` reports the top-k fully placed + // -- polled before the prefetch so we never launch a copy we would only drain again; the up-to-`p_eff - 1` + // prefetches already in flight (from earlier visits or priming) are drained after the loop. const auto consume = [&](offset_t i) -> bool { const offset_t o = forward ? i : (m - 1 - i); const int stage = static_cast(o % pe); @@ -1951,9 +1940,9 @@ private: // the boundary-crossing tile (see `all_below_cta` and the per-tile back logic). Early stop is not special-cased: // `total_candidates == num_kth` then makes every CTA `all_below_cta`. // - // Cache `early_stop`/`total_candidates` now, while every block is still tightly coupled to the pass loop's final - // `cluster.sync` -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). - const bool early_stopped = leader_state->early_stop(); + // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final + // `cluster.sync` + // -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). const offset_t total_candidates = leader_state->len; const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region @@ -1992,38 +1981,31 @@ private: offset_t running = cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive back prefix) bool tie_active = (num_back != out_offset_t{0}) && (cand_prefix < static_cast(num_back)); - // Exact, block-local front-complete flag: `front_local_inc` counts only this block's front placements, so a lane - // seeing `local + 1 >= my_front` means all `my_front` keys are emitted. Seeded true when this block owns no - // front. No remote read, no cross-CTA dependency. - bool front_done_seen = (my_front == out_offset_t{0}); - - // Uniform early-exit predicate: true once this block's ties are placed and all of its strictly-selected keys are - // emitted (it has no further output to contribute). Used between regions (skip the overflow stream); - // `process_tiles` folds the equivalent OR into its per-tile scan-reuse barrier instead of calling this. + // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and + // resolved its ties, so it has no further output to contribute. `front_local_cnt`/`back_local_cnt` are this + // block's own SMEM placement counters; the leading barrier makes the read block-wide (and resynchronizes lanes + // that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and before + // each streaming bulk copy -- never per tile. auto should_stop = [&]() -> bool { - if (tie_active) - { - return false; - } - return __syncthreads_or(static_cast(front_done_seen)) != 0; + __syncthreads(); + const bool front_done = temp_storage.front_local_cnt >= static_cast(my_front); + // Straddling/above CTAs finish the back when `tie_active` clears; an `all_below_cta` (which never clears it) + // finishes once all `my_cand_count` of its candidates are placed. + const bool back_done = !tie_active || (temp_storage.back_local_cnt >= my_cand_count); + return front_done && back_done; }; - // Sticky, block-uniform flag set by `process_tiles` when `should_stop()` first holds. It short-circuits the - // remaining tiles, the overflow stream (via `run_pass`'s `should_continue`), and every later region. - bool stop_now = false; // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. Strictly-selected keys // go to the front via a SMEM atomic (offset by `sel_prefix`); candidates go to the back (see the per-tile logic // below). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are unresolved -- // holds the boundary and scans directly. `running` carries across tiles/regions; `get_idx(pos)` is the - // segment-local index for the value payload. Each tile ends with one barrier that reduces the done flag and - // bails. + // segment-local index for the value payload. No per-tile early-exit or barrier here: the atomic front/back sub- + // paths write disjoint slots and may let lanes drift across tiles; only the lazy-scan `else` branch keeps its + // barriers. Early exit is decided at critical points (between regions, before each streaming bulk copy) via + // `should_stop`, which resynchronizes the drift. auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count, bool region_is_terminal) { constexpr int items = items_ic(); constexpr int tile = threads_per_block * items; - if (stop_now) - { - return; - } for (int tile_base = 0; tile_base < count; tile_base += tile) { key_t keys[items]; @@ -2058,15 +2040,9 @@ private: const out_offset_t out = static_cast(sel_prefix + local); block_keys_out[out] = keys[i]; write_value(out, get_idx(pos)); - front_done_seen = front_done_seen || (local + offset_t{1} >= static_cast(my_front)); } } - // Snapshot for the `all_below_cta` bail (relaxed atomic load -> no hazard vs the back atomics, no barrier). - // Monotonic and capped at `my_cand_count`, so observing prior/partial increments can only delay the bail, - // never trip it early. Unused by straddling/above CTAs (they bail on `!tie_active`). - const offset_t back_seen = all_below_cta ? back_local_load() : offset_t{0}; - // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: // * all_below_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. @@ -2166,19 +2142,11 @@ private: { tie_active = false; } + // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the + // next tile's `back_local_inc`. The other sub-paths write disjoint slots and need no per-tile barrier. + __syncthreads(); } } - - // One barrier per tile: reduce the done flag (and order `scan_storage` reuse against the next tile's scan). - // Back-complete is `!tie_active` (straddling/above) or `back_seen >= my_cand_count` (an `all_below_cta` that - // never clears `tie_active`), AND a front-complete sighting. The OR only fires once the whole output is - // placed (no false positive), so all lanes break together. - const bool back_complete = all_below_cta ? (back_seen >= my_cand_count) : !tie_active; - if (__syncthreads_or((back_complete && front_done_seen) ? 1 : 0) != 0) - { - stop_now = true; - break; - } } }; @@ -2301,27 +2269,28 @@ private: } }; - // Overflow chunks, in strict segment-index order (the tie-break scan demands it). The histogram leaves the - // streamer ping-ponging for L2 locality, but the final filter is a single ordered pass, so we drive the bulk-copy - // pipeline here in one fixed direction (`forward == !tie_reversed`). On the block-load path this folds each - // landed slot through `process_tiles` (reusing the TMA pipeline) instead of re-reading gmem; the generic fallback - // reads gmem chunk by chunk. The streamed unroll (`tie_break_items_streamed`) keeps a tile within a chunk, and - // `run_pass`'s `should_continue` breaks the stream once `process_tiles` reports the top-k fully placed. + // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA + // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream + // once the top-k is fully placed. // - // Slot reuse: the histogram leaves its last pass's `p_eff` turn-around chunks resident in the streaming slots, - // which (ping-pong) are exactly the first `p_eff` chunks of the *next* direction. Absent an early stop, exactly - // `num_passes` passes ran, so that direction is `num_passes % 2 == 0` (compile time); when it matches our scan - // direction we keep those chunks (`primed`) and only stream the rest, otherwise we re-prime (overwrite) the slots - // in our direction. An early stop makes the pass count -- hence the resident direction -- a runtime value, so we - // always re-prime then: it is still safe because on early stop every candidate wins (all CTAs are - // `all_below_cta`), so both the front and back writes here are order-independent. + // Direction/reuse: the histogram leaves its last pass's `p_eff` turn-around chunks resident, which (ping-pong) + // are the first `p_eff` chunks of the streamer's next direction -- `streamer.forward` still tracks that runtime + // continuation. Only the straddling CTA (still crossing the K-boundary) needs scan order for its index-ordered + // tie-break, so it forces `forward == !tie_reversed` and reuses the resident chunks only when the natural + // direction already matches (else re-primes). Every other CTA is order-independent (all_below wins every + // candidate, above places none, resolved CTAs have no back left), so it keeps the natural direction and always + // reuses the slots -- skipping the re-prime copies, early stop included. auto process_overflow = [&](bool reversed) { _CCCL_ASSERT(reversed == tie_reversed, "deterministic overflow walk must run in the tie-break scan direction"); - streamer.forward = !tie_reversed; - constexpr bool resident_dir_forward = (num_passes % 2) == 0; // streamer `forward` at filter entry (no early - // stop) - constexpr bool reuse_resident = (!tie_reversed) == resident_dir_forward; - streamer.primed = reuse_resident && !early_stopped; + if (tie_active && !all_below_cta) + { + streamer.primed = (streamer.forward == (!tie_reversed)); + streamer.forward = !tie_reversed; + } + else + { + streamer.primed = true; + } streamer.run_pass( // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. @@ -2397,9 +2366,10 @@ private: }, // No interleaved resident work: the deterministic filter folds its resident span separately. [] {}, - // Break the stream the moment the whole top-k is placed (set by `process_tiles` via `should_stop`). + // Checked before each refill bulk copy: break the stream once the whole top-k is placed. `should_stop`'s + // barrier also resynchronizes the lanes that drifted through the just-folded chunk's barrier-free tiles. [&] { - return !stop_now; + return !should_stop(); }); }; @@ -2520,19 +2490,19 @@ private: // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the // overflow re-stream -- mirroring the ascending path. process_tail_edge(true); - if (!stop_now && !should_stop()) + if (!should_stop()) { process_tail(true); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_resident(true); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_overflow(true); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_head_edge(true); } @@ -2540,19 +2510,19 @@ private: else { process_head_edge(false); - if (!stop_now && !should_stop()) + if (!should_stop()) { process_resident(false); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_overflow(false); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_tail(false); } - if (!stop_now && !should_stop()) + if (!should_stop()) { process_tail_edge(false); } From db38b622732eb131321f2802dd1fecb932ce5eba Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 2 Jul 2026 03:54:50 +0200 Subject: [PATCH 069/126] Simplify tail handling Always put the unaligned part into the tail edge buffer and do not force the remaining laigned tail to be resident. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 222 +++++------------- .../catch2_test_device_segmented_topk_keys.cu | 21 +- ...catch2_test_device_segmented_topk_pairs.cu | 25 +- 3 files changed, 94 insertions(+), 174 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 36b0abca988..a4a45ccac15 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -8,7 +8,7 @@ //! the multi-kernel + global histogram pipeline used by cub::DeviceTopK. Each //! cluster processes exactly one segment. //! -//! Histogram strategy (Pattern C): +//! Histogram strategy (block-private accumulation, then DSMEM merge into the leader): //! 1. Every block lays out `hist[num_buckets]` at the same offset in its own //! shared memory. Each block accumulates a block-private histogram using //! shared-space `red` reductions (cheap, SMEM-local): cta scope for @@ -418,7 +418,7 @@ struct agent_batched_topk_cluster // (ping-ponged) across radix passes; all are initialized once up front by `init_load_barriers`. ::cuda::std::uint64_t load_mbar[PipelineStages]; // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap)`, on rank 0) and - // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner when `full_slots == 1`), each + // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner whenever it is unaligned), each // strictly `< load_align_items` keys. Loaded once in the first pass and folded into every pass + the final filter. // Block-local (never reached through DSMEM). key_t edge_keys[2 * load_align_items]; @@ -615,8 +615,8 @@ struct agent_batched_topk_cluster } } - // Stage a small (`< load_align_items` items) unaligned run -- a boundary edge (head prefix / tail suffix) or a - // resident chunk's suffix -- from gmem `src` into SMEM `dst` and fold it into the first pass in one strided sweep: + // Stage a small (`< load_align_items` items) unaligned run -- a boundary edge (head prefix / tail suffix) -- from + // gmem `src` into SMEM `dst` and fold it into the first pass in one strided sweep: // each thread copies *and* folds the same indices it owns (`local % threads_per_block == threadIdx.x`), folding from // the just-loaded register rather than from SMEM. These runs cannot go through the aligned (16-byte-aligned dst, // guard-free) BlockLoadToShared path. Fusing the copy and the fold means no thread ever reads a key another thread @@ -1072,9 +1072,8 @@ private: const offset_t chunk_idx = chunk_index_of(overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); // Every chunk begins on a `load_align` boundary (zero prefix), so the guard-free aligned (TMA bulk) path applies. - // Only the global-last chunk can carry an unaligned suffix, and it only reaches the streamer when that tail was - // peeled (its suffix lives in `edge_keys`); streaming just the aligned bulk excludes it. For every interior chunk - // `bulk == count`. + // The global-last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk + // excludes it. For every interior chunk `bulk == count`. const auto split = agent.split_chunk(block_keys_base, chunk); _CCCL_ASSERT(split.prefix == offset_t{0}, "overflow streamer received a chunk with an unaligned start"); char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; @@ -1086,7 +1085,7 @@ private: // Rebuild the shared span for the chunk currently resident in `stage`'s slot without storing per-stage state: the // slot address is a pure function of `stage` and the length is recomputed from chunk index `o`, so there is no - // spillable `pending[]` array (see `bulk_span`). Returns the aligned bulk only (a peeled tail's suffix is + // spillable `pending[]` array (see `bulk_span`). Returns the aligned bulk only (the always-peeled tail suffix is // excluded). [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span stage_span(int stage, offset_t o) const { @@ -1454,20 +1453,19 @@ private: // tail of its block_tile and re-streams its overflow chunks from gmem each pass via `streamer`. // // Boundary edges (the unaligned head prefix on rank 0 and the unaligned tail suffix on the tail owner) cannot use - // the aligned TMA streamer. The head is always peeled into the persistent `edge_keys` buffer; the tail suffix is - // kept resident in its (partial, single-slot) tail chunk, peeled into `edge_keys` only when a single slot fits - // (`full_slots == 1`). Peeling the head removes the dual-boundary pressure, so streaming needs only `full_slots >= + // the aligned TMA streamer, so both are always peeled into the persistent `edge_keys` buffer. The tail chunk's + // aligned bulk is then a normal (possibly partial) aligned chunk that can be resident or streamed like any other; + // no partial tail chunk is ever kept resident. Peeling both boundaries means streaming needs only `full_slots >= // 1`. // - // `stream_slots` is right-sized: clamped to leave a resident slot for a non-peeled tail without exceeding the - // available slots; deep overflows still get the full `PipelineStages` depth. The generic fallback has no async - // pipeline (it re-reads overflow from gmem each pass and peels nothing), so it reserves no streaming slots. + // `stream_slots` is right-sized: clamped into `[1, full_slots]` when streaming; deep overflows still get the full + // `PipelineStages` depth. The generic fallback has no async pipeline (it re-reads overflow from gmem each pass and + // peels nothing), so it reserves no streaming slots. const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); [[maybe_unused]] const bool needs_streaming = my_chunks > full_slots; // Does this rank own the global tail, and does that tail carry an unaligned suffix? (block-load path only; the // generic fallback reads any trailing items straight from gmem and never peels.) - [[maybe_unused]] offset_t tail_local = offset_t{0}; [[maybe_unused]] offset_t tail_suffix_items = offset_t{0}; bool owns_suffix_tail = false; if constexpr (use_block_load_to_shared) @@ -1476,29 +1474,23 @@ private: // true for both the strided and blocked partitions). if (my_chunks > 0 && part.global_index(my_chunks - offset_t{1}) == chunks - offset_t{1}) { - tail_local = my_chunks - offset_t{1}; const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); tail_suffix_items = split_chunk(block_keys_base, tail_chunk).suffix; owns_suffix_tail = tail_suffix_items != offset_t{0}; } } - // Peel the suffix tail only when streaming and no resident slot can be spared for it (`full_slots == 1`); otherwise - // keep it resident. `reserved_resident` holds back one slot for a non-peeled tail; `stream_slots` is then clamped - // into `[1, full_slots - reserved_resident]`. - offset_t stream_slots = offset_t{0}; - bool peel_tail = false; - bool force_tail_resident = false; + // `peel_tail` mirrors `owns_suffix_tail` (always peel, per above). `stream_slots` clamps into `[1, full_slots]`. + offset_t stream_slots = offset_t{0}; + bool peel_tail = false; if constexpr (use_block_load_to_shared) { + peel_tail = owns_suffix_tail; if (needs_streaming) { - peel_tail = owns_suffix_tail && full_slots < offset_t{2}; - const offset_t reserved_resident = (owns_suffix_tail && !peel_tail) ? offset_t{1} : offset_t{0}; - const offset_t excess = my_chunks - full_slots; - const offset_t want_stream = (::cuda::std::min) (static_cast(PipelineStages), excess); - const offset_t max_stream = full_slots - reserved_resident; // >= 1 (full_slots >= 1; reserve <= f-1) - stream_slots = (::cuda::std::max) (offset_t{1}, (::cuda::std::min) (want_stream, max_stream)); + const offset_t excess = my_chunks - full_slots; + const offset_t want_stream = (::cuda::std::min) (static_cast(PipelineStages), excess); + stream_slots = (::cuda::std::max) (offset_t{1}, (::cuda::std::min) (want_stream, full_slots)); } } const offset_t resident_slots_cap = full_slots - stream_slots; @@ -1507,28 +1499,15 @@ private: // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. _CCCL_ASSERT(my_resident_chunks <= resident_slots_cap, "Dynamic shared memory block_tile is too small"); - // A non-peeled suffix tail is forced into the last resident slot iff it would otherwise stream (it falls in the - // overflow region). The overflow then begins one chunk earlier so the displaced middle chunk streams in its place. const offset_t overflow_count = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; - if constexpr (use_block_load_to_shared) - { - // Under `reverse_residency` the suffix tail (global-last chunk) is the top resident chunk by construction, so no - // forced-tail exception is needed (its suffix is appended to the resident span like any other resident tail). - force_tail_resident = !reverse_residency && owns_suffix_tail && !peel_tail && overflow_count > 0 - && (tail_local >= my_resident_chunks); - } // Rank-local base of the resident and overflow (streamed) chunk windows. Default: resident `[0, // my_resident_chunks)`, overflow the high-index rest. `reverse_residency` swaps them: resident `[overflow_count, - // my_chunks)`, overflow - // `[0, overflow_count)`. + // my_chunks)`, overflow `[0, overflow_count)`. const offset_t resident_base = reverse_residency ? overflow_count : offset_t{0}; - const offset_t overflow_base = - reverse_residency ? offset_t{0} : (force_tail_resident ? (my_resident_chunks - offset_t{1}) : my_resident_chunks); - _CCCL_ASSERT(!force_tail_resident || my_resident_chunks >= offset_t{1}, - "a forced-resident tail needs at least one resident slot"); + const offset_t overflow_base = reverse_residency ? offset_t{0} : my_resident_chunks; // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and - // for an aligned base); the peeled tail suffix lives on the tail owner only when `peel_tail`. + // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. [[maybe_unused]] const int head_edge_len = (cluster_rank == 0u) ? static_cast(head_items) : 0; [[maybe_unused]] const int tail_edge_len = peel_tail ? static_cast(tail_suffix_items) : 0; @@ -1602,19 +1581,19 @@ private: // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor // (`read_off`) mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed // `pending_spans` array that would anchor surrounding state to local memory. Every chunk begins on a - // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except a resident suffix - // tail, whose bulk is loaded here and whose suffix is appended afterwards. The read span is rebuilt from a - // rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot demote the reads to `LD`. + // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the global-last + // chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk here. The + // read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot + // demote the reads to `LD`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - // Resident slot -> rank-local chunk index. `resident_base + slot` (identity unless `reverse_residency` shifts - // the resident window to the high-index chunks), except the last slot holds a forced-resident tail. + // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `reverse_residency` shifts + // the resident window to the high-index chunks). const auto resident_local = [&](offset_t slot) -> offset_t { - return (force_tail_resident && slot == my_resident_chunks - offset_t{1}) - ? tail_local - : (resident_base + slot); + return resident_base + slot; }; - // Aligned bulk of the resident chunk in `slot` (its `count` minus any tail suffix); empty when it has none. + // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has + // none. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { const offset_t chunk_idx = part.global_index(resident_local(slot)); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); @@ -1630,8 +1609,7 @@ private: // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. int next_off = 0; - // Load every resident chunk's aligned bulk, densely packed in slot order, so the last slot's bulk (the tail - // bulk in the forced case) ends the packed region and its suffix can be appended right after it. + // Load every resident chunk's aligned bulk, densely packed in slot order. for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); @@ -1664,22 +1642,7 @@ private: } } - // Tail suffix of a resident (non-peeled) suffix tail: append right after the last (tail) bulk. Nothing - // follows, so its non-16 length is harmless. A peeled tail's suffix lives in `edge_keys` instead. - const offset_t last_slot = my_resident_chunks - offset_t{1}; - const auto last_chunk = get_chunk(part.global_index(resident_local(last_slot)), segment_size_u32, head_items); - const auto last_split = split_chunk(block_keys_base, last_chunk); - if (last_split.suffix > 0) - { - key_t* const edge_dst = reinterpret_cast(key_slots + next_off); - stage_and_fold_edge(edge_dst, - block_keys_base + last_chunk.offset + last_split.prefix + last_split.bulk, - static_cast(last_split.suffix), - add_first_pass); - next_off += static_cast(last_split.suffix) * int{sizeof(key_t)}; - } - - // The resident region is one contiguous span [bulks | tail suffix] for the later passes; the head prefix is + // The resident region is one contiguous span of aligned bulks for the later passes; both boundary edges are // folded separately from `edge_keys`. resident_keys = {reinterpret_cast(key_slots), static_cast<::cuda::std::size_t>(next_off / int{sizeof(key_t)})}; @@ -1711,7 +1674,8 @@ private: // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix - // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, `full_slots == 1`) trails the last chunk. The + // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. + // The // `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via // `fold_edges` after a barrier). if constexpr (use_block_load_to_shared) @@ -2160,43 +2124,34 @@ private: region_is_terminal); }; - // Resident-front extent (bulk path): the contiguous resident span minus the forced-resident tail chunk, visited - // separately after the overflow because it is the globally-last chunk. Only the ascending `force_tail_resident` - // case splits off a tail; under `reverse_residency` the tail stays folded into the span (`tail_count == 0`). - int tail_count = 0; - int front_count = span_size(resident_keys); - if constexpr (use_block_load_to_shared) - { - if (force_tail_resident) - { - tail_count = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).count; - front_count = front_count - tail_count; - } - } + // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the + // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never + // part of this span. + const int front_count = span_size(resident_keys); // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later - // region in the sweep (head/resident/overflow/tail/tail-edge ascending, reversed descending) carries work. A - // forced tail/edge after the last overflow chunk is not special-cased -- a stray direct scan there is still - // correct, just rare. + // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. const bool has_head = head_edge_len > 0; const bool has_resident = front_count > 0; const bool has_overflow = overflow_count > offset_t{0}; - const bool has_tail = tail_count > 0; const bool has_tail_edge = tail_edge_len > 0; [[maybe_unused]] const bool resident_terminal = - tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail && !has_tail_edge); - [[maybe_unused]] const bool tail_terminal = - tie_reversed ? (!has_resident && !has_overflow && !has_head) : (!has_tail_edge); + tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); [[maybe_unused]] const bool head_edge_terminal = - tie_reversed ? true : (!has_resident && !has_overflow && !has_tail && !has_tail_edge); + tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); [[maybe_unused]] const bool tail_edge_terminal = - tie_reversed ? (!has_tail && !has_resident && !has_overflow && !has_head) : true; + tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to // the high-index window under `reverse_residency`). The blocked partition packs the front contiguously, so - // element `pos` maps to `front_seg_base + pos`. - const offset_t front_seg_base = get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset; + // element `pos` maps to `front_seg_base + pos`. Guarded on `front_count > 0`: with a fully-streamed rank + // (`stream_slots == full_slots`) there are no resident chunks and `resident_base` would name an out-of-range + // local chunk under `reverse_residency`. + const offset_t front_seg_base = + (front_count > 0) + ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset + : offset_t{0}; auto process_resident = [&](bool reversed) { if constexpr (use_block_load_to_shared) @@ -2373,46 +2328,6 @@ private: }); }; - auto process_tail = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - if (!force_tail_resident) - { - return; - } - key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - key_t* const tptr = rfront + front_count; - const int tc = tail_count; - // The forced-resident tail is the globally-last chunk `chunks-1`; its segment-local base is that chunk's - // offset. - const offset_t tail_seg_base = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items).offset; - if (reversed) - { - process_flat( - [&](int pos) { - return tptr[tc - 1 - pos]; - }, - [&](int pos) { - return tail_seg_base + static_cast(tc - 1 - pos); - }, - tc, - tail_terminal); - } - else - { - process_flat( - [&](int pos) { - return tptr[pos]; - }, - [&](int pos) { - return tail_seg_base + static_cast(pos); - }, - tc, - tail_terminal); - } - } - }; - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. Processed // before every chunk (ascending) / after (descending) to keep global index order. `count == 0` is a no-op, so // non-head ranks skip it. @@ -2448,9 +2363,9 @@ private: } }; - // Peeled tail suffix edge (tail owner, `full_slots == 1`): the segment's highest indices, staged in `edge_keys`. - // Processed after every chunk (ascending) / before (descending). `tail_edge_len == 0` (tail kept resident or not - // owned here) makes it a no-op. + // Peeled tail suffix edge (tail owner): the segment's highest indices, staged in `edge_keys`. Processed after + // every chunk (ascending) / before (descending). `tail_edge_len == 0` (aligned tail or not owned here) makes it + // a no-op. auto process_tail_edge = [&](bool reversed) { if constexpr (use_block_load_to_shared) { @@ -2491,10 +2406,6 @@ private: // overflow re-stream -- mirroring the ascending path. process_tail_edge(true); if (!should_stop()) - { - process_tail(true); - } - if (!should_stop()) { process_resident(true); } @@ -2519,10 +2430,6 @@ private: process_overflow(false); } if (!should_stop()) - { - process_tail(false); - } - if (!should_stop()) { process_tail_edge(false); } @@ -2635,16 +2542,17 @@ private: const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) { - // Resident keys are densely packed in slot order, so a running cursor recovers per-chunk spans. The last - // slot holds the forced-resident tail when applicable. + // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers + // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into + // `edge_keys` and folded below), so iterate `split.bulk`, not `chunk.count`. key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); int cursor = 0; for (offset_t s = 0; s < my_resident_chunks; ++s) { - const offset_t rl = (force_tail_resident && s == my_resident_chunks - offset_t{1}) ? tail_local : s; - const auto chunk = get_chunk(part.global_index(rl), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - const int cc = chunk.count; + const auto chunk = get_chunk(part.global_index(resident_base + s), segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); + const offset_t base_off = chunk.offset + split.prefix; + const int cc = split.bulk; write_run( [&](int local) { return rk[cursor + local]; @@ -2700,10 +2608,10 @@ private: } } - // No final cluster barrier: both filter paths place output via block-local SMEM atomics into gmem, so the last - // cross-CTA DSMEM access is the combined scan's `prefix_pair` push, already fenced by its post-push - // `cluster.sync()` (and `early_stop` is cached pre-scan). With no shared-memory access to another block after the - // scan, a block can return without risking a "cluster target block not present" fault from a straggler. + // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into + // gmem, so the last cross-CTA DSMEM access is the combined scan's `prefix_pair` push, already fenced by its + // post-push `cluster.sync()` (and `early_stop` is cached pre-scan). With no shared-memory access to another block + // after the scan, a block can return without risking a "cluster target block not present" fault from a straggler. } // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 8fd8e77e501..0c552cbd057 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -548,17 +548,24 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned // `static_max_segment_size` is chosen to exceed the largest all-resident cluster coverage (~16 blocks worth of // resident SMEM), so the 1 Mi-element segments force the agent's gmem-streaming overflow path (including an // unaligned overflow tail via `- 31`), while the 128 Ki-element segment still runs fully resident under the same - // streaming-capable launch configuration. + // streaming-capable launch configuration. The `+ 1` / `- 4095` sizes make the global-last chunk a single item, i.e. + // a pure-suffix tail with an empty aligned bulk (`bulk == 0`) once `pad == 0` aligns the base, exercising the + // always-peeled tail edge on top of a zero-length resident/streamed tail chunk (resident `128 Ki + 1`, streamed + // `1 Mi - 4095`). constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; - constexpr auto direction = cub::detail::topk::select::max; - const int pad = GENERATE(0, 1, 3, 7); - const segment_size_t segment_size = - GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + constexpr auto direction = cub::detail::topk::select::max; + const int pad = GENERATE(0, 1, 3, 7); + const segment_size_t segment_size = GENERATE_COPY(values( + {static_max_segment_size, + static_max_segment_size - 31, + static_max_segment_size - 4095, + segment_size_t{128 * 1024}, + segment_size_t{128 * 1024 + 1}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index de87a2d199d..88f3073288b 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -630,11 +630,9 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments through a non- // Pair analog of the keys "large fixed-size unaligned segments" test: contiguous keys offset by `pad` take the // block-load path with an unaligned head edge, and the 1 Mi-element segments stream, so the value payloads exercise -// the boundary-edge value writes that the small and non-contiguous pair tests above do not. Which boundary the launch -// stresses depends on its resident capacity: the host/graph configs keep the tail suffix resident (`full_slots > 1`), -// covering the head edge and the resident-suffix write; the device-launch (`lid_1`) static config has a small -// `full_slots`, so the tail suffix is peeled into `edge_keys` and the persistent `tail_edge_len`/`process_tail_edge` -// path is covered there. +// the boundary-edge value writes that the small and non-contiguous pair tests above do not. An unaligned tail suffix +// is always peeled into `edge_keys` (like the head prefix), so every launch config that owns such a tail exercises the +// head edge plus the persistent `tail_edge_len`/`process_tail_edge` value writes. C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned segments", "[pairs][segmented][topk][device][cluster]", select_direction_list) @@ -651,11 +649,18 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned constexpr segment_index_t num_segments = 3; const int pad = GENERATE(0, 1, 3, 7); - const segment_size_t segment_size = - GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); - const segment_size_t num_items = num_segments * segment_size; + // The `+ 1` / `- 4095` sizes make the global-last chunk a single item, i.e. a pure-suffix tail with an empty aligned + // bulk (`bulk == 0`) once `pad == 0` aligns the base, exercising the always-peeled tail edge value writes on top of a + // zero-length resident/streamed tail chunk (resident `128 Ki + 1`, streamed `1 Mi - 4095`). + const segment_size_t segment_size = GENERATE_COPY(values( + {static_max_segment_size, + static_max_segment_size - 31, + static_max_segment_size - 4095, + segment_size_t{128 * 1024}, + segment_size_t{128 * 1024 + 1}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + const segment_size_t num_items = num_segments * segment_size; CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); From 7a47139b93b4c9f58b78497e0b08263e954ab36a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 2 Jul 2026 05:17:40 +0200 Subject: [PATCH 070/126] Replace cg with PTX Use aligned cluster barrier instructions among other things. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 112 +++++++++--------- .../dispatch_batched_topk_cluster.cuh | 2 +- 2 files changed, 57 insertions(+), 57 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index a4a45ccac15..b52d5a361eb 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -60,6 +60,7 @@ #include #include #include +#include #include #include #include @@ -76,8 +77,6 @@ #include -#include - CUB_NAMESPACE_BEGIN namespace detail::batched_topk_cluster @@ -188,7 +187,7 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, // Effective cluster blocks implied by a chunk count: a CTA joins the effective cluster iff it would own at least // `min_chunks_per_cta` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap -// differs (live `cluster.num_blocks()` on the device, max launchable blocks on the host). `min_chunks_per_cta` is +// differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_cta` is // `static_assert`ed positive, so the divide is well-defined. 64-bit math. [[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( ::cuda::std::uint64_t chunks, int min_chunks_per_cta, unsigned int cluster_blocks_cap) noexcept @@ -837,14 +836,11 @@ struct agent_batched_topk_cluster // --------------------------------------------------------------------------- // Main entry point // --------------------------------------------------------------------------- - // SM 9.0+ only. `_CG_HAS_CLUSTER_GROUP` keeps the body and the - // `process_impl` definition consistent across NVCC and clang-cuda/clangd; - // `NV_IF_TARGET` strips the call from NVCC's sub-SM-9.0 device passes. + // SM 9.0+ only. `NV_IF_TARGET` strips the call on NVCC's sub-SM-9.0 device passes, so `process_impl` (and the + // cluster/async PTX in it and its callees) is never ODR-used there and never reaches ptxas. _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() { -#if defined(_CG_HAS_CLUSTER_GROUP) NV_IF_TARGET(NV_PROVIDES_SM_90, (process_impl();)); -#endif } private: @@ -899,6 +895,16 @@ private: asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" : : "r"(remote), "r"(v) : "memory"); } + // Generic pointer to this CTA's `state` as seen in the CTA at cluster rank `rank` (reached over DSMEM) -- the PTX + // equivalent of cooperative_groups' `map_shared_rank`, via `mapa.u64` (generic-address form). + _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(unsigned int rank) + { + const ::cuda::std::uint64_t own = reinterpret_cast<::cuda::std::uint64_t>(&temp_storage.state); + ::cuda::std::uint64_t remote; + asm("mapa.u64 %0, %1, %2;" : "=l"(remote) : "l"(own), "r"(rank)); + return reinterpret_cast(remote); + } + // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. @@ -1301,15 +1307,19 @@ private: // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- - // Stripped on sub-SM-9.0 device passes; uses `cluster_group`, which is only - // declared when `_CG_HAS_CLUSTER_GROUP` is set. -#if defined(_CG_HAS_CLUSTER_GROUP) + // Cluster-wide barrier via PTX (replaces cooperative_groups' `cluster.sync()`): `.release` on arrive, `.acquire` on + // wait, both `.aligned` since every thread reaches it under a uniform branch. + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_barrier() + { + asm volatile("barrier.cluster.arrive.release.aligned;" : : : "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" : : : "memory"); + } + // Synchronize the segment's cluster. A single-CTA "cluster" keeps all state block-local, so `__syncthreads()` orders // it and the cluster-scoped barrier is unnecessary. `single_cta` is computed in `run()` from the collapsed cluster - // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw `cluster.num_blocks()`. - // It is per-segment uniform across the surviving block(s), so the branch is reached uniformly. - _CCCL_DEVICE _CCCL_FORCEINLINE static void - cluster_or_block_sync(::cooperative_groups::cluster_group& cluster, bool single_cta) + // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw cluster size. It is + // per-segment uniform across the surviving block(s), so the branch is reached uniformly. + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(bool single_cta) { if (single_cta) { @@ -1317,7 +1327,7 @@ private: } else { - cluster.sync(); + cluster_barrier(); } } @@ -1331,11 +1341,7 @@ private: // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( - ::cooperative_groups::cluster_group& cluster, - bool single_cta, - unsigned int cluster_rank, - unsigned int eff_cluster_blocks, - ::cuda::std::uint64_t packed) + bool single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) { if (single_cta) { @@ -1364,14 +1370,13 @@ private: } // TODO(cccl): idle ranks arrive here only to keep this barrier reachable; a sub-cluster mbarrier over the working // ranks would let them exit (see the pass loop). - cluster_or_block_sync(cluster, single_cta); + cluster_or_block_sync(single_cta); return temp_storage.prefix_pair; } template _CCCL_DEVICE _CCCL_FORCEINLINE void - run(::cooperative_groups::cluster_group& cluster, - num_segments_val_t segment_id, + run(num_segments_val_t segment_id, unsigned int cluster_rank, unsigned int cluster_blocks, segment_size_val_t segment_size, @@ -1386,7 +1391,7 @@ private: auto block_keys_in = d_key_segments_it[segment_id]; const auto segment_size_u32 = static_cast(segment_size); - // `cluster_blocks` is what `process_impl` runs at: the launched `cluster.num_blocks()`, or `1` when it collapsed a + // `cluster_blocks` is what `process_impl` runs at: the launched cluster size, or `1` when it collapsed a // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. @@ -1417,7 +1422,7 @@ private: // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_cta` each), <= the // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- - // but stay resident and still arrive at every `cluster.sync()` (a returned CTA would hang the barrier; see the + // but stay resident and still arrive at every cluster barrier (a returned CTA would hang the barrier; see the // TODOs at the barrier sites). Derived from this CTA's head-aligned `chunks` so it matches the partition exactly. // Stays at `cluster_blocks` for host-exact sizes (the dispatch already matched it) and on the single-CTA path. unsigned int eff_cluster_blocks = cluster_blocks; @@ -1444,11 +1449,10 @@ private: // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = - single_cta ? &temp_storage.state : cluster.map_shared_rank(&temp_storage.state, leader_rank); + state_t* leader_state = single_cta ? &temp_storage.state : map_state_to_rank(leader_rank); // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every - // `cluster.sync()` are reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) + // cluster barrier is reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) // keeps them all resident and streams nothing; an overflowing CTA reserves a round-robin streaming region at the // tail of its block_tile and re-streams its overflow chunks from gmem each pass via `streamer`. // @@ -1768,7 +1772,7 @@ private: // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` // are atomic at compatible scopes (see Step 1 dispatch). The // cluster-wide ordering before Step 3's leader read of `hist[]` - // is supplied by the `cluster.sync()` further below. + // is supplied by the cluster barrier further below. __syncthreads(); // Step 2: non-leader blocks fold their per-bucket raw counts into @@ -1790,9 +1794,9 @@ private: } } - // TODO(cccl): idle ranks arrive here only because `cluster.sync()` spans the whole launched cluster. An mbarrier - // over just the active ranks would let them exit and free their SM slots instead of spinning on this barrier. - cluster_or_block_sync(cluster, single_cta); + // TODO(cccl): idle ranks arrive here only because the cluster barrier spans the whole launched cluster. An + // mbarrier over just the active ranks would let them exit and free their SM slots instead of spinning here. + cluster_or_block_sync(single_cta); // Step 3: the leader prefix-scans the merged `hist` (raw counts) and // updates the cluster-shared `state`. Subsequent reads (end-of-pass @@ -1828,7 +1832,7 @@ private: reset_hist(); } // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. - cluster_or_block_sync(cluster, single_cta); + cluster_or_block_sync(single_cta); // End-of-pass splitter fold. Every block pulls the leader's just-published `result_pair` once through DSMEM (a // single naturally-aligned `u64`, ordered by the `cluster_or_block_sync()` above) and decodes both halves from @@ -1857,7 +1861,7 @@ private: } } } -# ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP +#ifndef CUB_DISABLE_CLUSTER_TOPK_EARLY_STOP // Early stop: the leader sets `early_stop` when the splitter bucket holds exactly the remaining `k` candidates, // so no further radix refinement can change the result. Every block decodes the same flag from the `pass_result` // it just loaded and breaks together; `last_pass`/`kth_key_bits_local` then match what the original @@ -1866,7 +1870,7 @@ private: { break; } -# endif +#endif } // ----------------------------------------------------------------------- @@ -1904,15 +1908,14 @@ private: // the boundary-crossing tile (see `all_below_cta` and the per-tile back logic). Early stop is not special-cased: // `total_candidates == num_kth` then makes every CTA `all_below_cta`. // - // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final - // `cluster.sync` - // -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). + // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster + // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). const offset_t total_candidates = leader_state->len; const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final - // `cluster.sync`) block-wide before they feed the scan and `front_count`. + // cluster barrier) block-wide before they feed the scan and `front_count`. __syncthreads(); const bool participates = !is_idle_rank && (cluster_rank != leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; @@ -1923,10 +1926,9 @@ private: const offset_t push_front = my_sel; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t pp = - combined_prefix_scan(cluster, single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(pp >> 32); - const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + const ::cuda::std::uint64_t pp = combined_prefix_scan(single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(pp >> 32); + const offset_t cand_prefix = static_cast(pp & 0xffffffffu); // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A // CTA is `all_below_cta` when all of its candidates sit at or below the K-boundary (`cand_prefix + my_cand_count @@ -2448,10 +2450,9 @@ private: const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t pp = - combined_prefix_scan(cluster, single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(pp >> 32); - const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + const ::cuda::std::uint64_t pp = combined_prefix_scan(single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(pp >> 32); + const offset_t cand_prefix = static_cast(pp & 0xffffffffu); if constexpr (is_keys_only) { @@ -2610,7 +2611,7 @@ private: // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into // gmem, so the last cross-CTA DSMEM access is the combined scan's `prefix_pair` push, already fenced by its - // post-push `cluster.sync()` (and `early_stop` is cached pre-scan). With no shared-memory access to another block + // post-push cluster barrier (and `early_stop` is cached pre-scan). With no shared-memory access to another block // after the scan, a block can return without risking a "cluster target block not present" fault from a straggler. } @@ -2706,11 +2707,11 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() { - ::cooperative_groups::cluster_group cluster = ::cooperative_groups::this_cluster(); - const unsigned int cluster_rank = cluster.block_rank(); + // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). + const unsigned int cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); // Runtime cluster blocks match the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). - const unsigned int cluster_blocks = cluster.num_blocks(); + const unsigned int cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) @@ -2796,18 +2797,17 @@ private: temp_storage.my_candidates = 0; } reset_hist(); - cluster_or_block_sync(cluster, single_cta); + cluster_or_block_sync(single_cta); [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( select_directions, segment_id, - [this, &cluster, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k](auto direction_tag) { + [this, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; - this->template run(cluster, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k); + this->template run(segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k); }); _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); } -#endif // _CG_HAS_CLUSTER_GROUP }; } // namespace detail::batched_topk_cluster diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh index d234e3f65f2..3c4d2fe5fab 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh @@ -76,7 +76,7 @@ template // Kernel entry points // ----------------------------------------------------------------------------- // Dynamic-cluster kernel for host launches; the agent reads the active cluster -// width via cooperative groups. +// width from the `%cluster_nctarank` PTX special register. template Date: Thu, 2 Jul 2026 05:48:02 +0200 Subject: [PATCH 071/126] Avoid bulk copy when ct max seg size is small If it fits a single chunk there isn't much latency hiding we can do. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index b52d5a361eb..e88efa9ac22 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -375,9 +375,17 @@ struct agent_batched_topk_cluster // resident range, so such types fall back to plain per-element loads/stores. static constexpr bool key_is_bulk_tileable = int{sizeof(key_t)} == int{alignof(key_t)} && LoadAlignBytes % int{sizeof(key_t)} == 0; + + // A statically-bounded size spanning at most one aligned chunk (`num_chunks = ceil((size-head_items)/chunk_items)`, + // `head_items < chunk_items`) has one working CTA and never overflows to gmem, so the bulk-copy pipeline has nothing + // to overlap or stream. Prefer the scalar fallback: it still stages the lone chunk resident and reads any unaligned + // tail inline, skipping the mbarrier/TMA/edge-peel setup. Only bounded sizes trip this; unbounded runtime sizes + // report the type maximum (`highest`) and keep the pipeline. + static constexpr bool single_chunk_bounded = + static_max_segment_size > 0 && static_max_segment_size <= ::cuda::std::int64_t{chunk_items}; static constexpr bool use_block_load_to_shared = THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v - && key_is_bulk_tileable; + && key_is_bulk_tileable && !single_chunk_bounded; // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram From 0b26e448fa2c5b07c7aeaa84b230c4db7b1e68c5 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 2 Jul 2026 05:58:52 +0200 Subject: [PATCH 072/126] Revert "Avoid bulk copy when ct max seg size is small" This reverts commit 1a8e1bb79f4a093454739daae6ae7fde44d0dced. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index e88efa9ac22..b52d5a361eb 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -375,17 +375,9 @@ struct agent_batched_topk_cluster // resident range, so such types fall back to plain per-element loads/stores. static constexpr bool key_is_bulk_tileable = int{sizeof(key_t)} == int{alignof(key_t)} && LoadAlignBytes % int{sizeof(key_t)} == 0; - - // A statically-bounded size spanning at most one aligned chunk (`num_chunks = ceil((size-head_items)/chunk_items)`, - // `head_items < chunk_items`) has one working CTA and never overflows to gmem, so the bulk-copy pipeline has nothing - // to overlap or stream. Prefer the scalar fallback: it still stages the lone chunk resident and reads any unaligned - // tail inline, skipping the mbarrier/TMA/edge-peel setup. Only bounded sizes trip this; unbounded runtime sizes - // report the type maximum (`highest`) and keep the pipeline. - static constexpr bool single_chunk_bounded = - static_max_segment_size > 0 && static_max_segment_size <= ::cuda::std::int64_t{chunk_items}; static constexpr bool use_block_load_to_shared = THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v - && key_is_bulk_tileable && !single_chunk_bounded; + && key_is_bulk_tileable; // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram From 94d5a647aa4074a055685edead557844444830b6 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 2 Jul 2026 15:32:28 +0200 Subject: [PATCH 073/126] Clean up big nested lambdas --- cub/cub/agent/agent_batched_topk_cluster.cuh | 933 +++++++++---------- 1 file changed, 440 insertions(+), 493 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index b52d5a361eb..687aa62ee95 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1374,6 +1374,396 @@ private: return temp_storage.prefix_pair; } + // Deterministic final-filter driver. The `run()`-local tie-break state (counts, prefixes, region extents, and the + // mutable `running`/`tie_active` scan cursor) is hoisted here so the per-region index-ordered sweeps are named member + // functions instead of a nest of `[&]` lambdas. Constructed once per `run()` in the deterministic branch. Kept an + // aggregate (no user constructor) so its members are initialized positionally at the single call site. + // + // Codegen: methods are `_CCCL_FORCEINLINE` and no SMEM key pointer is stored -- the resident window is carried as its + // 32-bit shared address (`resident_smem32`) and rebuilt with `__cvta_shared_to_generic` at use, so the reads stay + // `LDS` (see the `resident_smem32` note in `run`). + template + struct det_final_filter + { + using identify_candidates_op_t = + detail::topk::identify_candidates_op_t; + + agent_batched_topk_cluster& agent; + num_segments_val_t segment_id; + identify_candidates_op_t identify_op; + it_value_t block_keys_out; + key_it_t block_keys_in; + overflow_streamer& streamer; + chunk_partition part; + out_offset_t k; + out_offset_t num_back; + out_offset_t my_front; + offset_t sel_prefix; + offset_t cand_prefix; + offset_t my_cand_count; + offset_t resident_base; + offset_t my_resident_chunks; + offset_t segment_size_u32; + offset_t head_items; + offset_t front_seg_base; + ::cuda::std::uint32_t resident_smem32; + int front_count; + int head_edge_len; + int tail_edge_len; + bool all_below_cta; + bool resident_terminal; + bool head_edge_terminal; + bool tail_edge_terminal; + // Mutable per-invocation tie-break scan cursor. + offset_t running; + bool tie_active; + + // For each key written to `block_keys_out[pos]`, load the associated input value at the key's segment-local index + // `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences the null value iterators) + // in keys-only builds; `segment_id` is loop-invariant, so the per-segment iterators hoist out of the writes. + _CCCL_DEVICE _CCCL_FORCEINLINE void write_value(out_offset_t pos, offset_t seg_idx) const + { + if constexpr (!is_keys_only) + { + auto block_vals_in = agent.d_value_segments_it[segment_id]; + auto block_vals_out = agent.d_value_segments_out_it[segment_id]; + block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; + } + } + + // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and + // resolved its ties. The leading barrier makes the counter reads block-wide (and resynchronizes lanes that raced + // ahead through the barrier-free tiles). Polled only at critical points -- between regions and before each + // streaming bulk copy -- never per tile. + _CCCL_DEVICE _CCCL_FORCEINLINE bool should_stop() + { + __syncthreads(); + const bool front_done = agent.temp_storage.front_local_cnt >= static_cast(my_front); + // Straddling/above CTAs finish the back when `tie_active` clears; an `all_below_cta` (which never clears it) + // finishes once all `my_cand_count` of its candidates are placed. + const bool back_done = !tie_active || (agent.temp_storage.back_local_cnt >= my_cand_count); + return front_done && back_done; + } + + // Fold a flat scan position `pos` into its in-region index: forward, or (`Reversed`) counting down from the + // region's last element. The `FromSmem` local read and the segment-local value index share this one folded index. + template + [[nodiscard]] static _CCCL_DEVICE _CCCL_FORCEINLINE int fold_pos(int pos, int count) + { + return Reversed ? (count - 1 - pos) : pos; + } + + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. `FromSmem` selects the + // key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` + // (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. Strictly- + // selected keys go to the front via a SMEM atomic (offset by `sel_prefix`); candidates go to the back (see the + // per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are + // unresolved -- holds the boundary and scans directly. `running` carries across tiles/regions. No per-tile early- + // exit or barrier here except the lazy-scan `else` branch; early exit is decided at critical points via + // `should_stop`. The SMEM base is passed in (rebuilt via `__cvta_shared_to_generic` at the call site) rather than + // stored, so the reads stay `LDS`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + process_tiles([[maybe_unused]] const key_t* smem_src, offset_t seg_base, int count, bool region_is_terminal) + { + constexpr int items = Items; + constexpr int tile = threads_per_block * items; + for (int tile_base = 0; tile_base < count; tile_base += tile) + { + key_t keys[items]; + offset_t flags[items]; + detail::topk::candidate_class cls[items]; + bool valid[items]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + valid[i] = pos < count; + flags[i] = offset_t{0}; + if (valid[i]) + { + const int j = fold_pos(pos, count); + if constexpr (FromSmem) + { + keys[i] = smem_src[j]; + } + else + { + keys[i] = block_keys_in[static_cast(seg_base + static_cast(j))]; + } + cls[i] = identify_op(keys[i]); + flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + } + } + + // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`. + // The per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in + // here -- they always route through the back below, even on early stop. + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + const bool to_front = valid[i] && (cls[i] == detail::topk::candidate_class::selected); + if (to_front) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const offset_t local = agent.front_local_inc(); + const out_offset_t out = static_cast(sel_prefix + local); + const offset_t seg_idx = seg_base + static_cast(fold_pos(pos, count)); + block_keys_out[out] = keys[i]; + write_value(out, seg_idx); + } + } + + // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: + // * all_below_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. + // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. + // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the + // counter and, on the crossing tile only, overwrite the arrival slots in index order. + if (tie_active) + { + const bool terminal_tile = region_is_terminal && (tile_base + tile >= count); + if (all_below_cta) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = cand_prefix + agent.back_local_inc(); + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, seg_base + static_cast(fold_pos(pos, count))); + } + } + } + } + else if (terminal_tile) + { + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(agent.temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, seg_base + static_cast(fold_pos(pos, count))); + } + } + } + running += tile_total; + tie_active = false; + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = cand_prefix + agent.back_local_inc(); + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, seg_base + static_cast(fold_pos(pos, count))); + } + } + } + // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make + // the counter read race-free and block-uniform. + __syncthreads(); + const offset_t placed = agent.temp_storage.back_local_cnt; + if ((cand_prefix + placed) > static_cast(num_back)) + { + // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered + // winners (identical slot set, different candidate->slot mapping). + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(agent.temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (valid[i] && flags[i] != offset_t{0}) + { + const offset_t global_rank = running + excl[i]; + if (global_rank < static_cast(num_back)) + { + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = keys[i]; + write_value(out, seg_base + static_cast(fold_pos(pos, count))); + } + } + } + } + running = cand_prefix + placed; + if (running >= static_cast(num_back)) + { + tie_active = false; + } + // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the + // next tile's `back_local_inc`. The other sub-paths write disjoint slots and need no per-tile barrier. + __syncthreads(); + } + } + } + } + + // Resident-front region. Direction is the compile-time `reverse_residency` (== `tie_reversed` in deterministic + // mode): ascending walks the low-index window forward, descending walks the high-index window (`resident_base`) in + // reverse, so a single `process_tiles` call per span with the index folded at compile time replaces the old fwd/rev + // pair. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident() + { + if constexpr (use_block_load_to_shared) + { + // Whole contiguous resident span staged in SMEM; rebuild the base from its 32-bit address at use. + key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + process_tiles( + rfront, front_seg_base, front_count, resident_terminal); + } + else + { + const int rc = static_cast(my_resident_chunks); + for (int s = 0; s < rc; ++s) + { + const int local_slot = reverse_residency ? (rc - 1 - s) : s; + const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); + const auto chunk = agent.get_chunk(chunk_idx, segment_size_u32, head_items); + // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy + // per-tile boundary detection handles any boundary), so pass `false`. + process_tiles( + agent.slot_keys_unpadded(local_slot), chunk.offset, chunk.count, false); + } + } + } + + // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA + // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once + // the top-k is fully placed. Only the straddling CTA needs scan order, so it forces `forward == !tie_reversed` and + // reuses the resident turn-around chunks only when the natural direction already matches; every other CTA is + // order-independent and keeps the natural direction, reusing the slots and skipping the re-prime copies. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow() + { + if (tie_active && !all_below_cta) + { + streamer.primed = (streamer.forward == (!tie_reversed)); + streamer.forward = !tie_reversed; + } + else + { + streamer.primed = true; + } + + streamer.run_pass( + // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. + // `stage_span` rebuilds the slot pointer from its 32-bit shared address (spill-proof `LDS`) and returns only + // the aligned bulk (a peeled tail suffix is handled by `process_tail_edge`). + [&](int stage, offset_t o) { + const auto span = streamer.stage_span(stage, o); + const offset_t base_off = agent.get_chunk(streamer.chunk_index_of(o), segment_size_u32, head_items).offset; + // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == + // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving + // is one barrier on one tile. + process_tiles( + span.data(), base_off, static_cast(span.size()), false); + }, + // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). + [&](const auto& chunk) { + process_tiles( + nullptr, chunk.offset, chunk.count, false); + }, + // No interleaved resident work: the deterministic filter folds its resident span separately. + [] {}, + // Checked before each refill bulk copy: break the stream once the whole top-k is placed. `should_stop`'s + // barrier also resynchronizes the lanes that drifted through the just-folded chunk's barrier-free tiles. + [&] { + return !should_stop(); + }); + } + + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. In global- + // index order it is the leading region (ascending) / trailing region (descending). `head_edge_len == 0` is a no-op, + // so non-head ranks skip it. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge() + { + if constexpr (use_block_load_to_shared) + { + // Head prefix `[0, head_edge_len)` staged at `edge_keys`, so its segment-local base is 0. + process_tiles( + agent.temp_storage.edge_keys, offset_t{0}, head_edge_len, head_edge_terminal); + } + } + + // Peeled tail suffix edge (tail owner): the segment's highest indices, staged in `edge_keys`. In global-index order + // it is the trailing region (ascending) / leading region (descending). `tail_edge_len == 0` (aligned tail or not + // owned here) makes it a no-op. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge() + { + if constexpr (use_block_load_to_shared) + { + // Peeled tail suffix staged at `edge_keys + head_edge_cap`; its segment-local base is the last `tail_edge_len` + // indices of the segment. + const int tc = tail_edge_len; + const offset_t tail_base = segment_size_u32 - static_cast(tc); + process_tiles( + agent.temp_storage.edge_keys + head_edge_cap, tail_base, tc, tail_edge_terminal); + } + } + + // Drive the regions in global-index order (ascending, or descending under `tie_reversed`), bailing between regions + // once `should_stop` reports the whole top-k placed. + _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter() + { + if constexpr (tie_reversed) + { + // Descending global-index order: peeled suffix tail, resident high-index chunks, streamed low-index overflow, + // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the + // overflow re-stream -- mirroring the ascending path. + process_tail_edge(); + if (!should_stop()) + { + process_resident(); + } + if (!should_stop()) + { + process_overflow(); + } + if (!should_stop()) + { + process_head_edge(); + } + } + else + { + process_head_edge(); + if (!should_stop()) + { + process_resident(); + } + if (!should_stop()) + { + process_overflow(); + } + if (!should_stop()) + { + process_tail_edge(); + } + } + } + }; + template _CCCL_DEVICE _CCCL_FORCEINLINE void run(num_segments_val_t segment_id, @@ -1887,7 +2277,9 @@ private: // the sweeps below. The per-segment value iterators are derived *inside* the `is_keys_only` guard: in keys-only // builds the value iterators-of-iterators are `cub::NullType**` (null), so indexing them with `segment_id` here // would dereference a null pointer; `segment_id` is loop-invariant, so the compiler hoists these out of the writes. - const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { + // The deterministic path folds this logic into `det_final_filter::write_value`; this lambda now serves only the + // non-deterministic branch below, so it is unused (but still instantiated) in deterministic builds. + [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { if constexpr (!is_keys_only) { auto block_vals_in = d_value_segments_it[segment_id]; @@ -1944,188 +2336,6 @@ private: ? static_cast(num_selected - static_cast(sel_prefix)) : static_cast(push_front); - offset_t running = cand_prefix; // candidates owned by preceding CTAs (this CTA's exclusive back prefix) - bool tie_active = (num_back != out_offset_t{0}) && (cand_prefix < static_cast(num_back)); - - // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and - // resolved its ties, so it has no further output to contribute. `front_local_cnt`/`back_local_cnt` are this - // block's own SMEM placement counters; the leading barrier makes the read block-wide (and resynchronizes lanes - // that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and before - // each streaming bulk copy -- never per tile. - auto should_stop = [&]() -> bool { - __syncthreads(); - const bool front_done = temp_storage.front_local_cnt >= static_cast(my_front); - // Straddling/above CTAs finish the back when `tie_active` clears; an `all_below_cta` (which never clears it) - // finishes once all `my_cand_count` of its candidates are placed. - const bool back_done = !tie_active || (temp_storage.back_local_cnt >= my_cand_count); - return front_done && back_done; - }; - - // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. Strictly-selected keys - // go to the front via a SMEM atomic (offset by `sel_prefix`); candidates go to the back (see the per-tile logic - // below). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are unresolved -- - // holds the boundary and scans directly. `running` carries across tiles/regions; `get_idx(pos)` is the - // segment-local index for the value payload. No per-tile early-exit or barrier here: the atomic front/back sub- - // paths write disjoint slots and may let lanes drift across tiles; only the lazy-scan `else` branch keeps its - // barriers. Early exit is decided at critical points (between regions, before each streaming bulk copy) via - // `should_stop`, which resynchronizes the drift. - auto process_tiles = [&](auto items_ic, auto get_key, auto get_idx, int count, bool region_is_terminal) { - constexpr int items = items_ic(); - constexpr int tile = threads_per_block * items; - for (int tile_base = 0; tile_base < count; tile_base += tile) - { - key_t keys[items]; - offset_t flags[items]; - detail::topk::candidate_class cls[items]; - bool valid[items]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - valid[i] = pos < count; - flags[i] = offset_t{0}; - if (valid[i]) - { - keys[i] = get_key(pos); - cls[i] = identify_op(keys[i]); - flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; - } - } - - // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`. - // The per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in - // here -- they always route through the back below, even on early stop. - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - const bool to_front = valid[i] && (cls[i] == detail::topk::candidate_class::selected); - if (to_front) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const offset_t local = front_local_inc(); - const out_offset_t out = static_cast(sel_prefix + local); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } - } - - // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: - // * all_below_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. - // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. - // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the - // counter and, on the crossing tile only, overwrite the arrival slots in index order. - if (tie_active) - { - const bool terminal_tile = region_is_terminal && (tile_base + tile >= count); - if (all_below_cta) - { - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (valid[i] && flags[i] != offset_t{0}) - { - const offset_t global_rank = cand_prefix + back_local_inc(); - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } - } - } - } - else if (terminal_tile) - { - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (valid[i] && flags[i] != offset_t{0}) - { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } - } - } - running += tile_total; - tie_active = false; - } - else - { - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (valid[i] && flags[i] != offset_t{0}) - { - const offset_t global_rank = cand_prefix + back_local_inc(); - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } - } - } - // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make - // the counter read race-free and block-uniform. - __syncthreads(); - const offset_t placed = temp_storage.back_local_cnt; - if ((cand_prefix + placed) > static_cast(num_back)) - { - // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered - // winners (identical slot set, different candidate->slot mapping). - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (valid[i] && flags[i] != offset_t{0}) - { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = - static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, get_idx(pos)); - } - } - } - } - running = cand_prefix + placed; - if (running >= static_cast(num_back)) - { - tie_active = false; - } - // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the - // next tile's `back_local_inc`. The other sub-paths write disjoint slots and need no per-tile barrier. - __syncthreads(); - } - } - } - }; - - // Full-unroll flat scan for the resident and boundary-edge regions (one contiguous span; tile may exceed a - // chunk). - auto process_flat = [&](auto get_key, auto get_idx, int count, bool region_is_terminal) { - process_tiles(::cuda::std::integral_constant{}, - get_key, - get_idx, - count, - region_is_terminal); - }; - // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never // part of this span. @@ -2155,287 +2365,39 @@ private: ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset : offset_t{0}; - auto process_resident = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - const int fc = front_count; - if (reversed) - { - process_flat( - [&](int pos) { - return rfront[fc - 1 - pos]; - }, - [&](int pos) { - return front_seg_base + static_cast(fc - 1 - pos); - }, - fc, - resident_terminal); - } - else - { - process_flat( - [&](int pos) { - return rfront[pos]; - }, - [&](int pos) { - return front_seg_base + static_cast(pos); - }, - fc, - resident_terminal); - } - } - else - { - const int rc = static_cast(my_resident_chunks); - for (int s = 0; s < rc; ++s) - { - const int local_slot = reversed ? (rc - 1 - s) : s; - const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const int cc = chunk.count; - const offset_t base_off = chunk.offset; - key_t* const ck = slot_keys_unpadded(local_slot); - // Generic multi-chunk resident loop: never the terminal-tile fast path (the lazy per-tile boundary - // detection handles any boundary), so pass `false`. - if (reversed) - { - process_flat( - [&](int pos) { - return ck[cc - 1 - pos]; - }, - [&](int pos) { - return base_off + static_cast(cc - 1 - pos); - }, - cc, - false); - } - else - { - process_flat( - [&](int pos) { - return ck[pos]; - }, - [&](int pos) { - return base_off + static_cast(pos); - }, - cc, - false); - } - } - } - }; - - // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA - // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream - // once the top-k is fully placed. - // - // Direction/reuse: the histogram leaves its last pass's `p_eff` turn-around chunks resident, which (ping-pong) - // are the first `p_eff` chunks of the streamer's next direction -- `streamer.forward` still tracks that runtime - // continuation. Only the straddling CTA (still crossing the K-boundary) needs scan order for its index-ordered - // tie-break, so it forces `forward == !tie_reversed` and reuses the resident chunks only when the natural - // direction already matches (else re-primes). Every other CTA is order-independent (all_below wins every - // candidate, above places none, resolved CTAs have no back left), so it keeps the natural direction and always - // reuses the slots -- skipping the re-prime copies, early stop included. - auto process_overflow = [&](bool reversed) { - _CCCL_ASSERT(reversed == tie_reversed, "deterministic overflow walk must run in the tie-break scan direction"); - if (tie_active && !all_below_cta) - { - streamer.primed = (streamer.forward == (!tie_reversed)); - streamer.forward = !tie_reversed; - } - else - { - streamer.primed = true; - } - - streamer.run_pass( - // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. - // `stage_span` rebuilds the slot pointer from its 32-bit shared address (spill-proof `LDS`) and returns only - // the aligned bulk (a peeled tail suffix is handled by `process_tail_edge`). - [&](int stage, offset_t o) { - const auto span = streamer.stage_span(stage, o); - const key_t* const sm = span.data(); - const int cc = static_cast(span.size()); - const offset_t base_off = get_chunk(streamer.chunk_index_of(o), segment_size_u32, head_items).offset; - constexpr auto items = ::cuda::std::integral_constant{}; - // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == - // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving - // is one barrier on one tile. - if (reversed) - { - process_tiles( - items, - [&](int pos) { - return sm[cc - 1 - pos]; - }, - [&](int pos) { - return base_off + static_cast(cc - 1 - pos); - }, - cc, - false); - } - else - { - process_tiles( - items, - [&](int pos) { - return sm[pos]; - }, - [&](int pos) { - return base_off + static_cast(pos); - }, - cc, - false); - } - }, - // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). - [&](const auto& chunk) { - const offset_t base_off = chunk.offset; - const int cc = chunk.count; - constexpr auto items = ::cuda::std::integral_constant{}; - if (reversed) - { - process_tiles( - items, - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(cc - 1 - pos))]; - }, - [&](int pos) { - return base_off + static_cast(cc - 1 - pos); - }, - cc, - false); - } - else - { - process_tiles( - items, - [&](int pos) { - return block_keys_in[static_cast(base_off + static_cast(pos))]; - }, - [&](int pos) { - return base_off + static_cast(pos); - }, - cc, - false); - } - }, - // No interleaved resident work: the deterministic filter folds its resident span separately. - [] {}, - // Checked before each refill bulk copy: break the stream once the whole top-k is placed. `should_stop`'s - // barrier also resynchronizes the lanes that drifted through the just-folded chunk's barrier-free tiles. - [&] { - return !should_stop(); - }); - }; - - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. Processed - // before every chunk (ascending) / after (descending) to keep global index order. `count == 0` is a no-op, so - // non-head ranks skip it. - auto process_head_edge = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - key_t* const he = temp_storage.edge_keys; - const int hc = head_edge_len; - if (reversed) - { - process_flat( - [&](int pos) { - return he[hc - 1 - pos]; - }, - [&](int pos) { - return static_cast(hc - 1 - pos); - }, - hc, - head_edge_terminal); - } - else - { - process_flat( - [&](int pos) { - return he[pos]; - }, - [&](int pos) { - return static_cast(pos); - }, - hc, - head_edge_terminal); - } - } - }; - - // Peeled tail suffix edge (tail owner): the segment's highest indices, staged in `edge_keys`. Processed after - // every chunk (ascending) / before (descending). `tail_edge_len == 0` (aligned tail or not owned here) makes it - // a no-op. - auto process_tail_edge = [&](bool reversed) { - if constexpr (use_block_load_to_shared) - { - key_t* const te = temp_storage.edge_keys + head_edge_cap; - const int tc = tail_edge_len; - const offset_t tail_base = segment_size_u32 - static_cast(tc); - if (reversed) - { - process_flat( - [&](int pos) { - return te[tc - 1 - pos]; - }, - [&](int pos) { - return tail_base + static_cast(tc - 1 - pos); - }, - tc, - tail_edge_terminal); - } - else - { - process_flat( - [&](int pos) { - return te[pos]; - }, - [&](int pos) { - return tail_base + static_cast(pos); - }, - tc, - tail_edge_terminal); - } - } - }; - - if constexpr (tie_reversed) - { - // Descending global-index order: peeled suffix tail, resident high-index chunks, streamed low-index overflow, - // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the - // overflow re-stream -- mirroring the ascending path. - process_tail_edge(true); - if (!should_stop()) - { - process_resident(true); - } - if (!should_stop()) - { - process_overflow(true); - } - if (!should_stop()) - { - process_head_edge(true); - } - } - else - { - process_head_edge(false); - if (!should_stop()) - { - process_resident(false); - } - if (!should_stop()) - { - process_overflow(false); - } - if (!should_stop()) - { - process_tail_edge(false); - } - } + // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` + // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `tie_active` at whether this CTA + // straddles the K-boundary. + det_final_filter filt{ + *this, + segment_id, + identify_op, + block_keys_out, + block_keys_in, + streamer, + part, + k, + num_back, + my_front, + sel_prefix, + cand_prefix, + my_cand_count, + resident_base, + my_resident_chunks, + segment_size_u32, + head_items, + front_seg_base, + resident_smem32, + front_count, + head_edge_len, + tail_edge_len, + all_below_cta, + resident_terminal, + head_edge_terminal, + tail_edge_terminal, + cand_prefix, + (num_back != out_offset_t{0}) && (cand_prefix < static_cast(num_back))}; + filt.run_filter(); } else { @@ -2524,9 +2486,12 @@ private: } }; - // Iterate a contiguous run of `count` keys whose element `local` has segment-local index `base_off + local`. - // `get_key(local)` reads the key (from SMEM for resident chunks, from gmem for overflow chunks). - auto write_run = [&](auto get_key, offset_t base_off, int count) { + // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local + // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary + // edges); overflow chunks fold through the streamer's own indexed callback below. Materialize the key into a + // register first: `write_selected_idx` binds it by `const&` and reads it several times, so passing + // `smem[local]` directly would re-issue a narrow `LDS` per use instead of reusing the loaded value. + auto write_run = [&](const key_t* smem, offset_t base_off, int count) { const int iterations = ::cuda::ceil_div(count, threads_per_block); _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) @@ -2534,7 +2499,8 @@ private: const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < count) { - write_selected_idx(get_key(local), base_off + static_cast(local)); + const key_t key = smem[local]; + write_selected_idx(key, base_off + static_cast(local)); } } }; @@ -2554,12 +2520,7 @@ private: const auto split = split_chunk(block_keys_base, chunk); const offset_t base_off = chunk.offset + split.prefix; const int cc = split.bulk; - write_run( - [&](int local) { - return rk[cursor + local]; - }, - base_off, - cc); + write_run(rk + cursor, base_off, cc); cursor += cc; } } @@ -2569,13 +2530,7 @@ private: { const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); const offset_t base_off = chunk.offset; - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); - write_run( - [&](int local) { - return chunk_keys[local]; - }, - base_off, - chunk.count); + write_run(slot_keys_unpadded(static_cast(p)), base_off, chunk.count); } } // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the @@ -2584,21 +2539,13 @@ private: { if (head_edge_len > 0) { - write_run( - [&](int local) { - return temp_storage.edge_keys[local]; - }, - offset_t{0}, - head_edge_len); + write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len); } if (tail_edge_len > 0) { - write_run( - [&](int local) { - return temp_storage.edge_keys[head_edge_cap + local]; - }, - segment_size_u32 - static_cast(tail_edge_len), - tail_edge_len); + write_run(temp_storage.edge_keys + head_edge_cap, + segment_size_u32 - static_cast(tail_edge_len), + tail_edge_len); } } }; From 43cbd1f68efc5371f78cd3c39211e14e3cefba3c Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 5 Jul 2026 15:35:07 +0200 Subject: [PATCH 074/126] Unify backends into one dispatch/tuning/kernel --- cub/cub/agent/agent_batched_topk_cluster.cuh | 53 +- cub/cub/device/device_batched_topk.cuh | 162 ++- .../device/dispatch/dispatch_batched_topk.cuh | 999 +++++++++++++++++- .../dispatch_batched_topk_cluster.cuh | 797 -------------- .../dispatch/kernels/kernel_batched_topk.cuh | 284 ++++- .../dispatch/tuning/tuning_batched_topk.cuh | 476 ++++++++- .../tuning/tuning_batched_topk_cluster.cuh | 70 -- ...tch2_test_segmented_topk_cluster_layout.cu | 107 -- .../cub/api_docs/device_topk_requirements.rst | 38 +- 9 files changed, 1855 insertions(+), 1131 deletions(-) delete mode 100644 cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh delete mode 100644 cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh delete mode 100644 cub/test/catch2_test_segmented_topk_cluster_layout.cu diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 687aa62ee95..cf57faf19a6 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -54,6 +54,7 @@ #include #include +#include #include #include #include @@ -177,22 +178,20 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, // the single-CTA tuning threshold. Occupancy- and head-alignment-independent, so the host fast path and the device // collapse decision agree exactly. 64-bit math for wide segment-size types / loose bounds. [[nodiscard]] _CCCL_HOST_DEVICE constexpr bool single_cta_eligible( - ::cuda::std::uint64_t segment_size, - ::cuda::std::uint64_t block_tile_capacity, - int single_cta_max_segment_size) noexcept + ::cuda::std::uint64_t segment_size, ::cuda::std::uint64_t block_tile_capacity, int single_block_max_seg_size) noexcept { return segment_size <= block_tile_capacity - && segment_size <= static_cast<::cuda::std::uint64_t>(single_cta_max_segment_size); + && segment_size <= static_cast<::cuda::std::uint64_t>(single_block_max_seg_size); } // Effective cluster blocks implied by a chunk count: a CTA joins the effective cluster iff it would own at least -// `min_chunks_per_cta` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap -// differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_cta` is +// `min_chunks_per_block` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap +// differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_block` is // `static_assert`ed positive, so the divide is well-defined. 64-bit math. [[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( - ::cuda::std::uint64_t chunks, int min_chunks_per_cta, unsigned int cluster_blocks_cap) noexcept + ::cuda::std::uint64_t chunks, int min_chunks_per_block, unsigned int cluster_blocks_cap) noexcept { - const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_cta); + const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_block); return static_cast( (::cuda::std::clamp) (blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } @@ -214,8 +213,8 @@ template = 1, "min_chunks_per_cta must be positive"); + static constexpr int min_chunks_per_block = MinChunksPerBlock; + static_assert(min_chunks_per_block >= 1, "min_chunks_per_block must be positive"); // Enable the per-segment effective-single-CTA runtime path only when the host could not size the launch to each // segment's exact size: any per-segment sequence (`!is_single_value`) or a `deferred` single value. For host-exact @@ -300,7 +299,7 @@ struct agent_batched_topk_cluster // bound to trim predication/registers on sub-tile segments. Larger/unbounded types keep the full unroll (codegen // unchanged); the guard also keeps the rounds arithmetic in `int` range. static constexpr bool clamp_items_to_segment = - static_max_segment_size > 0 && static_max_segment_size < single_cta_max_segment_size; + static_max_segment_size > 0 && static_max_segment_size < single_block_max_seg_size; static constexpr int segment_rounds_ceil = clamp_items_to_segment ? static_cast( @@ -361,8 +360,10 @@ struct agent_batched_topk_cluster static_assert(LoadAlignBytes > 0); static_assert(ChunkBytes % LoadAlignBytes == 0, "ChunkBytes must be a multiple of LoadAlignBytes"); // The hybrid load relies on the aligned bulk-copy path being exact (no scalar guard), which requires the load - // alignment to be at least the bulk-copy minimum alignment. + // alignment to be a power of two and at least the bulk-copy minimum alignment. Mirrors `is_valid_cluster_policy` + // (checked by the dispatch on the whole policy); repeated here to also guard direct agent instantiations. static_assert(LoadAlignBytes >= detail::bulk_copy_min_align, "LoadAlignBytes must be >= bulk_copy_min_align"); + static_assert(::cuda::is_power_of_two(LoadAlignBytes), "LoadAlignBytes must be a power of two"); static_assert(chunk_items > 0); using decomposer_t = detail::identity_decomposer_t; @@ -1810,7 +1811,7 @@ private: // depends on the multiset of keys covered by the cluster. const offset_t chunks = num_chunks(segment_size_u32, head_items); - // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_cta` each), <= the + // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_block` each), <= the // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- // but stay resident and still arrive at every cluster barrier (a returned CTA would hang the barrier; see the // TODOs at the barrier sites). Derived from this CTA's head-aligned `chunks` so it matches the partition exactly. @@ -1821,7 +1822,7 @@ private: if (!single_cta) { eff_cluster_blocks = effective_cluster_blocks_from_chunks( - static_cast<::cuda::std::uint64_t>(chunks), min_chunks_per_cta, cluster_blocks); + static_cast<::cuda::std::uint64_t>(chunks), min_chunks_per_block, cluster_blocks); } } const bool is_idle_rank = cluster_rank >= eff_cluster_blocks; @@ -2667,11 +2668,14 @@ private: } const auto segment_size = static_cast(detail::params::get_param(segment_sizes, segment_id)); - const auto k_requested = static_cast(detail::params::get_param(k_param, segment_id)); - const auto k = - static_cast((::cuda::std::min) (static_cast(k_requested), segment_size)); - - if (k == 0) + // Clamp the requested `k` to the segment size in a 64-bit width holding both operands, *before* narrowing to + // `out_offset_t`. Clamping in a narrower type first would let a large "select all" `k` wrap to a small value and + // silently truncate the output (or wrap to 0 and skip the segment). + const auto k_clamped = + (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(detail::params::get_param(k_param, segment_id)), + static_cast<::cuda::std::uint64_t>(segment_size)); + + if (k_clamped == 0) { return; } @@ -2690,6 +2694,9 @@ private: return; } + // `k_clamped <= segment_size`, which now fits `out_offset_t`, so this narrowing is safe. + const auto k = static_cast(k_clamped); + // Select-all fast path: when `k` reaches the full segment, every element wins, so we skip the radix passes, // histogram, and output-ordering and just copy. Runs on the full launched cluster (before the effective-cluster // collapse); the decision is per-segment uniform, so the branch is cluster-uniform. @@ -2711,7 +2718,7 @@ private: const bool fits_single_cta = single_cta_eligible( static_cast<::cuda::std::uint64_t>(segment_size), static_cast<::cuda::std::uint64_t>(block_tile_capacity), - single_cta_max_segment_size); + single_block_max_seg_size); if (fits_single_cta) { if (cluster_rank != 0u) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 54261ac9590..3d8fa6f773c 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -82,7 +83,11 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( // 1. determinism and tie_break must be acknowledged together (both specified, or both omitted default) // 2. an explicit tie_break of prefer_smaller_index / prefer_larger_index fully pins the result set across GPUs and // therefore requires determinism::gpu_to_gpu (it cannot be paired with run_to_run or not_guaranteed) - // 3. this initial API surface only implements the fully opted-out configuration (non-deterministic, unsorted). + // 3. only `output_ordering::unsorted` is implemented. Given rules 1/2, this admits the five implemented + // (determinism, tie_break) combinations -- (not_guaranteed, unspecified), (run_to_run, unspecified), + // (gpu_to_gpu, {unspecified, prefer_smaller_index, prefer_larger_index}) -- while `sorted` / `stable_sorted` + // (and therefore the empty-env default, which resolves to `stable_sorted`) remain rejected. Deterministic + // requests route to the cluster backend on SM 9.0+; the opt-out configuration uses the arch+size crossover. // --------------------------------------------------------------------------- static_assert(!::cuda::std::execution::__queryable_with, "Determinism should be used inside cuda::execution::require to have an effect."); @@ -117,9 +122,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( constexpr bool tie_break_compatible_with_determinism = ::cuda::std::is_same_v || ::cuda::std::is_same_v; - constexpr bool is_non_deterministic_unsorted = - ::cuda::std::is_same_v - && ::cuda::std::is_same_v; + constexpr bool is_unsorted_output = + ::cuda::std::is_same_v; static_assert(determinism_and_tie_break_paired, "cub::DeviceBatchedTopK: determinism and tie_break requirements must be acknowledged together. Either " @@ -131,22 +135,23 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( "prefer_larger_index pins the result set across GPUs and therefore requires " "cuda::execution::determinism::gpu_to_gpu (it cannot be combined with run_to_run or not_guaranteed)."); static_assert( - !determinism_and_tie_break_paired || !tie_break_compatible_with_determinism || is_non_deterministic_unsorted, - "cub::DeviceBatchedTopK currently only implements non-deterministic, unsorted output. Request it " - "explicitly with cuda::execution::require(cuda::execution::determinism::not_guaranteed, " + !determinism_and_tie_break_paired || !tie_break_compatible_with_determinism || is_unsorted_output, + "cub::DeviceBatchedTopK currently only implements cuda::execution::output_ordering::unsorted output " + "(cuda::execution::output_ordering::sorted and stable_sorted are not yet implemented). Because the default " + "output ordering is stable_sorted, an empty (no-requirement) environment is rejected: request unsorted output " + "explicitly, e.g. cuda::execution::require(cuda::execution::determinism::not_guaranteed, " "cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)."); // --------------------------------------------------------------------------- // Resolve the (optionally tuned) policy selector from the environment. // --------------------------------------------------------------------------- - using key_t = cub::detail::it_value_t>; - using value_t = cub::detail::it_value_t>; - using default_policy_selector_t = batched_topk:: - policy_selector_from_types::highest>; + // A single tuning query on the whole `topk_policy`: a `tune`d selector returning `topk_policy` picks the backend + // (baseline vs cluster) and both sub-policies in one shot. Absent by default (the sentinel `no_override`), in which + // case the dispatch's automatic arch+size selector is used. Strip cv/ref so the override is a value type. using tuning_env_t = ::cuda::__call_result_or_t<::cuda::execution::__get_tuning_t, ::cuda::std::execution::env<>, EnvT>; - using policy_selector_t = ::cuda::std::execution:: - __query_result_or_t; + using selector_override_t = ::cuda::std::remove_cvref_t< + ::cuda::std::execution::__query_result_or_t>; // --------------------------------------------------------------------------- // Argument-annotation constraints surfaced at the call site. @@ -177,7 +182,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( // conservative 64-bit upper bound here. constexpr auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; - return batched_topk::dispatch( + return batched_topk::dispatch( d_temp_storage, temp_storage_bytes, d_keys_in, @@ -189,8 +197,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( ::cuda::args::constant{}, num_segments, total_num_items, - stream.get(), - policy_selector_t{}); + stream.get()); } //! @endcond } // namespace detail @@ -242,8 +249,9 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! //! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. ``num_segments`` must be //! a single value (``constant``, ``immediate``, or a plain integral), never a per-segment sequence. ``segment_sizes`` -//! must also carry a small compile-time upper bound (a ``constant`` or ``cuda::args::bounds()``), and tight -//! bounds on every parameter are encouraged. +//! must also carry a compile-time upper bound (a ``constant`` or ``cuda::args::bounds()``); the permitted +//! maximum is architecture-dependent (see *Current constraints* below), and tight bounds on every parameter are +//! encouraged. //! //! .. code-block:: c++ //! @@ -272,17 +280,19 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! This is an initial, intentionally restricted API surface. The following constraints are enforced at compile time //! (a ``static_assert`` fires if violated): //! -//! - **Small segments only.** Every segment must be processable by a single thread block (one worker per segment). -//! The *statically-known maximum* segment size (the upper bound of the ``segment_sizes`` annotation) must be small -//! enough that such a block fits within the shared-memory limit. Both uniform (fixed) and variable segment sizes are -//! supported as long as this maximum is honored. +//! - **Segment size is architecture-dependent.** On pre-Hopper GPUs (compute capability < 9.0) every segment must be +//! processable by a single thread block (one worker per segment): the *statically-known maximum* segment size (the +//! upper bound of the ``segment_sizes`` annotation) must be small enough that such a block fits within the +//! shared-memory limit. On Hopper and newer GPUs (compute capability >= 9.0) the thread-block-cluster backend also +//! handles larger segments that exceed this per-block limit. Both uniform (fixed) and variable segment sizes are +//! supported. //! - **Uniform number of segments.** ``num_segments`` must be a single value, never a per-segment sequence. -//! - **Explicit opt-out required for the output guarantees.** The deterministic, stable-sorted default contract -//! described in *Determinism, tie-breaking, and output ordering* below (and in :ref:`cub-topk-requirements`) is not -//! yet implemented. The caller must currently request non-deterministic, unsorted output explicitly by passing -//! ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, -//! cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)`` in the environment -//! (``determinism`` and ``tie_break`` must always be specified together). +//! - **Unsorted output required.** Only ``cuda::execution::output_ordering::unsorted`` is implemented; the sorted +//! orderings of the default contract described in *Determinism, tie-breaking, and output ordering* below (and hence +//! an empty, no-requirement environment, which defaults to ``stable_sorted``) are rejected at compile time. The +//! supported ``determinism`` / ``tie_break`` requirements depend on the architecture -- see the *Current support* +//! note below. ``determinism`` and ``tie_break`` must always be specified together, or both omitted to take the +//! default. //! //! Determinism, tie-breaking, and output ordering //! +++++++++++++++++++++++++++++++++++++++++++++++ @@ -303,12 +313,24 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! //! .. note:: //! -//! **Current support.** This release only implements the fully opted-out configuration, which must be requested -//! explicitly: ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, -//! cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)``. Any other combination -//! (including an empty, no-requirement environment) is rejected at compile time. In this configuration the -//! per-segment output is unordered and may be non-deterministic: if multiple items tie at the K-th position, the -//! subset of tied elements returned is not uniquely defined and may vary between runs. +//! **Current support.** Only the unsorted output ordering is implemented; +//! ``cuda::execution::output_ordering::unsorted`` must be requested explicitly (``sorted`` / ``stable_sorted``, and +//! thus an empty, no-requirement environment, are rejected at compile time). The supported selection requirements +//! and segment sizes differ by architecture: +//! +//! - **Pre-Hopper (compute capability < 9.0):** only the fully non-deterministic request +//! ``(determinism::not_guaranteed, tie_break::unspecified)`` is supported, and every segment must fit a single +//! thread block. Deterministic / tie-break requests and larger segments require the SM 9.0+ cluster backend and +//! are diagnosed at compile time (or, in relaxed builds, at runtime as ``cudaErrorNotSupported``). +//! - **Hopper and newer (compute capability >= 9.0):** all five acknowledged ``(determinism, tie_break)`` pairs are +//! supported -- ``(not_guaranteed, unspecified)``, ``(run_to_run, unspecified)``, and ``(gpu_to_gpu, +//! {unspecified, prefer_smaller_index, prefer_larger_index})`` -- and larger segments are handled by the cluster +//! backend. Among non-deterministic requests the baseline vs cluster backend is chosen by an architecture / +//! segment-size crossover. +//! +//! When ``determinism::not_guaranteed`` is requested the per-segment output may be non-deterministic: if multiple +//! items tie at the K-th position, the subset of tied elements returned is not uniquely defined and may vary between +//! runs. //! //! Usage Considerations //! ++++++++++++++++++++++++++ @@ -375,7 +397,8 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -387,8 +410,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -480,8 +506,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -579,8 +608,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -670,8 +702,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -786,8 +821,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -889,8 +927,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -997,8 +1038,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template ` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the + //! permitted maximum is architecture-dependent (see the *Current constraints* section). //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -1100,8 +1144,10 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, - //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `output_ordering::unsorted` (`sorted` / `stable_sorted`, and + //! thus an empty environment, are not yet supported). The selection requirements may be any acknowledged + //! `(determinism, tie_break)` pair: `(not_guaranteed, unspecified)`, `(run_to_run, unspecified)`, or `gpu_to_gpu` + //! with `unspecified` / `prefer_smaller_index` / `prefer_larger_index`. Deterministic requests require SM 9.0+. //! @endrst template +#include #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -32,16 +36,27 @@ #include #include +#include +#include +#include #include #include +#include #include +#include +#include #include #include +#include #include #include #include #include +#include + +#include + CUB_NAMESPACE_BEGIN namespace detail::batched_topk @@ -101,6 +116,10 @@ struct segment_size_to_tile_count_op // ----------------------------------------------------------------------------- // Segmented Top-K Dispatch // ----------------------------------------------------------------------------- +// +// NOTE: `baseline_dispatch` / `baseline_dispatch_with_env` are the legacy baseline-only entry points. The public API +// now routes through the unified `dispatch` (baseline + cluster) below and no longer reaches these; they are retained +// for now and must stay in sync with the baseline arm of `dispatch`. //! @param d_temp_storage Device-accessible allocation of temporary storage. When `nullptr`, the required allocation //! size is written to `temp_storage_bytes` and no work is done. @@ -128,14 +147,14 @@ template >, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>> + typename PolicySelector = baseline_policy_selector_from_types>, + it_value_t>, + ::cuda::std::int64_t, + ::cuda::args::__traits::highest>> #if _CCCL_HAS_CONCEPTS() - requires batched_topk_policy_selector + requires baseline_topk_policy_selector #endif // _CCCL_HAS_CONCEPTS() -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, @@ -245,7 +264,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( // TODO (elstehle): support number of segments provided by device-accessible iterator // Only uniform number of segments are supported (i.e., we need to resolve the number of segments on the host) static_assert(::cuda::args::__traits::is_single_value, - "Only uniform segment sizes are currently supported."); + "Only a uniform number of segments is currently supported."); if constexpr (any_small_segments) { @@ -330,7 +349,7 @@ template > -[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch_with_env( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, ValueInputItItT d_value_segments_it, @@ -343,13 +362,13 @@ template >, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>; + baseline_policy_selector_from_types>, + it_value_t>, + ::cuda::std::int64_t, + ::cuda::args::__traits::highest>; return detail::dispatch_with_env_and_tuning( env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { - return dispatch( + return baseline_dispatch( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -365,6 +384,956 @@ template ::highest` semantics: +// the compile-time bound for `constant`/bounded sequence arguments and the runtime value for a uniform +// `immediate`. For a per-segment sequence with only a static bound this can be the loose `numeric_limits::max()`. +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto runtime_max_segment_size(SegmentSizeParameterT segment_sizes) noexcept +{ + return ::cuda::args::__highest_(segment_sizes); +} + +// Host launches go through the single kernel symbol (`device_batched_topk_kernel`, see kernel_batched_topk.cuh). Only +// the CDP static-cluster kernel below remains a dedicated cluster symbol, because device-side launches cannot opt in to +// dynamic cluster dimensions. + +#ifdef CUB_RDC_ENABLED +// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so the +// triple-chevron launch from device code needs no `cudaFuncSetAttribute`. +template +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) + _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + ::cuda::std::uint32_t block_tile_capacity) +{ + using agent_t = batched_topk_cluster::agent_batched_topk_cluster< + ThreadsPerBlock, + HistogramItemsPerThread, + PipelineStages, + ChunkBytes, + LoadAlignBytes, + BitsPerPass, + TieBreakItemsPerThread, + SingleBlockMaxSegSize, + MinChunksPerBlock, + CopyItemsPerThread, + Determinism, + TieBreak, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` + // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } + + agent_t agent( + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + key_slots, + block_tile_capacity); + + agent.Process(); +} +#endif // CUB_RDC_ENABLED + +// CDP launch body, empty when CDP is disabled. Wrapped in a macro because +// `#ifdef` can't sit inside `NV_IF_TARGET`. +#ifndef CUB_RDC_ENABLED +// Without CDP/RDC a device-side launch is impossible; surface that instead of silently returning success (no-op). +# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH return cudaErrorNotSupported; +#else // CUB_RDC_ENABLED +# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ + auto static_kernel = detail::batched_topk::device_segmented_topk_cluster_kernel_static< \ + ThreadsPerBlock, \ + HistogramItemsPerThread, \ + PipelineStages, \ + ChunkBytes, \ + LoadAlignBytes, \ + BitsPerPass, \ + TieBreakItemsPerThread, \ + SingleBlockMaxSegSize, \ + MinChunksPerBlock, \ + CopyItemsPerThread, \ + Determinism, \ + TieBreak, \ + KeyInputItItT, \ + KeyOutputItItT, \ + ValueInputItItT, \ + ValueOutputItItT, \ + SegmentSizeParameterT, \ + KParameterT, \ + SelectDirectionParameterT, \ + NumSegmentsParameterT>; \ + if (const auto error = CubDebug( \ + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ + static_cast(grid_blocks), ThreadsPerBlock, dynamic_smem_bytes, stream) \ + .doit(static_kernel, \ + d_key_segments_it, \ + d_key_segments_out_it, \ + d_value_segments_it, \ + d_value_segments_out_it, \ + segment_sizes, \ + k_param, \ + select_directions, \ + num_segments, \ + block_tile_capacity))) \ + { \ + return error; \ + } +#endif // CUB_RDC_ENABLED + +// Cluster host-launch arm of the dispatch. Launches the single kernel symbol +// (`device_batched_topk_kernel`, passing an empty `baseline_kernel_args` and the resident `cluster_kernel_args`). +// `select_directions` arrives already wrapped; the cluster tuning is taken from `policy_getter` (the resolved-CC +// policy) and the requested `Determinism`/`TieBreak` from the `PolicySelector`. The CDP arm still launches the +// dedicated static-cluster kernel symbol. +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( + PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + cudaStream_t stream) +{ + constexpr auto Determinism = PolicySelector::determinism; + constexpr auto TieBreak = PolicySelector::tie_break; + // Use the cluster sub-policy for the *resolved* architecture (the one `dispatch_compute_cap` matched), i.e. exactly + // what the device kernel instantiates via `current_policy()`. Sourcing it from `policy_getter` keeps + // the host launch config (block size, shared-memory math, agent instantiation) in lock-step with the device policy + // per CC. `policy_getter()` is a constant expression in AOT builds, so `policy` is usable as a non-type template arg. + constexpr cluster_topk_policy policy = policy_getter().cluster; + constexpr int ThreadsPerBlock = policy.threads_per_block; + constexpr int HistogramItemsPerThread = policy.histogram_items_per_thread; + constexpr int PipelineStages = policy.pipeline_stages; + constexpr int ChunkBytes = policy.chunk_bytes; + constexpr int LoadAlignBytes = policy.load_align_bytes; + constexpr int BitsPerPass = policy.bits_per_pass; + constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; + constexpr int SingleBlockMaxSegSize = policy.single_block_max_seg_size; + constexpr int MinChunksPerBlock = policy.min_chunks_per_block; + constexpr int CopyItemsPerThread = policy.copy_items_per_thread; + + using key_it_t = it_value_t; + using key_t = it_value_t; + using layout_t = batched_topk_cluster::smem_block_tile_layout; + using agent_t = batched_topk_cluster::agent_batched_topk_cluster< + ThreadsPerBlock, + HistogramItemsPerThread, + PipelineStages, + ChunkBytes, + LoadAlignBytes, + BitsPerPass, + TieBreakItemsPerThread, + SingleBlockMaxSegSize, + MinChunksPerBlock, + CopyItemsPerThread, + Determinism, + TieBreak, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + // TODO: This should be taken care of in the public env-based interface. + // A tie-break preference is only meaningful once the result set itself is deterministic. + static_assert(Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed + || TieBreak == ::cuda::execution::tie_break::__tie_break_t::__unspecified, + "A tie-break preference requires a deterministic execution requirement"); + + // Validate the block_tile byte geometry (load alignment power-of-two / >= 16, chunk a multiple of it) in one place. + static_assert(is_valid_cluster_policy(policy)); + static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); + // Static-footprint estimate for the device-side CDP fallback, which cannot query `cudaFuncGetAttributes`. + // The host path instead uses the driver-reported `sharedSizeBytes` (see below), which is padding-aware. + constexpr int static_smem_bytes = static_cast(sizeof(typename agent_t::TempStorage)); + + const auto max_seg_size = runtime_max_segment_size(segment_sizes); + + // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. + size_t allocation_sizes[1] = {1}; + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + static_assert(::cuda::args::__traits::is_single_value, + "Number of segments must be resolved on the host."); + + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + const auto num_seg_val = detail::params::get_param(num_segments, num_segments_val_t{0}); + if (num_seg_val == 0) + { + return cudaSuccess; + } + + // A zero bound would drive `clusterDim.x = 0`, which the runtime rejects. + if (max_seg_size == 0) + { + return cudaSuccess; + } + + // Cluster launches require compute capability 9.0+. + int sm_version = 0; + if (const auto error = CubDebug(SmVersionUncached(sm_version))) + { + return error; + } + if (sm_version < 900) + { + return cudaErrorNotSupported; + } + + // Single kernel symbol; its cluster arm is selected device-side via `current_policy()`. The baseline + // arm is pruned per-arch, so no baseline symbol is emitted here. + constexpr auto dynamic_kernel = &device_batched_topk_kernel< + PolicySelector, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>; + + NV_IF_TARGET( + NV_IS_HOST, + ({ + // The launcher's `doit` carries the triple-chevron that makes NVCC emit `dynamic_kernel` for this TU, and + // performs the cluster launch via `cudaLaunchKernelEx`. The factory also wraps the pre-launch driver queries. + detail::TripleChevronFactory launcher_factory{}; + + // Opt in to non-portable cluster blocks (>8 on Hopper). + if (const auto error = launcher_factory.set_non_portable_cluster_allowed(dynamic_kernel)) + { + return error; + } + + // Config used only for the occupancy probe below; the final launch goes through `launcher_factory`. + // `clusterDim.x` is a placeholder since `cudaOccupancyMaxPotentialClusterSize` ignores it. + ::cudaLaunchAttribute cluster_attr{}; + cluster_attr.id = ::cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim.x = 1; + cluster_attr.val.clusterDim.y = 1; + cluster_attr.val.clusterDim.z = 1; + + ::cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(1, 1, 1); + cfg.blockDim = dim3(static_cast(ThreadsPerBlock), 1, 1); + cfg.dynamicSmemBytes = 0; + cfg.stream = stream; + cfg.attrs = &cluster_attr; + cfg.numAttrs = 1; + + // Resolve the per-block opt-in shared-memory budget and the kernel's static footprint from the driver so + // the dynamic-SMEM math below matches exactly what the launch permits. The opt-in budget + // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`) is the documented total per-block budget; the usable dynamic + // portion (`max_dynamic_smem_bytes`) is that budget minus the static footprint. + int device_id = 0; + if (const auto error = CubDebug(cudaGetDevice(&device_id))) + { + return error; + } + int max_smem_optin_bytes = 0; + if (const auto error = + CubDebug(cudaDeviceGetAttribute(&max_smem_optin_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id))) + { + return error; + } + // Use the driver-reported static footprint (`sharedSizeBytes`) rather than `sizeof(TempStorage)`: it reflects + // any padding the toolchain inserts to align the dynamic shared-memory section after the static one, so the + // derived dynamic sizes neither overshoot the budget nor conservatively drop the top table tier. + cudaFuncAttributes kernel_attrs{}; + if (const auto error = CubDebug(cudaFuncGetAttributes(&kernel_attrs, dynamic_kernel))) + { + return error; + } + // `cudaDevAttrMaxSharedMemoryPerBlockOptin` already excludes the driver's per-block reserved shared memory + // (opt-in == per-SM - reserved), so the dynamic budget is just the opt-in budget minus the static footprint; + // reserved must not be subtracted a second time. + const int nondynamic_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); + const int max_dynamic_smem_bytes = + (max_smem_optin_bytes > nondynamic_smem_bytes) ? max_smem_optin_bytes - nondynamic_smem_bytes : 0; + + // Raise the kernel's dynamic-SMEM opt-in lazily: occupancy queries and the launch must not request more than the + // currently configured `cudaFuncAttributeMaxDynamicSharedMemorySize`. The kernel's compile-time default already + // permits the portable 48 KiB total, i.e. that budget minus the static footprint. + constexpr int portable_total_smem_bytes = 48 * 1024; + int configured_dynamic_smem_limit = + (portable_total_smem_bytes > nondynamic_smem_bytes) ? portable_total_smem_bytes - nondynamic_smem_bytes : 0; + const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { + if (dynamic_smem_bytes <= configured_dynamic_smem_limit) + { + return cudaSuccess; + } + + if (const auto error = CubDebug(cudaFuncSetAttribute( + reinterpret_cast(dynamic_kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + dynamic_smem_bytes))) + { + return error; + } + configured_dynamic_smem_limit = dynamic_smem_bytes; + return cudaSuccess; + }; + + // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); + // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more + // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We + // pick the `C` that minimizes the number of waves, breaking ties toward the largest `C` (= smallest SMEM = most + // L1), which matches the profiled fast configs. We enumerate `C` analytically rather than discovering SMEM tiers + // via occupancy, so a register-limited occupancy (e.g. 1 CTA/SM) cannot collapse the candidate set. + const auto seg = static_cast<::cuda::std::uint64_t>(max_seg_size); + const auto chunk_items_u64 = static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); + const int max_block_tile_capacity = static_cast(layout_t::block_tile_capacity(max_dynamic_smem_bytes)); + if (max_block_tile_capacity <= 0) + { + // Not even one load-aligned chunk fits in the opt-in budget; the kernel cannot run. + return cudaErrorInvalidValue; + } + + // `S_res(items)`: smallest chunk-granular dynamic SMEM whose per-CTA capacity reaches `items`. + const auto smem_for_block_capacity = [&](::cuda::std::uint64_t items) { + const auto slots = ::cuda::ceil_div(items, chunk_items_u64); + return layout_t::base_padding_bytes + static_cast(slots) * layout_t::chunk_bytes; + }; + + // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (cap HW + // max). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest fully-resident `C`. + // Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. + // `numeric_limits::max()` for an unbounded deferred sequence); narrowing such a `C_lo` to `int` could wrap to + // a small (or negative) value and wrongly enter the resident branch instead of the oversize/streaming fallback. + const int c_full = static_cast( + (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), + ::cuda::ceil_div(seg, chunk_items_u64))); + const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); + // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than + // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. + const int desired_cluster_blocks = static_cast(batched_topk_cluster::effective_cluster_blocks_from_chunks( + ::cuda::ceil_div(seg, chunk_items_u64), + MinChunksPerBlock, + static_cast(max_supported_cluster_blocks))); + + int cluster_blocks = 0; + int dynamic_smem_sel = 0; + + if (batched_topk_cluster::single_cta_eligible( + seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity), SingleBlockMaxSegSize)) + { + // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's + // cluster-barrier-free path beats spreading it across more CTAs. `S_res(seg)` is within budget and one CTA is + // always launchable, so the occupancy probe is skipped (the shared `ensure_dynamic_smem_limit` below raises the + // opt-in for the selected SMEM). Larger fully-resident segments fall through to the wave-aware search below. + cluster_blocks = 1; + dynamic_smem_sel = smem_for_block_capacity(seg); + } + else if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) + { + // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every + // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. + // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. The + // upper bound is capped at the cluster blocks the max segment needs (`desired_cluster_blocks`), so the host + // never launches a wider cluster than necessary; at `min_chunks_per_block == 1` the cap equals `c_full`. + const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); + const int c_end = (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks)); + ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); + for (int c = c_begin; c <= c_end; ++c) + { + const auto per_block_items = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(c)); + const int s_res = smem_for_block_capacity(per_block_items); + if (s_res > max_dynamic_smem_bytes) + { + continue; // unreachable for c >= C_lo, but guards the SMEM budget regardless. + } + + if (const auto error = ensure_dynamic_smem_limit(s_res)) + { + return error; + } + + // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must + // be a multiple of the cluster blocks. The returned value is the device-wide clusters-per-wave (capacity), + // independent of grid size, and accounts for the static footprint and register pressure internally. + cluster_attr.val.clusterDim.x = static_cast(c); + cfg.gridDim = dim3(static_cast(c), 1, 1); + cfg.dynamicSmemBytes = static_cast(s_res); + int clusters_per_wave = 0; + if (const auto error = launcher_factory.max_active_clusters(clusters_per_wave, dynamic_kernel, &cfg)) + { + return error; + } + if (clusters_per_wave <= 0) + { + continue; // cluster blocks not launchable at this SMEM. + } + + const auto waves = ::cuda::ceil_div( + static_cast<::cuda::std::uint64_t>(num_seg_val), static_cast<::cuda::std::uint64_t>(clusters_per_wave)); + // Min waves; tie-break largest `C`. The loop ascends in `C`, so `<=` keeps the largest at equal waves. + if (cluster_blocks == 0 || waves <= best_waves) + { + best_waves = waves; + cluster_blocks = c; + dynamic_smem_sel = s_res; + } + } + + if (cluster_blocks == 0 && c_lo == 1) + { + // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, + // but `C_lo == 1` guarantees `S_res(seg)` fits the budget and one CTA is always launchable. + cluster_blocks = 1; + dynamic_smem_sel = smem_for_block_capacity(seg); + } + } + + if (cluster_blocks == 0) + { + // Oversize (`C_lo > HW max`) or nothing launchable in range: full residency is impossible, so maximize + // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. + if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) + { + return error; + } + cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` + cfg.gridDim = dim3(1, 1, 1); + cfg.dynamicSmemBytes = static_cast(max_dynamic_smem_bytes); + int hw_max_cluster_blocks = 0; + if (const auto error = launcher_factory.max_potential_cluster_size(hw_max_cluster_blocks, dynamic_kernel, &cfg)) + { + return error; + } + hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, max_supported_cluster_blocks); + if (hw_max_cluster_blocks <= 0) + { + return cudaErrorInvalidValue; + } + cluster_blocks = hw_max_cluster_blocks; + dynamic_smem_sel = max_dynamic_smem_bytes; + } + + // The launch needs `MaxDynamicSharedMemorySize >= dynamic_smem_sel`; the scan already raised the limit past the + // largest probed SMEM, so this is a no-op unless the selected config skipped the scan. + if (const auto error = ensure_dynamic_smem_limit(dynamic_smem_sel)) + { + return error; + } + + const int dynamic_smem_bytes = dynamic_smem_sel; + const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + + const auto grid_blocks = + static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); + if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + // The cluster dimension routes the host launch through `cudaLaunchKernelEx`; passing `dynamic_kernel` to `.doit` + // below forces NVCC to emit that kernel symbol for this TU. + if (const auto error = CubDebug( + launcher_factory(dim3(static_cast(grid_blocks), 1, 1), + dim3(static_cast(ThreadsPerBlock), 1, 1), + static_cast<::cuda::std::size_t>(dynamic_smem_bytes), + stream, + /*dependent_launch=*/false, + dim3(static_cast(cluster_blocks), 1, 1)) + .doit(dynamic_kernel, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + baseline_kernel_args{}, + cluster_kernel_args{static_cast<::cuda::std::uint32_t>(block_tile_capacity)}))) + { + return error; + } + }), + ({ + // CDP path: device-side launches cannot opt in to more than portable total SMEM or non-portable cluster blocks. + // Segments exceeding the portable resident coverage are still handled: the agent re-streams overflow from gmem. + constexpr int portable_total_smem_bytes = 48 * 1024; + constexpr int dynamic_smem_bytes = + (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; + + // The compile-time `ChunkBytes` is reused verbatim; the agent peels unaligned boundary edges into a tiny + // per-block buffer and re-streams overflow from gmem, so the only hard requirement is that one load-aligned chunk + // fits the worst-case portable SMEM block tile. + constexpr auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + static_assert(block_tile_capacity >= static_cast<::cuda::std::uint32_t>(layout_t::chunk_items), + "Portable SMEM is too small to fit even one load-aligned chunk for the device-launch (CDP) path"); + + const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) + * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); + if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + + CUB_TOPK_CLUSTER_DEVICE_LAUNCH + })); + + // Cluster launches can fail on the device while reporting success; sync. + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + return CubDebug(detail::DebugSyncStream(stream)); +} + +#undef CUB_TOPK_CLUSTER_DEVICE_LAUNCH + +// Baseline host-launch arm of the dispatch. Launches the single kernel symbol +// (`device_batched_topk_kernel`, packing the large-segment bookkeeping into `baseline_kernel_args` and passing an empty +// `cluster_kernel_args`). `select_directions` arrives already wrapped and the baseline tuning is taken from the +// `PolicySelector` (via `baseline_policy_selector_adaptor`). +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_baseline_arm( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + cudaStream_t stream) +{ + if constexpr (!PolicySelector::baseline_can_cover) + { + // The policy selector never routes to the baseline backend when it cannot cover the static max segment size, so + // this arm is pruned per-arch in AOT builds. Kept assert-free (no `find_smallest_covering_policy`) so it also + // compiles under runtime-policies mode, where both host arms are instantiated. + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return cudaErrorNotSupported; + } + else + { + using large_segment_tile_offset_t = LargeSegmentTileOffsetT; + using baseline_selector_t = baseline_policy_selector_adaptor; + + // Determine which one-worker-per-segment policy covers the segment-size range and k. + constexpr auto policy = find_smallest_covering_policy< + baseline_selector_t, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + large_segment_tile_offset_t>::policy; + constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; + constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; + + static constexpr int worker_per_segment_tile_size = + worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; + static constexpr bool any_small_segments = + ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; + static constexpr bool only_small_segments = + ::cuda::args::__traits::highest <= worker_per_segment_tile_size; + + // Allocation layout: + // only_small_segments: [0] dummy. + // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, + // [2] large-segment ids. + // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. + static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); + size_t allocation_sizes[allocations_array_size] = {1}; + + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + using counters_t = batched_topk_counters; + using segment_size_scan_offset_t = detail::choose_offset_t; + using segment_size_scan_input_op_t = + segment_size_to_tile_count_op; + static constexpr auto multi_worker_per_segment_tile_size = + multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; + const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; + // Transform iterator over [0, num_segments) producing the tile-count for each segment. + [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( + ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); + + if constexpr (!only_small_segments) + { + const auto num_segments_val = params::get_param(num_segments, 0); + allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); + if constexpr (any_small_segments) + { + allocation_sizes[1] = sizeof(counters_t); + allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); + } + else + { + // Query the temporary storage requirement of the segment-size transform-scan. + if (const auto error = CubDebug(detail::scan::dispatch( + nullptr, + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(nullptr), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(num_segments_val), + stream))) + { + return error; + } + } + } + + void* allocations[allocations_array_size] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + static_assert(::cuda::args::__traits::is_single_value, + "Only a uniform number of segments is currently supported."); + + if constexpr (any_small_segments) + { + if constexpr (!only_small_segments) + { + // Zero-initialize the counters struct read by the agent's atomics. + if (const auto error = CubDebug(cudaMemsetAsync(allocations[1], 0, sizeof(counters_t), stream))) + { + return error; + } + } + const int grid_dim = static_cast(params::get_param(num_segments, 0)); + constexpr int block_dim = worker_per_segment_policy.threads_per_block; + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid_dim, block_dim, 0, stream) + .doit( + device_batched_topk_kernel, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + baseline_kernel_args{ + only_small_segments ? nullptr : static_cast(allocations[1]), + only_small_segments ? nullptr : static_cast(allocations[2]), + only_small_segments ? nullptr : static_cast(allocations[0])}, + cluster_kernel_args{}))) + { + return error; + } + } + else + { + // No small segments: compute the per-segment tile offsets directly via a transform-scan over all segment sizes. + if (const auto error = CubDebug(detail::scan::dispatch( + allocations[1], + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(allocations[0]), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(params::get_param(num_segments, 0)), + stream))) + { + return error; + } + } + + return CubDebug(detail::DebugSyncStream(stream)); + } +} + +#if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) +// Returns true if at least one architecture in the compile target list (`CMAKE_CUDA_ARCHITECTURES`, exposed as +// `::cuda::__target_compute_capabilities()`) resolves to the `unsupported` backend for `PolicySelector` -- e.g. a +// deterministic request while a pre-SM90 target is present in the list. Used to turn a would-be runtime +// `cudaErrorNotSupported` into a compile-time diagnostic (see the static_assert in `dispatch`). +template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL bool any_target_cc_unsupported() +{ + bool any = false; + for (const auto cc : ::cuda::__target_compute_capabilities()) + { + any = any || (PolicySelector{}(cc).backend == topk_backend::unsupported); + } + return any; +} +#endif // _CCCL_CUDA_COMPILATION() && !CUB_DEFINE_RUNTIME_POLICIES && !NVRTC + +// Internal entry point: the single dispatch that replaces the standalone baseline / cluster dispatches. It resolves the +// runtime compute capability, then uses `dispatch_compute_cap` to pick, per architecture, the backend chosen by +// `policy_selector` (deterministic -> cluster; otherwise the arch+size crossover). Both host arms launch the same +// kernel symbol. `Determinism`/`TieBreak` are compile-time selection inputs; `Mode` lets a caller force a backend. +// +// A non-`no_override` `PolicySelectorOverride` (threaded through the tuning environment) fully replaces the automatic +// selector -- its `.backend` chooses the arm and its `.baseline`/`.cluster` carry the tunings -- so a benchmark or +// tuning test can pick the backend and its knobs in one selector. +template < + ::cuda::execution::determinism::__determinism_t Determinism = + ::cuda::execution::determinism::__determinism_t::__not_guaranteed, + ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, + backend_mode Mode = backend_mode::automatic, + class PolicySelectorOverride = no_override, + class KeyInputItItT, + class KeyOutputItItT, + class ValueInputItItT, + class ValueOutputItItT, + class SegmentSizeParameterT, + class KParameterT, + class SelectDirectionT, + class NumSegmentsParameterT, + class TotalNumItemsGuaranteeT> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionT select_direction, + NumSegmentsParameterT num_segments, + [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream) +{ + // The selection direction is a compile-time constant carried as `::cuda::args::constant`. Wrap it into the + // internal discrete param the kernel/agent expect (both host arms take the wrapped form). + auto select_directions = wrap_select_direction(select_direction); + using SelectDirectionParameterT = decltype(select_directions); + + using key_t = it_value_t>; + using value_t = it_value_t>; + using LargeSegmentTileOffsetT = typename ::cuda::args::__traits::element_type; + + constexpr ::cuda::std::int64_t max_k = ::cuda::args::__traits::highest; + constexpr ::cuda::std::int64_t static_max_seg = ::cuda::args::__traits::highest; + + // Baseline sub-policy the kernel will actually instantiate: the tuning override's `.baseline` when an override is + // present (its sub-policies are forwarded verbatim by `selector_override_adaptor`), otherwise the default + // baseline sub-selector. Coverage must be computed from this same policy so the backend decision (and the + // override adaptor's `baseline_can_cover` it borrows) matches the policy `find_smallest_covering_policy` resolves. + using launched_baseline_selector_t = + ::cuda::std::conditional_t<::cuda::std::is_same_v, + baseline_policy_selector_from_types, + baseline_policy_selector_adaptor>; + + // Assert-free coverage predicate (never instantiates `find_smallest_covering_policy`'s hard static_assert). + constexpr bool baseline_can_cover = baseline_can_cover_v< + launched_baseline_selector_t, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>; + + // Default automatic selector from the compile-time inputs; a non-`no_override` override replaces it wholesale. + using default_selector_t = policy_selector_from_types< + key_t, + value_t, + ::cuda::std::int64_t, + max_k, + static_max_seg, + Determinism, + TieBreak, + baseline_can_cover, + Mode>; + using selector_t = ::cuda::std::conditional_t<::cuda::std::is_same_v, + default_selector_t, + selector_override_adaptor>; + +#if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ + && !defined(_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // Strict mode (default): fail at compile time if the request cannot be served on *any* architecture in + // `CMAKE_CUDA_ARCHITECTURES` (e.g. a deterministic / large-segment request while a pre-SM90 target is present, since + // the cluster backend requires SM90+). This is the least-surprising UX for callers building the default multi-arch + // preset. Define `_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` to defer the diagnosis to runtime instead (the dispatch + // then returns `cudaErrorNotSupported` on unsupported devices); CUB's own tests and benchmarks do this so they can + // compile the full configuration space across all target architectures and skip at runtime where unsupported. + static_assert( + !any_target_cc_unsupported(), + "cub::DeviceBatchedTopK: the requested top-k configuration is not supported on at least one architecture in " + "CMAKE_CUDA_ARCHITECTURES (the deterministic / large-segment cluster backend requires SM90+). Remove the " + "unsupported architecture(s), relax the request, or define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the " + "diagnosis to runtime (cudaErrorNotSupported)."); +#endif // strict unsupported-arch check + + detail::TripleChevronFactory launcher_factory{}; + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + return detail::dispatch_compute_cap(selector_t{}, cc, [&](auto policy_getter) -> cudaError_t { + CUB_DETAIL_CONSTEXPR_ISH auto active_policy = policy_getter(); + if CUB_DETAIL_CONSTEXPR_ISH (active_policy.backend == topk_backend::baseline) + { + return launch_baseline_arm( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + stream); + } + else if CUB_DETAIL_CONSTEXPR_ISH (active_policy.backend == topk_backend::cluster) + { + return launch_cluster_arm( + policy_getter, + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + stream); + } + else + { + // Unsupported on this architecture (e.g. a deterministic request on pre-SM90). Report a positive temp-storage + // size so the two-phase protocol proceeds, then fail the launch explicitly. + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return cudaErrorNotSupported; + } + }); +} } // namespace detail::batched_topk CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh b/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh deleted file mode 100644 index c4ab4b26410..00000000000 --- a/cub/cub/device/dispatch/dispatch_batched_topk_cluster.cuh +++ /dev/null @@ -1,797 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -//! @file -//! Cluster-based batched top-k dispatch. -//! -//! Prototype that launches a grid of thread block clusters to compute a -//! segmented top-k. Each cluster processes one segment end-to-end: private -//! histograms are reduced into the leader block via DSMEM atomics, then every -//! block reads the merged histogram back through DSMEM, locally identifies the -//! k-th bucket, and refines its shared-memory cluster_tile across radix passes. -//! -//! Two kernels share the agent body: -//! * Host: no `__cluster_dims__`, launched via `cudaLaunchKernelExC` with the -//! cluster blocks chosen at runtime (up to 16 on Hopper). -//! * CDP: static `__cluster_dims__(max_portable_cluster_blocks, 1, 1)`, -//! gated on `CUB_RDC_ENABLED` so CDP-disabled builds skip emitting it -//! (same pattern as `dispatch_segmented_sort.cuh`). - -#pragma once - -#include - -#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) -# pragma GCC system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) -# pragma clang system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) -# pragma system_header -#endif // no system header - -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include - -CUB_NAMESPACE_BEGIN - -namespace detail::batched_topk_cluster -{ -// ----------------------------------------------------------------------------- -// Cluster-size / dynamic-SMEM selection -// ----------------------------------------------------------------------------- -// Tightest upper bound carried by the segment-size argument. Mirrors `args::__traits<>::highest` semantics: -// the compile-time bound for `constant`/bounded sequence arguments and the runtime value for a uniform -// `immediate`. For a per-segment sequence with only a static bound this can be the loose `numeric_limits::max()`. -template -[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto runtime_max_segment_size(SegmentSizeParameterT segment_sizes) noexcept -{ - return ::cuda::args::__highest_(segment_sizes); -} - -// ----------------------------------------------------------------------------- -// Kernel entry points -// ----------------------------------------------------------------------------- -// Dynamic-cluster kernel for host launches; the agent reads the active cluster -// width from the `%cluster_nctarank` PTX special register. -template -__launch_bounds__(ThreadsPerBlock, MinBlocksPerSm) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t block_tile_capacity) -{ - using agent_t = agent_batched_topk_cluster< - ThreadsPerBlock, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleCtaMaxSegmentSize, - MinChunksPerCta, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - __shared__ typename agent_t::TempStorage temp_storage; - extern __shared__ char topk_cluster_smem[]; - char* key_slots = topk_cluster_smem; - // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` - // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. - { - ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); - smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); - asm("" : "+r"(smem32)); - key_slots = static_cast(__cvta_shared_to_generic(smem32)); - } - - agent_t agent( - temp_storage, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - key_slots, - block_tile_capacity); - - agent.Process(); -} - -#ifdef CUB_RDC_ENABLED -// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so the -// triple-chevron launch from device code needs no `cudaFuncSetAttribute`. -template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) - _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t block_tile_capacity) -{ - using agent_t = agent_batched_topk_cluster< - ThreadsPerBlock, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleCtaMaxSegmentSize, - MinChunksPerCta, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - __shared__ typename agent_t::TempStorage temp_storage; - extern __shared__ char topk_cluster_smem[]; - char* key_slots = topk_cluster_smem; - // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` - // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. - { - ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); - smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); - asm("" : "+r"(smem32)); - key_slots = static_cast(__cvta_shared_to_generic(smem32)); - } - - agent_t agent( - temp_storage, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - key_slots, - block_tile_capacity); - - agent.Process(); -} -#endif // CUB_RDC_ENABLED - -// ----------------------------------------------------------------------------- -// Dispatch -// ----------------------------------------------------------------------------- -// Keys and key/value pairs; segments larger than the resident block tile are streamed from gmem by the agent. The host -// path picks `(cluster_blocks, dynamic_smem_bytes)` at runtime via a wave-aware cluster-size search (see below); CDP -// uses the static kernel at `max_portable_cluster_blocks` and portable SMEM. - -// CDP launch body, empty when CDP is disabled. Wrapped in a macro because -// `#ifdef` can't sit inside `NV_IF_TARGET`. -#ifndef CUB_RDC_ENABLED -# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH -#else // CUB_RDC_ENABLED -# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ - auto static_kernel = device_segmented_topk_cluster_kernel_static< \ - ThreadsPerBlock, \ - HistogramItemsPerThread, \ - PipelineStages, \ - ChunkBytes, \ - LoadAlignBytes, \ - BitsPerPass, \ - TieBreakItemsPerThread, \ - SingleCtaMaxSegmentSize, \ - MinChunksPerCta, \ - CopyItemsPerThread, \ - Determinism, \ - TieBreak, \ - KeyInputItItT, \ - KeyOutputItItT, \ - ValueInputItItT, \ - ValueOutputItItT, \ - SegmentSizeParameterT, \ - KParameterT, \ - SelectDirectionParameterT, \ - NumSegmentsParameterT>; \ - if (const auto error = CubDebug( \ - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ - static_cast(grid_blocks), ThreadsPerBlock, dynamic_smem_bytes, stream) \ - .doit(static_kernel, \ - d_key_segments_it, \ - d_key_segments_out_it, \ - d_value_segments_it, \ - d_value_segments_out_it, \ - segment_sizes, \ - k_param, \ - select_directions, \ - num_segments, \ - block_tile_capacity))) \ - { \ - return error; \ - } -#endif // CUB_RDC_ENABLED - -// `Determinism`/`TieBreak` carry the requested `cuda::execution` requirements down to the kernel and agent (the public -// env-based interface is handled separately). They default to the nondeterministic, racing-atomics behavior. -template < - ::cuda::execution::determinism::__determinism_t Determinism = - ::cuda::execution::determinism::__determinism_t::__not_guaranteed, - ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, - typename KeyInputItItT, - typename KeyOutputItItT, - typename ValueInputItItT, - typename ValueOutputItItT, - typename SegmentSizeParameterT, - typename KParameterT, - typename SelectDirectionT, - typename NumSegmentsParameterT, - typename PolicySelector = policy_selector> -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( - void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionT select_direction, - NumSegmentsParameterT num_segments, - cudaStream_t stream = nullptr, - [[maybe_unused]] PolicySelector policy_selector = {}) -{ - // The selection direction is a compile-time constant carried as `::cuda::args::constant`. Wrap it into - // the internal discrete param the kernel/agent expect, exactly as the baseline `batched_topk::dispatch` does. - auto select_directions = batched_topk::wrap_select_direction(select_direction); - using SelectDirectionParameterT = decltype(select_directions); - // Clusters are SM 9.0+ only and the tuning is currently identical across - // CCs, so pin the selector query to the minimum supported CC. - constexpr cluster_topk_policy policy = PolicySelector{}(::cuda::compute_capability{9, 0}); - constexpr int ThreadsPerBlock = policy.threads_per_block; - constexpr int MinBlocksPerSm = policy.min_blocks_per_sm; - constexpr int HistogramItemsPerThread = policy.histogram_items_per_thread; - constexpr int PipelineStages = policy.pipeline_stages; - constexpr int ChunkBytes = policy.chunk_bytes; - constexpr int LoadAlignBytes = policy.load_align_bytes; - constexpr int BitsPerPass = policy.bits_per_pass; - constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; - constexpr int SingleCtaMaxSegmentSize = policy.single_cta_max_segment_size; - constexpr int MinChunksPerCta = policy.min_chunks_per_cta; - constexpr int CopyItemsPerThread = policy.copy_items_per_thread; - - using key_it_t = it_value_t; - using key_t = it_value_t; - using layout_t = smem_block_tile_layout; - using agent_t = agent_batched_topk_cluster< - ThreadsPerBlock, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleCtaMaxSegmentSize, - MinChunksPerCta, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - // TODO: This should be taken care of in the public env-based interface. - // A tie-break preference is only meaningful once the result set itself is deterministic. - static_assert(Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed - || TieBreak == ::cuda::execution::tie_break::__tie_break_t::__unspecified, - "A tie-break preference requires a deterministic execution requirement"); - - static_assert(ChunkBytes % LoadAlignBytes == 0); - static_assert(LoadAlignBytes % int{sizeof(key_t)} == 0); - // Static-footprint estimate for the device-side CDP fallback, which cannot query `cudaFuncGetAttributes`. - // The host path instead uses the driver-reported `sharedSizeBytes` (see below), which is padding-aware. - constexpr int static_smem_bytes = static_cast(sizeof(typename agent_t::TempStorage)); - - const auto max_seg_size = runtime_max_segment_size(segment_sizes); - - // The harness expects temp_storage_bytes > 0. Allocate a single byte placeholder. - size_t allocation_sizes[1] = {1}; - void* allocations[1] = {}; - if (const auto error = - CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) - { - return error; - } - - if (d_temp_storage == nullptr) - { - return cudaSuccess; - } - - static_assert(::cuda::args::__traits::is_single_value, - "Number of segments must be resolved on the host."); - - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - const auto num_seg_val = detail::params::get_param(num_segments, num_segments_val_t{0}); - if (num_seg_val == 0) - { - return cudaSuccess; - } - - // A zero bound would drive `clusterDim.x = 0`, which the runtime rejects. - if (max_seg_size == 0) - { - return cudaSuccess; - } - - // Cluster launches require compute capability 9.0+. - int sm_version = 0; - if (const auto error = CubDebug(SmVersionUncached(sm_version))) - { - return error; - } - if (sm_version < 900) - { - return cudaErrorNotSupported; - } - - constexpr auto dynamic_kernel = &device_segmented_topk_cluster_kernel< - ThreadsPerBlock, - MinBlocksPerSm, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleCtaMaxSegmentSize, - MinChunksPerCta, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - NV_IF_TARGET( - NV_IS_HOST, - ({ - // The launcher's `doit` carries the triple-chevron that makes NVCC emit `dynamic_kernel` for this TU, and - // performs the cluster launch via `cudaLaunchKernelEx`. The factory also wraps the pre-launch driver queries. - detail::TripleChevronFactory launcher_factory{}; - - // Opt in to non-portable cluster blocks (>8 on Hopper). - if (const auto error = launcher_factory.set_non_portable_cluster_allowed(dynamic_kernel)) - { - return error; - } - - // Config used only for the occupancy probe below; the final launch goes through `launcher_factory`. - // `clusterDim.x` is a placeholder since `cudaOccupancyMaxPotentialClusterSize` ignores it. - ::cudaLaunchAttribute cluster_attr{}; - cluster_attr.id = ::cudaLaunchAttributeClusterDimension; - cluster_attr.val.clusterDim.x = 1; - cluster_attr.val.clusterDim.y = 1; - cluster_attr.val.clusterDim.z = 1; - - ::cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(1, 1, 1); - cfg.blockDim = dim3(static_cast(ThreadsPerBlock), 1, 1); - cfg.dynamicSmemBytes = 0; - cfg.stream = stream; - cfg.attrs = &cluster_attr; - cfg.numAttrs = 1; - - // Resolve the per-block opt-in shared-memory budget and the kernel's static footprint from the driver so - // the dynamic-SMEM math below matches exactly what the launch permits. The opt-in budget - // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`) is the documented total per-block budget; the usable dynamic - // portion (`max_dynamic_smem_bytes`) is that budget minus the static footprint. - int device_id = 0; - if (const auto error = CubDebug(cudaGetDevice(&device_id))) - { - return error; - } - int max_smem_optin_bytes = 0; - if (const auto error = - CubDebug(cudaDeviceGetAttribute(&max_smem_optin_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id))) - { - return error; - } - // Use the driver-reported static footprint (`sharedSizeBytes`) rather than `sizeof(TempStorage)`: it reflects - // any padding the toolchain inserts to align the dynamic shared-memory section after the static one, so the - // derived dynamic sizes neither overshoot the budget nor conservatively drop the top table tier. - cudaFuncAttributes kernel_attrs{}; - if (const auto error = CubDebug(cudaFuncGetAttributes(&kernel_attrs, dynamic_kernel))) - { - return error; - } - // `cudaDevAttrMaxSharedMemoryPerBlockOptin` already excludes the driver's per-block reserved shared memory - // (opt-in == per-SM - reserved), so the dynamic budget is just the opt-in budget minus the static footprint; - // reserved must not be subtracted a second time. - const int nondynamic_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); - const int max_dynamic_smem_bytes = - (max_smem_optin_bytes > nondynamic_smem_bytes) ? max_smem_optin_bytes - nondynamic_smem_bytes : 0; - - // Raise the kernel's dynamic-SMEM opt-in lazily: occupancy queries and the launch must not request more than the - // currently configured `cudaFuncAttributeMaxDynamicSharedMemorySize`. The kernel's compile-time default already - // permits the portable 48 KiB total, i.e. that budget minus the static footprint. - constexpr int portable_total_smem_bytes = 48 * 1024; - int configured_dynamic_smem_limit = - (portable_total_smem_bytes > nondynamic_smem_bytes) ? portable_total_smem_bytes - nondynamic_smem_bytes : 0; - const auto ensure_dynamic_smem_limit = [&](int dynamic_smem_bytes) { - if (dynamic_smem_bytes <= configured_dynamic_smem_limit) - { - return cudaSuccess; - } - - if (const auto error = CubDebug(cudaFuncSetAttribute( - reinterpret_cast(dynamic_kernel), - cudaFuncAttributeMaxDynamicSharedMemorySize, - dynamic_smem_bytes))) - { - return error; - } - configured_dynamic_smem_limit = dynamic_smem_bytes; - return cudaSuccess; - }; - - // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); - // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more - // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We - // pick the `C` that minimizes the number of waves, breaking ties toward the largest `C` (= smallest SMEM = most - // L1), which matches the profiled fast configs. We enumerate `C` analytically rather than discovering SMEM tiers - // via occupancy, so a register-limited occupancy (e.g. 1 CTA/SM) cannot collapse the candidate set. - const auto seg = static_cast<::cuda::std::uint64_t>(max_seg_size); - const auto chunk_items_u64 = static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); - const int max_block_tile_capacity = static_cast(layout_t::block_tile_capacity(max_dynamic_smem_bytes)); - if (max_block_tile_capacity <= 0) - { - // Not even one load-aligned chunk fits in the opt-in budget; the kernel cannot run. - return cudaErrorInvalidValue; - } - - // `S_res(items)`: smallest chunk-granular dynamic SMEM whose per-CTA capacity reaches `items`. - const auto smem_for_block_capacity = [&](::cuda::std::uint64_t items) { - const auto slots = ::cuda::ceil_div(items, chunk_items_u64); - return layout_t::base_padding_bytes + static_cast(slots) * layout_t::chunk_bytes; - }; - - // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (cap HW - // max). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest fully-resident `C`. - // Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. - // `numeric_limits::max()` for an unbounded deferred sequence); narrowing such a `C_lo` to `int` could wrap to - // a small (or negative) value and wrongly enter the resident branch instead of the oversize/streaming fallback. - const int c_full = static_cast( - (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), - ::cuda::ceil_div(seg, chunk_items_u64))); - const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); - // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than - // necessary). At `min_chunks_per_cta == 1` this equals `c_full`; a larger knob shrinks it. - const int desired_cluster_blocks = static_cast(effective_cluster_blocks_from_chunks( - ::cuda::ceil_div(seg, chunk_items_u64), - MinChunksPerCta, - static_cast(max_supported_cluster_blocks))); - - int cluster_blocks = 0; - int dynamic_smem_sel = 0; - - if (single_cta_eligible(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity), SingleCtaMaxSegmentSize)) - { - // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's - // cluster-barrier-free path beats spreading it across more CTAs. `S_res(seg)` is within budget and one CTA is - // always launchable, so the occupancy probe is skipped (the shared `ensure_dynamic_smem_limit` below raises the - // opt-in for the selected SMEM). Larger fully-resident segments fall through to the wave-aware search below. - cluster_blocks = 1; - dynamic_smem_sel = smem_for_block_capacity(seg); - } - else if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) - { - // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every - // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. - // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. The - // upper bound is capped at the cluster blocks the max segment needs (`desired_cluster_blocks`), so the host - // never launches a wider cluster than necessary; at `min_chunks_per_cta == 1` the cap equals `c_full`. - const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); - const int c_end = (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks)); - ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); - for (int c = c_begin; c <= c_end; ++c) - { - const auto per_block_items = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(c)); - const int s_res = smem_for_block_capacity(per_block_items); - if (s_res > max_dynamic_smem_bytes) - { - continue; // unreachable for c >= C_lo, but guards the SMEM budget regardless. - } - - if (const auto error = ensure_dynamic_smem_limit(s_res)) - { - return error; - } - - // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must - // be a multiple of the cluster blocks. The returned value is the device-wide clusters-per-wave (capacity), - // independent of grid size, and accounts for the static footprint and register pressure internally. - cluster_attr.val.clusterDim.x = static_cast(c); - cfg.gridDim = dim3(static_cast(c), 1, 1); - cfg.dynamicSmemBytes = static_cast(s_res); - int clusters_per_wave = 0; - if (const auto error = launcher_factory.max_active_clusters(clusters_per_wave, dynamic_kernel, &cfg)) - { - return error; - } - if (clusters_per_wave <= 0) - { - continue; // cluster blocks not launchable at this SMEM. - } - - const auto waves = ::cuda::ceil_div( - static_cast<::cuda::std::uint64_t>(num_seg_val), static_cast<::cuda::std::uint64_t>(clusters_per_wave)); - // Min waves; tie-break largest `C`. The loop ascends in `C`, so `<=` keeps the largest at equal waves. - if (cluster_blocks == 0 || waves <= best_waves) - { - best_waves = waves; - cluster_blocks = c; - dynamic_smem_sel = s_res; - } - } - - if (cluster_blocks == 0 && c_lo == 1) - { - // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, - // but `C_lo == 1` guarantees `S_res(seg)` fits the budget and one CTA is always launchable. - cluster_blocks = 1; - dynamic_smem_sel = smem_for_block_capacity(seg); - } - } - - if (cluster_blocks == 0) - { - // Oversize (`C_lo > HW max`) or nothing launchable in range: full residency is impossible, so maximize - // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. - if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) - { - return error; - } - cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` - cfg.gridDim = dim3(1, 1, 1); - cfg.dynamicSmemBytes = static_cast(max_dynamic_smem_bytes); - int hw_max_cluster_blocks = 0; - if (const auto error = launcher_factory.max_potential_cluster_size(hw_max_cluster_blocks, dynamic_kernel, &cfg)) - { - return error; - } - hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, max_supported_cluster_blocks); - if (hw_max_cluster_blocks <= 0) - { - return cudaErrorInvalidValue; - } - cluster_blocks = hw_max_cluster_blocks; - dynamic_smem_sel = max_dynamic_smem_bytes; - } - - // The launch needs `MaxDynamicSharedMemorySize >= dynamic_smem_sel`; the scan already raised the limit past the - // largest probed SMEM, so this is a no-op unless the selected config skipped the scan. - if (const auto error = ensure_dynamic_smem_limit(dynamic_smem_sel)) - { - return error; - } - - const int dynamic_smem_bytes = dynamic_smem_sel; - const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); - - const auto grid_blocks = - static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cluster_blocks); - if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) - { - return cudaErrorInvalidValue; - } - - // The cluster dimension routes the launch through `cudaLaunchKernelEx`; the sibling triple-chevron in - // `doit_host` forces NVCC to emit `dynamic_kernel` for this TU. - if (const auto error = CubDebug( - launcher_factory(dim3(static_cast(grid_blocks), 1, 1), - dim3(static_cast(ThreadsPerBlock), 1, 1), - static_cast<::cuda::std::size_t>(dynamic_smem_bytes), - stream, - /*dependent_launch=*/false, - dim3(static_cast(cluster_blocks), 1, 1)) - .doit(dynamic_kernel, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - block_tile_capacity))) - { - return error; - } - }), - ({ - // CDP path: device-side launches cannot opt in to more than portable - // total SMEM or non-portable cluster blocks. Segments that exceed the - // portable resident coverage are still handled: the agent re-streams the - // overflow chunks from gmem. - constexpr int portable_total_smem_bytes = 48 * 1024; - constexpr int dynamic_smem_bytes = - (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; - - // The compile-time `ChunkBytes` is reused verbatim for the device launch: the agent peels the unaligned boundary - // edges into a tiny per-block buffer, so streaming needs only a single resident-or-streaming slot, and segments - // exceeding the small portable-SMEM block tile are handled by re-streaming overflow from gmem. The only hard - // requirement is that at least one load-aligned chunk fits the worst-case portable SMEM block tile. - constexpr auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); - static_assert(block_tile_capacity >= static_cast<::cuda::std::uint32_t>(layout_t::chunk_items), - "Portable SMEM is too small to fit even one load-aligned chunk for the device-launch (CDP) path"); - - const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) - * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); - if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) - { - return cudaErrorInvalidValue; - } - - CUB_TOPK_CLUSTER_DEVICE_LAUNCH - })); - - // Cluster launches can fail on the device while reporting success; sync. - if (const auto error = CubDebug(cudaPeekAtLastError())) - { - return error; - } - - return CubDebug(detail::DebugSyncStream(stream)); -} - -#undef CUB_TOPK_CLUSTER_DEVICE_LAUNCH - -// Env-based dispatch that also handles temporary-storage allocation. This is usually done by the device-layer, but -// there is no public API for cluster segmented top-k yet. Mirrors `batched_topk::dispatch_with_env`: the cluster -// algorithm is single-phase (one kernel launch over a placeholder allocation), but routing it through the shared -// env-based machinery keeps the call shape close to the baseline backend (the cluster path takes no total-items -// guarantee -- it needs no gmem large-segment tiling) and lets it pick up the stream, memory resource, and tuning -// carried by the environment. -// -// `Determinism`/`TieBreak` are forwarded verbatim to `dispatch` (the public env-based interface is handled separately); -// they default to the nondeterministic, unspecified-tie-break behavior. -template < - ::cuda::execution::determinism::__determinism_t Determinism = - ::cuda::execution::determinism::__determinism_t::__not_guaranteed, - ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, - typename KeyInputItItT, - typename KeyOutputItItT, - typename ValueInputItItT, - typename ValueOutputItItT, - typename SegmentSizeParameterT, - typename KParameterT, - typename SelectDirectionT, - typename NumSegmentsParameterT, - typename EnvT = ::cuda::std::execution::env<>> -[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionT select_direction, - NumSegmentsParameterT num_segments, - EnvT env = {}) -{ - return detail::dispatch_with_env_and_tuning( - env, [&](auto policy_sel, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { - return dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_direction, - num_segments, - stream, - policy_sel); - }); -} -} // namespace detail::batched_topk_cluster - -CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 9c66be1cba3..044258e776f 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -17,21 +17,28 @@ #endif // no system header #include +#include #include #include +#include #include #include +#include + +#include CUB_NAMESPACE_BEGIN namespace detail::batched_topk { -// Given a policy_selector and a segment-size parameter, resolves the agent type to be instantiated by the kernel. -// Selects the smallest policy whose tile size still covers the upper bound on segment size AND whose instantiated -// agent's shared memory usage fits within the static shared memory limit (max_smem_per_block). +// Assert-free search shared by `find_smallest_covering_policy` and the backend coverage predicate. Returns the +// index of the smallest worker policy whose tile size still covers the upper bound on segment size AND whose +// instantiated agent's shared memory usage fits within the static shared memory limit (max_smem_per_block), or -1 if +// none does. Kept separate from `find_smallest_covering_policy` so callers can query coverage as a bool without +// tripping that trait's hard `static_assert`. template -struct find_smallest_covering_policy +struct find_covering_policy_index { private: struct policy_t @@ -40,7 +47,7 @@ private: multi_worker_policy multi_worker_per_segment_policy; }; static constexpr ::cuda::std::int64_t max_segment_size = ::cuda::args::__traits::highest; - static constexpr batched_topk_policy active_policy = current_policy(); + static constexpr baseline_topk_policy active_policy = current_policy(); template [[nodiscard]] static constexpr int find_index() @@ -77,15 +84,38 @@ private: } } - static constexpr int selected_index = find_index<0>(); +public: + static constexpr int value = find_index<0>(); +}; + +// True iff some one-worker-per-segment policy covers the statically-known maximum segment size within the shared-memory +// limit. Used by the backend selector to decide whether the baseline backend is viable at all (an oversize +// bound must route to the cluster backend instead of tripping the `static_assert` below). +template +inline constexpr bool baseline_can_cover_v = + find_covering_policy_index::value >= 0; + +// Resolves the agent type the kernel instantiates via the same covering-policy search as `find_covering_policy_index`, +// adding a hard `static_assert` when no policy covers the segment size within the shared-memory limit. +template +struct find_smallest_covering_policy +{ +private: + struct policy_t + { + worker_policy worker_per_segment_policy; + multi_worker_policy multi_worker_per_segment_policy; + }; + static constexpr baseline_topk_policy active_policy = current_policy(); + static constexpr int selected_index = + find_covering_policy_index::value; public: // TODO (elstehle): extend support for variable-size segments static_assert(selected_index >= 0, - "cub::DeviceBatchedTopK currently supports only segments small enough to be processed by a single " - "thread block (one worker per segment). No policy could cover the statically-known maximum segment " - "size within the shared-memory limit. Reduce the maximum segment size encoded in the segment-size " - "argument annotation."); + "cub::DeviceBatchedTopK: no baseline worker policy covers the statically-known maximum segment size " + "within the shared-memory limit. Reduce the maximum segment size encoded in the segment-size argument " + "annotation (larger segments are served by the SM 9.0+ cluster backend)."); static constexpr policy_t policy = { active_policy.worker_per_segment_policies[selected_index], active_policy.multi_worker_per_segment_policy}; @@ -113,7 +143,7 @@ template #if _CCCL_HAS_CONCEPTS() - requires batched_topk_policy_selector + requires baseline_topk_policy_selector #endif // _CCCL_HAS_CONCEPTS() __launch_bounds__(int( find_smallest_covering_policy< @@ -181,6 +211,238 @@ __launch_bounds__(int( // Process segments agent.Process(); } + +// ----------------------------------------------------------------------------- +// Single kernel symbol hosting both backends +// ----------------------------------------------------------------------------- +// There is exactly one kernel symbol per instantiation. Its body selects the active backend device-side via +// `current_policy()` (evaluated per `__CUDA_ARCH__` pass), so each target architecture compiles only +// the backend the selector picks for it -- honoring CUB's "one kernel per arch/problem" rule while still supporting a +// multi-architecture fatbin whose per-arch choice differs. The host still branches its launch configuration (grid, +// shared memory, cluster dimensions) per backend, but both host arms launch this same symbol. + +// Backend-specific kernel arguments. The unused struct is passed default-constructed (all-null / zero) to the arm the +// selector does not pick; passing it costs nothing (a few grid-constant scalars) and keeps a single kernel signature. +template +struct baseline_kernel_args +{ + batched_topk_counters* d_counters = nullptr; + NumSegmentsValueT* d_large_segments_ids = nullptr; + LargeSegmentTileOffsetT* d_large_segments_tile_offsets = nullptr; +}; + +struct cluster_kernel_args +{ + ::cuda::std::uint32_t block_tile_capacity = 0; +}; + +// ----------------------------------------------------------------------------- +// Launch-bounds helpers +// ----------------------------------------------------------------------------- +// The two backends use different `__launch_bounds__` shapes (baseline: just threads_per_block; cluster: threads plus a +// min-blocks-per-SM). We resolve both per architecture from the selected policy. `find_smallest_covering_policy` (which +// carries a hard `static_assert`) is only ever touched inside the `backend == baseline` branch, so an oversize bound +// routed to the cluster/unsupported backend never trips it. +_CCCL_EXEC_CHECK_DISABLE +template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL int topk_threads_per_block_helper() noexcept +{ + constexpr auto policy = current_policy(); + if constexpr (policy.backend == topk_backend::baseline) + { + return find_smallest_covering_policy, + SegmentSizeParameterT, + AgentParamsT...>::policy.worker_per_segment_policy.threads_per_block; + } + else if constexpr (policy.backend == topk_backend::cluster) + { + return policy.cluster.threads_per_block; + } + else + { + // unsupported: harmless positive default; the host never launches this arm. + return 128; + } +} + +_CCCL_EXEC_CHECK_DISABLE +template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL int topk_min_blocks_per_sm_helper() noexcept +{ + constexpr auto policy = current_policy(); + if constexpr (policy.backend == topk_backend::cluster) + { + return policy.cluster.min_blocks_per_sm; + } + else + { + // baseline / unsupported: no minimum-blocks constraint. + return 0; + } +} + +// Variable templates force constant evaluation of the helpers, otherwise nvcc reports a "bad attribute argument +// substitution" error on the `__launch_bounds__` below (same pattern as `transform_kernel`). +template +inline constexpr int topk_threads_per_block = + topk_threads_per_block_helper(); + +template +inline constexpr int topk_min_blocks_per_sm = + topk_min_blocks_per_sm_helper(); + +// ----------------------------------------------------------------------------- +// Global kernel entry point (single symbol for both backends) +// ----------------------------------------------------------------------------- +template +__launch_bounds__( + topk_threads_per_block< + PolicySelector, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>, + topk_min_blocks_per_sm< + PolicySelector, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>) + _CCCL_KERNEL_ATTRIBUTES void device_batched_topk_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + baseline_kernel_args::element_type, LargeSegmentTileOffsetT> + base_args, + [[maybe_unused]] cluster_kernel_args clus_args) +{ + constexpr auto policy = current_policy(); + + if constexpr (policy.backend == topk_backend::baseline) + { + using agent_t = typename find_smallest_covering_policy< + baseline_policy_selector_adaptor, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + LargeSegmentTileOffsetT>::agent_t; + + // Static assertions (constraints). + static_assert(agent_t::tile_size >= ::cuda::args::__traits::highest, + "Block size exceeds maximum segment size supported by SegmentSizeParameterT"); + static_assert(sizeof(typename agent_t::TempStorage) <= max_smem_per_block, + "Static shared memory per block must not exceed 48KB limit."); + + __shared__ typename agent_t::TempStorage temp_storage; + + agent_t agent( + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + base_args.d_counters, + base_args.d_large_segments_ids, + base_args.d_large_segments_tile_offsets); + + agent.Process(); + } + else if constexpr (policy.backend == topk_backend::cluster) + { + NV_IF_TARGET( + NV_PROVIDES_SM_90, + (using agent_t = batched_topk_cluster::agent_batched_topk_cluster< + policy.cluster.threads_per_block, + policy.cluster.histogram_items_per_thread, + policy.cluster.pipeline_stages, + policy.cluster.chunk_bytes, + policy.cluster.load_align_bytes, + policy.cluster.bits_per_pass, + policy.cluster.tie_break_items_per_thread, + policy.cluster.single_block_max_seg_size, + policy.cluster.min_chunks_per_block, + policy.cluster.copy_items_per_thread, + PolicySelector::determinism, + PolicySelector::tie_break, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same + // `load_align` alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves + // `base_padding_bytes` for this. + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } + + agent_t agent( + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + key_slots, + clus_args.block_tile_capacity); + + agent.Process();)); + } + else + { + // topk_backend::unsupported: the host arm returns cudaErrorNotSupported before launching, so this never runs. + return; + } +} } // namespace detail::batched_topk CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh index 6ea6a2fa7cd..540c6a4ad8a 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh @@ -16,20 +16,27 @@ #include #include #include +#include +#include #include +#include +#include #include #include +#include CUB_NAMESPACE_BEGIN namespace detail::batched_topk { +//! Sub-policy for the compaction epilogue shared by the baseline @ref DeviceBatchedTopK workers: it scans the radix +//! histogram and writes out the selected keys. struct epilogue_policy { - int items_per_thread; - BlockLoadAlgorithm load_algorithm; - BlockStoreAlgorithm store_algorithm; - BlockScanAlgorithm scan_algorithm; + int items_per_thread; //!< Keys each thread loads/stores per tile in the epilogue. + BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read keys back in the epilogue. + BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys. + BlockScanAlgorithm scan_algorithm; //!< Block scan algorithm used for the histogram prefix sum. _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const epilogue_policy& lhs, const epilogue_policy& rhs) { @@ -52,14 +59,16 @@ struct epilogue_policy #endif // !_CCCL_COMPILER(NVRTC) }; +//! Per-segment worker sub-policy for the baseline backend: one thread block cooperatively computes the top-k of a +//! single segment. @ref baseline_topk_policy holds several of these, ordered by decreasing tile size. struct worker_policy { - int threads_per_block; - int items_per_thread; - BlockLoadAlgorithm load_algorithm; - BlockStoreAlgorithm store_algorithm; + int threads_per_block; //!< Number of threads in a CUDA block. + int items_per_thread; //!< Keys each thread loads/processes per tile (with `threads_per_block` sets the tile size). + BlockLoadAlgorithm load_algorithm; //!< Block load algorithm used to read the segment's keys. + BlockStoreAlgorithm store_algorithm; //!< Block store algorithm used to write the selected keys. - epilogue_policy epilogue; + epilogue_policy epilogue; //!< Sub-policy for the compaction epilogue. _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const worker_policy& lhs, const worker_policy& rhs) { @@ -83,10 +92,12 @@ struct worker_policy #endif // !_CCCL_COMPILER(NVRTC) }; +//! Sub-policy for the baseline backend's multiple-blocks-per-segment worker path, used for segments too large for a +//! single worker block. struct multi_worker_policy { - int threads_per_block; - int items_per_thread; + int threads_per_block; //!< Number of threads in a CUDA block. + int items_per_thread; //!< Keys each thread loads/processes per tile. _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const multi_worker_policy& lhs, const multi_worker_policy& rhs) { @@ -107,28 +118,31 @@ struct multi_worker_policy #endif // _CCCL_HOSTED() }; -struct batched_topk_policy +//! Sub-policy for the baseline (worker-per-segment) backend of @ref DeviceBatchedTopK. +struct baseline_topk_policy { - // The list of per-segment agent policies is ordered by decreasing tile size. At compile time, the smallest policy - // whose tile size still covers the upper bound of the segment size is selected. + //! Per-segment worker policies ordered by decreasing tile size. At compile time the smallest policy whose tile size + //! still covers the upper bound of the segment size is selected. ::cuda::std::array worker_per_segment_policies; - multi_worker_policy multi_worker_per_segment_policy; + multi_worker_policy multi_worker_per_segment_policy; //!< Worker policy for segments too large for a single block. - _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API constexpr friend bool + operator==(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs) { return lhs.worker_per_segment_policies == rhs.worker_per_segment_policies && lhs.multi_worker_per_segment_policy == rhs.multi_worker_per_segment_policy; } - _CCCL_HOST_DEVICE_API constexpr friend bool operator!=(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API constexpr friend bool + operator!=(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs) { return !(lhs == rhs); } #if _CCCL_HOSTED() - friend ::std::ostream& operator<<(::std::ostream& os, const batched_topk_policy& p) + friend ::std::ostream& operator<<(::std::ostream& os, const baseline_topk_policy& p) { - os << "batched_topk_policy { .worker_per_segment_policies = { "; + os << "baseline_topk_policy { .worker_per_segment_policies = { "; for (::cuda::std::size_t i = 0; i < p.worker_per_segment_policies.size(); ++i) { if (i != 0) @@ -144,42 +158,424 @@ struct batched_topk_policy #if _CCCL_HAS_CONCEPTS() template -concept batched_topk_policy_selector = policy_selector; +concept baseline_topk_policy_selector = policy_selector; #endif // _CCCL_HAS_CONCEPTS() -struct policy_selector +// Default baseline sub-policy. A free function so both `baseline_policy_selector` and the combined `policy_selector` +// can build it inline. Tuning is currently CC-independent. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_baseline_policy() -> baseline_topk_policy +{ + constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE; + constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE; + constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS; + constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg}; + return baseline_topk_policy{ + {{ + worker_policy{256, 64, load_alg, store_alg, epilogue}, + worker_policy{256, 32, load_alg, store_alg, epilogue}, + worker_policy{256, 16, load_alg, store_alg, epilogue}, + worker_policy{256, 8, load_alg, store_alg, epilogue}, + worker_policy{256, 4, load_alg, store_alg, epilogue}, + worker_policy{128, 2, load_alg, store_alg, epilogue}, + }}, + multi_worker_policy{256, 64}}; +} + +// Largest maximum segment size (in keys) the baseline (worker-per-segment) backend can cover: the largest worker tile +// (threads_per_block * items_per_thread) in `policy`. A larger statically-known maximum segment size makes the baseline +// backend ineligible (the selector then picks the cluster backend where supported, otherwise `unsupported`). This is +// only the tile-based necessary condition; the exact predicate `baseline_can_cover_v` also checks the agent's +// shared-memory fit (which needs the concrete agent types). +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::int64_t +baseline_max_covered_segment_size(const baseline_topk_policy& policy) +{ + ::cuda::std::int64_t max_tile_size = 0; + for (const auto& worker : policy.worker_per_segment_policies) + { + const ::cuda::std::int64_t tile_size = ::cuda::std::int64_t{worker.threads_per_block} * worker.items_per_thread; + if (tile_size > max_tile_size) + { + max_tile_size = tile_size; + } + } + return max_tile_size; +} + +struct baseline_policy_selector { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> batched_topk_policy - { - constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE; - constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE; - constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS; - constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg}; - return batched_topk_policy{ - {{ - worker_policy{256, 64, load_alg, store_alg, epilogue}, - worker_policy{256, 32, load_alg, store_alg, epilogue}, - worker_policy{256, 16, load_alg, store_alg, epilogue}, - worker_policy{256, 8, load_alg, store_alg, epilogue}, - worker_policy{256, 4, load_alg, store_alg, epilogue}, - worker_policy{128, 2, load_alg, store_alg, epilogue}, - }}, - multi_worker_policy{256, 64}}; + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const + -> baseline_topk_policy + { + return make_baseline_policy(); } }; template +struct baseline_policy_selector_from_types +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> baseline_topk_policy + { + return baseline_policy_selector{}(cc); + } +}; + +#if _CCCL_HAS_CONCEPTS() +static_assert(baseline_topk_policy_selector); +#endif // _CCCL_HAS_CONCEPTS() + +//! Per-block execution shape for the thread-block-cluster backend of @ref DeviceBatchedTopK. The dispatch picks the +//! number of cluster blocks and the dynamic shared-memory block_tile capacity at runtime (occupancy / wave-aware), so +//! this policy carries only the per-block tuning knobs. +struct cluster_topk_policy +{ + // Fields grouped by kind and, across and within groups, ordered by kernel use: launch config, then load, then + // process, then output. + + int threads_per_block; //!< Number of threads in a CUDA block. + int min_blocks_per_sm; //!< Minimum resident blocks per SM, forwarded as the kernel launch-bounds occupancy hint. + + int min_chunks_per_block; //!< Minimum number of chunks a block must own to join a segment's effective cluster (the + //!< divisor mapping a segment's chunk count to its cluster width). Must be >= 1. + + int chunk_bytes; //!< Size in bytes of one block_tile chunk -- the granularity of the async-copy load pipeline. + int load_align_bytes; //!< Load / bulk-copy alignment in bytes. Must be a power of two and >= 16 + //!< (`detail::bulk_copy_min_align`), and `chunk_bytes` must be a multiple of it (see + //!< `is_valid_cluster_policy`). + int pipeline_stages; //!< Depth of the async-copy (mbarrier) pipeline that stages chunks into shared memory. + + int single_block_max_seg_size; //!< Largest segment size, in keys, still eligible for the single-block fast path + //!< (kept out of the byte-unit loading group above as it is measured in items). + int bits_per_pass; //!< Radix digit width per pass; each pass' histogram spans `1 << bits_per_pass` buckets. Together + //!< with `threads_per_block` this implicitly fixes the histogram block-scan's items per thread, + //!< `ceil_div(1 << bits_per_pass, threads_per_block)` buckets scanned per thread. + int histogram_items_per_thread; //!< Keys each thread accumulates per tile during the radix histogram passes. + int tie_break_items_per_thread; //!< Keys each thread processes per tile during the final tie-break / filter phase. + int copy_items_per_thread; //!< Keys each thread copies per tile on the select-all (k >= segment size) fast path. + + // Equality/streaming make this a regular type (required by the `policy_selector` concept / `dispatch_compute_cap`). + _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) + { + return lhs.threads_per_block == rhs.threads_per_block && lhs.min_blocks_per_sm == rhs.min_blocks_per_sm + && lhs.min_chunks_per_block == rhs.min_chunks_per_block && lhs.chunk_bytes == rhs.chunk_bytes + && lhs.load_align_bytes == rhs.load_align_bytes && lhs.pipeline_stages == rhs.pipeline_stages + && lhs.single_block_max_seg_size == rhs.single_block_max_seg_size && lhs.bits_per_pass == rhs.bits_per_pass + && lhs.histogram_items_per_thread == rhs.histogram_items_per_thread + && lhs.tie_break_items_per_thread == rhs.tie_break_items_per_thread + && lhs.copy_items_per_thread == rhs.copy_items_per_thread; + } + + _CCCL_HOST_DEVICE_API constexpr friend bool operator!=(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) + { + return !(lhs == rhs); + } + +#if _CCCL_HOSTED() + friend ::std::ostream& operator<<(::std::ostream& os, const cluster_topk_policy& p) + { + return os + << "cluster_topk_policy { .threads_per_block = " << p.threads_per_block + << ", .min_blocks_per_sm = " << p.min_blocks_per_sm << ", .min_chunks_per_block = " << p.min_chunks_per_block + << ", .chunk_bytes = " << p.chunk_bytes << ", .load_align_bytes = " << p.load_align_bytes + << ", .pipeline_stages = " << p.pipeline_stages + << ", .single_block_max_seg_size = " << p.single_block_max_seg_size << ", .bits_per_pass = " << p.bits_per_pass + << ", .histogram_items_per_thread = " << p.histogram_items_per_thread << ", .tie_break_items_per_thread = " + << p.tie_break_items_per_thread << ", .copy_items_per_thread = " << p.copy_items_per_thread << " }"; + } +#endif // _CCCL_HOSTED() +}; + +#if _CCCL_HAS_CONCEPTS() +template +concept cluster_topk_policy_selector = policy_selector; +#endif // _CCCL_HAS_CONCEPTS() + +// Default cluster sub-policy. A free function so both `cluster_policy_selector` and the combined `policy_selector` +// can build it inline. Tuning is currently CC-independent. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_cluster_policy() -> cluster_topk_policy +{ + return cluster_topk_policy{ + /*threads_per_block=*/512, + /*min_blocks_per_sm=*/1, + /*min_chunks_per_block=*/1, + /*chunk_bytes=*/16 * 1024, + /*load_align_bytes=*/128, + /*pipeline_stages=*/8, + /*single_block_max_seg_size=*/8 * 1024, + /*bits_per_pass=*/11, + /*histogram_items_per_thread=*/8, + /*tie_break_items_per_thread=*/8, + /*copy_items_per_thread=*/8}; +} + +// Hard constraints on the block_tile byte geometry. The aligned bulk-copy (TMA) load path addresses gmem/smem in +// `load_align_bytes`-sized, aligned units, so `load_align_bytes` must be a power of two and at least +// `bulk_copy_min_align` (16 B), and `chunk_bytes` must be a whole number of those units. `launch_cluster_arm` +// static-asserts this before instantiating the agent. +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_cluster_policy(cluster_topk_policy policy) -> bool +{ + return policy.load_align_bytes >= bulk_copy_min_align && ::cuda::is_power_of_two(policy.load_align_bytes) + && policy.chunk_bytes % policy.load_align_bytes == 0; +} + +static_assert(is_valid_cluster_policy(make_cluster_policy())); + +// Default selector for cluster-capable architectures (SM 9.0+). The tuning is currently identical across CCs. +struct cluster_policy_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy + { + return make_cluster_policy(); + } +}; + +#if _CCCL_HAS_CONCEPTS() +static_assert(cluster_topk_policy_selector); +#endif // _CCCL_HAS_CONCEPTS() + +// ----------------------------------------------------------------------------- +// Backend selection +// ----------------------------------------------------------------------------- +//! Backend algorithms for @ref DeviceBatchedTopK. Both backends are launched through a single kernel symbol; which one +//! runs is decided per architecture by `policy_selector` below, whose result also drives the device-side agent +//! selection (via `current_policy`). +enum class topk_backend +{ + baseline, //!< worker-per-segment backend (single thread block per segment) + cluster, //!< thread-block-cluster backend (SM 9.0+) + unsupported //!< no backend can serve the request on the target architecture; dispatch returns cudaErrorNotSupported +}; + +#if _CCCL_HOSTED() +[[nodiscard]] inline ::std::ostream& operator<<(::std::ostream& os, topk_backend backend) +{ + switch (backend) + { + case topk_backend::baseline: + return os << "baseline"; + case topk_backend::cluster: + return os << "cluster"; + default: + return os << "unsupported"; + } +} +#endif // _CCCL_HOSTED() + +//! The tuning policy for all backends of @ref DeviceBatchedTopK. It carries the selected backend plus both backends' +//! sub-policies; the kernel instantiates only the arm named by @p backend (chosen device-side via `current_policy`). +//! +//! This is a regular type: `detail::dispatch_compute_cap` (and the `policy_selector` concept) require the selector's +//! result to be `::cuda::std::regular`, hence the equality/streaming operators below. +struct topk_policy +{ + topk_backend backend; //!< Backend the dispatch selected, i.e. the kernel arm that runs. + baseline_topk_policy baseline; //!< Sub-policy used when @p backend is @p topk_backend::baseline. + cluster_topk_policy cluster; //!< Sub-policy used when @p backend is @p topk_backend::cluster. + + _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const topk_policy& lhs, const topk_policy& rhs) + { + return lhs.backend == rhs.backend && lhs.baseline == rhs.baseline && lhs.cluster == rhs.cluster; + } + + _CCCL_HOST_DEVICE_API constexpr friend bool operator!=(const topk_policy& lhs, const topk_policy& rhs) + { + return !(lhs == rhs); + } + +#if _CCCL_HOSTED() + friend ::std::ostream& operator<<(::std::ostream& os, const topk_policy& p) + { + return os << "topk_policy { .backend = " << p.backend << ", .baseline = " << p.baseline + << ", .cluster = " << p.cluster << " }"; + } +#endif // _CCCL_HOSTED() +}; + +// Sentinel meaning "no backend-selecting tuning override was provided", so the dispatch builds its automatic selector. +// Used as the default for `dispatch`'s `PolicySelectorOverride` and as the not-found result of the public +// API's `topk_policy` tuning query (a real empty type -- `void` cannot be a query default). +struct no_override +{}; + +// Crossover knobs (TODO: tune via SM100 benchmarks). +//! Clusters require SM 9.0+. +inline constexpr int cluster_min_cc_major = 9; +//! The cluster backend only consistently wins at/above SM 10.0 (measured on B200), and only for large segments. +inline constexpr int cluster_beneficial_min_cc_major = 10; +//! Smallest statically-known maximum segment size at which the cluster backend starts to win (measured on B200). This +//! is the backend crossover threshold and is intentionally part of the selector -- not a tunable policy field -- so +//! that tuning the cluster policy (e.g. its single-CTA threshold) does not silently shift which backend is chosen. +inline constexpr ::cuda::std::int64_t cluster_beneficial_min_segment_size = 8 * 1024; + +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool cluster_capable(::cuda::compute_capability cc) +{ + return cc >= ::cuda::compute_capability{cluster_min_cc_major, 0}; +} + +// How the backend is chosen. `automatic` applies the determinism/arch/size rules below; the `force_*` modes let a +// caller pin a specific backend (e.g. a tuning override that forces the cluster backend while still launching the +// single kernel symbol). +enum class backend_mode +{ + automatic, + force_baseline, + force_cluster, +}; + +// Field-based backend selector (like DeviceScan's / DeviceTransform's `policy_selector`): one selector that builds both +// sub-policies inline and makes the backend decision from plain value fields (keeping it a regular value type). +// `baseline_can_cover` is supplied by the dispatch, which alone knows the concrete agent types needed for the +// shared-memory fit, so this selector stays free of the agent-type-dependent `find_smallest_covering_policy` machinery. +struct policy_selector +{ + ::cuda::std::int64_t static_max_segment_size; + ::cuda::execution::determinism::__determinism_t determinism; + ::cuda::execution::tie_break::__tie_break_t tie_break; + bool baseline_can_cover; + backend_mode mode; + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> topk_policy + { + const auto baseline = make_baseline_policy(); + const auto cluster = make_cluster_policy(); + + // A deterministic result set (or a concrete tie-break preference) can only be honored by the cluster backend. + const bool deterministic = (determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) + || (tie_break != ::cuda::execution::tie_break::__tie_break_t::__unspecified); + + topk_backend backend = topk_backend::unsupported; + if (mode == backend_mode::force_cluster) + { + backend = cluster_capable(cc) ? topk_backend::cluster : topk_backend::unsupported; + } + else if (mode == backend_mode::force_baseline) + { + // The baseline backend cannot honor a deterministic result set / concrete tie-break preference, nor cover an + // oversize segment; reject (map to `unsupported`) in those cases rather than pinning a backend that cannot serve + // the request -- matching the hard constraints `selector_override_adaptor` enforces. + backend = (baseline_can_cover && !deterministic) ? topk_backend::baseline : topk_backend::unsupported; + } + else if (deterministic) + { + // Deterministic -> cluster (arch permitting), independent of the max segment size. + backend = cluster_capable(cc) ? topk_backend::cluster : topk_backend::unsupported; + } + else if (!baseline_can_cover) + { + // Oversize for the baseline backend: it must never be selected (its `find_smallest_covering_policy` would fail). + backend = cluster_capable(cc) ? topk_backend::cluster : topk_backend::unsupported; + } + else + { + // Baseline can cover: prefer the cluster backend only where it is beneficial, otherwise use the baseline. The + // size crossover is a fixed selector constant (not read from the tunable cluster policy), so tuning the cluster + // policy never shifts the backend choice. + const bool beneficial = cc >= ::cuda::compute_capability{cluster_beneficial_min_cc_major, 0} + && static_max_segment_size >= cluster_beneficial_min_segment_size; + backend = (cluster_capable(cc) && beneficial) ? topk_backend::cluster : topk_backend::baseline; + } + return topk_policy{backend, baseline, cluster}; + } +}; + +// Stateless selector built purely from the compile-time request facts; default-constructs the field-based +// `policy_selector` and delegates. This is the type threaded into the dispatch and kernel: `current_policy` / +// `dispatch_compute_cap` default-construct it, so the behavior must live in the type. The facts are re-exposed as +// static members so the tuning-override adaptor can borrow them. +template struct policy_selector_from_types +{ + static constexpr auto determinism = Determinism; + static constexpr auto tie_break = TieBreak; + static constexpr bool baseline_can_cover = BaselineCanCover; + static constexpr ::cuda::std::int64_t static_max_segment_size = StaticMaxSegSize; + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> topk_policy + { + return policy_selector{StaticMaxSegSize, Determinism, TieBreak, BaselineCanCover, Mode}(cc); + } +}; + +// Adapts a tune-provided selector (which only implements `operator() -> topk_policy`) to the full `PolicySelector` +// interface. The request facts (`baseline_can_cover` / `determinism` / `tie_break` / `static_max_segment_size`) are +// borrowed from the dispatch's automatic selector -- they are properties of the call, not the tuning. +// `baseline_can_cover` reflects the *default* baseline sub-policy's coverage (the adaptor lacks the concrete agent +// types needed to recompute it for an override's baseline sub-policy), so a tuning override that changes the baseline +// tile shape is still validated against the default's coverage. +// +// The override's sub-policies are forwarded verbatim and its `.backend` is honored, but validated against the chosen +// backend's hard constraints: an override selecting a backend that cannot serve the request is *rejected* (mapped to +// `unsupported`) rather than silently rerouted, surfacing the misconfiguration (compile error in strict mode, else +// cudaErrorNotSupported at runtime). Rejected cases: +// * baseline for a deterministic / tie-break request or a segment size it cannot cover, and +// * cluster on a pre-SM90 architecture. +template +struct selector_override_adaptor +{ + static constexpr auto determinism = Default::determinism; + static constexpr auto tie_break = Default::tie_break; + static constexpr bool baseline_can_cover = Default::baseline_can_cover; + static constexpr ::cuda::std::int64_t static_max_segment_size = Default::static_max_segment_size; + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> topk_policy + { + const auto overridden = Override{}(cc); + + constexpr bool deterministic = (determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) + || (tie_break != ::cuda::execution::tie_break::__tie_break_t::__unspecified); + + topk_backend backend = overridden.backend; + if (backend == topk_backend::baseline) + { + // The baseline backend cannot honor a deterministic / tie-break request, nor cover an oversize segment. + if (deterministic || !baseline_can_cover) + { + backend = topk_backend::unsupported; + } + } + else if (backend == topk_backend::cluster) + { + // The cluster backend requires SM 9.0+. + if (!cluster_capable(cc)) + { + backend = topk_backend::unsupported; + } + } + return topk_policy{backend, overridden.baseline, overridden.cluster}; + } +}; + +// Adapts a (combined) policy selector to a plain baseline policy selector (returns just the `.baseline` sub-policy), so +// the kernel can drive `find_smallest_covering_policy` from a single `PolicySelector` template parameter. +template +struct baseline_policy_selector_adaptor { [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const - -> batched_topk_policy + -> baseline_topk_policy { - return policy_selector{}(cc); + return PolicySelector{}(cc).baseline; } }; #if _CCCL_HAS_CONCEPTS() -static_assert(batched_topk_policy_selector); +static_assert( + detail::policy_selector, + topk_policy>); #endif // _CCCL_HAS_CONCEPTS() } // namespace detail::batched_topk diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh deleted file mode 100644 index 5bf9c7fc127..00000000000 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk_cluster.cuh +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -#pragma once - -#include - -#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) -# pragma GCC system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) -# pragma clang system_header -#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) -# pragma system_header -#endif // no system header - -#include - -CUB_NAMESPACE_BEGIN -namespace detail::batched_topk_cluster -{ -// Per-block execution shape. The dispatch picks the cluster blocks and the dynamic shared-memory block_tile capacity at -// runtime (occupancy/wave-aware), so the policy carries only the per-block tuning knobs. -struct cluster_topk_policy -{ - int threads_per_block; - int histogram_items_per_thread; - int pipeline_stages; - int chunk_bytes; - int load_align_bytes; - int bits_per_pass; - int min_blocks_per_sm; - int tie_break_items_per_thread; - int single_cta_max_segment_size; - int min_chunks_per_cta; - int copy_items_per_thread; -}; - -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_policy() -> cluster_topk_policy -{ - return cluster_topk_policy{ - /*threads_per_block=*/512, - /*histogram_items_per_thread=*/8, - /*pipeline_stages=*/8, - /*chunk_bytes=*/16 * 1024, - /*load_align_bytes=*/128, - /*bits_per_pass=*/11, - /*min_blocks_per_sm=*/1, - /*tie_break_items_per_thread=*/8, - /*single_cta_max_segment_size=*/8 * 1024, - /*min_chunks_per_cta=*/1, - /*copy_items_per_thread=*/8}; -} - -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_policy(cluster_topk_policy policy) -> bool -{ - return policy.chunk_bytes % policy.load_align_bytes == 0; -} - -static_assert(is_valid_policy(make_policy())); -// Default selector for cluster-capable architectures (SM 9.0+). The tuning is currently identical across CCs. -struct policy_selector -{ - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cluster_topk_policy - { - return make_policy(); - } -}; -} // namespace detail::batched_topk_cluster - -CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_segmented_topk_cluster_layout.cu deleted file mode 100644 index a943e172c15..00000000000 --- a/cub/test/catch2_test_segmented_topk_cluster_layout.cu +++ /dev/null @@ -1,107 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -#include - -#include - -#include - -namespace -{ -template -[[nodiscard]] constexpr SizeT host_ceil_div(SizeT numerator, SizeT denominator) -{ - return (numerator + denominator - 1) / denominator; -} - -template -void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) -{ - using layout_t = cub::detail::batched_topk_cluster::smem_block_tile_layout; - - const int usable_bytes = dynamic_smem_bytes - layout_t::base_padding_bytes; - REQUIRE(usable_bytes > 0); - - const int slots = usable_bytes / layout_t::chunk_bytes; - REQUIRE(slots > 0); - - const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); - REQUIRE(block_tile_capacity == static_cast(slots * layout_t::chunk_items)); - - const auto cluster_tile_capacity = - layout_t::template cluster_tile_capacity(cluster_blocks, block_tile_capacity); - REQUIRE(cluster_tile_capacity > 0); - // The head is an edge (static SMEM), not a reserved chunk, so coverage is the full physical capacity. - REQUIRE(cluster_tile_capacity == static_cast(cluster_blocks) * block_tile_capacity); - - // Worst-case resident chunks on any single rank, mirroring the agent: only the aligned region `[head_items, - // segment_size)` is chunked (`num_chunks`), then spread across `blocks` ranks (the `ceil_div` max for both the - // strided and blocked partitions). - const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { - using size_t = cuda::std::int64_t; - const size_t tail_items = segment_size - head_items; - const size_t chunks = host_ceil_div(tail_items, size_t{layout_t::chunk_items}); - return host_ceil_div(chunks, static_cast(blocks)); - }; - - const int heads[] = {0, 1, layout_t::chunk_items / 2, layout_t::chunk_items - 1}; - for (const int head_items : heads) - { - CAPTURE(c2h::type_name(), - ChunkBytes, - LoadAlignBytes, - dynamic_smem_bytes, - cluster_blocks, - slots, - block_tile_capacity, - cluster_tile_capacity, - head_items); - REQUIRE(max_rank_chunks(cluster_tile_capacity, head_items, cluster_blocks) <= slots); - } - - // Tightness: the full physical capacity fits resident (`slots` chunks per rank), one item beyond overflows. - CAPTURE( - c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, block_tile_capacity); - REQUIRE(max_rank_chunks(cluster_tile_capacity, 0, cluster_blocks) == slots); - REQUIRE(max_rank_chunks(cluster_tile_capacity + 1, 0, cluster_blocks) == slots + 1); -} - -template -void check_layout_matrix() -{ - constexpr int dynamic_smem_cases[] = {48 * 1024, 96 * 1024, 160 * 1024, 228 * 1024}; - constexpr int cluster_block_cases[] = {8, 16}; - - for (const int dynamic_smem_bytes : dynamic_smem_cases) - { - for (const int cluster_blocks : cluster_block_cases) - { - check_layout_case(dynamic_smem_bytes, cluster_blocks); - } - } -} -} // namespace - -TEST_CASE("Segmented TopK cluster SMEM layout exposes the full physical capacity (head is an edge, not a chunk)", - "[keys][segmented][topk][cluster][layout]") -{ - using default_policy = cub::detail::batched_topk_cluster::policy_selector; - constexpr auto policy = default_policy{}(cuda::compute_capability{9, 0}); - - using default_float_layout = - cub::detail::batched_topk_cluster::smem_block_tile_layout; - static_assert(default_float_layout::block_tile_capacity(0) == 0); - static_assert(default_float_layout::template cluster_tile_capacity(8, 0) == 0); - // No head reservation: a one-chunk cluster reports its full chunk as coverage. - static_assert(default_float_layout::template cluster_tile_capacity(1, default_float_layout::chunk_items) - == default_float_layout::chunk_items); - - check_layout_matrix(); - check_layout_matrix(); - check_layout_matrix(); - - check_layout_matrix(); - check_layout_matrix(); - check_layout_matrix(); -} diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index a3117f9cea2..d97fc2da9af 100644 --- a/docs/cub/api_docs/device_topk_requirements.rst +++ b/docs/cub/api_docs/device_topk_requirements.rst @@ -63,16 +63,34 @@ requires ``determinism::gpu_to_gpu``. See :ref:`cub-topk-set-membership` for the .. note:: - **Current support.** This initial API surface only implements the fully opted-out configuration. - For :cpp:struct:`cub::DeviceBatchedTopK` it must be requested **explicitly** as - ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, - cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)`` - (:cpp:struct:`cub::DeviceTopK` has no tie-break dimension yet and omits the ``tie_break`` token). - The algorithms ``static_assert`` for any other combination (including an empty, no-requirement - environment), so the deterministic default described above cannot yet be exercised in code. The - deterministic, tie-broken, and (stable-)sorted modes documented here define the committed long-term - contract and will become available (including as the no-requirement default) as those code paths - land. + **Current support.** Only the unsorted output ordering is implemented so far, so + ``cuda::execution::output_ordering::unsorted`` must be requested **explicitly**; ``sorted`` / + ``stable_sorted`` (and therefore an empty, no-requirement environment, which resolves to + ``stable_sorted``) ``static_assert`` at compile time. The supported *selection* requirements differ + by algorithm and, for :cpp:struct:`cub::DeviceBatchedTopK`, by architecture: + + * :cpp:struct:`cub::DeviceBatchedTopK`: + + - **Pre-Hopper (compute capability < 9.0):** only the fully non-deterministic request + ``(not_guaranteed, unspecified)`` is supported, and every segment must fit a single thread + block. Deterministic / tie-break requests and larger segments require the SM 9.0+ cluster + backend and are diagnosed at compile time (or at runtime as ``cudaErrorNotSupported`` in + relaxed builds). + - **Hopper and newer (compute capability >= 9.0):** every acknowledged ``(determinism, + tie_break)`` pair is supported -- ``(not_guaranteed, unspecified)``, ``(run_to_run, + unspecified)``, and ``(gpu_to_gpu, {unspecified, prefer_smaller_index, prefer_larger_index})`` + -- and segments too large for a single thread block are handled by the cluster backend. Among + non-deterministic requests the baseline vs cluster backend is chosen by an architecture / + segment-size crossover (the cluster backend is preferred for baseline-coverable segments only on + the newer architectures where it wins). + + * :cpp:struct:`cub::DeviceTopK` implements only the fully opted-out configuration + (``cuda::execution::determinism::not_guaranteed`` with unsorted output) and has no tie-break + dimension yet, so it omits the ``tie_break`` token. + + The remaining deterministic, tie-broken, and (stable-)sorted modes documented here define the + committed long-term contract and will become available (including as the no-requirement default) as + those code paths land. Requirements reference ---------------------- From b943a0c429185331b9af3b7d14224dbbdbc029ed Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 5 Jul 2026 15:36:42 +0200 Subject: [PATCH 075/126] Update tests to use the unified interface --- cub/test/CMakeLists.txt | 12 + .../catch2_test_device_batched_topk_api.cu | 5 + ...catch2_test_device_batched_topk_env_api.cu | 5 + ...st_device_segmented_topk_cluster_layout.cu | 108 ++++++++ .../catch2_test_device_segmented_topk_keys.cu | 204 +++++++-------- ...catch2_test_device_segmented_topk_pairs.cu | 234 +++++++++--------- cub/test/catch2_test_device_topk_common.cuh | 52 ++++ ...t_device_batched_topk_requirements_fail.cu | 6 + ...vice_batched_topk_unsupported_arch_fail.cu | 48 ++++ 9 files changed, 450 insertions(+), 224 deletions(-) create mode 100644 cub/test/catch2_test_device_segmented_topk_cluster_layout.cu create mode 100644 cub/test/test_device_batched_topk_unsupported_arch_fail.cu diff --git a/cub/test/CMakeLists.txt b/cub/test/CMakeLists.txt index ce46a86b93e..806e1ad5276 100644 --- a/cub/test/CMakeLists.txt +++ b/cub/test/CMakeLists.txt @@ -280,6 +280,18 @@ function( # Ensure that we test with assertions enabled target_compile_definitions(${test_target} PRIVATE CCCL_ENABLE_ASSERTIONS) + # The unsupported-arch xfail test omits the _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT define its siblings use, so it + # keeps the strict compile-time diagnostic. Pin it to 89-virtual (the highest arch without cluster support) so the + # deterministic request it issues has no viable backend and the static_assert fires regardless of the preset's archs. + if ( + "${test_src}" MATCHES "test_device_batched_topk_unsupported_arch_fail\\.cu" + ) + set_target_properties( + ${test_target} + PROPERTIES CUDA_ARCHITECTURES "89-virtual" + ) + endif() + # Disable clang-cuda >= 22 warnings regarding failed loop unrolling. if ( "${CMAKE_CUDA_COMPILER_ID}" STREQUAL "Clang" diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index 92d55da4ade..d54c3b79963 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -1,6 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) +// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See +// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index f0650bde1db..b3f90925e21 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -1,6 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) +// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See +// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include diff --git a/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu new file mode 100644 index 00000000000..1b13f8e29e5 --- /dev/null +++ b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include // smem_block_tile_layout +#include // batched_topk::cluster_policy_selector + +#include + +#include + +namespace +{ +template +[[nodiscard]] constexpr SizeT host_ceil_div(SizeT numerator, SizeT denominator) +{ + return (numerator + denominator - 1) / denominator; +} + +template +void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) +{ + using layout_t = cub::detail::batched_topk_cluster::smem_block_tile_layout; + + const int usable_bytes = dynamic_smem_bytes - layout_t::base_padding_bytes; + REQUIRE(usable_bytes > 0); + + const int slots = usable_bytes / layout_t::chunk_bytes; + REQUIRE(slots > 0); + + const auto block_tile_capacity = layout_t::block_tile_capacity(dynamic_smem_bytes); + REQUIRE(block_tile_capacity == static_cast(slots * layout_t::chunk_items)); + + const auto cluster_tile_capacity = + layout_t::template cluster_tile_capacity(cluster_blocks, block_tile_capacity); + REQUIRE(cluster_tile_capacity > 0); + // The head is an edge (static SMEM), not a reserved chunk, so coverage is the full physical capacity. + REQUIRE(cluster_tile_capacity == static_cast(cluster_blocks) * block_tile_capacity); + + // Worst-case resident chunks on any single rank, mirroring the agent: only the aligned region `[head_items, + // segment_size)` is chunked (`num_chunks`), then spread across `blocks` ranks (the `ceil_div` max for both the + // strided and blocked partitions). + const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { + using size_t = cuda::std::int64_t; + const size_t tail_items = segment_size - head_items; + const size_t chunks = host_ceil_div(tail_items, size_t{layout_t::chunk_items}); + return host_ceil_div(chunks, static_cast(blocks)); + }; + + const int heads[] = {0, 1, layout_t::chunk_items / 2, layout_t::chunk_items - 1}; + for (const int head_items : heads) + { + CAPTURE(c2h::type_name(), + ChunkBytes, + LoadAlignBytes, + dynamic_smem_bytes, + cluster_blocks, + slots, + block_tile_capacity, + cluster_tile_capacity, + head_items); + REQUIRE(max_rank_chunks(cluster_tile_capacity, head_items, cluster_blocks) <= slots); + } + + // Tightness: the full physical capacity fits resident (`slots` chunks per rank), one item beyond overflows. + CAPTURE( + c2h::type_name(), ChunkBytes, LoadAlignBytes, dynamic_smem_bytes, cluster_blocks, slots, block_tile_capacity); + REQUIRE(max_rank_chunks(cluster_tile_capacity, 0, cluster_blocks) == slots); + REQUIRE(max_rank_chunks(cluster_tile_capacity + 1, 0, cluster_blocks) == slots + 1); +} + +template +void check_layout_matrix() +{ + constexpr int dynamic_smem_cases[] = {48 * 1024, 96 * 1024, 160 * 1024, 228 * 1024}; + constexpr int cluster_block_cases[] = {8, 16}; + + for (const int dynamic_smem_bytes : dynamic_smem_cases) + { + for (const int cluster_blocks : cluster_block_cases) + { + check_layout_case(dynamic_smem_bytes, cluster_blocks); + } + } +} +} // namespace + +TEST_CASE("Segmented TopK cluster SMEM layout exposes the full physical capacity (head is an edge, not a chunk)", + "[keys][segmented][topk][cluster][layout]") +{ + using default_policy = cub::detail::batched_topk::cluster_policy_selector; + constexpr auto policy = default_policy{}(cuda::compute_capability{9, 0}); + + using default_float_layout = + cub::detail::batched_topk_cluster::smem_block_tile_layout; + static_assert(default_float_layout::block_tile_capacity(0) == 0); + static_assert(default_float_layout::template cluster_tile_capacity(8, 0) == 0); + // No head reservation: a one-chunk cluster reports its full chunk as coverage. + static_assert(default_float_layout::template cluster_tile_capacity(1, default_float_layout::chunk_items) + == default_float_layout::chunk_items); + + check_layout_matrix(); + check_layout_matrix(); + check_layout_matrix(); + + check_layout_matrix(); + check_layout_matrix(); + check_layout_matrix(); +} diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index b594a1d8d03..7a930225962 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -1,10 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) +// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See +// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include -#include #include #include @@ -49,14 +53,9 @@ struct heavy_tie_key_op } }; -enum class topk_backend -{ - baseline, - cluster, -}; - -inline constexpr topk_backend selected_backend = topk_backend::cluster; - +// All tests drive the algorithm through the public `cub::DeviceBatchedTopK` API, threading the requested +// determinism/tie-break into the environment via `require`. The dispatch selects the backend from the architecture and +// the statically-known maximum segment size (a deterministic request routes to the cluster backend on SM90+). template {}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted)}; + if constexpr (SelectDirection == cub::detail::topk::select::max) { - return cub::detail::batched_topk_cluster::dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - cuda::args::constant{}, - num_segments, - stream); + return cub::DeviceBatchedTopK::MaxKeys( + d_temp_storage, temp_storage_bytes, d_key_segments_it, d_key_segments_out_it, segment_sizes, k, num_segments, env); } else { - // Baseline backend routes through the public API; the cluster backend keeps using the lower-level dispatch above. - // The public API takes no total-items guarantee and always runs nondeterministic (the determinism-aware tests skip - // the deterministic combos for non-cluster backends). - auto env = cuda::std::execution::env{ - cuda::stream_ref{stream}, - cuda::execution::require(cuda::execution::determinism::not_guaranteed, - cuda::execution::tie_break::unspecified, - cuda::execution::output_ordering::unsorted)}; - if constexpr (SelectDirection == cub::detail::topk::select::max) - { - return cub::DeviceBatchedTopK::MaxKeys( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k, - num_segments, - env); - } - else - { - return cub::DeviceBatchedTopK::MinKeys( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k, - num_segments, - env); - } + return cub::DeviceBatchedTopK::MinKeys( + d_temp_storage, temp_storage_bytes, d_key_segments_it, d_key_segments_out_it, segment_sizes, k, num_segments, env); } } @@ -423,16 +387,10 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle a segment-size type narrower t using combo = c2h::get<1, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - constexpr auto direction = cub::detail::topk::select::max; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - + // Baseline-coverable segment sizes (statically bounded to 127 below): only the deterministic / tie-break requirements + // route to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(/*static_max_segment_size=*/127); + constexpr auto direction = cub::detail::topk::select::max; // Sizes fit both 8-bit types but sit far below the 512 threads a block launches; `127` probes the signed type's max. const seg_size_t segment_size = static_cast(GENERATE(values({3, 100, 127}))); const seg_size_t k = static_cast( @@ -479,6 +437,68 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle a segment-size type narrower t REQUIRE(expected_keys == keys_out_buffer); } + +// Automatic size-based backend crossover. The segment size is at the cluster crossover threshold +// (`cluster_beneficial_min_segment_size` == 8 Ki) yet still baseline-coverable (baseline top tile is 256*64 == 16 Ki), +// so on SM 10.0+ the automatic selector prefers the cluster backend purely on size, under a *non-deterministic* request +// -- a branch no other test hits (deterministic tests force cluster via the requirement, oversize tests via +// non-coverage). Below SM 10.0 the crossover does not apply, so the case is skipped. Checked vs a sorted reference. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys route a large baseline-coverable segment to the cluster backend on SM10+", + "[keys][segmented][topk][device][cluster]", + select_direction_list) +{ + // Gate on the compiled/JIT compute capability the dispatch resolves (cub::PtxVersion), not the physical device + // (cub::SmVersion): the size-based crossover fires from the selector evaluated at the resolved target, so a build + // whose highest applicable target is below SM 10.0 stays on the baseline backend even on an SM 10.0+ device. + int ptx_version = 0; + REQUIRE(cudaSuccess == cub::PtxVersion(ptx_version)); + constexpr int cluster_beneficial_min_ptx_version = 1000; // SM 10.0 + if (ptx_version < cluster_beneficial_min_ptx_version) + { + SKIP( + "The size-based baseline->cluster crossover only applies on SM 10.0+; below it the automatic selector stays on " + "the baseline backend (already covered by the small-segment tests)."); + } + + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + using key_t = cuda::std::uint8_t; // small key so the baseline backend can cover the 8 Ki segment + + constexpr auto direction = c2h::get<0, TestType>::value; + + constexpr segment_size_t static_max_segment_size = 8 * 1024; // == cluster_beneficial_min_segment_size + const segment_size_t segment_size = static_max_segment_size; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, segment_size_t{257}, static_max_segment_size})); + const segment_index_t num_segments = GENERATE_COPY(values({segment_index_t{1}, segment_index_t{5}})); + + CAPTURE(static_max_segment_size, segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // Defaulted determinism/tie-break (not_guaranteed / unspecified): the statically-known maximum segment size, not the + // requirement, is what routes this to the cluster backend. + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} #endif // TEST_TYPES == 0 #if TEST_TYPES == 2 @@ -549,15 +569,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - // `static_max_segment_size` is chosen to exceed the largest all-resident cluster coverage (~16 blocks worth of // resident SMEM), so the 1 Mi-element segments force the agent's gmem-streaming overflow path (including an // unaligned overflow tail via `- 31`), while the 128 Ki-element segment still runs fully resident under the same @@ -568,6 +579,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); constexpr auto direction = cub::detail::topk::select::max; const int pad = GENERATE(0, 1, 3, 7); @@ -625,15 +638,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 3 constexpr auto tie_break = combo::tie_break; constexpr auto direction = cub::detail::topk::select::max; - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - constexpr seg_size_t static_max_segment_size = 1024 * 1024; constexpr seg_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 2; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const int pad = GENERATE(0, 7); const seg_size_t segment_size = static_max_segment_size - 31; // unaligned -> forces streaming + unaligned tail edge @@ -708,15 +717,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-c using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - // The counting-iterator key source is non-contiguous, so the agent uses its generic overflow-streaming path rather // than BlockLoadToShared. `static_max_segment_size` exceeds the largest all-resident cluster coverage, so the 1 Mi // -element segments stream (incl. an unaligned `- 31` tail), while the 128 Ki-element segment validates the generic @@ -725,6 +725,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-c constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); constexpr auto direction = cub::detail::topk::select::max; const segment_size_t segment_size = @@ -876,15 +878,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch sizes a // wide (16-CTA) cluster and mixes every effective-width regime: streaming segments (the 1 Mi-element ones, one with // an unaligned `- 31` overflow tail), a fully-resident multi-CTA segment (96 Ki + 17), two *medium* segments that @@ -892,6 +885,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign // effective-cluster-width path), and a tiny segment that collapses onto a single CTA (257 elements). constexpr segment_size_t static_max_segment_size = 1100 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); constexpr auto direction = cub::detail::topk::select::max; const int pad = GENERATE(1, 3, 7); @@ -1173,21 +1168,14 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys handle heavy ties at the k-th boundar using combo = c2h::get<1, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - - using segment_size_t = cuda::std::int64_t; - using segment_index_t = cuda::std::int64_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; constexpr segment_size_t static_max_segment_size = 64 * 1024; constexpr segment_size_t static_max_k = 64 * 1024; constexpr segment_index_t num_segments = 3; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{257}, segment_size_t{4096}, segment_size_t{64 * 1024}})); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 43b8f6ff0cd..b39046a8676 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -1,10 +1,15 @@ // SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) +// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See +// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include -#include +#include // topk_policy / make_baseline_policy (cross-tuning test) #include #include @@ -16,6 +21,7 @@ #include #include #include +#include #include #include @@ -58,16 +64,9 @@ struct flag_intra_segment_duplicates template flag_intra_segment_duplicates(ItemItT, SegIdItT) -> flag_intra_segment_duplicates; -enum class topk_backend -{ - baseline, - cluster, -}; - -inline constexpr topk_backend selected_backend = topk_backend::cluster; - -// Routes the key-value (pairs) top-k to the baseline or cluster backend. `Determinism`/`TieBreak` are forwarded to the -// cluster dispatch to request the deterministic tie-break path; they default to the nondeterministic behavior. +// Routes the key-value (pairs) top-k through the public `cub::DeviceBatchedTopK` API, threading the requested +// determinism/tie-break into the environment via `require`. The dispatch selects the backend from the architecture and +// the statically-known maximum segment size (a deterministic request routes to the cluster backend on SM90+). template {}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted)}; + if constexpr (SelectDirection == cub::detail::topk::select::max) { - return cub::detail::batched_topk_cluster::dispatch( + return cub::DeviceBatchedTopK::MaxPairs( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -102,48 +106,22 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( d_value_segments_out_it, segment_sizes, k, - cuda::args::constant{}, num_segments, - stream); + env); } else { - // Baseline backend routes through the public API; the cluster backend keeps using the lower-level dispatch above. - // The public API takes no total-items guarantee and always runs nondeterministic (the determinism-aware tests skip - // the deterministic combos for non-cluster backends). - auto env = cuda::std::execution::env{ - cuda::stream_ref{stream}, - cuda::execution::require(cuda::execution::determinism::not_guaranteed, - cuda::execution::tie_break::unspecified, - cuda::execution::output_ordering::unsorted)}; - if constexpr (SelectDirection == cub::detail::topk::select::max) - { - return cub::DeviceBatchedTopK::MaxPairs( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - num_segments, - env); - } - else - { - return cub::DeviceBatchedTopK::MinPairs( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - num_segments, - env); - } + return cub::DeviceBatchedTopK::MinPairs( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + num_segments, + env); } } @@ -467,16 +445,10 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs handle a segment-size type narrower using combo = c2h::get<1, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - constexpr auto direction = cub::detail::topk::select::max; - - // Only the cluster backend honors determinism; elsewhere the deterministic combinations just rerun the - // nondeterministic path, so skip them (the base nondeterministic combo still runs). - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - + // Baseline-coverable segment sizes (statically bounded to 127 below): only the deterministic / tie-break requirements + // route to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(/*static_max_segment_size=*/127); + constexpr auto direction = cub::detail::topk::select::max; // Sizes fit both 8-bit types but sit far below the 512 threads a block launches; `127` probes the signed type's max. const seg_size_t segment_size = static_cast(GENERATE(values({3, 100, 127}))); const seg_size_t k = static_cast( @@ -594,6 +566,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments through a non- constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; + // Oversize segments always route to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t segment_size = GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); @@ -665,6 +639,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; + // Oversize segments always route to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const int pad = GENERATE(0, 1, 3, 7); // The `+ 1` / `- 4095` sizes make the global-last chunk a single item, i.e. a pure-suffix tail with an empty aligned @@ -743,15 +719,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments with a signed constexpr auto tie_break = combo::tie_break; constexpr auto direction = cub::detail::topk::select::max; - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - constexpr seg_size_t static_max_segment_size = 1024 * 1024; constexpr seg_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 2; + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const int pad = GENERATE(0, 7); const seg_size_t segment_size = static_max_segment_size - 31; // unaligned -> forces streaming + unaligned tail edge @@ -869,12 +841,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the select_direction_list, tie_break_pref_list) { - // Only the cluster backend provides the deterministic tie-break; the strict index-set check is meaningless elsewhere. - if constexpr (selected_backend != topk_backend::cluster) - { - SKIP("Deterministic tie-break is only provided by the cluster backend"); - } - using key_t = cuda::std::uint32_t; using val_t = cuda::std::int32_t; using segment_size_t = cuda::std::int64_t; @@ -888,6 +854,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the constexpr segment_size_t static_max_segment_size = 64 * 1024; constexpr segment_size_t static_max_k = 64 * 1024; constexpr segment_index_t num_segments = 2; + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{4096}, segment_size_t{64 * 1024}})); const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); @@ -943,19 +911,35 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the REQUIRE(ref == h_values_out); } -// A second valid cluster tuning for the gpu_to_gpu determinism test: starts from the default policy (for its valid -// launch configs) and overrides the shape knobs (block size, items-per-thread, pipeline depth, tie-break granularity). -struct alt_cluster_tuning +// Whole-`topk_policy` tuning overrides for the reproducibility test, threaded through the public API's single tuning +// query (`cuda::execution::tune`). Both force the cluster backend (which alone honors a deterministic request). +// `default_cluster_selector` uses the default cluster tuning; `alt_cluster_selector` starts from it and overrides the +// shape knobs (block size, items-per-thread, pipeline depth, tie-break granularity) to a second *valid* tuning, so the +// gpu_to_gpu config-independence check can compare two different launch configs. +struct default_cluster_selector { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const - -> cub::detail::batched_topk_cluster::cluster_topk_policy + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const + -> cub::detail::batched_topk::topk_policy { - auto policy = cub::detail::batched_topk_cluster::policy_selector{}(cc); - policy.threads_per_block = 256; - policy.histogram_items_per_thread = 2; - policy.pipeline_stages = 2; - policy.tie_break_items_per_thread = 2; - return policy; + return cub::detail::batched_topk::topk_policy{ + cub::detail::batched_topk::topk_backend::cluster, + cub::detail::batched_topk::make_baseline_policy(), + cub::detail::batched_topk::make_cluster_policy()}; + } +}; + +struct alt_cluster_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const + -> cub::detail::batched_topk::topk_policy + { + auto cluster = cub::detail::batched_topk::make_cluster_policy(); + cluster.threads_per_block = 256; + cluster.histogram_items_per_thread = 2; + cluster.pipeline_stages = 2; + cluster.tie_break_items_per_thread = 2; + return cub::detail::batched_topk::topk_policy{ + cub::detail::batched_topk::topk_backend::cluster, cub::detail::batched_topk::make_baseline_policy(), cluster}; } }; @@ -969,12 +953,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break select_direction_list, determinism_list) { - // Only the cluster backend honors determinism; the cross-tuning check is meaningless elsewhere. - if constexpr (selected_backend != topk_backend::cluster) - { - SKIP("Determinism is only provided by the cluster backend"); - } - using key_t = cuda::std::uint32_t; using val_t = cuda::std::int32_t; using segment_size_t = cuda::std::int64_t; @@ -987,6 +965,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break constexpr segment_size_t static_max_segment_size = 64 * 1024; constexpr segment_size_t static_max_k = 64 * 1024; constexpr segment_index_t num_segments = 2; + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{4096}, segment_size_t{64 * 1024}})); const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); @@ -1011,29 +991,56 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break c2h::device_vector values_out_a(num_segments * k, thrust::no_init); c2h::device_vector values_out_b(num_segments * k, thrust::no_init); - // Calls the dispatch directly (the launch wrapper does not thread a policy selector) so each run picks a tuning. + // Runs the same problem through the public `cub::DeviceBatchedTopK` API, tuning the whole `topk_policy` via + // `cuda::execution::tune` so each run can pick a different valid cluster tuning (the override forces the cluster + // backend). Drives the two-phase temp-storage protocol manually so both runs share this call site. const auto run_with_tuning = - [&](auto policy_selector, c2h::device_vector& keys_out, c2h::device_vector& values_out) { + [&](auto selector, c2h::device_vector& keys_out, c2h::device_vector& values_out) { auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune(selector)}; + size_t temp_bytes = 0; const auto invoke = [&](void* d_temp) { - return cub::detail::batched_topk_cluster::dispatch( - d_temp, - temp_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - cuda::args::immediate{segment_size, cuda::args::bounds()}, - cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::constant{}, - cuda::args::immediate{num_segments}, - cudaStream_t{0}, - policy_selector); + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + if constexpr (direction == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + } + else + { + return cub::DeviceBatchedTopK::MinPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + } }; REQUIRE(invoke(nullptr) == cudaSuccess); c2h::device_vector temp_storage(temp_bytes, thrust::no_init); @@ -1041,15 +1048,15 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break REQUIRE(cudaSuccess == cudaDeviceSynchronize()); }; - run_with_tuning(cub::detail::batched_topk_cluster::policy_selector{}, keys_out_a, values_out_a); - // run_to_run: same tuning twice. gpu_to_gpu: a different tuning for the stronger config-independent check. + run_with_tuning(default_cluster_selector{}, keys_out_a, values_out_a); + // run_to_run: same tuning twice. gpu_to_gpu: a different valid tuning for the stronger config-independent check. if constexpr (determinism == cuda::execution::determinism::__determinism_t::__gpu_to_gpu) { - run_with_tuning(alt_cluster_tuning{}, keys_out_b, values_out_b); + run_with_tuning(alt_cluster_selector{}, keys_out_b, values_out_b); } else { - run_with_tuning(cub::detail::batched_topk_cluster::policy_selector{}, keys_out_b, values_out_b); + run_with_tuning(default_cluster_selector{}, keys_out_b, values_out_b); } // Sanity: values still belong to their keys and no source index repeats within a segment. @@ -1088,17 +1095,12 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with mixed effective-width vari using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - - if constexpr (selected_backend != topk_backend::cluster - && determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed) - { - SKIP("Determinism requirements are only provided by the cluster backend"); - } - // The bound far exceeds the actual segments, so the launch picks a wide cluster while the segments stay small. constexpr segment_size_t static_max_segment_size = 1100 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr auto direction = cub::detail::topk::select::max; + // Oversize bound always routes to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); constexpr segment_size_t full_segment_size = 96 * 1024 + 17; // enough chunks to fill the launched cluster (no idle) diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index c6a9bcc122b..c2dfa4b3e2e 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -6,15 +6,67 @@ #include #include #include // topk::select::{min, max} +#include // make_baseline_policy +#include // cub::PtxVersion #include +#include +#include #include #include #include #include +// Low-level gate: skips the current test case when a request that must use the SM90+ cluster backend has no +// cluster-capable target in the build. With `_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` defined (as the top-k test +// sources do), such a request fails at runtime with cudaErrorNotSupported rather than at compile time. `needs_cluster` +// must be true when the caller knows the configuration requires the cluster backend; prefer +// skip_if_batched_topk_backend_unavailable(), which derives it from the request. +// +// Uses cub::PtxVersion (the compiled compute capability the dispatch resolves via PtxComputeCap), not cub::SmVersion +// (the physical device): the two diverge when compiling for a virtual architecture below the device (e.g. `89-virtual` +// on an SM120 GPU), where the cluster arm is never emitted and the dispatch returns cudaErrorNotSupported -- exactly +// what this skip must catch. +inline void skip_if_batched_topk_cluster_unavailable(bool needs_cluster) +{ + if (!needs_cluster) + { + return; + } + int ptx_version = 0; + REQUIRE(cudaSuccess == cub::PtxVersion(ptx_version)); + constexpr int cluster_min_ptx_version = 900; // SM 9.0 + if (ptx_version < cluster_min_ptx_version) + { + SKIP("This top-k configuration requires the SM90+ cluster backend (a deterministic / tie-break request or a " + "segment size beyond the baseline backend's capacity); the current build has no SM90+ target that can serve " + "it."); + } +} + +// Skips the current test case when the batched top-k configuration cannot run in the current build. The request needs +// the cluster backend if it is deterministic / has a concrete tie-break, or if `static_max_segment_size` exceeds the +// baseline backend's coverage; skips if that backend is unavailable. Deriving the size decision here (rather than a +// precomputed `oversize` bool) keeps the threshold in one place. Pass the same maximum segment size the test hands to +// the dispatch (its `cuda::args::bounds<...>` upper bound, or the type's maximum when unbounded). Note that `oversize` +// uses only the tile-size bound `baseline_max_covered_segment_size`; the dispatch's `baseline_can_cover_v` additionally +// checks the agent's shared-memory fit, so a borderline size the bound deems baseline-coverable could still route to +// the cluster backend (such a case would fail rather than skip if the cluster backend is unavailable). +template +void skip_if_batched_topk_backend_unavailable(cuda::std::int64_t static_max_segment_size) +{ + constexpr bool deterministic = Determinism != cuda::execution::determinism::__determinism_t::__not_guaranteed + || TieBreak != cuda::execution::tie_break::__tie_break_t::__unspecified; + const bool oversize = + static_max_segment_size + > cub::detail::batched_topk::baseline_max_covered_segment_size(cub::detail::batched_topk::make_baseline_policy()); + skip_if_batched_topk_cluster_unavailable(deterministic || oversize); +} + // Function object to generate monotonically non-decreasing values for small key types template struct inc_t diff --git a/cub/test/test_device_batched_topk_requirements_fail.cu b/cub/test/test_device_batched_topk_requirements_fail.cu index c9a00e9d31d..55d234c2f7b 100644 --- a/cub/test/test_device_batched_topk_requirements_fail.cu +++ b/cub/test/test_device_batched_topk_requirements_fail.cu @@ -3,6 +3,12 @@ // %PARAM% TEST_ERR err 0:1:2:3:4:5 +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) +// so only the requirement static_asserts under test fire, regardless of the target architecture (this test compiles for +// the preset's default archs, which include pre-SM90). See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/dispatch/dispatch_batched_topk.cuh. Precedes the CUB include below. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include #include diff --git a/cub/test/test_device_batched_topk_unsupported_arch_fail.cu b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu new file mode 100644 index 00000000000..61e269a294f --- /dev/null +++ b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +// Verifies the strict unsupported-architecture diagnostic in cub::DeviceBatchedTopK (see +// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh). A deterministic +// (gpu_to_gpu) request can only be served by the SM90+ cluster backend; when the translation unit targets an +// architecture that cannot run it (this target is pinned to a pre-SM90 arch via CUDA_ARCHITECTURES in CMakeLists.txt), +// the default (strict) mode must fail to compile rather than defer the failure to a runtime cudaErrorNotSupported. +// +// This is the only batched/segmented top-k test built *without* _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT, so it +// exercises the strict path that CUB's other top-k tests intentionally opt out of. + +int main() +{ + namespace ex = cuda::execution; + + int** d_keys_in = nullptr; + int** d_keys_out = nullptr; + auto segment_sizes = cuda::args::constant<8>{}; + auto k_arg = cuda::args::constant<3>{}; + auto num_segments = cuda::args::immediate{cuda::std::int64_t{2}}; + + // A deterministic result set routes to the cluster backend, which is unsupported on the pinned pre-SM90 target. + auto requirements = + ex::require(ex::determinism::gpu_to_gpu, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); + // expected-error {{"is not supported on at least one architecture"}} + + auto env = cuda::std::execution::env{requirements}; + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segments, env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; + } +} From 592279d16f6ef7b5deed65d74d2f6f61bec578e7 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 5 Jul 2026 15:37:55 +0200 Subject: [PATCH 076/126] Update benchmarks to use the unified interface Enable auto-tuning of the cluster backend. --- .../bench/segmented_topk/fixed/keys.cu | 149 +++++---- .../variable/indexed.cluster.cu | 32 ++ .../bench/segmented_topk/variable/indexed.cu | 186 ++--------- .../variable/indexed_common.cuh | 274 ++++++++++++++++ .../segmented_topk/variable/keys.cluster.cu | 32 ++ .../bench/segmented_topk/variable/keys.cu | 215 ++----------- .../segmented_topk/variable/keys_common.cuh | 298 ++++++++++++++++++ 7 files changed, 769 insertions(+), 417 deletions(-) create mode 100644 cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu create mode 100644 cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh create mode 100644 cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu create mode 100644 cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index f224ee378b9..37ea40eb095 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -1,13 +1,22 @@ // SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check so this benchmark compiles for the full +// configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) +// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include +#include #include #include -#include #include +#include +#include #include +#include #include #include #include @@ -20,34 +29,67 @@ // %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1 // %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32 // %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1 +// %RANGE% TUNE_BACKEND backend 0:1:1 enum class topk_backend { baseline, cluster, device, + automatic, }; -inline constexpr topk_backend selected_backend = topk_backend::baseline; +// Which backend this build benchmarks. `automatic` (the default) issues no `tune` override, leaving the choice to the +// library's arch/size selector. `baseline`/`cluster` force one of the two DeviceBatchedTopK backends via the `tune`d +// selector below; `device` is a reference that issues one `cub::DeviceTopK` call per segment. Autotuning sweeps the +// baseline and cluster backends (the `%RANGE% ... 0:1:1` above); only baseline knobs are exposed here, so the cluster +// backend uses its default sub-policy and `device` has no knobs. Override with -DTUNE_BACKEND=0/1 (force a backend), +// =2 (device reference), or =3 (automatic). +#ifndef TUNE_BACKEND +# if TUNE_BASE +# define TUNE_BACKEND 3 // automatic: the library's production selector, for base/benchmark builds +# else +# define TUNE_BACKEND 0 // force baseline when an actively-tuned variant does not sweep the backend +# endif +#endif + +inline constexpr topk_backend selected_backend = +#if TUNE_BACKEND == 0 + topk_backend::baseline; +#elif TUNE_BACKEND == 1 + topk_backend::cluster; +#elif TUNE_BACKEND == 2 + topk_backend::device; +#else + topk_backend::automatic; +#endif // Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; -// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. -static_assert(selected_backend == topk_backend::cluster +// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. The +// cluster backend implements them, and automatic honors them via the library's selector, so both allow non-defaults. +static_assert(selected_backend == topk_backend::cluster || selected_backend == topk_backend::automatic || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), - "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " - "selected_tie_break at their defaults for the baseline/device backends."); + "Only the cluster and automatic backends honor determinism/tie-break requirements; keep " + "selected_determinism and selected_tie_break at their defaults for the baseline/device backends."); -#if !TUNE_BASE -struct tuned_policy_selector +// Policy selector threaded through the public API's tuning environment when a concrete backend is forced (not +// `automatic`). Its `.backend` pins the backend for this build. In a TUNE_BASE build the forced backend uses the +// default sub-policies; otherwise the baseline knobs come from the TUNE_* macros. The cluster sub-policy is always the +// default (no cluster knobs are exposed here). +template +struct topk_backend_selector { - [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const - -> cub::detail::batched_topk::batched_topk_policy + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability cc) const + -> cub::detail::batched_topk::topk_policy { - // Single-entry policy chain driven by the tuning knobs. +#if TUNE_BASE + const auto baseline = + cub::detail::batched_topk::baseline_policy_selector_from_types{}(cc); +#else constexpr auto store_alg = cub::BLOCK_STORE_WARP_TRANSPOSE; # if TUNE_BLOCK_LOAD_ALGORITHM == 0 constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT; @@ -56,7 +98,7 @@ struct tuned_policy_selector # elif TUNE_BLOCK_LOAD_ALGORITHM == 2 constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE; # endif - return cub::detail::batched_topk::batched_topk_policy{{{ + const auto baseline = cub::detail::batched_topk::baseline_topk_policy{{{ cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, @@ -64,20 +106,25 @@ struct tuned_policy_selector cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, }}}; +#endif // TUNE_BASE + const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc); + constexpr auto backend = + (selected_backend == topk_backend::cluster) + ? cub::detail::batched_topk::topk_backend::cluster + : cub::detail::batched_topk::topk_backend::baseline; + return cub::detail::batched_topk::topk_policy{backend, baseline, cluster}; } }; -#endif // !TUNE_BASE -// Env-based dispatch over the selected backend. The cluster and baseline backends route through their respective -// `dispatch_with_env` entry points (temporary storage is allocated from the memory resource carried by `env`); the -// device backend issues one `cub::DeviceTopK::MaxKeys` per segment, reading the host-side segment sizes. +// Env-based dispatch over the selected backend. `automatic`/`baseline`/`cluster` all route through the public +// `cub::DeviceBatchedTopK` API (the latter two add a `tune`d `topk_backend_selector` that forces the backend; temp +// storage comes from the memory resource carried by `env`); the `device` backend issues one `cub::DeviceTopK::MaxKeys` +// per segment, reading the host-side segment sizes. template CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( @@ -85,26 +132,11 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( KeyOutputItItT d_keys_out, SegmentSizeParamT segment_sizes, KParamT k, - SelectDirectionParamT select_direction, NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items, [[maybe_unused]] const HostSegSizeT* h_segment_sizes, EnvT env) { - if constexpr (selected_backend == topk_backend::cluster) - { - return cub::detail::batched_topk_cluster::dispatch_with_env( - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - select_direction, - num_segments, - env); - } - else if constexpr (selected_backend == topk_backend::device) + if constexpr (selected_backend == topk_backend::device) { using num_segments_val_t = typename ::cuda::args::__traits::element_type; const auto num_segs = cub::detail::params::get_param(num_segments, num_segments_val_t{0}); @@ -135,17 +167,26 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( } else { - return cub::detail::batched_topk::dispatch_with_env( - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - select_direction, - num_segments, - total_num_items, - env); + // The determinism / tie-break / ordering requirement this benchmark issues; the library honors it whether or not we + // additionally force a backend. + auto req_env = cuda::std::execution::env{ + env, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted)}; + if constexpr (selected_backend == topk_backend::automatic) + { + // No `tune` override: the library's own selector picks the backend (arch/size crossover) -- the usual behavior. + return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, req_env); + } + else + { + using key_t = cub::detail::it_value_t>; + constexpr auto max_k = ::cuda::args::__traits::highest; + auto full_env = cuda::std::execution::env{ + req_env, cuda::execution::tune(topk_backend_selector{})}; + return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, full_env); + } } } @@ -160,7 +201,6 @@ void fixed_seg_size_topk_keys( const auto selected_elements = static_cast<::cuda::std::ptrdiff_t>(MaxNumSelected); const auto num_segments = ::cuda::std::max(1, (max_elements / segment_size)); const auto elements = num_segments * segment_size; - const auto total_num_items = ::cuda::args::immediate{static_cast<::cuda::std::int64_t>(elements)}; const bit_entropy entropy = str_to_entropy(state.get_string("Entropy")); // Skip workloads where k exceeds the segment size @@ -177,9 +217,8 @@ void fixed_seg_size_topk_keys( auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements); - auto segment_sizes = ::cuda::args::constant{}; - auto k = ::cuda::args::constant{}; - auto select_direction = ::cuda::args::constant{}; + auto segment_sizes = ::cuda::args::constant{}; + auto k = ::cuda::args::constant{}; state.add_element_count(elements, "NumElements"); state.add_element_count(segment_size, "SegmentSize"); @@ -193,15 +232,7 @@ void fixed_seg_size_topk_keys( caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env( - alloc, - launch -#if !TUNE_BASE - , - cuda::execution::tune(tuned_policy_selector{}) -#endif // !TUNE_BASE - ); - // TODO(bgruber): call the public API once available + auto env = cub_bench_env(alloc, launch); _CCCL_TRY_CUDA_API( batched_topk_keys, "batched topk failed", @@ -209,9 +240,7 @@ void fixed_seg_size_topk_keys( d_keys_out, segment_sizes, k, - select_direction, ::cuda::args::immediate{static_cast<::cuda::std::int64_t>(num_segments)}, - total_num_items, h_segment_sizes.data(), env); }); diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu b/cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu new file mode 100644 index 00000000000..40bbcc0c828 --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed.cluster.cu @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Cluster-backend variable-segment indexed (arg-top-k) benchmark. A base build runs the library's production +// (automatic) selector; a tuning variant forces the cluster backend and sweeps its per-block knobs plus the +// determinism / tie-break requirement. The baseline backend is benchmarked separately in the plain `indexed.cu`. + +// %RANGE% TUNE_CLUSTER_THREADS_PER_BLOCK tpb 128:1024:32 +// %RANGE% TUNE_CLUSTER_MIN_BLOCKS_PER_SM mbs 1:2:1 +// %RANGE% TUNE_CLUSTER_MIN_CHUNKS_PER_BLOCK mcb 1:2:1 +// %RANGE% TUNE_CLUSTER_CHUNK_KIB ckib 8:32:1 +// %RANGE% TUNE_CLUSTER_LOAD_ALIGN_BYTES_POW2 la 4:7:1 +// %RANGE% TUNE_CLUSTER_PIPELINE_STAGES ps 2:16:1 +// %RANGE% TUNE_CLUSTER_BITS_PER_PASS bpp 8:11:1 +// %RANGE% TUNE_CLUSTER_HIST_IPT hipt 1:24:1 +// %RANGE% TUNE_CLUSTER_TIEBREAK_IPT tipt 1:24:1 +// %RANGE% TUNE_CLUSTER_COPY_IPT cipt 1:24:1 +// %RANGE% TUNE_REQUIREMENT req 0:2:1 + +#ifndef TUNE_BACKEND +# if TUNE_BASE +# define TUNE_BACKEND 2 // automatic: the library's production selector, for base/benchmark builds +# else +# define TUNE_BACKEND 1 // cluster: the backend this file tunes +# endif +#endif + +#ifndef TUNE_REQUIREMENT +# define TUNE_REQUIREMENT 1 // deterministic + prefer-smaller-index (forward): the safe, least-surprising default +#endif + +#include "indexed_common.cuh" diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu index 9a36a7ccc1b..7583851dc0a 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu @@ -1,167 +1,25 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include - -#include - -#include "common.cuh" - -enum class topk_backend -{ - baseline, - cluster, -}; - -// Which backend computes the indexed (key + segment-local-index value) top-k. The cluster backend exercises the new -// key-value-pair path through `agent_batched_topk_cluster`; the baseline backend is kept for A/B comparison. -inline constexpr topk_backend selected_backend = topk_backend::cluster; - -// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). Only forwarded -// to the cluster backend -- the baseline backend does not implement these requirements. -inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; -inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; - -// The baseline backend ignores these requirements, so require the defaults there to avoid a silently ignored selection. -static_assert(selected_backend == topk_backend::cluster - || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed - && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), - "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " - "selected_tie_break at their defaults for the baseline backend."); - -template -CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_indexed( - KeyInputItItT d_keys_in, - KeyOutputItItT d_keys_out, - ValueInputItItT d_values_in, - ValueOutputItItT d_values_out, - SegmentSizeParamT segment_sizes, - KParamT k, - SelectDirectionParamT select_direction, - NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items, - EnvT env) -{ - if constexpr (selected_backend == topk_backend::cluster) - { - return cub::detail::batched_topk_cluster::dispatch_with_env( - d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, select_direction, num_segments, env); - } - else - { - return cub::detail::batched_topk::dispatch_with_env( - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k, - select_direction, - num_segments, - total_num_items, - env); - } -} - -// Indexed (arg-top-k) variant: each key carries a segment-local index as its value payload. The input values are -// produced by a counting iterator that restarts at 0 for every segment, so indices are not (pre-)materialized in global -// memory -template -void decode_style_variable_topk_indexed( - nvbench::state& state, nvbench::type_list, nvbench::enum_type>) -{ - if constexpr (K > MaxSegmentSize) - { - state.skip("K > MaxSegmentSize."); - return; - } - - const auto num_segments = static_cast(state.get_int64("NumSegments")); - const thrust::device_vector d_segment_sizes = generate( - static_cast(num_segments), - bit_entropy::_1_000, - static_cast(K), - static_cast(MaxSegmentSize)); - const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); - const auto output_elements = static_cast(num_segments) * K; - const auto total_num_items = cuda::args::immediate{static_cast(input_elements)}; - - auto in_keys_buffer = gen_data( - num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data())); - auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); - auto out_indices_buffer = thrust::device_vector(output_elements, thrust::no_init); - - auto segment_sizes_param = cuda::args::deferred_sequence{ - thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; - auto k_param = cuda::args::constant{}; - auto select_direction = cuda::args::constant{}; - auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; - - auto d_keys_in = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), - static_cast(MaxSegmentSize)); - auto d_keys_out = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), - static_cast(K)); - - // Input values: every segment maps to the same counting iterator starting at 0, so values are segment-local indices. - auto d_indices_in = cuda::make_constant_iterator(cuda::make_counting_iterator(IndexT{0})); - auto d_indices_out = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(out_indices_buffer.data())), - static_cast(K)); - - state.add_element_count(input_elements, "NumElements"); - state.add_global_memory_reads(input_elements, "InputKeys"); - state.add_global_memory_reads(num_segments, "SegmentSizes"); - state.add_global_memory_writes(output_elements, "OutputKeys"); - state.add_global_memory_writes(output_elements, "OutputIndices"); - - caching_allocator_t alloc; - state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env(alloc, launch); - // TODO(bgruber): call the public API once available - _CCCL_TRY_CUDA_API( - batched_topk_indexed, - "batched topk failed", - d_keys_in, - d_keys_out, - d_indices_in, - d_indices_out, - segment_sizes_param, - k_param, - select_direction, - num_segments_param, - total_num_items, - env); - }); -} - -// Index type is a compile-time axis: i32 for now, extensible to i64. -using index_type_list = nvbench::type_list; - -NVBENCH_BENCH_TYPES(decode_style_variable_topk_indexed, - NVBENCH_TYPE_AXES(key_type_list, index_type_list, max_segment_size_list, k_list)) - .set_name("decode_style_variable_topk_indexed") - .set_type_axes_names({"KeyT{ct}", "IndexT{ct}", "MaxSegmentSize{ct}", "K{ct}"}) - .add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32}) - .add_string_axis("Pattern", valid_patterns); +// Baseline-backend variable-segment indexed (arg-top-k) benchmark (the plain/default indexed benchmark). A base build +// runs the library's automatic selector; a tuning variant forces the baseline backend and sweeps its worker knobs. The +// baseline backend has no determinism / tie-break, so this build stays non-deterministic (TUNE_REQUIREMENT 0). The +// cluster backend and the requirement sweep live in `indexed.cluster.cu`. + +// %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1 +// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32 +// %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1 + +#ifndef TUNE_BACKEND +# if TUNE_BASE +# define TUNE_BACKEND 2 // automatic: the library's production selector, for base/benchmark builds +# else +# define TUNE_BACKEND 0 // baseline: the backend this file tunes +# endif +#endif + +#ifndef TUNE_REQUIREMENT +# define TUNE_REQUIREMENT 0 // baseline backend: non-deterministic only (a non-zero override trips the static_assert) +#endif + +#include "indexed_common.cuh" diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh new file mode 100644 index 00000000000..c8cfdebe1cd --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -0,0 +1,274 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Shared implementation for the variable-segment-size indexed (arg-top-k) benchmarks. The two thin TUs that include it +// -- `indexed.cu` (baseline) and `indexed.cluster.cu` (cluster) -- only set defaults and list their `%RANGE%` knobs; +// everything else (backend/requirement selection, the tuned `topk_backend_selector`, the dispatch wrapper, the nvbench +// body and its registration) lives here so the backends stay in lock-step. Each includer must define `TUNE_BACKEND` +// (0 baseline / 1 cluster / 2 automatic) and `TUNE_REQUIREMENT` (0 non-det / 1 det + prefer-smaller-index / 2 det + +// prefer-larger-index) before including. + +#pragma once + +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check so these benchmarks compile for the full +// configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) +// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "common.cuh" + +#ifndef TUNE_BACKEND +# error "indexed_common.cuh requires the includer to define TUNE_BACKEND (0 baseline / 1 cluster / 2 automatic)" +#endif +#ifndef TUNE_REQUIREMENT +# error \ + "indexed_common.cuh requires the includer to define TUNE_REQUIREMENT (0 non-det / 1 det+smaller / 2 det+larger)" +#endif + +enum class topk_backend +{ + baseline, + cluster, + automatic, +}; + +// Which backend this build benchmarks. `automatic` (the default) issues no `tune` override, leaving the choice to the +// library's arch/size selector -- convenient here since variable segment sizes can exceed the baseline backend's +// coverage (forcing baseline is only valid for coverable sizes). A tuning variant forces one concrete backend so its +// sub-policy knobs (below) take effect. +inline constexpr topk_backend selected_backend = +#if TUNE_BACKEND == 0 + topk_backend::baseline; +#elif TUNE_BACKEND == 1 + topk_backend::cluster; +#else + topk_backend::automatic; +#endif + +// The determinism / tie-break requirement this build issues. Only the cluster and automatic backends honor a +// deterministic result set / concrete tie-break preference; the baseline backend must stay on the non-deterministic +// path (enforced by the static_assert below). The three requirement values collapse the API's determinism/tie-break +// combinations to the distinct behaviors we benchmark: no guarantee, and deterministic with either index preference. +inline constexpr auto selected_determinism = +#if TUNE_REQUIREMENT == 0 + cuda::execution::determinism::__determinism_t::__not_guaranteed; +#else + cuda::execution::determinism::__determinism_t::__gpu_to_gpu; +#endif + +inline constexpr auto selected_tie_break = +#if TUNE_REQUIREMENT == 1 + cuda::execution::tie_break::__tie_break_t::__prefer_smaller_index; +#elif TUNE_REQUIREMENT == 2 + cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; +#else + cuda::execution::tie_break::__tie_break_t::__unspecified; +#endif + +// The baseline backend ignores determinism/tie-break, so require the defaults there to avoid a silently ignored +// selection. The cluster backend implements them, and automatic honors them via the library's selector, so both allow +// non-defaults. +static_assert(selected_backend == topk_backend::cluster || selected_backend == topk_backend::automatic + || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed + && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), + "Only the cluster and automatic backends honor determinism/tie-break requirements; keep TUNE_REQUIREMENT " + "at 0 (non-deterministic) for the baseline backend."); + +// Policy selector threaded through the public API's tuning environment when a concrete backend is forced (not +// `automatic`). Its `.backend` pins the backend for this build. In a base build both sub-policies are the defaults; +// in a tuning variant only the forced backend's sub-policy is driven by this build's TUNE_* macros (the other stays +// default), so each backend's benchmark sweeps only its own knobs. +template +struct topk_backend_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability cc) const + -> cub::detail::batched_topk::topk_policy + { +#if !TUNE_BASE && TUNE_BACKEND == 0 + constexpr auto store_alg = cub::BLOCK_STORE_WARP_TRANSPOSE; +# if TUNE_BLOCK_LOAD_ALGORITHM == 0 + constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT; +# elif TUNE_BLOCK_LOAD_ALGORITHM == 1 + constexpr auto load_alg = cub::BLOCK_LOAD_WARP_TRANSPOSE; +# elif TUNE_BLOCK_LOAD_ALGORITHM == 2 + constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE; +# endif + const auto baseline = cub::detail::batched_topk::baseline_topk_policy{{{ + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + }}}; +#else + const auto baseline = + cub::detail::batched_topk::baseline_policy_selector_from_types{}(cc); +#endif + +#if !TUNE_BASE && TUNE_BACKEND == 1 + const auto cluster = cub::detail::batched_topk::cluster_topk_policy{ + /*threads_per_block=*/TUNE_CLUSTER_THREADS_PER_BLOCK, + /*min_blocks_per_sm=*/TUNE_CLUSTER_MIN_BLOCKS_PER_SM, + /*min_chunks_per_block=*/TUNE_CLUSTER_MIN_CHUNKS_PER_BLOCK, + /*chunk_bytes=*/(TUNE_CLUSTER_CHUNK_KIB) * 1024, + /*load_align_bytes=*/(1 << (TUNE_CLUSTER_LOAD_ALIGN_BYTES_POW2)), + /*pipeline_stages=*/TUNE_CLUSTER_PIPELINE_STAGES, + /*single_block_max_seg_size=*/8 * 1024, + /*bits_per_pass=*/TUNE_CLUSTER_BITS_PER_PASS, + /*histogram_items_per_thread=*/TUNE_CLUSTER_HIST_IPT, + /*tie_break_items_per_thread=*/TUNE_CLUSTER_TIEBREAK_IPT, + /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT}; +#else + const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc); +#endif + + constexpr auto backend = + (selected_backend == topk_backend::cluster) + ? cub::detail::batched_topk::topk_backend::cluster + : cub::detail::batched_topk::topk_backend::baseline; + return cub::detail::batched_topk::topk_policy{backend, baseline, cluster}; + } +}; + +// Env-based dispatch over the selected backend. `automatic` routes through the public `cub::DeviceBatchedTopK` API with +// no backend override, so the library's own selector chooses; a forced backend routes through the same API but adds a +// `tune`d `topk_backend_selector` that pins the backend. Either way the determinism/tie-break/ordering requirement is +// layered on top of the benchmark environment (which carries the stream and the caching memory resource). +template +CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_indexed( + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + ValueInputItItT d_values_in, + ValueOutputItItT d_values_out, + SegmentSizeParamT segment_sizes, + KParamT k, + NumSegmentsParameterT num_segments, + EnvT env) +{ + auto req_env = cuda::std::execution::env{ + env, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted)}; + if constexpr (selected_backend == topk_backend::automatic) + { + // No `tune` override: the library's own selector picks the backend (arch/size crossover) -- the usual behavior. + return cub::DeviceBatchedTopK::MaxPairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, req_env); + } + else + { + using key_t = cub::detail::it_value_t>; + using value_t = cub::detail::it_value_t>; + constexpr auto max_k = ::cuda::args::__traits::highest; + auto full_env = cuda::std::execution::env{ + req_env, cuda::execution::tune(topk_backend_selector{})}; + return cub::DeviceBatchedTopK::MaxPairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, full_env); + } +} + +// Indexed (arg-top-k) variant: each key carries a segment-local index as its value payload. The input values are +// produced by a counting iterator that restarts at 0 for every segment, so indices are not (pre-)materialized in global +// memory +template +void decode_style_variable_topk_indexed( + nvbench::state& state, nvbench::type_list, nvbench::enum_type>) +{ + if constexpr (K > MaxSegmentSize) + { + state.skip("K > MaxSegmentSize."); + return; + } + + const auto num_segments = static_cast(state.get_int64("NumSegments")); + const thrust::device_vector d_segment_sizes = generate( + static_cast(num_segments), + bit_entropy::_1_000, + static_cast(K), + static_cast(MaxSegmentSize)); + const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); + const auto output_elements = static_cast(num_segments) * K; + + auto in_keys_buffer = gen_data( + num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data())); + auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); + auto out_indices_buffer = thrust::device_vector(output_elements, thrust::no_init); + + auto segment_sizes_param = cuda::args::deferred_sequence{ + thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; + auto k_param = cuda::args::constant{}; + auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + + auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), + static_cast(MaxSegmentSize)); + auto d_keys_out = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), + static_cast(K)); + + // Input values: every segment maps to the same counting iterator starting at 0, so values are segment-local indices. + auto d_indices_in = cuda::make_constant_iterator(cuda::make_counting_iterator(IndexT{0})); + auto d_indices_out = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(out_indices_buffer.data())), + static_cast(K)); + + state.add_element_count(input_elements, "NumElements"); + state.add_global_memory_reads(input_elements, "InputKeys"); + state.add_global_memory_reads(num_segments, "SegmentSizes"); + state.add_global_memory_writes(output_elements, "OutputKeys"); + state.add_global_memory_writes(output_elements, "OutputIndices"); + + caching_allocator_t alloc; + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + auto env = cub_bench_env(alloc, launch); + _CCCL_TRY_CUDA_API( + batched_topk_indexed, + "batched topk failed", + d_keys_in, + d_keys_out, + d_indices_in, + d_indices_out, + segment_sizes_param, + k_param, + num_segments_param, + env); + }); +} + +// Index type is a compile-time axis: i32 for now, extensible to i64. +using index_type_list = nvbench::type_list; + +NVBENCH_BENCH_TYPES(decode_style_variable_topk_indexed, + NVBENCH_TYPE_AXES(key_type_list, index_type_list, max_segment_size_list, k_list)) + .set_name("decode_style_variable_topk_indexed") + .set_type_axes_names({"KeyT{ct}", "IndexT{ct}", "MaxSegmentSize{ct}", "K{ct}"}) + .add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32}) + .add_string_axis("Pattern", valid_patterns); diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu new file mode 100644 index 00000000000..870dcab13c2 --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cluster.cu @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Cluster-backend variable-segment keys-only top-k benchmark. A base build runs the library's production (automatic) +// selector; a tuning variant forces the cluster backend and sweeps its per-block knobs. This build stays on the +// non-deterministic path (TUNE_REQUIREMENT 0) -- the determinism / tie-break requirement sweep lives in the indexed +// benchmark (`indexed.cluster.cu`). The baseline backend is benchmarked separately in the plain `keys.cu`. + +// %RANGE% TUNE_CLUSTER_THREADS_PER_BLOCK tpb 128:1024:32 +// %RANGE% TUNE_CLUSTER_MIN_BLOCKS_PER_SM mbs 1:2:1 +// %RANGE% TUNE_CLUSTER_MIN_CHUNKS_PER_BLOCK mcb 1:2:1 +// %RANGE% TUNE_CLUSTER_CHUNK_KIB ckib 8:32:1 +// %RANGE% TUNE_CLUSTER_LOAD_ALIGN_BYTES_POW2 la 4:7:1 +// %RANGE% TUNE_CLUSTER_PIPELINE_STAGES ps 2:16:1 +// %RANGE% TUNE_CLUSTER_BITS_PER_PASS bpp 8:11:1 +// %RANGE% TUNE_CLUSTER_HIST_IPT hipt 1:24:1 +// %RANGE% TUNE_CLUSTER_TIEBREAK_IPT tipt 1:24:1 +// %RANGE% TUNE_CLUSTER_COPY_IPT cipt 1:24:1 + +#ifndef TUNE_BACKEND +# if TUNE_BASE +# define TUNE_BACKEND 3 // automatic: the library's production selector, for base/benchmark builds +# else +# define TUNE_BACKEND 1 // cluster: the backend this file tunes +# endif +#endif + +#ifndef TUNE_REQUIREMENT +# define TUNE_REQUIREMENT 0 // cluster tuned non-deterministic; requirement sweep lives in indexed.cluster.cu +#endif + +#include "keys_common.cuh" diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys.cu b/cub/benchmarks/bench/segmented_topk/variable/keys.cu index c90be0bc834..f5a48a338c4 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -1,196 +1,25 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -#include -#include -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include - -#include -#include - -#include - -#include "common.cuh" - -enum class topk_backend -{ - baseline, - cluster, - device, -}; - -inline constexpr topk_backend selected_backend = topk_backend::baseline; - -// Determinism / tie-break requirement benchmarked by the cluster backend (a single combination for now). -inline constexpr auto selected_determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; -inline constexpr auto selected_tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; - -// The baseline/device backends ignore these requirements, so require the defaults there to avoid a silent mismatch. -static_assert(selected_backend == topk_backend::cluster - || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed - && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), - "Only the cluster backend implements determinism/tie-break requirements; keep selected_determinism and " - "selected_tie_break at their defaults for the baseline/device backends."); - -// Env-based dispatch over the selected backend. The cluster and baseline backends route through their respective -// `dispatch_with_env` entry points (temporary storage is allocated from the memory resource carried by `env`); the -// device backend issues one `cub::DeviceTopK::MaxKeys` per segment, reading the host-side segment sizes. -template -CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( - KeyInputItItT d_keys_in, - KeyOutputItItT d_keys_out, - SegmentSizeParamT segment_sizes, - KParamT k, - SelectDirectionParamT select_direction, - NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items, - [[maybe_unused]] const HostSegSizeT* h_segment_sizes, - EnvT env) -{ - if constexpr (selected_backend == topk_backend::cluster) - { - return cub::detail::batched_topk_cluster::dispatch_with_env( - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - select_direction, - num_segments, - env); - } - else if constexpr (selected_backend == topk_backend::device) - { - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - const auto num_segs = cub::detail::params::get_param(num_segments, num_segments_val_t{0}); - - // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement - // on top of the benchmark environment (which carries the stream and the caching memory resource). - auto seg_env = cuda::std::execution::env{ - env, - cuda::execution::require(cuda::execution::determinism::not_guaranteed, - cuda::execution::output_ordering::unsorted)}; - - for (num_segments_val_t i = 0; i < num_segs; ++i) - { - const auto k_value = cub::detail::params::get_param(k, i); - const auto seg_size = h_segment_sizes[i]; - if (const auto err = cub::DeviceTopK::MaxKeys( - d_keys_in[i], - d_keys_out[i], - static_cast(seg_size), - static_cast(k_value), - seg_env); - err != cudaSuccess) - { - return err; - } - } - return cudaSuccess; - } - else - { - return cub::detail::batched_topk::dispatch_with_env( - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - select_direction, - num_segments, - total_num_items, - env); - } -} - -template -void decode_style_variable_topk_keys( - nvbench::state& state, nvbench::type_list, nvbench::enum_type>) -{ - if constexpr (K > MaxSegmentSize) - { - state.skip("K > MaxSegmentSize."); - return; - } - - const auto num_segments = static_cast(state.get_int64("NumSegments")); - const thrust::device_vector d_segment_sizes = generate( - static_cast(num_segments), - bit_entropy::_1_000, - static_cast(K), - static_cast(MaxSegmentSize)); - const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); - const auto output_elements = static_cast(num_segments) * K; - const auto total_num_items = cuda::args::immediate{static_cast(input_elements)}; - - auto in_keys_buffer = gen_data( - num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data())); - auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); - - auto segment_sizes_param = cuda::args::deferred_sequence{ - thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; - auto k_param = cuda::args::constant{}; - auto select_direction = cuda::args::constant{}; - auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; - - auto d_keys_in = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), - static_cast(MaxSegmentSize)); - auto d_keys_out = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), - static_cast(K)); - - state.add_element_count(input_elements, "NumElements"); - state.add_global_memory_reads(input_elements, "InputKeys"); - state.add_global_memory_reads(num_segments, "SegmentSizes"); - state.add_global_memory_writes(output_elements, "OutputKeys"); - - // Host copy of segment sizes — consumed only by the per-segment device backend. - std::vector h_segment_sizes(static_cast(num_segments)); - thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), h_segment_sizes.begin()); - - caching_allocator_t alloc; - state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env(alloc, launch); - // TODO(bgruber): call the public API once available - _CCCL_TRY_CUDA_API( - batched_topk_keys, - "batched topk failed", - d_keys_in, - d_keys_out, - segment_sizes_param, - k_param, - select_direction, - num_segments_param, - total_num_items, - h_segment_sizes.data(), - env); - }); -} - -NVBENCH_BENCH_TYPES(decode_style_variable_topk_keys, NVBENCH_TYPE_AXES(key_type_list, max_segment_size_list, k_list)) - .set_name("decode_style_variable_topk_keys") - .set_type_axes_names({"KeyT{ct}", "MaxSegmentSize{ct}", "K{ct}"}) - .add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32}) - .add_string_axis("Pattern", valid_patterns); +// Baseline-backend variable-segment keys-only top-k benchmark (the plain/default keys benchmark). A base build runs the +// library's automatic selector; a tuning variant forces the baseline backend and sweeps its worker knobs. The baseline +// backend has no determinism / tie-break, so this build stays non-deterministic (TUNE_REQUIREMENT 0). The cluster +// backend is tuned separately in `keys.cluster.cu`; the `device` reference is reachable via -DTUNE_BACKEND=2. + +// %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1 +// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32 +// %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1 + +#ifndef TUNE_BACKEND +# if TUNE_BASE +# define TUNE_BACKEND 3 // automatic: the library's production selector, for base/benchmark builds +# else +# define TUNE_BACKEND 0 // baseline: the backend this file tunes +# endif +#endif + +#ifndef TUNE_REQUIREMENT +# define TUNE_REQUIREMENT 0 // baseline backend: non-deterministic only (a non-zero override trips the static_assert) +#endif + +#include "keys_common.cuh" diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh new file mode 100644 index 00000000000..a713cbdb5ad --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -0,0 +1,298 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// Shared implementation for the variable-segment-size keys-only top-k benchmarks. The two thin TUs that include it -- +// `keys.cu` (baseline) and `keys.cluster.cu` (cluster) -- only set defaults and list their `%RANGE%` knobs; everything +// else (backend/requirement selection, the tuned `topk_backend_selector`, the dispatch wrapper including the +// per-segment `device` reference, the nvbench body and its registration) lives here so the backends stay in lock-step. +// Each includer must define `TUNE_BACKEND` (0 baseline / 1 cluster / 2 device / 3 automatic) and `TUNE_REQUIREMENT` +// (0 non-det / 1 det + prefer-smaller-index / 2 det + prefer-larger-index) before including. + +#pragma once + +// Defer the unsupported-architecture diagnosis to the dispatch's runtime check so these benchmarks compile for the full +// configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) +// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. +#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include "common.cuh" + +#ifndef TUNE_BACKEND +# error "keys_common.cuh requires the includer to define TUNE_BACKEND (0 baseline/1 cluster/2 device/3 automatic)" +#endif +#ifndef TUNE_REQUIREMENT +# error "keys_common.cuh requires the includer to define TUNE_REQUIREMENT (0 non-det / 1 det+smaller / 2 det+larger)" +#endif + +enum class topk_backend +{ + baseline, + cluster, + device, + automatic, +}; + +// Which backend this build benchmarks. `automatic` (the default) issues no `tune` override, leaving the choice to the +// library's arch/size selector -- convenient here since variable segment sizes can exceed the baseline backend's +// coverage (forcing baseline is only valid for coverable sizes). A tuning variant forces one concrete backend so its +// sub-policy knobs (below) take effect; `device` issues one `cub::DeviceTopK` call per segment and has no knobs. +inline constexpr topk_backend selected_backend = +#if TUNE_BACKEND == 0 + topk_backend::baseline; +#elif TUNE_BACKEND == 1 + topk_backend::cluster; +#elif TUNE_BACKEND == 2 + topk_backend::device; +#else + topk_backend::automatic; +#endif + +// The determinism / tie-break requirement this build issues. Only the cluster and automatic backends honor a +// deterministic result set / concrete tie-break; the baseline and device backends must stay non-deterministic (enforced +// by the static_assert below). The three values cover the distinct behaviors we benchmark: no guarantee, and +// deterministic with either index preference. +inline constexpr auto selected_determinism = +#if TUNE_REQUIREMENT == 0 + cuda::execution::determinism::__determinism_t::__not_guaranteed; +#else + cuda::execution::determinism::__determinism_t::__gpu_to_gpu; +#endif + +inline constexpr auto selected_tie_break = +#if TUNE_REQUIREMENT == 1 + cuda::execution::tie_break::__tie_break_t::__prefer_smaller_index; +#elif TUNE_REQUIREMENT == 2 + cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; +#else + cuda::execution::tie_break::__tie_break_t::__unspecified; +#endif + +// The baseline/device backends ignore determinism/tie-break, so require the defaults there to avoid a silently ignored +// selection. The cluster backend implements them, and automatic honors them via the library's selector, so both allow +// non-defaults. +static_assert(selected_backend == topk_backend::cluster || selected_backend == topk_backend::automatic + || (selected_determinism == cuda::execution::determinism::__determinism_t::__not_guaranteed + && selected_tie_break == cuda::execution::tie_break::__tie_break_t::__unspecified), + "Only the cluster and automatic backends honor determinism/tie-break requirements; keep TUNE_REQUIREMENT " + "at 0 (non-deterministic) for the baseline/device backends."); + +// Policy selector threaded through the public API's tuning environment when a concrete backend is forced (not +// `automatic`). Its `.backend` pins the backend for this build. In a base build both sub-policies are the defaults; +// in a tuning variant only the forced backend's sub-policy is driven by this build's TUNE_* macros (the other stays +// default), so each backend's benchmark sweeps only its own knobs. +template +struct topk_backend_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability cc) const + -> cub::detail::batched_topk::topk_policy + { +#if !TUNE_BASE && TUNE_BACKEND == 0 + constexpr auto store_alg = cub::BLOCK_STORE_WARP_TRANSPOSE; +# if TUNE_BLOCK_LOAD_ALGORITHM == 0 + constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT; +# elif TUNE_BLOCK_LOAD_ALGORITHM == 1 + constexpr auto load_alg = cub::BLOCK_LOAD_WARP_TRANSPOSE; +# elif TUNE_BLOCK_LOAD_ALGORITHM == 2 + constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE; +# endif + const auto baseline = cub::detail::batched_topk::baseline_topk_policy{{{ + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg}, + }}}; +#else + const auto baseline = + cub::detail::batched_topk::baseline_policy_selector_from_types{}(cc); +#endif + +#if !TUNE_BASE && TUNE_BACKEND == 1 + const auto cluster = cub::detail::batched_topk::cluster_topk_policy{ + /*threads_per_block=*/TUNE_CLUSTER_THREADS_PER_BLOCK, + /*min_blocks_per_sm=*/TUNE_CLUSTER_MIN_BLOCKS_PER_SM, + /*min_chunks_per_block=*/TUNE_CLUSTER_MIN_CHUNKS_PER_BLOCK, + /*chunk_bytes=*/(TUNE_CLUSTER_CHUNK_KIB) * 1024, + /*load_align_bytes=*/(1 << (TUNE_CLUSTER_LOAD_ALIGN_BYTES_POW2)), + /*pipeline_stages=*/TUNE_CLUSTER_PIPELINE_STAGES, + /*single_block_max_seg_size=*/8 * 1024, + /*bits_per_pass=*/TUNE_CLUSTER_BITS_PER_PASS, + /*histogram_items_per_thread=*/TUNE_CLUSTER_HIST_IPT, + /*tie_break_items_per_thread=*/TUNE_CLUSTER_TIEBREAK_IPT, + /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT}; +#else + const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc); +#endif + + constexpr auto backend = + (selected_backend == topk_backend::cluster) + ? cub::detail::batched_topk::topk_backend::cluster + : cub::detail::batched_topk::topk_backend::baseline; + return cub::detail::batched_topk::topk_policy{backend, baseline, cluster}; + } +}; + +// Env-based dispatch over the selected backend. `automatic`/`baseline`/`cluster` all route through the public +// `cub::DeviceBatchedTopK` API (the latter two add a `tune`d `topk_backend_selector` that forces the backend; temp +// storage comes from the memory resource carried by `env`); the `device` backend issues one `cub::DeviceTopK::MaxKeys` +// per segment, reading the host-side segment sizes. +template +CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + SegmentSizeParamT segment_sizes, + KParamT k, + NumSegmentsParameterT num_segments, + [[maybe_unused]] const HostSegSizeT* h_segment_sizes, + EnvT env) +{ + if constexpr (selected_backend == topk_backend::device) + { + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + const auto num_segs = cub::detail::params::get_param(num_segments, num_segments_val_t{0}); + + // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement + // on top of the benchmark environment (which carries the stream and the caching memory resource). + auto seg_env = cuda::std::execution::env{ + env, + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::output_ordering::unsorted)}; + + for (num_segments_val_t i = 0; i < num_segs; ++i) + { + const auto k_value = cub::detail::params::get_param(k, i); + const auto seg_size = h_segment_sizes[i]; + if (const auto err = cub::DeviceTopK::MaxKeys( + d_keys_in[i], + d_keys_out[i], + static_cast(seg_size), + static_cast(k_value), + seg_env); + err != cudaSuccess) + { + return err; + } + } + return cudaSuccess; + } + else + { + // The determinism / tie-break / ordering requirement this benchmark issues; the library honors it whether or not we + // additionally force a backend. + auto req_env = cuda::std::execution::env{ + env, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted)}; + if constexpr (selected_backend == topk_backend::automatic) + { + // No `tune` override: the library's own selector picks the backend (arch/size crossover) -- the usual behavior. + return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, req_env); + } + else + { + using key_t = cub::detail::it_value_t>; + constexpr auto max_k = ::cuda::args::__traits::highest; + auto full_env = cuda::std::execution::env{ + req_env, cuda::execution::tune(topk_backend_selector{})}; + return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, full_env); + } + } +} + +template +void decode_style_variable_topk_keys( + nvbench::state& state, nvbench::type_list, nvbench::enum_type>) +{ + if constexpr (K > MaxSegmentSize) + { + state.skip("K > MaxSegmentSize."); + return; + } + + const auto num_segments = static_cast(state.get_int64("NumSegments")); + const thrust::device_vector d_segment_sizes = generate( + static_cast(num_segments), + bit_entropy::_1_000, + static_cast(K), + static_cast(MaxSegmentSize)); + const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); + const auto output_elements = static_cast(num_segments) * K; + + auto in_keys_buffer = gen_data( + num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data())); + auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); + + auto segment_sizes_param = cuda::args::deferred_sequence{ + thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; + auto k_param = cuda::args::constant{}; + auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + + auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), + static_cast(MaxSegmentSize)); + auto d_keys_out = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), + static_cast(K)); + + state.add_element_count(input_elements, "NumElements"); + state.add_global_memory_reads(input_elements, "InputKeys"); + state.add_global_memory_reads(num_segments, "SegmentSizes"); + state.add_global_memory_writes(output_elements, "OutputKeys"); + + // Host copy of segment sizes — consumed only by the per-segment device backend. + std::vector h_segment_sizes(static_cast(num_segments)); + thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), h_segment_sizes.begin()); + + caching_allocator_t alloc; + state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { + auto env = cub_bench_env(alloc, launch); + _CCCL_TRY_CUDA_API( + batched_topk_keys, + "batched topk failed", + d_keys_in, + d_keys_out, + segment_sizes_param, + k_param, + num_segments_param, + h_segment_sizes.data(), + env); + }); +} + +NVBENCH_BENCH_TYPES(decode_style_variable_topk_keys, NVBENCH_TYPE_AXES(key_type_list, max_segment_size_list, k_list)) + .set_name("decode_style_variable_topk_keys") + .set_type_axes_names({"KeyT{ct}", "MaxSegmentSize{ct}", "K{ct}"}) + .add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32}) + .add_string_axis("Pattern", valid_patterns); From d26c7b86c949cb7fc13d199f9f4f1845944cb852 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 5 Jul 2026 19:12:38 +0200 Subject: [PATCH 077/126] Avoid re-repriming streamer for filter Start ping-pong with the right direction from the start. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 42 ++-- ...catch2_test_device_segmented_topk_pairs.cu | 207 ++++++++++++++++++ 2 files changed, 234 insertions(+), 15 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index cf57faf19a6..5d10bde83c7 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1652,20 +1652,20 @@ private: // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once - // the top-k is fully placed. Only the straddling CTA needs scan order, so it forces `forward == !tie_reversed` and - // reuses the resident turn-around chunks only when the natural direction already matches; every other CTA is - // order-independent and keeps the natural direction, reusing the slots and skipping the re-prime copies. + // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the + // slots are already primed -- the histogram's first streaming pass primed the same persistent streamer, so + // `streamer.primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the + // generic fallback re-reads gmem each pass regardless). _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow() { - if (tie_active && !all_below_cta) - { - streamer.primed = (streamer.forward == (!tie_reversed)); - streamer.forward = !tie_reversed; - } - else - { - streamer.primed = true; - } + // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at + // `forward == !tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where the streaming + // direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary (`all_below_cta`, + // arrival-order atomics); or one with no back scan (`!tie_active`). Early stop is subsumed -- it makes every CTA + // `all_below_cta`, so the shorter (possibly mis-parity) pass count never reaches a straddler. + _CCCL_ASSERT(streamer.overflow_chunks == 0 || all_below_cta || !tie_active || streamer.forward == (!tie_reversed), + "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " + "the wrong streaming direction"); streamer.run_pass( // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. @@ -1920,6 +1920,17 @@ private: static_cast(stream_slots), my_chunks); + // Preselect the streamer's initial ping-pong direction. A streaming rank flips direction once per histogram pass, + // so the leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = + // (!tie_reversed) ^ (num_passes & 1)` makes that leftover `== !tie_reversed` -- exactly what the deterministic + // filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see + // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is + // moot). Non-deterministic filtering is order-independent, so leave its historical `forward` start untouched. + if constexpr (deterministic) + { + streamer.forward = (!tie_reversed) ^ ((num_passes & 1) != 0); + } + ::cuda::std::span resident_keys; // 32-bit shared-window address of `resident_keys.data()`. The resident span is read once per radix pass and in // the final filter; a 64-bit generic pointer kept live across that loop spills and reloads as a *generic* @@ -2089,9 +2100,10 @@ private: } } - // Fold the overflow chunks into the first-pass histogram. Forward direction; primes the streaming slots. The - // streamer reuses the resident load's stage mbarriers (all front-loaded at `run` entry); `wait_stage` provides - // the producer/consumer sync. + // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the streamer's initial + // direction (preselected above; the histogram is order-independent, so the direction only sets up the leftover + // parity for the final filter). The streamer reuses the resident load's stage mbarriers (all front-loaded at + // `run` entry); `wait_stage` provides the producer/consumer sync. streamer.process_pass(add_first_pass); const int resident_count = span_size(resident_keys); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index b39046a8676..84c26352449 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -911,6 +911,213 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the REQUIRE(ref == h_values_out); } +// The streaming tie-break regressions below need both `num_passes` parities so they exercise both ping-pong toggle +// counts. Enforce that the two widths actually straddle the parity against the real cluster `bits_per_pass` (rather +// than hard-coding "uint32 -> odd, uint64 -> even"): if tuning ever makes both widths share a parity, this fails to +// compile. +using streaming_tie_pair_key_types = c2h::type_list; +template +inline constexpr int cluster_num_passes = + cub::detail::topk::calc_num_passes(cub::detail::batched_topk::make_cluster_policy().bits_per_pass); +static_assert(cluster_num_passes % 2 != cluster_num_passes % 2, + "streaming_tie_pair_key_types must cover both num_passes parities"); + +// Deterministic tie-break *while streaming*: the value-observable (index-set) test for the preselected ping-pong +// direction. Heavy ties (keys in [0, 7]) straddle the k-th bucket while a 1 Mi segment (>> resident cluster coverage) +// forces the straddling CTA onto the overflow-streaming path, so a wrong preselected direction picks the wrong tied +// indices -- observable here (unlike keys-only, where all tied keys are equal) via the value payload checked against +// the host index reference. Two key widths cover both `num_passes` parities; `tie_break_pref_list` both tie-break +// directions. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the index-ordered top-k", + "[pairs][segmented][topk][device][cluster][determinism]", + streaming_tie_pair_key_types, + tie_break_pref_list) +{ + using key_t = c2h::get<0, TestType>; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = cub::detail::topk::select::max; + constexpr auto tie_break = c2h::get<1, TestType>::value; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr bool prefer_larger = tie_break == cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + + constexpr segment_size_t static_max_segment_size = 1024 * 1024; + constexpr segment_size_t static_max_k = 512 * 1024; + constexpr segment_index_t num_segments = 2; + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); + + const segment_size_t segment_size = static_max_segment_size - 31; // unaligned -> streaming + peeled tail edge + // With 8 near-uniform values each bucket holds ~segment_size/8 (~128 Ki) keys: k == 1 / 50 Ki straddle the top bucket + // (empty front), k == 300 Ki spans full buckets plus a mid-bucket straddle (non-empty front). + const segment_size_t k = + GENERATE_COPY(values({segment_size_t{1}, segment_size_t{50 * 1024}, segment_size_t{300 * 1024}})); + + CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, prefer_larger); + + // Few distinct key values -> many tied candidates in the k-th bucket. Contiguous -> BlockLoadToShared (streaming, as + // the 1 Mi segment exceeds resident coverage). + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer, key_t{0}, key_t{7}); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + // Values = global flattened index, so each selected value points back into the flattened input. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); + + // Values still belong to their keys, and no source index is selected twice. + c2h::device_vector expected_keys(keys_in_buffer); + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // The deterministic path must return *exactly* the index-ordered top-k. Compare per-segment selected index sets. + c2h::host_vector h_keys = keys_in_buffer; + const c2h::host_vector ref = + reference_deterministic_topk_indices(h_keys, num_segments, segment_size, k, direction, prefer_larger); + + c2h::host_vector h_values_out = values_out_buffer; + for (segment_index_t seg = 0; seg < num_segments; ++seg) + { + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + std::sort(seg_begin, seg_begin + k); + } + + REQUIRE(ref == h_values_out); +} + +// Maps a flattened element index to one of 8 values (host/device consistent), so a large segment has many duplicates +// and the k-th bucket holds a large tied set. Same twiddle as the keys test's heavy-tie op. +template +struct heavy_tie_key_op +{ + template + __host__ __device__ KeyT operator()(IndexT i) const + { + const unsigned h = static_cast(static_cast(i) * 2654435761u); + return static_cast(static_cast((h >> 13) & 7u)); + } +}; + +// Non-contiguous per-segment key iterator (like `counting_segment_keys_op`, but heavy-tie instead of identity): the +// transform iterator makes `use_block_load_to_shared` false, forcing the generic streaming path, while the heavy-tie +// op on the global index injects the boundary straddle. Segment `seg` -> keys at [seg*segment_size, ...). +template +struct heavy_tie_segment_keys_op +{ + SegmentSizeT segment_size; + + template + __host__ __device__ auto operator()(IndexT seg) const + { + return cuda::make_transform_iterator( + cuda::make_counting_iterator(static_cast(seg) * segment_size), heavy_tie_key_op{}); + } +}; + +// Generic-path counterpart of the deterministic streaming tie-break test above: a non-contiguous key iterator forces +// the agent's generic overflow-streaming path (not BlockLoadToShared) while still straddling the k-th bucket, so it +// covers the preselected ping-pong direction on the generic streamer for both `num_passes` parities and both tie-break +// directions. Value-observable: the value payload (global index) is checked against the host index reference. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the index-ordered top-k (generic path)", + "[pairs][segmented][topk][device][cluster][determinism]", + streaming_tie_pair_key_types, + tie_break_pref_list) +{ + using key_t = c2h::get<0, TestType>; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = cub::detail::topk::select::max; + constexpr auto tie_break = c2h::get<1, TestType>::value; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr bool prefer_larger = tie_break == cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + + constexpr segment_size_t static_max_segment_size = 1024 * 1024; + constexpr segment_size_t static_max_k = 512 * 1024; + constexpr segment_index_t num_segments = 2; + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + skip_if_batched_topk_backend_unavailable(static_max_segment_size); + + // Non-round size -> streaming (the generic fallback reads any trailing items straight from gmem; it never peels). + const segment_size_t segment_size = static_max_segment_size - 31; + const segment_size_t k = + GENERATE_COPY(values({segment_size_t{1}, segment_size_t{50 * 1024}, segment_size_t{300 * 1024}})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, prefer_larger); + + // Non-contiguous key input (transform iterator) -> generic streaming path; heavy-tie values straddle the k-th bucket. + auto d_keys_in = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), heavy_tie_segment_keys_op{segment_size}); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + // Values = global flattened index, so each selected value points back into the flattened input. + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + batched_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); + + // Materialize the (non-contiguous) key input for verification: same heavy-tie function of the global index. + c2h::host_vector h_keys(static_cast(num_items)); + heavy_tie_key_op key_op{}; + for (segment_size_t idx = 0; idx < num_items; ++idx) + { + h_keys[static_cast(idx)] = key_op(idx); + } + const c2h::device_vector keys_materialized = h_keys; + + // Values still belong to their keys, and no source index is selected twice. + REQUIRE(verify_pairs_consistency(keys_materialized, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // The deterministic path must return *exactly* the index-ordered top-k. Compare per-segment selected index sets. + const c2h::host_vector ref = + reference_deterministic_topk_indices(h_keys, num_segments, segment_size, k, direction, prefer_larger); + + c2h::host_vector h_values_out = values_out_buffer; + for (segment_index_t seg = 0; seg < num_segments; ++seg) + { + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + std::sort(seg_begin, seg_begin + k); + } + + REQUIRE(ref == h_values_out); +} + // Whole-`topk_policy` tuning overrides for the reproducibility test, threaded through the public API's single tuning // query (`cuda::execution::tune`). Both force the cluster backend (which alone honors a deterministic request). // `default_cluster_selector` uses the default cluster tuning; `alt_cluster_selector` starts from it and overrides the From e38149542d48f6f3e865861721bdc20323179cef Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 6 Jul 2026 01:14:47 +0200 Subject: [PATCH 078/126] Move CDP kernel to kernel header --- .../device/dispatch/dispatch_batched_topk.cuh | 92 +------------------ .../dispatch/kernels/kernel_batched_topk.cuh | 89 ++++++++++++++++++ 2 files changed, 92 insertions(+), 89 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 9627ce48f3d..004724dee4b 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -397,95 +397,9 @@ template return ::cuda::args::__highest_(segment_sizes); } -// Host launches go through the single kernel symbol (`device_batched_topk_kernel`, see kernel_batched_topk.cuh). Only -// the CDP static-cluster kernel below remains a dedicated cluster symbol, because device-side launches cannot opt in to -// dynamic cluster dimensions. - -#ifdef CUB_RDC_ENABLED -// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so the -// triple-chevron launch from device code needs no `cudaFuncSetAttribute`. -template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) - _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t block_tile_capacity) -{ - using agent_t = batched_topk_cluster::agent_batched_topk_cluster< - ThreadsPerBlock, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleBlockMaxSegSize, - MinChunksPerBlock, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - __shared__ typename agent_t::TempStorage temp_storage; - extern __shared__ char topk_cluster_smem[]; - char* key_slots = topk_cluster_smem; - // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` - // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. - { - ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); - smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); - asm("" : "+r"(smem32)); - key_slots = static_cast(__cvta_shared_to_generic(smem32)); - } - - agent_t agent( - temp_storage, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - key_slots, - block_tile_capacity); - - agent.Process(); -} -#endif // CUB_RDC_ENABLED +// Host launches go through the single kernel symbol (`device_batched_topk_kernel`); the CDP path uses a dedicated +// static-cluster kernel symbol (`device_segmented_topk_cluster_kernel_static`) because device-side launches cannot opt +// in to dynamic cluster dimensions. Both kernels live in kernel_batched_topk.cuh. // CDP launch body, empty when CDP is disabled. Wrapped in a macro because // `#ifdef` can't sit inside `NV_IF_TARGET`. diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 044258e776f..b557c1efd9a 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -20,6 +20,7 @@ #include #include #include +#include // max_portable_cluster_blocks (CDP static-cluster kernel) #include #include @@ -443,6 +444,94 @@ __launch_bounds__( return; } } + +#ifdef CUB_RDC_ENABLED +// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so a device-side (CDP) triple-chevron launch needs no +// `cudaFuncSetAttribute` (device-side launches cannot opt into dynamic cluster dimensions the way the host +// `device_batched_topk_kernel` does). It mirrors that kernel's cluster arm with a fixed cluster dim, and is consumed by +// the CDP arm of `launch_cluster_arm` (see `CUB_TOPK_CLUSTER_DEVICE_LAUNCH` in dispatch_batched_topk.cuh). +template +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) + _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + ::cuda::std::uint32_t block_tile_capacity) +{ + using agent_t = batched_topk_cluster::agent_batched_topk_cluster< + ThreadsPerBlock, + HistogramItemsPerThread, + PipelineStages, + ChunkBytes, + LoadAlignBytes, + BitsPerPass, + TieBreakItemsPerThread, + SingleBlockMaxSegSize, + MinChunksPerBlock, + CopyItemsPerThread, + Determinism, + TieBreak, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` + // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } + + agent_t agent( + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + key_slots, + block_tile_capacity); + + agent.Process(); +} +#endif // CUB_RDC_ENABLED } // namespace detail::batched_topk CUB_NAMESPACE_END From 53d73504e5fea9ba2d254bd938d3d23945432415 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 6 Jul 2026 14:15:58 +0200 Subject: [PATCH 079/126] Remove dead code and avoid repetition --- cub/cub/agent/agent_batched_topk_cluster.cuh | 173 ++++++------------- 1 file changed, 56 insertions(+), 117 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5d10bde83c7..d62498f3e76 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -102,7 +102,7 @@ struct alignas(16) cluster_topk_state // Decoders that operate on an already-loaded `result_pair`. The hot pass loop reads the leader's word once through // DSMEM and decodes both halves from that single value, so the splitter fold and the early-stop check never issue a - // second remote load. The instance accessors below delegate here, keeping the bit layout defined in one place. + // second remote load. `set_result` packs the same layout, keeping the bit layout defined in one place. [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static ::cuda::std::uint32_t kth_bucket_of(::cuda::std::uint64_t packed) { @@ -112,14 +112,6 @@ struct alignas(16) cluster_topk_state { return (packed >> 32) != ::cuda::std::uint64_t{0}; } - [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t kth_bucket() const - { - return kth_bucket_of(result_pair); - } - [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool early_stop() const - { - return early_stop_of(result_pair); - } _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void set_result(::cuda::std::uint32_t bucket, bool stop) { result_pair = static_cast<::cuda::std::uint64_t>(bucket) | (static_cast<::cuda::std::uint64_t>(stop) << 32); @@ -308,27 +300,23 @@ struct agent_batched_topk_cluster static constexpr int segment_rounds_floor = clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : 0; + // Clamp a per-thread unroll down to the segment's bounded round count (only when `clamp_items_to_segment`); + // larger/unbounded segments keep the full tuning width, and the guard keeps the rounds arithmetic in `int`. + [[nodiscard]] static constexpr int clamp_unroll(int rounds, int tuning) + { + return clamp_items_to_segment ? ::cuda::std::clamp(rounds, 1, tuning) : tuning; + } + // Two clamp flavors. The `floor` clamp pairs with an unpredicated main loop over full tiles plus a single // non-unrolled remainder loop (the chunk helper `for_each_chunk_key_impl` and the copy fast path); the `ceil` clamp - // keeps the whole resident segment inside one tile for fully-predicated loops (the deterministic filter). Each unroll - // is capped at its tuning value; non-clamped (large/unbounded) segments keep the full width, and the guard keeps the - // rounds arithmetic in `int`. + // keeps the whole resident segment inside one tile for fully-predicated loops (the deterministic filter). static constexpr int histogram_items_per_thread_clamped = - clamp_items_to_segment - ? ::cuda::std::clamp(segment_rounds_floor, 1, histogram_items_per_thread) - : histogram_items_per_thread; - // `ceil`-clamped for the deterministic filter's fully-predicated loops; `floor`-clamped for the non-deterministic - // path, which reuses the chunk helper. + clamp_unroll(segment_rounds_floor, histogram_items_per_thread); static constexpr int tie_break_items_per_thread_clamped = - clamp_items_to_segment - ? ::cuda::std::clamp(segment_rounds_ceil, 1, tie_break_items_per_thread) - : tie_break_items_per_thread; + clamp_unroll(segment_rounds_ceil, tie_break_items_per_thread); static constexpr int tie_break_items_per_thread_floor_clamped = - clamp_items_to_segment - ? ::cuda::std::clamp(segment_rounds_floor, 1, tie_break_items_per_thread) - : tie_break_items_per_thread; - static constexpr int copy_items_per_thread_clamped = - clamp_items_to_segment ? ::cuda::std::clamp(segment_rounds_floor, 1, copy_items_per_thread) : copy_items_per_thread; + clamp_unroll(segment_rounds_floor, tie_break_items_per_thread); + static constexpr int copy_items_per_thread_clamped = clamp_unroll(segment_rounds_floor, copy_items_per_thread); static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); static_assert(tie_break_items_per_thread_floor_clamped >= 1, @@ -429,53 +417,11 @@ struct agent_batched_topk_cluster // Split point of `edge_keys`: head edge in `[0, head_edge_cap)`, tail edge in `[head_edge_cap, 2 * head_edge_cap)`. static constexpr int head_edge_cap = load_align_items; - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span block_tile_buffer() const - { - const int slots = static_cast(block_tile_capacity / chunk_items); - return {key_slots, static_cast<::cuda::std::size_t>(slots * ChunkBytes)}; - } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE key_t* slot_keys_unpadded(int slot) const { return reinterpret_cast(key_slots + slot * ChunkBytes); } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span - available_block_tile_buffer(char* buffer_begin) const - { - const auto buffer = block_tile_buffer(); - char* const end = ::cuda::std::data(buffer) + static_cast(::cuda::std::size(buffer)); - _CCCL_ASSERT(buffer_begin >= ::cuda::std::data(buffer) && buffer_begin <= end, "Invalid block_tile buffer cursor"); - _CCCL_ASSERT(::cuda::is_aligned(buffer_begin, detail::LoadToSharedBufferAlignBytes()), - "block_tile buffer cursor must satisfy BlockLoadToShared's shared-memory alignment"); - return {buffer_begin, static_cast<::cuda::std::size_t>(end - buffer_begin)}; - } - - _CCCL_DEVICE _CCCL_FORCEINLINE void - append_contiguous_span(::cuda::std::span& merged, ::cuda::std::span next) const - { - const int next_count = static_cast(::cuda::std::size(next)); - _CCCL_ASSERT(static_cast<::cuda::std::size_t>(next_count) == ::cuda::std::size(next), - "Resident key span length must fit in int"); - if (next_count == 0) - { - return; - } - - const int merged_count = static_cast(::cuda::std::size(merged)); - _CCCL_ASSERT(static_cast<::cuda::std::size_t>(merged_count) == ::cuda::std::size(merged), - "Resident key span length must fit in int"); - if (merged_count == 0) - { - merged = next; - return; - } - - _CCCL_ASSERT(::cuda::std::data(merged) + merged_count == ::cuda::std::data(next), - "BlockLoadToShared returned non-contiguous resident key spans"); - merged = {::cuda::std::data(merged), static_cast<::cuda::std::size_t>(merged_count + next_count)}; - } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int span_size(::cuda::std::span keys) const { const int count = static_cast(::cuda::std::size(keys)); @@ -555,13 +501,6 @@ struct agent_batched_topk_cluster return {prefix, bulk, static_cast(chunk.count) - prefix - bulk}; } - template - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_aligned_chunk(PtrT base, const chunk_desc chunk) const - { - const auto s = split_chunk(base, chunk); - return s.prefix == 0 && s.suffix == 0; - } - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t num_rank_chunks(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { @@ -1694,74 +1633,74 @@ private: }); } - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged in `edge_keys`. In global- - // index order it is the leading region (ascending) / trailing region (descending). `head_edge_len == 0` is a no-op, - // so non-head ranks skip it. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge() + // Fold one persistent boundary edge (head prefix or peeled tail suffix), both staged in `edge_keys`. A no-op on the + // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge(const key_t* keys, offset_t seg_base, int count, bool terminal) { if constexpr (use_block_load_to_shared) { - // Head prefix `[0, head_edge_len)` staged at `edge_keys`, so its segment-local base is 0. process_tiles( - agent.temp_storage.edge_keys, offset_t{0}, head_edge_len, head_edge_terminal); + keys, seg_base, count, terminal); } } - // Peeled tail suffix edge (tail owner): the segment's highest indices, staged in `edge_keys`. In global-index order - // it is the trailing region (ascending) / leading region (descending). `tail_edge_len == 0` (aligned tail or not - // owned here) makes it a no-op. + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged at `edge_keys` (base 0). In + // global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or empty + // prefix is a no-op. + _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge() + { + process_edge(agent.temp_storage.edge_keys, offset_t{0}, head_edge_len, head_edge_terminal); + } + + // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap` (base + // `segment_size - tail_edge_len`). In global-index order it is the trailing region (ascending) / leading region + // (descending); an aligned or non-owned tail is a no-op. _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge() { - if constexpr (use_block_load_to_shared) - { - // Peeled tail suffix staged at `edge_keys + head_edge_cap`; its segment-local base is the last `tail_edge_len` - // indices of the segment. - const int tc = tail_edge_len; - const offset_t tail_base = segment_size_u32 - static_cast(tc); - process_tiles( - agent.temp_storage.edge_keys + head_edge_cap, tail_base, tc, tail_edge_terminal); - } + process_edge(agent.temp_storage.edge_keys + head_edge_cap, + segment_size_u32 - static_cast(tail_edge_len), + tail_edge_len, + tail_edge_terminal); } - // Drive the regions in global-index order (ascending, or descending under `tie_reversed`), bailing between regions - // once `should_stop` reports the whole top-k placed. + // Drive the four regions in global-index order (ascending, or descending under `tie_reversed`), bailing between + // regions once `should_stop` reports the whole top-k placed. The two orders share the resident->overflow middle and + // only swap which boundary edge leads: ascending is head, resident, overflow, tail; descending reverses the edges. + // Both visit resident before overflow, so `should_stop` can skip re-streaming the overflow once the top-k is + // placed; `reverse_residency` keeps the first-visited (high-index) chunks resident in the descending order so this + // holds. _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter() { - if constexpr (tie_reversed) - { - // Descending global-index order: peeled suffix tail, resident high-index chunks, streamed low-index overflow, - // peeled head prefix. With resident visited before overflow (`reverse_residency`), `should_stop` can skip the - // overflow re-stream -- mirroring the ascending path. - process_tail_edge(); - if (!should_stop()) - { - process_resident(); - } - if (!should_stop()) - { - process_overflow(); - } + const auto step = [&](auto&& region) { if (!should_stop()) { - process_head_edge(); + region(); } + }; + if constexpr (tie_reversed) + { + process_tail_edge(); } else { process_head_edge(); - if (!should_stop()) - { - process_resident(); - } - if (!should_stop()) + } + step([&] { + process_resident(); + }); + step([&] { + process_overflow(); + }); + step([&] { + if constexpr (tie_reversed) { - process_overflow(); + process_head_edge(); } - if (!should_stop()) + else { process_tail_edge(); } - } + }); } }; From ac027d28ba8a3e507216447447bc17ead8f22740 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 6 Jul 2026 15:47:18 +0200 Subject: [PATCH 080/126] Add emit_back_one function to avoid repetition --- cub/cub/agent/agent_batched_topk_cluster.cuh | 61 ++++++++++---------- 1 file changed, 29 insertions(+), 32 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index d62498f3e76..ac3599bc3c6 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1393,6 +1393,21 @@ private: return Reversed ? (count - 1 - pos) : pos; } + // Emit one back/tie candidate: if its `global_rank` (arrival- or scan-ordered by the caller) is a winner it lands + // in the top-k output at slot `k-1-global_rank`, with its value pulled from segment-local index `seg_idx`; ranks at + // or past `num_back` are losers and dropped (a no-op). Forced-inline with explicit args (no capture) so it folds + // into the hot back-placement loops with no extra live ranges; the caller keeps any counter side effects (e.g. + // `back_local_inc`) in the `global_rank` argument so they still run for every candidate. + _CCCL_DEVICE _CCCL_FORCEINLINE void emit_back_one(offset_t global_rank, const key_t& key, offset_t seg_idx) + { + if (global_rank < static_cast(num_back)) + { + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + block_keys_out[out] = key; + write_value(out, seg_idx); + } + } + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. `FromSmem` selects the // key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` // (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. Strictly- @@ -1469,14 +1484,10 @@ private: { if (valid[i] && flags[i] != offset_t{0}) { - const offset_t global_rank = cand_prefix + agent.back_local_inc(); - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, seg_base + static_cast(fold_pos(pos, count))); - } + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + emit_back_one(cand_prefix + agent.back_local_inc(), + keys[i], + seg_base + static_cast(fold_pos(pos, count))); } } } @@ -1490,14 +1501,9 @@ private: { if (valid[i] && flags[i] != offset_t{0}) { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, seg_base + static_cast(fold_pos(pos, count))); - } + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + emit_back_one( + running + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); } } running += tile_total; @@ -1510,14 +1516,10 @@ private: { if (valid[i] && flags[i] != offset_t{0}) { - const offset_t global_rank = cand_prefix + agent.back_local_inc(); - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, seg_base + static_cast(fold_pos(pos, count))); - } + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + emit_back_one(cand_prefix + agent.back_local_inc(), + keys[i], + seg_base + static_cast(fold_pos(pos, count))); } } // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make @@ -1536,14 +1538,9 @@ private: { if (valid[i] && flags[i] != offset_t{0}) { - const offset_t global_rank = running + excl[i]; - if (global_rank < static_cast(num_back)) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = keys[i]; - write_value(out, seg_base + static_cast(fold_pos(pos, count))); - } + const int pos = tile_base + static_cast(threadIdx.x) * items + i; + emit_back_one( + running + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); } } } From 692f05cf1ff2edf2fb3944b65e36005bcfc6f549 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 6 Jul 2026 19:35:47 +0200 Subject: [PATCH 081/126] Add sink removing repetition from write_selected[_idx] --- cub/cub/agent/agent_batched_topk_cluster.cuh | 152 ++++++++++--------- 1 file changed, 77 insertions(+), 75 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index ac3599bc3c6..231ad6a4cd9 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1044,10 +1044,10 @@ private: // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` resident in streaming // slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from gmem - // (fallback). The two entry points (`process_pass` / `process_pass_indexed`) differ only in whether the callable - // also gets the segment-local index, so share this loop. `mid()` runs once on a full pass -- after the first reload - // wave (`p_eff` visits) is issued but before it is waited on, overlapping the caller's resident-chunk work with - // those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with no unmatched barrier. + // (fallback). Both entry points delegate to `process_pass_dispatch`, which differs only in whether the callable + // also gets the segment-local index, so they share this loop. `mid()` runs once on a full pass -- after the first + // reload wave (`p_eff` visits) is issued but before it is waited on, overlapping the caller's resident-chunk work + // with those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with no unmatched barrier. // `should_continue()` is polled once after each consumed chunk (before its refill copy is issued); returning false // breaks the stream so the final filter bails once the top-k is placed. It must be block-uniform and is evaluated // uniformly by every lane, so it may contain a barrier (the deterministic filter's does, to read its placement @@ -1177,15 +1177,26 @@ private: } } - // Apply `f(key)` to every overflow key once, in the current ping-pong direction. See `run_pass` for the overlap - // semantics of `mid`. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers pass their - // clamped items-per-thread. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + // Fold every overflow key once in the current ping-pong direction. `Indexed` selects the callable shape: keys-only + // (`Indexed == false`) applies `f(key)`; the pair filter (`Indexed == true`) applies `f(key, seg_idx)` with each + // key's segment-local index, needed to fetch its value payload from gmem while still reusing the streamed overflow + // keys. The index math lives entirely inside the `if constexpr (Indexed)` arms, so it is elided in the keys-only + // instantiation. See `run_pass` for the overlap semantics of `mid`; `UnrollFactor` partially unrolls the generic + // (gmem fallback) fold loop. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_dispatch(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t o) { - agent.template for_each_chunk_key(stage_span(stage, o), f); + if constexpr (Indexed) + { + const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; + agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + } + else + { + agent.template for_each_chunk_key(stage_span(stage, o), f); + } }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1195,7 +1206,15 @@ private: const int local = j * threads_per_block + static_cast(threadIdx.x); if (local < chunk.count) { - f(block_keys_in[static_cast(chunk.offset + static_cast(local))]); + if constexpr (Indexed) + { + const offset_t seg_idx = chunk.offset + static_cast(local); + f(block_keys_in[static_cast(seg_idx)], seg_idx); + } + else + { + f(block_keys_in[static_cast(chunk.offset + static_cast(local))]); + } } } }, @@ -1205,6 +1224,14 @@ private: }); } + // `f(key)` over every overflow key. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers + // pass their clamped items-per-thread. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + { + process_pass_dispatch(static_cast(f), static_cast(mid)); + } + // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in // by the BlockLoadToShared pipeline (rather than already resident in SMEM). template @@ -1213,34 +1240,11 @@ private: process_pass(static_cast(f), [] {}); } - // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. The pair - // final filter needs that index to fetch each selected key's value payload from gmem, while still reusing the - // overflow keys from the streaming SMEM pipeline (block-load path) instead of re-reading them from gmem. + // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) { - run_pass( - [&](int stage, offset_t o) { - const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; - agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); - }, - [&](const auto& chunk) { - const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - _CCCL_PRAGMA_UNROLL(UnrollFactor) - for (int j = 0; j < iterations; ++j) - { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < chunk.count) - { - const offset_t seg_idx = chunk.offset + static_cast(local); - f(block_keys_in[static_cast(seg_idx)], seg_idx); - } - } - }, - static_cast(mid), - [] { - return true; - }); + process_pass_dispatch(static_cast(f), static_cast(mid)); } }; @@ -2365,24 +2369,42 @@ private: const offset_t sel_prefix = static_cast(pp >> 32); const offset_t cand_prefix = static_cast(pp & 0xffffffffu); - if constexpr (is_keys_only) - { - auto write_selected = [&](const key_t& key) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) + // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local + // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair + // mode each written key additionally pulls its value from gmem at `seg_idx`. Keys-only elides that write via the + // `if constexpr` below and drives the sink from index-free traversals, passing a dummy `seg_idx` that goes + // unread. + const auto sink = [&](const key_t& key, [[maybe_unused]] offset_t seg_idx) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); + block_keys_out[pos] = key; + if constexpr (!is_keys_only) { - const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); - block_keys_out[pos] = key; + write_value(pos, seg_idx); } - else if (res == detail::topk::candidate_class::candidate) + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); + if (back_pos < num_kth) { - const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); - if (back_pos < num_kth) + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + if constexpr (!is_keys_only) { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; + write_value(pos, seg_idx); } } + } + }; + + if constexpr (is_keys_only) + { + // Keys-only: no value payload, so drive the sink through the index-free traversal with a dummy `seg_idx`. + const auto write_selected = [&](const key_t& key) { + sink(key, offset_t{0}); }; // Fold the resident keys as the streamer's `mid` work so they overlap the first overflow reloads. Writes are // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. @@ -2411,35 +2433,15 @@ private: } else { - // Pair (key + value) path. Same block-local SMEM-atomic placement (`sel_prefix`/`cand_prefix`) and key reuse as - // the keys-only path; for each written key we additionally load its value from gmem at the segment-local index - // `seg_idx` and store it at the same slot (values are never streamed). Output order is not preserved, as the - // non-deterministic path permits. - auto write_selected_idx = [&](const key_t& key, offset_t seg_idx) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); - block_keys_out[pos] = key; - write_value(pos, seg_idx); - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); - if (back_pos < num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; - write_value(pos, seg_idx); - } - } - }; + // Pair (key + value) path: the shared `sink` above additionally stores each written key's value (fetched from + // gmem at its `seg_idx`). Unlike keys-only, the traversal must recover each key's segment-local index, so it + // folds resident/edge keys through `write_run` (below) and overflow keys through `process_pass_indexed`. // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary // edges); overflow chunks fold through the streamer's own indexed callback below. Materialize the key into a - // register first: `write_selected_idx` binds it by `const&` and reads it several times, so passing - // `smem[local]` directly would re-issue a narrow `LDS` per use instead of reusing the loaded value. + // register first: `sink` binds it by `const&` and reads it several times, so passing `smem[local]` directly + // would re-issue a narrow `LDS` per use instead of reusing the loaded value. auto write_run = [&](const key_t* smem, offset_t base_off, int count) { const int iterations = ::cuda::ceil_div(count, threads_per_block); _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) @@ -2449,7 +2451,7 @@ private: if (local < count) { const key_t key = smem[local]; - write_selected_idx(key, base_off + static_cast(local)); + sink(key, base_off + static_cast(local)); } } }; @@ -2501,7 +2503,7 @@ private: // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. - streamer.process_pass_indexed(write_selected_idx, fold_resident); + streamer.process_pass_indexed(sink, fold_resident); } } From d651595c920f6e33765a4f2871dadd8d1abadac3 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 01:21:08 +0200 Subject: [PATCH 082/126] Clean up variable/function naming --- cub/cub/agent/agent_batched_topk_cluster.cuh | 710 +++++++++--------- .../device/dispatch/dispatch_batched_topk.cuh | 2 +- ...catch2_test_device_segmented_topk_pairs.cu | 2 +- 3 files changed, 368 insertions(+), 346 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 231ad6a4cd9..4512dbeceba 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -108,13 +108,13 @@ struct alignas(16) cluster_topk_state { return static_cast<::cuda::std::uint32_t>(packed); } - [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static bool early_stop_of(::cuda::std::uint64_t packed) + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static bool is_early_stop(::cuda::std::uint64_t packed) { return (packed >> 32) != ::cuda::std::uint64_t{0}; } - _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void set_result(::cuda::std::uint32_t bucket, bool stop) + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void set_result(::cuda::std::uint32_t bucket, bool is_stop) { - result_pair = static_cast<::cuda::std::uint64_t>(bucket) | (static_cast<::cuda::std::uint64_t>(stop) << 32); + result_pair = static_cast<::cuda::std::uint64_t>(bucket) | (static_cast<::cuda::std::uint64_t>(is_stop) << 32); } }; @@ -169,7 +169,7 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, // Whether a segment takes the barrier-free single-CTA path: resident in one CTA (`<= block_tile_capacity`) and at/below // the single-CTA tuning threshold. Occupancy- and head-alignment-independent, so the host fast path and the device // collapse decision agree exactly. 64-bit math for wide segment-size types / loose bounds. -[[nodiscard]] _CCCL_HOST_DEVICE constexpr bool single_cta_eligible( +[[nodiscard]] _CCCL_HOST_DEVICE constexpr bool is_single_cta_eligible( ::cuda::std::uint64_t segment_size, ::cuda::std::uint64_t block_tile_capacity, int single_block_max_seg_size) noexcept { return segment_size <= block_tile_capacity @@ -240,23 +240,24 @@ struct agent_batched_topk_cluster using state_t = cluster_topk_state; using key_prefix_t = typename state_t::key_prefix_t; - // See the class comment: deterministic guarantees select the index-ordered scan; `tie_reversed` flips its order. - static constexpr bool deterministic = + // See the class comment: deterministic guarantees select the index-ordered scan; `is_tie_reversed` flips its order. + static constexpr bool need_determinism = Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed; - static constexpr bool tie_reversed = TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + static constexpr bool is_tie_reversed = + TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; // Push direction of the combined cross-CTA prefix scan (`combined_prefix_scan`). The leader must be *last* in scan // order so it derives its own (merged-away) counts from the predecessor sum: deterministic-prefer-smaller puts the // leader at the last effective rank and scans ascending; every other config (deterministic-prefer-larger and the // whole non-deterministic path) keeps the leader at rank 0 and scans descending, which makes rank 0 last. Matches the // `leader_rank` computed in `run`. - static constexpr bool scan_descending = !(deterministic && !tie_reversed); + static constexpr bool is_scan_descending = !(need_determinism && !is_tie_reversed); // The deterministic final scan visits chunks in global-index order and bails early (`should_stop`), so keeping the // *first-visited* chunks resident lets it skip re-reading the streamed overflow. Ascending visits low indices first // (the default low-resident split); descending (prefer-larger) visits high indices first, so we flip the split to // keep the high-index chunks resident (see `run`), restoring symmetry. Never set on the non-deterministic path. - static constexpr bool reverse_residency = deterministic && tie_reversed; + static constexpr bool is_residency_reversed = need_determinism && is_tie_reversed; // Segments at or below this size that also fit resident in one CTA take the single-CTA fast path (see // `process_impl`). @@ -290,21 +291,21 @@ struct agent_batched_topk_cluster // `ceil(static_max_segment_size / threads_per_block)` sweep rounds. We clamp each per-thread unroll down to that // bound to trim predication/registers on sub-tile segments. Larger/unbounded types keep the full unroll (codegen // unchanged); the guard also keeps the rounds arithmetic in `int` range. - static constexpr bool clamp_items_to_segment = + static constexpr bool should_clamp_items_to_segment = static_max_segment_size > 0 && static_max_segment_size < single_block_max_seg_size; static constexpr int segment_rounds_ceil = - clamp_items_to_segment + should_clamp_items_to_segment ? static_cast( ::cuda::ceil_div(static_max_segment_size, static_cast<::cuda::std::int64_t>(threads_per_block))) : 0; static constexpr int segment_rounds_floor = - clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : 0; + should_clamp_items_to_segment ? static_cast(static_max_segment_size / threads_per_block) : 0; - // Clamp a per-thread unroll down to the segment's bounded round count (only when `clamp_items_to_segment`); + // Clamp a per-thread unroll down to the segment's bounded round count (only when `should_clamp_items_to_segment`); // larger/unbounded segments keep the full tuning width, and the guard keeps the rounds arithmetic in `int`. [[nodiscard]] static constexpr int clamp_unroll(int rounds, int tuning) { - return clamp_items_to_segment ? ::cuda::std::clamp(rounds, 1, tuning) : tuning; + return should_clamp_items_to_segment ? ::cuda::std::clamp(rounds, 1, tuning) : tuning; } // Two clamp flavors. The `floor` clamp pairs with an unpredicated main loop over full tiles plus a single @@ -405,17 +406,18 @@ struct agent_batched_topk_cluster // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused // (ping-ponged) across radix passes; all are initialized once up front by `init_load_barriers`. ::cuda::std::uint64_t load_mbar[PipelineStages]; - // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap)`, on rank 0) and - // the peeled tail suffix (`[head_edge_cap, 2 * head_edge_cap)`, on the tail owner whenever it is unaligned), each - // strictly `< load_align_items` keys. Loaded once in the first pass and folded into every pass + the final filter. - // Block-local (never reached through DSMEM). + // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap_items)`, on rank + // 0) and the peeled tail suffix (`[head_edge_cap_items, 2 * head_edge_cap_items)`, on the tail owner whenever it is + // unaligned), each strictly `< load_align_items` keys. Loaded once in the first pass and folded into every pass + + // the final filter. Block-local (never reached through DSMEM). key_t edge_keys[2 * load_align_items]; }; // The `red.add.u64` on `prefix_pair`'s `.shared::cluster` address needs 8-byte alignment; the `uint64_t` member is // already at an 8-aligned struct offset, so guarding the struct's alignment covers the absolute (and peer) address. static_assert(alignof(_TempStorage) >= 8, "prefix_pair must be 8-byte aligned for the u64 DSMEM atomic"); - // Split point of `edge_keys`: head edge in `[0, head_edge_cap)`, tail edge in `[head_edge_cap, 2 * head_edge_cap)`. - static constexpr int head_edge_cap = load_align_items; + // Split point of `edge_keys`: head edge in `[0, head_edge_cap_items)`, tail edge in `[head_edge_cap_items, 2 * + // head_edge_cap_items)`. + static constexpr int head_edge_cap_items = load_align_items; [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE key_t* slot_keys_unpadded(int slot) const { @@ -483,11 +485,11 @@ struct agent_batched_topk_cluster template [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_split split_chunk(PtrT base, const chunk_desc chunk) const { - const auto la = static_cast<::cuda::std::uintptr_t>(load_align_bytes); + const auto align_bytes = static_cast<::cuda::std::uintptr_t>(load_align_bytes); const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); - const auto aligned_begin = ::cuda::round_up(begin, la); - const auto aligned_end = ::cuda::round_down(end, la); + const auto aligned_begin = ::cuda::round_up(begin, align_bytes); + const auto aligned_end = ::cuda::round_down(end, align_bytes); if (aligned_begin > aligned_end) { // The chunk lies strictly between two load_align boundaries (no aligned point inside): the whole chunk is an @@ -535,11 +537,11 @@ struct agent_batched_topk_cluster // non-empty rank gets the short remainder). // // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches - // ascending contiguous global-index ranges), so it is selected exactly when `deterministic` is set. + // ascending contiguous global-index ranges), so it is selected exactly when `need_determinism` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { - if constexpr (deterministic) + if constexpr (need_determinism) { const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_blocks)); const offset_t first = static_cast(cluster_rank) * chunks_per_cta; @@ -595,15 +597,15 @@ struct agent_batched_topk_cluster { key_t regs[Unroll]; _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) + for (int i = 0; i < Unroll; ++i) { - regs[t] = data[tile_base + t * threads_per_block + tid]; + regs[i] = data[tile_base + i * threads_per_block + tid]; } _CCCL_PRAGMA_UNROLL_FULL() - for (int t = 0; t < Unroll; ++t) + for (int i = 0; i < Unroll; ++i) { - const int local = tile_base + t * threads_per_block + tid; - apply(regs[t], local); + const int local = tile_base + i * threads_per_block + tid; + apply(regs[i], local); } } @@ -664,9 +666,9 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barriers() { _CCCL_PRAGMA_NOUNROLL() - for (int s = static_cast(threadIdx.x); s < PipelineStages; s += threads_per_block) + for (int stage = static_cast(threadIdx.x); stage < PipelineStages; stage += threads_per_block) { - ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[s], 1u); + ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[stage], 1u); } } @@ -676,8 +678,9 @@ struct agent_batched_topk_cluster // per stage, with a matching `wait_stage(stage)`. _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) { - // Only the block leader (see `load_leader`) drives the mbarrier, for a uniform branch and better mbarrier codegen. - if (!load_leader) + // Only the block leader (see `is_load_leader`) drives the mbarrier, for a uniform branch and better mbarrier + // codegen. + if (!is_load_leader) { return; } @@ -746,7 +749,7 @@ struct agent_batched_topk_cluster ::cuda::std::uint32_t load_phase{}; // The single block leader (warp 0's elected lane) that drives the bulk copies. Elected once at construction (the // block constructs convergently) and cached so the pipeline reuses it instead of re-electing per copy. - const bool load_leader = ::cuda::device::__block_elect_one(); + const bool is_load_leader = ::cuda::device::__block_elect_one(); _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -810,10 +813,10 @@ private: // Increment this block's own histogram bucket by one via its 32-bit shared address. The leader (rank 0) also receives // remote folds from the other cluster blocks (Step 2, `hist_fold_remote`), so its add must be cluster-scoped to be // mutually atomic with them; non-leaders only touch their own `hist` before the fold, so cta scope suffices. - _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket, bool leader) + _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket, bool is_leader) { const ::cuda::std::uint32_t addr = base32 + static_cast<::cuda::std::uint32_t>(bucket) * sizeof(offset_t); - if (leader) + if (is_leader) { asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); } @@ -879,11 +882,11 @@ private: return old; } - _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t v) + _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t count) { const ::cuda::std::uint32_t addr = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.num_strictly_selected)); - asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], %1;" : : "r"(addr), "r"(v) : "memory"); + asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], %1;" : : "r"(addr), "r"(count) : "memory"); } // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram @@ -941,18 +944,18 @@ private: // --------------------------------------------------------------------------- // Re-streams the per-rank "overflow" chunks (those that do not fit in the // resident SMEM region) from gmem through a small, fixed, round-robin set of - // `p_eff` (<= `PipelineStages`) streaming slots. The same object is reused for + // `eff_stages` (<= `PipelineStages`) streaming slots. The same object is reused for // every radix pass and the final filter. It ping-pongs the iteration order - // across calls so the `p_eff` turn-around chunks that one pass leaves resident + // across calls so the `eff_stages` turn-around chunks that one pass leaves resident // in the streaming slots are reused by the next pass with no reload; the - // remaining `overflow_chunks - p_eff` chunks are reloaded from gmem on each pass. - // The caller right-sizes the reservation to `p_eff = min(PipelineStages, excess)` + // remaining `overflow_chunks - eff_stages` chunks are reloaded from gmem on each pass. + // The caller right-sizes the reservation to `eff_stages = min(PipelineStages, excess)` // (where `excess = my_chunks - full_slots`), so a streaming rank always has - // `overflow_chunks = excess + p_eff` and reloads exactly `excess` chunks per - // pass - the reserved slots only ever buy reuse of those `p_eff` turn-around + // `overflow_chunks = excess + eff_stages` and reloads exactly `excess` chunks per + // pass - the reserved slots only ever buy reuse of those `eff_stages` turn-around // chunks, never a reload-free pass. The resident region is unaffected: it lives in the // slots `[0, resident_slots)`, the streaming region in `[stream_slot_base, - // stream_slot_base + p_eff)`. + // stream_slot_base + eff_stages)`. struct overflow_streamer { agent_batched_topk_cluster& agent; @@ -965,9 +968,9 @@ private: offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk int stream_slot_base; // SMEM slot index at which the streaming region begins offset_t overflow_chunks; // number of overflow chunks for this rank (M) - int p_eff; // streaming region size = reserved streaming slots (<= PipelineStages, <= M, >= 1) - bool forward = true; - bool primed = false; + int eff_stages; // streaming region size = reserved streaming slots (<= PipelineStages, <= M, >= 1) + bool is_forward = true; + bool is_primed = false; // Stage mbarriers are shared with the resident load (`agent.temp_storage.load_mbar`); stage `stage` targets slot // `stream_slot_base + stage`. `inflight_mask` bit `stage` is set only while a copy is in flight (issued, not yet @@ -999,12 +1002,12 @@ private: , stream_slot_base(stream_slot_base_) , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) { - // `p_eff` is the streaming region size the caller reserved at `[stream_slot_base, stream_slot_base + + // `eff_stages` is the streaming region size the caller reserved at `[stream_slot_base, stream_slot_base + // stream_slots)`; using all of it as pipeline stages keeps the ping-pong reuse maximal. It is `<= M` whenever // there is overflow (`M = excess + stream_slots >= stream_slots`); the `>= 1` floor only matters for the no-op // (`M == 0`) case, where the streamer never touches a slot. - p_eff = (stream_slots_ > 0) ? stream_slots_ : 1; - _CCCL_ASSERT(overflow_chunks == 0 || p_eff <= static_cast(overflow_chunks), + eff_stages = (stream_slots_ > 0) ? stream_slots_ : 1; + _CCCL_ASSERT(overflow_chunks == 0 || eff_stages <= static_cast(overflow_chunks), "streaming depth exceeds the overflow chunk count"); } @@ -1030,30 +1033,32 @@ private: } // Rebuild the shared span for the chunk currently resident in `stage`'s slot without storing per-stage state: the - // slot address is a pure function of `stage` and the length is recomputed from chunk index `o`, so there is no + // slot address is a pure function of `stage` and the length is recomputed from `overflow_idx`, so there is no // spillable `pending[]` array (see `bulk_span`). Returns the aligned bulk only (the always-peeled tail suffix is // excluded). - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span stage_span(int stage, offset_t o) const + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span + stage_span(int stage, offset_t overflow_idx) const { char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; - const auto chunk = agent.get_chunk(chunk_index_of(o), segment_size, head_items); + const auto chunk = agent.get_chunk(chunk_index_of(overflow_idx), segment_size, head_items); const auto split = agent.split_chunk(block_keys_base, chunk); return agent.bulk_span( {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), static_cast(split.bulk)}); } - // Shared driver for one overflow pass. `block_apply(stage, o)` folds the chunk for visit `o` resident in streaming - // slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from gmem - // (fallback). Both entry points delegate to `process_pass_dispatch`, which differs only in whether the callable - // also gets the segment-local index, so they share this loop. `mid()` runs once on a full pass -- after the first - // reload wave (`p_eff` visits) is issued but before it is waited on, overlapping the caller's resident-chunk work - // with those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with no unmatched barrier. - // `should_continue()` is polled once after each consumed chunk (before its refill copy is issued); returning false - // breaks the stream so the final filter bails once the top-k is placed. It must be block-uniform and is evaluated - // uniformly by every lane, so it may contain a barrier (the deterministic filter's does, to read its placement - // counters block-wide). The histogram and non-deterministic filter pass an always-true predicate (folding back to - // the unconditional loop). An early break can leave prefetches in flight, so the pass drains the remaining stages - // before returning (a full pass ends with an empty `inflight_mask`, so the drain is a no-op). + // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident + // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from + // gmem (fallback). Both entry points delegate to `process_pass_dispatch`, which differs only in whether the + // callable also gets the segment-local index, so they share this loop. `mid()` runs once on a full pass -- after + // the first reload wave (`eff_stages` visits) is issued but before it is waited on, overlapping the caller's + // resident-chunk work with those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with + // no unmatched barrier. `should_continue()` is polled once after each consumed chunk (before its refill copy is + // issued); returning false breaks the stream so the final filter bails once the top-k is placed. It must be + // block-uniform and is evaluated uniformly by every lane, so it may contain a barrier (the deterministic filter's + // does, to read its placement counters block-wide). The histogram and non-deterministic filter pass an always-true + // predicate (folding back to the unconditional loop). An early break can leave prefetches in flight, so the pass + // drains the remaining stages before returning (a full pass ends with an empty `inflight_mask`, so the drain is a + // no-op). template _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) @@ -1064,78 +1069,76 @@ private: return; } - const offset_t m = overflow_chunks; - if constexpr (use_block_load_to_shared) { - const offset_t pe = static_cast(p_eff); - // First ever call: prime the streaming slots. Subsequent calls inherit // the previous pass's resident tail, which (because the order - // ping-pongs) is exactly the first `p_eff` chunks of this direction. - if (!primed) + // ping-pongs) is exactly the first `eff_stages` chunks of this direction. + if (!is_primed) { // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else // the phase advances twice and a lagging thread misses the flip and spins forever. __syncthreads(); - for (int i = 0; i < p_eff; ++i) + for (int i = 0; i < eff_stages; ++i) { - const offset_t o = forward ? static_cast(i) : (m - 1 - static_cast(i)); - issue_load(static_cast(o % pe), o); + const offset_t overflow_idx = + is_forward ? static_cast(i) : (overflow_chunks - 1 - static_cast(i)); + issue_load(static_cast(overflow_idx % static_cast(eff_stages)), overflow_idx); } - primed = true; + is_primed = true; } - // Consume overflow visit `i`: wait for its slot, fold its keys via `block_apply`, then prefetch the chunk - // `p_eff` visits ahead into the slot just freed (a barrier guards the slot before the async copy can overwrite - // the data the block was just reading). Returns false once `should_continue()` reports the top-k fully placed - // -- polled before the prefetch so we never launch a copy we would only drain again; the up-to-`p_eff - 1` - // prefetches already in flight (from earlier visits or priming) are drained after the loop. + // Consume the `i`-th visit (its ping-pong-ordered position is `overflow_idx`): wait for its slot, fold its keys + // via `block_apply`, then prefetch the chunk `eff_stages` visits ahead into the slot just freed (a barrier + // guards the slot before the async copy can overwrite the data the block was just reading). Returns false once + // `should_continue()` reports the top-k fully placed -- polled before the prefetch so we never launch a copy we + // would only drain again; the up-to-`eff_stages - 1` prefetches already in flight (from earlier visits or + // priming) are drained after the loop. const auto consume = [&](offset_t i) -> bool { - const offset_t o = forward ? i : (m - 1 - i); - const int stage = static_cast(o % pe); + const offset_t overflow_idx = is_forward ? i : (overflow_chunks - 1 - i); + const int stage = static_cast(overflow_idx % static_cast(eff_stages)); if (inflight_mask & (::cuda::std::uint32_t{1} << stage)) { agent.wait_stage(stage); inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); } - block_apply(stage, o); + block_apply(stage, overflow_idx); if (!should_continue()) { return false; } - const offset_t ni = i + pe; - if (ni < m) + const offset_t next_step = i + static_cast(eff_stages); + if (next_step < overflow_chunks) { - const offset_t no = forward ? ni : (m - 1 - ni); + const offset_t next_overflow_idx = is_forward ? next_step : (overflow_chunks - 1 - next_step); __syncthreads(); - issue_load(stage, no); + issue_load(stage, next_overflow_idx); } return true; }; - // Phase 1: consume the first `p_eff` visits (the chunks reused from the previous pass, already resident in the - // streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. - bool stop = false; - const offset_t split = (::cuda::std::min) (pe, m); + // Phase 1: consume the first `eff_stages` visits (the chunks reused from the previous pass, already resident in + // the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. + bool is_stopped = false; + const offset_t split = (::cuda::std::min) (static_cast(eff_stages), overflow_chunks); for (offset_t i = 0; i < split; ++i) { if (!consume(i)) { - stop = true; + is_stopped = true; break; } } - if (!stop) + if (!is_stopped) { // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. mid(); // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). - for (offset_t i = split; i < m; ++i) + for (offset_t i = split; i < overflow_chunks; ++i) { if (!consume(i)) { @@ -1154,7 +1157,7 @@ private: agent.wait_stage(drain_stage); inflight_mask &= ~(::cuda::std::uint32_t{1} << drain_stage); } - forward = !forward; + is_forward = !is_forward; } else { @@ -1162,18 +1165,18 @@ private: // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no // SMEM reuse), with the walk still snaking for L2 locality. mid(); - for (offset_t i = 0; i < m; ++i) + for (offset_t i = 0; i < overflow_chunks; ++i) { - const offset_t o = forward ? i : (m - 1 - i); - const offset_t chunk_idx = chunk_index_of(o); - const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); + const offset_t overflow_idx = is_forward ? i : (overflow_chunks - 1 - i); + const offset_t chunk_idx = chunk_index_of(overflow_idx); + const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); generic_apply(chunk); if (!should_continue()) { break; } } - forward = !forward; + is_forward = !is_forward; } } @@ -1187,15 +1190,15 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_dispatch(F&& f, Mid&& mid) { run_pass( - [&](int stage, offset_t o) { + [&](int stage, offset_t overflow_idx) { if constexpr (Indexed) { - const offset_t base_off = agent.get_chunk(chunk_index_of(o), segment_size, head_items).offset; - agent.template for_each_chunk_key_indexed(stage_span(stage, o), base_off, f); + const offset_t base_off = agent.get_chunk(chunk_index_of(overflow_idx), segment_size, head_items).offset; + agent.template for_each_chunk_key_indexed(stage_span(stage, overflow_idx), base_off, f); } else { - agent.template for_each_chunk_key(stage_span(stage, o), f); + agent.template for_each_chunk_key(stage_span(stage, overflow_idx), f); } }, [&](const auto& chunk) { @@ -1260,12 +1263,12 @@ private: } // Synchronize the segment's cluster. A single-CTA "cluster" keeps all state block-local, so `__syncthreads()` orders - // it and the cluster-scoped barrier is unnecessary. `single_cta` is computed in `run()` from the collapsed cluster + // it and the cluster-scoped barrier is unnecessary. `is_single_cta` is computed in `run()` from the collapsed cluster // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw cluster size. It is // per-segment uniform across the surviving block(s), so the branch is reached uniformly. - _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(bool single_cta) + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(bool is_single_cta) { - if (single_cta) + if (is_single_cta) { __syncthreads(); } @@ -1276,29 +1279,30 @@ private: } // Combined 64-bit exclusive cross-CTA prefix scan over each working CTA's packed `(front_count << 32) | cand_count`. - // Each CTA pushes its counts into every successor's `prefix_pair` in `scan_descending` order; the leader is last and - // pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` + // Each CTA pushes its counts into every successor's `prefix_pair` in `is_scan_descending` order; the leader is last + // and pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` // and untouched until here, so a single post-push barrier suffices. Idle ranks and the leader pass `packed == 0`. - // Returns this CTA's exclusive prefix packed the same way; `single_cta` returns 0. + // Returns this CTA's exclusive prefix packed the same way; `is_single_cta` returns 0. // // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( - bool single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) + bool is_single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) { - if (single_cta) + if (is_single_cta) { return ::cuda::std::uint64_t{0}; } if (packed != ::cuda::std::uint64_t{0}) { - if constexpr (scan_descending) + if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() - for (unsigned int r = threadIdx.x; r < cluster_rank; r += threads_per_block) // lower ranks follow; leader last + for (unsigned int rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; + // leader last { - add_remote_prefix(r, packed); + add_remote_prefix(rank, packed); } } else @@ -1306,22 +1310,22 @@ private: // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. _CCCL_PRAGMA_NOUNROLL() - for (unsigned int r = cluster_rank + 1u + threadIdx.x; r < eff_cluster_blocks; r += threads_per_block) + for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < eff_cluster_blocks; rank += threads_per_block) { - add_remote_prefix(r, packed); + add_remote_prefix(rank, packed); } } } // TODO(cccl): idle ranks arrive here only to keep this barrier reachable; a sub-cluster mbarrier over the working // ranks would let them exit (see the pass loop). - cluster_or_block_sync(single_cta); + cluster_or_block_sync(is_single_cta); return temp_storage.prefix_pair; } // Deterministic final-filter driver. The `run()`-local tie-break state (counts, prefixes, region extents, and the - // mutable `running`/`tie_active` scan cursor) is hoisted here so the per-region index-ordered sweeps are named member - // functions instead of a nest of `[&]` lambdas. Constructed once per `run()` in the deterministic branch. Kept an - // aggregate (no user constructor) so its members are initialized positionally at the single call site. + // mutable `running`/`is_tie_active` scan cursor) is hoisted here so the per-region index-ordered sweeps are named + // member functions instead of a nest of `[&]` lambdas. Constructed once per `run()` in the deterministic branch. Kept + // an aggregate (no user constructor) so its members are initialized positionally at the single call site. // // Codegen: methods are `_CCCL_FORCEINLINE` and no SMEM key pointer is stored -- the resident window is carried as its // 32-bit shared address (`resident_smem32`) and rebuilt with `__cvta_shared_to_generic` at use, so the reads stay @@ -1352,15 +1356,15 @@ private: offset_t front_seg_base; ::cuda::std::uint32_t resident_smem32; int front_count; - int head_edge_len; - int tail_edge_len; - bool all_below_cta; - bool resident_terminal; - bool head_edge_terminal; - bool tail_edge_terminal; + int head_edge_len_items; + int tail_edge_len_items; + bool is_select_all_cand_cta; + bool is_resident_terminal; + bool is_head_edge_terminal; + bool is_tail_edge_terminal; // Mutable per-invocation tie-break scan cursor. offset_t running; - bool tie_active; + bool is_tie_active; // For each key written to `block_keys_out[pos]`, load the associated input value at the key's segment-local index // `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences the null value iterators) @@ -1382,11 +1386,11 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE bool should_stop() { __syncthreads(); - const bool front_done = agent.temp_storage.front_local_cnt >= static_cast(my_front); - // Straddling/above CTAs finish the back when `tie_active` clears; an `all_below_cta` (which never clears it) - // finishes once all `my_cand_count` of its candidates are placed. - const bool back_done = !tie_active || (agent.temp_storage.back_local_cnt >= my_cand_count); - return front_done && back_done; + const bool is_front_done = agent.temp_storage.front_local_cnt >= static_cast(my_front); + // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never + // clears it) finishes once all `my_cand_count` of its candidates are placed. + const bool is_back_done = !is_tie_active || (agent.temp_storage.back_local_cnt >= my_cand_count); + return is_front_done && is_back_done; } // Fold a flat scan position `pos` into its in-region index: forward, or (`Reversed`) counting down from the @@ -1432,23 +1436,23 @@ private: key_t keys[items]; offset_t flags[items]; detail::topk::candidate_class cls[items]; - bool valid[items]; + bool is_valid[items]; _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - valid[i] = pos < count; + is_valid[i] = pos < count; flags[i] = offset_t{0}; - if (valid[i]) + if (is_valid[i]) { - const int j = fold_pos(pos, count); + const int folded_pos = fold_pos(pos, count); if constexpr (FromSmem) { - keys[i] = smem_src[j]; + keys[i] = smem_src[folded_pos]; } else { - keys[i] = block_keys_in[static_cast(seg_base + static_cast(j))]; + keys[i] = block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; } cls[i] = identify_op(keys[i]); flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; @@ -1461,8 +1465,8 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - const bool to_front = valid[i] && (cls[i] == detail::topk::candidate_class::selected); - if (to_front) + const bool is_front_key = is_valid[i] && (cls[i] == detail::topk::candidate_class::selected); + if (is_front_key) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; const offset_t local = agent.front_local_inc(); @@ -1474,19 +1478,19 @@ private: } // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: - // * all_below_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. + // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the // counter and, on the crossing tile only, overwrite the arrival slots in index order. - if (tie_active) + if (is_tie_active) { - const bool terminal_tile = region_is_terminal && (tile_base + tile >= count); - if (all_below_cta) + const bool is_terminal_tile = region_is_terminal && (tile_base + tile >= count); + if (is_select_all_cand_cta) { _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && flags[i] != offset_t{0}) + if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; emit_back_one(cand_prefix + agent.back_local_inc(), @@ -1495,7 +1499,7 @@ private: } } } - else if (terminal_tile) + else if (is_terminal_tile) { offset_t excl[items]; offset_t tile_total = 0; @@ -1503,7 +1507,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && flags[i] != offset_t{0}) + if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; emit_back_one( @@ -1511,14 +1515,14 @@ private: } } running += tile_total; - tie_active = false; + is_tie_active = false; } else { _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && flags[i] != offset_t{0}) + if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; emit_back_one(cand_prefix + agent.back_local_inc(), @@ -1540,7 +1544,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - if (valid[i] && flags[i] != offset_t{0}) + if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; emit_back_one( @@ -1551,7 +1555,7 @@ private: running = cand_prefix + placed; if (running >= static_cast(num_back)) { - tie_active = false; + is_tie_active = false; } // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the // next tile's `back_local_inc`. The other sub-paths write disjoint slots and need no per-tile barrier. @@ -1561,30 +1565,30 @@ private: } } - // Resident-front region. Direction is the compile-time `reverse_residency` (== `tie_reversed` in deterministic - // mode): ascending walks the low-index window forward, descending walks the high-index window (`resident_base`) in - // reverse, so a single `process_tiles` call per span with the index folded at compile time replaces the old fwd/rev - // pair. + // Resident-front region. Direction is the compile-time `is_residency_reversed` (== `is_tie_reversed` in + // deterministic mode): ascending walks the low-index window forward, descending walks the high-index window + // (`resident_base`) in reverse, so a single `process_tiles` call per span with the index folded at compile time + // replaces the old fwd/rev pair. _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident() { if constexpr (use_block_load_to_shared) { // Whole contiguous resident span staged in SMEM; rebuild the base from its 32-bit address at use. key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - process_tiles( - rfront, front_seg_base, front_count, resident_terminal); + process_tiles( + rfront, front_seg_base, front_count, is_resident_terminal); } else { - const int rc = static_cast(my_resident_chunks); - for (int s = 0; s < rc; ++s) + const int resident_chunk_count = static_cast(my_resident_chunks); + for (int slot = 0; slot < resident_chunk_count; ++slot) { - const int local_slot = reverse_residency ? (rc - 1 - s) : s; + const int local_slot = is_residency_reversed ? (resident_chunk_count - 1 - slot) : slot; const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); const auto chunk = agent.get_chunk(chunk_idx, segment_size_u32, head_items); // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy // per-tile boundary detection handles any boundary), so pass `false`. - process_tiles( + process_tiles( agent.slot_keys_unpadded(local_slot), chunk.offset, chunk.count, false); } } @@ -1594,35 +1598,39 @@ private: // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the // slots are already primed -- the histogram's first streaming pass primed the same persistent streamer, so - // `streamer.primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the - // generic fallback re-reads gmem each pass regardless). + // `streamer.is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime + // (the generic fallback re-reads gmem each pass regardless). _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow() { // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at - // `forward == !tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where the streaming - // direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary (`all_below_cta`, - // arrival-order atomics); or one with no back scan (`!tie_active`). Early stop is subsumed -- it makes every CTA - // `all_below_cta`, so the shorter (possibly mis-parity) pass count never reaches a straddler. - _CCCL_ASSERT(streamer.overflow_chunks == 0 || all_below_cta || !tie_active || streamer.forward == (!tie_reversed), + // `is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where the + // streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary + // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is + // subsumed + // -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never reaches a + // straddler. + _CCCL_ASSERT(streamer.overflow_chunks == 0 || is_select_all_cand_cta || !is_tie_active + || streamer.is_forward == (!is_tie_reversed), "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " "the wrong streaming direction"); streamer.run_pass( - // Block-load: fold the chunk for overflow visit `o`, resident in streaming slot `stage`, straight from SMEM. + // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. // `stage_span` rebuilds the slot pointer from its 32-bit shared address (spill-proof `LDS`) and returns only // the aligned bulk (a peeled tail suffix is handled by `process_tail_edge`). - [&](int stage, offset_t o) { - const auto span = streamer.stage_span(stage, o); - const offset_t base_off = agent.get_chunk(streamer.chunk_index_of(o), segment_size_u32, head_items).offset; + [&](int stage, offset_t overflow_idx) { + const auto span = streamer.stage_span(stage, overflow_idx); + const offset_t base_off = + agent.get_chunk(streamer.chunk_index_of(overflow_idx), segment_size_u32, head_items).offset; // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving // is one barrier on one tile. - process_tiles( + process_tiles( span.data(), base_off, static_cast(span.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). [&](const auto& chunk) { - process_tiles( + process_tiles( nullptr, chunk.offset, chunk.count, false); }, // No interleaved resident work: the deterministic filter folds its resident span separately. @@ -1636,40 +1644,40 @@ private: // Fold one persistent boundary edge (head prefix or peeled tail suffix), both staged in `edge_keys`. A no-op on the // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge(const key_t* keys, offset_t seg_base, int count, bool terminal) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge(const key_t* keys, offset_t seg_base, int count, bool is_terminal) { if constexpr (use_block_load_to_shared) { - process_tiles( - keys, seg_base, count, terminal); + process_tiles( + keys, seg_base, count, is_terminal); } } - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len)`, staged at `edge_keys` (base 0). In - // global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or empty - // prefix is a no-op. + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len_items)`, staged at `edge_keys` (base + // 0). In global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or + // empty prefix is a no-op. _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge() { - process_edge(agent.temp_storage.edge_keys, offset_t{0}, head_edge_len, head_edge_terminal); + process_edge(agent.temp_storage.edge_keys, offset_t{0}, head_edge_len_items, is_head_edge_terminal); } - // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap` (base - // `segment_size - tail_edge_len`). In global-index order it is the trailing region (ascending) / leading region - // (descending); an aligned or non-owned tail is a no-op. + // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` + // (base `segment_size - tail_edge_len_items`). In global-index order it is the trailing region (ascending) / + // leading region (descending); an aligned or non-owned tail is a no-op. _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge() { - process_edge(agent.temp_storage.edge_keys + head_edge_cap, - segment_size_u32 - static_cast(tail_edge_len), - tail_edge_len, - tail_edge_terminal); + process_edge(agent.temp_storage.edge_keys + head_edge_cap_items, + segment_size_u32 - static_cast(tail_edge_len_items), + tail_edge_len_items, + is_tail_edge_terminal); } - // Drive the four regions in global-index order (ascending, or descending under `tie_reversed`), bailing between + // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between // regions once `should_stop` reports the whole top-k placed. The two orders share the resident->overflow middle and // only swap which boundary edge leads: ascending is head, resident, overflow, tail; descending reverses the edges. // Both visit resident before overflow, so `should_stop` can skip re-streaming the overflow once the top-k is - // placed; `reverse_residency` keeps the first-visited (high-index) chunks resident in the descending order so this - // holds. + // placed; `is_residency_reversed` keeps the first-visited (high-index) chunks resident in the descending order so + // this holds. _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter() { const auto step = [&](auto&& region) { @@ -1678,7 +1686,7 @@ private: region(); } }; - if constexpr (tie_reversed) + if constexpr (is_tie_reversed) { process_tail_edge(); } @@ -1693,7 +1701,7 @@ private: process_overflow(); }); step([&] { - if constexpr (tie_reversed) + if constexpr (is_tie_reversed) { process_head_edge(); } @@ -1726,7 +1734,7 @@ private: // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. - const bool single_cta = (cluster_blocks == 1u); + const bool is_single_cta = (cluster_blocks == 1u); // Per-block local copy of `kth_key_bits` so each key check hits the // block's own SMEM rather than DSMEM during the histogram loop. Built up one splitter digit per pass from the @@ -1759,7 +1767,7 @@ private: unsigned int eff_cluster_blocks = cluster_blocks; if constexpr (enable_runtime_single_cta) { - if (!single_cta) + if (!is_single_cta) { eff_cluster_blocks = effective_cluster_blocks_from_chunks( static_cast<::cuda::std::uint64_t>(chunks), min_chunks_per_block, cluster_blocks); @@ -1776,11 +1784,11 @@ private: // (`< eff_cluster_blocks`). The deterministic tie-break makes the leader the *last* CTA in scan order so it never // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. - const unsigned int leader_rank = (deterministic && !tie_reversed) ? (eff_cluster_blocks - 1u) : 0u; + const unsigned int leader_rank = (need_determinism && !is_tie_reversed) ? (eff_cluster_blocks - 1u) : 0u; // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = single_cta ? &temp_storage.state : map_state_to_rank(leader_rank); + state_t* leader_state = is_single_cta ? &temp_storage.state : map_state_to_rank(leader_rank); // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every // cluster barrier is reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) @@ -1815,12 +1823,13 @@ private: } } - // `peel_tail` mirrors `owns_suffix_tail` (always peel, per above). `stream_slots` clamps into `[1, full_slots]`. + // `should_peel_tail` mirrors `owns_suffix_tail` (always peel, per above). `stream_slots` clamps into `[1, + // full_slots]`. offset_t stream_slots = offset_t{0}; - bool peel_tail = false; + bool should_peel_tail = false; if constexpr (use_block_load_to_shared) { - peel_tail = owns_suffix_tail; + should_peel_tail = owns_suffix_tail; if (needs_streaming) { const offset_t excess = my_chunks - full_slots; @@ -1834,17 +1843,17 @@ private: // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. _CCCL_ASSERT(my_resident_chunks <= resident_slots_cap, "Dynamic shared memory block_tile is too small"); - const offset_t overflow_count = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; + const offset_t overflow_chunks = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; // Rank-local base of the resident and overflow (streamed) chunk windows. Default: resident `[0, - // my_resident_chunks)`, overflow the high-index rest. `reverse_residency` swaps them: resident `[overflow_count, - // my_chunks)`, overflow `[0, overflow_count)`. - const offset_t resident_base = reverse_residency ? overflow_count : offset_t{0}; - const offset_t overflow_base = reverse_residency ? offset_t{0} : my_resident_chunks; + // my_resident_chunks)`, overflow the high-index rest. `is_residency_reversed` swaps them: resident + // `[overflow_chunks, my_chunks)`, overflow `[0, overflow_chunks)`. + const offset_t resident_base = is_residency_reversed ? overflow_chunks : offset_t{0}; + const offset_t overflow_base = is_residency_reversed ? offset_t{0} : my_resident_chunks; // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. - [[maybe_unused]] const int head_edge_len = (cluster_rank == 0u) ? static_cast(head_items) : 0; - [[maybe_unused]] const int tail_edge_len = peel_tail ? static_cast(tail_suffix_items) : 0; + [[maybe_unused]] const int head_edge_len_items = (cluster_rank == 0u) ? static_cast(head_items) : 0; + [[maybe_unused]] const int tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. overflow_streamer streamer( @@ -1862,13 +1871,13 @@ private: // Preselect the streamer's initial ping-pong direction. A streaming rank flips direction once per histogram pass, // so the leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = - // (!tie_reversed) ^ (num_passes & 1)` makes that leftover `== !tie_reversed` -- exactly what the deterministic - // filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see + // (!is_tie_reversed) ^ (num_passes & 1)` makes that leftover `== !is_tie_reversed` -- exactly what the + // deterministic filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is - // moot). Non-deterministic filtering is order-independent, so leave its historical `forward` start untouched. - if constexpr (deterministic) + // moot). Non-deterministic filtering is order-independent, so leave its historical `is_forward` start untouched. + if constexpr (need_determinism) { - streamer.forward = (!tie_reversed) ^ ((num_passes & 1) != 0); + streamer.is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); } ::cuda::std::span resident_keys; @@ -1888,14 +1897,14 @@ private: { const int tid = static_cast(threadIdx.x); _CCCL_PRAGMA_NOUNROLL() - for (int local = tid; local < head_edge_len; local += threads_per_block) + for (int local = tid; local < head_edge_len_items; local += threads_per_block) { apply(temp_storage.edge_keys[local]); } _CCCL_PRAGMA_NOUNROLL() - for (int local = tid; local < tail_edge_len; local += threads_per_block) + for (int local = tid; local < tail_edge_len_items; local += threads_per_block) { - apply(temp_storage.edge_keys[head_edge_cap + local]); + apply(temp_storage.edge_keys[head_edge_cap_items + local]); } } }; @@ -1913,10 +1922,10 @@ private: const ::cuda::std::uint32_t hist_smem32 = hist_base32(); // Cluster-scope leader atomics are only needed to stay mutually atomic with the non-leaders' DSMEM folds; a lone // CTA has none, so it uses the cheaper CTA scope. - const bool leader_cluster_scope = (cluster_rank == leader_rank) && !single_cta; - auto add_first_pass = [&](const key_t& key) { + const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; + auto add_first_pass = [&](const key_t& key) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, leader_cluster_scope); + hist_inc(hist_smem32, bucket, is_leader_cluster_scope); }; if constexpr (use_block_load_to_shared) @@ -1925,16 +1934,16 @@ private: { // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor - // (`read_off`) mirrors the write cursor (`next_off`) as a running prefix sum, avoiding a dynamically-indexed - // `pending_spans` array that would anchor surrounding state to local memory. Every chunk begins on a - // `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the global-last - // chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk here. The - // read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor cannot - // demote the reads to `LD`. + // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a + // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk + // begins on a `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the + // global-last chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk + // here. The read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor + // cannot demote the reads to `LD`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `reverse_residency` shifts - // the resident window to the high-index chunks). + // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` + // shifts the resident window to the high-index chunks). const auto resident_local = [&](offset_t slot) -> offset_t { return resident_base + slot; }; @@ -1953,55 +1962,57 @@ private: // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. - int next_off = 0; + int next_off_bytes = 0; // Load every resident chunk's aligned bulk, densely packed in slot order. for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); - issue_bulk_copy(stage, key_slots + next_off, src); - next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + issue_bulk_copy(stage, key_slots + next_off_bytes, src); + next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } - // Read cursor trailing the write cursor: chunk `p`'s bulk was written at `read_off` (packed and consumed in - // the same order), and `bulk_src(p)` recomputes its length, so the read span needs no stored per-stage state. - int read_off = 0; - for (offset_t p = 0; p < my_resident_chunks; ++p) + // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed + // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs + // no stored per-stage state. + int read_off_bytes = 0; + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const int stage = static_cast(p % static_cast(prologue)); + const int stage = static_cast(local_chunk % static_cast(prologue)); wait_stage(stage); - const int read_len = static_cast(::cuda::std::size(bulk_src(p))); + const int read_len_items = static_cast(::cuda::std::size(bulk_src(local_chunk))); for_each_chunk_key( - bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off)), read_len}), + bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off_bytes)), + read_len_items}), add_first_pass); - read_off += read_len * int{sizeof(key_t)}; + read_off_bytes += read_len_items * int{sizeof(key_t)}; - const offset_t next_p = p + static_cast(prologue); - if (next_p < my_resident_chunks) + const offset_t next_local_chunk = local_chunk + static_cast(prologue); + if (next_local_chunk < my_resident_chunks) { - const auto src = bulk_src(next_p); + const auto src = bulk_src(next_local_chunk); // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads // leave the wait above would advance the phase twice, stranding a lagging waiter forever. __syncthreads(); - issue_bulk_copy(stage, key_slots + next_off, src); - next_off += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + issue_bulk_copy(stage, key_slots + next_off_bytes, src); + next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } } // The resident region is one contiguous span of aligned bulks for the later passes; both boundary edges are // folded separately from `edge_keys`. resident_keys = {reinterpret_cast(key_slots), - static_cast<::cuda::std::size_t>(next_off / int{sizeof(key_t)})}; + static_cast<::cuda::std::size_t>(next_off_bytes / int{sizeof(key_t)})}; resident_smem32 = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots)); } } else { - for (offset_t p = 0; p < my_resident_chunks; ++p) + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(resident_base + p); + const offset_t chunk_idx = part.global_index(resident_base + local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) @@ -2026,16 +2037,16 @@ private: // `fold_edges` after a barrier). if constexpr (use_block_load_to_shared) { - if (head_edge_len > 0) + if (head_edge_len_items > 0) { - stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len, add_first_pass); + stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len_items, add_first_pass); } - if (tail_edge_len > 0) + if (tail_edge_len_items > 0) { const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); - stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap, - block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len), - tail_edge_len, + stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, + block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), + tail_edge_len_items, add_first_pass); } } @@ -2067,12 +2078,12 @@ private: // Step 1: block-private histogram via shared-space `red` (see `hist_inc`): leader uses cluster scope to be // mutually atomic with the non-leaders' Step 2 DSMEM folds, non-leaders use the cheaper cta scope. const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool leader_cluster_scope = (cluster_rank == leader_rank) && !single_cta; + const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; auto add_hist = [&](const key_t& key) { if (identify_op(key) == detail::topk::candidate_class::candidate) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, leader_cluster_scope); + hist_inc(hist_smem32, bucket, is_leader_cluster_scope); } }; @@ -2084,17 +2095,17 @@ private: { // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value // spilled across the pass loop (see `resident_smem32`). - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); for_each_chunk_key( - {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); + {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); } else { - for (offset_t p = 0; p < my_resident_chunks; ++p) + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(resident_base + p); + const offset_t chunk_idx = part.global_index(resident_base + local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); } @@ -2129,17 +2140,18 @@ private: const ::cuda::std::uint32_t hist_smem32 = hist_base32(); for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) { - const offset_t v = temp_storage.hist[i]; - if (v != 0) + const offset_t bucket_count = temp_storage.hist[i]; + if (bucket_count != 0) { - hist_fold_remote(hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), v, leader_rank); + hist_fold_remote( + hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), bucket_count, leader_rank); } } } // TODO(cccl): idle ranks arrive here only because the cluster barrier spans the whole launched cluster. An // mbarrier over just the active ranks would let them exit and free their SM slots instead of spinning here. - cluster_or_block_sync(single_cta); + cluster_or_block_sync(is_single_cta); // Step 3: the leader prefix-scans the merged `hist` (raw counts) and // updates the cluster-shared `state`. Subsequent reads (end-of-pass @@ -2175,7 +2187,7 @@ private: reset_hist(); } // TODO(cccl): see the barrier above -- idle ranks arrive here only to keep the cluster barrier reachable. - cluster_or_block_sync(single_cta); + cluster_or_block_sync(is_single_cta); // End-of-pass splitter fold. Every block pulls the leader's just-published `result_pair` once through DSMEM (a // single naturally-aligned `u64`, ordered by the `cluster_or_block_sync()` above) and decodes both halves from @@ -2209,7 +2221,7 @@ private: // so no further radix refinement can change the result. Every block decodes the same flag from the `pass_result` // it just loaded and breaks together; `last_pass`/`kth_key_bits_local` then match what the original // top-of-next-pass break produced. - if (state_t::early_stop_of(pass_result)) + if (state_t::is_early_stop(pass_result)) { break; } @@ -2246,12 +2258,12 @@ private: // trailing digits as smaller and erroneously reject candidates that share the identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); - if constexpr (deterministic) + if constexpr (need_determinism) { // Deterministic tie-break: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on - // the boundary-crossing tile (see `all_below_cta` and the per-tile back logic). Early stop is not special-cased: - // `total_candidates == num_kth` then makes every CTA `all_below_cta`. + // the boundary-crossing tile (see `is_select_all_cand_cta` and the per-tile back logic). Early stop is not + // special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. // // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). @@ -2271,17 +2283,24 @@ private: const offset_t push_front = my_sel; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t pp = combined_prefix_scan(single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(pp >> 32); - const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + const ::cuda::std::uint64_t packed_prefix = + combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A - // CTA is `all_below_cta` when all of its candidates sit at or below the K-boundary (`cand_prefix + my_cand_count + // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary (`cand_prefix + + // my_cand_count // <= num_back`): every one wins, so the back places them with arrival-order SMEM atomics and skips the - // index-ordered scan. While `tie_active`, a non-`all_below_cta` CTA is the single boundary-crossing (straddling) - // CTA cluster-wide. - const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; - const bool all_below_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); + // index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the single boundary-crossing + // (straddling) CTA cluster-wide. + const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; + const bool is_select_all_cand_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); + // Mirror image: a CTA selects none of its candidates when the tie region is empty (`num_back == 0`) or all of its + // candidates sort strictly after the K-boundary (`cand_prefix >= num_back`). Such a CTA seeds `is_tie_active` + // false and skips the back placement entirely. + const bool is_select_no_cand_cta = + (num_back == out_offset_t{0}) || (cand_prefix >= static_cast(num_back)); // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. const out_offset_t my_front = @@ -2297,30 +2316,31 @@ private: // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. - const bool has_head = head_edge_len > 0; + const bool has_head = head_edge_len_items > 0; const bool has_resident = front_count > 0; - const bool has_overflow = overflow_count > offset_t{0}; - const bool has_tail_edge = tail_edge_len > 0; - [[maybe_unused]] const bool resident_terminal = - tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); - [[maybe_unused]] const bool head_edge_terminal = - tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); - [[maybe_unused]] const bool tail_edge_terminal = - tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; + const bool has_overflow = overflow_chunks > offset_t{0}; + const bool has_tail_edge = tail_edge_len_items > 0; + [[maybe_unused]] const bool is_resident_terminal = + is_tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); + [[maybe_unused]] const bool is_head_edge_terminal = + is_tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); + [[maybe_unused]] const bool is_tail_edge_terminal = + is_tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to - // the high-index window under `reverse_residency`). The blocked partition packs the front contiguously, so + // the high-index window under `is_residency_reversed`). The blocked partition packs the front contiguously, so // element `pos` maps to `front_seg_base + pos`. Guarded on `front_count > 0`: with a fully-streamed rank // (`stream_slots == full_slots`) there are no resident chunks and `resident_base` would name an out-of-range - // local chunk under `reverse_residency`. + // local chunk under `is_residency_reversed`. const offset_t front_seg_base = (front_count > 0) ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset : offset_t{0}; // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` - // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `tie_active` at whether this CTA - // straddles the K-boundary. + // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at + // `!is_select_no_cand_cta` -- true unless this is a select-no-candidates CTA (empty back region, or all of this + // CTA's candidates sort past the boundary). det_final_filter filt{ *this, segment_id, @@ -2342,14 +2362,14 @@ private: front_seg_base, resident_smem32, front_count, - head_edge_len, - tail_edge_len, - all_below_cta, - resident_terminal, - head_edge_terminal, - tail_edge_terminal, + head_edge_len_items, + tail_edge_len_items, + is_select_all_cand_cta, + is_resident_terminal, + is_head_edge_terminal, + is_tail_edge_terminal, cand_prefix, - (num_back != out_offset_t{0}) && (cand_prefix < static_cast(num_back))}; + !is_select_no_cand_cta}; filt.run_filter(); } else @@ -2365,9 +2385,10 @@ private: const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t pp = combined_prefix_scan(single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(pp >> 32); - const offset_t cand_prefix = static_cast(pp & 0xffffffffu); + const ::cuda::std::uint64_t packed_prefix = + combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair @@ -2411,17 +2432,17 @@ private: const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) { - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); for_each_chunk_key( - {rk, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); } else { - for (offset_t p = 0; p < my_resident_chunks; ++p) + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(p); + const offset_t chunk_idx = part.global_index(local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(p)); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); } @@ -2463,40 +2484,41 @@ private: // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into // `edge_keys` and folded below), so iterate `split.bulk`, not `chunk.count`. - key_t* const rk = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - int cursor = 0; - for (offset_t s = 0; s < my_resident_chunks; ++s) + key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + int cursor_items = 0; + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(resident_base + s), segment_size_u32, head_items); - const auto split = split_chunk(block_keys_base, chunk); - const offset_t base_off = chunk.offset + split.prefix; - const int cc = split.bulk; - write_run(rk + cursor, base_off, cc); - cursor += cc; + const auto chunk = + get_chunk(part.global_index(resident_base + local_chunk), segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); + const offset_t base_off = chunk.offset + split.prefix; + const int bulk_count_items = split.bulk; + write_run(resident_ptr + cursor_items, base_off, bulk_count_items); + cursor_items += bulk_count_items; } } else { - for (offset_t p = 0; p < my_resident_chunks; ++p) + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(p), segment_size_u32, head_items); + const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_u32, head_items); const offset_t base_off = chunk.offset; - write_run(slot_keys_unpadded(static_cast(p)), base_off, chunk.count); + write_run(slot_keys_unpadded(static_cast(local_chunk)), base_off, chunk.count); } } // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the - // peeled tail suffix at `segment_size - tail_edge_len`. Value fetched per key. + // peeled tail suffix at `segment_size - tail_edge_len_items`. Value fetched per key. if constexpr (use_block_load_to_shared) { - if (head_edge_len > 0) + if (head_edge_len_items > 0) { - write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len); + write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); } - if (tail_edge_len > 0) + if (tail_edge_len_items > 0) { - write_run(temp_storage.edge_keys + head_edge_cap, - segment_size_u32 - static_cast(tail_edge_len), - tail_edge_len); + write_run(temp_storage.edge_keys + head_edge_cap_items, + segment_size_u32 - static_cast(tail_edge_len_items), + tail_edge_len_items); } } }; @@ -2520,14 +2542,14 @@ private: unsigned int cluster_rank, unsigned int cluster_blocks) { - constexpr int copy_items = copy_items_per_thread_clamped; - const offset_t n = static_cast(segment_size); - const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; - const offset_t cluster_tids = cluster_blocks * static_cast(threads_per_block); - const offset_t step = cluster_tids * static_cast(copy_items); - const offset_t full_tiles = ::cuda::round_down(n, step); - auto keys_in_it = d_key_segments_it[segment_id]; - auto keys_out_it = d_key_segments_out_it[segment_id]; + constexpr int copy_items = copy_items_per_thread_clamped; + const offset_t num_items = static_cast(segment_size); + const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; + const offset_t cluster_threads = cluster_blocks * static_cast(threads_per_block); + const offset_t step = cluster_threads * static_cast(copy_items); + const offset_t full_tiles = ::cuda::round_down(num_items, step); + auto keys_in_it = d_key_segments_it[segment_id]; + auto keys_out_it = d_key_segments_out_it[segment_id]; // Per-segment value iterators (pairs only), hoisted once like the key iterators and reused by both loops. For // keys-only the value iterators-of-iterators are null, so the discarded `if constexpr` branch keeps the indexing // out of those builds (the variables stay unused there). @@ -2560,7 +2582,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) { - idx[j] = base + static_cast(j) * cluster_tids + cluster_tid; + idx[j] = base + static_cast(j) * cluster_threads + cluster_tid; } _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < copy_items; ++j) @@ -2592,13 +2614,13 @@ private: // Sub-tile remainder _CCCL_PRAGMA_NOUNROLL() - for (offset_t idx = full_tiles + cluster_tid; idx < n; idx += cluster_tids) + for (offset_t idx = full_tiles + cluster_tid; idx < num_items; idx += cluster_threads) { - const auto si = static_cast(idx); - keys_out_it[si] = keys_in_it[si]; + const auto seg_idx = static_cast(idx); + keys_out_it[seg_idx] = keys_in_it[seg_idx]; if constexpr (!is_keys_only) { - vals_out_it[si] = vals_in_it[si]; + vals_out_it[seg_idx] = vals_in_it[seg_idx]; } } } @@ -2665,7 +2687,7 @@ private: unsigned int eff_cluster_rank = cluster_rank; if constexpr (enable_runtime_single_cta) { - const bool fits_single_cta = single_cta_eligible( + const bool fits_single_cta = is_single_cta_eligible( static_cast<::cuda::std::uint64_t>(segment_size), static_cast<::cuda::std::uint64_t>(block_tile_capacity), single_block_max_seg_size); @@ -2679,7 +2701,7 @@ private: eff_cluster_rank = 0u; } } - const bool single_cta = (eff_cluster_blocks == 1u); + const bool is_single_cta = (eff_cluster_blocks == 1u); // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster @@ -2701,16 +2723,16 @@ private: temp_storage.my_candidates = 0; } reset_hist(); - cluster_or_block_sync(single_cta); + cluster_or_block_sync(is_single_cta); - [[maybe_unused]] const bool ok = detail::params::dispatch_discrete( + [[maybe_unused]] const bool is_ok = detail::params::dispatch_discrete( select_directions, segment_id, [this, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; this->template run(segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k); }); - _CCCL_ASSERT(ok, "Unsupported select direction for cluster top-k"); + _CCCL_ASSERT(is_ok, "Unsupported select direction for cluster top-k"); } }; } // namespace detail::batched_topk_cluster diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 004724dee4b..2ca2217ea7f 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -712,7 +712,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( int cluster_blocks = 0; int dynamic_smem_sel = 0; - if (batched_topk_cluster::single_cta_eligible( + if (batched_topk_cluster::is_single_cta_eligible( seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity), SingleBlockMaxSegSize)) { // Single-CTA fast path: the segment fits resident in one CTA and is small enough that the agent's diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 84c26352449..9898f7ea04f 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -624,7 +624,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments through a non- // block-load path with an unaligned head edge, and the 1 Mi-element segments stream, so the value payloads exercise // the boundary-edge value writes that the small and non-contiguous pair tests above do not. An unaligned tail suffix // is always peeled into `edge_keys` (like the head prefix), so every launch config that owns such a tail exercises the -// head edge plus the persistent `tail_edge_len`/`process_tail_edge` value writes. +// head edge plus the persistent `tail_edge_len_items`/`process_tail_edge` value writes. C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned segments", "[pairs][segmented][topk][device][cluster]", select_direction_list) From d38bf56837012e666e0fb76b0593547e445e4a2e Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 01:38:50 +0200 Subject: [PATCH 083/126] Trim comments a little --- cub/cub/agent/agent_batched_topk_cluster.cuh | 88 +++++++------------ .../dispatch/kernels/kernel_batched_topk.cuh | 3 - ...catch2_test_device_segmented_topk_pairs.cu | 8 +- 3 files changed, 36 insertions(+), 63 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 4512dbeceba..bc8d2e40d1b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -100,9 +100,7 @@ struct alignas(16) cluster_topk_state // change it). ::cuda::std::uint64_t result_pair; - // Decoders that operate on an already-loaded `result_pair`. The hot pass loop reads the leader's word once through - // DSMEM and decodes both halves from that single value, so the splitter fold and the early-stop check never issue a - // second remote load. `set_result` packs the same layout, keeping the bit layout defined in one place. + // Decode/encode an already-loaded `result_pair`, keeping the bit layout (see above) defined in one place. [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static ::cuda::std::uint32_t kth_bucket_of(::cuda::std::uint64_t packed) { @@ -898,10 +896,8 @@ private: // merge before invoking this. _CCCL_DEVICE _CCCL_FORCEINLINE void leader_identify_kth_bucket() { - // Capture `state.k` before the scan: this is the only legal window where - // every thread is guaranteed to read the previous pass's value. The - // owning thread overwrites `state.k` in the if-block below, so any read - // after that point would race with that write. + // Capture `state.k` before the scan: the owning thread overwrites it in the loop below, so any later read would + // race with that write. const out_offset_t target_k = temp_storage.state.k; offset_t hist_vals[buckets_per_thread]; @@ -928,11 +924,8 @@ private: const offset_t new_len = hist_vals[j]; temp_storage.state.len = new_len; temp_storage.state.k = new_k; - // Publish the splitter bucket and the early-stop flag (set when the bucket holds exactly the remaining `k`, so - // every candidate is part of the top-k and subsequent radix passes only redistribute the same items across - // finer buckets without changing the result); see `result_pair`. Every block pulls this packed word from the - // leader at the end of the pass and folds the bucket into its own `kth_key_bits_local`, so the full splitter - // key is never broadcast. + // Publish the splitter bucket and the early-stop flag (set when the bucket holds exactly the remaining `k`); + // see `result_pair`. temp_storage.state.set_result( static_cast<::cuda::std::uint32_t>(bucket), static_cast(new_len) == new_k); } @@ -942,20 +935,14 @@ private: // --------------------------------------------------------------------------- // Overflow streamer // --------------------------------------------------------------------------- - // Re-streams the per-rank "overflow" chunks (those that do not fit in the - // resident SMEM region) from gmem through a small, fixed, round-robin set of - // `eff_stages` (<= `PipelineStages`) streaming slots. The same object is reused for - // every radix pass and the final filter. It ping-pongs the iteration order - // across calls so the `eff_stages` turn-around chunks that one pass leaves resident - // in the streaming slots are reused by the next pass with no reload; the - // remaining `overflow_chunks - eff_stages` chunks are reloaded from gmem on each pass. - // The caller right-sizes the reservation to `eff_stages = min(PipelineStages, excess)` - // (where `excess = my_chunks - full_slots`), so a streaming rank always has - // `overflow_chunks = excess + eff_stages` and reloads exactly `excess` chunks per - // pass - the reserved slots only ever buy reuse of those `eff_stages` turn-around - // chunks, never a reload-free pass. The resident region is unaffected: it lives in the - // slots `[0, resident_slots)`, the streaming region in `[stream_slot_base, - // stream_slot_base + eff_stages)`. + // Re-streams a rank's "overflow" chunks (those that do not fit its resident SMEM region) from gmem through a fixed, + // round-robin set of `eff_stages` (<= `PipelineStages`) streaming slots, reused across every radix pass and the final + // filter. It ping-pongs the iteration order across calls so the `eff_stages` turn-around chunks that one pass leaves + // resident in the streaming slots are reused by the next with no reload; the remaining `overflow_chunks - eff_stages` + // are reloaded from gmem each pass. The caller sizes the reservation to `eff_stages = min(PipelineStages, excess)` + // (`excess = my_chunks - full_slots`), so a streaming rank reloads exactly `excess` chunks per pass -- the reserved + // slots only ever buy reuse of the turn-around chunks, never a reload-free pass. The resident region occupies slots + // `[0, resident_slots)`, the streaming region `[stream_slot_base, stream_slot_base + eff_stages)`. struct overflow_streamer { agent_batched_topk_cluster& agent; @@ -1048,17 +1035,15 @@ private: // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from - // gmem (fallback). Both entry points delegate to `process_pass_dispatch`, which differs only in whether the - // callable also gets the segment-local index, so they share this loop. `mid()` runs once on a full pass -- after - // the first reload wave (`eff_stages` visits) is issued but before it is waited on, overlapping the caller's - // resident-chunk work with those in-flight copies (skipped if phase 1 stops early); it must be block-uniform with - // no unmatched barrier. `should_continue()` is polled once after each consumed chunk (before its refill copy is - // issued); returning false breaks the stream so the final filter bails once the top-k is placed. It must be - // block-uniform and is evaluated uniformly by every lane, so it may contain a barrier (the deterministic filter's - // does, to read its placement counters block-wide). The histogram and non-deterministic filter pass an always-true - // predicate (folding back to the unconditional loop). An early break can leave prefetches in flight, so the pass - // drains the remaining stages before returning (a full pass ends with an empty `inflight_mask`, so the drain is a - // no-op). + // gmem (fallback). `mid()` runs once on a full pass -- after the first reload wave (`eff_stages` visits) is issued + // but before it is waited on, overlapping the caller's resident-chunk work with those in-flight copies (skipped if + // phase 1 stops early); it must be block-uniform with no unmatched barrier. `should_continue()` is polled once + // after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final + // filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the + // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (the deterministic + // filter's does, to read its placement counters block-wide); the histogram and non-deterministic filter pass an + // always-true predicate. An early break can leave prefetches in flight, so the pass drains the + // remaining stages before returning (a full pass ends with an empty `inflight_mask`, so the drain is a no-op). template _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) @@ -1606,9 +1591,8 @@ private: // `is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where the // streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is - // subsumed - // -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never reaches a - // straddler. + // subsumed -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never + // reaches a straddler. _CCCL_ASSERT(streamer.overflow_chunks == 0 || is_select_all_cand_cta || !is_tie_active || streamer.is_forward == (!is_tie_reversed), "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " @@ -2032,8 +2016,7 @@ private: // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. - // The - // `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via + // The `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via // `fold_edges` after a barrier). if constexpr (use_block_load_to_shared) { @@ -2237,13 +2220,11 @@ private: auto block_keys_out = d_key_segments_out_it[segment_id]; const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes - // For each key written to `block_keys_out[pos]`, the associated input value at the key's segment-local index - // `seg_idx` is loaded naively from gmem and written to `block_vals_out[pos]`. `seg_idx` is recomputed per region in - // the sweeps below. The per-segment value iterators are derived *inside* the `is_keys_only` guard: in keys-only - // builds the value iterators-of-iterators are `cub::NullType**` (null), so indexing them with `segment_id` here - // would dereference a null pointer; `segment_id` is loop-invariant, so the compiler hoists these out of the writes. - // The deterministic path folds this logic into `det_final_filter::write_value`; this lambda now serves only the - // non-deterministic branch below, so it is unused (but still instantiated) in deterministic builds. + // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index + // `seg_idx`. Deriving the per-segment value iterators *inside* the `is_keys_only` guard avoids indexing the null + // `cub::NullType**` value-iterators-of-iterators in keys-only builds; `segment_id` is loop-invariant, so they hoist + // out of the writes. Serves only the non-deterministic branch below (the deterministic path uses + // `det_final_filter::write_value`), so it is unused but still instantiated in deterministic builds. [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { if constexpr (!is_keys_only) { @@ -2289,11 +2270,10 @@ private: const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A - // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary (`cand_prefix + - // my_cand_count - // <= num_back`): every one wins, so the back places them with arrival-order SMEM atomics and skips the - // index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the single boundary-crossing - // (straddling) CTA cluster-wide. + // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary + // (`cand_prefix + my_cand_count <= num_back`): every one wins, so the back places them with arrival-order SMEM + // atomics and skips the index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the + // single boundary-crossing (straddling) CTA cluster-wide. const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; const bool is_select_all_cand_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); // Mirror image: a CTA selects none of its candidates when the tie region is empty (`num_back == 0`) or all of its diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index b557c1efd9a..35fc9cfa253 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -191,10 +191,8 @@ __launch_bounds__(int( static_assert(sizeof(typename agent_t::TempStorage) <= max_smem_per_block, "Static shared memory per block must not exceed 48KB limit."); - // Temporary storage allocation __shared__ typename agent_t::TempStorage temp_storage; - // Instantiate agent agent_t agent( temp_storage, d_key_segments_it, @@ -209,7 +207,6 @@ __launch_bounds__(int( d_large_segments_ids, d_large_segments_tile_offsets); - // Process segments agent.Process(); } diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 9898f7ea04f..926fd45adbb 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -338,14 +338,10 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small fixed-size segments" cuda::args::immediate{k, cuda::args::bounds()}, cuda::args::immediate{num_segments}); - // Verification: - // - We verify correct top-k selection through the keys - // - We verify that values were permuted along correctly by making sure values remain associated with their keys and - // making sure we do not duplicate values Verify values remain associated with their corresponding keys + // Verify values stayed associated with their keys. REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); - // Verify values don't appear more than once in the returned results - // This catches the case where we just returned a valid value multiple times + // Verify no value appears more than once (catches returning a valid value multiple times). REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); // Verify keys are returned correctly From 88cee8d68bd4344bad6dbeafd521b6a229dd50e8 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 04:20:52 +0200 Subject: [PATCH 084/126] Split up long functions Mostly SASS neutral up to a SASS reduction from process_tiles dedup. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 1060 ++++++++++-------- 1 file changed, 600 insertions(+), 460 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index bc8d2e40d1b..d9a26905a1d 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1470,8 +1470,10 @@ private: if (is_tie_active) { const bool is_terminal_tile = region_is_terminal && (tile_base + tile >= count); - if (is_select_all_cand_cta) - { + + // Arrival-order placement: each candidate grabs the next block-local back slot via an atomic. Used where the + // winning slot set is provisional (select-all, or the lazy path before the K-boundary is known). + const auto emit_arrival = [&] { _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { @@ -1483,9 +1485,10 @@ private: seg_base + static_cast(fold_pos(pos, count))); } } - } - else if (is_terminal_tile) - { + }; + // Index-ordered placement into slots based at `base`: a block scan assigns each candidate a deterministic + // slot by its in-tile rank. Returns the tile's candidate total. Used on the boundary-crossing tile. + const auto emit_indexed = [&](offset_t base) -> offset_t { offset_t excl[items]; offset_t tile_total = 0; block_scan_t(agent.temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); @@ -1495,47 +1498,34 @@ private: if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - emit_back_one( - running + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); + emit_back_one(base + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); } } - running += tile_total; + return tile_total; + }; + + if (is_select_all_cand_cta) + { + emit_arrival(); + } + else if (is_terminal_tile) + { + // Straddling CTA's last tile necessarily holds the boundary: scan it directly in index order. + running += emit_indexed(running); is_tie_active = false; } else { - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (is_valid[i] && flags[i] != offset_t{0}) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - emit_back_one(cand_prefix + agent.back_local_inc(), - keys[i], - seg_base + static_cast(fold_pos(pos, count))); - } - } + emit_arrival(); // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make // the counter read race-free and block-uniform. __syncthreads(); const offset_t placed = agent.temp_storage.back_local_cnt; if ((cand_prefix + placed) > static_cast(num_back)) { - // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered - // winners (identical slot set, different candidate->slot mapping). - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(agent.temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - if (is_valid[i] && flags[i] != offset_t{0}) - { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; - emit_back_one( - running + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); - } - } + // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered winners + // (identical slot set, different candidate->slot mapping). + emit_indexed(running); } running = cand_prefix + placed; if (running >= static_cast(num_back)) @@ -1697,6 +1687,530 @@ private: } }; + // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), + // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final + // filter; the deterministic filter folds the edges as separate index-ordered regions instead. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + fold_boundary_edges(int head_edge_len_items, int tail_edge_len_items, Apply&& apply) + { + if constexpr (use_block_load_to_shared) + { + const int tid = static_cast(threadIdx.x); + _CCCL_PRAGMA_NOUNROLL() + for (int local = tid; local < head_edge_len_items; local += threads_per_block) + { + apply(temp_storage.edge_keys[local]); + } + _CCCL_PRAGMA_NOUNROLL() + for (int local = tid; local < tail_edge_len_items; local += threads_per_block) + { + apply(temp_storage.edge_keys[head_edge_cap_items + local]); + } + } + } + + // Non-deterministic final filter: strictly-selected keys fill the front `[0, num_selected)`, the first `num_kth` + // candidates (arrival order) fill the back. The combined cross-CTA scan gives this block disjoint front/back bases + // (`sel_prefix`/`cand_prefix`); placement then uses block-local SMEM atomics since output order is not preserved. A + // perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( + num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int eff_cluster_blocks, + unsigned int leader_rank, + bool is_single_cta, + bool is_idle_rank, + out_offset_t k, + out_offset_t num_kth, + IdentifyOp identify_op, + KeyOutIt block_keys_out, + const key_t* block_keys_base, + const chunk_partition& part, + offset_t resident_base, + offset_t my_resident_chunks, + offset_t segment_size_u32, + offset_t head_items, + ::cuda::std::uint32_t resident_smem32, + ::cuda::std::span resident_keys, + int head_edge_len_items, + int tail_edge_len_items, + overflow_streamer& streamer) + { + // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index + // `seg_idx`. Deriving the per-segment value iterators *inside* the `is_keys_only` guard avoids indexing the null + // `cub::NullType**` value-iterators-of-iterators in keys-only builds; `segment_id` is loop-invariant, so they hoist + // out of the writes. + [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { + if constexpr (!is_keys_only) + { + auto block_vals_in = d_value_segments_it[segment_id]; + auto block_vals_out = d_value_segments_out_it[segment_id]; + block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; + } + }; + + __syncthreads(); + const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; + const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; + const ::cuda::std::uint64_t packed = + (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); + const ::cuda::std::uint64_t packed_prefix = + combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + + // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local + // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair + // mode each written key additionally pulls its value from gmem at `seg_idx`. Keys-only elides that write via the + // `if constexpr` below and drives the sink from index-free traversals, passing a dummy `seg_idx` that goes + // unread. + const auto sink = [&](const key_t& key, [[maybe_unused]] offset_t seg_idx) { + const auto res = identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); + block_keys_out[pos] = key; + if constexpr (!is_keys_only) + { + write_value(pos, seg_idx); + } + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); + if (back_pos < num_kth) + { + const out_offset_t pos = k - 1 - back_pos; + block_keys_out[pos] = key; + if constexpr (!is_keys_only) + { + write_value(pos, seg_idx); + } + } + } + }; + + if constexpr (is_keys_only) + { + // Keys-only: no value payload, so drive the sink through the index-free traversal with a dummy `seg_idx`. + const auto write_selected = [&](const key_t& key) { + sink(key, offset_t{0}); + }; + // Fold the resident keys as the streamer's `mid` work so they overlap the first overflow reloads. Writes are + // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) + { + key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + for_each_chunk_key( + {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + } + else + { + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + { + const offset_t chunk_idx = part.global_index(local_chunk); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); + for_each_chunk_key( + {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + } + } + // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). + fold_boundary_edges(head_edge_len_items, tail_edge_len_items, write_selected); + }; + streamer.process_pass(write_selected, fold_resident); + } + else + { + // Pair (key + value) path: the shared `sink` above additionally stores each written key's value (fetched from + // gmem at its `seg_idx`). Unlike keys-only, the traversal must recover each key's segment-local index, so it + // folds resident/edge keys through `write_run` (below) and overflow keys through `process_pass_indexed`. + + // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local + // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary + // edges); overflow chunks fold through the streamer's own indexed callback below. Materialize the key into a + // register first: `sink` binds it by `const&` and reads it several times, so passing `smem[local]` directly + // would re-issue a narrow `LDS` per use instead of reusing the loaded value. + auto write_run = [&](const key_t* smem, offset_t base_off, int count) { + const int iterations = ::cuda::ceil_div(count, threads_per_block); + _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) + for (int j = 0; j < iterations; ++j) + { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < count) + { + const key_t key = smem[local]; + sink(key, base_off + static_cast(local)); + } + } + }; + + // Fold the resident keys (and their values) as the streamer's `mid` work, exactly as in the keys-only path. + const auto fold_resident = [&] { + if constexpr (use_block_load_to_shared) + { + // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers + // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into + // `edge_keys` and folded below), so iterate `split.bulk`, not `chunk.count`. + key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + int cursor_items = 0; + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + { + const auto chunk = get_chunk(part.global_index(resident_base + local_chunk), segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); + const offset_t base_off = chunk.offset + split.prefix; + const int bulk_count_items = split.bulk; + write_run(resident_ptr + cursor_items, base_off, bulk_count_items); + cursor_items += bulk_count_items; + } + } + else + { + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + { + const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_u32, head_items); + const offset_t base_off = chunk.offset; + write_run(slot_keys_unpadded(static_cast(local_chunk)), base_off, chunk.count); + } + } + // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the + // peeled tail suffix at `segment_size - tail_edge_len_items`. Value fetched per key. + if constexpr (use_block_load_to_shared) + { + if (head_edge_len_items > 0) + { + write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); + } + if (tail_edge_len_items > 0) + { + write_run(temp_storage.edge_keys + head_edge_cap_items, + segment_size_u32 - static_cast(tail_edge_len_items), + tail_edge_len_items); + } + } + }; + + // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected + // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. + streamer.process_pass_indexed(sink, fold_resident); + } + } + + // Deterministic final filter: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); + // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on the + // single boundary-crossing (straddling) CTA. A combined cross-CTA scan gives this block its disjoint front/back + // bases and lets it detect whether all/none/some of its candidates win. The heavy lifting lives in + // `det_final_filter`; this member computes that struct's inputs and runs it. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( + num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int eff_cluster_blocks, + unsigned int leader_rank, + bool is_single_cta, + bool is_idle_rank, + out_offset_t k, + out_offset_t num_kth, + IdentifyOp identify_op, + KeyOutIt block_keys_out, + key_it_t block_keys_in, + const chunk_partition& part, + offset_t resident_base, + offset_t my_resident_chunks, + offset_t overflow_chunks, + offset_t segment_size_u32, + offset_t head_items, + state_t* leader_state, + ::cuda::std::uint32_t resident_smem32, + ::cuda::std::span resident_keys, + int head_edge_len_items, + int tail_edge_len_items, + overflow_streamer& streamer) + { + // Early stop is not special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. + // + // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster + // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). + const offset_t total_candidates = leader_state->len; + const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys + const out_offset_t num_selected = k - num_back; // front region + + // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final + // cluster barrier) block-wide before they feed the scan and `front_count`. + __syncthreads(); + const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; + const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; + // Front count pushed by this block: its strictly-selected count. Candidates always route through the back, so + // nothing folds into the front here. The leader and idle ranks push 0 -- the leader because its merged histogram + // cannot self-count (it derives its own front from the total below), idle ranks because they own nothing. + const offset_t push_front = my_sel; + const ::cuda::std::uint64_t packed = + (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); + const ::cuda::std::uint64_t packed_prefix = + combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so + // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A + // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary + // (`cand_prefix + my_cand_count <= num_back`): every one wins, so the back places them with arrival-order SMEM + // atomics and skips the index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the + // single boundary-crossing (straddling) CTA cluster-wide. + const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; + const bool is_select_all_cand_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); + // Mirror image: a CTA selects none of its candidates when the tie region is empty (`num_back == 0`) or all of its + // candidates sort strictly after the K-boundary (`cand_prefix >= num_back`). Such a CTA seeds `is_tie_active` + // false and skips the back placement entirely. + const bool is_select_no_cand_cta = + (num_back == out_offset_t{0}) || (cand_prefix >= static_cast(num_back)); + // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so + // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. + const out_offset_t my_front = + (cluster_rank == leader_rank) + ? static_cast(num_selected - static_cast(sel_prefix)) + : static_cast(push_front); + + // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the + // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never + // part of this span. + const int front_count = span_size(resident_keys); + + // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop + // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later + // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. + const bool has_head = head_edge_len_items > 0; + const bool has_resident = front_count > 0; + const bool has_overflow = overflow_chunks > offset_t{0}; + const bool has_tail_edge = tail_edge_len_items > 0; + [[maybe_unused]] const bool is_resident_terminal = + is_tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); + [[maybe_unused]] const bool is_head_edge_terminal = + is_tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); + [[maybe_unused]] const bool is_tail_edge_terminal = + is_tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; + + // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to + // the high-index window under `is_residency_reversed`). The blocked partition packs the front contiguously, so + // element `pos` maps to `front_seg_base + pos`. Guarded on `front_count > 0`: with a fully-streamed rank + // (`stream_slots == full_slots`) there are no resident chunks and `resident_base` would name an out-of-range + // local chunk under `is_residency_reversed`. + const offset_t front_seg_base = + (front_count > 0) + ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset + : offset_t{0}; + + // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` + // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at + // `!is_select_no_cand_cta` -- true unless this is a select-no-candidates CTA (empty back region, or all of this + // CTA's candidates sort past the boundary). + det_final_filter filt{ + *this, + segment_id, + identify_op, + block_keys_out, + block_keys_in, + streamer, + part, + k, + num_back, + my_front, + sel_prefix, + cand_prefix, + my_cand_count, + resident_base, + my_resident_chunks, + segment_size_u32, + head_items, + front_seg_base, + resident_smem32, + front_count, + head_edge_len_items, + tail_edge_len_items, + is_select_all_cand_cta, + is_resident_terminal, + is_head_edge_terminal, + is_tail_edge_terminal, + cand_prefix, + !is_select_no_cand_cta}; + filt.run_filter(); + } + + // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into + // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 + // needs no candidate filtering). Publishes the resident span as `resident_keys` and its 32-bit shared base address + // as `resident_smem32` for the later passes and the final filter; ends on the barrier that makes `edge_keys` and the + // histogram visible cluster-wide. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( + unsigned int cluster_rank, + unsigned int leader_rank, + bool is_single_cta, + const chunk_partition& part, + offset_t my_resident_chunks, + offset_t resident_base, + offset_t chunks, + offset_t segment_size_u32, + offset_t head_items, + key_it_t block_keys_in, + const key_t* block_keys_base, + int head_edge_len_items, + int tail_edge_len_items, + overflow_streamer& streamer, + ::cuda::std::span& resident_keys, + ::cuda::std::uint32_t& resident_smem32) + { + using extract_bin_op_t = detail::topk::extract_bin_op_t; + constexpr int total_bits = int{sizeof(key_t)} * 8; + + extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); + // Cluster-scope leader atomics are only needed to stay mutually atomic with the non-leaders' DSMEM folds; a lone + // CTA has none, so it uses the cheaper CTA scope. + const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; + auto add_first_pass = [&](const key_t& key) { + const int bucket = extract_op(key); + hist_inc(hist_smem32, bucket, is_leader_cluster_scope); + }; + + if constexpr (use_block_load_to_shared) + { + if (my_resident_chunks > 0) + { + // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). + // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor + // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a + // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk + // begins on a `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the + // global-last chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk + // here. The read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor + // cannot demote the reads to `LD`. + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); + + // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` + // shifts the resident window to the high-index chunks). + const auto resident_local = [&](offset_t slot) -> offset_t { + return resident_base + slot; + }; + // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has + // none. + const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { + const offset_t chunk_idx = part.global_index(resident_local(slot)); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto split = split_chunk(block_keys_base, chunk); + if (split.bulk == 0) + { + return {}; + } + return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; + }; + + // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives + // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. + int next_off_bytes = 0; + + // Load every resident chunk's aligned bulk, densely packed in slot order. + for (int stage = 0; stage < prologue; ++stage) + { + const auto src = bulk_src(static_cast(stage)); + issue_bulk_copy(stage, key_slots + next_off_bytes, src); + next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + } + + // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed + // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs + // no stored per-stage state. + int read_off_bytes = 0; + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + { + const int stage = static_cast(local_chunk % static_cast(prologue)); + wait_stage(stage); + const int read_len_items = static_cast(::cuda::std::size(bulk_src(local_chunk))); + for_each_chunk_key( + bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off_bytes)), + read_len_items}), + add_first_pass); + read_off_bytes += read_len_items * int{sizeof(key_t)}; + + const offset_t next_local_chunk = local_chunk + static_cast(prologue); + if (next_local_chunk < my_resident_chunks) + { + const auto src = bulk_src(next_local_chunk); + // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads + // leave the wait above would advance the phase twice, stranding a lagging waiter forever. + __syncthreads(); + issue_bulk_copy(stage, key_slots + next_off_bytes, src); + next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + } + } + + // The resident region is one contiguous span of aligned bulks for the later passes; both boundary edges are + // folded separately from `edge_keys`. + resident_keys = {reinterpret_cast(key_slots), + static_cast<::cuda::std::size_t>(next_off_bytes / int{sizeof(key_t)})}; + resident_smem32 = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots)); + } + } + else + { + for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + { + const offset_t chunk_idx = part.global_index(resident_base + local_chunk); + const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) + for (int j = 0; j < iterations; ++j) + { + const int local = j * threads_per_block + static_cast(threadIdx.x); + if (local < chunk.count) + { + const key_t key = + block_keys_in[static_cast(chunk.offset + static_cast(local))]; + chunk_keys[local] = key; + add_first_pass(key); + } + } + } + } + + // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see + // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix + // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. + // The `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via + // `fold_edges` after a barrier). + if constexpr (use_block_load_to_shared) + { + if (head_edge_len_items > 0) + { + stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len_items, add_first_pass); + } + if (tail_edge_len_items > 0) + { + const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); + stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, + block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), + tail_edge_len_items, + add_first_pass); + } + } + + // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the streamer's initial + // direction (preselected above; the histogram is order-independent, so the direction only sets up the leftover + // parity for the final filter). The streamer reuses the resident load's stage mbarriers (all front-loaded at + // `run` entry); `wait_stage` provides the producer/consumer sync. + streamer.process_pass(add_first_pass); + + const int resident_count = span_size(resident_keys); + _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, + "Dynamic shared memory block_tile is too small"); + __syncthreads(); + } + template _CCCL_DEVICE _CCCL_FORCEINLINE void run(num_segments_val_t segment_id, @@ -1873,24 +2387,8 @@ private: // cannot be re-anchored after the fact). ::cuda::std::uint32_t resident_smem32 = 0; - // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), - // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final - // filter; the deterministic filter folds the edges as separate index-ordered regions instead (see below). const auto fold_edges = [&](auto&& apply) { - if constexpr (use_block_load_to_shared) - { - const int tid = static_cast(threadIdx.x); - _CCCL_PRAGMA_NOUNROLL() - for (int local = tid; local < head_edge_len_items; local += threads_per_block) - { - apply(temp_storage.edge_keys[local]); - } - _CCCL_PRAGMA_NOUNROLL() - for (int local = tid; local < tail_edge_len_items; local += threads_per_block) - { - apply(temp_storage.edge_keys[head_edge_cap_items + local]); - } - } + fold_boundary_edges(head_edge_len_items, tail_edge_len_items, apply); }; // Front-load all stage mbarrier inits before any bulk copy issues; the barrier below orders them ahead of the first @@ -1901,150 +2399,23 @@ private: } __syncthreads(); - { - extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); - const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - // Cluster-scope leader atomics are only needed to stay mutually atomic with the non-leaders' DSMEM folds; a lone - // CTA has none, so it uses the cheaper CTA scope. - const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; - auto add_first_pass = [&](const key_t& key) { - const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, is_leader_cluster_scope); - }; - - if constexpr (use_block_load_to_shared) - { - if (my_resident_chunks > 0) - { - // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). - // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor - // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a - // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk - // begins on a `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the - // global-last chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk - // here. The read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor - // cannot demote the reads to `LD`. - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); - - // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` - // shifts the resident window to the high-index chunks). - const auto resident_local = [&](offset_t slot) -> offset_t { - return resident_base + slot; - }; - // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has - // none. - const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { - const offset_t chunk_idx = part.global_index(resident_local(slot)); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - const auto split = split_chunk(block_keys_base, chunk); - if (split.bulk == 0) - { - return {}; - } - return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; - }; - - // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives - // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. - int next_off_bytes = 0; - - // Load every resident chunk's aligned bulk, densely packed in slot order. - for (int stage = 0; stage < prologue; ++stage) - { - const auto src = bulk_src(static_cast(stage)); - issue_bulk_copy(stage, key_slots + next_off_bytes, src); - next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - } - - // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed - // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs - // no stored per-stage state. - int read_off_bytes = 0; - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) - { - const int stage = static_cast(local_chunk % static_cast(prologue)); - wait_stage(stage); - const int read_len_items = static_cast(::cuda::std::size(bulk_src(local_chunk))); - for_each_chunk_key( - bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off_bytes)), - read_len_items}), - add_first_pass); - read_off_bytes += read_len_items * int{sizeof(key_t)}; - - const offset_t next_local_chunk = local_chunk + static_cast(prologue); - if (next_local_chunk < my_resident_chunks) - { - const auto src = bulk_src(next_local_chunk); - // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads - // leave the wait above would advance the phase twice, stranding a lagging waiter forever. - __syncthreads(); - issue_bulk_copy(stage, key_slots + next_off_bytes, src); - next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; - } - } - - // The resident region is one contiguous span of aligned bulks for the later passes; both boundary edges are - // folded separately from `edge_keys`. - resident_keys = {reinterpret_cast(key_slots), - static_cast<::cuda::std::size_t>(next_off_bytes / int{sizeof(key_t)})}; - resident_smem32 = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots)); - } - } - else - { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) - { - const offset_t chunk_idx = part.global_index(resident_base + local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); - const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) - for (int j = 0; j < iterations; ++j) - { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < chunk.count) - { - const key_t key = - block_keys_in[static_cast(chunk.offset + static_cast(local))]; - chunk_keys[local] = key; - add_first_pass(key); - } - } - } - } - - // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see - // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix - // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. - // The `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via - // `fold_edges` after a barrier). - if constexpr (use_block_load_to_shared) - { - if (head_edge_len_items > 0) - { - stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len_items, add_first_pass); - } - if (tail_edge_len_items > 0) - { - const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); - stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, - block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), - tail_edge_len_items, - add_first_pass); - } - } - - // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the streamer's initial - // direction (preselected above; the histogram is order-independent, so the direction only sets up the leftover - // parity for the final filter). The streamer reuses the resident load's stage mbarriers (all front-loaded at - // `run` entry); `wait_stage` provides the producer/consumer sync. - streamer.process_pass(add_first_pass); - - const int resident_count = span_size(resident_keys); - _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, - "Dynamic shared memory block_tile is too small"); - __syncthreads(); - } + load_and_histogram_first_pass( + cluster_rank, + leader_rank, + is_single_cta, + part, + my_resident_chunks, + resident_base, + chunks, + segment_size_u32, + head_items, + block_keys_in, + block_keys_base, + head_edge_len_items, + tail_edge_len_items, + streamer, + resident_keys, + resident_smem32); for (int pass = 0; pass < num_passes; ++pass) { @@ -2220,20 +2591,6 @@ private: auto block_keys_out = d_key_segments_out_it[segment_id]; const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes - // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index - // `seg_idx`. Deriving the per-segment value iterators *inside* the `is_keys_only` guard avoids indexing the null - // `cub::NullType**` value-iterators-of-iterators in keys-only builds; `segment_id` is loop-invariant, so they hoist - // out of the writes. Serves only the non-deterministic branch below (the deterministic path uses - // `det_final_filter::write_value`), so it is unused but still instantiated in deterministic builds. - [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { - if constexpr (!is_keys_only) - { - auto block_vals_in = d_value_segments_it[segment_id]; - auto block_vals_out = d_value_segments_out_it[segment_id]; - block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; - } - }; - // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) // trailing digits as smaller and erroneously reject candidates that share the identified prefix. @@ -2241,272 +2598,55 @@ private: if constexpr (need_determinism) { - // Deterministic tie-break: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); - // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on - // the boundary-crossing tile (see `is_select_all_cand_cta` and the per-tile back logic). Early stop is not - // special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. - // - // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster - // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). - const offset_t total_candidates = leader_state->len; - const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys - const out_offset_t num_selected = k - num_back; // front region - - // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final - // cluster barrier) block-wide before they feed the scan and `front_count`. - __syncthreads(); - const bool participates = !is_idle_rank && (cluster_rank != leader_rank); - const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; - const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; - // Front count pushed by this block: its strictly-selected count. Candidates always route through the back, so - // nothing folds into the front here. The leader and idle ranks push 0 -- the leader because its merged histogram - // cannot self-count (it derives its own front from the total below), idle ranks because they own nothing. - const offset_t push_front = my_sel; - const ::cuda::std::uint64_t packed = - (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = - combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); - // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so - // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A - // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary - // (`cand_prefix + my_cand_count <= num_back`): every one wins, so the back places them with arrival-order SMEM - // atomics and skips the index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the - // single boundary-crossing (straddling) CTA cluster-wide. - const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; - const bool is_select_all_cand_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); - // Mirror image: a CTA selects none of its candidates when the tie region is empty (`num_back == 0`) or all of its - // candidates sort strictly after the K-boundary (`cand_prefix >= num_back`). Such a CTA seeds `is_tie_active` - // false and skips the back placement entirely. - const bool is_select_no_cand_cta = - (num_back == out_offset_t{0}) || (cand_prefix >= static_cast(num_back)); - // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so - // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. - const out_offset_t my_front = - (cluster_rank == leader_rank) - ? static_cast(num_selected - static_cast(sel_prefix)) - : static_cast(push_front); - - // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the - // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never - // part of this span. - const int front_count = span_size(resident_keys); - - // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop - // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later - // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. - const bool has_head = head_edge_len_items > 0; - const bool has_resident = front_count > 0; - const bool has_overflow = overflow_chunks > offset_t{0}; - const bool has_tail_edge = tail_edge_len_items > 0; - [[maybe_unused]] const bool is_resident_terminal = - is_tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); - [[maybe_unused]] const bool is_head_edge_terminal = - is_tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); - [[maybe_unused]] const bool is_tail_edge_terminal = - is_tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; - - // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to - // the high-index window under `is_residency_reversed`). The blocked partition packs the front contiguously, so - // element `pos` maps to `front_seg_base + pos`. Guarded on `front_count > 0`: with a fully-streamed rank - // (`stream_slots == full_slots`) there are no resident chunks and `resident_base` would name an out-of-range - // local chunk under `is_residency_reversed`. - const offset_t front_seg_base = - (front_count > 0) - ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset - : offset_t{0}; - - // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` - // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at - // `!is_select_no_cand_cta` -- true unless this is a select-no-candidates CTA (empty back region, or all of this - // CTA's candidates sort past the boundary). - det_final_filter filt{ - *this, + write_deterministic_topk( segment_id, + cluster_rank, + eff_cluster_blocks, + leader_rank, + is_single_cta, + is_idle_rank, + k, + num_kth, identify_op, block_keys_out, block_keys_in, - streamer, part, - k, - num_back, - my_front, - sel_prefix, - cand_prefix, - my_cand_count, resident_base, my_resident_chunks, + overflow_chunks, segment_size_u32, head_items, - front_seg_base, + leader_state, resident_smem32, - front_count, + resident_keys, head_edge_len_items, tail_edge_len_items, - is_select_all_cand_cta, - is_resident_terminal, - is_head_edge_terminal, - is_tail_edge_terminal, - cand_prefix, - !is_select_no_cand_cta}; - filt.run_filter(); + streamer); } else { - // Non-deterministic tie-break: strictly-selected keys fill the front `[0, num_selected)`, the first `num_kth` - // candidates (arrival order) fill the back. The same combined scan as the deterministic path gives this block - // disjoint front/back bases (`sel_prefix`/`cand_prefix`); placement then uses block-local SMEM atomics since - // output order is not preserved. A perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, - // not a correctness one. - __syncthreads(); - const bool participates = !is_idle_rank && (cluster_rank != leader_rank); - const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; - const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; - const ::cuda::std::uint64_t packed = - (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = - combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); - - // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local - // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair - // mode each written key additionally pulls its value from gmem at `seg_idx`. Keys-only elides that write via the - // `if constexpr` below and drives the sink from index-free traversals, passing a dummy `seg_idx` that goes - // unread. - const auto sink = [&](const key_t& key, [[maybe_unused]] offset_t seg_idx) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); - block_keys_out[pos] = key; - if constexpr (!is_keys_only) - { - write_value(pos, seg_idx); - } - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); - if (back_pos < num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; - if constexpr (!is_keys_only) - { - write_value(pos, seg_idx); - } - } - } - }; - - if constexpr (is_keys_only) - { - // Keys-only: no value payload, so drive the sink through the index-free traversal with a dummy `seg_idx`. - const auto write_selected = [&](const key_t& key) { - sink(key, offset_t{0}); - }; - // Fold the resident keys as the streamer's `mid` work so they overlap the first overflow reloads. Writes are - // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) - { - key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - for_each_chunk_key( - {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); - } - else - { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) - { - const offset_t chunk_idx = part.global_index(local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); - for_each_chunk_key( - {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); - } - } - // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_edges(write_selected); - }; - streamer.process_pass(write_selected, fold_resident); - } - else - { - // Pair (key + value) path: the shared `sink` above additionally stores each written key's value (fetched from - // gmem at its `seg_idx`). Unlike keys-only, the traversal must recover each key's segment-local index, so it - // folds resident/edge keys through `write_run` (below) and overflow keys through `process_pass_indexed`. - - // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local - // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary - // edges); overflow chunks fold through the streamer's own indexed callback below. Materialize the key into a - // register first: `sink` binds it by `const&` and reads it several times, so passing `smem[local]` directly - // would re-issue a narrow `LDS` per use instead of reusing the loaded value. - auto write_run = [&](const key_t* smem, offset_t base_off, int count) { - const int iterations = ::cuda::ceil_div(count, threads_per_block); - _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) - for (int j = 0; j < iterations; ++j) - { - const int local = j * threads_per_block + static_cast(threadIdx.x); - if (local < count) - { - const key_t key = smem[local]; - sink(key, base_off + static_cast(local)); - } - } - }; - - // Fold the resident keys (and their values) as the streamer's `mid` work, exactly as in the keys-only path. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) - { - // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers - // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into - // `edge_keys` and folded below), so iterate `split.bulk`, not `chunk.count`. - key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); - int cursor_items = 0; - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) - { - const auto chunk = - get_chunk(part.global_index(resident_base + local_chunk), segment_size_u32, head_items); - const auto split = split_chunk(block_keys_base, chunk); - const offset_t base_off = chunk.offset + split.prefix; - const int bulk_count_items = split.bulk; - write_run(resident_ptr + cursor_items, base_off, bulk_count_items); - cursor_items += bulk_count_items; - } - } - else - { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) - { - const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_u32, head_items); - const offset_t base_off = chunk.offset; - write_run(slot_keys_unpadded(static_cast(local_chunk)), base_off, chunk.count); - } - } - // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the - // peeled tail suffix at `segment_size - tail_edge_len_items`. Value fetched per key. - if constexpr (use_block_load_to_shared) - { - if (head_edge_len_items > 0) - { - write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); - } - if (tail_edge_len_items > 0) - { - write_run(temp_storage.edge_keys + head_edge_cap_items, - segment_size_u32 - static_cast(tail_edge_len_items), - tail_edge_len_items); - } - } - }; - - // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected - // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. - streamer.process_pass_indexed(sink, fold_resident); - } + write_nondeterministic_topk( + segment_id, + cluster_rank, + eff_cluster_blocks, + leader_rank, + is_single_cta, + is_idle_rank, + k, + num_kth, + identify_op, + block_keys_out, + block_keys_base, + part, + resident_base, + my_resident_chunks, + segment_size_u32, + head_items, + resident_smem32, + resident_keys, + head_edge_len_items, + tail_edge_len_items, + streamer); } // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into From d0458bfb9443e61300f529ded80c9ed120659ec3 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 05:04:59 +0200 Subject: [PATCH 085/126] Split up more long functions Slight SASS impact, need to benchmark. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 373 ++++++++++++------- 1 file changed, 228 insertions(+), 145 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index d9a26905a1d..5a83211105d 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -2211,82 +2211,92 @@ private: __syncthreads(); } - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - run(num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int cluster_blocks, - segment_size_val_t segment_size, - out_offset_t k) + // Per-segment, per-rank geometry computed once at the top of `run`: the head-aligned chunking, the effective + // (non-idle) cluster width and this rank's partition of it, the leader/idle roles, and the resident/streaming/edge + // split of this rank's chunks. All fields are loop-invariant for the segment; the streaming-only intermediates + // (`overflow_base`, `resident_slots_cap`, `stream_slots`, `my_chunks`) are kept because the `overflow_streamer` + // constructor needs them. + struct segment_layout { - using extract_bin_op_t = detail::topk::extract_bin_op_t; - using identify_candidates_op_t = - detail::topk::identify_candidates_op_t; - - constexpr int total_bits = int{sizeof(key_t)} * 8; - constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + offset_t segment_size_u32; + bool is_single_cta; + key_it_t block_keys_in; + const key_t* block_keys_base; + offset_t head_items; + offset_t chunks; + unsigned int eff_cluster_blocks; + bool is_idle_rank; + chunk_partition part; + unsigned int leader_rank; + state_t* leader_state; + offset_t my_chunks; + offset_t my_resident_chunks; + offset_t overflow_chunks; + offset_t resident_base; + offset_t overflow_base; + offset_t resident_slots_cap; + offset_t stream_slots; + int head_edge_len_items; + int tail_edge_len_items; + }; - auto block_keys_in = d_key_segments_it[segment_id]; - const auto segment_size_u32 = static_cast(segment_size); + _CCCL_DEVICE _CCCL_FORCEINLINE segment_layout compute_segment_layout( + num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int cluster_blocks, + segment_size_val_t segment_size) + { + segment_layout layout; + layout.block_keys_in = d_key_segments_it[segment_id]; + layout.segment_size_u32 = static_cast(segment_size); // `cluster_blocks` is what `process_impl` runs at: the launched cluster size, or `1` when it collapsed a // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. - const bool is_single_cta = (cluster_blocks == 1u); - - // Per-block local copy of `kth_key_bits` so each key check hits the - // block's own SMEM rather than DSMEM during the histogram loop. Built up one splitter digit per pass from the - // leader's published `kth_bucket` (see the pass loop), so the full key is never broadcast. - key_prefix_t kth_key_bits_local = {}; + layout.is_single_cta = (cluster_blocks == 1u); - // Tracks the highest pass count that actually executed. Without early - // stop this stays at `num_passes`; with early stop it captures the pass - // at which we broke out so the final filter can construct its identify - // operator at the matching radix level. - int last_pass = num_passes; - - const key_t* block_keys_base = nullptr; - offset_t head_items = 0; + layout.block_keys_base = nullptr; + layout.head_items = 0; if constexpr (use_block_load_to_shared) { - block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(block_keys_in); - head_items = aligned_head_items(block_keys_base, segment_size_u32); + layout.block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(layout.block_keys_in); + layout.head_items = aligned_head_items(layout.block_keys_base, layout.segment_size_u32); } // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only // depends on the multiset of keys covered by the cluster. - const offset_t chunks = num_chunks(segment_size_u32, head_items); + layout.chunks = num_chunks(layout.segment_size_u32, layout.head_items); // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_block` each), <= the // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- // but stay resident and still arrive at every cluster barrier (a returned CTA would hang the barrier; see the // TODOs at the barrier sites). Derived from this CTA's head-aligned `chunks` so it matches the partition exactly. // Stays at `cluster_blocks` for host-exact sizes (the dispatch already matched it) and on the single-CTA path. - unsigned int eff_cluster_blocks = cluster_blocks; + layout.eff_cluster_blocks = cluster_blocks; if constexpr (enable_runtime_single_cta) { - if (!is_single_cta) + if (!layout.is_single_cta) { - eff_cluster_blocks = effective_cluster_blocks_from_chunks( - static_cast<::cuda::std::uint64_t>(chunks), min_chunks_per_block, cluster_blocks); + layout.eff_cluster_blocks = effective_cluster_blocks_from_chunks( + static_cast<::cuda::std::uint64_t>(layout.chunks), min_chunks_per_block, cluster_blocks); } } - const bool is_idle_rank = cluster_rank >= eff_cluster_blocks; + layout.is_idle_rank = cluster_rank >= layout.eff_cluster_blocks; // Idle ranks own no chunks; `make_chunk_partition` assumes `rank < size`, so hand them an explicit empty partition. - const chunk_partition part = is_idle_rank ? chunk_partition{offset_t{0}, offset_t{1}, offset_t{0}} - : make_chunk_partition(chunks, cluster_rank, eff_cluster_blocks); - const offset_t my_chunks = part.count; + layout.part = layout.is_idle_rank ? chunk_partition{offset_t{0}, offset_t{1}, offset_t{0}} + : make_chunk_partition(layout.chunks, cluster_rank, layout.eff_cluster_blocks); + layout.my_chunks = layout.part.count; // Leader rank. The leader owns the cluster-merged histogram and the shared `state`, and is always a working rank // (`< eff_cluster_blocks`). The deterministic tie-break makes the leader the *last* CTA in scan order so it never // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. - const unsigned int leader_rank = (need_determinism && !is_tie_reversed) ? (eff_cluster_blocks - 1u) : 0u; + layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1u) : 0u; // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - state_t* leader_state = is_single_cta ? &temp_storage.state : map_state_to_rank(leader_rank); + layout.leader_state = layout.is_single_cta ? &temp_storage.state : map_state_to_rank(layout.leader_rank); // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every // cluster barrier is reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) @@ -2303,7 +2313,7 @@ private: // `PipelineStages` depth. The generic fallback has no async pipeline (it re-reads overflow from gmem each pass and // peels nothing), so it reserves no streaming slots. const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); - [[maybe_unused]] const bool needs_streaming = my_chunks > full_slots; + [[maybe_unused]] const bool needs_streaming = layout.my_chunks > full_slots; // Does this rank own the global tail, and does that tail carry an unaligned suffix? (block-load path only; the // generic fallback reads any trailing items straight from gmem and never peels.) @@ -2313,10 +2323,11 @@ private: { // This rank owns the global tail iff its last owned chunk is chunk `chunks-1` (its local index `my_chunks-1`, // true for both the strided and blocked partitions). - if (my_chunks > 0 && part.global_index(my_chunks - offset_t{1}) == chunks - offset_t{1}) + if (layout.my_chunks > 0 + && layout.part.global_index(layout.my_chunks - offset_t{1}) == layout.chunks - offset_t{1}) { - const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); - tail_suffix_items = split_chunk(block_keys_base, tail_chunk).suffix; + const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_u32, layout.head_items); + tail_suffix_items = split_chunk(layout.block_keys_base, tail_chunk).suffix; owns_suffix_tail = tail_suffix_items != offset_t{0}; } } @@ -2330,93 +2341,64 @@ private: should_peel_tail = owns_suffix_tail; if (needs_streaming) { - const offset_t excess = my_chunks - full_slots; + const offset_t excess = layout.my_chunks - full_slots; const offset_t want_stream = (::cuda::std::min) (static_cast(PipelineStages), excess); stream_slots = (::cuda::std::max) (offset_t{1}, (::cuda::std::min) (want_stream, full_slots)); } } - const offset_t resident_slots_cap = full_slots - stream_slots; - const offset_t my_resident_chunks = (::cuda::std::min) (my_chunks, resident_slots_cap); + layout.resident_slots_cap = full_slots - stream_slots; + layout.my_resident_chunks = (::cuda::std::min) (layout.my_chunks, layout.resident_slots_cap); // Resident chunks stay within the first `resident_slots_cap` slots; the streaming region occupies the slots // `[resident_slots_cap, full_slots)`, so both regions live inside the allocated block_tile buffer. - _CCCL_ASSERT(my_resident_chunks <= resident_slots_cap, "Dynamic shared memory block_tile is too small"); + _CCCL_ASSERT(layout.my_resident_chunks <= layout.resident_slots_cap, + "Dynamic shared memory block_tile is too small"); - const offset_t overflow_chunks = (my_chunks > my_resident_chunks) ? (my_chunks - my_resident_chunks) : offset_t{0}; + layout.overflow_chunks = + (layout.my_chunks > layout.my_resident_chunks) ? (layout.my_chunks - layout.my_resident_chunks) : offset_t{0}; // Rank-local base of the resident and overflow (streamed) chunk windows. Default: resident `[0, // my_resident_chunks)`, overflow the high-index rest. `is_residency_reversed` swaps them: resident // `[overflow_chunks, my_chunks)`, overflow `[0, overflow_chunks)`. - const offset_t resident_base = is_residency_reversed ? overflow_chunks : offset_t{0}; - const offset_t overflow_base = is_residency_reversed ? offset_t{0} : my_resident_chunks; + layout.resident_base = is_residency_reversed ? layout.overflow_chunks : offset_t{0}; + layout.overflow_base = is_residency_reversed ? offset_t{0} : layout.my_resident_chunks; + layout.stream_slots = stream_slots; // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. - [[maybe_unused]] const int head_edge_len_items = (cluster_rank == 0u) ? static_cast(head_items) : 0; - [[maybe_unused]] const int tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; - - // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. - overflow_streamer streamer( - *this, - block_keys_in, - block_keys_base, - segment_size_u32, - head_items, - part, - my_resident_chunks, - overflow_base, - static_cast(resident_slots_cap), - static_cast(stream_slots), - my_chunks); - - // Preselect the streamer's initial ping-pong direction. A streaming rank flips direction once per histogram pass, - // so the leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = - // (!is_tie_reversed) ^ (num_passes & 1)` makes that leftover `== !is_tie_reversed` -- exactly what the - // deterministic filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see - // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is - // moot). Non-deterministic filtering is order-independent, so leave its historical `is_forward` start untouched. - if constexpr (need_determinism) - { - streamer.is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); - } - - ::cuda::std::span resident_keys; - // 32-bit shared-window address of `resident_keys.data()`. The resident span is read once per radix pass and in - // the final filter; a 64-bit generic pointer kept live across that loop spills and reloads as a *generic* - // pointer (`LDL.64`), which demotes every key read from `LDS` to a generic `LD`. Carrying the base as a 32-bit - // shared address and rebuilding the pointer with `__cvta_shared_to_generic` at each use keeps the reads `LDS` - // even when the value spills (the cvta intrinsic re-confers shared provenance; a spilled 64-bit generic pointer - // cannot be re-anchored after the fact). - ::cuda::std::uint32_t resident_smem32 = 0; - - const auto fold_edges = [&](auto&& apply) { - fold_boundary_edges(head_edge_len_items, tail_edge_len_items, apply); - }; + layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_items) : 0; + layout.tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; + return layout; + } - // Front-load all stage mbarrier inits before any bulk copy issues; the barrier below orders them ahead of the first - // issue (see `init_load_barriers`). The generic fallback has no mbarriers but keeps the same pass-start barrier. - if constexpr (use_block_load_to_shared) - { - init_load_barriers(); - } - __syncthreads(); + // Radix histogram/scan passes: refine the splitter one `bits_per_pass` digit at a time until the top-k prefix is + // pinned down (or early stop fires). Folds `kth_key_bits_local` up digit by digit and returns the number of passes + // that actually ran (`last_pass`), which the final filter uses to size its identify operator. + template + _CCCL_DEVICE _CCCL_FORCEINLINE int run_radix_passes( + unsigned int cluster_rank, + unsigned int leader_rank, + bool is_single_cta, + bool is_idle_rank, + const chunk_partition& part, + offset_t my_resident_chunks, + offset_t resident_base, + offset_t segment_size_u32, + offset_t head_items, + state_t* leader_state, + int head_edge_len_items, + int tail_edge_len_items, + ::cuda::std::span resident_keys, + ::cuda::std::uint32_t resident_smem32, + overflow_streamer& streamer, + key_prefix_t& kth_key_bits_local) + { + using extract_bin_op_t = detail::topk::extract_bin_op_t; + using identify_candidates_op_t = + detail::topk::identify_candidates_op_t; - load_and_histogram_first_pass( - cluster_rank, - leader_rank, - is_single_cta, - part, - my_resident_chunks, - resident_base, - chunks, - segment_size_u32, - head_items, - block_keys_in, - block_keys_base, - head_edge_len_items, - tail_edge_len_items, - streamer, - resident_keys, - resident_smem32); + constexpr int total_bits = int{sizeof(key_t)} * 8; + constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + int last_pass = num_passes; for (int pass = 0; pass < num_passes; ++pass) { const bool is_first_pass = (pass == 0); @@ -2474,7 +2456,7 @@ private: // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's per-bucket counts (the source of its `num_strictly_selected` // and `my_candidates`, the cross-CTA scan inputs) inclusive of its edge candidates. - fold_edges(add_hist); + fold_boundary_edges(head_edge_len_items, tail_edge_len_items, add_hist); } // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` @@ -2581,6 +2563,107 @@ private: } #endif } + return last_pass; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + run(num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int cluster_blocks, + segment_size_val_t segment_size, + out_offset_t k) + { + using identify_candidates_op_t = + detail::topk::identify_candidates_op_t; + + constexpr int total_bits = int{sizeof(key_t)} * 8; + constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + + const segment_layout layout = compute_segment_layout(segment_id, cluster_rank, cluster_blocks, segment_size); + + // Per-block local copy of `kth_key_bits` so each key check hits the block's own SMEM rather than DSMEM during the + // histogram loop. Built up one splitter digit per pass from the leader's published `kth_bucket` (see + // `run_radix_passes`), so the full key is never broadcast. + key_prefix_t kth_key_bits_local = {}; + + // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. + overflow_streamer streamer( + *this, + layout.block_keys_in, + layout.block_keys_base, + layout.segment_size_u32, + layout.head_items, + layout.part, + layout.my_resident_chunks, + layout.overflow_base, + static_cast(layout.resident_slots_cap), + static_cast(layout.stream_slots), + layout.my_chunks); + + // Preselect the streamer's initial ping-pong direction. A streaming rank flips direction once per histogram pass, + // so the leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = + // (!is_tie_reversed) ^ (num_passes & 1)` makes that leftover `== !is_tie_reversed` -- exactly what the + // deterministic filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see + // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is + // moot). Non-deterministic filtering is order-independent, so leave its historical `is_forward` start untouched. + if constexpr (need_determinism) + { + streamer.is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); + } + + ::cuda::std::span resident_keys; + // 32-bit shared-window address of `resident_keys.data()`. The resident span is read once per radix pass and in + // the final filter; a 64-bit generic pointer kept live across that loop spills and reloads as a *generic* + // pointer (`LDL.64`), which demotes every key read from `LDS` to a generic `LD`. Carrying the base as a 32-bit + // shared address and rebuilding the pointer with `__cvta_shared_to_generic` at each use keeps the reads `LDS` + // even when the value spills (the cvta intrinsic re-confers shared provenance; a spilled 64-bit generic pointer + // cannot be re-anchored after the fact). + ::cuda::std::uint32_t resident_smem32 = 0; + + // Front-load all stage mbarrier inits before any bulk copy issues; the barrier below orders them ahead of the first + // issue (see `init_load_barriers`). The generic fallback has no mbarriers but keeps the same pass-start barrier. + if constexpr (use_block_load_to_shared) + { + init_load_barriers(); + } + __syncthreads(); + + load_and_histogram_first_pass( + cluster_rank, + layout.leader_rank, + layout.is_single_cta, + layout.part, + layout.my_resident_chunks, + layout.resident_base, + layout.chunks, + layout.segment_size_u32, + layout.head_items, + layout.block_keys_in, + layout.block_keys_base, + layout.head_edge_len_items, + layout.tail_edge_len_items, + streamer, + resident_keys, + resident_smem32); + + const int last_pass = run_radix_passes( + cluster_rank, + layout.leader_rank, + layout.is_single_cta, + layout.is_idle_rank, + layout.part, + layout.my_resident_chunks, + layout.resident_base, + layout.segment_size_u32, + layout.head_items, + layout.leader_state, + layout.head_edge_len_items, + layout.tail_edge_len_items, + resident_keys, + resident_smem32, + streamer, + kth_key_bits_local); // ----------------------------------------------------------------------- // Final filter pass: write the top-k keys for this segment. Strictly- @@ -2589,7 +2672,7 @@ private: // from each pass's bucket above), so no broadcast is needed here. // ----------------------------------------------------------------------- auto block_keys_out = d_key_segments_out_it[segment_id]; - const out_offset_t num_kth = leader_state->k; // remaining k after the radix passes + const out_offset_t num_kth = layout.leader_state->k; // remaining k after the radix passes // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) @@ -2601,26 +2684,26 @@ private: write_deterministic_topk( segment_id, cluster_rank, - eff_cluster_blocks, - leader_rank, - is_single_cta, - is_idle_rank, + layout.eff_cluster_blocks, + layout.leader_rank, + layout.is_single_cta, + layout.is_idle_rank, k, num_kth, identify_op, block_keys_out, - block_keys_in, - part, - resident_base, - my_resident_chunks, - overflow_chunks, - segment_size_u32, - head_items, - leader_state, + layout.block_keys_in, + layout.part, + layout.resident_base, + layout.my_resident_chunks, + layout.overflow_chunks, + layout.segment_size_u32, + layout.head_items, + layout.leader_state, resident_smem32, resident_keys, - head_edge_len_items, - tail_edge_len_items, + layout.head_edge_len_items, + layout.tail_edge_len_items, streamer); } else @@ -2628,24 +2711,24 @@ private: write_nondeterministic_topk( segment_id, cluster_rank, - eff_cluster_blocks, - leader_rank, - is_single_cta, - is_idle_rank, + layout.eff_cluster_blocks, + layout.leader_rank, + layout.is_single_cta, + layout.is_idle_rank, k, num_kth, identify_op, block_keys_out, - block_keys_base, - part, - resident_base, - my_resident_chunks, - segment_size_u32, - head_items, + layout.block_keys_base, + layout.part, + layout.resident_base, + layout.my_resident_chunks, + layout.segment_size_u32, + layout.head_items, resident_smem32, resident_keys, - head_edge_len_items, - tail_edge_len_items, + layout.head_edge_len_items, + layout.tail_edge_len_items, streamer); } From 19a9cae8b4973279e1099f39ef93b1d68de8b73f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 15:54:27 +0200 Subject: [PATCH 086/126] Use ceil_div and clamp --- cub/cub/agent/agent_batched_topk_cluster.cuh | 4 ++-- ...atch2_test_device_segmented_topk_cluster_layout.cu | 11 +++-------- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5a83211105d..3ebbfe79879 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -183,7 +183,7 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, { const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_block); return static_cast( - (::cuda::std::clamp) (blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); + ::cuda::std::clamp(blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } // Cluster top-k agent @@ -2343,7 +2343,7 @@ private: { const offset_t excess = layout.my_chunks - full_slots; const offset_t want_stream = (::cuda::std::min) (static_cast(PipelineStages), excess); - stream_slots = (::cuda::std::max) (offset_t{1}, (::cuda::std::min) (want_stream, full_slots)); + stream_slots = ::cuda::std::clamp(want_stream, offset_t{1}, full_slots); } } layout.resident_slots_cap = full_slots - stream_slots; diff --git a/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu index 1b13f8e29e5..71c8fde1f69 100644 --- a/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu +++ b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu @@ -4,18 +4,13 @@ #include // smem_block_tile_layout #include // batched_topk::cluster_policy_selector +#include #include #include namespace { -template -[[nodiscard]] constexpr SizeT host_ceil_div(SizeT numerator, SizeT denominator) -{ - return (numerator + denominator - 1) / denominator; -} - template void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) { @@ -42,8 +37,8 @@ void check_layout_case(int dynamic_smem_bytes, int cluster_blocks) const auto max_rank_chunks = [](cuda::std::int64_t segment_size, int head_items, int blocks) { using size_t = cuda::std::int64_t; const size_t tail_items = segment_size - head_items; - const size_t chunks = host_ceil_div(tail_items, size_t{layout_t::chunk_items}); - return host_ceil_div(chunks, static_cast(blocks)); + const size_t chunks = ::cuda::ceil_div(tail_items, size_t{layout_t::chunk_items}); + return ::cuda::ceil_div(chunks, static_cast(blocks)); }; const int heads[] = {0, 1, layout_t::chunk_items / 2, layout_t::chunk_items - 1}; From 75a88a88215c4ef3bfd2f0db06161814dda4fa1f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 7 Jul 2026 19:45:01 +0200 Subject: [PATCH 087/126] Use I32 offsets instead of U32 - Document the current restriction on segment sizes and add a compile-time check. - Change benchmarks to use I32 segment sizes instead of I64 since this is in-line with downstream users. --- .../bench/segmented_topk/fixed/keys.cu | 4 +- .../bench/segmented_topk/variable/common.cuh | 7 +- .../variable/indexed_common.cuh | 10 +- .../segmented_topk/variable/keys_common.cuh | 12 +- cub/cub/agent/agent_batched_topk.cuh | 12 +- cub/cub/agent/agent_batched_topk_cluster.cuh | 203 ++++++------- cub/cub/detail/segmented_params.cuh | 75 ++++- cub/cub/device/device_batched_topk.cuh | 222 ++++++++++---- .../device/dispatch/dispatch_batched_topk.cuh | 25 +- .../catch2_test_device_batched_topk_api.cu | 47 --- .../catch2_test_device_segmented_topk_keys.cu | 282 +++++++++++++++--- ...catch2_test_device_segmented_topk_pairs.cu | 6 +- cub/test/catch2_test_device_topk_common.cuh | 3 +- ...t_device_batched_topk_requirements_fail.cu | 75 ++++- 14 files changed, 695 insertions(+), 288 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index 37ea40eb095..cf0b139a785 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -227,8 +227,8 @@ void fixed_seg_size_topk_keys( state.add_global_memory_writes(selected_elements * num_segments, "OutputKeys"); // Host copy of segment sizes — all entries equal MaxSegmentSize for fixed-size segments. Consumed only by the - // per-segment device backend. - std::vector h_segment_sizes(num_segments, static_cast(MaxSegmentSize)); + // per-segment device backend. Segment sizes fit in a signed 32-bit integer (the library caps them at 2^21). + std::vector h_segment_sizes(num_segments, static_cast(MaxSegmentSize)); caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { diff --git a/cub/benchmarks/bench/segmented_topk/variable/common.cuh b/cub/benchmarks/bench/segmented_topk/variable/common.cuh index 7245cf9e194..6fa2a3eeff8 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/common.cuh @@ -19,6 +19,11 @@ #include +// Segment sizes are stored as signed 32-bit integers: the library caps supported segment sizes at 2^21 (about 2 +// million), so a wider type buys nothing here. The benchmarks always pass an explicit `cuda::args::bounds` (a bare +// int32 would be rejected, as its type maximum exceeds the cap). +using segment_size_t = cuda::std::int32_t; + namespace { enum class pattern_kind : int @@ -57,7 +62,7 @@ enum class pattern_kind : int template [[nodiscard]] thrust::device_vector -gen_data(int num_segments, pattern_kind pattern, const cuda::std::int64_t* d_seg_sizes) +gen_data(int num_segments, pattern_kind pattern, const segment_size_t* d_seg_sizes) { const auto num_keys = static_cast(num_segments) * static_cast(MaxSegmentSize); auto d_keys = thrust::device_vector{num_keys, thrust::no_init}; diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh index c8cfdebe1cd..2e92c2958a8 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -208,12 +208,12 @@ void decode_style_variable_topk_indexed( return; } - const auto num_segments = static_cast(state.get_int64("NumSegments")); - const thrust::device_vector d_segment_sizes = generate( + const auto num_segments = static_cast(state.get_int64("NumSegments")); + const thrust::device_vector d_segment_sizes = generate( static_cast(num_segments), bit_entropy::_1_000, - static_cast(K), - static_cast(MaxSegmentSize)); + static_cast(K), + static_cast(MaxSegmentSize)); const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); const auto output_elements = static_cast(num_segments) * K; @@ -242,7 +242,7 @@ void decode_style_variable_topk_indexed( state.add_element_count(input_elements, "NumElements"); state.add_global_memory_reads(input_elements, "InputKeys"); - state.add_global_memory_reads(num_segments, "SegmentSizes"); + state.add_global_memory_reads(num_segments, "SegmentSizes"); state.add_global_memory_writes(output_elements, "OutputKeys"); state.add_global_memory_writes(output_elements, "OutputIndices"); diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh index a713cbdb5ad..ce5f97408c3 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -241,12 +241,12 @@ void decode_style_variable_topk_keys( return; } - const auto num_segments = static_cast(state.get_int64("NumSegments")); - const thrust::device_vector d_segment_sizes = generate( + const auto num_segments = static_cast(state.get_int64("NumSegments")); + const thrust::device_vector d_segment_sizes = generate( static_cast(num_segments), bit_entropy::_1_000, - static_cast(K), - static_cast(MaxSegmentSize)); + static_cast(K), + static_cast(MaxSegmentSize)); const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end()); const auto output_elements = static_cast(num_segments) * K; @@ -268,11 +268,11 @@ void decode_style_variable_topk_keys( state.add_element_count(input_elements, "NumElements"); state.add_global_memory_reads(input_elements, "InputKeys"); - state.add_global_memory_reads(num_segments, "SegmentSizes"); + state.add_global_memory_reads(num_segments, "SegmentSizes"); state.add_global_memory_writes(output_elements, "OutputKeys"); // Host copy of segment sizes — consumed only by the per-segment device backend. - std::vector h_segment_sizes(static_cast(num_segments)); + std::vector h_segment_sizes(static_cast(num_segments)); thrust::copy(d_segment_sizes.begin(), d_segment_sizes.end(), h_segment_sizes.begin()); caching_allocator_t alloc; diff --git a/cub/cub/agent/agent_batched_topk.cuh b/cub/cub/agent/agent_batched_topk.cuh index d2a99cfc809..dc38b332de2 100644 --- a/cub/cub/agent/agent_batched_topk.cuh +++ b/cub/cub/agent/agent_batched_topk.cuh @@ -200,10 +200,12 @@ struct agent_batched_topk_worker_per_segment && ::cuda::args::__traits::lowest == tile_size; // Resolve Segment Parameters - const auto segment_size = params::get_param(segment_sizes, segment_id); + const auto segment_size = params::get_segment_size(segment_sizes, segment_id); if (!only_small_segments && segment_size > tile_size) { // Enqueue large segment + // TODO(topk): once the large-segment worker is wired up, skip enqueue when the effective k is 0 (nothing to + // select) so an empty/zero-k large segment does not schedule pointless work. if (threadIdx.x == 0u) { // Add to large segment queue @@ -216,8 +218,14 @@ struct agent_batched_topk_worker_per_segment else { // Process small segment - const auto k = (::cuda::std::min) (params::get_param(k_param, segment_id), + const auto k = (::cuda::std::min) (params::get_param(k_param, segment_id), static_cast(segment_size)); + // Nothing to select for an empty segment (including a negative size clamped to 0) or a zero k: skip it, leaving + // its output untouched. Also keeps the block primitive's `valid_items in [1, tile_items]` precondition. + if (k == 0) + { + return; + } const auto direction = select_directions.get_param(segment_id); // Determine padding key based on direction diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 3ebbfe79879..ad1161178a7 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -72,7 +72,6 @@ #include #include #include -#include #include #include @@ -178,11 +177,11 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, // `min_chunks_per_block` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap // differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_block` is // `static_assert`ed positive, so the divide is well-defined. 64-bit math. -[[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( - ::cuda::std::uint64_t chunks, int min_chunks_per_block, unsigned int cluster_blocks_cap) noexcept +[[nodiscard]] _CCCL_HOST_DEVICE constexpr int effective_cluster_blocks_from_chunks( + ::cuda::std::uint64_t chunks, int min_chunks_per_block, int cluster_blocks_cap) noexcept { const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_block); - return static_cast( + return static_cast( ::cuda::std::clamp(blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } @@ -233,8 +232,12 @@ struct agent_batched_topk_cluster using segment_size_val_t = typename ::cuda::args::__traits::element_type; using num_segments_val_t = typename ::cuda::args::__traits::element_type; - using offset_t = ::cuda::std::uint32_t; - using out_offset_t = ::cuda::std::uint32_t; + // Signed, so the indexing arithmetic below can never wrap on underflow. 32-bit comfortably covers every supported + // segment: the public entry caps the statically-known maximum segment size at 2^21, well within this width, so a + // value that actually exceeds its declared bound at runtime is a caller precondition violation (undefined behavior). + // The cross-CTA scan also packs two of these lanes into one `uint64_t`, which only holds for 32-bit lanes. + using offset_t = ::cuda::std::int32_t; + using out_offset_t = ::cuda::std::int32_t; using state_t = cluster_topk_state; using key_prefix_t = typename state_t::key_prefix_t; @@ -502,10 +505,11 @@ struct agent_batched_topk_cluster } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t - num_rank_chunks(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const + num_rank_chunks(offset_t chunks, int cluster_rank, int cluster_blocks) const { - return (cluster_rank < chunks) - ? static_cast((chunks - 1 - cluster_rank) / cluster_blocks + 1) + return (static_cast(cluster_rank) < chunks) + ? static_cast( + (chunks - 1 - static_cast(cluster_rank)) / static_cast(cluster_blocks) + 1) : offset_t{0}; } @@ -537,7 +541,7 @@ struct agent_batched_topk_cluster // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches // ascending contiguous global-index ranges), so it is selected exactly when `need_determinism` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition - make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const + make_chunk_partition(offset_t chunks, int cluster_rank, int cluster_blocks) const { if constexpr (need_determinism) { @@ -829,7 +833,7 @@ private: // it to rank 0 to form the leader's `shared::cluster` address (no 64-bit pointer, no memory descriptor). Cluster // scope makes it mutually atomic with the leader's `hist_inc` adds. _CCCL_DEVICE _CCCL_FORCEINLINE void - hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, unsigned int leader_rank) + hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, int leader_rank) { ::cuda::std::uint32_t remote; asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "r"(leader_rank)); @@ -838,7 +842,7 @@ private: // Generic pointer to this CTA's `state` as seen in the CTA at cluster rank `rank` (reached over DSMEM) -- the PTX // equivalent of cooperative_groups' `map_shared_rank`, via `mapa.u64` (generic-address form). - _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(unsigned int rank) + _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(int rank) { const ::cuda::std::uint64_t own = reinterpret_cast<::cuda::std::uint64_t>(&temp_storage.state); ::cuda::std::uint64_t remote; @@ -849,7 +853,7 @@ private: // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. - _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, ::cuda::std::uint64_t v) + _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(int target_rank, ::cuda::std::uint64_t v) { const ::cuda::std::uint32_t own = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.prefix_pair)); @@ -1272,8 +1276,8 @@ private: // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. - _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( - bool is_single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t + combined_prefix_scan(bool is_single_cta, int cluster_rank, int eff_cluster_blocks, ::cuda::std::uint64_t packed) { if (is_single_cta) { @@ -1284,8 +1288,9 @@ private: if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() - for (unsigned int rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; - // leader last + for (int rank = static_cast(threadIdx.x); rank < cluster_rank; rank += threads_per_block) // lower ranks + // follow; leader + // last { add_remote_prefix(rank, packed); } @@ -1295,7 +1300,8 @@ private: // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. _CCCL_PRAGMA_NOUNROLL() - for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < eff_cluster_blocks; rank += threads_per_block) + for (int rank = cluster_rank + 1 + static_cast(threadIdx.x); rank < eff_cluster_blocks; + rank += threads_per_block) { add_remote_prefix(rank, packed); } @@ -1336,7 +1342,7 @@ private: offset_t my_cand_count; offset_t resident_base; offset_t my_resident_chunks; - offset_t segment_size_u32; + offset_t segment_size_off; offset_t head_items; offset_t front_seg_base; ::cuda::std::uint32_t resident_smem32; @@ -1560,7 +1566,7 @@ private: { const int local_slot = is_residency_reversed ? (resident_chunk_count - 1 - slot) : slot; const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); - const auto chunk = agent.get_chunk(chunk_idx, segment_size_u32, head_items); + const auto chunk = agent.get_chunk(chunk_idx, segment_size_off, head_items); // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy // per-tile boundary detection handles any boundary), so pass `false`. process_tiles( @@ -1595,7 +1601,7 @@ private: [&](int stage, offset_t overflow_idx) { const auto span = streamer.stage_span(stage, overflow_idx); const offset_t base_off = - agent.get_chunk(streamer.chunk_index_of(overflow_idx), segment_size_u32, head_items).offset; + agent.get_chunk(streamer.chunk_index_of(overflow_idx), segment_size_off, head_items).offset; // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving // is one barrier on one tile. @@ -1641,7 +1647,7 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge() { process_edge(agent.temp_storage.edge_keys + head_edge_cap_items, - segment_size_u32 - static_cast(tail_edge_len_items), + segment_size_off - static_cast(tail_edge_len_items), tail_edge_len_items, is_tail_edge_terminal); } @@ -1717,9 +1723,9 @@ private: template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int eff_cluster_blocks, - unsigned int leader_rank, + int cluster_rank, + int eff_cluster_blocks, + int leader_rank, bool is_single_cta, bool is_idle_rank, out_offset_t k, @@ -1730,7 +1736,7 @@ private: const chunk_partition& part, offset_t resident_base, offset_t my_resident_chunks, - offset_t segment_size_u32, + offset_t segment_size_off, offset_t head_items, ::cuda::std::uint32_t resident_smem32, ::cuda::std::span resident_keys, @@ -1813,7 +1819,7 @@ private: for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { const offset_t chunk_idx = part.global_index(local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); @@ -1860,7 +1866,7 @@ private: int cursor_items = 0; for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(resident_base + local_chunk), segment_size_u32, head_items); + const auto chunk = get_chunk(part.global_index(resident_base + local_chunk), segment_size_off, head_items); const auto split = split_chunk(block_keys_base, chunk); const offset_t base_off = chunk.offset + split.prefix; const int bulk_count_items = split.bulk; @@ -1872,7 +1878,7 @@ private: { for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_u32, head_items); + const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_off, head_items); const offset_t base_off = chunk.offset; write_run(slot_keys_unpadded(static_cast(local_chunk)), base_off, chunk.count); } @@ -1888,7 +1894,7 @@ private: if (tail_edge_len_items > 0) { write_run(temp_storage.edge_keys + head_edge_cap_items, - segment_size_u32 - static_cast(tail_edge_len_items), + segment_size_off - static_cast(tail_edge_len_items), tail_edge_len_items); } } @@ -1908,9 +1914,9 @@ private: template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int eff_cluster_blocks, - unsigned int leader_rank, + int cluster_rank, + int eff_cluster_blocks, + int leader_rank, bool is_single_cta, bool is_idle_rank, out_offset_t k, @@ -1922,7 +1928,7 @@ private: offset_t resident_base, offset_t my_resident_chunks, offset_t overflow_chunks, - offset_t segment_size_u32, + offset_t segment_size_off, offset_t head_items, state_t* leader_state, ::cuda::std::uint32_t resident_smem32, @@ -2001,7 +2007,7 @@ private: // local chunk under `is_residency_reversed`. const offset_t front_seg_base = (front_count > 0) - ? get_chunk(part.global_index(resident_base), segment_size_u32, head_items).offset + ? get_chunk(part.global_index(resident_base), segment_size_off, head_items).offset : offset_t{0}; // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` @@ -2024,7 +2030,7 @@ private: my_cand_count, resident_base, my_resident_chunks, - segment_size_u32, + segment_size_off, head_items, front_seg_base, resident_smem32, @@ -2047,14 +2053,14 @@ private: // histogram visible cluster-wide. template _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( - unsigned int cluster_rank, - unsigned int leader_rank, + int cluster_rank, + int leader_rank, bool is_single_cta, const chunk_partition& part, offset_t my_resident_chunks, offset_t resident_base, offset_t chunks, - offset_t segment_size_u32, + offset_t segment_size_off, offset_t head_items, key_it_t block_keys_in, const key_t* block_keys_base, @@ -2100,7 +2106,7 @@ private: // none. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { const offset_t chunk_idx = part.global_index(resident_local(slot)); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); const auto split = split_chunk(block_keys_base, chunk); if (split.bulk == 0) { @@ -2160,7 +2166,7 @@ private: for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { const offset_t chunk_idx = part.global_index(resident_base + local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) @@ -2191,7 +2197,7 @@ private: } if (tail_edge_len_items > 0) { - const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_u32, head_items); + const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_off, head_items); stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), tail_edge_len_items, @@ -2218,16 +2224,16 @@ private: // constructor needs them. struct segment_layout { - offset_t segment_size_u32; + offset_t segment_size_off; bool is_single_cta; key_it_t block_keys_in; const key_t* block_keys_base; offset_t head_items; offset_t chunks; - unsigned int eff_cluster_blocks; + int eff_cluster_blocks; bool is_idle_rank; chunk_partition part; - unsigned int leader_rank; + int leader_rank; state_t* leader_state; offset_t my_chunks; offset_t my_resident_chunks; @@ -2241,31 +2247,28 @@ private: }; _CCCL_DEVICE _CCCL_FORCEINLINE segment_layout compute_segment_layout( - num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int cluster_blocks, - segment_size_val_t segment_size) + num_segments_val_t segment_id, int cluster_rank, int cluster_blocks, segment_size_val_t segment_size) { segment_layout layout; layout.block_keys_in = d_key_segments_it[segment_id]; - layout.segment_size_u32 = static_cast(segment_size); + layout.segment_size_off = static_cast(segment_size); // `cluster_blocks` is what `process_impl` runs at: the launched cluster size, or `1` when it collapsed a // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. - layout.is_single_cta = (cluster_blocks == 1u); + layout.is_single_cta = (cluster_blocks == 1); layout.block_keys_base = nullptr; layout.head_items = 0; if constexpr (use_block_load_to_shared) { layout.block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(layout.block_keys_in); - layout.head_items = aligned_head_items(layout.block_keys_base, layout.segment_size_u32); + layout.head_items = aligned_head_items(layout.block_keys_base, layout.segment_size_off); } // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only // depends on the multiset of keys covered by the cluster. - layout.chunks = num_chunks(layout.segment_size_u32, layout.head_items); + layout.chunks = num_chunks(layout.segment_size_off, layout.head_items); // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_block` each), <= the // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- @@ -2292,7 +2295,7 @@ private: // (`< eff_cluster_blocks`). The deterministic tie-break makes the leader the *last* CTA in scan order so it never // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. - layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1u) : 0u; + layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1) : 0; // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). @@ -2326,7 +2329,7 @@ private: if (layout.my_chunks > 0 && layout.part.global_index(layout.my_chunks - offset_t{1}) == layout.chunks - offset_t{1}) { - const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_u32, layout.head_items); + const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_off, layout.head_items); tail_suffix_items = split_chunk(layout.block_keys_base, tail_chunk).suffix; owns_suffix_tail = tail_suffix_items != offset_t{0}; } @@ -2364,7 +2367,7 @@ private: // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. - layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_items) : 0; + layout.head_edge_len_items = (cluster_rank == 0) ? static_cast(layout.head_items) : 0; layout.tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; return layout; } @@ -2374,14 +2377,14 @@ private: // that actually ran (`last_pass`), which the final filter uses to size its identify operator. template _CCCL_DEVICE _CCCL_FORCEINLINE int run_radix_passes( - unsigned int cluster_rank, - unsigned int leader_rank, + int cluster_rank, + int leader_rank, bool is_single_cta, bool is_idle_rank, const chunk_partition& part, offset_t my_resident_chunks, offset_t resident_base, - offset_t segment_size_u32, + offset_t segment_size_off, offset_t head_items, state_t* leader_state, int head_edge_len_items, @@ -2440,7 +2443,7 @@ private: for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { const offset_t chunk_idx = part.global_index(resident_base + local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_u32, head_items); + const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); @@ -2567,12 +2570,8 @@ private: } template - _CCCL_DEVICE _CCCL_FORCEINLINE void - run(num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int cluster_blocks, - segment_size_val_t segment_size, - out_offset_t k) + _CCCL_DEVICE _CCCL_FORCEINLINE void run( + num_segments_val_t segment_id, int cluster_rank, int cluster_blocks, segment_size_val_t segment_size, out_offset_t k) { using identify_candidates_op_t = detail::topk::identify_candidates_op_t; @@ -2592,7 +2591,7 @@ private: *this, layout.block_keys_in, layout.block_keys_base, - layout.segment_size_u32, + layout.segment_size_off, layout.head_items, layout.part, layout.my_resident_chunks, @@ -2637,7 +2636,7 @@ private: layout.my_resident_chunks, layout.resident_base, layout.chunks, - layout.segment_size_u32, + layout.segment_size_off, layout.head_items, layout.block_keys_in, layout.block_keys_base, @@ -2655,7 +2654,7 @@ private: layout.part, layout.my_resident_chunks, layout.resident_base, - layout.segment_size_u32, + layout.segment_size_off, layout.head_items, layout.leader_state, layout.head_edge_len_items, @@ -2697,7 +2696,7 @@ private: layout.resident_base, layout.my_resident_chunks, layout.overflow_chunks, - layout.segment_size_u32, + layout.segment_size_off, layout.head_items, layout.leader_state, resident_smem32, @@ -2723,7 +2722,7 @@ private: layout.part, layout.resident_base, layout.my_resident_chunks, - layout.segment_size_u32, + layout.segment_size_off, layout.head_items, resident_smem32, resident_keys, @@ -2740,15 +2739,13 @@ private: // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all( - num_segments_val_t segment_id, - segment_size_val_t segment_size, - unsigned int cluster_rank, - unsigned int cluster_blocks) + num_segments_val_t segment_id, segment_size_val_t segment_size, int cluster_rank, int cluster_blocks) { - constexpr int copy_items = copy_items_per_thread_clamped; - const offset_t num_items = static_cast(segment_size); - const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; - const offset_t cluster_threads = cluster_blocks * static_cast(threads_per_block); + constexpr int copy_items = copy_items_per_thread_clamped; + const offset_t num_items = static_cast(segment_size); + const offset_t cluster_tid = + static_cast(cluster_rank) * static_cast(threads_per_block) + threadIdx.x; + const offset_t cluster_threads = static_cast(cluster_blocks) * static_cast(threads_per_block); const offset_t step = cluster_threads * static_cast(copy_items); const offset_t full_tiles = ::cuda::round_down(num_items, step); auto keys_in_it = d_key_segments_it[segment_id]; @@ -2815,11 +2812,13 @@ private: } } - // Sub-tile remainder + // Sub-tile remainder. Iterate the tail relative to zero and offset by `full_tiles`, so every index stays in + // `[0, num_items)` without relying on `full_tiles + cluster_tid` staying within `offset_t`. + const offset_t tail_items = num_items - full_tiles; _CCCL_PRAGMA_NOUNROLL() - for (offset_t idx = full_tiles + cluster_tid; idx < num_items; idx += cluster_threads) + for (offset_t local = cluster_tid; local < tail_items; local += cluster_threads) { - const auto seg_idx = static_cast(idx); + const auto seg_idx = static_cast(full_tiles + local); keys_out_it[seg_idx] = keys_in_it[seg_idx]; if constexpr (!is_keys_only) { @@ -2830,19 +2829,21 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() { - // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). - const unsigned int cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); + // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). Cast the + // unsigned sregs to `int` once here so all downstream rank/block indexing stays signed (both always fit `int`). + const int cluster_rank = static_cast(::cuda::ptx::get_sreg_cluster_ctarank()); // Runtime cluster blocks match the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). - const unsigned int cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); - const auto segment_id = static_cast(blockIdx.x / cluster_blocks); + const int cluster_blocks = static_cast(::cuda::ptx::get_sreg_cluster_nctarank()); + const auto segment_id = static_cast(static_cast(blockIdx.x) / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) { return; } - const auto segment_size = static_cast(detail::params::get_param(segment_sizes, segment_id)); + const auto segment_size = + static_cast(detail::params::get_segment_size(segment_sizes, segment_id)); // Clamp the requested `k` to the segment size in a 64-bit width holding both operands, *before* narrowing to // `out_offset_t`. Clamping in a narrower type first would let a large "select all" `k` wrap to a small value and // silently truncate the output (or wrap to 0 and skip the segment). @@ -2855,19 +2856,9 @@ private: return; } - // Segments larger than the resident cluster_tile capacity are handled by re-streaming the overflow chunks from - // gmem (see `overflow_streamer`), so the only hard limit left is the 32-bit offset range used internally. - bool segment_fits_offset = true; - if constexpr (sizeof(segment_size_val_t) > sizeof(offset_t)) - { - segment_fits_offset = - segment_size <= static_cast(::cuda::std::numeric_limits::max()); - } - if (!segment_fits_offset) - { - _CCCL_ASSERT(false, "Segment exceeds the 32-bit offset range supported by cluster top-k"); - return; - } + // `segment_size` fits `offset_t` by precondition (see the `offset_t` definition above); an out-of-range runtime + // value is undefined behavior, not guarded here. Segments larger than the resident cluster_tile capacity are still + // handled -- the overflow chunks are re-streamed from gmem (see `overflow_streamer`). // `k_clamped <= segment_size`, which now fits `out_offset_t`, so this narrowing is safe. const auto k = static_cast(k_clamped); @@ -2886,8 +2877,8 @@ private: // served by rank 0 alone via the barrier-free path; the cluster's other CTAs exit immediately, freeing their SM // slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. Compiled out for // host-exact sizes, which the dispatch already sized to exact cluster blocks. - unsigned int eff_cluster_blocks = cluster_blocks; - unsigned int eff_cluster_rank = cluster_rank; + int eff_cluster_blocks = cluster_blocks; + int eff_cluster_rank = cluster_rank; if constexpr (enable_runtime_single_cta) { const bool fits_single_cta = is_single_cta_eligible( @@ -2896,15 +2887,15 @@ private: single_block_max_seg_size); if (fits_single_cta) { - if (cluster_rank != 0u) + if (cluster_rank != 0) { return; } - eff_cluster_blocks = 1u; - eff_cluster_rank = 0u; + eff_cluster_blocks = 1; + eff_cluster_rank = 0; } } - const bool is_single_cta = (eff_cluster_blocks == 1u); + const bool is_single_cta = (eff_cluster_blocks == 1); // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster diff --git a/cub/cub/detail/segmented_params.cuh b/cub/cub/detail/segmented_params.cuh index 18be7f651c7..8be71d5935f 100644 --- a/cub/cub/detail/segmented_params.cuh +++ b/cub/cub/detail/segmented_params.cuh @@ -14,8 +14,12 @@ #endif // no system header #include +#include +#include // indirectly_readable, random_access_iterator #include +#include // __cccl_is_integer_v #include +#include // cmp_less, cmp_greater_equal, cmp_less_equal #include #include @@ -23,6 +27,45 @@ CUB_NAMESPACE_BEGIN namespace detail::params { +// ===================================================================== +// Deferred handle requirements +// ===================================================================== + +// A `deferred` is read by dereferencing its handle (`*handle`, see the get_param overload below), so the handle must +// be indirectly readable (a pointer or other dereferenceable handle). Ranges/containers (span, array, ...) are +// rejected -- their bounds are ambiguous (values vs. size); use `deferred_sequence` for per-segment values. +template +inline constexpr bool __is_valid_deferred_handle_v = ::cuda::std::indirectly_readable<_Handle>; + +// A `deferred_sequence` is indexed per segment (`handle[index]`), so its handle must be a random-access iterator. +template +inline constexpr bool __is_valid_deferred_sequence_handle_v = ::cuda::std::random_access_iterator<_Handle>; + +// ===================================================================== +// Argument contract checks (debug-only) +// ===================================================================== + +// Shared debug-only contract check for the argument reads the host cannot validate at call time: a value read on the +// device in stream order from a `deferred`/`deferred_sequence` handle must lie within that argument's *effective* +// bounds -- the intersection of its static `cuda::args::bounds` and any runtime bounds, as computed by +// ::cuda::args::__lowest_/__highest_. A value outside them breaks the caller's promise and is otherwise undefined +// behavior. Compiled out when assertions are disabled (the bounds are not even computed). Gated on the argument's +// integer element type (not the read value type, which may be a proxy reference): `bool` and character types are +// excluded, matching the `cmp_*` comparators. +template +_CCCL_HOST_DEVICE constexpr void +__assert_param_in_bounds([[maybe_unused]] const _Arg& __arg, [[maybe_unused]] const _Value& __value) noexcept +{ + using __element_t = typename ::cuda::args::__traits<_Arg>::element_type; + if constexpr (::cuda::std::__cccl_is_integer_v<__element_t>) + { + const __element_t __checked = static_cast<__element_t>(__value); + _CCCL_ASSERT(::cuda::std::cmp_greater_equal(__checked, ::cuda::args::__lowest_(__arg)) + && ::cuda::std::cmp_less_equal(__checked, ::cuda::args::__highest_(__arg)), + "cub argument value is outside its declared bounds"); + } +} + // ===================================================================== // get_param — unified segment parameter access // ===================================================================== @@ -64,16 +107,40 @@ template [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(const ::cuda::args::deferred<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept { - // A single `deferred` wraps a device-accessible 1-element handle (pointer / iterator / 1-element span), not a value; - // the value is read on the device via element 0 (mirroring the `deferred_sequence` per-index read below). - return ::cuda::args::__unwrap(__arg)[0]; + // A single `deferred` wraps a handle to a device-side value (a pointer or input iterator), not the value itself; + // the value is read on the device by dereferencing the handle. + auto __value = *::cuda::args::__unwrap(__arg); + __assert_param_in_bounds(__arg, __value); + return __value; } template [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(const ::cuda::args::deferred_sequence<_Arg, _StaticBounds>& __arg, _SegmentIndexT __index) noexcept { - return ::cuda::args::__unwrap(__arg)[__index]; + auto __value = ::cuda::args::__unwrap(__arg)[__index]; + __assert_param_in_bounds(__arg, __value); + return __value; +} + +// Reads the size of segment `__index` as a non-negative count. A signed argument type with a negative static lower +// bound (e.g. an un-annotated `int16_t`, or an explicit `bounds` with a negative lower end) can legitimately produce a +// negative value; it is clamped up to 0 so a negative runtime size becomes an empty segment. A statically non-negative +// lower bound is trusted and left unclamped. For the deferred forms (whose values the host cannot see), `get_param` +// already checks the value against the argument's effective bounds in debug builds, so a value below a non-negative +// lower bound is caught there. +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_segment_size(const _Arg& __arg, _SegmentIndexT __index) noexcept +{ + auto __size = get_param(__arg, __index); + if constexpr (::cuda::std::cmp_less(::cuda::args::__traits<_Arg>::lowest, 0)) + { + return (::cuda::std::max) (__size, static_cast(0)); + } + else + { + return __size; + } } // ===================================================================== diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 3d8fa6f773c..49796a33b81 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -3,7 +3,7 @@ //! @file //! cub::DeviceBatchedTopK provides device-wide, parallel operations for finding the K largest (or smallest) items -//! from many (small) segments of unordered data items residing within device-accessible memory. +//! from many segments of unordered data items residing within device-accessible memory. #pragma once @@ -18,6 +18,7 @@ #endif // no system header #include +#include // detail::params::__is_valid_deferred_handle_v #include #include // topk::select::{min, max} #include @@ -31,9 +32,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -160,44 +163,118 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " "cuda::args::immediate), not a per-segment sequence."); + // The segment_sizes checks below are layered so that a single misuse reports a single, targeted diagnostic: each + // check is guarded by the validity of the ones before it (a later check "passes" once an earlier one has failed), + // so e.g. a raw pointer trips only the form check, not also the range check, and a scalar handed to `deferred` + // trips only the handle check, not also the negative-range check. + using seg_traits = ::cuda::args::__traits; + constexpr bool segment_sizes_is_form = + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v; + // The per-segment size is a count of items, so the type actually read as a size must be integral and not `bool`. A + // deferred handle or a sequence is read element-wise (dereferenced / indexed), so its `element_type` is the size; a + // plain value, `constant`, or `immediate` is used directly, so its `value_type` is -- using `value_type` there also + // rejects a handle wrongly wrapped as a single value (e.g. a pointer in `immediate`, whose `element_type` would hide + // the pointer behind its pointee). + using segment_size_value_t = + ::cuda::std::conditional_t; + constexpr bool segment_size_is_integral = + ::cuda::std::is_integral_v && !::cuda::std::is_same_v; + // A deferred / deferred_sequence argument wraps a handle read on the device (see cub::detail::params). The handle + // must match how it is read: a single `deferred` is dereferenced (`*handle`), so it must be a pointer or input + // iterator; a `deferred_sequence` is indexed per segment (`handle[index]`), so it must be a random-access iterator. + constexpr bool segment_sizes_is_deferred_single = seg_traits::is_deferred && seg_traits::is_single_value; + constexpr bool segment_sizes_is_deferred_seq = seg_traits::is_deferred && !seg_traits::is_single_value; + // Whether the deferred handle(s) are valid. Consulted only to gate the integral range block below; the per-check + // `static_assert`s below emit the actual per-misuse diagnostics. + constexpr bool segment_sizes_handle_ok = + (!segment_sizes_is_deferred_single || detail::params::__is_valid_deferred_handle_v) + && (!segment_sizes_is_deferred_seq + || detail::params::__is_valid_deferred_sequence_handle_v); + + static_assert(segment_sizes_is_form, + "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " + "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " + "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " + "cuda::args::deferred."); + static_assert(!segment_sizes_is_form || segment_size_is_integral, + "cub::DeviceBatchedTopK: segment_sizes must have an integral (non-bool) element type (the per-segment " + "size is a count of items)."); static_assert( - ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, - "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " - "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " - "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " - "cuda::args::deferred."); - static_assert(::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + !segment_sizes_is_form || !segment_size_is_integral || !segment_sizes_is_deferred_single + || detail::params::__is_valid_deferred_handle_v, + "cub::DeviceBatchedTopK: a cuda::args::deferred segment_sizes argument must wrap a pointer or other " + "dereferenceable handle to a single device-side integral segment size (it is read via *handle); a range " + "or container such as span / array is not accepted -- use cuda::args::deferred_sequence for per-segment " + "sizes."); + static_assert(!segment_sizes_is_form || !segment_size_is_integral || !segment_sizes_is_deferred_seq + || detail::params::__is_valid_deferred_sequence_handle_v, + "cub::DeviceBatchedTopK: a cuda::args::deferred_sequence segment_sizes argument must wrap a " + "random-access iterator (a pointer qualifies) over the per-segment integral segment sizes (it is " + "indexed per segment)."); + // Only the statically-known *maximum* segment size is constrained: it must not exceed 2^21 (about 2 million). + // Beyond that the streaming cluster backend is not competitive; larger segments are future work (a WIP multi-CTA + // baseline backend). So a segment-size type or bound whose maximum exceeds 2^21 (e.g. an un-annotated int32/uint32, + // or an explicit bound above 2^21) must carry a tighter compile-time `cuda::args::bounds`. The minimum is left + // unconstrained: a negative statically-known lower bound (a signed type with no non-negative bound, e.g. int16, or + // an explicit negative bound) is accepted, and the kernel clamps any negative runtime size up to 0 (see + // detail::params::get_segment_size). Gated on the checks above so an invalid form/element/handle reports just its + // own diagnostic. + if constexpr (segment_sizes_is_form && segment_size_is_integral && segment_sizes_handle_ok) + { + static_assert( + ::cuda::std::cmp_less_equal(seg_traits::highest, ::cuda::std::int64_t{1} << 21), + "cub::DeviceBatchedTopK: the statically-known maximum segment size exceeds the maximum currently supported " + "segment size (2^21, about 2 million). Give a segment-size type whose maximum is larger a compile-time upper " + "bound (cuda::args::bounds) not exceeding 2^21."); + } + constexpr bool k_is_form = ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v; + static_assert(k_is_form, "cub::DeviceBatchedTopK: k must be a cuda::args annotation or a plain integral value (taken as a " "uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap a per-segment k " "in cuda::args::deferred_sequence, or a single device-side value in cuda::args::deferred."); - static_assert( - ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, - "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " - "raw pointer or iterator is not accepted."); + constexpr bool num_segments_is_form = + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v; + static_assert(num_segments_is_form, + "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " + "raw pointer or iterator is not accepted."); - const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + // Only instantiate the dispatch once every argument passes its form/element checks above. An invalid argument has + // already failed a static_assert, so the else branch is unreachable in a well-formed program; guarding here keeps + // the diagnostic limited to that one targeted static_assert instead of a cascade of follow-on template errors from + // instantiating the device-side reads with an invalid argument (e.g. dereferencing a scalar `deferred` handle). + if constexpr (segment_sizes_is_form && segment_size_is_integral && segment_sizes_handle_ok && k_is_form + && num_segments_is_form) + { + const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); - // The total-number-of-items guarantee is intentionally not part of the initial public API surface. The dispatch - // only uses its element type to size internal large-segment offsets (the value itself is unused), so we pass a - // conservative 64-bit upper bound here. - constexpr auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; + // The total-number-of-items guarantee is intentionally not part of the initial public API surface. The dispatch + // only uses its element type to size internal large-segment offsets (the value itself is unused), so we pass a + // conservative 64-bit upper bound here. + constexpr auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; - return batched_topk::dispatch( - d_temp_storage, - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k, - ::cuda::args::constant{}, - num_segments, - total_num_items, - stream.get()); + return batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + ::cuda::args::constant{}, + num_segments, + total_num_items, + stream.get()); + } + else + { + return cudaErrorInvalidValue; + } } //! @endcond } // namespace detail @@ -236,7 +313,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! A plain integral value works too and is taken as a uniform ``immediate`` (no extra bounds). A pointer or iterator, //! by contrast, must be wrapped explicitly in ``deferred`` (single value) or ``deferred_sequence`` (per segment). //! Passing a raw pointer or iterator is rejected at compile time, because it would otherwise be misread as a single -//! value rather than a sequence. +//! value rather than a sequence. A plain integral (no bound) is still subject to each parameter's constraints: for +//! ``segment_sizes`` in particular, a plain signed integral such as ``int`` is rejected because its maximum exceeds +//! the supported maximum segment size of ``2^21`` (see *Which form each parameter accepts* and *Current constraints* +//! below). //! //! **How it is bounded.** A bound lets the algorithm reason about a value it does not know exactly: //! @@ -249,9 +329,12 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! //! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. ``num_segments`` must be //! a single value (``constant``, ``immediate``, or a plain integral), never a per-segment sequence. ``segment_sizes`` -//! must also carry a compile-time upper bound (a ``constant`` or ``cuda::args::bounds()``); the permitted -//! maximum is architecture-dependent (see *Current constraints* below), and tight bounds on every parameter are -//! encouraged. +//! must have a statically-known *maximum* not exceeding the supported ``2^21`` (about 2 million; see *Current +//! constraints* below): a type whose maximum already fits (a narrow type such as ``uint8_t``, ``int16_t``, or +//! ``uint16_t``) is accepted without an explicit bound, while a type whose maximum exceeds ``2^21`` (e.g. ``int32_t``, +//! ``uint32_t``, or ``int64_t``) must carry a compile-time upper bound (a ``constant`` or +//! ``cuda::args::bounds()``). A negative lower bound is allowed -- negative runtime sizes are clamped to an +//! empty segment (see *Current constraints*). Tight bounds on every parameter are encouraged. //! //! .. code-block:: c++ //! @@ -285,7 +368,16 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! upper bound of the ``segment_sizes`` annotation) must be small enough that such a block fits within the //! shared-memory limit. On Hopper and newer GPUs (compute capability >= 9.0) the thread-block-cluster backend also //! handles larger segments that exceed this per-block limit. Both uniform (fixed) and variable segment sizes are -//! supported. +//! supported. Independent of the architecture -- and independent of the integer type used for the ``segment_sizes`` +//! argument (a wider type such as ``int64_t`` does not raise it) -- an individual segment is currently limited to a +//! maximum of ``2^21`` (about 2 million) items, enforced at compile time from the statically-known maximum segment +//! size. Larger segments are future work. A type whose maximum already lies within it (a narrow type such as +//! ``uint8_t``, ``int16_t``, or ``uint16_t``) is accepted un-annotated; a type whose maximum exceeds ``2^21`` (e.g. +//! ``int32_t``, ``uint32_t``, or ``int64_t``) must carry a compile-time ``cuda::args::bounds`` whose upper end does +//! not exceed ``2^21``. A negative statically-known lower bound is accepted: a negative runtime size is treated as an +//! empty segment (clamped to 0). A non-negative lower bound is trusted -- an actual negative value there, like any +//! per-segment value outside its declared bound, is a caller precondition violation (undefined behavior; only +//! host-known values are checked, by assertions in debug builds). //! - **Uniform number of segments.** ``num_segments`` must be a single value, never a per-segment sequence. //! - **Unsorted output required.** Only ``cuda::execution::output_ordering::unsorted`` is implemented; the sorted //! orderings of the default contract described in *Determinism, tie-breaking, and output ordering* below (and hence @@ -397,8 +489,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -493,8 +588,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -595,8 +693,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -689,8 +790,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -808,8 +912,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -914,8 +1021,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -1025,8 +1135,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! @@ -1131,8 +1244,11 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a compile-time upper bound; the - //! permitted maximum is architecture-dependent (see the *Current constraints* section). + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Its statically-known maximum must not exceed the + //! supported `2^21` (about 2 million; see the *Current constraints* section): a type whose maximum already + //! fits (a narrow type such as `int16_t` or `uint16_t`) needs no explicit bound, while a type whose maximum exceeds + //! 2^21 (e.g. `int32_t`, `uint32_t`, `int64_t`) must carry a compile-time upper bound. A negative lower bound is + //! allowed; negative runtime sizes are clamped to an empty segment. //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 2ca2217ea7f..75cb8ad6367 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -109,7 +109,7 @@ struct segment_size_to_tile_count_op _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr TotalNumItemsValueType operator()(SegmentIndexT segment_id) const { return static_cast( - ::cuda::ceil_div(params::get_param(segment_sizes, segment_id), large_segment_agent_tile_size)); + ::cuda::ceil_div(params::get_segment_size(segment_sizes, segment_id), large_segment_agent_tile_size)); } }; @@ -559,8 +559,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return cudaSuccess; } - // A zero bound would drive `clusterDim.x = 0`, which the runtime rejects. - if (max_seg_size == 0) + // No work to launch when the tightest known upper bound on any segment size (static or runtime) is non-positive: + // every segment is empty (e.g. a uniform negative size, which the kernel would clamp to 0). Also avoids + // `clusterDim.x = 0` and a negative maximum feeding the unsigned sizing math below. + if (max_seg_size <= 0) { return cudaSuccess; } @@ -704,10 +706,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. - const int desired_cluster_blocks = static_cast(batched_topk_cluster::effective_cluster_blocks_from_chunks( - ::cuda::ceil_div(seg, chunk_items_u64), - MinChunksPerBlock, - static_cast(max_supported_cluster_blocks))); + const int desired_cluster_blocks = batched_topk_cluster::effective_cluster_blocks_from_chunks( + ::cuda::ceil_div(seg, chunk_items_u64), MinChunksPerBlock, max_supported_cluster_blocks); int cluster_blocks = 0; int dynamic_smem_sel = 0; @@ -1009,6 +1009,13 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_baseline_arm( return cudaSuccess; } + // No work to launch when the tightest known upper bound on any segment size is non-positive (mirrors the cluster + // arm): every segment is empty (e.g. a uniform host-known negative size, which the agent clamps to 0 device-side). + if (runtime_max_segment_size(segment_sizes) <= 0) + { + return cudaSuccess; + } + static_assert(::cuda::args::__traits::is_single_value, "Only a uniform number of segments is currently supported."); @@ -1195,6 +1202,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( "diagnosis to runtime (cudaErrorNotSupported)."); #endif // strict unsupported-arch check + // The supported maximum segment size (2^21) is enforced at compile time at the public entry; a statically negative + // lower bound is allowed and negative runtime sizes are clamped to 0 (see detail::params::get_segment_size). A + // per-segment value outside its declared bound is a caller error (asserted for host-known values, otherwise UB). + detail::TripleChevronFactory launcher_factory{}; ::cuda::compute_capability cc{}; if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index d54c3b79963..58ad6fa903c 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -273,50 +273,3 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); REQUIRE(keys_out == expected_result_set); } - -// The temporary storage size requirement must not assume a particular base-pointer alignment (the public contract -// states that no special alignment is required). Over-allocate by one byte and offset the base pointer. -C2H_TEST("cub::DeviceBatchedTopK::MaxKeys handles a misaligned temporary storage pointer", "[batched_topk][device]") -{ - constexpr int num_segments = 2; - constexpr int segment_size = 8; - constexpr int k = 3; - - auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; - auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); - - auto d_keys_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); - auto d_keys_out = - cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); - - constexpr auto segment_sizes = cuda::args::constant{}; - constexpr auto k_arg = cuda::args::constant{}; - auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{cuda::execution::require( - cuda::execution::determinism::not_guaranteed, - cuda::execution::tie_break::unspecified, - cuda::execution::output_ordering::unsorted)}; - - size_t temp_storage_bytes = 0; - auto error = cub::DeviceBatchedTopK::MaxKeys( - nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); - REQUIRE(error == cudaSuccess); - - // Allocate one extra byte and offset the base pointer by one to misalign it. - thrust::device_vector temp_storage(temp_storage_bytes + 1, thrust::no_init); - error = cub::DeviceBatchedTopK::MaxKeys( - thrust::raw_pointer_cast(temp_storage.data()) + 1, - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k_arg, - num_segs, - env); - REQUIRE(error == cudaSuccess); - - thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); - thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); - REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); -} diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 7a930225962..551d9873304 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -14,8 +14,10 @@ #include #include #include +#include #include #include +#include #include #include @@ -25,6 +27,8 @@ #include #include #include +#include +#include #include "catch2_test_device_topk_common.cuh" #include "catch2_test_launch_helper.h" @@ -137,38 +141,6 @@ using key_types = using select_direction_list = c2h::enum_type_list; -// Segment-size argument form for the single-value (uniform) fixed-size tests. The host-value forms (`immediate` and an -// un-annotated raw scalar, which auto-wraps as an `immediate` with loose bounds) are distributed across the TEST_TYPES -// variants so each variant compiles only one extra dispatch instantiation (balancing build time). The un-annotated form -// reports a `numeric_limits::max()` upper bound, so the types_2 variant also exercises the loose-bound -// oversize/streaming fallback (and its int-narrowing overflow guard) on the cluster backend. The remaining public forms -// get their own dedicated tests below (`constant` needs a compile-time size; `deferred` needs a device-accessible -// handle), and `deferred_sequence` is covered by the variable-size tests. -enum class seg_size_arg -{ - immediate_form, - unannotated_form, -}; - -template -auto make_segment_size_arg(SizeT segment_size) -{ - if constexpr (Form == seg_size_arg::immediate_form) - { - return cuda::args::immediate{segment_size, cuda::args::bounds()}; - } - else - { - return segment_size; // un-annotated: auto-wrapped as an immediate with loose bounds - } -} - -#if TEST_TYPES == 2 -inline constexpr seg_size_arg fixed_seg_size_arg = seg_size_arg::unannotated_form; -#else -inline constexpr seg_size_arg fixed_seg_size_arg = seg_size_arg::immediate_form; -#endif - // A (determinism, tie-break) requirement pair, used as a single compile-time test axis. Not a full cross product: a // tie-break preference is only meaningful with a deterministic requirement and is `gpu_to_gpu` by definition. template @@ -256,12 +228,10 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", // Copy input for verification c2h::device_vector expected_keys(keys_in_buffer); - // Run the top-k algorithm. The segment-size argument form (immediate / un-annotated) is selected per TEST_TYPES so - // the suite covers each form without any single variant compiling all of them. batched_topk_keys( d_keys_in, d_keys_out, - make_segment_size_arg(segment_size), + cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, cuda::args::immediate{num_segments}); // Prepare expected results @@ -623,8 +593,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned } // Streaming counterpart of the narrow-segment-size regression. Streaming needs segments larger than the resident -// cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- same width as the -// (unsigned) internal `offset_t` but signed -- to exercise the streaming path's index arithmetic across det/non-det. +// cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- the same signed +// `int32_t` as the internal `offset_t` -- to exercise the streaming path's index arithmetic across det/non-det. C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 32-bit segment-size type", "[keys][segmented][topk][device][cluster][determinism]", det_tie_combos) @@ -1298,3 +1268,241 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys accept unwrapped (plain integral) k a REQUIRE(expected_keys == keys_out_buffer); } + +// The following tests exercise public-API behavior that is independent of key type and launch mechanism, so they call +// the API directly (not through the templated launch wrapper) and are compiled once for the whole param matrix. +#if TEST_TYPES == 0 && TEST_LAUNCH == 0 +// The temporary storage size requirement must not assume a particular base-pointer alignment (the public contract +// states that no special alignment is required). Over-allocate by one byte and offset the base pointer. +C2H_TEST("DeviceBatchedTopK::MaxKeys handles a misaligned temporary storage pointer", "[keys][segmented][topk][device]") +{ + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + // Allocate one extra byte and offset the base pointer by one to misalign it. + thrust::device_vector temp_storage(temp_storage_bytes + 1, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()) + 1, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); +} + +// The supported maximum segment size (2^21, about 2 million) is enforced at compile time from the statically-known +// maximum. An un-annotated segment size is accepted only when its element type's maximum already fits: a narrow type +// (e.g. uint16, [0, 65535]) qualifies with no cuda::args::bounds, while a type whose max exceeds 2^21 (int32, uint32, +// int64) needs one. Verify the un-annotated narrow path compiles and runs end-to-end. +C2H_TEST("DeviceBatchedTopK::MaxKeys accepts an un-annotated narrow-unsigned segment size", + "[keys][segmented][topk][device]") +{ + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + // Un-annotated narrow-unsigned segment size: a bare integral value (no cuda::args wrapper), taken as a uniform + // immediate. With no wrapper/bounds the static maximum is the type maximum (uint16 -> 65535), which must itself fall + // not exceed 2^21 to be accepted; the static_assert below pins that exposed bound. Any narrow-unsigned type + // qualifies the same way (e.g. uint8, [0, 255]); uint16 is representative. + auto segment_sizes = cuda::std::uint16_t{segment_size}; + static_assert(cuda::args::__traits::highest == 65535, + "expected the un-annotated uint16 segment size to expose its full type range as the static bound"); + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); +} + +// A negative statically-known lower bound (here an explicit `bounds<-8, ...>`, matching e.g. a bare `int16_t`) is +// accepted: the kernel clamps a negative runtime segment size up to 0, so that segment is treated as empty (skipped, +// no output) instead of indexing with a negative count. A non-negative lower bound is instead trusted (no clamp). +// Exercise the clamp with a mixed batch -- segment 0 declares a negative size, segment 1 a normal one. Small size + +// no determinism requirement keeps this on the baseline backend (the cluster backend is covered separately below). +C2H_TEST("DeviceBatchedTopK::MaxKeys clamps a negative segment size to an empty segment (baseline backend)", + "[keys][segmented][topk][device]") +{ + using seg_size_t = cuda::std::int16_t; // negative-capable lower bound -> clamp path + constexpr int num_segments = 2; + constexpr int k = 3; + constexpr int stride = 8; + constexpr int sentinel = -12345; + + // Segment 0: declared size -1 -> clamped to 0 -> skipped. Segment 1: 8 real keys, top-3 max = {9, 8, 7}. + auto d_segment_sizes = thrust::device_vector{seg_size_t{-1}, seg_size_t{stride}}; + auto keys_in = thrust::device_vector{0, 0, 0, 0, 0, 0, 0, 0, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, sentinel); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), stride); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + auto segment_sizes = + cuda::args::deferred_sequence{thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<-8, 100>()}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + + // Segment 0's output slots are left untouched (still the sentinel); segment 1 holds its top-3. + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{sentinel, sentinel, sentinel, 9, 8, 7}); +} + +// Cluster-backend counterpart of the clamp test above: a deterministic requirement forces the SM90+ cluster backend +// even for small segments, so this pins the cluster agent's own empty-segment early-out on a clamped negative size. +C2H_TEST("DeviceBatchedTopK::MaxKeys clamps a negative segment size to an empty segment (cluster backend)", + "[keys][segmented][topk][device][cluster][determinism]") +{ + using seg_size_t = cuda::std::int16_t; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__prefer_smaller_index; + constexpr int num_segments = 2; + constexpr int k = 3; + constexpr int stride = 8; + constexpr int sentinel = -12345; + constexpr cuda::std::int64_t static_max_segment_size = 100; + + skip_if_batched_topk_backend_unavailable(static_max_segment_size); + + // Segment 0: declared size -1 -> clamped to 0 -> skipped. Segment 1: 8 real keys, top-3 max = {9, 8, 7}. + auto d_segment_sizes = thrust::device_vector{seg_size_t{-1}, seg_size_t{stride}}; + auto keys_in = thrust::device_vector{0, 0, 0, 0, 0, 0, 0, 0, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, sentinel); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), stride); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + auto segment_sizes = + cuda::args::deferred_sequence{thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<-8, 100>()}; + + batched_topk_keys( + d_keys_in, + d_keys_out, + segment_sizes, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}); + + // Segment 0's output slots are left untouched (still the sentinel); segment 1 holds its top-3. + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{sentinel, sentinel, sentinel, 9, 8, 7}); +} + +// A uniform (host-known) negative segment size means every segment is empty. Routed to the cluster dispatch by a +// deterministic requirement, this must be recognized on the host from the non-positive upper bound and skipped without +// launching (the `max_seg_size <= 0` guard), leaving the output untouched -- rather than casting the negative maximum +// to unsigned and sizing an enormous launch. +C2H_TEST("DeviceBatchedTopK::MaxKeys treats a uniform negative segment size as no work (cluster backend)", + "[keys][segmented][topk][device][cluster][determinism]") +{ + using seg_size_t = cuda::std::int16_t; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__prefer_smaller_index; + constexpr int num_segments = 2; + constexpr int k = 3; + constexpr int stride = 8; + constexpr int sentinel = -12345; + constexpr cuda::std::int64_t static_max_segment_size = 100; + + skip_if_batched_topk_backend_unavailable(static_max_segment_size); + + auto keys_in = thrust::device_vector{0, 9, 3, 2, 1, 8, 7, 4, /**/ 5, 6, 1, 0, 3, 2, 8, 7}; + auto keys_out = thrust::device_vector(num_segments * k, sentinel); + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), stride); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{seg_size_t{-1}, cuda::args::bounds<-8, 100>()}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}); + + // No segment had any items, so the whole output is left untouched. + REQUIRE(keys_out == thrust::device_vector(num_segments * k, sentinel)); +} +#endif // TEST_TYPES == 0 && TEST_LAUNCH == 0 diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 926fd45adbb..58b6358c460 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -698,9 +698,9 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with large fixed-size unaligned } // Streaming counterpart of the narrow-segment-size regression (pairs). Streaming needs segments larger than the -// resident cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- same width as -// the (unsigned) internal `offset_t` but signed -- to exercise the streaming path's index arithmetic (including the -// value gather) across det/non-det. +// resident cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use the same signed 32-bit type as +// the internal `offset_t` to exercise the streaming path's index arithmetic (including the value gather) across +// det/non-det. C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream large segments with a signed 32-bit segment-size type", "[pairs][segmented][topk][device][cluster][determinism]", det_tie_pair_combos) diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index c2dfa4b3e2e..b802d7ea14d 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -50,7 +50,8 @@ inline void skip_if_batched_topk_cluster_unavailable(bool needs_cluster) // the cluster backend if it is deterministic / has a concrete tie-break, or if `static_max_segment_size` exceeds the // baseline backend's coverage; skips if that backend is unavailable. Deriving the size decision here (rather than a // precomputed `oversize` bool) keeps the threshold in one place. Pass the same maximum segment size the test hands to -// the dispatch (its `cuda::args::bounds<...>` upper bound, or the type's maximum when unbounded). Note that `oversize` +// the dispatch: its `cuda::args::bounds<...>` upper bound (or, for an un-annotated narrow type, that type's maximum; +// a type whose maximum exceeds 2^21 no longer compiles without a bound). Note that `oversize` // uses only the tile-size bound `baseline_max_covered_segment_size`; the dispatch's `baseline_can_cover_v` additionally // checks the agent's shared-memory fit, so a borderline size the bound deems baseline-coverable could still route to // the cluster backend (such a case would fail rather than skip if the cluster backend is unavailable). diff --git a/cub/test/test_device_batched_topk_requirements_fail.cu b/cub/test/test_device_batched_topk_requirements_fail.cu index 55d234c2f7b..b6109c8c893 100644 --- a/cub/test/test_device_batched_topk_requirements_fail.cu +++ b/cub/test/test_device_batched_topk_requirements_fail.cu @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// %PARAM% TEST_ERR err 0:1:2:3:4:5 +// %PARAM% TEST_ERR err 0:1:2:3:4:5:6:7:8:9:10:11:12:13 // Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) // so only the requirement static_asserts under test fire, regardless of the target architecture (this test compiles for @@ -18,14 +18,24 @@ #include #include #include +#include #include -// Verifies that cub::DeviceBatchedTopK rejects, at compile time, the determinism and tie_break requirement -// combinations that the public contract marks as ill-formed (see docs/cub/api_docs/device_topk_requirements.rst): +// Verifies that cub::DeviceBatchedTopK rejects, at compile time, requests the public contract marks as ill-formed: // * determinism and tie_break must be acknowledged together, both specified or both omitted to take the default // * an explicit tie_break of prefer_smaller_index or prefer_larger_index pins the result set across GPUs and so // requires determinism::gpu_to_gpu (it cannot be paired with not_guaranteed or run_to_run) +// * an explicit segment_sizes compile-time upper bound above the maximum supported segment size (2^21) +// * an un-annotated segment_sizes argument whose element type max exceeds 2^21 (a compile-time bound is required): +// int64, uint32, and int32 all need one; only a type whose max already fits (e.g. int16, uint16) is accepted bare +// * a deferred segment_sizes argument wrapping a scalar rather than a dereferenceable pointer / handle +// * a deferred segment_sizes argument wrapping a range / container (even a static-extent-1 span): ranges are not +// accepted, a deferred handle must be a dereferenceable pointer / handle (sizes belong in deferred_sequence) +// * a deferred_sequence segment_sizes argument wrapping a non-random-access handle (a span is a range, not an +// iterator); a deferred_sequence handle is indexed per segment, so it must be a random-access iterator +// * a single-value wrapper (immediate) around a pointer / handle rather than an integral value: its value_type is +// not integral, so it is rejected before the pointer could be misread as a size // Each variant exercises one rejected cell of that contract table. int main() @@ -33,34 +43,71 @@ int main() namespace ex = cuda::execution; // Per-segment key iterators (iterator-of-iterators). The keys-only entry point ignores value iterators. - int** d_keys_in = nullptr; - int** d_keys_out = nullptr; - auto segment_sizes = cuda::args::constant<8>{}; - auto k_arg = cuda::args::constant<3>{}; - auto num_segments = cuda::args::immediate{cuda::std::int64_t{2}}; + int** d_keys_in = nullptr; + int** d_keys_out = nullptr; + auto k_arg = cuda::args::constant<3>{}; + auto num_segments = cuda::args::immediate{cuda::std::int64_t{2}}; -#if TEST_ERR == 0 // determinism specified without a paired tie_break +#if TEST_ERR == 6 // explicit segment-size upper bound above 2^21 (unsupported by the streaming cluster backend) + auto segment_sizes = cuda::args::immediate{ + cuda::std::int64_t{8}, cuda::args::bounds()}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-6 {{"exceeds the maximum currently supported segment size"}} +#elif TEST_ERR == 7 // un-annotated wide segment-size type (int64 type max > 2^21): a compile-time bound is required + auto segment_sizes = cuda::args::immediate{cuda::std::int64_t{8}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-7 {{"exceeds the maximum currently supported segment size"}} +#elif TEST_ERR == 8 // un-annotated unsigned segment-size type (uint32 type max > 2^21): a compile-time bound too + auto segment_sizes = cuda::args::immediate{cuda::std::uint32_t{8}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-8 {{"exceeds the maximum currently supported segment size"}} +#elif TEST_ERR == 9 // un-annotated int32 segment-size type (type max 2^31 - 1 > 2^21): a compile-time bound is required + auto segment_sizes = cuda::args::immediate{cuda::std::int32_t{8}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-9 {{"exceeds the maximum currently supported segment size"}} +#elif TEST_ERR == 10 // deferred wrapping a scalar (not a dereferenceable pointer/handle): rejected + auto segment_sizes = cuda::args::deferred{cuda::std::int32_t{8}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-10 {{"must wrap a pointer or other dereferenceable handle"}} +#elif TEST_ERR == 11 // deferred wrapping a range/container: rejected even at static extent 1 (still a range, not a + // handle) + cuda::std::uint16_t seg_storage[1]{}; + auto segment_sizes = cuda::args::deferred{cuda::std::span{seg_storage}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-11 {{"must wrap a pointer or other dereferenceable handle"}} +#elif TEST_ERR == 12 // deferred_sequence wrapping a non-random-access handle (a span is a range, not an iterator) + auto segment_sizes = cuda::args::deferred_sequence{cuda::std::span{}}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-12 {{"must wrap a random-access iterator"}} +#elif TEST_ERR == 13 // single-value wrapper (immediate) around a pointer/handle rather than an integral value + auto segment_sizes = cuda::args::immediate{static_cast(nullptr), cuda::args::bounds<0, 8>()}; + auto requirements = ex::require(ex::output_ordering::unsorted); + // expected-error-13 {{"must have an integral \(non-bool\) element type"}} +#else + auto segment_sizes = cuda::args::constant<8>{}; +# if TEST_ERR == 0 // determinism specified without a paired tie_break auto requirements = ex::require(ex::determinism::not_guaranteed, ex::output_ordering::unsorted); // expected-error-0 {{"must be acknowledged together"}} -#elif TEST_ERR == 1 // tie_break specified without a paired determinism +# elif TEST_ERR == 1 // tie_break specified without a paired determinism auto requirements = ex::require(ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); // expected-error-1 {{"must be acknowledged together"}} -#elif TEST_ERR == 2 // explicit tie_break with not_guaranteed +# elif TEST_ERR == 2 // explicit tie_break with not_guaranteed auto requirements = ex::require(ex::determinism::not_guaranteed, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); // expected-error-2 {{"pins the result set across GPUs and therefore requires"}} -#elif TEST_ERR == 3 // explicit tie_break with not_guaranteed +# elif TEST_ERR == 3 // explicit tie_break with not_guaranteed auto requirements = ex::require(ex::determinism::not_guaranteed, ex::tie_break::prefer_larger_index, ex::output_ordering::unsorted); // expected-error-3 {{"pins the result set across GPUs and therefore requires"}} -#elif TEST_ERR == 4 // explicit tie_break with run_to_run +# elif TEST_ERR == 4 // explicit tie_break with run_to_run auto requirements = ex::require(ex::determinism::run_to_run, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); // expected-error-4 {{"pins the result set across GPUs and therefore requires"}} -#elif TEST_ERR == 5 // explicit tie_break with run_to_run +# elif TEST_ERR == 5 // explicit tie_break with run_to_run auto requirements = ex::require(ex::determinism::run_to_run, ex::tie_break::prefer_larger_index, ex::output_ordering::unsorted); // expected-error-5 {{"pins the result set across GPUs and therefore requires"}} +# endif #endif auto env = cuda::std::execution::env{requirements}; From ac18c3851f83326cdac898aae2c7a0040c1ca474 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 9 Jul 2026 16:28:30 +0200 Subject: [PATCH 088/126] Go back to unsigned offsets in the agent --- cub/cub/agent/agent_batched_topk_cluster.cuh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index ad1161178a7..c71ce726ad4 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -232,12 +232,12 @@ struct agent_batched_topk_cluster using segment_size_val_t = typename ::cuda::args::__traits::element_type; using num_segments_val_t = typename ::cuda::args::__traits::element_type; - // Signed, so the indexing arithmetic below can never wrap on underflow. 32-bit comfortably covers every supported - // segment: the public entry caps the statically-known maximum segment size at 2^21, well within this width, so a - // value that actually exceeds its declared bound at runtime is a caller precondition violation (undefined behavior). - // The cross-CTA scan also packs two of these lanes into one `uint64_t`, which only holds for 32-bit lanes. - using offset_t = ::cuda::std::int32_t; - using out_offset_t = ::cuda::std::int32_t; + // 32-bit covers every supported segment: the public entry caps the statically-known maximum segment size at 2^21, so + // a runtime value exceeding its declared bound is a caller precondition violation (undefined behavior). Unsigned + // because all offsets are non-negative (segment sizes are clamped to >= 0 upstream; ranks/blocks are `int`s cast in + // at the boundaries). The cross-CTA scan also packs two lanes into one `uint64_t`, which needs 32-bit lanes. + using offset_t = ::cuda::std::uint32_t; + using out_offset_t = ::cuda::std::uint32_t; using state_t = cluster_topk_state; using key_prefix_t = typename state_t::key_prefix_t; From 40144aaefc02f07379f1f92ed6235fa7b25d74d0 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 9 Jul 2026 17:56:50 +0200 Subject: [PATCH 089/126] Go back to unsigned integers for indices internally --- cub/cub/agent/agent_batched_topk_cluster.cuh | 115 +++++++++--------- .../device/dispatch/dispatch_batched_topk.cuh | 6 +- 2 files changed, 64 insertions(+), 57 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index c71ce726ad4..b1fd6355773 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -177,11 +177,11 @@ inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, // `min_chunks_per_block` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap // differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_block` is // `static_assert`ed positive, so the divide is well-defined. 64-bit math. -[[nodiscard]] _CCCL_HOST_DEVICE constexpr int effective_cluster_blocks_from_chunks( - ::cuda::std::uint64_t chunks, int min_chunks_per_block, int cluster_blocks_cap) noexcept +[[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( + ::cuda::std::uint64_t chunks, int min_chunks_per_block, unsigned int cluster_blocks_cap) noexcept { const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_block); - return static_cast( + return static_cast( ::cuda::std::clamp(blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } @@ -234,8 +234,8 @@ struct agent_batched_topk_cluster // 32-bit covers every supported segment: the public entry caps the statically-known maximum segment size at 2^21, so // a runtime value exceeding its declared bound is a caller precondition violation (undefined behavior). Unsigned - // because all offsets are non-negative (segment sizes are clamped to >= 0 upstream; ranks/blocks are `int`s cast in - // at the boundaries). The cross-CTA scan also packs two lanes into one `uint64_t`, which needs 32-bit lanes. + // because all offsets, ranks, and block counts are non-negative (segment sizes are clamped to >= 0 upstream). The + // cross-CTA scan also packs two lanes into one `uint64_t`, which needs 32-bit lanes. using offset_t = ::cuda::std::uint32_t; using out_offset_t = ::cuda::std::uint32_t; using state_t = cluster_topk_state; @@ -505,11 +505,10 @@ struct agent_batched_topk_cluster } [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t - num_rank_chunks(offset_t chunks, int cluster_rank, int cluster_blocks) const + num_rank_chunks(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { - return (static_cast(cluster_rank) < chunks) - ? static_cast( - (chunks - 1 - static_cast(cluster_rank)) / static_cast(cluster_blocks) + 1) + return (cluster_rank < chunks) + ? static_cast((chunks - 1 - cluster_rank) / cluster_blocks + 1) : offset_t{0}; } @@ -541,7 +540,7 @@ struct agent_batched_topk_cluster // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches // ascending contiguous global-index ranges), so it is selected exactly when `need_determinism` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition - make_chunk_partition(offset_t chunks, int cluster_rank, int cluster_blocks) const + make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { if constexpr (need_determinism) { @@ -833,7 +832,7 @@ private: // it to rank 0 to form the leader's `shared::cluster` address (no 64-bit pointer, no memory descriptor). Cluster // scope makes it mutually atomic with the leader's `hist_inc` adds. _CCCL_DEVICE _CCCL_FORCEINLINE void - hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, int leader_rank) + hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, unsigned int leader_rank) { ::cuda::std::uint32_t remote; asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "r"(leader_rank)); @@ -842,7 +841,7 @@ private: // Generic pointer to this CTA's `state` as seen in the CTA at cluster rank `rank` (reached over DSMEM) -- the PTX // equivalent of cooperative_groups' `map_shared_rank`, via `mapa.u64` (generic-address form). - _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(int rank) + _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(unsigned int rank) { const ::cuda::std::uint64_t own = reinterpret_cast<::cuda::std::uint64_t>(&temp_storage.state); ::cuda::std::uint64_t remote; @@ -853,7 +852,7 @@ private: // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. - _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(int target_rank, ::cuda::std::uint64_t v) + _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, ::cuda::std::uint64_t v) { const ::cuda::std::uint32_t own = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.prefix_pair)); @@ -1276,8 +1275,8 @@ private: // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. - _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t - combined_prefix_scan(bool is_single_cta, int cluster_rank, int eff_cluster_blocks, ::cuda::std::uint64_t packed) + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( + bool is_single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) { if (is_single_cta) { @@ -1288,9 +1287,8 @@ private: if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() - for (int rank = static_cast(threadIdx.x); rank < cluster_rank; rank += threads_per_block) // lower ranks - // follow; leader - // last + for (unsigned int rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; + // leader last { add_remote_prefix(rank, packed); } @@ -1300,8 +1298,7 @@ private: // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. _CCCL_PRAGMA_NOUNROLL() - for (int rank = cluster_rank + 1 + static_cast(threadIdx.x); rank < eff_cluster_blocks; - rank += threads_per_block) + for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < eff_cluster_blocks; rank += threads_per_block) { add_remote_prefix(rank, packed); } @@ -1723,9 +1720,9 @@ private: template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( num_segments_val_t segment_id, - int cluster_rank, - int eff_cluster_blocks, - int leader_rank, + unsigned int cluster_rank, + unsigned int eff_cluster_blocks, + unsigned int leader_rank, bool is_single_cta, bool is_idle_rank, out_offset_t k, @@ -1914,9 +1911,9 @@ private: template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( num_segments_val_t segment_id, - int cluster_rank, - int eff_cluster_blocks, - int leader_rank, + unsigned int cluster_rank, + unsigned int eff_cluster_blocks, + unsigned int leader_rank, bool is_single_cta, bool is_idle_rank, out_offset_t k, @@ -2053,8 +2050,8 @@ private: // histogram visible cluster-wide. template _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( - int cluster_rank, - int leader_rank, + unsigned int cluster_rank, + unsigned int leader_rank, bool is_single_cta, const chunk_partition& part, offset_t my_resident_chunks, @@ -2230,10 +2227,10 @@ private: const key_t* block_keys_base; offset_t head_items; offset_t chunks; - int eff_cluster_blocks; + unsigned int eff_cluster_blocks; bool is_idle_rank; chunk_partition part; - int leader_rank; + unsigned int leader_rank; state_t* leader_state; offset_t my_chunks; offset_t my_resident_chunks; @@ -2247,7 +2244,10 @@ private: }; _CCCL_DEVICE _CCCL_FORCEINLINE segment_layout compute_segment_layout( - num_segments_val_t segment_id, int cluster_rank, int cluster_blocks, segment_size_val_t segment_size) + num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int cluster_blocks, + segment_size_val_t segment_size) { segment_layout layout; layout.block_keys_in = d_key_segments_it[segment_id]; @@ -2256,7 +2256,7 @@ private: // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. - layout.is_single_cta = (cluster_blocks == 1); + layout.is_single_cta = (cluster_blocks == 1u); layout.block_keys_base = nullptr; layout.head_items = 0; @@ -2295,7 +2295,7 @@ private: // (`< eff_cluster_blocks`). The deterministic tie-break makes the leader the *last* CTA in scan order so it never // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. - layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1) : 0; + layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1u) : 0u; // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). @@ -2367,7 +2367,7 @@ private: // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. - layout.head_edge_len_items = (cluster_rank == 0) ? static_cast(layout.head_items) : 0; + layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_items) : 0; layout.tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; return layout; } @@ -2377,8 +2377,8 @@ private: // that actually ran (`last_pass`), which the final filter uses to size its identify operator. template _CCCL_DEVICE _CCCL_FORCEINLINE int run_radix_passes( - int cluster_rank, - int leader_rank, + unsigned int cluster_rank, + unsigned int leader_rank, bool is_single_cta, bool is_idle_rank, const chunk_partition& part, @@ -2570,8 +2570,12 @@ private: } template - _CCCL_DEVICE _CCCL_FORCEINLINE void run( - num_segments_val_t segment_id, int cluster_rank, int cluster_blocks, segment_size_val_t segment_size, out_offset_t k) + _CCCL_DEVICE _CCCL_FORCEINLINE void + run(num_segments_val_t segment_id, + unsigned int cluster_rank, + unsigned int cluster_blocks, + segment_size_val_t segment_size, + out_offset_t k) { using identify_candidates_op_t = detail::topk::identify_candidates_op_t; @@ -2739,13 +2743,15 @@ private: // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all( - num_segments_val_t segment_id, segment_size_val_t segment_size, int cluster_rank, int cluster_blocks) + num_segments_val_t segment_id, + segment_size_val_t segment_size, + unsigned int cluster_rank, + unsigned int cluster_blocks) { - constexpr int copy_items = copy_items_per_thread_clamped; - const offset_t num_items = static_cast(segment_size); - const offset_t cluster_tid = - static_cast(cluster_rank) * static_cast(threads_per_block) + threadIdx.x; - const offset_t cluster_threads = static_cast(cluster_blocks) * static_cast(threads_per_block); + constexpr int copy_items = copy_items_per_thread_clamped; + const offset_t num_items = static_cast(segment_size); + const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; + const offset_t cluster_threads = cluster_blocks * static_cast(threads_per_block); const offset_t step = cluster_threads * static_cast(copy_items); const offset_t full_tiles = ::cuda::round_down(num_items, step); auto keys_in_it = d_key_segments_it[segment_id]; @@ -2829,13 +2835,12 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() { - // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). Cast the - // unsigned sregs to `int` once here so all downstream rank/block indexing stays signed (both always fit `int`). - const int cluster_rank = static_cast(::cuda::ptx::get_sreg_cluster_ctarank()); + // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). + const unsigned int cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); // Runtime cluster blocks match the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). - const int cluster_blocks = static_cast(::cuda::ptx::get_sreg_cluster_nctarank()); - const auto segment_id = static_cast(static_cast(blockIdx.x) / cluster_blocks); + const unsigned int cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); + const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) { @@ -2877,8 +2882,8 @@ private: // served by rank 0 alone via the barrier-free path; the cluster's other CTAs exit immediately, freeing their SM // slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. Compiled out for // host-exact sizes, which the dispatch already sized to exact cluster blocks. - int eff_cluster_blocks = cluster_blocks; - int eff_cluster_rank = cluster_rank; + unsigned int eff_cluster_blocks = cluster_blocks; + unsigned int eff_cluster_rank = cluster_rank; if constexpr (enable_runtime_single_cta) { const bool fits_single_cta = is_single_cta_eligible( @@ -2887,15 +2892,15 @@ private: single_block_max_seg_size); if (fits_single_cta) { - if (cluster_rank != 0) + if (cluster_rank != 0u) { return; } - eff_cluster_blocks = 1; - eff_cluster_rank = 0; + eff_cluster_blocks = 1u; + eff_cluster_rank = 0u; } } - const bool is_single_cta = (eff_cluster_blocks == 1); + const bool is_single_cta = (eff_cluster_blocks == 1u); // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 75cb8ad6367..9969361b039 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -706,8 +706,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. - const int desired_cluster_blocks = batched_topk_cluster::effective_cluster_blocks_from_chunks( - ::cuda::ceil_div(seg, chunk_items_u64), MinChunksPerBlock, max_supported_cluster_blocks); + const int desired_cluster_blocks = ::cuda::narrow(batched_topk_cluster::effective_cluster_blocks_from_chunks( + ::cuda::ceil_div(seg, chunk_items_u64), + MinChunksPerBlock, + ::cuda::narrow(max_supported_cluster_blocks))); int cluster_blocks = 0; int dynamic_smem_sel = 0; From 44e2a4320ab030e5967bd66eedc7febe8f3b2f57 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 01:21:45 +0200 Subject: [PATCH 090/126] Improve latency hiding via split-barrier --- cub/cub/agent/agent_batched_topk_cluster.cuh | 110 ++++++++++++++----- 1 file changed, 82 insertions(+), 28 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index b1fd6355773..38895387486 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1242,14 +1242,26 @@ private: // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- - // Cluster-wide barrier via PTX (replaces cooperative_groups' `cluster.sync()`): `.release` on arrive, `.acquire` on - // wait, both `.aligned` since every thread reaches it under a uniform branch. - _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_barrier() + // Split halves of the cluster-wide barrier (replaces cooperative_groups' `cluster.sync()`). `arrive` releases this + // CTA's prior writes and signals arrival; `wait` acquires all CTAs' writes once every CTA has arrived. Both are + // `.aligned` since every thread reaches them under a uniform branch. Issuing independent, block-local work between + // them hides the cluster-arrival latency. + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_arrive() { asm volatile("barrier.cluster.arrive.release.aligned;" : : : "memory"); + } + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_wait() + { asm volatile("barrier.cluster.wait.acquire.aligned;" : : : "memory"); } + // Cluster-wide barrier: arrive immediately followed by wait, no work window exploited. + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_barrier() + { + cluster_arrive(); + cluster_wait(); + } + // Synchronize the segment's cluster. A single-CTA "cluster" keeps all state block-local, so `__syncthreads()` orders // it and the cluster-scoped barrier is unnecessary. `is_single_cta` is computed in `run()` from the collapsed cluster // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw cluster size. It is @@ -1266,6 +1278,29 @@ private: } } + // Split form of `cluster_or_block_sync`, for overlapping the cluster-arrival latency with independent block-local + // work: `arrive`, then the work, then `wait`. A single-CTA "cluster" has no arrival latency to hide and no cross-CTA + // state, so it takes the whole `__syncthreads()` at `arrive` and its `wait` is a no-op. In the multi-CTA case the + // work between the two must touch only block-local state (a cross-CTA read before `wait` may miss a peer's writes). + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_arrive(bool is_single_cta) + { + if (is_single_cta) + { + __syncthreads(); + } + else + { + cluster_arrive(); + } + } + _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_wait(bool is_single_cta) + { + if (!is_single_cta) + { + cluster_wait(); + } + } + // Combined 64-bit exclusive cross-CTA prefix scan over each working CTA's packed `(front_count << 32) | cand_count`. // Each CTA pushes its counts into every successor's `prefix_pair` in `is_scan_descending` order; the leader is last // and pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` @@ -2046,8 +2081,9 @@ private: // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 // needs no candidate filtering). Publishes the resident span as `resident_keys` and its 32-bit shared base address - // as `resident_smem32` for the later passes and the final filter; ends on the barrier that makes `edge_keys` and the - // histogram visible cluster-wide. + // as `resident_smem32` for the later passes and the final filter; ends on a `__syncthreads()` that makes `edge_keys` + // and the block-local histogram visible within the block. Cross-CTA visibility of the zeroed histogram comes later, + // from the deferred initial cluster wait in `run_radix_passes` (this whole load runs in that arrive->wait window). template _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( unsigned int cluster_rank, @@ -2462,12 +2498,20 @@ private: fold_boundary_edges(head_edge_len_items, tail_edge_len_items, add_hist); } - // Local barrier is enough: all Step 1 / Step 2 writes to `hist[]` - // are atomic at compatible scopes (see Step 1 dispatch). The - // cluster-wide ordering before Step 3's leader read of `hist[]` - // is supplied by the cluster barrier further below. + // Local barrier is enough here: all Step 1 / Step 2 writes to `hist[]` are atomic at compatible scopes (see Step + // 1 dispatch). The cluster-wide ordering before Step 3's leader read of the merged `hist[]` is supplied by the + // split post-fold cluster wait further below. __syncthreads(); + // First pass only: complete the deferred initial cluster barrier here. `process_impl` issued its matching + // `cluster_arrive` right after zeroing `state`/`hist`, so the cluster-arrival latency overlapped the fused + // first-pass load + histogram. The wait must precede the first cross-CTA access -- the Step-2 folds just below, + // which need the leader's `hist` visibly zeroed. + if (is_first_pass) + { + cluster_or_block_wait(is_single_cta); + } + // Step 2: non-leader blocks fold their per-bucket raw counts into // the leader's `hist` via cluster-scope DSMEM atomics (see // `hist_fold_remote`). The @@ -2488,26 +2532,21 @@ private: } } + // Split the post-fold cluster barrier: arrive once the folds are released, then overlap the non-leaders' + // own-histogram scan (block-local) with the cluster-arrival latency, and only wait before the leader reads the + // merged `hist`. // TODO(cccl): idle ranks arrive here only because the cluster barrier spans the whole launched cluster. An // mbarrier over just the active ranks would let them exit and free their SM slots instead of spinning here. - cluster_or_block_sync(is_single_cta); + cluster_or_block_arrive(is_single_cta); - // Step 3: the leader prefix-scans the merged `hist` (raw counts) and - // updates the cluster-shared `state`. Subsequent reads (end-of-pass - // fold, last filter) all observe these writes after the next cluster sync. - // - // In parallel, each non-leader exclusive-scans its *own* (un-merged) histogram into registers (the leader was - // otherwise the only block doing useful work here). Once the leader publishes `kth_bucket` below, the lane that - // owns it reads its exclusive prefix (= this block's keys strictly above the splitter this pass -> accumulated - // into `num_strictly_selected`) and its raw bucket count (-> `my_candidates`). Keeping the scan in registers lets - // `hist` reset on the normal schedule (the regs survive the reset and the next cluster sync). + // Step 3 (non-leader half, run in the arrive->wait window): each non-leader exclusive-scans its *own* (un-merged) + // histogram into registers -- the leader is otherwise the only block doing useful work at Step 3. Once the leader + // publishes `kth_bucket` below, the lane owning it reads its exclusive prefix (this block's keys strictly above + // the splitter this pass -> `num_strictly_selected`) and its raw bucket count (-> `my_candidates`). Keeping the + // scan in registers lets `hist` reset on the normal schedule (the regs survive the reset and the next sync). offset_t local_prefixes[buckets_per_thread]{}; offset_t local_hist_vals[buckets_per_thread]{}; - if (cluster_rank == leader_rank) - { - leader_identify_kth_bucket(); - } - else if (!is_idle_rank) + if (cluster_rank != leader_rank && !is_idle_rank) { _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < buckets_per_thread; ++j) @@ -2518,6 +2557,16 @@ private: block_scan_t(temp_storage.scan_storage).ExclusiveSum(local_hist_vals, local_prefixes); } + cluster_or_block_wait(is_single_cta); + + // Step 3 (leader half): the leader prefix-scans the merged `hist` (raw counts, all folds now visible) and updates + // the cluster-shared `state`. Subsequent reads (end-of-pass fold, last filter) observe these writes after the + // next cluster sync. + if (cluster_rank == leader_rank) + { + leader_identify_kth_bucket(); + } + if (pass + 1 < num_passes) { // Unconditional barrier: every working rank just read `hist` (the leader via its BlockScan, non-leaders via the @@ -2912,9 +2961,10 @@ private: temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; temp_storage.state.result_pair = 0; - // Front-load every counter the final filter relies on so the first cluster barrier below publishes the zeros to - // all ranks (the combined scan's DSMEM pushes then add into an already-zeroed `prefix_pair`, needing only a - // post-push barrier). `my_candidates` is zeroed too so idle/leader ranks (which never write it) read 0. + // Front-load every counter the final filter relies on so the initial cluster barrier below (arrive here, wait in + // the first pass) publishes the zeros to all ranks (the combined scan's DSMEM pushes then add into an + // already-zeroed `prefix_pair`, needing only a post-push barrier). `my_candidates` is zeroed too so idle/leader + // ranks (which never write it) read 0. temp_storage.prefix_pair = 0; temp_storage.front_local_cnt = 0; temp_storage.back_local_cnt = 0; @@ -2922,7 +2972,11 @@ private: temp_storage.my_candidates = 0; } reset_hist(); - cluster_or_block_sync(is_single_cta); + // Only arrive here; the matching wait is deferred to just before the first cross-CTA fold in `run_radix_passes`' + // first pass, so the cluster-arrival latency overlaps the fused first-pass load + histogram. Safe because nothing + // between here and that wait touches another rank's DSMEM (the first-pass histogram writes only block-local + // `hist`; `leader_state`/`prefix_pair` are untouched until later). + cluster_or_block_arrive(is_single_cta); [[maybe_unused]] const bool is_ok = detail::params::dispatch_discrete( select_directions, From e9ed89465bc25bdf217ba66be216480d91166c88 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 07:16:16 +0200 Subject: [PATCH 091/126] Clean up - Remove lots of unneeded inline assembly (some is needed for codegen) - Remove redundant span usage - Manually inline basic single-line functions for readability - Simplify functions with libcu++ helpers --- cub/cub/agent/agent_batched_topk_cluster.cuh | 500 +++++++------------ 1 file changed, 189 insertions(+), 311 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 38895387486..7578568550f 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -11,8 +11,9 @@ //! Histogram strategy (block-private accumulation, then DSMEM merge into the leader): //! 1. Every block lays out `hist[num_buckets]` at the same offset in its own //! shared memory. Each block accumulates a block-private histogram using -//! shared-space `red` reductions (cheap, SMEM-local): cta scope for -//! non-leaders, cluster scope for the leader (see `hist_inc`). +//! cluster-scope shared-space `red` reductions (cheap, SMEM-local), so the +//! leader's adds stay mutually atomic with the DSMEM folds below (see +//! `hist_inc`). //! 2. After a cluster-wide barrier, every non-leader block walks its //! histogram and folds its bucket counts into the leader block's `hist` //! via cluster-scope DSMEM atomics. The leader's `hist` therefore plays @@ -56,10 +57,11 @@ #include #include #include -#include #include #include #include +#include +#include #include #include #include @@ -331,6 +333,10 @@ struct agent_batched_topk_cluster static constexpr int load_align_items = smem_layout_t::load_align_items; static constexpr int slot_alignment = smem_layout_t::slot_alignment; + // (pointer, count) carrier for a contiguous run of SMEM-staged keys, used only where the length must travel with the + // base across call boundaries (the resident window and the streamer's per-stage view); elsewhere: `key_t*` + count. + using smem_keys_t = ::cuda::std::span; + // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a // time: clamp items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit // to <= chunk granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. @@ -420,49 +426,17 @@ struct agent_batched_topk_cluster // head_edge_cap_items)`. static constexpr int head_edge_cap_items = load_align_items; - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE key_t* slot_keys_unpadded(int slot) const - { - return reinterpret_cast(key_slots + slot * ChunkBytes); - } - - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int span_size(::cuda::std::span keys) const - { - const int count = static_cast(::cuda::std::size(keys)); - _CCCL_ASSERT(static_cast<::cuda::std::size_t>(count) == ::cuda::std::size(keys), - "Resident key span length must fit in int"); - return count; - } - - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t - aligned_head_items(const key_t* base, offset_t segment_size) const - { - const auto base_addr = reinterpret_cast<::cuda::std::uintptr_t>(base); - const auto rem = base_addr % static_cast<::cuda::std::uintptr_t>(load_align_bytes); - const auto bytes = - (rem == 0) ? ::cuda::std::uintptr_t{0} : static_cast<::cuda::std::uintptr_t>(load_align_bytes) - rem; - const auto items = static_cast(bytes / static_cast<::cuda::std::uintptr_t>(sizeof(key_t))); - return (::cuda::std::min) (items, segment_size); - } - - // Number of aligned chunks covering the segment. The unaligned head prefix (`head_items`) is handled as a separate - // edge, not a chunk, so the aligned region `[head_items, segment_size)` chunks uniformly at `chunk_items` (a segment - // lying entirely before the first boundary has no chunks). `head_items == 0` (aligned base or generic fallback) is - // then plain uniform chunking from offset 0. - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t num_chunks(offset_t segment_size, offset_t head_items) const - { - const offset_t aligned_items = (segment_size > head_items) ? (segment_size - head_items) : offset_t{0}; - return static_cast(::cuda::ceil_div(aligned_items, offset_t{chunk_items})); - } - struct chunk_desc { offset_t offset; int count; }; - // Chunk `chunk_idx` of the aligned region: it begins on a `load_align` boundary (`head_items + chunk_idx * - // chunk_items`, with `head_items` itself aligning the base), so every chunk has a zero prefix and only the last chunk - // can carry an unaligned suffix (the segment's trailing `< load_align_items` items). + // Chunk `chunk_idx` of the aligned region, beginning on a `load_align` boundary (`head_items + chunk_idx * + // chunk_items`, with `head_items` aligning the base and `chunk_items` a multiple of `load_align_items`): so every + // chunk has a zero prefix and only the last chunk can carry an unaligned suffix (its trailing `< load_align_items` + // items). A chunk's aligned bulk (what the guard-free TMA path copies) is `round_down(count, load_align_items)` + // (`== count` for interior chunks); the tail's `count - bulk` remainder is peeled into `edge_keys`. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_desc get_chunk(offset_t chunk_idx, offset_t segment_size, offset_t head_items) const { @@ -471,47 +445,6 @@ struct agent_batched_topk_cluster return {offset, static_cast((::cuda::std::min) (remaining, offset_t{chunk_items}))}; } - // Splits a chunk into its unaligned front edge, aligned interior (bulk), and unaligned back edge, relative to the - // gmem base. The interior begins and ends on a `load_align` boundary so it can be loaded with the aligned, - // guard-free BlockLoadToShared path; the edges (each `< load_align_items` items) are staged via - // `stage_and_fold_edge`. The head prefix is peeled as a separate edge before chunking, so a `get_chunk` chunk begins - // on a boundary (zero prefix) and only the last chunk (tail) carries a nonzero suffix. - struct chunk_split - { - offset_t prefix; - offset_t bulk; - offset_t suffix; - }; - - template - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_split split_chunk(PtrT base, const chunk_desc chunk) const - { - const auto align_bytes = static_cast<::cuda::std::uintptr_t>(load_align_bytes); - const auto begin = reinterpret_cast<::cuda::std::uintptr_t>(base + chunk.offset); - const auto end = begin + static_cast<::cuda::std::uintptr_t>(chunk.count) * sizeof(key_t); - const auto aligned_begin = ::cuda::round_up(begin, align_bytes); - const auto aligned_end = ::cuda::round_down(end, align_bytes); - if (aligned_begin > aligned_end) - { - // The chunk lies strictly between two load_align boundaries (no aligned point inside): the whole chunk is an - // unaligned edge, attributed entirely to the front edge. `get_chunk` chunks begin on a boundary (the head is - // peeled separately), so this only guards a degenerate sub-`load_align` chunk. A tail always begins on a - // boundary, so it takes the `aligned_begin <= aligned_end` path below and its unaligned remainder is the suffix. - return {static_cast(chunk.count), offset_t{0}, offset_t{0}}; - } - const offset_t prefix = static_cast((aligned_begin - begin) / sizeof(key_t)); - const offset_t bulk = static_cast((aligned_end - aligned_begin) / sizeof(key_t)); - return {prefix, bulk, static_cast(chunk.count) - prefix - bulk}; - } - - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t - num_rank_chunks(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const - { - return (cluster_rank < chunks) - ? static_cast((chunks - 1 - cluster_rank) / cluster_blocks + 1) - : offset_t{0}; - } - // Assignment of the cluster's global chunk indices `[0, chunks)` to its CTAs. A rank owns `count` chunks; its i-th // owned chunk has global index `global_index(i) = first + i * stride`. The single mapping point lets the rest of the // agent (resident load, streamer, per-pass scans) stay agnostic to the layout chosen by `make_chunk_partition`. @@ -551,9 +484,13 @@ struct agent_batched_topk_cluster } else { - return {static_cast(cluster_rank), - static_cast(cluster_blocks), - num_rank_chunks(chunks, cluster_rank, cluster_blocks)}; + // Strided count: rank `cluster_rank` owns the global indices `cluster_rank, cluster_rank + cluster_blocks, ...` + // that stay `< chunks`, i.e. `ceil((chunks - cluster_rank) / cluster_blocks)`, written as + // `(chunks - 1 - cluster_rank) / cluster_blocks + 1` for exact integer arithmetic. The `cluster_rank >= chunks` + // guard (rank owns nothing) keeps the subtraction from underflowing the unsigned `offset_t`. + const offset_t count = + (cluster_rank < chunks) ? static_cast((chunks - 1 - cluster_rank) / cluster_blocks + 1) : offset_t{0}; + return {static_cast(cluster_rank), static_cast(cluster_blocks), count}; } } @@ -587,7 +524,7 @@ struct agent_batched_topk_cluster // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped (floor) items per // thread, so the sub-tile remainder is handled by a single non-unrolled fused block-stride loop bounded by the count. template - _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* data, int chunk_count, Apply&& apply) const + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* keys, int chunk_count, Apply&& apply) const { constexpr int tile = Unroll * threads_per_block; const int tid = static_cast(threadIdx.x); @@ -600,7 +537,7 @@ struct agent_batched_topk_cluster _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < Unroll; ++i) { - regs[i] = data[tile_base + i * threads_per_block + tid]; + regs[i] = keys[tile_base + i * threads_per_block + tid]; } _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < Unroll; ++i) @@ -614,14 +551,14 @@ struct agent_batched_topk_cluster _CCCL_PRAGMA_NOUNROLL() for (int local = full_tiles + tid; local < chunk_count; local += threads_per_block) { - apply(data[local], local); + apply(keys[local], local); } } template - _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(::cuda::std::span chunk_keys, F&& f) const + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(const key_t* chunk_keys, int count, F&& f) const { - for_each_chunk_key_impl(::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int) { + for_each_chunk_key_impl(chunk_keys, count, [&](const key_t& key, int) { f(key); }); } @@ -631,26 +568,11 @@ struct agent_batched_topk_cluster // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. template _CCCL_DEVICE _CCCL_FORCEINLINE void - for_each_chunk_key_indexed(::cuda::std::span chunk_keys, offset_t base_off, F&& f) const - { - for_each_chunk_key_impl( - ::cuda::std::data(chunk_keys), span_size(chunk_keys), [&](const key_t& key, int local) { - f(key, base_off + static_cast(local)); - }); - } - - // A bulk in the block_tile as a 32-bit shared address + length. A spilled 32-bit shared address (rebuilt with - // `__cvta_shared_to_generic`) keeps the key reads `LDS`; a spilled 64-bit generic pointer would demote them to `LD`. - struct shared_bulk - { - ::cuda::std::uint32_t smem32; - int len; - }; - - // Rebuild a bulk's span from its 32-bit shared address at the point of use (spill-proof `LDS`). - [[nodiscard]] static _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span bulk_span(shared_bulk b) + for_each_chunk_key_indexed(const key_t* chunk_keys, int count, offset_t base_off, F&& f) const { - return {reinterpret_cast(__cvta_shared_to_generic(b.smem32)), static_cast<::cuda::std::size_t>(b.len)}; + for_each_chunk_key_impl(chunk_keys, count, [&](const key_t& key, int local) { + f(key, base_off + static_cast(local)); + }); } // --------------------------------------------------------------------------- @@ -797,34 +719,28 @@ private: } // --------------------------------------------------------------------------- - // Block-private histogram atomics (shared-space `red` via cuda::ptx-style inline PTX) + // Block-private histogram atomics (shared-space `red` via inline PTX) // --------------------------------------------------------------------------- - // A builtin `atomicAdd(&temp_storage.hist[bucket], 1)` compiles to a generic atomic (`ATOM.E`) whose 64-bit base is - // spilled and reloaded (`LDL.64`) at every update across this huge agent. A shared-space `red` instead addresses with - // the 32-bit shared address (no base to spill). Shared atomics only allow cta/cluster scope, which is exactly right: - // every writer of a given `hist` is in the same cluster. `red` (no return) matches the discarded `atomicAdd` result. + // The per-key histogram update is the hottest path. A builtin `atomicAdd(&temp_storage.hist[bucket], 1)` lowers to + // the same warp-aggregated shared atomic (`ATOMS.POPC.INC.32`), but recomputes `&hist[bucket]` from the 64-bit + // generic base of `temp_storage` on every key; the inline `red` adds a pre-hoisted 32-bit shared base instead + // (measured ~3-4% fewer total instructions across this huge agent). - // The 32-bit shared address of `hist[0]`. Hoisted once per histogram region so the per-key address math is a pure - // 32-bit add; recomputing it per key would reload the 64-bit generic base of `temp_storage` from the stack. + // The 32-bit shared address of `hist[0]`, hoisted once per histogram region so the per-key update stays a pure + // 32-bit add. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t hist_base32() const { return static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(temp_storage.hist)); } - // Increment this block's own histogram bucket by one via its 32-bit shared address. The leader (rank 0) also receives - // remote folds from the other cluster blocks (Step 2, `hist_fold_remote`), so its add must be cluster-scoped to be - // mutually atomic with them; non-leaders only touch their own `hist` before the fold, so cta scope suffices. - _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket, bool is_leader) + // Increment this block's own histogram bucket by one, at cluster scope so it stays mutually atomic with the DSMEM + // folds peers push into the leader's `hist` (Step 2, `hist_fold_remote`). Cluster and cta scope lower to identical + // shared-atomic SASS here, so applying it unconditionally (non-leaders, lone CTA) costs nothing and drops the leader + // branch. + _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket) { const ::cuda::std::uint32_t addr = base32 + static_cast<::cuda::std::uint32_t>(bucket) * sizeof(offset_t); - if (is_leader) - { - asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); - } - else - { - asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); - } + asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); } // Step 2: a non-leader folds one bucket into the leader's histogram through DSMEM. `own_bucket_addr32` is the @@ -861,35 +777,6 @@ private: asm volatile("red.relaxed.cluster.shared::cluster.add.u64 [%0], %1;" : : "r"(remote), "l"(v) : "memory"); } - // Block-local SMEM atomics for the final filter: `front_local_inc`/`back_local_inc` hand out this block's next - // front/back output slot (pre-increment value), `add_local_selected` accumulates its strictly-selected count across - // passes. Each addresses the 32-bit shared slot directly (like `hist_inc`) to dodge the generic-atomic base spill of - // `atomicAdd(&shared)`, at cta scope since every writer of a block's counter lives in that block. - _CCCL_DEVICE _CCCL_FORCEINLINE offset_t front_local_inc() - { - const ::cuda::std::uint32_t addr = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.front_local_cnt)); - offset_t old; - asm volatile("atom.relaxed.cta.shared::cta.add.u32 %0, [%1], 1;" : "=r"(old) : "r"(addr) : "memory"); - return old; - } - - _CCCL_DEVICE _CCCL_FORCEINLINE offset_t back_local_inc() - { - const ::cuda::std::uint32_t addr = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)); - offset_t old; - asm volatile("atom.relaxed.cta.shared::cta.add.u32 %0, [%1], 1;" : "=r"(old) : "r"(addr) : "memory"); - return old; - } - - _CCCL_DEVICE _CCCL_FORCEINLINE void add_local_selected(offset_t count) - { - const ::cuda::std::uint32_t addr = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.num_strictly_selected)); - asm volatile("red.relaxed.cta.shared::cta.add.u32 [%0], %1;" : : "r"(addr), "r"(count) : "memory"); - } - // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram // plus identification of the bucket holding the k-th item. Each thread of // the leader block contributes `buckets_per_thread` consecutive buckets in @@ -1001,39 +888,33 @@ private: "streaming depth exceeds the overflow chunk count"); } - _CCCL_DEVICE _CCCL_FORCEINLINE offset_t chunk_index_of(offset_t overflow_idx) const - { - return part.global_index(overflow_base + overflow_idx); - } - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) { - const offset_t chunk_idx = chunk_index_of(overflow_idx); + const offset_t chunk_idx = part.global_index(overflow_base + overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - // Every chunk begins on a `load_align` boundary (zero prefix), so the guard-free aligned (TMA bulk) path applies. - // The global-last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk - // excludes it. For every interior chunk `bulk == count`. - const auto split = agent.split_chunk(block_keys_base, chunk); - _CCCL_ASSERT(split.prefix == offset_t{0}, "overflow streamer received a chunk with an unaligned start"); + // Every chunk begins on a `load_align` boundary, so the guard-free aligned (TMA bulk) path applies. The global- + // last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk excludes + // it. For every interior chunk `bulk == count`. + const offset_t bulk = + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); + _CCCL_ASSERT(::cuda::is_aligned(block_keys_base + chunk.offset, load_align_bytes), + "overflow streamer received a chunk with an unaligned start"); char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; - const ::cuda::std::span src{ - block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(split.bulk)}; + const ::cuda::std::span src{block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; agent.issue_bulk_copy(stage, dst, src); inflight_mask |= (::cuda::std::uint32_t{1} << stage); } - // Rebuild the shared span for the chunk currently resident in `stage`'s slot without storing per-stage state: the - // slot address is a pure function of `stage` and the length is recomputed from `overflow_idx`, so there is no - // spillable `pending[]` array (see `bulk_span`). Returns the aligned bulk only (the always-peeled tail suffix is - // excluded). - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span - stage_span(int stage, offset_t overflow_idx) const + // Shared-memory view of the chunk currently resident in `stage`'s slot without storing per-stage state: the slot + // index is a pure function of `stage` and the length is recomputed from `overflow_idx`, so there is no spillable + // `pending[]` array. Returns the aligned bulk only (the always-peeled tail suffix is excluded). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE smem_keys_t stage_span(int stage, offset_t overflow_idx) const { - char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; - const auto chunk = agent.get_chunk(chunk_index_of(overflow_idx), segment_size, head_items); - const auto split = agent.split_chunk(block_keys_base, chunk); - return agent.bulk_span( - {static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(dst)), static_cast(split.bulk)}); + const auto chunk = agent.get_chunk(part.global_index(overflow_base + overflow_idx), segment_size, head_items); + return smem_keys_t( + ::cuda::ptr_rebind(agent.key_slots + (stream_slot_base + stage) * ChunkBytes), + static_cast( + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); } // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident @@ -1156,7 +1037,7 @@ private: for (offset_t i = 0; i < overflow_chunks; ++i) { const offset_t overflow_idx = is_forward ? i : (overflow_chunks - 1 - i); - const offset_t chunk_idx = chunk_index_of(overflow_idx); + const offset_t chunk_idx = part.global_index(overflow_base + overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); generic_apply(chunk); if (!should_continue()) @@ -1181,12 +1062,16 @@ private: [&](int stage, offset_t overflow_idx) { if constexpr (Indexed) { - const offset_t base_off = agent.get_chunk(chunk_index_of(overflow_idx), segment_size, head_items).offset; - agent.template for_each_chunk_key_indexed(stage_span(stage, overflow_idx), base_off, f); + const offset_t base_off = + agent.get_chunk(part.global_index(overflow_base + overflow_idx), segment_size, head_items).offset; + const auto keys = stage_span(stage, overflow_idx); + agent.template for_each_chunk_key_indexed( + keys.data(), static_cast(keys.size()), base_off, f); } else { - agent.template for_each_chunk_key(stage_span(stage, overflow_idx), f); + const auto keys = stage_span(stage, overflow_idx); + agent.template for_each_chunk_key(keys.data(), static_cast(keys.size()), f); } }, [&](const auto& chunk) { @@ -1350,9 +1235,8 @@ private: // member functions instead of a nest of `[&]` lambdas. Constructed once per `run()` in the deterministic branch. Kept // an aggregate (no user constructor) so its members are initialized positionally at the single call site. // - // Codegen: methods are `_CCCL_FORCEINLINE` and no SMEM key pointer is stored -- the resident window is carried as its - // 32-bit shared address (`resident_smem32`) and rebuilt with `__cvta_shared_to_generic` at use, so the reads stay - // `LDS` (see the `resident_smem32` note in `run`). + // Codegen: methods are `_CCCL_FORCEINLINE`; the resident window is carried as a `smem_keys_t` view + // (`resident_front`). template struct det_final_filter { @@ -1377,7 +1261,7 @@ private: offset_t segment_size_off; offset_t head_items; offset_t front_seg_base; - ::cuda::std::uint32_t resident_smem32; + smem_keys_t resident_front; int front_count; int head_edge_len_items; int tail_edge_len_items; @@ -1416,19 +1300,11 @@ private: return is_front_done && is_back_done; } - // Fold a flat scan position `pos` into its in-region index: forward, or (`Reversed`) counting down from the - // region's last element. The `FromSmem` local read and the segment-local value index share this one folded index. - template - [[nodiscard]] static _CCCL_DEVICE _CCCL_FORCEINLINE int fold_pos(int pos, int count) - { - return Reversed ? (count - 1 - pos) : pos; - } - // Emit one back/tie candidate: if its `global_rank` (arrival- or scan-ordered by the caller) is a winner it lands // in the top-k output at slot `k-1-global_rank`, with its value pulled from segment-local index `seg_idx`; ranks at // or past `num_back` are losers and dropped (a no-op). Forced-inline with explicit args (no capture) so it folds // into the hot back-placement loops with no extra live ranges; the caller keeps any counter side effects (e.g. - // `back_local_inc`) in the `global_rank` argument so they still run for every candidate. + // the `back_local_cnt` atomic) in the `global_rank` argument so they still run for every candidate. _CCCL_DEVICE _CCCL_FORCEINLINE void emit_back_one(offset_t global_rank, const key_t& key, offset_t seg_idx) { if (global_rank < static_cast(num_back)) @@ -1446,8 +1322,7 @@ private: // per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are // unresolved -- holds the boundary and scans directly. `running` carries across tiles/regions. No per-tile early- // exit or barrier here except the lazy-scan `else` branch; early exit is decided at critical points via - // `should_stop`. The SMEM base is passed in (rebuilt via `__cvta_shared_to_generic` at the call site) rather than - // stored, so the reads stay `LDS`. + // `should_stop`. template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles([[maybe_unused]] const key_t* smem_src, offset_t seg_base, int count, bool region_is_terminal) @@ -1468,7 +1343,9 @@ private: flags[i] = offset_t{0}; if (is_valid[i]) { - const int folded_pos = fold_pos(pos, count); + // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` + // local read and the segment-local value index. + const int folded_pos = Reversed ? (count - 1 - pos) : pos; if constexpr (FromSmem) { keys[i] = smem_src[folded_pos]; @@ -1492,9 +1369,9 @@ private: if (is_front_key) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - const offset_t local = agent.front_local_inc(); + const offset_t local = atomicAdd(&agent.temp_storage.front_local_cnt, offset_t{1}); const out_offset_t out = static_cast(sel_prefix + local); - const offset_t seg_idx = seg_base + static_cast(fold_pos(pos, count)); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); block_keys_out[out] = keys[i]; write_value(out, seg_idx); } @@ -1518,9 +1395,9 @@ private: if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - emit_back_one(cand_prefix + agent.back_local_inc(), + emit_back_one(cand_prefix + atomicAdd(&agent.temp_storage.back_local_cnt, offset_t{1}), keys[i], - seg_base + static_cast(fold_pos(pos, count))); + seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); } } }; @@ -1536,7 +1413,8 @@ private: if (is_valid[i] && flags[i] != offset_t{0}) { const int pos = tile_base + static_cast(threadIdx.x) * items + i; - emit_back_one(base + excl[i], keys[i], seg_base + static_cast(fold_pos(pos, count))); + emit_back_one( + base + excl[i], keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); } } return tile_total; @@ -1571,7 +1449,8 @@ private: is_tie_active = false; } // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the - // next tile's `back_local_inc`. The other sub-paths write disjoint slots and need no per-tile barrier. + // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no per-tile + // barrier. __syncthreads(); } } @@ -1586,10 +1465,9 @@ private: { if constexpr (use_block_load_to_shared) { - // Whole contiguous resident span staged in SMEM; rebuild the base from its 32-bit address at use. - key_t* const rfront = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + // Whole contiguous resident span staged in SMEM. process_tiles( - rfront, front_seg_base, front_count, is_resident_terminal); + resident_front.data(), front_seg_base, front_count, is_resident_terminal); } else { @@ -1602,7 +1480,7 @@ private: // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy // per-tile boundary detection handles any boundary), so pass `false`. process_tiles( - agent.slot_keys_unpadded(local_slot), chunk.offset, chunk.count, false); + ::cuda::ptr_rebind(agent.key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); } } } @@ -1628,20 +1506,22 @@ private: streamer.run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. - // `stage_span` rebuilds the slot pointer from its 32-bit shared address (spill-proof `LDS`) and returns only - // the aligned bulk (a peeled tail suffix is handled by `process_tail_edge`). + // `stage_span` returns the slot's aligned-bulk view (a peeled tail suffix is handled by `process_tail_edge`). [&](int stage, offset_t overflow_idx) { - const auto span = streamer.stage_span(stage, overflow_idx); + const auto keys = streamer.stage_span(stage, overflow_idx); const offset_t base_off = - agent.get_chunk(streamer.chunk_index_of(overflow_idx), segment_size_off, head_items).offset; + agent + .get_chunk(streamer.part.global_index(streamer.overflow_base + overflow_idx), segment_size_off, head_items) + .offset; // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving // is one barrier on one tile. process_tiles( - span.data(), base_off, static_cast(span.size()), false); + keys.data(), base_off, static_cast(keys.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). [&](const auto& chunk) { + // FromSmem=false: `smem_src` is unread, so pass nullptr. process_tiles( nullptr, chunk.offset, chunk.count, false); }, @@ -1655,13 +1535,15 @@ private: } // Fold one persistent boundary edge (head prefix or peeled tail suffix), both staged in `edge_keys`. A no-op on the - // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. + // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. An edge spans fewer + // than `load_align_items` items (it fits in its `edge_keys` slot), so it runs `process_tiles` at unroll factor 1 + // (non-unrolled) instead of the wider tie-break unroll the potentially large resident/overflow regions use. _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge(const key_t* keys, offset_t seg_base, int count, bool is_terminal) { if constexpr (use_block_load_to_shared) { - process_tiles( - keys, seg_base, count, is_terminal); + _CCCL_ASSERT(count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed>(keys, seg_base, count, is_terminal); } } @@ -1770,8 +1652,7 @@ private: offset_t my_resident_chunks, offset_t segment_size_off, offset_t head_items, - ::cuda::std::uint32_t resident_smem32, - ::cuda::std::span resident_keys, + smem_keys_t resident_keys, int head_edge_len_items, int tail_edge_len_items, overflow_streamer& streamer) @@ -1809,8 +1690,9 @@ private: const auto res = identify_op(key); if (res == detail::topk::candidate_class::selected) { - const out_offset_t pos = static_cast(sel_prefix + front_local_inc()); - block_keys_out[pos] = key; + const out_offset_t pos = + static_cast(sel_prefix + atomicAdd(&temp_storage.front_local_cnt, offset_t{1})); + block_keys_out[pos] = key; if constexpr (!is_keys_only) { write_value(pos, seg_idx); @@ -1818,7 +1700,8 @@ private: } else if (res == detail::topk::candidate_class::candidate) { - const out_offset_t back_pos = static_cast(cand_prefix + back_local_inc()); + const out_offset_t back_pos = + static_cast(cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1})); if (back_pos < num_kth) { const out_offset_t pos = k - 1 - back_pos; @@ -1842,9 +1725,8 @@ private: const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) { - key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); for_each_chunk_key( - {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, write_selected); + resident_keys.data(), static_cast(resident_keys.size()), write_selected); } else { @@ -1852,9 +1734,10 @@ private: { const offset_t chunk_idx = part.global_index(local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( - {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, write_selected); + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + static_cast(chunk.count), + write_selected); } } // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). @@ -1887,21 +1770,32 @@ private: } }; + // Boundary edges span fewer than `load_align_items` items, so fold them with a plain block-stride loop rather + // than the tiled/unrolled `write_run` used for the (potentially large) resident bulks. + const auto write_edge = [&](const key_t* smem, offset_t base_off, int count) { + for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) + { + const key_t key = smem[local]; + sink(key, base_off + static_cast(local)); + } + }; + // Fold the resident keys (and their values) as the streamer's `mid` work, exactly as in the keys-only path. const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) { // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into - // `edge_keys` and folded below), so iterate `split.bulk`, not `chunk.count`. - key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); + // `edge_keys` and folded below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not + // `chunk.count`. + key_t* const resident_ptr = resident_keys.data(); int cursor_items = 0; for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) { const auto chunk = get_chunk(part.global_index(resident_base + local_chunk), segment_size_off, head_items); - const auto split = split_chunk(block_keys_base, chunk); - const offset_t base_off = chunk.offset + split.prefix; - const int bulk_count_items = split.bulk; + const offset_t base_off = chunk.offset; + const int bulk_count_items = static_cast( + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); write_run(resident_ptr + cursor_items, base_off, bulk_count_items); cursor_items += bulk_count_items; } @@ -1912,7 +1806,8 @@ private: { const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_off, head_items); const offset_t base_off = chunk.offset; - write_run(slot_keys_unpadded(static_cast(local_chunk)), base_off, chunk.count); + write_run( + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), base_off, chunk.count); } } // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the @@ -1921,13 +1816,13 @@ private: { if (head_edge_len_items > 0) { - write_run(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); + write_edge(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); } if (tail_edge_len_items > 0) { - write_run(temp_storage.edge_keys + head_edge_cap_items, - segment_size_off - static_cast(tail_edge_len_items), - tail_edge_len_items); + write_edge(temp_storage.edge_keys + head_edge_cap_items, + segment_size_off - static_cast(tail_edge_len_items), + tail_edge_len_items); } } }; @@ -1963,8 +1858,7 @@ private: offset_t segment_size_off, offset_t head_items, state_t* leader_state, - ::cuda::std::uint32_t resident_smem32, - ::cuda::std::span resident_keys, + smem_keys_t resident_keys, int head_edge_len_items, int tail_edge_len_items, overflow_streamer& streamer) @@ -2016,7 +1910,7 @@ private: // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never // part of this span. - const int front_count = span_size(resident_keys); + const int front_count = static_cast(resident_keys.size()); // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later @@ -2065,7 +1959,7 @@ private: segment_size_off, head_items, front_seg_base, - resident_smem32, + resident_keys, front_count, head_edge_len_items, tail_edge_len_items, @@ -2080,15 +1974,12 @@ private: // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 - // needs no candidate filtering). Publishes the resident span as `resident_keys` and its 32-bit shared base address - // as `resident_smem32` for the later passes and the final filter; ends on a `__syncthreads()` that makes `edge_keys` - // and the block-local histogram visible within the block. Cross-CTA visibility of the zeroed histogram comes later, - // from the deferred initial cluster wait in `run_radix_passes` (this whole load runs in that arrive->wait window). + // needs no candidate filtering). Publishes the resident span as `resident_keys` for the later passes and the final + // filter; ends on a `__syncthreads()` that makes `edge_keys` and the block-local histogram visible within the block. + // Cross-CTA visibility of the zeroed histogram comes later, from the deferred initial cluster wait in + // `run_radix_passes` (this whole load runs in that arrive->wait window). template _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( - unsigned int cluster_rank, - unsigned int leader_rank, - bool is_single_cta, const chunk_partition& part, offset_t my_resident_chunks, offset_t resident_base, @@ -2100,20 +1991,16 @@ private: int head_edge_len_items, int tail_edge_len_items, overflow_streamer& streamer, - ::cuda::std::span& resident_keys, - ::cuda::std::uint32_t& resident_smem32) + smem_keys_t& resident_keys) { using extract_bin_op_t = detail::topk::extract_bin_op_t; constexpr int total_bits = int{sizeof(key_t)} * 8; extract_bin_op_t extract_op(0, total_bits, decomposer_t{}); const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - // Cluster-scope leader atomics are only needed to stay mutually atomic with the non-leaders' DSMEM folds; a lone - // CTA has none, so it uses the cheaper CTA scope. - const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; - auto add_first_pass = [&](const key_t& key) { + auto add_first_pass = [&](const key_t& key) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, is_leader_cluster_scope); + hist_inc(hist_smem32, bucket); }; if constexpr (use_block_load_to_shared) @@ -2124,10 +2011,9 @@ private: // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk - // begins on a `load_align` boundary (zero prefix), so its whole `count` is the aligned bulk - except the - // global-last chunk, whose unaligned suffix is always peeled into `edge_keys`, leaving only its aligned bulk - // here. The read span is rebuilt from a rooted 32-bit shared address (see `bulk_span`) so a spilled cursor - // cannot demote the reads to `LD`. + // begins on a `load_align` boundary (zero prefix), so its aligned bulk is `round_down(count, + // load_align_items)`: the whole `count` for interior chunks, and for the global-last chunk its `count` minus + // the unaligned suffix that is always peeled into `edge_keys`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` @@ -2140,12 +2026,13 @@ private: const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { const offset_t chunk_idx = part.global_index(resident_local(slot)); const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); - const auto split = split_chunk(block_keys_base, chunk); - if (split.bulk == 0) + const offset_t bulk = + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); + if (bulk == 0) { return {}; } - return {block_keys_base + chunk.offset + split.prefix, static_cast<::cuda::std::size_t>(split.bulk)}; + return {block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; }; // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives @@ -2170,9 +2057,7 @@ private: wait_stage(stage); const int read_len_items = static_cast(::cuda::std::size(bulk_src(local_chunk))); for_each_chunk_key( - bulk_span({static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots + read_off_bytes)), - read_len_items}), - add_first_pass); + ::cuda::ptr_rebind(key_slots + read_off_bytes), read_len_items, add_first_pass); read_off_bytes += read_len_items * int{sizeof(key_t)}; const offset_t next_local_chunk = local_chunk + static_cast(prologue); @@ -2187,11 +2072,9 @@ private: } } - // The resident region is one contiguous span of aligned bulks for the later passes; both boundary edges are + // The resident region is one contiguous run of aligned bulks for the later passes; both boundary edges are // folded separately from `edge_keys`. - resident_keys = {reinterpret_cast(key_slots), - static_cast<::cuda::std::size_t>(next_off_bytes / int{sizeof(key_t)})}; - resident_smem32 = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(key_slots)); + resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), next_off_bytes / int{sizeof(key_t)}); } } else @@ -2200,7 +2083,7 @@ private: { const offset_t chunk_idx = part.global_index(resident_base + local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); + key_t* const chunk_keys = ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) @@ -2244,7 +2127,7 @@ private: // `run` entry); `wait_stage` provides the producer/consumer sync. streamer.process_pass(add_first_pass); - const int resident_count = span_size(resident_keys); + const int resident_count = static_cast(resident_keys.size()); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, "Dynamic shared memory block_tile is too small"); __syncthreads(); @@ -2289,9 +2172,10 @@ private: layout.block_keys_in = d_key_segments_it[segment_id]; layout.segment_size_off = static_cast(segment_size); // `cluster_blocks` is what `process_impl` runs at: the launched cluster size, or `1` when it collapsed a - // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()`, keeps `state`/atomics block-local, - // and uses CTA-scope histogram atomics (no cross-rank DSMEM folds to be mutually atomic with). For wider clusters, - // `eff_cluster_blocks` (below) further excludes ranks that receive no chunks; they stay resident but idle. + // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()` and keeps `state`/atomics block-local + // (no cross-rank DSMEM folds); its histogram increments use the same unconditional cluster-scope `hist_inc`, which + // is identical SASS at cluster size 1. For wider clusters, `eff_cluster_blocks` (below) further excludes ranks that + // receive no chunks; they stay resident but idle. layout.is_single_cta = (cluster_blocks == 1u); layout.block_keys_base = nullptr; @@ -2299,12 +2183,21 @@ private: if constexpr (use_block_load_to_shared) { layout.block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(layout.block_keys_in); - layout.head_items = aligned_head_items(layout.block_keys_base, layout.segment_size_off); - } - // The generic fallback does not use BlockLoadToShared's alignment hint or peeling path, so it can keep a simple - // uniform chunking (`head_items == 0`). The two chunkings may assign keys to CTAs differently, but top-k only - // depends on the multiset of keys covered by the cluster. - layout.chunks = num_chunks(layout.segment_size_off, layout.head_items); + // Items from `block_keys_base` to the next `load_align_bytes` boundary (0 when already aligned), clamped to the + // segment. `load_align_bytes` is a multiple of `sizeof(key_t)`, so the aligned pointer lands on a key boundary + // and the pointer difference is already an exact item count; narrow it to 32-bit `offset_t` right away. + const offset_t head_to_boundary = + static_cast(::cuda::align_up(layout.block_keys_base, load_align_bytes) - layout.block_keys_base); + layout.head_items = (::cuda::std::min) (head_to_boundary, layout.segment_size_off); + } + // Number of aligned chunks covering the segment: the unaligned head prefix (`head_items`) is a separate edge, not a + // chunk, so the aligned region `[head_items, segment_size)` chunks uniformly at `chunk_items` (a segment lying + // entirely before the first boundary yields `head_items == segment_size` and zero chunks). The generic fallback + // skips the alignment/peeling path and keeps `head_items == 0`; the two chunkings may assign keys to CTAs + // differently, but top-k only depends on the multiset of keys the cluster covers. + _CCCL_ASSERT(layout.head_items <= layout.segment_size_off, "head prefix cannot exceed the segment"); + layout.chunks = + static_cast(::cuda::ceil_div(layout.segment_size_off - layout.head_items, offset_t{chunk_items})); // Effective cluster blocks: the CTAs that actually receive chunks (at least `min_chunks_per_block` each), <= the // launched `cluster_blocks`. Ranks at or beyond it are idle -- they own no chunks, fold nothing, and never lead -- @@ -2366,8 +2259,10 @@ private: && layout.part.global_index(layout.my_chunks - offset_t{1}) == layout.chunks - offset_t{1}) { const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_off, layout.head_items); - tail_suffix_items = split_chunk(layout.block_keys_base, tail_chunk).suffix; - owns_suffix_tail = tail_suffix_items != offset_t{0}; + tail_suffix_items = + static_cast(tail_chunk.count) + - ::cuda::round_down(static_cast(tail_chunk.count), static_cast(load_align_items)); + owns_suffix_tail = tail_suffix_items != offset_t{0}; } } @@ -2425,8 +2320,7 @@ private: state_t* leader_state, int head_edge_len_items, int tail_edge_len_items, - ::cuda::std::span resident_keys, - ::cuda::std::uint32_t resident_smem32, + smem_keys_t resident_keys, overflow_streamer& streamer, key_prefix_t& kth_key_bits_local) { @@ -2450,15 +2344,13 @@ private: identify_candidates_op_t identify_op(&kth_key_bits_local, pass, total_bits, decomposer_t{}); extract_bin_op_t extract_op(pass, total_bits, decomposer_t{}); - // Step 1: block-private histogram via shared-space `red` (see `hist_inc`): leader uses cluster scope to be - // mutually atomic with the non-leaders' Step 2 DSMEM folds, non-leaders use the cheaper cta scope. + // Step 1: block-private histogram via shared-space `red` at cluster scope (see `hist_inc`). const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - const bool is_leader_cluster_scope = (cluster_rank == leader_rank) && !is_single_cta; auto add_hist = [&](const key_t& key) { if (identify_op(key) == detail::topk::candidate_class::candidate) { const int bucket = extract_op(key); - hist_inc(hist_smem32, bucket, is_leader_cluster_scope); + hist_inc(hist_smem32, bucket); } }; @@ -2468,11 +2360,8 @@ private: const auto fold_resident_hist = [&] { if constexpr (use_block_load_to_shared) { - // Rebuild the resident pointer from its 32-bit shared address so the reads stay `LDS` even if the value - // spilled across the pass loop (see `resident_smem32`). - key_t* const resident_ptr = reinterpret_cast(__cvta_shared_to_generic(resident_smem32)); for_each_chunk_key( - {resident_ptr, static_cast<::cuda::std::size_t>(span_size(resident_keys))}, add_hist); + resident_keys.data(), static_cast(resident_keys.size()), add_hist); } else { @@ -2480,9 +2369,10 @@ private: { const offset_t chunk_idx = part.global_index(resident_base + local_chunk); const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); - key_t* const chunk_keys = slot_keys_unpadded(static_cast(local_chunk)); for_each_chunk_key( - {chunk_keys, static_cast<::cuda::std::size_t>(chunk.count)}, add_hist); + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + static_cast(chunk.count), + add_hist); } } }; @@ -2599,7 +2489,7 @@ private: if (static_cast(threadIdx.x) == owner) { const int slot = bucket - owner * buckets_per_thread; - add_local_selected(local_prefixes[slot]); + atomicAdd(&temp_storage.num_strictly_selected, local_prefixes[slot]); temp_storage.my_candidates = local_hist_vals[slot]; } } @@ -2664,14 +2554,9 @@ private: streamer.is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); } - ::cuda::std::span resident_keys; - // 32-bit shared-window address of `resident_keys.data()`. The resident span is read once per radix pass and in - // the final filter; a 64-bit generic pointer kept live across that loop spills and reloads as a *generic* - // pointer (`LDL.64`), which demotes every key read from `LDS` to a generic `LD`. Carrying the base as a 32-bit - // shared address and rebuilding the pointer with `__cvta_shared_to_generic` at each use keeps the reads `LDS` - // even when the value spills (the cvta intrinsic re-confers shared provenance; a spilled 64-bit generic pointer - // cannot be re-anchored after the fact). - ::cuda::std::uint32_t resident_smem32 = 0; + // Contiguous resident-key window staged in SMEM (block-load path); read once per radix pass and in the final + // filter. Rebound to the real resident window by `load_and_histogram_first_pass`. + smem_keys_t resident_keys; // Front-load all stage mbarrier inits before any bulk copy issues; the barrier below orders them ahead of the first // issue (see `init_load_barriers`). The generic fallback has no mbarriers but keeps the same pass-start barrier. @@ -2682,9 +2567,6 @@ private: __syncthreads(); load_and_histogram_first_pass( - cluster_rank, - layout.leader_rank, - layout.is_single_cta, layout.part, layout.my_resident_chunks, layout.resident_base, @@ -2696,8 +2578,7 @@ private: layout.head_edge_len_items, layout.tail_edge_len_items, streamer, - resident_keys, - resident_smem32); + resident_keys); const int last_pass = run_radix_passes( cluster_rank, @@ -2713,7 +2594,6 @@ private: layout.head_edge_len_items, layout.tail_edge_len_items, resident_keys, - resident_smem32, streamer, kth_key_bits_local); @@ -2752,7 +2632,6 @@ private: layout.segment_size_off, layout.head_items, layout.leader_state, - resident_smem32, resident_keys, layout.head_edge_len_items, layout.tail_edge_len_items, @@ -2777,7 +2656,6 @@ private: layout.my_resident_chunks, layout.segment_size_off, layout.head_items, - resident_smem32, resident_keys, layout.head_edge_len_items, layout.tail_edge_len_items, From 45b80d2b34a72a57cc39a70ea107cbbda1f3ae19 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 07:33:56 +0200 Subject: [PATCH 092/126] Revert unnecessary changes to API tests. --- .../catch2_test_device_batched_topk_api.cu | 52 +++++++++++++++++-- ...catch2_test_device_batched_topk_env_api.cu | 5 -- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index 58ad6fa903c..92d55da4ade 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -1,11 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) -// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See -// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT - #include "insert_nested_NVTX_range_guard.h" #include @@ -273,3 +268,50 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); REQUIRE(keys_out == expected_result_set); } + +// The temporary storage size requirement must not assume a particular base-pointer alignment (the public contract +// states that no special alignment is required). Over-allocate by one byte and offset the base pointer. +C2H_TEST("cub::DeviceBatchedTopK::MaxKeys handles a misaligned temporary storage pointer", "[batched_topk][device]") +{ + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + // Allocate one extra byte and offset the base pointer by one to misalign it. + thrust::device_vector temp_storage(temp_storage_bytes + 1, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()) + 1, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); +} diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index b3f90925e21..f0650bde1db 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -1,11 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) -// so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See -// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT - #include "insert_nested_NVTX_range_guard.h" #include From 7e51a2885e8cb84edb5753f054e37c53e5d33a8d Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 07:39:06 +0200 Subject: [PATCH 093/126] Remove redundant helper --- cub/cub/agent/agent_batched_topk_cluster.cuh | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 7578568550f..68bfdd7e7a9 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -151,17 +151,6 @@ struct smem_block_tile_layout } }; -// True for `cuda::args` segment-size arguments whose exact per-segment value the host cannot know at dispatch time, so -// the launch is sized for a static upper bound: the `deferred` forms. Combined with `!is_single_value` (any -// per-segment sequence) this gates the runtime effective-single-CTA path; host-exact `immediate`/`constant` singles, -// which the dispatch already sizes exactly, are excluded. -template -inline constexpr bool __is_deferred_arg_v = false; -template -inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred<_Arg, _Bounds>> = true; -template -inline constexpr bool __is_deferred_arg_v<::cuda::args::deferred_sequence<_Arg, _Bounds>> = true; - // ----------------------------------------------------------------------------- // Occupancy-free cluster-blocks arithmetic (shared by host dispatch and device agent) // ----------------------------------------------------------------------------- @@ -277,7 +266,8 @@ struct agent_batched_topk_cluster // `immediate`/`constant` singles the dispatch already picks the matching cluster blocks, so the logic is compiled // out. static constexpr bool enable_runtime_single_cta = - !::cuda::args::__traits::is_single_value || __is_deferred_arg_v; + !::cuda::args::__traits::is_single_value + || ::cuda::args::__traits::is_deferred; static constexpr int threads_per_block = ThreadsPerBlock; static constexpr int histogram_items_per_thread = HistogramItemsPerThread; From fe003dafbf48806d915176179cc6e8fb03392d87 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 15:30:53 +0200 Subject: [PATCH 094/126] Add defensive asserts --- cub/cub/agent/agent_batched_topk_cluster.cuh | 107 +++++++++++++++++-- 1 file changed, 96 insertions(+), 11 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 68bfdd7e7a9..38f7b6cd6b4 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -430,7 +430,9 @@ struct agent_batched_topk_cluster [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_desc get_chunk(offset_t chunk_idx, offset_t segment_size, offset_t head_items) const { - const offset_t offset = head_items + chunk_idx * offset_t{chunk_items}; + const offset_t offset = head_items + chunk_idx * offset_t{chunk_items}; + // Callers pass a chunk index `< chunks`, so the chunk is non-empty and unsigned `remaining` below can't underflow. + _CCCL_ASSERT(offset < segment_size, "get_chunk: chunk index lies past the segment end"); const offset_t remaining = segment_size - offset; return {offset, static_cast((::cuda::std::min) (remaining, offset_t{chunk_items}))}; } @@ -446,6 +448,7 @@ struct agent_batched_topk_cluster [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t global_index(offset_t local) const { + _CCCL_ASSERT(local < count, "global_index: rank-local chunk index out of range"); return first + local * stride; } }; @@ -465,6 +468,8 @@ struct agent_batched_topk_cluster [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const { + // Idle ranks never reach here -- `compute_segment_layout` gives them an explicit empty partition. + _CCCL_ASSERT(cluster_rank < cluster_blocks, "make_chunk_partition assumes a working rank (rank < cluster_blocks)"); if constexpr (need_determinism) { const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_blocks)); @@ -499,6 +504,7 @@ struct agent_batched_topk_cluster template _CCCL_DEVICE _CCCL_FORCEINLINE void stage_and_fold_edge(key_t* dst, const key_t* src, int count, Apply&& apply) const { + _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) { const key_t key = src[local]; @@ -591,6 +597,7 @@ struct agent_batched_topk_cluster // per stage, with a matching `wait_stage(stage)`. _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) { + _CCCL_ASSERT(stage >= 0 && stage < PipelineStages, "pipeline stage index out of range"); // Only the block leader (see `is_load_leader`) drives the mbarrier, for a uniform branch and better mbarrier // codegen. if (!is_load_leader) @@ -598,6 +605,9 @@ struct agent_batched_topk_cluster return; } const int num_bytes = static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + // A source is one chunk's aligned bulk (`round_down(count, load_align_items)`): a whole load-align unit, <= a slot. + _CCCL_ASSERT(num_bytes % load_align_bytes == 0, "a bulk copy size must be a whole load-alignment unit"); + _CCCL_ASSERT(num_bytes <= ChunkBytes, "a bulk copy must not exceed one chunk slot"); if (num_bytes > 0) { // The TMA path only requires `bulk_copy_min_align` (>= 16), but the aligned base + `load_align`-multiple bulk @@ -633,6 +643,7 @@ struct agent_batched_topk_cluster // Wait for stage `stage`'s copy to land (all threads spin on its current parity), then flip that stage's parity bit. _CCCL_DEVICE _CCCL_FORCEINLINE void wait_stage(int stage) { + _CCCL_ASSERT(stage >= 0 && stage < PipelineStages, "pipeline stage index out of range"); const ::cuda::std::uint32_t parity = (load_phase >> stage) & 1u; while (!::cuda::ptx::mbarrier_try_wait_parity(&temp_storage.load_mbar[stage], parity)) { @@ -729,6 +740,7 @@ private: // branch. _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket) { + _CCCL_ASSERT(bucket >= 0 && bucket < num_buckets, "histogram bucket index out of range"); const ::cuda::std::uint32_t addr = base32 + static_cast<::cuda::std::uint32_t>(bucket) * sizeof(offset_t); asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); } @@ -779,6 +791,9 @@ private: // Capture `state.k` before the scan: the owning thread overwrites it in the loop below, so any later read would // race with that write. const out_offset_t target_k = temp_storage.state.k; + // Every pass keeps the remaining `k` inside the chosen bucket, so `1 <= target_k <= len`. + _CCCL_ASSERT(target_k >= 1, "leader scan: remaining k must stay positive"); + _CCCL_ASSERT(target_k <= temp_storage.state.len, "leader scan: remaining k cannot exceed the candidate count"); offset_t hist_vals[buckets_per_thread]; offset_t prefixes[buckets_per_thread]; @@ -835,7 +850,7 @@ private: offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk int stream_slot_base; // SMEM slot index at which the streaming region begins offset_t overflow_chunks; // number of overflow chunks for this rank (M) - int eff_stages; // streaming region size = reserved streaming slots (<= PipelineStages, <= M, >= 1) + int eff_stages; // streaming region size = reserved streaming slots (<= PipelineStages; <= M when M>0, else 1) bool is_forward = true; bool is_primed = false; @@ -869,17 +884,28 @@ private: , stream_slot_base(stream_slot_base_) , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) { - // `eff_stages` is the streaming region size the caller reserved at `[stream_slot_base, stream_slot_base + - // stream_slots)`; using all of it as pipeline stages keeps the ping-pong reuse maximal. It is `<= M` whenever - // there is overflow (`M = excess + stream_slots >= stream_slots`); the `>= 1` floor only matters for the no-op - // (`M == 0`) case, where the streamer never touches a slot. + // Use the whole reserved streaming region as pipeline stages for maximal ping-pong reuse; the `1` floor covers + // the no-op (`M == 0`) case, where the streamer never touches a slot (asserted `<= overflow_chunks` below). eff_stages = (stream_slots_ > 0) ? stream_slots_ : 1; + // Both chunk windows must lie inside this rank's `part.count` (resident first, streamed from `overflow_base`). + _CCCL_ASSERT(my_chunks_ == part.count && resident_chunks <= my_chunks_ && overflow_base <= my_chunks_ + && overflow_chunks <= my_chunks_ - overflow_base, + "overflow streamer chunk windows escape the rank's partition"); + // The streaming region is carved from the tile's slots and capped by the pipeline depth. + _CCCL_ASSERT(stream_slot_base >= 0 && stream_slots_ >= 0 && stream_slots_ <= PipelineStages + && static_cast(stream_slot_base) + static_cast(stream_slots_) + <= agent.block_tile_capacity / static_cast(chunk_items), + "overflow streamer streaming slots escape the block tile or pipeline"); _CCCL_ASSERT(overflow_chunks == 0 || eff_stages <= static_cast(overflow_chunks), "streaming depth exceeds the overflow chunk count"); } _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) { + _CCCL_ASSERT(overflow_idx < overflow_chunks, "overflow chunk index out of range"); + _CCCL_ASSERT(stage >= 0 && stage < eff_stages, "overflow stage index exceeds the reserved streaming slots"); + _CCCL_ASSERT((inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, + "cannot issue a load into a streaming stage that is still in flight"); const offset_t chunk_idx = part.global_index(overflow_base + overflow_idx); const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); // Every chunk begins on a `load_align` boundary, so the guard-free aligned (TMA bulk) path applies. The global- @@ -922,6 +948,7 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) { + _CCCL_ASSERT(inflight_mask == ::cuda::std::uint32_t{0}, "an overflow pass must begin with no copies in flight"); if (overflow_chunks == 0) { mid(); @@ -1194,6 +1221,8 @@ private: } if (packed != ::cuda::std::uint64_t{0}) { + // Only working (non-leader) ranks reach here with a nonzero count (see above); idle ranks / the leader push 0. + _CCCL_ASSERT(cluster_rank < eff_cluster_blocks, "a nonzero prefix count must come from a working rank"); if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() @@ -1270,6 +1299,8 @@ private: { if constexpr (!is_keys_only) { + _CCCL_ASSERT(pos < k && seg_idx < segment_size_off, + "value write must land in the top-k output and read inside the segment"); auto block_vals_in = agent.d_value_segments_it[segment_id]; auto block_vals_out = agent.d_value_segments_out_it[segment_id]; block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; @@ -1283,6 +1314,10 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE bool should_stop() { __syncthreads(); + // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. + _CCCL_ASSERT(agent.temp_storage.front_local_cnt <= static_cast(my_front) + && agent.temp_storage.back_local_cnt <= my_cand_count, + "final-filter placement counters exceeded this CTA's assigned work"); const bool is_front_done = agent.temp_storage.front_local_cnt >= static_cast(my_front); // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never // clears it) finishes once all `my_cand_count` of its candidates are placed. @@ -1317,6 +1352,12 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles([[maybe_unused]] const key_t* smem_src, offset_t seg_base, int count, bool region_is_terminal) { + // A non-negative span inside the segment, with a valid SMEM source when read from SMEM (the generic overflow + // fallback passes `nullptr` with `FromSmem == false`). + _CCCL_ASSERT( + count >= 0 && seg_base <= segment_size_off && static_cast(count) <= segment_size_off - seg_base + && (!FromSmem || count == 0 || smem_src != nullptr), + "process_tiles: region must be a valid in-segment source span"); constexpr int items = Items; constexpr int tile = threads_per_block * items; for (int tile_base = 0; tile_base < count; tile_base += tile) @@ -1532,7 +1573,7 @@ private: { if constexpr (use_block_load_to_shared) { - _CCCL_ASSERT(count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); + _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed>(keys, seg_base, count, is_terminal); } } @@ -1654,6 +1695,8 @@ private: [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { if constexpr (!is_keys_only) { + _CCCL_ASSERT(pos < k && seg_idx < segment_size_off, + "value write must land in the top-k output and read inside the segment"); auto block_vals_in = d_value_segments_it[segment_id]; auto block_vals_out = d_value_segments_out_it[segment_id]; block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; @@ -1670,6 +1713,10 @@ private: combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); const offset_t sel_prefix = static_cast(packed_prefix >> 32); const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. + _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) + <= static_cast<::cuda::std::uint64_t>(k), + "selected prefix must fit before the tie-back output region"); // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair @@ -1858,6 +1905,9 @@ private: // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). const offset_t total_candidates = leader_state->len; + // Guards the unsigned `k - num_back` below: the remaining tie count fits both `k` and the splitter bucket. + _CCCL_ASSERT(num_kth <= k && static_cast(num_kth) <= total_candidates, + "remaining k must fit within both the top-k and the splitter bucket"); const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region @@ -1877,6 +1927,9 @@ private: combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); const offset_t sel_prefix = static_cast(packed_prefix >> 32); const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + // Guards the leader's remainder subtractions below (`total_candidates - cand_prefix`, `num_selected - sel_prefix`). + _CCCL_ASSERT(cand_prefix <= total_candidates && sel_prefix <= static_cast(num_selected), + "cross-CTA prefixes must stay within their candidate and selected totals"); // This block's own candidate count: non-leaders hold it in `my_cand`; the leader is last in scan order, so // `cand_prefix` already sums every other block's candidates and `total_candidates - cand_prefix` is its own. A // CTA is `is_select_all_cand_cta` when all of its candidates sit at or below the K-boundary @@ -2022,6 +2075,9 @@ private: { return {}; } + // Chunks start after the aligned head on an alignment-multiple stride (mirrors the streamer's `issue_load`). + _CCCL_ASSERT(::cuda::is_aligned(block_keys_base + chunk.offset, load_align_bytes), + "resident loader received a chunk with an unaligned start"); return {block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; }; @@ -2062,6 +2118,9 @@ private: } } + // Read and write cursors sum the same aligned bulks in the same order, so they meet on a load-align boundary. + _CCCL_ASSERT(read_off_bytes == next_off_bytes && next_off_bytes % load_align_bytes == 0, + "resident bulk read/write cursors must meet on a load-alignment boundary"); // The resident region is one contiguous run of aligned bulks for the later passes; both boundary edges are // folded separately from `edge_keys`. resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), next_off_bytes / int{sizeof(key_t)}); @@ -2103,7 +2162,9 @@ private: } if (tail_edge_len_items > 0) { + _CCCL_ASSERT(chunks > offset_t{0}, "a peeled tail edge requires at least one aligned chunk"); const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_off, head_items); + _CCCL_ASSERT(tail_edge_len_items <= tail_chunk.count, "peeled tail length cannot exceed its source chunk"); stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), tail_edge_len_items, @@ -2173,12 +2234,17 @@ private: if constexpr (use_block_load_to_shared) { layout.block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(layout.block_keys_in); + // `run` is reached only for `0 < k < segment_size`, so a bulk-load segment is non-empty and has a real base. + _CCCL_ASSERT(layout.block_keys_base != nullptr, "a bulk-load segment must have a valid input pointer"); // Items from `block_keys_base` to the next `load_align_bytes` boundary (0 when already aligned), clamped to the // segment. `load_align_bytes` is a multiple of `sizeof(key_t)`, so the aligned pointer lands on a key boundary // and the pointer difference is already an exact item count; narrow it to 32-bit `offset_t` right away. const offset_t head_to_boundary = static_cast(::cuda::align_up(layout.block_keys_base, load_align_bytes) - layout.block_keys_base); layout.head_items = (::cuda::std::min) (head_to_boundary, layout.segment_size_off); + // Distance to the next alignment boundary is < one alignment unit, so the head prefix fits its `edge_keys` slot. + _CCCL_ASSERT(layout.head_items < static_cast(load_align_items), + "the unaligned head prefix must be shorter than one load-alignment unit"); } // Number of aligned chunks covering the segment: the unaligned head prefix (`head_items`) is a separate edge, not a // chunk, so the aligned region `[head_items, segment_size)` chunks uniformly at `chunk_items` (a segment lying @@ -2203,6 +2269,8 @@ private: static_cast<::cuda::std::uint64_t>(layout.chunks), min_chunks_per_block, cluster_blocks); } } + _CCCL_ASSERT(layout.eff_cluster_blocks >= 1u && layout.eff_cluster_blocks <= cluster_blocks, + "effective cluster blocks must stay within [1, launched cluster blocks]"); layout.is_idle_rank = cluster_rank >= layout.eff_cluster_blocks; // Idle ranks own no chunks; `make_chunk_partition` assumes `rank < size`, so hand them an explicit empty partition. @@ -2215,6 +2283,7 @@ private: // needs its own (merged-away) local candidate count: prefer-smallest scans ascending by rank (leader = last // effective rank), prefer-largest scans descending (leader = rank 0). The nondeterministic path keeps rank 0. layout.leader_rank = (need_determinism && !is_tie_reversed) ? (layout.eff_cluster_blocks - 1u) : 0u; + _CCCL_ASSERT(layout.leader_rank < layout.eff_cluster_blocks, "leader must be a working rank"); // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). @@ -2234,6 +2303,10 @@ private: // `stream_slots` is right-sized: clamped into `[1, full_slots]` when streaming; deep overflows still get the full // `PipelineStages` depth. The generic fallback has no async pipeline (it re-reads overflow from gmem each pass and // peels nothing), so it reserves no streaming slots. + // Dispatch guarantees >= 1 whole chunk slot of capacity, so the streaming clamp below is well-defined. + _CCCL_ASSERT(block_tile_capacity >= static_cast(chunk_items) + && block_tile_capacity % static_cast(chunk_items) == offset_t{0}, + "block tile capacity must be a positive whole number of chunk slots"); const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); [[maybe_unused]] const bool needs_streaming = layout.my_chunks > full_slots; @@ -2466,6 +2539,7 @@ private: if (!is_idle_rank) { const int bucket = static_cast(state_t::kth_bucket_of(pass_result)); + _CCCL_ASSERT(bucket >= 0 && bucket < num_buckets, "published splitter bucket index is out of range"); detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); last_pass = pass + 1; @@ -2512,6 +2586,9 @@ private: constexpr int total_bits = int{sizeof(key_t)} * 8; constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + // `process_impl` handles `k == 0` and select-all, so the radix path sees a strict `0 < k < segment_size`. + _CCCL_ASSERT(k > out_offset_t{0} && static_cast(k) < segment_size, + "radix path requires 0 < k < segment size"); const segment_layout layout = compute_segment_layout(segment_id, cluster_rank, cluster_blocks, segment_size); // Per-block local copy of `kth_key_bits` so each key check hits the block's own SMEM rather than DSMEM during the @@ -2595,6 +2672,10 @@ private: // ----------------------------------------------------------------------- auto block_keys_out = d_key_segments_out_it[segment_id]; const out_offset_t num_kth = layout.leader_state->k; // remaining k after the radix passes + // Each pass keeps `0 < num_kth <= k` inside the splitter bucket (same leader read as `num_kth`, equally safe). + _CCCL_ASSERT( + num_kth > out_offset_t{0} && num_kth <= k && static_cast(num_kth) <= layout.leader_state->len, + "radix passes produced an invalid remaining-k count"); // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) @@ -2757,7 +2838,9 @@ private: // Runtime cluster blocks match the launch attribute the dispatch passed // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). const unsigned int cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); - const auto segment_id = static_cast(blockIdx.x / cluster_blocks); + _CCCL_ASSERT(cluster_blocks > 0u && cluster_rank < cluster_blocks, + "hardware cluster rank must lie within a non-empty cluster"); + const auto segment_id = static_cast(blockIdx.x / cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) { @@ -2766,6 +2849,9 @@ private: const auto segment_size = static_cast(detail::params::get_segment_size(segment_sizes, segment_id)); + // Precondition: sizes are clamped >= 0 and capped at 2^21, so a value exceeding 32-bit `offset_t` is a violation. + _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(segment_size) <= ::cuda::std::uint64_t{0xffffffffu}, + "segment size must be non-negative and fit the 32-bit cluster offset type"); // Clamp the requested `k` to the segment size in a 64-bit width holding both operands, *before* narrowing to // `out_offset_t`. Clamping in a narrower type first would let a large "select all" `k` wrap to a small value and // silently truncate the output (or wrap to 0 and skip the segment). @@ -2778,9 +2864,8 @@ private: return; } - // `segment_size` fits `offset_t` by precondition (see the `offset_t` definition above); an out-of-range runtime - // value is undefined behavior, not guarded here. Segments larger than the resident cluster_tile capacity are still - // handled -- the overflow chunks are re-streamed from gmem (see `overflow_streamer`). + // Segments larger than the resident cluster_tile capacity are still handled -- the overflow chunks are re-streamed + // from gmem (see `overflow_streamer`). // `k_clamped <= segment_size`, which now fits `out_offset_t`, so this narrowing is safe. const auto k = static_cast(k_clamped); From fbeddc52cdb4ba06078a293d7859e369f3c0565a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 16:08:51 +0200 Subject: [PATCH 095/126] Move barriers up the call chain Cosmetic change that should improve overview. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 34 +++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 38f7b6cd6b4..85b17b85240 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1308,12 +1308,11 @@ private: } // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and - // resolved its ties. The leading barrier makes the counter reads block-wide (and resynchronizes lanes that raced - // ahead through the barrier-free tiles). Polled only at critical points -- between regions and before each - // streaming bulk copy -- never per tile. + // resolved its ties. Callers must `__syncthreads()` first (the counter reads are block-wide and must resynchronize + // lanes that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and + // before each streaming bulk copy -- never per tile. _CCCL_DEVICE _CCCL_FORCEINLINE bool should_stop() { - __syncthreads(); // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. _CCCL_ASSERT(agent.temp_storage.front_local_cnt <= static_cast(my_front) && agent.temp_storage.back_local_cnt <= my_cand_count, @@ -1558,9 +1557,10 @@ private: }, // No interleaved resident work: the deterministic filter folds its resident span separately. [] {}, - // Checked before each refill bulk copy: break the stream once the whole top-k is placed. `should_stop`'s - // barrier also resynchronizes the lanes that drifted through the just-folded chunk's barrier-free tiles. + // Break the stream once the whole top-k is placed. The barrier makes the counter reads block-wide and resyncs + // lanes that drifted through the just-folded chunk's barrier-free tiles (polled before each refill copy). [&] { + __syncthreads(); return !should_stop(); }); } @@ -1606,6 +1606,9 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter() { const auto step = [&](auto&& region) { + // Barrier before polling: makes the placement-counter reads block-wide and resynchronizes lanes that raced + // ahead through the previous region's barrier-free tiles. + __syncthreads(); if (!should_stop()) { region(); @@ -1703,7 +1706,6 @@ private: } }; - __syncthreads(); const bool participates = !is_idle_rank && (cluster_rank != leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; @@ -1911,9 +1913,6 @@ private: const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region - // Publish the last pass's `num_strictly_selected`/`my_candidates` (written by the owning lane after the final - // cluster barrier) block-wide before they feed the scan and `front_count`. - __syncthreads(); const bool participates = !is_idle_rank && (cluster_rank != leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; @@ -2018,9 +2017,9 @@ private: // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 // needs no candidate filtering). Publishes the resident span as `resident_keys` for the later passes and the final - // filter; ends on a `__syncthreads()` that makes `edge_keys` and the block-local histogram visible within the block. - // Cross-CTA visibility of the zeroed histogram comes later, from the deferred initial cluster wait in - // `run_radix_passes` (this whole load runs in that arrive->wait window). + // filter. The caller's `__syncthreads()` makes `edge_keys` and the block-local histogram block-visible; cross-CTA + // visibility of the zeroed histogram comes later, from the deferred initial cluster wait in `run_radix_passes` (this + // whole load runs in that arrive->wait window). template _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( const chunk_partition& part, @@ -2152,8 +2151,6 @@ private: // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. - // The `__syncthreads()` below publishes `edge_keys` for the later passes and the final filter (which read it via - // `fold_edges` after a barrier). if constexpr (use_block_load_to_shared) { if (head_edge_len_items > 0) @@ -2181,7 +2178,6 @@ private: const int resident_count = static_cast(resident_keys.size()); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, "Dynamic shared memory block_tile is too small"); - __syncthreads(); } // Per-segment, per-rank geometry computed once at the top of `run`: the head-aligned chunking, the effective @@ -2647,6 +2643,9 @@ private: streamer, resident_keys); + // Publish the first pass's staged `edge_keys` and per-rank histogram block-wide before the radix passes read them. + __syncthreads(); + const int last_pass = run_radix_passes( cluster_rank, layout.leader_rank, @@ -2682,6 +2681,9 @@ private: // trailing digits as smaller and erroneously reject candidates that share the identified prefix. identify_candidates_op_t identify_op(&kth_key_bits_local, last_pass, total_bits, decomposer_t{}); + // Publish the final pass's per-rank `num_strictly_selected`/`my_candidates` (written by one lane after the last + // cluster barrier) block-wide before the final-filter scan below reads them. + __syncthreads(); if constexpr (need_determinism) { write_deterministic_topk( From 7148ce429948c6292e58c353b8648a921b3568cc Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 17:07:41 +0200 Subject: [PATCH 096/126] Use member variables to avoid many arguments --- cub/cub/agent/agent_batched_topk_cluster.cuh | 548 ++++++++----------- 1 file changed, 219 insertions(+), 329 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 85b17b85240..8d26a78cad4 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -505,7 +505,7 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void stage_and_fold_edge(key_t* dst, const key_t* src, int count, Apply&& apply) const { _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); - for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) + for (int local = tid; local < count; local += threads_per_block) { const key_t key = src[local]; dst[local] = key; @@ -523,7 +523,6 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* keys, int chunk_count, Apply&& apply) const { constexpr int tile = Unroll * threads_per_block; - const int tid = static_cast(threadIdx.x); const int full_tiles = ::cuda::round_down(chunk_count, tile); _CCCL_PRAGMA_NOUNROLL() @@ -585,7 +584,7 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void init_load_barriers() { _CCCL_PRAGMA_NOUNROLL() - for (int stage = static_cast(threadIdx.x); stage < PipelineStages; stage += threads_per_block) + for (int stage = tid; stage < PipelineStages; stage += threads_per_block) { ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[stage], 1u); } @@ -654,6 +653,34 @@ struct agent_batched_topk_cluster struct TempStorage : Uninitialized<_TempStorage> {}; + // Per-segment, per-rank geometry computed once by `compute_segment_layout` at the top of `run`: the head-aligned + // chunking, the effective (non-idle) cluster width and this rank's partition of it, the leader/idle roles, and the + // resident/streaming/edge split of this rank's chunks. All fields are loop-invariant for the segment; the + // streaming-only intermediates (`overflow_base`, `resident_slots_cap`, `stream_slots`, `my_chunks`) are kept because + // the `overflow_streamer` constructor needs them. + struct segment_layout + { + offset_t segment_size_off; + key_it_t block_keys_in; + const key_t* block_keys_base; + offset_t head_items; + offset_t chunks; + unsigned int eff_cluster_blocks; + bool is_idle_rank; + chunk_partition part; + unsigned int leader_rank; + state_t* leader_state; + offset_t my_chunks; + offset_t my_resident_chunks; + offset_t overflow_chunks; + offset_t resident_base; + offset_t overflow_base; + offset_t resident_slots_cap; + offset_t stream_slots; + int head_edge_len_items; + int tail_edge_len_items; + }; + // --------------------------------------------------------------------------- // Per-thread members // --------------------------------------------------------------------------- @@ -674,6 +701,22 @@ struct agent_batched_topk_cluster // The single block leader (warp 0's elected lane) that drives the bulk copies. Elected once at construction (the // block constructs convergently) and cached so the pipeline reuses it instead of re-electing per copy. const bool is_load_leader = ::cuda::device::__block_elect_one(); + // This thread's block-local index, signed for the agent's block-stride loops and cached to avoid re-casting + // `threadIdx.x` at each site. + const int tid = static_cast(threadIdx.x); + // Effective cluster geometry the radix/scan/filter path runs on, set once by `init_effective_cluster`: the launched + // cluster, or rank 0 alone (`cluster_blocks == 1`) when a small segment collapsed onto a single resident CTA. The + // select-all fast path predates the collapse and uses the raw hardware sregs instead. + unsigned int cluster_rank = 0u; + unsigned int cluster_blocks = 1u; + bool is_single_cta = true; + // Per-segment constants for this block, set once and then read directly by the helpers instead of being threaded + // through their signatures. `segment_id`/`segment_size`/`k` are set in `process_impl`; `layout` is filled by + // `compute_segment_layout` on the radix path (the select-all fast path and redundant CTAs never touch it). + num_segments_val_t segment_id{}; + segment_size_val_t segment_size{}; + out_offset_t k{}; + segment_layout layout{}; _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, @@ -713,7 +756,7 @@ struct agent_batched_topk_cluster private: _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() { - for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + for (int i = tid; i < num_buckets; i += threads_per_block) { temp_storage.hist[i] = 0; } @@ -801,7 +844,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < buckets_per_thread; ++j) { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + const int bucket = tid * buckets_per_thread + j; hist_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; } @@ -812,7 +855,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < buckets_per_thread; ++j) { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + const int bucket = tid * buckets_per_thread + j; if (bucket < num_buckets && prefixes[j] < target_k && prefixes[j] + hist_vals[j] >= target_k) { const out_offset_t new_k = target_k - static_cast(prefixes[j]); @@ -1096,7 +1139,7 @@ private: _CCCL_PRAGMA_UNROLL(UnrollFactor) for (int j = 0; j < iterations; ++j) { - const int local = j * threads_per_block + static_cast(threadIdx.x); + const int local = j * threads_per_block + agent.tid; if (local < chunk.count) { if constexpr (Indexed) @@ -1165,8 +1208,8 @@ private: } // Synchronize the segment's cluster. A single-CTA "cluster" keeps all state block-local, so `__syncthreads()` orders - // it and the cluster-scoped barrier is unnecessary. `is_single_cta` is computed in `run()` from the collapsed cluster - // blocks `process_impl` passes in (1 when a small segment collapsed onto rank 0), not the raw cluster size. It is + // it and the cluster-scoped barrier is unnecessary. `is_single_cta` is the agent's effective-cluster member (see + // `init_effective_cluster`): 1 when a small segment collapsed onto rank 0, not the raw cluster size. It is // per-segment uniform across the surviving block(s), so the branch is reached uniformly. _CCCL_DEVICE _CCCL_FORCEINLINE static void cluster_or_block_sync(bool is_single_cta) { @@ -1212,8 +1255,7 @@ private: // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. - _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan( - bool is_single_cta, unsigned int cluster_rank, unsigned int eff_cluster_blocks, ::cuda::std::uint64_t packed) + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan(::cuda::std::uint64_t packed) { if (is_single_cta) { @@ -1222,7 +1264,7 @@ private: if (packed != ::cuda::std::uint64_t{0}) { // Only working (non-leader) ranks reach here with a nonzero count (see above); idle ranks / the leader push 0. - _CCCL_ASSERT(cluster_rank < eff_cluster_blocks, "a nonzero prefix count must come from a working rank"); + _CCCL_ASSERT(cluster_rank < layout.eff_cluster_blocks, "a nonzero prefix count must come from a working rank"); if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() @@ -1237,7 +1279,8 @@ private: // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. _CCCL_PRAGMA_NOUNROLL() - for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < eff_cluster_blocks; rank += threads_per_block) + for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < layout.eff_cluster_blocks; + rank += threads_per_block) { add_remote_prefix(rank, packed); } @@ -1368,7 +1411,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < items; ++i) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const int pos = tile_base + agent.tid * items + i; is_valid[i] = pos < count; flags[i] = offset_t{0}; if (is_valid[i]) @@ -1398,7 +1441,7 @@ private: const bool is_front_key = is_valid[i] && (cls[i] == detail::topk::candidate_class::selected); if (is_front_key) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const int pos = tile_base + agent.tid * items + i; const offset_t local = atomicAdd(&agent.temp_storage.front_local_cnt, offset_t{1}); const out_offset_t out = static_cast(sel_prefix + local); const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); @@ -1424,7 +1467,7 @@ private: { if (is_valid[i] && flags[i] != offset_t{0}) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const int pos = tile_base + agent.tid * items + i; emit_back_one(cand_prefix + atomicAdd(&agent.temp_storage.back_local_cnt, offset_t{1}), keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); @@ -1442,7 +1485,7 @@ private: { if (is_valid[i] && flags[i] != offset_t{0}) { - const int pos = tile_base + static_cast(threadIdx.x) * items + i; + const int pos = tile_base + agent.tid * items + i; emit_back_one( base + excl[i], keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); } @@ -1650,7 +1693,6 @@ private: { if constexpr (use_block_load_to_shared) { - const int tid = static_cast(threadIdx.x); _CCCL_PRAGMA_NOUNROLL() for (int local = tid; local < head_edge_len_items; local += threads_per_block) { @@ -1670,25 +1712,10 @@ private: // perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( - num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int eff_cluster_blocks, - unsigned int leader_rank, - bool is_single_cta, - bool is_idle_rank, - out_offset_t k, out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, - const key_t* block_keys_base, - const chunk_partition& part, - offset_t resident_base, - offset_t my_resident_chunks, - offset_t segment_size_off, - offset_t head_items, smem_keys_t resident_keys, - int head_edge_len_items, - int tail_edge_len_items, overflow_streamer& streamer) { // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index @@ -1698,7 +1725,7 @@ private: [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { if constexpr (!is_keys_only) { - _CCCL_ASSERT(pos < k && seg_idx < segment_size_off, + _CCCL_ASSERT(pos < k && seg_idx < layout.segment_size_off, "value write must land in the top-k output and read inside the segment"); auto block_vals_in = d_value_segments_it[segment_id]; auto block_vals_out = d_value_segments_out_it[segment_id]; @@ -1706,15 +1733,14 @@ private: } }; - const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = - combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) <= static_cast<::cuda::std::uint64_t>(k), @@ -1769,10 +1795,10 @@ private: } else { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); + const offset_t chunk_idx = layout.part.global_index(local_chunk); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); for_each_chunk_key( ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), static_cast(chunk.count), @@ -1780,7 +1806,7 @@ private: } } // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_boundary_edges(head_edge_len_items, tail_edge_len_items, write_selected); + fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, write_selected); }; streamer.process_pass(write_selected, fold_resident); } @@ -1800,7 +1826,7 @@ private: _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) { - const int local = j * threads_per_block + static_cast(threadIdx.x); + const int local = j * threads_per_block + tid; if (local < count) { const key_t key = smem[local]; @@ -1812,7 +1838,7 @@ private: // Boundary edges span fewer than `load_align_items` items, so fold them with a plain block-stride loop rather // than the tiled/unrolled `write_run` used for the (potentially large) resident bulks. const auto write_edge = [&](const key_t* smem, offset_t base_off, int count) { - for (int local = static_cast(threadIdx.x); local < count; local += threads_per_block) + for (int local = tid; local < count; local += threads_per_block) { const key_t key = smem[local]; sink(key, base_off + static_cast(local)); @@ -1829,9 +1855,10 @@ private: // `chunk.count`. key_t* const resident_ptr = resident_keys.data(); int cursor_items = 0; - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(resident_base + local_chunk), segment_size_off, head_items); + const auto chunk = get_chunk( + layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); const offset_t base_off = chunk.offset; const int bulk_count_items = static_cast( ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); @@ -1841,9 +1868,10 @@ private: } else { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const auto chunk = get_chunk(part.global_index(local_chunk), segment_size_off, head_items); + const auto chunk = + get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); const offset_t base_off = chunk.offset; write_run( ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), base_off, chunk.count); @@ -1853,15 +1881,15 @@ private: // peeled tail suffix at `segment_size - tail_edge_len_items`. Value fetched per key. if constexpr (use_block_load_to_shared) { - if (head_edge_len_items > 0) + if (layout.head_edge_len_items > 0) { - write_edge(temp_storage.edge_keys, offset_t{0}, head_edge_len_items); + write_edge(temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items); } - if (tail_edge_len_items > 0) + if (layout.tail_edge_len_items > 0) { write_edge(temp_storage.edge_keys + head_edge_cap_items, - segment_size_off - static_cast(tail_edge_len_items), - tail_edge_len_items); + layout.segment_size_off - static_cast(layout.tail_edge_len_items), + layout.tail_edge_len_items); } } }; @@ -1879,41 +1907,24 @@ private: // `det_final_filter`; this member computes that struct's inputs and runs it. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( - num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int eff_cluster_blocks, - unsigned int leader_rank, - bool is_single_cta, - bool is_idle_rank, - out_offset_t k, out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, - key_it_t block_keys_in, - const chunk_partition& part, - offset_t resident_base, - offset_t my_resident_chunks, - offset_t overflow_chunks, - offset_t segment_size_off, - offset_t head_items, - state_t* leader_state, smem_keys_t resident_keys, - int head_edge_len_items, - int tail_edge_len_items, overflow_streamer& streamer) { // Early stop is not special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. // // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). - const offset_t total_candidates = leader_state->len; + const offset_t total_candidates = layout.leader_state->len; // Guards the unsigned `k - num_back` below: the remaining tie count fits both `k` and the splitter bucket. _CCCL_ASSERT(num_kth <= k && static_cast(num_kth) <= total_candidates, "remaining k must fit within both the top-k and the splitter bucket"); const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region - const bool participates = !is_idle_rank && (cluster_rank != leader_rank); + const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; // Front count pushed by this block: its strictly-selected count. Candidates always route through the back, so @@ -1922,10 +1933,9 @@ private: const offset_t push_front = my_sel; const ::cuda::std::uint64_t packed = (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = - combined_prefix_scan(is_single_cta, cluster_rank, eff_cluster_blocks, packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // Guards the leader's remainder subtractions below (`total_candidates - cand_prefix`, `num_selected - sel_prefix`). _CCCL_ASSERT(cand_prefix <= total_candidates && sel_prefix <= static_cast(num_selected), "cross-CTA prefixes must stay within their candidate and selected totals"); @@ -1935,7 +1945,7 @@ private: // (`cand_prefix + my_cand_count <= num_back`): every one wins, so the back places them with arrival-order SMEM // atomics and skips the index-ordered scan. While `is_tie_active`, a non-`is_select_all_cand_cta` CTA is the // single boundary-crossing (straddling) CTA cluster-wide. - const offset_t my_cand_count = (cluster_rank == leader_rank) ? (total_candidates - cand_prefix) : my_cand; + const offset_t my_cand_count = (cluster_rank == layout.leader_rank) ? (total_candidates - cand_prefix) : my_cand; const bool is_select_all_cand_cta = (cand_prefix + my_cand_count) <= static_cast(num_back); // Mirror image: a CTA selects none of its candidates when the tie region is empty (`num_back == 0`) or all of its // candidates sort strictly after the K-boundary (`cand_prefix >= num_back`). Such a CTA seeds `is_tie_active` @@ -1945,7 +1955,7 @@ private: // This block's own front size: non-leaders know it directly (`push_front`); the leader is last in scan order, so // `sel_prefix` already sums every other block's front and `num_selected - sel_prefix` is the remainder it owns. const out_offset_t my_front = - (cluster_rank == leader_rank) + (cluster_rank == layout.leader_rank) ? static_cast(num_selected - static_cast(sel_prefix)) : static_cast(push_front); @@ -1957,10 +1967,10 @@ private: // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. - const bool has_head = head_edge_len_items > 0; + const bool has_head = layout.head_edge_len_items > 0; const bool has_resident = front_count > 0; - const bool has_overflow = overflow_chunks > offset_t{0}; - const bool has_tail_edge = tail_edge_len_items > 0; + const bool has_overflow = layout.overflow_chunks > offset_t{0}; + const bool has_tail_edge = layout.tail_edge_len_items > 0; [[maybe_unused]] const bool is_resident_terminal = is_tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); [[maybe_unused]] const bool is_head_edge_terminal = @@ -1975,7 +1985,7 @@ private: // local chunk under `is_residency_reversed`. const offset_t front_seg_base = (front_count > 0) - ? get_chunk(part.global_index(resident_base), segment_size_off, head_items).offset + ? get_chunk(layout.part.global_index(layout.resident_base), layout.segment_size_off, layout.head_items).offset : offset_t{0}; // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` @@ -1987,24 +1997,24 @@ private: segment_id, identify_op, block_keys_out, - block_keys_in, + layout.block_keys_in, streamer, - part, + layout.part, k, num_back, my_front, sel_prefix, cand_prefix, my_cand_count, - resident_base, - my_resident_chunks, - segment_size_off, - head_items, + layout.resident_base, + layout.my_resident_chunks, + layout.segment_size_off, + layout.head_items, front_seg_base, resident_keys, front_count, - head_edge_len_items, - tail_edge_len_items, + layout.head_edge_len_items, + layout.tail_edge_len_items, is_select_all_cand_cta, is_resident_terminal, is_head_edge_terminal, @@ -2021,19 +2031,8 @@ private: // visibility of the zeroed histogram comes later, from the deferred initial cluster wait in `run_radix_passes` (this // whole load runs in that arrive->wait window). template - _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass( - const chunk_partition& part, - offset_t my_resident_chunks, - offset_t resident_base, - offset_t chunks, - offset_t segment_size_off, - offset_t head_items, - key_it_t block_keys_in, - const key_t* block_keys_base, - int head_edge_len_items, - int tail_edge_len_items, - overflow_streamer& streamer, - smem_keys_t& resident_keys) + _CCCL_DEVICE _CCCL_FORCEINLINE void + load_and_histogram_first_pass(overflow_streamer& streamer, smem_keys_t& resident_keys) { using extract_bin_op_t = detail::topk::extract_bin_op_t; constexpr int total_bits = int{sizeof(key_t)} * 8; @@ -2047,7 +2046,7 @@ private: if constexpr (use_block_load_to_shared) { - if (my_resident_chunks > 0) + if (layout.my_resident_chunks > 0) { // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor @@ -2056,18 +2055,18 @@ private: // begins on a `load_align` boundary (zero prefix), so its aligned bulk is `round_down(count, // load_align_items)`: the whole `count` for interior chunks, and for the global-last chunk its `count` minus // the unaligned suffix that is always peeled into `edge_keys`. - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(my_resident_chunks)); + const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` // shifts the resident window to the high-index chunks). const auto resident_local = [&](offset_t slot) -> offset_t { - return resident_base + slot; + return layout.resident_base + slot; }; // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has // none. const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { - const offset_t chunk_idx = part.global_index(resident_local(slot)); - const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); + const offset_t chunk_idx = layout.part.global_index(resident_local(slot)); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); const offset_t bulk = ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); if (bulk == 0) @@ -2075,9 +2074,9 @@ private: return {}; } // Chunks start after the aligned head on an alignment-multiple stride (mirrors the streamer's `issue_load`). - _CCCL_ASSERT(::cuda::is_aligned(block_keys_base + chunk.offset, load_align_bytes), + _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), "resident loader received a chunk with an unaligned start"); - return {block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; + return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; }; // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives @@ -2096,7 +2095,7 @@ private: // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs // no stored per-stage state. int read_off_bytes = 0; - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const int stage = static_cast(local_chunk % static_cast(prologue)); wait_stage(stage); @@ -2106,7 +2105,7 @@ private: read_off_bytes += read_len_items * int{sizeof(key_t)}; const offset_t next_local_chunk = local_chunk + static_cast(prologue); - if (next_local_chunk < my_resident_chunks) + if (next_local_chunk < layout.my_resident_chunks) { const auto src = bulk_src(next_local_chunk); // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads @@ -2127,20 +2126,20 @@ private: } else { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(resident_base + local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); + const offset_t chunk_idx = layout.part.global_index(layout.resident_base + local_chunk); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); key_t* const chunk_keys = ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes); const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); _CCCL_PRAGMA_UNROLL(histogram_items_per_thread_clamped) for (int j = 0; j < iterations; ++j) { - const int local = j * threads_per_block + static_cast(threadIdx.x); + const int local = j * threads_per_block + tid; if (local < chunk.count) { const key_t key = - block_keys_in[static_cast(chunk.offset + static_cast(local))]; + layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]; chunk_keys[local] = key; add_first_pass(key); } @@ -2153,19 +2152,21 @@ private: // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. if constexpr (use_block_load_to_shared) { - if (head_edge_len_items > 0) + if (layout.head_edge_len_items > 0) { - stage_and_fold_edge(temp_storage.edge_keys, block_keys_base, head_edge_len_items, add_first_pass); + stage_and_fold_edge(temp_storage.edge_keys, layout.block_keys_base, layout.head_edge_len_items, add_first_pass); } - if (tail_edge_len_items > 0) + if (layout.tail_edge_len_items > 0) { - _CCCL_ASSERT(chunks > offset_t{0}, "a peeled tail edge requires at least one aligned chunk"); - const auto tail_chunk = get_chunk(chunks - offset_t{1}, segment_size_off, head_items); - _CCCL_ASSERT(tail_edge_len_items <= tail_chunk.count, "peeled tail length cannot exceed its source chunk"); - stage_and_fold_edge(temp_storage.edge_keys + head_edge_cap_items, - block_keys_base + tail_chunk.offset + (tail_chunk.count - tail_edge_len_items), - tail_edge_len_items, - add_first_pass); + _CCCL_ASSERT(layout.chunks > offset_t{0}, "a peeled tail edge requires at least one aligned chunk"); + const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_off, layout.head_items); + _CCCL_ASSERT(layout.tail_edge_len_items <= tail_chunk.count, + "peeled tail length cannot exceed its source chunk"); + stage_and_fold_edge( + temp_storage.edge_keys + head_edge_cap_items, + layout.block_keys_base + tail_chunk.offset + (tail_chunk.count - layout.tail_edge_len_items), + layout.tail_edge_len_items, + add_first_pass); } } @@ -2180,50 +2181,15 @@ private: "Dynamic shared memory block_tile is too small"); } - // Per-segment, per-rank geometry computed once at the top of `run`: the head-aligned chunking, the effective - // (non-idle) cluster width and this rank's partition of it, the leader/idle roles, and the resident/streaming/edge - // split of this rank's chunks. All fields are loop-invariant for the segment; the streaming-only intermediates - // (`overflow_base`, `resident_slots_cap`, `stream_slots`, `my_chunks`) are kept because the `overflow_streamer` - // constructor needs them. - struct segment_layout - { - offset_t segment_size_off; - bool is_single_cta; - key_it_t block_keys_in; - const key_t* block_keys_base; - offset_t head_items; - offset_t chunks; - unsigned int eff_cluster_blocks; - bool is_idle_rank; - chunk_partition part; - unsigned int leader_rank; - state_t* leader_state; - offset_t my_chunks; - offset_t my_resident_chunks; - offset_t overflow_chunks; - offset_t resident_base; - offset_t overflow_base; - offset_t resident_slots_cap; - offset_t stream_slots; - int head_edge_len_items; - int tail_edge_len_items; - }; - - _CCCL_DEVICE _CCCL_FORCEINLINE segment_layout compute_segment_layout( - num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int cluster_blocks, - segment_size_val_t segment_size) + // Fill the `layout` member for this block's `segment_id`/`segment_size`. Called once at the top of `run`. + _CCCL_DEVICE _CCCL_FORCEINLINE void compute_segment_layout() { - segment_layout layout; layout.block_keys_in = d_key_segments_it[segment_id]; layout.segment_size_off = static_cast(segment_size); - // `cluster_blocks` is what `process_impl` runs at: the launched cluster size, or `1` when it collapsed a - // small segment onto rank 0. A lone CTA routes barriers to `__syncthreads()` and keeps `state`/atomics block-local - // (no cross-rank DSMEM folds); its histogram increments use the same unconditional cluster-scope `hist_inc`, which - // is identical SASS at cluster size 1. For wider clusters, `eff_cluster_blocks` (below) further excludes ranks that + // A lone CTA (`is_single_cta`) routes barriers to `__syncthreads()` and keeps `state`/atomics block-local (no + // cross-rank DSMEM folds); its histogram increments still use the unconditional cluster-scope `hist_inc`, which is + // identical SASS at cluster size 1. For wider clusters, `eff_cluster_blocks` (below) further excludes ranks that // receive no chunks; they stay resident but idle. - layout.is_single_cta = (cluster_blocks == 1u); layout.block_keys_base = nullptr; layout.head_items = 0; @@ -2259,7 +2225,7 @@ private: layout.eff_cluster_blocks = cluster_blocks; if constexpr (enable_runtime_single_cta) { - if (!layout.is_single_cta) + if (!is_single_cta) { layout.eff_cluster_blocks = effective_cluster_blocks_from_chunks( static_cast<::cuda::std::uint64_t>(layout.chunks), min_chunks_per_block, cluster_blocks); @@ -2283,7 +2249,7 @@ private: // DSMEM pointer into the leader block's shared memory. The Step 2 histogram fold reaches the leader's `hist` // through a `mapa`-formed `shared::cluster` address instead (see `hist_fold_remote`). - layout.leader_state = layout.is_single_cta ? &temp_storage.state : map_state_to_rank(layout.leader_rank); + layout.leader_state = is_single_cta ? &temp_storage.state : map_state_to_rank(layout.leader_rank); // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every // cluster barrier is reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) @@ -2359,29 +2325,14 @@ private: // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_items) : 0; layout.tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; - return layout; } // Radix histogram/scan passes: refine the splitter one `bits_per_pass` digit at a time until the top-k prefix is // pinned down (or early stop fires). Folds `kth_key_bits_local` up digit by digit and returns the number of passes // that actually ran (`last_pass`), which the final filter uses to size its identify operator. template - _CCCL_DEVICE _CCCL_FORCEINLINE int run_radix_passes( - unsigned int cluster_rank, - unsigned int leader_rank, - bool is_single_cta, - bool is_idle_rank, - const chunk_partition& part, - offset_t my_resident_chunks, - offset_t resident_base, - offset_t segment_size_off, - offset_t head_items, - state_t* leader_state, - int head_edge_len_items, - int tail_edge_len_items, - smem_keys_t resident_keys, - overflow_streamer& streamer, - key_prefix_t& kth_key_bits_local) + _CCCL_DEVICE _CCCL_FORCEINLINE int + run_radix_passes(smem_keys_t resident_keys, overflow_streamer& streamer, key_prefix_t& kth_key_bits_local) { using extract_bin_op_t = detail::topk::extract_bin_op_t; using identify_candidates_op_t = @@ -2424,10 +2375,10 @@ private: } else { - for (offset_t local_chunk = 0; local_chunk < my_resident_chunks; ++local_chunk) + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const offset_t chunk_idx = part.global_index(resident_base + local_chunk); - const auto chunk = get_chunk(chunk_idx, segment_size_off, head_items); + const offset_t chunk_idx = layout.part.global_index(layout.resident_base + local_chunk); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); for_each_chunk_key( ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), static_cast(chunk.count), @@ -2444,7 +2395,7 @@ private: // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's per-bucket counts (the source of its `num_strictly_selected` // and `my_candidates`, the cross-CTA scan inputs) inclusive of its edge candidates. - fold_boundary_edges(head_edge_len_items, tail_edge_len_items, add_hist); + fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, add_hist); } // Local barrier is enough here: all Step 1 / Step 2 writes to `hist[]` are atomic at compatible scopes (see Step @@ -2467,16 +2418,16 @@ private: // leader skips this to avoid double-counting its own contribution; // idle ranks (`>= eff_cluster_blocks`) have an all-zero histogram, so // they skip the fold entirely (the loop would only read zeros). - if (cluster_rank != leader_rank && !is_idle_rank) + if (cluster_rank != layout.leader_rank && !layout.is_idle_rank) { const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - for (int i = static_cast(threadIdx.x); i < num_buckets; i += threads_per_block) + for (int i = tid; i < num_buckets; i += threads_per_block) { const offset_t bucket_count = temp_storage.hist[i]; if (bucket_count != 0) { hist_fold_remote( - hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), bucket_count, leader_rank); + hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), bucket_count, layout.leader_rank); } } } @@ -2495,12 +2446,12 @@ private: // scan in registers lets `hist` reset on the normal schedule (the regs survive the reset and the next sync). offset_t local_prefixes[buckets_per_thread]{}; offset_t local_hist_vals[buckets_per_thread]{}; - if (cluster_rank != leader_rank && !is_idle_rank) + if (cluster_rank != layout.leader_rank && !layout.is_idle_rank) { _CCCL_PRAGMA_UNROLL_FULL() for (int j = 0; j < buckets_per_thread; ++j) { - const int bucket = static_cast(threadIdx.x) * buckets_per_thread + j; + const int bucket = tid * buckets_per_thread + j; local_hist_vals[j] = (bucket < num_buckets) ? temp_storage.hist[bucket] : offset_t{0}; } block_scan_t(temp_storage.scan_storage).ExclusiveSum(local_hist_vals, local_prefixes); @@ -2511,7 +2462,7 @@ private: // Step 3 (leader half): the leader prefix-scans the merged `hist` (raw counts, all folds now visible) and updates // the cluster-shared `state`. Subsequent reads (end-of-pass fold, last filter) observe these writes after the // next cluster sync. - if (cluster_rank == leader_rank) + if (cluster_rank == layout.leader_rank) { leader_identify_kth_bucket(); } @@ -2531,8 +2482,8 @@ private: // that one load: working threads fold the `kth_bucket` digit into their own `kth_key_bits_local` (so the full // splitter key is reconstructed locally without a broadcast), and every block -- including idle ranks -- reuses // the same value for the uniform early-stop check below, so the whole cluster breaks together. - const ::cuda::std::uint64_t pass_result = leader_state->result_pair; - if (!is_idle_rank) + const ::cuda::std::uint64_t pass_result = layout.leader_state->result_pair; + if (!layout.is_idle_rank) { const int bucket = static_cast(state_t::kth_bucket_of(pass_result)); _CCCL_ASSERT(bucket >= 0 && bucket < num_buckets, "published splitter bucket index is out of range"); @@ -2543,10 +2494,10 @@ private: // the scan above. Accumulate the strictly-selected count and overwrite `my_candidates` with this pass's // splitter-bucket count (the last pass's value is what the filter reads, however the loop exits). The leader's // `hist` is merged, so it derives its own counts from the scan total instead (see the filter). - if (cluster_rank != leader_rank) + if (cluster_rank != layout.leader_rank) { const int owner = bucket / buckets_per_thread; - if (static_cast(threadIdx.x) == owner) + if (tid == owner) { const int slot = bucket - owner * buckets_per_thread; atomicAdd(&temp_storage.num_strictly_selected, local_prefixes[slot]); @@ -2569,12 +2520,7 @@ private: } template - _CCCL_DEVICE _CCCL_FORCEINLINE void - run(num_segments_val_t segment_id, - unsigned int cluster_rank, - unsigned int cluster_blocks, - segment_size_val_t segment_size, - out_offset_t k) + _CCCL_DEVICE _CCCL_FORCEINLINE void run() { using identify_candidates_op_t = detail::topk::identify_candidates_op_t; @@ -2585,7 +2531,7 @@ private: // `process_impl` handles `k == 0` and select-all, so the radix path sees a strict `0 < k < segment_size`. _CCCL_ASSERT(k > out_offset_t{0} && static_cast(k) < segment_size, "radix path requires 0 < k < segment size"); - const segment_layout layout = compute_segment_layout(segment_id, cluster_rank, cluster_blocks, segment_size); + compute_segment_layout(); // Per-block local copy of `kth_key_bits` so each key check hits the block's own SMEM rather than DSMEM during the // histogram loop. Built up one splitter digit per pass from the leader's published `kth_bucket` (see @@ -2629,39 +2575,12 @@ private: } __syncthreads(); - load_and_histogram_first_pass( - layout.part, - layout.my_resident_chunks, - layout.resident_base, - layout.chunks, - layout.segment_size_off, - layout.head_items, - layout.block_keys_in, - layout.block_keys_base, - layout.head_edge_len_items, - layout.tail_edge_len_items, - streamer, - resident_keys); + load_and_histogram_first_pass(streamer, resident_keys); // Publish the first pass's staged `edge_keys` and per-rank histogram block-wide before the radix passes read them. __syncthreads(); - const int last_pass = run_radix_passes( - cluster_rank, - layout.leader_rank, - layout.is_single_cta, - layout.is_idle_rank, - layout.part, - layout.my_resident_chunks, - layout.resident_base, - layout.segment_size_off, - layout.head_items, - layout.leader_state, - layout.head_edge_len_items, - layout.tail_edge_len_items, - resident_keys, - streamer, - kth_key_bits_local); + const int last_pass = run_radix_passes(resident_keys, streamer, kth_key_bits_local); // ----------------------------------------------------------------------- // Final filter pass: write the top-k keys for this segment. Strictly- @@ -2686,53 +2605,11 @@ private: __syncthreads(); if constexpr (need_determinism) { - write_deterministic_topk( - segment_id, - cluster_rank, - layout.eff_cluster_blocks, - layout.leader_rank, - layout.is_single_cta, - layout.is_idle_rank, - k, - num_kth, - identify_op, - block_keys_out, - layout.block_keys_in, - layout.part, - layout.resident_base, - layout.my_resident_chunks, - layout.overflow_chunks, - layout.segment_size_off, - layout.head_items, - layout.leader_state, - resident_keys, - layout.head_edge_len_items, - layout.tail_edge_len_items, - streamer); + write_deterministic_topk(num_kth, identify_op, block_keys_out, resident_keys, streamer); } else { - write_nondeterministic_topk( - segment_id, - cluster_rank, - layout.eff_cluster_blocks, - layout.leader_rank, - layout.is_single_cta, - layout.is_idle_rank, - k, - num_kth, - identify_op, - block_keys_out, - layout.block_keys_base, - layout.part, - layout.resident_base, - layout.my_resident_chunks, - layout.segment_size_off, - layout.head_items, - resident_keys, - layout.head_edge_len_items, - layout.tail_edge_len_items, - streamer); + write_nondeterministic_topk(num_kth, identify_op, block_keys_out, resident_keys, streamer); } // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into @@ -2741,17 +2618,16 @@ private: // after the scan, a block can return without risking a "cluster target block not present" fault from a straggler. } - // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). - _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all( - num_segments_val_t segment_id, - segment_size_val_t segment_size, - unsigned int cluster_rank, - unsigned int cluster_blocks) + // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). Runs before + // the effective-cluster collapse, so it takes the raw hardware rank/size (the `cluster_*` members are not set yet) + // and reads the `segment_id`/`segment_size` members `process_impl` filled in. + _CCCL_DEVICE _CCCL_FORCEINLINE void + copy_segment_select_all(unsigned int hw_cluster_rank, unsigned int hw_cluster_blocks) { constexpr int copy_items = copy_items_per_thread_clamped; const offset_t num_items = static_cast(segment_size); - const offset_t cluster_tid = cluster_rank * static_cast(threads_per_block) + threadIdx.x; - const offset_t cluster_threads = cluster_blocks * static_cast(threads_per_block); + const offset_t cluster_tid = hw_cluster_rank * static_cast(threads_per_block) + threadIdx.x; + const offset_t cluster_threads = hw_cluster_blocks * static_cast(threads_per_block); const offset_t step = cluster_threads * static_cast(copy_items); const offset_t full_tiles = ::cuda::round_down(num_items, step); auto keys_in_it = d_key_segments_it[segment_id]; @@ -2833,24 +2709,56 @@ private: } } + // Set the effective-cluster members (`cluster_rank`/`cluster_blocks`/`is_single_cta`) from the hardware cluster and + // this segment's size. For a per-segment (deferred) size argument the launch is sized for the maximum segment, so a + // small segment that fits resident in one CTA and is at/below the single-CTA tuning threshold is served by rank 0 + // alone via the barrier-free path; the cluster's other CTAs are redundant and this returns false so they exit, + // freeing their SM slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. + // Compiled out for host-exact sizes, which the dispatch already sized to exact cluster blocks (returns true for all). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool + init_effective_cluster(unsigned int hw_cluster_rank, unsigned int hw_cluster_blocks) + { + cluster_rank = hw_cluster_rank; + cluster_blocks = hw_cluster_blocks; + if constexpr (enable_runtime_single_cta) + { + const bool fits_single_cta = is_single_cta_eligible( + static_cast<::cuda::std::uint64_t>(segment_size), + static_cast<::cuda::std::uint64_t>(block_tile_capacity), + single_block_max_seg_size); + if (fits_single_cta) + { + if (hw_cluster_rank != 0u) + { + return false; + } + cluster_rank = 0u; + cluster_blocks = 1u; + } + } + is_single_cta = (cluster_blocks == 1u); + return true; + } + _CCCL_DEVICE _CCCL_FORCEINLINE void process_impl() { - // Cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). - const unsigned int cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); - // Runtime cluster blocks match the launch attribute the dispatch passed - // to `cudaLaunchKernelExC` (or the kernel's `__cluster_dims__` on CDP). - const unsigned int cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); - _CCCL_ASSERT(cluster_blocks > 0u && cluster_rank < cluster_blocks, + // Hardware cluster rank/size from the PTX special registers (replaces cooperative_groups' `this_cluster()`). The + // radix path runs on the *effective* geometry (the `cluster_rank`/`cluster_blocks` members set by + // `init_effective_cluster` below); only the pre-collapse steps here (segment id, select-all fast path) use the raw + // hardware values. Runtime cluster blocks match the launch attribute the dispatch passed to `cudaLaunchKernelExC` + // (or the kernel's `__cluster_dims__` on CDP). + const unsigned int hw_cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); + const unsigned int hw_cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); + _CCCL_ASSERT(hw_cluster_blocks > 0u && hw_cluster_rank < hw_cluster_blocks, "hardware cluster rank must lie within a non-empty cluster"); - const auto segment_id = static_cast(blockIdx.x / cluster_blocks); + segment_id = static_cast(blockIdx.x / hw_cluster_blocks); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) { return; } - const auto segment_size = - static_cast(detail::params::get_segment_size(segment_sizes, segment_id)); + segment_size = static_cast(detail::params::get_segment_size(segment_sizes, segment_id)); // Precondition: sizes are clamped >= 0 and capped at 2^21, so a value exceeding 32-bit `offset_t` is a violation. _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(segment_size) <= ::cuda::std::uint64_t{0xffffffffu}, "segment size must be non-negative and fit the 32-bit cluster offset type"); @@ -2870,48 +2778,32 @@ private: // from gmem (see `overflow_streamer`). // `k_clamped <= segment_size`, which now fits `out_offset_t`, so this narrowing is safe. - const auto k = static_cast(k_clamped); + k = static_cast(k_clamped); // Select-all fast path: when `k` reaches the full segment, every element wins, so we skip the radix passes, // histogram, and output-ordering and just copy. Runs on the full launched cluster (before the effective-cluster - // collapse); the decision is per-segment uniform, so the branch is cluster-uniform. + // collapse), so it uses the raw hardware rank/size; the decision is per-segment uniform, so the branch is + // cluster-uniform. if (static_cast(k) == segment_size) { - copy_segment_select_all(segment_id, segment_size, cluster_rank, cluster_blocks); + copy_segment_select_all(hw_cluster_rank, hw_cluster_blocks); return; } - // Effective cluster blocks/rank. For a per-segment (deferred) size argument the launch is sized for the maximum - // segment, so a small segment that fits resident in one CTA and is at/below the single-CTA tuning threshold is - // served by rank 0 alone via the barrier-free path; the cluster's other CTAs exit immediately, freeing their SM - // slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. Compiled out for - // host-exact sizes, which the dispatch already sized to exact cluster blocks. - unsigned int eff_cluster_blocks = cluster_blocks; - unsigned int eff_cluster_rank = cluster_rank; - if constexpr (enable_runtime_single_cta) + // Collapse the launched cluster onto the effective geometry the radix path runs on (sets the `cluster_rank`/ + // `cluster_blocks`/`is_single_cta` members). Returns false on the CTAs the collapse makes redundant, which exit + // here. + if (!init_effective_cluster(hw_cluster_rank, hw_cluster_blocks)) { - const bool fits_single_cta = is_single_cta_eligible( - static_cast<::cuda::std::uint64_t>(segment_size), - static_cast<::cuda::std::uint64_t>(block_tile_capacity), - single_block_max_seg_size); - if (fits_single_cta) - { - if (cluster_rank != 0u) - { - return; - } - eff_cluster_blocks = 1u; - eff_cluster_rank = 0u; - } + return; } - const bool is_single_cta = (eff_cluster_blocks == 1u); // Every block's thread 0 initializes its local `state`. Only the // leader's copy is semantically read (non-leaders reach the cluster // state through `leader_state`), but mirroring the writes everywhere // keeps every block's unconditional `state.k` load safe under // compute-sanitizer. - if (threadIdx.x == 0) + if (tid == 0) { temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; @@ -2933,12 +2825,10 @@ private: // `hist`; `leader_state`/`prefix_pair` are untouched until later). cluster_or_block_arrive(is_single_cta); - [[maybe_unused]] const bool is_ok = detail::params::dispatch_discrete( - select_directions, - segment_id, - [this, segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k](auto direction_tag) { + [[maybe_unused]] const bool is_ok = + detail::params::dispatch_discrete(select_directions, segment_id, [this](auto direction_tag) { constexpr detail::topk::select Direction = decltype(direction_tag)::value; - this->template run(segment_id, eff_cluster_rank, eff_cluster_blocks, segment_size, k); + this->template run(); }); _CCCL_ASSERT(is_ok, "Unsupported select direction for cluster top-k"); } From 0a35ad745dc30509de6eb7de2ed0686d8533d5ed Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 17:51:44 +0200 Subject: [PATCH 097/126] Flatten overflow streamer into agent --- cub/cub/agent/agent_batched_topk_cluster.cuh | 648 +++++++++---------- 1 file changed, 303 insertions(+), 345 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 8d26a78cad4..ff1062fd50d 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -324,7 +324,8 @@ struct agent_batched_topk_cluster static constexpr int slot_alignment = smem_layout_t::slot_alignment; // (pointer, count) carrier for a contiguous run of SMEM-staged keys, used only where the length must travel with the - // base across call boundaries (the resident window and the streamer's per-stage view); elsewhere: `key_t*` + count. + // base across call boundaries (the resident window and the overflow stream's per-stage view); elsewhere: `key_t*` + + // count. using smem_keys_t = ::cuda::std::span; // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a @@ -400,7 +401,7 @@ struct agent_batched_topk_cluster offset_t num_strictly_selected; offset_t my_candidates; typename block_scan_t::TempStorage scan_storage; - // One mbarrier handle per pipeline stage, shared by the resident load and the overflow streamer and reused + // One mbarrier handle per pipeline stage, shared by the resident load and the overflow stream and reused // (ping-ponged) across radix passes; all are initialized once up front by `init_load_barriers`. ::cuda::std::uint64_t load_mbar[PipelineStages]; // Persistent unaligned boundary edges (block-load path only): the head prefix (`[0, head_edge_cap_items)`, on rank @@ -439,7 +440,8 @@ struct agent_batched_topk_cluster // Assignment of the cluster's global chunk indices `[0, chunks)` to its CTAs. A rank owns `count` chunks; its i-th // owned chunk has global index `global_index(i) = first + i * stride`. The single mapping point lets the rest of the - // agent (resident load, streamer, per-pass scans) stay agnostic to the layout chosen by `make_chunk_partition`. + // agent (resident load, overflow stream, per-pass scans) stay agnostic to the layout chosen by + // `make_chunk_partition`. struct chunk_partition { offset_t first; // global index of this rank's first owned chunk @@ -455,7 +457,7 @@ struct agent_batched_topk_cluster // Decides which global chunks a cluster rank owns. Both layouts keep chunk 0 on rank 0 (which also stages the head // edge) and the tail (chunk `chunks-1`) on a single rank, and leave the per-chunk alignment, the resident/streaming - // split, and the streamer ping-pong untouched, because all of those depend only on the global chunk index, not on + // split, and the streaming ping-pong untouched, because all of those depend only on the global chunk index, not on // which rank owns it. // // * Strided (default): chunk `i` goes to rank `i % cluster_blocks`, so each CTA walks `first, first+S, first+2S, @@ -657,7 +659,7 @@ struct agent_batched_topk_cluster // chunking, the effective (non-idle) cluster width and this rank's partition of it, the leader/idle roles, and the // resident/streaming/edge split of this rank's chunks. All fields are loop-invariant for the segment; the // streaming-only intermediates (`overflow_base`, `resident_slots_cap`, `stream_slots`, `my_chunks`) are kept because - // the `overflow_streamer` constructor needs them. + // the overflow-streaming helpers (`init_overflow_stream`, `run_pass`, ...) read them. struct segment_layout { offset_t segment_size_off; @@ -696,7 +698,7 @@ struct agent_batched_topk_cluster char* key_slots; offset_t block_tile_capacity; // Per-thread mbarrier phase parity, one bit per pipeline stage (see `wait_stage`); the resident load and the - // overflow streamer keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. + // overflow stream keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. ::cuda::std::uint32_t load_phase{}; // The single block leader (warp 0's elected lane) that drives the bulk copies. Elected once at construction (the // block constructs convergently) and cached so the pipeline reuses it instead of re-electing per copy. @@ -718,6 +720,15 @@ struct agent_batched_topk_cluster out_offset_t k{}; segment_layout layout{}; + // Overflow-streaming slot geometry + ping-pong cursor, set once per segment by `init_overflow_stream` and carried + // across every radix pass and the final filter (the "Overflow streaming" section documents the reuse scheme). Inert + // when this rank has no overflow chunks (`init_overflow_stream` still runs, but streaming then touches no slot). + int stream_slot_base = 0; + int stream_stages = 1; + bool stream_is_forward{true}; + bool stream_is_primed{false}; + ::cuda::std::uint32_t stream_inflight_mask{0}; + _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( TempStorage& temp_storage_, KeyInputItItT d_key_segments_it_, @@ -871,318 +882,287 @@ private: } // --------------------------------------------------------------------------- - // Overflow streamer + // Overflow streaming // --------------------------------------------------------------------------- // Re-streams a rank's "overflow" chunks (those that do not fit its resident SMEM region) from gmem through a fixed, - // round-robin set of `eff_stages` (<= `PipelineStages`) streaming slots, reused across every radix pass and the final - // filter. It ping-pongs the iteration order across calls so the `eff_stages` turn-around chunks that one pass leaves - // resident in the streaming slots are reused by the next with no reload; the remaining `overflow_chunks - eff_stages` - // are reloaded from gmem each pass. The caller sizes the reservation to `eff_stages = min(PipelineStages, excess)` - // (`excess = my_chunks - full_slots`), so a streaming rank reloads exactly `excess` chunks per pass -- the reserved - // slots only ever buy reuse of the turn-around chunks, never a reload-free pass. The resident region occupies slots - // `[0, resident_slots)`, the streaming region `[stream_slot_base, stream_slot_base + eff_stages)`. - struct overflow_streamer + // round-robin set of `stream_stages` (<= `PipelineStages`) streaming slots, reused across every radix pass and the + // final filter. It ping-pongs the iteration order across calls (`stream_is_forward`) so the `stream_stages` + // turn-around chunks one pass leaves resident in the streaming slots are reused by the next with no reload; the + // remaining `overflow_chunks - stream_stages` are reloaded from gmem each pass. `compute_segment_layout` sizes the + // reservation `stream_slots = min(PipelineStages, excess)` (`excess = my_chunks - full_slots`) that + // `init_overflow_stream` adopts as `stream_stages`, so a streaming rank reloads exactly `excess` chunks per pass -- + // the reserved slots only ever buy reuse of the turn-around chunks, never a reload-free pass. The resident region + // occupies slots `[0, resident_slots)`, the streaming region `[stream_slot_base, stream_slot_base + stream_stages)`. + // All geometry comes from `layout`; the only per-stage state + // is one bit each of `stream_inflight_mask` (set while a copy is in flight) and the `load_phase` parity, so a + // spillable per-stage array is avoided (the read span is recomputed on demand by `stage_span`). + + // Adopt the streaming-slot window (`stream_slots`, sized by `compute_segment_layout`) and reset the ping-pong/priming + // cursor for this segment. `run` may then flip `stream_is_forward` for the deterministic filter's entry parity. + _CCCL_DEVICE _CCCL_FORCEINLINE void init_overflow_stream() { - agent_batched_topk_cluster& agent; - key_it_t block_keys_in; - const key_t* block_keys_base; // unwrapped contiguous base (pipeline path only; null otherwise) - offset_t segment_size; - offset_t head_items; - chunk_partition part; // rank -> global chunk index mapping (strided or blocked) - offset_t resident_chunks; // number of rank-local chunks kept resident - offset_t overflow_base; // rank-local chunk index of the first overflow (streamed) chunk - int stream_slot_base; // SMEM slot index at which the streaming region begins - offset_t overflow_chunks; // number of overflow chunks for this rank (M) - int eff_stages; // streaming region size = reserved streaming slots (<= PipelineStages; <= M when M>0, else 1) - bool is_forward = true; - bool is_primed = false; - - // Stage mbarriers are shared with the resident load (`agent.temp_storage.load_mbar`); stage `stage` targets slot - // `stream_slot_base + stage`. `inflight_mask` bit `stage` is set only while a copy is in flight (issued, not yet - // waited). The slot/stage mapping is fixed, so the read span is recomputed on demand by `stage_span` rather than - // held in a spillable per-stage array; the only per-stage state is one bit each of `inflight_mask` (here) and the - // agent's `load_phase` parity. - ::cuda::std::uint32_t inflight_mask = 0; - - _CCCL_DEVICE _CCCL_FORCEINLINE overflow_streamer( - agent_batched_topk_cluster& agent_, - key_it_t block_keys_in_, - const key_t* block_keys_base_, - offset_t segment_size_, - offset_t head_items_, - chunk_partition part_, - offset_t resident_chunks_, - offset_t overflow_base_, - int stream_slot_base_, - int stream_slots_, - offset_t my_chunks_) - : agent(agent_) - , block_keys_in(block_keys_in_) - , block_keys_base(block_keys_base_) - , segment_size(segment_size_) - , head_items(head_items_) - , part(part_) - , resident_chunks(resident_chunks_) - , overflow_base(overflow_base_) - , stream_slot_base(stream_slot_base_) - , overflow_chunks((my_chunks_ > resident_chunks_) ? (my_chunks_ - resident_chunks_) : offset_t{0}) - { - // Use the whole reserved streaming region as pipeline stages for maximal ping-pong reuse; the `1` floor covers - // the no-op (`M == 0`) case, where the streamer never touches a slot (asserted `<= overflow_chunks` below). - eff_stages = (stream_slots_ > 0) ? stream_slots_ : 1; - // Both chunk windows must lie inside this rank's `part.count` (resident first, streamed from `overflow_base`). - _CCCL_ASSERT(my_chunks_ == part.count && resident_chunks <= my_chunks_ && overflow_base <= my_chunks_ - && overflow_chunks <= my_chunks_ - overflow_base, - "overflow streamer chunk windows escape the rank's partition"); - // The streaming region is carved from the tile's slots and capped by the pipeline depth. - _CCCL_ASSERT(stream_slot_base >= 0 && stream_slots_ >= 0 && stream_slots_ <= PipelineStages - && static_cast(stream_slot_base) + static_cast(stream_slots_) - <= agent.block_tile_capacity / static_cast(chunk_items), - "overflow streamer streaming slots escape the block tile or pipeline"); - _CCCL_ASSERT(overflow_chunks == 0 || eff_stages <= static_cast(overflow_chunks), - "streaming depth exceeds the overflow chunk count"); - } + stream_slot_base = static_cast(layout.resident_slots_cap); + // Use the whole reserved streaming region as pipeline stages for maximal ping-pong reuse; the `1` floor covers the + // no-op (`overflow_chunks == 0`) case, where streaming never touches a slot (asserted below). + stream_stages = (layout.stream_slots > offset_t{0}) ? static_cast(layout.stream_slots) : 1; + stream_is_forward = true; + stream_is_primed = false; + stream_inflight_mask = 0; + // Both chunk windows must lie inside this rank's `part.count` (resident first, streamed from `overflow_base`). + _CCCL_ASSERT(layout.my_chunks == layout.part.count && layout.my_resident_chunks <= layout.my_chunks + && layout.overflow_base <= layout.my_chunks + && layout.overflow_chunks <= layout.my_chunks - layout.overflow_base, + "overflow stream chunk windows escape the rank's partition"); + // The streaming region is carved from the tile's slots and capped by the pipeline depth. + _CCCL_ASSERT(stream_slot_base >= 0 && layout.stream_slots <= static_cast(PipelineStages) + && static_cast(stream_slot_base) + layout.stream_slots + <= block_tile_capacity / static_cast(chunk_items), + "overflow stream slots escape the block tile or pipeline"); + _CCCL_ASSERT(layout.overflow_chunks == 0 || stream_stages <= static_cast(layout.overflow_chunks), + "streaming depth exceeds the overflow chunk count"); + } - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) - { - _CCCL_ASSERT(overflow_idx < overflow_chunks, "overflow chunk index out of range"); - _CCCL_ASSERT(stage >= 0 && stage < eff_stages, "overflow stage index exceeds the reserved streaming slots"); - _CCCL_ASSERT((inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, - "cannot issue a load into a streaming stage that is still in flight"); - const offset_t chunk_idx = part.global_index(overflow_base + overflow_idx); - const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - // Every chunk begins on a `load_align` boundary, so the guard-free aligned (TMA bulk) path applies. The global- - // last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk excludes - // it. For every interior chunk `bulk == count`. - const offset_t bulk = - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); - _CCCL_ASSERT(::cuda::is_aligned(block_keys_base + chunk.offset, load_align_bytes), - "overflow streamer received a chunk with an unaligned start"); - char* const dst = agent.key_slots + (stream_slot_base + stage) * ChunkBytes; - const ::cuda::std::span src{block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; - agent.issue_bulk_copy(stage, dst, src); - inflight_mask |= (::cuda::std::uint32_t{1} << stage); - } + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) + { + _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); + _CCCL_ASSERT(stage >= 0 && stage < stream_stages, "overflow stage index exceeds the reserved streaming slots"); + _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, + "cannot issue a load into a streaming stage that is still in flight"); + const offset_t chunk_idx = layout.part.global_index(layout.overflow_base + overflow_idx); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); + // Every chunk begins on a `load_align` boundary, so the guard-free aligned (TMA bulk) path applies. The global- + // last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk excludes + // it. For every interior chunk `bulk == count`. + const offset_t bulk = + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); + _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), + "overflow stream received a chunk with an unaligned start"); + char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + const ::cuda::std::span src{ + layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; + issue_bulk_copy(stage, dst, src); + stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); + } + + // Shared-memory view of the chunk currently resident in `stage`'s slot: the slot index is a pure function of `stage` + // and the length is recomputed from `overflow_idx` (no spillable per-stage array). Returns the aligned bulk only (the + // always-peeled tail suffix is excluded). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE smem_keys_t stage_span(int stage, offset_t overflow_idx) const + { + const auto chunk = get_chunk( + layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items); + return smem_keys_t( + ::cuda::ptr_rebind(key_slots + (stream_slot_base + stage) * ChunkBytes), + static_cast(::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); + } - // Shared-memory view of the chunk currently resident in `stage`'s slot without storing per-stage state: the slot - // index is a pure function of `stage` and the length is recomputed from `overflow_idx`, so there is no spillable - // `pending[]` array. Returns the aligned bulk only (the always-peeled tail suffix is excluded). - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE smem_keys_t stage_span(int stage, offset_t overflow_idx) const + // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident + // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from + // gmem (fallback). `mid()` runs once on a full pass -- after the first reload wave (`stream_stages` visits) is issued + // but before it is waited on, overlapping the caller's resident-chunk work with those in-flight copies (skipped if + // phase 1 stops early); it must be block-uniform with no unmatched barrier. `should_continue()` is polled once + // after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final + // filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the + // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (the deterministic + // filter's does, to read its placement counters block-wide); the histogram and non-deterministic filter pass an + // always-true predicate. An early break can leave prefetches in flight, so the pass drains the + // remaining stages before returning (a full pass ends with an empty `stream_inflight_mask`, so the drain is a no-op). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) + { + _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, + "an overflow pass must begin with no copies in flight"); + if (layout.overflow_chunks == 0) { - const auto chunk = agent.get_chunk(part.global_index(overflow_base + overflow_idx), segment_size, head_items); - return smem_keys_t( - ::cuda::ptr_rebind(agent.key_slots + (stream_slot_base + stage) * ChunkBytes), - static_cast( - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); + mid(); + return; } - // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident - // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from - // gmem (fallback). `mid()` runs once on a full pass -- after the first reload wave (`eff_stages` visits) is issued - // but before it is waited on, overlapping the caller's resident-chunk work with those in-flight copies (skipped if - // phase 1 stops early); it must be block-uniform with no unmatched barrier. `should_continue()` is polled once - // after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final - // filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the - // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (the deterministic - // filter's does, to read its placement counters block-wide); the histogram and non-deterministic filter pass an - // always-true predicate. An early break can leave prefetches in flight, so the pass drains the - // remaining stages before returning (a full pass ends with an empty `inflight_mask`, so the drain is a no-op). - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) + if constexpr (use_block_load_to_shared) { - _CCCL_ASSERT(inflight_mask == ::cuda::std::uint32_t{0}, "an overflow pass must begin with no copies in flight"); - if (overflow_chunks == 0) + // First ever call: prime the streaming slots. Subsequent calls inherit the previous pass's resident tail, which + // (because the order ping-pongs) is exactly the first `stream_stages` chunks of this direction. + if (!stream_is_primed) { - mid(); - return; + // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else + // the phase advances twice and a lagging thread misses the flip and spins forever. + __syncthreads(); + for (int i = 0; i < stream_stages; ++i) + { + const offset_t overflow_idx = + stream_is_forward ? static_cast(i) : (layout.overflow_chunks - 1 - static_cast(i)); + issue_load(static_cast(overflow_idx % static_cast(stream_stages)), overflow_idx); + } + stream_is_primed = true; } - if constexpr (use_block_load_to_shared) - { - // First ever call: prime the streaming slots. Subsequent calls inherit - // the previous pass's resident tail, which (because the order - // ping-pongs) is exactly the first `eff_stages` chunks of this direction. - if (!is_primed) + // Consume the `i`-th visit (its ping-pong-ordered position is `overflow_idx`): wait for its slot, fold its keys + // via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the slot just freed (a barrier + // guards the slot before the async copy can overwrite the data the block was just reading). Returns false once + // `should_continue()` reports the top-k fully placed -- polled before the prefetch so we never launch a copy we + // would only drain again; the up-to-`stream_stages - 1` prefetches already in flight (from earlier visits or + // priming) are drained after the loop. + const auto consume = [&](offset_t i) -> bool { + const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); + const int stage = static_cast(overflow_idx % static_cast(stream_stages)); + if (stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) { - // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else - // the phase advances twice and a lagging thread misses the flip and spins forever. - __syncthreads(); - for (int i = 0; i < eff_stages; ++i) - { - const offset_t overflow_idx = - is_forward ? static_cast(i) : (overflow_chunks - 1 - static_cast(i)); - issue_load(static_cast(overflow_idx % static_cast(eff_stages)), overflow_idx); - } - is_primed = true; + wait_stage(stage); + stream_inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); } + block_apply(stage, overflow_idx); - // Consume the `i`-th visit (its ping-pong-ordered position is `overflow_idx`): wait for its slot, fold its keys - // via `block_apply`, then prefetch the chunk `eff_stages` visits ahead into the slot just freed (a barrier - // guards the slot before the async copy can overwrite the data the block was just reading). Returns false once - // `should_continue()` reports the top-k fully placed -- polled before the prefetch so we never launch a copy we - // would only drain again; the up-to-`eff_stages - 1` prefetches already in flight (from earlier visits or - // priming) are drained after the loop. - const auto consume = [&](offset_t i) -> bool { - const offset_t overflow_idx = is_forward ? i : (overflow_chunks - 1 - i); - const int stage = static_cast(overflow_idx % static_cast(eff_stages)); - if (inflight_mask & (::cuda::std::uint32_t{1} << stage)) - { - agent.wait_stage(stage); - inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); - } - block_apply(stage, overflow_idx); - - if (!should_continue()) - { - return false; - } - - const offset_t next_step = i + static_cast(eff_stages); - if (next_step < overflow_chunks) - { - const offset_t next_overflow_idx = is_forward ? next_step : (overflow_chunks - 1 - next_step); - __syncthreads(); - issue_load(stage, next_overflow_idx); - } - return true; - }; - - // Phase 1: consume the first `eff_stages` visits (the chunks reused from the previous pass, already resident in - // the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. - bool is_stopped = false; - const offset_t split = (::cuda::std::min) (static_cast(eff_stages), overflow_chunks); - for (offset_t i = 0; i < split; ++i) + if (!should_continue()) { - if (!consume(i)) - { - is_stopped = true; - break; - } + return false; } - if (!is_stopped) + const offset_t next_step = i + static_cast(stream_stages); + if (next_step < layout.overflow_chunks) { - // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. - mid(); - - // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). - for (offset_t i = split; i < overflow_chunks; ++i) - { - if (!consume(i)) - { - break; - } - } + const offset_t next_overflow_idx = stream_is_forward ? next_step : (layout.overflow_chunks - 1 - next_step); + __syncthreads(); + issue_load(stage, next_overflow_idx); } + return true; + }; - // Drain prefetches still in flight before returning: an early break leaves outstanding bulk copies whose - // mbarriers were never waited, and they must complete before the block can exit (their slots are never read). - // `inflight_mask` is block-uniform (set/cleared under uniform control flow), so the trip count and each - // collective `wait_stage` are uniform across the block. - while (inflight_mask != ::cuda::std::uint32_t{0}) + // Phase 1: consume the first `stream_stages` visits (the chunks reused from the previous pass, already resident + // in the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. + bool is_stopped = false; + const offset_t split = (::cuda::std::min) (static_cast(stream_stages), layout.overflow_chunks); + for (offset_t i = 0; i < split; ++i) + { + if (!consume(i)) { - const int drain_stage = __ffs(static_cast(inflight_mask)) - 1; - agent.wait_stage(drain_stage); - inflight_mask &= ~(::cuda::std::uint32_t{1} << drain_stage); + is_stopped = true; + break; } - is_forward = !is_forward; } - else + + if (!is_stopped) { - // Generic fallback: no async SMEM pipeline, so resident work cannot hide load latency here. Fold the resident - // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no - // SMEM reuse), with the walk still snaking for L2 locality. + // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. mid(); - for (offset_t i = 0; i < overflow_chunks; ++i) + + // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). + for (offset_t i = split; i < layout.overflow_chunks; ++i) { - const offset_t overflow_idx = is_forward ? i : (overflow_chunks - 1 - i); - const offset_t chunk_idx = part.global_index(overflow_base + overflow_idx); - const auto chunk = agent.get_chunk(chunk_idx, segment_size, head_items); - generic_apply(chunk); - if (!should_continue()) + if (!consume(i)) { break; } } - is_forward = !is_forward; } - } - // Fold every overflow key once in the current ping-pong direction. `Indexed` selects the callable shape: keys-only - // (`Indexed == false`) applies `f(key)`; the pair filter (`Indexed == true`) applies `f(key, seg_idx)` with each - // key's segment-local index, needed to fetch its value payload from gmem while still reusing the streamed overflow - // keys. The index math lives entirely inside the `if constexpr (Indexed)` arms, so it is elided in the keys-only - // instantiation. See `run_pass` for the overlap semantics of `mid`; `UnrollFactor` partially unrolls the generic - // (gmem fallback) fold loop. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_dispatch(F&& f, Mid&& mid) + // Drain prefetches still in flight before returning: an early break leaves outstanding bulk copies whose + // mbarriers were never waited, and they must complete before the block can exit (their slots are never read). + // `stream_inflight_mask` is block-uniform (set/cleared under uniform control flow), so the trip count and each + // collective `wait_stage` are uniform across the block. + while (stream_inflight_mask != ::cuda::std::uint32_t{0}) + { + const int drain_stage = __ffs(static_cast(stream_inflight_mask)) - 1; + wait_stage(drain_stage); + stream_inflight_mask &= ~(::cuda::std::uint32_t{1} << drain_stage); + } + stream_is_forward = !stream_is_forward; + } + else { - run_pass( - [&](int stage, offset_t overflow_idx) { - if constexpr (Indexed) - { - const offset_t base_off = - agent.get_chunk(part.global_index(overflow_base + overflow_idx), segment_size, head_items).offset; - const auto keys = stage_span(stage, overflow_idx); - agent.template for_each_chunk_key_indexed( - keys.data(), static_cast(keys.size()), base_off, f); - } - else - { - const auto keys = stage_span(stage, overflow_idx); - agent.template for_each_chunk_key(keys.data(), static_cast(keys.size()), f); - } - }, - [&](const auto& chunk) { - const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); - _CCCL_PRAGMA_UNROLL(UnrollFactor) - for (int j = 0; j < iterations; ++j) + // Generic fallback: no async SMEM pipeline, so resident work cannot hide load latency here. Fold the resident + // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no + // SMEM reuse), with the walk still snaking for L2 locality. + mid(); + for (offset_t i = 0; i < layout.overflow_chunks; ++i) + { + const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); + const offset_t chunk_idx = layout.part.global_index(layout.overflow_base + overflow_idx); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); + generic_apply(chunk); + if (!should_continue()) + { + break; + } + } + stream_is_forward = !stream_is_forward; + } + } + + // Fold every overflow key once in the current ping-pong direction. `Indexed` selects the callable shape: keys-only + // (`Indexed == false`) applies `f(key)`; the pair filter (`Indexed == true`) applies `f(key, seg_idx)` with each + // key's segment-local index, needed to fetch its value payload from gmem while still reusing the streamed overflow + // keys. The index math lives entirely inside the `if constexpr (Indexed)` arms, so it is elided in the keys-only + // instantiation. See `run_pass` for the overlap semantics of `mid`; `UnrollFactor` partially unrolls the generic + // (gmem fallback) fold loop. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_dispatch(F&& f, Mid&& mid) + { + run_pass( + [&](int stage, offset_t overflow_idx) { + if constexpr (Indexed) + { + const offset_t base_off = + get_chunk( + layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items) + .offset; + const auto keys = stage_span(stage, overflow_idx); + for_each_chunk_key_indexed(keys.data(), static_cast(keys.size()), base_off, f); + } + else + { + const auto keys = stage_span(stage, overflow_idx); + for_each_chunk_key(keys.data(), static_cast(keys.size()), f); + } + }, + [&](const auto& chunk) { + const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); + _CCCL_PRAGMA_UNROLL(UnrollFactor) + for (int j = 0; j < iterations; ++j) + { + const int local = j * threads_per_block + tid; + if (local < chunk.count) { - const int local = j * threads_per_block + agent.tid; - if (local < chunk.count) + if constexpr (Indexed) { - if constexpr (Indexed) - { - const offset_t seg_idx = chunk.offset + static_cast(local); - f(block_keys_in[static_cast(seg_idx)], seg_idx); - } - else - { - f(block_keys_in[static_cast(chunk.offset + static_cast(local))]); - } + const offset_t seg_idx = chunk.offset + static_cast(local); + f(layout.block_keys_in[static_cast(seg_idx)], seg_idx); + } + else + { + f(layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); } } - }, - static_cast(mid), - [] { - return true; - }); - } + } + }, + static_cast(mid), + [] { + return true; + }); + } - // `f(key)` over every overflow key. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers - // pass their clamped items-per-thread. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) - { - process_pass_dispatch(static_cast(f), static_cast(mid)); - } + // `f(key)` over every overflow key. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers + // pass their clamped items-per-thread. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + { + process_pass_dispatch(static_cast(f), static_cast(mid)); + } - // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in - // by the BlockLoadToShared pipeline (rather than already resident in SMEM). - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) - { - process_pass(static_cast(f), [] {}); - } + // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in + // by the BlockLoadToShared pipeline (rather than already resident in SMEM). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) + { + process_pass(static_cast(f), [] {}); + } - // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) - { - process_pass_dispatch(static_cast(f), static_cast(mid)); - } - }; + // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) + { + process_pass_dispatch(static_cast(f), static_cast(mid)); + } // ------------------------------------------------------------------------- // Per-direction implementation @@ -1310,7 +1290,6 @@ private: identify_candidates_op_t identify_op; it_value_t block_keys_out; key_it_t block_keys_in; - overflow_streamer& streamer; chunk_partition part; out_offset_t k; out_offset_t num_back; @@ -1561,33 +1540,32 @@ private: // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the - // slots are already primed -- the histogram's first streaming pass primed the same persistent streamer, so - // `streamer.is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime - // (the generic fallback re-reads gmem each pass regardless). + // slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so + // `agent.stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no + // re-prime (the generic fallback re-reads gmem each pass regardless). _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow() { // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at - // `is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where the - // streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary + // `stream_is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where + // the streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is // subsumed -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never // reaches a straddler. - _CCCL_ASSERT(streamer.overflow_chunks == 0 || is_select_all_cand_cta || !is_tie_active - || streamer.is_forward == (!is_tie_reversed), + _CCCL_ASSERT(agent.layout.overflow_chunks == 0 || is_select_all_cand_cta || !is_tie_active + || agent.stream_is_forward == (!is_tie_reversed), "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " "the wrong streaming direction"); - streamer.run_pass( + agent.run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. // `stage_span` returns the slot's aligned-bulk view (a peeled tail suffix is handled by `process_tail_edge`). [&](int stage, offset_t overflow_idx) { - const auto keys = streamer.stage_span(stage, overflow_idx); + const auto keys = agent.stage_span(stage, overflow_idx); const offset_t base_off = - agent - .get_chunk(streamer.part.global_index(streamer.overflow_base + overflow_idx), segment_size_off, head_items) + agent.get_chunk(part.global_index(agent.layout.overflow_base + overflow_idx), segment_size_off, head_items) .offset; // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == - // false`): a stray terminal direct scan here would need the streamer to flag its last chunk, and the saving + // false`): a stray terminal direct scan here would need the stream to flag its last chunk, and the saving // is one barrier on one tile. process_tiles( keys.data(), base_off, static_cast(keys.size()), false); @@ -1712,11 +1690,7 @@ private: // perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( - out_offset_t num_kth, - IdentifyOp identify_op, - KeyOutIt block_keys_out, - smem_keys_t resident_keys, - overflow_streamer& streamer) + out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) { // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index // `seg_idx`. Deriving the per-segment value iterators *inside* the `is_keys_only` guard avoids indexing the null @@ -1785,7 +1759,7 @@ private: const auto write_selected = [&](const key_t& key) { sink(key, offset_t{0}); }; - // Fold the resident keys as the streamer's `mid` work so they overlap the first overflow reloads. Writes are + // Fold the resident keys as the overflow pass's `mid` work so they overlap the first overflow reloads. Writes are // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) @@ -1808,7 +1782,7 @@ private: // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, write_selected); }; - streamer.process_pass(write_selected, fold_resident); + process_pass(write_selected, fold_resident); } else { @@ -1818,7 +1792,7 @@ private: // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary - // edges); overflow chunks fold through the streamer's own indexed callback below. Materialize the key into a + // edges); overflow chunks fold through the indexed overflow-pass callback below. Materialize the key into a // register first: `sink` binds it by `const&` and reads it several times, so passing `smem[local]` directly // would re-issue a narrow `LDS` per use instead of reusing the loaded value. auto write_run = [&](const key_t* smem, offset_t base_off, int count) { @@ -1845,7 +1819,7 @@ private: } }; - // Fold the resident keys (and their values) as the streamer's `mid` work, exactly as in the keys-only path. + // Fold the resident keys (and their values) as the overflow pass's `mid` work, exactly as in the keys-only path. const auto fold_resident = [&] { if constexpr (use_block_load_to_shared) { @@ -1895,8 +1869,8 @@ private: }; // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected - // key's value at index `seg_idx`. The resident keys above fold in as the streamer's `mid` work. - streamer.process_pass_indexed(sink, fold_resident); + // key's value at index `seg_idx`. The resident keys above fold in as the overflow pass's `mid` work. + process_pass_indexed(sink, fold_resident); } } @@ -1907,11 +1881,7 @@ private: // `det_final_filter`; this member computes that struct's inputs and runs it. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( - out_offset_t num_kth, - IdentifyOp identify_op, - KeyOutIt block_keys_out, - smem_keys_t resident_keys, - overflow_streamer& streamer) + out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) { // Early stop is not special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. // @@ -1998,7 +1968,6 @@ private: identify_op, block_keys_out, layout.block_keys_in, - streamer, layout.part, k, num_back, @@ -2031,8 +2000,7 @@ private: // visibility of the zeroed histogram comes later, from the deferred initial cluster wait in `run_radix_passes` (this // whole load runs in that arrive->wait window). template - _CCCL_DEVICE _CCCL_FORCEINLINE void - load_and_histogram_first_pass(overflow_streamer& streamer, smem_keys_t& resident_keys) + _CCCL_DEVICE _CCCL_FORCEINLINE void load_and_histogram_first_pass(smem_keys_t& resident_keys) { using extract_bin_op_t = detail::topk::extract_bin_op_t; constexpr int total_bits = int{sizeof(key_t)} * 8; @@ -2048,8 +2016,9 @@ private: { if (layout.my_resident_chunks > 0) { - // Stage mbarriers and the `load_phase` parity are shared with the streamer (no per-chunk token array needed). - // Chunks are written densely in slot order from offset 0 and read back in the same order, so the read cursor + // Stage mbarriers and the `load_phase` parity are shared with the overflow stream (no per-chunk token array + // needed). Chunks are written densely in slot order from offset 0 and read back in the same order, so the read + // cursor // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk // begins on a `load_align` boundary (zero prefix), so its aligned bulk is `round_down(count, @@ -2073,7 +2042,7 @@ private: { return {}; } - // Chunks start after the aligned head on an alignment-multiple stride (mirrors the streamer's `issue_load`). + // Chunks start after the aligned head on an alignment-multiple stride (mirrors `issue_load`). _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), "resident loader received a chunk with an unaligned start"); return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; @@ -2170,11 +2139,11 @@ private: } } - // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the streamer's initial + // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the stream's initial // direction (preselected above; the histogram is order-independent, so the direction only sets up the leftover - // parity for the final filter). The streamer reuses the resident load's stage mbarriers (all front-loaded at + // parity for the final filter). The overflow stream reuses the resident load's stage mbarriers (all front-loaded at // `run` entry); `wait_stage` provides the producer/consumer sync. - streamer.process_pass(add_first_pass); + process_pass(add_first_pass); const int resident_count = static_cast(resident_keys.size()); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -2254,10 +2223,10 @@ private: // Resident vs. streaming split, decided independently per CTA (CTAs need not agree -- cross-CTA traffic and every // cluster barrier is reached uniformly). A CTA whose chunks fit its resident slots (`my_chunks <= full_slots`) // keeps them all resident and streams nothing; an overflowing CTA reserves a round-robin streaming region at the - // tail of its block_tile and re-streams its overflow chunks from gmem each pass via `streamer`. + // tail of its block_tile and re-streams its overflow chunks from gmem each pass via the overflow stream. // // Boundary edges (the unaligned head prefix on rank 0 and the unaligned tail suffix on the tail owner) cannot use - // the aligned TMA streamer, so both are always peeled into the persistent `edge_keys` buffer. The tail chunk's + // the aligned TMA stream, so both are always peeled into the persistent `edge_keys` buffer. The tail chunk's // aligned bulk is then a normal (possibly partial) aligned chunk that can be resident or streamed like any other; // no partial tail chunk is ever kept resident. Peeling both boundaries means streaming needs only `full_slots >= // 1`. @@ -2331,8 +2300,7 @@ private: // pinned down (or early stop fires). Folds `kth_key_bits_local` up digit by digit and returns the number of passes // that actually ran (`last_pass`), which the final filter uses to size its identify operator. template - _CCCL_DEVICE _CCCL_FORCEINLINE int - run_radix_passes(smem_keys_t resident_keys, overflow_streamer& streamer, key_prefix_t& kth_key_bits_local) + _CCCL_DEVICE _CCCL_FORCEINLINE int run_radix_passes(smem_keys_t resident_keys, key_prefix_t& kth_key_bits_local) { using extract_bin_op_t = detail::topk::extract_bin_op_t; using identify_candidates_op_t = @@ -2364,9 +2332,9 @@ private: } }; - // Resident-chunk histogram, deferred into the streamer so it overlaps the streamer's in-flight first reload - // wave (see `process_pass`). The histogram is order-independent, so folding resident keys between the - // streamer's load issue and its wait does not change the result. + // Resident-chunk histogram, deferred into the overflow stream so it overlaps the stream's in-flight first + // reload wave (see `process_pass`). The histogram is order-independent, so folding resident keys between the + // stream's load issue and its wait does not change the result. const auto fold_resident_hist = [&] { if constexpr (use_block_load_to_shared) { @@ -2390,7 +2358,7 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. - streamer.process_pass(add_hist, fold_resident_hist); + process_pass(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's per-bucket counts (the source of its `num_strictly_selected` @@ -2538,29 +2506,19 @@ private: // `run_radix_passes`), so the full key is never broadcast. key_prefix_t kth_key_bits_local = {}; - // Persistent streamer for the overflow chunks; a no-op (constructs nothing) when this rank has no overflow. - overflow_streamer streamer( - *this, - layout.block_keys_in, - layout.block_keys_base, - layout.segment_size_off, - layout.head_items, - layout.part, - layout.my_resident_chunks, - layout.overflow_base, - static_cast(layout.resident_slots_cap), - static_cast(layout.stream_slots), - layout.my_chunks); + // Size the persistent overflow-streaming window for this segment; a no-op when this rank has no overflow. + init_overflow_stream(); - // Preselect the streamer's initial ping-pong direction. A streaming rank flips direction once per histogram pass, - // so the leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = + // Preselect the streaming ping-pong direction. A streaming rank flips direction once per histogram pass, so the + // leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = // (!is_tie_reversed) ^ (num_passes & 1)` makes that leftover `== !is_tie_reversed` -- exactly what the // deterministic filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is - // moot). Non-deterministic filtering is order-independent, so leave its historical `is_forward` start untouched. + // moot). Non-deterministic filtering is order-independent, so leave its historical `stream_is_forward` start + // untouched. if constexpr (need_determinism) { - streamer.is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); + stream_is_forward = (!is_tie_reversed) ^ ((num_passes & 1) != 0); } // Contiguous resident-key window staged in SMEM (block-load path); read once per radix pass and in the final @@ -2575,12 +2533,12 @@ private: } __syncthreads(); - load_and_histogram_first_pass(streamer, resident_keys); + load_and_histogram_first_pass(resident_keys); // Publish the first pass's staged `edge_keys` and per-rank histogram block-wide before the radix passes read them. __syncthreads(); - const int last_pass = run_radix_passes(resident_keys, streamer, kth_key_bits_local); + const int last_pass = run_radix_passes(resident_keys, kth_key_bits_local); // ----------------------------------------------------------------------- // Final filter pass: write the top-k keys for this segment. Strictly- @@ -2605,11 +2563,11 @@ private: __syncthreads(); if constexpr (need_determinism) { - write_deterministic_topk(num_kth, identify_op, block_keys_out, resident_keys, streamer); + write_deterministic_topk(num_kth, identify_op, block_keys_out, resident_keys); } else { - write_nondeterministic_topk(num_kth, identify_op, block_keys_out, resident_keys, streamer); + write_nondeterministic_topk(num_kth, identify_op, block_keys_out, resident_keys); } // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into @@ -2775,7 +2733,7 @@ private: } // Segments larger than the resident cluster_tile capacity are still handled -- the overflow chunks are re-streamed - // from gmem (see `overflow_streamer`). + // from gmem (see the "Overflow streaming" section). // `k_clamped <= segment_size`, which now fits `out_offset_t`, so this narrowing is safe. k = static_cast(k_clamped); From c95ef855461f9d0cde600085c789fa51ff1d90e2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 18:53:58 +0200 Subject: [PATCH 098/126] Flatten deterministic filter class into agent Lowers complexity due to the tight coupling. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 680 ++++++++++--------- 1 file changed, 341 insertions(+), 339 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index ff1062fd50d..b42e6ccac4b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -245,10 +245,11 @@ struct agent_batched_topk_cluster // `leader_rank` computed in `run`. static constexpr bool is_scan_descending = !(need_determinism && !is_tie_reversed); - // The deterministic final scan visits chunks in global-index order and bails early (`should_stop`), so keeping the - // *first-visited* chunks resident lets it skip re-reading the streamed overflow. Ascending visits low indices first - // (the default low-resident split); descending (prefer-larger) visits high indices first, so we flip the split to - // keep the high-index chunks resident (see `run`), restoring symmetry. Never set on the non-deterministic path. + // The deterministic final scan visits chunks in global-index order and bails early (`final_filter_should_stop`), so + // keeping the *first-visited* chunks resident lets it skip re-reading the streamed overflow. Ascending visits low + // indices first (the default low-resident split); descending (prefer-larger) visits high indices first, so we flip + // the split to keep the high-index chunks resident (see `run`), restoring symmetry. Never set on the + // non-deterministic path. static constexpr bool is_residency_reversed = need_determinism && is_tie_reversed; // Segments at or below this size that also fit resident in one CTA take the single-CTA fast path (see @@ -1272,40 +1273,27 @@ private: return temp_storage.prefix_pair; } - // Deterministic final-filter driver. The `run()`-local tie-break state (counts, prefixes, region extents, and the - // mutable `running`/`is_tie_active` scan cursor) is hoisted here so the per-region index-ordered sweeps are named - // member functions instead of a nest of `[&]` lambdas. Constructed once per `run()` in the deterministic branch. Kept - // an aggregate (no user constructor) so its members are initialized positionally at the single call site. - // - // Codegen: methods are `_CCCL_FORCEINLINE`; the resident window is carried as a `smem_keys_t` view - // (`resident_front`). + // Deterministic final-filter state: the `run()`-local tie-break values (counts, prefixes, region extents) plus the + // mutable `running`/`is_tie_active` scan cursor, bundled into one POD so the flattened filter methods below take a + // single `st` argument instead of ~16. A deterministic-path-only local of `write_deterministic_topk`; never + // instantiated on the non-deterministic path. Segment-invariant inputs (`segment_id`, `k`, and the `layout.*` + // geometry) are read from the agent directly rather than duplicated here. `resident_front` carries the resident + // window as a `smem_keys_t` view. template - struct det_final_filter + struct det_filter_state { - using identify_candidates_op_t = - detail::topk::identify_candidates_op_t; + using identify_op_t = detail::topk::identify_candidates_op_t; - agent_batched_topk_cluster& agent; - num_segments_val_t segment_id; - identify_candidates_op_t identify_op; + identify_op_t identify_op; it_value_t block_keys_out; - key_it_t block_keys_in; - chunk_partition part; - out_offset_t k; out_offset_t num_back; out_offset_t my_front; offset_t sel_prefix; offset_t cand_prefix; offset_t my_cand_count; - offset_t resident_base; - offset_t my_resident_chunks; - offset_t segment_size_off; - offset_t head_items; offset_t front_seg_base; smem_keys_t resident_front; int front_count; - int head_edge_len_items; - int tail_edge_len_items; bool is_select_all_cand_cta; bool is_resident_terminal; bool is_head_edge_terminal; @@ -1313,354 +1301,379 @@ private: // Mutable per-invocation tie-break scan cursor. offset_t running; bool is_tie_active; + }; - // For each key written to `block_keys_out[pos]`, load the associated input value at the key's segment-local index - // `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences the null value iterators) - // in keys-only builds; `segment_id` is loop-invariant, so the per-segment iterators hoist out of the writes. - _CCCL_DEVICE _CCCL_FORCEINLINE void write_value(out_offset_t pos, offset_t seg_idx) const + // --------------------------------------------------------------------------- + // Deterministic final filter + // --------------------------------------------------------------------------- + // The per-region index-ordered sweeps below are named agent methods over a `det_filter_state` (`st`), driven by + // `run_filter`; `write_deterministic_topk` computes `st` and calls `run_filter(st)`. See `write_deterministic_topk` + // for the front/back placement scheme. All methods are `_CCCL_FORCEINLINE`. + + // For each key written to `block_keys_out[pos]`, load the associated input value at the key's segment-local index + // `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences the null value iterators) + // in keys-only builds; `segment_id` is loop-invariant, so the per-segment iterators hoist out of the writes. + _CCCL_DEVICE _CCCL_FORCEINLINE void final_filter_write_value(out_offset_t pos, offset_t seg_idx) + { + if constexpr (!is_keys_only) { - if constexpr (!is_keys_only) - { - _CCCL_ASSERT(pos < k && seg_idx < segment_size_off, - "value write must land in the top-k output and read inside the segment"); - auto block_vals_in = agent.d_value_segments_it[segment_id]; - auto block_vals_out = agent.d_value_segments_out_it[segment_id]; - block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; - } + _CCCL_ASSERT(pos < k && seg_idx < layout.segment_size_off, + "value write must land in the top-k output and read inside the segment"); + auto block_vals_in = d_value_segments_it[segment_id]; + auto block_vals_out = d_value_segments_out_it[segment_id]; + block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; } + } - // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and - // resolved its ties. Callers must `__syncthreads()` first (the counter reads are block-wide and must resynchronize - // lanes that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and - // before each streaming bulk copy -- never per tile. - _CCCL_DEVICE _CCCL_FORCEINLINE bool should_stop() - { - // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. - _CCCL_ASSERT(agent.temp_storage.front_local_cnt <= static_cast(my_front) - && agent.temp_storage.back_local_cnt <= my_cand_count, - "final-filter placement counters exceeded this CTA's assigned work"); - const bool is_front_done = agent.temp_storage.front_local_cnt >= static_cast(my_front); - // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never - // clears it) finishes once all `my_cand_count` of its candidates are placed. - const bool is_back_done = !is_tie_active || (agent.temp_storage.back_local_cnt >= my_cand_count); - return is_front_done && is_back_done; - } + // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and + // resolved its ties. Callers must `__syncthreads()` first (the counter reads are block-wide and must resynchronize + // lanes that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and + // before each streaming bulk copy -- never per tile. + template + _CCCL_DEVICE _CCCL_FORCEINLINE bool final_filter_should_stop(const det_filter_state& st) + { + // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. + _CCCL_ASSERT(temp_storage.front_local_cnt <= static_cast(st.my_front) + && temp_storage.back_local_cnt <= st.my_cand_count, + "final-filter placement counters exceeded this CTA's assigned work"); + const bool is_front_done = temp_storage.front_local_cnt >= static_cast(st.my_front); + // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never + // clears it) finishes once all `my_cand_count` of its candidates are placed. + const bool is_back_done = !st.is_tie_active || (temp_storage.back_local_cnt >= st.my_cand_count); + return is_front_done && is_back_done; + } - // Emit one back/tie candidate: if its `global_rank` (arrival- or scan-ordered by the caller) is a winner it lands - // in the top-k output at slot `k-1-global_rank`, with its value pulled from segment-local index `seg_idx`; ranks at - // or past `num_back` are losers and dropped (a no-op). Forced-inline with explicit args (no capture) so it folds - // into the hot back-placement loops with no extra live ranges; the caller keeps any counter side effects (e.g. - // the `back_local_cnt` atomic) in the `global_rank` argument so they still run for every candidate. - _CCCL_DEVICE _CCCL_FORCEINLINE void emit_back_one(offset_t global_rank, const key_t& key, offset_t seg_idx) + // Emit one back/tie candidate: if its `global_rank` (arrival- or scan-ordered by the caller) is a winner it lands + // in the top-k output at slot `k-1-global_rank`, with its value pulled from segment-local index `seg_idx`; ranks at + // or past `num_back` are losers and dropped (a no-op). Forced-inline so it folds into the hot back-placement loops; + // the caller keeps any counter side effects (e.g. the `back_local_cnt` atomic) in the `global_rank` argument so they + // still run for every candidate. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + emit_back_one(const det_filter_state& st, offset_t global_rank, const key_t& key, offset_t seg_idx) + { + if (global_rank < static_cast(st.num_back)) { - if (global_rank < static_cast(num_back)) - { - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - block_keys_out[out] = key; - write_value(out, seg_idx); - } + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + st.block_keys_out[out] = key; + final_filter_write_value(out, seg_idx); } + } - // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. `FromSmem` selects the - // key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` - // (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. Strictly- - // selected keys go to the front via a SMEM atomic (offset by `sel_prefix`); candidates go to the back (see the - // per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are - // unresolved -- holds the boundary and scans directly. `running` carries across tiles/regions. No per-tile early- - // exit or barrier here except the lazy-scan `else` branch; early exit is decided at critical points via - // `should_stop`. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - process_tiles([[maybe_unused]] const key_t* smem_src, offset_t seg_base, int count, bool region_is_terminal) + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. `FromSmem` selects the + // key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` + // (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. Strictly- + // selected keys go to the front via a SMEM atomic (offset by `st.sel_prefix`); candidates go to the back (see the + // per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are + // unresolved -- holds the boundary and scans directly. `st.running` carries across tiles/regions. No per-tile early- + // exit or barrier here except the lazy-scan `else` branch; early exit is decided at critical points via + // `final_filter_should_stop`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles( + det_filter_state& st, + [[maybe_unused]] const key_t* smem_src, + offset_t seg_base, + int count, + bool region_is_terminal) + { + // The SMEM-source clause is exempt when reading from gmem (`FromSmem == false`, the generic overflow fallback + // passes `nullptr`) or for an empty tile. + _CCCL_ASSERT(count >= 0 && seg_base <= layout.segment_size_off + && static_cast(count) <= layout.segment_size_off - seg_base + && (!FromSmem || count == 0 || smem_src != nullptr), + "process_tiles: region must be a valid in-segment source span"); + constexpr int items = Items; + constexpr int tile = threads_per_block * items; + for (int tile_base = 0; tile_base < count; tile_base += tile) { - // A non-negative span inside the segment, with a valid SMEM source when read from SMEM (the generic overflow - // fallback passes `nullptr` with `FromSmem == false`). - _CCCL_ASSERT( - count >= 0 && seg_base <= segment_size_off && static_cast(count) <= segment_size_off - seg_base - && (!FromSmem || count == 0 || smem_src != nullptr), - "process_tiles: region must be a valid in-segment source span"); - constexpr int items = Items; - constexpr int tile = threads_per_block * items; - for (int tile_base = 0; tile_base < count; tile_base += tile) + key_t keys[items]; + offset_t flags[items]; + detail::topk::candidate_class cls[items]; + bool is_valid[items]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) { - key_t keys[items]; - offset_t flags[items]; - detail::topk::candidate_class cls[items]; - bool is_valid[items]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + const int pos = tile_base + tid * items + i; + is_valid[i] = pos < count; + flags[i] = offset_t{0}; + if (is_valid[i]) { - const int pos = tile_base + agent.tid * items + i; - is_valid[i] = pos < count; - flags[i] = offset_t{0}; - if (is_valid[i]) + // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` + // local read and the segment-local value index. + const int folded_pos = Reversed ? (count - 1 - pos) : pos; + if constexpr (FromSmem) { - // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` - // local read and the segment-local value index. - const int folded_pos = Reversed ? (count - 1 - pos) : pos; - if constexpr (FromSmem) - { - keys[i] = smem_src[folded_pos]; - } - else - { - keys[i] = block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; - } - cls[i] = identify_op(keys[i]); - flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + keys[i] = smem_src[folded_pos]; } - } - - // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `sel_prefix`. - // The per-block slices (disjoint by `sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in - // here -- they always route through the back below, even on early stop. - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) - { - const bool is_front_key = is_valid[i] && (cls[i] == detail::topk::candidate_class::selected); - if (is_front_key) + else { - const int pos = tile_base + agent.tid * items + i; - const offset_t local = atomicAdd(&agent.temp_storage.front_local_cnt, offset_t{1}); - const out_offset_t out = static_cast(sel_prefix + local); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - block_keys_out[out] = keys[i]; - write_value(out, seg_idx); + keys[i] = + layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; } + cls[i] = st.identify_op(keys[i]); + flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; } + } - // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: - // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. - // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. - // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the - // counter and, on the crossing tile only, overwrite the arrival slots in index order. - if (is_tie_active) + // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `st.sel_prefix`. + // The per-block slices (disjoint by `st.sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in + // here -- they always route through the back below, even on early stop. + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + const bool is_front_key = is_valid[i] && (cls[i] == detail::topk::candidate_class::selected); + if (is_front_key) { - const bool is_terminal_tile = region_is_terminal && (tile_base + tile >= count); + const int pos = tile_base + tid * items + i; + const offset_t local = atomicAdd(&temp_storage.front_local_cnt, offset_t{1}); + const out_offset_t out = static_cast(st.sel_prefix + local); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + st.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); + } + } - // Arrival-order placement: each candidate grabs the next block-local back slot via an atomic. Used where the - // winning slot set is provisional (select-all, or the lazy path before the K-boundary is known). - const auto emit_arrival = [&] { - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: + // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. + // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. + // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the + // counter and, on the crossing tile only, overwrite the arrival slots in index order. + if (st.is_tie_active) + { + const bool is_terminal_tile = region_is_terminal && (tile_base + tile >= count); + + // Arrival-order placement: each candidate grabs the next block-local back slot via an atomic. Used where the + // winning slot set is provisional (select-all, or the lazy path before the K-boundary is known). + const auto emit_arrival = [&] { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (is_valid[i] && flags[i] != offset_t{0}) { - if (is_valid[i] && flags[i] != offset_t{0}) - { - const int pos = tile_base + agent.tid * items + i; - emit_back_one(cand_prefix + atomicAdd(&agent.temp_storage.back_local_cnt, offset_t{1}), - keys[i], - seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); - } + const int pos = tile_base + tid * items + i; + emit_back_one(st, + st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}), + keys[i], + seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); } - }; - // Index-ordered placement into slots based at `base`: a block scan assigns each candidate a deterministic - // slot by its in-tile rank. Returns the tile's candidate total. Used on the boundary-crossing tile. - const auto emit_indexed = [&](offset_t base) -> offset_t { - offset_t excl[items]; - offset_t tile_total = 0; - block_scan_t(agent.temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + } + }; + // Index-ordered placement into slots based at `base`: a block scan assigns each candidate a deterministic + // slot by its in-tile rank. Returns the tile's candidate total. Used on the boundary-crossing tile. + const auto emit_indexed = [&](offset_t base) -> offset_t { + offset_t excl[items]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items; ++i) + { + if (is_valid[i] && flags[i] != offset_t{0}) { - if (is_valid[i] && flags[i] != offset_t{0}) - { - const int pos = tile_base + agent.tid * items + i; - emit_back_one( - base + excl[i], keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); - } + const int pos = tile_base + tid * items + i; + emit_back_one( + st, base + excl[i], keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); } - return tile_total; - }; - - if (is_select_all_cand_cta) - { - emit_arrival(); } - else if (is_terminal_tile) + return tile_total; + }; + + if (st.is_select_all_cand_cta) + { + emit_arrival(); + } + else if (is_terminal_tile) + { + // Straddling CTA's last tile necessarily holds the boundary: scan it directly in index order. + st.running += emit_indexed(st.running); + st.is_tie_active = false; + } + else + { + emit_arrival(); + // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make + // the counter read race-free and block-uniform. + __syncthreads(); + const offset_t placed = temp_storage.back_local_cnt; + if ((st.cand_prefix + placed) > static_cast(st.num_back)) { - // Straddling CTA's last tile necessarily holds the boundary: scan it directly in index order. - running += emit_indexed(running); - is_tie_active = false; + // Boundary tile: overwrite this tile's arrival slots `{k-1-st.running, ...}` with the index-ordered winners + // (identical slot set, different candidate->slot mapping). + emit_indexed(st.running); } - else + st.running = st.cand_prefix + placed; + if (st.running >= static_cast(st.num_back)) { - emit_arrival(); - // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make - // the counter read race-free and block-uniform. - __syncthreads(); - const offset_t placed = agent.temp_storage.back_local_cnt; - if ((cand_prefix + placed) > static_cast(num_back)) - { - // Boundary tile: overwrite this tile's arrival slots `{k-1-running, ...}` with the index-ordered winners - // (identical slot set, different candidate->slot mapping). - emit_indexed(running); - } - running = cand_prefix + placed; - if (running >= static_cast(num_back)) - { - is_tie_active = false; - } - // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the - // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no per-tile - // barrier. - __syncthreads(); + st.is_tie_active = false; } + // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the + // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no per-tile + // barrier. + __syncthreads(); } } } + } - // Resident-front region. Direction is the compile-time `is_residency_reversed` (== `is_tie_reversed` in - // deterministic mode): ascending walks the low-index window forward, descending walks the high-index window - // (`resident_base`) in reverse, so a single `process_tiles` call per span with the index folded at compile time - // replaces the old fwd/rev pair. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident() + // Resident-front region. Direction is the compile-time `is_residency_reversed` (== `is_tie_reversed` in + // deterministic mode): ascending walks the low-index window forward, descending walks the high-index window + // (`resident_base`) in reverse, so a single `process_tiles` call per span with the index folded at compile time + // replaces the old fwd/rev pair. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident(det_filter_state& st) + { + if constexpr (use_block_load_to_shared) { - if constexpr (use_block_load_to_shared) + // Whole contiguous resident span staged in SMEM. + process_tiles( + st, st.resident_front.data(), st.front_seg_base, st.front_count, st.is_resident_terminal); + } + else + { + const int resident_chunk_count = static_cast(layout.my_resident_chunks); + for (int slot = 0; slot < resident_chunk_count; ++slot) { - // Whole contiguous resident span staged in SMEM. + const int local_slot = is_residency_reversed ? (resident_chunk_count - 1 - slot) : slot; + const offset_t chunk_idx = layout.part.global_index(layout.resident_base + static_cast(local_slot)); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); + // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy + // per-tile boundary detection handles any boundary), so pass `false`. process_tiles( - resident_front.data(), front_seg_base, front_count, is_resident_terminal); - } - else - { - const int resident_chunk_count = static_cast(my_resident_chunks); - for (int slot = 0; slot < resident_chunk_count; ++slot) - { - const int local_slot = is_residency_reversed ? (resident_chunk_count - 1 - slot) : slot; - const offset_t chunk_idx = part.global_index(resident_base + static_cast(local_slot)); - const auto chunk = agent.get_chunk(chunk_idx, segment_size_off, head_items); - // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy - // per-tile boundary detection handles any boundary), so pass `false`. - process_tiles( - ::cuda::ptr_rebind(agent.key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); - } + st, ::cuda::ptr_rebind(key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); } } + } + + // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA + // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once + // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the + // slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so + // `stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the + // generic fallback re-reads gmem each pass regardless). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow(det_filter_state& st) + { + // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at + // `stream_is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where + // the streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary + // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is + // subsumed -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never + // reaches a straddler. + _CCCL_ASSERT(layout.overflow_chunks == 0 || st.is_select_all_cand_cta || !st.is_tie_active + || stream_is_forward == (!is_tie_reversed), + "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " + "the wrong streaming direction"); - // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA - // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once - // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the - // slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so - // `agent.stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no - // re-prime (the generic fallback re-reads gmem each pass regardless). - _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow() + run_pass( + // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. + // `stage_span` returns the slot's aligned-bulk view (a peeled tail suffix is handled by `process_tail_edge`). + [&](int stage, offset_t overflow_idx) { + const auto keys = stage_span(stage, overflow_idx); + const offset_t base_off = + get_chunk( + layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items) + .offset; + // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == + // false`): a stray terminal direct scan here would need the stream to flag its last chunk, and the saving + // is one barrier on one tile. + process_tiles( + st, keys.data(), base_off, static_cast(keys.size()), false); + }, + // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). + [&](const auto& chunk) { + // FromSmem=false: `smem_src` is unread, so pass nullptr. + process_tiles( + st, nullptr, chunk.offset, chunk.count, false); + }, + // No interleaved resident work: the deterministic filter folds its resident span separately. + [] {}, + // Break the stream once the whole top-k is placed. The barrier makes the counter reads block-wide and resyncs + // lanes that drifted through the just-folded chunk's barrier-free tiles (polled before each refill copy). + [&] { + __syncthreads(); + return !final_filter_should_stop(st); + }); + } + + // Fold one persistent boundary edge (head prefix or peeled tail suffix), both staged in `edge_keys`. A no-op on the + // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. An edge spans fewer + // than `load_align_items` items (it fits in its `edge_keys` slot), so it runs `process_tiles` at unroll factor 1 + // (non-unrolled) instead of the wider tie-break unroll the potentially large resident/overflow regions use. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + process_edge(det_filter_state& st, const key_t* keys, offset_t seg_base, int count, bool is_terminal) + { + if constexpr (use_block_load_to_shared) { - // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at - // `stream_is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where - // the streaming direction is moot: no overflow (`overflow_chunks == 0`); a CTA entirely at/below the boundary - // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is - // subsumed -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never - // reaches a straddler. - _CCCL_ASSERT(agent.layout.overflow_chunks == 0 || is_select_all_cand_cta || !is_tie_active - || agent.stream_is_forward == (!is_tie_reversed), - "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " - "the wrong streaming direction"); - - agent.run_pass( - // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. - // `stage_span` returns the slot's aligned-bulk view (a peeled tail suffix is handled by `process_tail_edge`). - [&](int stage, offset_t overflow_idx) { - const auto keys = agent.stage_span(stage, overflow_idx); - const offset_t base_off = - agent.get_chunk(part.global_index(agent.layout.overflow_base + overflow_idx), segment_size_off, head_items) - .offset; - // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == - // false`): a stray terminal direct scan here would need the stream to flag its last chunk, and the saving - // is one barrier on one tile. - process_tiles( - keys.data(), base_off, static_cast(keys.size()), false); - }, - // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). - [&](const auto& chunk) { - // FromSmem=false: `smem_src` is unread, so pass nullptr. - process_tiles( - nullptr, chunk.offset, chunk.count, false); - }, - // No interleaved resident work: the deterministic filter folds its resident span separately. - [] {}, - // Break the stream once the whole top-k is placed. The barrier makes the counter reads block-wide and resyncs - // lanes that drifted through the just-folded chunk's barrier-free tiles (polled before each refill copy). - [&] { - __syncthreads(); - return !should_stop(); - }); + _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed>(st, keys, seg_base, count, is_terminal); } + } - // Fold one persistent boundary edge (head prefix or peeled tail suffix), both staged in `edge_keys`. A no-op on the - // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. An edge spans fewer - // than `load_align_items` items (it fits in its `edge_keys` slot), so it runs `process_tiles` at unroll factor 1 - // (non-unrolled) instead of the wider tie-break unroll the potentially large resident/overflow regions use. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge(const key_t* keys, offset_t seg_base, int count, bool is_terminal) - { - if constexpr (use_block_load_to_shared) + // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len_items)`, staged at `edge_keys` (base + // 0). In global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or + // empty prefix is a no-op. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge(det_filter_state& st) + { + process_edge(st, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, st.is_head_edge_terminal); + } + + // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` + // (base `segment_size - tail_edge_len_items`). In global-index order it is the trailing region (ascending) / + // leading region (descending); an aligned or non-owned tail is a no-op. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge(det_filter_state& st) + { + process_edge(st, + temp_storage.edge_keys + head_edge_cap_items, + layout.segment_size_off - static_cast(layout.tail_edge_len_items), + layout.tail_edge_len_items, + st.is_tail_edge_terminal); + } + + // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between + // regions once `final_filter_should_stop` reports the whole top-k placed. The two orders share the resident->overflow + // middle and only swap which boundary edge leads: ascending is head, resident, overflow, tail; descending reverses + // the edges. Both visit resident before overflow, so the stop check can skip re-streaming the overflow once the + // top-k is placed; `is_residency_reversed` keeps the first-visited (high-index) chunks resident in the descending + // order so this holds. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter(det_filter_state& st) + { + const auto step = [&](auto&& region) { + // Barrier before polling: makes the placement-counter reads block-wide and resynchronizes lanes that raced + // ahead through the previous region's barrier-free tiles. + __syncthreads(); + if (!final_filter_should_stop(st)) { - _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed>(keys, seg_base, count, is_terminal); + region(); } - } - - // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len_items)`, staged at `edge_keys` (base - // 0). In global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or - // empty prefix is a no-op. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge() + }; + if constexpr (is_tie_reversed) { - process_edge(agent.temp_storage.edge_keys, offset_t{0}, head_edge_len_items, is_head_edge_terminal); + process_tail_edge(st); } - - // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` - // (base `segment_size - tail_edge_len_items`). In global-index order it is the trailing region (ascending) / - // leading region (descending); an aligned or non-owned tail is a no-op. - _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge() + else { - process_edge(agent.temp_storage.edge_keys + head_edge_cap_items, - segment_size_off - static_cast(tail_edge_len_items), - tail_edge_len_items, - is_tail_edge_terminal); + process_head_edge(st); } - - // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between - // regions once `should_stop` reports the whole top-k placed. The two orders share the resident->overflow middle and - // only swap which boundary edge leads: ascending is head, resident, overflow, tail; descending reverses the edges. - // Both visit resident before overflow, so `should_stop` can skip re-streaming the overflow once the top-k is - // placed; `is_residency_reversed` keeps the first-visited (high-index) chunks resident in the descending order so - // this holds. - _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter() - { - const auto step = [&](auto&& region) { - // Barrier before polling: makes the placement-counter reads block-wide and resynchronizes lanes that raced - // ahead through the previous region's barrier-free tiles. - __syncthreads(); - if (!should_stop()) - { - region(); - } - }; + step([&] { + process_resident(st); + }); + step([&] { + process_overflow(st); + }); + step([&] { if constexpr (is_tie_reversed) { - process_tail_edge(); + process_head_edge(st); } else { - process_head_edge(); + process_tail_edge(st); } - step([&] { - process_resident(); - }); - step([&] { - process_overflow(); - }); - step([&] { - if constexpr (is_tie_reversed) - { - process_head_edge(); - } - else - { - process_tail_edge(); - } - }); - } - }; + }); + } // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final @@ -1877,8 +1890,8 @@ private: // Deterministic final filter: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on the // single boundary-crossing (straddling) CTA. A combined cross-CTA scan gives this block its disjoint front/back - // bases and lets it detect whether all/none/some of its candidates win. The heavy lifting lives in - // `det_final_filter`; this member computes that struct's inputs and runs it. + // bases and lets it detect whether all/none/some of its candidates win. This member computes the `det_filter_state` + // inputs and hands them to `run_filter`, which drives the per-region sweeps. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) @@ -1958,39 +1971,28 @@ private: ? get_chunk(layout.part.global_index(layout.resident_base), layout.segment_size_off, layout.head_items).offset : offset_t{0}; - // Positional aggregate init in declaration order. The last two initializers seed the tie-break cursor: `running` - // at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at + // Positional aggregate init in `det_filter_state` declaration order. The last two initializers seed the tie-break + // cursor: `running` at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at // `!is_select_no_cand_cta` -- true unless this is a select-no-candidates CTA (empty back region, or all of this // CTA's candidates sort past the boundary). - det_final_filter filt{ - *this, - segment_id, + det_filter_state st{ identify_op, block_keys_out, - layout.block_keys_in, - layout.part, - k, num_back, my_front, sel_prefix, cand_prefix, my_cand_count, - layout.resident_base, - layout.my_resident_chunks, - layout.segment_size_off, - layout.head_items, front_seg_base, resident_keys, front_count, - layout.head_edge_len_items, - layout.tail_edge_len_items, is_select_all_cand_cta, is_resident_terminal, is_head_edge_terminal, is_tail_edge_terminal, cand_prefix, !is_select_no_cand_cta}; - filt.run_filter(); + run_filter(st); } // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into @@ -2513,7 +2515,7 @@ private: // leftover after the compile-time `num_passes` passes is `initial ^ (num_passes & 1)`; choosing `initial = // (!is_tie_reversed) ^ (num_passes & 1)` makes that leftover `== !is_tie_reversed` -- exactly what the // deterministic filter's straddling CTA needs to reuse its resident turn-around chunks with no re-prime (see - // `det_final_filter::process_overflow`; early exit runs fewer passes but then has no straddler, so direction is + // `process_overflow`; early exit runs fewer passes but then has no straddler, so direction is // moot). Non-deterministic filtering is order-independent, so leave its historical `stream_is_forward` start // untouched. if constexpr (need_determinism) From 545c8f522c726b0c61d2f416db80767d8dab5927 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 19:34:16 +0200 Subject: [PATCH 099/126] Flatten write_nondeterministic_topk lambda --- cub/cub/agent/agent_batched_topk_cluster.cuh | 341 ++++++++++--------- 1 file changed, 183 insertions(+), 158 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index b42e6ccac4b..0058cfb826e 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1310,9 +1310,10 @@ private: // `run_filter`; `write_deterministic_topk` computes `st` and calls `run_filter(st)`. See `write_deterministic_topk` // for the front/back placement scheme. All methods are `_CCCL_FORCEINLINE`. - // For each key written to `block_keys_out[pos]`, load the associated input value at the key's segment-local index - // `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences the null value iterators) - // in keys-only builds; `segment_id` is loop-invariant, so the per-segment iterators hoist out of the writes. + // Shared by both final filters: for each key written to `block_keys_out[pos]`, load the associated input value at the + // key's segment-local index `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences + // the null value iterators) in keys-only builds; `segment_id` is loop-invariant, so the per-segment iterators hoist + // out of the writes. _CCCL_DEVICE _CCCL_FORCEINLINE void final_filter_write_value(out_offset_t pos, offset_t seg_idx) { if constexpr (!is_keys_only) @@ -1697,193 +1698,217 @@ private: } } - // Non-deterministic final filter: strictly-selected keys fill the front `[0, num_selected)`, the first `num_kth` - // candidates (arrival order) fill the back. The combined cross-CTA scan gives this block disjoint front/back bases - // (`sel_prefix`/`cand_prefix`); placement then uses block-local SMEM atomics since output order is not preserved. A - // perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. + // --------------------------------------------------------------------------- + // Non-deterministic final filter + // --------------------------------------------------------------------------- + // The sweep helpers below are named agent methods over a `nondet_filter_state` (`st`); `k`, `segment_id`, and + // `layout.*` are read from the agent. `write_nondeterministic_topk` computes `st` and drives the sweeps (see + // `nondet_place` for the per-key front/back placement). All methods are `_CCCL_FORCEINLINE`. template - _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( - out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) + struct nondet_filter_state { - // Store the value for a key written to `block_keys_out[pos]`, fetched from gmem at its segment-local index - // `seg_idx`. Deriving the per-segment value iterators *inside* the `is_keys_only` guard avoids indexing the null - // `cub::NullType**` value-iterators-of-iterators in keys-only builds; `segment_id` is loop-invariant, so they hoist - // out of the writes. - [[maybe_unused]] const auto write_value = [&](out_offset_t pos, offset_t seg_idx) { + IdentifyOp identify_op; + KeyOutIt block_keys_out; + out_offset_t num_kth; + offset_t sel_prefix; + offset_t cand_prefix; + smem_keys_t resident_keys; + }; + + // Classify one key and place it: strictly-selected keys go to the front (`sel_prefix` + a block-local atomic); the + // first `num_kth` candidates (arrival order) go to the back (`cand_prefix` + a block-local atomic), later candidates + // are dropped. In pair mode each written key additionally pulls its value from gmem at `seg_idx`; keys-only elides + // that and passes a dummy `seg_idx`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + nondet_place(const nondet_filter_state& st, const key_t& key, [[maybe_unused]] offset_t seg_idx) + { + const auto res = st.identify_op(key); + if (res == detail::topk::candidate_class::selected) + { + const out_offset_t pos = + static_cast(st.sel_prefix + atomicAdd(&temp_storage.front_local_cnt, offset_t{1})); + st.block_keys_out[pos] = key; if constexpr (!is_keys_only) { - _CCCL_ASSERT(pos < k && seg_idx < layout.segment_size_off, - "value write must land in the top-k output and read inside the segment"); - auto block_vals_in = d_value_segments_it[segment_id]; - auto block_vals_out = d_value_segments_out_it[segment_id]; - block_vals_out[pos] = block_vals_in[static_cast(seg_idx)]; + final_filter_write_value(pos, seg_idx); } - }; - - const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); - const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; - const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; - const ::cuda::std::uint64_t packed = - (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); - // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. - _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) - <= static_cast<::cuda::std::uint64_t>(k), - "selected prefix must fit before the tie-back output region"); - - // Placement sink shared by both value modes: strictly-selected keys go to the front (`sel_prefix` + a block-local - // atomic), the first `num_kth` candidates (arrival order) to the back; output order is not preserved. In pair - // mode each written key additionally pulls its value from gmem at `seg_idx`. Keys-only elides that write via the - // `if constexpr` below and drives the sink from index-free traversals, passing a dummy `seg_idx` that goes - // unread. - const auto sink = [&](const key_t& key, [[maybe_unused]] offset_t seg_idx) { - const auto res = identify_op(key); - if (res == detail::topk::candidate_class::selected) + } + else if (res == detail::topk::candidate_class::candidate) + { + const out_offset_t back_pos = + static_cast(st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1})); + if (back_pos < st.num_kth) { - const out_offset_t pos = - static_cast(sel_prefix + atomicAdd(&temp_storage.front_local_cnt, offset_t{1})); - block_keys_out[pos] = key; + const out_offset_t pos = k - 1 - back_pos; + st.block_keys_out[pos] = key; if constexpr (!is_keys_only) { - write_value(pos, seg_idx); + final_filter_write_value(pos, seg_idx); } } - else if (res == detail::topk::candidate_class::candidate) + } + } + + // Fold a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local index + // `base_off + local` (pair mode needs that index for the value fetch). Materialize each key into a register first: + // `nondet_place` reads it more than once, so indexing `smem[local]` per use would re-issue a narrow `LDS`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + nondet_write_run(const nondet_filter_state& st, const key_t* smem, offset_t base_off, int count) + { + const int iterations = ::cuda::ceil_div(count, threads_per_block); + _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) + for (int j = 0; j < iterations; ++j) + { + const int local = j * threads_per_block + tid; + if (local < count) { - const out_offset_t back_pos = - static_cast(cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1})); - if (back_pos < num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - block_keys_out[pos] = key; - if constexpr (!is_keys_only) - { - write_value(pos, seg_idx); - } - } + const key_t key = smem[local]; + nondet_place(st, key, base_off + static_cast(local)); } - }; + } + } + + // Fold a boundary edge (fewer than `load_align_items` items) with a plain block-stride loop rather than the + // tiled/unrolled `nondet_write_run` used for the potentially large resident bulks. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_write_edge( + const nondet_filter_state& st, const key_t* smem, offset_t base_off, int count) + { + for (int local = tid; local < count; local += threads_per_block) + { + const key_t key = smem[local]; + nondet_place(st, key, base_off + static_cast(local)); + } + } + // Fold this rank's resident keys (and, in pair mode, their values) as the overflow pass's `mid` work so they overlap + // the first overflow reloads. Writes are order-independent atomics and resident SMEM slots are disjoint from the + // streaming slots, so this never races the overflow apply. Keys-only drives the index-free traversal with a dummy + // `seg_idx`; pair mode recovers each key's segment-local index via `nondet_write_run`/`nondet_write_edge`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_fold_resident(const nondet_filter_state& st) + { if constexpr (is_keys_only) { - // Keys-only: no value payload, so drive the sink through the index-free traversal with a dummy `seg_idx`. const auto write_selected = [&](const key_t& key) { - sink(key, offset_t{0}); + nondet_place(st, key, offset_t{0}); }; - // Fold the resident keys as the overflow pass's `mid` work so they overlap the first overflow reloads. Writes are - // order-independent atomics, and resident SMEM slots are disjoint from streaming slots, so `mid` never races. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) + if constexpr (use_block_load_to_shared) + { + for_each_chunk_key( + st.resident_keys.data(), static_cast(st.resident_keys.size()), write_selected); + } + else + { + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { + const offset_t chunk_idx = layout.part.global_index(local_chunk); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); for_each_chunk_key( - resident_keys.data(), static_cast(resident_keys.size()), write_selected); - } - else - { - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const offset_t chunk_idx = layout.part.global_index(local_chunk); - const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); - for_each_chunk_key( - ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), - static_cast(chunk.count), - write_selected); - } + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + static_cast(chunk.count), + write_selected); } - // Scan the persistent boundary edges alongside the resident keys (order-independent atomic writes). - fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, write_selected); - }; - process_pass(write_selected, fold_resident); + } + // Order-independent atomic writes, so the persistent boundary edges fold in alongside the resident keys. + fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, write_selected); } else { - // Pair (key + value) path: the shared `sink` above additionally stores each written key's value (fetched from - // gmem at its `seg_idx`). Unlike keys-only, the traversal must recover each key's segment-local index, so it - // folds resident/edge keys through `write_run` (below) and overflow keys through `process_pass_indexed`. - - // Iterate a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local - // index `base_off + local`. Every source folded here is SMEM (resident slots and the persistent boundary - // edges); overflow chunks fold through the indexed overflow-pass callback below. Materialize the key into a - // register first: `sink` binds it by `const&` and reads it several times, so passing `smem[local]` directly - // would re-issue a narrow `LDS` per use instead of reusing the loaded value. - auto write_run = [&](const key_t* smem, offset_t base_off, int count) { - const int iterations = ::cuda::ceil_div(count, threads_per_block); - _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) - for (int j = 0; j < iterations; ++j) - { - const int local = j * threads_per_block + tid; - if (local < count) - { - const key_t key = smem[local]; - sink(key, base_off + static_cast(local)); - } - } - }; - - // Boundary edges span fewer than `load_align_items` items, so fold them with a plain block-stride loop rather - // than the tiled/unrolled `write_run` used for the (potentially large) resident bulks. - const auto write_edge = [&](const key_t* smem, offset_t base_off, int count) { - for (int local = tid; local < count; local += threads_per_block) + if constexpr (use_block_load_to_shared) + { + // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers per-chunk + // spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into `edge_keys` and folded + // below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not `chunk.count`. + key_t* const resident_ptr = st.resident_keys.data(); + int cursor_items = 0; + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const key_t key = smem[local]; - sink(key, base_off + static_cast(local)); + const auto chunk = get_chunk( + layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); + const int bulk_count_items = static_cast( + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); + nondet_write_run(st, resident_ptr + cursor_items, chunk.offset, bulk_count_items); + cursor_items += bulk_count_items; } - }; - - // Fold the resident keys (and their values) as the overflow pass's `mid` work, exactly as in the keys-only path. - const auto fold_resident = [&] { - if constexpr (use_block_load_to_shared) + } + else + { + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers - // per-chunk spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into - // `edge_keys` and folded below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not - // `chunk.count`. - key_t* const resident_ptr = resident_keys.data(); - int cursor_items = 0; - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const auto chunk = get_chunk( - layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); - const offset_t base_off = chunk.offset; - const int bulk_count_items = static_cast( - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); - write_run(resident_ptr + cursor_items, base_off, bulk_count_items); - cursor_items += bulk_count_items; - } + const auto chunk = + get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); + nondet_write_run(st, + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + chunk.offset, + chunk.count); } - else + } + // Fold the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the + // peeled tail suffix at `segment_size - tail_edge_len_items`. + if constexpr (use_block_load_to_shared) + { + if (layout.head_edge_len_items > 0) { - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const auto chunk = - get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); - const offset_t base_off = chunk.offset; - write_run( - ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), base_off, chunk.count); - } + nondet_write_edge(st, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items); } - // Scan the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the - // peeled tail suffix at `segment_size - tail_edge_len_items`. Value fetched per key. - if constexpr (use_block_load_to_shared) + if (layout.tail_edge_len_items > 0) { - if (layout.head_edge_len_items > 0) - { - write_edge(temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items); - } - if (layout.tail_edge_len_items > 0) - { - write_edge(temp_storage.edge_keys + head_edge_cap_items, - layout.segment_size_off - static_cast(layout.tail_edge_len_items), - layout.tail_edge_len_items); - } + nondet_write_edge(st, + temp_storage.edge_keys + head_edge_cap_items, + layout.segment_size_off - static_cast(layout.tail_edge_len_items), + layout.tail_edge_len_items); } - }; + } + } + } + + // Non-deterministic final filter driver. The combined cross-CTA scan gives this block disjoint front/back bases + // (`sel_prefix`/`cand_prefix`); `nondet_place` then places each key with block-local SMEM atomics (output order not + // preserved). A perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( + out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) + { + const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); + const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; + const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; + const ::cuda::std::uint64_t packed = + (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); + const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. + _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) + <= static_cast<::cuda::std::uint64_t>(k), + "selected prefix must fit before the tie-back output region"); + + nondet_filter_state st{ + identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, resident_keys}; - // Overflow chunks: reuse the streamed keys (the generic fallback re-reads from gmem) and fetch each selected - // key's value at index `seg_idx`. The resident keys above fold in as the overflow pass's `mid` work. - process_pass_indexed(sink, fold_resident); + // `f` folds each overflow key; `mid` folds the resident keys + boundary edges, overlapping the first overflow + // reloads. Keys-only passes a dummy `seg_idx`; pair mode recovers each key's index for its value fetch. + if constexpr (is_keys_only) + { + process_pass( + [&](const key_t& key) { + nondet_place(st, key, offset_t{0}); + }, + [&] { + nondet_fold_resident(st); + }); + } + else + { + process_pass_indexed( + [&](const key_t& key, offset_t seg_idx) { + nondet_place(st, key, seg_idx); + }, + [&] { + nondet_fold_resident(st); + }); } } From 63a7fafd98cf646e5d11da0c83ed2d14d821c5cd Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 21:36:39 +0200 Subject: [PATCH 100/126] Avoid bool/enum arrays Leave it to the compiler to hoist the comparisons out of unrolled loops. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 126 +++++++++---------- 1 file changed, 60 insertions(+), 66 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 0058cfb826e..9e033c7f77c 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1344,32 +1344,15 @@ private: return is_front_done && is_back_done; } - // Emit one back/tie candidate: if its `global_rank` (arrival- or scan-ordered by the caller) is a winner it lands - // in the top-k output at slot `k-1-global_rank`, with its value pulled from segment-local index `seg_idx`; ranks at - // or past `num_back` are losers and dropped (a no-op). Forced-inline so it folds into the hot back-placement loops; - // the caller keeps any counter side effects (e.g. the `back_local_cnt` atomic) in the `global_rank` argument so they - // still run for every candidate. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - emit_back_one(const det_filter_state& st, offset_t global_rank, const key_t& key, offset_t seg_idx) - { - if (global_rank < static_cast(st.num_back)) - { - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - st.block_keys_out[out] = key; - final_filter_write_value(out, seg_idx); - } - } - - // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * Items`. `FromSmem` selects the - // key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` - // (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. Strictly- - // selected keys go to the front via a SMEM atomic (offset by `st.sel_prefix`); candidates go to the back (see the - // per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are - // unresolved -- holds the boundary and scans directly. `st.running` carries across tiles/regions. No per-tile early- - // exit or barrier here except the lazy-scan `else` branch; early exit is decided at critical points via - // `final_filter_should_stop`. - template + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` + // selects the key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem + // `block_keys_in` (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. + // Strictly-selected keys go to the front via a SMEM atomic (offset by `st.sel_prefix`); candidates go to the back + // (see the per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are + // unresolved -- holds the boundary and scans directly. `st.running` carries across tiles/regions. The only explicit + // per-tile `__syncthreads()` are in the lazy-scan `else` branch (block scans elsewhere synchronize internally); early + // exit is decided at critical points via `final_filter_should_stop`. + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles( det_filter_state& st, [[maybe_unused]] const key_t* smem_src, @@ -1383,21 +1366,20 @@ private: && static_cast(count) <= layout.segment_size_off - seg_base && (!FromSmem || count == 0 || smem_src != nullptr), "process_tiles: region must be a valid in-segment source span"); - constexpr int items = Items; - constexpr int tile = threads_per_block * items; - for (int tile_base = 0; tile_base < count; tile_base += tile) + constexpr int tile_size = threads_per_block * ItemsPerThread; + for (int tile_base = 0; tile_base < count; tile_base += tile_size) { - key_t keys[items]; - offset_t flags[items]; - detail::topk::candidate_class cls[items]; - bool is_valid[items]; + key_t keys[ItemsPerThread]; + offset_t flags[ItemsPerThread]; _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + for (int i = 0; i < ItemsPerThread; ++i) { - const int pos = tile_base + tid * items + i; - is_valid[i] = pos < count; - flags[i] = offset_t{0}; - if (is_valid[i]) + const int pos = tile_base + tid * ItemsPerThread + i; + // Zero-init doubles as the out-of-range/non-candidate mask -- only in-range candidates set it below, so the + // back-placement loops gate on `flags[i]` alone. Validity (`pos < count`) and the selected class are cheap + // recomputed on the loaded key, not cached per item. + flags[i] = offset_t{0}; + if (pos < count) { // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` // local read and the segment-local value index. @@ -1411,21 +1393,20 @@ private: keys[i] = layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; } - cls[i] = st.identify_op(keys[i]); - flags[i] = (cls[i] == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; + flags[i] = (st.identify_op(keys[i]) == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; } } // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `st.sel_prefix`. // The per-block slices (disjoint by `st.sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in - // here -- they always route through the back below, even on early stop. + // here -- they always route through the back below, even on early stop. The `&&` short-circuits so `identify_op` + // only touches in-range (loaded) keys. _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + for (int i = 0; i < ItemsPerThread; ++i) { - const bool is_front_key = is_valid[i] && (cls[i] == detail::topk::candidate_class::selected); - if (is_front_key) + const int pos = tile_base + tid * ItemsPerThread + i; + if (pos < count && st.identify_op(keys[i]) == detail::topk::candidate_class::selected) { - const int pos = tile_base + tid * items + i; const offset_t local = atomicAdd(&temp_storage.front_local_cnt, offset_t{1}); const out_offset_t out = static_cast(st.sel_prefix + local); const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); @@ -1434,45 +1415,58 @@ private: } } - // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths: + // Back/tie placement (only while this CTA still has unresolved ties). A candidate whose global rank is below + // `num_back` lands in the top-k output at slot `k-1-rank`; higher ranks are dropped. Three block-uniform + // sub-paths assign that rank: // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the // counter and, on the crossing tile only, overwrite the arrival slots in index order. if (st.is_tie_active) { - const bool is_terminal_tile = region_is_terminal && (tile_base + tile >= count); + const bool is_terminal_tile = region_is_terminal && (tile_base + tile_size >= count); - // Arrival-order placement: each candidate grabs the next block-local back slot via an atomic. Used where the - // winning slot set is provisional (select-all, or the lazy path before the K-boundary is known). + // Arrival-order placement: each candidate grabs the next back rank via an atomic on `back_local_cnt` (added to + // `cand_prefix`); the atomic must advance for every candidate, win or lose. Arrival order is acceptable because + // either every candidate wins (select-all) or these slots are overwritten in index order later (lazy path). const auto emit_arrival = [&] { _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + for (int i = 0; i < ItemsPerThread; ++i) { - if (is_valid[i] && flags[i] != offset_t{0}) + if (flags[i] != offset_t{0}) { - const int pos = tile_base + tid * items + i; - emit_back_one(st, - st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}), - keys[i], - seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); + const int pos = tile_base + tid * ItemsPerThread + i; + const offset_t global_rank = st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}); + if (global_rank < static_cast(st.num_back)) + { + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + st.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); + } } } }; - // Index-ordered placement into slots based at `base`: a block scan assigns each candidate a deterministic - // slot by its in-tile rank. Returns the tile's candidate total. Used on the boundary-crossing tile. + // Index-ordered placement: a block scan assigns each candidate a deterministic rank (`base` plus its in-tile + // position). Returns the tile's candidate total. Used on the boundary-crossing tile. const auto emit_indexed = [&](offset_t base) -> offset_t { - offset_t excl[items]; + offset_t excl[ItemsPerThread]; offset_t tile_total = 0; block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < items; ++i) + for (int i = 0; i < ItemsPerThread; ++i) { - if (is_valid[i] && flags[i] != offset_t{0}) + if (flags[i] != offset_t{0}) { - const int pos = tile_base + tid * items + i; - emit_back_one( - st, base + excl[i], keys[i], seg_base + static_cast(Reversed ? (count - 1 - pos) : pos)); + const offset_t global_rank = base + excl[i]; + if (global_rank < static_cast(st.num_back)) + { + const int pos = tile_base + tid * ItemsPerThread + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + st.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); + } } } return tile_total; @@ -1507,8 +1501,8 @@ private: st.is_tie_active = false; } // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the - // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no per-tile - // barrier. + // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no additional + // explicit barrier. __syncthreads(); } } From 302b852f15ab04bba4fdcc86e8007a8e77d8d302 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Fri, 10 Jul 2026 22:01:47 +0200 Subject: [PATCH 101/126] Flatten emit_[arrival|indexed] lambdas --- cub/cub/agent/agent_batched_topk_cluster.cuh | 125 +++++++++++-------- 1 file changed, 72 insertions(+), 53 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 9e033c7f77c..5606e90190b 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1344,6 +1344,72 @@ private: return is_front_done && is_back_done; } + // Arrival-order back placement for one tile: each candidate advances `back_local_cnt` (win or lose); the returned + // count offset by `cand_prefix` is its back rank. A winner (rank below `num_back`) lands at output slot `k-1-rank`; + // higher ranks are dropped. Arrival order is acceptable because either every candidate wins (select-all) or, on the + // lazy path, only the boundary-crossing tile is later overwritten in index order (earlier tiles are all winners). + template + _CCCL_DEVICE _CCCL_FORCEINLINE void emit_arrival( + det_filter_state& st, + const key_t (&keys)[ItemsPerThread], + const offset_t (&flags)[ItemsPerThread], + offset_t seg_base, + int count, + int tile_base) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + if (flags[i] != offset_t{0}) + { + const int pos = tile_base + tid * ItemsPerThread + i; + const offset_t global_rank = st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}); + if (global_rank < static_cast(st.num_back)) + { + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + st.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); + } + } + } + } + + // Index-ordered back placement for one tile: a block scan gives each candidate a deterministic rank (`base` plus the + // count of preceding candidates in the tile). Placement rule matches `emit_arrival`. Returns the tile's candidate + // total. Used where the K-boundary falls in this tile (terminal tile, or the lazy path's crossing tile). + template + _CCCL_DEVICE _CCCL_FORCEINLINE offset_t emit_indexed( + det_filter_state& st, + const key_t (&keys)[ItemsPerThread], + offset_t (&flags)[ItemsPerThread], + offset_t seg_base, + int count, + int tile_base, + offset_t base) + { + offset_t excl[ItemsPerThread]; + offset_t tile_total = 0; + block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + if (flags[i] != offset_t{0}) + { + const offset_t global_rank = base + excl[i]; + if (global_rank < static_cast(st.num_back)) + { + const int pos = tile_base + tid * ItemsPerThread + i; + const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + st.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); + } + } + } + return tile_total; + } + // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` // selects the key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem // `block_keys_in` (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. @@ -1415,9 +1481,8 @@ private: } } - // Back/tie placement (only while this CTA still has unresolved ties). A candidate whose global rank is below - // `num_back` lands in the top-k output at slot `k-1-rank`; higher ranks are dropped. Three block-uniform - // sub-paths assign that rank: + // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths pick the + // candidate rank (see `emit_arrival`/`emit_indexed` for the shared rank->slot rule): // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the @@ -1426,65 +1491,19 @@ private: { const bool is_terminal_tile = region_is_terminal && (tile_base + tile_size >= count); - // Arrival-order placement: each candidate grabs the next back rank via an atomic on `back_local_cnt` (added to - // `cand_prefix`); the atomic must advance for every candidate, win or lose. Arrival order is acceptable because - // either every candidate wins (select-all) or these slots are overwritten in index order later (lazy path). - const auto emit_arrival = [&] { - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < ItemsPerThread; ++i) - { - if (flags[i] != offset_t{0}) - { - const int pos = tile_base + tid * ItemsPerThread + i; - const offset_t global_rank = st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}); - if (global_rank < static_cast(st.num_back)) - { - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - st.block_keys_out[out] = keys[i]; - final_filter_write_value(out, seg_idx); - } - } - } - }; - // Index-ordered placement: a block scan assigns each candidate a deterministic rank (`base` plus its in-tile - // position). Returns the tile's candidate total. Used on the boundary-crossing tile. - const auto emit_indexed = [&](offset_t base) -> offset_t { - offset_t excl[ItemsPerThread]; - offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < ItemsPerThread; ++i) - { - if (flags[i] != offset_t{0}) - { - const offset_t global_rank = base + excl[i]; - if (global_rank < static_cast(st.num_back)) - { - const int pos = tile_base + tid * ItemsPerThread + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - st.block_keys_out[out] = keys[i]; - final_filter_write_value(out, seg_idx); - } - } - } - return tile_total; - }; - if (st.is_select_all_cand_cta) { - emit_arrival(); + emit_arrival(st, keys, flags, seg_base, count, tile_base); } else if (is_terminal_tile) { // Straddling CTA's last tile necessarily holds the boundary: scan it directly in index order. - st.running += emit_indexed(st.running); + st.running += emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); st.is_tie_active = false; } else { - emit_arrival(); + emit_arrival(st, keys, flags, seg_base, count, tile_base); // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make // the counter read race-free and block-uniform. __syncthreads(); @@ -1493,7 +1512,7 @@ private: { // Boundary tile: overwrite this tile's arrival slots `{k-1-st.running, ...}` with the index-ordered winners // (identical slot set, different candidate->slot mapping). - emit_indexed(st.running); + emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); } st.running = st.cand_prefix + placed; if (st.running >= static_cast(st.num_back)) From 71dfeb0523aeef1c2327c651900fca6df5364fbb Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 01:37:46 +0200 Subject: [PATCH 102/126] Unify candidate and selected path --- cub/cub/agent/agent_batched_topk_cluster.cuh | 164 ++++++++++--------- 1 file changed, 89 insertions(+), 75 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5606e90190b..1053f58d6fd 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1326,6 +1326,22 @@ private: } } + // Per-item placement code stored in `flags[]` by the load/classify pass and reused by `place_tile`/`emit_indexed`, so + // each key is classified exactly once. + static constexpr offset_t flag_none = 0; // rejected or out-of-range: no placement, absent from the tie scan + static constexpr offset_t flag_candidate = 1; // tie candidate: routed to the back, counted by the boundary scan + static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front + + // Shared output write for the deterministic placement paths (`place_tile`, `emit_indexed`): store the key at + // `block_keys_out[out]`, then (pairs builds) copy its value from segment-local index `seg_idx`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + final_filter_place(det_filter_state& st, out_offset_t out, const key_t& key, offset_t seg_idx) + { + st.block_keys_out[out] = key; + final_filter_write_value(out, seg_idx); + } + // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and // resolved its ties. Callers must `__syncthreads()` first (the counter reads are block-wide and must resynchronize // lanes that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and @@ -1344,66 +1360,81 @@ private: return is_front_done && is_back_done; } - // Arrival-order back placement for one tile: each candidate advances `back_local_cnt` (win or lose); the returned - // count offset by `cand_prefix` is its back rank. A winner (rank below `num_back`) lands at output slot `k-1-rank`; - // higher ranks are dropped. Arrival order is acceptable because either every candidate wins (select-all) or, on the - // lazy path, only the boundary-crossing tile is later overwritten in index order (earlier tiles are all winners). + // Single-pass placement for one tile: both regions fill forward -- front `[0, num_selected)`, back `[num_selected, + // k)` -- so a key's class only selects the counter and base offset (below), and one uniform `out < k` guard drops + // losing candidates while always accepting selected keys (whose `out` is `< num_selected`). The candidate atomic + // advances win or lose; arrival order is fine because either every candidate wins (select-all) or the lazy crossing + // tile is overwritten in index order later. `do_arrival == false` (terminal tile only) skips candidates here and + // leaves them to `emit_indexed`; the lazy path passes `true` and its crossing tile is later overwritten by + // `emit_indexed`. template - _CCCL_DEVICE _CCCL_FORCEINLINE void emit_arrival( + _CCCL_DEVICE _CCCL_FORCEINLINE void place_tile( det_filter_state& st, const key_t (&keys)[ItemsPerThread], const offset_t (&flags)[ItemsPerThread], offset_t seg_base, int count, - int tile_base) + int tile_base, + bool do_arrival) { + const offset_t num_selected = static_cast(k) - static_cast(st.num_back); // front region size _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; ++i) { - if (flags[i] != offset_t{0}) + const bool is_cand = flags[i] == flag_candidate; + if (flags[i] == flag_none || (is_cand && !do_arrival)) { - const int pos = tile_base + tid * ItemsPerThread + i; - const offset_t global_rank = st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1}); - if (global_rank < static_cast(st.num_back)) - { - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - st.block_keys_out[out] = keys[i]; - final_filter_write_value(out, seg_idx); - } + continue; + } + offset_t* const counter = is_cand ? &temp_storage.back_local_cnt : &temp_storage.front_local_cnt; + const offset_t base = is_cand ? (num_selected + st.cand_prefix) : st.sel_prefix; + const out_offset_t out = static_cast(base + atomicAdd(counter, offset_t{1})); + if (out < k) + { + const int pos = tile_base + tid * ItemsPerThread + i; + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + final_filter_place(st, out, keys[i], seg_idx); } } } - // Index-ordered back placement for one tile: a block scan gives each candidate a deterministic rank (`base` plus the - // count of preceding candidates in the tile). Placement rule matches `emit_arrival`. Returns the tile's candidate - // total. Used where the K-boundary falls in this tile (terminal tile, or the lazy path's crossing tile). + // Index-ordered back placement for one tile: a block scan over the candidate mask gives each candidate a + // deterministic rank (`base` plus the count of preceding candidates in the tile). A winner (rank below `num_back`) + // lands at forward output slot `num_selected + rank`, matching `place_tile`. Returns the tile's candidate total. Used + // where the K-boundary falls in this tile (terminal tile, or the lazy path's crossing tile). template _CCCL_DEVICE _CCCL_FORCEINLINE offset_t emit_indexed( det_filter_state& st, const key_t (&keys)[ItemsPerThread], - offset_t (&flags)[ItemsPerThread], + const offset_t (&flags)[ItemsPerThread], offset_t seg_base, int count, int tile_base, offset_t base) { + const offset_t num_selected = static_cast(k) - static_cast(st.num_back); // front region size + // `flags[]` is 3-valued (none/candidate/selected); the scan must count candidates only, so reduce to a 0/1 mask. + offset_t cand[ItemsPerThread]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + cand[i] = (flags[i] == flag_candidate) ? offset_t{1} : offset_t{0}; + } offset_t excl[ItemsPerThread]; offset_t tile_total = 0; - block_scan_t(temp_storage.scan_storage).ExclusiveSum(flags, excl, tile_total); + block_scan_t(temp_storage.scan_storage).ExclusiveSum(cand, excl, tile_total); _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; ++i) { - if (flags[i] != offset_t{0}) + if (cand[i] != offset_t{0}) { const offset_t global_rank = base + excl[i]; if (global_rank < static_cast(st.num_back)) { const int pos = tile_base + tid * ItemsPerThread + i; - const out_offset_t out = static_cast(k - 1) - static_cast(global_rank); + const out_offset_t out = static_cast(num_selected + global_rank); const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - st.block_keys_out[out] = keys[i]; - final_filter_write_value(out, seg_idx); + final_filter_place(st, out, keys[i], seg_idx); } } } @@ -1441,10 +1472,9 @@ private: for (int i = 0; i < ItemsPerThread; ++i) { const int pos = tile_base + tid * ItemsPerThread + i; - // Zero-init doubles as the out-of-range/non-candidate mask -- only in-range candidates set it below, so the - // back-placement loops gate on `flags[i]` alone. Validity (`pos < count`) and the selected class are cheap - // recomputed on the loaded key, not cached per item. - flags[i] = offset_t{0}; + // Classify each in-range key once into `flags[i]` (see `flag_*`); `place_tile`/`emit_indexed` read it back + // without re-running `identify_op`. Out-of-range lanes stay `flag_none` and are skipped everywhere. + flags[i] = flag_none; if (pos < count) { // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` @@ -1459,59 +1489,42 @@ private: keys[i] = layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; } - flags[i] = (st.identify_op(keys[i]) == detail::topk::candidate_class::candidate) ? offset_t{1} : offset_t{0}; - } - } - - // Strictly-selected keys go to this block's front slice via a block-local SMEM atomic offset by `st.sel_prefix`. - // The per-block slices (disjoint by `st.sel_prefix`) together fill `[0, num_selected)`. Candidates never fold in - // here -- they always route through the back below, even on early stop. The `&&` short-circuits so `identify_op` - // only touches in-range (loaded) keys. - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < ItemsPerThread; ++i) - { - const int pos = tile_base + tid * ItemsPerThread + i; - if (pos < count && st.identify_op(keys[i]) == detail::topk::candidate_class::selected) - { - const offset_t local = atomicAdd(&temp_storage.front_local_cnt, offset_t{1}); - const out_offset_t out = static_cast(st.sel_prefix + local); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - st.block_keys_out[out] = keys[i]; - final_filter_write_value(out, seg_idx); + const auto cls = st.identify_op(keys[i]); + flags[i] = (cls == detail::topk::candidate_class::candidate) + ? flag_candidate + : (cls == detail::topk::candidate_class::selected ? flag_selected : flag_none); } } - // Back/tie placement (only while this CTA still has unresolved ties). Three block-uniform sub-paths pick the - // candidate rank (see `emit_arrival`/`emit_indexed` for the shared rank->slot rule): - // * is_select_all_cand_cta -- every candidate wins: arrival-order SMEM atomics, no scan, no barrier. + // Placement. `place_tile` routes selected keys to the front on every tile and, on arrival tiles, candidates to + // the back in arrival order. Two block-uniform sub-paths need a block scan instead and run through + // `emit_indexed`: // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. - // * other tiles -- straddling CTA, boundary not yet known: place in arrival order, then `B1` to read the - // counter and, on the crossing tile only, overwrite the arrival slots in index order. - if (st.is_tie_active) + // * lazy crossing tile -- boundary not yet known: candidates land in arrival order first, then `B1` reads the + // counter and, on the tile that crosses `num_back`, overwrites those slots in index + // order. + // A select-all CTA never scans (every candidate wins), so it places entirely via `place_tile`'s arrival route. + const bool is_terminal_tile = st.is_tie_active && region_is_terminal && (tile_base + tile_size >= count); + const bool do_arrival = st.is_tie_active && (st.is_select_all_cand_cta || !is_terminal_tile); + place_tile(st, keys, flags, seg_base, count, tile_base, do_arrival); + + if (st.is_tie_active && !st.is_select_all_cand_cta) { - const bool is_terminal_tile = region_is_terminal && (tile_base + tile_size >= count); - - if (st.is_select_all_cand_cta) - { - emit_arrival(st, keys, flags, seg_base, count, tile_base); - } - else if (is_terminal_tile) + if (is_terminal_tile) { - // Straddling CTA's last tile necessarily holds the boundary: scan it directly in index order. st.running += emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); st.is_tie_active = false; } else { - emit_arrival(st, keys, flags, seg_base, count, tile_base); - // B1: order the arrival global writes ahead of the index-order overwrite (same boundary slots) and make - // the counter read race-free and block-uniform. + // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary + // slots) and make the counter read race-free and block-uniform. __syncthreads(); const offset_t placed = temp_storage.back_local_cnt; if ((st.cand_prefix + placed) > static_cast(st.num_back)) { - // Boundary tile: overwrite this tile's arrival slots `{k-1-st.running, ...}` with the index-ordered winners - // (identical slot set, different candidate->slot mapping). + // Crossing tile: overwrite this tile's arrival slots `{num_selected+st.running, ...}` with the + // index-ordered winners (identical slot set, different candidate->slot mapping). emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); } st.running = st.cand_prefix + placed; @@ -1519,9 +1532,9 @@ private: { st.is_tie_active = false; } - // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the - // next tile's `back_local_cnt` atomic. The other sub-paths write disjoint slots and need no additional - // explicit barrier. + // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the next + // tile's `back_local_cnt` atomic. The other paths write disjoint slots and need no additional explicit + // barrier. __syncthreads(); } } @@ -1925,11 +1938,12 @@ private: } } - // Deterministic final filter: strictly-selected keys fill the front via a SMEM atomic (offset by `sel_prefix`); - // candidates fill the back -- arrival-order atomics away from the K-boundary, an index-ordered BlockScan only on the - // single boundary-crossing (straddling) CTA. A combined cross-CTA scan gives this block its disjoint front/back - // bases and lets it detect whether all/none/some of its candidates win. This member computes the `det_filter_state` - // inputs and hands them to `run_filter`, which drives the per-region sweeps. + // Deterministic final filter: both regions fill forward -- strictly-selected keys into the front `[0, num_selected)` + // via a SMEM atomic (offset by `sel_prefix`), candidates into the back `[num_selected, k)`. Candidate placement uses + // arrival-order atomics, with an index-ordered BlockScan only on the single boundary-crossing (straddling) CTA. A + // combined cross-CTA scan gives this block its disjoint front/back bases and lets it detect whether all/none/some of + // its candidates win. This member computes the `det_filter_state` inputs and hands them to `run_filter`, which drives + // the per-region sweeps. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) From 7dc328479b95a47d329d88806cac14f6763cd506 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 03:31:04 +0200 Subject: [PATCH 103/126] Unify det and non-det filters --- cub/cub/agent/agent_batched_topk_cluster.cuh | 589 ++++++++----------- 1 file changed, 238 insertions(+), 351 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 1053f58d6fd..6e7ea59b236 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -304,18 +304,14 @@ struct agent_batched_topk_cluster // Two clamp flavors. The `floor` clamp pairs with an unpredicated main loop over full tiles plus a single // non-unrolled remainder loop (the chunk helper `for_each_chunk_key_impl` and the copy fast path); the `ceil` clamp - // keeps the whole resident segment inside one tile for fully-predicated loops (the deterministic filter). + // keeps the whole resident segment inside one tile for fully-predicated loops (the final filters' `process_tiles`). static constexpr int histogram_items_per_thread_clamped = clamp_unroll(segment_rounds_floor, histogram_items_per_thread); static constexpr int tie_break_items_per_thread_clamped = clamp_unroll(segment_rounds_ceil, tie_break_items_per_thread); - static constexpr int tie_break_items_per_thread_floor_clamped = - clamp_unroll(segment_rounds_floor, tie_break_items_per_thread); static constexpr int copy_items_per_thread_clamped = clamp_unroll(segment_rounds_floor, copy_items_per_thread); static_assert(histogram_items_per_thread_clamped >= 1, "histogram_items_per_thread_clamped must be positive"); static_assert(tie_break_items_per_thread_clamped >= 1, "tie_break_items_per_thread_clamped must be positive"); - static_assert(tie_break_items_per_thread_floor_clamped >= 1, - "tie_break_items_per_thread_floor_clamped must be positive"); static_assert(copy_items_per_thread_clamped >= 1, "copy_items_per_thread_clamped must be positive"); static constexpr int num_buckets = 1 << bits_per_pass; @@ -329,9 +325,9 @@ struct agent_batched_topk_cluster // count. using smem_keys_t = ::cuda::std::span; - // Tie-break unroll for the deterministic filter's streamed overflow, which feeds `process_tiles` one chunk slot at a - // time: clamp items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit - // to <= chunk granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. + // Tie-break unroll for either filter's streamed overflow, which feeds `process_tiles` one chunk slot at a time: clamp + // items so the tile (`threads_per_block * items`) stays within a chunk, bounding the per-tile early-exit to <= chunk + // granularity. Resident/edge regions keep the full `tie_break_items_per_thread_clamped`. Floors at 1. static constexpr int tie_break_items_streamed = ::cuda::std::clamp(chunk_items / threads_per_block, 1, tie_break_items_per_thread_clamped); static_assert(tie_break_items_streamed >= 1, "tie_break_items_streamed must be positive"); @@ -561,18 +557,6 @@ struct agent_batched_topk_cluster }); } - // Like `for_each_chunk_key`, but also hands `f` each key's segment-local index `base_off + local`, where `base_off` - // is the segment-local offset of the chunk's first element. The pair path uses that index to fetch the key's value - // payload from gmem, so overflow keys can be reused from the streaming SMEM pipeline instead of re-read from gmem. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - for_each_chunk_key_indexed(const key_t* chunk_keys, int count, offset_t base_off, F&& f) const - { - for_each_chunk_key_impl(chunk_keys, count, [&](const key_t& key, int local) { - f(key, base_off + static_cast(local)); - }); - } - // --------------------------------------------------------------------------- // Async bulk-copy pipeline helpers (raw mbarrier + cp.async.bulk via cuda::ptx) // --------------------------------------------------------------------------- @@ -1090,31 +1074,16 @@ private: } } - // Fold every overflow key once in the current ping-pong direction. `Indexed` selects the callable shape: keys-only - // (`Indexed == false`) applies `f(key)`; the pair filter (`Indexed == true`) applies `f(key, seg_idx)` with each - // key's segment-local index, needed to fetch its value payload from gmem while still reusing the streamed overflow - // keys. The index math lives entirely inside the `if constexpr (Indexed)` arms, so it is elided in the keys-only - // instantiation. See `run_pass` for the overlap semantics of `mid`; `UnrollFactor` partially unrolls the generic - // (gmem fallback) fold loop. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_dispatch(F&& f, Mid&& mid) + // Apply `f(key)` to every overflow key once in the current ping-pong direction. `UnrollFactor` partially unrolls the + // generic (gmem fallback) fold loop; callers pass their clamped items-per-thread. See `run_pass` for the overlap + // semantics of `mid`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) { run_pass( [&](int stage, offset_t overflow_idx) { - if constexpr (Indexed) - { - const offset_t base_off = - get_chunk( - layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items) - .offset; - const auto keys = stage_span(stage, overflow_idx); - for_each_chunk_key_indexed(keys.data(), static_cast(keys.size()), base_off, f); - } - else - { - const auto keys = stage_span(stage, overflow_idx); - for_each_chunk_key(keys.data(), static_cast(keys.size()), f); - } + const auto keys = stage_span(stage, overflow_idx); + for_each_chunk_key(keys.data(), static_cast(keys.size()), f); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1124,15 +1093,7 @@ private: const int local = j * threads_per_block + tid; if (local < chunk.count) { - if constexpr (Indexed) - { - const offset_t seg_idx = chunk.offset + static_cast(local); - f(layout.block_keys_in[static_cast(seg_idx)], seg_idx); - } - else - { - f(layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); - } + f(layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); } } }, @@ -1142,14 +1103,6 @@ private: }); } - // `f(key)` over every overflow key. `UnrollFactor` partially unrolls the generic (gmem fallback) fold loop; callers - // pass their clamped items-per-thread. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) - { - process_pass_dispatch(static_cast(f), static_cast(mid)); - } - // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in // by the BlockLoadToShared pipeline (rather than already resident in SMEM). template @@ -1158,13 +1111,6 @@ private: process_pass(static_cast(f), [] {}); } - // Like `process_pass`, but applies `f(key, seg_idx)` where `seg_idx` is the key's segment-local index. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass_indexed(F&& f, Mid&& mid) - { - process_pass_dispatch(static_cast(f), static_cast(mid)); - } - // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- @@ -1275,7 +1221,7 @@ private: // Deterministic final-filter state: the `run()`-local tie-break values (counts, prefixes, region extents) plus the // mutable `running`/`is_tie_active` scan cursor, bundled into one POD so the flattened filter methods below take a - // single `st` argument instead of ~16. A deterministic-path-only local of `write_deterministic_topk`; never + // single `state` argument instead of ~16. A deterministic-path-only local of `write_deterministic_topk`; never // instantiated on the non-deterministic path. Segment-invariant inputs (`segment_id`, `k`, and the `layout.*` // geometry) are read from the agent directly rather than duplicated here. `resident_front` carries the resident // window as a `smem_keys_t` view. @@ -1304,11 +1250,12 @@ private: }; // --------------------------------------------------------------------------- - // Deterministic final filter + // Final-filter placement (shared) + deterministic filter // --------------------------------------------------------------------------- - // The per-region index-ordered sweeps below are named agent methods over a `det_filter_state` (`st`), driven by - // `run_filter`; `write_deterministic_topk` computes `st` and calls `run_filter(st)`. See `write_deterministic_topk` - // for the front/back placement scheme. All methods are `_CCCL_FORCEINLINE`. + // The classify/place helpers below (`place_one`, `place_tile`, `process_tiles`, ...) are shared by both filters. The + // deterministic-only sweeps that follow are named agent methods over a `det_filter_state` (`state`), driven by + // `run_filter`; `write_deterministic_topk` computes `state` and calls `run_filter(state)`. See + // `write_deterministic_topk` for the front/back placement scheme. All methods are `_CCCL_FORCEINLINE`. // Shared by both final filters: for each key written to `block_keys_out[pos]`, load the associated input value at the // key's segment-local index `seg_idx` from gmem and store it at the same slot. Compiled out (and never dereferences @@ -1332,14 +1279,29 @@ private: static constexpr offset_t flag_candidate = 1; // tie candidate: routed to the back, counted by the boundary scan static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front - // Shared output write for the deterministic placement paths (`place_tile`, `emit_indexed`): store the key at - // `block_keys_out[out]`, then (pairs builds) copy its value from segment-local index `seg_idx`. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - final_filter_place(det_filter_state& st, out_offset_t out, const key_t& key, offset_t seg_idx) - { - st.block_keys_out[out] = key; - final_filter_write_value(out, seg_idx); + // Shared per-key arrival placement core (called by `place_tile` for both final filters): route one key to its region + // by class -- selected to the front (base `sel_prefix`), candidate to the back (base `num_selected + cand_prefix`, + // where `num_selected = k - `) -- via a block-local SMEM atomic under the uniform `out < k` guard + // that drops losing candidates while always accepting selected keys. Pairs builds also copy the key's value from its + // segment-local index `seg_idx`. (Index-ordered deterministic tie resolution goes through `emit_indexed` instead.) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void place_one( + KeyOutIt block_keys_out, + bool is_cand, + const key_t& key, + offset_t seg_idx, + offset_t num_selected, + offset_t sel_prefix, + offset_t cand_prefix) + { + offset_t* const counter = is_cand ? &temp_storage.back_local_cnt : &temp_storage.front_local_cnt; + const offset_t base = is_cand ? (num_selected + cand_prefix) : sel_prefix; + const out_offset_t out = static_cast(base + atomicAdd(counter, offset_t{1})); + if (out < k) + { + block_keys_out[out] = key; + final_filter_write_value(out, seg_idx); + } } // Uniform "all placed" predicate: true once this block has emitted all `my_front` strictly-selected keys and @@ -1347,29 +1309,27 @@ private: // lanes that raced ahead through the barrier-free tiles). Polled only at critical points -- between regions and // before each streaming bulk copy -- never per tile. template - _CCCL_DEVICE _CCCL_FORCEINLINE bool final_filter_should_stop(const det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE bool final_filter_should_stop(const det_filter_state& state) { // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. - _CCCL_ASSERT(temp_storage.front_local_cnt <= static_cast(st.my_front) - && temp_storage.back_local_cnt <= st.my_cand_count, + _CCCL_ASSERT(temp_storage.front_local_cnt <= static_cast(state.my_front) + && temp_storage.back_local_cnt <= state.my_cand_count, "final-filter placement counters exceeded this CTA's assigned work"); - const bool is_front_done = temp_storage.front_local_cnt >= static_cast(st.my_front); + const bool is_front_done = temp_storage.front_local_cnt >= static_cast(state.my_front); // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never // clears it) finishes once all `my_cand_count` of its candidates are placed. - const bool is_back_done = !st.is_tie_active || (temp_storage.back_local_cnt >= st.my_cand_count); + const bool is_back_done = !state.is_tie_active || (temp_storage.back_local_cnt >= state.my_cand_count); return is_front_done && is_back_done; } - // Single-pass placement for one tile: both regions fill forward -- front `[0, num_selected)`, back `[num_selected, - // k)` -- so a key's class only selects the counter and base offset (below), and one uniform `out < k` guard drops - // losing candidates while always accepting selected keys (whose `out` is `< num_selected`). The candidate atomic - // advances win or lose; arrival order is fine because either every candidate wins (select-all) or the lazy crossing - // tile is overwritten in index order later. `do_arrival == false` (terminal tile only) skips candidates here and - // leaves them to `emit_indexed`; the lazy path passes `true` and its crossing tile is later overwritten by - // `emit_indexed`. - template + // Arrival-order placement for one tile: each in-range key is routed by `place_one` (both regions fill forward). The + // candidate atomic advances win or lose; arrival order is fine because either every candidate wins (select-all) or + // the lazy crossing tile is overwritten in index order later. `do_arrival == false` (terminal tile only) skips + // candidates here and leaves them to `emit_indexed`; the lazy path passes `true` and its crossing tile is later + // overwritten by `emit_indexed`. + template _CCCL_DEVICE _CCCL_FORCEINLINE void place_tile( - det_filter_state& st, + State& state, const key_t (&keys)[ItemsPerThread], const offset_t (&flags)[ItemsPerThread], offset_t seg_base, @@ -1377,7 +1337,7 @@ private: int tile_base, bool do_arrival) { - const offset_t num_selected = static_cast(k) - static_cast(st.num_back); // front region size + const offset_t num_selected = static_cast(k) - static_cast(state.num_back); // front region size _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; ++i) { @@ -1386,15 +1346,9 @@ private: { continue; } - offset_t* const counter = is_cand ? &temp_storage.back_local_cnt : &temp_storage.front_local_cnt; - const offset_t base = is_cand ? (num_selected + st.cand_prefix) : st.sel_prefix; - const out_offset_t out = static_cast(base + atomicAdd(counter, offset_t{1})); - if (out < k) - { - const int pos = tile_base + tid * ItemsPerThread + i; - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - final_filter_place(st, out, keys[i], seg_idx); - } + const int pos = tile_base + tid * ItemsPerThread + i; + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + place_one(state.block_keys_out, is_cand, keys[i], seg_idx, num_selected, state.sel_prefix, state.cand_prefix); } } @@ -1402,9 +1356,9 @@ private: // deterministic rank (`base` plus the count of preceding candidates in the tile). A winner (rank below `num_back`) // lands at forward output slot `num_selected + rank`, matching `place_tile`. Returns the tile's candidate total. Used // where the K-boundary falls in this tile (terminal tile, or the lazy path's crossing tile). - template + template _CCCL_DEVICE _CCCL_FORCEINLINE offset_t emit_indexed( - det_filter_state& st, + State& state, const key_t (&keys)[ItemsPerThread], const offset_t (&flags)[ItemsPerThread], offset_t seg_base, @@ -1412,7 +1366,7 @@ private: int tile_base, offset_t base) { - const offset_t num_selected = static_cast(k) - static_cast(st.num_back); // front region size + const offset_t num_selected = static_cast(k) - static_cast(state.num_back); // front region size // `flags[]` is 3-valued (none/candidate/selected); the scan must count candidates only, so reduce to a 0/1 mask. offset_t cand[ItemsPerThread]; _CCCL_PRAGMA_UNROLL_FULL() @@ -1429,33 +1383,36 @@ private: if (cand[i] != offset_t{0}) { const offset_t global_rank = base + excl[i]; - if (global_rank < static_cast(st.num_back)) + if (global_rank < static_cast(state.num_back)) { - const int pos = tile_base + tid * ItemsPerThread + i; - const out_offset_t out = static_cast(num_selected + global_rank); - const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - final_filter_place(st, out, keys[i], seg_idx); + const int pos = tile_base + tid * ItemsPerThread + i; + const out_offset_t out = static_cast(num_selected + global_rank); + const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); + state.block_keys_out[out] = keys[i]; + final_filter_write_value(out, seg_idx); } } } return tile_total; } - // Process a flat span of `count` keys in scan order, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` - // selects the key source: the SMEM buffer `smem_src` (indexed by the folded in-region position) or gmem - // `block_keys_in` (indexed by the segment-local index `seg_base + folded`). `Reversed` walks the span high-to-low. - // Strictly-selected keys go to the front via a SMEM atomic (offset by `st.sel_prefix`); candidates go to the back - // (see the per-tile logic). `region_is_terminal` marks this region as the CTA's last, so its last tile -- if ties are - // unresolved -- holds the boundary and scans directly. `st.running` carries across tiles/regions. The only explicit - // per-tile `__syncthreads()` are in the lazy-scan `else` branch (block scans elsewhere synchronize internally); early - // exit is decided at critical points via `final_filter_should_stop`. - template + // Classify and place each of `count` keys in scan order, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` + // picks the key source: `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` (indexed by the + // segment-local index `seg_base + folded`); `Reversed` walks the span high-to-low. Both final filters share this via + // `Deterministic`: + // * `true` (a `det_filter_state`): resolves the boundary ties in segment-index order via `emit_indexed`, so callers + // walk the regions in global-index order and carry `state.running`/`state.is_tie_active` across tiles; the + // lazy-scan `else` branch holds the only explicit per-tile `__syncthreads()`; `region_is_terminal` marks the + // CTA's last. + // * `false` (a `nondet_filter_state`): all arrival-order placement (`place_tile` with `do_arrival == true`); no + // scan/tie-state/barrier, `Reversed`/`region_is_terminal` unused. + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles( - det_filter_state& st, + State& state, [[maybe_unused]] const key_t* smem_src, offset_t seg_base, int count, - bool region_is_terminal) + [[maybe_unused]] bool region_is_terminal) { // The SMEM-source clause is exempt when reading from gmem (`FromSmem == false`, the generic overflow fallback // passes `nullptr`) or for an empty tile. @@ -1489,53 +1446,62 @@ private: keys[i] = layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; } - const auto cls = st.identify_op(keys[i]); + const auto cls = state.identify_op(keys[i]); flags[i] = (cls == detail::topk::candidate_class::candidate) ? flag_candidate : (cls == detail::topk::candidate_class::selected ? flag_selected : flag_none); } } - // Placement. `place_tile` routes selected keys to the front on every tile and, on arrival tiles, candidates to - // the back in arrival order. Two block-uniform sub-paths need a block scan instead and run through - // `emit_indexed`: - // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. - // * lazy crossing tile -- boundary not yet known: candidates land in arrival order first, then `B1` reads the - // counter and, on the tile that crosses `num_back`, overwrites those slots in index - // order. - // A select-all CTA never scans (every candidate wins), so it places entirely via `place_tile`'s arrival route. - const bool is_terminal_tile = st.is_tie_active && region_is_terminal && (tile_base + tile_size >= count); - const bool do_arrival = st.is_tie_active && (st.is_select_all_cand_cta || !is_terminal_tile); - place_tile(st, keys, flags, seg_base, count, tile_base, do_arrival); - - if (st.is_tie_active && !st.is_select_all_cand_cta) + if constexpr (!Deterministic) { - if (is_terminal_tile) - { - st.running += emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); - st.is_tie_active = false; - } - else + // Non-deterministic: place every tile in arrival order (selected to the front, the first `num_back` candidates + // to the back via `place_one`'s `out < k` guard). No index-ordered scan, tie-state, or per-tile barrier. + place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/true); + } + else + { + // Placement. `place_tile` routes selected keys to the front on every tile and, on arrival tiles, candidates to + // the back in arrival order. Two block-uniform sub-paths need a block scan instead and run through + // `emit_indexed`: + // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. + // * lazy crossing tile -- boundary not yet known: candidates land in arrival order first, then `B1` reads the + // counter and, on the tile that crosses `num_back`, overwrites those slots in index + // order. + // A select-all CTA never scans (every candidate wins), so it places entirely via `place_tile`'s arrival route. + const bool is_terminal_tile = state.is_tie_active && region_is_terminal && (tile_base + tile_size >= count); + const bool do_arrival = state.is_tie_active && (state.is_select_all_cand_cta || !is_terminal_tile); + place_tile(state, keys, flags, seg_base, count, tile_base, do_arrival); + + if (state.is_tie_active && !state.is_select_all_cand_cta) { - // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary - // slots) and make the counter read race-free and block-uniform. - __syncthreads(); - const offset_t placed = temp_storage.back_local_cnt; - if ((st.cand_prefix + placed) > static_cast(st.num_back)) + if (is_terminal_tile) { - // Crossing tile: overwrite this tile's arrival slots `{num_selected+st.running, ...}` with the - // index-ordered winners (identical slot set, different candidate->slot mapping). - emit_indexed(st, keys, flags, seg_base, count, tile_base, st.running); + state.running += emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); + state.is_tie_active = false; } - st.running = st.cand_prefix + placed; - if (st.running >= static_cast(st.num_back)) + else { - st.is_tie_active = false; + // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary + // slots) and make the counter read race-free and block-uniform. + __syncthreads(); + const offset_t placed = temp_storage.back_local_cnt; + if ((state.cand_prefix + placed) > static_cast(state.num_back)) + { + // Crossing tile: overwrite this tile's arrival slots `{num_selected+state.running, ...}` with the + // index-ordered winners (identical slot set, different candidate->slot mapping). + emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); + } + state.running = state.cand_prefix + placed; + if (state.running >= static_cast(state.num_back)) + { + state.is_tie_active = false; + } + // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the next + // tile's `back_local_cnt` atomic. The other paths write disjoint slots and need no additional explicit + // barrier. + __syncthreads(); } - // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the next - // tile's `back_local_cnt` atomic. The other paths write disjoint slots and need no additional explicit - // barrier. - __syncthreads(); } } } @@ -1546,13 +1512,16 @@ private: // (`resident_base`) in reverse, so a single `process_tiles` call per span with the index folded at compile time // replaces the old fwd/rev pair. template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident(det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident(det_filter_state& state) { if constexpr (use_block_load_to_shared) { // Whole contiguous resident span staged in SMEM. - process_tiles( - st, st.resident_front.data(), st.front_seg_base, st.front_count, st.is_resident_terminal); + process_tiles( + state, state.resident_front.data(), state.front_seg_base, state.front_count, state.is_resident_terminal); } else { @@ -1564,8 +1533,11 @@ private: const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); // Generic multi-chunk resident loop reads the chunk's SMEM slot; never the terminal-tile fast path (the lazy // per-tile boundary detection handles any boundary), so pass `false`. - process_tiles( - st, ::cuda::ptr_rebind(key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); + process_tiles( + state, ::cuda::ptr_rebind(key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); } } } @@ -1577,7 +1549,7 @@ private: // `stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the // generic fallback re-reads gmem each pass regardless). template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow(det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow(det_filter_state& state) { // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at // `stream_is_forward == !is_tie_reversed` after the full `num_passes`. The escape disjuncts are the cases where @@ -1585,7 +1557,7 @@ private: // (`is_select_all_cand_cta`, arrival-order atomics); or one with no back scan (`!is_tie_active`). Early stop is // subsumed -- it makes every CTA `is_select_all_cand_cta`, so the shorter (possibly mis-parity) pass count never // reaches a straddler. - _CCCL_ASSERT(layout.overflow_chunks == 0 || st.is_select_all_cand_cta || !st.is_tie_active + _CCCL_ASSERT(layout.overflow_chunks == 0 || state.is_select_all_cand_cta || !state.is_tie_active || stream_is_forward == (!is_tie_reversed), "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " "the wrong streaming direction"); @@ -1602,14 +1574,16 @@ private: // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == // false`): a stray terminal direct scan here would need the stream to flag its last chunk, and the saving // is one barrier on one tile. - process_tiles( - st, keys.data(), base_off, static_cast(keys.size()), false); + process_tiles( + state, keys.data(), base_off, static_cast(keys.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). [&](const auto& chunk) { // FromSmem=false: `smem_src` is unread, so pass nullptr. - process_tiles( - st, nullptr, chunk.offset, chunk.count, false); + process_tiles(state, nullptr, chunk.offset, chunk.count, false); }, // No interleaved resident work: the deterministic filter folds its resident span separately. [] {}, @@ -1617,7 +1591,7 @@ private: // lanes that drifted through the just-folded chunk's barrier-free tiles (polled before each refill copy). [&] { __syncthreads(); - return !final_filter_should_stop(st); + return !final_filter_should_stop(state); }); } @@ -1626,13 +1600,14 @@ private: // than `load_align_items` items (it fits in its `edge_keys` slot), so it runs `process_tiles` at unroll factor 1 // (non-unrolled) instead of the wider tie-break unroll the potentially large resident/overflow regions use. template - _CCCL_DEVICE _CCCL_FORCEINLINE void - process_edge(det_filter_state& st, const key_t* keys, offset_t seg_base, int count, bool is_terminal) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge( + det_filter_state& state, const key_t* keys, offset_t seg_base, int count, bool is_terminal) { if constexpr (use_block_load_to_shared) { _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed>(st, keys, seg_base, count, is_terminal); + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true>( + state, keys, seg_base, count, is_terminal); } } @@ -1640,22 +1615,22 @@ private: // 0). In global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or // empty prefix is a no-op. template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge(det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge(det_filter_state& state) { - process_edge(st, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, st.is_head_edge_terminal); + process_edge(state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, state.is_head_edge_terminal); } // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` // (base `segment_size - tail_edge_len_items`). In global-index order it is the trailing region (ascending) / // leading region (descending); an aligned or non-owned tail is a no-op. template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge(det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge(det_filter_state& state) { - process_edge(st, + process_edge(state, temp_storage.edge_keys + head_edge_cap_items, layout.segment_size_off - static_cast(layout.tail_edge_len_items), layout.tail_edge_len_items, - st.is_tail_edge_terminal); + state.is_tail_edge_terminal); } // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between @@ -1665,46 +1640,46 @@ private: // top-k is placed; `is_residency_reversed` keeps the first-visited (high-index) chunks resident in the descending // order so this holds. template - _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter(det_filter_state& st) + _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter(det_filter_state& state) { const auto step = [&](auto&& region) { // Barrier before polling: makes the placement-counter reads block-wide and resynchronizes lanes that raced // ahead through the previous region's barrier-free tiles. __syncthreads(); - if (!final_filter_should_stop(st)) + if (!final_filter_should_stop(state)) { region(); } }; if constexpr (is_tie_reversed) { - process_tail_edge(st); + process_tail_edge(state); } else { - process_head_edge(st); + process_head_edge(state); } step([&] { - process_resident(st); + process_resident(state); }); step([&] { - process_overflow(st); + process_overflow(state); }); step([&] { if constexpr (is_tie_reversed) { - process_head_edge(st); + process_head_edge(state); } else { - process_tail_edge(st); + process_tail_edge(state); } }); } // Fold the persistent boundary edges into `apply` (head prefix on rank 0; peeled tail suffix on the tail owner), - // reading the keys already staged in `edge_keys`. Used by the histogram passes and the non-deterministic final - // filter; the deterministic filter folds the edges as separate index-ordered regions instead. + // reading the keys already staged in `edge_keys`. Used by the histogram passes; both final filters fold the edges + // through `process_tiles` instead (as separate regions for the deterministic filter, via `nondet_fold_resident`). template _CCCL_DEVICE _CCCL_FORCEINLINE void fold_boundary_edges(int head_edge_len_items, int tail_edge_len_items, Apply&& apply) @@ -1727,173 +1702,83 @@ private: // --------------------------------------------------------------------------- // Non-deterministic final filter // --------------------------------------------------------------------------- - // The sweep helpers below are named agent methods over a `nondet_filter_state` (`st`); `k`, `segment_id`, and - // `layout.*` are read from the agent. `write_nondeterministic_topk` computes `st` and drives the sweeps (see - // `nondet_place` for the per-key front/back placement). All methods are `_CCCL_FORCEINLINE`. + // Order-independent counterpart of the deterministic filter: every region is swept by + // `process_tiles<..., Deterministic=false>` (arrival-order placement, no scan/tie-state/barrier). Since order is + // irrelevant, `write_nondeterministic_topk` folds the resident keys as the overflow stream's `mid` work (overlapping + // the first reloads) and never bails early. template struct nondet_filter_state { IdentifyOp identify_op; KeyOutIt block_keys_out; - out_offset_t num_kth; + out_offset_t num_back; offset_t sel_prefix; offset_t cand_prefix; smem_keys_t resident_keys; }; - // Classify one key and place it: strictly-selected keys go to the front (`sel_prefix` + a block-local atomic); the - // first `num_kth` candidates (arrival order) go to the back (`cand_prefix` + a block-local atomic), later candidates - // are dropped. In pair mode each written key additionally pulls its value from gmem at `seg_idx`; keys-only elides - // that and passes a dummy `seg_idx`. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - nondet_place(const nondet_filter_state& st, const key_t& key, [[maybe_unused]] offset_t seg_idx) - { - const auto res = st.identify_op(key); - if (res == detail::topk::candidate_class::selected) - { - const out_offset_t pos = - static_cast(st.sel_prefix + atomicAdd(&temp_storage.front_local_cnt, offset_t{1})); - st.block_keys_out[pos] = key; - if constexpr (!is_keys_only) - { - final_filter_write_value(pos, seg_idx); - } - } - else if (res == detail::topk::candidate_class::candidate) - { - const out_offset_t back_pos = - static_cast(st.cand_prefix + atomicAdd(&temp_storage.back_local_cnt, offset_t{1})); - if (back_pos < st.num_kth) - { - const out_offset_t pos = k - 1 - back_pos; - st.block_keys_out[pos] = key; - if constexpr (!is_keys_only) - { - final_filter_write_value(pos, seg_idx); - } - } - } - } - - // Fold a contiguous run of `count` keys staged in SMEM at `smem`, whose element `local` has segment-local index - // `base_off + local` (pair mode needs that index for the value fetch). Materialize each key into a register first: - // `nondet_place` reads it more than once, so indexing `smem[local]` per use would re-issue a narrow `LDS`. + // Fold this rank's resident keys (and boundary edges) through `process_tiles`, run as the overflow pass's `mid` work + // so they overlap the first overflow reloads. Placement uses order-independent atomics into slots disjoint from the + // streaming slots, so this never races the overflow apply. Each chunk's `chunk.offset` is passed as the region's + // segment base (used for the pair value fetch; elided in keys-only). template - _CCCL_DEVICE _CCCL_FORCEINLINE void - nondet_write_run(const nondet_filter_state& st, const key_t* smem, offset_t base_off, int count) + _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_fold_resident(nondet_filter_state& state) { - const int iterations = ::cuda::ceil_div(count, threads_per_block); - _CCCL_PRAGMA_UNROLL(tie_break_items_per_thread_clamped) - for (int j = 0; j < iterations; ++j) + if constexpr (use_block_load_to_shared) { - const int local = j * threads_per_block + tid; - if (local < count) + // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers per-chunk + // spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into `edge_keys` and folded + // below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not `chunk.count`. + key_t* const resident_ptr = state.resident_keys.data(); + int cursor_items = 0; + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const key_t key = smem[local]; - nondet_place(st, key, base_off + static_cast(local)); + const auto chunk = get_chunk( + layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); + const int bulk_count_items = static_cast( + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); + process_tiles( + state, resident_ptr + cursor_items, chunk.offset, bulk_count_items, false); + cursor_items += bulk_count_items; } - } - } - - // Fold a boundary edge (fewer than `load_align_items` items) with a plain block-stride loop rather than the - // tiled/unrolled `nondet_write_run` used for the potentially large resident bulks. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_write_edge( - const nondet_filter_state& st, const key_t* smem, offset_t base_off, int count) - { - for (int local = tid; local < count; local += threads_per_block) - { - const key_t key = smem[local]; - nondet_place(st, key, base_off + static_cast(local)); - } - } - - // Fold this rank's resident keys (and, in pair mode, their values) as the overflow pass's `mid` work so they overlap - // the first overflow reloads. Writes are order-independent atomics and resident SMEM slots are disjoint from the - // streaming slots, so this never races the overflow apply. Keys-only drives the index-free traversal with a dummy - // `seg_idx`; pair mode recovers each key's segment-local index via `nondet_write_run`/`nondet_write_edge`. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_fold_resident(const nondet_filter_state& st) - { - if constexpr (is_keys_only) - { - const auto write_selected = [&](const key_t& key) { - nondet_place(st, key, offset_t{0}); - }; - if constexpr (use_block_load_to_shared) + // Persistent boundary edges (unroll factor 1: an edge spans fewer than `load_align_items` items): head prefix at + // `edge_keys` (segment index 0), peeled tail suffix at `edge_keys + head_edge_cap_items` (segment index + // `segment_size - tail_edge_len_items`). + if (layout.head_edge_len_items > 0) { - for_each_chunk_key( - st.resident_keys.data(), static_cast(st.resident_keys.size()), write_selected); + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false>( + state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, false); } - else + if (layout.tail_edge_len_items > 0) { - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const offset_t chunk_idx = layout.part.global_index(local_chunk); - const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); - for_each_chunk_key( - ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), - static_cast(chunk.count), - write_selected); - } + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false>( + state, + temp_storage.edge_keys + head_edge_cap_items, + layout.segment_size_off - static_cast(layout.tail_edge_len_items), + layout.tail_edge_len_items, + false); } - // Order-independent atomic writes, so the persistent boundary edges fold in alongside the resident keys. - fold_boundary_edges(layout.head_edge_len_items, layout.tail_edge_len_items, write_selected); } else { - if constexpr (use_block_load_to_shared) - { - // Resident keys are densely packed in slot order (aligned bulks only), so a running cursor recovers per-chunk - // spans. Only the global-last chunk can be partial (its unaligned suffix is peeled into `edge_keys` and folded - // below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not `chunk.count`. - key_t* const resident_ptr = st.resident_keys.data(); - int cursor_items = 0; - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const auto chunk = get_chunk( - layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); - const int bulk_count_items = static_cast( - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); - nondet_write_run(st, resident_ptr + cursor_items, chunk.offset, bulk_count_items); - cursor_items += bulk_count_items; - } - } - else - { - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) - { - const auto chunk = - get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); - nondet_write_run(st, - ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), - chunk.offset, - chunk.count); - } - } - // Fold the persistent boundary edges with their segment-local indices: the head prefix starts at index 0, the - // peeled tail suffix at `segment_size - tail_edge_len_items`. - if constexpr (use_block_load_to_shared) + // Generic fallback: each resident chunk is staged in its own `key_slots` slot (indexed by `local_chunk`); no + // edges are peeled (boundary items are read as ordinary chunks). + for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - if (layout.head_edge_len_items > 0) - { - nondet_write_edge(st, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items); - } - if (layout.tail_edge_len_items > 0) - { - nondet_write_edge(st, - temp_storage.edge_keys + head_edge_cap_items, - layout.segment_size_off - static_cast(layout.tail_edge_len_items), - layout.tail_edge_len_items); - } + const auto chunk = get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); + process_tiles( + state, + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + chunk.offset, + chunk.count, + false); } } } // Non-deterministic final filter driver. The combined cross-CTA scan gives this block disjoint front/back bases - // (`sel_prefix`/`cand_prefix`); `nondet_place` then places each key with block-local SMEM atomics (output order not - // preserved). A perf change over the old cluster-wide `out_cnt`/`out_back_cnt` DSMEM atomics, not a correctness one. + // (`sel_prefix`/`cand_prefix`); overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and + // place into block-local SMEM atomics. `should_continue` is always true: every key must be classified. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) @@ -1911,31 +1796,33 @@ private: <= static_cast<::cuda::std::uint64_t>(k), "selected prefix must fit before the tie-back output region"); - nondet_filter_state st{ + nondet_filter_state state{ identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, resident_keys}; - // `f` folds each overflow key; `mid` folds the resident keys + boundary edges, overlapping the first overflow - // reloads. Keys-only passes a dummy `seg_idx`; pair mode recovers each key's index for its value fetch. - if constexpr (is_keys_only) - { - process_pass( - [&](const key_t& key) { - nondet_place(st, key, offset_t{0}); - }, - [&] { - nondet_fold_resident(st); - }); - } - else - { - process_pass_indexed( - [&](const key_t& key, offset_t seg_idx) { - nondet_place(st, key, seg_idx); - }, - [&] { - nondet_fold_resident(st); - }); - } + run_pass( + // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. + [&](int stage, offset_t overflow_idx) { + const auto keys = stage_span(stage, overflow_idx); + const offset_t base_off = + get_chunk( + layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items) + .offset; + process_tiles( + state, keys.data(), base_off, static_cast(keys.size()), false); + }, + // Generic fallback: read the overflow chunk straight from gmem. + [&](const auto& chunk) { + process_tiles( + state, nullptr, chunk.offset, chunk.count, false); + }, + // Fold the resident keys + boundary edges, overlapping the first overflow reloads. + [&] { + nondet_fold_resident(state); + }, + // Never break early: every key must be classified (selected/candidate keys placed, the rest dropped in-tile). + [] { + return true; + }); } // Deterministic final filter: both regions fill forward -- strictly-selected keys into the front `[0, num_selected)` @@ -2027,7 +1914,7 @@ private: // cursor: `running` at this CTA's exclusive back prefix (candidates owned by preceding CTAs), `is_tie_active` at // `!is_select_no_cand_cta` -- true unless this is a select-no-candidates CTA (empty back region, or all of this // CTA's candidates sort past the boundary). - det_filter_state st{ + det_filter_state state{ identify_op, block_keys_out, num_back, @@ -2044,7 +1931,7 @@ private: is_tail_edge_terminal, cand_prefix, !is_select_no_cand_cta}; - run_filter(st); + run_filter(state); } // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into From f0d1799c38b0ccdf52bea752e2350010b7c52e69 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 04:51:16 +0200 Subject: [PATCH 104/126] Prime output rank counters with offsets --- cub/cub/agent/agent_batched_topk_cluster.cuh | 91 ++++++++++++-------- 1 file changed, 57 insertions(+), 34 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 6e7ea59b236..53adb5fceba 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1279,24 +1279,19 @@ private: static constexpr offset_t flag_candidate = 1; // tie candidate: routed to the back, counted by the boundary scan static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front - // Shared per-key arrival placement core (called by `place_tile` for both final filters): route one key to its region - // by class -- selected to the front (base `sel_prefix`), candidate to the back (base `num_selected + cand_prefix`, - // where `num_selected = k - `) -- via a block-local SMEM atomic under the uniform `out < k` guard - // that drops losing candidates while always accepting selected keys. Pairs builds also copy the key's value from its - // segment-local index `seg_idx`. (Index-ordered deterministic tie resolution goes through `emit_indexed` instead.) + // Shared per-key arrival placement core (called by `place_tile` for both final filters): route one key by class to + // the front (selected) or back (candidate) counter, each primed with its region base by the final-filter drivers so + // the SMEM atomic returns the absolute output slot with no per-key offset. The uniform `out < k` guard drops losing + // candidates but always accepts selected keys; pairs builds also copy the key's value from `seg_idx`. (Deterministic + // index-ordered tie resolution goes through `emit_indexed` instead.) template - _CCCL_DEVICE _CCCL_FORCEINLINE void place_one( - KeyOutIt block_keys_out, - bool is_cand, - const key_t& key, - offset_t seg_idx, - offset_t num_selected, - offset_t sel_prefix, - offset_t cand_prefix) + _CCCL_DEVICE _CCCL_FORCEINLINE void + place_one(KeyOutIt block_keys_out, bool is_cand, const key_t& key, offset_t seg_idx) { offset_t* const counter = is_cand ? &temp_storage.back_local_cnt : &temp_storage.front_local_cnt; - const offset_t base = is_cand ? (num_selected + cand_prefix) : sel_prefix; - const out_offset_t out = static_cast(base + atomicAdd(counter, offset_t{1})); + const out_offset_t out = static_cast(atomicAdd(counter, offset_t{1})); + // The front counter never overruns the selected region, so only losing candidates may exceed `k`. + _CCCL_ASSERT(is_cand || out < k, "a strictly-selected key must always land within the top-k output"); if (out < k) { block_keys_out[out] = key; @@ -1311,14 +1306,17 @@ private: template _CCCL_DEVICE _CCCL_FORCEINLINE bool final_filter_should_stop(const det_filter_state& state) { - // Each region is visited at most once, so neither counter can pass this CTA's exact selected/candidate totals. - _CCCL_ASSERT(temp_storage.front_local_cnt <= static_cast(state.my_front) - && temp_storage.back_local_cnt <= state.my_cand_count, + // Counters run absolute (primed with the region base), so completion compares against each region's end. Every + // region is visited at most once, so neither counter can pass its end. + const offset_t num_selected = static_cast(k) - static_cast(state.num_back); + const offset_t front_end = state.sel_prefix + static_cast(state.my_front); + const offset_t back_end = num_selected + state.cand_prefix + state.my_cand_count; + _CCCL_ASSERT(temp_storage.front_local_cnt <= front_end && temp_storage.back_local_cnt <= back_end, "final-filter placement counters exceeded this CTA's assigned work"); - const bool is_front_done = temp_storage.front_local_cnt >= static_cast(state.my_front); + const bool is_front_done = temp_storage.front_local_cnt >= front_end; // Straddling/above CTAs finish the back when `is_tie_active` clears; an `is_select_all_cand_cta` (which never // clears it) finishes once all `my_cand_count` of its candidates are placed. - const bool is_back_done = !state.is_tie_active || (temp_storage.back_local_cnt >= state.my_cand_count); + const bool is_back_done = !state.is_tie_active || (temp_storage.back_local_cnt >= back_end); return is_front_done && is_back_done; } @@ -1337,7 +1335,6 @@ private: int tile_base, bool do_arrival) { - const offset_t num_selected = static_cast(k) - static_cast(state.num_back); // front region size _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; ++i) { @@ -1348,7 +1345,7 @@ private: } const int pos = tile_base + tid * ItemsPerThread + i; const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); - place_one(state.block_keys_out, is_cand, keys[i], seg_idx, num_selected, state.sel_prefix, state.cand_prefix); + place_one(state.block_keys_out, is_cand, keys[i], seg_idx); } } @@ -1485,19 +1482,24 @@ private: // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary // slots) and make the counter read race-free and block-uniform. __syncthreads(); - const offset_t placed = temp_storage.back_local_cnt; - if ((state.cand_prefix + placed) > static_cast(state.num_back)) + // `back_local_cnt` is primed with the back base (`num_selected + cand_prefix`), so subtracting the front + // region size recovers this CTA's candidate-rank reached so far (`cand_prefix` plus the candidates placed). + const offset_t num_selected = static_cast(k) - static_cast(state.num_back); + _CCCL_ASSERT(temp_storage.back_local_cnt >= num_selected, + "back counter must stay at or above its primed base (no unsigned underflow)"); + const offset_t reached = temp_storage.back_local_cnt - num_selected; + if (reached > static_cast(state.num_back)) { // Crossing tile: overwrite this tile's arrival slots `{num_selected+state.running, ...}` with the // index-ordered winners (identical slot set, different candidate->slot mapping). emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); } - state.running = state.cand_prefix + placed; + state.running = reached; if (state.running >= static_cast(state.num_back)) { state.is_tie_active = false; } - // Trailing barrier for the lazy-scan path only: separate this tile's `placed` read (B1 above) from the next + // Trailing barrier for the lazy-scan path only: separate this tile's counter read (B1 above) from the next // tile's `back_local_cnt` atomic. The other paths write disjoint slots and need no additional explicit // barrier. __syncthreads(); @@ -1799,6 +1801,15 @@ private: nondet_filter_state state{ identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, resident_keys}; + // Prime the placement counters with this CTA's absolute region bases so `place_one` returns final output slots with + // no per-key offset; the barrier publishes both writes before the first placement. + if (tid == 0) + { + temp_storage.front_local_cnt = sel_prefix; + temp_storage.back_local_cnt = static_cast(k - num_kth) + cand_prefix; + } + __syncthreads(); + run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. [&](int stage, offset_t overflow_idx) { @@ -1826,8 +1837,9 @@ private: } // Deterministic final filter: both regions fill forward -- strictly-selected keys into the front `[0, num_selected)` - // via a SMEM atomic (offset by `sel_prefix`), candidates into the back `[num_selected, k)`. Candidate placement uses - // arrival-order atomics, with an index-ordered BlockScan only on the single boundary-crossing (straddling) CTA. A + // via a SMEM atomic (primed with this CTA's `sel_prefix`), candidates into the back `[num_selected, k)`. Candidate + // placement uses arrival-order atomics, with an index-ordered BlockScan only on the single boundary-crossing + // (straddling) CTA. A // combined cross-CTA scan gives this block its disjoint front/back bases and lets it detect whether all/none/some of // its candidates win. This member computes the `det_filter_state` inputs and hands them to `run_filter`, which drives // the per-region sweeps. @@ -1880,6 +1892,10 @@ private: (cluster_rank == layout.leader_rank) ? static_cast(num_selected - static_cast(sel_prefix)) : static_cast(push_front); + // This CTA's front slots `[sel_prefix, sel_prefix + my_front)` (the range its primed `front_local_cnt` walks) must + // stay within the selected region `[0, num_selected)`. + _CCCL_ASSERT(sel_prefix + static_cast(my_front) <= static_cast(num_selected), + "this CTA's front slots must fit within the selected region"); // Resident-front extent (bulk path): the whole contiguous resident span. The unaligned tail suffix (the // globally-last chunk's) is always peeled into `edge_keys` and folded by `process_tail_edge`, so it is never @@ -1931,6 +1947,15 @@ private: is_tail_edge_terminal, cand_prefix, !is_select_no_cand_cta}; + + // Prime the placement counters with this CTA's absolute region bases so `place_one` returns final output slots with + // no per-key offset; the barrier publishes both writes before `run_filter`'s first placement. + if (tid == 0) + { + temp_storage.front_local_cnt = sel_prefix; + temp_storage.back_local_cnt = static_cast(num_selected) + cand_prefix; + } + __syncthreads(); run_filter(state); } @@ -2707,13 +2732,11 @@ private: temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; temp_storage.state.result_pair = 0; - // Front-load every counter the final filter relies on so the initial cluster barrier below (arrive here, wait in - // the first pass) publishes the zeros to all ranks (the combined scan's DSMEM pushes then add into an - // already-zeroed `prefix_pair`, needing only a post-push barrier). `my_candidates` is zeroed too so idle/leader - // ranks (which never write it) read 0. + // Front-load the cross-pass accumulators so the initial cluster barrier below (arrive here, wait in the first + // pass) publishes the zeros to all ranks: the combined scan's DSMEM pushes add into an already-zeroed + // `prefix_pair` (needing only a post-push barrier), and idle/leader ranks that never write `my_candidates` read + // 0. (The placement counters are primed with region bases in the final-filter drivers, not here.) temp_storage.prefix_pair = 0; - temp_storage.front_local_cnt = 0; - temp_storage.back_local_cnt = 0; temp_storage.num_strictly_selected = 0; temp_storage.my_candidates = 0; } From 22d9ee9b13655dfaa0480edcf03caebd82103220 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 06:01:49 +0200 Subject: [PATCH 105/126] Split scan atomics since 64b ones seem to spin --- cub/cub/agent/agent_batched_topk_cluster.cuh | 208 +++++++++++-------- 1 file changed, 121 insertions(+), 87 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 53adb5fceba..e764bc43661 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -26,9 +26,10 @@ //! folds the bucket into its own local splitter key. //! //! The final filter places each block's output through per-CTA shared-memory -//! atomics, seeded by a single combined 64-bit cross-CTA prefix scan (each -//! block's selected-front and candidate-back base offsets); no cluster-wide -//! output cursor is kept in `state`. +//! atomics, seeded by a cross-CTA prefix scan of two independent 32-bit DSMEM +//! reductions (each block's selected-front and candidate-back base offsets) +//! fused directly into the placement counters; no cluster-wide output cursor is +//! kept in `state`. #pragma once @@ -225,8 +226,8 @@ struct agent_batched_topk_cluster // 32-bit covers every supported segment: the public entry caps the statically-known maximum segment size at 2^21, so // a runtime value exceeding its declared bound is a caller precondition violation (undefined behavior). Unsigned - // because all offsets, ranks, and block counts are non-negative (segment sizes are clamped to >= 0 upstream). The - // cross-CTA scan also packs two lanes into one `uint64_t`, which needs 32-bit lanes. + // because all offsets, ranks, and block counts are non-negative (segment sizes are clamped to >= 0 upstream). + // `prime_placement_counters` also returns its two prefix lanes packed into one `uint64_t`, which needs 32-bit lanes. using offset_t = ::cuda::std::uint32_t; using out_offset_t = ::cuda::std::uint32_t; using state_t = cluster_topk_state; @@ -238,7 +239,7 @@ struct agent_batched_topk_cluster static constexpr bool is_tie_reversed = TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; - // Push direction of the combined cross-CTA prefix scan (`combined_prefix_scan`). The leader must be *last* in scan + // Push direction of the cross-CTA prefix scan (`prime_placement_counters`). The leader must be *last* in scan // order so it derives its own (merged-away) counts from the predecessor sum: deterministic-prefer-smaller puts the // leader at the last effective rank and scans ascending; every other config (deterministic-prefer-larger and the // whole non-deterministic path) keeps the leader at rank 0 and scans descending, which makes rank 0 last. Matches the @@ -383,16 +384,16 @@ struct agent_batched_topk_cluster // sync the non-leader blocks fold their bucket counts into the leader's // `hist` through DSMEM atomics. `state` is meaningful only in the leader // block; the other blocks reach it exclusively through the DSMEM mapping. - // `prefix_pair` packs this block's exclusive cross-CTA scan result (high = `sel_prefix`, low = `cand_prefix`); peers - // add into it through DSMEM (`add_remote_prefix`), so it must sit at an identical offset in every block's storage. - // The remaining scalars are block-local: `front_local_cnt`/`back_local_cnt` hand out output slots via SMEM atomics, - // `num_strictly_selected` accumulates this block's strictly-selected count across passes, and `my_candidates` holds - // the last pass's splitter-bucket count. + // `front_local_cnt`/`back_local_cnt` are the final-filter output-slot counters, but they first serve as the cross-CTA + // scan accumulators: peers add their selected/candidate counts into them through DSMEM (`add_remote_prefix`) while + // this block seeds its own back base locally, leaving each counter primed with this block's absolute region base + // (front = `sel_prefix`, back = `num_selected + cand_prefix`). Because peers reach them over DSMEM they must sit at + // an identical offset in every block's storage. `num_strictly_selected` accumulates this block's strictly-selected + // count across passes, and `my_candidates` holds the last pass's splitter-bucket count. struct _TempStorage { offset_t hist[num_buckets]; state_t state; - ::cuda::std::uint64_t prefix_pair; offset_t front_local_cnt; offset_t back_local_cnt; offset_t num_strictly_selected; @@ -407,9 +408,6 @@ struct agent_batched_topk_cluster // the final filter. Block-local (never reached through DSMEM). key_t edge_keys[2 * load_align_items]; }; - // The `red.add.u64` on `prefix_pair`'s `.shared::cluster` address needs 8-byte alignment; the `uint64_t` member is - // already at an 8-aligned struct offset, so guarding the struct's alignment covers the absolute (and peer) address. - static_assert(alignof(_TempStorage) >= 8, "prefix_pair must be 8-byte aligned for the u64 DSMEM atomic"); // Split point of `edge_keys`: head edge in `[0, head_edge_cap_items)`, tail edge in `[head_edge_cap_items, 2 * // head_edge_cap_items)`. static constexpr int head_edge_cap_items = load_align_items; @@ -773,15 +771,24 @@ private: return static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(temp_storage.hist)); } - // Increment this block's own histogram bucket by one, at cluster scope so it stays mutually atomic with the DSMEM - // folds peers push into the leader's `hist` (Step 2, `hist_fold_remote`). Cluster and cta scope lower to identical - // shared-atomic SASS here, so applying it unconditionally (non-leaders, lone CTA) costs nothing and drops the leader - // branch. + // Cluster-scope `red.add` of `v` into this CTA's own shared `u32` at 32-bit `.shared::cta` address `addr` (e.g. from + // `__cvta_generic_to_shared`). Cluster scope -- not the address space -- is what keeps it mutually atomic with peers' + // DSMEM `red.add`s into the same location (`hist_fold_remote`, `add_remote_prefix`), so this CTA's own address needs + // no `mapa`; it lowers to the same shared-atomic SASS as cta scope. `hist_inc` passes a compile-time `1`, which + // current ptxas still folds into the warp-aggregated `ATOMS.POPC.INC.32` (verified in SASS) despite the register + // operand. + _CCCL_DEVICE _CCCL_FORCEINLINE static void add_local_shared_cluster(::cuda::std::uint32_t addr, offset_t v) + { + asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], %1;" : : "r"(addr), "r"(v) : "memory"); + } + + // Increment this block's own histogram bucket by one. Applied unconditionally (non-leaders and the lone CTA + // included): the cluster-scope add costs nothing there and it drops the leader branch. _CCCL_DEVICE _CCCL_FORCEINLINE void hist_inc(::cuda::std::uint32_t base32, int bucket) { _CCCL_ASSERT(bucket >= 0 && bucket < num_buckets, "histogram bucket index out of range"); const ::cuda::std::uint32_t addr = base32 + static_cast<::cuda::std::uint32_t>(bucket) * sizeof(offset_t); - asm volatile("red.relaxed.cluster.shared::cta.add.u32 [%0], 1;" : : "r"(addr) : "memory"); + add_local_shared_cluster(addr, offset_t{1}); } // Step 2: a non-leader folds one bucket into the leader's histogram through DSMEM. `own_bucket_addr32` is the @@ -806,16 +813,37 @@ private: return reinterpret_cast(remote); } - // Adds the packed 64-bit `v` to the `prefix_pair` of the CTA at cluster rank `target_rank` through DSMEM (mirrors - // `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add`). Exact because the two 32-bit lanes - // never carry into each other. Drives the combined cross-CTA selected/candidate prefix scan. - _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned int target_rank, ::cuda::std::uint64_t v) + // Adds `push_front`/`push_cand` into the front/back placement counters of the CTA at cluster rank `target_rank` + // through DSMEM (mirrors `hist_fold_remote`: `mapa` to `target_rank`, then a cluster-scope `red.add` per counter). + // Two independent 32-bit reductions rather than one 64-bit one: `red.add.u64` on shared memory is emulated with a + // CAS spin loop, whereas `red.add.u32` is a native shared-memory atomic. Drives the cross-CTA selected/candidate + // prefix scan (see `prime_placement_counters`). + _CCCL_DEVICE _CCCL_FORCEINLINE void + add_remote_prefix(unsigned int target_rank, offset_t push_front, offset_t push_cand) + { + const ::cuda::std::uint32_t own_front = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.front_local_cnt)); + const ::cuda::std::uint32_t own_back = + static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)); + ::cuda::std::uint32_t remote_front, remote_back; + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote_front) : "r"(own_front), "r"(target_rank)); + asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote_back) : "r"(own_back), "r"(target_rank)); + asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" + : + : "r"(remote_front), "r"(push_front) + : "memory"); + asm volatile("red.relaxed.cluster.shared::cluster.add.u32 [%0], %1;" + : + : "r"(remote_back), "r"(push_cand) + : "memory"); + } + + // Folds the back region base (`num_selected`) into this CTA's own `back_local_cnt`, in parallel with the peers' + // cluster-scope candidate-prefix pushes into the same counter (`add_remote_prefix`). + _CCCL_DEVICE _CCCL_FORCEINLINE void seed_local_back(offset_t v) { - const ::cuda::std::uint32_t own = - static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.prefix_pair)); - ::cuda::std::uint32_t remote; - asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own), "r"(target_rank)); - asm volatile("red.relaxed.cluster.shared::cluster.add.u64 [%0], %1;" : : "r"(remote), "l"(v) : "memory"); + add_local_shared_cluster(static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.back_local_cnt)), + v); } // Parallel prefix sum (cub::BlockScan) over the leader's merged histogram @@ -1173,24 +1201,37 @@ private: } } - // Combined 64-bit exclusive cross-CTA prefix scan over each working CTA's packed `(front_count << 32) | cand_count`. - // Each CTA pushes its counts into every successor's `prefix_pair` in `is_scan_descending` order; the leader is last - // and pushes to nobody, so it ends up holding the full predecessor sum. `prefix_pair` is pre-zeroed in `process_impl` - // and untouched until here, so a single post-push barrier suffices. Idle ranks and the leader pass `packed == 0`. - // Returns this CTA's exclusive prefix packed the same way; `is_single_cta` returns 0. + // Exclusive cross-CTA prefix scan fused with priming the final-filter placement counters. Each working CTA pushes its + // `push_front`/`push_cand` counts into every successor's front/back counter in `is_scan_descending` order (the leader + // is last and pushes to nobody, so it holds the full predecessor sum) and folds its own back-region base + // `num_selected` into its own back counter. All are commutative `red.add`s into the counters zeroed in `process_impl` + // (so a single post-push barrier suffices), leaving `front_local_cnt = sel_prefix` and + // `back_local_cnt = num_selected + cand_prefix` -- the absolute output-slot bases `place_one` expects. // - // The successor pushes are lane-parallel: the `red.add` reductions are commutative and target distinct remote ranks, - // so each thread owns a strided slice of the successor range (one push per thread for the usual small cluster). All - // threads see the same CTA-uniform `packed`, so the guard and the post-push barrier stay uniform. - _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t combined_prefix_scan(::cuda::std::uint64_t packed) + // Returns this CTA's exclusive prefix packed as `(sel_prefix << 32) | cand_prefix` for the driver's region math; + // `is_single_cta` yields 0 (front stays 0, back is just `num_selected`). The successor pushes are lane-parallel (each + // thread owns a strided slice); all threads see CTA-uniform counts, so the guard and the barrier stay uniform. + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t + prime_placement_counters(offset_t push_front, offset_t push_cand, offset_t num_selected) { if (is_single_cta) { + // No peers: the front base is 0 (untouched since init) and the back base is just `num_selected`. + if (tid == 0) + { + temp_storage.back_local_cnt = num_selected; + } + __syncthreads(); return ::cuda::std::uint64_t{0}; } - if (packed != ::cuda::std::uint64_t{0}) + // Fold this CTA's back base in parallel with the peers pushing their candidate counts into the same counter. + if (tid == 0) + { + seed_local_back(num_selected); + } + if (push_front != offset_t{0} || push_cand != offset_t{0}) { - // Only working (non-leader) ranks reach here with a nonzero count (see above); idle ranks / the leader push 0. + // Only working (non-leader) ranks reach here with a nonzero count; idle ranks / the leader push 0. _CCCL_ASSERT(cluster_rank < layout.eff_cluster_blocks, "a nonzero prefix count must come from a working rank"); if constexpr (is_scan_descending) { @@ -1198,7 +1239,7 @@ private: for (unsigned int rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; // leader last { - add_remote_prefix(rank, packed); + add_remote_prefix(rank, push_front, push_cand); } } else @@ -1209,14 +1250,22 @@ private: for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < layout.eff_cluster_blocks; rank += threads_per_block) { - add_remote_prefix(rank, packed); + add_remote_prefix(rank, push_front, push_cand); } } } // TODO(cccl): idle ranks arrive here only to keep this barrier reachable; a sub-cluster mbarrier over the working // ranks would let them exit (see the pass loop). cluster_or_block_sync(is_single_cta); - return temp_storage.prefix_pair; + // The local seed guarantees the back counter carries at least `num_selected` (postcondition of the fused prime). + _CCCL_ASSERT(temp_storage.back_local_cnt >= num_selected, + "back counter must include the seeded region base after the scan"); + const offset_t sel_prefix = temp_storage.front_local_cnt; + const offset_t cand_prefix = temp_storage.back_local_cnt - num_selected; + // Every thread must finish snapshotting the primed bases before any lane's `place_one` mutates the same counters: + // the leading boundary edge in each driver has no barrier of its own before its first placement atomic. + __syncthreads(); + return (static_cast<::cuda::std::uint64_t>(sel_prefix) << 32) | static_cast<::cuda::std::uint64_t>(cand_prefix); } // Deterministic final-filter state: the `run()`-local tie-break values (counts, prefixes, region extents) plus the @@ -1280,10 +1329,10 @@ private: static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front // Shared per-key arrival placement core (called by `place_tile` for both final filters): route one key by class to - // the front (selected) or back (candidate) counter, each primed with its region base by the final-filter drivers so - // the SMEM atomic returns the absolute output slot with no per-key offset. The uniform `out < k` guard drops losing - // candidates but always accepts selected keys; pairs builds also copy the key's value from `seg_idx`. (Deterministic - // index-ordered tie resolution goes through `emit_indexed` instead.) + // the front (selected) or back (candidate) counter, each primed with its region base by `prime_placement_counters` + // so the SMEM atomic returns the absolute output slot with no per-key offset. The uniform `out < k` guard drops + // losing candidates but always accepts selected keys; pairs builds also copy the key's value from `seg_idx`. + // (Deterministic index-ordered tie resolution goes through `emit_indexed` instead.) template _CCCL_DEVICE _CCCL_FORCEINLINE void place_one(KeyOutIt block_keys_out, bool is_cand, const key_t& key, offset_t seg_idx) @@ -1778,7 +1827,7 @@ private: } } - // Non-deterministic final filter driver. The combined cross-CTA scan gives this block disjoint front/back bases + // Non-deterministic final filter driver. `prime_placement_counters` gives this block disjoint front/back bases // (`sel_prefix`/`cand_prefix`); overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and // place into block-local SMEM atomics. `should_continue` is always true: every key must be classified. template @@ -1788,9 +1837,10 @@ private: const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; - const ::cuda::std::uint64_t packed = - (static_cast<::cuda::std::uint64_t>(my_sel) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); + // The scan doubles as counter priming: it leaves this CTA's placement counters holding its absolute region bases + // (front = `sel_prefix`, back = `num_selected + cand_prefix`), so no explicit priming follows. + const offset_t num_selected = static_cast(k) - static_cast(num_kth); + const ::cuda::std::uint64_t packed_prefix = prime_placement_counters(my_sel, my_cand, num_selected); const offset_t sel_prefix = static_cast(packed_prefix >> 32); const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. @@ -1801,15 +1851,6 @@ private: nondet_filter_state state{ identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, resident_keys}; - // Prime the placement counters with this CTA's absolute region bases so `place_one` returns final output slots with - // no per-key offset; the barrier publishes both writes before the first placement. - if (tid == 0) - { - temp_storage.front_local_cnt = sel_prefix; - temp_storage.back_local_cnt = static_cast(k - num_kth) + cand_prefix; - } - __syncthreads(); - run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. [&](int stage, offset_t overflow_idx) { @@ -1839,10 +1880,9 @@ private: // Deterministic final filter: both regions fill forward -- strictly-selected keys into the front `[0, num_selected)` // via a SMEM atomic (primed with this CTA's `sel_prefix`), candidates into the back `[num_selected, k)`. Candidate // placement uses arrival-order atomics, with an index-ordered BlockScan only on the single boundary-crossing - // (straddling) CTA. A - // combined cross-CTA scan gives this block its disjoint front/back bases and lets it detect whether all/none/some of - // its candidates win. This member computes the `det_filter_state` inputs and hands them to `run_filter`, which drives - // the per-region sweeps. + // (straddling) CTA. `prime_placement_counters` gives this block its disjoint front/back bases and lets it detect + // whether all/none/some of its candidates win. This member computes the `det_filter_state` inputs and hands them to + // `run_filter`, which drives the per-region sweeps. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) @@ -1865,11 +1905,12 @@ private: // nothing folds into the front here. The leader and idle ranks push 0 -- the leader because its merged histogram // cannot self-count (it derives its own front from the total below), idle ranks because they own nothing. const offset_t push_front = my_sel; - const ::cuda::std::uint64_t packed = - (static_cast<::cuda::std::uint64_t>(push_front) << 32) | static_cast<::cuda::std::uint64_t>(my_cand); - const ::cuda::std::uint64_t packed_prefix = combined_prefix_scan(packed); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + // The scan doubles as counter priming: it leaves this CTA's placement counters holding its absolute region bases + // (front = `sel_prefix`, back = `num_selected + cand_prefix`), so no explicit priming follows. + const ::cuda::std::uint64_t packed_prefix = + prime_placement_counters(push_front, my_cand, static_cast(num_selected)); + const offset_t sel_prefix = static_cast(packed_prefix >> 32); + const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); // Guards the leader's remainder subtractions below (`total_candidates - cand_prefix`, `num_selected - sel_prefix`). _CCCL_ASSERT(cand_prefix <= total_candidates && sel_prefix <= static_cast(num_selected), "cross-CTA prefixes must stay within their candidate and selected totals"); @@ -1947,15 +1988,6 @@ private: is_tail_edge_terminal, cand_prefix, !is_select_no_cand_cta}; - - // Prime the placement counters with this CTA's absolute region bases so `place_one` returns final output slots with - // no per-key offset; the barrier publishes both writes before `run_filter`'s first placement. - if (tid == 0) - { - temp_storage.front_local_cnt = sel_prefix; - temp_storage.back_local_cnt = static_cast(num_selected) + cand_prefix; - } - __syncthreads(); run_filter(state); } @@ -2537,9 +2569,10 @@ private: } // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into - // gmem, so the last cross-CTA DSMEM access is the combined scan's `prefix_pair` push, already fenced by its - // post-push cluster barrier (and `early_stop` is cached pre-scan). With no shared-memory access to another block - // after the scan, a block can return without risking a "cluster target block not present" fault from a straggler. + // gmem, so the last cross-CTA DSMEM access is the scan's counter push in `prime_placement_counters`, already fenced + // by its post-push cluster barrier (and `early_stop` is cached pre-scan). With no shared-memory access to another + // block after the scan, a block can return without risking a "cluster target block not present" fault from a + // straggler. } // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). Runs before @@ -2732,11 +2765,12 @@ private: temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; temp_storage.state.result_pair = 0; - // Front-load the cross-pass accumulators so the initial cluster barrier below (arrive here, wait in the first - // pass) publishes the zeros to all ranks: the combined scan's DSMEM pushes add into an already-zeroed - // `prefix_pair` (needing only a post-push barrier), and idle/leader ranks that never write `my_candidates` read - // 0. (The placement counters are primed with region bases in the final-filter drivers, not here.) - temp_storage.prefix_pair = 0; + // Front-load the scan accumulators so the initial cluster barrier below (arrive here, wait in the first pass) + // publishes the zeros to all ranks: the final filter's `prime_placement_counters` then adds into already-zeroed + // `front_local_cnt`/`back_local_cnt` (its own local seed + peers' DSMEM pushes, needing only a post-push + // barrier). `my_candidates` is zeroed too so idle/leader ranks that never write it read 0. + temp_storage.front_local_cnt = 0; + temp_storage.back_local_cnt = 0; temp_storage.num_strictly_selected = 0; temp_storage.my_candidates = 0; } @@ -2744,7 +2778,7 @@ private: // Only arrive here; the matching wait is deferred to just before the first cross-CTA fold in `run_radix_passes`' // first pass, so the cluster-arrival latency overlaps the fused first-pass load + histogram. Safe because nothing // between here and that wait touches another rank's DSMEM (the first-pass histogram writes only block-local - // `hist`; `leader_state`/`prefix_pair` are untouched until later). + // `hist`; `leader_state` and the scan counters are untouched until later). cluster_or_block_arrive(is_single_cta); [[maybe_unused]] const bool is_ok = From dab74837c83bafc9e17872771a0d71acc78b84e1 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 17:35:35 +0200 Subject: [PATCH 106/126] Try to fix non-det perf regression from unification - load striped from smem as we do not need to scan - Bonus: Early exit from filter possible before loading more keys from gmem --- cub/cub/agent/agent_batched_topk_cluster.cuh | 147 ++++++++++++------ .../catch2_test_device_segmented_topk_keys.cu | 14 +- 2 files changed, 110 insertions(+), 51 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index e764bc43661..eca71d5579f 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -971,15 +971,16 @@ private: // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from - // gmem (fallback). `mid()` runs once on a full pass -- after the first reload wave (`stream_stages` visits) is issued - // but before it is waited on, overlapping the caller's resident-chunk work with those in-flight copies (skipped if - // phase 1 stops early); it must be block-uniform with no unmatched barrier. `should_continue()` is polled once - // after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final + // gmem (fallback). `mid()` runs at most once per pass, positioned to overlap the caller's resident-chunk work with + // in-flight copies: after the first reload wave (`stream_stages` visits) is issued but before it is waited on + // (block-load path), or before the gmem loop (fallback). A phase-1 early stop on the block-load path skips it, since + // nothing is then left to place; it must be block-uniform with no unmatched barrier. `should_continue()` is polled + // once after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final // filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the - // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (the deterministic - // filter's does, to read its placement counters block-wide); the histogram and non-deterministic filter pass an - // always-true predicate. An early break can leave prefetches in flight, so the pass drains the - // remaining stages before returning (a full pass ends with an empty `stream_inflight_mask`, so the drain is a no-op). + // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (both final filters' + // do, to read their placement counters block-wide); the histogram passes an always-true predicate. An early break + // can leave prefetches in flight, so the pass drains the remaining stages before returning (a full pass ends with an + // empty `stream_inflight_mask`, so the drain is a no-op). template _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) @@ -1057,6 +1058,8 @@ private: if (!is_stopped) { // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. + // Skipped on an early stop: the stream only breaks once this CTA's whole contribution is placed, so no resident + // key can still be required (the non-deterministic filter folds its resident keys as `mid`). mid(); // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). @@ -1328,6 +1331,25 @@ private: static constexpr offset_t flag_candidate = 1; // tie candidate: routed to the back, counted by the boundary scan static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front + // Map thread `tid`'s item slot `i` to its position within the current tile. `Blocked` is mandatory for the + // deterministic filter: `emit_indexed`'s `BlockScan` ranks candidates in blocked order, which is what makes the tie + // ranks segment-index-ordered. The order-independent non-deterministic filter picks striped so consecutive threads + // hit consecutive addresses (no SMEM bank conflicts on the staged reads, coalesced loads in the gmem fallback). + // `process_tiles` and `place_tile` must use the same `Blocked` so `keys[i]`/`flags[i]` denote the same element; + // `emit_indexed` (deterministic-only) is intrinsically blocked. + template + _CCCL_DEVICE _CCCL_FORCEINLINE int tile_item_pos(int tile_base, int i) const + { + if constexpr (Blocked) + { + return tile_base + tid * ItemsPerThread + i; + } + else + { + return tile_base + i * threads_per_block + tid; + } + } + // Shared per-key arrival placement core (called by `place_tile` for both final filters): route one key by class to // the front (selected) or back (candidate) counter, each primed with its region base by `prime_placement_counters` // so the SMEM atomic returns the absolute output slot with no per-key offset. The uniform `out < k` guard drops @@ -1374,7 +1396,7 @@ private: // the lazy crossing tile is overwritten in index order later. `do_arrival == false` (terminal tile only) skips // candidates here and leaves them to `emit_indexed`; the lazy path passes `true` and its crossing tile is later // overwritten by `emit_indexed`. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void place_tile( State& state, const key_t (&keys)[ItemsPerThread], @@ -1392,7 +1414,7 @@ private: { continue; } - const int pos = tile_base + tid * ItemsPerThread + i; + const int pos = tile_item_pos(tile_base, i); const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); place_one(state.block_keys_out, is_cand, keys[i], seg_idx); } @@ -1431,7 +1453,8 @@ private: const offset_t global_rank = base + excl[i]; if (global_rank < static_cast(state.num_back)) { - const int pos = tile_base + tid * ItemsPerThread + i; + // `emit_indexed` is deterministic-only, and its `BlockScan` ranks assume a blocked arrangement. + const int pos = tile_item_pos(tile_base, i); const out_offset_t out = static_cast(num_selected + global_rank); const offset_t seg_idx = seg_base + static_cast(Reversed ? (count - 1 - pos) : pos); state.block_keys_out[out] = keys[i]; @@ -1442,17 +1465,17 @@ private: return tile_total; } - // Classify and place each of `count` keys in scan order, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` + // Classify and place each of `count` keys, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` // picks the key source: `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` (indexed by the - // segment-local index `seg_base + folded`); `Reversed` walks the span high-to-low. Both final filters share this via - // `Deterministic`: - // * `true` (a `det_filter_state`): resolves the boundary ties in segment-index order via `emit_indexed`, so callers - // walk the regions in global-index order and carry `state.running`/`state.is_tie_active` across tiles; the - // lazy-scan `else` branch holds the only explicit per-tile `__syncthreads()`; `region_is_terminal` marks the + // segment-local index `seg_base + folded`); `Reversed` walks the span high-to-low. `Blocked` selects the intra-tile + // thread arrangement (see `tile_item_pos`). Both final filters share this via `Deterministic`: + // * `true` (a `det_filter_state`, blocked): resolves the boundary ties in segment-index order via `emit_indexed`, + // so callers walk the regions in global-index order and carry `state.running`/`state.is_tie_active` across tiles; + // the lazy-scan `else` branch holds the only explicit per-tile `__syncthreads()`; `region_is_terminal` marks the // CTA's last. - // * `false` (a `nondet_filter_state`): all arrival-order placement (`place_tile` with `do_arrival == true`); no - // scan/tie-state/barrier, `Reversed`/`region_is_terminal` unused. - template + // * `false` (a `nondet_filter_state`, striped): all arrival-order placement (`place_tile` with `do_arrival == + // true`); no scan/tie-state/barrier, `Reversed`/`region_is_terminal` unused. + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles( State& state, [[maybe_unused]] const key_t* smem_src, @@ -1460,6 +1483,9 @@ private: int count, [[maybe_unused]] bool region_is_terminal) { + // The deterministic filter's index-ordered `emit_indexed` scan only produces segment-index-ordered ranks in a + // blocked arrangement; only the order-independent non-deterministic filter may pick striped. + static_assert(!Deterministic || Blocked, "the deterministic filter requires a blocked arrangement"); // The SMEM-source clause is exempt when reading from gmem (`FromSmem == false`, the generic overflow fallback // passes `nullptr`) or for an empty tile. _CCCL_ASSERT(count >= 0 && seg_base <= layout.segment_size_off @@ -1474,7 +1500,7 @@ private: _CCCL_PRAGMA_UNROLL_FULL() for (int i = 0; i < ItemsPerThread; ++i) { - const int pos = tile_base + tid * ItemsPerThread + i; + const int pos = tile_item_pos(tile_base, i); // Classify each in-range key once into `flags[i]` (see `flag_*`); `place_tile`/`emit_indexed` read it back // without re-running `identify_op`. Out-of-range lanes stay `flag_none` and are skipped everywhere. flags[i] = flag_none; @@ -1503,7 +1529,7 @@ private: { // Non-deterministic: place every tile in arrival order (selected to the front, the first `num_back` candidates // to the back via `place_one`'s `out < k` guard). No index-ordered scan, tie-state, or per-tile barrier. - place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/true); + place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/true); } else { @@ -1517,7 +1543,7 @@ private: // A select-all CTA never scans (every candidate wins), so it places entirely via `place_tile`'s arrival route. const bool is_terminal_tile = state.is_tie_active && region_is_terminal && (tile_base + tile_size >= count); const bool do_arrival = state.is_tie_active && (state.is_select_all_cand_cta || !is_terminal_tile); - place_tile(state, keys, flags, seg_base, count, tile_base, do_arrival); + place_tile(state, keys, flags, seg_base, count, tile_base, do_arrival); if (state.is_tie_active && !state.is_select_all_cand_cta) { @@ -1571,7 +1597,8 @@ private: process_tiles( + /*Deterministic=*/true, + /*Blocked=*/true>( state, state.resident_front.data(), state.front_seg_base, state.front_count, state.is_resident_terminal); } else @@ -1587,7 +1614,8 @@ private: process_tiles( + /*Deterministic=*/true, + /*Blocked=*/true>( state, ::cuda::ptr_rebind(key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); } } @@ -1625,8 +1653,11 @@ private: // The multi-chunk overflow stream stays on the lazy per-tile boundary detection (`region_is_terminal == // false`): a stray terminal direct scan here would need the stream to flag its last chunk, and the saving // is one barrier on one tile. - process_tiles( - state, keys.data(), base_off, static_cast(keys.size()), false); + process_tiles(state, keys.data(), base_off, static_cast(keys.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). [&](const auto& chunk) { @@ -1634,7 +1665,8 @@ private: process_tiles(state, nullptr, chunk.offset, chunk.count, false); + /*Deterministic=*/true, + /*Blocked=*/true>(state, nullptr, chunk.offset, chunk.count, false); }, // No interleaved resident work: the deterministic filter folds its resident span separately. [] {}, @@ -1657,7 +1689,7 @@ private: if constexpr (use_block_load_to_shared) { _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true>( + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true, /*Blocked=*/true>( state, keys, seg_base, count, is_terminal); } } @@ -1756,7 +1788,7 @@ private: // Order-independent counterpart of the deterministic filter: every region is swept by // `process_tiles<..., Deterministic=false>` (arrival-order placement, no scan/tie-state/barrier). Since order is // irrelevant, `write_nondeterministic_topk` folds the resident keys as the overflow stream's `mid` work (overlapping - // the first reloads) and never bails early. + // the first reloads) and bails out of the overflow stream once its contribution is fully placed. template struct nondet_filter_state { @@ -1765,6 +1797,7 @@ private: out_offset_t num_back; offset_t sel_prefix; offset_t cand_prefix; + offset_t my_front; // this CTA's front region size (see `write_deterministic_topk`'s `my_front`) smem_keys_t resident_keys; }; @@ -1788,8 +1821,11 @@ private: layout.part.global_index(layout.resident_base + local_chunk), layout.segment_size_off, layout.head_items); const int bulk_count_items = static_cast( ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items))); - process_tiles( - state, resident_ptr + cursor_items, chunk.offset, bulk_count_items, false); + process_tiles(state, resident_ptr + cursor_items, chunk.offset, bulk_count_items, false); cursor_items += bulk_count_items; } // Persistent boundary edges (unroll factor 1: an edge spans fewer than `load_align_items` items): head prefix at @@ -1797,12 +1833,12 @@ private: // `segment_size - tail_edge_len_items`). if (layout.head_edge_len_items > 0) { - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false>( + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false, /*Blocked=*/false>( state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, false); } if (layout.tail_edge_len_items > 0) { - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false>( + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/false, /*Deterministic=*/false, /*Blocked=*/false>( state, temp_storage.edge_keys + head_edge_cap_items, layout.segment_size_off - static_cast(layout.tail_edge_len_items), @@ -1817,7 +1853,11 @@ private: for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const auto chunk = get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); - process_tiles( + process_tiles( state, ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), chunk.offset, @@ -1829,7 +1869,8 @@ private: // Non-deterministic final filter driver. `prime_placement_counters` gives this block disjoint front/back bases // (`sel_prefix`/`cand_prefix`); overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and - // place into block-local SMEM atomics. `should_continue` is always true: every key must be classified. + // place into block-local SMEM atomics. `run_pass` breaks the stream once this CTA's contribution is fully placed + // (all its selected keys in the front, and its back counter reached the region end `k`). template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) @@ -1847,9 +1888,13 @@ private: _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) <= static_cast<::cuda::std::uint64_t>(k), "selected prefix must fit before the tie-back output region"); + // This CTA's front region size (mirrors `write_deterministic_topk`): the leader is last in scan order and derives + // it from the total (`num_selected - sel_prefix`) since its merged histogram can't self-count; others place + // `my_sel`. + const offset_t my_front = (cluster_rank == layout.leader_rank) ? (num_selected - sel_prefix) : my_sel; nondet_filter_state state{ - identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, resident_keys}; + identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, my_front, resident_keys}; run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. @@ -1859,21 +1904,35 @@ private: get_chunk( layout.part.global_index(layout.overflow_base + overflow_idx), layout.segment_size_off, layout.head_items) .offset; - process_tiles( - state, keys.data(), base_off, static_cast(keys.size()), false); + process_tiles(state, keys.data(), base_off, static_cast(keys.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem. [&](const auto& chunk) { - process_tiles( - state, nullptr, chunk.offset, chunk.count, false); + process_tiles(state, nullptr, chunk.offset, chunk.count, false); }, // Fold the resident keys + boundary edges, overlapping the first overflow reloads. [&] { nondet_fold_resident(state); }, - // Never break early: every key must be classified (selected/candidate keys placed, the rest dropped in-tile). - [] { - return true; + // Break the stream once this CTA's whole contribution is placed: its front is full (`front_local_cnt` reached + // `sel_prefix + my_front`) and its back is full (`back_local_cnt >= k`, so further candidates fail `place_one`'s + // `out < k` guard). The barrier makes the counter reads block-wide, resyncing lanes that raced through the + // barrier-free tiles. The stop needs both conditions, so a CTA whose candidates never fill its back to `k` + // simply never stops early. + [&] { + __syncthreads(); + const offset_t front_end = state.sel_prefix + state.my_front; + _CCCL_ASSERT(temp_storage.front_local_cnt <= front_end, + "front counter must stay within this CTA's front region"); + return !(temp_storage.front_local_cnt >= front_end && temp_storage.back_local_cnt >= static_cast(k)); }); } diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 551d9873304..709eb4c7366 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -119,9 +119,9 @@ using max_num_k_list = c2h::enum_type_list; using key_types = c2h::type_list; // clang-format on #elif TEST_TYPES == 1 @@ -130,10 +130,10 @@ using key_types = c2h::type_list; using key_types = c2h::type_list; +#if TEST_BF_T() + , bfloat16_t +#endif // TEST_BF_T() +>; // clang-format on #endif From dadf4e9443c1348a6dad424fbb15bb6eb6d60d0a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 18:25:54 +0200 Subject: [PATCH 107/126] Try to use striped loads in the deterministic filter Only the straddling CTA needs blocked loads for the potential scan until it found the k-boundary and can switch to striped loads for the remaining strictly selected keys. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 288 +++++++++++-------- 1 file changed, 173 insertions(+), 115 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index eca71d5579f..d6352cd5ef2 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1331,12 +1331,13 @@ private: static constexpr offset_t flag_candidate = 1; // tie candidate: routed to the back, counted by the boundary scan static constexpr offset_t flag_selected = 2; // strictly selected: routed to the front - // Map thread `tid`'s item slot `i` to its position within the current tile. `Blocked` is mandatory for the - // deterministic filter: `emit_indexed`'s `BlockScan` ranks candidates in blocked order, which is what makes the tie - // ranks segment-index-ordered. The order-independent non-deterministic filter picks striped so consecutive threads - // hit consecutive addresses (no SMEM bank conflicts on the staged reads, coalesced loads in the gmem fallback). - // `process_tiles` and `place_tile` must use the same `Blocked` so `keys[i]`/`flags[i]` denote the same element; - // `emit_indexed` (deterministic-only) is intrinsically blocked. + // Map thread `tid`'s item slot `i` to its position within the current tile. Blocked is used by the straddling + // deterministic CTA while it resolves ties (`process_tiles` phase A): `emit_indexed`'s `BlockScan` ranks candidates + // in blocked order, which makes the tie ranks segment-index-ordered. Striped (consecutive threads hit consecutive + // addresses: no SMEM bank conflicts on the staged reads, coalesced loads in the gmem fallback) is used everywhere + // else -- the whole non-deterministic filter, every non-straddling deterministic CTA, and the straddling CTA's tiles + // past the boundary (phase B). `classify_tile` and `place_tile` must use the same `Blocked` so `keys[i]`/`flags[i]` + // denote the same element; `emit_indexed` is intrinsically blocked. template _CCCL_DEVICE _CCCL_FORCEINLINE int tile_item_pos(int tile_base, int i) const { @@ -1465,16 +1466,57 @@ private: return tile_total; } - // Classify and place each of `count` keys, tiled by `threads_per_block * ItemsPerThread`. `FromSmem` - // picks the key source: `smem_src` (indexed by the folded in-region position) or gmem `block_keys_in` (indexed by the - // segment-local index `seg_base + folded`); `Reversed` walks the span high-to-low. `Blocked` selects the intra-tile - // thread arrangement (see `tile_item_pos`). Both final filters share this via `Deterministic`: - // * `true` (a `det_filter_state`, blocked): resolves the boundary ties in segment-index order via `emit_indexed`, - // so callers walk the regions in global-index order and carry `state.running`/`state.is_tie_active` across tiles; - // the lazy-scan `else` branch holds the only explicit per-tile `__syncthreads()`; `region_is_terminal` marks the - // CTA's last. - // * `false` (a `nondet_filter_state`, striped): all arrival-order placement (`place_tile` with `do_arrival == - // true`); no scan/tie-state/barrier, `Reversed`/`region_is_terminal` unused. + // Load and 3-way classify one tile's items into `keys`/`flags` (see `flag_*`), re-run by `place_tile`/`emit_indexed` + // without touching `identify_op` again. `Blocked` picks the thread->element arrangement (`tile_item_pos`); the source + // is `smem_src` (folded in-region position) or gmem `block_keys_in`. Out-of-range lanes stay `flag_none`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void classify_tile( + State& state, + [[maybe_unused]] const key_t* smem_src, + offset_t seg_base, + int count, + int tile_base, + key_t (&keys)[ItemsPerThread], + offset_t (&flags)[ItemsPerThread]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + const int pos = tile_item_pos(tile_base, i); + flags[i] = flag_none; + if (pos < count) + { + // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` + // local read and the segment-local value index. + const int folded_pos = Reversed ? (count - 1 - pos) : pos; + if constexpr (FromSmem) + { + keys[i] = smem_src[folded_pos]; + } + else + { + keys[i] = layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; + } + const auto cls = state.identify_op(keys[i]); + flags[i] = (cls == detail::topk::candidate_class::candidate) + ? flag_candidate + : (cls == detail::topk::candidate_class::selected ? flag_selected : flag_none); + } + } + } + + // Classify (`classify_tile`) and place each of `count` keys, tiled by `threads_per_block * ItemsPerThread`. + // `FromSmem` picks the key source; `Reversed` walks the span high-to-low. Both final filters share this via + // `Deterministic`: + // * `false` (a `nondet_filter_state`, always striped): every tile placed in arrival order (`place_tile` with + // `do_arrival == true`); no scan/tie-state/barrier, `region_is_terminal` unused. + // * `true` (a `det_filter_state`): resolves the boundary ties in segment-index order via `emit_indexed`, so callers + // walk the regions in global-index order and carry `state.running`/`state.is_tie_active` across tiles; + // `region_is_terminal` marks the CTA's last. Only the boundary-straddling CTA resolves ties and is dispatched + // `Blocked`; other deterministic CTAs never scan and are dispatched striped (see `write_deterministic_topk`). + // `Blocked` is the *entry* arrangement (`tile_item_pos`). The blocked deterministic path is two-phase: it loads + // blocked only while ties remain (phase A, `emit_indexed`'s scan needs it) and switches once to striped for the + // remaining strictly-selected/rejected keys (phase B). template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tiles( State& state, @@ -1483,9 +1525,6 @@ private: int count, [[maybe_unused]] bool region_is_terminal) { - // The deterministic filter's index-ordered `emit_indexed` scan only produces segment-index-ordered ranks in a - // blocked arrangement; only the order-independent non-deterministic filter may pick striped. - static_assert(!Deterministic || Blocked, "the deterministic filter requires a blocked arrangement"); // The SMEM-source clause is exempt when reading from gmem (`FromSmem == false`, the generic overflow fallback // passes `nullptr`) or for an empty tile. _CCCL_ASSERT(count >= 0 && seg_base <= layout.segment_size_off @@ -1493,92 +1532,99 @@ private: && (!FromSmem || count == 0 || smem_src != nullptr), "process_tiles: region must be a valid in-segment source span"); constexpr int tile_size = threads_per_block * ItemsPerThread; - for (int tile_base = 0; tile_base < count; tile_base += tile_size) + int tile_base = 0; + + if constexpr (Deterministic && Blocked) { - key_t keys[ItemsPerThread]; - offset_t flags[ItemsPerThread]; - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < ItemsPerThread; ++i) + // Blocked is dispatched only for the boundary-straddling CTA, which always has ties and is never select-all. + _CCCL_ASSERT(!state.is_select_all_cand_cta, "blocked deterministic path is only for the straddling CTA"); + + // Phase A: while ties remain, load blocked (`emit_indexed`'s scan needs it) and resolve them. A region entered + // past the crossing (its ties already placed) has `is_tie_active == false`, so phase A is skipped and everything + // falls through to phase B. + for (; state.is_tie_active && tile_base < count; tile_base += tile_size) { - const int pos = tile_item_pos(tile_base, i); - // Classify each in-range key once into `flags[i]` (see `flag_*`); `place_tile`/`emit_indexed` read it back - // without re-running `identify_op`. Out-of-range lanes stay `flag_none` and are skipped everywhere. - flags[i] = flag_none; - if (pos < count) + key_t keys[ItemsPerThread]; + offset_t flags[ItemsPerThread]; + classify_tile( + state, smem_src, seg_base, count, tile_base, keys, flags); + + // The boundary tile places its candidates via `emit_indexed` (arrival placement skipped, `do_arrival == + // false`); every other tile uses `place_tile`'s arrival route: + // * terminal tile -- the CTA's last tile necessarily holds the boundary: scan it directly. + // * lazy crossing tile -- boundary not yet known: candidates land in arrival order first, then `B1` reads the + // counter and, on the tile that crosses `num_back`, overwrites those slots in index + // order. + const bool is_terminal_tile = region_is_terminal && (tile_base + tile_size >= count); + place_tile(state, keys, flags, seg_base, count, tile_base, !is_terminal_tile); + + if (is_terminal_tile) + { + state.running += emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); + state.is_tie_active = false; + } + else { - // In-region index: forward, or (`Reversed`) counting down from the last element. Shared by the `FromSmem` - // local read and the segment-local value index. - const int folded_pos = Reversed ? (count - 1 - pos) : pos; - if constexpr (FromSmem) + // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary + // slots) and make the counter read race-free and block-uniform. + __syncthreads(); + // `back_local_cnt` is primed with the back base (`num_selected + cand_prefix`), so subtracting the front + // region size recovers this CTA's candidate-rank reached so far (`cand_prefix` plus placed candidates). + const offset_t num_selected = static_cast(k) - static_cast(state.num_back); + _CCCL_ASSERT(temp_storage.back_local_cnt >= num_selected, + "back counter must stay at or above its primed base (no unsigned underflow)"); + const offset_t reached = temp_storage.back_local_cnt - num_selected; + if (reached > static_cast(state.num_back)) { - keys[i] = smem_src[folded_pos]; + // Crossing tile: overwrite this tile's arrival slots `{num_selected+state.running, ...}` with the + // index-ordered winners (identical slot set, different candidate->slot mapping). + emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); } - else + state.running = reached; + if (state.running >= static_cast(state.num_back)) { - keys[i] = - layout.block_keys_in[static_cast(seg_base + static_cast(folded_pos))]; + state.is_tie_active = false; } - const auto cls = state.identify_op(keys[i]); - flags[i] = (cls == detail::topk::candidate_class::candidate) - ? flag_candidate - : (cls == detail::topk::candidate_class::selected ? flag_selected : flag_none); + // Trailing barrier for the lazy-scan path only: separate this tile's counter read (B1 above) from the next + // tile's `back_local_cnt` atomic. Other paths write disjoint slots and need no explicit barrier. + __syncthreads(); } } - if constexpr (!Deterministic) + // Phase B: ties are placed, so any remaining candidate-class keys are losers (dropped via `do_arrival == false`); + // only strictly-selected keys are placed, to the front in arrival order. Load striped (coalesced/conflict-free); + // no scan or barrier -- `front_local_cnt` is order-independent, so the A->B switch needs none either. + for (; tile_base < count; tile_base += tile_size) { - // Non-deterministic: place every tile in arrival order (selected to the front, the first `num_back` candidates - // to the back via `place_one`'s `out < k` guard). No index-ordered scan, tie-state, or per-tile barrier. - place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/true); + key_t keys[ItemsPerThread]; + offset_t flags[ItemsPerThread]; + classify_tile( + state, smem_src, seg_base, count, tile_base, keys, flags); + place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/false); } - else + } + else + { + // Single striped pass: the non-deterministic filter (always) or a deterministic CTA that never scans. + for (; tile_base < count; tile_base += tile_size) { - // Placement. `place_tile` routes selected keys to the front on every tile and, on arrival tiles, candidates to - // the back in arrival order. Two block-uniform sub-paths need a block scan instead and run through - // `emit_indexed`: - // * terminal tile -- the straddling CTA's last tile necessarily holds the boundary: scan it directly. - // * lazy crossing tile -- boundary not yet known: candidates land in arrival order first, then `B1` reads the - // counter and, on the tile that crosses `num_back`, overwrites those slots in index - // order. - // A select-all CTA never scans (every candidate wins), so it places entirely via `place_tile`'s arrival route. - const bool is_terminal_tile = state.is_tie_active && region_is_terminal && (tile_base + tile_size >= count); - const bool do_arrival = state.is_tie_active && (state.is_select_all_cand_cta || !is_terminal_tile); - place_tile(state, keys, flags, seg_base, count, tile_base, do_arrival); - - if (state.is_tie_active && !state.is_select_all_cand_cta) + key_t keys[ItemsPerThread]; + offset_t flags[ItemsPerThread]; + classify_tile( + state, smem_src, seg_base, count, tile_base, keys, flags); + if constexpr (Deterministic) { - if (is_terminal_tile) - { - state.running += emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); - state.is_tie_active = false; - } - else - { - // B1: order the arrival global writes (from `place_tile`) ahead of the index-order overwrite (same boundary - // slots) and make the counter read race-free and block-uniform. - __syncthreads(); - // `back_local_cnt` is primed with the back base (`num_selected + cand_prefix`), so subtracting the front - // region size recovers this CTA's candidate-rank reached so far (`cand_prefix` plus the candidates placed). - const offset_t num_selected = static_cast(k) - static_cast(state.num_back); - _CCCL_ASSERT(temp_storage.back_local_cnt >= num_selected, - "back counter must stay at or above its primed base (no unsigned underflow)"); - const offset_t reached = temp_storage.back_local_cnt - num_selected; - if (reached > static_cast(state.num_back)) - { - // Crossing tile: overwrite this tile's arrival slots `{num_selected+state.running, ...}` with the - // index-ordered winners (identical slot set, different candidate->slot mapping). - emit_indexed(state, keys, flags, seg_base, count, tile_base, state.running); - } - state.running = reached; - if (state.running >= static_cast(state.num_back)) - { - state.is_tie_active = false; - } - // Trailing barrier for the lazy-scan path only: separate this tile's counter read (B1 above) from the next - // tile's `back_local_cnt` atomic. The other paths write disjoint slots and need no additional explicit - // barrier. - __syncthreads(); - } + // Non-straddling CTA: a select-all CTA places its (all-winning) candidates in arrival order; a select-no or + // already-resolved CTA has none. Never the index-ordered scan. + _CCCL_ASSERT(state.is_select_all_cand_cta || !state.is_tie_active, + "striped deterministic tile must never need index-ordered tie resolution"); + const bool do_arrival = state.is_tie_active && state.is_select_all_cand_cta; + place_tile(state, keys, flags, seg_base, count, tile_base, do_arrival); + } + else + { + // Selected to the front, the first `num_back` candidates to the back via `place_one`'s `out < k` guard. + place_tile(state, keys, flags, seg_base, count, tile_base, /*do_arrival=*/true); } } } @@ -1588,7 +1634,7 @@ private: // deterministic mode): ascending walks the low-index window forward, descending walks the high-index window // (`resident_base`) in reverse, so a single `process_tiles` call per span with the index folded at compile time // replaces the old fwd/rev pair. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_resident(det_filter_state& state) { if constexpr (use_block_load_to_shared) @@ -1598,7 +1644,7 @@ private: /*FromSmem=*/true, /*Reversed=*/is_residency_reversed, /*Deterministic=*/true, - /*Blocked=*/true>( + Blocked>( state, state.resident_front.data(), state.front_seg_base, state.front_count, state.is_resident_terminal); } else @@ -1615,7 +1661,7 @@ private: /*FromSmem=*/true, /*Reversed=*/is_residency_reversed, /*Deterministic=*/true, - /*Blocked=*/true>( + Blocked>( state, ::cuda::ptr_rebind(key_slots + local_slot * ChunkBytes), chunk.offset, chunk.count, false); } } @@ -1627,7 +1673,7 @@ private: // slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so // `stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the // generic fallback re-reads gmem each pass regardless). - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_overflow(det_filter_state& state) { // Guards the straddling CTA (the only rank needing scan order): `run` preselected the direction so it enters at @@ -1657,7 +1703,7 @@ private: /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true, - /*Blocked=*/true>(state, keys.data(), base_off, static_cast(keys.size()), false); + Blocked>(state, keys.data(), base_off, static_cast(keys.size()), false); }, // Generic fallback: read the overflow chunk straight from gmem (full count; the fallback never peels a tail). [&](const auto& chunk) { @@ -1666,7 +1712,7 @@ private: /*FromSmem=*/false, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true, - /*Blocked=*/true>(state, nullptr, chunk.offset, chunk.count, false); + Blocked>(state, nullptr, chunk.offset, chunk.count, false); }, // No interleaved resident work: the deterministic filter folds its resident span separately. [] {}, @@ -1682,14 +1728,14 @@ private: // generic fallback (which reads boundary items straight from gmem) and for a zero-length edge. An edge spans fewer // than `load_align_items` items (it fits in its `edge_keys` slot), so it runs `process_tiles` at unroll factor 1 // (non-unrolled) instead of the wider tie-break unroll the potentially large resident/overflow regions use. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_edge( det_filter_state& state, const key_t* keys, offset_t seg_base, int count, bool is_terminal) { if constexpr (use_block_load_to_shared) { _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); - process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true, /*Blocked=*/true>( + process_tiles<1, /*FromSmem=*/true, /*Reversed=*/is_tie_reversed, /*Deterministic=*/true, Blocked>( state, keys, seg_base, count, is_terminal); } } @@ -1697,23 +1743,25 @@ private: // Head prefix edge (rank 0): the segment's lowest indices `[0, head_edge_len_items)`, staged at `edge_keys` (base // 0). In global-index order it is the leading region (ascending) / trailing region (descending); a non-head rank or // empty prefix is a no-op. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge(det_filter_state& state) { - process_edge(state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, state.is_head_edge_terminal); + process_edge( + state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, state.is_head_edge_terminal); } // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` // (base `segment_size - tail_edge_len_items`). In global-index order it is the trailing region (ascending) / // leading region (descending); an aligned or non-owned tail is a no-op. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void process_tail_edge(det_filter_state& state) { - process_edge(state, - temp_storage.edge_keys + head_edge_cap_items, - layout.segment_size_off - static_cast(layout.tail_edge_len_items), - layout.tail_edge_len_items, - state.is_tail_edge_terminal); + process_edge( + state, + temp_storage.edge_keys + head_edge_cap_items, + layout.segment_size_off - static_cast(layout.tail_edge_len_items), + layout.tail_edge_len_items, + state.is_tail_edge_terminal); } // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between @@ -1722,7 +1770,7 @@ private: // the edges. Both visit resident before overflow, so the stop check can skip re-streaming the overflow once the // top-k is placed; `is_residency_reversed` keeps the first-visited (high-index) chunks resident in the descending // order so this holds. - template + template _CCCL_DEVICE _CCCL_FORCEINLINE void run_filter(det_filter_state& state) { const auto step = [&](auto&& region) { @@ -1736,26 +1784,26 @@ private: }; if constexpr (is_tie_reversed) { - process_tail_edge(state); + process_tail_edge(state); } else { - process_head_edge(state); + process_head_edge(state); } step([&] { - process_resident(state); + process_resident(state); }); step([&] { - process_overflow(state); + process_overflow(state); }); step([&] { if constexpr (is_tie_reversed) { - process_head_edge(state); + process_head_edge(state); } else { - process_tail_edge(state); + process_tail_edge(state); } }); } @@ -2047,7 +2095,17 @@ private: is_tail_edge_terminal, cand_prefix, !is_select_no_cand_cta}; - run_filter(state); + // Arrangement dispatch: only the boundary-straddling CTA (neither select-all nor select-no candidates) resolves + // ties via `emit_indexed`, so it enters `Blocked` (then switches to striped past the boundary; see + // `process_tiles`). Every other CTA places purely via arrival-order atomics and loads striped throughout. + if (!is_select_all_cand_cta && !is_select_no_cand_cta) + { + run_filter(state); + } + else + { + run_filter(state); + } } // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into From 6b9eeae83c87589f91b3d23eaf77d9130b256b7f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 22:08:50 +0200 Subject: [PATCH 108/126] Use elected leader thread also outside TMA --- cub/cub/agent/agent_batched_topk_cluster.cuh | 29 +++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index d6352cd5ef2..5b5ab8922a9 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -19,8 +19,9 @@ //! via cluster-scope DSMEM atomics. The leader's `hist` therefore plays //! a dual role: its own block-private histogram in step 1, then the //! cluster-merged histogram after the second cluster sync. -//! 3. The leader's thread 0 prefix-scans `hist`, identifies the bucket of -//! the k-th key, and updates the cluster-shared `state`. Every block +//! 3. The leader block prefix-scans `hist` (a block-wide scan) and identifies +//! the bucket of the k-th key; its owning lane updates the cluster-shared +//! `state`. Every block //! reads the leader's packed `state.result_pair` (splitter bucket plus //! early-stop flag) from the leader via DSMEM at the end of each pass and //! folds the bucket into its own local splitter key. @@ -582,9 +583,9 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void issue_bulk_copy(int stage, char* dst, ::cuda::std::span src) { _CCCL_ASSERT(stage >= 0 && stage < PipelineStages, "pipeline stage index out of range"); - // Only the block leader (see `is_load_leader`) drives the mbarrier, for a uniform branch and better mbarrier + // Only the block leader (see `is_block_leader`) drives the mbarrier, for a uniform branch and better mbarrier // codegen. - if (!is_load_leader) + if (!is_block_leader) { return; } @@ -683,9 +684,11 @@ struct agent_batched_topk_cluster // Per-thread mbarrier phase parity, one bit per pipeline stage (see `wait_stage`); the resident load and the // overflow stream keep their per-stage issue/wait calls balanced so each bit tracks its mbarrier's phase. ::cuda::std::uint32_t load_phase{}; - // The single block leader (warp 0's elected lane) that drives the bulk copies. Elected once at construction (the - // block constructs convergently) and cached so the pipeline reuses it instead of re-electing per copy. - const bool is_load_leader = ::cuda::device::__block_elect_one(); + // The single block leader (warp 0's elected lane) for all block-wide single-thread work: the bulk copies and the + // arbitrary `threadIdx.x == 0` picks (state init, counter seeding). Elected once at construction (convergent there) + // and cached, so callers reuse it instead of re-electing; the `elect.sync` predicate also keeps each guarded branch + // on a uniform data path. + const bool is_block_leader = ::cuda::device::__block_elect_one(); // This thread's block-local index, signed for the agent's block-stride loops and cached to avoid re-casting // `threadIdx.x` at each site. const int tid = static_cast(threadIdx.x); @@ -1220,7 +1223,7 @@ private: if (is_single_cta) { // No peers: the front base is 0 (untouched since init) and the back base is just `num_selected`. - if (tid == 0) + if (is_block_leader) { temp_storage.back_local_cnt = num_selected; } @@ -1228,7 +1231,7 @@ private: return ::cuda::std::uint64_t{0}; } // Fold this CTA's back base in parallel with the peers pushing their candidate counts into the same counter. - if (tid == 0) + if (is_block_leader) { seed_local_back(num_selected); } @@ -2872,12 +2875,12 @@ private: return; } - // Every block's thread 0 initializes its local `state`. Only the - // leader's copy is semantically read (non-leaders reach the cluster - // state through `leader_state`), but mirroring the writes everywhere + // Every block's elected leader initializes its local `state`. Only the + // leader block's copy is semantically read (non-leaders reach the cluster + // state through `leader_state`), but mirroring the write in every block // keeps every block's unconditional `state.k` load safe under // compute-sanitizer. - if (tid == 0) + if (is_block_leader) { temp_storage.state.len = static_cast(segment_size); temp_storage.state.k = k; From 1866bc8f1949f374a429d7673c0e3c8ff28c121a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 22:57:18 +0200 Subject: [PATCH 109/126] More clean up --- cub/cub/agent/agent_batched_topk_cluster.cuh | 202 +++++++++++-------- 1 file changed, 114 insertions(+), 88 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 5b5ab8922a9..310f42b48a2 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -170,11 +170,11 @@ struct smem_block_tile_layout // `min_chunks_per_block` chunks, clamped to `[1, cluster_blocks_cap]`. Identical arithmetic on both sides; only the cap // differs (the live cluster size on the device, max launchable blocks on the host). `min_chunks_per_block` is // `static_assert`ed positive, so the divide is well-defined. 64-bit math. -[[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned int effective_cluster_blocks_from_chunks( - ::cuda::std::uint64_t chunks, int min_chunks_per_block, unsigned int cluster_blocks_cap) noexcept +[[nodiscard]] _CCCL_HOST_DEVICE constexpr unsigned effective_cluster_blocks_from_chunks( + ::cuda::std::uint64_t chunks, int min_chunks_per_block, unsigned cluster_blocks_cap) noexcept { const auto blocks = chunks / static_cast<::cuda::std::uint64_t>(min_chunks_per_block); - return static_cast( + return static_cast( ::cuda::std::clamp(blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); } @@ -228,7 +228,6 @@ struct agent_batched_topk_cluster // 32-bit covers every supported segment: the public entry caps the statically-known maximum segment size at 2^21, so // a runtime value exceeding its declared bound is a caller precondition violation (undefined behavior). Unsigned // because all offsets, ranks, and block counts are non-negative (segment sizes are clamped to >= 0 upstream). - // `prime_placement_counters` also returns its two prefix lanes packed into one `uint64_t`, which needs 32-bit lanes. using offset_t = ::cuda::std::uint32_t; using out_offset_t = ::cuda::std::uint32_t; using state_t = cluster_topk_state; @@ -464,7 +463,7 @@ struct agent_batched_topk_cluster // The blocked layout is required by the deterministic tie-break (its cross-CTA scan assumes CTA-rank order matches // ascending contiguous global-index ranges), so it is selected exactly when `need_determinism` is set. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE chunk_partition - make_chunk_partition(offset_t chunks, unsigned int cluster_rank, unsigned int cluster_blocks) const + make_chunk_partition(offset_t chunks, unsigned cluster_rank, unsigned cluster_blocks) const { // Idle ranks never reach here -- `compute_segment_layout` gives them an explicit empty partition. _CCCL_ASSERT(cluster_rank < cluster_blocks, "make_chunk_partition assumes a working rank (rank < cluster_blocks)"); @@ -651,10 +650,10 @@ struct agent_batched_topk_cluster const key_t* block_keys_base; offset_t head_items; offset_t chunks; - unsigned int eff_cluster_blocks; + unsigned eff_cluster_blocks; bool is_idle_rank; chunk_partition part; - unsigned int leader_rank; + unsigned leader_rank; state_t* leader_state; offset_t my_chunks; offset_t my_resident_chunks; @@ -695,9 +694,9 @@ struct agent_batched_topk_cluster // Effective cluster geometry the radix/scan/filter path runs on, set once by `init_effective_cluster`: the launched // cluster, or rank 0 alone (`cluster_blocks == 1`) when a small segment collapsed onto a single resident CTA. The // select-all fast path predates the collapse and uses the raw hardware sregs instead. - unsigned int cluster_rank = 0u; - unsigned int cluster_blocks = 1u; - bool is_single_cta = true; + unsigned cluster_rank = 0u; + unsigned cluster_blocks = 1u; + bool is_single_cta = true; // Per-segment constants for this block, set once and then read directly by the helpers instead of being threaded // through their signatures. `segment_id`/`segment_size`/`k` are set in `process_impl`; `layout` is filled by // `compute_segment_layout` on the radix path (the select-all fast path and redundant CTAs never touch it). @@ -799,7 +798,7 @@ private: // it to rank 0 to form the leader's `shared::cluster` address (no 64-bit pointer, no memory descriptor). Cluster // scope makes it mutually atomic with the leader's `hist_inc` adds. _CCCL_DEVICE _CCCL_FORCEINLINE void - hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, unsigned int leader_rank) + hist_fold_remote(::cuda::std::uint32_t own_bucket_addr32, offset_t v, unsigned leader_rank) { ::cuda::std::uint32_t remote; asm("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(own_bucket_addr32), "r"(leader_rank)); @@ -808,7 +807,7 @@ private: // Generic pointer to this CTA's `state` as seen in the CTA at cluster rank `rank` (reached over DSMEM) -- the PTX // equivalent of cooperative_groups' `map_shared_rank`, via `mapa.u64` (generic-address form). - _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(unsigned int rank) + _CCCL_DEVICE _CCCL_FORCEINLINE state_t* map_state_to_rank(unsigned rank) { const ::cuda::std::uint64_t own = reinterpret_cast<::cuda::std::uint64_t>(&temp_storage.state); ::cuda::std::uint64_t remote; @@ -821,8 +820,7 @@ private: // Two independent 32-bit reductions rather than one 64-bit one: `red.add.u64` on shared memory is emulated with a // CAS spin loop, whereas `red.add.u32` is a native shared-memory atomic. Drives the cross-CTA selected/candidate // prefix scan (see `prime_placement_counters`). - _CCCL_DEVICE _CCCL_FORCEINLINE void - add_remote_prefix(unsigned int target_rank, offset_t push_front, offset_t push_cand) + _CCCL_DEVICE _CCCL_FORCEINLINE void add_remote_prefix(unsigned target_rank, offset_t push_front, offset_t push_cand) { const ::cuda::std::uint32_t own_front = static_cast<::cuda::std::uint32_t>(__cvta_generic_to_shared(&temp_storage.front_local_cnt)); @@ -1214,10 +1212,16 @@ private: // (so a single post-push barrier suffices), leaving `front_local_cnt = sel_prefix` and // `back_local_cnt = num_selected + cand_prefix` -- the absolute output-slot bases `place_one` expects. // - // Returns this CTA's exclusive prefix packed as `(sel_prefix << 32) | cand_prefix` for the driver's region math; - // `is_single_cta` yields 0 (front stays 0, back is just `num_selected`). The successor pushes are lane-parallel (each - // thread owns a strided slice); all threads see CTA-uniform counts, so the guard and the barrier stay uniform. - _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint64_t + // Returns this CTA's exclusive front/back prefixes for the driver's region math (the non-deterministic path uses only + // `sel_prefix`); `is_single_cta` yields both zero (front stays 0, back is just `num_selected`). The successor pushes + // are lane-parallel (each thread owns a strided slice); all threads see CTA-uniform counts, so the guard and the + // barrier stay uniform. + struct cross_cta_prefixes + { + offset_t sel_prefix; // exclusive selected-front prefix across CTAs + offset_t cand_prefix; // exclusive candidate-back prefix across CTAs (candidate-rank space) + }; + _CCCL_DEVICE _CCCL_FORCEINLINE cross_cta_prefixes prime_placement_counters(offset_t push_front, offset_t push_cand, offset_t num_selected) { if (is_single_cta) @@ -1228,7 +1232,7 @@ private: temp_storage.back_local_cnt = num_selected; } __syncthreads(); - return ::cuda::std::uint64_t{0}; + return cross_cta_prefixes{offset_t{0}, offset_t{0}}; } // Fold this CTA's back base in parallel with the peers pushing their candidate counts into the same counter. if (is_block_leader) @@ -1242,8 +1246,8 @@ private: if constexpr (is_scan_descending) { _CCCL_PRAGMA_NOUNROLL() - for (unsigned int rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; - // leader last + for (unsigned rank = threadIdx.x; rank < cluster_rank; rank += threads_per_block) // lower ranks follow; + // leader last { add_remote_prefix(rank, push_front, push_cand); } @@ -1253,7 +1257,7 @@ private: // Higher ranks follow. Stops at `eff_cluster_blocks` since idle ranks own nothing; the leader at the last // effective rank is last. _CCCL_PRAGMA_NOUNROLL() - for (unsigned int rank = cluster_rank + 1u + threadIdx.x; rank < layout.eff_cluster_blocks; + for (unsigned rank = cluster_rank + 1u + threadIdx.x; rank < layout.eff_cluster_blocks; rank += threads_per_block) { add_remote_prefix(rank, push_front, push_cand); @@ -1266,14 +1270,26 @@ private: // The local seed guarantees the back counter carries at least `num_selected` (postcondition of the fused prime). _CCCL_ASSERT(temp_storage.back_local_cnt >= num_selected, "back counter must include the seeded region base after the scan"); - const offset_t sel_prefix = temp_storage.front_local_cnt; - const offset_t cand_prefix = temp_storage.back_local_cnt - num_selected; + const cross_cta_prefixes prefixes{temp_storage.front_local_cnt, temp_storage.back_local_cnt - num_selected}; // Every thread must finish snapshotting the primed bases before any lane's `place_one` mutates the same counters: // the leading boundary edge in each driver has no barrier of its own before its first placement atomic. __syncthreads(); - return (static_cast<::cuda::std::uint64_t>(sel_prefix) << 32) | static_cast<::cuda::std::uint64_t>(cand_prefix); + return prefixes; } + // Which region of the deterministic sweep carries this CTA's last placement work -- the "terminal" region whose last + // tile may hold the K-boundary and so runs the direct tie scan (all earlier regions, plus the lazy overflow/generic + // loops, detect the crossing per tile instead). Exactly one region is terminal. `overflow` names an overflow-terminal + // sweep; it is never queried (overflow uses lazy detection), and each block-load region compares against its own + // enumerator. + enum class terminal_region + { + head_edge, + resident, + overflow, + tail_edge + }; + // Deterministic final-filter state: the `run()`-local tie-break values (counts, prefixes, region extents) plus the // mutable `running`/`is_tie_active` scan cursor, bundled into one POD so the flattened filter methods below take a // single `state` argument instead of ~16. A deterministic-path-only local of `write_deterministic_topk`; never @@ -1296,9 +1312,7 @@ private: smem_keys_t resident_front; int front_count; bool is_select_all_cand_cta; - bool is_resident_terminal; - bool is_head_edge_terminal; - bool is_tail_edge_terminal; + terminal_region terminal; // Mutable per-invocation tie-break scan cursor. offset_t running; bool is_tie_active; @@ -1648,7 +1662,11 @@ private: /*Reversed=*/is_residency_reversed, /*Deterministic=*/true, Blocked>( - state, state.resident_front.data(), state.front_seg_base, state.front_count, state.is_resident_terminal); + state, + state.resident_front.data(), + state.front_seg_base, + state.front_count, + state.terminal == terminal_region::resident); } else { @@ -1750,7 +1768,11 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void process_head_edge(det_filter_state& state) { process_edge( - state, temp_storage.edge_keys, offset_t{0}, layout.head_edge_len_items, state.is_head_edge_terminal); + state, + temp_storage.edge_keys, + offset_t{0}, + layout.head_edge_len_items, + state.terminal == terminal_region::head_edge); } // Peeled tail suffix edge (tail owner): the segment's highest indices, staged at `edge_keys + head_edge_cap_items` @@ -1764,7 +1786,7 @@ private: temp_storage.edge_keys + head_edge_cap_items, layout.segment_size_off - static_cast(layout.tail_edge_len_items), layout.tail_edge_len_items, - state.is_tail_edge_terminal); + state.terminal == terminal_region::tail_edge); } // Drive the four regions in global-index order (ascending, or descending under `is_tie_reversed`), bailing between @@ -1847,7 +1869,6 @@ private: KeyOutIt block_keys_out; out_offset_t num_back; offset_t sel_prefix; - offset_t cand_prefix; offset_t my_front; // this CTA's front region size (see `write_deterministic_topk`'s `my_front`) smem_keys_t resident_keys; }; @@ -1918,25 +1939,28 @@ private: } } - // Non-deterministic final filter driver. `prime_placement_counters` gives this block disjoint front/back bases - // (`sel_prefix`/`cand_prefix`); overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and - // place into block-local SMEM atomics. `run_pass` breaks the stream once this CTA's contribution is fully placed - // (all its selected keys in the front, and its back counter reached the region end `k`). + // Non-deterministic final filter driver. `prime_placement_counters` primes this block's front/back placement counters + // to their absolute bases (only `sel_prefix` is read back here; the back base is used in place via `back_local_cnt`); + // overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and place into block-local SMEM + // atomics. `run_pass` breaks the stream once this CTA's contribution is fully placed (all its selected keys in the + // front, and its back counter reached the region end `k`). template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( - out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) + out_offset_t num_tie_winners, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) { const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); const offset_t my_sel = participates ? temp_storage.num_strictly_selected : offset_t{0}; const offset_t my_cand = participates ? temp_storage.my_candidates : offset_t{0}; // The scan doubles as counter priming: it leaves this CTA's placement counters holding its absolute region bases // (front = `sel_prefix`, back = `num_selected + cand_prefix`), so no explicit priming follows. - const offset_t num_selected = static_cast(k) - static_cast(num_kth); - const ::cuda::std::uint64_t packed_prefix = prime_placement_counters(my_sel, my_cand, num_selected); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); - // The selected region has size `k - num_kth`, so the selected prefix leaves room for the `num_kth` tie-back slots. - _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_kth) + const offset_t num_selected = static_cast(k) - static_cast(num_tie_winners); + // Only the front prefix is needed here: back placement uses the primed `back_local_cnt` directly and the early stop + // tests it against `k`, so the candidate prefix is discarded. + const offset_t sel_prefix = prime_placement_counters(my_sel, my_cand, num_selected).sel_prefix; + // This CTA's exclusive selected prefix must leave room for the `num_tie_winners` tie-back slots. It is `<=`, not + // `==`: `sel_prefix` reaches the full front total `num_selected` (so the sum hits `k`) only for the last CTA in + // scan order (the leader); every earlier CTA sits strictly below it. + _CCCL_ASSERT(static_cast<::cuda::std::uint64_t>(sel_prefix) + static_cast<::cuda::std::uint64_t>(num_tie_winners) <= static_cast<::cuda::std::uint64_t>(k), "selected prefix must fit before the tie-back output region"); // This CTA's front region size (mirrors `write_deterministic_topk`): the leader is last in scan order and derives @@ -1945,7 +1969,7 @@ private: const offset_t my_front = (cluster_rank == layout.leader_rank) ? (num_selected - sel_prefix) : my_sel; nondet_filter_state state{ - identify_op, block_keys_out, num_kth, sel_prefix, cand_prefix, my_front, resident_keys}; + identify_op, block_keys_out, num_tie_winners, sel_prefix, my_front, resident_keys}; run_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. @@ -1995,17 +2019,18 @@ private: // `run_filter`, which drives the per-region sweeps. template _CCCL_DEVICE _CCCL_FORCEINLINE void write_deterministic_topk( - out_offset_t num_kth, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) + out_offset_t num_tie_winners, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) { - // Early stop is not special-cased: `total_candidates == num_kth` then makes every CTA `is_select_all_cand_cta`. + // Early stop is not special-cased: `total_candidates == num_tie_winners` then makes every CTA + // `is_select_all_cand_cta`. // // Cache `total_candidates` now, while every block is still tightly coupled to the pass loop's final cluster // barrier -- a post-scan re-read of `leader_state` could touch an already-returned leader (barrier gone). const offset_t total_candidates = layout.leader_state->len; // Guards the unsigned `k - num_back` below: the remaining tie count fits both `k` and the splitter bucket. - _CCCL_ASSERT(num_kth <= k && static_cast(num_kth) <= total_candidates, + _CCCL_ASSERT(num_tie_winners <= k && static_cast(num_tie_winners) <= total_candidates, "remaining k must fit within both the top-k and the splitter bucket"); - const out_offset_t num_back = num_kth; // all candidates go to the back; the front holds only selected keys + const out_offset_t num_back = num_tie_winners; // all candidates go to the back; the front holds only selected keys const out_offset_t num_selected = k - num_back; // front region const bool participates = !layout.is_idle_rank && (cluster_rank != layout.leader_rank); @@ -2017,10 +2042,9 @@ private: const offset_t push_front = my_sel; // The scan doubles as counter priming: it leaves this CTA's placement counters holding its absolute region bases // (front = `sel_prefix`, back = `num_selected + cand_prefix`), so no explicit priming follows. - const ::cuda::std::uint64_t packed_prefix = - prime_placement_counters(push_front, my_cand, static_cast(num_selected)); - const offset_t sel_prefix = static_cast(packed_prefix >> 32); - const offset_t cand_prefix = static_cast(packed_prefix & 0xffffffffu); + const auto prefixes = prime_placement_counters(push_front, my_cand, static_cast(num_selected)); + const offset_t sel_prefix = prefixes.sel_prefix; + const offset_t cand_prefix = prefixes.cand_prefix; // Guards the leader's remainder subtractions below (`total_candidates - cand_prefix`, `num_selected - sel_prefix`). _CCCL_ASSERT(cand_prefix <= total_candidates && sel_prefix <= static_cast(num_selected), "cross-CTA prefixes must stay within their candidate and selected totals"); @@ -2053,19 +2077,27 @@ private: // part of this span. const int front_count = static_cast(resident_keys.size()); - // Terminal-region flags for the last-tile direct-scan gate (block-load regions only; the generic resident loop - // and overflow stream pass `false` and stay on lazy per-tile detection). A region is terminal when no later - // region in the sweep (head/resident/overflow/tail-edge ascending, reversed descending) carries work. + // The terminal region (see `terminal_region`): the last one carrying work in sweep order (ascending + // head->resident->overflow->tail, reversed tail->resident->overflow->head). const bool has_head = layout.head_edge_len_items > 0; const bool has_resident = front_count > 0; const bool has_overflow = layout.overflow_chunks > offset_t{0}; const bool has_tail_edge = layout.tail_edge_len_items > 0; - [[maybe_unused]] const bool is_resident_terminal = - is_tie_reversed ? (!has_overflow && !has_head) : (!has_overflow && !has_tail_edge); - [[maybe_unused]] const bool is_head_edge_terminal = - is_tie_reversed ? true : (!has_resident && !has_overflow && !has_tail_edge); - [[maybe_unused]] const bool is_tail_edge_terminal = - is_tie_reversed ? (!has_resident && !has_overflow && !has_head) : true; + terminal_region terminal; + if constexpr (is_tie_reversed) + { + terminal = has_head ? terminal_region::head_edge + : has_overflow ? terminal_region::overflow + : has_resident ? terminal_region::resident + : terminal_region::tail_edge; + } + else + { + terminal = has_tail_edge ? terminal_region::tail_edge + : has_overflow ? terminal_region::overflow + : has_resident ? terminal_region::resident + : terminal_region::head_edge; + } // Segment-local base of the resident-front span (its lowest-index resident chunk; `resident_base` shifts it to // the high-index window under `is_residency_reversed`). The blocked partition packs the front contiguously, so @@ -2093,9 +2125,7 @@ private: resident_keys, front_count, is_select_all_cand_cta, - is_resident_terminal, - is_head_edge_terminal, - is_tail_edge_terminal, + terminal, cand_prefix, !is_select_no_cand_cta}; // Arrangement dispatch: only the boundary-straddling CTA (neither select-all nor select-no candidates) resolves @@ -2359,10 +2389,10 @@ private: const offset_t full_slots = block_tile_capacity / static_cast(chunk_items); [[maybe_unused]] const bool needs_streaming = layout.my_chunks > full_slots; - // Does this rank own the global tail, and does that tail carry an unaligned suffix? (block-load path only; the - // generic fallback reads any trailing items straight from gmem and never peels.) - [[maybe_unused]] offset_t tail_suffix_items = offset_t{0}; - bool owns_suffix_tail = false; + // The unaligned suffix of the global tail chunk, peeled into `edge_keys` when this rank owns it (block-load path + // only; the generic fallback reads any trailing items straight from gmem and never peels). Stays 0 for every other + // rank and for an aligned tail, so it doubles as the "peel this tail" predicate. + offset_t tail_suffix_items = offset_t{0}; if constexpr (use_block_load_to_shared) { // This rank owns the global tail iff its last owned chunk is chunk `chunks-1` (its local index `my_chunks-1`, @@ -2374,17 +2404,13 @@ private: tail_suffix_items = static_cast(tail_chunk.count) - ::cuda::round_down(static_cast(tail_chunk.count), static_cast(load_align_items)); - owns_suffix_tail = tail_suffix_items != offset_t{0}; } } - // `should_peel_tail` mirrors `owns_suffix_tail` (always peel, per above). `stream_slots` clamps into `[1, - // full_slots]`. + // Streaming-slot reservation: clamped into `[1, full_slots]` when this rank overflows, else 0. offset_t stream_slots = offset_t{0}; - bool should_peel_tail = false; if constexpr (use_block_load_to_shared) { - should_peel_tail = owns_suffix_tail; if (needs_streaming) { const offset_t excess = layout.my_chunks - full_slots; @@ -2411,7 +2437,7 @@ private: // Persistent boundary-edge lengths: the head prefix lives on rank 0 (`head_items` is 0 on the generic fallback and // for an aligned base); the peeled tail suffix lives on the tail owner whenever it is unaligned. layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_items) : 0; - layout.tail_edge_len_items = should_peel_tail ? static_cast(tail_suffix_items) : 0; + layout.tail_edge_len_items = static_cast(tail_suffix_items); } // Radix histogram/scan passes: refine the splitter one `bits_per_pass` digit at a time until the top-k prefix is @@ -2660,16 +2686,17 @@ private: // ----------------------------------------------------------------------- // Final filter pass: write the top-k keys for this segment. Strictly- - // selected keys go to the front; the `num_kth` tied candidates fill the + // selected keys go to the front; the `num_tie_winners` tied candidates fill the // back. `kth_key_bits_local` already holds the full splitter key (folded // from each pass's bucket above), so no broadcast is needed here. // ----------------------------------------------------------------------- - auto block_keys_out = d_key_segments_out_it[segment_id]; - const out_offset_t num_kth = layout.leader_state->k; // remaining k after the radix passes - // Each pass keeps `0 < num_kth <= k` inside the splitter bucket (same leader read as `num_kth`, equally safe). - _CCCL_ASSERT( - num_kth > out_offset_t{0} && num_kth <= k && static_cast(num_kth) <= layout.leader_state->len, - "radix passes produced an invalid remaining-k count"); + auto block_keys_out = d_key_segments_out_it[segment_id]; + const out_offset_t num_tie_winners = layout.leader_state->k; // remaining k after the radix passes + // Each pass keeps `0 < num_tie_winners <= k` inside the splitter bucket (same leader read as `num_tie_winners`, + // equally safe). + _CCCL_ASSERT(num_tie_winners > out_offset_t{0} && num_tie_winners <= k + && static_cast(num_tie_winners) <= layout.leader_state->len, + "radix passes produced an invalid remaining-k count"); // `last_pass` controls how many radix levels of `kth_key_bits_local` are significant. After an early-stop break, // only the first `last_pass` digits of the splitter were folded; comparing all bits would treat the (still-zero) @@ -2681,11 +2708,11 @@ private: __syncthreads(); if constexpr (need_determinism) { - write_deterministic_topk(num_kth, identify_op, block_keys_out, resident_keys); + write_deterministic_topk(num_tie_winners, identify_op, block_keys_out, resident_keys); } else { - write_nondeterministic_topk(num_kth, identify_op, block_keys_out, resident_keys); + write_nondeterministic_topk(num_tie_winners, identify_op, block_keys_out, resident_keys); } // No cluster barrier after the final filter pass: both filter paths place output via block-local SMEM atomics into @@ -2698,8 +2725,7 @@ private: // Copies an entire segment `input[i] -> output[i]` for the select-all fast path (`k >= segment_size`). Runs before // the effective-cluster collapse, so it takes the raw hardware rank/size (the `cluster_*` members are not set yet) // and reads the `segment_id`/`segment_size` members `process_impl` filled in. - _CCCL_DEVICE _CCCL_FORCEINLINE void - copy_segment_select_all(unsigned int hw_cluster_rank, unsigned int hw_cluster_blocks) + _CCCL_DEVICE _CCCL_FORCEINLINE void copy_segment_select_all(unsigned hw_cluster_rank, unsigned hw_cluster_blocks) { constexpr int copy_items = copy_items_per_thread_clamped; const offset_t num_items = static_cast(segment_size); @@ -2793,7 +2819,7 @@ private: // freeing their SM slots. The decision is per-segment uniform across the block, so a redundant CTA returns whole. // Compiled out for host-exact sizes, which the dispatch already sized to exact cluster blocks (returns true for all). [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool - init_effective_cluster(unsigned int hw_cluster_rank, unsigned int hw_cluster_blocks) + init_effective_cluster(unsigned hw_cluster_rank, unsigned hw_cluster_blocks) { cluster_rank = hw_cluster_rank; cluster_blocks = hw_cluster_blocks; @@ -2824,8 +2850,8 @@ private: // `init_effective_cluster` below); only the pre-collapse steps here (segment id, select-all fast path) use the raw // hardware values. Runtime cluster blocks match the launch attribute the dispatch passed to `cudaLaunchKernelExC` // (or the kernel's `__cluster_dims__` on CDP). - const unsigned int hw_cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); - const unsigned int hw_cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); + const unsigned hw_cluster_rank = ::cuda::ptx::get_sreg_cluster_ctarank(); + const unsigned hw_cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); _CCCL_ASSERT(hw_cluster_blocks > 0u && hw_cluster_rank < hw_cluster_blocks, "hardware cluster rank must lie within a non-empty cluster"); segment_id = static_cast(blockIdx.x / hw_cluster_blocks); From 65390a6ac1c93d4e8a34cb5fe726b1089e2a9e7f Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sat, 11 Jul 2026 23:44:43 +0200 Subject: [PATCH 110/126] Make sure no unrolling decision is left to the compiler --- cub/cub/agent/agent_batched_topk_cluster.cuh | 53 ++++++++++++++++++-- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 310f42b48a2..cbbc9d827a5 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -502,6 +502,7 @@ struct agent_batched_topk_cluster _CCCL_DEVICE _CCCL_FORCEINLINE void stage_and_fold_edge(key_t* dst, const key_t* src, int count, Apply&& apply) const { _CCCL_ASSERT(count >= 0 && count <= head_edge_cap_items, "a boundary edge must fit in its edge_keys slot"); + _CCCL_PRAGMA_NOUNROLL() for (int local = tid; local < count; local += threads_per_block) { const key_t key = src[local]; @@ -629,6 +630,7 @@ struct agent_batched_topk_cluster { _CCCL_ASSERT(stage >= 0 && stage < PipelineStages, "pipeline stage index out of range"); const ::cuda::std::uint32_t parity = (load_phase >> stage) & 1u; + _CCCL_PRAGMA_NOUNROLL() while (!::cuda::ptx::mbarrier_try_wait_parity(&temp_storage.load_mbar[stage], parity)) { } @@ -752,9 +754,21 @@ struct agent_batched_topk_cluster private: _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() { - for (int i = tid; i < num_buckets; i += threads_per_block) + // The `num_buckets / threads_per_block` full strided rounds are in range for every thread, so they unroll with no + // bounds check; the `< threads_per_block` leftover is at most one more guarded write per thread, compiled out when + // `num_buckets` is a multiple of `threads_per_block`. + constexpr int full_rounds = num_buckets / threads_per_block; + _CCCL_PRAGMA_UNROLL_FULL() + for (int r = 0; r < full_rounds; ++r) + { + temp_storage.hist[r * threads_per_block + tid] = 0; + } + if constexpr (num_buckets % threads_per_block != 0) { - temp_storage.hist[i] = 0; + if (const int i = full_rounds * threads_per_block + tid; i < num_buckets) + { + temp_storage.hist[i] = 0; + } } } @@ -1003,6 +1017,7 @@ private: // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else // the phase advances twice and a lagging thread misses the flip and spins forever. __syncthreads(); + _CCCL_PRAGMA_NOUNROLL() for (int i = 0; i < stream_stages; ++i) { const offset_t overflow_idx = @@ -1047,6 +1062,7 @@ private: // in the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. bool is_stopped = false; const offset_t split = (::cuda::std::min) (static_cast(stream_stages), layout.overflow_chunks); + _CCCL_PRAGMA_NOUNROLL() for (offset_t i = 0; i < split; ++i) { if (!consume(i)) @@ -1064,6 +1080,7 @@ private: mid(); // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). + _CCCL_PRAGMA_NOUNROLL() for (offset_t i = split; i < layout.overflow_chunks; ++i) { if (!consume(i)) @@ -1077,6 +1094,7 @@ private: // mbarriers were never waited, and they must complete before the block can exit (their slots are never read). // `stream_inflight_mask` is block-uniform (set/cleared under uniform control flow), so the trip count and each // collective `wait_stage` are uniform across the block. + _CCCL_PRAGMA_NOUNROLL() while (stream_inflight_mask != ::cuda::std::uint32_t{0}) { const int drain_stage = __ffs(static_cast(stream_inflight_mask)) - 1; @@ -1091,6 +1109,7 @@ private: // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no // SMEM reuse), with the walk still snaking for L2 locality. mid(); + _CCCL_PRAGMA_NOUNROLL() for (offset_t i = 0; i < layout.overflow_chunks; ++i) { const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); @@ -1559,6 +1578,7 @@ private: // Phase A: while ties remain, load blocked (`emit_indexed`'s scan needs it) and resolve them. A region entered // past the crossing (its ties already placed) has `is_tie_active == false`, so phase A is skipped and everything // falls through to phase B. + _CCCL_PRAGMA_NOUNROLL() for (; state.is_tie_active && tile_base < count; tile_base += tile_size) { key_t keys[ItemsPerThread]; @@ -1611,6 +1631,7 @@ private: // Phase B: ties are placed, so any remaining candidate-class keys are losers (dropped via `do_arrival == false`); // only strictly-selected keys are placed, to the front in arrival order. Load striped (coalesced/conflict-free); // no scan or barrier -- `front_local_cnt` is order-independent, so the A->B switch needs none either. + _CCCL_PRAGMA_NOUNROLL() for (; tile_base < count; tile_base += tile_size) { key_t keys[ItemsPerThread]; @@ -1623,6 +1644,7 @@ private: else { // Single striped pass: the non-deterministic filter (always) or a deterministic CTA that never scans. + _CCCL_PRAGMA_NOUNROLL() for (; tile_base < count; tile_base += tile_size) { key_t keys[ItemsPerThread]; @@ -1671,6 +1693,7 @@ private: else { const int resident_chunk_count = static_cast(layout.my_resident_chunks); + _CCCL_PRAGMA_NOUNROLL() for (int slot = 0; slot < resident_chunk_count; ++slot) { const int local_slot = is_residency_reversed ? (resident_chunk_count - 1 - slot) : slot; @@ -1887,6 +1910,7 @@ private: // below), so iterate the aligned bulk (`round_down(count, load_align_items)`), not `chunk.count`. key_t* const resident_ptr = state.resident_keys.data(); int cursor_items = 0; + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const auto chunk = get_chunk( @@ -1922,6 +1946,7 @@ private: { // Generic fallback: each resident chunk is staged in its own `key_slots` slot (indexed by `local_chunk`); no // edges are peeled (boundary items are read as ordinary chunks). + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const auto chunk = get_chunk(layout.part.global_index(local_chunk), layout.segment_size_off, layout.head_items); @@ -2201,6 +2226,7 @@ private: int next_off_bytes = 0; // Load every resident chunk's aligned bulk, densely packed in slot order. + _CCCL_PRAGMA_NOUNROLL() for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); @@ -2212,6 +2238,7 @@ private: // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs // no stored per-stage state. int read_off_bytes = 0; + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const int stage = static_cast(local_chunk % static_cast(prologue)); @@ -2243,6 +2270,7 @@ private: } else { + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const offset_t chunk_idx = layout.part.global_index(layout.resident_base + local_chunk); @@ -2454,6 +2482,7 @@ private: constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); int last_pass = num_passes; + _CCCL_PRAGMA_NOUNROLL() for (int pass = 0; pass < num_passes; ++pass) { const bool is_first_pass = (pass == 0); @@ -2487,6 +2516,7 @@ private: } else { + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const offset_t chunk_idx = layout.part.global_index(layout.resident_base + local_chunk); @@ -2533,14 +2563,28 @@ private: if (cluster_rank != layout.leader_rank && !layout.is_idle_rank) { const ::cuda::std::uint32_t hist_smem32 = hist_base32(); - for (int i = tid; i < num_buckets; i += threads_per_block) - { + // Same split as `reset_hist`: unroll the full strided rounds, then (only when there's a remainder) at most one + // more guarded fold. + const auto fold_bucket = [&](int i) { const offset_t bucket_count = temp_storage.hist[i]; if (bucket_count != 0) { hist_fold_remote( hist_smem32 + static_cast<::cuda::std::uint32_t>(i) * sizeof(offset_t), bucket_count, layout.leader_rank); } + }; + constexpr int full_rounds = num_buckets / threads_per_block; + _CCCL_PRAGMA_UNROLL_FULL() + for (int r = 0; r < full_rounds; ++r) + { + fold_bucket(r * threads_per_block + tid); + } + if constexpr (num_buckets % threads_per_block != 0) + { + if (const int i = full_rounds * threads_per_block + tid; i < num_buckets) + { + fold_bucket(i); + } } } @@ -2759,6 +2803,7 @@ private: } }(); + _CCCL_PRAGMA_NOUNROLL() for (offset_t base = 0; base < full_tiles; base += step) { offset_t idx[copy_items]; From 355e107b63f30e407df6db4e70f201cface8a328 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 12 Jul 2026 02:15:36 +0200 Subject: [PATCH 111/126] Revert forcing nounroll for some loops due to perf --- cub/cub/agent/agent_batched_topk_cluster.cuh | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index cbbc9d827a5..67e9b019729 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -304,7 +304,7 @@ struct agent_batched_topk_cluster } // Two clamp flavors. The `floor` clamp pairs with an unpredicated main loop over full tiles plus a single - // non-unrolled remainder loop (the chunk helper `for_each_chunk_key_impl` and the copy fast path); the `ceil` clamp + // remainder loop (the chunk helper `for_each_chunk_key_impl` and the copy fast path); the `ceil` clamp // keeps the whole resident segment inside one tile for fully-predicated loops (the final filters' `process_tiles`). static constexpr int histogram_items_per_thread_clamped = clamp_unroll(segment_rounds_floor, histogram_items_per_thread); @@ -516,7 +516,7 @@ struct agent_batched_topk_cluster // by one unrolled loop, then handed to `apply` by a second. Splitting the loads out matters for the histogram passes: // `apply`'s SMEM atomics can't be proven disjoint from the SMEM key reads, so a fused loop would interleave each load // with its atomic instead of hoisting the whole load wave ahead. `Unroll` is the caller's clamped (floor) items per - // thread, so the sub-tile remainder is handled by a single non-unrolled fused block-stride loop bounded by the count. + // thread, so the sub-tile remainder is handled by a single fused block-stride loop bounded by the count. template _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key_impl(const key_t* keys, int chunk_count, Apply&& apply) const { @@ -540,8 +540,7 @@ struct agent_batched_topk_cluster } } - // Sub-tile remainder - _CCCL_PRAGMA_NOUNROLL() + // Sub-tile remainder. No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (int local = full_tiles + tid; local < chunk_count; local += threads_per_block) { apply(keys[local], local); @@ -1017,7 +1016,7 @@ private: // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else // the phase advances twice and a lagging thread misses the flip and spins forever. __syncthreads(); - _CCCL_PRAGMA_NOUNROLL() + // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (int i = 0; i < stream_stages; ++i) { const offset_t overflow_idx = @@ -2226,7 +2225,7 @@ private: int next_off_bytes = 0; // Load every resident chunk's aligned bulk, densely packed in slot order. - _CCCL_PRAGMA_NOUNROLL() + // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (int stage = 0; stage < prologue; ++stage) { const auto src = bulk_src(static_cast(stage)); @@ -2482,7 +2481,7 @@ private: constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); int last_pass = num_passes; - _CCCL_PRAGMA_NOUNROLL() + // No unroll pragma on purpose: the compiler's partial/full auto-unroll beats forced nounroll here. for (int pass = 0; pass < num_passes; ++pass) { const bool is_first_pass = (pass == 0); @@ -2845,7 +2844,7 @@ private: // Sub-tile remainder. Iterate the tail relative to zero and offset by `full_tiles`, so every index stays in // `[0, num_items)` without relying on `full_tiles + cluster_tid` staying within `offset_t`. const offset_t tail_items = num_items - full_tiles; - _CCCL_PRAGMA_NOUNROLL() + // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (offset_t local = cluster_tid; local < tail_items; local += cluster_threads) { const auto seg_idx = static_cast(full_tiles + local); From 960d6a2b021f5ab15be50b03c7fdd6f5f2198e5e Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 12 Jul 2026 03:14:56 +0200 Subject: [PATCH 112/126] Flatten more lambdas --- cub/cub/agent/agent_batched_topk_cluster.cuh | 119 +++++++++---------- 1 file changed, 59 insertions(+), 60 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 67e9b019729..e4b28351a78 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -983,6 +983,39 @@ private: static_cast(::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); } + // Consume the `i`-th visit on the block-load path (its ping-pong-ordered position is `overflow_idx`): wait for its + // slot, fold its keys via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the just-freed + // slot (a barrier orders the block's read ahead of the overwriting copy). Returns false once `should_continue()` + // reports the top-k fully placed -- polled before the prefetch so we never launch a copy we would only drain again; + // the prefetches already in flight are drained after `run_pass`'s consume loops. + template + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool + consume_overflow_visit(offset_t i, BlockApply&& block_apply, Continue&& should_continue) + { + const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); + const int stage = static_cast(overflow_idx % static_cast(stream_stages)); + if (stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) + { + wait_stage(stage); + stream_inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); + } + block_apply(stage, overflow_idx); + + if (!should_continue()) + { + return false; + } + + const offset_t next_step = i + static_cast(stream_stages); + if (next_step < layout.overflow_chunks) + { + const offset_t next_overflow_idx = stream_is_forward ? next_step : (layout.overflow_chunks - 1 - next_step); + __syncthreads(); + issue_load(stage, next_overflow_idx); + } + return true; + } + // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from // gmem (fallback). `mid()` runs at most once per pass, positioned to overlap the caller's resident-chunk work with @@ -1026,37 +1059,6 @@ private: stream_is_primed = true; } - // Consume the `i`-th visit (its ping-pong-ordered position is `overflow_idx`): wait for its slot, fold its keys - // via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the slot just freed (a barrier - // guards the slot before the async copy can overwrite the data the block was just reading). Returns false once - // `should_continue()` reports the top-k fully placed -- polled before the prefetch so we never launch a copy we - // would only drain again; the up-to-`stream_stages - 1` prefetches already in flight (from earlier visits or - // priming) are drained after the loop. - const auto consume = [&](offset_t i) -> bool { - const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); - const int stage = static_cast(overflow_idx % static_cast(stream_stages)); - if (stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) - { - wait_stage(stage); - stream_inflight_mask &= ~(::cuda::std::uint32_t{1} << stage); - } - block_apply(stage, overflow_idx); - - if (!should_continue()) - { - return false; - } - - const offset_t next_step = i + static_cast(stream_stages); - if (next_step < layout.overflow_chunks) - { - const offset_t next_overflow_idx = stream_is_forward ? next_step : (layout.overflow_chunks - 1 - next_step); - __syncthreads(); - issue_load(stage, next_overflow_idx); - } - return true; - }; - // Phase 1: consume the first `stream_stages` visits (the chunks reused from the previous pass, already resident // in the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. bool is_stopped = false; @@ -1064,7 +1066,7 @@ private: _CCCL_PRAGMA_NOUNROLL() for (offset_t i = 0; i < split; ++i) { - if (!consume(i)) + if (!consume_overflow_visit(i, block_apply, should_continue)) { is_stopped = true; break; @@ -1082,7 +1084,7 @@ private: _CCCL_PRAGMA_NOUNROLL() for (offset_t i = split; i < layout.overflow_chunks; ++i) { - if (!consume(i)) + if (!consume_overflow_visit(i, block_apply, should_continue)) { break; } @@ -2165,6 +2167,25 @@ private: } } + // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has none. + // Slot -> rank-local chunk index is `resident_base + slot` (identity unless `is_residency_reversed` shifts the + // resident window to the high-index chunks). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span resident_bulk_src(offset_t slot) const + { + const offset_t chunk_idx = layout.part.global_index(layout.resident_base + slot); + const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); + const offset_t bulk = + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); + if (bulk == 0) + { + return {}; + } + // Chunks start after the aligned head on an alignment-multiple stride (mirrors `issue_load`). + _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), + "resident loader received a chunk with an unaligned start"); + return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; + } + // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 // needs no candidate filtering). Publishes the resident span as `resident_keys` for the later passes and the final @@ -2198,28 +2219,6 @@ private: // the unaligned suffix that is always peeled into `edge_keys`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); - // Resident slot -> rank-local chunk index (`resident_base + slot`; identity unless `is_residency_reversed` - // shifts the resident window to the high-index chunks). - const auto resident_local = [&](offset_t slot) -> offset_t { - return layout.resident_base + slot; - }; - // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has - // none. - const auto bulk_src = [&](offset_t slot) -> ::cuda::std::span { - const offset_t chunk_idx = layout.part.global_index(resident_local(slot)); - const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); - const offset_t bulk = - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); - if (bulk == 0) - { - return {}; - } - // Chunks start after the aligned head on an alignment-multiple stride (mirrors `issue_load`). - _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), - "resident loader received a chunk with an unaligned start"); - return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; - }; - // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. int next_off_bytes = 0; @@ -2228,21 +2227,21 @@ private: // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (int stage = 0; stage < prologue; ++stage) { - const auto src = bulk_src(static_cast(stage)); + const auto src = resident_bulk_src(static_cast(stage)); issue_bulk_copy(stage, key_slots + next_off_bytes, src); next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; } // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed - // and consumed in the same order), and `bulk_src(local_chunk)` recomputes its length, so the read span needs - // no stored per-stage state. + // and consumed in the same order), and `resident_bulk_src(local_chunk)` recomputes its length, so the read span + // needs no stored per-stage state. int read_off_bytes = 0; _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { const int stage = static_cast(local_chunk % static_cast(prologue)); wait_stage(stage); - const int read_len_items = static_cast(::cuda::std::size(bulk_src(local_chunk))); + const int read_len_items = static_cast(::cuda::std::size(resident_bulk_src(local_chunk))); for_each_chunk_key( ::cuda::ptr_rebind(key_slots + read_off_bytes), read_len_items, add_first_pass); read_off_bytes += read_len_items * int{sizeof(key_t)}; @@ -2250,7 +2249,7 @@ private: const offset_t next_local_chunk = local_chunk + static_cast(prologue); if (next_local_chunk < layout.my_resident_chunks) { - const auto src = bulk_src(next_local_chunk); + const auto src = resident_bulk_src(next_local_chunk); // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads // leave the wait above would advance the phase twice, stranding a lagging waiter forever. __syncthreads(); From cf4272daa74c67a0962f7d0f25eae1ecf2641715 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Sun, 12 Jul 2026 19:37:41 +0200 Subject: [PATCH 113/126] Improve TMA latency hiding - Prime streaming pipeline already while still iterating through the resident pipeline. - Avoid runtime division/remainder where possible. Add comments/TODOs where not easily avoidable. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 391 ++++++++++++++----- 1 file changed, 285 insertions(+), 106 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index e4b28351a78..33525fd08b7 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -278,6 +278,13 @@ struct agent_batched_topk_cluster static constexpr int tie_break_items_per_thread = TieBreakItemsPerThread; static constexpr int copy_items_per_thread = CopyItemsPerThread; // select-all copy fast path + // Radix passes for this key type (compile-time). The overflow stream flips direction each pass, and `run` preselects + // the initial direction so the leftover after `num_passes` matches the deterministic filter (see `run`). The first + // wave is always primed in that initial direction, so `first_wave_is_forward` -- the only direction + // `first_wave_chunk_for_stage` needs -- is a compile-time constant (non-deterministic keeps the forward default). + static constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); + static constexpr bool first_wave_is_forward = !need_determinism || ((!is_tie_reversed) ^ ((num_passes & 1) != 0)); + // Static upper bound on segment size: exact for constant/immediate sizes, the type maximum for runtime sizes. static constexpr ::cuda::std::int64_t static_max_segment_size = ::cuda::args::__traits::highest; @@ -467,6 +474,13 @@ struct agent_batched_topk_cluster { // Idle ranks never reach here -- `compute_segment_layout` gives them an explicit empty partition. _CCCL_ASSERT(cluster_rank < cluster_blocks, "make_chunk_partition assumes a working rank (rank < cluster_blocks)"); + // TODO(cccl): both branches divide by `cluster_blocks`, the launched cluster dimension (runtime, non-power-of-two). + // Critical path: runs at kernel entry and gates the first TMA copy (the partition names which chunks to load), so + // it is one-time but has nothing earlier to hide it and cannot be relocated after priming -- only removing the + // divide helps. For host-exact (immediate/constant) sizes the effective cluster equals the launch + // (`enable_runtime_single_cta == false`), so this divisor is host-known: precompute a `cuda::fast_mod_div` on the + // host and pass it in to turn the divide into a multiply-shift. Highest-priority runtime divide here and likely + // worth doing (out-of-scope for now); N/A when the effective cluster is device-collapsed (deferred/sequence sizes). if constexpr (need_determinism) { const offset_t chunks_per_cta = ::cuda::ceil_div(chunks, static_cast(cluster_blocks)); @@ -949,25 +963,49 @@ private: "streaming depth exceeds the overflow chunk count"); } - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_load(int stage, offset_t overflow_idx) + // gmem aligned bulk of the chunk at rank-local index `window_base + index` (the resident or overflow window): its + // `count` rounded down to a `load_align` multiple, as a span from `block_keys_base`; empty when it has none. Every + // chunk starts on a `load_align` boundary (guard-free TMA path); the global-last chunk's unaligned suffix is always + // peeled into `edge_keys`, so the aligned bulk excludes it (`bulk == count` for every interior chunk). Shared source + // computation for both issue paths (`issue_resident_copy`, `issue_stream_copy`). + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span + chunk_bulk_src(offset_t window_base, offset_t index) const + { + const auto chunk = + get_chunk(layout.part.global_index(window_base + index), layout.segment_size_off, layout.head_items); + const offset_t bulk = + ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); + if (bulk == 0) + { + return {}; + } + _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), + "a chunk source must start on a load-alignment boundary"); + return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; + } + + // Issue one resident chunk's aligned bulk on stage `stage` into the densely-packed block_tile at `off_bytes`, + // advancing `off_bytes` by the bytes written. Counterpart of `issue_stream_copy` for the resident window: the + // destination is the running packed cursor (vs a fixed slot) and there is no inflight bookkeeping (resident stages + // are driven directly, not through the stream ring). + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_resident_copy(int stage, offset_t chunk_index, int& off_bytes) + { + const auto src = chunk_bulk_src(layout.resident_base, chunk_index); + issue_bulk_copy(stage, key_slots + off_bytes, src); + off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + } + + // Issue one overflow chunk's aligned bulk on stage `stage` into its fixed streaming slot and mark the stage in + // flight. Counterpart of `issue_resident_copy` for the overflow window (fixed slot vs packed cursor; the stream ring + // records in-flight stages in `stream_inflight_mask`). + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_stream_copy(int stage, offset_t overflow_idx) { _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); _CCCL_ASSERT(stage >= 0 && stage < stream_stages, "overflow stage index exceeds the reserved streaming slots"); _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, "cannot issue a load into a streaming stage that is still in flight"); - const offset_t chunk_idx = layout.part.global_index(layout.overflow_base + overflow_idx); - const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); - // Every chunk begins on a `load_align` boundary, so the guard-free aligned (TMA bulk) path applies. The global- - // last chunk's unaligned suffix is always peeled into `edge_keys`, so streaming just its aligned bulk excludes - // it. For every interior chunk `bulk == count`. - const offset_t bulk = - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); - _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), - "overflow stream received a chunk with an unaligned start"); char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - const ::cuda::std::span src{ - layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; - issue_bulk_copy(stage, dst, src); + issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); } @@ -983,6 +1021,65 @@ private: static_cast(::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); } + // Chunk the first overflow wave assigns to `stage` (`stage < stream_stages`). The wave is primed in the compile-time + // `first_wave_is_forward` direction. Forward: stage `s` -> chunk `s`. Reverse: the last `stream_stages` chunks, each + // landing in stage `chunk % stream_stages`, so `s` takes the one in that window congruent to it. This + // `stage == chunk % stream_stages` invariant lets both the eager (interleaved) and post-loop primes address a + // stage's slot with no stored mapping. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t first_wave_chunk_for_stage(int stage) const + { + _CCCL_ASSERT(stage >= 0 && stage < stream_stages, "first-wave stage index out of range"); + _CCCL_ASSERT(static_cast(stream_stages) <= layout.overflow_chunks, + "first wave assumes at least `stream_stages` overflow chunks"); + // Only ever evaluated while priming the first wave, which `run` always issues in this preselected direction. + _CCCL_ASSERT(stream_is_forward == first_wave_is_forward, "first-wave prime ran in an unexpected stream direction"); + if constexpr (first_wave_is_forward) + { + return static_cast(stage); + } + else + { + // Perf note: runtime modulo by `stream_stages` (a runtime, not compile-time-constant, divisor; not host-known -- + // it derives from the per-CTA `stream_slots`). Cold: reached only while priming the first wave (`stream_stages` + // times, once per streaming CTA), feeding async TMA copies that mostly overlap sibling copies (fallback prime) or + // resident folds (interleaved prime). Lowest priority; flagged so a future bottleneck search finds every `/`,`%`. + const offset_t stages = static_cast(stream_stages); + const offset_t base = layout.overflow_chunks - stages; // window of the last `stream_stages` chunks + const offset_t chunk = base + (static_cast(stage) + stages - base % stages) % stages; + _CCCL_ASSERT(chunk % stages == static_cast(stage), + "reverse first-wave chunk must map back to its stage"); + return chunk; + } + } + + // Prime the streaming slots: issue this direction's first `stream_stages` overflow copies, arming the reload + // pipeline. Idempotent (`stream_is_primed`) and a no-op without overflow -- the first effective call per segment + // primes; later passes inherit the previous pass's resident turn-around chunks (exactly this direction's first + // `stream_stages` chunks, since the order ping-pongs) with no re-prime. The only path that actually issues here has + // no resident chunks (`load_and_histogram_first_pass` interleaves the prime whenever it has some) and already follows + // the kernel-init barrier, so the leading `__syncthreads()` is effectively redundant today; it is kept so the prime + // stays phase-safe if a future caller ever runs it right after a resident load (every thread must clear that load's + // final `wait_stage` before these copies re-arm the shared mbarriers, else the phase advances twice and a lagging + // thread spins forever). + _CCCL_DEVICE _CCCL_FORCEINLINE void prime_overflow_stream() + { + if (stream_is_primed || layout.overflow_chunks == 0) + { + return; + } + // Reaching here is the once-per-segment prime: the guard rules out a re-prime, and the resident load drives the + // shared stages via `issue_resident_copy` (never the stream mask), so the stream must be fully drained here. + _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, + "overflow stream priming must start from a fully drained stream"); + __syncthreads(); + // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. + for (int stage = 0; stage < stream_stages; ++stage) + { + issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + } + stream_is_primed = true; + } + // Consume the `i`-th visit on the block-load path (its ping-pong-ordered position is `overflow_idx`): wait for its // slot, fold its keys via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the just-freed // slot (a barrier orders the block's read ahead of the overwriting copy). Returns false once `should_continue()` @@ -990,10 +1087,14 @@ private: // the prefetches already in flight are drained after `run_pass`'s consume loops. template [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool - consume_overflow_visit(offset_t i, BlockApply&& block_apply, Continue&& should_continue) + consume_overflow_visit(offset_t i, int& stage, BlockApply&& block_apply, Continue&& should_continue) { const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); - const int stage = static_cast(overflow_idx % static_cast(stream_stages)); + // `stage` is a rolling ring pointer owned by the caller and stepped here for the next visit, replacing a per-visit + // `overflow_idx % stream_stages` (a runtime, non-power-of-two divisor). The assert pins it to that closed form so + // a later edit cannot desync the two. + _CCCL_ASSERT(stage == static_cast(overflow_idx % static_cast(stream_stages)), + "overflow stage desynced from its rolling ring pointer"); if (stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) { wait_stage(stage); @@ -1011,7 +1112,19 @@ private: { const offset_t next_overflow_idx = stream_is_forward ? next_step : (layout.overflow_chunks - 1 - next_step); __syncthreads(); - issue_load(stage, next_overflow_idx); + issue_stream_copy(stage, next_overflow_idx); + } + // Step the ring pointer one slot in the consume direction, wrapping with a compare-and-select (no modulo). + if (stream_is_forward) + { + if (++stage == stream_stages) + { + stage = 0; + } + } + else + { + stage = (stage == 0) ? stream_stages - 1 : stage - 1; } return true; } @@ -1032,8 +1145,13 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) { - _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, - "an overflow pass must begin with no copies in flight"); + // A pass begins drained, except the first, where `load_and_histogram_first_pass` early-primed the first wave (all + // `stream_stages` stages in flight) to overlap its edge work; phase 1 waits on them. Later passes re-enter drained: + // the reused turn-around chunks are resident, not in flight. + // 64-bit shift keeps a full 32-stage mask defined; the `uint32` mask promotes to match. + _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0} + || stream_inflight_mask == ((::cuda::std::uint64_t{1} << stream_stages) - ::cuda::std::uint64_t{1}), + "an overflow pass must begin drained or holding only the early-primed first wave"); if (layout.overflow_chunks == 0) { mid(); @@ -1042,31 +1160,26 @@ private: if constexpr (use_block_load_to_shared) { - // First ever call: prime the streaming slots. Subsequent calls inherit the previous pass's resident tail, which - // (because the order ping-pongs) is exactly the first `stream_stages` chunks of this direction. - if (!stream_is_primed) - { - // Wait for all threads to leave the resident load's final wait before re-arming its shared mbarriers; else - // the phase advances twice and a lagging thread misses the flip and spins forever. - __syncthreads(); - // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. - for (int i = 0; i < stream_stages; ++i) - { - const offset_t overflow_idx = - stream_is_forward ? static_cast(i) : (layout.overflow_chunks - 1 - static_cast(i)); - issue_load(static_cast(overflow_idx % static_cast(stream_stages)), overflow_idx); - } - stream_is_primed = true; - } - - // Phase 1: consume the first `stream_stages` visits (the chunks reused from the previous pass, already resident - // in the streaming slots), which issues the prefetch loads for this pass's reload wave into the freed slots. + // Idempotent prime; `load_and_histogram_first_pass` normally primes earlier, so this is usually a no-op. + prime_overflow_stream(); + + // Phase 1: consume the first `stream_stages` visits and issue this pass's reload wave into the freed slots. Their + // slots hold the early-primed wave (still in flight) on the first pass and the previous pass's reused turn-around + // chunks (already resident) on later passes; `consume_overflow_visit` waits only when the stage is in flight. + // Rolling ring pointer carried across both consume loops so no visit pays a runtime `% stream_stages`: forward + // starts at slot 0, reverse at the last chunk's slot (so this modulo runs only on reverse passes). That seed is + // the only runtime modulo left in these consume loops in release builds (the per-visit ones became the rolling + // step; a debug-only assert still evaluates one): a runtime, not compile-time-constant, divisor, at most once per + // pass -- cold (on pass 0 it overlaps the in-flight first wave; later passes enter drained, so it is a cheap + // one-off). Lowest priority. `consume_overflow_visit` then steps the pointer with compare-select. + int stage = + stream_is_forward ? 0 : static_cast((layout.overflow_chunks - 1) % static_cast(stream_stages)); bool is_stopped = false; const offset_t split = (::cuda::std::min) (static_cast(stream_stages), layout.overflow_chunks); _CCCL_PRAGMA_NOUNROLL() for (offset_t i = 0; i < split; ++i) { - if (!consume_overflow_visit(i, block_apply, should_continue)) + if (!consume_overflow_visit(i, stage, block_apply, should_continue)) { is_stopped = true; break; @@ -1084,7 +1197,7 @@ private: _CCCL_PRAGMA_NOUNROLL() for (offset_t i = split; i < layout.overflow_chunks; ++i) { - if (!consume_overflow_visit(i, block_apply, should_continue)) + if (!consume_overflow_visit(i, stage, block_apply, should_continue)) { break; } @@ -1155,8 +1268,8 @@ private: }); } - // Overload with no interleaved work, for the fused first pass where the resident keys are still being streamed in - // by the BlockLoadToShared pipeline (rather than already resident in SMEM). + // Overload with no interleaved (`mid`) work: the fused first pass folds its resident keys during the load itself, so + // pass 0 only folds the overflow chunks here and has no resident work left to overlap. template _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) { @@ -2167,25 +2280,6 @@ private: } } - // Aligned bulk of the resident chunk in `slot` (its `count` minus any peeled tail suffix); empty when it has none. - // Slot -> rank-local chunk index is `resident_base + slot` (identity unless `is_residency_reversed` shifts the - // resident window to the high-index chunks). - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span resident_bulk_src(offset_t slot) const - { - const offset_t chunk_idx = layout.part.global_index(layout.resident_base + slot); - const auto chunk = get_chunk(chunk_idx, layout.segment_size_off, layout.head_items); - const offset_t bulk = - ::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)); - if (bulk == 0) - { - return {}; - } - // Chunks start after the aligned head on an alignment-multiple stride (mirrors `issue_load`). - _CCCL_ASSERT(::cuda::is_aligned(layout.block_keys_base + chunk.offset, load_align_bytes), - "resident loader received a chunk with an unaligned start"); - return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; - } - // Fused first pass: load this rank's resident chunks into the block_tile, stage the persistent boundary edges into // `edge_keys`, and fold every key (resident + edges + overflow) into pass 0's histogram in the same sweep (pass 0 // needs no candidate filtering). Publishes the resident span as `resident_keys` for the later passes and the final @@ -2219,29 +2313,83 @@ private: // the unaligned suffix that is always peeled into `edge_keys`. const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); + // Prime the overflow stream while the resident load still runs to hide its latency, over two disjoint stage + // sets. The shared stages `[0, prologue)` the resident load cycles through are re-armed as they drain (tail + // loop below); the stages `[prologue, stream_stages)` it never touches (non-empty only when the resident region + // is smaller than the streaming reservation) are free from the start and primed up front. Together they cover + // every stream stage, so only the no-resident-chunks path leaves priming to the post-loop call. + const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; + // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. int next_off_bytes = 0; - // Load every resident chunk's aligned bulk, densely packed in slot order. + // Rotate the resident load's mbarrier-stage assignment so that, when active, the tail (the last `prologue` + // chunks, which have no resident successor) drains stages 0,1,...,prologue-1 in order (without the rotation it + // would start at `my_resident_chunks % prologue`). The stream's first wave reuses stages `[0, stream_stages)` + // as those drain, so this packs every stream prime into the earliest `stream_stages` tail folds -- maximizing + // overlap of the in-flight overflow copies with the remaining resident/edge work -- and leaves the slot-less + // idle folds (`stage >= stream_stages`) at the end. Only relabels which mbarrier a chunk uses (the packing is + // unchanged); `stage_rot` is 0 (identity) except when all of: there is overflow to prime, there are idle folds + // (`stream_stages < prologue`), the pipeline fully cycles (`prologue == PipelineStages`), and the tail is + // misaligned (`tail_align = my_resident_chunks % prologue != 0`) -- then it is `prologue - tail_align`. + // Perf note: `% prologue` is a runtime, non-power-of-two modulo. Critical path (one-time; seeds the resident + // stage rotation before the first resident copy) but TMA-dwarfed, with nothing earlier to overlap it; + // `prologue` is device-derived, so only removing the divide -- not relocation or a host `fast_mod_div` -- could + // help. Low priority. Flagged so a future bottleneck search finds every runtime `/`,`%`. + const offset_t tail_align = layout.my_resident_chunks % static_cast(prologue); + const offset_t stage_rot = + (interleave_overflow_prime && stream_stages < prologue && tail_align != offset_t{0}) + ? static_cast(prologue) - tail_align + : offset_t{0}; + + // Load every resident chunk's aligned bulk, densely packed in slot order. `stage` is a second, unit-stride + // rolling counter: it just increments and wraps to zero (a compare-and-select, no `+ stage_rot` and no runtime + // `% prologue`, which would be costly as `prologue` is not a compile-time power of two). The per-iteration + // assert pins it to the closed form `(local_chunk + stage_rot) % prologue` so a later edit cannot desync them. // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. - for (int stage = 0; stage < prologue; ++stage) { - const auto src = resident_bulk_src(static_cast(stage)); - issue_bulk_copy(stage, key_slots + next_off_bytes, src); - next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + int stage = static_cast(stage_rot); + for (offset_t local_chunk = 0; local_chunk < static_cast(prologue); ++local_chunk) + { + _CCCL_ASSERT(static_cast(stage) == (local_chunk + stage_rot) % static_cast(prologue), + "resident issue stage desynced from its rolling counter"); + issue_resident_copy(stage, local_chunk, next_off_bytes); + if (++stage == prologue) + { + stage = 0; + } + } + } + + // Prime the stream stages the resident load never uses (`[prologue, stream_stages)`; empty unless the resident + // region is smaller than the streaming reservation). Their mbarriers are fresh -- no resident drain to wait + // for -- so these copies start now, concurrent with the resident bulks. The shared stages `[0, prologue)` are + // primed as the resident load frees them in the tail loop below. + // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. + if (interleave_overflow_prime) + { + for (int stage = prologue; stage < stream_stages; ++stage) + { + issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + } } // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed - // and consumed in the same order), and `resident_bulk_src(local_chunk)` recomputes its length, so the read span - // needs no stored per-stage state. + // and consumed in the same order), and `chunk_bulk_src` recomputes its length, so the read span needs no + // stored per-stage state. int read_off_bytes = 0; + // Same unit-stride rolling counter as the issue loop: increment and select-zero on wrap, no runtime + // `% prologue`. Starts at `stage_rot`; the assert pins it to `(local_chunk + stage_rot) % prologue`. + int stage = static_cast(stage_rot); _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - const int stage = static_cast(local_chunk % static_cast(prologue)); + _CCCL_ASSERT(static_cast(stage) == (local_chunk + stage_rot) % static_cast(prologue), + "resident read stage desynced from its rolling counter"); wait_stage(stage); - const int read_len_items = static_cast(::cuda::std::size(resident_bulk_src(local_chunk))); + const int read_len_items = + static_cast(::cuda::std::size(chunk_bulk_src(layout.resident_base, local_chunk))); for_each_chunk_key( ::cuda::ptr_rebind(key_slots + read_off_bytes), read_len_items, add_first_pass); read_off_bytes += read_len_items * int{sizeof(key_t)}; @@ -2249,12 +2397,24 @@ private: const offset_t next_local_chunk = local_chunk + static_cast(prologue); if (next_local_chunk < layout.my_resident_chunks) { - const auto src = resident_bulk_src(next_local_chunk); // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads // leave the wait above would advance the phase twice, stranding a lagging waiter forever. __syncthreads(); - issue_bulk_copy(stage, key_slots + next_off_bytes, src); - next_off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + issue_resident_copy(stage, next_local_chunk, next_off_bytes); + } + else if (interleave_overflow_prime && stage < stream_stages) + { + // Shared stage `[0, prologue)` with no resident successor: re-arm it with its first-wave chunk + // (`first_wave_chunk_for_stage`, direction-aware). Same phase-safety barrier as the resident re-arm; the + // copy targets the disjoint stream region (`stream_slot_base` onward), never the densely-packed resident + // data the later tail folds still read. Stages `>= stream_stages` are extra resident pipeline depth with no + // stream slot, so they idle -- `stage_rot` places them last so all primes issue first. + __syncthreads(); + issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + } + if (++stage == prologue) + { + stage = 0; } } @@ -2264,6 +2424,40 @@ private: // The resident region is one contiguous run of aligned bulks for the later passes; both boundary edges are // folded separately from `edge_keys`. resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), next_off_bytes / int{sizeof(key_t)}); + + // The up-front and tail primes together armed every stream stage, so the stream is primed; the post-loop call + // no-ops. + if (interleave_overflow_prime) + { + _CCCL_ASSERT(stream_inflight_mask == ((::cuda::std::uint64_t{1} << stream_stages) - ::cuda::std::uint64_t{1}), + "interleaved prime must arm exactly the first wave across all stream stages"); + stream_is_primed = true; + } + } + + // Prime the first overflow wave (a no-op when the resident load above already interleaved it; still primes the + // no-resident-chunks path) so the edge staging/folding below -- independent of the streaming slots -- runs under + // the in-flight TMA copies. + prime_overflow_stream(); + + // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see + // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix + // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. + if (layout.head_edge_len_items > 0) + { + stage_and_fold_edge(temp_storage.edge_keys, layout.block_keys_base, layout.head_edge_len_items, add_first_pass); + } + if (layout.tail_edge_len_items > 0) + { + _CCCL_ASSERT(layout.chunks > offset_t{0}, "a peeled tail edge requires at least one aligned chunk"); + const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_off, layout.head_items); + _CCCL_ASSERT(layout.tail_edge_len_items <= tail_chunk.count, + "peeled tail length cannot exceed its source chunk"); + stage_and_fold_edge( + temp_storage.edge_keys + head_edge_cap_items, + layout.block_keys_base + tail_chunk.offset + (tail_chunk.count - layout.tail_edge_len_items), + layout.tail_edge_len_items, + add_first_pass); } } else @@ -2290,33 +2484,10 @@ private: } } - // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see - // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix - // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. - if constexpr (use_block_load_to_shared) - { - if (layout.head_edge_len_items > 0) - { - stage_and_fold_edge(temp_storage.edge_keys, layout.block_keys_base, layout.head_edge_len_items, add_first_pass); - } - if (layout.tail_edge_len_items > 0) - { - _CCCL_ASSERT(layout.chunks > offset_t{0}, "a peeled tail edge requires at least one aligned chunk"); - const auto tail_chunk = get_chunk(layout.chunks - offset_t{1}, layout.segment_size_off, layout.head_items); - _CCCL_ASSERT(layout.tail_edge_len_items <= tail_chunk.count, - "peeled tail length cannot exceed its source chunk"); - stage_and_fold_edge( - temp_storage.edge_keys + head_edge_cap_items, - layout.block_keys_base + tail_chunk.offset + (tail_chunk.count - layout.tail_edge_len_items), - layout.tail_edge_len_items, - add_first_pass); - } - } - - // Fold the overflow chunks into the first-pass histogram, priming the streaming slots in the stream's initial - // direction (preselected above; the histogram is order-independent, so the direction only sets up the leftover - // parity for the final filter). The overflow stream reuses the resident load's stage mbarriers (all front-loaded at - // `run` entry); `wait_stage` provides the producer/consumer sync. + // Fold the overflow chunks into the first-pass histogram. On the block-load path the stream was primed above (in + // its preselected initial direction; the histogram is order-independent, so the direction only sets up the leftover + // parity for the final filter) and reuses the resident load's stage mbarriers (all front-loaded at `run` entry), + // with `wait_stage` for producer/consumer sync; the generic fallback re-reads each overflow chunk from gmem. process_pass(add_first_pass); const int resident_count = static_cast(resident_keys.size()); @@ -2774,7 +2945,6 @@ private: const offset_t cluster_tid = hw_cluster_rank * static_cast(threads_per_block) + threadIdx.x; const offset_t cluster_threads = hw_cluster_blocks * static_cast(threads_per_block); const offset_t step = cluster_threads * static_cast(copy_items); - const offset_t full_tiles = ::cuda::round_down(num_items, step); auto keys_in_it = d_key_segments_it[segment_id]; auto keys_out_it = d_key_segments_out_it[segment_id]; // Per-segment value iterators (pairs only), hoisted once like the key iterators and reused by both loops. For @@ -2801,8 +2971,12 @@ private: } }(); + // Full tiles: advance while a whole `step`-sized tile still fits (a grid-stride bound, not a precomputed count), so + // `base` ends at `round_down(num_items, step)` with no division. Each tile batches unrolled loads then stores for + // memory-level parallelism; the remainder below is a plain grid-stride tail. + offset_t base = 0; _CCCL_PRAGMA_NOUNROLL() - for (offset_t base = 0; base < full_tiles; base += step) + for (; base + step <= num_items; base += step) { offset_t idx[copy_items]; key_t keys[copy_items]; @@ -2840,13 +3014,13 @@ private: } } - // Sub-tile remainder. Iterate the tail relative to zero and offset by `full_tiles`, so every index stays in - // `[0, num_items)` without relying on `full_tiles + cluster_tid` staying within `offset_t`. - const offset_t tail_items = num_items - full_tiles; + // Sub-tile remainder. Iterate the tail relative to zero and offset by `base`, so every index stays in + // `[0, num_items)` without relying on `base + cluster_tid` staying within `offset_t`. + const offset_t tail_items = num_items - base; // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. for (offset_t local = cluster_tid; local < tail_items; local += cluster_threads) { - const auto seg_idx = static_cast(full_tiles + local); + const auto seg_idx = static_cast(base + local); keys_out_it[seg_idx] = keys_in_it[seg_idx]; if constexpr (!is_keys_only) { @@ -2897,7 +3071,12 @@ private: const unsigned hw_cluster_blocks = ::cuda::ptx::get_sreg_cluster_nctarank(); _CCCL_ASSERT(hw_cluster_blocks > 0u && hw_cluster_rank < hw_cluster_blocks, "hardware cluster rank must lie within a non-empty cluster"); - segment_id = static_cast(blockIdx.x / hw_cluster_blocks); + // Segment id is the cluster's grid coordinate read straight from `%clusterid`, not derived as + // `blockIdx.x / hw_cluster_blocks`: the hardware value needs no division and stays correct even if a launch ever + // used non-uniform cluster sizes (the derived form assumes every cluster holds `hw_cluster_blocks` CTAs). The + // dispatch launches one cluster per segment (`gridDim.x = num_segments * clusterDim.x`), so `clusterid.x` is + // exactly the segment id. + segment_id = static_cast(::cuda::ptx::get_sreg_clusterid_x()); if (segment_id >= detail::params::get_param(num_segments, num_segments_val_t{0})) { From d7fe9c21129a10d850fc8909f42d0e6ebb5329a7 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 13 Jul 2026 15:49:38 +0200 Subject: [PATCH 114/126] Hopefully revert perf regression from previous commit --- cub/cub/agent/agent_batched_topk_cluster.cuh | 103 +++++++++++-------- 1 file changed, 60 insertions(+), 43 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 33525fd08b7..6032d5c24d9 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -2314,62 +2314,76 @@ private: const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); // Prime the overflow stream while the resident load still runs to hide its latency, over two disjoint stage - // sets. The shared stages `[0, prologue)` the resident load cycles through are re-armed as they drain (tail - // loop below); the stages `[prologue, stream_stages)` it never touches (non-empty only when the resident region - // is smaller than the streaming reservation) are free from the start and primed up front. Together they cover - // every stream stage, so only the no-resident-chunks path leaves priming to the post-loop call. + // sets: the `prologue` mbarrier stages the resident load cycles through (re-armed for the stream as they drain + // in the tail loop below) and the stages it never touches (fresh, primed up front). Together they cover every + // stream stage, so only the no-resident-chunks path leaves priming to the post-loop call. const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. int next_off_bytes = 0; - // Rotate the resident load's mbarrier-stage assignment so that, when active, the tail (the last `prologue` - // chunks, which have no resident successor) drains stages 0,1,...,prologue-1 in order (without the rotation it - // would start at `my_resident_chunks % prologue`). The stream's first wave reuses stages `[0, stream_stages)` - // as those drain, so this packs every stream prime into the earliest `stream_stages` tail folds -- maximizing - // overlap of the in-flight overflow copies with the remaining resident/edge work -- and leaves the slot-less - // idle folds (`stage >= stream_stages`) at the end. Only relabels which mbarrier a chunk uses (the packing is - // unchanged); `stage_rot` is 0 (identity) except when all of: there is overflow to prime, there are idle folds - // (`stream_stages < prologue`), the pipeline fully cycles (`prologue == PipelineStages`), and the tail is - // misaligned (`tail_align = my_resident_chunks % prologue != 0`) -- then it is `prologue - tail_align`. - // Perf note: `% prologue` is a runtime, non-power-of-two modulo. Critical path (one-time; seeds the resident - // stage rotation before the first resident copy) but TMA-dwarfed, with nothing earlier to overlap it; - // `prologue` is device-derived, so only removing the divide -- not relocation or a host `fast_mod_div` -- could - // help. Low priority. Flagged so a future bottleneck search finds every runtime `/`,`%`. + // Place the resident load's `prologue`-wide mbarrier window so the stream's first wave is primed in forward + // consume order (stage `s`, holding chunk `s`, first), removing the ordering inversion where the up-front prime + // would otherwise issue the last-consumed chunks first. Two complementary knobs, at most one non-zero + // (`stream_stages > prologue` vs `<=`); both only relabel which mbarrier a chunk uses (the packing and the + // stream ring invariant `stage == chunk % stream_stages` are untouched): + // * `stage_base` shifts the window to `[stage_base, stage_base + prologue)`, leaving the low stream stages + // `[0, stage_base)` free to prime up front. Forward first wave with `stream_stages > prologue` only (which + // forces `prologue == my_resident_chunks`, so there is no reload). Reverse keeps it 0 (its first-consumed + // stage is data-dependent, so no static low-stage reservation helps). + // * `stage_rot` rotates the resident cycle (mod `prologue`) so a misaligned reload tail still drains stages + // 0,1,...,prologue-1 in order. Only when the stream stages fit the resident cycle (`stream_stages <= + // prologue`) and the tail is misaligned (`my_resident_chunks % prologue != 0`, implying `prologue == + // PipelineStages`). Its `% prologue` is a one-time runtime, non-power-of-two modulo, TMA-dwarfed. + int stage_base = 0; + if constexpr (first_wave_is_forward) + { + if (interleave_overflow_prime && stream_stages > prologue) + { + stage_base = stream_stages - prologue; + } + } const offset_t tail_align = layout.my_resident_chunks % static_cast(prologue); const offset_t stage_rot = - (interleave_overflow_prime && stream_stages < prologue && tail_align != offset_t{0}) + (interleave_overflow_prime && stream_stages <= prologue && tail_align != offset_t{0}) ? static_cast(prologue) - tail_align : offset_t{0}; // Load every resident chunk's aligned bulk, densely packed in slot order. `stage` is a second, unit-stride - // rolling counter: it just increments and wraps to zero (a compare-and-select, no `+ stage_rot` and no runtime - // `% prologue`, which would be costly as `prologue` is not a compile-time power of two). The per-iteration - // assert pins it to the closed form `(local_chunk + stage_rot) % prologue` so a later edit cannot desync them. + // rolling counter over `[stage_base, stage_base + prologue)`: it increments and select-wraps back to + // `stage_base` (no `+ stage_rot` and no runtime `% prologue`, costly as `prologue` is not a power of two). The + // per-iteration assert pins it to the closed form so a later edit cannot desync them. // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. { - int stage = static_cast(stage_rot); + int stage = stage_base + static_cast(stage_rot); for (offset_t local_chunk = 0; local_chunk < static_cast(prologue); ++local_chunk) { - _CCCL_ASSERT(static_cast(stage) == (local_chunk + stage_rot) % static_cast(prologue), - "resident issue stage desynced from its rolling counter"); + _CCCL_ASSERT( + stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), + "resident issue stage desynced from its rolling counter"); issue_resident_copy(stage, local_chunk, next_off_bytes); - if (++stage == prologue) + if (++stage == stage_base + prologue) { - stage = 0; + stage = stage_base; } } } - // Prime the stream stages the resident load never uses (`[prologue, stream_stages)`; empty unless the resident - // region is smaller than the streaming reservation). Their mbarriers are fresh -- no resident drain to wait - // for -- so these copies start now, concurrent with the resident bulks. The shared stages `[0, prologue)` are - // primed as the resident load frees them in the tail loop below. + // Prime the stream stages the resident load never uses -- the complement of `[stage_base, stage_base + + // prologue)` within `[0, stream_stages)`. Their mbarriers are fresh (no resident drain to wait for), so these + // copies start now, concurrent with the resident bulks; the shared resident stages are primed as the tail frees + // them (read loop below). Forward the free set is the low `[0, stage_base)` (first-consumed chunks); reverse it + // is the high `[prologue, stream_stages)`; at most one loop is non-empty (both are empty when the whole stream + // reservation overlaps the resident window, i.e. `stage_base == 0 && stream_stages <= prologue`). // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. if (interleave_overflow_prime) { - for (int stage = prologue; stage < stream_stages; ++stage) + for (int stage = 0; stage < stage_base; ++stage) + { + issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + } + for (int stage = stage_base + prologue; stage < stream_stages; ++stage) { issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); } @@ -2379,14 +2393,16 @@ private: // and consumed in the same order), and `chunk_bulk_src` recomputes its length, so the read span needs no // stored per-stage state. int read_off_bytes = 0; - // Same unit-stride rolling counter as the issue loop: increment and select-zero on wrap, no runtime - // `% prologue`. Starts at `stage_rot`; the assert pins it to `(local_chunk + stage_rot) % prologue`. - int stage = static_cast(stage_rot); + // Same unit-stride rolling counter as the issue loop over `[stage_base, stage_base + prologue)`: increment and + // select-wrap to `stage_base`, no runtime `% prologue`. Starts at `stage_base + stage_rot`; the assert pins it + // to the closed form. + int stage = stage_base + static_cast(stage_rot); _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) { - _CCCL_ASSERT(static_cast(stage) == (local_chunk + stage_rot) % static_cast(prologue), - "resident read stage desynced from its rolling counter"); + _CCCL_ASSERT( + stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), + "resident read stage desynced from its rolling counter"); wait_stage(stage); const int read_len_items = static_cast(::cuda::std::size(chunk_bulk_src(layout.resident_base, local_chunk))); @@ -2404,17 +2420,18 @@ private: } else if (interleave_overflow_prime && stage < stream_stages) { - // Shared stage `[0, prologue)` with no resident successor: re-arm it with its first-wave chunk - // (`first_wave_chunk_for_stage`, direction-aware). Same phase-safety barrier as the resident re-arm; the - // copy targets the disjoint stream region (`stream_slot_base` onward), never the densely-packed resident - // data the later tail folds still read. Stages `>= stream_stages` are extra resident pipeline depth with no - // stream slot, so they idle -- `stage_rot` places them last so all primes issue first. + // Shared resident stage (in `[stage_base, stage_base + prologue)`) with no resident successor: re-arm it + // with its first-wave chunk (`first_wave_chunk_for_stage`, direction-aware). Same phase-safety barrier as + // the resident re-arm; the copy targets the disjoint stream region (`stream_slot_base` onward), never the + // densely-packed resident data the later tail folds still read. Stages `>= stream_stages` are extra + // resident pipeline depth with no stream slot, so they idle -- `stage_rot` places them last so all primes + // issue first. __syncthreads(); issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); } - if (++stage == prologue) + if (++stage == stage_base + prologue) { - stage = 0; + stage = stage_base; } } From c9f087836aff04750f550afe8a6ec93fa9f0d93c Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 13 Jul 2026 17:40:49 +0200 Subject: [PATCH 115/126] Improve testing Add two tuning parameters that artificially restric the amount of CTAs per cluster and the amount of chunk slots and therefore dynamic shared memory per CTA. This allows us to test specific scenarios/code-paths deterministically without relying on assumptions or duplication of the dynamic dispatch. While we will keep these parameters out of tuning for now, they could be interesting for users once the tuning becomes public. E.g. a user might have a concurrent kernel that needs a specific smem/L1 carveout for the kernel to run in parallel. --- .../variable/indexed_common.cuh | 4 +- .../segmented_topk/variable/keys_common.cuh | 4 +- .../device/dispatch/dispatch_batched_topk.cuh | 114 +++- .../dispatch/kernels/kernel_batched_topk.cuh | 8 +- .../dispatch/tuning/tuning_batched_topk.cuh | 36 +- cub/cub/util_device.cuh | 3 - .../catch2_test_device_segmented_topk_keys.cu | 578 ++++++++++++++---- ...catch2_test_device_segmented_topk_pairs.cu | 521 +++++++++++++--- cub/test/catch2_test_device_topk_common.cuh | 60 ++ 9 files changed, 1094 insertions(+), 234 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh index 2e92c2958a8..86ae3d3c73d 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -137,7 +137,9 @@ struct topk_backend_selector /*bits_per_pass=*/TUNE_CLUSTER_BITS_PER_PASS, /*histogram_items_per_thread=*/TUNE_CLUSTER_HIST_IPT, /*tie_break_items_per_thread=*/TUNE_CLUSTER_TIEBREAK_IPT, - /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT}; + /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; #else const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc); #endif diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh index ce5f97408c3..7237e70fe11 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -144,7 +144,9 @@ struct topk_backend_selector /*bits_per_pass=*/TUNE_CLUSTER_BITS_PER_PASS, /*histogram_items_per_thread=*/TUNE_CLUSTER_HIST_IPT, /*tie_break_items_per_thread=*/TUNE_CLUSTER_TIEBREAK_IPT, - /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT}; + /*copy_items_per_thread=*/TUNE_CLUSTER_COPY_IPT, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; #else const auto cluster = cub::detail::batched_topk::cluster_policy_selector{}(cc); #endif diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 9969361b039..e88febbccaa 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -419,6 +419,7 @@ template SingleBlockMaxSegSize, \ MinChunksPerBlock, \ CopyItemsPerThread, \ + cdp_cluster_blocks, \ Determinism, \ TieBreak, \ KeyInputItItT, \ @@ -493,7 +494,17 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( constexpr int TieBreakItemsPerThread = policy.tie_break_items_per_thread; constexpr int SingleBlockMaxSegSize = policy.single_block_max_seg_size; constexpr int MinChunksPerBlock = policy.min_chunks_per_block; + constexpr int MaxBlocksPerCluster = policy.max_blocks_per_cluster; + constexpr int MaxChunkSlotsPerBlock = policy.max_chunk_slots_per_block; constexpr int CopyItemsPerThread = policy.copy_items_per_thread; + static_assert(MaxBlocksPerCluster >= 0, + "max_blocks_per_cluster must be 0 (unrestricted) or a positive cluster width"); + static_assert(MaxChunkSlotsPerBlock >= 0, "max_chunk_slots_per_block must be 0 (unrestricted) or a positive count"); + + // The effective launched-cluster-width cap (`cluster_cap`) folds `MaxBlocksPerCluster` (0 = unrestricted) into the + // hardware ceiling. The host arm derives that ceiling from the runtime (below) rather than a compile-time constant, + // so a future device that supports wider (non-portable) clusters is not artificially limited; the CDP arm is + // compile-time and bounded by the portable ceiling (device launches cannot opt into non-portable clusters). using key_it_t = it_value_t; using key_t = it_value_t; @@ -624,7 +635,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( // Resolve the per-block opt-in shared-memory budget and the kernel's static footprint from the driver so // the dynamic-SMEM math below matches exactly what the launch permits. The opt-in budget // (`cudaDevAttrMaxSharedMemoryPerBlockOptin`) is the documented total per-block budget; the usable dynamic - // portion (`max_dynamic_smem_bytes`) is that budget minus the static footprint. + // portion (`hw_dynamic_smem_bytes`) is that budget minus the static footprint (the policy slot cap may narrow it + // further into `max_dynamic_smem_bytes`). int device_id = 0; if (const auto error = CubDebug(cudaGetDevice(&device_id))) { @@ -648,8 +660,19 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( // (opt-in == per-SM - reserved), so the dynamic budget is just the opt-in budget minus the static footprint; // reserved must not be subtracted a second time. const int nondynamic_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); - const int max_dynamic_smem_bytes = + const int hw_dynamic_smem_bytes = (max_smem_optin_bytes > nondynamic_smem_bytes) ? max_smem_optin_bytes - nondynamic_smem_bytes : 0; + // Optional policy cap on resident chunk slots per block (`max_chunk_slots_per_block == 0` -> unrestricted, i.e. + // the full hardware budget). Expressed as the SMEM those slots need (`base_padding + slots * chunk_bytes`), then + // clamped to the hardware budget: a cap the hardware cannot satisfy is a no-op (hardware wins). Fewer slots + // lowers every CTA's resident dynamic shared-memory request, so a smaller segment overflows into streaming -- + // useful to leave shared memory free for a concurrent kernel (or to reach the streaming / schedule paths at a + // small footprint in tests). A cap below one slot trips the `max_block_tile_capacity <= 0` guard below. + const int max_dynamic_smem_bytes = + (MaxChunkSlotsPerBlock == 0) + ? hw_dynamic_smem_bytes + : (::cuda::std::min) (hw_dynamic_smem_bytes, + layout_t::base_padding_bytes + MaxChunkSlotsPerBlock * layout_t::chunk_bytes); // Raise the kernel's dynamic-SMEM opt-in lazily: occupancy queries and the launch must not request more than the // currently configured `cudaFuncAttributeMaxDynamicSharedMemorySize`. The kernel's compile-time default already @@ -674,6 +697,29 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return cudaSuccess; }; + // Hardware cluster-width ceiling, taken from the runtime rather than a compile-time constant so a future device + // with wider (non-portable) clusters is not artificially capped (the same reason the dynamic-SMEM budget above is + // queried, not hardcoded). Probed at zero dynamic SMEM to get the architectural/kernel ceiling independent of the + // per-`C` SMEM chosen below; each candidate width is still validated against its own SMEM by the resident scan's + // `max_active_clusters` probe, and the oversize branch re-probes at its (larger) launch SMEM. Requires the + // non-portable opt-in already set above so the probe can report widths beyond the portable ceiling. + cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` + cfg.gridDim = dim3(1, 1, 1); + cfg.dynamicSmemBytes = 0; + int hw_cluster_ceiling = 0; + if (const auto error = launcher_factory.max_potential_cluster_size(hw_cluster_ceiling, dynamic_kernel, &cfg)) + { + return error; + } + if (hw_cluster_ceiling <= 0) + { + return cudaErrorInvalidValue; + } + // `MaxBlocksPerCluster == 0` -> the full hardware ceiling; a non-zero knob narrows it, clamped to that ceiling. A + // cap narrower than a segment needs pushes it into the oversize/streaming fallback below (cap 1 -> single-CTA). + const int cluster_cap = + (MaxBlocksPerCluster == 0) ? hw_cluster_ceiling : (::cuda::std::min) (MaxBlocksPerCluster, hw_cluster_ceiling); + // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We @@ -695,21 +741,18 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return layout_t::base_padding_bytes + static_cast(slots) * layout_t::chunk_bytes; }; - // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (cap HW - // max). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest fully-resident `C`. - // Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. + // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (capped at + // `cluster_cap`). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest resident + // `C`. Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. // `numeric_limits::max()` for an unbounded deferred sequence); narrowing such a `C_lo` to `int` could wrap to // a small (or negative) value and wrongly enter the resident branch instead of the oversize/streaming fallback. const int c_full = static_cast( - (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks), - ::cuda::ceil_div(seg, chunk_items_u64))); + (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(cluster_cap), ::cuda::ceil_div(seg, chunk_items_u64))); const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. const int desired_cluster_blocks = ::cuda::narrow(batched_topk_cluster::effective_cluster_blocks_from_chunks( - ::cuda::ceil_div(seg, chunk_items_u64), - MinChunksPerBlock, - ::cuda::narrow(max_supported_cluster_blocks))); + ::cuda::ceil_div(seg, chunk_items_u64), MinChunksPerBlock, ::cuda::narrow(cluster_cap))); int cluster_blocks = 0; int dynamic_smem_sel = 0; @@ -724,15 +767,20 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( cluster_blocks = 1; dynamic_smem_sel = smem_for_block_capacity(seg); } - else if (c_lo <= static_cast<::cuda::std::uint64_t>(max_supported_cluster_blocks)) + else if (c_lo <= static_cast<::cuda::std::uint64_t>(cluster_cap)) { - // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= HW max`, so every + // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= cluster_cap`, so every // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. - // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. `C = 1` is handled above. The - // upper bound is capped at the cluster blocks the max segment needs (`desired_cluster_blocks`), so the host - // never launches a wider cluster than necessary; at `min_chunks_per_block == 1` the cap equals `c_full`. + // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. The upper bound is capped at the + // cluster blocks the max segment needs (`desired_cluster_blocks`), so the host never launches a wider cluster + // than necessary; at `min_chunks_per_block == 1` the cap equals `c_full`. `C_end` is additionally clamped to + // `cluster_cap`, which empties the range only at `cluster_cap == 1` (`c_begin == 2 > 1`) -- reachable here for + // a one-CTA-resident segment whose single-CTA fast path is disabled. Then `c_lo == 1` and the fallback below + // takes the single CTA, so the width cap is never exceeded. const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); - const int c_end = (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks)); + const int c_end = + (::cuda::std::min) (cluster_cap, + (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks))); ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); for (int c = c_begin; c <= c_end; ++c) { @@ -786,7 +834,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( if (cluster_blocks == 0) { - // Oversize (`C_lo > HW max`) or nothing launchable in range: full residency is impossible, so maximize + // Oversize (`C_lo > cluster_cap`) or nothing launchable in range: full residency is impossible, so maximize // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) { @@ -800,7 +848,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( { return error; } - hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, max_supported_cluster_blocks); + hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, cluster_cap); if (hw_max_cluster_blocks <= 0) { return cudaErrorInvalidValue; @@ -851,11 +899,31 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( } }), ({ - // CDP path: device-side launches cannot opt in to more than portable total SMEM or non-portable cluster blocks. - // Segments exceeding the portable resident coverage are still handled: the agent re-streams overflow from gmem. + // CDP path: device-side launches cannot opt in to more than portable total SMEM or non-portable cluster blocks, + // so both geometry knobs start from the portable ceiling and are only ever narrowed. Segments exceeding the + // portable resident coverage are still handled: the agent re-streams overflow from gmem. + // + // Cluster width: the portable ceiling (device launches cannot opt into non-portable clusters, so unlike the host + // arm there is no runtime ceiling to query), narrowed when the policy's `max_blocks_per_cluster` knob is tighter + // (a knob above the portable ceiling is naturally a no-op here). Compile-time because the static kernel's + // `__cluster_dims__` is a compile-time attribute. + constexpr int cdp_cluster_blocks = + (MaxBlocksPerCluster == 0) + ? max_portable_cluster_blocks + : (::cuda::std::min) (MaxBlocksPerCluster, max_portable_cluster_blocks); + static_assert(cdp_cluster_blocks >= 1, "device-launch (CDP) cluster width must be at least one CTA"); + + // Dynamic SMEM: portable budget (total minus static footprint), further narrowed by the policy's + // `max_chunk_slots_per_block` cap (same slots -> bytes conversion as the host arm; a cap above the portable + // budget is a no-op). Fewer resident slots just forces the agent to stream more from gmem. constexpr int portable_total_smem_bytes = 48 * 1024; - constexpr int dynamic_smem_bytes = + constexpr int portable_dynamic_smem_bytes = (portable_total_smem_bytes > static_smem_bytes) ? portable_total_smem_bytes - static_smem_bytes : 0; + constexpr int dynamic_smem_bytes = + (MaxChunkSlotsPerBlock == 0) + ? portable_dynamic_smem_bytes + : (::cuda::std::min) (portable_dynamic_smem_bytes, + layout_t::base_padding_bytes + MaxChunkSlotsPerBlock * layout_t::chunk_bytes); // The compile-time `ChunkBytes` is reused verbatim; the agent peels unaligned boundary edges into a tiny // per-block buffer and re-streams overflow from gmem, so the only hard requirement is that one load-aligned chunk @@ -864,8 +932,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( static_assert(block_tile_capacity >= static_cast<::cuda::std::uint32_t>(layout_t::chunk_items), "Portable SMEM is too small to fit even one load-aligned chunk for the device-launch (CDP) path"); - const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) - * static_cast<::cuda::std::uint64_t>(max_portable_cluster_blocks); + const auto grid_blocks = + static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cdp_cluster_blocks); if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) { return cudaErrorInvalidValue; diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 35fc9cfa253..fc109ff8bc7 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -20,7 +20,7 @@ #include #include #include -#include // max_portable_cluster_blocks (CDP static-cluster kernel) +#include #include #include @@ -447,6 +447,9 @@ __launch_bounds__( // `cudaFuncSetAttribute` (device-side launches cannot opt into dynamic cluster dimensions the way the host // `device_batched_topk_kernel` does). It mirrors that kernel's cluster arm with a fixed cluster dim, and is consumed by // the CDP arm of `launch_cluster_arm` (see `CUB_TOPK_CLUSTER_DEVICE_LAUNCH` in dispatch_batched_topk.cuh). +// `ClusterBlocks` is the compile-time cluster width the CDP arm picks: `<= max_portable_cluster_blocks` (device +// launches can't opt into non-portable cluster sizes), and further narrowed when the policy's `max_blocks_per_cluster` +// cap is tighter. template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(max_portable_cluster_blocks, 1, 1) +__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterBlocks, 1, 1) _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, diff --git a/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh b/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh index 540c6a4ad8a..2d881015837 100644 --- a/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh @@ -224,9 +224,10 @@ struct baseline_policy_selector_from_types static_assert(baseline_topk_policy_selector); #endif // _CCCL_HAS_CONCEPTS() -//! Per-block execution shape for the thread-block-cluster backend of @ref DeviceBatchedTopK. The dispatch picks the -//! number of cluster blocks and the dynamic shared-memory block_tile capacity at runtime (occupancy / wave-aware), so -//! this policy carries only the per-block tuning knobs. +//! Execution shape for the thread-block-cluster backend of @ref DeviceBatchedTopK. The dispatch picks the number of +//! cluster blocks and the dynamic shared-memory block_tile capacity at runtime (occupancy / wave-aware), so this policy +//! mostly carries per-block tuning knobs; the two trailing `max_*` fields are optional launch-geometry caps that bound +//! that runtime choice. struct cluster_topk_policy { // Fields grouped by kind and, across and within groups, ordered by kernel use: launch config, then load, then @@ -253,6 +254,23 @@ struct cluster_topk_policy int tie_break_items_per_thread; //!< Keys each thread processes per tile during the final tie-break / filter phase. int copy_items_per_thread; //!< Keys each thread copies per tile on the select-all (k >= segment size) fast path. + // Launch-geometry caps that bound the otherwise heuristic / hardware-derived cluster width and resident shared-memory + // footprint. Both default to 0 (= unrestricted) and are deliberately not auto-tuned. Beyond deterministically + // steering tests onto the streaming / cluster paths at a small footprint, they let a caller trade top-k throughput + // for resources it wants to leave free -- e.g. capping resident slots to fit a shared-memory carveout reserved for a + // concurrently running kernel, or narrowing the cluster width to co-schedule other work. The algorithm stays correct + // at any cap: a segment that no longer fits resident simply streams the remainder from global memory. + int max_blocks_per_cluster; //!< Upper bound on the launched cluster width (CTAs per segment); 0 = unrestricted (the + //!< hardware cluster-width ceiling, queried from the runtime for a host launch and the + //!< portable ceiling for a device (CDP) launch). Non-zero is additionally clamped to that + //!< same ceiling. A cap narrower than a segment needs pushes it into the streaming + //!< fallback (cap 1 -> single-CTA streaming). + int max_chunk_slots_per_block; //!< Upper bound on resident chunk slots per block; 0 = unrestricted (the full + //!< shared-memory budget: the hardware opt-in budget for a host launch, the portable + //!< 48 KiB budget for a device (CDP) launch). A smaller cap shrinks each CTA's + //!< resident capacity (and thus its dynamic shared-memory request), so a smaller + //!< segment overflows into the streaming path. + // Equality/streaming make this a regular type (required by the `policy_selector` concept / `dispatch_compute_cap`). _CCCL_HOST_DEVICE_API constexpr friend bool operator==(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) { @@ -262,7 +280,9 @@ struct cluster_topk_policy && lhs.single_block_max_seg_size == rhs.single_block_max_seg_size && lhs.bits_per_pass == rhs.bits_per_pass && lhs.histogram_items_per_thread == rhs.histogram_items_per_thread && lhs.tie_break_items_per_thread == rhs.tie_break_items_per_thread - && lhs.copy_items_per_thread == rhs.copy_items_per_thread; + && lhs.copy_items_per_thread == rhs.copy_items_per_thread + && lhs.max_blocks_per_cluster == rhs.max_blocks_per_cluster + && lhs.max_chunk_slots_per_block == rhs.max_chunk_slots_per_block; } _CCCL_HOST_DEVICE_API constexpr friend bool operator!=(const cluster_topk_policy& lhs, const cluster_topk_policy& rhs) @@ -280,7 +300,9 @@ struct cluster_topk_policy << ", .pipeline_stages = " << p.pipeline_stages << ", .single_block_max_seg_size = " << p.single_block_max_seg_size << ", .bits_per_pass = " << p.bits_per_pass << ", .histogram_items_per_thread = " << p.histogram_items_per_thread << ", .tie_break_items_per_thread = " - << p.tie_break_items_per_thread << ", .copy_items_per_thread = " << p.copy_items_per_thread << " }"; + << p.tie_break_items_per_thread << ", .copy_items_per_thread = " << p.copy_items_per_thread + << ", .max_blocks_per_cluster = " << p.max_blocks_per_cluster + << ", .max_chunk_slots_per_block = " << p.max_chunk_slots_per_block << " }"; } #endif // _CCCL_HOSTED() }; @@ -305,7 +327,9 @@ concept cluster_topk_policy_selector = policy_selector; /*bits_per_pass=*/11, /*histogram_items_per_thread=*/8, /*tie_break_items_per_thread=*/8, - /*copy_items_per_thread=*/8}; + /*copy_items_per_thread=*/8, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; } // Hard constraints on the block_tile byte geometry. The aligned bulk-copy (TMA) load path addresses gmem/smem in diff --git a/cub/cub/util_device.cuh b/cub/cub/util_device.cuh index 7b5eca2e7b3..c77a4417cf1 100644 --- a/cub/cub/util_device.cuh +++ b/cub/cub/util_device.cuh @@ -603,9 +603,6 @@ inline constexpr int bulk_copy_min_align = 16; // Largest thread-block cluster width guaranteed on every SM 9.0+ device. inline constexpr int max_portable_cluster_blocks = 8; -// CUDA's hardware ceiling on thread-block cluster width (Hopper supports up to 16). -inline constexpr int max_supported_cluster_blocks = 16; - //! @brief Returns the alignment needed for the shared memory destination buffer of BlockLoadToShared. //! @tparam T //! Value type to be loaded. diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 709eb4c7366..4280297e7b3 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -9,6 +9,7 @@ #include "insert_nested_NVTX_range_guard.h" #include +#include // topk_policy / make_{baseline,cluster}_policy (cluster-cap tuning test) #include #include @@ -24,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -524,6 +526,343 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with a deferred (device-resident } #endif // TEST_TYPES == 2 +#if TEST_TYPES == 1 +template +struct cast_to_key_op +{ + template + __host__ __device__ KeyT operator()(T x) const + { + return static_cast(x); + } +}; + +// Yields, for each segment, a *non-contiguous* iterator over that segment's keys (an integral counting iterator cast to +// the key type). Feeding the cluster top-k a non-contiguous key iterator makes `use_block_load_to_shared` false, so the +// agent takes its generic (non-BlockLoadToShared) overflow-streaming path. Segment `seg` produces keys +// [seg * segment_size, (seg + 1) * segment_size), so the flattened input equals the identity sequence. +template +struct counting_segment_keys_op +{ + SegmentSizeT segment_size; + + template + __host__ __device__ auto operator()(IndexT seg) const + { + return cuda::make_transform_iterator( + cuda::make_counting_iterator(static_cast(seg) * segment_size), cast_to_key_op{}); + } +}; +#endif // TEST_TYPES == 1 + +// The cluster-path coverage below pins the launch geometry through `cluster_tuning_selector` (a whole-`topk_policy` +// tune override), which is a direct-API construct -- so this whole block is `TEST_LAUNCH == 0` (built once) rather than +// going through the launch-wrapper macro. These tests exercise agent-internal cluster paths (multi-CTA scan / cluster +// barriers / gmem streaming / stage schedule / idle ranks) that are identical across launch modes, so host launch is +// sufficient; they run at a racecheck-tiny footprint instead of the ~1 Mi segments these paths used to require (small +// segments otherwise collapse to the single-CTA fast path). The chunk stride is shrunk to 512 B (128 floats), so every +// segment still spans several chunks per block rather than degenerating to one. On a device with a very small +// shared-memory budget the real resident capacity could dip below the pinned slot cap, but that only forces *more* +// streaming, so the paths under test stay covered. +#if TEST_TYPES == 1 && TEST_LAUNCH == 0 +// Runs the direct-API cluster top-k twice (temp-size query, then the real call) and syncs, requiring success at each +// step. `Direction` selects Min/Max at compile time. Factored out because the tune-override tests here cannot use the +// launch-wrapper macro (it owns the env) and would otherwise repeat this boilerplate. +template +void run_cluster_topk_keys( + KeyInItT d_keys_in, KeyOutItT d_keys_out, SegSizesT seg_sizes, KParamT k_param, NumSegT num_seg, EnvT env) +{ + std::size_t temp_bytes = 0; + auto run = [&](void* d_temp) { + if constexpr (Direction == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxKeys( + d_temp, temp_bytes, d_keys_in, d_keys_out, seg_sizes, k_param, num_seg, env); + } + else + { + return cub::DeviceBatchedTopK::MinKeys( + d_temp, temp_bytes, d_keys_in, d_keys_out, seg_sizes, k_param, num_seg, env); + } + }; + REQUIRE(cudaSuccess == run(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == run(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); +} + +// Builds the require/tune env pinning `Selector` as the whole-policy tune override. +template +auto make_cluster_tune_env(Selector selector) +{ + return cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune(selector)}; +} + +// A small multi-CTA segment that stays *fully resident* across a forced 2-CTA cluster: `single_block_max_seg_size = 0` +// disables the single-CTA fast path so even this tiny segment fans out, and cap 2 pins the width. This exercises the +// real cross-CTA prefix scan (`prime_placement_counters` / remote `red.add`), the cluster barriers, and the DSMEM +// histogram fold -- machinery that previously only ran on 64 Ki+ segments too large for `compute-sanitizer racecheck`. +// Non-deterministic here (striped load + early-stop driver, the path unique to keys); the deterministic driver's use of +// the same scan is verified against an index reference by the pairs "deterministic cross-CTA scan" test. No overflow +// (verifies the scan in isolation). +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys run a small multi-CTA segment through the cross-CTA scan", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + constexpr auto direction = cub::detail::topk::select::max; + skip_if_batched_topk_cluster_unavailable(true); // the tune override forces the SM90+ cluster backend + + // 2048 floats = 16 chunks; a 2-CTA cluster holds 8 chunks each -> fully resident, no streaming. + constexpr segment_size_t static_max_segment_size = 2048; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + constexpr segment_size_t segment_size = static_max_segment_size; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k})); + + CAPTURE(segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // Force a 2-CTA cluster with the single-CTA fast path disabled; slots stay unrestricted (the segment is resident). + auto env = + make_cluster_tune_env(cluster_tuning_selector<2, 0, 0, cluster_test_chunk_bytes>{}); + run_cluster_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + env); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Cluster-width cap (`cluster_cap_list`) plus a resident-slot cap force a tiny segment to overflow and stream: cap 1 +// streams inside one CTA (barrier-free), cap 2 across a fixed 2-CTA cluster (cross-CTA scan while streaming). The +// resident capacity is `slots * chunk_items` (host-known), so 4 slots x 128 floats = 512 resident floats and the ~1.5 K +// segment spills several overflow chunks. Both `stream_stages` and `prologue` are `min(PipelineStages, .)` (of the +// per-CTA overflow count and the 4 resident chunks respectively), so sweeping `PipelineStages` {2,4,8} across the two +// caps (cap 1 -> ~8 overflow chunks/CTA, cap 2 -> ~2) spans the `stream_stages` <, ==, > `prologue` trichotomy and +// reaches the `stage_base` up-front prime (the `>` case). The 4-slot cap leaves an aligned resident tail here; the +// complementary misaligned-tail `stage_rot` reorder is covered by the dedicated test below (reached through the slot +// cap, not a distinct `PipelineStages`). An unaligned size (`- 31`) + `pad` also covers the peeled overflow tail edge. +// Non-deterministic (forward first wave only); the pairs schedule test streams the deterministic combos, which +// additionally reach the reverse first wave. Verified against a sorted reference. +using cluster_cap_list = c2h::enum_type_list; +using stage_list = c2h::enum_type_list; + +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream a tiny oversize segment across the pipeline-stage schedule", + "[keys][segmented][topk][device][cluster]", + cluster_cap_list, + stage_list) +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr int cluster_cap = c2h::get<0, TestType>::value; + constexpr int stages = c2h::get<1, TestType>::value; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + constexpr auto direction = cub::detail::topk::select::max; + skip_if_batched_topk_cluster_unavailable(true); + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + + const int pad = GENERATE(0, 7); + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31})); // unaligned -> peeled overflow tail + // edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(cluster_cap, stages, pad, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(pad + num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(num_segments * segment_size, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + // cap + 4 resident slots -> a tiny segment overflows and streams; `stages` varies the pipeline depth. + auto env = make_cluster_tune_env( + cluster_tuning_selector{}); + run_cluster_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + env); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Misaligned-tail stage rotation (`stage_rot`): the interleaved overflow prime rotates the resident mbarrier-stage +// assignment when the resident chunk count is not a multiple of the pipeline depth. Reuses the schedule sweep's exact +// cap-1 / stages-2 kernel (same `cluster_tuning_selector<1, 4, 0, .., 2>` type and same size/k bounds -> no new +// instantiation); only the *runtime* segment size differs -- 5 chunks here vs 12 there. With 4 slots and one chunk of +// excess over them, 1 stream slot is reserved, leaving 3 resident chunks, and `3 % prologue(2) == 1` misaligns the tail +// (whereas the sweep's 12-chunk segment leaves an aligned 2-chunk tail). Single-CTA (cap 1) streaming, +// non-deterministic (forward first wave); the pairs stage-rotation test streams the deterministic combos, which +// additionally rotate the reverse first wave. Verified against a sorted reference. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream a tiny oversize segment with a misaligned resident tail", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + constexpr auto direction = cub::detail::topk::select::max; + skip_if_batched_topk_cluster_unavailable(true); + + // Bounds match the schedule sweep (same kernel); the runtime segment is sized to 5 chunks (640 floats) so, at 4 slots + // minus 1 reserved stream slot, only 3 resident chunks remain -> a misaligned reload tail (`3 % prologue(2) == 1`). + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + constexpr segment_size_t segment_size = 640; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k})); + + CAPTURE(segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // Same selector as the schedule sweep's cap-1 / stages-2 case; the 5-chunk runtime segment overflows 4 slots and + // leaves 3 resident chunks -> a misaligned reload tail that triggers `stage_rot`. + auto env = make_cluster_tune_env( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/2>{}); + run_cluster_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + env); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Generic (non-BlockLoadToShared) overflow streaming at a tiny footprint: the non-contiguous counting-iterator key +// source makes `use_block_load_to_shared` false, and the slot cap forces a ~1.5 K segment to overflow, so the agent +// takes its generic gmem-streaming path. Non-deterministic (the load path is determinism-independent; the deterministic +// generic streamer is covered by the pairs generic index-order test). Verified against the identity sequence. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream a tiny oversize segment through a non-contiguous key iterator", + "[keys][segmented][topk][device][cluster]") +{ + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + constexpr auto direction = cub::detail::topk::select::max; + skip_if_batched_topk_cluster_unavailable(true); + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + constexpr segment_size_t segment_size = static_max_segment_size; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(segment_size, k, num_segments); + + // Non-contiguous input: segment `seg` is the counting iterator [seg * segment_size, (seg + 1) * segment_size). + auto d_keys_in = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), counting_segment_keys_op{segment_size}); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto env = make_cluster_tune_env( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes>{}); + run_cluster_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}, + env); + + // The flattened input is the identity sequence, so build the expected keys directly and reuse the standard + // sort + compact verification. + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::sequence(expected_keys.begin(), expected_keys.end()); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +#endif // TEST_TYPES == 1 && TEST_LAUNCH == 0 + +// Big-segment general coverage (launch wrapper -> all launch modes, no tune override). The tiny tune-override tests +// above pin geometry to reach the multi-CTA / streaming / schedule / idle paths at a racecheck-tiny footprint; these +// keep real 1 Mi-scale coverage so the index/offset arithmetic (and the streaming path under device/graph launch) is +// exercised at scale. Their sweeps are trimmed (fewer sizes/k/pads) since the fine-grained path coverage now lives in +// the tiny tests. #if TEST_TYPES == 1 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned segments", "[keys][segmented][topk][device][cluster][determinism]", @@ -533,35 +872,23 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - // Requirement under test. Random keys rarely tie at the boundary, so this mainly exercises the blocked-partition + - // streaming/edge interaction; the key result is invariant to the requirement, so all combinations match the - // reference. using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - // `static_max_segment_size` is chosen to exceed the largest all-resident cluster coverage (~16 blocks worth of - // resident SMEM), so the 1 Mi-element segments force the agent's gmem-streaming overflow path (including an - // unaligned overflow tail via `- 31`), while the 128 Ki-element segment still runs fully resident under the same - // streaming-capable launch configuration. The `+ 1` / `- 4095` sizes make the global-last chunk a single item, i.e. - // a pure-suffix tail with an empty aligned bulk (`bulk == 0`) once `pad == 0` aligns the base, exercising the - // always-peeled tail edge on top of a zero-length resident/streamed tail chunk (resident `128 Ki + 1`, streamed - // `1 Mi - 4095`). + // 1 Mi segments overflow the resident cluster and stream (one unaligned `- 31` tail; `- 4095` makes the global-last + // chunk a single item -> pure-suffix tail with empty aligned bulk once `pad == 0`). The requirement is swept because + // the deterministic (blocked-load) path has distinct large-offset arithmetic. constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 3; - // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. skip_if_batched_topk_backend_unavailable(static_max_segment_size); - constexpr auto direction = cub::detail::topk::select::max; - const int pad = GENERATE(0, 1, 3, 7); - const segment_size_t segment_size = GENERATE_COPY(values( - {static_max_segment_size, - static_max_segment_size - 31, - static_max_segment_size - 4095, - segment_size_t{128 * 1024}, - segment_size_t{128 * 1024 + 1}})); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); + constexpr auto direction = cub::detail::topk::select::max; + const int pad = GENERATE(0, 7); + const segment_size_t segment_size = + GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, static_max_segment_size - 4095})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); CAPTURE(pad, static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); @@ -586,21 +913,19 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large fixed-size unaligned fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); compact_sorted_keys_to_topk(expected_keys, segment_size, k); - fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); REQUIRE(expected_keys == keys_out_buffer); } -// Streaming counterpart of the narrow-segment-size regression. Streaming needs segments larger than the resident -// cluster coverage (>128 Ki), which 8/16-bit types can't represent, so we use signed 32-bit -- the same signed -// `int32_t` as the internal `offset_t` -- to exercise the streaming path's index arithmetic across det/non-det. +// Big streaming with a signed 32-bit segment-size type (same width as the internal `offset_t`, but signed): the point +// is that the large-offset arithmetic stays correct on a signed type at real 1 Mi scale. C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 32-bit segment-size type", "[keys][segmented][topk][device][cluster][determinism]", det_tie_combos) { using key_t = float; - using seg_size_t = int; // signed 32-bit: same width as offset_t, but signed + using seg_size_t = int; // signed 32-bit using segment_index_t = cuda::std::int64_t; using combo = c2h::get<0, TestType>; @@ -611,7 +936,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 3 constexpr seg_size_t static_max_segment_size = 1024 * 1024; constexpr seg_size_t static_max_k = 4 * 1024; constexpr segment_index_t num_segments = 2; - // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. skip_if_batched_topk_backend_unavailable(static_max_segment_size); const int pad = GENERATE(0, 7); @@ -647,34 +971,90 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments with a signed 3 REQUIRE(expected_keys == keys_out_buffer); } -template -struct cast_to_key_op +// Big variable-size batch: a loose 1.1 Mi bound sizes a wide cluster while the actual segments mix a streaming 1 Mi +// segment, a fully-resident multi-CTA segment, and small segments that leave surplus CTAs idle -- large-offset +// coverage of every effective-width regime in one launch (the tiny idle-rank test above covers the idle path +// racecheck-clean; this covers it at scale). +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", + "[keys][segmented][topk][device][cluster][determinism]", + det_tie_combos) { - template - __host__ __device__ KeyT operator()(T x) const - { - return static_cast(x); - } -}; + using key_t = float; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; -// Yields, for each segment, a *non-contiguous* iterator over that segment's keys (an integral counting iterator -// cast to the key type). Feeding the cluster top-k a non-contiguous key iterator makes `use_block_load_to_shared` -// false, so the agent takes its generic (non-BlockLoadToShared) overflow-streaming path. Segment `seg` produces -// keys [seg * segment_size, (seg + 1) * segment_size), so the flattened input equals the identity sequence and the -// expected top-k is exact. -template -struct counting_segment_keys_op -{ - SegmentSizeT segment_size; + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr segment_size_t static_max_segment_size = 1100 * 1024; + constexpr segment_size_t static_max_k = 4 * 1024; + skip_if_batched_topk_backend_unavailable(static_max_segment_size); - template - __host__ __device__ auto operator()(IndexT seg) const - { - return cuda::make_transform_iterator( - cuda::make_counting_iterator(static_cast(seg) * segment_size), cast_to_key_op{}); - } -}; + constexpr auto direction = cub::detail::topk::select::max; + const int pad = GENERATE(1, 7); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k})); + + constexpr segment_size_t big_segment_size = 1024 * 1024; + c2h::host_vector h_segment_offsets{ + 0, + big_segment_size, + big_segment_size + (big_segment_size - 31), + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17), + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257, + big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257 + (12 * 1024 + 1)}; + c2h::device_vector segment_offsets = h_segment_offsets; + const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); + const segment_size_t num_items = h_segment_offsets.back(); + + auto segment_offsets_it = thrust::raw_pointer_cast(segment_offsets.data()); + auto segment_size_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), segment_size_op{segment_offsets_it}); + + CAPTURE(pad, static_max_segment_size, static_max_k, k, num_segments, num_items, direction); + + auto compacted_output_sizes_it = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), + get_output_size_op{segment_offsets.cbegin(), cuda::constant_iterator(k), num_segments}); + c2h::device_vector compacted_offsets(num_segments + 1, thrust::no_init); + thrust::exclusive_scan( + compacted_output_sizes_it, compacted_output_sizes_it + num_segments + 1, compacted_offsets.begin()); + segment_size_t total_output_size = compacted_offsets.back(); + + c2h::device_vector keys_in_buffer(pad + num_items, thrust::no_init); + c2h::device_vector keys_out_buffer(total_output_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()) + pad; + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_offsets.cbegin()); + auto d_keys_out = + cuda::make_permutation_iterator(cuda::make_counting_iterator(d_keys_out_ptr), compacted_offsets.cbegin()); + + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); + segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); + expected_keys = compact_to_topk_batched(expected_keys, segment_offsets, k); + + segmented_sort_keys( + keys_out_buffer, num_segments, compacted_offsets.cbegin(), compacted_offsets.cbegin() + 1, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Big generic (non-BlockLoadToShared) streaming under default tuning: the non-contiguous counting-iterator key source +// makes `use_block_load_to_shared` false, so the dispatch routes these 1 Mi segments to the cluster backend's generic +// overflow-streaming path (vs. the tiny forced-cluster generic test above, this exercises the real default routing at +// large-offset scale). The 1 Mi - 31 segment streams with an unaligned tail edge; the 128 Ki segment validates the +// generic resident path (no streaming) through the same code. Keeping totals below 2^24 makes every key an exact float. C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-contiguous key iterator", "[keys][segmented][topk][device][cluster][determinism]", det_tie_combos) @@ -683,35 +1063,26 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-c using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - // Requirement under test (the key result is invariant to it; exercises the generic streaming path). using combo = c2h::get<0, TestType>; constexpr auto determinism = combo::determinism; constexpr auto tie_break = combo::tie_break; - // The counting-iterator key source is non-contiguous, so the agent uses its generic overflow-streaming path rather - // than BlockLoadToShared. `static_max_segment_size` exceeds the largest all-resident cluster coverage, so the 1 Mi - // -element segments stream (incl. an unaligned `- 31` tail), while the 128 Ki-element segment validates the generic - // resident path (no streaming) through the same code. Keeping the largest total below 2^24 makes every key an exact - // float. + constexpr auto direction = cub::detail::topk::select::max; + constexpr segment_size_t static_max_segment_size = 1024 * 1024; constexpr segment_size_t static_max_k = 4 * 1024; - constexpr segment_index_t num_segments = 3; - // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + constexpr segment_index_t num_segments = 2; skip_if_batched_topk_backend_unavailable(static_max_segment_size); - constexpr auto direction = cub::detail::topk::select::max; - const segment_size_t segment_size = - GENERATE_COPY(values({static_max_segment_size, static_max_segment_size - 31, segment_size_t{128 * 1024}})); - const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2, max_k})); - const segment_size_t num_items = num_segments * segment_size; + const segment_size_t segment_size = GENERATE_COPY(values({static_max_segment_size - 31, segment_size_t{128 * 1024}})); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + const segment_size_t num_items = num_segments * segment_size; CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); - // Non-contiguous input: segment `seg` is the counting iterator [seg * segment_size, (seg + 1) * segment_size). auto d_keys_in = cuda::make_transform_iterator( cuda::make_counting_iterator(segment_index_t{0}), counting_segment_keys_op{segment_size}); - // Output is a real buffer (the output iterator stays contiguous; only the input drives the streaming path). c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); @@ -723,8 +1094,6 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys stream large segments through a non-c cuda::args::immediate{k, cuda::args::bounds()}, cuda::args::immediate{num_segments}); - // The flattened input is the identity sequence, so build the expected keys directly and reuse the standard - // sort + compact verification. c2h::device_vector expected_keys(num_items, thrust::no_init); thrust::sequence(expected_keys.begin(), expected_keys.end()); @@ -835,44 +1204,38 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } -#if TEST_TYPES == 1 -C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unaligned segments", - "[keys][segmented][topk][device][cluster][determinism]", - det_tie_combos) +#if TEST_TYPES == 1 && TEST_LAUNCH == 0 +// Small variable-size batch under a forced wide (4-CTA) cluster with the single-CTA fast path disabled. One per-batch +// launch sizes the cluster from the loose bound (4096 -> 4 CTAs), but each segment derives its own effective width from +// its actual chunk count, so the smaller segments leave surplus CTAs idle -- exercising the effective-cluster-width +// arithmetic and the `rank >= eff_cluster_blocks` idle early-out at a racecheck-tiny footprint (previously only +// reachable on a 64 Ki+ mixed batch). Non-deterministic (the effective-width arithmetic is determinism-independent; the +// deterministic idle path is covered by the pairs mixed-effective-width test). Verified against a sorted reference. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys leave surplus cluster CTAs idle on small variable-size segments", + "[keys][segmented][topk][device][cluster]") { using key_t = float; using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - // Requirement under test (the key result is invariant to it; exercises mixed resident/streaming segments). - using combo = c2h::get<0, TestType>; - constexpr auto determinism = combo::determinism; - constexpr auto tie_break = combo::tie_break; - // `static_max_segment_size` exceeds the largest all-resident cluster coverage, so a single per-segment launch sizes a - // wide (16-CTA) cluster and mixes every effective-width regime: streaming segments (the 1 Mi-element ones, one with - // an unaligned `- 31` overflow tail), a fully-resident multi-CTA segment (96 Ki + 17), two *medium* segments that - // exceed the single-CTA threshold but need only a few chunks (so surplus cluster CTAs go idle -- the runtime - // effective-cluster-width path), and a tiny segment that collapses onto a single CTA (257 elements). - constexpr segment_size_t static_max_segment_size = 1100 * 1024; - constexpr segment_size_t static_max_k = 4 * 1024; - // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. - skip_if_batched_topk_backend_unavailable(static_max_segment_size); - - constexpr auto direction = cub::detail::topk::select::max; - const int pad = GENERATE(1, 3, 7); - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2, static_max_k})); - - constexpr segment_size_t big_segment_size = 1024 * 1024; - constexpr segment_size_t med_segment_a = 12 * 1024 + 1; // a few chunks, unaligned tail -> most cluster CTAs idle - constexpr segment_size_t med_segment_b = 40 * 1024; // more chunks, still well below the 16-CTA launch width - c2h::host_vector h_segment_offsets{ - 0, - big_segment_size, - big_segment_size + (big_segment_size - 31), - big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17), - big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257, - big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257 + med_segment_a, - big_segment_size + (big_segment_size - 31) + (96 * 1024 + 17) + 257 + med_segment_a + med_segment_b}; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__not_guaranteed; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + constexpr auto direction = cub::detail::topk::select::max; + skip_if_batched_topk_cluster_unavailable(true); + + constexpr segment_size_t static_max_segment_size = 4096; // loose bound -> 4-CTA physical cluster (128 floats/chunk) + constexpr segment_size_t static_max_k = 512; + const int pad = GENERATE(0, 7); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k})); + + // Mixed effective widths under the fixed 4-CTA launch: 4096 -> full 4, 1024/512 -> 4, 256 -> 2 (2 idle), 128 -> 1 + // (3 idle). + c2h::host_vector h_segment_offsets{0}; + for (const segment_size_t s : + {segment_size_t{4096}, segment_size_t{128}, segment_size_t{512}, segment_size_t{256}, segment_size_t{1024}}) + { + h_segment_offsets.push_back(h_segment_offsets.back() + s); + } c2h::device_vector segment_offsets = h_segment_offsets; const segment_index_t num_segments = static_cast(h_segment_offsets.size() - 1); const segment_size_t num_items = h_segment_offsets.back(); @@ -881,7 +1244,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign auto segment_size_it = cuda::make_transform_iterator( cuda::make_counting_iterator(segment_index_t{0}), segment_size_op{segment_offsets_it}); - CAPTURE(pad, static_max_segment_size, static_max_k, k, num_segments, num_items, direction); + CAPTURE(pad, static_max_segment_size, static_max_k, k, num_segments, num_items); // Each output segment holds exactly min(k, segment_size[i]) items, tightly packed. auto compacted_output_sizes_it = cuda::make_transform_iterator( @@ -906,12 +1269,15 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign c2h::device_vector expected_keys(num_items, thrust::no_init); thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); - batched_topk_keys( + auto env = make_cluster_tune_env( + cluster_tuning_selector<4, /*slots=*/0, /*single_block=*/0, cluster_test_chunk_bytes>{}); + run_cluster_topk_keys( d_keys_in, d_keys_out, cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}); + cuda::args::immediate{num_segments}, + env); segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); expected_keys = compact_to_topk_batched(expected_keys, segment_offsets, k); @@ -922,7 +1288,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with large variable-size unalign REQUIRE(expected_keys == keys_out_buffer); } -#endif // TEST_TYPES == 1 +#endif // TEST_TYPES == 1 && TEST_LAUNCH == 0 C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per-segment k", "[keys][segmented][topk][device]", diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 58b6358c460..af15cbe0947 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -514,6 +514,79 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs handle a segment-size type narrower // every launch id, including device launch (`lid_1`): the CDP static config's small resident capacity is what streams // big segments and peels the unaligned tail edge, so it is the path that must cover them. #if TEST_TYPES == 1 +// Launch-wrapper-compatible pairs cluster dispatch that also pins a whole-`topk_policy` tune override (`Selector`). The +// env (require + tune) is built internally from the threaded stream, so -- unlike a direct-API call that owns its env +// -- a test using it runs under every launch mode. Used only by the interface-level determinism/reproducibility test +// below (the path-pinning tune tests stay host-only direct-API); `CUB_RUNTIME_FUNCTION` so the CDP wrapper can invoke +// it device-side, confirming the tune override and the determinism guarantee hold when invoked from device code. +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_cluster_topk_pairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + cudaStream_t stream = nullptr) +{ + auto env = cuda::std::execution::env{ + cuda::stream_ref{stream}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune(Selector{})}; + if constexpr (SelectDirection == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + num_segments, + env); + } + else + { + return cub::DeviceBatchedTopK::MinPairs( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + num_segments, + env); + } +} + +DECLARE_TMPL_LAUNCH_WRAPPER( + dispatch_cluster_topk_pairs, + cluster_topk_pairs, + ESCAPE_LIST(typename Selector, + cub::detail::topk::select SelectDirection, + cuda::execution::determinism::__determinism_t Determinism, + cuda::execution::tie_break::__tie_break_t TieBreak), + ESCAPE_LIST(Selector, SelectDirection, Determinism, TieBreak)); + template struct cast_to_key_op { @@ -907,6 +980,125 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the REQUIRE(ref == h_values_out); } +# if TEST_LAUNCH == 0 +// Tiny multi-CTA resident cross-CTA scan (pairs): `single_block = 0` disables the single-CTA fast path and cap 2 pins a +// 2-CTA cluster, so a small *fully-resident* segment (no overflow) runs the real cross-CTA prefix scan +// (`prime_placement_counters` / remote `red.add`), cluster barriers, and DSMEM histogram fold that previously only ran +// on 64 Ki+ segments too large for `compute-sanitizer racecheck`. Heavy ties (keys in [0, 7]) fill the k-th bucket so +// the deterministic scan/tie-break drives the result; the value payload (global index) is checked against the host +// index reference, so a wrong scan result is observable. Direct-API, so built once for `TEST_LAUNCH == 0`. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs run a tiny multi-CTA segment through the deterministic cross-CTA scan", + "[pairs][segmented][topk][device][cluster][determinism]", + select_direction_list, + tie_break_pref_list) +{ + using key_t = cuda::std::uint32_t; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr auto direction = c2h::get<0, TestType>::value; + constexpr auto tie_break = c2h::get<1, TestType>::value; + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr bool prefer_larger = tie_break == cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + skip_if_batched_topk_cluster_unavailable(true); + + constexpr segment_size_t static_max_segment_size = 2048; + constexpr segment_size_t static_max_k = 1024; + constexpr segment_index_t num_segments = 2; + + const segment_size_t segment_size = static_max_segment_size; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, static_max_k / 2})); + const segment_size_t num_items = num_segments * segment_size; + + CAPTURE(segment_size, k, num_segments, direction, prefer_larger); + + // Few distinct key values -> many tied candidates in the k-th bucket. Contiguous -> resident BlockLoadToShared path. + c2h::device_vector keys_in_buffer(num_items, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer, key_t{0}, key_t{7}); + + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + // Force the cluster backend and a resident 2-CTA cluster (single-CTA fast path disabled) through the tune query. + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune(cluster_tuning_selector<2, /*slots=*/0, /*single_block=*/0, cluster_test_chunk_bytes>{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + if constexpr (direction == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + } + else + { + return cub::DeviceBatchedTopK::MinPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + } + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + // Values still belong to their keys, and no source index is selected twice. + c2h::device_vector expected_keys(keys_in_buffer); + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + // The deterministic path must return *exactly* the index-ordered top-k. Compare per-segment selected index sets. + c2h::host_vector h_keys = keys_in_buffer; + const c2h::host_vector ref = + reference_deterministic_topk_indices(h_keys, num_segments, segment_size, k, direction, prefer_larger); + + c2h::host_vector h_values_out = values_out_buffer; + for (segment_index_t seg = 0; seg < num_segments; ++seg) + { + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + std::sort(seg_begin, seg_begin + k); + } + + REQUIRE(ref == h_values_out); +} +# endif // TEST_LAUNCH == 0 + // The streaming tie-break regressions below need both `num_passes` parities so they exercise both ping-pong toggle // counts. Enforce that the two widths actually straddle the parity against the real cluster `bits_per_pass` (rather // than hard-coding "uint32 -> odd, uint64 -> even"): if tuning ever makes both widths share a parity, this fails to @@ -946,10 +1138,9 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the skip_if_batched_topk_backend_unavailable(static_max_segment_size); const segment_size_t segment_size = static_max_segment_size - 31; // unaligned -> streaming + peeled tail edge - // With 8 near-uniform values each bucket holds ~segment_size/8 (~128 Ki) keys: k == 1 / 50 Ki straddle the top bucket - // (empty front), k == 300 Ki spans full buckets plus a mid-bucket straddle (non-empty front). - const segment_size_t k = - GENERATE_COPY(values({segment_size_t{1}, segment_size_t{50 * 1024}, segment_size_t{300 * 1024}})); + // With 8 near-uniform values each bucket holds ~segment_size/8 (~128 Ki) keys: k == 1 straddles the top bucket (empty + // front), k == 300 Ki spans full buckets plus a mid-bucket straddle (non-empty front). + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, segment_size_t{300 * 1024}})); CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, prefer_larger); @@ -1057,9 +1248,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the // Non-round size -> streaming (the generic fallback reads any trailing items straight from gmem; it never peels). const segment_size_t segment_size = static_max_segment_size - 31; - const segment_size_t k = - GENERATE_COPY(values({segment_size_t{1}, segment_size_t{50 * 1024}, segment_size_t{300 * 1024}})); - const segment_size_t num_items = num_segments * segment_size; + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, segment_size_t{300 * 1024}})); + const segment_size_t num_items = num_segments * segment_size; CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, prefer_larger); @@ -1114,43 +1304,26 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the REQUIRE(ref == h_values_out); } -// Whole-`topk_policy` tuning overrides for the reproducibility test, threaded through the public API's single tuning -// query (`cuda::execution::tune`). Both force the cluster backend (which alone honors a deterministic request). -// `default_cluster_selector` uses the default cluster tuning; `alt_cluster_selector` starts from it and overrides the -// shape knobs (block size, items-per-thread, pipeline depth, tie-break granularity) to a second *valid* tuning, so the -// gpu_to_gpu config-independence check can compare two different launch configs. -struct default_cluster_selector -{ - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const - -> cub::detail::batched_topk::topk_policy - { - return cub::detail::batched_topk::topk_policy{ - cub::detail::batched_topk::topk_backend::cluster, - cub::detail::batched_topk::make_baseline_policy(), - cub::detail::batched_topk::make_cluster_policy()}; - } -}; - -struct alt_cluster_selector -{ - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const - -> cub::detail::batched_topk::topk_policy - { - auto cluster = cub::detail::batched_topk::make_cluster_policy(); - cluster.threads_per_block = 256; - cluster.histogram_items_per_thread = 2; - cluster.pipeline_stages = 2; - cluster.tie_break_items_per_thread = 2; - return cub::detail::batched_topk::topk_policy{ - cub::detail::batched_topk::topk_backend::cluster, cub::detail::batched_topk::make_baseline_policy(), cluster}; - } -}; - +# if TEST_LAUNCH != 2 // Reproducibility with an *unspecified* tie-break (which tied candidate wins is an implementation detail). We run twice // and require the same selected index set, per each requirement's contract: `run_to_run` only promises repeated runs of // the *same* config agree (so both runs share a tuning); `gpu_to_gpu` must be config-independent (so the second run // uses a different valid tuning), mirroring the reduce/scan deterministic tests. Within-top-k order is unspecified, so // we compare sorted sets. `gpu_to_gpu` cannot be checked more strictly without a second device. +// +// Both configs reuse tiny `cluster_tuning_selector` shapes already instantiated by the tests above (no extra kernels): +// config A is the 2-CTA fully-resident shape (P1's cross-CTA-scan tuning); config B is the single-CTA streaming shape +// (cap 1 from the schedule sweep). The CTA-count/streaming contrast is the second valid config the gpu_to_gpu check +// needs, while forcing the cluster backend (which alone honors a deterministic request) at a racecheck-tiny footprint. +// +// Unlike the other tune-override tests (which pin agent-internal paths identical across launch modes and so stay +// `TEST_LAUNCH == 0`), this one asserts an *interface-level* contract -- a deterministic request stays reproducible -- +// so it is routed through the `cluster_topk_pairs` launch wrapper and also built for device launch (CDP), confirming +// the tune override and the determinism guarantee hold when invoked from device code. Graph launch (`TEST_LAUNCH == 2`) +// re-runs the same host dispatch arm, so it is skipped to avoid the extra kernel instantiations. +using repro_config_a = cluster_tuning_selector<2, /*slots=*/0, /*single_block=*/0, cluster_test_chunk_bytes>; +using repro_config_b = + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/4>; C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break is reproducible", "[pairs][segmented][topk][device][cluster][determinism]", select_direction_list, @@ -1165,13 +1338,16 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break constexpr auto determinism = c2h::get<1, TestType>::value; constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; - constexpr segment_size_t static_max_segment_size = 64 * 1024; - constexpr segment_size_t static_max_k = 64 * 1024; + // Tiny footprint: config B (cap 1, 4 resident slots = 512 keys) overflows this 2048-key segment and streams, while + // config A (cap 2, unrestricted slots) keeps it resident across a 2-CTA cluster -- the streaming-vs-resident contrast + // the gpu_to_gpu config-independence check needs, at a racecheck-tiny size. + constexpr segment_size_t static_max_segment_size = 2048; + constexpr segment_size_t static_max_k = 1024; constexpr segment_index_t num_segments = 2; - // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + // A deterministic request always routes to the SM90+ cluster backend (the tune override forces it at this tiny size). skip_if_batched_topk_backend_unavailable(static_max_segment_size); - const segment_size_t segment_size = GENERATE_COPY(values({segment_size_t{4096}, segment_size_t{64 * 1024}})); + const segment_size_t segment_size = static_max_segment_size; const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k / 2})); const segment_size_t num_items = num_segments * segment_size; @@ -1196,7 +1372,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break // Runs the same problem through the public `cub::DeviceBatchedTopK` API, tuning the whole `topk_policy` via // `cuda::execution::tune` so each run can pick a different valid cluster tuning (the override forces the cluster - // backend). Drives the two-phase temp-storage protocol manually so both runs share this call site. + // backend). `selector` is only used for its type; the `cluster_topk_pairs` wrapper drives the two-phase + // temp-storage protocol (and the host/device-launch switch) so both runs share this call site. const auto run_with_tuning = [&](auto selector, c2h::device_vector& keys_out, c2h::device_vector& values_out) { auto d_keys_out = @@ -1204,62 +1381,25 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); - auto env = cuda::std::execution::env{ - cuda::stream_ref{cudaStream_t{0}}, - cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, - cuda::execution::tie_break::__tie_break_holder_t{}, - cuda::execution::output_ordering::unsorted), - cuda::execution::tune(selector)}; - - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { - auto seg_sizes = - cuda::args::immediate{segment_size, cuda::args::bounds()}; - auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - if constexpr (direction == cub::detail::topk::select::max) - { - return cub::DeviceBatchedTopK::MaxPairs( - d_temp, - temp_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - seg_sizes, - k_param, - cuda::args::immediate{num_segments}, - env); - } - else - { - return cub::DeviceBatchedTopK::MinPairs( - d_temp, - temp_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - seg_sizes, - k_param, - cuda::args::immediate{num_segments}, - env); - } - }; - REQUIRE(invoke(nullptr) == cudaSuccess); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); - REQUIRE(invoke(thrust::raw_pointer_cast(temp_storage.data())) == cudaSuccess); - REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + cluster_topk_pairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); }; - run_with_tuning(default_cluster_selector{}, keys_out_a, values_out_a); + run_with_tuning(repro_config_a{}, keys_out_a, values_out_a); // run_to_run: same tuning twice. gpu_to_gpu: a different valid tuning for the stronger config-independent check. if constexpr (determinism == cuda::execution::determinism::__determinism_t::__gpu_to_gpu) { - run_with_tuning(alt_cluster_selector{}, keys_out_b, values_out_b); + run_with_tuning(repro_config_b{}, keys_out_b, values_out_b); } else { - run_with_tuning(default_cluster_selector{}, keys_out_b, values_out_b); + run_with_tuning(repro_config_a{}, keys_out_b, values_out_b); } // Sanity: values still belong to their keys and no source index repeats within a segment. @@ -1279,6 +1419,203 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break } REQUIRE(h_values_a == h_values_b); } +# endif // TEST_LAUNCH != 2 + +# if TEST_LAUNCH == 0 +// Cluster-width cap + tiny streaming (pairs): force the cluster backend, cap the launch to 1 or 2 CTAs, and cap the +// resident slots so a tiny segment overflows and streams -- deterministically reaching the overflow/streaming path +// (single-CTA at cap 1, a fixed 2-CTA cluster at cap 2) with the value payload carried along, at a small footprint +// rather than the 1 Mi a hardware-derived wide cluster would otherwise need. Both `stream_stages` and `prologue` are +// `min(PipelineStages, .)`, so sweeping `PipelineStages` {2,4,8} across the two caps (cap 1 -> ~8 overflow chunks/CTA, +// cap 2 -> ~2) spans the `stream_stages` <, ==, > `prologue` trichotomy and reaches the `stage_base` up-front prime +// (the `>` case). The 4-slot cap leaves an aligned resident tail here; the complementary misaligned-tail `stage_rot` +// reorder is covered by the dedicated test below (reached through the slot cap, not a distinct `PipelineStages`). +// Swept over `det_tie_pair_combos` (both drivers); the deterministic combos additionally reach the reverse first-wave +// direction. Direct-API, so built once for `TEST_LAUNCH == 0`. +using cluster_cap_list = c2h::enum_type_list; +using stage_list = c2h::enum_type_list; + +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment across the pipeline-stage schedule", + "[pairs][segmented][topk][device][cluster][determinism]", + cluster_cap_list, + stage_list, + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + constexpr int cluster_cap = c2h::get<0, TestType>::value; + constexpr int stages = c2h::get<1, TestType>::value; + using combo = c2h::get<2, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + // Forcing the cluster backend needs a cluster-capable target; the guard skips when absent. + skip_if_batched_topk_cluster_unavailable(true); + + const segment_size_t segment_size = static_max_segment_size - 31; // unaligned -> peeled overflow tail edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(cluster_cap, stages, static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // Force the cluster backend at the requested width cap, cap resident slots to 4, and set the pipeline depth, all + // through the public API's single tuning query. + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune( + cluster_tuning_selector{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Deterministic counterpart of the keys misaligned-tail (`stage_rot`) test: reuses the schedule sweep's cap-1 / +// stages-2 kernel (same `cluster_tuning_selector<1, 4, 0, .., 2>` type and same size/k bounds -> no new instantiation); +// only the *runtime* segment size differs. Sized to 5 chunks so, at 4 slots minus 1 reserved stream slot, 3 resident +// chunks remain and `3 % prologue(2) == 1` misaligns the tail. Swept over `det_tie_pair_combos` so the deterministic +// combos flip the first-wave direction, rotating both the forward and the reverse `first_wave_chunk_for_stage` +// mapping. The unaligned size (`609 = 5 chunks - 31`) also peels the unaligned tail edge under the rotation. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with a misaligned resident tail", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + skip_if_batched_topk_cluster_unavailable(true); + + const segment_size_t segment_size = 640 - 31; // 5 chunks, unaligned -> misaligned tail + peeled overflow tail edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // Same selector as the schedule sweep's cap-1 / stages-2 case; the 5-chunk runtime segment overflows 4 slots and + // leaves 3 resident chunks -> a misaligned reload tail that triggers `stage_rot`. + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/2>{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} +# endif // TEST_LAUNCH == 0 // Effective-cluster-width coverage for pairs: the large `deferred_sequence` bound sizes a wide (16-CTA) launch, but the // actual segments are small, so a single launch mixes a fully-resident multi-CTA segment (96 Ki + 17, every rank diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index b802d7ea14d..b7c633fe060 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -68,6 +68,66 @@ void skip_if_batched_topk_backend_unavailable(cuda::std::int64_t static_max_segm skip_if_batched_topk_cluster_unavailable(deterministic || oversize); } +// Whole-`topk_policy` tuning override that forces the cluster backend and pins otherwise heuristic / hardware-derived +// launch geometry, so a *small* segment can deterministically reach paths that normally only large segments hit. Shared +// by the segmented keys and pairs cluster tests (identical there, so it lives here). The levers (a `0`/`-1` sentinel +// keeps the production default): +// * `MaxBlocksPerCluster` -- cap on the launched cluster width. A cap narrower than a segment needs forces the +// oversize/streaming fallback (cap 1 -> single-CTA streaming; cap 2 -> a fixed 2-CTA +// cluster with a real cross-CTA scan / barriers) instead of the hardware-derived width. +// * `MaxChunkSlotsPerBlock` -- cap on resident chunk slots per block. Resident capacity is then `slots * chunk_items` +// (host-known), so a tiny segment above it overflows into streaming -- how a test +// reaches the streaming / stage-schedule paths at a racecheck-tiny footprint. Combined +// with a segment size that overflows to a non-divisor resident-chunk count, it also +// reaches the misaligned-tail (`stage_rot`) prime rotation. +// * `SingleBlockMaxSegSize` -- pin the single-CTA-fastpath threshold; `0` disables it so even a small resident +// segment fans out across the cluster (multi-CTA scan / idle-rank paths). +// * `ChunkBytes` -- shrink the chunk (== slot) stride so a small segment still spans several chunks. +// * `PipelineStages` -- vary the streaming pipeline depth relative to the resident chunk count, sweeping the +// `stream_stages` <, ==, > `prologue` stage-schedule trichotomy. +// * `MinChunksPerBlock` -- raise the chunks-per-block divisor to collapse the effective cluster below the +// physical width (idle-rank early-out). +template +struct cluster_tuning_selector +{ + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const + -> cub::detail::batched_topk::topk_policy + { + auto cluster = cub::detail::batched_topk::make_cluster_policy(); + cluster.max_blocks_per_cluster = MaxBlocksPerCluster; + if constexpr (MaxChunkSlotsPerBlock != 0) + { + cluster.max_chunk_slots_per_block = MaxChunkSlotsPerBlock; + } + if constexpr (SingleBlockMaxSegSize >= 0) + { + cluster.single_block_max_seg_size = SingleBlockMaxSegSize; + } + if constexpr (ChunkBytes != 0) + { + cluster.chunk_bytes = ChunkBytes; + } + if constexpr (PipelineStages != 0) + { + cluster.pipeline_stages = PipelineStages; + } + if constexpr (MinChunksPerBlock != 0) + { + cluster.min_chunks_per_block = MinChunksPerBlock; + } + return cub::detail::batched_topk::topk_policy{ + cub::detail::batched_topk::topk_backend::cluster, cub::detail::batched_topk::make_baseline_policy(), cluster}; + } +}; + +// 512 B chunk stride -> 128 floats per chunk (== per slot); small enough that a tiny segment spans several chunks. +inline constexpr int cluster_test_chunk_bytes = 512; + // Function object to generate monotonically non-decreasing values for small key types template struct inc_t From 10391be645709c9b6fcb001f20a3df4a7dc20231 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 13 Jul 2026 20:02:44 +0200 Subject: [PATCH 116/126] Try fixing perf regression again --- cub/cub/agent/agent_batched_topk_cluster.cuh | 115 +++++++++++-------- 1 file changed, 67 insertions(+), 48 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 6032d5c24d9..9be9fcc7c09 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -984,20 +984,18 @@ private: return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; } - // Issue one resident chunk's aligned bulk on stage `stage` into the densely-packed block_tile at `off_bytes`, - // advancing `off_bytes` by the bytes written. Counterpart of `issue_stream_copy` for the resident window: the - // destination is the running packed cursor (vs a fixed slot) and there is no inflight bookkeeping (resident stages - // are driven directly, not through the stream ring). - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_resident_copy(int stage, offset_t chunk_index, int& off_bytes) + // Issue resident chunk `chunk_index`'s aligned bulk on stage `stage` into its slot at the fixed stride + // `chunk_index * ChunkBytes` (every resident chunk but the segment's last is a full `chunk_items` slot, so the dense + // packing needs no running write cursor). Like `issue_stream_copy` but without the stream ring's inflight mask. + _CCCL_DEVICE _CCCL_FORCEINLINE void issue_resident_copy(int stage, offset_t chunk_index) { - const auto src = chunk_bulk_src(layout.resident_base, chunk_index); - issue_bulk_copy(stage, key_slots + off_bytes, src); - off_bytes += static_cast(::cuda::std::size(src)) * int{sizeof(key_t)}; + issue_bulk_copy( + stage, key_slots + static_cast(chunk_index) * ChunkBytes, chunk_bulk_src(layout.resident_base, chunk_index)); } - // Issue one overflow chunk's aligned bulk on stage `stage` into its fixed streaming slot and mark the stage in - // flight. Counterpart of `issue_resident_copy` for the overflow window (fixed slot vs packed cursor; the stream ring - // records in-flight stages in `stream_inflight_mask`). + // Issue one overflow chunk's aligned bulk on stage `stage` into its streaming slot and mark the stage in flight. + // Counterpart of `issue_resident_copy` for the overflow window: a ring slot keyed by `stage` (vs. the resident + // window's chunk-indexed slot) whose in-flight stages are tracked in `stream_inflight_mask`. _CCCL_DEVICE _CCCL_FORCEINLINE void issue_stream_copy(int stage, offset_t overflow_idx) { _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); @@ -2304,13 +2302,13 @@ private: if (layout.my_resident_chunks > 0) { // Stage mbarriers and the `load_phase` parity are shared with the overflow stream (no per-chunk token array - // needed). Chunks are written densely in slot order from offset 0 and read back in the same order, so the read - // cursor - // (`read_off_bytes`) mirrors the write cursor (`next_off_bytes`) as a running prefix sum, avoiding a - // dynamically-indexed `pending_spans` array that would anchor surrounding state to local memory. Every chunk + // needed). Chunks are written densely in slot order from offset 0 and read back in the same order; the fixed + // slot stride (`chunk_index * ChunkBytes`) positions both the write and the fold with no running cursor or + // dynamically-indexed `pending_spans` array (which would anchor surrounding state to local memory). Every chunk // begins on a `load_align` boundary (zero prefix), so its aligned bulk is `round_down(count, - // load_align_items)`: the whole `count` for interior chunks, and for the global-last chunk its `count` minus - // the unaligned suffix that is always peeled into `edge_keys`. + // load_align_items)`: the whole `count` for interior chunks (`== chunk_items`, a full slot), and for the + // global-last chunk its `count` minus the unaligned suffix that is always peeled into `edge_keys` (the only + // chunk that can be short -- and, folded in ascending order, always the last resident chunk). const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); // Prime the overflow stream while the resident load still runs to hide its latency, over two disjoint stage @@ -2319,10 +2317,6 @@ private: // stream stage, so only the no-resident-chunks path leaves priming to the post-loop call. const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; - // Bulks are densely packed from the start of the block_tile. The head prefix is no longer kept here (it lives - // in `edge_keys`), so there is no reserved front gap: the write cursor starts at offset 0. - int next_off_bytes = 0; - // Place the resident load's `prologue`-wide mbarrier window so the stream's first wave is primed in forward // consume order (stage `s`, holding chunk `s`, first), removing the ordering inversion where the up-front prime // would otherwise issue the last-consumed chunks first. Two complementary knobs, at most one non-zero @@ -2362,7 +2356,7 @@ private: _CCCL_ASSERT( stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), "resident issue stage desynced from its rolling counter"); - issue_resident_copy(stage, local_chunk, next_off_bytes); + issue_resident_copy(stage, local_chunk); if (++stage == stage_base + prologue) { stage = stage_base; @@ -2389,10 +2383,8 @@ private: } } - // Read cursor trailing the write cursor: chunk `local_chunk`'s bulk was written at `read_off_bytes` (packed - // and consumed in the same order), and `chunk_bulk_src` recomputes its length, so the read span needs no - // stored per-stage state. - int read_off_bytes = 0; + // Chunk `local_chunk` was written at its fixed slot offset `local_chunk * ChunkBytes` (same stride the issue + // path uses), and `chunk_bulk_src` recomputes its length, so the fold needs no stored per-chunk state. // Same unit-stride rolling counter as the issue loop over `[stage_base, stage_base + prologue)`: increment and // select-wrap to `stage_base`, no runtime `% prologue`. Starts at `stage_base + stage_rot`; the assert pins it // to the closed form. @@ -2406,28 +2398,50 @@ private: wait_stage(stage); const int read_len_items = static_cast(::cuda::std::size(chunk_bulk_src(layout.resident_base, local_chunk))); + // Fixed-stride packing invariant: only the segment's global-last chunk can be short, and (folded ascending) + // it is always the last resident chunk, so every earlier resident chunk fills its slot. + _CCCL_ASSERT(local_chunk + offset_t{1} == layout.my_resident_chunks || read_len_items == chunk_items, + "only the last resident chunk may carry a short aligned bulk"); for_each_chunk_key( - ::cuda::ptr_rebind(key_slots + read_off_bytes), read_len_items, add_first_pass); - read_off_bytes += read_len_items * int{sizeof(key_t)}; + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + read_len_items, + add_first_pass); const offset_t next_local_chunk = local_chunk + static_cast(prologue); + // Inline both re-arm dispatches (resident successor vs. first-wave stream prime) so each only computes its + // `(dst, src)`; the single (force-inlined) bulk copy is then issued once, uniformly, below. + char* dst = nullptr; + ::cuda::std::span src = {}; + bool rearm = false; if (next_local_chunk < layout.my_resident_chunks) { - // Phase safety, not data safety (the target offset is fresh): re-arming this stage before all threads - // leave the wait above would advance the phase twice, stranding a lagging waiter forever. - __syncthreads(); - issue_resident_copy(stage, next_local_chunk, next_off_bytes); + // Resident successor: densely packed at its fixed slot offset (constant `ChunkBytes` stride). + dst = key_slots + static_cast(next_local_chunk) * ChunkBytes; + src = chunk_bulk_src(layout.resident_base, next_local_chunk); + rearm = true; } else if (interleave_overflow_prime && stage < stream_stages) { // Shared resident stage (in `[stage_base, stage_base + prologue)`) with no resident successor: re-arm it - // with its first-wave chunk (`first_wave_chunk_for_stage`, direction-aware). Same phase-safety barrier as - // the resident re-arm; the copy targets the disjoint stream region (`stream_slot_base` onward), never the - // densely-packed resident data the later tail folds still read. Stages `>= stream_stages` are extra - // resident pipeline depth with no stream slot, so they idle -- `stage_rot` places them last so all primes - // issue first. + // with its first-wave chunk (`first_wave_chunk_for_stage`, direction-aware) into the disjoint stream region + // (`stream_slot_base` onward), never the densely-packed resident data the later tail folds still read. + // Stages `>= stream_stages` are extra resident pipeline depth with no stream slot, so they idle + // (`stage_rot` places them last so all primes issue first). + const offset_t overflow_idx = first_wave_chunk_for_stage(stage); + _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); + _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, + "cannot issue a load into a streaming stage that is still in flight"); + dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + src = chunk_bulk_src(layout.overflow_base, overflow_idx); + stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); + rearm = true; + } + if (rearm) + { + // Phase safety, not data safety (the target is fresh/disjoint): re-arming this stage before all threads + // leave the wait above would advance the phase twice, stranding a lagging waiter forever. __syncthreads(); - issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + issue_bulk_copy(stage, dst, src); } if (++stage == stage_base + prologue) { @@ -2435,12 +2449,17 @@ private: } } - // Read and write cursors sum the same aligned bulks in the same order, so they meet on a load-align boundary. - _CCCL_ASSERT(read_off_bytes == next_off_bytes && next_off_bytes % load_align_bytes == 0, - "resident bulk read/write cursors must meet on a load-alignment boundary"); - // The resident region is one contiguous run of aligned bulks for the later passes; both boundary edges are - // folded separately from `edge_keys`. - resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), next_off_bytes / int{sizeof(key_t)}); + // Resident region extent for the later passes (one contiguous run of aligned bulks; both boundary edges live in + // `edge_keys`). Inferred once here rather than carried as a loop cursor: the fixed slot stride makes it + // `(my_resident_chunks - 1)` full `chunk_items` slots plus the last chunk's (possibly short) aligned bulk. The + // sum is a `load_align_items` multiple because each term is. + const offset_t last_bulk_items = static_cast( + ::cuda::std::size(chunk_bulk_src(layout.resident_base, layout.my_resident_chunks - offset_t{1}))); + const int resident_items = static_cast( + (layout.my_resident_chunks - offset_t{1}) * static_cast(chunk_items) + last_bulk_items); + _CCCL_ASSERT(resident_items % load_align_items == 0, + "the resident region must end on a load-alignment boundary"); + resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), resident_items); // The up-front and tail primes together armed every stream stage, so the stream is primed; the post-loop call // no-ops. @@ -2588,10 +2607,10 @@ private: // tail of its block_tile and re-streams its overflow chunks from gmem each pass via the overflow stream. // // Boundary edges (the unaligned head prefix on rank 0 and the unaligned tail suffix on the tail owner) cannot use - // the aligned TMA stream, so both are always peeled into the persistent `edge_keys` buffer. The tail chunk's - // aligned bulk is then a normal (possibly partial) aligned chunk that can be resident or streamed like any other; - // no partial tail chunk is ever kept resident. Peeling both boundaries means streaming needs only `full_slots >= - // 1`. + // the aligned TMA stream, so both are always peeled into the persistent `edge_keys` buffer. Only that unaligned + // suffix is excluded; the tail chunk's aligned bulk is then a normal aligned chunk (possibly shorter than a full + // slot) that can be resident or streamed like any other. Peeling both boundaries means streaming needs only + // `full_slots >= 1`. // // `stream_slots` is right-sized: clamped into `[1, full_slots]` when streaming; deep overflows still get the full // `PipelineStages` depth. The generic fallback has no async pipeline (it re-reads overflow from gmem each pass and From b9a790f4dee283c1a0d9576fa84b80908de6b5e2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Mon, 13 Jul 2026 22:53:35 +0200 Subject: [PATCH 117/126] Next try --- cub/cub/agent/agent_batched_topk_cluster.cuh | 84 ++++++++++---------- 1 file changed, 41 insertions(+), 43 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 9be9fcc7c09..a77da83e692 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -967,7 +967,7 @@ private: // `count` rounded down to a `load_align` multiple, as a span from `block_keys_base`; empty when it has none. Every // chunk starts on a `load_align` boundary (guard-free TMA path); the global-last chunk's unaligned suffix is always // peeled into `edge_keys`, so the aligned bulk excludes it (`bulk == count` for every interior chunk). Shared source - // computation for both issue paths (`issue_resident_copy`, `issue_stream_copy`). + // computation for both the resident and overflow-stream bulk-copy issues. [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::span chunk_bulk_src(offset_t window_base, offset_t index) const { @@ -984,29 +984,6 @@ private: return {layout.block_keys_base + chunk.offset, static_cast<::cuda::std::size_t>(bulk)}; } - // Issue resident chunk `chunk_index`'s aligned bulk on stage `stage` into its slot at the fixed stride - // `chunk_index * ChunkBytes` (every resident chunk but the segment's last is a full `chunk_items` slot, so the dense - // packing needs no running write cursor). Like `issue_stream_copy` but without the stream ring's inflight mask. - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_resident_copy(int stage, offset_t chunk_index) - { - issue_bulk_copy( - stage, key_slots + static_cast(chunk_index) * ChunkBytes, chunk_bulk_src(layout.resident_base, chunk_index)); - } - - // Issue one overflow chunk's aligned bulk on stage `stage` into its streaming slot and mark the stage in flight. - // Counterpart of `issue_resident_copy` for the overflow window: a ring slot keyed by `stage` (vs. the resident - // window's chunk-indexed slot) whose in-flight stages are tracked in `stream_inflight_mask`. - _CCCL_DEVICE _CCCL_FORCEINLINE void issue_stream_copy(int stage, offset_t overflow_idx) - { - _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); - _CCCL_ASSERT(stage >= 0 && stage < stream_stages, "overflow stage index exceeds the reserved streaming slots"); - _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, - "cannot issue a load into a streaming stage that is still in flight"); - char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); - stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); - } - // Shared-memory view of the chunk currently resident in `stage`'s slot: the slot index is a pure function of `stage` // and the length is recomputed from `overflow_idx` (no spillable per-stage array). Returns the aligned bulk only (the // always-peeled tail suffix is excluded). @@ -1066,14 +1043,19 @@ private: return; } // Reaching here is the once-per-segment prime: the guard rules out a re-prime, and the resident load drives the - // shared stages via `issue_resident_copy` (never the stream mask), so the stream must be fully drained here. + // shared stages via its own bulk copies (never the stream mask), so the stream must be fully drained here. _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, "overflow stream priming must start from a fully drained stream"); __syncthreads(); - // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. + // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. + _CCCL_PRAGMA_NOUNROLL() for (int stage = 0; stage < stream_stages; ++stage) { - issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + const offset_t overflow_idx = first_wave_chunk_for_stage(stage); + _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); + char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); + stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); } stream_is_primed = true; } @@ -1109,8 +1091,13 @@ private: if (next_step < layout.overflow_chunks) { const offset_t next_overflow_idx = stream_is_forward ? next_step : (layout.overflow_chunks - 1 - next_step); + _CCCL_ASSERT(next_overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); __syncthreads(); - issue_stream_copy(stage, next_overflow_idx); + _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, + "cannot issue a load into a streaming stage that is still in flight"); + char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, next_overflow_idx)); + stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); } // Step the ring pointer one slot in the consume direction, wrapping with a compare-and-select (no modulo). if (stream_is_forward) @@ -2348,15 +2335,18 @@ private: // rolling counter over `[stage_base, stage_base + prologue)`: it increments and select-wraps back to // `stage_base` (no `+ stage_rot` and no runtime `% prologue`, costly as `prologue` is not a power of two). The // per-iteration assert pins it to the closed form so a later edit cannot desync them. - // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. + // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. { int stage = stage_base + static_cast(stage_rot); + _CCCL_PRAGMA_NOUNROLL() for (offset_t local_chunk = 0; local_chunk < static_cast(prologue); ++local_chunk) { _CCCL_ASSERT( stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), "resident issue stage desynced from its rolling counter"); - issue_resident_copy(stage, local_chunk); + issue_bulk_copy(stage, + key_slots + static_cast(local_chunk) * ChunkBytes, + chunk_bulk_src(layout.resident_base, local_chunk)); if (++stage == stage_base + prologue) { stage = stage_base; @@ -2364,22 +2354,30 @@ private: } } - // Prime the stream stages the resident load never uses -- the complement of `[stage_base, stage_base + - // prologue)` within `[0, stream_stages)`. Their mbarriers are fresh (no resident drain to wait for), so these - // copies start now, concurrent with the resident bulks; the shared resident stages are primed as the tail frees - // them (read loop below). Forward the free set is the low `[0, stage_base)` (first-consumed chunks); reverse it - // is the high `[prologue, stream_stages)`; at most one loop is non-empty (both are empty when the whole stream - // reservation overlaps the resident window, i.e. `stage_base == 0 && stream_stages <= prologue`). - // No unroll pragma on purpose: the compiler's partial auto-unroll beats forced nounroll here. + // Prime the stream stages the resident load never uses: the complement of the resident window + // `[stage_base, stage_base + prologue)` within `[0, stream_stages)` (`stream_stages - prologue` stages, reached + // by mapping `i` around the window). Their mbarriers are fresh (no resident drain to wait for), so these copies + // start now, concurrent with the resident bulks; the shared resident stages are primed as the tail frees them + // (read loop below). `stage_base` steers this free set to the stages consumed first (low `[0, stage_base)` when + // the first wave runs forward, high `[prologue, stream_stages)` when reverse), matching up-front prime order to + // consume order. Empty (non-positive trip count) when the reservation fits the resident window + // (`stream_stages <= prologue`). + // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. if (interleave_overflow_prime) { - for (int stage = 0; stage < stage_base; ++stage) + _CCCL_PRAGMA_NOUNROLL() + for (int i = 0; i < stream_stages - prologue; ++i) { - issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); - } - for (int stage = stage_base + prologue; stage < stream_stages; ++stage) - { - issue_stream_copy(stage, first_wave_chunk_for_stage(stage)); + const int stage = (i < stage_base) ? i : i + prologue; + _CCCL_ASSERT(stage < stage_base || stage >= stage_base + prologue, + "up-front prime stage must lie outside the resident window"); + const offset_t overflow_idx = first_wave_chunk_for_stage(stage); + _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); + _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, + "cannot issue a load into a streaming stage that is still in flight"); + char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); + stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); } } From cc127604c59cd3661d24d4532347ee1346f696b2 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 14 Jul 2026 01:00:30 +0200 Subject: [PATCH 118/126] Cut out dead streaming prime call --- cub/cub/agent/agent_batched_topk_cluster.cuh | 20 +- .../device/dispatch/dispatch_batched_topk.cuh | 235 +++++++++--------- 2 files changed, 133 insertions(+), 122 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index a77da83e692..807313d7436 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -1028,14 +1028,12 @@ private: } // Prime the streaming slots: issue this direction's first `stream_stages` overflow copies, arming the reload - // pipeline. Idempotent (`stream_is_primed`) and a no-op without overflow -- the first effective call per segment - // primes; later passes inherit the previous pass's resident turn-around chunks (exactly this direction's first - // `stream_stages` chunks, since the order ping-pongs) with no re-prime. The only path that actually issues here has - // no resident chunks (`load_and_histogram_first_pass` interleaves the prime whenever it has some) and already follows - // the kernel-init barrier, so the leading `__syncthreads()` is effectively redundant today; it is kept so the prime - // stays phase-safe if a future caller ever runs it right after a resident load (every thread must clear that load's - // final `wait_stage` before these copies re-arm the shared mbarriers, else the phase advances twice and a lagging - // thread spins forever). + // pipeline. Idempotent (`stream_is_primed`) and a no-op without overflow. Single caller + // `load_and_histogram_first_pass`: its resident path interleaves the prime inline (so this call then no-ops), leaving + // only the no-resident-chunks path to actually issue here. That path already follows the kernel-init barrier, so the + // leading `__syncthreads()` is redundant today; it is kept so the prime stays phase-safe if a future caller runs it + // right after a resident load (every thread must clear that load's final `wait_stage` before these copies re-arm the + // shared mbarriers, else the phase advances twice and a lagging thread spins forever). _CCCL_DEVICE _CCCL_FORCEINLINE void prime_overflow_stream() { if (stream_is_primed || layout.overflow_chunks == 0) @@ -1145,8 +1143,10 @@ private: if constexpr (use_block_load_to_shared) { - // Idempotent prime; `load_and_histogram_first_pass` normally primes earlier, so this is usually a no-op. - prime_overflow_stream(); + // The stream is always primed before the first `run_pass`: `load_and_histogram_first_pass` primes every streaming + // rank on the block-load path (interleaved with its resident load, or directly when it has no resident chunks), + // and `stream_is_primed` then carries across all radix passes and the final filter with no re-prime. + _CCCL_ASSERT(stream_is_primed, "run_pass requires the overflow stream primed by load_and_histogram_first_pass"); // Phase 1: consume the first `stream_stages` visits and issue this pass's reload wave into the freed slots. Their // slots hold the early-primed wave (still in flight) on the first pass and the previous pass's reused turn-around diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index e88febbccaa..4ffe89c8738 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -697,35 +697,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return cudaSuccess; }; - // Hardware cluster-width ceiling, taken from the runtime rather than a compile-time constant so a future device - // with wider (non-portable) clusters is not artificially capped (the same reason the dynamic-SMEM budget above is - // queried, not hardcoded). Probed at zero dynamic SMEM to get the architectural/kernel ceiling independent of the - // per-`C` SMEM chosen below; each candidate width is still validated against its own SMEM by the resident scan's - // `max_active_clusters` probe, and the oversize branch re-probes at its (larger) launch SMEM. Requires the - // non-portable opt-in already set above so the probe can report widths beyond the portable ceiling. - cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` - cfg.gridDim = dim3(1, 1, 1); - cfg.dynamicSmemBytes = 0; - int hw_cluster_ceiling = 0; - if (const auto error = launcher_factory.max_potential_cluster_size(hw_cluster_ceiling, dynamic_kernel, &cfg)) - { - return error; - } - if (hw_cluster_ceiling <= 0) - { - return cudaErrorInvalidValue; - } - // `MaxBlocksPerCluster == 0` -> the full hardware ceiling; a non-zero knob narrows it, clamped to that ceiling. A - // cap narrower than a segment needs pushes it into the oversize/streaming fallback below (cap 1 -> single-CTA). - const int cluster_cap = - (MaxBlocksPerCluster == 0) ? hw_cluster_ceiling : (::cuda::std::min) (MaxBlocksPerCluster, hw_cluster_ceiling); - - // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); - // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs more - // SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). We - // pick the `C` that minimizes the number of waves, breaking ties toward the largest `C` (= smallest SMEM = most - // L1), which matches the profiled fast configs. We enumerate `C` analytically rather than discovering SMEM tiers - // via occupancy, so a register-limited occupancy (e.g. 1 CTA/SM) cannot collapse the candidate set. + // Segment/capacity basics that gate the launch shape. Computed before any occupancy query so the single-CTA fast + // path below can be taken without one -- that driver query otherwise dominates the runtime of tiny launches. const auto seg = static_cast<::cuda::std::uint64_t>(max_seg_size); const auto chunk_items_u64 = static_cast<::cuda::std::uint64_t>(layout_t::chunk_items); const int max_block_tile_capacity = static_cast(layout_t::block_tile_capacity(max_dynamic_smem_bytes)); @@ -741,18 +714,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return layout_t::base_padding_bytes + static_cast(slots) * layout_t::chunk_bytes; }; - // `C_full`: at the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (capped at - // `cluster_cap`). `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest resident - // `C`. Both are computed and compared in 64-bit, because `max_seg_size` may be a loose bound (e.g. - // `numeric_limits::max()` for an unbounded deferred sequence); narrowing such a `C_lo` to `int` could wrap to - // a small (or negative) value and wrongly enter the resident branch instead of the oversize/streaming fallback. - const int c_full = static_cast( - (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(cluster_cap), ::cuda::ceil_div(seg, chunk_items_u64))); + // `C_lo`: at the largest SMEM each CTA holds `max_block_tile_capacity`, the smallest resident `C`. Computed in + // 64-bit because `max_seg_size` may be a loose bound (e.g. `numeric_limits::max()` for an unbounded deferred + // sequence); narrowing such a `C_lo` to `int` could wrap and wrongly enter the resident branch below. const auto c_lo = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(max_block_tile_capacity)); - // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than - // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. - const int desired_cluster_blocks = ::cuda::narrow(batched_topk_cluster::effective_cluster_blocks_from_chunks( - ::cuda::ceil_div(seg, chunk_items_u64), MinChunksPerBlock, ::cuda::narrow(cluster_cap))); int cluster_blocks = 0; int dynamic_smem_sel = 0; @@ -767,94 +732,140 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( cluster_blocks = 1; dynamic_smem_sel = smem_for_block_capacity(seg); } - else if (c_lo <= static_cast<::cuda::std::uint64_t>(cluster_cap)) + else { - // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= cluster_cap`, so every - // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. - // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. The upper bound is capped at the - // cluster blocks the max segment needs (`desired_cluster_blocks`), so the host never launches a wider cluster - // than necessary; at `min_chunks_per_block == 1` the cap equals `c_full`. `C_end` is additionally clamped to - // `cluster_cap`, which empties the range only at `cluster_cap == 1` (`c_begin == 2 > 1`) -- reachable here for - // a one-CTA-resident segment whose single-CTA fast path is disabled. Then `c_lo == 1` and the fallback below - // takes the single CTA, so the width cap is never exceeded. - const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); - const int c_end = - (::cuda::std::min) (cluster_cap, - (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks))); - ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); - for (int c = c_begin; c <= c_end; ++c) + // Hardware cluster-width ceiling, taken from the runtime rather than a compile-time constant so a future device + // with wider (non-portable) clusters is not artificially capped (the same reason the dynamic-SMEM budget above + // is queried, not hardcoded). Probed at zero dynamic SMEM to get the architectural/kernel ceiling independent + // of the per-`C` SMEM chosen below; each candidate width is still validated against its own SMEM by the + // resident scan's `max_active_clusters` probe, and the oversize branch re-probes at its (larger) launch SMEM. + // Requires the non-portable opt-in already set above so the probe can report widths beyond the portable + // ceiling. + cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` + cfg.gridDim = dim3(1, 1, 1); + cfg.dynamicSmemBytes = 0; + int hw_cluster_ceiling = 0; + if (const auto error = launcher_factory.max_potential_cluster_size(hw_cluster_ceiling, dynamic_kernel, &cfg)) { - const auto per_block_items = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(c)); - const int s_res = smem_for_block_capacity(per_block_items); - if (s_res > max_dynamic_smem_bytes) + return error; + } + if (hw_cluster_ceiling <= 0) + { + return cudaErrorInvalidValue; + } + // `MaxBlocksPerCluster == 0` -> the full hardware ceiling; a non-zero knob narrows it, clamped to that ceiling. + // A cap narrower than a segment needs pushes it into the oversize/streaming fallback below. + const int cluster_cap = + (MaxBlocksPerCluster == 0) + ? hw_cluster_ceiling + : (::cuda::std::min) (MaxBlocksPerCluster, hw_cluster_ceiling); + + // Wave-aware cluster-blocks selection. The free variable is the cluster blocks `C` (one cluster per segment); + // each `C` is paired with the smallest dynamic SMEM that keeps a segment fully resident. A smaller `C` needs + // more SMEM (fewer clusters-per-wave, less L1); a larger `C` needs less SMEM (more clusters-per-wave, more L1). + // We pick the `C` that minimizes waves, breaking ties toward the largest `C` (= smallest SMEM = most L1), which + // matches the profiled fast configs. `C` is enumerated analytically rather than by discovering SMEM tiers via + // occupancy, so a register-limited occupancy (e.g. 1 CTA/SM) cannot collapse the candidate set. `C_full`: at + // the 1-chunk SMEM each CTA holds `chunk_items`, so full residency needs this many CTAs (capped at + // `cluster_cap`). + const int c_full = static_cast( + (::cuda::std::min) (static_cast<::cuda::std::uint64_t>(cluster_cap), ::cuda::ceil_div(seg, chunk_items_u64))); + // Cluster blocks the max segment actually needs (shared with the device so the launch is never wider than + // necessary). At `min_chunks_per_block == 1` this equals `c_full`; a larger knob shrinks it. + const int desired_cluster_blocks = + ::cuda::narrow(batched_topk_cluster::effective_cluster_blocks_from_chunks( + ::cuda::ceil_div(seg, chunk_items_u64), MinChunksPerBlock, ::cuda::narrow(cluster_cap))); + + if (c_lo <= static_cast<::cuda::std::uint64_t>(cluster_cap)) + { + // Full residency is achievable. `seg <= C_lo * max_block_tile_capacity` with `C_lo <= cluster_cap`, so every + // per-CTA capacity (and thus its slot count and SMEM bytes) below stays well within `int` -- no overflow. + // Scan `C` in `[max(C_lo, 2), C_end]`, minimize waves, tie-break largest `C`. The upper bound is capped at + // the cluster blocks the max segment needs (`desired_cluster_blocks`), so the host never launches a wider + // cluster than necessary; at `min_chunks_per_block == 1` the cap equals `c_full`. `C_end` is additionally + // clamped to `cluster_cap`, which empties the range only at `cluster_cap == 1` (`c_begin == 2 > 1`) -- + // reachable here for a one-CTA-resident segment whose single-CTA fast path is disabled. Then `c_lo == 1` and + // the fallback below takes the single CTA, so the width cap is never exceeded. + const int c_begin = (::cuda::std::max) (2, static_cast(c_lo)); + const int c_end = + (::cuda::std::min) (cluster_cap, + (::cuda::std::max) (c_begin, (::cuda::std::min) (c_full, desired_cluster_blocks))); + ::cuda::std::uint64_t best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); + for (int c = c_begin; c <= c_end; ++c) { - continue; // unreachable for c >= C_lo, but guards the SMEM budget regardless. + const auto per_block_items = ::cuda::ceil_div(seg, static_cast<::cuda::std::uint64_t>(c)); + const int s_res = smem_for_block_capacity(per_block_items); + if (s_res > max_dynamic_smem_bytes) + { + continue; // unreachable for c >= C_lo, but guards the SMEM budget regardless. + } + + if (const auto error = ensure_dynamic_smem_limit(s_res)) + { + return error; + } + + // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must + // be a multiple of the cluster blocks. The returned value is the device-wide clusters-per-wave (capacity), + // independent of grid size, and accounts for the static footprint and register pressure internally. + cluster_attr.val.clusterDim.x = static_cast(c); + cfg.gridDim = dim3(static_cast(c), 1, 1); + cfg.dynamicSmemBytes = static_cast(s_res); + int clusters_per_wave = 0; + if (const auto error = launcher_factory.max_active_clusters(clusters_per_wave, dynamic_kernel, &cfg)) + { + return error; + } + if (clusters_per_wave <= 0) + { + continue; // cluster blocks not launchable at this SMEM. + } + + const auto waves = ::cuda::ceil_div( + static_cast<::cuda::std::uint64_t>(num_seg_val), static_cast<::cuda::std::uint64_t>(clusters_per_wave)); + // Min waves; tie-break largest `C`. The loop ascends in `C`, so `<=` keeps the largest at equal waves. + if (cluster_blocks == 0 || waves <= best_waves) + { + best_waves = waves; + cluster_blocks = c; + dynamic_smem_sel = s_res; + } } - if (const auto error = ensure_dynamic_smem_limit(s_res)) + if (cluster_blocks == 0 && c_lo == 1) { - return error; + // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, + // but `C_lo == 1` guarantees `S_res(seg)` fits the budget and one CTA is always launchable. + cluster_blocks = 1; + dynamic_smem_sel = smem_for_block_capacity(seg); } + } - // `cudaOccupancyMaxActiveClusters` needs the cluster dimension and the matching dynamic SMEM; the grid must - // be a multiple of the cluster blocks. The returned value is the device-wide clusters-per-wave (capacity), - // independent of grid size, and accounts for the static footprint and register pressure internally. - cluster_attr.val.clusterDim.x = static_cast(c); - cfg.gridDim = dim3(static_cast(c), 1, 1); - cfg.dynamicSmemBytes = static_cast(s_res); - int clusters_per_wave = 0; - if (const auto error = launcher_factory.max_active_clusters(clusters_per_wave, dynamic_kernel, &cfg)) + if (cluster_blocks == 0) + { + // Oversize (`C_lo > cluster_cap`) or nothing launchable in range: full residency is impossible, so maximize + // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. + if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) { return error; } - if (clusters_per_wave <= 0) + cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` + cfg.gridDim = dim3(1, 1, 1); + cfg.dynamicSmemBytes = static_cast(max_dynamic_smem_bytes); + int hw_max_cluster_blocks = 0; + if (const auto error = + launcher_factory.max_potential_cluster_size(hw_max_cluster_blocks, dynamic_kernel, &cfg)) { - continue; // cluster blocks not launchable at this SMEM. + return error; } - - const auto waves = ::cuda::ceil_div( - static_cast<::cuda::std::uint64_t>(num_seg_val), static_cast<::cuda::std::uint64_t>(clusters_per_wave)); - // Min waves; tie-break largest `C`. The loop ascends in `C`, so `<=` keeps the largest at equal waves. - if (cluster_blocks == 0 || waves <= best_waves) + hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, cluster_cap); + if (hw_max_cluster_blocks <= 0) { - best_waves = waves; - cluster_blocks = c; - dynamic_smem_sel = s_res; + return cudaErrorInvalidValue; } + cluster_blocks = hw_max_cluster_blocks; + dynamic_smem_sel = max_dynamic_smem_bytes; } - - if (cluster_blocks == 0 && c_lo == 1) - { - // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, - // but `C_lo == 1` guarantees `S_res(seg)` fits the budget and one CTA is always launchable. - cluster_blocks = 1; - dynamic_smem_sel = smem_for_block_capacity(seg); - } - } - - if (cluster_blocks == 0) - { - // Oversize (`C_lo > cluster_cap`) or nothing launchable in range: full residency is impossible, so maximize - // residency with the largest launchable cluster at the largest SMEM and let the agent stream the overflow. - if (const auto error = ensure_dynamic_smem_limit(max_dynamic_smem_bytes)) - { - return error; - } - cluster_attr.val.clusterDim.x = 1; // ignored by `max_potential_cluster_size` - cfg.gridDim = dim3(1, 1, 1); - cfg.dynamicSmemBytes = static_cast(max_dynamic_smem_bytes); - int hw_max_cluster_blocks = 0; - if (const auto error = launcher_factory.max_potential_cluster_size(hw_max_cluster_blocks, dynamic_kernel, &cfg)) - { - return error; - } - hw_max_cluster_blocks = (::cuda::std::min) (hw_max_cluster_blocks, cluster_cap); - if (hw_max_cluster_blocks <= 0) - { - return cudaErrorInvalidValue; - } - cluster_blocks = hw_max_cluster_blocks; - dynamic_smem_sel = max_dynamic_smem_bytes; } // The launch needs `MaxDynamicSharedMemorySize >= dynamic_smem_sel`; the scan already raised the limit past the From a92bdd34a2999cb36e95ff3468ac103f74fb7098 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 14 Jul 2026 03:40:09 +0200 Subject: [PATCH 119/126] Clean up --- cub/cub/agent/agent_batched_topk_cluster.cuh | 219 ++++++++++--------- 1 file changed, 113 insertions(+), 106 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 807313d7436..2a1843db4ae 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -54,7 +54,6 @@ #include #include -#include #include #include @@ -73,11 +72,12 @@ #include #include #include +#include #include #include +#include #include #include -#include #include @@ -370,7 +370,7 @@ struct agent_batched_topk_cluster int{sizeof(key_t)} == int{alignof(key_t)} && LoadAlignBytes % int{sizeof(key_t)} == 0; static constexpr bool use_block_load_to_shared = THRUST_NS_QUALIFIER::is_trivially_relocatable_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v - && key_is_bulk_tileable; + && ::cuda::std::__can_to_address && key_is_bulk_tileable; // --------------------------------------------------------------------------- // Block-scan used by the leader block to prefix-sum its merged histogram @@ -561,11 +561,11 @@ struct agent_batched_topk_cluster } } - template - _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(const key_t* chunk_keys, int count, F&& f) const + template + _CCCL_DEVICE _CCCL_FORCEINLINE void for_each_chunk_key(const key_t* chunk_keys, int count, KeyOp&& key_op) const { for_each_chunk_key_impl(chunk_keys, count, [&](const key_t& key, int) { - f(key); + key_op(key); }); } @@ -657,7 +657,7 @@ struct agent_batched_topk_cluster // chunking, the effective (non-idle) cluster width and this rank's partition of it, the leader/idle roles, and the // resident/streaming/edge split of this rank's chunks. All fields are loop-invariant for the segment; the // streaming-only intermediates (`overflow_base`, `resident_slots_cap`, `stream_slots`, `my_chunks`) are kept because - // the overflow-streaming helpers (`init_overflow_stream`, `run_pass`, ...) read them. + // the overflow-streaming helpers (`init_overflow_stream`, `run_overflow_pass`, ...) read them. struct segment_layout { offset_t segment_size_off; @@ -1062,7 +1062,7 @@ private: // slot, fold its keys via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the just-freed // slot (a barrier orders the block's read ahead of the overwriting copy). Returns false once `should_continue()` // reports the top-k fully placed -- polled before the prefetch so we never launch a copy we would only drain again; - // the prefetches already in flight are drained after `run_pass`'s consume loops. + // the prefetches already in flight are drained after `run_overflow_pass`'s consume loops. template [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool consume_overflow_visit(offset_t i, int& stage, BlockApply&& block_apply, Continue&& should_continue) @@ -1114,19 +1114,19 @@ private: // Shared driver for one overflow pass. `block_apply(stage, overflow_idx)` folds the chunk `overflow_idx` resident // in streaming slot `stage` (block-load path); `generic_apply(chunk)` folds an overflow chunk read straight from - // gmem (fallback). `mid()` runs at most once per pass, positioned to overlap the caller's resident-chunk work with - // in-flight copies: after the first reload wave (`stream_stages` visits) is issued but before it is waited on - // (block-load path), or before the gmem loop (fallback). A phase-1 early stop on the block-load path skips it, since - // nothing is then left to place; it must be block-uniform with no unmatched barrier. `should_continue()` is polled - // once after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the final - // filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the + // gmem (fallback). `overlap_work()` runs at most once per pass, positioned to overlap the caller's resident-chunk + // work with in-flight copies: after the first reload wave (`stream_stages` visits) is issued but before it is waited + // on (block-load path), or before the gmem loop (fallback). A phase-1 early stop on the block-load path skips it, + // since nothing is then left to place; it must be block-uniform with no unmatched barrier. `should_continue()` is + // polled once after each consumed chunk (before its refill copy is issued); returning false breaks the stream so the + // final filter bails once the top-k is placed. Its result must be block-uniform (all lanes break together, else the // post-break barrier deadlocks) and it is evaluated by every lane, so it may contain a barrier (both final filters' // do, to read their placement counters block-wide); the histogram passes an always-true predicate. An early break // can leave prefetches in flight, so the pass drains the remaining stages before returning (a full pass ends with an // empty `stream_inflight_mask`, so the drain is a no-op). - template - _CCCL_DEVICE _CCCL_FORCEINLINE void - run_pass(BlockApply&& block_apply, GenericApply&& generic_apply, Mid&& mid, Continue&& should_continue) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run_overflow_pass( + BlockApply&& block_apply, GenericApply&& generic_apply, OverlapWork&& overlap_work, Continue&& should_continue) { // A pass begins drained, except the first, where `load_and_histogram_first_pass` early-primed the first wave (all // `stream_stages` stages in flight) to overlap its edge work; phase 1 waits on them. Later passes re-enter drained: @@ -1137,16 +1137,17 @@ private: "an overflow pass must begin drained or holding only the early-primed first wave"); if (layout.overflow_chunks == 0) { - mid(); + overlap_work(); return; } if constexpr (use_block_load_to_shared) { - // The stream is always primed before the first `run_pass`: `load_and_histogram_first_pass` primes every streaming - // rank on the block-load path (interleaved with its resident load, or directly when it has no resident chunks), - // and `stream_is_primed` then carries across all radix passes and the final filter with no re-prime. - _CCCL_ASSERT(stream_is_primed, "run_pass requires the overflow stream primed by load_and_histogram_first_pass"); + // The stream is always primed before the first `run_overflow_pass`: `load_and_histogram_first_pass` primes every + // streaming rank on the block-load path (interleaved with its resident load, or directly when it has no resident + // chunks), and `stream_is_primed` then carries across all radix passes and the final filter with no re-prime. + _CCCL_ASSERT(stream_is_primed, + "run_overflow_pass requires the overflow stream primed by load_and_histogram_first_pass"); // Phase 1: consume the first `stream_stages` visits and issue this pass's reload wave into the freed slots. Their // slots hold the early-primed wave (still in flight) on the first pass and the previous pass's reused turn-around @@ -1173,12 +1174,12 @@ private: if (!is_stopped) { - // The reload wave is now in flight; run the caller's resident-chunk work to hide its latency before waiting. - // Skipped on an early stop: the stream only breaks once this CTA's whole contribution is placed, so no resident - // key can still be required (the non-deterministic filter folds its resident keys as `mid`). - mid(); + // The reload wave is now in flight; run the caller's overlap work to hide its latency before waiting. Skipped + // on an early stop: the stream only breaks once this CTA's whole contribution is placed, so no resident key can + // still be required (the non-deterministic filter folds its resident keys as `overlap_work`). + overlap_work(); - // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `mid`). + // Phase 2: consume the remaining visits (their loads were issued in phase 1 and overlapped `overlap_work`). _CCCL_PRAGMA_NOUNROLL() for (offset_t i = split; i < layout.overflow_chunks; ++i) { @@ -1207,7 +1208,7 @@ private: // Generic fallback: no async SMEM pipeline, so resident work cannot hide load latency here. Fold the resident // chunks first (preserving the prior ordering), then read the overflow keys straight from gmem each pass (no // SMEM reuse), with the walk still snaking for L2 locality. - mid(); + overlap_work(); _CCCL_PRAGMA_NOUNROLL() for (offset_t i = 0; i < layout.overflow_chunks; ++i) { @@ -1224,16 +1225,16 @@ private: } } - // Apply `f(key)` to every overflow key once in the current ping-pong direction. `UnrollFactor` partially unrolls the - // generic (gmem fallback) fold loop; callers pass their clamped items-per-thread. See `run_pass` for the overlap - // semantics of `mid`. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f, Mid&& mid) + // Apply `fold_key_op(key)` to every overflow key once in the current ping-pong direction. `UnrollFactor` partially + // unrolls the generic (gmem fallback) fold loop; callers pass their clamped items-per-thread. See `run_overflow_pass` + // for the overlap semantics of `overlap_work`. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void fold_overflow_keys(FoldKeyOp&& fold_key_op, OverlapWork&& overlap_work) { - run_pass( + run_overflow_pass( [&](int stage, offset_t overflow_idx) { const auto keys = stage_span(stage, overflow_idx); - for_each_chunk_key(keys.data(), static_cast(keys.size()), f); + for_each_chunk_key(keys.data(), static_cast(keys.size()), fold_key_op); }, [&](const auto& chunk) { const int iterations = ::cuda::ceil_div(chunk.count, threads_per_block); @@ -1243,24 +1244,17 @@ private: const int local = j * threads_per_block + tid; if (local < chunk.count) { - f(layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); + fold_key_op( + layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); } } }, - static_cast(mid), + ::cuda::std::forward(overlap_work), [] { return true; }); } - // Overload with no interleaved (`mid`) work: the fused first pass folds its resident keys during the load itself, so - // pass 0 only folds the overflow chunks here and has no resident work left to overlap. - template - _CCCL_DEVICE _CCCL_FORCEINLINE void process_pass(F&& f) - { - process_pass(static_cast(f), [] {}); - } - // ------------------------------------------------------------------------- // Per-direction implementation // ------------------------------------------------------------------------- @@ -1811,9 +1805,9 @@ private: } // Overflow chunks. On the block-load path each landed slot folds through `process_tiles` (reusing the TMA - // pipeline); the generic fallback reads gmem chunk by chunk. `run_pass`'s `should_continue` breaks the stream once - // the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), and the - // slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so + // pipeline); the generic fallback reads gmem chunk by chunk. `run_overflow_pass`'s `should_continue` breaks the + // stream once the top-k is fully placed. The streaming direction is already correct on entry (preselected in `run`), + // and the slots are already primed -- the histogram's first streaming pass primed the same persistent stream, so // `stream_is_primed` carries in as `true` and this pass reuses the resident turn-around chunks with no re-prime (the // generic fallback re-reads gmem each pass regardless). template @@ -1830,7 +1824,7 @@ private: "preselected ping-pong parity mismatch: the straddling CTA entered the deterministic filter with " "the wrong streaming direction"); - run_pass( + run_overflow_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. // `stage_span` returns the slot's aligned-bulk view (a peeled tail suffix is handled by `process_tail_edge`). [&](int stage, offset_t overflow_idx) { @@ -1982,8 +1976,8 @@ private: // --------------------------------------------------------------------------- // Order-independent counterpart of the deterministic filter: every region is swept by // `process_tiles<..., Deterministic=false>` (arrival-order placement, no scan/tie-state/barrier). Since order is - // irrelevant, `write_nondeterministic_topk` folds the resident keys as the overflow stream's `mid` work (overlapping - // the first reloads) and bails out of the overflow stream once its contribution is fully placed. + // irrelevant, `write_nondeterministic_topk` folds the resident keys as the overflow stream's `overlap_work` + // (overlapping the first reloads) and bails out of the overflow stream once its contribution is fully placed. template struct nondet_filter_state { @@ -1995,10 +1989,10 @@ private: smem_keys_t resident_keys; }; - // Fold this rank's resident keys (and boundary edges) through `process_tiles`, run as the overflow pass's `mid` work - // so they overlap the first overflow reloads. Placement uses order-independent atomics into slots disjoint from the - // streaming slots, so this never races the overflow apply. Each chunk's `chunk.offset` is passed as the region's - // segment base (used for the pair value fetch; elided in keys-only). + // Fold this rank's resident keys (and boundary edges) through `process_tiles`, run as the overflow pass's + // `overlap_work` so they overlap the first overflow reloads. Placement uses order-independent atomics into slots + // disjoint from the streaming slots, so this never races the overflow apply. Each chunk's `chunk.offset` is passed as + // the region's segment base (used for the pair value fetch; elided in keys-only). template _CCCL_DEVICE _CCCL_FORCEINLINE void nondet_fold_resident(nondet_filter_state& state) { @@ -2065,9 +2059,9 @@ private: // Non-deterministic final filter driver. `prime_placement_counters` primes this block's front/back placement counters // to their absolute bases (only `sel_prefix` is read back here; the back base is used in place via `back_local_cnt`); - // overflow keys then stream through `run_pass` (resident keys folded as its `mid`) and place into block-local SMEM - // atomics. `run_pass` breaks the stream once this CTA's contribution is fully placed (all its selected keys in the - // front, and its back counter reached the region end `k`). + // overflow keys then stream through `run_overflow_pass` (resident keys folded as its `overlap_work`) and place into + // block-local SMEM atomics. `run_overflow_pass` breaks the stream once this CTA's contribution is fully placed (all + // its selected keys in the front, and its back counter reached the region end `k`). template _CCCL_DEVICE _CCCL_FORCEINLINE void write_nondeterministic_topk( out_offset_t num_tie_winners, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) @@ -2095,7 +2089,7 @@ private: nondet_filter_state state{ identify_op, block_keys_out, num_tie_winners, sel_prefix, my_front, resident_keys}; - run_pass( + run_overflow_pass( // Block-load: fold the chunk `overflow_idx`, resident in streaming slot `stage`, straight from SMEM. [&](int stage, offset_t overflow_idx) { const auto keys = stage_span(stage, overflow_idx); @@ -2331,54 +2325,61 @@ private: ? static_cast(prologue) - tail_align : offset_t{0}; - // Load every resident chunk's aligned bulk, densely packed in slot order. `stage` is a second, unit-stride - // rolling counter over `[stage_base, stage_base + prologue)`: it increments and select-wraps back to - // `stage_base` (no `+ stage_rot` and no runtime `% prologue`, costly as `prologue` is not a power of two). The - // per-iteration assert pins it to the closed form so a later edit cannot desync them. + // Issue every up-front bulk copy from one TMA site, over two disjoint stage sets. The first `prologue` + // iterations load the resident chunks, densely packed in slot order over the mbarrier window + // `[stage_base, stage_base + prologue)` (`stage_res`: unit-stride, select-wrap to `stage_base`, avoiding a + // `% prologue`, a real division since `prologue` is a runtime value). The remaining `stream_stages - prologue` + // iterations (only when + // interleaving) prime the complement stages the resident load never touches with their first-wave chunk into + // the disjoint stream region; those mbarriers are fresh, so the copies run concurrently with the resident bulks + // (the shared resident stages are primed as the read loop below frees them). `stage_base` steers the free set + // to the first-consumed stages (low for a forward first wave, high for reverse), matching prime order to + // consume order. + // // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. + const int upfront_copies = (interleave_overflow_prime && stream_stages > prologue) ? stream_stages : prologue; + int stage_res = stage_base + static_cast(stage_rot); + _CCCL_PRAGMA_NOUNROLL() + for (int idx = 0; idx < upfront_copies; ++idx) { - int stage = stage_base + static_cast(stage_rot); - _CCCL_PRAGMA_NOUNROLL() - for (offset_t local_chunk = 0; local_chunk < static_cast(prologue); ++local_chunk) + // Each branch selects only its `dst` and chunk source; the shared `chunk_bulk_src` and bulk copy issue once, + // uniformly, below. + int stage = 0; + char* dst = nullptr; + offset_t src_base = 0; + offset_t src_index = 0; + if (idx < prologue) { + const offset_t local_chunk = static_cast(idx); _CCCL_ASSERT( - stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), + stage_res == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), "resident issue stage desynced from its rolling counter"); - issue_bulk_copy(stage, - key_slots + static_cast(local_chunk) * ChunkBytes, - chunk_bulk_src(layout.resident_base, local_chunk)); - if (++stage == stage_base + prologue) + stage = stage_res; + dst = key_slots + idx * ChunkBytes; + src_base = layout.resident_base; + src_index = local_chunk; + if (++stage_res == stage_base + prologue) { - stage = stage_base; + stage_res = stage_base; } } - } - - // Prime the stream stages the resident load never uses: the complement of the resident window - // `[stage_base, stage_base + prologue)` within `[0, stream_stages)` (`stream_stages - prologue` stages, reached - // by mapping `i` around the window). Their mbarriers are fresh (no resident drain to wait for), so these copies - // start now, concurrent with the resident bulks; the shared resident stages are primed as the tail frees them - // (read loop below). `stage_base` steers this free set to the stages consumed first (low `[0, stage_base)` when - // the first wave runs forward, high `[prologue, stream_stages)` when reverse), matching up-front prime order to - // consume order. Empty (non-positive trip count) when the reservation fits the resident window - // (`stream_stages <= prologue`). - // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. - if (interleave_overflow_prime) - { - _CCCL_PRAGMA_NOUNROLL() - for (int i = 0; i < stream_stages - prologue; ++i) + else { - const int stage = (i < stage_base) ? i : i + prologue; + // Map the complement index `idx - prologue` around the resident window: low stages `[0, stage_base)` map to + // themselves, the rest skip the window (`+ prologue`, i.e. back to `idx`). + stage = (idx - prologue < stage_base) ? idx - prologue : idx; _CCCL_ASSERT(stage < stage_base || stage >= stage_base + prologue, "up-front prime stage must lie outside the resident window"); const offset_t overflow_idx = first_wave_chunk_for_stage(stage); _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, "cannot issue a load into a streaming stage that is still in flight"); - char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); + dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + src_base = layout.overflow_base; + src_index = overflow_idx; stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); } + issue_bulk_copy(stage, dst, chunk_bulk_src(src_base, src_index)); } // Chunk `local_chunk` was written at its fixed slot offset `local_chunk * ChunkBytes` (same stride the issue @@ -2406,17 +2407,20 @@ private: add_first_pass); const offset_t next_local_chunk = local_chunk + static_cast(prologue); - // Inline both re-arm dispatches (resident successor vs. first-wave stream prime) so each only computes its - // `(dst, src)`; the single (force-inlined) bulk copy is then issued once, uniformly, below. - char* dst = nullptr; - ::cuda::std::span src = {}; - bool rearm = false; + // Inline both re-arm dispatches (resident successor vs. first-wave stream prime) so each only selects its + // `dst` and chunk source (`src_base`/`src_index`); the shared `chunk_bulk_src` and the single (force-inlined) + // bulk copy then run once, uniformly, below. + char* dst = nullptr; + offset_t src_base = 0; + offset_t src_index = 0; + bool rearm = false; if (next_local_chunk < layout.my_resident_chunks) { // Resident successor: densely packed at its fixed slot offset (constant `ChunkBytes` stride). - dst = key_slots + static_cast(next_local_chunk) * ChunkBytes; - src = chunk_bulk_src(layout.resident_base, next_local_chunk); - rearm = true; + dst = key_slots + static_cast(next_local_chunk) * ChunkBytes; + src_base = layout.resident_base; + src_index = next_local_chunk; + rearm = true; } else if (interleave_overflow_prime && stage < stream_stages) { @@ -2429,8 +2433,9 @@ private: _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, "cannot issue a load into a streaming stage that is still in flight"); - dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - src = chunk_bulk_src(layout.overflow_base, overflow_idx); + dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + src_base = layout.overflow_base; + src_index = overflow_idx; stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); rearm = true; } @@ -2439,7 +2444,7 @@ private: // Phase safety, not data safety (the target is fresh/disjoint): re-arming this stage before all threads // leave the wait above would advance the phase twice, stranding a lagging waiter forever. __syncthreads(); - issue_bulk_copy(stage, dst, src); + issue_bulk_copy(stage, dst, chunk_bulk_src(src_base, src_index)); } if (++stage == stage_base + prologue) { @@ -2522,7 +2527,9 @@ private: // its preselected initial direction; the histogram is order-independent, so the direction only sets up the leftover // parity for the final filter) and reuses the resident load's stage mbarriers (all front-loaded at `run` entry), // with `wait_stage` for producer/consumer sync; the generic fallback re-reads each overflow chunk from gmem. - process_pass(add_first_pass); + // No overlap work: this fused first pass already folded its resident keys during the load, so only the overflow + // chunks remain to fold here. + fold_overflow_keys(add_first_pass, [] {}); const int resident_count = static_cast(resident_keys.size()); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, @@ -2543,7 +2550,7 @@ private: layout.head_items = 0; if constexpr (use_block_load_to_shared) { - layout.block_keys_base = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(layout.block_keys_in); + layout.block_keys_base = ::cuda::std::to_address(layout.block_keys_in); // `run` is reached only for `0 < k < segment_size`, so a bulk-load segment is non-empty and has a real base. _CCCL_ASSERT(layout.block_keys_base != nullptr, "a bulk-load segment must have a valid input pointer"); // Items from `block_keys_base` to the next `load_align_bytes` boundary (0 when already aligned), clamped to the @@ -2709,8 +2716,8 @@ private: }; // Resident-chunk histogram, deferred into the overflow stream so it overlaps the stream's in-flight first - // reload wave (see `process_pass`). The histogram is order-independent, so folding resident keys between the - // stream's load issue and its wait does not change the result. + // reload wave (see `fold_overflow_keys`). The histogram is order-independent, so folding resident keys between + // the stream's load issue and its wait does not change the result. const auto fold_resident_hist = [&] { if constexpr (use_block_load_to_shared) { @@ -2735,7 +2742,7 @@ private: // Re-stream the overflow chunks into this pass's histogram, overlapping the resident-chunk histogram with the // first wave of reload bulk copies. Ping-pongs direction and reuses the turn-around chunks left resident by the // previous pass. - process_pass(add_hist, fold_resident_hist); + fold_overflow_keys(add_hist, fold_resident_hist); // Fold the persistent boundary edges (loaded once in the first pass) into this pass's histogram, alongside the // resident and overflow keys. Keeps every owner's per-bucket counts (the source of its `num_strictly_selected` From 9abc4eff335434421440f3a0651c59fa00e8220b Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 14 Jul 2026 08:05:48 +0200 Subject: [PATCH 120/126] First pass opimization Also add testing to increase coverage due to an uncovered bug found during implementation. --- cub/cub/agent/agent_batched_topk_cluster.cuh | 477 +++++++++--------- ...catch2_test_device_segmented_topk_pairs.cu | 91 ++++ 2 files changed, 328 insertions(+), 240 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index 2a1843db4ae..cfb2feb294f 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -280,8 +280,8 @@ struct agent_batched_topk_cluster // Radix passes for this key type (compile-time). The overflow stream flips direction each pass, and `run` preselects // the initial direction so the leftover after `num_passes` matches the deterministic filter (see `run`). The first - // wave is always primed in that initial direction, so `first_wave_is_forward` -- the only direction - // `first_wave_chunk_for_stage` needs -- is a compile-time constant (non-deterministic keeps the forward default). + // wave is always primed in that initial direction, so `first_wave_is_forward` -- the only direction the first-wave + // stage->chunk mapping needs -- is a compile-time constant (non-deterministic keeps the forward default). static constexpr int num_passes = detail::topk::calc_num_passes(bits_per_pass); static constexpr bool first_wave_is_forward = !need_determinism || ((!is_tie_reversed) ^ ((num_passes & 1) != 0)); @@ -727,6 +727,10 @@ struct agent_batched_topk_cluster int stream_stages = 1; bool stream_is_forward{true}; bool stream_is_primed{false}; + // In-flight streaming copies, one bit per stage. Maintained only by the radix-pass consume path + // (`consume_overflow_visit`/`run_overflow_pass`), where phase-1 visits hit resident turn-around chunks and must not + // wait. The histogram first pass owns no persistent stream state: it waits its own primes in ring order and leaves + // this zero, so every `run_overflow_pass` enters drained. ::cuda::std::uint32_t stream_inflight_mask{0}; _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_cluster( @@ -996,68 +1000,6 @@ private: static_cast(::cuda::round_down(static_cast(chunk.count), static_cast(load_align_items)))); } - // Chunk the first overflow wave assigns to `stage` (`stage < stream_stages`). The wave is primed in the compile-time - // `first_wave_is_forward` direction. Forward: stage `s` -> chunk `s`. Reverse: the last `stream_stages` chunks, each - // landing in stage `chunk % stream_stages`, so `s` takes the one in that window congruent to it. This - // `stage == chunk % stream_stages` invariant lets both the eager (interleaved) and post-loop primes address a - // stage's slot with no stored mapping. - [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE offset_t first_wave_chunk_for_stage(int stage) const - { - _CCCL_ASSERT(stage >= 0 && stage < stream_stages, "first-wave stage index out of range"); - _CCCL_ASSERT(static_cast(stream_stages) <= layout.overflow_chunks, - "first wave assumes at least `stream_stages` overflow chunks"); - // Only ever evaluated while priming the first wave, which `run` always issues in this preselected direction. - _CCCL_ASSERT(stream_is_forward == first_wave_is_forward, "first-wave prime ran in an unexpected stream direction"); - if constexpr (first_wave_is_forward) - { - return static_cast(stage); - } - else - { - // Perf note: runtime modulo by `stream_stages` (a runtime, not compile-time-constant, divisor; not host-known -- - // it derives from the per-CTA `stream_slots`). Cold: reached only while priming the first wave (`stream_stages` - // times, once per streaming CTA), feeding async TMA copies that mostly overlap sibling copies (fallback prime) or - // resident folds (interleaved prime). Lowest priority; flagged so a future bottleneck search finds every `/`,`%`. - const offset_t stages = static_cast(stream_stages); - const offset_t base = layout.overflow_chunks - stages; // window of the last `stream_stages` chunks - const offset_t chunk = base + (static_cast(stage) + stages - base % stages) % stages; - _CCCL_ASSERT(chunk % stages == static_cast(stage), - "reverse first-wave chunk must map back to its stage"); - return chunk; - } - } - - // Prime the streaming slots: issue this direction's first `stream_stages` overflow copies, arming the reload - // pipeline. Idempotent (`stream_is_primed`) and a no-op without overflow. Single caller - // `load_and_histogram_first_pass`: its resident path interleaves the prime inline (so this call then no-ops), leaving - // only the no-resident-chunks path to actually issue here. That path already follows the kernel-init barrier, so the - // leading `__syncthreads()` is redundant today; it is kept so the prime stays phase-safe if a future caller runs it - // right after a resident load (every thread must clear that load's final `wait_stage` before these copies re-arm the - // shared mbarriers, else the phase advances twice and a lagging thread spins forever). - _CCCL_DEVICE _CCCL_FORCEINLINE void prime_overflow_stream() - { - if (stream_is_primed || layout.overflow_chunks == 0) - { - return; - } - // Reaching here is the once-per-segment prime: the guard rules out a re-prime, and the resident load drives the - // shared stages via its own bulk copies (never the stream mask), so the stream must be fully drained here. - _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, - "overflow stream priming must start from a fully drained stream"); - __syncthreads(); - // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. - _CCCL_PRAGMA_NOUNROLL() - for (int stage = 0; stage < stream_stages; ++stage) - { - const offset_t overflow_idx = first_wave_chunk_for_stage(stage); - _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); - char* const dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - issue_bulk_copy(stage, dst, chunk_bulk_src(layout.overflow_base, overflow_idx)); - stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); - } - stream_is_primed = true; - } - // Consume the `i`-th visit on the block-load path (its ping-pong-ordered position is `overflow_idx`): wait for its // slot, fold its keys via `block_apply`, then prefetch the chunk `stream_stages` visits ahead into the just-freed // slot (a barrier orders the block's read ahead of the overwriting copy). Returns false once `should_continue()` @@ -1128,13 +1070,9 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void run_overflow_pass( BlockApply&& block_apply, GenericApply&& generic_apply, OverlapWork&& overlap_work, Continue&& should_continue) { - // A pass begins drained, except the first, where `load_and_histogram_first_pass` early-primed the first wave (all - // `stream_stages` stages in flight) to overlap its edge work; phase 1 waits on them. Later passes re-enter drained: - // the reused turn-around chunks are resident, not in flight. - // 64-bit shift keeps a full 32-stage mask defined; the `uint32` mask promotes to match. - _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0} - || stream_inflight_mask == ((::cuda::std::uint64_t{1} << stream_stages) - ::cuda::std::uint64_t{1}), - "an overflow pass must begin drained or holding only the early-primed first wave"); + // Every pass begins drained: `load_and_histogram_first_pass` fully drains the first wave before returning, and each + // pass reuses the prior pass's turn-around chunks resident (not in flight). + _CCCL_ASSERT(stream_inflight_mask == ::cuda::std::uint32_t{0}, "an overflow pass must begin drained"); if (layout.overflow_chunks == 0) { overlap_work(); @@ -1149,15 +1087,14 @@ private: _CCCL_ASSERT(stream_is_primed, "run_overflow_pass requires the overflow stream primed by load_and_histogram_first_pass"); - // Phase 1: consume the first `stream_stages` visits and issue this pass's reload wave into the freed slots. Their - // slots hold the early-primed wave (still in flight) on the first pass and the previous pass's reused turn-around - // chunks (already resident) on later passes; `consume_overflow_visit` waits only when the stage is in flight. + // Phase 1: consume the first `stream_stages` visits and issue this pass's reload wave into the freed slots. Every + // pass enters drained, so these slots always hold the prior pass's reused turn-around chunks (already resident); + // `consume_overflow_visit` issues but never waits here, and phase 2 waits on the reloads it launches. // Rolling ring pointer carried across both consume loops so no visit pays a runtime `% stream_stages`: forward // starts at slot 0, reverse at the last chunk's slot (so this modulo runs only on reverse passes). That seed is // the only runtime modulo left in these consume loops in release builds (the per-visit ones became the rolling // step; a debug-only assert still evaluates one): a runtime, not compile-time-constant, divisor, at most once per - // pass -- cold (on pass 0 it overlaps the in-flight first wave; later passes enter drained, so it is a cheap - // one-off). Lowest priority. `consume_overflow_visit` then steps the pointer with compare-select. + // pass and cold. Lowest priority. `consume_overflow_visit` then steps the pointer with compare-select. int stage = stream_is_forward ? 0 : static_cast((layout.overflow_chunks - 1) % static_cast(stream_stages)); bool is_stopped = false; @@ -2280,182 +2217,259 @@ private: if constexpr (use_block_load_to_shared) { - if (layout.my_resident_chunks > 0) + // Resident load geometry (meaningful only when this rank owns resident chunks). Chunks are written densely in + // slot order from offset 0 at the fixed stride `chunk_index * ChunkBytes` and read back in the same order, so the + // fold needs no stored per-chunk state. `prologue_stages` shared mbarrier stages carry the resident pipeline; + // `stage_base`/`stage_rot` only relabel which stage each resident chunk uses (the packing and the ring invariant + // `stage == chunk % stream_stages` are untouched) so the interleaved first-wave stream prime lands in forward + // consume order and a misaligned reload tail still drains stages 0..prologue_stages-1 in order: + // * `stage_base` shifts the window to `[stage_base, stage_base + prologue_stages)`, freeing the low stream + // stages for + // the up-front prime. Forward first wave with `stream_stages > prologue_stages` only (which forces + // `prologue_stages == my_resident_chunks`, so there is no reload). Reverse keeps it 0 (its first-consumed + // stage is data-dependent). + // * `stage_rot` rotates the resident cycle (mod `prologue_stages`), only when the stream fits the resident cycle + // (`stream_stages <= prologue_stages`) and the tail is misaligned. Its one-time `% prologue_stages` is + // TMA-dwarfed. + const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; + const int prologue_stages = + layout.my_resident_chunks > 0 + ? (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)) + : 0; + int stage_base = 0; + if constexpr (first_wave_is_forward) { - // Stage mbarriers and the `load_phase` parity are shared with the overflow stream (no per-chunk token array - // needed). Chunks are written densely in slot order from offset 0 and read back in the same order; the fixed - // slot stride (`chunk_index * ChunkBytes`) positions both the write and the fold with no running cursor or - // dynamically-indexed `pending_spans` array (which would anchor surrounding state to local memory). Every chunk - // begins on a `load_align` boundary (zero prefix), so its aligned bulk is `round_down(count, - // load_align_items)`: the whole `count` for interior chunks (`== chunk_items`, a full slot), and for the - // global-last chunk its `count` minus the unaligned suffix that is always peeled into `edge_keys` (the only - // chunk that can be short -- and, folded in ascending order, always the last resident chunk). - const int prologue = (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)); - - // Prime the overflow stream while the resident load still runs to hide its latency, over two disjoint stage - // sets: the `prologue` mbarrier stages the resident load cycles through (re-armed for the stream as they drain - // in the tail loop below) and the stages it never touches (fresh, primed up front). Together they cover every - // stream stage, so only the no-resident-chunks path leaves priming to the post-loop call. - const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; - - // Place the resident load's `prologue`-wide mbarrier window so the stream's first wave is primed in forward - // consume order (stage `s`, holding chunk `s`, first), removing the ordering inversion where the up-front prime - // would otherwise issue the last-consumed chunks first. Two complementary knobs, at most one non-zero - // (`stream_stages > prologue` vs `<=`); both only relabel which mbarrier a chunk uses (the packing and the - // stream ring invariant `stage == chunk % stream_stages` are untouched): - // * `stage_base` shifts the window to `[stage_base, stage_base + prologue)`, leaving the low stream stages - // `[0, stage_base)` free to prime up front. Forward first wave with `stream_stages > prologue` only (which - // forces `prologue == my_resident_chunks`, so there is no reload). Reverse keeps it 0 (its first-consumed - // stage is data-dependent, so no static low-stage reservation helps). - // * `stage_rot` rotates the resident cycle (mod `prologue`) so a misaligned reload tail still drains stages - // 0,1,...,prologue-1 in order. Only when the stream stages fit the resident cycle (`stream_stages <= - // prologue`) and the tail is misaligned (`my_resident_chunks % prologue != 0`, implying `prologue == - // PipelineStages`). Its `% prologue` is a one-time runtime, non-power-of-two modulo, TMA-dwarfed. - int stage_base = 0; - if constexpr (first_wave_is_forward) + if (interleave_overflow_prime && layout.my_resident_chunks > 0 && stream_stages > prologue_stages) { - if (interleave_overflow_prime && stream_stages > prologue) - { - stage_base = stream_stages - prologue; - } + stage_base = stream_stages - prologue_stages; } - const offset_t tail_align = layout.my_resident_chunks % static_cast(prologue); - const offset_t stage_rot = - (interleave_overflow_prime && stream_stages <= prologue && tail_align != offset_t{0}) - ? static_cast(prologue) - tail_align - : offset_t{0}; - - // Issue every up-front bulk copy from one TMA site, over two disjoint stage sets. The first `prologue` - // iterations load the resident chunks, densely packed in slot order over the mbarrier window - // `[stage_base, stage_base + prologue)` (`stage_res`: unit-stride, select-wrap to `stage_base`, avoiding a - // `% prologue`, a real division since `prologue` is a runtime value). The remaining `stream_stages - prologue` - // iterations (only when - // interleaving) prime the complement stages the resident load never touches with their first-wave chunk into - // the disjoint stream region; those mbarriers are fresh, so the copies run concurrently with the resident bulks - // (the shared resident stages are primed as the read loop below frees them). `stage_base` steers the free set - // to the first-consumed stages (low for a forward first wave, high for reverse), matching prime order to - // consume order. + } + const offset_t tail_align = + prologue_stages > 0 ? layout.my_resident_chunks % static_cast(prologue_stages) : offset_t{0}; + const offset_t stage_rot = + (interleave_overflow_prime && prologue_stages > 0 && stream_stages <= prologue_stages + && tail_align != offset_t{0}) + ? static_cast(prologue_stages) - tail_align + : offset_t{0}; + + // First-wave stage -> chunk mapping, shared by the up-front prime and the merged loop's resident re-arms below. + // The wave occupies the last `stream_stages` overflow chunks, each landing in stage `chunk % stream_stages`. + // Forward keeps stage `s` -> chunk `s`. Reverse walks the window `[overflow_chunks - stream_stages, + // overflow_chunks)`; since +1 in stage is +1 in chunk there (with wrap), each prime site seeds a ring counter + // once (one modulo) at its first stage and steps it unit-stride, replacing the per-copy modulo. + const offset_t fw_stages = static_cast(stream_stages); + const offset_t fw_base = layout.overflow_chunks - fw_stages; + + // The up-front prime and the merged loop's resident re-arms below share one first-wave ring counter, seeded once + // here in prime order and stepped unit-stride (`fw_base`/`fw_stages`: +1 with wrap). `prime_complement` is the + // sub-case A predicate -- the wave is wider than the resident window, so the up-front loop primes the complement + // stages the resident load never touches, walking the counter to `fw(stage_base)` for the merged loop to continue + // from. Forward primes upward from stage 0; reverse's complement starts at stage `prologue_stages`. Without a + // complement (forward, or reverse sub-case B) the merged loop consumes the counter first, so it must already sit + // at `fw(stage_base)`: 0 for forward, `fw(0)` for reverse. + const bool prime_complement = interleave_overflow_prime && stream_stages > prologue_stages; + offset_t first_wave_chunk = + first_wave_is_forward + ? offset_t{0} + : fw_base + + ((prime_complement ? static_cast(prologue_stages) : offset_t{0}) + fw_stages + - fw_base % fw_stages) + % fw_stages; + + { + // Issue every up-front bulk copy from one TMA site, over two disjoint stage sets. The first `prologue_stages` + // iterations load the resident chunks over the mbarrier window `[stage_base, stage_base + prologue_stages)` + // (`resident_stage`: unit-stride, select-wrap to `stage_base`, avoiding a `% prologue_stages`, a real division + // since `prologue_stages` is a runtime value). The remaining `stream_stages - prologue_stages` iterations (only + // when interleaving) prime the complement stages the resident load never touches with their first-wave chunk + // into the disjoint stream region; those mbarriers are fresh, so the copies run concurrently with the resident + // bulks (the shared resident stages are primed as the merged fold loop below frees them). `stage_base` steers + // the free set to the first-consumed stages (low for a forward first wave, high for reverse), matching prime + // order to consume order. With no resident chunks `prologue_stages == 0`, so every iteration takes the + // complement path and this loop alone arms the whole first wave (stages `0..stream_stages-1`). // // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. - const int upfront_copies = (interleave_overflow_prime && stream_stages > prologue) ? stream_stages : prologue; - int stage_res = stage_base + static_cast(stage_rot); + _CCCL_ASSERT(!interleave_overflow_prime || stream_is_forward == first_wave_is_forward, + "first-wave prime ran in an unexpected stream direction"); + _CCCL_ASSERT(!interleave_overflow_prime || static_cast(stream_stages) <= layout.overflow_chunks, + "first wave assumes at least `stream_stages` overflow chunks"); + const int upfront_copies = prime_complement ? stream_stages : prologue_stages; + int resident_stage = stage_base + static_cast(stage_rot); _CCCL_PRAGMA_NOUNROLL() for (int idx = 0; idx < upfront_copies; ++idx) { - // Each branch selects only its `dst` and chunk source; the shared `chunk_bulk_src` and bulk copy issue once, - // uniformly, below. + // Each branch selects only its destination slot and chunk source; the shared `chunk_bulk_src` and bulk copy + // issue once, uniformly, below. int stage = 0; - char* dst = nullptr; + int dst_slot = 0; offset_t src_base = 0; offset_t src_index = 0; - if (idx < prologue) + if (idx < prologue_stages) { const offset_t local_chunk = static_cast(idx); _CCCL_ASSERT( - stage_res == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), + resident_stage + == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue_stages)), "resident issue stage desynced from its rolling counter"); - stage = stage_res; - dst = key_slots + idx * ChunkBytes; + stage = resident_stage; + dst_slot = idx; src_base = layout.resident_base; src_index = local_chunk; - if (++stage_res == stage_base + prologue) + if (++resident_stage == stage_base + prologue_stages) { - stage_res = stage_base; + resident_stage = stage_base; } } else { - // Map the complement index `idx - prologue` around the resident window: low stages `[0, stage_base)` map to - // themselves, the rest skip the window (`+ prologue`, i.e. back to `idx`). - stage = (idx - prologue < stage_base) ? idx - prologue : idx; - _CCCL_ASSERT(stage < stage_base || stage >= stage_base + prologue, + // Map the complement index `idx - prologue_stages` around the resident window: low stages `[0, stage_base)` + // map to themselves, the rest skip the window (`+ prologue_stages`, i.e. back to `idx`). + stage = (idx - prologue_stages < stage_base) ? idx - prologue_stages : idx; + _CCCL_ASSERT(stage < stage_base || stage >= stage_base + prologue_stages, "up-front prime stage must lie outside the resident window"); - const offset_t overflow_idx = first_wave_chunk_for_stage(stage); - _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); - _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, - "cannot issue a load into a streaming stage that is still in flight"); - dst = key_slots + (stream_slot_base + stage) * ChunkBytes; + _CCCL_ASSERT( + first_wave_chunk < layout.overflow_chunks && first_wave_chunk % fw_stages == static_cast(stage), + "first-wave chunk must map back to its stage"); + dst_slot = stream_slot_base + stage; src_base = layout.overflow_base; - src_index = overflow_idx; - stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); + src_index = first_wave_chunk; + first_wave_chunk = + (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; } - issue_bulk_copy(stage, dst, chunk_bulk_src(src_base, src_index)); + issue_bulk_copy(stage, key_slots + dst_slot * ChunkBytes, chunk_bulk_src(src_base, src_index)); } + } - // Chunk `local_chunk` was written at its fixed slot offset `local_chunk * ChunkBytes` (same stride the issue - // path uses), and `chunk_bulk_src` recomputes its length, so the fold needs no stored per-chunk state. - // Same unit-stride rolling counter as the issue loop over `[stage_base, stage_base + prologue)`: increment and - // select-wrap to `stage_base`, no runtime `% prologue`. Starts at `stage_base + stage_rot`; the assert pins it - // to the closed form. - int stage = stage_base + static_cast(stage_rot); - _CCCL_PRAGMA_NOUNROLL() - for (offset_t local_chunk = 0; local_chunk < layout.my_resident_chunks; ++local_chunk) + // Fold this rank's chunks into the first-pass histogram through one loop and one bulk-copy site: the resident + // chunks first (`chunk < my_resident_chunks`), then the overflow chunks. Each iteration folds exactly one chunk + // and optionally refills one stage's pipeline; the `fold_is_resident` x `refill_kind` classification below names + // the five states it can be in. The stream is fully primed before the overflow phase (the up-front loop above, + // plus this loop's resident re-arms) and the histogram never early-stops, so overflow visits only wait/fold/ + // prefetch and the pass ends drained. Edges are folded after. `resident_stage` cycles the resident window + // `[stage_base, stage_base + prologue_stages)`, `streaming_stage` the stream ring `[0, stream_stages)`; both step + // unit-stride with a select-wrap (no runtime modulo). + enum class refill_kind + { + none, + resident, + stream + }; + const offset_t total_chunks = layout.my_resident_chunks + layout.overflow_chunks; + int resident_stage = stage_base + static_cast(stage_rot); + int streaming_stage = + (interleave_overflow_prime && !first_wave_is_forward) + ? static_cast((layout.overflow_chunks - offset_t{1}) % static_cast(stream_stages)) + : 0; + // `first_wave_chunk` carries over from the up-front prime, which left it at `fw(stage_base)` (sub-case A walked + // it there through the complement; otherwise it was seeded there directly). The resident re-arms below fire on + // consecutive iterations starting at `stage_base` -- the last `prologue_stages` resident chunks occupy the window + // `[stage_base, stage_base + prologue_stages)` in +1 order -- so they consume the counter unit-stride from here. + _CCCL_PRAGMA_NOUNROLL() + for (offset_t chunk = offset_t{0}; chunk < total_chunks; ++chunk) + { + // Fold axis: pick the chunk this iteration folds, the stage it occupies, and step that pipeline's cursor. + bool fold_is_resident = true; + int stage = 0; + offset_t fold_chunk = 0; // resident: local chunk index; stream: ping-pong overflow chunk index + offset_t overflow_visit = 0; // stream only: 0-based position in the overflow sequence + if (chunk < layout.my_resident_chunks) { + fold_is_resident = true; + stage = resident_stage; + fold_chunk = chunk; _CCCL_ASSERT( - stage == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue)), + stage == stage_base + static_cast((chunk + stage_rot) % static_cast(prologue_stages)), "resident read stage desynced from its rolling counter"); - wait_stage(stage); - const int read_len_items = - static_cast(::cuda::std::size(chunk_bulk_src(layout.resident_base, local_chunk))); - // Fixed-stride packing invariant: only the segment's global-last chunk can be short, and (folded ascending) - // it is always the last resident chunk, so every earlier resident chunk fills its slot. - _CCCL_ASSERT(local_chunk + offset_t{1} == layout.my_resident_chunks || read_len_items == chunk_items, - "only the last resident chunk may carry a short aligned bulk"); - for_each_chunk_key( - ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), - read_len_items, - add_first_pass); - - const offset_t next_local_chunk = local_chunk + static_cast(prologue); - // Inline both re-arm dispatches (resident successor vs. first-wave stream prime) so each only selects its - // `dst` and chunk source (`src_base`/`src_index`); the shared `chunk_bulk_src` and the single (force-inlined) - // bulk copy then run once, uniformly, below. - char* dst = nullptr; - offset_t src_base = 0; - offset_t src_index = 0; - bool rearm = false; - if (next_local_chunk < layout.my_resident_chunks) - { - // Resident successor: densely packed at its fixed slot offset (constant `ChunkBytes` stride). - dst = key_slots + static_cast(next_local_chunk) * ChunkBytes; - src_base = layout.resident_base; - src_index = next_local_chunk; - rearm = true; - } - else if (interleave_overflow_prime && stage < stream_stages) - { - // Shared resident stage (in `[stage_base, stage_base + prologue)`) with no resident successor: re-arm it - // with its first-wave chunk (`first_wave_chunk_for_stage`, direction-aware) into the disjoint stream region - // (`stream_slot_base` onward), never the densely-packed resident data the later tail folds still read. - // Stages `>= stream_stages` are extra resident pipeline depth with no stream slot, so they idle - // (`stage_rot` places them last so all primes issue first). - const offset_t overflow_idx = first_wave_chunk_for_stage(stage); - _CCCL_ASSERT(overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); - _CCCL_ASSERT((stream_inflight_mask & (::cuda::std::uint32_t{1} << stage)) == 0, - "cannot issue a load into a streaming stage that is still in flight"); - dst = key_slots + (stream_slot_base + stage) * ChunkBytes; - src_base = layout.overflow_base; - src_index = overflow_idx; - stream_inflight_mask |= (::cuda::std::uint32_t{1} << stage); - rearm = true; - } - if (rearm) + resident_stage = (resident_stage + 1 == stage_base + prologue_stages) ? stage_base : resident_stage + 1; + } + else + { + fold_is_resident = false; + overflow_visit = chunk - layout.my_resident_chunks; + stage = streaming_stage; + fold_chunk = first_wave_is_forward ? overflow_visit : (layout.overflow_chunks - offset_t{1} - overflow_visit); + _CCCL_ASSERT(stage == static_cast(fold_chunk % static_cast(stream_stages)), + "overflow stage desynced from its rolling ring pointer"); + // First-pass direction is the compile-time `first_wave_is_forward` (`stream_is_forward` flips only after). + if constexpr (first_wave_is_forward) { - // Phase safety, not data safety (the target is fresh/disjoint): re-arming this stage before all threads - // leave the wait above would advance the phase twice, stranding a lagging waiter forever. - __syncthreads(); - issue_bulk_copy(stage, dst, chunk_bulk_src(src_base, src_index)); + streaming_stage = (streaming_stage + 1 == stream_stages) ? 0 : streaming_stage + 1; } - if (++stage == stage_base + prologue) + else { - stage = stage_base; + streaming_stage = (streaming_stage == 0) ? stream_stages - 1 : streaming_stage - 1; } } + // Refill axis (flat over the four refilling states; the remaining two fall through to `none`): a resident stage + // refills its resident successor, or -- lacking one -- re-arms its shared stream slot with the first overflow + // wave; a stream stage prefetches the chunk `stream_stages` visits ahead. `refill_chunk` is a resident local + // chunk index for `resident`, an overflow chunk index for `stream`. + refill_kind refill = refill_kind::none; + offset_t refill_chunk = 0; + if (fold_is_resident && chunk + static_cast(prologue_stages) < layout.my_resident_chunks) + { + refill = refill_kind::resident; + refill_chunk = chunk + static_cast(prologue_stages); + } + else if (fold_is_resident && interleave_overflow_prime && stage < stream_stages) + { + _CCCL_ASSERT(first_wave_chunk % fw_stages == static_cast(stage), + "first-wave chunk must map back to its stage"); + refill = refill_kind::stream; + refill_chunk = first_wave_chunk; + first_wave_chunk = + (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; + } + else if (!fold_is_resident && overflow_visit + static_cast(stream_stages) < layout.overflow_chunks) + { + const offset_t next_step = overflow_visit + static_cast(stream_stages); + refill = refill_kind::stream; + refill_chunk = first_wave_is_forward ? next_step : (layout.overflow_chunks - offset_t{1} - next_step); + } + _CCCL_ASSERT(refill != refill_kind::stream || refill_chunk < layout.overflow_chunks, + "overflow chunk index out of range"); + + // Fold: wait the classified slot, then histogram it. Resident chunks sit at their dense slot, stream chunks at + // this stage's stream slot. + const int fold_slot = fold_is_resident ? static_cast(fold_chunk) : stream_slot_base + stage; + const offset_t fold_base = fold_is_resident ? layout.resident_base : layout.overflow_base; + wait_stage(stage); + const key_t* fold_src = ::cuda::ptr_rebind(key_slots + fold_slot * ChunkBytes); + const int fold_len = static_cast(::cuda::std::size(chunk_bulk_src(fold_base, fold_chunk))); + // Only the segment's global-last chunk can be short; folded ascending it is the last resident chunk (overflow + // chunks fold after and may also be short), so every earlier resident chunk fills its slot. + _CCCL_ASSERT(!fold_is_resident || chunk + offset_t{1} == layout.my_resident_chunks || fold_len == chunk_items, + "only the last resident chunk may carry a short aligned bulk"); + for_each_chunk_key(fold_src, fold_len, add_first_pass); + + // Refill: issue this stage's next copy. A resident successor lands at its dense slot; either stream case (a + // first-wave re-arm or a reload) lands at this stage's stream slot. + if (refill != refill_kind::none) + { + const bool refill_resident = refill == refill_kind::resident; + const int refill_slot = refill_resident ? static_cast(refill_chunk) : stream_slot_base + stage; + const offset_t refill_base = refill_resident ? layout.resident_base : layout.overflow_base; + // Phase safety: re-arming before all threads clear the wait above would advance the mbarrier phase twice, + // stranding a lagging waiter. + __syncthreads(); + issue_bulk_copy(stage, key_slots + refill_slot * ChunkBytes, chunk_bulk_src(refill_base, refill_chunk)); + } + } + + if (interleave_overflow_prime) + { + // Every visit above waited its own prefetch, so the first pass ends with all copies complete; the direction + // flip and `stream_is_primed` carry to the radix passes and final filter, which reuse the persistent stream. + stream_is_forward = !stream_is_forward; + stream_is_primed = true; + } + + if (layout.my_resident_chunks > 0) + { // Resident region extent for the later passes (one contiguous run of aligned bulks; both boundary edges live in - // `edge_keys`). Inferred once here rather than carried as a loop cursor: the fixed slot stride makes it - // `(my_resident_chunks - 1)` full `chunk_items` slots plus the last chunk's (possibly short) aligned bulk. The - // sum is a `load_align_items` multiple because each term is. + // `edge_keys`): `(my_resident_chunks - 1)` full `chunk_items` slots plus the last chunk's (possibly short) + // aligned bulk. The sum is a `load_align_items` multiple because each term is. const offset_t last_bulk_items = static_cast( ::cuda::std::size(chunk_bulk_src(layout.resident_base, layout.my_resident_chunks - offset_t{1}))); const int resident_items = static_cast( @@ -2463,25 +2477,12 @@ private: _CCCL_ASSERT(resident_items % load_align_items == 0, "the resident region must end on a load-alignment boundary"); resident_keys = smem_keys_t(::cuda::ptr_rebind(key_slots), resident_items); - - // The up-front and tail primes together armed every stream stage, so the stream is primed; the post-loop call - // no-ops. - if (interleave_overflow_prime) - { - _CCCL_ASSERT(stream_inflight_mask == ((::cuda::std::uint64_t{1} << stream_stages) - ::cuda::std::uint64_t{1}), - "interleaved prime must arm exactly the first wave across all stream stages"); - stream_is_primed = true; - } } - // Prime the first overflow wave (a no-op when the resident load above already interleaved it; still primes the - // no-resident-chunks path) so the edge staging/folding below -- independent of the streaming slots -- runs under - // the in-flight TMA copies. - prime_overflow_stream(); - - // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass in the same sweep (see - // `stage_and_fold_edge`: each thread folds keys it just wrote, so no barrier is needed here). The head prefix - // (rank 0) precedes chunk 0; the peeled tail suffix (tail owner, always when unaligned) trails the last chunk. + // Stage the persistent boundary edges into `edge_keys` and fold them into the first pass, after the whole + // resident+overflow loop (the histogram is order-independent). `stage_and_fold_edge` folds keys each thread just + // wrote, so no barrier is needed here. The head prefix (rank 0) precedes chunk 0; the peeled tail suffix (tail + // owner, always when unaligned) trails the last chunk. if (layout.head_edge_len_items > 0) { stage_and_fold_edge(temp_storage.edge_keys, layout.block_keys_base, layout.head_edge_len_items, add_first_pass); @@ -2521,15 +2522,11 @@ private: } } } - } - // Fold the overflow chunks into the first-pass histogram. On the block-load path the stream was primed above (in - // its preselected initial direction; the histogram is order-independent, so the direction only sets up the leftover - // parity for the final filter) and reuses the resident load's stage mbarriers (all front-loaded at `run` entry), - // with `wait_stage` for producer/consumer sync; the generic fallback re-reads each overflow chunk from gmem. - // No overlap work: this fused first pass already folded its resident keys during the load, so only the overflow - // chunks remain to fold here. - fold_overflow_keys(add_first_pass, [] {}); + // Fold the overflow chunks straight from gmem (the block-load path folds them in the merged loop above, reusing + // the resident load's stage mbarriers). No overlap work: this pass already folded its resident keys above. + fold_overflow_keys(add_first_pass, [] {}); + } const int resident_count = static_cast(resident_keys.size()); _CCCL_ASSERT(resident_count == 0 || static_cast(resident_count) <= block_tile_capacity, diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index af15cbe0947..20391d1cf27 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -1615,6 +1615,97 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with REQUIRE(expected_keys == keys_out_buffer); } + +// Forward first-wave `stage_base` prime (`stream_stages > prologue` *with* a resident chunk present). Reuses the +// schedule sweep's cap-1 / stages-4 kernel (same `cluster_tuning_selector<1, 4, 0, .., 4>` type and static bounds -> no +// new instantiation); only the runtime segment size differs. Sized to 7 chunks so, at 4 slots, streaming reserves 3 (= +// min(stages, excess=3)) and 1 resident chunk remains: `stream_stages(3) > prologue(1)` yet `my_resident > 0`, so a +// forward wave sets `stage_base = 2` and the merged loop's stream re-arm must seed its ring counter at +// `fw(stage_base)`, not `fw(0)`. The schedule sweep never reaches this (its `>` configs all collapse to `my_resident == +// 0` pure streaming, which guards `stage_base` back to 0). Non-deterministic combos drive the forward wave (the new +// coverage); the deterministic combos exercise the already-working reverse counterpart (`stage_base == 0`). +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with a resident chunk below a wider stream", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + skip_if_batched_topk_cluster_unavailable(true); + + const segment_size_t segment_size = 896 - 31; // 7 chunks (128 items each), unaligned -> peeled overflow tail edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/4>{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} # endif // TEST_LAUNCH == 0 // Effective-cluster-width coverage for pairs: the large `deferred_sequence` bound sizes a wide (16-CTA) launch, but the From 983a8aa4338929c2965e7a5a5f92611515fd116a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 14 Jul 2026 08:46:28 +0200 Subject: [PATCH 121/126] Fix pipeline ordering --- cub/cub/agent/agent_batched_topk_cluster.cuh | 249 ++++++++++++------ ...catch2_test_device_segmented_topk_pairs.cu | 186 ++++++++++++- 2 files changed, 349 insertions(+), 86 deletions(-) diff --git a/cub/cub/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh index cfb2feb294f..ec4f22e7083 100644 --- a/cub/cub/agent/agent_batched_topk_cluster.cuh +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -2219,74 +2219,100 @@ private: { // Resident load geometry (meaningful only when this rank owns resident chunks). Chunks are written densely in // slot order from offset 0 at the fixed stride `chunk_index * ChunkBytes` and read back in the same order, so the - // fold needs no stored per-chunk state. `prologue_stages` shared mbarrier stages carry the resident pipeline; - // `stage_base`/`stage_rot` only relabel which stage each resident chunk uses (the packing and the ring invariant - // `stage == chunk % stream_stages` are untouched) so the interleaved first-wave stream prime lands in forward - // consume order and a misaligned reload tail still drains stages 0..prologue_stages-1 in order: - // * `stage_base` shifts the window to `[stage_base, stage_base + prologue_stages)`, freeing the low stream - // stages for - // the up-front prime. Forward first wave with `stream_stages > prologue_stages` only (which forces - // `prologue_stages == my_resident_chunks`, so there is no reload). Reverse keeps it 0 (its first-consumed - // stage is data-dependent). - // * `stage_rot` rotates the resident cycle (mod `prologue_stages`), only when the stream fits the resident cycle - // (`stream_stages <= prologue_stages`) and the tail is misaligned. Its one-time `% prologue_stages` is - // TMA-dwarfed. + // fold needs no stored per-chunk state. `prologue_stages` shared mbarrier stages carry the resident pipeline. The + // stream ring (`stream_stages` stages) and the resident window share the same mbarrier array; the first wave is + // primed in consume order so the earliest-consumed copies issue first (see the direction-aware layout below). const bool interleave_overflow_prime = layout.overflow_chunks > offset_t{0}; const int prologue_stages = layout.my_resident_chunks > 0 ? (::cuda::std::min) (PipelineStages, static_cast(layout.my_resident_chunks)) : 0; + + // First wave = the `stream_stages` overflow chunks resident when the overflow phase begins (forward chunks + // `[0, stream_stages)`, reverse chunks `[fw_base, overflow_chunks)`); a chunk's stage is always `chunk % + // stream_stages`. + const offset_t fw_stages = static_cast(stream_stages); + const offset_t fw_base = layout.overflow_chunks - fw_stages; + + // Distinct mbarrier stages in flight: the stream ring and the resident window share mbarriers, so the cycle spans + // whichever is wider -- sub-case A (`stream_stages > prologue_stages`) nests the resident window in the stream + // ring; sub-case B nests the stream ring in the resident window. + const int stage_cycle = (::cuda::std::max) (stream_stages, prologue_stages); + + // Reverse consumes the first wave as stages descending from `reverse_first_stage`; forward ascending from 0. We + // prime in that same order. The `stream_stages - prologue_stages` early-consumed stages are the "complement", + // primed up-front onto fresh mbarriers for maximal overlap; the `prologue_stages` late-consumed stages are the + // "resident window", primed as the resident fold frees their shared mbarriers. + const int reverse_first_stage = + interleave_overflow_prime ? static_cast((layout.overflow_chunks - offset_t{1}) % fw_stages) : 0; + + // `stage_base` places the resident window on the late-consumed stages. Forward: shift up by the early-consumed + // count `stream_stages - prologue_stages` (nonzero only when `stream_stages > prologue_stages`, which forces + // `prologue_stages == my_resident_chunks`, so no reload). Reverse: the late-consumed stages form the cyclic + // window starting at `(reverse_first_stage + 1) % stage_cycle` (== `overflow_chunks % stream_stages` for sub-case + // A). int stage_base = 0; if constexpr (first_wave_is_forward) { - if (interleave_overflow_prime && layout.my_resident_chunks > 0 && stream_stages > prologue_stages) + if (interleave_overflow_prime && prologue_stages > 0 && stream_stages > prologue_stages) { stage_base = stream_stages - prologue_stages; } } + else + { + if (prologue_stages > 0) + { + stage_base = static_cast( + (static_cast(reverse_first_stage) + offset_t{1}) % static_cast(stage_cycle)); + } + } + + // Resident chunk -> mbarrier-stage rotation, chosen so the last `prologue_stages` resident chunks free the + // resident window in consume order as they fold. Forward walks the window `+1` (ascending); `stage_rot` corrects + // a misaligned reload tail so the last chunks still land `[stage_base, stage_base + prologue_stages)` in order. + // Reverse walks the window `-1` (descending) from its top; `reverse_cycle_seed` seeds the cycle position so chunk + // `my_resident-1` frees `stage_base`. const offset_t tail_align = prologue_stages > 0 ? layout.my_resident_chunks % static_cast(prologue_stages) : offset_t{0}; const offset_t stage_rot = - (interleave_overflow_prime && prologue_stages > 0 && stream_stages <= prologue_stages + (first_wave_is_forward && interleave_overflow_prime && prologue_stages > 0 && stream_stages <= prologue_stages && tail_align != offset_t{0}) ? static_cast(prologue_stages) - tail_align : offset_t{0}; + const int reverse_cycle_seed = + prologue_stages > 0 + ? static_cast((layout.my_resident_chunks - offset_t{1}) % static_cast(prologue_stages)) + : 0; - // First-wave stage -> chunk mapping, shared by the up-front prime and the merged loop's resident re-arms below. - // The wave occupies the last `stream_stages` overflow chunks, each landing in stage `chunk % stream_stages`. - // Forward keeps stage `s` -> chunk `s`. Reverse walks the window `[overflow_chunks - stream_stages, - // overflow_chunks)`; since +1 in stage is +1 in chunk there (with wrap), each prime site seeds a ring counter - // once (one modulo) at its first stage and steps it unit-stride, replacing the per-copy modulo. - const offset_t fw_stages = static_cast(stream_stages); - const offset_t fw_base = layout.overflow_chunks - fw_stages; - - // The up-front prime and the merged loop's resident re-arms below share one first-wave ring counter, seeded once - // here in prime order and stepped unit-stride (`fw_base`/`fw_stages`: +1 with wrap). `prime_complement` is the - // sub-case A predicate -- the wave is wider than the resident window, so the up-front loop primes the complement - // stages the resident load never touches, walking the counter to `fw(stage_base)` for the merged loop to continue - // from. Forward primes upward from stage 0; reverse's complement starts at stage `prologue_stages`. Without a - // complement (forward, or reverse sub-case B) the merged loop consumes the counter first, so it must already sit - // at `fw(stage_base)`: 0 for forward, `fw(0)` for reverse. + // Shared first-wave ring counter, seeded once in prime (== consume) order and stepped unit-stride across both the + // up-front prime and the merged fold loop's resident re-arms: forward `+1` from chunk 0, reverse `-1` from + // `overflow_chunks - 1` (== `fw(reverse_first_stage)`). `prime_complement` marks sub-case A, where the up-front + // loop primes the early-consumed complement stages first and then hands the counter to the re-arms. const bool prime_complement = interleave_overflow_prime && stream_stages > prologue_stages; - offset_t first_wave_chunk = - first_wave_is_forward - ? offset_t{0} - : fw_base - + ((prime_complement ? static_cast(prologue_stages) : offset_t{0}) + fw_stages - - fw_base % fw_stages) - % fw_stages; + // Sub-case A implies `prologue_stages < PipelineStages`, which forces `prologue_stages == my_resident_chunks`: + // the resident fold then never reloads a resident successor (`chunk + prologue_stages >= my_resident_chunks` + // always), so every last-`prologue_stages` fold re-arms a shared stream stage. + _CCCL_ASSERT(!prime_complement || layout.my_resident_chunks == static_cast(prologue_stages), + "sub-case A (stream wider than the resident window) admits no resident reload"); + offset_t first_wave_chunk = first_wave_is_forward ? offset_t{0} : (layout.overflow_chunks - offset_t{1}); + + // Reverse resident chunk -> mbarrier stage (assertion + cursor cross-check): descending cycle position + // `reverse_cycle_seed - chunk` (Euclidean) mapped into the cyclic window. Only valid when `prologue_stages > 0`. + [[maybe_unused]] const auto reverse_resident_stage = [&](offset_t chunk_index) -> int { + const offset_t cycle_len = static_cast(prologue_stages); + const offset_t cycle_pos = + (static_cast(reverse_cycle_seed) + cycle_len - chunk_index % cycle_len) % cycle_len; + const int stage_index = stage_base + static_cast(cycle_pos); + return stage_index >= stage_cycle ? stage_index - stage_cycle : stage_index; + }; { - // Issue every up-front bulk copy from one TMA site, over two disjoint stage sets. The first `prologue_stages` - // iterations load the resident chunks over the mbarrier window `[stage_base, stage_base + prologue_stages)` - // (`resident_stage`: unit-stride, select-wrap to `stage_base`, avoiding a `% prologue_stages`, a real division - // since `prologue_stages` is a runtime value). The remaining `stream_stages - prologue_stages` iterations (only - // when interleaving) prime the complement stages the resident load never touches with their first-wave chunk - // into the disjoint stream region; those mbarriers are fresh, so the copies run concurrently with the resident - // bulks (the shared resident stages are primed as the merged fold loop below frees them). `stage_base` steers - // the free set to the first-consumed stages (low for a forward first wave, high for reverse), matching prime - // order to consume order. With no resident chunks `prologue_stages == 0`, so every iteration takes the - // complement path and this loop alone arms the whole first wave (stages `0..stream_stages-1`). + // Issue every up-front bulk copy from one TMA site, in first-wave consume order. The first `prologue_stages` + // iterations load resident chunks into their dense slots over the resident-window mbarriers; the rest (sub-case + // A) prime the early-consumed complement onto fresh stream mbarriers that overlap the resident bulks (the + // resident-window stages are re-armed later, by the fold loop). No resident chunks -> this loop arms the whole + // wave. // // Force nounroll: auto-unroll replicates the full TMA issue per iteration and bloats the load path. _CCCL_ASSERT(!interleave_overflow_prime || stream_is_forward == first_wave_is_forward, @@ -2294,7 +2320,10 @@ private: _CCCL_ASSERT(!interleave_overflow_prime || static_cast(stream_stages) <= layout.overflow_chunks, "first wave assumes at least `stream_stages` overflow chunks"); const int upfront_copies = prime_complement ? stream_stages : prologue_stages; - int resident_stage = stage_base + static_cast(stage_rot); + // Resident-load mbarrier cursor (see the rotation note above): forward walks `+1` from `stage_base`, reverse + // walks the cycle position `-1` from `reverse_cycle_seed`; both select-wrap, no runtime modulo. + [[maybe_unused]] int resident_stage = first_wave_is_forward ? stage_base + static_cast(stage_rot) : 0; + [[maybe_unused]] int resident_cycle_pos = reverse_cycle_seed; _CCCL_PRAGMA_NOUNROLL() for (int idx = 0; idx < upfront_copies; ++idx) { @@ -2307,25 +2336,49 @@ private: if (idx < prologue_stages) { const offset_t local_chunk = static_cast(idx); - _CCCL_ASSERT( - resident_stage - == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue_stages)), - "resident issue stage desynced from its rolling counter"); - stage = resident_stage; + if constexpr (first_wave_is_forward) + { + _CCCL_ASSERT( + resident_stage + == stage_base + static_cast((local_chunk + stage_rot) % static_cast(prologue_stages)), + "resident issue stage desynced from its rolling counter"); + stage = resident_stage; + if (++resident_stage == stage_base + prologue_stages) + { + resident_stage = stage_base; + } + } + else + { + stage = stage_base + resident_cycle_pos; + if (stage >= stage_cycle) + { + stage -= stage_cycle; + } + _CCCL_ASSERT(stage == reverse_resident_stage(local_chunk), + "resident issue stage desynced from its rolling counter"); + resident_cycle_pos = (resident_cycle_pos == 0) ? prologue_stages - 1 : resident_cycle_pos - 1; + } dst_slot = idx; src_base = layout.resident_base; src_index = local_chunk; - if (++resident_stage == stage_base + prologue_stages) - { - resident_stage = stage_base; - } } else { - // Map the complement index `idx - prologue_stages` around the resident window: low stages `[0, stage_base)` - // map to themselves, the rest skip the window (`+ prologue_stages`, i.e. back to `idx`). - stage = (idx - prologue_stages < stage_base) ? idx - prologue_stages : idx; - _CCCL_ASSERT(stage < stage_base || stage >= stage_base + prologue_stages, + if constexpr (first_wave_is_forward) + { + // Map the complement index `idx - prologue_stages` around the resident window: low stages `[0, + // stage_base)` map to themselves, the rest skip the window (`+ prologue_stages`, i.e. back to `idx`). + stage = (idx - prologue_stages < stage_base) ? idx - prologue_stages : idx; + } + else + { + // Reverse complement stages descend from `reverse_first_stage`; the ring counter already walks them, so + // the stage is just its chunk's stage. + stage = static_cast(first_wave_chunk % fw_stages); + } + _CCCL_ASSERT(prologue_stages == 0 + || static_cast((stage - stage_base + stage_cycle) % stage_cycle) >= prologue_stages, "up-front prime stage must lie outside the resident window"); _CCCL_ASSERT( first_wave_chunk < layout.overflow_chunks && first_wave_chunk % fw_stages == static_cast(stage), @@ -2333,8 +2386,16 @@ private: dst_slot = stream_slot_base + stage; src_base = layout.overflow_base; src_index = first_wave_chunk; - first_wave_chunk = - (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; + if constexpr (first_wave_is_forward) + { + first_wave_chunk = + (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; + } + else + { + first_wave_chunk = + (first_wave_chunk == fw_base) ? fw_base + fw_stages - offset_t{1} : first_wave_chunk - offset_t{1}; + } } issue_bulk_copy(stage, key_slots + dst_slot * ChunkBytes, chunk_bulk_src(src_base, src_index)); } @@ -2345,25 +2406,24 @@ private: // and optionally refills one stage's pipeline; the `fold_is_resident` x `refill_kind` classification below names // the five states it can be in. The stream is fully primed before the overflow phase (the up-front loop above, // plus this loop's resident re-arms) and the histogram never early-stops, so overflow visits only wait/fold/ - // prefetch and the pass ends drained. Edges are folded after. `resident_stage` cycles the resident window - // `[stage_base, stage_base + prologue_stages)`, `streaming_stage` the stream ring `[0, stream_stages)`; both step - // unit-stride with a select-wrap (no runtime modulo). + // prefetch and the pass ends drained. Edges are folded after. The resident cursor cycles the resident window and + // `streaming_stage` the stream ring `[0, stream_stages)`, both stepping unit-stride with a select-wrap (no + // runtime modulo); forward ascends, reverse descends. enum class refill_kind { none, resident, stream }; - const offset_t total_chunks = layout.my_resident_chunks + layout.overflow_chunks; - int resident_stage = stage_base + static_cast(stage_rot); - int streaming_stage = - (interleave_overflow_prime && !first_wave_is_forward) - ? static_cast((layout.overflow_chunks - offset_t{1}) % static_cast(stream_stages)) - : 0; - // `first_wave_chunk` carries over from the up-front prime, which left it at `fw(stage_base)` (sub-case A walked - // it there through the complement; otherwise it was seeded there directly). The resident re-arms below fire on - // consecutive iterations starting at `stage_base` -- the last `prologue_stages` resident chunks occupy the window - // `[stage_base, stage_base + prologue_stages)` in +1 order -- so they consume the counter unit-stride from here. + const offset_t total_chunks = layout.my_resident_chunks + layout.overflow_chunks; + [[maybe_unused]] int resident_stage = first_wave_is_forward ? stage_base + static_cast(stage_rot) : 0; + [[maybe_unused]] int resident_cycle_pos = reverse_cycle_seed; + int streaming_stage = first_wave_is_forward ? 0 : static_cast(reverse_first_stage); // reverse 1st-consumed + // stage + // `first_wave_chunk` carries over from the up-front prime in consume order: sub-case A walked it through the + // early-consumed complement, leaving it at the first late-consumed chunk; otherwise it still sits at its seed. + // The last `prologue_stages` resident chunks free the resident window in consume order as they fold, so their + // re-arms consume the counter unit-stride from here (forward `+1`, reverse `-1`). _CCCL_PRAGMA_NOUNROLL() for (offset_t chunk = offset_t{0}; chunk < total_chunks; ++chunk) { @@ -2374,13 +2434,26 @@ private: offset_t overflow_visit = 0; // stream only: 0-based position in the overflow sequence if (chunk < layout.my_resident_chunks) { - fold_is_resident = true; - stage = resident_stage; - fold_chunk = chunk; - _CCCL_ASSERT( - stage == stage_base + static_cast((chunk + stage_rot) % static_cast(prologue_stages)), - "resident read stage desynced from its rolling counter"); - resident_stage = (resident_stage + 1 == stage_base + prologue_stages) ? stage_base : resident_stage + 1; + fold_chunk = chunk; + if constexpr (first_wave_is_forward) + { + stage = resident_stage; + _CCCL_ASSERT( + stage == stage_base + static_cast((chunk + stage_rot) % static_cast(prologue_stages)), + "resident read stage desynced from its rolling counter"); + resident_stage = (resident_stage + 1 == stage_base + prologue_stages) ? stage_base : resident_stage + 1; + } + else + { + stage = stage_base + resident_cycle_pos; + if (stage >= stage_cycle) + { + stage -= stage_cycle; + } + _CCCL_ASSERT(stage == reverse_resident_stage(chunk), + "resident read stage desynced from its rolling counter"); + resident_cycle_pos = (resident_cycle_pos == 0) ? prologue_stages - 1 : resident_cycle_pos - 1; + } } else { @@ -2418,8 +2491,16 @@ private: "first-wave chunk must map back to its stage"); refill = refill_kind::stream; refill_chunk = first_wave_chunk; - first_wave_chunk = - (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; + if constexpr (first_wave_is_forward) + { + first_wave_chunk = + (first_wave_chunk + offset_t{1} == fw_base + fw_stages) ? fw_base : first_wave_chunk + offset_t{1}; + } + else + { + first_wave_chunk = + (first_wave_chunk == fw_base) ? fw_base + fw_stages - offset_t{1} : first_wave_chunk - offset_t{1}; + } } else if (!fold_is_resident && overflow_visit + static_cast(stream_stages) < layout.overflow_chunks) { diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 20391d1cf27..e1b1f66c195 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -1622,8 +1622,9 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with // min(stages, excess=3)) and 1 resident chunk remains: `stream_stages(3) > prologue(1)` yet `my_resident > 0`, so a // forward wave sets `stage_base = 2` and the merged loop's stream re-arm must seed its ring counter at // `fw(stage_base)`, not `fw(0)`. The schedule sweep never reaches this (its `>` configs all collapse to `my_resident == -// 0` pure streaming, which guards `stage_base` back to 0). Non-deterministic combos drive the forward wave (the new -// coverage); the deterministic combos exercise the already-working reverse counterpart (`stage_base == 0`). +// 0` pure streaming, which guards `stage_base` back to 0). The combos that select a forward first wave drive the new +// coverage; the reverse ones here land on `stage_base == 0` (`overflow_chunks(6) % stream_stages(3) == 0`), the +// simplest reverse case (the wrapping non-zero reverse `stage_base` is pinned by the C = 5 test below). C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with a resident chunk below a wider stream", "[pairs][segmented][topk][device][cluster][determinism]", det_tie_pair_combos) @@ -1706,6 +1707,187 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with REQUIRE(expected_keys == keys_out_buffer); } + +// Reverse first-wave with a non-zero cyclic `stage_base` and a wrapping resident window. Reuses the schedule sweep's +// cap-1 / slots-4 / stages-4 kernel (no new instantiation); only the runtime segment size differs. Sized to 5 chunks so +// the single CTA reserves 1 stream slot (`excess = 1`), leaving 3 resident chunks: sub-case B (`stream_stages(1) <= +// prologue(3)`) with `overflow_chunks = 2`, so a reverse first wave (`first_wave_is_forward == false`) sets +// `stage_base = (s0+1) % stage_cycle = 1` over `stage_cycle = 3`, giving the cyclic window `[1,4) mod 3`. Resident +// chunks 0,1,2 then map to stages 0,2,1 (descending cycle position from `reverse_cycle_seed`), and the `-1` re-arm must +// free stream stage 0 in consume order. The `stage_base == 0` reverse case is covered by the wider-stream test above (C +// = 7); this pins the wrapping, non-zero-`stage_base` reverse path. Only the deterministic combos whose (tie-break, +// pass-parity) select a reverse first wave exercise it; the rest re-drive forward paths. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with a wrapping reverse resident window", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + skip_if_batched_topk_cluster_unavailable(true); + + const segment_size_t segment_size = 640 - 31; // 5 chunks (128 items each), unaligned -> peeled overflow tail edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/4>{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} + +// Partial final wave: `overflow_chunks % stream_stages != 0`, so the streaming reload branch wraps its ring on an +// uneven tail. Reuses the misaligned-tail test's cap-1 / slots-4 / stages-2 kernel (no new instantiation). Sized to 7 +// chunks so, at 4 slots / 2 stages, streaming reserves 2 (`excess = 3`, `min(stages,excess) = 2`) and 2 resident chunks +// remain: sub-case B with `overflow_chunks = 5`, `stream_stages = 2` -> `5 % 2 == 1`. The other tune tests all keep +// `overflow_chunks` a multiple of `stream_stages`, so this is the only tiny check of the uneven-tail reload for both +// directions (deterministic combos drive reverse, non-deterministic forward). +C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with a partial final overflow wave", + "[pairs][segmented][topk][device][cluster][determinism]", + det_tie_pair_combos) +{ + using key_t = float; + using val_t = cuda::std::int32_t; + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using combo = c2h::get<0, TestType>; + constexpr auto determinism = combo::determinism; + constexpr auto tie_break = combo::tie_break; + constexpr auto direction = cub::detail::topk::select::max; + + constexpr segment_size_t static_max_segment_size = 1536; + constexpr segment_size_t static_max_k = 512; + constexpr segment_index_t num_segments = 2; + skip_if_batched_topk_cluster_unavailable(true); + + const segment_size_t segment_size = 896 - 63; // 7 chunks (128 items each), unaligned -> peeled overflow tail edge + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k})); + + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + auto values_in_it = cuda::make_counting_iterator(val_t{0}); + auto d_values_in = cuda::make_strided_iterator(cuda::make_counting_iterator(values_in_it), segment_size); + c2h::device_vector values_out_buffer(num_segments * k, thrust::no_init); + auto d_values_out_ptr = thrust::raw_pointer_cast(values_out_buffer.data()); + auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_values_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + auto env = cuda::std::execution::env{ + cuda::stream_ref{cudaStream_t{0}}, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, + cuda::execution::output_ordering::unsorted), + cuda::execution::tune( + cluster_tuning_selector<1, /*slots=*/4, /*single_block=*/0, cluster_test_chunk_bytes, /*stages=*/2>{})}; + + auto seg_sizes = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; + + size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp, + temp_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + seg_sizes, + k_param, + cuda::args::immediate{num_segments}, + env); + }; + REQUIRE(cudaSuccess == invoke(nullptr)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_buffer, values_out_buffer) == true); + REQUIRE(verify_unique_indices(values_out_buffer, num_segments, k) == true); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} # endif // TEST_LAUNCH == 0 // Effective-cluster-width coverage for pairs: the large `deferred_sequence` bound sizes a wide (16-CTA) launch, but the From ad6ddc3b8ce59645191ae3259b4b69d56179410a Mon Sep 17 00:00:00 2001 From: pauleonix Date: Tue, 14 Jul 2026 09:29:58 +0200 Subject: [PATCH 122/126] Revert .clangd change --- .clangd | 1 - 1 file changed, 1 deletion(-) diff --git a/.clangd b/.clangd index 079f1b7b8b8..2200fd689ba 100644 --- a/.clangd +++ b/.clangd @@ -73,7 +73,6 @@ CompileFlags: Add: - -x - cuda - - --offload-arch=sm_90 - -Wno-unknown-cuda-version - -Wno-pragma-system-header-outside-header - --no-cuda-version-check From 1f9d81b0e7859d6168a7f6394ca9839b544052af Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 15 Jul 2026 01:24:29 +0200 Subject: [PATCH 123/126] Implement CodeRabbit feedback --- .../bench/segmented_topk/fixed/keys.cu | 22 +-- .../variable/indexed_common.cuh | 26 ++- .../segmented_topk/variable/keys_common.cuh | 24 ++- cub/cub/agent/agent_batched_topk.cuh | 173 +++++++++--------- cub/cub/device/device_batched_topk.cuh | 53 ++++-- .../device/dispatch/dispatch_batched_topk.cuh | 151 ++++++++------- .../catch2_test_device_segmented_topk_keys.cu | 97 +++++++++- ...catch2_test_device_segmented_topk_pairs.cu | 78 ++++---- ...vice_batched_topk_unsupported_arch_fail.cu | 7 +- .../cub/api_docs/device_topk_requirements.rst | 2 +- 10 files changed, 377 insertions(+), 256 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index cf0b139a785..a9f8a1a8614 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -143,7 +143,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement // on top of the benchmark environment (which carries the stream and the caching memory resource). - auto seg_env = cuda::std::execution::env{ + const auto seg_env = cuda::std::execution::env{ env, cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; @@ -169,7 +169,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { // The determinism / tie-break / ordering requirement this benchmark issues; the library honors it whether or not we // additionally force a backend. - auto req_env = cuda::std::execution::env{ + const auto req_env = cuda::std::execution::env{ env, cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, cuda::execution::tie_break::__tie_break_holder_t{}, @@ -183,7 +183,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { using key_t = cub::detail::it_value_t>; constexpr auto max_k = ::cuda::args::__traits::highest; - auto full_env = cuda::std::execution::env{ + const auto full_env = cuda::std::execution::env{ req_env, cuda::execution::tune(topk_backend_selector{})}; return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, full_env); } @@ -212,13 +212,13 @@ void fixed_seg_size_topk_keys( thrust::device_vector in_keys_buffer = generate(elements, entropy); thrust::device_vector out_keys_buffer(selected_elements * num_segments, thrust::no_init); - auto d_keys_in_ptr = thrust::raw_pointer_cast(in_keys_buffer.data()); - auto d_keys_out_ptr = thrust::raw_pointer_cast(out_keys_buffer.data()); - auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); - auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements); + const auto d_keys_in_ptr = thrust::raw_pointer_cast(in_keys_buffer.data()); + const auto d_keys_out_ptr = thrust::raw_pointer_cast(out_keys_buffer.data()); + const auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + const auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements); - auto segment_sizes = ::cuda::args::constant{}; - auto k = ::cuda::args::constant{}; + const auto segment_sizes = ::cuda::args::constant{}; + const auto k = ::cuda::args::constant{}; state.add_element_count(elements, "NumElements"); state.add_element_count(segment_size, "SegmentSize"); @@ -228,11 +228,11 @@ void fixed_seg_size_topk_keys( // Host copy of segment sizes — all entries equal MaxSegmentSize for fixed-size segments. Consumed only by the // per-segment device backend. Segment sizes fit in a signed 32-bit integer (the library caps them at 2^21). - std::vector h_segment_sizes(num_segments, static_cast(MaxSegmentSize)); + const std::vector h_segment_sizes(num_segments, static_cast(MaxSegmentSize)); caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env(alloc, launch); + const auto env = cub_bench_env(alloc, launch); _CCCL_TRY_CUDA_API( batched_topk_keys, "batched topk failed", diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh index 86ae3d3c73d..438b776a3e8 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -43,6 +43,12 @@ # error \ "indexed_common.cuh requires the includer to define TUNE_REQUIREMENT (0 non-det / 1 det+smaller / 2 det+larger)" #endif +#if TUNE_BACKEND < 0 || TUNE_BACKEND > 2 +# error "indexed_common.cuh: TUNE_BACKEND must be 0 (baseline), 1 (cluster), or 2 (automatic)" +#endif +#if TUNE_REQUIREMENT < 0 || TUNE_REQUIREMENT > 2 +# error "indexed_common.cuh: TUNE_REQUIREMENT must be 0 (non-det), 1 (det+smaller), or 2 (det+larger)" +#endif enum class topk_backend { @@ -174,7 +180,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_indexed( NumSegmentsParameterT num_segments, EnvT env) { - auto req_env = cuda::std::execution::env{ + const auto req_env = cuda::std::execution::env{ env, cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, cuda::execution::tie_break::__tie_break_holder_t{}, @@ -190,7 +196,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_indexed( using key_t = cub::detail::it_value_t>; using value_t = cub::detail::it_value_t>; constexpr auto max_k = ::cuda::args::__traits::highest; - auto full_env = cuda::std::execution::env{ + const auto full_env = cuda::std::execution::env{ req_env, cuda::execution::tune(topk_backend_selector{})}; return cub::DeviceBatchedTopK::MaxPairs( d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, full_env); @@ -224,21 +230,21 @@ void decode_style_variable_topk_indexed( auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); auto out_indices_buffer = thrust::device_vector(output_elements, thrust::no_init); - auto segment_sizes_param = cuda::args::deferred_sequence{ + const auto segment_sizes_param = cuda::args::deferred_sequence{ thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; - auto k_param = cuda::args::constant{}; - auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + const auto k_param = cuda::args::constant{}; + const auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; - auto d_keys_in = cuda::make_strided_iterator( + const auto d_keys_in = cuda::make_strided_iterator( cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), static_cast(MaxSegmentSize)); - auto d_keys_out = cuda::make_strided_iterator( + const auto d_keys_out = cuda::make_strided_iterator( cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), static_cast(K)); // Input values: every segment maps to the same counting iterator starting at 0, so values are segment-local indices. - auto d_indices_in = cuda::make_constant_iterator(cuda::make_counting_iterator(IndexT{0})); - auto d_indices_out = cuda::make_strided_iterator( + const auto d_indices_in = cuda::make_constant_iterator(cuda::make_counting_iterator(IndexT{0})); + const auto d_indices_out = cuda::make_strided_iterator( cuda::make_counting_iterator(thrust::raw_pointer_cast(out_indices_buffer.data())), static_cast(K)); @@ -250,7 +256,7 @@ void decode_style_variable_topk_indexed( caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env(alloc, launch); + const auto env = cub_bench_env(alloc, launch); _CCCL_TRY_CUDA_API( batched_topk_indexed, "batched topk failed", diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh index 7237e70fe11..3b0f554c513 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -47,6 +47,12 @@ #ifndef TUNE_REQUIREMENT # error "keys_common.cuh requires the includer to define TUNE_REQUIREMENT (0 non-det / 1 det+smaller / 2 det+larger)" #endif +#if TUNE_BACKEND < 0 || TUNE_BACKEND > 3 +# error "keys_common.cuh: TUNE_BACKEND must be 0 (baseline), 1 (cluster), 2 (device), or 3 (automatic)" +#endif +#if TUNE_REQUIREMENT < 0 || TUNE_REQUIREMENT > 2 +# error "keys_common.cuh: TUNE_REQUIREMENT must be 0 (non-det), 1 (det+smaller), or 2 (det+larger)" +#endif enum class topk_backend { @@ -186,7 +192,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( // The per-segment device backend uses the unsorted / not-guaranteed-determinism fast path. Layer the requirement // on top of the benchmark environment (which carries the stream and the caching memory resource). - auto seg_env = cuda::std::execution::env{ + const auto seg_env = cuda::std::execution::env{ env, cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; @@ -212,7 +218,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { // The determinism / tie-break / ordering requirement this benchmark issues; the library honors it whether or not we // additionally force a backend. - auto req_env = cuda::std::execution::env{ + const auto req_env = cuda::std::execution::env{ env, cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, cuda::execution::tie_break::__tie_break_holder_t{}, @@ -226,7 +232,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t batched_topk_keys( { using key_t = cub::detail::it_value_t>; constexpr auto max_k = ::cuda::args::__traits::highest; - auto full_env = cuda::std::execution::env{ + const auto full_env = cuda::std::execution::env{ req_env, cuda::execution::tune(topk_backend_selector{})}; return cub::DeviceBatchedTopK::MaxKeys(d_keys_in, d_keys_out, segment_sizes, k, num_segments, full_env); } @@ -256,15 +262,15 @@ void decode_style_variable_topk_keys( num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data())); auto out_keys_buffer = thrust::device_vector(output_elements, thrust::no_init); - auto segment_sizes_param = cuda::args::deferred_sequence{ + const auto segment_sizes_param = cuda::args::deferred_sequence{ thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; - auto k_param = cuda::args::constant{}; - auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + const auto k_param = cuda::args::constant{}; + const auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; - auto d_keys_in = cuda::make_strided_iterator( + const auto d_keys_in = cuda::make_strided_iterator( cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), static_cast(MaxSegmentSize)); - auto d_keys_out = cuda::make_strided_iterator( + const auto d_keys_out = cuda::make_strided_iterator( cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())), static_cast(K)); @@ -279,7 +285,7 @@ void decode_style_variable_topk_keys( caching_allocator_t alloc; state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) { - auto env = cub_bench_env(alloc, launch); + const auto env = cub_bench_env(alloc, launch); _CCCL_TRY_CUDA_API( batched_topk_keys, "batched topk failed", diff --git a/cub/cub/agent/agent_batched_topk.cuh b/cub/cub/agent/agent_batched_topk.cuh index dc38b332de2..e7fd55157a1 100644 --- a/cub/cub/agent/agent_batched_topk.cuh +++ b/cub/cub/agent/agent_batched_topk.cuh @@ -220,112 +220,113 @@ struct agent_batched_topk_worker_per_segment // Process small segment const auto k = (::cuda::std::min) (params::get_param(k_param, segment_id), static_cast(segment_size)); - // Nothing to select for an empty segment (including a negative size clamped to 0) or a zero k: skip it, leaving - // its output untouched. Also keeps the block primitive's `valid_items in [1, tile_items]` precondition. - if (k == 0) + // Nothing to select for an empty segment (including a negative size clamped to 0) or a zero k: skip the block + // work, leaving its output untouched (also keeps the block primitive's `valid_items in [1, tile_items]` + // precondition). We must not `return` here -- the large-segment epilogue below is unconditional participation, so + // bailing out would drop this block from the retirement count and stall the epilogue scan. + if (k != 0) { - return; - } - const auto direction = select_directions.get_param(segment_id); - - // Determine padding key based on direction - const key_t padding_key = - (direction == detail::topk::select::max) - ? ::cuda::std::numeric_limits::lowest() - : (::cuda::std::numeric_limits::max)(); - - // Dereference iterator-of-iterators to get the segment specific iterator - auto block_keys_in = d_key_segments_it[segment_id]; - - // Load Keys - key_t thread_keys[items_per_thread]; - if constexpr (is_full_tile) - { - // No padding needed - block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys); - } - else - { - // Potentially partial final load with padding - // TODO (elstehle): explore whether a runtime check for segment_size == tile_size improves performance - block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys, segment_size); - } + const auto direction = select_directions.get_param(segment_id); - // Load Values (if applicable) - [[maybe_unused]] value_t thread_values[items_per_thread]; + // Determine padding key based on direction + const key_t padding_key = + (direction == detail::topk::select::max) + ? ::cuda::std::numeric_limits::lowest() + : (::cuda::std::numeric_limits::max)(); - if constexpr (!is_keys_only) - { - __syncthreads(); - auto block_vals_in = d_value_segments_it[segment_id]; + // Dereference iterator-of-iterators to get the segment specific iterator + auto block_keys_in = d_key_segments_it[segment_id]; + // Load Keys + key_t thread_keys[items_per_thread]; if constexpr (is_full_tile) { // No padding needed - block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values); + block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys); } else { // Potentially partial final load with padding // TODO (elstehle): explore whether a runtime check for segment_size == tile_size improves performance - block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values, segment_size); + block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys, segment_size); } - } - __syncthreads(); - - // Perform Block Top-K - if constexpr (is_keys_only) - { - const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( - select_directions, segment_id, [this, &thread_keys, k, segment_size](auto direction_tag) { - if constexpr (decltype(direction_tag)::value == detail::topk::select::max) - { - block_topk_t(temp_storage.topk).template max_keys(thread_keys, k, segment_size); - } - else - { - block_topk_t(temp_storage.topk).template min_keys(thread_keys, k, segment_size); - } - }); - _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); - } - else - { - // Pass both keys and values - const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( - select_directions, segment_id, [this, &thread_keys, &thread_values, k, segment_size](auto direction_tag) { - if constexpr (decltype(direction_tag)::value == detail::topk::select::max) - { - block_topk_t(temp_storage.topk) - .template max_pairs(thread_keys, thread_values, k, segment_size); - } - else - { - block_topk_t(temp_storage.topk) - .template min_pairs(thread_keys, thread_values, k, segment_size); - } - }); - _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); - } + // Load Values (if applicable) + [[maybe_unused]] value_t thread_values[items_per_thread]; - __syncthreads(); + if constexpr (!is_keys_only) + { + __syncthreads(); + auto block_vals_in = d_value_segments_it[segment_id]; + + if constexpr (is_full_tile) + { + // No padding needed + block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values); + } + else + { + // Potentially partial final load with padding + // TODO (elstehle): explore whether a runtime check for segment_size == tile_size improves performance + block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values, segment_size); + } + } - auto block_keys_out = d_key_segments_out_it[segment_id]; + __syncthreads(); - block_store_keys_t(temp_storage.store_keys) - .Store(block_keys_out, - thread_keys, - k // Only store K items - ); + // Perform Block Top-K + if constexpr (is_keys_only) + { + const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( + select_directions, segment_id, [this, &thread_keys, k, segment_size](auto direction_tag) { + if constexpr (decltype(direction_tag)::value == detail::topk::select::max) + { + block_topk_t(temp_storage.topk).template max_keys(thread_keys, k, segment_size); + } + else + { + block_topk_t(temp_storage.topk).template min_keys(thread_keys, k, segment_size); + } + }); + _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); + } + else + { + // Pass both keys and values + const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( + select_directions, segment_id, [this, &thread_keys, &thread_values, k, segment_size](auto direction_tag) { + if constexpr (decltype(direction_tag)::value == detail::topk::select::max) + { + block_topk_t(temp_storage.topk) + .template max_pairs(thread_keys, thread_values, k, segment_size); + } + else + { + block_topk_t(temp_storage.topk) + .template min_pairs(thread_keys, thread_values, k, segment_size); + } + }); + _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); + } - if constexpr (!is_keys_only) - { __syncthreads(); - auto block_vals_out = d_value_segments_out_it[segment_id]; - block_store_vals_t(temp_storage.store_vals).Store(block_vals_out, thread_values, k); - } + auto block_keys_out = d_key_segments_out_it[segment_id]; + + block_store_keys_t(temp_storage.store_keys) + .Store(block_keys_out, + thread_keys, + k // Only store K items + ); + + if constexpr (!is_keys_only) + { + __syncthreads(); + auto block_vals_out = d_value_segments_out_it[segment_id]; + + block_store_vals_t(temp_storage.store_vals).Store(block_vals_out, thread_values, k); + } + } // if (k != 0) } // Epilogue: Scan queued large segment sizes (in tiles not elements) for load balancing search in the large segment diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 49796a33b81..478aacf4a82 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -159,10 +159,13 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( // --------------------------------------------------------------------------- // Argument-annotation constraints surfaced at the call site. // --------------------------------------------------------------------------- - static_assert(::cuda::args::__traits::is_single_value, - "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " - "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " - "cuda::args::immediate), not a per-segment sequence."); + static_assert( + ::cuda::args::__traits::is_single_value + && !::cuda::args::__traits::is_deferred, + "cub::DeviceBatchedTopK requires a single (uniform) number of segments resolved on the host: pass " + "num_segments as a host-known single value (e.g. cuda::args::constant or cuda::args::immediate). A " + "per-segment sequence is not a meaningful segment count; a single deferred (device-resident) value is " + "meaningful but not yet supported."); // The segment_sizes checks below are layered so that a single misuse reports a single, targeted diagnostic: each // check is guarded by the validity of the ones before it (a later check "passes" once an earlier one has failed), // so e.g. a raw pointer trips only the form check, not also the range check, and a scalar handed to `deferred` @@ -328,13 +331,16 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! as narrow, lying within the compile-time range and only tightening it further. //! //! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. ``num_segments`` must be -//! a single value (``constant``, ``immediate``, or a plain integral), never a per-segment sequence. ``segment_sizes`` +//! a single, non-negative value known on the host (``constant``, ``immediate``, or a plain integral). +//! ``segment_sizes`` //! must have a statically-known *maximum* not exceeding the supported ``2^21`` (about 2 million; see *Current //! constraints* below): a type whose maximum already fits (a narrow type such as ``uint8_t``, ``int16_t``, or //! ``uint16_t``) is accepted without an explicit bound, while a type whose maximum exceeds ``2^21`` (e.g. ``int32_t``, //! ``uint32_t``, or ``int64_t``) must carry a compile-time upper bound (a ``constant`` or -//! ``cuda::args::bounds()``). A negative lower bound is allowed -- negative runtime sizes are clamped to an -//! empty segment (see *Current constraints*). Tight bounds on every parameter are encouraged. +//! ``cuda::args::bounds()``). A negative statically-known lower bound is allowed: negative runtime sizes are +//! clamped to an empty segment (size 0). A non-negative lower bound is trusted -- passing an actual value outside its +//! declared bound (for instance a negative value under a non-negative bound) is a caller precondition violation +//! (undefined behavior). Tight bounds on every parameter are encouraged. //! //! .. code-block:: c++ //! @@ -374,11 +380,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! size. Larger segments are future work. A type whose maximum already lies within it (a narrow type such as //! ``uint8_t``, ``int16_t``, or ``uint16_t``) is accepted un-annotated; a type whose maximum exceeds ``2^21`` (e.g. //! ``int32_t``, ``uint32_t``, or ``int64_t``) must carry a compile-time ``cuda::args::bounds`` whose upper end does -//! not exceed ``2^21``. A negative statically-known lower bound is accepted: a negative runtime size is treated as an -//! empty segment (clamped to 0). A non-negative lower bound is trusted -- an actual negative value there, like any -//! per-segment value outside its declared bound, is a caller precondition violation (undefined behavior; only -//! host-known values are checked, by assertions in debug builds). -//! - **Uniform number of segments.** ``num_segments`` must be a single value, never a per-segment sequence. +//! not exceed ``2^21`` (see *Which form each parameter accepts* above for how the lower bound and out-of-bound values +//! are handled). +//! - **Uniform number of segments.** ``num_segments`` must be a single value (``constant``, ``immediate``, or a plain +//! integral) resolved on the host. A ``deferred`` (device-resident) count is not supported at this time. //! - **Unsorted output required.** Only ``cuda::execution::output_ordering::unsorted`` is implemented; the sorted //! orderings of the default contract described in *Determinism, tie-breaking, and output ordering* below (and hence //! an empty, no-requirement environment, which defaults to ``stable_sorted``) are rejected at compile time. The @@ -501,7 +506,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -600,7 +606,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -705,7 +712,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -802,7 +810,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -924,7 +933,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -1033,7 +1043,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -1147,7 +1158,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst @@ -1256,7 +1268,8 @@ struct DeviceBatchedTopK //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! The number of segments, given as a `cuda::args` annotation or a plain integral value. Must be + //! non-negative. //! //! @param[in] env //! @rst diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 4ffe89c8738..463bf4f2f32 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -172,7 +172,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch( using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; // Wrap the raw enum into the internal discrete param type - auto select_directions = wrap_select_direction(select_direction); + const auto select_directions = wrap_select_direction(select_direction); using SelectDirectionParameterT = decltype(select_directions); // Helper that determines (a) whether there's any one-worker-per-segment policy supporting the range of segment @@ -263,8 +263,17 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch( // TODO (elstehle): support number of segments provided by device-accessible iterator // Only uniform number of segments are supported (i.e., we need to resolve the number of segments on the host) - static_assert(::cuda::args::__traits::is_single_value, - "Only a uniform number of segments is currently supported."); + static_assert(::cuda::args::__traits::is_single_value + && !::cuda::args::__traits::is_deferred, + "Only a host-known uniform number of segments is currently supported; a deferred (device-resident) " + "count is not yet supported (see the TODO above)."); + + // No work to launch when there are no segments; otherwise the grid launches below would use `grid_dim = 0`, an + // invalid launch configuration. + if (params::get_param(num_segments, 0) == 0) + { + return cudaSuccess; + } if constexpr (any_small_segments) { @@ -407,44 +416,44 @@ template // Without CDP/RDC a device-side launch is impossible; surface that instead of silently returning success (no-op). # define CUB_TOPK_CLUSTER_DEVICE_LAUNCH return cudaErrorNotSupported; #else // CUB_RDC_ENABLED -# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ - auto static_kernel = detail::batched_topk::device_segmented_topk_cluster_kernel_static< \ - ThreadsPerBlock, \ - HistogramItemsPerThread, \ - PipelineStages, \ - ChunkBytes, \ - LoadAlignBytes, \ - BitsPerPass, \ - TieBreakItemsPerThread, \ - SingleBlockMaxSegSize, \ - MinChunksPerBlock, \ - CopyItemsPerThread, \ - cdp_cluster_blocks, \ - Determinism, \ - TieBreak, \ - KeyInputItItT, \ - KeyOutputItItT, \ - ValueInputItItT, \ - ValueOutputItItT, \ - SegmentSizeParameterT, \ - KParameterT, \ - SelectDirectionParameterT, \ - NumSegmentsParameterT>; \ - if (const auto error = CubDebug( \ - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ - static_cast(grid_blocks), ThreadsPerBlock, dynamic_smem_bytes, stream) \ - .doit(static_kernel, \ - d_key_segments_it, \ - d_key_segments_out_it, \ - d_value_segments_it, \ - d_value_segments_out_it, \ - segment_sizes, \ - k_param, \ - select_directions, \ - num_segments, \ - block_tile_capacity))) \ - { \ - return error; \ +# define CUB_TOPK_CLUSTER_DEVICE_LAUNCH \ + const auto static_kernel = detail::batched_topk::device_segmented_topk_cluster_kernel_static< \ + ThreadsPerBlock, \ + HistogramItemsPerThread, \ + PipelineStages, \ + ChunkBytes, \ + LoadAlignBytes, \ + BitsPerPass, \ + TieBreakItemsPerThread, \ + SingleBlockMaxSegSize, \ + MinChunksPerBlock, \ + CopyItemsPerThread, \ + cdp_cluster_blocks, \ + Determinism, \ + TieBreak, \ + KeyInputItItT, \ + KeyOutputItItT, \ + ValueInputItItT, \ + ValueOutputItItT, \ + SegmentSizeParameterT, \ + KParameterT, \ + SelectDirectionParameterT, \ + NumSegmentsParameterT>; \ + if (const auto error = CubDebug( \ + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( \ + static_cast(grid_blocks), ThreadsPerBlock, dynamic_smem_bytes, stream) \ + .doit(static_kernel, \ + d_key_segments_it, \ + d_key_segments_out_it, \ + d_value_segments_it, \ + d_value_segments_out_it, \ + segment_sizes, \ + k_param, \ + select_directions, \ + num_segments, \ + block_tile_capacity))) \ + { \ + return error; \ } #endif // CUB_RDC_ENABLED @@ -560,23 +569,10 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( return cudaSuccess; } - static_assert(::cuda::args::__traits::is_single_value, - "Number of segments must be resolved on the host."); - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - const auto num_seg_val = detail::params::get_param(num_segments, num_segments_val_t{0}); - if (num_seg_val == 0) - { - return cudaSuccess; - } - - // No work to launch when the tightest known upper bound on any segment size (static or runtime) is non-positive: - // every segment is empty (e.g. a uniform negative size, which the kernel would clamp to 0). Also avoids - // `clusterDim.x = 0` and a negative maximum feeding the unsigned sizing math below. - if (max_seg_size <= 0) - { - return cudaSuccess; - } + // `num_segments > 0` and `max_seg_size > 0` here: the generic `dispatch` returns for the empty-batch cases (no + // segments, or a non-positive max segment size) before selecting an arm. + const auto num_seg_val = detail::params::get_param(num_segments, num_segments_val_t{0}); // Cluster launches require compute capability 9.0+. int sm_version = 0; @@ -1090,15 +1086,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_baseline_arm( return cudaSuccess; } - // No work to launch when the tightest known upper bound on any segment size is non-positive (mirrors the cluster - // arm): every segment is empty (e.g. a uniform host-known negative size, which the agent clamps to 0 device-side). - if (runtime_max_segment_size(segment_sizes) <= 0) - { - return cudaSuccess; - } - - static_assert(::cuda::args::__traits::is_single_value, - "Only a uniform number of segments is currently supported."); + // `num_segments > 0` and the max segment size > 0 here: the generic `dispatch` returns for the empty-batch cases + // (no segments, or a non-positive max segment size) before selecting an arm. if constexpr (any_small_segments) { @@ -1178,7 +1167,7 @@ template } return any; } -#endif // _CCCL_CUDA_COMPILATION() && !CUB_DEFINE_RUNTIME_POLICIES && !NVRTC +#endif // _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) // Internal entry point: the single dispatch that replaces the standalone baseline / cluster dispatches. It resolves the // runtime compute capability, then uses `dispatch_compute_cap` to pick, per architecture, the backend chosen by @@ -1217,9 +1206,27 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, cudaStream_t stream) { + // Both arms resolve `num_segments` on the host via detail::params::get_param (allocation sizing, grid extent, + // empty-batch guard), so it must be a host-known single value; device-side counts are future work (see TODOs below). + static_assert(::cuda::args::__traits::is_single_value + && !::cuda::args::__traits::is_deferred, + "cub::DeviceBatchedTopK requires a host-known uniform number of segments (constant, immediate, or a " + "plain integral): a per-segment sequence is not a meaningful segment count, and a single deferred " + "(device-resident) value is meaningful but not yet supported (resolve the count on the host)."); + + // No work to launch when the batch is empty: no segments, or the tightest known max segment size is non-positive + // (every segment empty; e.g. a uniform negative size the kernel clamps to 0). A negative `num_segments` is out of + // contract (non-negative is a documented precondition), hence `== 0` not `<= 0`. Guard only the actual launch: the + // query pass (`d_temp_storage == nullptr`) must fall through so the chosen arm still reports `temp_storage_bytes`. + if (d_temp_storage != nullptr + && (detail::params::get_param(num_segments, 0) == 0 || runtime_max_segment_size(segment_sizes) <= 0)) + { + return cudaSuccess; + } + // The selection direction is a compile-time constant carried as `::cuda::args::constant`. Wrap it into the // internal discrete param the kernel/agent expect (both host arms take the wrapped form). - auto select_directions = wrap_select_direction(select_direction); + const auto select_directions = wrap_select_direction(select_direction); using SelectDirectionParameterT = decltype(select_directions); using key_t = it_value_t>; @@ -1281,11 +1288,15 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( "CMAKE_CUDA_ARCHITECTURES (the deterministic / large-segment cluster backend requires SM90+). Remove the " "unsupported architecture(s), relax the request, or define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the " "diagnosis to runtime (cudaErrorNotSupported)."); -#endif // strict unsupported-arch check +#endif // _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + // && !defined(_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) // The supported maximum segment size (2^21) is enforced at compile time at the public entry; a statically negative // lower bound is allowed and negative runtime sizes are clamped to 0 (see detail::params::get_segment_size). A - // per-segment value outside its declared bound is a caller error (asserted for host-known values, otherwise UB). + // per-segment value outside its declared bound is a caller error (UB): the statically declared bounds are validated + // at compile time, while the argument values are bounds-checked only by assertions active in assertion-enabled (e.g. + // debug) builds -- host-side for a host-known immediate value and device-side for values read from a deferred / + // deferred_sequence handle. detail::TripleChevronFactory launcher_factory{}; ::cuda::compute_capability cc{}; diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 4280297e7b3..272352b358b 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -73,7 +74,7 @@ template CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( void* d_temp_storage, - size_t& temp_storage_bytes, + cuda::std::size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, SegmentSizeParamT segment_sizes, @@ -578,8 +579,8 @@ template temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == run(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); } @@ -1662,8 +1663,8 @@ C2H_TEST("DeviceBatchedTopK::MaxKeys handles a misaligned temporary storage poin cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; - size_t temp_storage_bytes = 0; - auto error = cub::DeviceBatchedTopK::MaxKeys( + cuda::std::size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); REQUIRE(error == cudaSuccess); @@ -1718,8 +1719,8 @@ C2H_TEST("DeviceBatchedTopK::MaxKeys accepts an un-annotated narrow-unsigned seg cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; - size_t temp_storage_bytes = 0; - auto error = cub::DeviceBatchedTopK::MaxKeys( + cuda::std::size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); REQUIRE(error == cudaSuccess); @@ -1773,8 +1774,8 @@ C2H_TEST("DeviceBatchedTopK::MaxKeys clamps a negative segment size to an empty cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; - size_t temp_storage_bytes = 0; - auto error = cub::DeviceBatchedTopK::MaxKeys( + cuda::std::size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); REQUIRE(error == cudaSuccess); @@ -1871,4 +1872,80 @@ C2H_TEST("DeviceBatchedTopK::MaxKeys treats a uniform negative segment size as n // No segment had any items, so the whole output is left untouched. REQUIRE(keys_out == thrust::device_vector(num_segments * k, sentinel)); } + +// Zero segments is a no-op, but with a positive segment-size bound the `max_seg_size <= 0` guard does not fire, so the +// dispatch must elide the launch from `num_segments == 0` alone -- otherwise the grid would use `grid_dim == 0`, an +// invalid launch configuration. Small size + no determinism keeps this on the baseline arm; the cluster arm is below. +C2H_TEST("DeviceBatchedTopK::MaxKeys treats zero segments as no work (baseline backend)", + "[keys][segmented][topk][device]") +{ + constexpr int k = 3; + constexpr int stride = 8; + constexpr int sentinel = -12345; + auto keys_in = thrust::device_vector{}; + // Canary output slots that must stay untouched: with zero segments nothing may be written. + auto keys_out = thrust::device_vector(k, sentinel); + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), stride); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + auto segment_sizes = cuda::args::immediate{cuda::std::int16_t{stride}, cuda::args::bounds<0, 100>()}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{0}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + cuda::std::size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + REQUIRE(keys_out == thrust::device_vector(k, sentinel)); +} + +// Cluster-backend counterpart: a deterministic requirement forces the SM90+ cluster arm, whose own `num_seg_val == 0` +// guard must likewise elide the launch (positive size bound -> the `max_seg_size <= 0` guard does not apply). +C2H_TEST("DeviceBatchedTopK::MaxKeys treats zero segments as no work (cluster backend)", + "[keys][segmented][topk][device][cluster][determinism]") +{ + constexpr auto determinism = cuda::execution::determinism::__determinism_t::__gpu_to_gpu; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__prefer_smaller_index; + constexpr int k = 3; + constexpr int stride = 8; + constexpr int sentinel = -12345; + constexpr cuda::std::int64_t static_max_segment_size = 100; + + skip_if_batched_topk_backend_unavailable(static_max_segment_size); + + auto keys_in = thrust::device_vector{}; + // Canary output slots that must stay untouched: with zero segments nothing may be written. + auto keys_out = thrust::device_vector(k, sentinel); + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), stride); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{cuda::std::int16_t{stride}, cuda::args::bounds<0, 100>()}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{0}}); + + REQUIRE(keys_out == thrust::device_vector(k, sentinel)); +} #endif // TEST_TYPES == 0 && TEST_LAUNCH == 0 diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index c988ee51161..6683ca06ca0 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include #include @@ -81,7 +83,7 @@ template CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( void* d_temp_storage, - size_t& temp_storage_bytes, + cuda::std::size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, ValueInputItItT d_value_segments_it, @@ -252,8 +254,8 @@ bool verify_unique_indices(const c2h::device_vector& values_compacted, flag_intra_segment_duplicates flag_op{sorted_values.cbegin(), segment_ids.cbegin()}; - auto num_duplicates = - thrust::count_if(cuda::make_counting_iterator(size_t{0}), cuda::make_counting_iterator(num_items - 1), flag_op); + auto num_duplicates = thrust::count_if( + cuda::make_counting_iterator(cuda::std::size_t{0}), cuda::make_counting_iterator(num_items - 1), flag_op); return num_duplicates == 0; } @@ -533,7 +535,7 @@ template CUB_RUNTIME_FUNCTION static cudaError_t dispatch_cluster_topk_pairs( void* d_temp_storage, - size_t& temp_storage_bytes, + cuda::std::size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, KeyOutputItItT d_key_segments_out_it, ValueInputItItT d_value_segments_it, @@ -874,15 +876,15 @@ c2h::host_vector reference_deterministic_topk_indices( return reverse_index ? static_cast(last_index - idx) : idx; }; - c2h::host_vector selected(static_cast(num_segments * k)); - std::vector> pairs(static_cast(segment_size)); + c2h::host_vector selected(static_cast(num_segments * k)); + std::vector> pairs(static_cast(segment_size)); for (SegSizeT seg = 0; seg < num_segments; ++seg) { const SegSizeT base = seg * segment_size; for (SegSizeT i = 0; i < segment_size; ++i) { - const IndexT idx = static_cast(base + i); - pairs[static_cast(i)] = {h_keys[static_cast(base + i)], encode(idx)}; + const IndexT idx = static_cast(base + i); + pairs[static_cast(i)] = {h_keys[static_cast(base + i)], encode(idx)}; } if (want_max) { @@ -892,10 +894,10 @@ c2h::host_vector reference_deterministic_topk_indices( { std::partial_sort(pairs.begin(), pairs.begin() + k, pairs.end()); } - const auto seg_begin = selected.begin() + static_cast(seg * k); + const auto seg_begin = selected.begin() + static_cast(seg * k); for (SegSizeT i = 0; i < k; ++i) { - seg_begin[i] = encode(pairs[static_cast(i)].second); + seg_begin[i] = encode(pairs[static_cast(i)].second); } std::sort(seg_begin, seg_begin + k); } @@ -974,7 +976,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break returns the c2h::host_vector h_values_out = values_out_buffer; for (segment_index_t seg = 0; seg < num_segments; ++seg) { - const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); std::sort(seg_begin, seg_begin + k); } @@ -1043,8 +1045,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs run a tiny multi-CTA segment through cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { if constexpr (direction == cub::detail::topk::select::max) { return cub::DeviceBatchedTopK::MaxPairs( @@ -1075,7 +1077,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs run a tiny multi-CTA segment through } }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -1092,7 +1094,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs run a tiny multi-CTA segment through c2h::host_vector h_values_out = values_out_buffer; for (segment_index_t seg = 0; seg < num_segments; ++seg) { - const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); std::sort(seg_begin, seg_begin + k); } @@ -1186,7 +1188,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the c2h::host_vector h_values_out = values_out_buffer; for (segment_index_t seg = 0; seg < num_segments; ++seg) { - const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); std::sort(seg_begin, seg_begin + k); } @@ -1279,11 +1281,11 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the cuda::args::immediate{num_segments}); // Materialize the (non-contiguous) key input for verification: same heavy-tie function of the global index. - c2h::host_vector h_keys(static_cast(num_items)); + c2h::host_vector h_keys(static_cast(num_items)); heavy_tie_key_op key_op{}; for (segment_size_t idx = 0; idx < num_items; ++idx) { - h_keys[static_cast(idx)] = key_op(idx); + h_keys[static_cast(idx)] = key_op(idx); } const c2h::device_vector keys_materialized = h_keys; @@ -1298,7 +1300,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic tie-break streams the c2h::host_vector h_values_out = values_out_buffer; for (segment_index_t seg = 0; seg < num_segments; ++seg) { - const auto seg_begin = h_values_out.begin() + static_cast(seg * k); + const auto seg_begin = h_values_out.begin() + static_cast(seg * k); std::sort(seg_begin, seg_begin + k); } @@ -1407,14 +1409,18 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs deterministic unspecified tie-break c2h::device_vector expected_keys(keys_in_buffer); REQUIRE(verify_pairs_consistency(expected_keys, keys_out_a, values_out_a) == true); REQUIRE(verify_unique_indices(values_out_a, num_segments, k) == true); + // The second run uses config B (the streaming path under gpu_to_gpu); validate its pairing/uniqueness too so a + // B-only key/value mismatch cannot hide behind the index-set comparison below. + REQUIRE(verify_pairs_consistency(expected_keys, keys_out_b, values_out_b) == true); + REQUIRE(verify_unique_indices(values_out_b, num_segments, k) == true); // Determinism: both runs must select the same per-segment index set. Sort each segment (within-top-k order is free). c2h::host_vector h_values_a = values_out_a; c2h::host_vector h_values_b = values_out_b; for (segment_index_t seg = 0; seg < num_segments; ++seg) { - const auto a_begin = h_values_a.begin() + static_cast(seg * k); - const auto b_begin = h_values_b.begin() + static_cast(seg * k); + const auto a_begin = h_values_a.begin() + static_cast(seg * k); + const auto b_begin = h_values_b.begin() + static_cast(seg * k); std::sort(a_begin, a_begin + k); std::sort(b_begin, b_begin + k); } @@ -1497,8 +1503,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment acros cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { return cub::DeviceBatchedTopK::MaxPairs( d_temp, temp_bytes, @@ -1512,7 +1518,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment acros env); }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -1588,8 +1594,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { return cub::DeviceBatchedTopK::MaxPairs( d_temp, temp_bytes, @@ -1603,7 +1609,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with env); }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -1680,8 +1686,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { return cub::DeviceBatchedTopK::MaxPairs( d_temp, temp_bytes, @@ -1695,7 +1701,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with env); }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -1772,8 +1778,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { return cub::DeviceBatchedTopK::MaxPairs( d_temp, temp_bytes, @@ -1787,7 +1793,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with env); }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); @@ -1861,8 +1867,8 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_param = cuda::args::immediate{k, cuda::args::bounds()}; - size_t temp_bytes = 0; - const auto invoke = [&](void* d_temp) { + cuda::std::size_t temp_bytes = 0; + const auto invoke = [&](void* d_temp) { return cub::DeviceBatchedTopK::MaxPairs( d_temp, temp_bytes, @@ -1876,7 +1882,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs stream a tiny oversize segment with env); }; REQUIRE(cudaSuccess == invoke(nullptr)); - c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); REQUIRE(cudaSuccess == invoke(thrust::raw_pointer_cast(temp_storage.data()))); REQUIRE(cudaSuccess == cudaDeviceSynchronize()); diff --git a/cub/test/test_device_batched_topk_unsupported_arch_fail.cu b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu index 61e269a294f..aee6d6dfe16 100644 --- a/cub/test/test_device_batched_topk_unsupported_arch_fail.cu +++ b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -37,9 +38,9 @@ int main() ex::require(ex::determinism::gpu_to_gpu, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); // expected-error {{"is not supported on at least one architecture"}} - auto env = cuda::std::execution::env{requirements}; - size_t temp_storage_bytes = 0; - auto error = cub::DeviceBatchedTopK::MaxKeys( + auto env = cuda::std::execution::env{requirements}; + cuda::std::size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segments, env); if (error != cudaSuccess) { diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index d97fc2da9af..1f2bdad810b 100644 --- a/docs/cub/api_docs/device_topk_requirements.rst +++ b/docs/cub/api_docs/device_topk_requirements.rst @@ -174,7 +174,7 @@ execution environment alongside other properties such as a stream: cuda::execution::require( cuda::execution::determinism::gpu_to_gpu, cuda::execution::tie_break::prefer_smaller_index, - cuda::execution::output_ordering::sorted), + cuda::execution::output_ordering::unsorted), stream_ref}; .. _cub-topk-set-membership: From 563d2467b033fb269bb5908246e53f7850218f51 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 15 Jul 2026 02:07:22 +0200 Subject: [PATCH 124/126] Remove dead legacy dispatch functions --- .../device/dispatch/dispatch_batched_topk.cuh | 282 ------------------ .../dispatch/kernels/kernel_batched_topk.cuh | 80 ----- 2 files changed, 362 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 463bf4f2f32..da68abbe61c 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -113,287 +112,6 @@ struct segment_size_to_tile_count_op } }; -// ----------------------------------------------------------------------------- -// Segmented Top-K Dispatch -// ----------------------------------------------------------------------------- -// -// NOTE: `baseline_dispatch` / `baseline_dispatch_with_env` are the legacy baseline-only entry points. The public API -// now routes through the unified `dispatch` (baseline + cluster) below and no longer reaches these; they are retained -// for now and must stay in sync with the baseline arm of `dispatch`. - -//! @param d_temp_storage Device-accessible allocation of temporary storage. When `nullptr`, the required allocation -//! size is written to `temp_storage_bytes` and no work is done. -//! @param temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation -//! @param d_key_segments_it d_key_segments_it[segment_index] -> iterator to the input sequence of key data for segment -//! `segment_index` -//! @param d_key_segments_out_it d_key_segments_out_it[segment_index] -> iterator to the output sequence of key data for -//! segment `segment_index` -//! @param d_value_segments_it d_value_segments_it[segment_index] -> iterator to the input sequence of associated value -//! items for segment `segment_index`. When cub::NullType**, only keys are provided. -//! @param d_value_segments_out_it d_value_segments_out_it[segment_index] -> iterator to the output sequence of -//! associated value items for segment `segment_index` -//! @param segment_sizes Parameter providing segment sizes for each segment -//! @param k Parameter providing K for each segment -//! @param select_directions Parameter providing the selection direction for each segment -//! @param num_segments Number of segments -//! @param total_num_items_guarantee Allows the user to provide a guarantee on the upper bound of the total number of -//! items -template >, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>> -#if _CCCL_HAS_CONCEPTS() - requires baseline_topk_policy_selector -#endif // _CCCL_HAS_CONCEPTS() -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch( - void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionT select_direction, - NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr, - [[maybe_unused]] PolicySelector policy_selector = {}) -{ - using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; - - // Wrap the raw enum into the internal discrete param type - const auto select_directions = wrap_select_direction(select_direction); - using SelectDirectionParameterT = decltype(select_directions); - - // Helper that determines (a) whether there's any one-worker-per-segment policy supporting the range of segment - // sizes and k, and (b) if so, which set of one-worker-per-segment policies to use - constexpr auto policy = find_smallest_covering_policy< - PolicySelector, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - large_segment_tile_offset_t>::policy; - constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; - constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; - - static constexpr int worker_per_segment_tile_size = - worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; - static constexpr bool any_small_segments = - ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; - static constexpr bool only_small_segments = - ::cuda::args::__traits::highest <= worker_per_segment_tile_size; - - // Allocation layout: - // only_small_segments: [0] dummy. - // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, - // [2] large-segment ids. - // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. - static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); - size_t allocation_sizes[allocations_array_size] = {1}; - - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - using counters_t = batched_topk_counters; - using segment_size_scan_offset_t = detail::choose_offset_t; - using segment_size_scan_input_op_t = - segment_size_to_tile_count_op; - static constexpr auto multi_worker_per_segment_tile_size = - multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; - const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; - // Transform iterator over [0, num_segments) producing the tile-count for each segment. - [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( - ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); - - if constexpr (!only_small_segments) - { - const auto num_segments_val = params::get_param(num_segments, 0); - // Scan output - allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); - if constexpr (any_small_segments) - { - allocation_sizes[1] = sizeof(counters_t); - // Large segment ids for indirectly accessing the large segment parameters - allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); - } - else - { - // Query the temporary storage requirement of the segment-size transform-scan. - if (const auto error = CubDebug(detail::scan::dispatch( - nullptr, - allocation_sizes[1], - segment_size_scan_input_it, - static_cast(nullptr), - ::cuda::std::plus<>{}, - detail::InputValue(large_segment_tile_offset_t{0}), - static_cast(num_segments_val), - stream))) - { - return error; - } - } - } - - // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) - void* allocations[allocations_array_size] = {}; - if (const auto error = - CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) - { - return error; - } - - if (d_temp_storage == nullptr) - { - return cudaSuccess; - } - - // TODO (elstehle): support number of segments provided by device-accessible iterator - // Only uniform number of segments are supported (i.e., we need to resolve the number of segments on the host) - static_assert(::cuda::args::__traits::is_single_value - && !::cuda::args::__traits::is_deferred, - "Only a host-known uniform number of segments is currently supported; a deferred (device-resident) " - "count is not yet supported (see the TODO above)."); - - // No work to launch when there are no segments; otherwise the grid launches below would use `grid_dim = 0`, an - // invalid launch configuration. - if (params::get_param(num_segments, 0) == 0) - { - return cudaSuccess; - } - - if constexpr (any_small_segments) - { - if constexpr (!only_small_segments) - { - // Zero-initialize the counters struct that holds the large-segment queue length and the block retirement - // counter; both are read by the agent's atomic operations and must start at 0. - if (const auto error = CubDebug(cudaMemsetAsync(allocations[1], 0, sizeof(counters_t), stream))) - { - return error; - } - } - const int grid_dim = static_cast(params::get_param(num_segments, 0)); - constexpr int block_dim = worker_per_segment_policy.threads_per_block; - if (const auto error = CubDebug( - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid_dim, block_dim, 0, stream) - .doit( - device_segmented_topk_kernel< - PolicySelector, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - large_segment_tile_offset_t>, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - only_small_segments ? nullptr : static_cast(allocations[1]), - only_small_segments ? nullptr : static_cast(allocations[2]), - only_small_segments ? nullptr : static_cast(allocations[0])))) - { - return error; - } - } - else - { - // No small segments: the small-kernel epilogue (which would otherwise produce the per-segment tile offsets) does - // not run. Compute the per-segment tile offsets directly via a transform-scan over all segment sizes. - // The large segment agent will either consume these offsets directly (segment_id -> tile offset) or, when going - // through the large-segment queue, via a transform iterator over `d_large_segments_ids` (level of indirection). - if (const auto error = CubDebug(detail::scan::dispatch( - allocations[1], - allocation_sizes[1], - segment_size_scan_input_it, - static_cast(allocations[0]), - ::cuda::std::plus<>{}, - detail::InputValue(large_segment_tile_offset_t{0}), - static_cast(params::get_param(num_segments, 0)), - stream))) - { - return error; - } - } - - if constexpr (!only_small_segments) - { - // TODO (elstehle): support larger number of segments through multiple kernel launches - // Depending on any_small_segments, we need to either: - // - Indirectly get the large segment parameters via the queued large segment IDs - // - Directly take the segment parameters since all segments are large - } - return CubDebug(detail::DebugSyncStream(stream)); -} -// Env-based dispatch function handling memory allocation as well. This is usually done by the device-layer, but there -// is no public API for segmented topk yet. -template > -[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t baseline_dispatch_with_env( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, - EnvT env = {}) -{ - using default_policy_selector = - baseline_policy_selector_from_types>, - it_value_t>, - ::cuda::std::int64_t, - ::cuda::args::__traits::highest>; - return detail::dispatch_with_env_and_tuning( - env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { - return baseline_dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - total_num_items_guarantee, - stream, - policy_selector); - }); -} - // ----------------------------------------------------------------------------- // Dispatch (both backends behind one kernel symbol) // ----------------------------------------------------------------------------- diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index fc109ff8bc7..59772de6687 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -130,86 +130,6 @@ public: using agent_t = agent_batched_topk_worker_per_segment; }; -// ----------------------------------------------------------------------------- -// Global Kernel Entry Point -// ----------------------------------------------------------------------------- -template -#if _CCCL_HAS_CONCEPTS() - requires baseline_topk_policy_selector -#endif // _CCCL_HAS_CONCEPTS() -__launch_bounds__(int( - find_smallest_covering_policy< - PolicySelector, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>::policy.worker_per_segment_policy.threads_per_block)) - _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_kernel( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - batched_topk_counters::element_type>* d_counters, - typename ::cuda::args::__traits::element_type* d_large_segments_ids, - LargeSegmentTileOffsetT* d_large_segments_tile_offsets) -{ - using agent_t = typename find_smallest_covering_policy< - PolicySelector, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>::agent_t; - - // Static Assertions (Constraints) - static_assert(agent_t::tile_size >= ::cuda::args::__traits::highest, - "Block size exceeds maximum segment size supported by SegmentSizeParameterT"); - static_assert(sizeof(typename agent_t::TempStorage) <= max_smem_per_block, - "Static shared memory per block must not exceed 48KB limit."); - - __shared__ typename agent_t::TempStorage temp_storage; - - agent_t agent( - temp_storage, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - d_counters, - d_large_segments_ids, - d_large_segments_tile_offsets); - - agent.Process(); -} - // ----------------------------------------------------------------------------- // Single kernel symbol hosting both backends // ----------------------------------------------------------------------------- From fc46a4affadd09e8d0a072f32462e716c5168d7b Mon Sep 17 00:00:00 2001 From: pauleonix Date: Wed, 15 Jul 2026 03:03:00 +0200 Subject: [PATCH 125/126] Make CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERR public --- .../bench/segmented_topk/fixed/keys.cu | 6 ++--- .../variable/indexed_common.cuh | 6 ++--- .../segmented_topk/variable/keys_common.cuh | 6 ++--- cub/cub/device/device_batched_topk.cuh | 22 ++++++++++++++++- .../device/dispatch/dispatch_batched_topk.cuh | 24 +++++++++---------- cub/test/CMakeLists.txt | 2 +- .../catch2_test_device_segmented_topk_keys.cu | 4 ++-- ...catch2_test_device_segmented_topk_pairs.cu | 4 ++-- cub/test/catch2_test_device_topk_common.cuh | 2 +- ...t_device_batched_topk_requirements_fail.cu | 7 +++--- ...vice_batched_topk_unsupported_arch_fail.cu | 4 ++-- .../cub/api_docs/device_topk_requirements.rst | 4 ++-- 12 files changed, 55 insertions(+), 36 deletions(-) diff --git a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index a9f8a1a8614..1a70bc95908 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -3,9 +3,9 @@ // Defer the unsupported-architecture diagnosis to the dispatch's runtime check so this benchmark compiles for the full // configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) -// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in -// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// across all target architectures, including pre-SM90. See CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/device_batched_topk.cuh. Must precede the CUB includes below. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include #include diff --git a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh index 438b776a3e8..5f5cf6ab1ad 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -12,9 +12,9 @@ // Defer the unsupported-architecture diagnosis to the dispatch's runtime check so these benchmarks compile for the full // configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) -// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in -// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// across all target architectures, including pre-SM90. See CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/device_batched_topk.cuh. Must precede the CUB includes below. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include #include diff --git a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh index 3b0f554c513..460fa0c023b 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -12,9 +12,9 @@ // Defer the unsupported-architecture diagnosis to the dispatch's runtime check so these benchmarks compile for the full // configuration space (including deterministic / large-segment requests, which only the SM90+ cluster backend serves) -// across all target architectures, including pre-SM90. See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in -// cub/device/dispatch/dispatch_batched_topk.cuh. Must precede the CUB includes below. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// across all target architectures, including pre-SM90. See CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in +// cub/device/device_batched_topk.cuh. Must precede the CUB includes below. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include #include diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 478aacf4a82..7e539b7553b 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -41,6 +41,25 @@ #include #include +#ifdef _CCCL_DOXYGEN_INVOKED // Only parse this during doxygen passes: + +//! @def CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +//! +//! Specific to cub::DeviceBatchedTopK (it has no effect on any other CUB algorithm). +//! +//! By default, cub::DeviceBatchedTopK fails at compile time (via `static_assert`) when the requested configuration +//! cannot be served on *every* compute capability the translation unit is being compiled for. Some requests (a +//! deterministic result, or a segment too large for the single-block backend) require the SM90+ cluster backend, so +//! they cannot compile when a pre-SM90 compute capability is among the targets. +//! +//! Define this macro (before including any CUB header) to suppress that compile-time check and defer the diagnosis to +//! runtime instead: on a device that cannot serve the request, dispatch returns `cudaErrorNotSupported`. This is +//! useful when a single translation unit must compile the full configuration space across a mix of compute +//! capabilities and decide what is runnable at runtime. CUB's own tests and benchmarks define it for this reason. +# define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + +#endif // _CCCL_DOXYGEN_INVOKED + CUB_NAMESPACE_BEGIN namespace detail @@ -418,7 +437,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! - **Pre-Hopper (compute capability < 9.0):** only the fully non-deterministic request //! ``(determinism::not_guaranteed, tie_break::unspecified)`` is supported, and every segment must fit a single //! thread block. Deterministic / tie-break requests and larger segments require the SM 9.0+ cluster backend and -//! are diagnosed at compile time (or, in relaxed builds, at runtime as ``cudaErrorNotSupported``). +//! are diagnosed at compile time (or, when ``CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT`` is defined, deferred to +//! runtime as ``cudaErrorNotSupported``). //! - **Hopper and newer (compute capability >= 9.0):** all five acknowledged ``(determinism, tie_break)`` pairs are //! supported -- ``(not_guaranteed, unspecified)``, ``(run_to_run, unspecified)``, and ``(gpu_to_gpu, //! {unspecified, prefer_smaller_index, prefer_larger_index})`` -- and larger segments are handled by the cluster diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index da68abbe61c..477377fa4b8 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -871,7 +871,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_baseline_arm( } #if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) -// Returns true if at least one architecture in the compile target list (`CMAKE_CUDA_ARCHITECTURES`, exposed as +// Returns true if at least one architecture this translation unit targets (the compile target list exposed as // `::cuda::__target_compute_capabilities()`) resolves to the `unsupported` backend for `PolicySelector` -- e.g. a // deterministic request while a pre-SM90 target is present in the list. Used to turn a would-be runtime // `cudaErrorNotSupported` into a compile-time diagnostic (see the static_assert in `dispatch`). @@ -993,21 +993,21 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( selector_override_adaptor>; #if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ - && !defined(_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - // Strict mode (default): fail at compile time if the request cannot be served on *any* architecture in - // `CMAKE_CUDA_ARCHITECTURES` (e.g. a deterministic / large-segment request while a pre-SM90 target is present, since - // the cluster backend requires SM90+). This is the least-surprising UX for callers building the default multi-arch - // preset. Define `_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` to defer the diagnosis to runtime instead (the dispatch - // then returns `cudaErrorNotSupported` on unsupported devices); CUB's own tests and benchmarks do this so they can - // compile the full configuration space across all target architectures and skip at runtime where unsupported. + && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // Strict mode (default): fail at compile time if the request cannot be served on *any* architecture this translation + // unit targets (e.g. a deterministic / large-segment request while a pre-SM90 target is present, since the cluster + // backend requires SM90+). This is the least-surprising UX for callers whose build targets multiple architectures. + // Define `CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` to defer the diagnosis to runtime instead (the dispatch then + // returns `cudaErrorNotSupported` on unsupported devices); CUB's own tests and benchmarks do this so they can compile + // the full configuration space across all target architectures and skip at runtime where unsupported. static_assert( !any_target_cc_unsupported(), - "cub::DeviceBatchedTopK: the requested top-k configuration is not supported on at least one architecture in " - "CMAKE_CUDA_ARCHITECTURES (the deterministic / large-segment cluster backend requires SM90+). Remove the " - "unsupported architecture(s), relax the request, or define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the " + "cub::DeviceBatchedTopK: the requested top-k configuration is not supported on at least one architecture this " + "translation unit targets (the deterministic / large-segment cluster backend requires SM90+). Remove the " + "unsupported architecture(s), relax the request, or define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the " "diagnosis to runtime (cudaErrorNotSupported)."); #endif // _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) - // && !defined(_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) // The supported maximum segment size (2^21) is enforced at compile time at the public entry; a statically negative // lower bound is allowed and negative runtime sizes are clamped to 0 (see detail::params::get_segment_size). A diff --git a/cub/test/CMakeLists.txt b/cub/test/CMakeLists.txt index 091fa8e89fb..b2537a25915 100644 --- a/cub/test/CMakeLists.txt +++ b/cub/test/CMakeLists.txt @@ -294,7 +294,7 @@ function( # Ensure that we test with assertions enabled target_compile_definitions(${test_target} PRIVATE CCCL_ENABLE_ASSERTIONS) - # The unsupported-arch xfail test omits the _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT define its siblings use, so it + # The unsupported-arch xfail test omits the CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT define its siblings use, so it # keeps the strict compile-time diagnostic. Pin it to 89-virtual (the highest arch without cluster support) so the # deterministic request it issues has no viable backend and the static_assert fires regardless of the preset's archs. if ( diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 272352b358b..464ff228301 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -3,8 +3,8 @@ // Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) // so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See -// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/device_batched_topk.cuh. Precedes CUB includes. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include "insert_nested_NVTX_range_guard.h" diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 6683ca06ca0..c1247712859 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -3,8 +3,8 @@ // Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) // so this test compiles across all target architectures, including pre-SM90, for the full configuration space. See -// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh. Precedes CUB includes. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/device_batched_topk.cuh. Precedes CUB includes. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include "insert_nested_NVTX_range_guard.h" diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index b1edd681075..87bdc636318 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -20,7 +20,7 @@ #include // Low-level gate: skips the current test case when a request that must use the SM90+ cluster backend has no -// cluster-capable target in the build. With `_CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` defined (as the top-k test +// cluster-capable target in the build. With `CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT` defined (as the top-k test // sources do), such a request fails at runtime with cudaErrorNotSupported rather than at compile time. `needs_cluster` // must be true when the caller knows the configuration requires the cluster backend; prefer // skip_if_batched_topk_backend_unavailable(), which derives it from the request. diff --git a/cub/test/test_device_batched_topk_requirements_fail.cu b/cub/test/test_device_batched_topk_requirements_fail.cu index b6109c8c893..cbf885988c3 100644 --- a/cub/test/test_device_batched_topk_requirements_fail.cu +++ b/cub/test/test_device_batched_topk_requirements_fail.cu @@ -4,10 +4,9 @@ // %PARAM% TEST_ERR err 0:1:2:3:4:5:6:7:8:9:10:11:12:13 // Defer the unsupported-architecture diagnosis to the dispatch's runtime check (not a compile-time static_assert) -// so only the requirement static_asserts under test fire, regardless of the target architecture (this test compiles for -// the preset's default archs, which include pre-SM90). See _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in -// cub/device/dispatch/dispatch_batched_topk.cuh. Precedes the CUB include below. -#define _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT +// so only the requirement static_asserts under test fire, regardless of which architectures this test is compiled for. +// See CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/device_batched_topk.cuh. Precedes the CUB include below. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT #include diff --git a/cub/test/test_device_batched_topk_unsupported_arch_fail.cu b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu index aee6d6dfe16..054fc81f065 100644 --- a/cub/test/test_device_batched_topk_unsupported_arch_fail.cu +++ b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu @@ -15,12 +15,12 @@ #include // Verifies the strict unsupported-architecture diagnostic in cub::DeviceBatchedTopK (see -// _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/dispatch/dispatch_batched_topk.cuh). A deterministic +// CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT in cub/device/device_batched_topk.cuh). A deterministic // (gpu_to_gpu) request can only be served by the SM90+ cluster backend; when the translation unit targets an // architecture that cannot run it (this target is pinned to a pre-SM90 arch via CUDA_ARCHITECTURES in CMakeLists.txt), // the default (strict) mode must fail to compile rather than defer the failure to a runtime cudaErrorNotSupported. // -// This is the only batched/segmented top-k test built *without* _CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT, so it +// This is the only batched/segmented top-k test built *without* CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT, so it // exercises the strict path that CUB's other top-k tests intentionally opt out of. int main() diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index 1f2bdad810b..7e643c7d02b 100644 --- a/docs/cub/api_docs/device_topk_requirements.rst +++ b/docs/cub/api_docs/device_topk_requirements.rst @@ -74,8 +74,8 @@ requires ``determinism::gpu_to_gpu``. See :ref:`cub-topk-set-membership` for the - **Pre-Hopper (compute capability < 9.0):** only the fully non-deterministic request ``(not_guaranteed, unspecified)`` is supported, and every segment must fit a single thread block. Deterministic / tie-break requests and larger segments require the SM 9.0+ cluster - backend and are diagnosed at compile time (or at runtime as ``cudaErrorNotSupported`` in - relaxed builds). + backend and are diagnosed at compile time (or, when ``CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT`` + is defined, deferred to runtime as ``cudaErrorNotSupported``). - **Hopper and newer (compute capability >= 9.0):** every acknowledged ``(determinism, tie_break)`` pair is supported -- ``(not_guaranteed, unspecified)``, ``(run_to_run, unspecified)``, and ``(gpu_to_gpu, {unspecified, prefer_smaller_index, prefer_larger_index})`` From 5d86cdc68044c0f4af634047c0b4d1f9845dcb11 Mon Sep 17 00:00:00 2001 From: pauleonix Date: Thu, 16 Jul 2026 07:17:38 +0200 Subject: [PATCH 126/126] Fix issues with Clang as host compiler --- .clang-format | 2 + .../device/dispatch/dispatch_batched_topk.cuh | 3 + .../dispatch/kernels/kernel_batched_topk.cuh | 246 ++++++++++-------- .../cuda/std/__cccl/cuda_capabilities.h | 22 ++ 4 files changed, 161 insertions(+), 112 deletions(-) diff --git a/.clang-format b/.clang-format index bfcbb33c711..36b10810e07 100644 --- a/.clang-format +++ b/.clang-format @@ -40,7 +40,9 @@ AttributeMacros: [ '_CCCL_TYPE_VISIBILITY_HIDDEN', '_CCCL_VISIBILITY_HIDDEN', '_CCCL_LAUNCH_BOUNDS', + '_CCCL_LAUNCH_BOUNDS_CLUSTER', '_CCCL_BLOCK_SIZE', + '_CCCL_CLUSTER_DIMS', 'CUB_RUNTIME_FUNCTION', 'THRUST_RUNTIME_FUNCTION', 'CCCL_DEPRECATED', diff --git a/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cub/cub/device/dispatch/dispatch_batched_topk.cuh index 477377fa4b8..69cfe334f4f 100644 --- a/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -147,6 +147,7 @@ template MinChunksPerBlock, \ CopyItemsPerThread, \ cdp_cluster_blocks, \ + policy.min_blocks_per_sm, \ Determinism, \ TieBreak, \ KeyInputItItT, \ @@ -657,6 +658,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t launch_cluster_arm( static_assert(block_tile_capacity >= static_cast<::cuda::std::uint32_t>(layout_t::chunk_items), "Portable SMEM is too small to fit even one load-aligned chunk for the device-launch (CDP) path"); + // Explicit-cluster semantics (`_CCCL_CLUSTER_DIMS`): the launch grid counts blocks, so one cluster spans every + // `cdp_cluster_blocks` blocks; the cluster width itself is fixed by the kernel attributes. const auto grid_blocks = static_cast<::cuda::std::uint64_t>(num_seg_val) * static_cast<::cuda::std::uint64_t>(cdp_cluster_blocks); if (grid_blocks > static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())) diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 59772de6687..4c38acc557b 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -158,9 +158,9 @@ struct cluster_kernel_args // Launch-bounds helpers // ----------------------------------------------------------------------------- // The two backends use different `__launch_bounds__` shapes (baseline: just threads_per_block; cluster: threads plus a -// min-blocks-per-SM). We resolve both per architecture from the selected policy. `find_smallest_covering_policy` (which -// carries a hard `static_assert`) is only ever touched inside the `backend == baseline` branch, so an oversize bound -// routed to the cluster/unsupported backend never trips it. +// min-blocks-per-SM and an optional max-blocks-per-cluster cap). We resolve all three per architecture from the +// selected policy. `find_smallest_covering_policy` (which carries a hard `static_assert`) is only ever touched inside +// the `backend == baseline` branch, so an oversize bound routed to the cluster/unsupported backend never trips it. _CCCL_EXEC_CHECK_DISABLE template [[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL int topk_threads_per_block_helper() noexcept @@ -184,7 +184,7 @@ template +template [[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL int topk_min_blocks_per_sm_helper() noexcept { constexpr auto policy = current_policy(); @@ -199,19 +199,43 @@ template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL int topk_max_blocks_per_cluster_helper() noexcept +{ + constexpr auto policy = current_policy(); + if constexpr (policy.backend == topk_backend::cluster) + { + return policy.cluster.max_blocks_per_cluster; + } + else + { + // baseline / unsupported: not a cluster launch, so no cluster-width cap. + return 0; + } +} + // Variable templates force constant evaluation of the helpers, otherwise nvcc reports a "bad attribute argument // substitution" error on the `__launch_bounds__` below (same pattern as `transform_kernel`). template inline constexpr int topk_threads_per_block = topk_threads_per_block_helper(); -template -inline constexpr int topk_min_blocks_per_sm = - topk_min_blocks_per_sm_helper(); +template +inline constexpr int topk_min_blocks_per_sm = topk_min_blocks_per_sm_helper(); + +template +inline constexpr int topk_max_blocks_per_cluster = topk_max_blocks_per_cluster_helper(); // ----------------------------------------------------------------------------- // Global kernel entry point (single symbol for both backends) // ----------------------------------------------------------------------------- +// Launch bounds: only `topk_threads_per_block` takes the full kernel type list (its baseline branch runs the +// covering-policy search); min/max-blocks depend on `PolicySelector` alone. The parentheses around +// `topk_threads_per_block<...>` hide its template commas from the fixed-arity `_CCCL_LAUNCH_BOUNDS_CLUSTER(a, b, c)`. template -__launch_bounds__( - topk_threads_per_block< - PolicySelector, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>, - topk_min_blocks_per_sm< - PolicySelector, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>) - _CCCL_KERNEL_ATTRIBUTES void device_batched_topk_kernel( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - baseline_kernel_args::element_type, LargeSegmentTileOffsetT> - base_args, - [[maybe_unused]] cluster_kernel_args clus_args) +_CCCL_LAUNCH_BOUNDS_CLUSTER((topk_threads_per_block), + topk_min_blocks_per_sm, + topk_max_blocks_per_cluster) _CCCL_KERNEL_ATTRIBUTES void +device_batched_topk_kernel( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + baseline_kernel_args::element_type, LargeSegmentTileOffsetT> + base_args, + [[maybe_unused]] cluster_kernel_args clus_args) { constexpr auto policy = current_policy(); @@ -303,7 +315,7 @@ __launch_bounds__( } else if constexpr (policy.backend == topk_backend::cluster) { - NV_IF_TARGET( + NV_IF_ELSE_TARGET( NV_PROVIDES_SM_90, (using agent_t = batched_topk_cluster::agent_batched_topk_cluster< policy.cluster.threads_per_block, @@ -353,7 +365,9 @@ __launch_bounds__( key_slots, clus_args.block_tile_capacity); - agent.Process();)); + agent.Process();), + // Cluster-policy kernels are only ever launched on SM90+, so the sub-SM90 device pass is unreachable at runtime. + (_CCCL_UNREACHABLE();)); } else { @@ -363,13 +377,12 @@ __launch_bounds__( } #ifdef CUB_RDC_ENABLED -// CDP-only static-cluster kernel: compile-time `__cluster_dims__` so a device-side (CDP) triple-chevron launch needs no -// `cudaFuncSetAttribute` (device-side launches cannot opt into dynamic cluster dimensions the way the host -// `device_batched_topk_kernel` does). It mirrors that kernel's cluster arm with a fixed cluster dim, and is consumed by -// the CDP arm of `launch_cluster_arm` (see `CUB_TOPK_CLUSTER_DEVICE_LAUNCH` in dispatch_batched_topk.cuh). -// `ClusterBlocks` is the compile-time cluster width the CDP arm picks: `<= max_portable_cluster_blocks` (device -// launches can't opt into non-portable cluster sizes), and further narrowed when the policy's `max_blocks_per_cluster` -// cap is tighter. +// CDP-only static-cluster kernel: a compile-time `__cluster_dims__` (via `_CCCL_CLUSTER_DIMS`) lets a device-side (CDP) +// triple-chevron launch skip the `cudaFuncSetAttribute` dynamic-cluster-dim call the host `device_batched_topk_kernel` +// relies on (device launches can't make it). Mirrors that kernel's cluster arm; consumed by +// `CUB_TOPK_CLUSTER_DEVICE_LAUNCH` in dispatch_batched_topk.cuh. `ClusterBlocks` is the cluster width (<= +// `max_portable_cluster_blocks`, narrowed by the policy's `max_blocks_per_cluster`); `MinBlocksPerSm` forwards the +// `min_blocks_per_sm` hint via `_CCCL_LAUNCH_BOUNDS`, active under EWP but stubbed under RDC (#902). template -__launch_bounds__(ThreadsPerBlock) __cluster_dims__(ClusterBlocks, 1, 1) - _CCCL_KERNEL_ATTRIBUTES void device_segmented_topk_cluster_kernel_static( - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - ::cuda::std::uint32_t block_tile_capacity) +_CCCL_CLUSTER_DIMS(ClusterBlocks, 1, 1) _CCCL_LAUNCH_BOUNDS(ThreadsPerBlock, MinBlocksPerSm) +_CCCL_KERNEL_ATTRIBUTES void +device_segmented_topk_cluster_kernel_static( + [[maybe_unused]] KeyInputItItT d_key_segments_it, + [[maybe_unused]] KeyOutputItItT d_key_segments_out_it, + [[maybe_unused]] ValueInputItItT d_value_segments_it, + [[maybe_unused]] ValueOutputItItT d_value_segments_out_it, + [[maybe_unused]] SegmentSizeParameterT segment_sizes, + [[maybe_unused]] KParameterT k_param, + [[maybe_unused]] SelectDirectionParameterT select_directions, + [[maybe_unused]] NumSegmentsParameterT num_segments, + [[maybe_unused]] ::cuda::std::uint32_t block_tile_capacity) { - using agent_t = batched_topk_cluster::agent_batched_topk_cluster< - ThreadsPerBlock, - HistogramItemsPerThread, - PipelineStages, - ChunkBytes, - LoadAlignBytes, - BitsPerPass, - TieBreakItemsPerThread, - SingleBlockMaxSegSize, - MinChunksPerBlock, - CopyItemsPerThread, - Determinism, - TieBreak, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT>; - - __shared__ typename agent_t::TempStorage temp_storage; - extern __shared__ char topk_cluster_smem[]; - char* key_slots = topk_cluster_smem; - // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` - // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes` for this. - { - ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); - smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); - asm("" : "+r"(smem32)); - key_slots = static_cast(__cvta_shared_to_generic(smem32)); - } - - agent_t agent( - temp_storage, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k_param, - select_directions, - num_segments, - key_slots, - block_tile_capacity); - - agent.Process(); + // The agent's cluster/async PTX only assembles on SM90+, so gate the whole body on `NV_PROVIDES_SM_90`: host and + // sub-SM90 device passes emit an empty kernel and never instantiate the agent (the CDP arm only launches this on + // SM90+). Mirrors the cluster arm of `device_batched_topk_kernel`. + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_90, + (using agent_t = batched_topk_cluster::agent_batched_topk_cluster< + ThreadsPerBlock, + HistogramItemsPerThread, + PipelineStages, + ChunkBytes, + LoadAlignBytes, + BitsPerPass, + TieBreakItemsPerThread, + SingleBlockMaxSegSize, + MinChunksPerBlock, + CopyItemsPerThread, + Determinism, + TieBreak, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT>; + + __shared__ typename agent_t::TempStorage temp_storage; + extern __shared__ char topk_cluster_smem[]; + char* key_slots = topk_cluster_smem; + // Align the base up to `slot_alignment` (>= load_align) so every bulk-copy destination gets the same `load_align` + // alignment the gmem sources have (peak TMA throughput on Hopper). The layout reserves `base_padding_bytes`. + { + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(key_slots); + smem32 = ::cuda::round_up(smem32, static_cast<::cuda::std::uint32_t>(agent_t::slot_alignment)); + asm("" : "+r"(smem32)); + key_slots = static_cast(__cvta_shared_to_generic(smem32)); + } + + agent_t agent( + temp_storage, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k_param, + select_directions, + num_segments, + key_slots, + block_tile_capacity); + + agent.Process();), + // Cluster-policy kernels are only ever launched on SM90+, so the sub-SM90 device pass is unreachable at runtime. + (_CCCL_UNREACHABLE();)); } #endif // CUB_RDC_ENABLED } // namespace detail::batched_topk diff --git a/libcudacxx/include/cuda/std/__cccl/cuda_capabilities.h b/libcudacxx/include/cuda/std/__cccl/cuda_capabilities.h index 2182bf5f96f..2938738636c 100644 --- a/libcudacxx/include/cuda/std/__cccl/cuda_capabilities.h +++ b/libcudacxx/include/cuda/std/__cccl/cuda_capabilities.h @@ -109,6 +109,28 @@ # define _CCCL_BLOCK_SIZE(_NTID, _NCTA_PER_CLUSTER) #endif // ^^^ no __block_size__ attribute ^^^ +// __cluster_dims__ attribute (nvcc, nvrtc, and clang-cuda 22+, hopper+): pins the cluster dimension at compile time +// with explicit-cluster launch semantics, so unlike _CCCL_BLOCK_SIZE it also works for device-side (CDP) launches. +// This attribute should be used only for cluster launches on sm90+. +#if (_CCCL_CUDA_COMPILER(NVCC) || _CCCL_CUDA_COMPILER(NVRTC) || _CCCL_CUDA_COMPILER(CLANG, >=, 22)) \ + && _CCCL_PTX_ARCH() >= 900 +# define _CCCL_CLUSTER_DIMS(...) __cluster_dims__(__VA_ARGS__) +#else // ^^^ has __cluster_dims__ attribute ^^^ / vvv no __cluster_dims__ attribute vvv +# define _CCCL_CLUSTER_DIMS(...) +#endif // ^^^ no __cluster_dims__ attribute ^^^ + +// `_CCCL_LAUNCH_BOUNDS` plus the optional third `maxBlocksPerCluster` operand, which nvcc, nvrtc, and clang-cuda 18+ +// (clang gained `.maxclusterrank` in 18) accept on the SM90+ device pass but reject on host/pre-SM90 passes. So emit it +// only there and fall back to the 2-argument form elsewhere -- those passes never launch clusters. +#if (_CCCL_CUDA_COMPILER(NVCC) || _CCCL_CUDA_COMPILER(NVRTC) || _CCCL_CUDA_COMPILER(CLANG, >=, 18)) \ + && _CCCL_PTX_ARCH() >= 900 +# define _CCCL_LAUNCH_BOUNDS_CLUSTER(_MAX_THREADS_PER_BLOCK, _MIN_BLOCKS_PER_SM, _MAX_BLOCKS_PER_CLUSTER) \ + _CCCL_LAUNCH_BOUNDS(_MAX_THREADS_PER_BLOCK, _MIN_BLOCKS_PER_SM, _MAX_BLOCKS_PER_CLUSTER) +#else // ^^^ SM90+ device pass ^^^ / vvv host / pre-SM90 / non-CUDA pass vvv +# define _CCCL_LAUNCH_BOUNDS_CLUSTER(_MAX_THREADS_PER_BLOCK, _MIN_BLOCKS_PER_SM, _MAX_BLOCKS_PER_CLUSTER) \ + _CCCL_LAUNCH_BOUNDS(_MAX_THREADS_PER_BLOCK, _MIN_BLOCKS_PER_SM) +#endif // SM90+ CUDA device pass (nvcc / nvrtc / clang-cuda) + #if _CCCL_HAS_CDP() # ifdef CUDA_FORCE_CDP1_IF_SUPPORTED # error "CUDA Dynamic Parallelism 1 is no longer supported. Please undefine CUDA_FORCE_CDP1_IF_SUPPORTED."