From e794e7578932be59bb34c67928720edb9485d677 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 17 Aug 2026 16:05:06 -0700 Subject: [PATCH 1/8] Use cuda::stream_ref in libcudf_streaming --- .../benchmarks/bench_pack.cpp | 27 ++++++++++--------- .../benchmarks/bench_partition.cpp | 10 +++---- .../benchmarks/bench_shuffle.cpp | 14 +++++----- .../streaming/bench_streaming_shuffle.cpp | 4 +-- .../benchmarks/streaming/data_generator.hpp | 2 +- .../benchmarks/streaming/ndsh/join.cpp | 4 +-- .../benchmarks/streaming/ndsh/q03.cpp | 2 +- .../benchmarks/streaming/ndsh/utils.hpp | 7 +++-- .../benchmarks/utils/random_data.cu | 8 +++--- .../benchmarks/utils/random_data.hpp | 6 ++--- .../examples/example_shuffle.cpp | 2 +- .../detail/approx_distinct_count.hpp | 6 ++--- .../detail/device_bloom_filter.hpp | 11 ++++---- .../include/cudf_streaming/parquet.hpp | 2 +- .../cudf_streaming/partition_utils.hpp | 8 +++--- .../include/cudf_streaming/table_chunk.hpp | 10 +++---- .../include/cudf_streaming/utils.hpp | 10 +++---- .../src/approx_distinct_count.cpp | 7 +++-- cpp/libcudf_streaming/src/bloom_filter.cpp | 2 +- .../src/detail/approx_distinct_count.cu | 8 +++--- .../src/detail/device_bloom_filter.cu | 10 +++---- cpp/libcudf_streaming/src/parquet.cpp | 4 +-- cpp/libcudf_streaming/src/partition_utils.cpp | 13 ++++----- cpp/libcudf_streaming/src/table_chunk.cpp | 6 ++--- cpp/libcudf_streaming/src/utils.cpp | 26 ++++++++---------- .../streaming/base_streaming_fixture.hpp | 2 +- .../tests/streaming/test_channel_metadata.cpp | 5 ++-- .../tests/streaming/test_cudf_utils.cpp | 2 +- .../tests/streaming/test_read_parquet.cpp | 9 ++++--- .../tests/streaming/test_table_chunk.cpp | 21 ++++++++------- cpp/libcudf_streaming/tests/test_shuffler.cpp | 6 ++--- cpp/libcudf_streaming/tests/utils.hpp | 12 ++++----- .../cudf_streaming/channel_metadata.pxd | 1 - .../cudf_streaming/channel_metadata.pyx | 2 +- .../cudf_streaming/partition_utils.pyx | 20 +++++++------- .../cudf_streaming/stream_ref.pxd | 11 ++++++++ .../cudf_streaming/table_chunk.pxd | 4 +-- .../cudf_streaming/table_chunk.pyx | 7 ++--- 38 files changed, 161 insertions(+), 150 deletions(-) create mode 100644 python/cudf_streaming/cudf_streaming/stream_ref.pxd diff --git a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp index 7cfbc0ab3a02..a685f231de90 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp @@ -11,11 +11,12 @@ #include #include -#include #include #include #include +#include + #include #include #include @@ -38,7 +39,7 @@ void run_pack(benchmark::State& state, std::size_t table_size_mb, rmm::device_async_resource_ref table_mr, rmm::device_async_resource_ref pack_mr, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { auto const table_size_bytes = table_size_mb * MB; @@ -49,12 +50,12 @@ void run_pack(benchmark::State& state, // Warm up auto warm_up = cudf::pack(table.view(), stream, pack_mr); - stream.synchronize(); + stream.sync(); for (auto _ : state) { auto packed = cudf::pack(table.view(), stream, pack_mr); benchmark::DoNotOptimize(packed); - stream.synchronize(); + stream.sync(); } state.SetBytesProcessed(static_cast(state.iterations()) * @@ -70,7 +71,7 @@ static void BM_Pack_device(benchmark::State& state) { auto const table_size_mb = static_cast(state.range(0)); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -91,7 +92,7 @@ static void BM_Pack_pinned(benchmark::State& state) auto const table_size_mb = static_cast(state.range(0)); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{ @@ -115,7 +116,7 @@ void run_chunked_pack(benchmark::State& state, std::size_t table_size, rmm::device_async_resource_ref table_mr, rmm::device_async_resource_ref pack_mr, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // Calculate number of rows for a single-column table of the desired size auto const nrows = rapidsmpf::safe_cast(table_size / sizeof(random_data_t)); @@ -150,13 +151,13 @@ void run_chunked_pack(benchmark::State& state, { run_packer(); - stream.synchronize(); + stream.sync(); } for (auto _ : state) { run_packer(); benchmark::DoNotOptimize(destination); - stream.synchronize(); + stream.sync(); } state.SetBytesProcessed(static_cast(state.iterations()) * @@ -178,7 +179,7 @@ static void BM_ChunkedPack_device(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40)}; @@ -203,7 +204,7 @@ static void BM_ChunkedPack_pinned(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) @@ -247,7 +248,7 @@ static void BM_ChunkedPack_fixed_table_device(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -271,7 +272,7 @@ static void BM_ChunkedPack_fixed_table_pinned(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) diff --git a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp index b4f348faabbe..73981ef8998f 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -11,11 +11,12 @@ #include -#include #include #include #include +#include + #include #include @@ -23,8 +24,7 @@ #include // Helper function to create a table with a single int column -std::unique_ptr create_int_table(cudf::size_type num_rows, - rmm::cuda_stream_view stream) +std::unique_ptr create_int_table(cudf::size_type num_rows, cuda::stream_ref stream) { auto data = rmm::device_buffer(rapidsmpf::safe_cast(num_rows) * sizeof(std::int32_t), stream); @@ -45,7 +45,7 @@ static void BM_PartitionAndPack(benchmark::State& state) int const num_partitions = state.range(1); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; // Get total GPU memory cudaDeviceProp prop; @@ -94,7 +94,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) int num_rows = int(local_size / std::int64_t{sizeof(std::int32_t)} / std::int64_t{num_partitions}); - rmm::cuda_stream_view stream = rmm::cuda_stream_default; + cuda::stream_ref stream = cuda::stream_ref{}; // Get total GPU memory cudaDeviceProp prop; diff --git a/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp index ddbc7c11a38a..de17023fec96 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp @@ -259,7 +259,7 @@ void barrier(std::shared_ptr& comm) rapidsmpf::Duration do_run(rapidsmpf::shuffler::PartID const total_num_partitions, std::shared_ptr& comm, ArgumentParser const& args, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, std::shared_ptr statistics, auto&& shuffle_insert_fn) @@ -296,7 +296,7 @@ rapidsmpf::Duration do_run(rapidsmpf::shuffler::PartID const total_num_partition output_partitions.emplace_back(std::move(output_partition)); } } - stream.synchronize(); + stream.sync(); } auto const elapsed = rapidsmpf::Clock::now() - t0_elapsed; @@ -330,7 +330,7 @@ template >> std::vector generate_input_partitions(ArgumentParser const& args, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, TransformFn&& transform_fn) { @@ -350,7 +350,7 @@ std::vector generate_input_partitions(ArgumentParser const& ar random_table(num_columns, num_local_rows, min_val, max_val, stream, br->device_mr()); input_partitions.emplace_back(transform_fn(std::move(table))); } - stream.synchronize(); + stream.sync(); return input_partitions; } @@ -396,7 +396,7 @@ void do_insert(rapidsmpf::shuffler::Shuffler& shuffler, */ rapidsmpf::Duration run_hash_partition_inline(std::shared_ptr& comm, ArgumentParser const& args, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, std::shared_ptr statistics) { @@ -437,7 +437,7 @@ rapidsmpf::Duration run_hash_partition_inline(std::shared_ptr& comm, ArgumentParser const& args, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, std::shared_ptr statistics) { @@ -565,7 +565,7 @@ int main(int argc, char** argv) args.pprint(*comm); - rmm::cuda_stream_view stream = cudf::get_default_stream(); + cuda::stream_ref stream = cudf::get_default_stream(); // Print benchmark/hardware info. { diff --git a/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp index 1d7f1604498e..54e5dfd6d758 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -215,7 +215,7 @@ rapidsmpf::streaming::Actor consumer(std::shared_ptr ctx, std::shared_ptr comm, ArgumentParser const& args, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { constexpr std::int32_t min_val = 0; constexpr std::int32_t max_val = 10; @@ -349,7 +349,7 @@ int main(int argc, char** argv) auto& stat_enabled_mr = br->device_mr_adaptor(); rmm::mr::set_current_device_resource(stat_enabled_mr); - rmm::cuda_stream_view stream = cudf::get_default_stream(); + cuda::stream_ref stream = cudf::get_default_stream(); // Print benchmark/hardware info. { diff --git a/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp b/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp index 7333e63934c7..d0c1619d96b5 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp @@ -45,7 +45,7 @@ using cudf_streaming::table_chunk; * and sent, and the channel has been drained. */ inline Actor random_table_generator(std::shared_ptr ctx, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, std::shared_ptr ch_out, std::uint64_t num_blocks, cudf::size_type ncolumns, diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index dcfc3b3ae2c6..acfa11ba80de 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -19,7 +19,7 @@ #include #include -#include +#include #include #include @@ -204,7 +204,7 @@ streaming::Message inner_join_chunk(std::shared_ptr ctx, cudf::hash_join& joiner, cudf::table_view build_carrier, std::vector right_on, - rmm::cuda_stream_view build_stream, + cuda::stream_ref build_stream, CudaEvent* build_event, CudaEvent* tmp_event diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp index 783014e8ce0e..663f7425d745 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp @@ -232,7 +232,7 @@ rapidsmpf::streaming::Actor top_k_by(std::shared_ptrexecutor()->schedule(); std::vector> partials; - std::vector chunk_streams; + std::vector chunk_streams; while (true) { auto msg = co_await ch_in->receive(); if (msg.empty()) { break; } diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp index b5e28e2fdc63..bc8876b4e03c 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp @@ -13,9 +13,8 @@ #include #include -#include - #include +#include #include #include @@ -95,7 +94,7 @@ namespace detail { * @return Filter expression with proper lifetime management */ template -std::unique_ptr make_date_filter(rmm::cuda_stream_view stream, +std::unique_ptr make_date_filter(cuda::stream_ref stream, cuda::std::chrono::year_month_day date, std::string const& column_name, cudf::ast::ast_operator op) @@ -135,7 +134,7 @@ std::unique_ptr make_date_filter(rmm::cuda_stream_view s */ template std::unique_ptr make_date_range_filter( - rmm::cuda_stream_view stream, + cuda::stream_ref stream, cuda::std::chrono::year_month_day start_date, cuda::std::chrono::year_month_day end_date, std::string const& column_name) diff --git a/cpp/libcudf_streaming/benchmarks/utils/random_data.cu b/cpp/libcudf_streaming/benchmarks/utils/random_data.cu index 823a248a03aa..a11e2e74b942 100644 --- a/cpp/libcudf_streaming/benchmarks/utils/random_data.cu +++ b/cpp/libcudf_streaming/benchmarks/utils/random_data.cu @@ -26,7 +26,7 @@ rmm::device_uvector random_device_vector(std::size_t nelem, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { // Fill vector with random data. @@ -51,7 +51,7 @@ rmm::device_uvector random_device_vector(std::size_t nelem, std::unique_ptr random_column(cudf::size_type nrows, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { auto vec = @@ -63,7 +63,7 @@ cudf::table random_table(cudf::size_type ncolumns, cudf::size_type nrows, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { std::vector> cols; @@ -85,7 +85,7 @@ void random_fill(rapidsmpf::Buffer& buffer, rmm::device_async_resource_ref mr) std::numeric_limits::max(), buffer.stream(), mr); - buffer.write_access([&](std::byte* buffer_data, rmm::cuda_stream_view stream) { + buffer.write_access([&](std::byte* buffer_data, cuda::stream_ref stream) { RAPIDSMPF_CUDA_TRY( rapidsmpf::cuda_memcpy_async(buffer_data, vec.data(), buffer.size, stream)); }); diff --git a/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp b/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp index 6cdae39a56aa..bfd44eb4b6e3 100644 --- a/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp +++ b/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp @@ -49,7 +49,7 @@ std::size_t constexpr random_table_size_lower_bound(cudf::size_type ncolumns, cu rmm::device_uvector random_device_vector(std::size_t nelem, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr); /** @@ -70,7 +70,7 @@ rmm::device_uvector random_device_vector(std::size_t nelem, std::unique_ptr random_column(cudf::size_type nrows, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr); /** @@ -93,7 +93,7 @@ cudf::table random_table(cudf::size_type ncolumns, cudf::size_type nrows, std::int32_t min_val, std::int32_t max_val, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr); /** diff --git a/cpp/libcudf_streaming/examples/example_shuffle.cpp b/cpp/libcudf_streaming/examples/example_shuffle.cpp index 91fadb98eb40..53a954b59b81 100644 --- a/cpp/libcudf_streaming/examples/example_shuffle.cpp +++ b/cpp/libcudf_streaming/examples/example_shuffle.cpp @@ -45,7 +45,7 @@ int main(int argc, char** argv) std::make_shared(MPI_COMM_WORLD, progress_thread, log); // We will use the same stream, memory, and buffer resource throughout the example. - rmm::cuda_stream_view stream = cudf::get_default_stream(); + cuda::stream_ref stream = cudf::get_default_stream(); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); auto br = rapidsmpf::BufferResource::create(mr); diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp index 2f384ad2fa52..8d21ce9b8c1e 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include @@ -20,7 +20,7 @@ namespace cudf_streaming::detail { * @param value Value to set @p data to. * @param stream CUDA stream for kernel launches and memory operations. */ -void set_value(std::uint64_t* data, std::uint64_t value, rmm::cuda_stream_view stream); +void set_value(std::uint64_t* data, std::uint64_t value, cuda::stream_ref stream); /** * @brief Add the value in @p left into @p right. @@ -31,6 +31,6 @@ void set_value(std::uint64_t* data, std::uint64_t value, rmm::cuda_stream_view s * @param right Array to add into. * @param stream CUDA stream for kernel launches and memory operations. */ -void add_values(std::uint64_t const* left, std::uint64_t* right, rmm::cuda_stream_view stream); +void add_values(std::uint64_t const* left, std::uint64_t* right, cuda::stream_ref stream); } // namespace cudf_streaming::detail diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index 435087c60ffb..64c6b59f6ba1 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -6,11 +6,12 @@ #pragma once #include -#include #include #include #include +#include + #include #include #include @@ -62,7 +63,7 @@ struct device_bloom_filter { * filter size. */ static std::unique_ptr storage(std::size_t filter_size, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr); /** @@ -88,7 +89,7 @@ struct device_bloom_filter { * @param mr Memory resource for allocations. */ void add(cudf::table_view const& values_to_hash, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr); /** @@ -99,7 +100,7 @@ struct device_bloom_filter { * * @throws std::logic_error If `other` is not compatible with this filter. */ - void merge(device_bloom_filter const& other, rmm::cuda_stream_view stream); + void merge(device_bloom_filter const& other, cuda::stream_ref stream); /** * @brief Return a mask of which rows are contained in the filter. @@ -111,7 +112,7 @@ struct device_bloom_filter { * @return Mask vector to be used for filtering the table. */ [[nodiscard]] rmm::device_uvector contains(cudf::table_view const& values, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const; /** diff --git a/cpp/libcudf_streaming/include/cudf_streaming/parquet.hpp b/cpp/libcudf_streaming/include/cudf_streaming/parquet.hpp index 3e2ba6b13ab1..2315492bf702 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/parquet.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/parquet.hpp @@ -23,7 +23,7 @@ namespace cudf_streaming { * @brief Filter ast expression with lifetime/stream management. */ struct filter { - rmm::cuda_stream_view stream; ///< Stream the filter's scalars are valid on. + cuda::stream_ref stream; ///< Stream the filter's scalars are valid on. cudf::ast::expression& filter; ///< Filter expression. rapidsmpf::OwningWrapper owner{}; ///< Owner of all objects in the filter. }; diff --git a/cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp b/cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp index 5dd5787bde95..f9ca36ab0466 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp @@ -46,7 +46,7 @@ partition_and_split( int num_partitions, cudf::hash_id hash_function, std::uint32_t seed, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); @@ -78,7 +78,7 @@ partition_and_pack( int num_partitions, cudf::hash_id hash_function, std::uint32_t seed, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); @@ -104,7 +104,7 @@ partition_and_pack( [[nodiscard]] std::unordered_map split_and_pack( cudf::table_view const& table, std::vector const& splits, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); @@ -134,7 +134,7 @@ partition_and_pack( */ [[nodiscard]] std::unique_ptr unpack_and_concat( std::vector&& partitions, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES); diff --git a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp index 7099f755fd4f..8f016e2d5214 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include @@ -62,7 +62,7 @@ class table_chunk { * @param table Device-resident table. * @param stream The CUDA stream on which the table was created. */ - table_chunk(std::unique_ptr table, rmm::cuda_stream_view stream); + table_chunk(std::unique_ptr table, cuda::stream_ref stream); /** * @brief Construct a table_chunk from a device table view. @@ -91,7 +91,7 @@ class table_chunk { * is therefore not spillable. */ table_chunk(cudf::table_view table_view, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::OwningWrapper&& owner, exclusive_view exclusive_view); @@ -130,7 +130,7 @@ class table_chunk { * * @return The CUDA stream view. */ - [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; + [[nodiscard]] cuda::stream_ref stream() const noexcept; /** * @brief Number of bytes allocated for the data in the specified memory type. @@ -312,7 +312,7 @@ class table_chunk { std::array data_alloc_size_ = {}; std::size_t make_available_cost_; // For now, only device memory cost is tracked. - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; bool is_spillable_; }; diff --git a/cpp/libcudf_streaming/include/cudf_streaming/utils.hpp b/cpp/libcudf_streaming/include/cudf_streaming/utils.hpp index 354bc73df436..0fb50bbaad0b 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/utils.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/utils.hpp @@ -25,7 +25,7 @@ namespace cudf_streaming { */ std::string str(cudf::column_view col, cudf::size_type index, - rmm::cuda_stream_view stream = cudf::get_default_stream(), + cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** @@ -37,7 +37,7 @@ std::string str(cudf::column_view col, * @return A string representation of all elements in the column. */ std::string str(cudf::column_view col, - rmm::cuda_stream_view stream = cudf::get_default_stream(), + cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** @@ -49,7 +49,7 @@ std::string str(cudf::column_view col, * @return A string representation of all rows in the table. */ std::string str(cudf::table_view tbl, - rmm::cuda_stream_view stream = cudf::get_default_stream(), + cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** @@ -59,7 +59,7 @@ std::string str(cudf::table_view tbl, * @param stream CUDA stream used for device memory operations and kernel launches. * @return The estimated memory usage of the column. */ -std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_stream_view stream); +std::size_t estimated_memory_usage(cudf::column_view const& col, cuda::stream_ref stream); /** * @brief Estimate the memory usage of a table. @@ -68,6 +68,6 @@ std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_strea * @param stream CUDA stream used for device memory operations and kernel launches. * @return The estimated memory usage of the table. */ -std::size_t estimated_memory_usage(cudf::table_view const& tbl, rmm::cuda_stream_view stream); +std::size_t estimated_memory_usage(cudf::table_view const& tbl, cuda::stream_ref stream); } // namespace cudf_streaming diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index af4134c09f27..ef56f22d5780 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -13,8 +13,7 @@ #include #include -#include - +#include #include #include @@ -141,7 +140,7 @@ rapidsmpf::streaming::Actor cardinality_estimator::estimate( tag_, [precision = precision_, sketch_bytes, row_count_offset](rapidsmpf::Buffer const* left, rapidsmpf::Buffer* right) { - right->write_access([&](std::byte* out, rmm::cuda_stream_view stream) { + right->write_access([&](std::byte* out, cuda::stream_ref stream) { auto sketch = cudf::approx_distinct_count({reinterpret_cast(out), sketch_bytes}, precision, @@ -160,7 +159,7 @@ rapidsmpf::streaming::Actor cardinality_estimator::estimate( } auto const [distinct_count, row_count] = - storage->write_access([&](std::byte* data, rmm::cuda_stream_view stream) { + storage->write_access([&](std::byte* data, cuda::stream_ref stream) { auto sketch = cudf::approx_distinct_count({reinterpret_cast(data), sketch_bytes}, precision_, diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index c3411c25a6cb..3025bfcf2671 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -91,7 +91,7 @@ rapidsmpf::streaming::Actor bloom_filter::build( tag, [filter_size = filter_size_, seed = seed_](rapidsmpf::Buffer const* left, rapidsmpf::Buffer* right) { - right->write_access([&](std::byte* out_bytes, rmm::cuda_stream_view stream) { + right->write_access([&](std::byte* out_bytes, cuda::stream_ref stream) { auto const in = cudf_streaming::detail::device_bloom_filter::view(filter_size, seed, left->data()); cudf_streaming::detail::device_bloom_filter(filter_size, seed, out_bytes) diff --git a/cpp/libcudf_streaming/src/detail/approx_distinct_count.cu b/cpp/libcudf_streaming/src/detail/approx_distinct_count.cu index af1926155a7f..3f8bf9916c03 100644 --- a/cpp/libcudf_streaming/src/detail/approx_distinct_count.cu +++ b/cpp/libcudf_streaming/src/detail/approx_distinct_count.cu @@ -24,15 +24,15 @@ __global__ void add_values_kernel(std::uint64_t const* left, std::uint64_t* righ } // namespace -void set_value(std::uint64_t* data, std::uint64_t value, rmm::cuda_stream_view stream) +void set_value(std::uint64_t* data, std::uint64_t value, cuda::stream_ref stream) { - set_value_kernel<<<1, 1, 0, stream.value()>>>(data, value); + set_value_kernel<<<1, 1, 0, stream.get()>>>(data, value); RAPIDSMPF_CUDA_TRY(cudaPeekAtLastError()); } -void add_values(std::uint64_t const* left, std::uint64_t* right, rmm::cuda_stream_view stream) +void add_values(std::uint64_t const* left, std::uint64_t* right, cuda::stream_ref stream) { - add_values_kernel<<<1, 1, 0, stream.value()>>>(left, right); + add_values_kernel<<<1, 1, 0, stream.get()>>>(left, right); RAPIDSMPF_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index ad3a2217ba13..1103577aa504 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -35,11 +35,11 @@ #include #include -#include #include #include #include +#include #include #include @@ -94,7 +94,7 @@ device_bloom_filter const device_bloom_filter::view(std::size_t filter_size, } std::unique_ptr device_bloom_filter::storage(std::size_t filter_size, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { return std::make_unique( @@ -102,7 +102,7 @@ std::unique_ptr device_bloom_filter::storage(std::size_t fil } void device_bloom_filter::add(cudf::table_view const& values_to_hash, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { RAPIDSMPF_NVTX_FUNC_RANGE(); @@ -115,7 +115,7 @@ void device_bloom_filter::add(cudf::table_view const& values_to_hash, filter_ref.add_async(hash_view.begin(), hash_view.end(), stream); } -void device_bloom_filter::merge(device_bloom_filter const& other, rmm::cuda_stream_view stream) +void device_bloom_filter::merge(device_bloom_filter const& other, cuda::stream_ref stream) { RAPIDSMPF_NVTX_FUNC_RANGE(); RAPIDSMPF_EXPECTS(num_blocks_ == other.num_blocks_, "Mismatching number of blocks in filters"); @@ -127,7 +127,7 @@ void device_bloom_filter::merge(device_bloom_filter const& other, rmm::cuda_stre } rmm::device_uvector device_bloom_filter::contains(cudf::table_view const& values, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) const { RAPIDSMPF_NVTX_FUNC_RANGE(); diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index d4395d766e95..8880ae96418e 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -12,7 +12,7 @@ #include #include -#include +#include #include #include @@ -195,7 +195,7 @@ class FileCache { * @return Message representing the read chunk. */ rapidsmpf::streaming::Message read_parquet_chunk(std::shared_ptr ctx, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, cudf::io::parquet_reader_options options, std::uint64_t sequence_number) { diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index 99458b77fee9..0481472bd838 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -12,9 +12,10 @@ #include #include -#include #include +#include + #include #include #include @@ -36,7 +37,7 @@ std::pair, std::unique_ptr> partition int num_partitions, cudf::hash_id hash_function, std::uint32_t seed, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking) { @@ -75,7 +76,7 @@ std::unordered_map partition int num_partitions, cudf::hash_id hash_function, std::uint32_t seed, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking) { @@ -102,7 +103,7 @@ std::unordered_map partition std::unordered_map split_and_pack( cudf::table_view const& table, std::vector const& splits, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking) { @@ -128,7 +129,7 @@ std::unordered_map split_and } std::unique_ptr unpack_and_concat(std::vector&& partitions, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br, rapidsmpf::AllowOverbooking allow_overbooking) { @@ -151,7 +152,7 @@ std::unique_ptr unpack_and_concat(std::vector unpacked; std::vector references; - std::vector packed_data_streams; + std::vector packed_data_streams; unpacked.reserve(partitions.size()); references.reserve(partitions.size()); packed_data_streams.reserve(partitions.size()); diff --git a/cpp/libcudf_streaming/src/table_chunk.cpp b/cpp/libcudf_streaming/src/table_chunk.cpp index 3a0563fa1a20..113babb589cb 100644 --- a/cpp/libcudf_streaming/src/table_chunk.cpp +++ b/cpp/libcudf_streaming/src/table_chunk.cpp @@ -20,7 +20,7 @@ namespace cudf_streaming { -table_chunk::table_chunk(std::unique_ptr table, rmm::cuda_stream_view stream) +table_chunk::table_chunk(std::unique_ptr table, cuda::stream_ref stream) : table_{std::move(table)}, stream_{stream}, is_spillable_{true} { RAPIDSMPF_EXPECTS(table_ != nullptr, "table pointer cannot be null", std::invalid_argument); @@ -31,7 +31,7 @@ table_chunk::table_chunk(std::unique_ptr table, rmm::cuda_stream_vi } table_chunk::table_chunk(cudf::table_view table_view, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::OwningWrapper&& owner, exclusive_view exclusive_view) : owner_{std::move(owner)}, @@ -93,7 +93,7 @@ table_chunk& table_chunk::operator=(table_chunk&& other) noexcept return *this; } -rmm::cuda_stream_view table_chunk::stream() const noexcept { return stream_; } +cuda::stream_ref table_chunk::stream() const noexcept { return stream_; } std::size_t table_chunk::data_alloc_size(rapidsmpf::MemoryType mem_type) const { diff --git a/cpp/libcudf_streaming/src/utils.cpp b/cpp/libcudf_streaming/src/utils.cpp index 75bad965c73a..ba366cb9dd4d 100644 --- a/cpp/libcudf_streaming/src/utils.cpp +++ b/cpp/libcudf_streaming/src/utils.cpp @@ -34,7 +34,7 @@ struct str_cudf_column_scalar_fn { requires(cudf::is_numeric()) std::string operator()(cudf::column_view col, cudf::size_type index, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { std::unique_ptr scalar = cudf::get_element(col, index, stream, mr); @@ -48,7 +48,7 @@ struct str_cudf_column_scalar_fn { requires(!cudf::is_numeric()) std::string operator()(cudf::column_view /* col */, cudf::size_type /* index */, - rmm::cuda_stream_view /* stream */, + cuda::stream_ref /* stream */, rmm::device_async_resource_ref /* mr */ ) { @@ -59,7 +59,7 @@ struct str_cudf_column_scalar_fn { struct cudf_column_data_size_fn { template requires(cudf::is_fixed_width()) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + std::size_t operator()(cudf::column_view const& col, cuda::stream_ref) { return rapidsmpf::safe_cast(col.size()) * cudf::size_of(col.type()) + bitmask_size(col); @@ -68,7 +68,7 @@ struct cudf_column_data_size_fn { // string type specialization template requires(std::is_same_v) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view stream) + std::size_t operator()(cudf::column_view const& col, cuda::stream_ref stream) { cudf::strings_column_view sv(col); return rapidsmpf::safe_cast(sv.chars_size(stream)) + bitmask_size(col); @@ -77,7 +77,7 @@ struct cudf_column_data_size_fn { // compound type specialization except string template requires(!std::is_same_v && cudf::is_compound()) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + std::size_t operator()(cudf::column_view const& col, cuda::stream_ref) { // compound types (except string) ie. list, dict, structs dont have a // content::data buffer. Data is stored in children columns. So, just return the @@ -86,7 +86,7 @@ struct cudf_column_data_size_fn { } template - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) + std::size_t operator()(cudf::column_view const& col, cuda::stream_ref) { RAPIDSMPF_FAIL("not implemented for type: " + cudf::type_to_name(col.type())); } @@ -101,15 +101,13 @@ struct cudf_column_data_size_fn { std::string str(cudf::column_view col, cudf::size_type index, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref mr) { return cudf::type_dispatcher(col.type(), str_cudf_column_scalar_fn{}, col, index, stream, mr); } -std::string str(cudf::column_view col, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +std::string str(cudf::column_view col, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { std::stringstream ss; ss << "Column(["; @@ -121,9 +119,7 @@ std::string str(cudf::column_view col, return ss.str(); } -std::string str(cudf::table_view tbl, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +std::string str(cudf::table_view tbl, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { std::stringstream ss; ss << "Table(["; @@ -137,7 +133,7 @@ std::string str(cudf::table_view tbl, return ss.str(); } -std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_stream_view stream) +std::size_t estimated_memory_usage(cudf::column_view const& col, cuda::stream_ref stream) { return std::transform_reduce( col.child_begin(), @@ -147,7 +143,7 @@ std::size_t estimated_memory_usage(cudf::column_view const& col, rmm::cuda_strea [&stream](cudf::column_view const& child) { return estimated_memory_usage(child, stream); }); } -std::size_t estimated_memory_usage(cudf::table_view const& tbl, rmm::cuda_stream_view stream) +std::size_t estimated_memory_usage(cudf::table_view const& tbl, cuda::stream_ref stream) { return std::transform_reduce( tbl.begin(), tbl.end(), std::size_t{0}, std::plus{}, [&stream](cudf::column_view const& col) { diff --git a/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp index 8a849bd22d98..d7786e5074db 100644 --- a/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/libcudf_streaming/tests/streaming/base_streaming_fixture.hpp @@ -51,7 +51,7 @@ class BaseStreamingFixture : public ::testing::Test { std::move(options), GlobalEnvironment->comm_->logger(), br); } - rmm::cuda_stream_view stream; + cuda::stream_ref stream; rmm::mr::cuda_memory_resource mr_cuda; std::shared_ptr br; std::shared_ptr ctx; diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp index ef2c88f94037..16e1b0a7fe7d 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -12,9 +12,10 @@ #include #include -#include #include +#include + #include #include @@ -160,7 +161,7 @@ TEST_F(StreamingChannelMetadata, MessageRoundTrip) class StreamingChannelMetadataGPU : public ::testing::Test { protected: - rmm::cuda_stream_view stream{cudf::get_default_stream()}; + cuda::stream_ref stream{cudf::get_default_stream()}; std::shared_ptr br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); diff --git a/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp b/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp index a9e4b5b5219f..b3392693eb00 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_cudf_utils.cpp @@ -15,7 +15,7 @@ class BaseEstimatedMemoryUsageTest : public ::testing::Test { protected: void SetUp() override { stream = cudf::get_default_stream(); } - rmm::cuda_stream_view stream; + cuda::stream_ref stream; }; /** diff --git a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp index bb943b9f5ea9..97b134d2d1fa 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp @@ -25,9 +25,10 @@ #include #include -#include #include +#include + #include #include #include @@ -191,9 +192,9 @@ TEST_P(StreamingReadParquetParams, ReadParquet) if (filter_expr != nullptr) { auto expected_options = options; expected_options.set_filter(filter_expr->filter); - filter_expr->stream.synchronize(); + filter_expr->stream.sync(); auto expected = cudf::io::read_parquet(expected_options).tbl; - filter_expr->stream.synchronize(); + filter_expr->stream.sync(); return expected; } else { return cudf::io::read_parquet(options).tbl; @@ -239,7 +240,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) // May as well check on all ranks, so we also mildly exercise the allgather. auto gathered_packed_data = allgather.wait_and_extract(rapidsmpf::coll::AllGather::Ordered::YES); auto result = cudf_streaming::unpack_and_concat( - std::move(gathered_packed_data), rmm::cuda_stream_default, br.get()); + std::move(gathered_packed_data), cuda::stream_ref{}, br.get()); EXPECT_EQ(result->num_rows(), expected->num_rows()); EXPECT_EQ(result->num_columns(), expected->num_columns()); EXPECT_EQ(result->num_columns(), 1); diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp index 35719072c67b..0e51e4e232af 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -16,9 +16,10 @@ #include -#include #include +#include + #include #include @@ -55,7 +56,7 @@ class StreamingTableChunk : public BaseStreamingFixture, options, GlobalEnvironment->comm_->logger(), br); } - rmm::cuda_stream_view stream; + cuda::stream_ref stream; rmm::mr::cuda_memory_resource mr_cuda; std::shared_ptr br; std::shared_ptr ctx; @@ -69,7 +70,7 @@ TEST_F(StreamingTableChunk, FromTable) cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); table_chunk chunk{std::make_unique(expect), stream}; - EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_EQ(chunk.stream().get(), stream.get()); EXPECT_TRUE(chunk.is_available()); EXPECT_TRUE(chunk.is_spillable()); EXPECT_EQ(chunk.make_available_cost(), 0); @@ -102,7 +103,7 @@ TEST_F(StreamingTableChunk, TableChunkOwner) return table_chunk{expect, stream, rapidsmpf::OwningWrapper(new int, deleter), exclusive_view}; }; auto check_chunk = [&](table_chunk const& chunk, bool is_spillable) { - EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_EQ(chunk.stream().get(), stream.get()); EXPECT_TRUE(chunk.is_available()); EXPECT_EQ(chunk.is_spillable(), is_spillable); EXPECT_EQ(chunk.make_available_cost(), 0); @@ -152,7 +153,7 @@ TEST_F(StreamingTableChunk, FromPackedDataOnDevice) std::move(packed_columns.metadata), br->move(std::move(packed_columns.gpu_data), stream)); table_chunk chunk{std::move(packed_data)}; - EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_EQ(chunk.stream().get(), stream.get()); // chunk was created from packed data on device, so it is available and make available // cost is 0. EXPECT_TRUE(chunk.is_available()); @@ -202,7 +203,7 @@ TEST_P(StreamingTableChunk, FromPackedDataOn) std::move(gpu_data_in_spill_memory)); table_chunk chunk{std::move(packed_data)}; - EXPECT_EQ(chunk.stream().value(), stream.value()); + EXPECT_EQ(chunk.stream().get(), stream.get()); EXPECT_FALSE(chunk.is_available()); EXPECT_TRUE(chunk.is_spillable()); EXPECT_THROW(std::ignore = chunk.table_view(), std::invalid_argument); @@ -277,7 +278,7 @@ TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) table_chunk dev_chunk{std::make_unique(expect), stream}; EXPECT_TRUE(dev_chunk.is_available()); EXPECT_TRUE(dev_chunk.is_spillable()); - EXPECT_EQ(dev_chunk.stream().value(), stream.value()); + EXPECT_EQ(dev_chunk.stream().get(), stream.get()); EXPECT_EQ(dev_chunk.make_available_cost(), 0); { auto cd = get_content_description(dev_chunk); @@ -293,7 +294,7 @@ TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) auto host_copy = dev_chunk.copy(host_res); EXPECT_FALSE(host_copy.is_available()); EXPECT_TRUE(host_copy.is_spillable()); - EXPECT_EQ(host_copy.stream().value(), stream.value()); + EXPECT_EQ(host_copy.stream().get(), stream.get()); EXPECT_GT(host_copy.make_available_cost(), 0); { auto cd = get_content_description(host_copy); @@ -308,7 +309,7 @@ TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) auto host_copy2 = host_copy.copy(host_res2); EXPECT_FALSE(host_copy2.is_available()); EXPECT_TRUE(host_copy2.is_spillable()); - EXPECT_EQ(host_copy2.stream().value(), stream.value()); + EXPECT_EQ(host_copy2.stream().get(), stream.get()); EXPECT_EQ(host_copy2.make_available_cost(), host_copy.make_available_cost()); { auto cd = get_content_description(host_copy2); @@ -324,7 +325,7 @@ TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) auto dev_back = host_copy2.make_available(dev_res); EXPECT_TRUE(dev_back.is_available()); EXPECT_TRUE(dev_back.is_spillable()); - EXPECT_EQ(dev_back.stream().value(), stream.value()); + EXPECT_EQ(dev_back.stream().get(), stream.get()); EXPECT_EQ(dev_back.make_available_cost(), 0); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(dev_back.table_view(), expect); { diff --git a/cpp/libcudf_streaming/tests/test_shuffler.cpp b/cpp/libcudf_streaming/tests/test_shuffler.cpp index 8a9274ed39c9..c78235f5023d 100644 --- a/cpp/libcudf_streaming/tests/test_shuffler.cpp +++ b/cpp/libcudf_streaming/tests/test_shuffler.cpp @@ -54,7 +54,7 @@ void test_shuffler(std::shared_ptr const& comm, std::size_t total_num_rows, std::int64_t seed, cudf::hash_id hash_fn, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br) { // To expose unexpected deadlocks, we use a 30s timeout. In a normal run, the @@ -155,7 +155,7 @@ class MemoryLimits_NumPartition std::size_t total_num_rows; std::int64_t seed = 42; cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; - rmm::cuda_stream_view stream; + cuda::stream_ref stream; std::shared_ptr br; std::unique_ptr shuffler; }; @@ -205,7 +205,7 @@ class ConcurrentShuffleTest : public cudf::test::BaseFixtureWithParam br; }; diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index b49a6d0ee92e..930dc5682894 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -11,10 +11,10 @@ #include #include -#include #include #include +#include #include #include @@ -165,7 +165,7 @@ template */ [[nodiscard]] inline rapidsmpf::PackedData generate_packed_data(int n_elements, int offset, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource& br) { auto values = iota_vector(n_elements, offset); @@ -191,7 +191,7 @@ template inline void validate_packed_data(rapidsmpf::PackedData&& packed_data, int n_elements, int offset, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource& br) { auto const& metadata = *packed_data.metadata; @@ -237,19 +237,19 @@ class DelayedMemoryResource { RAPIDSMPF_FATAL("synchronous deallocation not supported"); } - void* allocate(rmm::cuda_stream_view stream, + void* allocate(cuda::stream_ref stream, std::size_t size, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { void* ptr = upstream_.allocate(stream, size, alignment); if (size > 0) { RAPIDSMPF_CUDA_TRY( - cudaLaunchHostFunc(stream.value(), sleep_on_stream, new std::chrono::milliseconds(delay_))); + cudaLaunchHostFunc(stream.get(), sleep_on_stream, new std::chrono::milliseconds(delay_))); } return ptr; } - void deallocate(rmm::cuda_stream_view stream, + void deallocate(cuda::stream_ref stream, void* ptr, std::size_t size, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd index f59e8c8c8bef..1721f1f78af5 100644 --- a/python/cudf_streaming/cudf_streaming/channel_metadata.pxd +++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pxd @@ -9,7 +9,6 @@ from libcpp.utility cimport pair from libcpp.vector cimport vector from pylibcudf.libcudf.types cimport null_order as cpp_null_order from pylibcudf.libcudf.types cimport order as cpp_order -from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rapidsmpf._detail.exception_handling cimport ex_handler diff --git a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx index a40a4b1fee98..70987ed0f719 100644 --- a/python/cudf_streaming/cudf_streaming/channel_metadata.pyx +++ b/python/cudf_streaming/cudf_streaming/channel_metadata.pyx @@ -190,7 +190,7 @@ cdef class Ordering: Buffer resource to associate with the returned table chunk. """ cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() - cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) + cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().get()) tbl = Table.from_table_view_of_arbitrary( chunk.table_view(), owner=self, stream=stream ) diff --git a/python/cudf_streaming/cudf_streaming/partition_utils.pyx b/python/cudf_streaming/cudf_streaming/partition_utils.pyx index 9e81e8737afd..fa5d215e24e1 100644 --- a/python/cudf_streaming/cudf_streaming/partition_utils.pyx +++ b/python/cudf_streaming/cudf_streaming/partition_utils.pyx @@ -14,13 +14,13 @@ from pylibcudf.libcudf.table.table cimport table as cpp_table from pylibcudf.libcudf.table.table_view cimport table_view from pylibcudf.libcudf.types cimport size_type from pylibcudf.table cimport Table -from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rmm.librmm.device_buffer cimport device_buffer from rmm.pylibrmm.stream cimport Stream from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.buffer_resource cimport BufferResource, cpp_BufferResource from rapidsmpf.memory.packed_data cimport PackedData, cpp_PackedData +from cudf_streaming.stream_ref cimport stream_ref cdef extern from "" nogil: @@ -34,7 +34,7 @@ cdef extern from "" nogil: int num_partitions, int hash_function, uint32_t seed, - cuda_stream_view stream, + stream_ref stream, cpp_BufferResource* br, ) except +ex_handler @@ -42,7 +42,7 @@ cdef extern from "" nogil: "cudf_streaming::split_and_pack"( const table_view& table, const vector[size_type] &splits, - cuda_stream_view stream, + stream_ref stream, cpp_BufferResource* br, ) except +ex_handler @@ -86,7 +86,7 @@ cpdef object partition_and_pack( pylibcudf.contiguous_split.pack cudf_streaming.partition_utils.split_and_pack """ - cdef cuda_stream_view _stream = stream.view() + cdef stream_ref _stream = stream_ref(stream.view().value()) cdef cpp_BufferResource* _br = br.ptr() cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash) cdef unordered_map[uint32_t, cpp_PackedData] _ret @@ -148,7 +148,7 @@ cpdef object split_and_pack( pylibcudf.copying.split cudf_streaming.partition_utils.partition_and_pack """ - cdef cuda_stream_view _stream = stream.view() + cdef stream_ref _stream = stream_ref(stream.view().value()) cdef cpp_BufferResource* _br = br.ptr() cdef vector[size_type] _splits = tuple(splits) cdef unordered_map[uint32_t, cpp_PackedData] _ret @@ -185,7 +185,7 @@ cdef extern from "" nogil: cdef unique_ptr[cpp_table] cpp_unpack_and_concat \ "cudf_streaming::unpack_and_concat"( vector[cpp_PackedData] partition, - cuda_stream_view stream, + stream_ref stream, cpp_BufferResource* br, ) except +ex_handler @@ -241,7 +241,7 @@ cpdef object unpack_and_concat( -------- cudf_streaming.partition_utils.partition_and_pack """ - cdef cuda_stream_view _stream = stream.view() + cdef stream_ref _stream = stream_ref(stream.view().value()) cdef cpp_BufferResource* _br = br.ptr() cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) cdef unique_ptr[cpp_table] _ret @@ -263,7 +263,7 @@ cdef extern from *: std::unique_ptr cpp_packed_data_from_buffers( std::unique_ptr> metadata, std::unique_ptr gpu_data, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rapidsmpf::BufferResource* br ) { return std::make_unique( @@ -274,7 +274,7 @@ cdef extern from *: unique_ptr[cpp_PackedData] cpp_packed_data_from_buffers( unique_ptr[vector[uint8_t]] metadata, unique_ptr[device_buffer] gpu_data, - cuda_stream_view stream, + stream_ref stream, cpp_BufferResource* br, ) except +ex_handler nogil @@ -317,7 +317,7 @@ cpdef object packed_data_from_cudf_packed_columns( """ if packed_columns is None or stream is None or br is None: raise TypeError("Arguments must not be None") - cdef cuda_stream_view _stream = stream.view() + cdef stream_ref _stream = stream_ref(stream.view().value()) cdef cpp_BufferResource* _br = br.ptr() cdef PackedData ret = PackedData.__new__(PackedData) with nogil: diff --git a/python/cudf_streaming/cudf_streaming/stream_ref.pxd b/python/cudf_streaming/cudf_streaming/stream_ref.pxd new file mode 100644 index 000000000000..8d27f14ee162 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/stream_ref.pxd @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from cuda.bindings.cyruntime cimport cudaStream_t + + +cdef extern from "" namespace "cuda" nogil: + cdef cppclass stream_ref: + stream_ref() noexcept + stream_ref(cudaStream_t) noexcept + cudaStream_t get() noexcept diff --git a/python/cudf_streaming/cudf_streaming/table_chunk.pxd b/python/cudf_streaming/cudf_streaming/table_chunk.pxd index f35a8dfbfb70..87b423e0ac05 100644 --- a/python/cudf_streaming/cudf_streaming/table_chunk.pxd +++ b/python/cudf_streaming/cudf_streaming/table_chunk.pxd @@ -8,7 +8,6 @@ from libcpp.memory cimport unique_ptr from libcpp.utility cimport pair from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view from pylibcudf.libcudf.types cimport size_type -from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rmm.pylibrmm.stream cimport Stream from rapidsmpf._detail.exception_handling cimport ex_handler @@ -17,12 +16,13 @@ from rapidsmpf.memory.buffer_resource cimport (BufferResource, cpp_BufferResource) from rapidsmpf.memory.memory_reservation cimport cpp_MemoryReservation from rapidsmpf.memory.packed_data cimport cpp_PackedData +from cudf_streaming.stream_ref cimport stream_ref cdef extern from "" nogil: cdef cppclass cpp_TableChunk "cudf_streaming::table_chunk": cpp_TableChunk(unique_ptr[cpp_PackedData]) except +ex_handler - cuda_stream_view stream() noexcept + stream_ref stream() noexcept size_t data_alloc_size(MemoryType mem_type) except +ex_handler bool_t is_available() noexcept size_t make_available_cost() noexcept diff --git a/python/cudf_streaming/cudf_streaming/table_chunk.pyx b/python/cudf_streaming/cudf_streaming/table_chunk.pyx index 3ae38093a1c0..dda313812a44 100644 --- a/python/cudf_streaming/cudf_streaming/table_chunk.pyx +++ b/python/cudf_streaming/cudf_streaming/table_chunk.pyx @@ -9,6 +9,7 @@ from libcpp.memory cimport make_unique, unique_ptr from libcpp.utility cimport move from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view from pylibcudf.table cimport Table +from cudf_streaming.stream_ref cimport stream_ref from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.buffer_resource cimport (BufferResource, @@ -46,7 +47,7 @@ cdef extern from * nogil: std::unique_ptr cpp_from_table_view_with_owner( cudf::table_view view, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, PyObject *owner, void(*py_deleter)(void *), bool exclusive_view @@ -182,7 +183,7 @@ cdef class TableChunk: persists even when the chunk is transferred through Channels. """ - cdef cuda_stream_view _stream = stream.view() + cdef stream_ref _stream = stream_ref(stream.view().value()) cdef cpp_table_view view = table.view() return TableChunk.from_handle( cpp_from_table_view_with_owner( @@ -313,7 +314,7 @@ cdef class TableChunk: The CUDA stream. """ return Stream._from_cudaStream_t( - deref(self.handle_ptr()).stream().value() + deref(self.handle_ptr()).stream().get() ) def data_alloc_size(self, mem_type=None): From eca1474cb2d783d05db4371dd15c563267808329 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 17 Aug 2026 16:59:18 -0700 Subject: [PATCH 2/8] Address stream_ref streaming review comments --- cpp/libcudf_streaming/benchmarks/bench_pack.cpp | 14 +++++++------- .../benchmarks/bench_partition.cpp | 6 +++--- .../benchmarks/streaming/ndsh/join.cpp | 2 +- .../benchmarks/streaming/ndsh/utils.hpp | 2 +- .../detail/approx_distinct_count.hpp | 2 +- .../cudf_streaming/detail/device_bloom_filter.hpp | 2 +- .../include/cudf_streaming/table_chunk.hpp | 2 +- .../src/approx_distinct_count.cpp | 2 +- .../src/detail/device_bloom_filter.cu | 2 +- cpp/libcudf_streaming/src/parquet.cpp | 2 +- cpp/libcudf_streaming/src/partition_utils.cpp | 2 +- .../tests/streaming/test_channel_metadata.cpp | 2 +- .../tests/streaming/test_read_parquet.cpp | 4 ++-- .../tests/streaming/test_table_chunk.cpp | 2 +- cpp/libcudf_streaming/tests/utils.hpp | 2 +- .../cudf_streaming/cudf_streaming/stream_ref.pxd | 2 +- 16 files changed, 25 insertions(+), 25 deletions(-) diff --git a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp index a685f231de90..adc2fd95f4be 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -71,7 +71,7 @@ static void BM_Pack_device(benchmark::State& state) { auto const table_size_mb = static_cast(state.range(0)); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -92,7 +92,7 @@ static void BM_Pack_pinned(benchmark::State& state) auto const table_size_mb = static_cast(state.range(0)); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{ @@ -179,7 +179,7 @@ static void BM_ChunkedPack_device(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40)}; @@ -204,7 +204,7 @@ static void BM_ChunkedPack_pinned(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) @@ -248,7 +248,7 @@ static void BM_ChunkedPack_fixed_table_device(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -272,7 +272,7 @@ static void BM_ChunkedPack_fixed_table_pinned(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) diff --git a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp index 73981ef8998f..421c43bdbcd7 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -45,7 +45,7 @@ static void BM_PartitionAndPack(benchmark::State& state) int const num_partitions = state.range(1); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; // Get total GPU memory cudaDeviceProp prop; @@ -94,7 +94,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) int num_rows = int(local_size / std::int64_t{sizeof(std::int32_t)} / std::int64_t{num_partitions}); - cuda::stream_ref stream = cuda::stream_ref{}; + cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; // Get total GPU memory cudaDeviceProp prop; diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index acfa11ba80de..a354f2d523bc 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -19,7 +19,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp index bc8876b4e03c..abb37de58b83 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp index 8d21ce9b8c1e..cc3136934fce 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index 64c6b59f6ba1..85a536afd0a6 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -10,7 +10,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp index 8f016e2d5214..cfb8d8ca5251 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index ef56f22d5780..3dcbca000539 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -13,7 +13,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 1103577aa504..8bf85e770727 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -39,7 +39,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index 8880ae96418e..d166251a096b 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -12,7 +12,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index 0481472bd838..f717eda99d8d 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -14,7 +14,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp index 16e1b0a7fe7d..d369a6427e66 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -14,7 +14,7 @@ #include -#include +#include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp index 97b134d2d1fa..30288fe0f90a 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp @@ -27,7 +27,7 @@ #include -#include +#include #include #include @@ -240,7 +240,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) // May as well check on all ranks, so we also mildly exercise the allgather. auto gathered_packed_data = allgather.wait_and_extract(rapidsmpf::coll::AllGather::Ordered::YES); auto result = cudf_streaming::unpack_and_concat( - std::move(gathered_packed_data), cuda::stream_ref{}, br.get()); + std::move(gathered_packed_data), cuda::stream_ref{cudaStreamLegacy}, br.get()); EXPECT_EQ(result->num_rows(), expected->num_rows()); EXPECT_EQ(result->num_columns(), expected->num_columns()); EXPECT_EQ(result->num_columns(), 1); diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp index 0e51e4e232af..a1cfb8ce6c74 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -18,7 +18,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index 930dc5682894..58e324278b30 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/python/cudf_streaming/cudf_streaming/stream_ref.pxd b/python/cudf_streaming/cudf_streaming/stream_ref.pxd index 8d27f14ee162..98140b8b6ee4 100644 --- a/python/cudf_streaming/cudf_streaming/stream_ref.pxd +++ b/python/cudf_streaming/cudf_streaming/stream_ref.pxd @@ -4,7 +4,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t -cdef extern from "" namespace "cuda" nogil: +cdef extern from "" namespace "cuda" nogil: cdef cppclass stream_ref: stream_ref() noexcept stream_ref(cudaStream_t) noexcept From 9f1a4eab9c31db44396bd587f0fa4fa365c00f94 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 19 Aug 2026 11:51:50 -0700 Subject: [PATCH 3/8] Adapt streaming joins for cuda stream_ref --- .../benchmarks/bench_partition.cpp | 4 +- .../benchmarks/streaming/ndsh/concatenate.cpp | 10 ++-- .../benchmarks/streaming/ndsh/join.cpp | 14 ++--- .../streaming/ndsh/parquet_writer.cpp | 5 +- .../benchmarks/streaming/ndsh/q03.cpp | 4 +- .../cudf_streaming/detail/stream_adapter.hpp | 53 +++++++++++++++++++ .../src/approx_distinct_count.cpp | 3 +- cpp/libcudf_streaming/src/bloom_filter.cpp | 7 +-- .../src/channel_metadata.cpp | 3 +- cpp/libcudf_streaming/src/parquet.cpp | 5 +- cpp/libcudf_streaming/src/partition_utils.cpp | 3 +- cpp/libcudf_streaming/tests/utils.hpp | 2 +- 12 files changed, 89 insertions(+), 24 deletions(-) create mode 100644 cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp diff --git a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp index 421c43bdbcd7..9de743e194e4 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -74,7 +74,7 @@ static void BM_PartitionAndPack(benchmark::State& state) stream, br.get()); benchmark::DoNotOptimize(pack_partitions); - CUDF_CUDA_TRY(cudaStreamSynchronize(stream)); + stream.sync(); } // Set metrics @@ -125,7 +125,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) br.get()); benchmark::DoNotOptimize(pack_partitions); } - CUDF_CUDA_TRY(cudaStreamSynchronize(stream)); + stream.sync(); } // Set metrics diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp index d93b44c53767..252bf1704221 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -55,15 +56,16 @@ streaming::Actor concatenate(std::shared_ptr ctx, views.reserve(messages.size()); for (auto&& msg : messages) { auto chunk = co_await msg.release().make_available(ctx); - cuda_stream_join(concat_stream, chunk.stream(), &event); + cudf_streaming::detail::cuda_stream_join(concat_stream, chunk.stream(), &event); views.push_back(chunk.table_view()); chunks.push_back(std::move(chunk)); } auto result = std::make_unique( cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream); - cuda_stream_join(chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), - std::ranges::single_view(concat_stream), - &event); + cudf_streaming::detail::cuda_stream_join( + chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), + std::ranges::single_view(concat_stream), + &event); chunks.clear(); co_await ch_out->send(cudf_streaming::to_message(0, std::move(result))); } diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index a354f2d523bc..dea73e181d8a 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include @@ -63,7 +64,7 @@ coro::task broadcast(std::shared_ptr ctx auto msg = co_await ch_in->receive(); if (msg.empty()) { break; } auto chunk = co_await msg.release().make_available(ctx); - cuda_stream_join(gather_stream, chunk.stream(), &event); + cudf_streaming::detail::cuda_stream_join(gather_stream, chunk.stream(), &event); views.push_back(chunk.table_view()); chunks.push_back(std::move(chunk)); } @@ -74,9 +75,10 @@ coro::task broadcast(std::shared_ptr ctx auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr()); // So that deallocation of the consitutent tables is stream-ordered wrt the // concatenation. - cuda_stream_join(chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), - std::ranges::single_view(gather_stream), - &event); + cudf_streaming::detail::cuda_stream_join( + chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), + std::ranges::single_view(gather_stream), + &event); co_return to_message( 0, std::make_unique(std::move(result), gather_stream)); } @@ -176,7 +178,7 @@ streaming::Message semi_join_chunk(std::shared_ptr ctx, auto result_table = std::make_unique(std::move(result_columns)); // Deallocation of the join indices will happen on chunk_stream, so add stream dep - cuda_stream_join(left_chunk.stream(), chunk_stream); + cudf_streaming::detail::cuda_stream_join(left_chunk.stream(), chunk_stream); return to_message( sequence, std::make_unique(std::move(result_table), chunk_stream)); @@ -242,7 +244,7 @@ streaming::Message inner_join_chunk(std::shared_ptr ctx, std::back_inserter(result_columns)); // Deallocation of the join indices will happen on build_stream, so add stream dep // This also ensure deallocation of the hash_join object waits for completion. - cuda_stream_join(build_stream, chunk_stream, tmp_event); + cudf_streaming::detail::cuda_stream_join(build_stream, chunk_stream, tmp_event); return to_message(sequence, std::make_unique( std::make_unique(std::move(result_columns)), chunk_stream)); diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp index d6989b7b739d..b4dc5a5ca185 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -51,9 +52,9 @@ rapidsmpf::streaming::Actor write_parquet(std::shared_ptr(table.num_columns()) == column_names.size(), "Mismatching number of column names and chunk columns"); - cuda_stream_join(write_stream, chunk.stream(), &event); + cudf_streaming::detail::cuda_stream_join(write_stream, chunk.stream(), &event); writer.write(table); - cuda_stream_join(chunk.stream(), write_stream, &event); + cudf_streaming::detail::cuda_stream_join(chunk.stream(), write_stream, &event); } writer.close(); } diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp index 663f7425d745..32b52ade8395 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp @@ -26,6 +26,7 @@ #include #include +#include #include #include #include @@ -251,7 +252,8 @@ rapidsmpf::streaming::Actor top_k_by(std::shared_ptr 0, "No chunks to sort"); auto out_stream = chunk_streams.front(); rapidsmpf::CudaEvent event; - rapidsmpf::cuda_stream_join(std::ranges::single_view{out_stream}, chunk_streams, &event); + cudf_streaming::detail::cuda_stream_join( + std::ranges::single_view{out_stream}, chunk_streams, &event); std::vector views; std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp new file mode 100644 index 000000000000..e296f56538bd --- /dev/null +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include +#include + +#include +#include + +namespace cudf_streaming::detail { + +inline rmm::cuda_stream_view as_rmm_stream(rmm::cuda_stream_view stream) { return stream; } + +inline rmm::cuda_stream_view as_rmm_stream(cuda::stream_ref stream) +{ + return rmm::cuda_stream_view{stream.get()}; +} + +template +auto as_rmm_streams(Range&& streams) +{ + return std::forward(streams) | + std::views::transform([](auto&& stream) { return as_rmm_stream(stream); }); +} + +template + requires(!std::ranges::range && !std::ranges::range) +void cuda_stream_join(Downstream downstream, + Upstream upstream, + rapidsmpf::CudaEvent* event = nullptr) +{ + rapidsmpf::cuda_stream_join(as_rmm_stream(downstream), as_rmm_stream(upstream), event); +} + +template + requires(std::ranges::range && std::ranges::range) +void cuda_stream_join(DownstreamRange&& downstreams, + UpstreamRange&& upstreams, + rapidsmpf::CudaEvent* event = nullptr) +{ + rapidsmpf::cuda_stream_join(as_rmm_streams(std::forward(downstreams)), + as_rmm_streams(std::forward(upstreams)), + event); +} + +} // namespace cudf_streaming::detail diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index 3dcbca000539..e8b1c0d48bb7 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -115,7 +116,7 @@ rapidsmpf::streaming::Actor cardinality_estimator::estimate( auto const table = column_indices.empty() ? chunk.table_view() : chunk.table_view().select(column_indices); sketch.add(table, chunk.stream()); - rapidsmpf::cuda_stream_join(sketch_stream, chunk.stream(), &add_event); + cudf_streaming::detail::cuda_stream_join(sketch_stream, chunk.stream(), &add_event); reservation.clear(); if (ch_sampled != nullptr) { co_await ch_sampled->send( diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 3025bfcf2671..300a54902a3d 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -78,7 +79,7 @@ rapidsmpf::streaming::Actor bloom_filter::build( // kernels doing that concurrently because the updates are atomic. build_event.stream_wait(chunk.stream()); filter.add(chunk.table_view(), chunk.stream(), mr); - rapidsmpf::cuda_stream_join(filter_stream, chunk.stream(), &event); + cudf_streaming::detail::cuda_stream_join(filter_stream, chunk.stream(), &event); } if (comm_->nranks() > 1) { auto reducer = rapidsmpf::streaming::AllReduce( @@ -130,7 +131,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( ctx_, -rapidsmpf::safe_cast(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE))); auto chunk_stream = chunk.stream(); - rapidsmpf::cuda_stream_join(chunk_stream, stream, &event); + cudf_streaming::detail::cuda_stream_join(chunk_stream, stream, &event); // Reservation for the mask construction and guess at output size. auto res = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE) ->reserve_or_wait(rapidsmpf::safe_cast(chunk.table_view().num_rows()) @@ -142,7 +143,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( 0); auto mask = filter.contains(chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr()); - rapidsmpf::cuda_stream_join(stream, chunk_stream, &event); + cudf_streaming::detail::cuda_stream_join(stream, chunk_stream, &event); RAPIDSMPF_EXPECTS(mask.size() == static_cast(chunk.table_view().num_rows()), "Invalid mask size"); auto mask_view = cudf::column_view{cudf::data_type{cudf::type_id::BOOL8}, diff --git a/cpp/libcudf_streaming/src/channel_metadata.cpp b/cpp/libcudf_streaming/src/channel_metadata.cpp index fe31d81f5e94..26cfc5de05f4 100644 --- a/cpp/libcudf_streaming/src/channel_metadata.cpp +++ b/cpp/libcudf_streaming/src/channel_metadata.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -69,7 +70,7 @@ bool ordering::boundaries_aligned_with(ordering const& other, rapidsmpf::BufferR auto const lhs = boundaries->table_view(); auto const rhs = other.boundaries->table_view(); auto const stream = boundaries->stream(); - rapidsmpf::cuda_stream_join(stream, other.boundaries->stream()); + cudf_streaming::detail::cuda_stream_join(stream, other.boundaries->stream()); for (cudf::size_type i = 0; i < lhs.num_columns(); ++i) { auto eq = cudf::binary_operation(lhs.column(i), rhs.column(i), diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index d166251a096b..c046768f22ae 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -319,7 +320,7 @@ rapidsmpf::streaming::Actor read_parquet(std::shared_ptrfilter); // Let's just join all the possible streams here rather than inducing cross-stream // deps in the tasks - rapidsmpf::cuda_stream_join( + cudf_streaming::detail::cuda_stream_join( std::ranges::transform_view( std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool()->get_pool_size()), [&](auto i) { return ctx->br()->stream_pool()->get_stream(i); }), @@ -398,7 +399,7 @@ rapidsmpf::streaming::Actor read_parquet(std::shared_ptrstream), std::ranges::transform_view( std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool()->get_pool_size()), diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index f717eda99d8d..c788c9fb43a1 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -176,7 +177,7 @@ std::unique_ptr unpack_and_concat(std::vectorset_stream(stream); } diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index 58e324278b30..570a97a2911e 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -207,7 +207,7 @@ inline void validate_packed_data(rapidsmpf::PackedData&& packed_data, auto res = br.reserve_or_fail(packed_data.data->size, rapidsmpf::MemoryType::HOST); auto data_on_host = br.move_to_host_buffer(std::move(packed_data.data), res); - RAPIDSMPF_CUDA_TRY(cudaStreamSynchronize(stream)); + stream.sync(); EXPECT_EQ(metadata, data_on_host->copy_to_uint8_vector()); } From b531833548e36c048da77559b1feede940f3e8ab Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Sat, 22 Aug 2026 19:32:23 -0700 Subject: [PATCH 4/8] Use cuda stream_ref header in libcudf_streaming --- cpp/libcudf_streaming/benchmarks/bench_pack.cpp | 2 +- cpp/libcudf_streaming/benchmarks/bench_partition.cpp | 2 +- cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp | 2 +- cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp | 2 +- .../include/cudf_streaming/detail/approx_distinct_count.hpp | 2 +- .../include/cudf_streaming/detail/device_bloom_filter.hpp | 2 +- cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp | 2 +- cpp/libcudf_streaming/src/approx_distinct_count.cpp | 2 +- cpp/libcudf_streaming/src/detail/device_bloom_filter.cu | 2 +- cpp/libcudf_streaming/src/parquet.cpp | 2 +- cpp/libcudf_streaming/src/partition_utils.cpp | 2 +- cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp | 2 +- cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp | 2 +- cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp | 2 +- cpp/libcudf_streaming/tests/utils.hpp | 2 +- 15 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp index adc2fd95f4be..0ba915b6f3d2 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp index 9de743e194e4..6312ccf2a2d1 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index dea73e181d8a..a1bb45c9fb51 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -20,7 +20,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp index abb37de58b83..bc8876b4e03c 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp index cc3136934fce..8d21ce9b8c1e 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index 85a536afd0a6..64c6b59f6ba1 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -10,7 +10,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp index cfb8d8ca5251..8f016e2d5214 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index e8b1c0d48bb7..0a57eb3ecb29 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 8bf85e770727..1103577aa504 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -39,7 +39,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index c046768f22ae..e1fa2867b4de 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -13,7 +13,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index c788c9fb43a1..42428c1e7f62 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -15,7 +15,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp index d369a6427e66..16e1b0a7fe7d 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -14,7 +14,7 @@ #include -#include +#include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp index 30288fe0f90a..e5c140adfae0 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp @@ -27,7 +27,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp index a1cfb8ce6c74..0e51e4e232af 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -18,7 +18,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index 570a97a2911e..7fc137804f39 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include From 808be8443852095cb82ff61dfc56005752ac790f Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 24 Aug 2026 10:15:40 -0700 Subject: [PATCH 5/8] Fix streaming stream_ref synchronization --- cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp index 16e1b0a7fe7d..bb82a91a7d12 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -168,7 +168,7 @@ class StreamingChannelMetadataGPU : public ::testing::Test { std::shared_ptr make_chunk(std::vector vals) { rmm::device_buffer buf(vals.data(), vals.size() * sizeof(int32_t), stream); - stream.synchronize(); + stream.sync(); auto col = std::make_unique(cudf::data_type{cudf::type_id::INT32}, static_cast(vals.size()), std::move(buf), From 9029757b34530947e4353b268a3c571b751658ce Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 26 Aug 2026 15:56:28 -0500 Subject: [PATCH 6/8] Use canonical CUDA streams in libcudf streaming --- cpp/libcudf_streaming/benchmarks/bench_pack.cpp | 14 +++++++------- .../benchmarks/bench_partition.cpp | 6 +++--- .../benchmarks/streaming/ndsh/join.cpp | 2 +- .../benchmarks/streaming/ndsh/utils.hpp | 2 +- .../detail/approx_distinct_count.hpp | 2 +- .../cudf_streaming/detail/device_bloom_filter.hpp | 2 +- .../cudf_streaming/detail/stream_adapter.hpp | 2 +- .../include/cudf_streaming/table_chunk.hpp | 2 +- .../src/approx_distinct_count.cpp | 2 +- .../src/detail/device_bloom_filter.cu | 2 +- cpp/libcudf_streaming/src/parquet.cpp | 2 +- cpp/libcudf_streaming/src/partition_utils.cpp | 2 +- .../tests/streaming/test_channel_metadata.cpp | 2 +- .../tests/streaming/test_read_parquet.cpp | 4 ++-- .../tests/streaming/test_table_chunk.cpp | 2 +- cpp/libcudf_streaming/tests/utils.hpp | 2 +- 16 files changed, 25 insertions(+), 25 deletions(-) diff --git a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp index 0ba915b6f3d2..59f6b0641133 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -71,7 +71,7 @@ static void BM_Pack_device(benchmark::State& state) { auto const table_size_mb = static_cast(state.range(0)); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -92,7 +92,7 @@ static void BM_Pack_pinned(benchmark::State& state) auto const table_size_mb = static_cast(state.range(0)); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); // Create memory resources rmm::mr::pool_memory_resource pool_mr{ @@ -179,7 +179,7 @@ static void BM_ChunkedPack_device(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40)}; @@ -204,7 +204,7 @@ static void BM_ChunkedPack_pinned(benchmark::State& state) // Bounce buffer size: max(1MB, table_size / 10) auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) @@ -248,7 +248,7 @@ static void BM_ChunkedPack_fixed_table_device(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); // Create memory resources rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_async_memory_resource{}, @@ -272,7 +272,7 @@ static void BM_ChunkedPack_fixed_table_pinned(benchmark::State& state) auto const bounce_buffer_size = static_cast(state.range(0)) * MB; constexpr std::size_t table_size_bytes = 1024 * MB; - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); rmm::mr::pool_memory_resource pool_mr{ rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) diff --git a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp index 6312ccf2a2d1..6fdcb03c068f 100644 --- a/cpp/libcudf_streaming/benchmarks/bench_partition.cpp +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -45,7 +45,7 @@ static void BM_PartitionAndPack(benchmark::State& state) int const num_partitions = state.range(1); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); // Get total GPU memory cudaDeviceProp prop; @@ -94,7 +94,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) int num_rows = int(local_size / std::int64_t{sizeof(std::int32_t)} / std::int64_t{num_partitions}); - cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy}; + cuda::stream_ref stream = cudf::get_default_stream(); // Get total GPU memory cudaDeviceProp prop; diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index a1bb45c9fb51..dea73e181d8a 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -20,7 +20,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp index bc8876b4e03c..abb37de58b83 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp index 8d21ce9b8c1e..cc3136934fce 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/approx_distinct_count.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index 64c6b59f6ba1..85a536afd0a6 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -10,7 +10,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp index e296f56538bd..3b272581b705 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp @@ -6,7 +6,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp index 8f016e2d5214..cfb8d8ca5251 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/table_chunk.hpp @@ -9,7 +9,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index 0a57eb3ecb29..e8b1c0d48bb7 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 6dcd3e52f7dd..c960d9a1e98f 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -39,7 +39,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index e1fa2867b4de..c046768f22ae 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -13,7 +13,7 @@ #include #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index 42428c1e7f62..c788c9fb43a1 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -15,7 +15,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp index bb82a91a7d12..1a8044502e19 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_channel_metadata.cpp @@ -14,7 +14,7 @@ #include -#include +#include #include diff --git a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp index e5c140adfae0..d42874b15400 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_read_parquet.cpp @@ -27,7 +27,7 @@ #include -#include +#include #include #include @@ -240,7 +240,7 @@ TEST_P(StreamingReadParquetParams, ReadParquet) // May as well check on all ranks, so we also mildly exercise the allgather. auto gathered_packed_data = allgather.wait_and_extract(rapidsmpf::coll::AllGather::Ordered::YES); auto result = cudf_streaming::unpack_and_concat( - std::move(gathered_packed_data), cuda::stream_ref{cudaStreamLegacy}, br.get()); + std::move(gathered_packed_data), cudf::get_default_stream(), br.get()); EXPECT_EQ(result->num_rows(), expected->num_rows()); EXPECT_EQ(result->num_columns(), expected->num_columns()); EXPECT_EQ(result->num_columns(), 1); diff --git a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp index 0e51e4e232af..a1cfb8ce6c74 100644 --- a/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp +++ b/cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp @@ -18,7 +18,7 @@ #include -#include +#include #include #include diff --git a/cpp/libcudf_streaming/tests/utils.hpp b/cpp/libcudf_streaming/tests/utils.hpp index 7fc137804f39..570a97a2911e 100644 --- a/cpp/libcudf_streaming/tests/utils.hpp +++ b/cpp/libcudf_streaming/tests/utils.hpp @@ -14,7 +14,7 @@ #include #include -#include +#include #include #include From 1a658f5eaec0577ecd32b1e069ae9fd7d69e430e Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 28 Aug 2026 08:45:08 +0000 Subject: [PATCH 7/8] Remove temporary stream adapter now rapidsmpf uses cuda::stream_ref --- .../benchmarks/streaming/ndsh/concatenate.cpp | 5 +- .../benchmarks/streaming/ndsh/join.cpp | 9 ++-- .../streaming/ndsh/parquet_writer.cpp | 5 +- .../benchmarks/streaming/ndsh/q03.cpp | 4 +- .../cudf_streaming/detail/stream_adapter.hpp | 53 ------------------- .../src/approx_distinct_count.cpp | 3 +- cpp/libcudf_streaming/src/bloom_filter.cpp | 7 ++- .../src/channel_metadata.cpp | 3 +- cpp/libcudf_streaming/src/parquet.cpp | 5 +- cpp/libcudf_streaming/src/partition_utils.cpp | 3 +- 10 files changed, 17 insertions(+), 80 deletions(-) delete mode 100644 cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp index 252bf1704221..57d631d9d9a6 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include @@ -56,13 +55,13 @@ streaming::Actor concatenate(std::shared_ptr ctx, views.reserve(messages.size()); for (auto&& msg : messages) { auto chunk = co_await msg.release().make_available(ctx); - cudf_streaming::detail::cuda_stream_join(concat_stream, chunk.stream(), &event); + rapidsmpf::cuda_stream_join(concat_stream, chunk.stream(), &event); views.push_back(chunk.table_view()); chunks.push_back(std::move(chunk)); } auto result = std::make_unique( cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream); - cudf_streaming::detail::cuda_stream_join( + rapidsmpf::cuda_stream_join( chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), std::ranges::single_view(concat_stream), &event); diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index dea73e181d8a..412462247e04 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -16,7 +16,6 @@ #include #include -#include #include #include @@ -64,7 +63,7 @@ coro::task broadcast(std::shared_ptr ctx auto msg = co_await ch_in->receive(); if (msg.empty()) { break; } auto chunk = co_await msg.release().make_available(ctx); - cudf_streaming::detail::cuda_stream_join(gather_stream, chunk.stream(), &event); + rapidsmpf::cuda_stream_join(gather_stream, chunk.stream(), &event); views.push_back(chunk.table_view()); chunks.push_back(std::move(chunk)); } @@ -75,7 +74,7 @@ coro::task broadcast(std::shared_ptr ctx auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr()); // So that deallocation of the consitutent tables is stream-ordered wrt the // concatenation. - cudf_streaming::detail::cuda_stream_join( + rapidsmpf::cuda_stream_join( chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), std::ranges::single_view(gather_stream), &event); @@ -178,7 +177,7 @@ streaming::Message semi_join_chunk(std::shared_ptr ctx, auto result_table = std::make_unique(std::move(result_columns)); // Deallocation of the join indices will happen on chunk_stream, so add stream dep - cudf_streaming::detail::cuda_stream_join(left_chunk.stream(), chunk_stream); + rapidsmpf::cuda_stream_join(left_chunk.stream(), chunk_stream); return to_message( sequence, std::make_unique(std::move(result_table), chunk_stream)); @@ -244,7 +243,7 @@ streaming::Message inner_join_chunk(std::shared_ptr ctx, std::back_inserter(result_columns)); // Deallocation of the join indices will happen on build_stream, so add stream dep // This also ensure deallocation of the hash_join object waits for completion. - cudf_streaming::detail::cuda_stream_join(build_stream, chunk_stream, tmp_event); + rapidsmpf::cuda_stream_join(build_stream, chunk_stream, tmp_event); return to_message(sequence, std::make_unique( std::make_unique(std::move(result_columns)), chunk_stream)); diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp index b4dc5a5ca185..c18fdef44932 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp @@ -8,7 +8,6 @@ #include #include -#include #include #include @@ -52,9 +51,9 @@ rapidsmpf::streaming::Actor write_parquet(std::shared_ptr(table.num_columns()) == column_names.size(), "Mismatching number of column names and chunk columns"); - cudf_streaming::detail::cuda_stream_join(write_stream, chunk.stream(), &event); + rapidsmpf::cuda_stream_join(write_stream, chunk.stream(), &event); writer.write(table); - cudf_streaming::detail::cuda_stream_join(chunk.stream(), write_stream, &event); + rapidsmpf::cuda_stream_join(chunk.stream(), write_stream, &event); } writer.close(); } diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp index 9d0aa9c4673b..48cc3d07b7f0 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp @@ -26,7 +26,6 @@ #include #include -#include #include #include #include @@ -252,8 +251,7 @@ rapidsmpf::streaming::Actor top_k_by(std::shared_ptr 0, "No chunks to sort"); auto out_stream = chunk_streams.front(); rapidsmpf::CudaEvent event; - cudf_streaming::detail::cuda_stream_join( - std::ranges::single_view{out_stream}, chunk_streams, &event); + rapidsmpf::cuda_stream_join(std::ranges::single_view{out_stream}, chunk_streams, &event); std::vector views; std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp deleted file mode 100644 index 3b272581b705..000000000000 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/stream_adapter.hpp +++ /dev/null @@ -1,53 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include - -#include - -#include -#include - -#include -#include - -namespace cudf_streaming::detail { - -inline rmm::cuda_stream_view as_rmm_stream(rmm::cuda_stream_view stream) { return stream; } - -inline rmm::cuda_stream_view as_rmm_stream(cuda::stream_ref stream) -{ - return rmm::cuda_stream_view{stream.get()}; -} - -template -auto as_rmm_streams(Range&& streams) -{ - return std::forward(streams) | - std::views::transform([](auto&& stream) { return as_rmm_stream(stream); }); -} - -template - requires(!std::ranges::range && !std::ranges::range) -void cuda_stream_join(Downstream downstream, - Upstream upstream, - rapidsmpf::CudaEvent* event = nullptr) -{ - rapidsmpf::cuda_stream_join(as_rmm_stream(downstream), as_rmm_stream(upstream), event); -} - -template - requires(std::ranges::range && std::ranges::range) -void cuda_stream_join(DownstreamRange&& downstreams, - UpstreamRange&& upstreams, - rapidsmpf::CudaEvent* event = nullptr) -{ - rapidsmpf::cuda_stream_join(as_rmm_streams(std::forward(downstreams)), - as_rmm_streams(std::forward(upstreams)), - event); -} - -} // namespace cudf_streaming::detail diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index e8b1c0d48bb7..3dcbca000539 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -11,7 +11,6 @@ #include #include -#include #include #include @@ -116,7 +115,7 @@ rapidsmpf::streaming::Actor cardinality_estimator::estimate( auto const table = column_indices.empty() ? chunk.table_view() : chunk.table_view().select(column_indices); sketch.add(table, chunk.stream()); - cudf_streaming::detail::cuda_stream_join(sketch_stream, chunk.stream(), &add_event); + rapidsmpf::cuda_stream_join(sketch_stream, chunk.stream(), &add_event); reservation.clear(); if (ch_sampled != nullptr) { co_await ch_sampled->send( diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 7bbbf0fe3621..5a19c24e36d2 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -7,7 +7,6 @@ #include #include -#include #include #include @@ -79,7 +78,7 @@ rapidsmpf::streaming::Actor bloom_filter::build( // kernels doing that concurrently because the updates are atomic. build_event.stream_wait(chunk.stream()); filter.add(chunk.table_view(), chunk.stream(), mr); - cudf_streaming::detail::cuda_stream_join(filter_stream, chunk.stream(), &event); + rapidsmpf::cuda_stream_join(filter_stream, chunk.stream(), &event); } if (comm_->nranks() > 1) { auto reducer = rapidsmpf::streaming::AllReduce( @@ -131,7 +130,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( ctx_, -rapidsmpf::safe_cast(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE))); auto chunk_stream = chunk.stream(); - cudf_streaming::detail::cuda_stream_join(chunk_stream, stream, &event); + rapidsmpf::cuda_stream_join(chunk_stream, stream, &event); // Reservation for the mask construction and guess at output size. auto res = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE) ->reserve_or_wait(rapidsmpf::safe_cast(chunk.table_view().num_rows()) @@ -143,7 +142,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( 0); auto mask = filter.contains(chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr()); - cudf_streaming::detail::cuda_stream_join(stream, chunk_stream, &event); + rapidsmpf::cuda_stream_join(stream, chunk_stream, &event); RAPIDSMPF_EXPECTS(mask.size() == static_cast(chunk.table_view().num_rows()), "Invalid mask size"); auto mask_view = cudf::column_view{cudf::data_type{cudf::type_id::BOOL8}, diff --git a/cpp/libcudf_streaming/src/channel_metadata.cpp b/cpp/libcudf_streaming/src/channel_metadata.cpp index 26cfc5de05f4..fe31d81f5e94 100644 --- a/cpp/libcudf_streaming/src/channel_metadata.cpp +++ b/cpp/libcudf_streaming/src/channel_metadata.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include @@ -70,7 +69,7 @@ bool ordering::boundaries_aligned_with(ordering const& other, rapidsmpf::BufferR auto const lhs = boundaries->table_view(); auto const rhs = other.boundaries->table_view(); auto const stream = boundaries->stream(); - cudf_streaming::detail::cuda_stream_join(stream, other.boundaries->stream()); + rapidsmpf::cuda_stream_join(stream, other.boundaries->stream()); for (cudf::size_type i = 0; i < lhs.num_columns(); ++i) { auto eq = cudf::binary_operation(lhs.column(i), rhs.column(i), diff --git a/cpp/libcudf_streaming/src/parquet.cpp b/cpp/libcudf_streaming/src/parquet.cpp index c046768f22ae..d166251a096b 100644 --- a/cpp/libcudf_streaming/src/parquet.cpp +++ b/cpp/libcudf_streaming/src/parquet.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include @@ -320,7 +319,7 @@ rapidsmpf::streaming::Actor read_parquet(std::shared_ptrfilter); // Let's just join all the possible streams here rather than inducing cross-stream // deps in the tasks - cudf_streaming::detail::cuda_stream_join( + rapidsmpf::cuda_stream_join( std::ranges::transform_view( std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool()->get_pool_size()), [&](auto i) { return ctx->br()->stream_pool()->get_stream(i); }), @@ -399,7 +398,7 @@ rapidsmpf::streaming::Actor read_parquet(std::shared_ptrstream), std::ranges::transform_view( std::ranges::iota_view(std::size_t{0}, ctx->br()->stream_pool()->get_pool_size()), diff --git a/cpp/libcudf_streaming/src/partition_utils.cpp b/cpp/libcudf_streaming/src/partition_utils.cpp index c788c9fb43a1..f717eda99d8d 100644 --- a/cpp/libcudf_streaming/src/partition_utils.cpp +++ b/cpp/libcudf_streaming/src/partition_utils.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include @@ -177,7 +176,7 @@ std::unique_ptr unpack_and_concat(std::vectorset_stream(stream); } From bae2ece0e75ab9b1dd96881f550e6f38f21e8752 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 28 Aug 2026 08:53:15 +0000 Subject: [PATCH 8/8] Use explicit cuda::stream_ref when obtaining stream from stream pool The RMM stream pool has not yet been migrated to return cuda::stream_ref, so we must explicitly ask for a stream_ref when obtaining fresh streams. --- .../benchmarks/streaming/ndsh/concatenate.cpp | 2 +- .../benchmarks/streaming/ndsh/join.cpp | 2 +- cpp/libcudf_streaming/src/approx_distinct_count.cpp | 12 ++++++------ cpp/libcudf_streaming/src/bloom_filter.cpp | 10 +++++----- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp index 57d631d9d9a6..16382695a981 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp @@ -31,7 +31,7 @@ streaming::Actor concatenate(std::shared_ptr ctx, CudaEvent event; std::vector messages; ctx->logger()->print("Concatenate"); - auto concat_stream = ctx->br()->stream_pool()->get_stream(); + cuda::stream_ref concat_stream = ctx->br()->stream_pool()->get_stream(); while (!ch_out->is_shutdown()) { co_await ctx->executor()->schedule(); auto msg = co_await ch_in->receive(); diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp index 412462247e04..6cb43c146a2b 100644 --- a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -58,7 +58,7 @@ coro::task broadcast(std::shared_ptr ctx if (comm->nranks() == 1) { std::vector chunks; std::vector views; - auto gather_stream = ctx->br()->stream_pool()->get_stream(); + cuda::stream_ref gather_stream = ctx->br()->stream_pool()->get_stream(); while (true) { auto msg = co_await ch_in->receive(); if (msg.empty()) { break; } diff --git a/cpp/libcudf_streaming/src/approx_distinct_count.cpp b/cpp/libcudf_streaming/src/approx_distinct_count.cpp index 3dcbca000539..2cff62acd9e8 100644 --- a/cpp/libcudf_streaming/src/approx_distinct_count.cpp +++ b/cpp/libcudf_streaming/src/approx_distinct_count.cpp @@ -75,16 +75,16 @@ rapidsmpf::streaming::Actor cardinality_estimator::estimate( if (ch_sampled != nullptr) { co_await ch_sampled->shutdown_metadata(); } co_await ch_out->shutdown_metadata(); - auto const& br = ctx_->br(); - auto const sketch_stream = br->stream_pool()->get_stream(); - auto const sketch_bytes = cudf::approx_distinct_count::sketch_bytes(precision_); - auto const row_count_offset = sketch_bytes; - auto const storage_bytes = sketch_bytes + sizeof(std::uint64_t); + auto const& br = ctx_->br(); + cuda::stream_ref const sketch_stream = br->stream_pool()->get_stream(); + auto const sketch_bytes = cudf::approx_distinct_count::sketch_bytes(precision_); + auto const row_count_offset = sketch_bytes; + auto const storage_bytes = sketch_bytes + sizeof(std::uint64_t); auto reservation = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE)->reserve_or_wait(storage_bytes, 0); auto buf = rmm::device_buffer( storage_bytes, cudf::approx_distinct_count::sketch_alignment(), sketch_stream, br->device_mr()); - RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(buf.data(), 0, storage_bytes, sketch_stream)); + RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(buf.data(), 0, storage_bytes, sketch_stream.get())); reservation.clear(); rapidsmpf::CudaEvent init_event; rapidsmpf::CudaEvent add_event; diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 5a19c24e36d2..6383304cece7 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -48,13 +48,13 @@ rapidsmpf::streaming::Actor bloom_filter::build( co_await ctx_->executor()->schedule(); co_await ch_in->shutdown_metadata(); co_await ch_out->shutdown_metadata(); - auto const& br = ctx_->br(); - auto mr = br->device_mr(); - auto filter_stream = br->stream_pool()->get_stream(); + auto const& br = ctx_->br(); + auto mr = br->device_mr(); + cuda::stream_ref const filter_stream = br->stream_pool()->get_stream(); rapidsmpf::CudaEvent event; auto storage = cudf_streaming::detail::device_bloom_filter::storage(filter_size_, filter_stream, mr); - RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream)); + RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream.get())); auto filter = cudf_streaming::detail::device_bloom_filter(filter_size_, seed_, storage->data()); rapidsmpf::CudaEvent build_event; build_event.record(filter_stream); @@ -117,7 +117,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( auto storage = (co_await bloom_filter->receive()).release(); RAPIDSMPF_EXPECTS((co_await bloom_filter->receive()).empty(), "Bloom filter channel contained more than one message"); - auto stream = storage.stream(); + cuda::stream_ref const stream = storage.stream(); rapidsmpf::CudaEvent event; auto filter = cudf_streaming::detail::device_bloom_filter(filter_size_, seed_, storage.data()); auto meta = co_await ch_in->receive_metadata();