Adopt CUDA stream compatibility accessors

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

@@ -19,6 +19,8 @@
 #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 +92,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);
+                            rmm::cuda_stream_view stream = cuda::stream_ref{
+                              cudaStream_t{cudaStreamDefault}});
 
   /**
    * @brief Get the callback for the user solution

cpp/include/cuopt/mathematical_optimization/optimization_problem_solution.hpp

@@ -65,7 +65,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 +77,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 +88,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 +154,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 +167,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 +180,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 +193,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 +205,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 +218,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 +230,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 +245,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 +260,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 +406,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/pdlp/solver_settings.hpp

@@ -8,6 +8,8 @@
 #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 +153,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);
+                                   rmm::cuda_stream_view stream = cuda::stream_ref{
+                                     cudaStream_t{cudaStreamDefault}});
 
   /**
    * @brief Set an initial dual solution.
@@ -165,7 +168,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);
+                                 rmm::cuda_stream_view stream = cuda::stream_ref{
+                                   cudaStream_t{cudaStreamDefault}});
 
   /** TODO batch mode: tmp
    * @brief Set an initial step size.
@@ -200,11 +204,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/solver_settings.hpp

@@ -10,6 +10,8 @@
 #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>
@@ -52,10 +54,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);
+                                        rmm::cuda_stream_view 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);
+                                      rmm::cuda_stream_view 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 +86,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);
+                                rmm::cuda_stream_view stream = cuda::stream_ref{
+                                  cudaStream_t{cudaStreamDefault}});
   void set_mip_callback(internals::base_solution_callback_t* callback = nullptr,
                         void* user_data                               = nullptr);
 

cpp/include/cuopt/mathematical_optimization/utilities/segmented_sum_handler.cuh

@@ -22,7 +22,7 @@ struct segmented_sum_handler_t {
                             i_t problem_size)
   {
     cub::DeviceSegmentedReduce::Sum(
-      nullptr, byte_needed_, input, output, batch_size, problem_size, stream_view_);
+      nullptr, byte_needed_, input, output, batch_size, problem_size, stream_view_.get());
 
     segmented_sum_storage_.resize(byte_needed_, stream_view_);
 
@@ -32,7 +32,7 @@ struct segmented_sum_handler_t {
                                     output,
                                     batch_size,
                                     problem_size,
-                                    stream_view_);
+                                    stream_view_.get());
   }
 
   template <typename InputIteratorT, typename ReductionOpT>
@@ -51,9 +51,9 @@ struct segmented_sum_handler_t {
                                        problem_size,
                                        reduction_op,
                                        initial_value,
-                                       stream_view_.value());
+                                       stream_view_.get());
 
-    segmented_sum_storage_.resize(byte_needed_, stream_view_.value());
+    segmented_sum_storage_.resize(byte_needed_, stream_view_.get());
 
     cub::DeviceSegmentedReduce::Reduce(segmented_sum_storage_.data(),
                                        byte_needed_,
@@ -63,7 +63,7 @@ struct segmented_sum_handler_t {
                                        problem_size,
                                        reduction_op,
                                        initial_value,
-                                       stream_view_.value());
+                                       stream_view_.get());
   }
 
   size_t byte_needed_;