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

cpp/docs/DEVELOPER_GUIDE.md

@@ -134,7 +134,7 @@ Similar to a `rmm::device_vector`, allocates a contiguous set of elements in dev
 key differences:
 - As an optimization, elements are uninitialized and no synchronization occurs at construction.
 This limits the types `T` to trivially copyable types.
-- All operations are stream ordered (i.e., they accept a `cuda_stream_view` specifying the stream
+- All operations are stream ordered (i.e., they accept a `cuda::stream_ref` specifying the stream
 on which the operation is performed).
 
 ## Namespaces

cpp/include/cuopt/mathematical_optimization/cpu_pdlp_warm_start_data.hpp

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include <cuda/stream>
 #include <cuopt/export.hpp>
 #include <cuopt/mathematical_optimization/pdlp/pdlp_warm_start_data.hpp>
 
@@ -114,12 +115,12 @@ struct pdlp_warm_start_data_t;
 // Convert GPU → CPU warmstart (D2H copy)
 template <typename i_t, typename f_t>
 cpu_pdlp_warm_start_data_t<i_t, f_t> convert_to_cpu_warmstart(
-  const pdlp_warm_start_data_t<i_t, f_t>& gpu_data, rmm::cuda_stream_view stream);
+  const pdlp_warm_start_data_t<i_t, f_t>& gpu_data, cuda::stream_ref stream);
 
 // Convert CPU → GPU warmstart (H2D copy)
 template <typename i_t, typename f_t>
 pdlp_warm_start_data_t<i_t, f_t> convert_to_gpu_warmstart(
-  const cpu_pdlp_warm_start_data_t<i_t, f_t>& cpu_data, rmm::cuda_stream_view stream);
+  const cpu_pdlp_warm_start_data_t<i_t, f_t>& cpu_data, cuda::stream_ref stream);
 
 }  // namespace CUOPT_EXPORT mathematical_optimization
 }  // namespace cuopt

cpp/include/cuopt/mathematical_optimization/mip/solver_settings.hpp

@@ -19,6 +19,7 @@
 #include <cuopt/mathematical_optimization/pdlp/solver_settings.hpp>
 #include <cuopt/mathematical_optimization/utilities/internals.hpp>
 
+#include <cuda/stream>
 #include <raft/core/device_span.hpp>
 #include <rmm/device_uvector.hpp>
 
@@ -90,7 +91,8 @@ class mip_solver_settings_t {
    */
   void add_initial_solution(const f_t* initial_solution,
                             i_t size,
-                            rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                            cuda::stream_ref stream = cuda::stream_ref{
+                              cudaStream_t{cudaStreamDefault}});
 
   /**
    * @brief Get the callback for the user solution

cpp/include/cuopt/mathematical_optimization/mip/solver_solution.hpp

@@ -13,7 +13,7 @@
 #include <cuopt/mathematical_optimization/mip/solver_stats.hpp>
 #include <cuopt/mathematical_optimization/utilities/internals.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 
 #include <raft/core/handle.hpp>
@@ -52,8 +52,8 @@ class mip_solution_t : public base_solution_t {
 
   mip_solution_t(mip_termination_status_t termination_status,
                  solver_stats_t<i_t, f_t> stats,
-                 rmm::cuda_stream_view stream_view);
-  mip_solution_t(const cuopt::logic_error& error_status, rmm::cuda_stream_view stream_view);
+                 cuda::stream_ref stream_view);
+  mip_solution_t(const cuopt::logic_error& error_status, cuda::stream_ref stream_view);
 
   bool is_mip() const override { return true; }
   const rmm::device_uvector<f_t>& get_solution() const;
@@ -76,7 +76,7 @@ class mip_solution_t : public base_solution_t {
   i_t get_num_simplex_iterations() const;
   const std::vector<std::string>& get_variable_names() const;
   const std::vector<rmm::device_uvector<f_t>>& get_solution_pool() const;
-  void write_to_sol_file(std::string_view filename, rmm::cuda_stream_view stream_view) const;
+  void write_to_sol_file(std::string_view filename, cuda::stream_ref stream_view) const;
   void log_detailed_summary() const;
   void log_summary() const;
 

cpp/include/cuopt/mathematical_optimization/optimization_problem.hpp

@@ -11,6 +11,7 @@
 #include <cuopt/mathematical_optimization/optimization_problem_interface.hpp>
 #include <cuopt/mathematical_optimization/utilities/internals.hpp>
 
+#include <cuda/stream>
 #include <raft/core/device_span.hpp>
 #include <raft/core/handle.hpp>
 #include <rmm/device_uvector.hpp>
@@ -350,7 +351,7 @@ class optimization_problem_t : public optimization_problem_interface_t<i_t, f_t>
    * @tparam other_f_t  Target floating-point type (e.g. float when this is double)
    */
   template <typename other_f_t>
-  optimization_problem_t<i_t, other_f_t> convert_to_other_prec(rmm::cuda_stream_view stream) const;
+  optimization_problem_t<i_t, other_f_t> convert_to_other_prec(cuda::stream_ref stream) const;
 
   /**
    * @brief Returns nullptr since this is already a GPU problem.
@@ -388,7 +389,7 @@ class optimization_problem_t : public optimization_problem_interface_t<i_t, f_t>
 
  private:
   raft::handle_t const* handle_ptr_{nullptr};
-  rmm::cuda_stream_view stream_view_;
+  cuda::stream_ref stream_view_;
 
   problem_category_t problem_category_ = problem_category_t::LP;
   bool maximize_{false};

cpp/include/cuopt/mathematical_optimization/optimization_problem_solution.hpp

@@ -15,7 +15,6 @@
 #include <cuopt/mathematical_optimization/utilities/cython_types.hpp>
 
 #include <raft/core/copy.hpp>
-#include <rmm/cuda_stream_view.hpp>
 
 namespace cuopt {
 namespace CUOPT_EXPORT mathematical_optimization {
@@ -65,7 +64,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
                solution_.get_primal_solution().data(),
                solution_.get_primal_solution().size(),
                stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -77,7 +76,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
                solution_.get_dual_solution().data(),
                solution_.get_dual_solution().size(),
                stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -88,7 +87,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     auto stream = reduced_cost.stream();
     std::vector<f_t> result(reduced_cost.size());
     raft::copy(result.data(), reduced_cost.data(), reduced_cost.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -154,7 +153,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
                ws.current_primal_solution_.data(),
                ws.current_primal_solution_.size(),
                stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -167,7 +166,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     std::vector<f_t> result(ws.current_dual_solution_.size());
     raft::copy(
       result.data(), ws.current_dual_solution_.data(), ws.current_dual_solution_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -180,7 +179,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     std::vector<f_t> result(ws.initial_primal_average_.size());
     raft::copy(
       result.data(), ws.initial_primal_average_.data(), ws.initial_primal_average_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -193,7 +192,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     std::vector<f_t> result(ws.initial_dual_average_.size());
     raft::copy(
       result.data(), ws.initial_dual_average_.data(), ws.initial_dual_average_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -205,7 +204,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     auto stream = ws.current_ATY_.stream();
     std::vector<f_t> result(ws.current_ATY_.size());
     raft::copy(result.data(), ws.current_ATY_.data(), ws.current_ATY_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -218,7 +217,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     std::vector<f_t> result(ws.sum_primal_solutions_.size());
     raft::copy(
       result.data(), ws.sum_primal_solutions_.data(), ws.sum_primal_solutions_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -230,7 +229,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
     auto stream = ws.sum_dual_solutions_.stream();
     std::vector<f_t> result(ws.sum_dual_solutions_.size());
     raft::copy(result.data(), ws.sum_dual_solutions_.data(), ws.sum_dual_solutions_.size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -245,7 +244,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
                ws.last_restart_duality_gap_primal_solution_.data(),
                ws.last_restart_duality_gap_primal_solution_.size(),
                stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -260,7 +259,7 @@ class gpu_lp_solution_t : public lp_solution_interface_t<i_t, f_t> {
                ws.last_restart_duality_gap_dual_solution_.data(),
                ws.last_restart_duality_gap_dual_solution_.size(),
                stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 
@@ -406,7 +405,7 @@ class gpu_mip_solution_t : public mip_solution_interface_t<i_t, f_t> {
     std::vector<f_t> result(solution_.get_solution().size());
     raft::copy(
       result.data(), solution_.get_solution().data(), solution_.get_solution().size(), stream);
-    stream.synchronize();
+    stream.sync();
     return result;
   }
 

cpp/include/cuopt/mathematical_optimization/optimization_problem_solution_interface.hpp

@@ -13,7 +13,6 @@
 #include <cuopt/mathematical_optimization/pdlp/pdlp_warm_start_data.hpp>
 #include <cuopt/mathematical_optimization/pdlp/solver_solution.hpp>  // For pdlp_termination_status_t
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuopt/mathematical_optimization/utilities/cython_types.hpp>

cpp/include/cuopt/mathematical_optimization/pdlp/pdlp_warm_start_data.hpp

@@ -9,6 +9,7 @@
 
 #include <cuopt/export.hpp>
 
+#include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 
 #include <span>
@@ -67,7 +68,7 @@ struct pdlp_warm_start_data_t {
 
   // Copy constructor using the view version for the cython_solver
   pdlp_warm_start_data_t(const pdlp_warm_start_data_view_t<i_t, f_t>& other,
-                         rmm::cuda_stream_view stream_view);
+                         cuda::stream_ref stream_view);
 
   // Copy constructor for when copying the solver_settings object in the PDLP object
   pdlp_warm_start_data_t(const pdlp_warm_start_data_t<i_t, f_t>& other);

cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <cuopt/mathematical_optimization/constants.h>
+#include <cuda/stream>
 #include <cuopt/export.hpp>
 #include <cuopt/mathematical_optimization/cpu_pdlp_warm_start_data.hpp>
 #include <cuopt/mathematical_optimization/pdlp/pdlp_hyper_params.cuh>
@@ -151,7 +152,8 @@ class pdlp_solver_settings_t {
    */
   void set_initial_primal_solution(const f_t* initial_primal_solution,
                                    i_t size,
-                                   rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                                   cuda::stream_ref stream = cuda::stream_ref{
+                                     cudaStream_t{cudaStreamDefault}});
 
   /**
    * @brief Set an initial dual solution.
@@ -165,7 +167,8 @@ class pdlp_solver_settings_t {
    */
   void set_initial_dual_solution(const f_t* initial_dual_solution,
                                  i_t size,
-                                 rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                                 cuda::stream_ref stream = cuda::stream_ref{
+                                   cudaStream_t{cudaStreamDefault}});
 
   /** TODO batch mode: tmp
    * @brief Set an initial step size.
@@ -200,11 +203,12 @@ class pdlp_solver_settings_t {
    * @param constraint_mapping Constraints indices to scatter to in case the new
    * problem has less constraints
    */
-  void set_pdlp_warm_start_data(pdlp_warm_start_data_t<i_t, f_t>& pdlp_warm_start_data_view,
-                                const rmm::device_uvector<i_t>& var_mapping =
-                                  rmm::device_uvector<i_t>{0, rmm::cuda_stream_default},
-                                const rmm::device_uvector<i_t>& constraint_mapping =
-                                  rmm::device_uvector<i_t>{0, rmm::cuda_stream_default});
+  void set_pdlp_warm_start_data(
+    pdlp_warm_start_data_t<i_t, f_t>& pdlp_warm_start_data_view,
+    const rmm::device_uvector<i_t>& var_mapping =
+      rmm::device_uvector<i_t>{0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}},
+    const rmm::device_uvector<i_t>& constraint_mapping = rmm::device_uvector<i_t>{
+      0, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}});
 
   // Same but for the Cython interface
   void set_pdlp_warm_start_data(const f_t* current_primal_solution,

cpp/include/cuopt/mathematical_optimization/pdlp/solver_solution.hpp

@@ -14,7 +14,7 @@
 #include <cuopt/mathematical_optimization/pdlp/solver_settings.hpp>
 #include <cuopt/mathematical_optimization/utilities/internals.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
+#include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 
 #include <raft/core/handle.hpp>
@@ -105,7 +105,7 @@ class optimization_problem_solution_t : public base_solution_t {
    * @param[in] stream_view An rmm view to a stream. All computations will go through this stream
    */
   optimization_problem_solution_t(pdlp_termination_status_t termination_status_,
-                                  rmm::cuda_stream_view stream_view);
+                                  cuda::stream_ref stream_view);
 
   /**
    * @brief Construct an optimization problem solution that serves as PDLP solver output
@@ -115,8 +115,7 @@ class optimization_problem_solution_t : public base_solution_t {
    * 'Optimal', 'PrimalInfeasible', 'DualInfeasible', 'TimeLimit'
    * @param[in] stream_view An rmm view to a stream. All computations will go through this stream
    */
-  optimization_problem_solution_t(cuopt::logic_error error_status_,
-                                  rmm::cuda_stream_view stream_view);
+  optimization_problem_solution_t(cuopt::logic_error error_status_, cuda::stream_ref stream_view);
   /**
    * @brief Construct an optimization problem solution that serves as PDLP solver output
    *
@@ -271,7 +270,7 @@ class optimization_problem_solution_t : public base_solution_t {
    * @param stream_view Non-owning stream view object
    */
   void write_to_file(std::string_view filename,
-                     rmm::cuda_stream_view stream_view,
+                     cuda::stream_ref stream_view,
                      bool generate_variable_values = true);
 
   /**
@@ -281,7 +280,7 @@ class optimization_problem_solution_t : public base_solution_t {
    * @param filename Name of the output file
    * @param stream_view Non-owning stream view object
    */
-  void write_to_sol_file(std::string_view filename, rmm::cuda_stream_view stream_view) const;
+  void write_to_sol_file(std::string_view filename, cuda::stream_ref stream_view) const;
 
   /**
    * @brief Copy solution from another solution object

cpp/include/cuopt/mathematical_optimization/solver_settings.hpp

@@ -10,9 +10,9 @@
 #include <cuopt/export.hpp>
 #include <cuopt/mathematical_optimization/pdlp/pdlp_warm_start_data.hpp>
 
+#include <cuda/stream>
 #include <raft/core/device_span.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuopt/mathematical_optimization/constants.h>
@@ -52,10 +52,12 @@ class solver_settings_t {
 
   void set_initial_pdlp_primal_solution(const f_t* initial_primal_solution,
                                         i_t size,
-                                        rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                                        cuda::stream_ref stream = cuda::stream_ref{
+                                          cudaStream_t{cudaStreamDefault}});
   void set_initial_pdlp_dual_solution(const f_t* initial_dual_solution,
                                       i_t size,
-                                      rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                                      cuda::stream_ref stream = cuda::stream_ref{
+                                        cudaStream_t{cudaStreamDefault}});
   void set_pdlp_warm_start_data(const f_t* current_primal_solution,
                                 const f_t* current_dual_solution,
                                 const f_t* initial_primal_average,
@@ -82,7 +84,8 @@ class solver_settings_t {
   // MIP Settings
   void add_initial_mip_solution(const f_t* initial_solution,
                                 i_t size,
-                                rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+                                cuda::stream_ref stream = cuda::stream_ref{
+                                  cudaStream_t{cudaStreamDefault}});
   void set_mip_callback(internals::base_solution_callback_t* callback = nullptr,
                         void* user_data                               = nullptr);