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/benchmarks/bench/segmented_topk/fixed/keys.cu b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu index 77bd997b742..1a70bc95908 100644 --- a/cub/benchmarks/bench/segmented_topk/fixed/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/fixed/keys.cu @@ -1,25 +1,95 @@ // 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/device_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 // %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, +}; + +// 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. 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 " + "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; @@ -28,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}, @@ -36,9 +106,89 @@ 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. `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). + const 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. + 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{}, + 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; + 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); + } + } +} template void fixed_seg_size_topk_keys( @@ -51,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 @@ -63,14 +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{}; - auto select_direction = ::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"); @@ -78,29 +226,22 @@ 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. Consumed only by the + // per-segment device backend. Segment sizes fit in a signed 32-bit integer (the library caps them at 2^21). + 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 -#if !TUNE_BASE - , - cuda::execution::tune(tuned_policy_selector{}) -#endif // !TUNE_BASE - ); - // TODO(bgruber): call the public API once available + const auto env = cub_bench_env(alloc, launch); _CCCL_TRY_CUDA_API( - cub::detail::batched_topk::dispatch_with_env, + batched_topk_keys, "batched topk failed", d_keys_in, d_keys_out, - static_cast(nullptr), - static_cast(nullptr), 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/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.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 8cb72ce26f7..7583851dc0a 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/indexed.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed.cu @@ -1,99 +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 "common.cuh" - -// 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( - cub::detail::batched_topk::dispatch_with_env, - "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..5f5cf6ab1ad --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/indexed_common.cuh @@ -0,0 +1,282 @@ +// 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/device_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 +#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 +{ + 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, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; +#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) +{ + 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{}, + 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; + 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); + } +} + +// 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); + + const auto segment_sizes_param = cuda::args::deferred_sequence{ + thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; + const auto k_param = cuda::args::constant{}; + const auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + + const auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), + static_cast(MaxSegmentSize)); + 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. + 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)); + + 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) { + const 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 f5c6e4e4949..f5a48a338c4 100644 --- a/cub/benchmarks/bench/segmented_topk/variable/keys.cu +++ b/cub/benchmarks/bench/segmented_topk/variable/keys.cu @@ -1,83 +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 "common.cuh" - -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"); - - 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( - cub::detail::batched_topk::dispatch_with_env, - "batched topk failed", - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes_param, - k_param, - select_direction, - num_segments_param, - total_num_items, - 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..460fa0c023b --- /dev/null +++ b/cub/benchmarks/bench/segmented_topk/variable/keys_common.cuh @@ -0,0 +1,306 @@ +// 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/device_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 +#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 +{ + 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, + /*max_blocks_per_cluster=*/0, + /*max_chunk_slots_per_block=*/0}; +#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). + const 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. + 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{}, + 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; + 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); + } + } +} + +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); + + const auto segment_sizes_param = cuda::args::deferred_sequence{ + thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()}; + const auto k_param = cuda::args::constant{}; + const auto num_segments_param = cuda::args::immediate{static_cast(num_segments)}; + + const auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())), + static_cast(MaxSegmentSize)); + const 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) { + const 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); diff --git a/cub/cub/agent/agent_batched_topk.cuh b/cub/cub/agent/agent_batched_topk.cuh index d2a99cfc809..e7fd55157a1 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,108 +218,115 @@ 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)); - 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 + // 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) { - // 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(); + // Load Values (if applicable) + [[maybe_unused]] value_t thread_values[items_per_thread]; - // 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"); - } - - __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/agent/agent_batched_topk_cluster.cuh b/cub/cub/agent/agent_batched_topk_cluster.cuh new file mode 100644 index 00000000000..ec4f22e7083 --- /dev/null +++ b/cub/cub/agent/agent_batched_topk_cluster.cuh @@ -0,0 +1,3281 @@ +// 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 (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 +//! 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 +//! a dual role: its own block-private histogram in step 1, then the +//! cluster-merged histogram after the second cluster sync. +//! 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. +//! +//! The final filter places each block's output through per-CTA shared-memory +//! 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 + +#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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk_cluster +{ +// ----------------------------------------------------------------------------- +// 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; + // 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; + + // 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) + { + return static_cast<::cuda::std::uint32_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 is_stop) + { + result_pair = static_cast<::cuda::std::uint64_t>(bucket) | (static_cast<::cuda::std::uint64_t>(is_stop) << 32); + } +}; + +// 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 +{ + 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)}; + // 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 / ChunkBytes; + return static_cast<::cuda::std::uint32_t>(slots * chunk_items); + } + + template + [[nodiscard]] _CCCL_HOST_DEVICE static constexpr SizeT + cluster_tile_capacity(int cluster_blocks, ::cuda::std::uint32_t physical_block_tile_capacity) noexcept + { + return static_cast(cluster_blocks) * static_cast(physical_block_tile_capacity); + } +}; + +// ----------------------------------------------------------------------------- +// 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 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 + && 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_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 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( + ::cuda::std::clamp(blocks, ::cuda::std::uint64_t{1}, static_cast<::cuda::std::uint64_t>(cluster_blocks_cap))); +} + +// Cluster top-k agent +// ----------------------------------------------------------------------------- +// 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. +// +// `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 +struct agent_batched_topk_cluster +{ + // --------------------------------------------------------------------------- + // Types / constants + // --------------------------------------------------------------------------- + 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::args::__traits::element_type; + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + + // 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). + 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; + + // 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 is_tie_reversed = + TieBreak == ::cuda::execution::tie_break::__tie_break_t::__prefer_larger_index; + + // 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 + // `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 (`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 + // `process_impl`). + static constexpr int single_block_max_seg_size = SingleBlockMaxSegSize; + + // A CTA joins a segment's effective cluster only if it would own at least this many chunks (see `run`). At 1 the + // effective cluster is just the CTAs that receive any chunk. Must be positive: it is the divisor in + // `effective_cluster_blocks_from_chunks` and a zero-chunk CTA has no work. + 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 + // `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 + || ::cuda::args::__traits::is_deferred; + + 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 + + // 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 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)); + + // 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`) 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 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 = + 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 = + 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 `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 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 + // 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); + static constexpr int tie_break_items_per_thread_clamped = + clamp_unroll(segment_rounds_ceil, 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(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; + 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 overflow stream's per-stage view); elsewhere: `key_t*` + + // count. + using smem_keys_t = ::cuda::std::span; + + // 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"); + + 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"); + // `__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"); + 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"); + // The hybrid load relies on the aligned bulk-copy path being exact (no scalar guard), which requires the load + // 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; + // 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 + && ::cuda::std::__can_to_address && key_is_bulk_tileable; + + // --------------------------------------------------------------------------- + // Block-scan used by the leader block to prefix-sum its merged histogram + // --------------------------------------------------------------------------- + // 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 + // --------------------------------------------------------------------------- + // 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. + // `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; + 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 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 + // 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]; + }; + // 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; + + struct chunk_desc + { + offset_t offset; + int count; + }; + + // 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 + { + 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}))}; + } + + // 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, 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 + 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 + { + _CCCL_ASSERT(local < count, "global_index: rank-local chunk index out of range"); + return first + local * stride; + } + }; + + // 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 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, + // ...`. + // * 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 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 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)"); + // 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)); + 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 + { + // 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}; + } + } + + // 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 + // 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. + 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"); + _CCCL_PRAGMA_NOUNROLL() + for (int local = tid; local < count; local += threads_per_block) + { + const key_t key = src[local]; + dst[local] = key; + apply(key); + } + } + + // 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: + // `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 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 + { + constexpr int tile = Unroll * threads_per_block; + 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 i = 0; i < Unroll; ++i) + { + regs[i] = keys[tile_base + i * threads_per_block + tid]; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Unroll; ++i) + { + const int local = tile_base + i * threads_per_block + tid; + apply(regs[i], local); + } + } + + // 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); + } + } + + 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) { + key_op(key); + }); + } + + // --------------------------------------------------------------------------- + // 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: 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() + { + _CCCL_PRAGMA_NOUNROLL() + for (int stage = tid; stage < PipelineStages; stage += threads_per_block) + { + ::cuda::ptx::mbarrier_init(&temp_storage.load_mbar[stage], 1u); + } + } + + // 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). + // 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) + { + _CCCL_ASSERT(stage >= 0 && stage < PipelineStages, "pipeline stage index out of range"); + // Only the block leader (see `is_block_leader`) drives the mbarrier, for a uniform branch and better mbarrier + // codegen. + if (!is_block_leader) + { + 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 + // 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( + ::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) + { + _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)) + { + } + load_phase ^= (::cuda::std::uint32_t{1} << stage); + } + + 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-streaming helpers (`init_overflow_stream`, `run_overflow_pass`, ...) read 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 eff_cluster_blocks; + bool is_idle_rank; + chunk_partition part; + unsigned 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 + // --------------------------------------------------------------------------- + _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; + 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 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) 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); + // 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 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). + num_segments_val_t segment_id{}; + segment_size_val_t segment_size{}; + 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}; + // 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( + 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_, + NumSegmentsParameterT num_segments_, + char* key_slots_, + 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_) + , 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_) + , num_segments(num_segments_) + , key_slots(key_slots_) + , block_tile_capacity(block_tile_capacity_) + {} + + // --------------------------------------------------------------------------- + // Main entry point + // --------------------------------------------------------------------------- + // 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() + { + NV_IF_TARGET(NV_PROVIDES_SM_90, (process_impl();)); + } + +private: + _CCCL_DEVICE _CCCL_FORCEINLINE void reset_hist() + { + // 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) + { + if (const int i = full_rounds * threads_per_block + tid; i < num_buckets) + { + temp_storage.hist[i] = 0; + } + } + } + + // --------------------------------------------------------------------------- + // Block-private histogram atomics (shared-space `red` via inline PTX) + // --------------------------------------------------------------------------- + // 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 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)); + } + + // 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); + 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 + // 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, 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)); + 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 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 `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 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) + { + 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 + // 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() + { + // 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]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < 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}; + } + + block_scan_t(temp_storage.scan_storage).ExclusiveSum(hist_vals, prefixes); + + // 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 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]); + 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`. + temp_storage.state.set_result( + static_cast<::cuda::std::uint32_t>(bucket), static_cast(new_len) == new_k); + } + } + } + + // --------------------------------------------------------------------------- + // 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 `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() + { + 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"); + } + + // 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 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 + { + 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)}; + } + + // 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)))); + } + + // 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_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) + { + const offset_t overflow_idx = stream_is_forward ? i : (layout.overflow_chunks - 1 - i); + // `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); + 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); + _CCCL_ASSERT(next_overflow_idx < layout.overflow_chunks, "overflow chunk index out of range"); + __syncthreads(); + _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) + { + if (++stage == stream_stages) + { + stage = 0; + } + } + else + { + stage = (stage == 0) ? stream_stages - 1 : stage - 1; + } + 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). `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_overflow_pass( + BlockApply&& block_apply, GenericApply&& generic_apply, OverlapWork&& overlap_work, Continue&& should_continue) + { + // 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(); + return; + } + + if constexpr (use_block_load_to_shared) + { + // 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. 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 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; + 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, stage, block_apply, should_continue)) + { + is_stopped = true; + break; + } + } + + if (!is_stopped) + { + // 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 `overlap_work`). + _CCCL_PRAGMA_NOUNROLL() + for (offset_t i = split; i < layout.overflow_chunks; ++i) + { + if (!consume_overflow_visit(i, stage, block_apply, should_continue)) + { + 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). + // `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; + wait_stage(drain_stage); + stream_inflight_mask &= ~(::cuda::std::uint32_t{1} << drain_stage); + } + stream_is_forward = !stream_is_forward; + } + else + { + // 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. + overlap_work(); + _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); + 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; + } + } + + // 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_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()), fold_key_op); + }, + [&](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) + { + fold_key_op( + layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]); + } + } + }, + ::cuda::std::forward(overlap_work), + [] { + return true; + }); + } + + // ------------------------------------------------------------------------- + // Per-direction implementation + // ------------------------------------------------------------------------- + // 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 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) + { + if (is_single_cta) + { + __syncthreads(); + } + else + { + cluster_barrier(); + } + } + + // 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(); + } + } + + // 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. + // + // 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) + { + // No peers: the front base is 0 (untouched since init) and the back base is just `num_selected`. + if (is_block_leader) + { + temp_storage.back_local_cnt = num_selected; + } + __syncthreads(); + 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) + { + 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; 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) + { + _CCCL_PRAGMA_NOUNROLL() + 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); + } + } + else + { + // 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 rank = cluster_rank + 1u + threadIdx.x; rank < layout.eff_cluster_blocks; + rank += threads_per_block) + { + 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); + // 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 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 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 + // 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_filter_state + { + using identify_op_t = detail::topk::identify_candidates_op_t; + + identify_op_t identify_op; + it_value_t block_keys_out; + 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 front_seg_base; + smem_keys_t resident_front; + int front_count; + bool is_select_all_cand_cta; + terminal_region terminal; + // Mutable per-invocation tie-break scan cursor. + offset_t running; + bool is_tie_active; + }; + + // --------------------------------------------------------------------------- + // Final-filter placement (shared) + deterministic filter + // --------------------------------------------------------------------------- + // 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 + // 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) + { + _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)]; + } + } + + // 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 + + // 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 + { + 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 + // 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* const counter = is_cand ? &temp_storage.back_local_cnt : &temp_storage.front_local_cnt; + 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; + 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 + // before each streaming bulk copy -- never per tile. + template + _CCCL_DEVICE _CCCL_FORCEINLINE bool final_filter_should_stop(const det_filter_state& state) + { + // 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 >= 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 >= back_end); + return is_front_done && is_back_done; + } + + // 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( + State& state, + const key_t (&keys)[ItemsPerThread], + const offset_t (&flags)[ItemsPerThread], + offset_t seg_base, + int count, + int tile_base, + bool do_arrival) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + const bool is_cand = flags[i] == flag_candidate; + if (flags[i] == flag_none || (is_cand && !do_arrival)) + { + continue; + } + 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); + } + } + + // 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( + State& state, + const key_t (&keys)[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(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() + 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(cand, excl, tile_total); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + if (cand[i] != offset_t{0}) + { + const offset_t global_rank = base + excl[i]; + if (global_rank < static_cast(state.num_back)) + { + // `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]; + final_filter_write_value(out, seg_idx); + } + } + } + return tile_total; + } + + // 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, + [[maybe_unused]] const key_t* smem_src, + offset_t seg_base, + int count, + [[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. + _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 tile_size = threads_per_block * ItemsPerThread; + int tile_base = 0; + + if constexpr (Deterministic && Blocked) + { + // 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. + _CCCL_PRAGMA_NOUNROLL() + for (; state.is_tie_active && tile_base < count; tile_base += tile_size) + { + 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 + { + // 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)) + { + // 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. Other paths write disjoint slots and need no explicit barrier. + __syncthreads(); + } + } + + // 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]; + 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 + { + // 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]; + offset_t flags[ItemsPerThread]; + classify_tile( + state, smem_src, seg_base, count, tile_base, keys, flags); + if constexpr (Deterministic) + { + // 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); + } + } + } + } + + // 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& state) + { + if constexpr (use_block_load_to_shared) + { + // Whole contiguous resident span staged in SMEM. + process_tiles( + state, + state.resident_front.data(), + state.front_seg_base, + state.front_count, + state.terminal == terminal_region::resident); + } + 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; + 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( + state, ::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_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 + _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 + // 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 || 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"); + + 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) { + 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(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(state, 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(state); + }); + } + + // 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& 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>( + state, keys, seg_base, count, is_terminal); + } + } + + // 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& state) + { + process_edge( + 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` + // (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& 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.terminal == terminal_region::tail_edge); + } + + // 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& 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(state)) + { + region(); + } + }; + if constexpr (is_tie_reversed) + { + process_tail_edge(state); + } + else + { + process_head_edge(state); + } + step([&] { + process_resident(state); + }); + step([&] { + process_overflow(state); + }); + step([&] { + if constexpr (is_tie_reversed) + { + process_head_edge(state); + } + else + { + 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; 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) + { + if constexpr (use_block_load_to_shared) + { + _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 + // --------------------------------------------------------------------------- + // 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 `overlap_work` + // (overlapping the first reloads) and bails out of the overflow stream once its contribution is fully placed. + template + struct nondet_filter_state + { + IdentifyOp identify_op; + KeyOutIt block_keys_out; + out_offset_t num_back; + offset_t sel_prefix; + offset_t my_front; // this CTA's front region size (see `write_deterministic_topk`'s `my_front`) + smem_keys_t resident_keys; + }; + + // 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) + { + 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 = 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( + 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; + } + // 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) + { + 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, /*Blocked=*/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); + } + } + else + { + // 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); + process_tiles( + state, + ::cuda::ptr_rebind(key_slots + static_cast(local_chunk) * ChunkBytes), + chunk.offset, + chunk.count, + false); + } + } + } + + // 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_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) + { + 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_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 + // 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_tie_winners, sel_prefix, my_front, resident_keys}; + + 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); + 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); + }, + // 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)); + }); + } + + // 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. `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_tie_winners, IdentifyOp identify_op, KeyOutIt block_keys_out, smem_keys_t resident_keys) + { + // 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_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_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); + 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; + // 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 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"); + // 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 == 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` + // 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 == 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 + // part of this span. + const int front_count = static_cast(resident_keys.size()); + + // 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; + 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 + // 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(layout.part.global_index(layout.resident_base), layout.segment_size_off, layout.head_items).offset + : offset_t{0}; + + // 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_filter_state state{ + identify_op, + block_keys_out, + num_back, + my_front, + sel_prefix, + cand_prefix, + my_cand_count, + front_seg_base, + resident_keys, + front_count, + is_select_all_cand_cta, + terminal, + cand_prefix, + !is_select_no_cand_cta}; + // 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 + // `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. 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(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(); + auto add_first_pass = [&](const key_t& key) { + const int bucket = extract_op(key); + hist_inc(hist_smem32, bucket); + }; + + if constexpr (use_block_load_to_shared) + { + // 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. 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 && 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 = + (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; + + // 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; + // 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, 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, + "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; + // 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) + { + // 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; + int dst_slot = 0; + offset_t src_base = 0; + offset_t src_index = 0; + if (idx < prologue_stages) + { + const offset_t local_chunk = static_cast(idx); + 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; + } + else + { + 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), + "first-wave chunk must map back to its stage"); + dst_slot = stream_slot_base + stage; + src_base = layout.overflow_base; + src_index = first_wave_chunk; + 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)); + } + } + + // 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. 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; + [[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) + { + // 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_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 + { + 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) + { + streaming_stage = (streaming_stage + 1 == stream_stages) ? 0 : streaming_stage + 1; + } + else + { + 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; + 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) + { + 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`): `(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); + } + + // 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); + } + 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 + { + _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); + 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 + tid; + if (local < chunk.count) + { + const key_t key = + layout.block_keys_in[static_cast(chunk.offset + static_cast(local))]; + chunk_keys[local] = key; + add_first_pass(key); + } + } + } + + // 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, + "Dynamic shared memory block_tile is too small"); + } + + // 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() + { + layout.block_keys_in = d_key_segments_it[segment_id]; + layout.segment_size_off = static_cast(segment_size); + // 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.block_keys_base = nullptr; + layout.head_items = 0; + if constexpr (use_block_load_to_shared) + { + 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 + // 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 + // 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 -- + // 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. + layout.eff_cluster_blocks = cluster_blocks; + if constexpr (enable_runtime_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); + } + } + _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. + 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. + 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`). + 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`) + // 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 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. 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 + // 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; + + // 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`, + // true for both the strided and blocked partitions). + 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_off, layout.head_items); + tail_suffix_items = + static_cast(tail_chunk.count) + - ::cuda::round_down(static_cast(tail_chunk.count), static_cast(load_align_items)); + } + } + + // Streaming-slot reservation: clamped into `[1, full_slots]` when this rank overflows, else 0. + offset_t stream_slots = offset_t{0}; + if constexpr (use_block_load_to_shared) + { + if (needs_streaming) + { + 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::clamp(want_stream, offset_t{1}, full_slots); + } + } + 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(layout.my_resident_chunks <= layout.resident_slots_cap, + "Dynamic shared memory block_tile is too small"); + + 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)`. + 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. + layout.head_edge_len_items = (cluster_rank == 0u) ? static_cast(layout.head_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 + // 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, 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; + + 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; + // 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); + + // `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) + { + 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` at cluster scope (see `hist_inc`). + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); + 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); + } + }; + + // Resident-chunk histogram, deferred into the overflow stream so it overlaps the stream's in-flight first + // 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) + { + for_each_chunk_key( + resident_keys.data(), static_cast(resident_keys.size()), add_hist); + } + 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); + 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), + 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 turn-around chunks left resident by the + // previous pass. + 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` + // and `my_candidates`, the cross-CTA scan inputs) inclusive of its edge candidates. + 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 + // 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 + // 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 != layout.leader_rank && !layout.is_idle_rank) + { + const ::cuda::std::uint32_t hist_smem32 = hist_base32(); + // 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); + } + } + } + + // 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_arrive(is_single_cta); + + // 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 != layout.leader_rank && !layout.is_idle_rank) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < 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); + } + + 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 == layout.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 + // 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. + 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 + // 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 = 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"); + detail::topk::set_kth_key_bits(kth_key_bits_local, pass, bucket); + 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 != layout.leader_rank) + { + const int owner = bucket / buckets_per_thread; + if (tid == owner) + { + const int slot = bucket - owner * buckets_per_thread; + atomicAdd(&temp_storage.num_strictly_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, + // 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::is_early_stop(pass_result)) + { + break; + } +#endif + } + return last_pass; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void run() + { + 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); + + // `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"); + 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 + // `run_radix_passes`), so the full key is never broadcast. + key_prefix_t kth_key_bits_local = {}; + + // Size the persistent overflow-streaming window for this segment; a no-op when this rank has no overflow. + init_overflow_stream(); + + // 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 + // `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) + { + 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 + // 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. + if constexpr (use_block_load_to_shared) + { + init_load_barriers(); + } + __syncthreads(); + + 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, kth_key_bits_local); + + // ----------------------------------------------------------------------- + // Final filter pass: write the top-k keys for this segment. Strictly- + // 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_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) + // 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(num_tie_winners, identify_op, block_keys_out, resident_keys); + } + else + { + 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 + // 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 + // 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 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); + 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); + 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). + [[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{}; + } + }(); + + // 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 (; base + step <= num_items; base += step) + { + 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] = base + static_cast(j) * cluster_threads + cluster_tid; + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + keys[j] = keys_in_it[static_cast(idx[j])]; + } + if constexpr (!is_keys_only) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + vals[j] = vals_in_it[static_cast(idx[j])]; + } + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + keys_out_it[static_cast(idx[j])] = keys[j]; + } + if constexpr (!is_keys_only) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < copy_items; ++j) + { + vals_out_it[static_cast(idx[j])] = vals[j]; + } + } + } + + // 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(base + local); + keys_out_it[seg_idx] = keys_in_it[seg_idx]; + if constexpr (!is_keys_only) + { + vals_out_it[seg_idx] = vals_in_it[seg_idx]; + } + } + } + + // 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 hw_cluster_rank, unsigned 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() + { + // 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 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 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})) + { + return; + } + + 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). + 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; + } + + // Segments larger than the resident cluster_tile capacity are still handled -- the overflow chunks are re-streamed + // 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); + + // 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), 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(hw_cluster_rank, hw_cluster_blocks); + return; + } + + // 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)) + { + return; + } + + // 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 (is_block_leader) + { + temp_storage.state.len = static_cast(segment_size); + temp_storage.state.k = k; + temp_storage.state.result_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; + } + reset_hist(); + // 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` and the scan counters 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](auto direction_tag) { + constexpr detail::topk::select Direction = decltype(direction_tag)::value; + this->template run(); + }); + _CCCL_ASSERT(is_ok, "Unsupported select direction for cluster top-k"); + } +}; +} // namespace detail::batched_topk_cluster + +CUB_NAMESPACE_END diff --git a/cub/cub/detail/launcher/cuda_runtime.cuh b/cub/cub/detail/launcher/cuda_runtime.cuh index ab8e3a95cdb..2cf8dbb374b 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,34 @@ 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)); + } + + // 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/detail/segmented_params.cuh b/cub/cub/detail/segmented_params.cuh index 0ece8092d0c..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,14 +107,40 @@ 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 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 54261ac9590..7e539b7553b 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,12 +32,34 @@ #include #include #include +#include #include #include +#include +#include #include #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 @@ -82,7 +105,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 +144,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,66 +157,146 @@ 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. // --------------------------------------------------------------------------- - 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::__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, + ::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` + // 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( + !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(), - policy_selector_t{}); + 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 @@ -229,7 +335,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: //! @@ -241,9 +350,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`` -//! must also carry a small compile-time upper bound (a ``constant`` or ``cuda::args::bounds()``), and tight -//! bounds on every parameter are encouraged. +//! 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 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++ //! @@ -272,17 +388,27 @@ 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. -//! - **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). +//! - **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. 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`` (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 +//! 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 +429,25 @@ 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, 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 +//! 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 +514,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 small compile-time maximum. + //! 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). //! @@ -383,12 +526,15 @@ 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 - //! **[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). 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). //! @@ -476,12 +626,15 @@ 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 - //! **[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). 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). //! @@ -575,12 +732,15 @@ 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 - //! **[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). 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). //! @@ -666,12 +830,15 @@ 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 - //! **[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). 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). //! @@ -782,12 +953,15 @@ 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 - //! **[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). 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). //! @@ -885,12 +1063,15 @@ 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 - //! **[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). 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). //! @@ -993,12 +1178,15 @@ 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 - //! **[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). 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). //! @@ -1096,12 +1288,15 @@ 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 - //! **[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 +35,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 @@ -94,48 +108,92 @@ 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)); } }; // ----------------------------------------------------------------------------- -// Segmented Top-K Dispatch +// Dispatch (both backends behind one kernel symbol) // ----------------------------------------------------------------------------- +// 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); +} -//! @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 batched_topk_policy_selector -#endif // _CCCL_HAS_CONCEPTS() -CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( +// 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`. +#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 \ + 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, \ + policy.min_blocks_per_sm, \ + 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, @@ -143,24 +201,113 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( ValueInputItItT d_value_segments_it, ValueOutputItItT d_value_segments_out_it, SegmentSizeParameterT segment_sizes, - KParameterT k, - SelectDirectionT select_direction, + KParameterT k_param, + SelectDirectionParameterT select_directions, NumSegmentsParameterT num_segments, - [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream = nullptr, - [[maybe_unused]] PolicySelector policy_selector = {}) + cudaStream_t stream) { - using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; + 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 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"); - // Wrap the raw enum into the internal discrete param type - auto select_directions = wrap_select_direction(select_direction); - using SelectDirectionParameterT = decltype(select_directions); + // 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; + 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; + } - // 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< + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + // `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 invoking this launch 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; + 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, - SegmentSizeParameterT, KeyInputItItT, KeyOutputItItT, ValueInputItItT, @@ -169,187 +316,725 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( 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))) + 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; } - } - } - // 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))) + // 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 (`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))) + { + 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 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 + // 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; + }; + + // 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)); + 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_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)); + + int cluster_blocks = 0; + int dynamic_smem_sel = 0; + + 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 + // 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 + { + // 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. + 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) + { + 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 > 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 + // 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, + // 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 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 + // 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"); + + // 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())) + { + 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; } - if (d_temp_storage == nullptr) + 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) { - return cudaSuccess; + // 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; - // 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."); + // 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 (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))) + 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) { - return error; + 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; + } } } - 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])))) + + void* allocations[allocations_array_size] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) { 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))) + + if (d_temp_storage == nullptr) { - return error; + return cudaSuccess; + } + + // `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 invoking this launch arm. + + 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 constexpr (!only_small_segments) +#if _CCCL_CUDA_COMPILATION() && !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) +// 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`). +template +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL bool any_target_cc_unsupported() +{ + bool any = false; + for (const auto cc : ::cuda::__target_compute_capabilities()) { - // 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 + any = any || (PolicySelector{}(cc).backend == topk_backend::unsupported); } - return CubDebug(detail::DebugSyncStream(stream)); + return any; } -// 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 dispatch_with_env( +#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 +// `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, - SelectDirectionParameterT select_directions, + SelectDirectionT select_direction, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, - EnvT env = {}) + [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream) { - using default_policy_selector = - 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( + // 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)."); + + // 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). + const 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 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 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) + + // 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 (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{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + // Empty batch = no work to launch: no segments, or a non-positive tightest max segment size (every segment empty, + // e.g. a uniform negative size clamped to 0). `== 0` not `<= 0` for `num_segments`, as a negative count is out of + // contract. Consulted only on the launch (`d_temp_storage != nullptr`) of a *supported* arm below: the query pass + // falls through to size `temp_storage_bytes`, and the unsupported arm ignores it so an unavailable request still + // fails with cudaErrorNotSupported rather than being masked into success. + const auto empty_batch_no_launch = [&] { + return d_temp_storage != nullptr + && (detail::params::get_param(num_segments, 0) == 0 || runtime_max_segment_size(segment_sizes) <= 0); + }; + + 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) + { + if (empty_batch_no_launch()) + { + return cudaSuccess; + } + return launch_baseline_arm( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -360,10 +1045,40 @@ template ( + 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 diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 9c66be1cba3..4c38acc557b 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -17,21 +17,29 @@ #endif // no system header #include +#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 +48,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 +85,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}; @@ -100,8 +131,111 @@ public: }; // ----------------------------------------------------------------------------- -// Global Kernel Entry Point +// 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 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 +{ + 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; + } +} + +// Third `__launch_bounds__` argument (`maxBlocksPerCluster`): the cluster policy's `max_blocks_per_cluster` cap. The +// host arm launches a dynamic cluster width, so this is the only compile-time width hint `ptxas` sees, and +// `launch_cluster_arm` clamps the launch to `<= max_blocks_per_cluster`. `0` disables the cap. +_CCCL_EXEC_CHECK_DISABLE +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_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 -#if _CCCL_HAS_CONCEPTS() - requires batched_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) +_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(); + + 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_ELSE_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();), + // Cluster-policy kernels are only ever launched on SM90+, so the sub-SM90 device pass is unreachable at runtime. + (_CCCL_UNREACHABLE();)); + } + else + { + // topk_backend::unsupported: the host arm returns cudaErrorNotSupported before launching, so this never runs. + return; + } +} + +#ifdef CUB_RDC_ENABLED +// 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 +_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 = 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."); - - // Temporary storage allocation - __shared__ typename agent_t::TempStorage temp_storage; - - // Instantiate agent - 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); - - // Process segments - 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 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 48f976514bd..2e7c54ac1bb 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 friend constexpr 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 friend constexpr 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 friend constexpr 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 friend constexpr bool operator==(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr 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 friend constexpr bool operator!=(const batched_topk_policy& lhs, const batched_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr 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,448 @@ 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) { - [[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}}; + ::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 + -> 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() + +//! 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 + // 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. + + // 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 friend constexpr 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 + && 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 friend constexpr 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 + << ", .max_blocks_per_cluster = " << p.max_blocks_per_cluster + << ", .max_chunk_slots_per_block = " << p.max_chunk_slots_per_block << " }"; + } +#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, + /*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 +// `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 friend constexpr 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 friend constexpr 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/util_device.cuh b/cub/cub/util_device.cuh index d25b0ee09dc..98df257fd93 100644 --- a/cub/cub/util_device.cuh +++ b/cub/cub/util_device.cuh @@ -524,6 +524,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; } @@ -626,6 +628,9 @@ 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; + //! @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/CMakeLists.txt b/cub/test/CMakeLists.txt index 865785835a2..3dfe9bd89ff 100644 --- a/cub/test/CMakeLists.txt +++ b/cub/test/CMakeLists.txt @@ -294,6 +294,25 @@ 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. Its two %PARAM% variants pin their own arch lists (independent of the + # preset) so the diagnostic fires regardless: `arch_0` a single pre-SM90 arch (89-virtual, the highest without cluster + # support), and `arch_1` a mixed list adding a cluster-capable 90-virtual to show the diagnostic fires when *any* + # targeted arch is unsupported. + if ("${test_target}" MATCHES "batched_topk_unsupported_arch_fail\\.arch_0$") + set_target_properties( + ${test_target} + PROPERTIES CUDA_ARCHITECTURES "89-virtual" + ) + elseif ( + "${test_target}" MATCHES "batched_topk_unsupported_arch_fail\\.arch_1$" + ) + set_target_properties( + ${test_target} + PROPERTIES CUDA_ARCHITECTURES "89-virtual;90-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_segmented_topk_cluster_layout.cu b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu new file mode 100644 index 00000000000..71c8fde1f69 --- /dev/null +++ b/cub/test/catch2_test_device_segmented_topk_cluster_layout.cu @@ -0,0 +1,103 @@ +// 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 + +#include + +namespace +{ +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 = ::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}; + 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 8e1d6ba29f3..b5842091574 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -1,20 +1,37 @@ // 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/device_batched_topk.cuh. Precedes CUB includes. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include +#include // topk_policy / make_{baseline,cluster}_policy (cluster-cap tuning test) #include +#include #include #include +#include #include +#include +#include +#include #include #include #include +#include +#include #include #include +#include +#include +#include +#include #include "catch2_test_device_topk_common.cuh" #include "catch2_test_launch_helper.h" @@ -30,7 +47,26 @@ struct is_minus_zero } }; +// 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 +{ + 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)); + } +}; + +// 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 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, @@ -48,8 +84,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( { 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::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, cuda::execution::output_ordering::unsorted)}; if constexpr (SelectDirection == cub::detail::topk::select::max) { @@ -63,123 +99,1040 @@ 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 SelectDirection, SelectDirection); +// %PARAM% TEST_LAUNCH lid 0:1:2 +DECLARE_TMPL_LAUNCH_WRAPPER( + dispatch_batched_topk_keys, + batched_topk_keys, + ESCAPE_LIST( + cub::detail::topk::select SelectDirection, + 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(SelectDirection, Determinism, TieBreak)); + +// Wrapper-test companion to expect_batched_topk_unsupported_and_skip: when the request's backend is unavailable in this +// build, dispatch it directly (host), verify the runtime cudaErrorNotSupported, and skip the correctness checks; +// otherwise return so the caller runs its normal batched_topk_keys<...> launch + checks. Pass the same trailing +// arguments (and Direction / Determinism / TieBreak) as that launch. +template +void skip_unless_batched_topk_keys_supported(cuda::std::int64_t static_max_segment_size, Args... args) +{ + if (batched_topk_backend_unavailable(static_max_segment_size)) + { + expect_batched_topk_unsupported_and_skip([&](void* d_temp_storage, cuda::std::size_t& temp_storage_bytes) { + return dispatch_batched_topk_keys(d_temp_storage, temp_storage_bytes, args...); + }); + } +} + +// Total segment size +using max_segment_size_list = c2h::enum_type_list; + +// Segment size: static, uniform +using max_num_k_list = c2h::enum_type_list; + +// %PARAM% TEST_TYPES types 0:1:2 + +#if TEST_TYPES == 0 +using key_types = + c2h::type_list; +// clang-format on +#elif TEST_TYPES == 1 +using key_types = c2h::type_list; +#elif TEST_TYPES == 2 +using key_types = + c2h::type_list; +// clang-format on +#endif + +// Selection direction is a compile-time option; cover both as a static test axis. +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, + max_segment_size_list, + max_num_k_list, + 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>; + + // Statically constrained maximum segment size and k + 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; + + // Selection direction comes from the compile-time test axis. + constexpr auto direction = c2h::get<3, TestType>::value; + + // Generate segment size + constexpr segment_size_t min_segment_size = 1; + constexpr auto max_segment_size = static_max_segment_size; + const segment_size_t segment_size = GENERATE_COPY(values({min_segment_size, segment_size_t{3}, max_segment_size}), + take(4, random(min_segment_size, max_segment_size))); + const segment_size_t max_k = (cuda::std::min) (static_max_k, segment_size); + + // Skip invalid combinations + if (segment_size > max_segment_size) + { + SKIP("The given segment size may not exceed the maximum segment size, we statically constrained the algorithm on."); + } + + // Set the k value + const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k}), take(3, random(segment_size_t{1}, max_k))); + + // Generate number of segments + const segment_index_t num_segments = GENERATE_COPY( + values({segment_index_t{1}, segment_index_t{42}}), take(4, random(segment_index_t{1}, segment_index_t{1000}))); + + // Capture test parameters + CAPTURE(c2h::type_name(), + c2h::type_name(), + c2h::type_name(), + static_max_segment_size, + static_max_k, + segment_size, + k, + num_segments, + direction); + + // Prepare input & output + 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); + + // Copy input for verification + 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}); + // Prepare expected results + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + + // Since the results of top-k are unordered, sort output segments before comparison. + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + 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}); + + 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); +} + +// 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}); + + 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; + // 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); + + // Baseline-coverable segment sizes (bounded to 127); only a deterministic / tie-break requirement routes to the SM90+ + // cluster backend. + skip_unless_batched_topk_keys_supported( + /*static_max_segment_size=*/127, + 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}); + 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}); + + // 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); +} + +// 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 +// `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}); + + 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 +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. The tune override forces the +// SM90+ cluster backend, so where no SM90+ target can serve it this verifies the runtime cudaErrorNotSupported and +// skips the correctness checks instead. +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) +{ + const auto dispatch = [&](void* d_temp, cuda::std::size_t& temp_bytes) { + 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); + } + }; + if (batched_topk_cluster_backend_unavailable(/*needs_cluster=*/true)) + { + expect_batched_topk_unsupported_and_skip(dispatch); + } + cuda::std::size_t temp_bytes = 0; + REQUIRE(cudaSuccess == dispatch(nullptr, temp_bytes)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == dispatch(thrust::raw_pointer_cast(temp_storage.data()), temp_bytes)); + 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; + + // 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; + + 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; + + // 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; + + 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]", + det_tie_combos) +{ + using key_t = float; + 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; + // 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; + + 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); + + 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()); + + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_keys_supported( + static_max_segment_size, d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + batched_topk_keys(d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + + 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); +} + +// 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 + 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 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()); + + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_keys_supported( + static_max_segment_size, d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + batched_topk_keys(d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + + 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); +} + +// 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) +{ + using key_t = float; + 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 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 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}); -// Total segment size -using max_segment_size_list = c2h::enum_type_list; + CAPTURE(pad, static_max_segment_size, static_max_k, k, num_segments, num_items, direction); -// Segment size: static, uniform -using max_num_k_list = c2h::enum_type_list; + 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(); -// %PARAM% TEST_TYPES types 0:1:2 + 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); -#if TEST_TYPES == 0 -using key_types = - c2h::type_list; -// clang-format on -#elif TEST_TYPES == 1 -using key_types = c2h::type_list; -#elif TEST_TYPES == 2 -using key_types = - c2h::type_list; -// clang-format on -#endif + 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()); -// Selection direction is a compile-time option; cover both as a static test axis. -using select_direction_list = - c2h::enum_type_list; + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::copy(keys_in_buffer.cbegin() + pad, keys_in_buffer.cend(), expected_keys.begin()); -C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", - "[keys][segmented][topk][device]", - key_types, - max_segment_size_list, - max_num_k_list, - select_direction_list) + const auto seg_arg = + cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_keys_supported( + static_max_segment_size, d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + batched_topk_keys(d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + + 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) { + using key_t = float; using segment_size_t = cuda::std::int64_t; using segment_index_t = cuda::std::int64_t; - using key_t = c2h::get<0, TestType>; - - // Statically constrained maximum segment size and k - 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; + 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; - // Selection direction comes from the compile-time test axis. - constexpr auto direction = c2h::get<3, 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 = 2; - // Generate segment size - constexpr segment_size_t min_segment_size = 1; - constexpr auto max_segment_size = static_max_segment_size; - const segment_size_t segment_size = GENERATE_COPY(values({min_segment_size, segment_size_t{3}, max_segment_size}), - take(4, random(min_segment_size, max_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; - // Skip invalid combinations - if (segment_size > max_segment_size) - { - SKIP("The given segment size may not exceed the maximum segment size, we statically constrained the algorithm on."); - } - - // Set the k value - const segment_size_t k = GENERATE_COPY(values({segment_size_t{1}, max_k}), take(3, random(segment_size_t{1}, max_k))); - - // Generate number of segments - const segment_index_t num_segments = GENERATE_COPY( - values({segment_index_t{1}, segment_index_t{42}}), take(4, random(segment_index_t{1}, segment_index_t{1000}))); + CAPTURE(static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); - // Capture test parameters - CAPTURE(c2h::type_name(), - c2h::type_name(), - c2h::type_name(), - static_max_segment_size, - static_max_k, - segment_size, - k, - num_segments, - direction); + auto d_keys_in = cuda::make_transform_iterator( + cuda::make_counting_iterator(segment_index_t{0}), counting_segment_keys_op{segment_size}); - // Prepare input & output - 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); - // Copy input for verification - c2h::device_vector expected_keys(keys_in_buffer); + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_keys_supported( + static_max_segment_size, d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + batched_topk_keys(d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + + c2h::device_vector expected_keys(num_items, thrust::no_init); + thrust::sequence(expected_keys.begin(), expected_keys.end()); - // Run the top-k algorithm - 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}); - // Prepare expected results fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); compact_sorted_keys_to_topk(expected_keys, segment_size, k); - // Since the results of top-k are unordered, sort output segments 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 small variable-size segments", "[keys][segmented][topk][device]", @@ -279,6 +1232,91 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment REQUIRE(expected_keys == keys_out_buffer); } +#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; + + 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; + + 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(); + + 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); + + // 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()); + 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()); + + 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}, + 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); + + 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 && TEST_LAUNCH == 0 + C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per-segment k", "[keys][segmented][topk][device]", key_types, @@ -473,6 +1511,71 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with variable-size segments and REQUIRE(expected_keys == keys_out_buffer); } +#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. 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 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][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; + 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); + + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_keys_supported( + static_max_segment_size, d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + batched_topk_keys(d_keys_in, d_keys_out, seg_arg, k_arg, ns_arg); + + 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 // Regression test: top-k must preserve -0.0f in the output (not normalize to +0.0f). C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys preserve -0.0f in output", @@ -558,3 +1661,341 @@ 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)}; + + 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); + + // 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)}; + + const auto dispatch = [&](void* d_temp_storage, cuda::std::size_t& temp_storage_bytes) { + return cub::DeviceBatchedTopK::MaxKeys( + d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + }; + + // The exposed static max (65535) exceeds the baseline backend's per-tile coverage, so this config routes to the SM90+ + // cluster backend (unavailable pre-SM90). + if (batched_topk_backend_unavailable( + /*static_max_segment_size=*/cuda::args::__traits::highest)) + { + expect_batched_topk_unsupported_and_skip(dispatch); + } + + cuda::std::size_t temp_storage_bytes = 0; + auto error = dispatch(nullptr, temp_storage_bytes); + REQUIRE(error == cudaSuccess); + + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + error = dispatch(thrust::raw_pointer_cast(temp_storage.data()), temp_storage_bytes); + 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. A small size +// and no determinism requirement route this to the baseline backend (the deterministic-requirement counterpart, which +// routes to the cluster backend, is below). +C2H_TEST("DeviceBatchedTopK::MaxKeys clamps a negative segment size to an empty segment (no determinism requirement)", + "[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)}; + + 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); + + // 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}); +} + +// Deterministic-requirement counterpart of the clamp test above: a gpu_to_gpu requirement routes to the SM90+ cluster +// backend even for small segments, pinning that backend's own empty-segment early-out on a clamped negative size. +C2H_TEST("DeviceBatchedTopK::MaxKeys clamps a negative segment size to an empty segment (deterministic requirement)", + "[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; + + // 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>()}; + + skip_unless_batched_topk_keys_supported( + static_max_segment_size, + d_keys_in, + d_keys_out, + segment_sizes, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}); + 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. With a deterministic requirement (which +// routes to the cluster backend), this must be recognized on the host from the non-positive runtime maximum segment +// size and skipped without launching, leaving the output untouched -- rather than casting the negative maximum to +// unsigned and sizing an enormous launch. Where that backend is unavailable the request fails with +// cudaErrorNotSupported even though it is a no-op (the wrapper verifies that, then skips). +C2H_TEST("DeviceBatchedTopK::MaxKeys treats a uniform negative segment size as no work (deterministic requirement)", + "[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; + + 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); + + skip_unless_batched_topk_keys_supported( + static_max_segment_size, + 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}}); + 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)); +} + +// 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. A small size and no determinism requirement route this to the baseline backend; the +// deterministic-requirement counterpart is below. +C2H_TEST("DeviceBatchedTopK::MaxKeys treats zero segments as no work (no determinism requirement)", + "[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)); +} + +// Deterministic-requirement counterpart: a gpu_to_gpu requirement routes to the SM90+ cluster backend. A positive size +// bound means the `max_seg_size <= 0` guard does not apply, so the launch must be elided from `num_segments == 0` +// alone -- by the dispatch's empty-batch guard, before the cluster arm launches. Where that backend is unavailable the +// request fails with cudaErrorNotSupported even though it is a no-op (the wrapper verifies that, then skips). +C2H_TEST("DeviceBatchedTopK::MaxKeys treats zero segments as no work (deterministic requirement)", + "[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; + + 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); + + skip_unless_batched_topk_keys_supported( + static_max_segment_size, + 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}}); + 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 17fc28848a7..2502f986b45 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -1,19 +1,35 @@ // 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/device_batched_topk.cuh. Precedes CUB includes. +#define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT + #include "insert_nested_NVTX_range_guard.h" #include +#include // topk_policy / make_baseline_policy (cross-tuning test) #include #include #include #include +#include #include #include #include +#include +#include #include +#include +#include + +#include +#include +#include +#include #include "catch2_test_device_topk_common.cuh" #include "catch2_test_launch_helper.h" @@ -51,40 +67,46 @@ struct flag_intra_segment_duplicates template flag_intra_segment_duplicates(ItemItT, SegIdItT) -> flag_intra_segment_duplicates; +// 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 CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_keys_in, - KeyOutputItItT d_keys_out, - ValueInputItItT d_values_in, - ValueOutputItItT d_values_out, - SegmentSizeParamT segment_sizes, - KParamT k, + cuda::std::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::not_guaranteed, - cuda::execution::tie_break::unspecified, + cuda::execution::require(cuda::execution::determinism::__determinism_holder_t{}, + cuda::execution::tie_break::__tie_break_holder_t{}, cuda::execution::output_ordering::unsorted)}; if constexpr (SelectDirection == cub::detail::topk::select::max) { return cub::DeviceBatchedTopK::MaxPairs( d_temp_storage, temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, segment_sizes, k, num_segments, @@ -95,10 +117,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( return cub::DeviceBatchedTopK::MinPairs( d_temp_storage, temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, segment_sizes, k, num_segments, @@ -108,7 +130,33 @@ 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 SelectDirection, SelectDirection); + dispatch_batched_topk_pairs, + batched_topk_pairs, + ESCAPE_LIST( + cub::detail::topk::select SelectDirection, + 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(SelectDirection, Determinism, TieBreak)); + +// Wrapper-test companion to expect_batched_topk_unsupported_and_skip: when the request's backend is unavailable in this +// build, dispatch it directly (host), verify the runtime cudaErrorNotSupported, and skip the correctness checks; +// otherwise return so the caller runs its normal batched_topk_pairs<...> launch + checks. Pass the same trailing +// arguments (and Direction / Determinism / TieBreak) as that launch. +template +void skip_unless_batched_topk_pairs_supported(cuda::std::int64_t static_max_segment_size, Args... args) +{ + if (batched_topk_backend_unavailable(static_max_segment_size)) + { + expect_batched_topk_unsupported_and_skip([&](void* d_temp_storage, cuda::std::size_t& temp_storage_bytes) { + return dispatch_batched_topk_pairs(d_temp_storage, temp_storage_bytes, args...); + }); + } +} // Total segment size using max_segment_size_list = c2h::enum_type_list; @@ -147,6 +195,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, @@ -206,8 +273,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; } @@ -293,14 +360,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 @@ -313,6 +376,1558 @@ 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}); + + 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; + // 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); + + // Baseline-coverable segment sizes (bounded to 127); only a deterministic / tie-break requirement routes to the SM90+ + // cluster backend. + const auto seg_arg = cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + /*static_max_segment_size=*/127, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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 +// 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, + cuda::std::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)); + +// 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. The path-pinning cluster tests below hand it an env carrying a +// `cluster_tuning_selector` tune override, which forces the SM90+ cluster backend; where no SM90+ target can serve it +// this verifies the runtime cudaErrorNotSupported and skips the correctness checks instead. +template +void run_cluster_topk_pairs( + KeyInItT d_keys_in, + KeyOutItT d_keys_out, + ValueInItT d_values_in, + ValueOutItT d_values_out, + SegSizesT seg_sizes, + KParamT k_param, + NumSegT num_seg, + EnvT env) +{ + const auto dispatch = [&](void* d_temp, cuda::std::size_t& temp_bytes) { + 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, num_seg, 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, num_seg, env); + } + }; + if (batched_topk_cluster_backend_unavailable(/*needs_cluster=*/true)) + { + expect_batched_topk_unsupported_and_skip(dispatch); + } + cuda::std::size_t temp_bytes = 0; + REQUIRE(cudaSuccess == dispatch(nullptr, temp_bytes)); + c2h::device_vector temp_storage(temp_bytes, thrust::no_init); + REQUIRE(cudaSuccess == dispatch(thrust::raw_pointer_cast(temp_storage.data()), temp_bytes)); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); +} + +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); + + // Oversize segments always route to the SM90+ cluster backend. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs(d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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); +} + +// 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. 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_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) +{ + 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); + // 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); + + // 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()); + + // Oversize segments always route to the SM90+ cluster backend. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs(d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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); +} + +// 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 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) +{ + 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; + + 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()); + + // Oversize segments always route to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + 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; + +// 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) +{ + 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); + + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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); +} + +# 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + // 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 +// 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; + 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 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); + + // 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); + + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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; + // 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{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); + + // Deterministic and oversize: this configuration always routes to the SM90+ cluster backend. + const auto seg_arg = + cuda::args::immediate{segment_size, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + // 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); +} + +# 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, + determinism_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 determinism = c2h::get<1, TestType>::value; + constexpr auto tie_break = cuda::execution::tie_break::__tie_break_t::__unspecified; + + // 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; + 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; + + 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); + + // 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). `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 = + 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); + + 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}); + }; + + // Deterministic request always routes to the SM90+ cluster backend (the tune override forces it at this tiny size). + { + auto d_keys_out_a = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out_a.data())), k); + auto d_values_out_a = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out_a.data())), k); + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, + d_keys_in, + d_keys_out_a, + d_values_in, + d_values_out_a, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + cuda::args::immediate{k, cuda::args::bounds()}, + cuda::args::immediate{num_segments}); + } + + 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(repro_config_b{}, keys_out_b, values_out_b); + } + else + { + 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. + 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); + std::sort(a_begin, a_begin + k); + std::sort(b_begin, b_begin + k); + } + 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + 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); +} + +// 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). 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) +{ + 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + 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); +} + +// 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + 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; + 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()}; + + run_cluster_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_sizes, k_param, cuda::args::immediate{num_segments}, env); + + 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 +// 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; + // 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); + + // Oversize bound always routes to the SM90+ cluster backend, regardless of the determinism / tie-break requirement. + const auto seg_arg = + cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}; + const auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + const auto ns_arg = cuda::args::immediate{num_segments}; + skip_unless_batched_topk_pairs_supported( + static_max_segment_size, d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + batched_topk_pairs( + d_keys_in, d_keys_out, d_values_in, d_values_out, seg_arg, k_arg, ns_arg); + + 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", "[pairs][segmented][topk][device]", key_types, diff --git a/cub/test/catch2_test_device_topk_common.cuh b/cub/test/catch2_test_device_topk_common.cuh index 68552951ec9..a9abcdf4811 100644 --- a/cub/test/catch2_test_device_topk_common.cuh +++ b/cub/test/catch2_test_device_topk_common.cuh @@ -6,15 +6,147 @@ #include #include #include // topk::select::{min, max} +#include // make_baseline_policy +#include // cub::PtxVersion +#include #include +#include +#include #include +#include +#include #include #include #include +// Low-level predicate: true 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 +// degrades to a runtime cudaErrorNotSupported rather than a compile-time error; callers dispatch it, verify that error +// (via expect_batched_topk_unsupported_and_skip), and skip the result-correctness checks. `needs_cluster` must be true +// when the caller knows the configuration requires the cluster backend; prefer 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 must catch. +inline bool batched_topk_cluster_backend_unavailable(bool needs_cluster) +{ + if (!needs_cluster) + { + return false; + } + int ptx_version = 0; + REQUIRE(cudaSuccess == cub::PtxVersion(ptx_version)); + constexpr int cluster_min_ptx_version = 900; // SM 9.0 + return ptx_version < cluster_min_ptx_version; +} + +// True when the batched top-k configuration cannot run in the current build (see +// batched_topk_cluster_backend_unavailable). 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. 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, 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). +template +bool 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()); + return batched_topk_cluster_backend_unavailable(deterministic || oversize); +} + +// Runs the two-phase direct-API dispatch `dispatch(d_temp_storage, temp_storage_bytes)` (temp-size query, then launch) +// for a request whose backend is unavailable in this build, and verifies it degraded gracefully at runtime: the query +// still succeeds (returns a placeholder size), while the launch must return cudaErrorNotSupported +// (CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT defers the would-be compile-time diagnostic to runtime). Then skips the +// result-correctness checks, which need a device that can run the request. Callers gate this on +// batched_topk_backend_unavailable() (or batched_topk_cluster_backend_unavailable() when the request is already known +// to require the cluster backend); `dispatch` forwards to the direct-API entry point (host-side; the runtime error is +// backend-selection driven and independent of the launch variant under test). +template +void expect_batched_topk_unsupported_and_skip(DispatchFn&& dispatch) +{ + cuda::std::size_t temp_storage_bytes = 0; + REQUIRE(dispatch(nullptr, temp_storage_bytes) == cudaSuccess); + c2h::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + REQUIRE(dispatch(thrust::raw_pointer_cast(temp_storage.data()), temp_storage_bytes) == cudaErrorNotSupported); + SKIP("This batched top-k request routes to the SM90+ cluster backend, which the compute capability this dispatch " + "resolved to does not provide; the dispatch reported cudaErrorNotSupported at runtime " + "(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) as expected. Skipping the result-correctness checks."); +} + +// 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 @@ -314,9 +446,10 @@ c2h::device_vector compact_to_topk_batched( cuda::make_counting_iterator(0), get_output_size_op{d_offsets.cbegin(), k_it, static_cast(num_segments)}); - // 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); diff --git a/cub/test/test_device_batched_topk_requirements_fail.cu b/cub/test/test_device_batched_topk_requirements_fail.cu index c9a00e9d31d..cbf885988c3 100644 --- a/cub/test/test_device_batched_topk_requirements_fail.cu +++ b/cub/test/test_device_batched_topk_requirements_fail.cu @@ -1,7 +1,12 @@ // 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 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 @@ -12,14 +17,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() @@ -27,34 +42,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}; 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..2247c7bf712 --- /dev/null +++ b/cub/test/test_device_batched_topk_unsupported_arch_fail.cu @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// %PARAM% CUB_TEST_TOPK_UNSUPPORTED_ARCH arch 0:1 + +#include + +#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/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, the default (strict) mode must fail to compile rather than defer the failure to a runtime cudaErrorNotSupported. +// 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. +// +// The two %PARAM% variants compile this identical source against different architecture lists (pinned per-variant in +// CMakeLists.txt, independent of the preset): variant 0 pins a single pre-SM90 arch, variant 1 a mixed pre-/post-SM90 +// list. The latter documents the "fail if *any* target architecture is unsupported" contract -- the static_assert +// (guarded by any_target_cc_unsupported, which scans every targeted compute capability) fires because at least one +// target is unsupported, even though the co-pinned SM90+ target could serve the request. CUB_TEST_TOPK_UNSUPPORTED_ARCH +// only spawns the two variants; the architecture list itself is a CMake target property. + +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}; + 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) + { + std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; + } +} diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index a3117f9cea2..7e643c7d02b 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, 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})`` + -- 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 ---------------------- @@ -156,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: 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." 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...); }