Skip to content

[libcu++] Use CCCL Runtime in CUB metric pre-runs pt 1 #9802

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pciolkosz wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[libcu++] Use CCCL Runtime in CUB metric pre-runs pt 1 #9802

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
13e2c3f
Use CCCL Runtime in CUB metric pre-runs
pciolkosz Jul 10, 2026
86b7cf4
Preserve benchmark temp-storage cache
pciolkosz Jul 14, 2026
68fac22
Address benchmark runtime feedback
pciolkosz Jul 15, 2026
4ef8c07
Merge branch 'main' into adopt-cccl-runtime-2800-set6
pciolkosz Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 23 additions & 15 deletions cub/benchmarks/bench/reduce/by_key.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,11 @@

#include <cub/device/device_reduce.cuh>

#include <cuda/buffer>
#include <cuda/std/execution>
#include <cuda/std/functional>
#include <cuda/stream>

#include <look_back_helper.cuh>
#include <nvbench_helper.cuh>

Expand Down Expand Up @@ -40,19 +45,22 @@ static void reduce_by_key(nvbench::state& state, nvbench::type_list<KeyT, ValueT
constexpr std::size_t min_segment_size = 1;
const std::size_t max_segment_size = static_cast<std::size_t>(state.get_int64("MaxSegSize"));

thrust::device_vector<OffsetT> num_runs_out(1);
thrust::device_vector<ValueT> in_vals(elements);
thrust::device_vector<ValueT> out_vals(elements);
thrust::device_vector<KeyT> out_keys(elements);
thrust::device_vector<KeyT> in_keys = generate.uniform.key_segments(elements, min_segment_size, max_segment_size);
const auto stream = get_stream_ref(state);
const auto device = stream.device();
caching_allocator_t alloc;

const KeyT* d_in_keys = thrust::raw_pointer_cast(in_keys.data());
KeyT* d_out_keys = thrust::raw_pointer_cast(out_keys.data());
const ValueT* d_in_vals = thrust::raw_pointer_cast(in_vals.data());
ValueT* d_out_vals = thrust::raw_pointer_cast(out_vals.data());
OffsetT* d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data());
auto num_runs_out = cuda::make_buffer<OffsetT>(stream, pinned_memory_resource(), 1, cuda::no_init);
const auto in_vals = cuda::make_device_buffer<ValueT>(stream, device, elements, ValueT{});
auto out_vals = cuda::make_device_buffer<ValueT>(stream, device, elements, cuda::no_init);
auto out_keys = cuda::make_device_buffer<KeyT>(stream, device, elements, cuda::no_init);
const auto in_keys =
generate.uniform.key_segments(elements, min_segment_size, max_segment_size).device_buffer<KeyT>(stream, device);

caching_allocator_t alloc;
const KeyT* d_in_keys = in_keys.data();
KeyT* d_out_keys = out_keys.data();
const ValueT* d_in_vals = in_vals.data();
ValueT* d_out_vals = out_vals.data();
OffsetT* d_num_runs_out = num_runs_out.data();

// Run once to get the number of runs for reporting
_CCCL_TRY_CUDA_API(
Expand All @@ -65,8 +73,8 @@ static void reduce_by_key(nvbench::state& state, nvbench::type_list<KeyT, ValueT
d_num_runs_out,
reduction_op_t{},
static_cast<OffsetT>(elements),
alloc);
cudaDeviceSynchronize();
cub_bench_env(alloc, stream));
stream.sync();
const OffsetT num_runs = num_runs_out[0];

state.add_element_count(elements);
Expand All @@ -79,9 +87,9 @@ static void reduce_by_key(nvbench::state& state, nvbench::type_list<KeyT, ValueT
state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
auto env = cub_bench_env(
alloc,
launch
get_stream_ref(launch)
#if !TUNE_BASE
,
,
cuda::execution::tune(bench_reduce_by_key_policy_selector{})
#endif // !TUNE_BASE
);
Expand Down
43 changes: 29 additions & 14 deletions cub/benchmarks/bench/run_length_encode/encode.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@

#include <cub/device/device_run_length_encode.cuh>

#include <cuda/buffer>
#include <cuda/std/execution>
#include <cuda/stream>

#include <look_back_helper.cuh>
#include <nvbench_helper.cuh>

Expand Down Expand Up @@ -42,20 +46,32 @@ static void rle(nvbench::state& state, nvbench::type_list<T, OffsetT, RunLengthT
constexpr std::size_t min_segment_size = 1;
const std::size_t max_segment_size = static_cast<std::size_t>(state.get_int64("MaxSegSize"));

thrust::device_vector<offset_t> num_runs_out(1);
thrust::device_vector<RunLengthT> out_counts(elements);
thrust::device_vector<T> out_keys(elements);
thrust::device_vector<T> in_keys = generate.uniform.key_segments(elements, min_segment_size, max_segment_size);
const auto stream = get_stream_ref(state);
const auto device = stream.device();
caching_allocator_t alloc;

auto num_runs_out = cuda::make_buffer<offset_t>(stream, pinned_memory_resource(), 1, cuda::no_init);
auto out_counts = cuda::make_device_buffer<RunLengthT>(stream, device, elements, cuda::no_init);
Comment on lines +53 to +54

@bernhardmgruber bernhardmgruber Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Q: Why do we use two different allocation functions to get a device buffer? Why does num_runs_out need to have a different memory resource?

@pciolkosz pciolkosz Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Its read on host later, so it need to be pinned memory

auto out_keys = cuda::make_device_buffer<T>(stream, device, elements, cuda::no_init);
const auto in_keys =
generate.uniform.key_segments(elements, min_segment_size, max_segment_size).device_buffer<T>(stream, device);

const T* d_in_keys = thrust::raw_pointer_cast(in_keys.data());
T* d_out_keys = thrust::raw_pointer_cast(out_keys.data());
RunLengthT* d_out_counts = thrust::raw_pointer_cast(out_counts.data());
offset_t* d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data());
const T* d_in_keys = in_keys.data();
T* d_out_keys = out_keys.data();
RunLengthT* d_out_counts = out_counts.data();
offset_t* d_num_runs_out = num_runs_out.data();

// Run once to get num_runs for memory accounting
(void) cub::DeviceRunLengthEncode::Encode(
d_in_keys, d_out_keys, d_out_counts, d_num_runs_out, static_cast<OffsetT>(elements));
cudaDeviceSynchronize();
_CCCL_TRY_CUDA_API(
cub::DeviceRunLengthEncode::Encode,
"Encode failed",
d_in_keys,
d_out_keys,
d_out_counts,
d_num_runs_out,
static_cast<OffsetT>(elements),
cub_bench_env(alloc, stream));
stream.sync();
Comment thread
coderabbitai[bot] marked this conversation as resolved.
const offset_t num_runs = num_runs_out[0];

state.add_element_count(elements);
Expand All @@ -64,13 +80,12 @@ static void rle(nvbench::state& state, nvbench::type_list<T, OffsetT, RunLengthT
state.add_global_memory_writes<RunLengthT>(num_runs);
state.add_global_memory_writes<offset_t>(1);

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
get_stream_ref(launch)
#if !TUNE_BASE
,
,
cuda::execution::tune(bench_encode_policy_selector{})
#endif // !TUNE_BASE
);
Comment on lines 84 to 91

@bernhardmgruber bernhardmgruber Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Q: Is memory_resource caching allocations? So if I allocate 1234 bytes, then free 1234 bytes, then allocate again. The second allocation will not have ANY overhead?

In other words, we want the first warmup run to do all allocations, and all following runs not have any allocation overhead on the host.

Q: Should we maybe run and compare a benchmark to be sure?

@pciolkosz pciolkosz Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There was a small ~2% regression, probably from the overhead to call to the driver, the memory itself was cached in the mempool.

I reverted back to the cached allocator

Expand Down
45 changes: 27 additions & 18 deletions cub/benchmarks/bench/run_length_encode/non_trivial_runs.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@

#include <cub/device/device_run_length_encode.cuh>

#include <cuda/buffer>
#include <cuda/std/execution>
#include <cuda/stream>

#include <look_back_helper.cuh>
#include <nvbench_helper.cuh>

Expand Down Expand Up @@ -43,23 +47,23 @@ static void rle(nvbench::state& state, nvbench::type_list<T, OffsetT, RunLengthT
constexpr std::size_t min_segment_size = 1;
const std::size_t max_segment_size = static_cast<std::size_t>(state.get_int64("MaxSegSize"));

thrust::device_vector<offset_t> num_runs_out(1);
thrust::device_vector<offset_t> out_offsets(elements);
thrust::device_vector<RunLengthT> out_lengths(elements);
thrust::device_vector<T> in_keys = generate.uniform.key_segments(elements, min_segment_size, max_segment_size);
const auto stream = get_stream_ref(state);
const auto device = stream.device();
caching_allocator_t alloc;

auto num_runs_out = cuda::make_buffer<offset_t>(stream, pinned_memory_resource(), 1, cuda::no_init);
auto out_offsets = cuda::make_device_buffer<offset_t>(stream, device, elements, cuda::no_init);
auto out_lengths = cuda::make_device_buffer<RunLengthT>(stream, device, elements, cuda::no_init);
const auto in_keys =
generate.uniform.key_segments(elements, min_segment_size, max_segment_size).device_buffer<T>(stream, device);

const T* d_in_keys = thrust::raw_pointer_cast(in_keys.data());
offset_t* d_out_offsets = thrust::raw_pointer_cast(out_offsets.data());
RunLengthT* d_out_lengths = thrust::raw_pointer_cast(out_lengths.data());
offset_t* d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data());
const T* d_in_keys = in_keys.data();
offset_t* d_out_offsets = out_offsets.data();
RunLengthT* d_out_lengths = out_lengths.data();
offset_t* d_num_runs_out = num_runs_out.data();

{
// Run once to get num_runs for memory accounting
auto memory_env = cuda::std::execution::env{
#if !TUNE_BASE
cuda::execution::tune(bench_rle_policy_selector{})
#endif // !TUNE_BASE
};
_CCCL_TRY_CUDA_API(
cub::DeviceRunLengthEncode::NonTrivialRuns,
"NonTrivialRuns failed",
Expand All @@ -68,8 +72,14 @@ static void rle(nvbench::state& state, nvbench::type_list<T, OffsetT, RunLengthT
d_out_lengths,
d_num_runs_out,
static_cast<OffsetT>(elements),
memory_env);
cudaDeviceSynchronize();
cub_bench_env(alloc,
stream
#if !TUNE_BASE
,
cuda::execution::tune(bench_rle_policy_selector{})
#endif // !TUNE_BASE
));
stream.sync();
}
const OffsetT num_runs = num_runs_out[0];

Expand All @@ -79,13 +89,12 @@ static void rle(nvbench::state& state, nvbench::type_list<T, OffsetT, RunLengthT
state.add_global_memory_writes<OffsetT>(num_runs);
state.add_global_memory_writes<OffsetT>(1);

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
get_stream_ref(launch)
#if !TUNE_BASE
,
,
cuda::execution::tune(bench_rle_policy_selector{})
#endif // !TUNE_BASE
);
Expand Down
32 changes: 20 additions & 12 deletions cub/benchmarks/bench/select/unique.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@

#include <cub/device/device_select.cuh>

#include <cuda/std/algorithm>
#include <cuda/buffer>
#include <cuda/std/execution>
#include <cuda/std/functional>
#include <cuda/stream>

#include <limits>

Expand Down Expand Up @@ -43,13 +46,18 @@ static void unique(nvbench::state& state, nvbench::type_list<T, InPlace>)
const auto elements = state.get_int64("Elements{io}");
const auto max_segment_size = state.get_int64("MaxSegSize");

thrust::device_vector<T> in = generate.uniform.key_segments(elements, /* min_segmented_size */ 1, max_segment_size);
thrust::device_vector<T> out(elements, thrust::no_init);
thrust::device_vector<offset_t> num_unique_out(1);
const auto stream = get_stream_ref(state);
const auto device = stream.device();
caching_allocator_t alloc;

auto in = generate.uniform.key_segments(elements, /* min_segmented_size */ 1, max_segment_size)
.device_buffer<T>(stream, device);
auto out = cuda::make_device_buffer<T>(stream, device, elements, cuda::no_init);
auto num_unique_out = cuda::make_buffer<offset_t>(stream, pinned_memory_resource(), 1, cuda::no_init);

T* d_in = thrust::raw_pointer_cast(in.data());
T* d_out = thrust::raw_pointer_cast(out.data());
offset_t* d_num_unique = thrust::raw_pointer_cast(num_unique_out.data());
T* d_in = in.data();
T* d_out = out.data();
offset_t* d_num_unique = num_unique_out.data();

// Get number of unique elements for metrics
_CCCL_TRY_CUDA_API(
Expand All @@ -59,22 +67,22 @@ static void unique(nvbench::state& state, nvbench::type_list<T, InPlace>)
d_out,
d_num_unique,
static_cast<offset_t>(elements),
::cuda::std::equal_to<>{});
cudaDeviceSynchronize();
::cuda::std::equal_to<>{},
cub_bench_env(alloc, stream));
stream.sync();
const offset_t num_unique = num_unique_out[0];

state.add_element_count(elements);
state.add_global_memory_reads<T>(elements);
state.add_global_memory_writes<T>(num_unique);
state.add_global_memory_writes<offset_t>(1);

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
get_stream_ref(launch)
#if !TUNE_BASE
,
,
cuda::execution::tune(bench_policy_selector<T>{})
#endif // !TUNE_BASE
);
Expand Down
40 changes: 25 additions & 15 deletions cub/benchmarks/bench/select/unique_by_key.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,11 @@

#include <cub/device/device_select.cuh>

#include <cuda/buffer>
#include <cuda/std/execution>
#include <cuda/std/functional>
#include <cuda/stream>

#include <look_back_helper.cuh>
#include <nvbench_helper.cuh>

Expand Down Expand Up @@ -50,17 +55,22 @@ static void select(nvbench::state& state, nvbench::type_list<KeyT, ValueT, Offse
constexpr std::size_t min_segment_size = 1;
const std::size_t max_segment_size = static_cast<std::size_t>(state.get_int64("MaxSegSize"));

thrust::device_vector<OffsetT> num_runs_out(1);
thrust::device_vector<ValueT> in_vals(elements);
thrust::device_vector<ValueT> out_vals(elements);
thrust::device_vector<KeyT> out_keys(elements);
thrust::device_vector<KeyT> in_keys = generate.uniform.key_segments(elements, min_segment_size, max_segment_size);
const auto stream = get_stream_ref(state);
const auto device = stream.device();
caching_allocator_t alloc;

auto num_runs_out = cuda::make_buffer<OffsetT>(stream, pinned_memory_resource(), 1, cuda::no_init);
const auto in_vals = cuda::make_device_buffer<ValueT>(stream, device, elements, ValueT{});
auto out_vals = cuda::make_device_buffer<ValueT>(stream, device, elements, cuda::no_init);
auto out_keys = cuda::make_device_buffer<KeyT>(stream, device, elements, cuda::no_init);
const auto in_keys =
generate.uniform.key_segments(elements, min_segment_size, max_segment_size).device_buffer<KeyT>(stream, device);

const KeyT* d_in_keys = thrust::raw_pointer_cast(in_keys.data());
KeyT* d_out_keys = thrust::raw_pointer_cast(out_keys.data());
const ValueT* d_in_vals = thrust::raw_pointer_cast(in_vals.data());
ValueT* d_out_vals = thrust::raw_pointer_cast(out_vals.data());
OffsetT* d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data());
const KeyT* d_in_keys = in_keys.data();
KeyT* d_out_keys = out_keys.data();
const ValueT* d_in_vals = in_vals.data();
ValueT* d_out_vals = out_vals.data();
OffsetT* d_num_runs_out = num_runs_out.data();

const auto num_items = static_cast<OffsetT>(elements);

Expand All @@ -74,8 +84,9 @@ static void select(nvbench::state& state, nvbench::type_list<KeyT, ValueT, Offse
d_out_vals,
d_num_runs_out,
num_items,
equality_op_t{});
_CCCL_TRY_CUDA_API(cudaDeviceSynchronize, "Sync failed");
equality_op_t{},
cub_bench_env(alloc, stream));
stream.sync();
const OffsetT num_runs = num_runs_out[0];

state.add_element_count(elements);
Expand All @@ -85,13 +96,12 @@ static void select(nvbench::state& state, nvbench::type_list<KeyT, ValueT, Offse
state.add_global_memory_writes<KeyT>(num_runs);
state.add_global_memory_writes<OffsetT>(1);

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
get_stream_ref(launch)
#if !TUNE_BASE
,
,
cuda::execution::tune(bench_unique_by_key_policy_selector{})
#endif // !TUNE_BASE
);
Expand Down
Loading
Loading