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

c/src/core/c_api.cpp

@@ -15,7 +15,7 @@
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 #include <rapids_logger/logger.hpp>
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/mr/cuda_async_memory_resource.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
@@ -136,7 +136,7 @@ extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream)
 {
   return cuvs::core::translate_exceptions([=] {
     auto res_ptr = reinterpret_cast<raft::resources*>(res);
-    raft::resource::set_cuda_stream(*res_ptr, static_cast<rmm::cuda_stream_view>(stream));
+    raft::resource::set_cuda_stream(*res_ptr, static_cast<cuda::stream_ref>(stream));
   });
 }
 

cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -18,7 +18,7 @@
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/failure_callback_resource_adaptor.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
@@ -122,8 +122,8 @@ class configured_raft_resources {
    */
   explicit configured_raft_resources(const std::shared_ptr<shared_raft_resources>& shared_res)
     : shared_res_{shared_res},
-      res_{std::make_unique<raft::device_resources>(
-        rmm::cuda_stream_view(get_stream_from_global_pool()))}
+      res_{
+        std::make_unique<raft::device_resources>(cuda::stream_ref(get_stream_from_global_pool()))}
   {
     raft::resource::set_large_workspace_resource(
       *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()});

cpp/include/cuvs/neighbors/scann.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -17,7 +17,6 @@
 #include <raft/core/resource/stream_view.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/integer_utils.hpp>
-#include <rmm/cuda_stream_view.hpp>
 
 #include <cmath>
 #include <cuvs/core/export.hpp>

cpp/include/cuvs/neighbors/vamana.hpp

@@ -16,7 +16,6 @@
 #include <raft/core/resource/stream_view.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/integer_utils.hpp>
-#include <rmm/cuda_stream_view.hpp>
 
 #include <cuvs/core/export.hpp>
 #include <optional>

cpp/internal/cuvs_internal/neighbors/refine_helper.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -15,7 +15,7 @@
 #include <raft/core/resources.hpp>
 #include <raft/random/rng.cuh>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 
 namespace cuvs::neighbors {
@@ -127,7 +127,7 @@ class RefineHelper {
  public:
   RefineInputs<IdxT> p;
   const raft::resources& handle_;
-  rmm::cuda_stream_view stream_;
+  cuda::stream_ref stream_;
 
   raft::device_matrix<DataT, IdxT, row_major> dataset;
   raft::device_matrix<DataT, IdxT, row_major> queries;

cpp/src/neighbors/detail/ann_utils.cuh

@@ -18,7 +18,7 @@
 #include <raft/util/cudart_utils.hpp>
 #include <raft/util/integer_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
@@ -224,7 +224,7 @@ HDI constexpr auto mapping<int8_t>::operator()(const float& x) const -> int8_t
  * @param[in] n_bytes
  */
 template <typename T, typename IdxT>
-inline void memzero(T* ptr, IdxT n_elems, rmm::cuda_stream_view stream)
+inline void memzero(T* ptr, IdxT n_elems, cuda::stream_ref stream)
 {
   switch (check_pointer_residency(ptr)) {
     case pointer_residency::host_and_device:
@@ -298,7 +298,7 @@ void block_copy(const IdxT* in_offsets,
                 const T* in_data,
                 T* out_data,
                 IdxT n_mult,
-                rmm::cuda_stream_view stream)
+                cuda::stream_ref stream)
 {
   IdxT in_size;
   update_host(&in_size, in_offsets + n_blocks, 1, stream);
@@ -325,7 +325,7 @@ void block_copy(const IdxT* in_offsets,
  * @param stream
  */
 template <typename T, typename IdxT>
-void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, rmm::cuda_stream_view stream)
+void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, cuda::stream_ref stream)
 {
   dim3 threads(128, 1, 1);
   dim3 blocks(raft::ceildiv<IdxT>(len_a * len_b, threads.x), 1, 1);
@@ -370,7 +370,7 @@ void copy_selected(IdxT n_rows,
                    IdxT ld_src,
                    T* dst,
                    IdxT ld_dst,
-                   rmm::cuda_stream_view stream)
+                   cuda::stream_ref stream)
 {
   switch (check_pointer_residency(src, dst, row_ids)) {
     case pointer_residency::host_and_device:
@@ -403,8 +403,7 @@ void copy_selected(IdxT n_rows,
  * the main stream itself is returned with `false`, and the caller should treat prefetch as a
  * no-op (no overlap is possible on a single stream).
  */
-inline auto get_prefetch_stream(raft::resources const& res)
-  -> std::pair<rmm::cuda_stream_view, bool>
+inline auto get_prefetch_stream(raft::resources const& res) -> std::pair<cuda::stream_ref, bool>
 {
   if (res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) &&
       raft::resource::get_stream_pool_size(res) >= 1) {
@@ -600,7 +599,7 @@ struct batch_load_iterator {
     batch(raft::resources const& res,
           MdspanT input_view,
           size_type batch_size,
-          rmm::cuda_stream_view copy_stream,
+          cuda::stream_ref copy_stream,
           rmm::device_async_resource_ref mr,
           bool prefetch,
           bool initialize,
@@ -806,7 +805,7 @@ struct batch_load_iterator {
                                     copy_stream_.get()));
     }
 
-    rmm::cuda_stream_view copy_stream_;
+    cuda::stream_ref copy_stream_;
     raft::resources const* res_;
     MdspanT input_view_;
     element_type* source_;
@@ -860,7 +859,7 @@ struct batch_load_iterator {
   batch_load_iterator(raft::resources const& res,
                       MdspanT input_view,
                       size_type batch_size,
-                      rmm::cuda_stream_view copy_stream,
+                      cuda::stream_ref copy_stream,
                       rmm::device_async_resource_ref mr,
                       bool prefetch       = false,
                       bool initialize     = true,
@@ -876,7 +875,7 @@ struct batch_load_iterator {
   batch_load_iterator(raft::resources const& res,
                       MdspanT input_view,
                       size_type batch_size,
-                      rmm::cuda_stream_view copy_stream,
+                      cuda::stream_ref copy_stream,
                       bool prefetch       = false,
                       bool initialize     = true,
                       bool host_writeback = false)
@@ -1025,7 +1024,7 @@ class batch_load_iterator_dyn {
                           IdxT n_rows,
                           IdxT row_width,
                           size_type batch_size,
-                          rmm::cuda_stream_view copy_stream,
+                          cuda::stream_ref copy_stream,
                           rmm::device_async_resource_ref mr,
                           bool prefetch       = false,
                           bool initialize     = true,
@@ -1050,7 +1049,7 @@ class batch_load_iterator_dyn {
                           IdxT n_rows,
                           IdxT row_width,
                           size_type batch_size,
-                          rmm::cuda_stream_view copy_stream,
+                          cuda::stream_ref copy_stream,
                           bool prefetch       = false,
                           bool initialize     = true,
                           bool host_writeback = false)
@@ -1156,7 +1155,7 @@ class batch_load_iterator_dyn {
                         IdxT n_rows,
                         IdxT row_width,
                         size_type batch_size,
-                        rmm::cuda_stream_view copy_stream,
+                        cuda::stream_ref copy_stream,
                         rmm::device_async_resource_ref mr,
                         bool prefetch,
                         bool initialize,
@@ -1219,7 +1218,7 @@ auto make_batch_load_iterator(raft::resources const& res,
                               detail::type_identity_t<IdxT> n_rows,
                               detail::type_identity_t<IdxT> row_width,
                               size_t batch_size,
-                              rmm::cuda_stream_view copy_stream,
+                              cuda::stream_ref copy_stream,
                               rmm::device_async_resource_ref mr,
                               bool prefetch       = false,
                               bool initialize     = true,
@@ -1244,7 +1243,7 @@ auto make_batch_load_iterator(raft::resources const& res,
                               detail::type_identity_t<IdxT> n_rows,
                               detail::type_identity_t<IdxT> row_width,
                               size_t batch_size,
-                              rmm::cuda_stream_view copy_stream,
+                              cuda::stream_ref copy_stream,
                               bool prefetch       = false,
                               bool initialize     = true,
                               bool host_writeback = false) -> batch_load_iterator_dyn<T, IdxT>

cpp/src/neighbors/detail/cagra/compute_distance.hpp

@@ -9,6 +9,7 @@
 #include "hashmap.hpp"
 #include "utils.hpp"
 
+#include <cuda/stream>
 #include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/common.hpp>
@@ -211,9 +212,8 @@ struct dataset_descriptor_host {
   // Codebook type is determined by DataT for VPQ (always half for now)
 
   struct state {
-    using ready_t = std::tuple<dev_descriptor_t*, rmm::cuda_stream_view>;
-    using init_f =
-      std::tuple<std::function<void(dev_descriptor_t*, rmm::cuda_stream_view)>, size_t>;
+    using ready_t = std::tuple<dev_descriptor_t*, cuda::stream_ref>;
+    using init_f  = std::tuple<std::function<void(dev_descriptor_t*, cuda::stream_ref)>, size_t>;
 
     std::mutex mutex;
     std::atomic<bool> ready;  // Not sure if std::holds_alternative is thread-safe
@@ -235,7 +235,7 @@ struct dataset_descriptor_host {
       RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(ready_event));
     }
 
-    void eval(rmm::cuda_stream_view stream)
+    void eval(cuda::stream_ref stream)
     {
       std::lock_guard<std::mutex> lock(mutex);
       if (std::holds_alternative<init_f>(value)) {
@@ -249,7 +249,7 @@ struct dataset_descriptor_host {
       }
     }
 
-    auto get(rmm::cuda_stream_view stream) -> dev_descriptor_t*
+    auto get(cuda::stream_ref stream) -> dev_descriptor_t*
     {
       if (!ready.load(std::memory_order_acquire)) { eval(stream); }
       // value is immutable at this point.
@@ -288,12 +288,12 @@ struct dataset_descriptor_host {
   /**
    * Return the device pointer, possibly evaluating it in the given thread.
    */
-  [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) const -> const dev_descriptor_t*
+  [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) const -> const dev_descriptor_t*
   {
     return value_->get(stream);
   }
 
-  [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) -> dev_descriptor_t*
+  [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) -> dev_descriptor_t*
   {
     return value_->get(stream);
   }

cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh

@@ -6,6 +6,7 @@
 
 #include "compute_distance_standard.hpp"
 
+#include <cuda/stream>
 #include <cuvs/distance/distance.hpp>
 #include <raft/util/pow2_utils.cuh>
 
@@ -131,23 +132,23 @@ standard_descriptor_spec<Metric, TeamSize, DatasetBlockDim, DataT, IndexT, Dista
   RAFT_EXPECTS(Metric != cuvs::distance::DistanceType::CosineExpanded || dataset_norms != nullptr,
                "Dataset norms must be provided for CosineExpanded metric");
 
-  return host_type{desc_type{ptr, size, dim, ld, dataset_norms},
-                   [=](dataset_descriptor_base_t<DataT, IndexT, DistanceT>* dev_ptr,
-                       rmm::cuda_stream_view stream) {
-                     standard_dataset_descriptor_init_kernel<Metric,
-                                                             TeamSize,
-                                                             DatasetBlockDim,
-                                                             DataT,
-                                                             IndexT,
-                                                             DistanceT>
-                       <<<1, 1, 0, stream.get()>>>(dev_ptr, ptr, size, dim, ld, dataset_norms);
-                     RAFT_CUDA_TRY(cudaPeekAtLastError());
-                   },
-                   Metric,
-                   DatasetBlockDim,
-                   false,  // is_vpq
-                   0,      // pq_bits
-                   0};     // pq_len
+  return host_type{
+    desc_type{ptr, size, dim, ld, dataset_norms},
+    [=](dataset_descriptor_base_t<DataT, IndexT, DistanceT>* dev_ptr, cuda::stream_ref stream) {
+      standard_dataset_descriptor_init_kernel<Metric,
+                                              TeamSize,
+                                              DatasetBlockDim,
+                                              DataT,
+                                              IndexT,
+                                              DistanceT>
+        <<<1, 1, 0, stream.get()>>>(dev_ptr, ptr, size, dim, ld, dataset_norms);
+      RAFT_CUDA_TRY(cudaPeekAtLastError());
+    },
+    Metric,
+    DatasetBlockDim,
+    false,  // is_vpq
+    0,      // pq_bits
+    0};     // pq_len
 }
 
 }  // namespace cuvs::neighbors::cagra::detail

cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh

@@ -8,6 +8,7 @@
 #include "compute_distance_vpq.hpp"
 #include "packed_type.hpp"
 
+#include <cuda/stream>
 #include <cuvs/distance/distance.hpp>
 #include <raft/util/pow2_utils.cuh>
 
@@ -221,8 +222,7 @@ vpq_descriptor_spec<Metric,
   return host_type{
     desc_type{
       encoded_dataset_ptr, encoded_dataset_dim, vq_code_book_ptr, pq_code_book_ptr, size, dim},
-    [=](dataset_descriptor_base_t<DataT, IndexT, DistanceT>* dev_ptr,
-        rmm::cuda_stream_view stream) {
+    [=](dataset_descriptor_base_t<DataT, IndexT, DistanceT>* dev_ptr, cuda::stream_ref stream) {
       vpq_dataset_descriptor_init_kernel<Metric,
                                          TeamSize,
                                          DatasetBlockDim,

cpp/src/neighbors/detail/cagra/search_plan.cuh

@@ -7,6 +7,7 @@
 
 #include "hashmap.hpp"
 
+#include <cuda/stream>
 #include <cuvs/neighbors/common.hpp>
 #include <neighbors/detail/cagra/compute_distance-ext.cuh>
 #include <raft/core/resource/cuda_stream.hpp>
@@ -38,8 +39,8 @@ namespace cuvs::neighbors::cagra::detail {
 template <typename T>
 struct lightweight_uvector {
  private:
-  using raft_res_type = const raft::resources*;
-  using rmm_res_type  = std::tuple<rmm::device_async_resource_ref, rmm::cuda_stream_view>;
+  using raft_res_type            = const raft::resources*;
+  using rmm_res_type             = std::tuple<rmm::device_async_resource_ref, cuda::stream_ref>;
   static constexpr size_t kAlign = 256;
 
   std::variant<raft_res_type, rmm_res_type> res_;
@@ -75,14 +76,14 @@ struct lightweight_uvector {
     size_ = new_size;
   }
 
-  void resize(size_t new_size, rmm::cuda_stream_view stream)
+  void resize(size_t new_size, cuda::stream_ref stream)
   {
     if (new_size == size_) { return; }
     if (std::holds_alternative<raft_res_type>(res_)) {
       auto& h = std::get<raft_res_type>(res_);
       res_    = rmm_res_type{raft::resource::get_workspace_resource_ref(*h), stream};
     } else {
-      std::get<rmm::cuda_stream_view>(std::get<rmm_res_type>(res_)) = stream;
+      std::get<cuda::stream_ref>(std::get<rmm_res_type>(res_)) = stream;
     }
     resize(new_size);
   }