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,7 +69,7 @@ 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_));
@@ -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)}

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,7 +97,7 @@ 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),
@@ -123,10 +124,7 @@ 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),

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,7 +70,7 @@ 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.get()),
@@ -94,7 +95,7 @@ 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.get()),

cpp/include/raft/core/detail/nvtx.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -8,8 +8,6 @@
 #include <raft/core/detail/macros.hpp>
 #include <raft/core/detail/nvtx_range_stack.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
-
 #ifdef NVTX_ENABLED
 
 #include <nvtx3/nvToolsExt.h>

cpp/include/raft/core/device_container_policy.hpp

@@ -20,11 +20,11 @@
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/stream>
 #include <thrust/device_ptr.h>
 
 namespace RAFT_EXPORT raft {
@@ -43,11 +43,10 @@ class device_reference {
 
  private:
   std::conditional_t<std::is_const<T>::value, const_pointer, pointer> ptr_;
-  rmm::cuda_stream_view stream_;
+  cuda::stream_ref stream_;
 
  public:
-  device_reference(thrust::device_ptr<T> ptr, rmm::cuda_stream_view stream)
-    : ptr_{ptr}, stream_{stream}
+  device_reference(thrust::device_ptr<T> ptr, cuda::stream_ref stream) : ptr_{ptr}, stream_{stream}
   {
   }
 
@@ -103,12 +102,12 @@ class device_uvector {
   /**
    * @brief Ctor that accepts a size, stream and an optional mr.
    */
-  explicit device_uvector(std::size_t size, rmm::cuda_stream_view stream) : data_{size, stream} {}
+  explicit device_uvector(std::size_t size, cuda::stream_ref stream) : data_{size, stream} {}
   /**
    * @brief Ctor that accepts a size, stream and a memory resource.
    */
   explicit device_uvector(std::size_t size,
-                          rmm::cuda_stream_view stream,
+                          cuda::stream_ref stream,
                           rmm::device_async_resource_ref mr)
     : data_{size, stream, mr}
   {

cpp/include/raft/core/device_resources.hpp

@@ -76,7 +76,7 @@ class device_resources : public resources {
    * @param[in] allocation_limit the total amount of memory in bytes available to the temporary
    *            workspace resources.
    */
-  device_resources(rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread},
+  device_resources(cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread},
                    std::shared_ptr<rmm::cuda_stream_pool> stream_pool          = {nullptr},
                    std::optional<raft::mr::device_resource> workspace_resource = std::nullopt,
                    std::optional<std::size_t> allocation_limit                 = std::nullopt)
@@ -119,7 +119,7 @@ class device_resources : public resources {
    * @param[in] stream stream to synchronize
    * @param[in] location the call site to blame for the errors; leave at its default
    */
-  void sync_stream(rmm::cuda_stream_view stream,
+  void sync_stream(cuda::stream_ref stream,
                    std::source_location location = std::source_location::current()) const
   {
     resource::sync_stream(*this, stream, location);
@@ -138,7 +138,7 @@ class device_resources : public resources {
   /**
    * @brief returns main stream on the current container
    */
-  rmm::cuda_stream_view get_stream() const { return resource::get_cuda_stream(*this); }
+  cuda::stream_ref get_stream() const { return resource::get_cuda_stream(*this); }
 
   /**
    * @brief returns whether stream pool was initialized on the current container
@@ -159,23 +159,23 @@ class device_resources : public resources {
   /**
    * @brief return stream from pool
    */
-  rmm::cuda_stream_view get_stream_from_stream_pool() const
+  cuda::stream_ref get_stream_from_stream_pool() const
   {
     return resource::get_stream_from_stream_pool(*this);
   }
 
   /**
    * @brief return stream from pool at index
    */
-  rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const
+  cuda::stream_ref get_stream_from_stream_pool(std::size_t stream_idx) const
   {
     return resource::get_stream_from_stream_pool(*this, stream_idx);
   }
 
   /**
    * @brief return stream from pool if size > 0, else main stream on current container
    */
-  rmm::cuda_stream_view get_next_usable_stream() const
+  cuda::stream_ref get_next_usable_stream() const
   {
     return resource::get_next_usable_stream(*this);
   }
@@ -185,7 +185,7 @@ class device_resources : public resources {
    *
    * @param[in] stream_idx the required index of the stream in the stream pool if available
    */
-  rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const
+  cuda::stream_ref get_next_usable_stream(std::size_t stream_idx) const
   {
     return resource::get_next_usable_stream(*this, stream_idx);
   }

cpp/include/raft/core/dry_run_resources.hpp

@@ -16,7 +16,6 @@
 #include <raft/mr/host_device_resource.hpp>
 #include <raft/mr/host_memory_resource.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 

cpp/include/raft/core/handle.hpp

@@ -43,7 +43,7 @@ class handle_t : public raft::device_resources {
    * @param[in] workspace_resource an optional resource used by some functions for allocating
    *            temporary workspaces.
    */
-  handle_t(rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread},
+  handle_t(cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread},
            std::shared_ptr<rmm::cuda_stream_pool> stream_pool          = {nullptr},
            std::optional<raft::mr::device_resource> workspace_resource = std::nullopt)
     : device_resources{stream_view, stream_pool, std::move(workspace_resource)}

cpp/include/raft/core/interruptible.hpp

@@ -12,7 +12,7 @@
 #include <raft/core/error.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 
 #include <atomic>
 #include <memory>
@@ -76,7 +76,7 @@ class interruptible {
    * thread before the currently captured work has been finished.
    * @throw raft::cuda_error if another CUDA error happens.
    */
-  static inline void synchronize(rmm::cuda_stream_view stream,
+  static inline void synchronize(cuda::stream_ref stream,
                                  std::source_location location = std::source_location::current())
   {
     get_token()->synchronize_impl(cudaStreamQuery, stream.get(), "cudaStreamQuery", location);

cpp/include/raft/core/memory_stats_resources.hpp

@@ -14,7 +14,6 @@
 #include <raft/mr/host_memory_resource.hpp>
 #include <raft/mr/statistics_adaptor.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 

cpp/include/raft/core/memory_tracking_resources.hpp

@@ -17,7 +17,6 @@
 #include <raft/mr/resource_monitor.hpp>
 #include <raft/mr/statistics_adaptor.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 

cpp/include/raft/core/resource/cublas_handle.hpp

@@ -17,7 +17,7 @@ namespace resource {
 
 class cublas_resource : public resource {
  public:
-  cublas_resource(rmm::cuda_stream_view stream)
+  cublas_resource(cuda::stream_ref stream)
   {
     RAFT_CUBLAS_TRY_NO_THROW(cublasCreate(&cublas_res));
     RAFT_CUBLAS_TRY_NO_THROW(cublasSetStream(cublas_res, stream.get()));
@@ -38,12 +38,12 @@ class cublas_resource : public resource {
  */
 class cublas_resource_factory : public resource_factory {
  public:
-  cublas_resource_factory(rmm::cuda_stream_view stream) : stream_(stream) {}
+  cublas_resource_factory(cuda::stream_ref stream) : stream_(stream) {}
   resource_type get_resource_type() override { return resource_type::CUBLAS_HANDLE; }
   resource* make_resource() override { return new cublas_resource(stream_); }
 
  private:
-  rmm::cuda_stream_view stream_;
+  cuda::stream_ref stream_;
 };
 
 /**

cpp/include/raft/core/resource/cuda_stream.hpp

@@ -11,8 +11,6 @@
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
-
 #include <cuda/stream>
 #include <cuda_runtime.h>
 
@@ -22,7 +20,7 @@ namespace RAFT_EXPORT raft {
 namespace resource {
 class cuda_stream_resource : public resource {
  public:
-  cuda_stream_resource(rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread})
+  cuda_stream_resource(cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread})
     : stream(stream_view)
   {
   }
@@ -31,7 +29,7 @@ class cuda_stream_resource : public resource {
   ~cuda_stream_resource() override {}
 
  private:
-  rmm::cuda_stream_view stream;
+  cuda::stream_ref stream;
 };
 
 /**
@@ -40,43 +38,41 @@ class cuda_stream_resource : public resource {
  */
 class cuda_stream_resource_factory : public resource_factory {
  public:
-  cuda_stream_resource_factory(
-    rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread})
+  cuda_stream_resource_factory(cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread})
     : stream(stream_view)
   {
   }
   resource_type get_resource_type() override { return resource_type::CUDA_STREAM_VIEW; }
   resource* make_resource() override { return new cuda_stream_resource(stream); }
 
  private:
-  rmm::cuda_stream_view stream;
+  cuda::stream_ref stream;
 };
 
 /**
  * @defgroup resource_cuda_stream CUDA stream resource functions
  * @{
  */
 /**
- * Load a rmm::cuda_stream_view from a resources instance (and populate it on the res
+ * Load a cuda::stream_ref from a resources instance (and populate it on the res
  * if needed).
  * @param res raft res object for managing resources
  * @return
  */
-inline rmm::cuda_stream_view get_cuda_stream(resources const& res)
+inline cuda::stream_ref get_cuda_stream(resources const& res)
 {
   if (!res.has_resource_factory(resource_type::CUDA_STREAM_VIEW)) {
     res.ensure_default_factory(std::make_shared<cuda_stream_resource_factory>());
   }
-  return *res.get_resource<rmm::cuda_stream_view>(resource_type::CUDA_STREAM_VIEW);
+  return *res.get_resource<cuda::stream_ref>(resource_type::CUDA_STREAM_VIEW);
 };
 
 /**
- * Load a rmm::cuda_stream_view from a resources instance (and populate it on the res
- * if needed).
+ * Set a cuda::stream_ref on a resources instance.
  * @param[in] res raft resources object for managing resources
- * @param[in] stream_view cuda stream view
+ * @param[in] stream_view cuda stream reference
  */
-inline void set_cuda_stream(resources& res, rmm::cuda_stream_view stream_view)
+inline void set_cuda_stream(resources& res, cuda::stream_ref stream_view)
 {
   res.add_resource_factory(std::make_shared<cuda_stream_resource_factory>(stream_view));
 };
@@ -90,7 +86,7 @@ inline void set_cuda_stream(resources& res, rmm::cuda_stream_view stream_view)
  * synchronizing on behalf of a caller, in which case forward the caller's location.
  */
 inline void sync_stream(const resources& res,
-                        rmm::cuda_stream_view stream,
+                        cuda::stream_ref stream,
                         std::source_location location = std::source_location::current())
 {
   if (raft::resource::get_dry_run_flag(res)) { return; }