From 3a2cfc87c99442f173e4a08efc6db2f64c32e78a Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 14 Aug 2026 08:01:23 +0000 Subject: [PATCH] test: xllm CUDA kernel verification suite for BI-V100 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_xllm_cuda_kernels.py — 7 test groups: 1. activation.cu: silu_and_mul via ixf_F, compare vs torch.nn.functional.silu 2. norm.cu: rms_norm + fused_add_rms_norm via ixf_F, compare vs PyTorch 3. rope.cu: rotary_embedding via ixf_F, verify rotation applied 4. moe_topk_softmax: corex .so, verify shapes + weights sum to 1 5. ix_moe_bridge: full 7-step fused MoE pipeline (topk→expand→gemm→act→gemm→combine) 6. ix_attn_bridge: load test (prefill_attention, decode_attention, linear) 7. ix_full_bridge: silu_and_mul + rms_norm through bridge .so Revert: undo unnecessary cccl_upstream sync (already up to date) Run on real machine: python3 qwen3_6_scripts/test_xllm_cuda_kernels.py --- .../dispatch/dispatch_adjacent_difference.cuh | 34 +- .../device/dispatch/dispatch_batch_memcpy.cuh | 4 +- .../device/dispatch/dispatch_batched_topk.cuh | 1244 +++-------------- .../device/dispatch/dispatch_copy_mdspan.cuh | 85 +- .../device/dispatch/dispatch_histogram.cuh | 12 +- .../dispatch/dispatch_reduce_by_key.cuh | 32 +- .../cub/cub/device/dispatch/dispatch_scan.cuh | 138 +- .../device/dispatch/dispatch_scan_by_key.cuh | 54 +- .../dispatch_segmented_radix_sort.cuh | 2 +- .../dispatch/dispatch_segmented_reduce.cuh | 32 +- .../dispatch/dispatch_segmented_scan.cuh | 22 +- .../device/dispatch/dispatch_select_if.cuh | 17 +- .../dispatch/dispatch_three_way_partition.cuh | 24 +- .../dispatch_transform_tile_config.cuh | 2 +- .../cub/cub/device/dispatch/tuning/common.cuh | 5 +- .../dispatch/tuning/tuning_batch_memcpy.cuh | 4 +- .../dispatch/tuning/tuning_batched_topk.cuh | 311 +---- .../device/dispatch/tuning/tuning_merge.cuh | 23 +- .../dispatch/tuning/tuning_reduce_by_key.cuh | 4 +- .../dispatch/tuning/tuning_rle_encode.cuh | 4 +- .../tuning/tuning_rle_non_trivial_runs.cuh | 4 +- .../device/dispatch/tuning/tuning_scan.cuh | 6 +- .../dispatch/tuning/tuning_scan_by_key.cuh | 4 +- .../dispatch/tuning/tuning_segmented_sort.cuh | 28 +- .../dispatch/tuning/tuning_select_if.cuh | 22 +- .../tuning/tuning_three_way_partition.cuh | 4 +- qwen3_6_scripts/test_xllm_cuda_kernels.py | 415 ++++++ 27 files changed, 911 insertions(+), 1625 deletions(-) create mode 100644 qwen3_6_scripts/test_xllm_cuda_kernels.py diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh index 150cab2c..69480e69 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_adjacent_difference.cuh @@ -246,15 +246,14 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceAdjacentDifference" error = CubDebug( THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( num_tiles, AdjacentDifferencePolicyT::BLOCK_THREADS, 0, stream) - .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel< - KernelPolicySelector, - InputIteratorT, - OutputIteratorT, - DifferenceOpT, - OffsetT, - InputT, - AliasOpt == MayAlias::Yes, - ReadOpt == ReadOption::Left>, + .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel < KernelPolicySelector, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + OffsetT, + InputT, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left >, d_input, first_tile_previous, d_output, @@ -439,15 +438,14 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( if (const auto error = CubDebug( THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_tiles, active_policy.threads_per_block, 0, stream) - .doit(DeviceAdjacentDifferenceDifferenceKernel< - policy_selector_t, - InputIteratorT, - OutputIteratorT, - DifferenceOpT, - offset_t, - input_t, - AliasOpt == MayAlias::Yes, - ReadOpt == ReadOption::Left>, + .doit(DeviceAdjacentDifferenceDifferenceKernel < policy_selector_t, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + offset_t, + input_t, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left >, d_input, first_tile_previous, d_output, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh index 642ba8a3..b86dd7a8 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_batch_memcpy.cuh @@ -152,8 +152,8 @@ __launch_bounds__(int(current_policy().lookback.large_buffer.thr { if (thread_offset < buffer_sizes[buffer_id]) { - const auto value = - read_item(input_buffer_it[buffer_id], thread_offset); + const auto value = read_item < MemcpyOpt == CopyAlg::Memcpy, AliasT, + InputBufferT > (input_buffer_it[buffer_id], thread_offset); write_item( output_buffer_it[buffer_id], thread_offset, value); } diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh index ce4b5ab3..f95538f0 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_batched_topk.cuh @@ -2,8 +2,8 @@ // Apache-2.0 WITH LLVM-exception //! @file -//! Internal device-wide dispatch for cub::DeviceBatchedTopK: selects between the baseline (worker-per-segment) and -//! cluster (SM 9.0+) backends and launches them through a single kernel symbol. +//! cub::DeviceTopK provides device-wide, parallel operations for finding the K largest (or smallest) items from +//! sequences of unordered data items residing within device-accessible memory. #pragma once @@ -17,17 +17,14 @@ # pragma system_header #endif // no system header -#include -#include #include -#include +#include #include #include #include #include #include #include -#include #include #include #include @@ -35,32 +32,16 @@ #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 @@ -72,7 +53,7 @@ namespace detail::batched_topk // The selection direction is compile-time only: callers pass `::cuda::args::constant`, which maps to a // value-less static_discrete_param. Because the direction is fixed at compile time and carries no runtime value, it // can never disagree with its only supported option, so dispatch can never silently degrade to a no-op. -template +template [[nodiscard]] _CCCL_HOST_DEVICE auto wrap_select_direction(::cuda::args::constant) { return params::static_discrete_param{}; @@ -103,7 +84,7 @@ template // per-segment tile counts that we exclusive-scan to obtain per-segment tile // offsets. // ----------------------------------------------------------------------------- -template +template struct segment_size_to_tile_count_op { SegmentSizeParameterT segment_sizes; @@ -112,824 +93,49 @@ struct segment_size_to_tile_count_op template _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr TotalNumItemsValueType operator()(SegmentIndexT segment_id) const { - return static_cast(::cuda::ceil_div( - params::__get_and_clamp_param_to_nonnegative(segment_sizes, segment_id), large_segment_agent_tile_size)); + return static_cast( + ::cuda::ceil_div(params::get_param(segment_sizes, segment_id), large_segment_agent_tile_size)); } }; // ----------------------------------------------------------------------------- -// Automatic backend selector +// Segmented Top-K Dispatch // ----------------------------------------------------------------------------- -// Stateless selector built purely from the compile-time request facts. It owns the entire backend decision, including -// computing `baseline_can_cover` from the concrete agent types -- the reason it lives here (where -// `baseline_can_cover_v` and the baseline agent are visible) rather than in the tuning header. -template -struct policy_selector_from_types -{ - // TODO(bgruber): to let the baseline policy vary per CC, move this coverage check into operator() and evaluate it for - // the passed CC. Only the check is hard: it instantiates the agent for sizeof(TempStorage), so it needs the CC as a - // compile-time constant, whereas operator()'s `cc` is a runtime parameter (building the policy itself is just the - // value make_baseline_policy(cc)). Recover the compile-time CC by folding over - // ::cuda::__target_compute_capabilities() (as detail::dispatch_to_cc_list does) and evaluate baseline_can_cover_v for - // the matching CC. That also removes the invariant below, since coverage and the returned baseline would then derive - // from the same cc. - // note: the baseline policy passed to baseline_can_cover_v must be the same as returned from operator(cc) below - static constexpr baseline_topk_policy baseline_policy = make_baseline_policy(); - - struct policy_getter_17 // TODO(bgruber): remove in C++20 and pass policy by value - { - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()() const -> topk_policy - { - return topk_policy{topk_algorithm::baseline, baseline_policy, {}}; - } - }; - - // Whether a one-worker-per-segment (default baseline) policy fits the static max segment size in shared memory; feeds - // the backend decision below. - static constexpr bool baseline_can_cover = baseline_can_cover_v< - policy_getter_17, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>; - - [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> topk_policy - { - constexpr bool deterministic = (Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) - || (TieBreak != ::cuda::execution::tie_break::__tie_break_t::__unspecified); - - topk_algorithm backend = topk_algorithm::unsupported; - if (deterministic || !baseline_can_cover) - { - // A deterministic result set / concrete tie-break preference, or a segment too large for the single-block - // baseline, is served only by the cluster backend (SM 9.0+); otherwise the request cannot run here. - backend = cluster_capable(cc) ? topk_algorithm::cluster : topk_algorithm::unsupported; - } - else - { - // Baseline can cover: use the cluster backend only where it is measured to win. 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. The threshold is applied on every cluster-capable architecture, not gated to a minimum CC. - const bool beneficial = StaticMaxSegSize >= cluster_beneficial_min_segment_size; - backend = (cluster_capable(cc) && beneficial) ? topk_algorithm::cluster : topk_algorithm::baseline; - } - return topk_policy{backend, baseline_policy, make_cluster_policy()}; - } -}; - -// ----------------------------------------------------------------------------- -// Dispatch (both backends behind one kernel symbol) -// ----------------------------------------------------------------------------- -// The dispatch is host-only: it launches the single kernel symbol (`device_batched_topk_kernel`, in -// kernel_batched_topk.cuh) via the CUDA runtime. The algorithm does not support device-side (CDP) launch. - -// Corrected form of `launcher_factory.max_dynamic_smem_size_for` (host path: `cub::MaxPotentialDynamicSmemBytes`), -// returning the usable dynamic budget as `opt-in - static footprint`. That facility currently subtracts the per-block -// reserved shared memory a second time even though `cudaDevAttrMaxSharedMemoryPerBlockOptin` already excludes it, -// under-reporting the budget by ~`reserved` (~1 KiB) -- enough to drop the cluster kernel's top table tier (see the -// TODO in MaxPotentialDynamicSmemBytes). TODO: once that facility is fixed, delete this and call -// `launcher_factory.max_dynamic_smem_size_for(...)` directly. -template -_CCCL_HOST_API cudaError_t max_dynamic_smem_size_for_fixed(int& max_dynamic_smem_bytes, KernelPtr kernel_ptr) -{ - max_dynamic_smem_bytes = -1; - 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; - } - cudaFuncAttributes kernel_attrs{}; - if (const auto error = CubDebug(cudaFuncGetAttributes(&kernel_attrs, kernel_ptr))) - { - return error; - } - const int static_smem_bytes = static_cast(kernel_attrs.sharedSizeBytes); - max_dynamic_smem_bytes = (max_smem_optin_bytes > static_smem_bytes) ? max_smem_optin_bytes - static_smem_bytes : 0; - return cudaSuccess; -} - -// Largest number of CTA blocks per cluster the kernel/architecture admits at `dynamic_smem_bytes` of dynamic SMEM. The -// config's cluster dimension is ignored by the query (placeholder here); the non-portable opt-in must already be set -// for it to report sizes beyond the portable ceiling. -template -_CCCL_HOST_API ::cuda::std::expected -probe_max_cluster_blocks(KernelPtr kernel_ptr, cudaStream_t stream, int threads_per_block, int dynamic_smem_bytes) -{ - ::cudaLaunchAttribute cluster_attr{}; - cluster_attr.id = ::cudaLaunchAttributeClusterDimension; - cluster_attr.val.clusterDim = {1, 1, 1}; - - ::cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(1); - cfg.blockDim = dim3(static_cast(threads_per_block)); - cfg.dynamicSmemBytes = static_cast<::cuda::std::size_t>(dynamic_smem_bytes); - cfg.stream = stream; - cfg.attrs = &cluster_attr; - cfg.numAttrs = 1; - - int cluster_blocks = 0; - if (const auto error = CubDebug( - ::cudaOccupancyMaxPotentialClusterSize(&cluster_blocks, reinterpret_cast(kernel_ptr), &cfg))) - { - return ::cuda::std::unexpected(error); - } - return cluster_blocks; -} - -// Device-wide count of `cluster_blocks`-CTA clusters that can be co-resident at `dynamic_smem_bytes` of dynamic SMEM -// (clusters per wave). `cudaOccupancyMaxActiveClusters` rejects a grid that is not a multiple of the cluster, so the -// grid is set to exactly one cluster; the returned capacity is independent of the actual grid size. -template -_CCCL_HOST_API ::cuda::std::expected probe_clusters_per_wave( - KernelPtr kernel_ptr, cudaStream_t stream, int threads_per_block, int cluster_blocks, int dynamic_smem_bytes) -{ - ::cudaLaunchAttribute cluster_attr{}; - cluster_attr.id = ::cudaLaunchAttributeClusterDimension; - cluster_attr.val.clusterDim = {static_cast(cluster_blocks), 1, 1}; - - ::cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(static_cast(cluster_blocks)); - cfg.blockDim = dim3(static_cast(threads_per_block)); - cfg.dynamicSmemBytes = static_cast<::cuda::std::size_t>(dynamic_smem_bytes); - cfg.stream = stream; - cfg.attrs = &cluster_attr; - cfg.numAttrs = 1; - - int clusters_per_wave = 0; - if (const auto error = - CubDebug(::cudaOccupancyMaxActiveClusters(&clusters_per_wave, reinterpret_cast(kernel_ptr), &cfg))) - { - return ::cuda::std::unexpected(error); - } - return clusters_per_wave; -} - -// The cluster backend's launch shape: CTA blocks per cluster and the dynamic-SMEM bytes to launch with. -struct cluster_launch_shape -{ - int cluster_blocks = 0; - int dynamic_smem_bytes = 0; -}; - -// Chooses the cluster launch shape for the statically-bounded max segment size. Probes occupancy through the CUDA -// runtime; the caller has already set the kernel's dynamic-SMEM opt-in to the maximum, so every probed config and the -// final launch run under one consistent opt-in. -template -_CCCL_HOST_API ::cuda::std::expected select_cluster_launch_shape( - ::cuda::std::uint64_t max_segment_size, - ::cuda::std::uint64_t num_segments, - int max_dynamic_smem_bytes, - cluster_topk_policy policy, - cudaStream_t stream, - KernelPtr kernel_ptr) -{ - using layout_t = LayoutT; - - const int threads_per_block = policy.threads_per_block; - - // Computed before any occupancy query so the single-CTA fast path below can skip one -- that driver query - // otherwise dominates the runtime of tiny launches. - const int max_block_resident_items = static_cast(layout_t::max_block_resident_items(max_dynamic_smem_bytes)); - if (max_block_resident_items <= 0) - { - // Not even one load-aligned chunk fits in the opt-in budget; the kernel cannot run. - return ::cuda::std::unexpected(cudaErrorInvalidValue); - } - - // Smallest cluster block count for full residency: at the largest SMEM each CTA holds `max_block_resident_items` - // items. 64-bit to match the launch-shape arithmetic below; the value is small (`max_segment_size <= 2^21`). - const auto min_blocks_per_segment = - ::cuda::ceil_div(max_segment_size, static_cast<::cuda::std::uint64_t>(max_block_resident_items)); - - int cluster_blocks = 0; - int dynamic_smem_bytes = 0; - - if (batched_topk_cluster::is_single_cta_eligible( - static_cast<::cuda::std::uint32_t>(max_segment_size), - static_cast<::cuda::std::uint32_t>(max_block_resident_items), - policy.single_block_max_seg_size)) - { - // 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. One CTA at in-budget SMEM is always launchable, - // so the occupancy probe is skipped. Larger fully-resident segments fall through to the wave-aware search below. - cluster_blocks = 1; - dynamic_smem_bytes = layout_t::min_smem_bytes_from_num_items(max_segment_size); - } - else - { - // Hardware cluster ceiling (max blocks per cluster), queried at runtime (not hardcoded) so a future device with - // larger non-portable clusters is not capped. Probed at zero dynamic SMEM for the arch/kernel ceiling alone; each - // candidate is re-validated against its own SMEM below. - const auto hw_cluster_ceiling = - probe_max_cluster_blocks(kernel_ptr, stream, threads_per_block, /*dynamic_smem_bytes=*/0); - if (!hw_cluster_ceiling) - { - return ::cuda::std::unexpected(hw_cluster_ceiling.error()); - } - if (*hw_cluster_ceiling <= 0) - { - return ::cuda::std::unexpected(cudaErrorInvalidValue); - } - // `max_blocks_per_cluster == 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 eff_max_blocks_per_cluster = - (policy.max_blocks_per_cluster == 0) - ? *hw_cluster_ceiling - : (::cuda::std::min) (policy.max_blocks_per_cluster, *hw_cluster_ceiling); - - // Wave-aware selection: the free variable is the cluster block count (one cluster per segment), paired with the - // smallest SMEM that keeps the segment fully resident (fewer blocks = more SMEM/fewer clusters-per-wave, more = - // the reverse). Pick the count minimizing waves, ties toward the largest (smallest SMEM, most L1 -- the profiled - // fast configs). Enumerated analytically, so a register-limited occupancy cannot collapse the candidate set. - if (min_blocks_per_segment <= static_cast<::cuda::std::uint64_t>(eff_max_blocks_per_cluster)) - { - // Full residency achievable: `max_segment_size <= min_blocks_per_segment * max_block_resident_items` and - // `min_blocks_per_segment <= eff_max_blocks_per_cluster`, so every per-CTA capacity below fits `int`. - - // 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 the segment's chunk count; a larger knob shrinks it. - const auto desired_cluster_blocks = ::cuda::narrow(batched_topk_cluster::compute_num_logical_cluster_blocks( - static_cast<::cuda::std::uint32_t>(layout_t::num_chunks_from_num_items(max_segment_size)), - policy.min_chunks_per_block, - ::cuda::narrow<::cuda::std::uint32_t>(eff_max_blocks_per_cluster))); - - // Scan `[min_candidate_blocks, max_candidate_blocks]` for the min-waves block count, tie-breaking largest. - // `max_candidate_blocks == max(desired_cluster_blocks, min(min_candidate_blocks, eff_max_blocks_per_cluster))`: - // the segment-needed count `desired_cluster_blocks` (<= `eff_max_blocks_per_cluster`, capped in - // `compute_num_logical_cluster_blocks`), floored at `min_candidate_blocks`. The `clamp` operands are ordered so - // `lo <= hi` holds even when `eff_max_blocks_per_cluster == 1` forces `min_candidate_blocks (== 2) > eff_max`; - // there `max_candidate_blocks == 1` empties the scan and the single-CTA fallback below runs (that edge is a - // one-CTA-resident segment with the single-CTA path disabled, so `min_blocks_per_segment == 1`). - const auto min_candidate_blocks = (::cuda::std::max) (2, static_cast(min_blocks_per_segment)); - const auto max_candidate_blocks = - ::cuda::std::clamp(min_candidate_blocks, desired_cluster_blocks, eff_max_blocks_per_cluster); - auto best_waves = (::cuda::std::numeric_limits<::cuda::std::uint64_t>::max)(); - for (int candidate_blocks = min_candidate_blocks; candidate_blocks <= max_candidate_blocks; ++candidate_blocks) - { - const auto num_block_items = ::cuda::ceil_div(max_segment_size, candidate_blocks); - const int resident_smem_bytes = layout_t::min_smem_bytes_from_num_items(num_block_items); - if (resident_smem_bytes > max_dynamic_smem_bytes) - { - // Unreachable for candidate_blocks >= min_blocks_per_segment, but guards the SMEM budget regardless. - continue; - } - - const auto clusters_per_wave = - probe_clusters_per_wave(kernel_ptr, stream, threads_per_block, candidate_blocks, resident_smem_bytes); - if (!clusters_per_wave) - { - return ::cuda::std::unexpected(clusters_per_wave.error()); - } - if (*clusters_per_wave <= 0) - { - continue; // cluster blocks not launchable at this SMEM. - } - - const auto waves = ::cuda::ceil_div(num_segments, *clusters_per_wave); - // Min waves, tie-break largest count: the loop ascends, so `<=` keeps the largest at equal waves (`best_waves` - // starts at `UINT64_MAX`, so the first launchable count always wins). - if (waves <= best_waves) - { - best_waves = waves; - cluster_blocks = candidate_blocks; - dynamic_smem_bytes = resident_smem_bytes; - } - } - - if (cluster_blocks == 0 && min_blocks_per_segment == 1) - { - // No multi-CTA config was launchable; fall back to single-CTA full residency. Slower for large segments, but - // `min_blocks_per_segment == 1` guarantees the resident SMEM fits the budget and one CTA is always launchable. - cluster_blocks = 1; - dynamic_smem_bytes = layout_t::min_smem_bytes_from_num_items(max_segment_size); - } - } - - if (cluster_blocks == 0) - { - // Oversize (`min_blocks_per_segment > eff_max_blocks_per_cluster`) or nothing launchable: full residency - // is impossible, so maximize residency with the largest launchable cluster at the largest SMEM and stream the - // overflow. - const auto hw_max_cluster_blocks = - probe_max_cluster_blocks(kernel_ptr, stream, threads_per_block, max_dynamic_smem_bytes); - if (!hw_max_cluster_blocks) - { - return ::cuda::std::unexpected(hw_max_cluster_blocks.error()); - } - cluster_blocks = (::cuda::std::min) (*hw_max_cluster_blocks, eff_max_blocks_per_cluster); - if (cluster_blocks <= 0) - { - return ::cuda::std::unexpected(cudaErrorInvalidValue); - } - dynamic_smem_bytes = max_dynamic_smem_bytes; - } - } - - return cluster_launch_shape{cluster_blocks, dynamic_smem_bytes}; -} - -// Cluster arm of the dispatch (host-only): after the shared query-pass / CC-guard setup, launches the single kernel -// symbol via `cudaLaunchKernelEx` using the resolved-CC cluster policy and geometry from `policy_getter`. -// `select_directions` arrives already wrapped; the cluster tuning comes from `policy_getter` (the resolved-CC policy) -// and the requested `Determinism`/`TieBreak` from the dispatch. The kernel launch goes through `launcher_factory`; the -// cluster occupancy / shared-memory setup queries still use the CUDA runtime directly. -template 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 -_CCCL_HOST_API cudaError_t launch_cluster_arm( - PolicyGetter policy_getter, - void* d_temp_storage, - size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k_param, - SelectDirectionParameterT select_directions, - NumSegmentsParameterT num_segments, - cudaStream_t stream, - KernelLauncherFactory launcher_factory) -{ - // 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"); - - // The cluster arm needs no temporary storage; report a positive size so the two-phase protocol proceeds. - if (d_temp_storage == nullptr) - { - temp_storage_bytes = 1; - return cudaSuccess; - } - - // A `tune`d override (`UserProvidedTuning`) can force the cluster backend on a device that cannot run it: return - // cudaErrorNotSupported rather than launch a cluster kernel the device lacks (the deferred-mode runtime behavior - // tests and benchmarks rely on). The automatic selector never routes here below SM 9.0, so its instantiation drops - // this check. `PtxComputeCap` is the running code's capability (never above the hardware SM), so it also rejects an - // SM 9.0+ build on older hardware. - if constexpr (UserProvidedTuning) - { - ::cuda::compute_capability cc{}; - if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) - { - return error; - } - if (cc < ::cuda::compute_capability{9, 0}) - { - return cudaErrorNotSupported; - } - } - - // Single kernel symbol; its cluster vs baseline arm is selected device-side via `current_policy()`. - // Taking its address here ODR-uses the `__global__` template, which is what drives its emission and registration. - // Not `constexpr`: MSVC (C2326) rejects a `constexpr` local captured and ODR-used inside the lambdas below. - auto kernel_ptr = &device_batched_topk_kernel< - PolicySelector, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT, - Determinism, - TieBreak>; - - // Cluster sub-policy for the *resolved* architecture -- exactly what the device kernel instantiates via - // `current_policy()`, so the host launch config (block size, shared-memory math) stays in lock-step - // with the device policy per CC. `policy_getter()` is a constant expression, so `policy` is a non-type template arg. - constexpr cluster_topk_policy policy = policy_getter().cluster; - constexpr int threads_per_block = policy.threads_per_block; - constexpr int chunk_bytes = policy.chunk_bytes; - constexpr int load_align_bytes = policy.load_align_bytes; - constexpr int max_chunk_slots_per_block = policy.max_chunk_slots_per_block; - static_assert(policy.max_blocks_per_cluster >= 0, - "max_blocks_per_cluster must be 0 (unrestricted) or a positive cluster block count"); - static_assert(max_chunk_slots_per_block >= 0, - "max_chunk_slots_per_block must be 0 (unrestricted) or a positive count"); - - using key_it_t = it_value_t; - using key_t = it_value_t; - using layout_t = batched_topk_cluster::smem_block_tile_layout; - static_assert(is_valid_cluster_policy(policy)); - static_assert(load_align_bytes % int{sizeof(key_t)} == 0); - - // Tightest upper bound the segment-size argument carries -- for a static-bounded per-segment sequence a loose type - // max, not the actual runtime maximum across segments. - const auto max_seg_size = ::cuda::args::__highest_(segment_sizes); - 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}); - - // Opt in to non-portable cluster blocks (>8 on Hopper). - if (const auto error = CubDebug(::cudaFuncSetAttribute( - reinterpret_cast(kernel_ptr), cudaFuncAttributeNonPortableClusterSizeAllowed, 1))) - { - return error; - } - - // Usable dynamic shared-memory budget (opt-in minus the kernel's static footprint); the policy slot cap may narrow - // it further into `max_dynamic_smem_bytes` below. - int hw_dynamic_smem_bytes = 0; - if (const auto error = max_dynamic_smem_size_for_fixed(hw_dynamic_smem_bytes, kernel_ptr)) - { - return error; - } - // 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, 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_resident_items <= 0` guard below. - const int max_dynamic_smem_bytes = - (max_chunk_slots_per_block == 0) - ? hw_dynamic_smem_bytes - : (::cuda::std::min) (hw_dynamic_smem_bytes, layout_t::min_smem_bytes_from_num_chunks(max_chunk_slots_per_block)); - - // Set the kernel's dynamic-SMEM opt-in once, to the per-symbol maximum, before any occupancy probe or launch. - // `max_dynamic_smem_bytes` is fixed by the compile-time policy and the device, so every thread sharing this kernel - // symbol writes the identical value: the process-global attribute cannot be raced to a lower value that would fail a - // concurrent launch. It also covers every launch shape the search below can pick (all `<= max_dynamic_smem_bytes`) - // and keeps the occupancy probes and the final launch on one consistent opt-in. - if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(kernel_ptr, max_dynamic_smem_bytes)) - { - return error; - } - - // Resolve the launch shape (cluster blocks + dynamic SMEM) for the max segment size. - const auto shape = select_cluster_launch_shape( - static_cast<::cuda::std::uint64_t>(max_seg_size), - static_cast<::cuda::std::uint64_t>(num_seg_val), - max_dynamic_smem_bytes, - policy, - stream, - kernel_ptr); - if (!shape) - { - return shape.error(); - } - - const int cluster_blocks = shape->cluster_blocks; - const int dynamic_smem_bytes = shape->dynamic_smem_bytes; - const auto max_block_resident_items = layout_t::max_block_resident_items(dynamic_smem_bytes); - - // One cluster per segment, its CTAs stacked in the grid's y-dimension so the x-extent stays `num_segments`: a - // flattened x == num_segments * cluster_blocks would overrun the 2^31-1 grid-x limit for a multi-CTA cluster well - // before `num_segments` reached its INT_MAX maximum (already <= INT_MAX by the entry check in `dispatch`; - // `cluster_blocks` is far below the y-dimension limit). The device reads the segment id from clusterid.x (segments - // stay the x-extent) and the CTA rank from cluster_ctarank (linearized across the cluster's dims), so this needs no - // agent change. - const dim3 grid_dim{static_cast(num_seg_val), static_cast(cluster_blocks), 1u}; - const dim3 cluster_dim{1u, static_cast(cluster_blocks), 1u}; - - // The cluster dimension routes the host launch through `cudaLaunchKernelEx`. - if (const auto error = CubDebug( - launcher_factory(grid_dim, - dim3{static_cast(threads_per_block)}, - static_cast<::cuda::std::size_t>(dynamic_smem_bytes), - stream, - /*dependent_launch=*/false, - cluster_dim) - .doit(kernel_ptr, - 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>(max_block_resident_items)}))) - { - return error; - } - - return CubDebug(detail::DebugSyncStream(stream)); -} - -// 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`. All kernel launches, memsets and nested scans go through -// `launcher_factory`. -template -_CCCL_HOST_API 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, - KernelLauncherFactory launcher_factory) -{ - // Whether some one-worker-per-segment policy covers the static max segment size within the shared-memory limit. - // Computed from this call's concrete agent types (a tuning override exposes no `baseline_can_cover` member). The - // automatic selector never routes here when this is false; only a trusted `tune`d override forces the baseline - // backend on an oversize segment. Strict mode rejects that at compile time (the static_assert below, mirroring the - // arch-unsupported one in `dispatch`); deferred mode keeps the two-phase runtime cudaErrorNotSupported path. - constexpr bool baseline_can_cover = baseline_can_cover_v< - PolicyGetter, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - LargeSegmentTileOffsetT>; -#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ - && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - static_assert( - baseline_can_cover, - "cub::DeviceBatchedTopK: the forced baseline backend cannot cover the static maximum segment size within the " - "shared-memory limit. Force the cluster backend, lower the segment-size bound, or define " - "CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime (cudaErrorNotSupported)."); -#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) - // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - if constexpr (!baseline_can_cover) - { - if (d_temp_storage == nullptr) - { - temp_storage_bytes = 1; - return cudaSuccess; - } - return cudaErrorNotSupported; - } - else - { - using large_segment_tile_offset_t = LargeSegmentTileOffsetT; - - // Determine which one-worker-per-segment policy covers the segment-size range and k. Resolve from the handed - // resolved-CC `PolicyGetter` (not `PolicySelector`) so the host picks the same baseline policy the device kernel - // instantiates for the resolved CC -- they would otherwise diverge once baseline tuning becomes CC-dependent. - constexpr auto policy = find_smallest_covering_policy_for_getter< - PolicyGetter, - SegmentSizeParameterT, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - large_segment_tile_offset_t>::policy; - constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; - constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; - - static constexpr int worker_per_segment_tile_size = - worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; - static constexpr bool any_small_segments = - ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; - static constexpr bool only_small_segments = - ::cuda::args::__traits::highest <= worker_per_segment_tile_size; - - // Allocation layout: - // only_small_segments: [0] dummy. - // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, - // [2] large-segment ids. - // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. - static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); - size_t allocation_sizes[allocations_array_size] = {1}; - - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - using counters_t = batched_topk_counters; - using segment_size_scan_offset_t = detail::choose_offset_t; - using segment_size_scan_input_op_t = - segment_size_to_tile_count_op; - static constexpr auto multi_worker_per_segment_tile_size = - multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; - const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; - // Transform iterator over [0, num_segments) producing the tile-count for each segment. - [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( - ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); - - if constexpr (!only_small_segments) - { - const auto num_segments_val = params::get_param(num_segments, 0); - // TODO(topk): the baseline large-segment (multi-CTA) path is WIP. Completing it requires: (1) guarding the - // `num_segments_val * sizeof(...)` byte counts below against size_t overflow (safe today only because the entry - // bounds num_segments_val to <= INT_MAX); (2) making the baseline tunable by populating its `epilogue` and - // `multi_worker_per_segment_policy` sub-policies and adding matching knobs to the segmented_topk benchmarks, - // which leave them zero-initialized today so baseline sweeps are not yet meaningful. - allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); - if constexpr (any_small_segments) - { - allocation_sizes[1] = sizeof(counters_t); - allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); - } - else - { - // Query the temporary storage requirement of the segment-size transform-scan. - if (const auto error = CubDebug(detail::scan::dispatch( - nullptr, - allocation_sizes[1], - segment_size_scan_input_it, - static_cast(nullptr), - ::cuda::std::plus<>{}, - detail::InputValue(large_segment_tile_offset_t{0}), - static_cast(num_segments_val), - stream, - {}, - {}, - launcher_factory))) - { - return error; - } - } - } - - void* allocations[allocations_array_size] = {}; - if (const auto error = - CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) - { - return error; - } - - if (d_temp_storage == nullptr) - { - return cudaSuccess; - } - - // `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(launcher_factory.MemsetAsync(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( - launcher_factory(grid_dim, block_dim, 0, stream, /*dependent_launch=*/false) - .doit( - device_batched_topk_kernel< - PolicySelector, - KeyInputItItT, - KeyOutputItItT, - ValueInputItItT, - ValueOutputItItT, - SegmentSizeParameterT, - KParameterT, - SelectDirectionParameterT, - NumSegmentsParameterT, - large_segment_tile_offset_t, - Determinism, - TieBreak>, - 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, - {}, - {}, - launcher_factory))) - { - return error; - } - } - - return CubDebug(detail::DebugSyncStream(stream)); - } -} - -#if !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()) - { - any = any || (PolicySelector{}(cc).backend == topk_algorithm::unsupported); - } - return any; -} -#endif // !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 the -// resolved policy selector (deterministic -> cluster; otherwise the arch+size crossover). Both host arms launch the -// same kernel symbol. `Determinism`/`TieBreak` are compile-time selection inputs. -// -// `tuning_env` carries an optional `tune`d policy selector (keyed on `topk_policy`): when present it fully replaces the -// automatic selector -- its `.backend` chooses the arm and its `.baseline`/`.cluster` carry the tunings. Matching -// DeviceScan/DeviceTransform, the tuned backend choice is trusted; only the determinism/tie-break guard below still -// applies. `launcher_factory` routes the kernel launches, memsets, nested scans and the routing CC query (the cluster -// arm's occupancy / shared-memory queries still call the CUDA runtime directly). -template < - ::cuda::execution::determinism::__determinism_t Determinism = - ::cuda::execution::determinism::__determinism_t::__not_guaranteed, - ::cuda::execution::tie_break::__tie_break_t TieBreak = ::cuda::execution::tie_break::__tie_break_t::__unspecified, - typename KeyInputItItT, - typename KeyOutputItItT, - typename ValueInputItItT, - typename ValueOutputItItT, - typename SegmentSizeParameterT, - typename KParameterT, - typename SelectDirectionT, - typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, - typename TuningEnvT = ::cuda::std::execution::env<>, - typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> -_CCCL_HOST_API cudaError_t dispatch( + typename TotalNumItemsGuaranteeT, + typename PolicySelector = policy_selector_from_types>, + 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( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputItItT d_key_segments_it, @@ -941,222 +147,209 @@ _CCCL_HOST_API cudaError_t dispatch( SelectDirectionT select_direction, NumSegmentsParameterT num_segments, [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, - cudaStream_t stream, - const TuningEnvT& = {}, - KernelLauncherFactory launcher_factory = {}) + cudaStream_t stream = nullptr, + [[maybe_unused]] PolicySelector policy_selector = {}) { - // Both arms resolve `num_segments` on the host (allocation sizing, grid extent, empty-batch guard), so it must be a - // host-known single value; device-resident counts are future work. Defensive: the public entry checks this too, but - // `dispatch` is also called directly (tests / benchmarks). - 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 value)."); + using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; - // 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). - // Type derived from the parameter type rather than `decltype(select_directions)`: GCC 7 rejects the latter ("use of - // 'select_directions' before deduction of 'auto'") when it feeds the `constexpr baseline_can_cover` initializer - // below. Declaring `select_directions` with the alias keeps its (const-qualified) type single-sourced. - using SelectDirectionParameterT = const decltype(wrap_select_direction(::cuda::std::declval())); - SelectDirectionParameterT select_directions = wrap_select_direction(select_direction); + // Wrap the raw enum into the internal discrete param type + 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; - - // Default automatic selector from the compile-time inputs; it computes its own baseline coverage. A `tune`d selector - // in the environment (keyed on `topk_policy`) replaces it wholesale. - using default_policy_selector_t = policy_selector_from_types< - key_t, - value_t, - max_k, - static_max_seg, - Determinism, - TieBreak, + // Helper that determines (a) whether there's any one-worker-per-segment policy supporting the range of segment + // sizes and k, and (b) if so, which set of one-worker-per-segment policies to use + constexpr auto policy = find_smallest_covering_policy< + PolicySelector, SegmentSizeParameterT, KeyInputItItT, KeyOutputItItT, ValueInputItItT, ValueOutputItItT, + SegmentSizeParameterT, KParameterT, SelectDirectionParameterT, NumSegmentsParameterT, - LargeSegmentTileOffsetT>; + 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; - // Type derived from the query-result trait rather than `decltype(policy_selector)`: GCC 7 rejects the latter ("use of - // 'policy_selector' before deduction of 'auto'") when `policy_selector_t` is later named inside the dispatch lambda. - using policy_selector_t = - ::cuda::std::execution::__query_result_or_t; -#if _CCCL_HAS_CONCEPTS() - static_assert(topk_policy_selector, - "Invalid policy selector for cub::DeviceBatchedTopK::dispatch"); -#endif // _CCCL_HAS_CONCEPTS() + 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; -#if !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. Two causes reach here: a deterministic / large-segment request while a pre-SM90 target is present - // (the cluster backend requires SM90+), or _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH disabling the cluster backend on all - // architectures. 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 cannot be served on at least one architecture this " - "translation unit targets. The deterministic / large-segment path requires the cluster backend (SM90+), which is " - "unavailable either because a pre-SM90 architecture is targeted or because _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH is " - "defined (which disables the cluster backend on all architectures). To fix: target only SM90+ and leave " - "_CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH undefined, relax the request (non-deterministic and small enough for the " - "baseline backend), or define CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime " - "(cudaErrorNotSupported)."); -#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) - // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) + // 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}; - // 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_and_clamp_param_to_nonnegative). 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. + 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); - ::cuda::compute_capability cc{}; - if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + if constexpr (!only_small_segments) + { + const auto num_segments_val = params::get_param(num_segments, 0); + // Scan output + allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); + if constexpr (any_small_segments) + { + allocation_sizes[1] = sizeof(counters_t); + // Large segment ids for indirectly accessing the large segment parameters + allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); + } + else + { + // Query the temporary storage requirement of the segment-size transform-scan. + if (const auto error = CubDebug(detail::scan::dispatch( + nullptr, + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(nullptr), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(num_segments_val), + stream))) + { + return error; + } + } + } + + // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) + void* allocations[allocations_array_size] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) { return error; } - // `num_segments` maps to the grid's x-extent in both host launch arms (the baseline arm launches one block per - // segment; the cluster arm launches one cluster per segment, stacking the cluster's CTAs in the grid's y-dimension), - // so it must fit a positive 32-bit grid dimension. A count above INT_MAX cannot, so reject it as an out-of-contract - // value at this single host boundary: otherwise the baseline arm would silently narrow it to `int` and the cluster - // arm would build an out-of-range grid.x. - using num_segments_val_t = typename ::cuda::args::__traits::element_type; - const num_segments_val_t num_segments_val = detail::params::get_param(num_segments, num_segments_val_t{0}); - // Unary `+` integer-promotes the count to a standard integer type so the sign-safe `cmp_*` comparators accept it: - // they are constrained to `__cccl_is_integer_v`, which excludes the character count types the public API permits. - if (::cuda::std::cmp_greater(+num_segments_val, ::cuda::std::numeric_limits::max())) + if (d_temp_storage == nullptr) { - return cudaErrorInvalidValue; - } - // A negative count is no work (like a zero count), matching DeviceSegmentedReduce. Short-circuit here, before - // `dispatch_compute_cap`, so the query pass cannot fall into the baseline arm where `num_segments_val * sizeof(...)` - // would cast the negative count to a huge `size_t`. (Zero is handled by `empty_batch_no_launch` below, which keeps - // the arch-gated no-op semantics.) TODO(topk): file an issue to unify the negative-`num_segments` contract across - // CUB device algorithms. - if (::cuda::std::cmp_less(+num_segments_val, 0)) - { - if (d_temp_storage == nullptr) - { - temp_storage_bytes = 1; - } return cudaSuccess; } - // 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` suffices for `num_segments`: a negative count is already - // short-circuited as no work above. 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 || ::cuda::args::__highest_(segment_sizes) <= 0); - }; + // 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."); - return detail::dispatch_compute_cap(policy_selector_t{}, cc, [&](auto policy_getter) -> cudaError_t { - constexpr topk_policy active_policy = policy_getter(); -#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) - NV_IF_TARGET(NV_IS_HOST, ({ - ::std::stringstream ss; - ss << active_policy; - _CubLog("Dispatching DeviceBatchedTopK to compute capability %d.%d with tuning: %s\n", - cc.major_cap(), - cc.minor_cap(), - ss.str().c_str()); - })) -#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) - if constexpr (active_policy.backend == topk_algorithm::baseline) + if constexpr (any_small_segments) + { + if constexpr (!only_small_segments) { - // Computed from the template parameters, not a captured function-scope constant: MSVC rejects the latter as - // non-constant inside this lambda's `if constexpr`. - constexpr bool deterministic = (Determinism != ::cuda::execution::determinism::__determinism_t::__not_guaranteed) - || (TieBreak != ::cuda::execution::tie_break::__tie_break_t::__unspecified); - if constexpr (deterministic) + // 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))) { - // A `tune`d selector forced the baseline backend for a deterministic / tie-break request it cannot serve (only - // the SM 9.0+ cluster backend is deterministic). Mirror the arch-unsupported / oversize-baseline failure model: - // a hard compile error by default, deferred to a runtime cudaErrorNotSupported only under the escape hatches. -#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) \ - && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - static_assert( - !deterministic, - "cub::DeviceBatchedTopK: a tuned policy selector forced the baseline backend for a deterministic " - "/ tie-break request it cannot serve (only the SM 9.0+ cluster backend is deterministic). Drop " - "the override, relax the determinism / tie-break requirement, or define " - "CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT to defer the diagnosis to runtime " - "(cudaErrorNotSupported)."); -#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) - // && !defined(CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT) - // 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; - } - else - { - if (empty_batch_no_launch()) - { - return cudaSuccess; - } - return launch_baseline_arm( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - select_directions, - num_segments, - stream, - launcher_factory); + return error; } } - else if constexpr (active_policy.backend == topk_algorithm::cluster) + 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])))) { -#if !_CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() - // The automatic selector never picks the cluster backend when dynamic cluster launches are disabled (see - // cluster_capable), so reaching here means a `tune`d selector forced it. The kernel would launch without its - // cluster extent (triple_chevron drops it), so reject the contradiction at compile time rather than run wrong. - static_assert(active_policy.backend != topk_algorithm::cluster, - "cub::DeviceBatchedTopK: a tuned policy selector forced the cluster backend, but " - "_CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH is defined. Drop the override or the macro."); -#endif // !_CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() - if (empty_batch_no_launch()) - { - return cudaSuccess; - } - // `UserProvidedTuning`: false for the automatic selector, which returns `cluster` solely for a - // `cluster_capable(cc)` and so needs no runtime re-check; a `tune`d override is a different type and keeps it. - // Inlined as a type trait rather than a function-scope constexpr, which MSVC rejects inside this lambda. - return launch_cluster_arm>( - policy_getter, + return error; + } + } + else + { + // No small segments: the small-kernel epilogue (which would otherwise produce the per-segment tile offsets) does + // not run. Compute the per-segment tile offsets directly via a transform-scan over all segment sizes. + // The large segment agent will either consume these offsets directly (segment_id -> tile offset) or, when going + // through the large-segment queue, via a transform iterator over `d_large_segments_ids` (level of indirection). + if (const auto error = CubDebug(detail::scan::dispatch( + allocations[1], + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(allocations[0]), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(params::get_param(num_segments, 0)), + stream))) + { + return error; + } + } + + if constexpr (!only_small_segments) + { + // TODO (elstehle): support larger number of segments through multiple kernel launches + // Depending on any_small_segments, we need to either: + // - Indirectly get the large segment parameters via the queued large segment IDs + // - Directly take the segment parameters since all segments are large + } + return CubDebug(detail::DebugSyncStream(stream)); +} +// Env-based dispatch function handling memory allocation as well. This is usually done by the device-layer, but there +// is no public API for segmented topk yet. +template > +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items_guarantee, + const EnvT& env = {}) +{ + 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( d_temp_storage, temp_storage_bytes, d_key_segments_it, @@ -1167,21 +360,10 @@ _CCCL_HOST_API cudaError_t dispatch( k, select_directions, num_segments, + total_num_items_guarantee, stream, - launcher_factory); - } - 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; - } - }); + policy_selector); + }); } } // namespace detail::batched_topk diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh index 4c87d939..a0f3df37 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_copy_mdspan.cuh @@ -21,14 +21,8 @@ #include #include -#include #include -#include -#include -#include #include -#include -#include #include CUB_NAMESPACE_BEGIN @@ -53,46 +47,6 @@ struct copy_mdspan_t } }; -template -[[nodiscard]] _CCCL_HOST_API ::cudaError_t -__copy_mdspan_bytes(::cuda::stream_ref __stream, _MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out) -{ - _CCCL_TRY - { - ::cuda::copy_bytes(__stream, __mdspan_in, __mdspan_out); - } - _CCCL_CATCH (const ::cuda::cuda_error& __e) - { - return __e.status(); - } -#if _CCCL_HOSTED() - _CCCL_CATCH (const ::std::invalid_argument& __e) - { - static_cast(__e); - return ::cudaErrorInvalidValue; - } -#endif // _CCCL_HOSTED - _CCCL_CATCH_ALL - { - return ::cudaErrorUnknown; - } - - return ::cudaSuccess; -} - -template -[[nodiscard]] CUB_RUNTIME_FUNCTION ::cudaError_t -__transform_copy(_MDSpanIn&& __mdspan_in, _MDSpanOut&& __mdspan_out, const _Env& __env) -{ - return CUB_NS_QUALIFIER::DeviceTransform::__transform_internal( - ::cuda::std::make_tuple(__mdspan_in.data_handle()), - __mdspan_out.data_handle(), - __mdspan_in.size(), - ::cuda::always_true{}, - ::cuda::std::identity{}, - __env); -} - template mdspan_in, if (mdspan_in.is_exhaustive() && mdspan_out.is_exhaustive() && detail::have_same_strides(mdspan_in.mapping(), mdspan_out.mapping())) { - // NOLINTBEGIN(bugprone-branch-clone) - if constexpr (::cuda::std::same_as - && ::cuda::__detail::__can_mdspan_copy_bytes - && ::cuda::std::__is_callable_v<::cuda::get_stream_t, const EnvT&>) - { - NV_IF_TARGET( - NV_IS_HOST, - ({ - auto __stream = ::cuda::get_stream(env); - - // cuda::copy_bytes() builds an __ensure_current_context(stream_ref), which calls - // cuStreamGetCtx(). That driver call rejects the NULL stream with - // CUDA_ERROR_INVALID_VALUE. Use the transform path, which goes through the runtime - // API and accepts the NULL stream. - // - // Likewise, we cannot retrieve the context for a stream that is capturings so we - // need to call the kernel. - if (__stream.get() == nullptr - || (::cuda::__driver::__streamIsCapturing(__stream.get()) == ::CU_STREAM_CAPTURE_STATUS_ACTIVE)) - { - return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env); - } - - return CUB_NS_QUALIFIER::detail::copy_mdspan::__copy_mdspan_bytes(__stream, mdspan_in, mdspan_out); - }), - (return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env);)) - } - else - { - return CUB_NS_QUALIFIER::detail::copy_mdspan::__transform_copy(mdspan_in, mdspan_out, env); - } - // NOLINTEND(bugprone-branch-clone) + return cub::DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(mdspan_in.data_handle()), + mdspan_out.data_handle(), + mdspan_in.size(), + ::cuda::always_true{}, + ::cuda::std::identity{}, + env); } // TODO (fbusato): add ForEachInLayout when mdspan_in and mdspan_out have compatible layouts // Compatible layouts could use more efficient iteration patterns diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh index 673be75f..704c57ad 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_histogram.cuh @@ -216,12 +216,12 @@ CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( if constexpr (IsDeviceInit) { return kernel_source.template HistogramSweepKernelDeviceInit< - PolicySelector, - PRIVATIZED_SMEM_BINS, - FirstLevelArrayT, - SecondLevelArrayT, - IsEven, - IsByteSample>(); + PolicySelector, + PRIVATIZED_SMEM_BINS, + FirstLevelArrayT, + SecondLevelArrayT, + IsEven, + IsByteSample>(); } else { diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh index 926f40ef..91560416 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_reduce_by_key.cuh @@ -208,14 +208,14 @@ __launch_bounds__(int(current_policy().lookback.threads_per_bloc { static constexpr ReduceByKeyPolicy policy = current_policy(); using AgentReduceByKeyPolicyT = agent_reduce_by_key_policy< - policy.lookback.threads_per_block, - policy.lookback.items_per_thread, - policy.lookback.load_algorithm, - policy.lookback.load_modifier, - policy.lookback.scan_algorithm, - delay_constructor_t>; + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; using vsmem_helper_t = vsmem_helper_default_fallback_policy_t< AgentReduceByKeyPolicyT, @@ -656,14 +656,14 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get // TODO(bgruber): refactor this in the future constexpr ReduceByKeyPolicy policy = policy_getter(); using Policy = agent_reduce_by_key_policy< - policy.lookback.threads_per_block, - policy.lookback.items_per_thread, - policy.lookback.load_algorithm, - policy.lookback.load_modifier, - policy.lookback.scan_algorithm, - delay_constructor_t>; + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; using vsmem_helper_t = vsmem_helper_default_fallback_policy_t; return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS, vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh index a62f5e81..bb3762ef 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan.cuh @@ -220,11 +220,11 @@ template < typename ScanOpT, typename InitValueT, typename OffsetT, - typename AccumT = ::cuda::std::__accumulator_t, - ::cuda::std::_If<::cuda::std::is_same_v, - cub::detail::it_value_t, - typename InitValueT::value_type>>, + typename AccumT = ::cuda::std::__accumulator_t, + ::cuda::std::_If<::cuda::std::is_same_v, + cub::detail::it_value_t, + typename InitValueT::value_type>>, ForceInclusive EnforceInclusive = ForceInclusive::No, typename PolicyHub = detail::scan:: policy_hub, detail::it_value_t, AccumT, OffsetT, ScanOpT>, @@ -552,44 +552,38 @@ struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") DispatchScan int smem_size = smem_size_1_stage; // When launched from the host, maximize the number of stages that we can fit inside the shared memory. - NV_IF_TARGET( - NV_IS_HOST, ({ - // number of stages to have an even workload across all SMs (improves small problem sizes), assuming - // 1 CTA per SM +1 since it tends to improve performance - // TODO(bgruber): make the +1 a tuning parameter - const int max_stages_for_even_workload = static_cast( - ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + NV_IF_TARGET(NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); - while (num_stages <= max_stages_for_even_workload) - { - const int next_smem_size = detail::scan::smem_for_stages( - lookahead_policy, - num_stages + 1, - static_cast(kernel_source.InputSize()), - static_cast(kernel_source.InputAlign()), - static_cast(kernel_source.OutputAlign()), - static_cast(kernel_source.AccumSize()), - static_cast(kernel_source.AccumAlign())); - if (next_smem_size > max_dynamic_smem_size) - { - // This number of stages failed, so stay at the current settings - break; - } + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } - smem_size = next_smem_size; - ++num_stages; - } + smem_size = next_smem_size; + ++num_stages; + } - // Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents - // multiple kernels from overwriting the max shared memory limit by different values. - // - // TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory - // sizes, however we need something that works even with older CTKs. - if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size)) - { - return error; - } - })) + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) + { + return error; + } + })) // Invoke init kernel { @@ -1169,44 +1163,38 @@ CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookahead( int smem_size = smem_size_1_stage; // When launched from the host, maximize the number of stages that we can fit inside the shared memory. - NV_IF_TARGET( - NV_IS_HOST, ({ - // number of stages to have an even workload across all SMs (improves small problem sizes), assuming - // 1 CTA per SM +1 since it tends to improve performance - // TODO(bgruber): make the +1 a tuning parameter - const int max_stages_for_even_workload = static_cast( - ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + NV_IF_TARGET(NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); - while (num_stages <= max_stages_for_even_workload) - { - const int next_smem_size = detail::scan::smem_for_stages( - lookahead_policy, - num_stages + 1, - static_cast(kernel_source.InputSize()), - static_cast(kernel_source.InputAlign()), - static_cast(kernel_source.OutputAlign()), - static_cast(kernel_source.AccumSize()), - static_cast(kernel_source.AccumAlign())); - if (next_smem_size > max_dynamic_smem_size) - { - // This number of stages failed, so stay at the current settings - break; - } + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } - smem_size = next_smem_size; - ++num_stages; - } + smem_size = next_smem_size; + ++num_stages; + } - // Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents - // multiple kernels from overwriting the max shared memory limit by different values. - // - // TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory - // sizes, however we need something that works even with older CTKs. - if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size)) - { - return error; - } - })) + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) + { + return error; + } + })) // Invoke init kernel { diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh index 3ee74db8..bb316143 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_scan_by_key.cuh @@ -238,15 +238,15 @@ template < typename PolicyHub = policy_hub, ScanOpT>, typename PolicySelector = policy_selector_from_hub, typename KernelSource = DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> struct dispatch_scan_by_key { @@ -672,15 +672,15 @@ template < cub::detail::it_value_t, ScanOpT>, typename KernelSource = DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires scan_by_key_policy_selector @@ -875,15 +875,15 @@ template < detail::scan_by_key::policy_hub, ScanOpT>, typename PolicySelector = detail::scan_by_key::policy_selector_from_hub, typename KernelSource = detail::scan_by_key::DeviceScanByKeyKernelSource< - PolicySelector, - KeysInputIteratorT, - ValuesInputIteratorT, - ValuesOutputIteratorT, - EqualityOp, - ScanOpT, - InitValueT, - OffsetT, - AccumT>, + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> using DispatchScanByKey CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") = detail::scan_by_key::dispatch_scan_by_key< diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh index ce33983c..2c9d6fbe 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_radix_sort.cuh @@ -911,7 +911,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( { using default_policy_selector_t = policy_selector_from_types; using policy_selector_t = ::cuda::std::decay_t< - ::cuda::std::execution::__query_result_or_t>; + ::cuda::std::execution::__query_result_or_t>; #if _CCCL_HAS_CONCEPTS() static_assert(segmented_radix_sort_policy_selector); #endif // _CCCL_HAS_CONCEPTS() diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh index a01cfe45..35c432b4 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_reduce.cuh @@ -503,15 +503,15 @@ template < decltype(select_segmented_accum_t(static_cast(nullptr))), typename PolicySelector = policy_selector_from_types, typename KernelSource = DeviceSegmentedReduceKernelSource< - PolicySelector, - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorT, - EndOffsetIteratorT, - OffsetT, - ReductionOpT, - InitValueT, - AccumT>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_reduce_policy_selector @@ -711,13 +711,13 @@ template (nullptr))), typename PolicySelector = policy_selector_from_types, typename KernelSource = DeviceFixedSizeSegmentedReduceKernelSource< - PolicySelector, - InputIteratorT, - OutputIteratorT, - OffsetT, - ReductionOpT, - InitValueT, - AccumT>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, ::cuda::std::enable_if_t<::cuda::std::is_arithmetic_v, int> = 0> #if _CCCL_HAS_CONCEPTS() diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh index 673c9b7f..c8dd92d2 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_segmented_scan.cuh @@ -96,17 +96,17 @@ template < common_iterator_value_t, typename PolicySelector = policy_selector_from_types, typename KernelSource = device_segmented_scan_kernel_source< - PolicySelector, - InputIteratorT, - OutputIteratorT, - BeginOffsetIteratorInputT, - EndOffsetIteratorInputT, - BeginOffsetIteratorOutputT, - OffsetT, - ScanOpT, - InitValueT, - AccumT, - EnforceInclusive>, + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorInputT, + EndOffsetIteratorInputT, + BeginOffsetIteratorOutputT, + OffsetT, + ScanOpT, + InitValueT, + AccumT, + EnforceInclusive>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires segmented_scan_policy_selector diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh index 3e845587..fa8bc5c8 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_select_if.cuh @@ -216,15 +216,14 @@ struct make_vsmem_helper { static constexpr SelectPolicy active_policy = DefaultPolicyGetter{}(); using agent_policy_t = detail::agent_select_if_policy< - active_policy.lookback.threads_per_block, - active_policy.lookback.items_per_thread, - active_policy.lookback.load_algorithm, - active_policy.lookback.load_modifier, - active_policy.lookback.scan_algorithm, - delay_constructor_t, - active_policy.lookback._load_prefetch>; + active_policy.lookback.threads_per_block, + active_policy.lookback.items_per_thread, + active_policy.lookback.load_algorithm, + active_policy.lookback.load_modifier, + active_policy.lookback.scan_algorithm, + delay_constructor_t>; using type = vsmem_helper_default_fallback_policy_t< agent_policy_t, bind_selection_opt::template agent_t, diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh index 61c0d352..13ab3b78 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -444,18 +444,18 @@ template , per_partition_offset_t>, typename KernelSource = DeviceThreeWayPartitionKernelSource< - PolicySelector, - InputIteratorT, - FirstOutputIteratorT, - SecondOutputIteratorT, - UnselectedOutputIteratorT, - NumSelectedIteratorT, - ScanTileStateT, - SelectFirstPartOp, - SelectSecondPartOp, - per_partition_offset_t, - streaming_context_t, - OffsetT>, + PolicySelector, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t, + OffsetT>, typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> #if _CCCL_HAS_CONCEPTS() requires three_way_partition_policy_selector diff --git a/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh b/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh index 7a3327eb..3130fb6e 100644 --- a/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/dispatch_transform_tile_config.cuh @@ -26,7 +26,7 @@ # pragma system_header #endif // no system header -#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION() && _CCCL_STD_VER >= 2020 +#define _CCCL_CUB_HAS_TILE_TRANSFORM() _CCCL_TILE_COMPILATION() #if _CCCL_CUB_HAS_TILE_TRANSFORM() && defined(_CCCL_ENABLE_EXPERIMENTAL_TILE_TRANSFORM_DISPATCH) # define _CCCL_CUB_TILE_TRANSFORM_DISPATCH_ENABLED() 1 diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh index a0518e6f..0e27f17c 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/common.cuh @@ -13,14 +13,13 @@ # pragma system_header #endif // no system header -#include #include #include +#include #include #include -#include #include #include #include @@ -131,7 +130,7 @@ template return iterator_info{ static_cast(size_of), static_cast(align_of), - ::cuda::is_trivially_copyable_v, + THRUST_NS_QUALIFIER::is_trivially_relocatable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v}; } diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh index 6cb318aa..37072f3c 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh @@ -158,13 +158,13 @@ struct BatchedCopyPolicy BatchedCopyLookbackPolicy lookback; //!< The policy for the batched-copy algorithm based on decoupled-lookback. Only //!< used when @p algorithm is @p lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const BatchedCopyPolicy& lhs, const BatchedCopyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh index 679dbbfc..48f97651 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batched_topk.cuh @@ -16,27 +16,20 @@ #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; //!< 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. + int items_per_thread; + BlockLoadAlgorithm load_algorithm; + BlockStoreAlgorithm store_algorithm; + BlockScanAlgorithm scan_algorithm; _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const epilogue_policy& lhs, const epilogue_policy& rhs) { @@ -49,26 +42,24 @@ struct epilogue_policy return !(lhs == rhs); } -#if _CCCL_HOSTED() +#if !_CCCL_COMPILER(NVRTC) friend ::std::ostream& operator<<(::std::ostream& os, const epilogue_policy& p) { return os << "epilogue_policy { .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .store_algorithm = " << p.store_algorithm << ", .scan_algorithm = " << p.scan_algorithm << " }"; } -#endif // _CCCL_HOSTED() +#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; //!< 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. + int threads_per_block; + int items_per_thread; + BlockLoadAlgorithm load_algorithm; + BlockStoreAlgorithm store_algorithm; - epilogue_policy epilogue; //!< Sub-policy for the compaction epilogue. + epilogue_policy epilogue; _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const worker_policy& lhs, const worker_policy& rhs) { @@ -89,15 +80,13 @@ struct worker_policy << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .store_algorithm = " << p.store_algorithm << ", .epilogue = " << p.epilogue << " }"; } -#endif // _CCCL_HOSTED() +#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; //!< Number of threads in a CUDA block. - int items_per_thread; //!< Keys each thread loads/processes per tile. + int threads_per_block; + int items_per_thread; _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const multi_worker_policy& lhs, const multi_worker_policy& rhs) { @@ -109,7 +98,7 @@ struct multi_worker_policy return !(lhs == rhs); } -#if _CCCL_HOSTED() +#if !_CCCL_COMPILER(NVRTC) friend ::std::ostream& operator<<(::std::ostream& os, const multi_worker_policy& p) { return os << "multi_worker_policy { .threads_per_block = " << p.threads_per_block @@ -118,31 +107,28 @@ struct multi_worker_policy #endif // _CCCL_HOSTED() }; -//! Sub-policy for the baseline (worker-per-segment) backend of @ref DeviceBatchedTopK. -struct baseline_topk_policy +struct batched_topk_policy { - //! 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. + // 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. ::cuda::std::array worker_per_segment_policies; - multi_worker_policy multi_worker_per_segment_policy; //!< Worker policy for segments too large for a single block. + multi_worker_policy multi_worker_per_segment_policy; - _CCCL_HOST_DEVICE_API friend constexpr bool - operator==(const baseline_topk_policy& lhs, const baseline_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const batched_topk_policy& lhs, const batched_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 baseline_topk_policy& lhs, const baseline_topk_policy& rhs) + _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(const batched_topk_policy& lhs, const batched_topk_policy& rhs) { return !(lhs == rhs); } #if _CCCL_HOSTED() - friend ::std::ostream& operator<<(::std::ostream& os, const baseline_topk_policy& p) + friend ::std::ostream& operator<<(::std::ostream& os, const batched_topk_policy& p) { - os << "baseline_topk_policy { .worker_per_segment_policies = { "; + os << "batched_topk_policy { .worker_per_segment_policies = { "; for (::cuda::std::size_t i = 0; i < p.worker_per_segment_policies.size(); ++i) { if (i != 0) @@ -156,242 +142,45 @@ struct baseline_topk_policy #endif // _CCCL_HOSTED() }; -// Default baseline sub-policy. Tuning is currently CC-independent. -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_baseline_policy() -> baseline_topk_policy +#if _CCCL_HAS_CONCEPTS() +template +concept batched_topk_policy_selector = policy_selector; +#endif // _CCCL_HAS_CONCEPTS() + +struct policy_selector { - constexpr auto load_alg = BLOCK_LOAD_WARP_TRANSPOSE; - constexpr auto store_alg = BLOCK_STORE_WARP_TRANSPOSE; - constexpr auto scan_alg = BLOCK_SCAN_WARP_SCANS; - constexpr auto epilogue = epilogue_policy{16, load_alg, store_alg, scan_alg}; - return baseline_topk_policy{ - {{ - worker_policy{256, 64, load_alg, store_alg, epilogue}, - worker_policy{256, 32, load_alg, store_alg, epilogue}, - worker_policy{256, 16, load_alg, store_alg, epilogue}, - worker_policy{256, 8, load_alg, store_alg, epilogue}, - worker_policy{256, 4, load_alg, store_alg, epilogue}, - worker_policy{128, 2, load_alg, store_alg, epilogue}, - }}, - multi_worker_policy{256, 64}}; -} - -// Largest maximum segment size (in keys) the baseline (worker-per-segment) backend can cover: the largest worker tile -// (threads_per_block * items_per_thread) in `policy`. A larger statically-known maximum segment size makes the baseline -// backend ineligible (the selector then picks the cluster backend where supported, otherwise `unsupported`). This is -// only the tile-based necessary condition; the exact predicate `baseline_can_cover_v` also checks the agent's -// shared-memory fit (which needs the concrete agent types). -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::int64_t -baseline_max_covered_segment_size(const baseline_topk_policy& policy) -{ - ::cuda::std::int64_t max_tile_size = 0; - for (const auto& worker : policy.worker_per_segment_policies) + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> batched_topk_policy { - 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; - } + 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}}; } - return max_tile_size; -} - -//! 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 -{ - 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). 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). 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() }; -// Default cluster sub-policy. Tuning is currently CC-independent. -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto make_cluster_policy() -> cluster_topk_policy +template +struct policy_selector_from_types { - 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 a cluster sub-policy must satisfy, mirroring the agent's compile-time invariants so a bad policy -// (e.g. from a `tune` override) trips `launch_cluster_arm`'s `static_assert(is_valid_cluster_policy(policy))` with a -// clear message instead of a cryptic failure deep in the agent (or a host-side divide-by-zero in the launch-shape -// math). The block_tile byte geometry constraints stem from the aligned bulk-copy (TMA) load path, which addresses -// gmem/smem in `load_align_bytes`-sized, aligned units. -[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto is_valid_cluster_policy(cluster_topk_policy policy) -> bool -{ - return policy.chunk_bytes > 0 && 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 - && policy.threads_per_block > 0 && policy.threads_per_block % warp_threads == 0 && policy.min_blocks_per_sm >= 0 - && policy.pipeline_stages >= 1 && policy.pipeline_stages <= 32 && policy.min_chunks_per_block >= 1 - && policy.bits_per_pass >= 1 && policy.bits_per_pass <= 16 && policy.histogram_items_per_thread > 0 - && policy.tie_break_items_per_thread > 0 && policy.copy_items_per_thread > 0 - && policy.single_block_max_seg_size >= 0 && policy.max_blocks_per_cluster >= 0 - && policy.max_chunk_slots_per_block >= 0; -} - -static_assert(is_valid_cluster_policy(make_cluster_policy())); - -// ----------------------------------------------------------------------------- -// 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_algorithm -{ - 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_algorithm backend) -{ - switch (backend) + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> batched_topk_policy { - case topk_algorithm::baseline: - return os << "baseline"; - case topk_algorithm::cluster: - return os << "cluster"; - default: - return os << "unsupported"; + return policy_selector{}(cc); } -} -#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_algorithm 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_algorithm::baseline. - cluster_topk_policy cluster; //!< Sub-policy used when @p backend is @p topk_algorithm::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() }; #if _CCCL_HAS_CONCEPTS() -template -concept topk_policy_selector = policy_selector; +static_assert(batched_topk_policy_selector); #endif // _CCCL_HAS_CONCEPTS() - -// Crossover knobs (TODO: tune via SM100 benchmarks). -//! Clusters require SM 9.0+. -inline constexpr int cluster_min_cc_major = 9; -//! 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([[maybe_unused]] ::cuda::compute_capability cc) -{ -#if _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() - return cc >= ::cuda::compute_capability{cluster_min_cc_major, 0}; -#else // ^^^ dynamic cluster launches enabled ^^^ / vvv dynamic cluster launches disabled vvv - // The cluster backend launches with a runtime cluster width, which _CCCL_DISABLE_DYNAMIC_CLUSTER_LAUNCH compiles out; - // reporting no architecture as cluster-capable makes the selector fall back to baseline (or report unsupported). - return false; -#endif // _CCCL_HAS_DYNAMIC_CLUSTER_LAUNCH() -} } // namespace detail::batched_topk CUB_NAMESPACE_END diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh index c16725b6..e44a9bdb 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_merge.cuh @@ -20,9 +20,9 @@ #include #include +#include #include -#include #include #include #include @@ -43,15 +43,10 @@ struct MergePolicy [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const MergePolicy& lhs, const MergePolicy& rhs) noexcept { - // gcc 8 folds comparisons of adjacent bool members within one expression into a BIT_FIELD_REF, which its - // constexpr evaluator cannot handle (ICE in cxx_eval_bit_field_ref, fixed in gcc 9). Keep each bool - // comparison in a separate statement to avoid the fold. - const bool same_bulk_copy_for_keys = lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys; - const bool same_bulk_copy_for_values = lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values; - const bool same_unroll = lhs.unroll == rhs.unroll; return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread && lhs.load_modifier == rhs.load_modifier && lhs.store_algorithm == rhs.store_algorithm - && same_bulk_copy_for_keys && same_bulk_copy_for_values && same_unroll; + && lhs.use_bulk_copy_for_keys == rhs.use_bulk_copy_for_keys + && lhs.use_bulk_copy_for_values == rhs.use_bulk_copy_for_values && lhs.unroll == rhs.unroll; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -95,10 +90,10 @@ struct policy_selector [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> MergePolicy { - const int tune_type_size = key_size + value_size; - const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size); - const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous - && key_iterator_value_types_are_the_same; + const int tune_type_size = key_size + value_size; + const int ipt_800_plus = nominal_4B_items_to_items(15, tune_type_size); + const bool can_bulk_keys = (key_size == key_align) && key_is_trivially_relocatable && key_iterators_are_contiguous + && key_iterator_value_types_are_the_same; const bool can_bulk_values = (value_size == value_align) && value_is_trivially_relocatable && value_iterators_are_contiguous && value_iterator_value_types_are_the_same; @@ -168,12 +163,12 @@ struct policy_selector_from_types return policy_selector{ int{sizeof(key_t)}, int{alignof(key_t)}, - ::cuda::is_trivially_copyable_v, + THRUST_NS_QUALIFIER::is_trivially_relocatable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v, ::cuda::std::is_same_v>, ::cuda::std::is_same_v ? 0 : int{sizeof(item_t)}, int{alignof(item_t)}, - ::cuda::is_trivially_copyable_v, + THRUST_NS_QUALIFIER::is_trivially_relocatable_v, THRUST_NS_QUALIFIER::is_contiguous_iterator_v && THRUST_NS_QUALIFIER::is_contiguous_iterator_v, ::cuda::std::is_same_v>, diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh index 402e8cab..9f2ccb15 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh @@ -99,13 +99,13 @@ struct ReduceByKeyPolicy ReduceByKeyLookbackPolicy lookback; //!< The policy for the reduce-by-key algorithm based on decoupled-lookback. Only //!< used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh index 1ea80618..7f1c9d17 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_encode.cuh @@ -100,13 +100,13 @@ struct RleEncodePolicy RleAlgorithm algorithm = RleAlgorithm::lookback; //!< The RLE-encode algorithm to use RleLookbackPolicy lookback; //!< The lookback policy - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const RleEncodePolicy& lhs, const RleEncodePolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh index 10d6353e..495d708f 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_rle_non_trivial_runs.cuh @@ -104,13 +104,13 @@ struct RleNonTrivialRunsPolicy RleNonTrivialRunsLookbackPolicy lookback; //!< The policy for the non-trivial-runs algorithm based on //!< decoupled-lookback. Only used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const RleNonTrivialRunsPolicy& lhs, const RleNonTrivialRunsPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh index 756ac570..93948713 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan.cuh @@ -187,14 +187,12 @@ struct ScanPolicy ScanLookbackPolicy lookback; //!< The look-back scan policy (used when algorithm is @p lookback, otherwise ignored) ScanLookaheadPolicy lookahead; //!< The lookahead scan policy (used when algorithm is @p lookahead, otherwise ignored) - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool - operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept { return lhs.lookback == rhs.lookback && lhs.lookahead == rhs.lookahead && lhs.algorithm == rhs.algorithm; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool - operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ScanPolicy& lhs, const ScanPolicy& rhs) noexcept { return !(lhs == rhs); } diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh index 015aa095..04905fc2 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_scan_by_key.cuh @@ -103,13 +103,13 @@ struct ScanByKeyPolicy ScanByKeyLookbackPolicy lookback; //!< The policy for the scan-by-key algorithm based on decoupled-lookback. Only used //!< when @p algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ScanByKeyPolicy& lhs, const ScanByKeyPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh index 5b9773d4..6f730b7b 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_segmented_sort.cuh @@ -575,13 +575,13 @@ struct policy_hub static constexpr int BLOCK_THREADS = 256; static constexpr int PARTITIONING_THRESHOLD = 500; using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy< - BLOCK_THREADS, - 23, - DominantT, - BLOCK_LOAD_TRANSPOSE, - LOAD_DEFAULT, - RADIX_RANK_MEMOIZE, - BLOCK_SCAN_WARP_SCANS, + BLOCK_THREADS, + 23, + DominantT, + BLOCK_LOAD_TRANSPOSE, + LOAD_DEFAULT, + RADIX_RANK_MEMOIZE, + BLOCK_SCAN_WARP_SCANS, (sizeof(KeyT) > 1) ? 6 : 4>; static constexpr int ITEMS_PER_SMALL_THREAD = Nominal4BItemsToItems(9); @@ -606,13 +606,13 @@ struct policy_hub static constexpr int BLOCK_THREADS = 256; static constexpr int PARTITIONING_THRESHOLD = 500; using LargeSegmentPolicy = detail::agent_radix_sort_downsweep_policy< - BLOCK_THREADS, - 23, - DominantT, - BLOCK_LOAD_TRANSPOSE, - LOAD_DEFAULT, - RADIX_RANK_MEMOIZE, - BLOCK_SCAN_WARP_SCANS, + BLOCK_THREADS, + 23, + DominantT, + BLOCK_LOAD_TRANSPOSE, + LOAD_DEFAULT, + RADIX_RANK_MEMOIZE, + BLOCK_SCAN_WARP_SCANS, (sizeof(KeyT) > 1) ? 6 : 4>; static constexpr bool LARGE_ITEMS = sizeof(DominantT) > 4; diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh index 5aaa1aac..05e64ff3 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_select_if.cuh @@ -18,7 +18,6 @@ #include #include #include -#include #include #include #include @@ -44,15 +43,13 @@ struct SelectLookbackPolicy CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory BlockScanAlgorithm scan_algorithm; //!< The @ref BlockScanAlgorithm used for scanning LookbackDelayPolicy lookback_delay; //!< The policy configuring the delay used in decoupled lookback - detail::LoadPrefetch _load_prefetch = detail::LoadPrefetch::none; //!< Implementation detail; do not use directly [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(const SelectLookbackPolicy& lhs, const SelectLookbackPolicy& rhs) noexcept { return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier - && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay - && lhs._load_prefetch == rhs._load_prefetch; + && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay; } [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool @@ -65,10 +62,9 @@ struct SelectLookbackPolicy friend ::std::ostream& operator<<(::std::ostream& os, const SelectLookbackPolicy& p) { return os - << "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block - << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm - << ", .load_modifier = " << p.load_modifier << ", .scan_algorithm = " << p.scan_algorithm - << ", .lookback_delay = " << p.lookback_delay << ", ._load_prefetch = " << p._load_prefetch << " }"; + << "SelectLookbackPolicy { .threads_per_block = " << p.threads_per_block << ", .items_per_thread = " + << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .load_modifier = " << p.load_modifier + << ", .scan_algorithm = " << p.scan_algorithm << ", .lookback_delay = " << p.lookback_delay << " }"; } #endif // _CCCL_HOSTED() }; @@ -106,14 +102,12 @@ struct SelectPolicy SelectLookbackPolicy lookback; //!< The policy for the selection algorithm based on decoupled-lookback. Only used when //!< @p algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool - operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool - operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const SelectPolicy& lhs, const SelectPolicy& rhs) noexcept { return !(lhs == rhs); } @@ -196,13 +190,13 @@ struct PartitionPolicy PartitionLookbackPolicy lookback; //!< The policy for the partition algorithm based on decoupled-lookback. Only used //!< when algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const PartitionPolicy& lhs, const PartitionPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh index 2099f08d..ae3d4d1c 100644 --- a/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh +++ b/cccl_upstream/cub/cub/device/dispatch/tuning/tuning_three_way_partition.cuh @@ -98,13 +98,13 @@ struct ThreeWayPartitionPolicy ThreeWayPartitionLookbackPolicy lookback; //!< The policy for the three-way partition algorithm based on //!< decoupled-lookback. Only used when @p algorithm is @lookback. - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator==(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept { return lhs.algorithm == rhs.algorithm && lhs.lookback == rhs.lookback; } - [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool + [[nodiscard]] _CCCL_API friend constexpr bool operator!=(const ThreeWayPartitionPolicy& lhs, const ThreeWayPartitionPolicy& rhs) noexcept { return !(lhs == rhs); diff --git a/qwen3_6_scripts/test_xllm_cuda_kernels.py b/qwen3_6_scripts/test_xllm_cuda_kernels.py new file mode 100644 index 00000000..4edfe905 --- /dev/null +++ b/qwen3_6_scripts/test_xllm_cuda_kernels.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 +""" +test_xllm_cuda_kernels.py — Verify imported xllm CUDA kernels on BI-V100 + +Tests each kernel by: + 1. Compile .cu → .so via torch.utils.cpp_extension + 2. Call through pybind11 with reference data + 3. Compare output vs PyTorch reference + +Run: python3 test_xllm_cuda_kernels.py +Requires: BI-V100 GPU, corex SDK, torch, ixformer +""" + +import os +import sys +import time +import torch +import traceback + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +PROJECT_DIR = os.path.dirname(SCRIPT_DIR) +CUDA_DIR = os.path.join(PROJECT_DIR, "ex_engine", "xllm_kernels", "cuda") +HEADER_DIR = os.path.join(CUDA_DIR, "headers") +MOE_DIR = os.path.join(CUDA_DIR, "moe") + +results = [] + +def report(name, status, detail=""): + sym = "✓" if status == "PASS" else "✗" if status == "FAIL" else "⊘" + results.append((name, status, detail)) + print(f" {sym} {name}: {status} {detail}") + + +def try_compile_cu(name, cu_file, extra_sources=None, extra_include=None): + """Try to compile a .cu file using torch.utils.cpp_extension.""" + try: + from torch.utils.cpp_extension import load + import glob + + sources = [cu_file] + if extra_sources: + sources.extend(extra_sources) + + extra_cflags = ["-O2", "-std=c++17"] + extra_cuda_cflags = [] + include_dirs = [HEADER_DIR] + if extra_include: + include_dirs.extend(extra_include) + + extra_ldflags = [] + try: + import ixformer + ixf_dir = os.path.dirname(ixformer.__file__) + for so in glob.glob(os.path.join(ixf_dir, "*.so")): + extra_ldflags.append(so) + extra_ldflags.append(f"-Wl,-rpath,{ixf_dir}") + except ImportError: + pass + + corex_lib = "/usr/local/corex/lib64" + if os.path.isdir(corex_lib): + extra_ldflags.append(f"-Wl,-rpath,{corex_lib}") + extra_ldflags.append(f"-L{corex_lib}") + include_dirs.append("/usr/local/corex/include") + + mod = load( + name=name, + sources=sources, + extra_cflags=extra_cflags, + extra_cuda_cflags=extra_cuda_cflags, + extra_ldflags=extra_ldflags, + extra_include_paths=include_dirs, + verbose=False, + ) + return mod + except Exception as e: + return str(e) + + +# ========================================================================= +# Test 1: activation.cu — silu_and_mul +# ========================================================================= +def test_activation(): + cu_file = os.path.join(CUDA_DIR, "activation.cu") + if not os.path.isfile(cu_file): + report("activation.cu", "SKIP", "file not found") + return + + # Test via ixformer.functions (already compiled in base image) + try: + import ixformer.functions as ixf_F + x = torch.randn(4, 256, dtype=torch.float16, device="cuda") + out = torch.empty(4, 128, dtype=torch.float16, device="cuda") + ixf_F.silu_and_mul(x, out) + + # Reference + gate, up = x.float().chunk(2, dim=-1) + ref = (torch.sigmoid(gate) * gate * up).half() # silu(gate) * up — wait, silu = x*sigmoid(x) + ref2 = (torch.nn.functional.silu(gate) * up).half() + + err = (out.float() - ref2.float()).abs().max().item() + report("activation.cu (silu_and_mul via ixf_F)", "PASS", f"max_err={err:.6f}") + except Exception as e: + report("activation.cu (silu_and_mul via ixf_F)", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Test 2: norm.cu — rms_norm, fused_add_rms_norm +# ========================================================================= +def test_norm(): + try: + import ixformer.functions as ixf_F + + hidden = 2048 + eps = 1e-6 + + # rms_norm + x = torch.randn(4, hidden, dtype=torch.float16, device="cuda") + w = torch.ones(hidden, dtype=torch.float16, device="cuda") + out = torch.empty_like(x) + ixf_F.rms_norm(x, w, out, eps) + + # Reference + x_f = x.float() + rms = torch.sqrt(x_f.pow(2).mean(-1, keepdim=True) + eps) + ref = (x_f / rms).half() + err = (out.float() - ref.float()).abs().max().item() + report("norm.cu (rms_norm via ixf_F)", "PASS", f"max_err={err:.6f}") + + # fused_add_rms_norm + inp = torch.randn(4, hidden, dtype=torch.float16, device="cuda") + res = torch.randn(4, hidden, dtype=torch.float16, device="cuda") + res_orig = res.clone() + ixf_F.fused_add_rms_norm(inp, res, w, eps) + # After: inp = rms_norm(inp + res_orig), res = inp + res_orig + combined = (inp.float() + res_orig.float()) + rms2 = torch.sqrt(combined.pow(2).mean(-1, keepdim=True) + eps) + # inp should now be normalized + report("norm.cu (fused_add_rms_norm via ixf_F)", "PASS", "ran without error") + except Exception as e: + report("norm.cu", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Test 3: rope.cu — rotary_embedding +# ========================================================================= +def test_rope(): + try: + import ixformer.functions as ixf_F + + head_size = 256 + rotary_dim = 64 # partial_rotary_factor=0.25 + max_pos = 1024 + num_heads = 6 + seq_len = 8 + + # Build cos_sin_cache + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim)) + t = torch.arange(max_pos, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1).cuda() + + positions = torch.arange(seq_len, dtype=torch.long, device="cuda") + q = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda") + k = torch.randn(seq_len, num_heads * head_size, dtype=torch.float16, device="cuda") + + q_orig = q.clone() + ixf_F.vllm_rotary_embedding_neox(positions, q, k, head_size, cos_sin_cache, True) + + # Verify something changed in the rotary dims + diff = (q.float() - q_orig.float()).abs().sum().item() + report("rope.cu (rotary_embedding via ixf_F)", "PASS", f"q_diff={diff:.2f}") + except Exception as e: + report("rope.cu", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Test 4: MoE topk_softmax +# ========================================================================= +def test_moe_topk(): + try: + # Try our prebuilt corex_moe_topk_softmax.so + sys.path.insert(0, os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10")) + try: + from vllm import corex_moe_topk_softmax + mod = corex_moe_topk_softmax + except ImportError: + import importlib.util + so_path = os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", + "corex_moe_topk_softmax.so") + if os.path.isfile(so_path): + spec = importlib.util.spec_from_file_location("corex_moe_topk_softmax", so_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + else: + report("moe_topk_softmax", "SKIP", "no .so found") + return + + num_tokens = 8 + num_experts = 256 + top_k = 8 + + gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device="cuda") + w, ids = mod.moe_topk_softmax(gating, top_k, True) + + # Reference + topk_logits, topk_ids_ref = torch.topk(gating, top_k, dim=-1) + topk_w_ref = torch.softmax(topk_logits, dim=-1) + topk_w_ref = topk_w_ref / topk_w_ref.sum(-1, keepdim=True) + + # Check shapes + assert w.shape == (num_tokens, top_k), f"weight shape {w.shape}" + assert ids.shape == (num_tokens, top_k), f"ids shape {ids.shape}" + + # Check weights sum to ~1 + w_sum = w.sum(-1) + w_sum_err = (w_sum - 1.0).abs().max().item() + report("moe_topk_softmax", "PASS", f"shape OK, weight_sum_err={w_sum_err:.6f}") + except Exception as e: + report("moe_topk_softmax", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Test 5: ix_moe_bridge — full fused MoE pipeline +# ========================================================================= +def test_ix_moe_bridge(): + # Try loading the bridge + so_paths = [ + os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_moe_bridge.so"), + os.path.join(SCRIPT_DIR, "ix_moe_bridge.so"), + ] + bridge = None + for p in so_paths: + if os.path.isfile(p): + try: + import importlib.util + spec = importlib.util.spec_from_file_location("ix_moe_bridge", p) + bridge = importlib.util.module_from_spec(spec) + spec.loader.exec_module(bridge) + break + except Exception: + pass + + if bridge is None: + report("ix_moe_bridge", "SKIP", "no prebuilt .so — run build_ix_moe_bridge.sh first") + return + + fns = [x for x in dir(bridge) if not x.startswith("_")] + report("ix_moe_bridge (load)", "PASS", f"functions: {fns}") + + # Test topk_softmax + try: + gating = torch.randn(4, 256, dtype=torch.float32, device="cuda") + w, ids = bridge.topk_softmax(gating, 8, True) + assert w.shape == (4, 8) + report("ix_moe_bridge.topk_softmax", "PASS", f"shape={w.shape}") + except Exception as e: + report("ix_moe_bridge.topk_softmax", "FAIL", str(e)[:120]) + + # Test moe_gen_idx + try: + expert_ids = torch.randint(0, 256, (32,), dtype=torch.int32, device="cuda") + results_list = bridge.moe_gen_idx(expert_ids, 256) + assert len(results_list) == 4 + report("ix_moe_bridge.moe_gen_idx", "PASS", f"got {len(results_list)} tensors") + except Exception as e: + report("ix_moe_bridge.moe_gen_idx", "FAIL", str(e)[:120]) + + # Test fused_moe_forward (full pipeline) + try: + T, H, E, I = 4, 2048, 256, 128 # TP-sharded: I = moe_intermediate_size / tp_size + hidden = torch.randn(T, H, dtype=torch.float16, device="cuda") + logits = torch.randn(T, E, dtype=torch.float32, device="cuda") + w13 = torch.randn(E, 2*I, H, dtype=torch.float16, device="cuda") * 0.01 + w2 = torch.randn(E, H, I, dtype=torch.float16, device="cuda") * 0.01 + out = bridge.fused_moe_forward(hidden, logits, w13, w2, 8, E, True) + assert out.shape == (T, H), f"output shape {out.shape}" + nan_count = torch.isnan(out).sum().item() + report("ix_moe_bridge.fused_moe_forward", "PASS", + f"shape={out.shape}, nans={nan_count}") + except Exception as e: + report("ix_moe_bridge.fused_moe_forward", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Test 6: ix_attn_bridge — attention functions +# ========================================================================= +def test_ix_attn_bridge(): + so_paths = [ + os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_attn_bridge.so"), + os.path.join(SCRIPT_DIR, "ix_attn_bridge.so"), + ] + bridge = None + for p in so_paths: + if os.path.isfile(p): + try: + import importlib.util + spec = importlib.util.spec_from_file_location("ix_attn_bridge", p) + bridge = importlib.util.module_from_spec(spec) + spec.loader.exec_module(bridge) + break + except Exception: + pass + + if bridge is None: + report("ix_attn_bridge", "SKIP", "no prebuilt .so — run build_ix_attn_bridge.sh first") + return + + fns = [x for x in dir(bridge) if not x.startswith("_")] + report("ix_attn_bridge (load)", "PASS", f"functions: {fns}") + + +# ========================================================================= +# Test 7: ix_full_bridge — basic ops bridge +# ========================================================================= +def test_ix_full_bridge(): + so_paths = [ + os.path.join(SCRIPT_DIR, "prebuilt", "corex-3.2.3-ivcore10", "ix_full_bridge.so"), + ] + bridge = None + for p in so_paths: + if os.path.isfile(p): + try: + import importlib.util + spec = importlib.util.spec_from_file_location("ix_full_bridge", p) + bridge = importlib.util.module_from_spec(spec) + spec.loader.exec_module(bridge) + break + except Exception: + pass + + if bridge is None: + report("ix_full_bridge", "SKIP", "no prebuilt .so") + return + + fns = [x for x in dir(bridge) if not x.startswith("_")] + report("ix_full_bridge (load)", "PASS", f"functions: {fns}") + + # Test silu_and_mul + try: + x = torch.randn(4, 256, dtype=torch.float16, device="cuda") + out = torch.empty(4, 128, dtype=torch.float16, device="cuda") + bridge.silu_and_mul(x, out) + report("ix_full_bridge.silu_and_mul", "PASS", f"shape={out.shape}") + except Exception as e: + report("ix_full_bridge.silu_and_mul", "FAIL", str(e)[:120]) + + # Test rms_norm + try: + x = torch.randn(4, 2048, dtype=torch.float16, device="cuda") + w = torch.ones(2048, dtype=torch.float16, device="cuda") + out = torch.empty_like(x) + bridge.rms_norm(out, x, w, 1e-6) + report("ix_full_bridge.rms_norm", "PASS", f"shape={out.shape}") + except Exception as e: + report("ix_full_bridge.rms_norm", "FAIL", str(e)[:120]) + + +# ========================================================================= +# Main +# ========================================================================= +if __name__ == "__main__": + print("=" * 60) + print(" xllm CUDA kernel verification on BI-V100") + print("=" * 60) + print() + + if not torch.cuda.is_available(): + print("ERROR: CUDA not available") + sys.exit(1) + + dev = torch.cuda.get_device_name(0) + print(f"GPU: {dev}") + print(f"CUDA kernels: {CUDA_DIR}") + print(f"MOE kernels: {MOE_DIR}") + print() + + t0 = time.time() + + print("[1/7] activation (silu_and_mul)") + test_activation() + + print("[2/7] norm (rms_norm, fused_add_rms_norm)") + test_norm() + + print("[3/7] rope (rotary_embedding)") + test_rope() + + print("[4/7] MoE topk_softmax") + test_moe_topk() + + print("[5/7] ix_moe_bridge (full fused MoE)") + test_ix_moe_bridge() + + print("[6/7] ix_attn_bridge (attention)") + test_ix_attn_bridge() + + print("[7/7] ix_full_bridge (basic ops)") + test_ix_full_bridge() + + elapsed = time.time() - t0 + print() + print("=" * 60) + passed = sum(1 for _, s, _ in results if s == "PASS") + failed = sum(1 for _, s, _ in results if s == "FAIL") + skipped = sum(1 for _, s, _ in results if s == "SKIP") + print(f" {passed} PASS {failed} FAIL {skipped} SKIP ({elapsed:.1f}s)") + print("=" * 60) + + if failed > 0: + print("\nFAILED tests:") + for name, s, detail in results: + if s == "FAIL": + print(f" ✗ {name}: {detail}") + sys.exit(1)