Migrate stream APIs from rmm::cuda_stream_view to cuda::stream_ref

cpp/bench/prims/common/benchmark.hpp

@@ -15,12 +15,13 @@
 
 #include <rmm/cuda_device.hpp>
 #include <rmm/cuda_stream.hpp>
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
+#include <cuda/stream>
+
 #include <benchmark/benchmark.h>
 
 #include <memory>
@@ -59,7 +60,7 @@ struct using_pool_memory_res {
 struct cuda_event_timer {
  private:
   ::benchmark::State* state_;
-  rmm::cuda_stream_view stream_;
+  cuda::stream_ref stream_;
   cudaEvent_t start_;
   cudaEvent_t stop_;
 
@@ -68,13 +69,13 @@ struct cuda_event_timer {
    * @param state  the benchmark::State whose timer we are going to update.
    * @param stream CUDA stream we are measuring time on.
    */
-  cuda_event_timer(::benchmark::State& state, rmm::cuda_stream_view stream)
+  cuda_event_timer(::benchmark::State& state, cuda::stream_ref stream)
     : state_(&state), stream_(stream)
   {
     RAFT_CUDA_TRY(cudaEventCreate(&start_));
     RAFT_CUDA_TRY(cudaEventCreate(&stop_));
     raft::interruptible::synchronize(stream_);
-    RAFT_CUDA_TRY(cudaEventRecord(start_, stream_));
+    RAFT_CUDA_TRY(cudaEventRecord(start_, stream_.get()));
   }
   cuda_event_timer() = delete;
 
@@ -85,7 +86,7 @@ struct cuda_event_timer {
    */
   ~cuda_event_timer()
   {
-    RAFT_CUDA_TRY_NO_THROW(cudaEventRecord(stop_, stream_));
+    RAFT_CUDA_TRY_NO_THROW(cudaEventRecord(stop_, stream_.get()));
     raft::interruptible::synchronize(stop_);
     float milliseconds = 0.0f;
     RAFT_CUDA_TRY_NO_THROW(cudaEventElapsedTime(&milliseconds, start_, stop_));
@@ -102,7 +103,7 @@ class fixture {
 
  public:
   raft::device_resources handle;
-  rmm::cuda_stream_view stream;
+  cuda::stream_ref stream;
 
   explicit fixture(bool use_pool_memory_resource = false)
     : stream{resource::get_cuda_stream(handle)}
@@ -137,7 +138,7 @@ class fixture {
   /** The helper that writes zeroes to some buffer in GPU memory to flush the L2 cache.  */
   void flush_L2_cache()
   {
-    RAFT_CUDA_TRY(cudaMemsetAsync(scratch_buf_.data(), 0, scratch_buf_.size(), stream));
+    RAFT_CUDA_TRY(cudaMemsetAsync(scratch_buf_.data(), 0, scratch_buf_.size(), stream.get()));
   }
 
   /**

cpp/bench/prims/core/memory_tracking.cu

@@ -10,7 +10,6 @@
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resources.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <unistd.h>

cpp/bench/prims/linalg/add.cu

@@ -22,7 +22,7 @@ struct add : public fixture {
   void run_benchmark(::benchmark::State& state) override
   {
     loop_on_state(state, [this]() {
-      raft::linalg::add(ptr0.data(), ptr0.data(), ptr1.data(), params.len, stream);
+      raft::linalg::add(ptr0.data(), ptr0.data(), ptr1.data(), params.len, stream.get());
     });
   }
 

cpp/bench/prims/linalg/map_then_reduce.cu

@@ -27,7 +27,8 @@ struct map_then_reduce : public fixture {
   void run_benchmark(::benchmark::State& state) override
   {
     loop_on_state(state, [this]() {
-      raft::linalg::mapThenSumReduce(out.data(), params.len, Identity<T>(), stream, in.data());
+      raft::linalg::mapThenSumReduce(
+        out.data(), params.len, Identity<T>(), stream.get(), in.data());
     });
   }
 

cpp/bench/prims/linalg/matrix_vector_op.cu

@@ -54,7 +54,7 @@ struct mat_vec_op : public fixture {
                                                      params.cols,
                                                      params.rows,
                                                      OpT{},
-                                                     stream);
+                                                     stream.get());
           } else {
             raft::linalg::matrixVectorOp<true, false>(out.data() + params.outAlignOffset,
                                                       in.data() + params.inAlignOffset,
@@ -63,7 +63,7 @@ struct mat_vec_op : public fixture {
                                                       params.cols,
                                                       params.rows,
                                                       OpT{},
-                                                      stream);
+                                                      stream.get());
           }
         } else {
           if (params.bcastAlongRows) {
@@ -74,7 +74,7 @@ struct mat_vec_op : public fixture {
                                                       params.cols,
                                                       params.rows,
                                                       OpT{},
-                                                      stream);
+                                                      stream.get());
           } else {
             raft::linalg::matrixVectorOp<false, false>(out.data() + params.outAlignOffset,
                                                        in.data() + params.inAlignOffset,
@@ -83,7 +83,7 @@ struct mat_vec_op : public fixture {
                                                        params.cols,
                                                        params.rows,
                                                        OpT{},
-                                                       stream);
+                                                       stream.get());
           }
         }
       } else {
@@ -95,15 +95,15 @@ struct mat_vec_op : public fixture {
                                                      params.cols,
                                                      params.rows,
                                                      OpT{},
-                                                     stream);
+                                                     stream.get());
           } else {
             raft::linalg::matrixVectorOp<true, false>(out.data() + params.outAlignOffset,
                                                       in.data() + params.inAlignOffset,
                                                       vec1.data(),
                                                       params.cols,
                                                       params.rows,
                                                       OpT{},
-                                                      stream);
+                                                      stream.get());
           }
         } else {
           if (params.bcastAlongRows) {
@@ -113,15 +113,15 @@ struct mat_vec_op : public fixture {
                                                       params.cols,
                                                       params.rows,
                                                       OpT{},
-                                                      stream);
+                                                      stream.get());
           } else {
             raft::linalg::matrixVectorOp<false, false>(out.data() + params.outAlignOffset,
                                                        in.data() + params.inAlignOffset,
                                                        vec1.data(),
                                                        params.cols,
                                                        params.rows,
                                                        OpT{},
-                                                       stream);
+                                                       stream.get());
           }
         }
       }

cpp/bench/prims/linalg/reduce.cu

@@ -28,10 +28,10 @@ struct reduce : public fixture {
     loop_on_state(state, [this]() {
       if (along_rows) {
         raft::linalg::reduce<true, true>(
-          out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream);
+          out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream.get());
       } else {
         raft::linalg::reduce<true, false>(
-          out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream);
+          out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream.get());
       }
     });
   }

cpp/bench/prims/linalg/reduce_cols_by_key.cu

@@ -42,8 +42,14 @@ struct reduce_cols_by_key : public fixture {
     state.SetLabel(label_stream.str());
 
     loop_on_state(state, [this]() {
-      raft::linalg::reduce_cols_by_key(
-        in.data(), keys.data(), out.data(), params.rows, params.cols, params.keys, stream, false);
+      raft::linalg::reduce_cols_by_key(in.data(),
+                                       keys.data(),
+                                       out.data(),
+                                       params.rows,
+                                       params.cols,
+                                       params.keys,
+                                       stream.get(),
+                                       false);
     });
   }
 

cpp/bench/prims/linalg/reduce_rows_by_key.cu

@@ -41,7 +41,7 @@ struct reduce_rows_by_key : public fixture {
                                        params.cols,
                                        params.keys,
                                        out.data(),
-                                       stream,
+                                       stream.get(),
                                        false);
     });
   }

cpp/bench/prims/linalg/transpose.cu

@@ -52,7 +52,7 @@ struct TransposeBench : public fixture {
                               output_view.data_handle(),
                               params.rows,
                               params.cols,
-                              handle.get_stream());
+                              handle.get_stream().get());
     });
   }
 

cpp/bench/prims/matrix/select_k.cu

@@ -49,7 +49,7 @@ struct selection : public fixture {
       out_dists_(p.batch_size * p.k, stream),
       out_ids_(p.batch_size * p.k, stream)
   {
-    raft::sparse::iota_fill(in_ids_.data(), IdxT(p.batch_size), IdxT(p.len), stream);
+    raft::sparse::iota_fill(in_ids_.data(), IdxT(p.batch_size), IdxT(p.len), stream.get());
     raft::random::RngState state{42};
 
     KeyT min_value = -1.0;

cpp/bench/prims/random/make_blobs.cu

@@ -42,7 +42,7 @@ struct make_blobs : public fixture {
                                params.rows,
                                params.cols,
                                params.clusters,
-                               this->stream,
+                               this->stream.get(),
                                params.row_major);
     });
   }

cpp/bench/prims/random/permute.cu

@@ -45,7 +45,7 @@ struct permute : public fixture {
                             params.cols,
                             params.rows,
                             params.rowMajor,
-                            stream,
+                            stream.get(),
                             123456ULL);
     });
   }
@@ -97,7 +97,7 @@ struct permute_perms_only : public fixture {
                             IntType(0),
                             IntType(n_rows),
                             true,
-                            stream,
+                            stream.get(),
                             123456ULL);
       bytes_processed += size_t(n_rows) * sizeof(IntType);
     });

cpp/bench/prims/sparse/convert_csr.cu

@@ -56,7 +56,7 @@ struct bench_base : public fixture {
       // col_ind is over-dimensioned because nnz is unknown at this point
       col_ind(p.num_rows * p.num_cols, stream)
   {
-    init_adj(adj.data(), p.num_rows, p.num_cols, p.divisor, stream);
+    init_adj(adj.data(), p.num_rows, p.num_cols, p.divisor, stream.get());
 
     std::vector<index_t> row_ind_host(p.num_rows);
     for (size_t i = 0; i < row_ind_host.size(); ++i) {

cpp/bench/prims/sparse/select_k_csr.cu

@@ -83,7 +83,7 @@ struct SelectKCsrTest : public fixture {
                                                  1,
                                                  nnz,
                                                  1,
-                                                 stream,
+                                                 stream.get(),
                                                  false,
                                                  nullptr,
                                                  nullptr,

cpp/bench/prims/sparse/svds.cu

@@ -254,7 +254,7 @@ class svds_bench_base : public fixture {
                        CUBLAS_OP_N,
                        one,
                        zero,
-                       stream);
+                       stream.get());
 
     raft::linalg::gemm(handle,
                        Vt.data_handle(),
@@ -268,7 +268,7 @@ class svds_bench_base : public fixture {
                        CUBLAS_OP_T,
                        one,
                        zero,
-                       stream);
+                       stream.get());
 
     std::vector<value_t> h_utu(static_cast<std::size_t>(params.k) * params.k);
     std::vector<value_t> h_vvt(static_cast<std::size_t>(params.k) * params.k);

cpp/bench/prims/util/fast_int_div.cu

@@ -75,20 +75,20 @@ struct fast_int_div_bench : public fixture {
                                   h_numerators.data(),
                                   h_numerators.size() * sizeof(IntT),
                                   cudaMemcpyHostToDevice,
-                                  stream));
+                                  stream.get()));
     RAFT_CUDA_TRY(cudaMemcpyAsync(d_divisors.data(),
                                   h_divisors.data(),
                                   h_divisors.size() * sizeof(divisor_t),
                                   cudaMemcpyHostToDevice,
-                                  stream));
-    stream.synchronize();
+                                  stream.get()));
+    stream.sync();
   }
 
   void run_benchmark(::benchmark::State& state) override
   {
     const auto* divisors = static_cast<const divisor_t*>(d_divisors.data());
     loop_on_state(state, [this, divisors]() {
-      divmod_kernel<IntT, divisor_t><<<kBlocks, kThreads, 0, stream>>>(
+      divmod_kernel<IntT, divisor_t><<<kBlocks, kThreads, 0, stream.get()>>>(
         d_numerators.data(), kNumNumerators, divisors, kNumDivisors, out_d.data());
       RAFT_CUDA_TRY(cudaPeekAtLastError());
     });

cpp/include/raft/comms/detail/mpi_comms.hpp

@@ -12,9 +12,10 @@
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_scalar.hpp>
 
+#include <cuda/stream>
+
 #include <mpi.h>
 #include <nccl.h>
 
@@ -96,14 +97,14 @@ constexpr MPI_Op get_mpi_op(const op_t op)
 
 class mpi_comms : public comms_iface {
  public:
-  mpi_comms(MPI_Comm comm, const bool owns_mpi_comm, rmm::cuda_stream_view stream)
+  mpi_comms(MPI_Comm comm, const bool owns_mpi_comm, cuda::stream_ref stream)
     : owns_mpi_comm_(owns_mpi_comm),
       mpi_comm_(comm),
       size_(0),
       rank_(1),
       status_(stream),
       next_request_id_(0),
-      stream_(stream)
+      stream_(stream.get())
   {
     int mpi_is_initialized = 0;
     RAFT_MPI_TRY(MPI_Initialized(&mpi_is_initialized));
@@ -123,18 +124,15 @@ class mpi_comms : public comms_iface {
     initialize();
   }
 
-  mpi_comms(MPI_Comm mpi_comm,
-            bool owns_mpi_comm,
-            ncclComm_t nccl_comm,
-            rmm::cuda_stream_view stream)
+  mpi_comms(MPI_Comm mpi_comm, bool owns_mpi_comm, ncclComm_t nccl_comm, cuda::stream_ref stream)
     : owns_mpi_comm_(owns_mpi_comm),
       mpi_comm_(mpi_comm),
       nccl_comm_(nccl_comm),
       size_(0),
       rank_(1),
       status_(stream),
       next_request_id_(0),
-      stream_(stream)
+      stream_(stream.get())
   {
     int mpi_is_initialized = 0;
     RAFT_MPI_TRY(MPI_Initialized(&mpi_is_initialized));

cpp/include/raft/comms/detail/std_comms.hpp

@@ -16,6 +16,7 @@
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/stream>
 #include <cuda_runtime.h>
 #include <thrust/iterator/zip_iterator.h>
 
@@ -69,10 +70,10 @@ class std_comms : public comms_iface {
             ucx_objects_t ucx_objects,
             int num_ranks,
             int rank,
-            rmm::cuda_stream_view stream,
+            cuda::stream_ref stream,
             bool subcomms_ucp = true)
     : nccl_comm_(nccl_comm),
-      stream_(stream),
+      stream_(stream.get()),
       status_(stream),
       num_ranks_(num_ranks),
       rank_(rank),
@@ -94,10 +95,10 @@ class std_comms : public comms_iface {
   std_comms(const ncclComm_t nccl_comm,
             int num_ranks,
             int rank,
-            rmm::cuda_stream_view stream,
+            cuda::stream_ref stream,
             bool own_nccl_comm = false)
     : nccl_comm_(nccl_comm),
-      stream_(stream),
+      stream_(stream.get()),
       status_(stream),
       num_ranks_(num_ranks),
       rank_(rank),