dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

Browse Source 
CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0
This commit is contained in:
EngineX CI
2026-07-30 09:35:51 +00:00
parent b4d01f481e
commit 56fd68e7dd
8871 changed files with 1454674 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
cccl_upstream/cub/benchmarks/bench/segmented_topk/fixed/keys.cu Normal file
View File

@@ -0,0 +1,122 @@
// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include <cub/detail/choose_offset.cuh>
#include <cub/device/dispatch/dispatch_batched_topk.cuh>
#include <cuda/argument>
#include <cuda/iterator>
#include <nvbench_helper.cuh>
// %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1
// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
// %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1
#if !TUNE_BASE
struct tuned_policy_selector
{
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-> cub::detail::batched_topk::batched_topk_policy
{
// Single-entry policy chain driven by the tuning knobs.
constexpr auto store_alg = cub::BLOCK_STORE_WARP_TRANSPOSE;
# if TUNE_BLOCK_LOAD_ALGORITHM == 0
constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT;
# elif TUNE_BLOCK_LOAD_ALGORITHM == 1
constexpr auto load_alg = cub::BLOCK_LOAD_WARP_TRANSPOSE;
# elif TUNE_BLOCK_LOAD_ALGORITHM == 2
constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE;
# endif
return cub::detail::batched_topk::batched_topk_policy{{{
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
cub::detail::batched_topk::worker_policy{TUNE_THREADS_PER_BLOCK, TUNE_ITEMS_PER_THREAD, load_alg, store_alg},
}}};
}
};
#endif // !TUNE_BASE
template <typename KeyT, int MaxSegmentSize, int MaxNumSelected>
void fixed_seg_size_topk_keys(
nvbench::state& state,
nvbench::type_list<KeyT, nvbench::enum_type<MaxSegmentSize>, nvbench::enum_type<MaxNumSelected>>)
{
// Retrieve axis parameters
const auto max_elements = static_cast<size_t>(state.get_int64("Elements{io}"));
const auto segment_size = static_cast<::cuda::std::ptrdiff_t>(MaxSegmentSize);
const auto selected_elements = static_cast<::cuda::std::ptrdiff_t>(MaxNumSelected);
const auto num_segments = ::cuda::std::max<std::size_t>(1, (max_elements / segment_size));
const auto elements = num_segments * segment_size;
const auto total_num_items = ::cuda::args::immediate{static_cast<::cuda::std::int64_t>(elements)};
const bit_entropy entropy = str_to_entropy(state.get_string("Entropy"));
// Skip workloads where k exceeds the segment size
if (selected_elements >= segment_size)
{
state.skip("Skipping workload where K >= SegmentSize.");
return;
}
thrust::device_vector<KeyT> in_keys_buffer = generate(elements, entropy);
thrust::device_vector<KeyT> out_keys_buffer(selected_elements * num_segments, thrust::no_init);
auto d_keys_in_ptr = thrust::raw_pointer_cast(in_keys_buffer.data());
auto d_keys_out_ptr = thrust::raw_pointer_cast(out_keys_buffer.data());
auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size);
auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), selected_elements);
auto segment_sizes = ::cuda::args::constant<MaxSegmentSize>{};
auto k = ::cuda::args::constant<MaxNumSelected>{};
auto select_direction = ::cuda::args::constant<cub::detail::topk::select::max>{};
state.add_element_count(elements, "NumElements");
state.add_element_count(segment_size, "SegmentSize");
state.add_element_count(selected_elements, "NumSelectedElements");
state.add_global_memory_reads<KeyT>(elements, "InputKeys");
state.add_global_memory_writes<KeyT>(selected_elements * num_segments, "OutputKeys");
caching_allocator_t alloc;
state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
auto env = cub_bench_env(
alloc,
launch
#if !TUNE_BASE
,
cuda::execution::tune(tuned_policy_selector{})
#endif // !TUNE_BASE
);
// TODO(bgruber): call the public API once available
_CCCL_TRY_CUDA_API(
cub::detail::batched_topk::dispatch_with_env,
"batched topk failed",
d_keys_in,
d_keys_out,
static_cast<cub::NullType**>(nullptr),
static_cast<cub::NullType**>(nullptr),
segment_sizes,
k,
select_direction,
::cuda::args::immediate{static_cast<::cuda::std::int64_t>(num_segments)},
total_num_items,
env);
});
}
using key_type_list = nvbench::type_list<float>;
using segment_size_type_list = nvbench::type_list<uint32_t>;
using out_offset_type_list = nvbench::type_list<uint32_t>;
using segment_size_ = nvbench::type_list<uint32_t>;
using out_offset_type_list = nvbench::type_list<uint32_t>;
using small_segment_size_list = nvbench::enum_type_list<64, 128, 256, 512, 1024>;
using small_k_list = nvbench::enum_type_list<8, 16, 32, 128, 512, 1024>;
NVBENCH_BENCH_TYPES(fixed_seg_size_topk_keys, NVBENCH_TYPE_AXES(key_type_list, small_segment_size_list, small_k_list))
.set_name("small")
.set_type_axes_names({"KeyT{ct}", "MaxSegmentSize{ct}", "MaxNumSelected{ct}"})
.add_int64_power_of_two_axis("Elements{io}", nvbench::range(28, 28, 4))
.add_string_axis("Entropy", {"1.000", "0.544", "0.201", "0.000"});

175
cccl_upstream/cub/benchmarks/bench/segmented_topk/variable/common.cuh Normal file
View File

@@ -0,0 +1,175 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#pragma once
#include <thrust/device_vector.h>
#include <thrust/tabulate.h>
#include <cuda/random>
#include <cuda/std/algorithm>
#include <cuda/std/cmath>
#include <cuda/std/cstdint>
#include <cuda/std/random>
#include <cstddef>
#include <stdexcept>
#include <string>
#include <vector>
#include <nvbench_helper.cuh>
namespace
{
enum class pattern_kind : int
{
random = 0,
quantized_random,
relu_quantized,
tie_heavy,
pivot_tie
};
[[nodiscard]] pattern_kind string_to_pattern(const std::string& pattern)
{
if (pattern == "random")
{
return pattern_kind::random;
}
if (pattern == "quantized_random")
{
return pattern_kind::quantized_random;
}
if (pattern == "relu_quantized")
{
return pattern_kind::relu_quantized;
}
if (pattern == "tie_heavy")
{
return pattern_kind::tie_heavy;
}
if (pattern == "pivot_tie")
{
return pattern_kind::pivot_tie;
}
throw std::runtime_error("Invalid Pattern axis value: " + pattern);
}
template <int MaxSegmentSize, int K>
[[nodiscard]] thrust::device_vector<float>
gen_data(int num_segments, pattern_kind pattern, const cuda::std::int64_t* d_seg_sizes)
{
const auto num_keys = static_cast<std::size_t>(num_segments) * static_cast<std::size_t>(MaxSegmentSize);
auto d_keys = thrust::device_vector<float>{num_keys, thrust::no_init};
// gt_count == "greater-than count": number of 2.0 values placed at the tail of each segment's live region.
constexpr int gt_count = cuda::std::max(1, cuda::std::min(K / 4, MaxSegmentSize / 8));
thrust::tabulate(d_keys.begin(), d_keys.end(), [pattern, d_seg_sizes] __device__(std::size_t idx) -> float {
auto quantize = [](float base) -> float {
const auto r = cuda::std::rint(base);
const auto scaled_fr = cuda::std::rint((base - r) * 32.0f);
return r + (scaled_fr / 32.0f);
};
auto random_value = [](unsigned long long idx) -> float {
cuda::pcg64 rng(42);
rng.discard(idx);
cuda::std::normal_distribution<float> normal(0.f, 1.f);
return normal(rng);
};
const auto j = static_cast<int>(idx % MaxSegmentSize);
switch (pattern)
{
// ##
// ####
// ########
// ############
// ################
// ######################
// ##############################
// ------------------------------
// -3 0 3
case pattern_kind::random:
return random_value(idx);
// |
// |
// | | |
// | | | | |
// | | | | | | |
// ----------------------------
// -3 0 3
case pattern_kind::quantized_random:
return quantize(random_value(idx));
// |
// |
// |
// |
// |
// | |
// | | |
// | | | | |
// | | | | | | |
// ----------------------------
// 0 3
case pattern_kind::relu_quantized:
return quantize(cuda::std::max(random_value(idx), 0.f));
// | | | | | | | |
// | | | | | | | |
// | | | | | | | |
// --------------------------------
// 0/64 63/64
case pattern_kind::tie_heavy:
return static_cast<float>(j % 64) / 64.f;
// |
// |
// |
// |
// |
// |
// |
// | |
// ----------------------------
// 1.0 2.0
case pattern_kind::pivot_tie: {
const auto seg_size = static_cast<int>(d_seg_sizes[idx / MaxSegmentSize]);
return (j >= seg_size - gt_count) ? 2.f : 1.f;
}
default:
_CCCL_UNREACHABLE();
}
});
return d_keys;
}
} // namespace
const std::vector<std::string> valid_patterns = {
"random", "quantized_random", "relu_quantized", "tie_heavy", "pivot_tie"};
using key_type_list = nvbench::type_list<float>;
using max_segment_size_list = nvbench::enum_type_list< //
512,
1024,
2048,
4096,
8192
#if 0 // need these, waiting for implementation to catch up
,
16384,
32768,
65536,
131072,
262144,
524288,
1048576
#endif
>;
using k_list = nvbench::enum_type_list<512, 1024, 2048>;

99
cccl_upstream/cub/benchmarks/bench/segmented_topk/variable/indexed.cu Normal file
View File

@@ -0,0 +1,99 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include <cub/detail/choose_offset.cuh>
#include <cub/device/dispatch/dispatch_batched_topk.cuh>
#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <cuda/argument>
#include <cuda/iterator>
#include <cuda/std/cstdint>
#include <nvbench_helper.cuh>
#include "common.cuh"
// Indexed (arg-top-k) variant: each key carries a segment-local index as its value payload. The input values are
// produced by a counting iterator that restarts at 0 for every segment, so indices are not (pre-)materialized in global
// memory
template <typename KeyT, typename IndexT, int MaxSegmentSize, int K>
void decode_style_variable_topk_indexed(
nvbench::state& state, nvbench::type_list<KeyT, IndexT, nvbench::enum_type<MaxSegmentSize>, nvbench::enum_type<K>>)
{
if constexpr (K > MaxSegmentSize)
{
state.skip("K > MaxSegmentSize.");
return;
}
const auto num_segments = static_cast<int>(state.get_int64("NumSegments"));
const thrust::device_vector<cuda::std::int64_t> d_segment_sizes = generate(
static_cast<std::size_t>(num_segments),
bit_entropy::_1_000,
static_cast<cuda::std::int64_t>(K),
static_cast<cuda::std::int64_t>(MaxSegmentSize));
const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end());
const auto output_elements = static_cast<std::size_t>(num_segments) * K;
const auto total_num_items = cuda::args::immediate{static_cast<cuda::std::int64_t>(input_elements)};
auto in_keys_buffer = gen_data<MaxSegmentSize, K>(
num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data()));
auto out_keys_buffer = thrust::device_vector<KeyT>(output_elements, thrust::no_init);
auto out_indices_buffer = thrust::device_vector<IndexT>(output_elements, thrust::no_init);
auto segment_sizes_param = cuda::args::deferred_sequence{
thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()};
auto k_param = cuda::args::constant<K>{};
auto select_direction = cuda::args::constant<cub::detail::topk::select::max>{};
auto num_segments_param = cuda::args::immediate{static_cast<cuda::std::int64_t>(num_segments)};
auto d_keys_in = cuda::make_strided_iterator(
cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())),
static_cast<cuda::std::ptrdiff_t>(MaxSegmentSize));
auto d_keys_out = cuda::make_strided_iterator(
cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())),
static_cast<cuda::std::ptrdiff_t>(K));
// Input values: every segment maps to the same counting iterator starting at 0, so values are segment-local indices.
auto d_indices_in = cuda::make_constant_iterator(cuda::make_counting_iterator(IndexT{0}));
auto d_indices_out = cuda::make_strided_iterator(
cuda::make_counting_iterator(thrust::raw_pointer_cast(out_indices_buffer.data())),
static_cast<cuda::std::ptrdiff_t>(K));
state.add_element_count(input_elements, "NumElements");
state.add_global_memory_reads<KeyT>(input_elements, "InputKeys");
state.add_global_memory_reads<cuda::std::int64_t>(num_segments, "SegmentSizes");
state.add_global_memory_writes<KeyT>(output_elements, "OutputKeys");
state.add_global_memory_writes<IndexT>(output_elements, "OutputIndices");
caching_allocator_t alloc;
state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
auto env = cub_bench_env(alloc, launch);
// TODO(bgruber): call the public API once available
_CCCL_TRY_CUDA_API(
cub::detail::batched_topk::dispatch_with_env,
"batched topk failed",
d_keys_in,
d_keys_out,
d_indices_in,
d_indices_out,
segment_sizes_param,
k_param,
select_direction,
num_segments_param,
total_num_items,
env);
});
}
// Index type is a compile-time axis: i32 for now, extensible to i64.
using index_type_list = nvbench::type_list<cuda::std::int32_t>;
NVBENCH_BENCH_TYPES(decode_style_variable_topk_indexed,
NVBENCH_TYPE_AXES(key_type_list, index_type_list, max_segment_size_list, k_list))
.set_name("decode_style_variable_topk_indexed")
.set_type_axes_names({"KeyT{ct}", "IndexT{ct}", "MaxSegmentSize{ct}", "K{ct}"})
.add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32})
.add_string_axis("Pattern", valid_patterns);

83
cccl_upstream/cub/benchmarks/bench/segmented_topk/variable/keys.cu Normal file
View File

@@ -0,0 +1,83 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include <cub/detail/choose_offset.cuh>
#include <cub/device/dispatch/dispatch_batched_topk.cuh>
#include <thrust/device_vector.h>
#include <thrust/reduce.h>
#include <cuda/argument>
#include <cuda/iterator>
#include <nvbench_helper.cuh>
#include "common.cuh"
template <typename KeyT, int MaxSegmentSize, int K>
void decode_style_variable_topk_keys(
nvbench::state& state, nvbench::type_list<KeyT, nvbench::enum_type<MaxSegmentSize>, nvbench::enum_type<K>>)
{
if constexpr (K > MaxSegmentSize)
{
state.skip("K > MaxSegmentSize.");
return;
}
const auto num_segments = static_cast<int>(state.get_int64("NumSegments"));
const thrust::device_vector<cuda::std::int64_t> d_segment_sizes = generate(
static_cast<std::size_t>(num_segments),
bit_entropy::_1_000,
static_cast<cuda::std::int64_t>(K),
static_cast<cuda::std::int64_t>(MaxSegmentSize));
const auto input_elements = thrust::reduce(d_segment_sizes.begin(), d_segment_sizes.end());
const auto output_elements = static_cast<std::size_t>(num_segments) * K;
const auto total_num_items = cuda::args::immediate{static_cast<cuda::std::int64_t>(input_elements)};
auto in_keys_buffer = gen_data<MaxSegmentSize, K>(
num_segments, string_to_pattern(state.get_string("Pattern")), thrust::raw_pointer_cast(d_segment_sizes.data()));
auto out_keys_buffer = thrust::device_vector<KeyT>(output_elements, thrust::no_init);
auto segment_sizes_param = cuda::args::deferred_sequence{
thrust::raw_pointer_cast(d_segment_sizes.data()), cuda::args::bounds<1, MaxSegmentSize>()};
auto k_param = cuda::args::constant<K>{};
auto select_direction = cuda::args::constant<cub::detail::topk::select::max>{};
auto num_segments_param = cuda::args::immediate{static_cast<cuda::std::int64_t>(num_segments)};
auto d_keys_in = cuda::make_strided_iterator(
cuda::make_counting_iterator(thrust::raw_pointer_cast(in_keys_buffer.data())),
static_cast<cuda::std::ptrdiff_t>(MaxSegmentSize));
auto d_keys_out = cuda::make_strided_iterator(
cuda::make_counting_iterator(thrust::raw_pointer_cast(out_keys_buffer.data())),
static_cast<cuda::std::ptrdiff_t>(K));
state.add_element_count(input_elements, "NumElements");
state.add_global_memory_reads<KeyT>(input_elements, "InputKeys");
state.add_global_memory_reads<cuda::std::int64_t>(num_segments, "SegmentSizes");
state.add_global_memory_writes<KeyT>(output_elements, "OutputKeys");
caching_allocator_t alloc;
state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
auto env = cub_bench_env(alloc, launch);
// TODO(bgruber): call the public API once available
_CCCL_TRY_CUDA_API(
cub::detail::batched_topk::dispatch_with_env,
"batched topk failed",
d_keys_in,
d_keys_out,
static_cast<cub::NullType**>(nullptr),
static_cast<cub::NullType**>(nullptr),
segment_sizes_param,
k_param,
select_direction,
num_segments_param,
total_num_items,
env);
});
}
NVBENCH_BENCH_TYPES(decode_style_variable_topk_keys, NVBENCH_TYPE_AXES(key_type_list, max_segment_size_list, k_list))
.set_name("decode_style_variable_topk_keys")
.set_type_axes_names({"KeyT{ct}", "MaxSegmentSize{ct}", "K{ct}"})
.add_int64_axis("NumSegments", {1, 2, 4, 8, 16, 32})
.add_string_axis("Pattern", valid_patterns);