Use cuda::stream_ref in libcudf_streaming

cpp/libcudf_streaming/benchmarks/bench_pack.cpp

@@ -11,11 +11,12 @@
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/cuda_device.hpp>
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/mr/cuda_async_memory_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
+#include <cuda/stream>
+
 #include <benchmark/benchmark.h>
 #include <rapidsmpf/memory/cuda_memcpy_async.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
@@ -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<std::int64_t>(state.iterations()) *
@@ -70,7 +71,7 @@ static void BM_Pack_device(benchmark::State& state)
 {
   auto const table_size_mb = static_cast<std::size_t>(state.range(0));
 
-  rmm::cuda_stream_view stream = rmm::cuda_stream_default;
+  cuda::stream_ref stream = cudf::get_default_stream();
 
   // 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<std::size_t>(state.range(0));
 
-  rmm::cuda_stream_view stream = rmm::cuda_stream_default;
+  cuda::stream_ref stream = cudf::get_default_stream();
 
   // 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<cudf::size_type>(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<std::int64_t>(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 = cudf::get_default_stream();
 
   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 = cudf::get_default_stream();
 
   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<std::size_t>(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 = cudf::get_default_stream();
 
   // 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<std::size_t>(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 = cudf::get_default_stream();
 
   rmm::mr::pool_memory_resource pool_mr{
     rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40)

cpp/libcudf_streaming/benchmarks/bench_partition.cpp

@@ -11,20 +11,20 @@
 
 #include <cudf_streaming/partition_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
+#include <cuda/stream>
+
 #include <benchmark/benchmark.h>
 #include <rapidsmpf/utils/misc.hpp>
 
 #include <memory>
 #include <vector>
 
 // Helper function to create a table with a single int column
-std::unique_ptr<cudf::table> create_int_table(cudf::size_type num_rows,
-                                              rmm::cuda_stream_view stream)
+std::unique_ptr<cudf::table> create_int_table(cudf::size_type num_rows, cuda::stream_ref stream)
 {
   auto data =
     rmm::device_buffer(rapidsmpf::safe_cast<std::size_t>(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 = cudf::get_default_stream();
 
   // Get total GPU memory
   cudaDeviceProp prop;
@@ -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
@@ -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 = cudf::get_default_stream();
 
   // Get total GPU memory
   cudaDeviceProp prop;
@@ -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

cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp

@@ -259,7 +259,7 @@ void barrier(std::shared_ptr<rapidsmpf::Communicator>& comm)
 rapidsmpf::Duration do_run(rapidsmpf::shuffler::PartID const total_num_partitions,
                            std::shared_ptr<rapidsmpf::Communicator>& comm,
                            ArgumentParser const& args,
-                           rmm::cuda_stream_view stream,
+                           cuda::stream_ref stream,
                            rapidsmpf::BufferResource* br,
                            std::shared_ptr<rapidsmpf::Statistics> 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 <typename TransformFn,
           typename InputPartitionsT =
             std::remove_reference_t<std::invoke_result_t<TransformFn, cudf::table&&>>>
 std::vector<InputPartitionsT> 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<InputPartitionsT> 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<rapidsmpf::Communicator>& comm,
                                               ArgumentParser const& args,
-                                              rmm::cuda_stream_view stream,
+                                              cuda::stream_ref stream,
                                               rapidsmpf::BufferResource* br,
                                               std::shared_ptr<rapidsmpf::Statistics> statistics)
 {
@@ -437,7 +437,7 @@ rapidsmpf::Duration run_hash_partition_inline(std::shared_ptr<rapidsmpf::Communi
 rapidsmpf::Duration run_hash_partition_with_datagen(
   std::shared_ptr<rapidsmpf::Communicator>& comm,
   ArgumentParser const& args,
-  rmm::cuda_stream_view stream,
+  cuda::stream_ref stream,
   rapidsmpf::BufferResource* br,
   std::shared_ptr<rapidsmpf::Statistics> 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.
   {

cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp

@@ -215,7 +215,7 @@ rapidsmpf::streaming::Actor consumer(std::shared_ptr<rapidsmpf::streaming::Conte
 rapidsmpf::Duration run(std::shared_ptr<rapidsmpf::streaming::Context> ctx,
                         std::shared_ptr<rapidsmpf::Communicator> 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.
   {

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<Context> ctx,
-                                    rmm::cuda_stream_view stream,
+                                    cuda::stream_ref stream,
                                     std::shared_ptr<Channel> ch_out,
                                     std::uint64_t num_blocks,
                                     cudf::size_type ncolumns,

cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp

@@ -31,7 +31,7 @@ streaming::Actor concatenate(std::shared_ptr<streaming::Context> ctx,
   CudaEvent event;
   std::vector<streaming::Message> 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();
@@ -55,15 +55,16 @@ streaming::Actor concatenate(std::shared_ptr<streaming::Context> ctx,
     views.reserve(messages.size());
     for (auto&& msg : messages) {
       auto chunk = co_await msg.release<cudf_streaming::table_chunk>().make_available(ctx);
-      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_streaming::table_chunk>(
       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);
+    rapidsmpf::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)));
   }

cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp

@@ -19,7 +19,7 @@
 #include <cudf_streaming/partition_utils.hpp>
 #include <cudf_streaming/table_chunk.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/cuda_event.hpp>
@@ -58,12 +58,12 @@ coro::task<streaming::Message> broadcast(std::shared_ptr<streaming::Context> ctx
   if (comm->nranks() == 1) {
     std::vector<cudf_streaming::table_chunk> chunks;
     std::vector<cudf::table_view> 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; }
       auto chunk = co_await msg.release<cudf_streaming::table_chunk>().make_available(ctx);
-      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));
     }
@@ -74,9 +74,10 @@ coro::task<streaming::Message> broadcast(std::shared_ptr<streaming::Context> 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);
+      rapidsmpf::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<cudf_streaming::table_chunk>(std::move(result), gather_stream));
     }
@@ -176,7 +177,7 @@ streaming::Message semi_join_chunk(std::shared_ptr<streaming::Context> ctx,
 
   auto result_table = std::make_unique<cudf::table>(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);
+  rapidsmpf::cuda_stream_join(left_chunk.stream(), chunk_stream);
 
   return to_message(
     sequence, std::make_unique<cudf_streaming::table_chunk>(std::move(result_table), chunk_stream));
@@ -204,7 +205,7 @@ streaming::Message inner_join_chunk(std::shared_ptr<streaming::Context> ctx,
                                     cudf::hash_join& joiner,
                                     cudf::table_view build_carrier,
                                     std::vector<cudf::size_type> right_on,
-                                    rmm::cuda_stream_view build_stream,
+                                    cuda::stream_ref build_stream,
                                     CudaEvent* build_event,
                                     CudaEvent* tmp_event
 
@@ -242,7 +243,7 @@ streaming::Message inner_join_chunk(std::shared_ptr<streaming::Context> 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);
+  rapidsmpf::cuda_stream_join(build_stream, chunk_stream, tmp_event);
   return to_message(sequence,
                     std::make_unique<cudf_streaming::table_chunk>(
                       std::make_unique<cudf::table>(std::move(result_columns)), chunk_stream));

cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp

@@ -51,9 +51,9 @@ rapidsmpf::streaming::Actor write_parquet(std::shared_ptr<rapidsmpf::streaming::
     table = chunk.table_view();
     RAPIDSMPF_EXPECTS(static_cast<std::size_t>(table.num_columns()) == column_names.size(),
                       "Mismatching number of column names and chunk columns");
-    cuda_stream_join(write_stream, chunk.stream(), &event);
+    rapidsmpf::cuda_stream_join(write_stream, chunk.stream(), &event);
     writer.write(table);
-    cuda_stream_join(chunk.stream(), write_stream, &event);
+    rapidsmpf::cuda_stream_join(chunk.stream(), write_stream, &event);
   }
   writer.close();
 }

cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp

@@ -232,7 +232,7 @@ rapidsmpf::streaming::Actor top_k_by(std::shared_ptr<rapidsmpf::streaming::Conte
 
   co_await ctx->executor()->schedule();
   std::vector<std::unique_ptr<cudf::table>> partials;
-  std::vector<rmm::cuda_stream_view> chunk_streams;
+  std::vector<cuda::stream_ref> chunk_streams;
   while (true) {
     auto msg = co_await ch_in->receive();
     if (msg.empty()) { break; }