Adopt CUDA stream compatibility accessors

cpp/bench/sg/arima_loglikelihood.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,6 +14,7 @@
 
 #include <rmm/device_uvector.hpp>
 
+#include <cuda/stream>
 #include <thrust/execution_policy.h>
 #include <thrust/for_each.h>
 #include <thrust/iterator/counting_iterator.h>
@@ -33,9 +34,9 @@ class ArimaLoglikelihood : public TsFixtureRandom<DataT> {
   ArimaLoglikelihood(const std::string& name, const ArimaParams& p)
     : TsFixtureRandom<DataT>(name, p.data),
       order(p.order),
-      param(0, rmm::cuda_stream_default),
-      loglike(0, rmm::cuda_stream_default),
-      temp_mem(0, rmm::cuda_stream_default)
+      param(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}),
+      loglike(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}),
+      temp_mem(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}})
   {
   }
 
@@ -45,7 +46,7 @@ class ArimaLoglikelihood : public TsFixtureRandom<DataT> {
     using MLCommon::Bench::CudaEventTimer;
 
     auto& handle  = *this->handle;
-    auto stream   = handle.get_stream();
+    auto stream   = handle.get_stream().get();
     auto counting = thrust::make_counting_iterator(0);
 
     // Generate random parameters

cpp/bench/sg/benchmark.cuh

@@ -14,7 +14,7 @@
 #include <raft/core/handle.hpp>
 #include <raft/util/cudart_utils.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <rmm/cuda_stream_pool.hpp>
 
 #include <cuda_runtime.h>
 

cpp/bench/sg/dataset.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,6 +14,8 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <cuda/stream>
+
 #include <fstream>
 #include <iostream>
 #include <sstream>
@@ -65,7 +67,11 @@ struct RegressionParams {
  */
 template <typename D, typename L, typename IdxT = int>
 struct Dataset {
-  Dataset() : X(0, rmm::cuda_stream_default), y(0, rmm::cuda_stream_default) {}
+  Dataset()
+    : X(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}),
+      y(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}})
+  {
+  }
   /** input data */
   rmm::device_uvector<D> X;
   /** labels or output associated with each row of input data */
@@ -97,7 +103,7 @@ struct Dataset {
   void blobs(const raft::handle_t& handle, const DatasetParams& p, const BlobsParams& b)
   {
     const auto& handle_impl = handle;
-    auto stream             = handle_impl.get_stream();
+    auto stream             = handle_impl.get_stream().get();
     auto cublas_handle      = handle_impl.get_cublas_handle();
 
     // Make blobs will generate labels of type IdxT which has to be an integer
@@ -139,7 +145,7 @@ struct Dataset {
   {
     ASSERT(!isClassification(), "make_regression: is only for regression problems!");
     const auto& handle_impl = handle;
-    auto stream             = handle_impl.get_stream();
+    auto stream             = handle_impl.get_stream().get();
     auto cublas_handle      = handle_impl.get_cublas_handle();
     auto cusolver_handle    = handle_impl.get_cusolver_dn_handle();
 

cpp/bench/sg/dataset_ts.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -10,6 +10,8 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <cuda/stream>
+
 namespace ML {
 namespace Bench {
 
@@ -26,7 +28,7 @@ struct TimeSeriesParams {
  */
 template <typename DataT>
 struct TimeSeriesDataset {
-  TimeSeriesDataset() : X(0, rmm::cuda_stream_default) {}
+  TimeSeriesDataset() : X(0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}) {}
 
   /** input data */
   rmm::device_uvector<DataT> X;
@@ -44,7 +46,7 @@ struct TimeSeriesDataset {
               DataT sigma = 1)
   {
     raft::random::Rng gpu_gen(p.seed, raft::random::GenPhilox);
-    gpu_gen.normal(X.data(), p.batch_size * p.n_obs, mu, sigma, handle.get_stream());
+    gpu_gen.normal(X.data(), p.batch_size * p.n_obs, mu, sigma, handle.get_stream().get());
   }
 };
 

cpp/src/arima/batched_arima.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -43,7 +43,7 @@ void pack(raft::handle_t& handle,
           int batch_size,
           double* param_vec)
 {
-  const auto stream = handle.get_stream();
+  const auto stream = handle.get_stream().get();
   params.pack(order, batch_size, param_vec, stream);
 }
 
@@ -53,7 +53,7 @@ void unpack(raft::handle_t& handle,
             int batch_size,
             const double* param_vec)
 {
-  const auto stream = handle.get_stream();
+  const auto stream = handle.get_stream().get();
   params.unpack(order, batch_size, param_vec, stream);
 }
 
@@ -64,7 +64,7 @@ void batched_diff(raft::handle_t& handle,
                   int n_obs,
                   const ARIMAOrder& order)
 {
-  const auto stream = handle.get_stream();
+  const auto stream = handle.get_stream().get();
   MLCommon::TimeSeries::prepare_data(
     d_y_diff, d_y, batch_size, n_obs, order.d, order.D, order.s, stream);
 }
@@ -80,7 +80,7 @@ struct is_missing {
 bool detect_missing(raft::handle_t& handle, const double* d_y, int n_elem)
 {
   return thrust::any_of(
-    thrust::cuda::par.on(handle.get_stream()), d_y, d_y + n_elem, is_missing<double>());
+    thrust::cuda::par.on(handle.get_stream().get()), d_y, d_y + n_elem, is_missing<double>());
 }
 
 void predict(raft::handle_t& handle,
@@ -101,7 +101,7 @@ void predict(raft::handle_t& handle,
              double* d_upper)
 {
   raft::common::nvtx::range fun_scope(__func__);
-  const auto stream = handle.get_stream();
+  const auto stream = handle.get_stream().get();
 
   bool diff     = order.need_diff() && pre_diff && level == 0;
   int num_steps = std::max(end - n_obs, 0);
@@ -356,7 +356,7 @@ void conditional_sum_of_squares(raft::handle_t& handle,
                                 int truncate)
 {
   raft::common::nvtx::range fun_scope(__func__);
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   int n_phi     = order.n_phi();
   int n_theta   = order.n_theta();
@@ -412,7 +412,7 @@ void batched_loglike(raft::handle_t& handle,
 {
   raft::common::nvtx::range fun_scope(__func__);
 
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   double* d_pred = arima_mem.pred;
 
@@ -485,7 +485,7 @@ void batched_loglike(raft::handle_t& handle,
   raft::common::nvtx::range fun_scope(__func__);
 
   // unpack parameters
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   ARIMAParams<double> params = {arima_mem.params_mu,
                                 arima_mem.params_beta,
@@ -527,7 +527,7 @@ void batched_loglike_grad(raft::handle_t& handle,
                           int truncate)
 {
   raft::common::nvtx::range fun_scope(__func__);
-  auto stream   = handle.get_stream();
+  auto stream   = handle.get_stream().get();
   auto counting = thrust::make_counting_iterator(0);
   int N         = order.complexity();
 
@@ -601,7 +601,7 @@ void information_criterion(raft::handle_t& handle,
                            int ic_type)
 {
   raft::common::nvtx::range fun_scope(__func__);
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   /* Compute log-likelihood in d_ic */
   batched_loglike(
@@ -674,7 +674,7 @@ void _arma_least_squares(raft::handle_t& handle,
                          double* d_mu = nullptr)
 {
   const auto& handle_impl = handle;
-  auto stream             = handle_impl.get_stream();
+  auto stream             = handle_impl.get_stream().get();
   auto cublas_handle      = handle_impl.get_cublas_handle();
   auto counting           = thrust::make_counting_iterator(0);
 
@@ -956,7 +956,7 @@ void estimate_x0(raft::handle_t& handle,
 {
   raft::common::nvtx::range fun_scope(__func__);
   const auto& handle_impl = handle;
-  auto stream             = handle_impl.get_stream();
+  auto stream             = handle_impl.get_stream().get();
   auto cublas_handle      = handle_impl.get_cublas_handle();
 
   /// TODO: solve exogenous coefficients with only valid rows instead of interpolation?

cpp/src/arima/batched_kalman.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -850,7 +850,7 @@ void _lyapunov_wrapper(raft::handle_t& handle,
                        int r)
 {
   if (r <= 5) {
-    auto stream       = handle.get_stream();
+    auto stream       = handle.get_stream().get();
     auto cublasHandle = handle.get_cublas_handle();
     int batch_size    = ML::narrow_cast<int>(A.batches());
     int r2            = r * r;
@@ -909,7 +909,7 @@ void _batched_kalman_filter(raft::handle_t& handle,
                             double* d_upper)
 {
   const size_t batch_size = Zb.batches();
-  auto stream             = handle.get_stream();
+  auto stream             = handle.get_stream().get();
   auto cublasHandle       = handle.get_cublas_handle();
 
   auto counting = thrust::make_counting_iterator(0);
@@ -1152,7 +1152,7 @@ void init_batched_kalman_matrices(raft::handle_t& handle,
 {
   raft::common::nvtx::range fun_scope(__func__);
 
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   // Note: Z is unused yet but kept to avoid reintroducing it later when
   // adding support for exogeneous variables
@@ -1265,7 +1265,7 @@ void batched_kalman_filter(raft::handle_t& handle,
   raft::common::nvtx::range fun_scope(__func__);
 
   auto cublasHandle = handle.get_cublas_handle();
-  auto stream       = handle.get_stream();
+  auto stream       = handle.get_stream().get();
 
   // see (3.18) in TSA by D&K
   int rd = order.rd();
@@ -1324,7 +1324,7 @@ void batched_jones_transform(raft::handle_t& handle,
                              double* h_Tparams)
 {
   int N                       = order.complexity();
-  auto stream                 = handle.get_stream();
+  auto stream                 = handle.get_stream().get();
   double* d_params            = arima_mem.d_params;
   double* d_Tparams           = arima_mem.d_Tparams;
   ARIMAParams<double> params  = {arima_mem.params_mu,

cpp/src/datasets/make_arima.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -23,7 +23,7 @@ inline void make_arima_helper(const raft::handle_t& handle,
                               DataT intercept_scale,
                               uint64_t seed)
 {
-  auto stream = handle.get_stream();
+  auto stream = handle.get_stream().get();
 
   MLCommon::Random::make_arima(
     out, batch_size, n_obs, order, stream, scale, noise_scale, intercept_scale, seed);

cpp/src/datasets/make_blobs.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -31,7 +31,7 @@ void make_blobs(const raft::handle_t& handle,
                            n_rows,
                            n_cols,
                            n_clusters,
-                           handle.get_stream(),
+                           handle.get_stream().get(),
                            row_major,
                            centers,
                            cluster_std,
@@ -62,7 +62,7 @@ void make_blobs(const raft::handle_t& handle,
                            n_rows,
                            n_cols,
                            n_clusters,
-                           handle.get_stream(),
+                           handle.get_stream().get(),
                            row_major,
                            centers,
                            cluster_std,
@@ -93,7 +93,7 @@ void make_blobs(const raft::handle_t& handle,
                            n_rows,
                            n_cols,
                            n_clusters,
-                           handle.get_stream(),
+                           handle.get_stream().get(),
                            row_major,
                            centers,
                            cluster_std,
@@ -124,7 +124,7 @@ void make_blobs(const raft::handle_t& handle,
                            n_rows,
                            n_cols,
                            n_clusters,
-                           handle.get_stream(),
+                           handle.get_stream().get(),
                            row_major,
                            centers,
                            cluster_std,

cpp/src/datasets/make_regression.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -28,7 +28,7 @@ void make_regression_helper(const raft::handle_t& handle,
                             uint64_t seed)
 {
   const auto& handle_impl            = handle;
-  cudaStream_t stream                = handle_impl.get_stream();
+  cudaStream_t stream                = handle_impl.get_stream().get();
   cublasHandle_t cublas_handle       = handle_impl.get_cublas_handle();
   cusolverDnHandle_t cusolver_handle = handle_impl.get_cusolver_dn_handle();