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 = cuda::stream_ref{cudaStreamLegacy};
 
   // 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 = cuda::stream_ref{cudaStreamLegacy};
 
   // 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 = cuda::stream_ref{cudaStreamLegacy};
 
   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{cudaStreamLegacy};
 
   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 = cuda::stream_ref{cudaStreamLegacy};
 
   // 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 = cuda::stream_ref{cudaStreamLegacy};
 
   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 = 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});
 
-  rmm::cuda_stream_view stream = rmm::cuda_stream_default;
+  cuda::stream_ref stream = cuda::stream_ref{cudaStreamLegacy};
 
   // Get total GPU memory
   cudaDeviceProp prop;

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/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>
@@ -204,7 +204,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
 

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; }

cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp

@@ -13,9 +13,8 @@
 #include <cudf_streaming/parquet.hpp>
 #include <cudf_streaming/table_chunk.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
-
 #include <cuda/memory_resource>
+#include <cuda/stream>
 
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/memory/buffer_resource.hpp>
@@ -95,7 +94,7 @@ namespace detail {
  * @return Filter expression with proper lifetime management
  */
 template <typename timestamp_type>
-std::unique_ptr<cudf_streaming::filter> make_date_filter(rmm::cuda_stream_view stream,
+std::unique_ptr<cudf_streaming::filter> 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<cudf_streaming::filter> make_date_filter(rmm::cuda_stream_view s
  */
 template <typename timestamp_type>
 std::unique_ptr<cudf_streaming::filter> 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)

cpp/libcudf_streaming/benchmarks/utils/random_data.cu

@@ -26,7 +26,7 @@
 rmm::device_uvector<std::int32_t> 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<std::int32_t> random_device_vector(std::size_t nelem,
 std::unique_ptr<cudf::column> 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<std::unique_ptr<cudf::column>> cols;
@@ -85,7 +85,7 @@ void random_fill(rapidsmpf::Buffer& buffer, rmm::device_async_resource_ref mr)
                                       std::numeric_limits<std::int32_t>::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));
       });

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<std::int32_t> 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<std::int32_t> random_device_vector(std::size_t nelem,
 std::unique_ptr<cudf::column> 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);
 
 /**

cpp/libcudf_streaming/examples/example_shuffle.cpp

@@ -45,7 +45,7 @@ int main(int argc, char** argv)
     std::make_shared<rapidsmpf::MPI>(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);