Update APIs - SOCP support

cpp/include/cuopt/mathematical_optimization/utilities/barrier_cache.hpp

@@ -23,6 +23,8 @@ void destroy_iteration_data(iteration_data_t<int, double>* data);
 void apply_barrier_linear_objective(iteration_data_t<int, double>& data,
                                     double const* barrier_c,
                                     int n);
+
+void apply_barrier_rhs(iteration_data_t<int, double>& data, double const* barrier_b, int m);
 }  // namespace cuopt::mathematical_optimization::barrier
 
 namespace cuopt {
@@ -34,8 +36,8 @@ struct barrier_transform_t;
  * @brief GPU solve cache owned by DataModel when CUOPT_SEQUENCE_SOLVE is enabled.
  *
  * After an Optimal full solve, holds iteration_data_t and the user-barrier transform.
- * update_linear_objective crushes the new linear objective and sets c_dirty so the next Solve
- * reuses that workspace (skip convert/presolve/scaling).
+ * The update APIs crush new user data into that workspace and mark the cache dirty so the
+ * next Solve reuses it (skip convert/presolve/scaling).
  */
 class barrier_cache_t {
  public:
@@ -64,15 +66,25 @@ class barrier_cache_t {
   void store_transform(std::unique_ptr<barrier_transform_t> transform);
   [[nodiscard]] barrier_transform_t* transform();
   [[nodiscard]] barrier_transform_t const* transform() const;
-  void set_c_dirty(bool dirty);
-  [[nodiscard]] bool c_dirty() const;
+  /** True when an update API has staged new data that the next Solve should reuse. */
+  [[nodiscard]] bool dirty() const;
+  void mark_clean();
+
+  /** True when the last update_rhs made a row presolve dropped as empty infeasible. */
+  [[nodiscard]] bool rhs_infeasible() const;
 
   /**
-   * Crush the input linear objective into cached iteration_data_t.c / d_c_ and set c_dirty.
+   * Crush the input linear objective into cached iteration_data_t.c / d_c_ and mark dirty.
    * Requires a stored transform and iteration_data from an Optimal solve.
    */
   void update_linear_objective(double const* c, int n);
 
+  /**
+   * Crush the input constraint RHS into cached iteration_data_t.b / d_b_ and mark dirty.
+   * Requires a stored transform and iteration_data from an Optimal solve.
+   */
+  void update_rhs(double const* b, int m);
+
  private:
   barrier_cache_t(std::unique_ptr<rmm::cuda_stream> stream, std::unique_ptr<raft::handle_t> handle);
 

cpp/src/barrier/CMakeLists.txt

@@ -9,6 +9,7 @@ set(BARRIER_SRC_FILES
   ${CMAKE_CURRENT_SOURCE_DIR}/barrier_cache.cu
   ${CMAKE_CURRENT_SOURCE_DIR}/device_sparse_matrix.cu
   ${CMAKE_CURRENT_SOURCE_DIR}/pinned_host_allocator.cu
+  ${CMAKE_CURRENT_SOURCE_DIR}/scaling_gpu.cu
   )
 
 set(BARRIER_SRC_FILES ${BARRIER_SRC_FILES} PARENT_SCOPE)

cpp/src/barrier/barrier.cu

@@ -445,11 +445,44 @@ class barrier_reduce_helper_t {
 template <typename i_t, typename f_t>
 class iteration_data_t {
  public:
+  /** The device Q to factorize: adopted from the scaling when already there, uploaded otherwise. */
+  static device_csc_matrix_t<i_t, f_t> make_device_Q(
+    std::shared_ptr<device_csc_matrix_t<i_t, f_t>> scaled_device_Q,
+    const lp_problem_t<i_t, f_t>& lp,
+    const csc_matrix_t<i_t, f_t>& Qin,
+    cuda::stream_ref stream)
+  {
+    const bool has_entries = Qin.n > 0 && Qin.col_start[Qin.n] > 0;
+    if (scaled_device_Q && has_entries) {
+      device_csc_matrix_t<i_t, f_t> dQ(std::move(*scaled_device_Q));
+      // All that is missing is the slack padding create_Q applies on host, which only extends
+      // col_start with the final nz.
+      const i_t old_n = dQ.n;
+      cuopt_assert(old_n <= Qin.n, "device Q has more columns than the host Q");
+      dQ.m = dQ.n = Qin.n;
+      if (old_n < Qin.n) {
+        dQ.col_start.resize(Qin.n + 1, stream);
+        thrust::fill(rmm::exec_policy(stream),
+                     dQ.col_start.begin() + old_n + 1,
+                     dQ.col_start.end(),
+                     dQ.nz_max);
+      }
+      return dQ;
+    }
+    if (has_entries) { return device_csc_matrix_t<i_t, f_t>(Qin, stream); }
+    // Keep an empty but correctly shaped Q so device views are never zero-sized/uninitialized.
+    device_csc_matrix_t<i_t, f_t> empty(stream);
+    empty.reset_empty(lp.num_cols, lp.num_cols, stream);
+    return empty;
+  }
+
   iteration_data_t(const lp_problem_t<i_t, f_t>& lp,
                    i_t num_upper_bounds,
                    const std::vector<i_t>& direct_free_variables,
                    const csc_matrix_t<i_t, f_t>& Qin,
-                   const simplex_solver_settings_t<i_t, f_t>& settings)
+                   const simplex_solver_settings_t<i_t, f_t>& settings,
+                   std::shared_ptr<device_csc_matrix_t<i_t, f_t>> scaled_device_A,
+                   std::shared_ptr<device_csc_matrix_t<i_t, f_t>> scaled_device_Q)
     : upper_bounds(num_upper_bounds),
       c(lp.objective),
       b(lp.rhs),
@@ -473,7 +506,6 @@ class iteration_data_t {
       inv_diag(lp.num_cols),
       inv_sqrt_diag(lp.num_cols),
       AD(lp.num_cols, lp.num_rows, 0),
-      AT(lp.num_rows, lp.num_cols, 0),
       ADAT(lp.num_rows, lp.num_rows, 0),
       // augmented(lp.num_cols + lp.num_rows, lp.num_cols + lp.num_rows, 0),
       A_dense(lp.num_rows, 0),
@@ -482,17 +514,28 @@ class iteration_data_t {
       Hchol(0, 0),
       A(lp.A),
       Q(Qin),
-      cusparse_Q_view_(lp.handle_ptr, Q),
-      cusparse_view_(lp.handle_ptr, lp.A),
+      // Q is stored fully symmetric, so CSC(Q) is CSR(Q) and both descriptors borrow the one
+      // device_Q_csc_, which is declared earlier and so is already built.
+      cusparse_Q_view_(lp.handle_ptr, device_Q_csc_, device_Q_csc_),
       cusparse_info_(nullptr),
+      // Borrows device_A_csc_ / device_AT_csc_, both of which are declared before it and so are
+      // already built.
+      cusparse_view_(lp.handle_ptr, device_A_csc_, device_AT_csc_),
       device_AD(lp.num_cols, lp.num_rows, 0, lp.handle_ptr->get_stream()),
       device_A(lp.num_cols, lp.num_rows, 0, lp.handle_ptr->get_stream()),
       device_ADAT(lp.num_rows, lp.num_rows, 0, lp.handle_ptr->get_stream()),
       device_augmented(
         lp.num_cols + lp.num_rows, lp.num_cols + lp.num_rows, 0, lp.handle_ptr->get_stream()),
-      device_A_csc_(lp.handle_ptr->get_stream()),
-      device_Q_csc_(lp.handle_ptr->get_stream()),
-      device_AT_csc_(lp.handle_ptr->get_stream()),
+      // Take over the scaled A when the scaling already left it on device, so it is neither
+      // downloaded there nor uploaded again here.
+      device_A_csc_(scaled_device_A
+                      ? std::move(*scaled_device_A)
+                      : device_csc_matrix_t<i_t, f_t>(lp.A, lp.handle_ptr->get_stream())),
+      device_Q_csc_(
+        make_device_Q(std::move(scaled_device_Q), lp, Qin, lp.handle_ptr->get_stream())),
+      device_AT_csc_(typename device_csc_matrix_t<i_t, f_t>::transposed_t{},
+                     device_A_csc_,
+                     lp.handle_ptr->get_stream()),
       d_original_A_values(0, lp.handle_ptr->get_stream()),
       d_inv_diag_prime(0, lp.handle_ptr->get_stream()),
       d_flag_buffer(0, lp.handle_ptr->get_stream()),
@@ -800,11 +843,17 @@ class iteration_data_t {
 
     if (settings.concurrent_halt != nullptr && *settings.concurrent_halt == 1) { return; }
 
-    {
+    // AD only feeds ADAT; the augmented path derives A^T on device instead.
+    if (!use_augmented) {
       raft::common::nvtx::range scope("Barrier: LP Data: AD matrix setup");
-      // Copy A into AD
-      AD = lp.A;
-      if (!use_augmented && n_dense_columns > 0) {
+      if (n_dense_columns == 0) {
+        // AD is A, which is already on device; only AD's dimensions are read on the host, so its
+        // entries are never materialised here.
+        AD.m = lp.A.m;
+        AD.n = lp.A.n;
+      } else {
+        // Copy A into AD
+        AD = lp.A;
         cols_to_remove.resize(lp.num_cols, 0);
         for (i_t k : dense_columns_unordered) {
           cols_to_remove[k] = 1;
@@ -830,34 +879,29 @@ class iteration_data_t {
           A_dense.from_sparse(lp.A, j, k++);
         }
       }
-
-      AD.transpose(AT);
-    }
-
-    if (use_augmented) {
-      raft::common::nvtx::range scope("Barrier: augmented: device CSC upload");
-      device_A_csc_.copy(A, handle_ptr->get_stream());
-      device_AT_csc_.copy(AT, handle_ptr->get_stream());
-      if (Q.n > 0 && Q.col_start[Q.n] > 0) {
-        device_Q_csc_.copy(Q, handle_ptr->get_stream());
-      } else {
-        // Keep an empty but correctly shaped Q so device views are never zero-sized/uninitialized.
-        device_Q_csc_.reset_empty(A.n, A.n, handle_ptr->get_stream());
-      }
     }
 
     // device_AD / device_A / ADAT path is only used when forming ADAT (!use_augmented).
     if (!use_augmented) {
       raft::common::nvtx::range scope("Barrier: LP Data: device AD path");
-      device_AD.copy(AD, handle_ptr->get_stream());
-      d_original_A_values.resize(device_AD.x.size(), handle_ptr->get_stream());
-      raft::copy(d_original_A_values.data(),
-                 device_AD.x.data(),
-                 device_AD.x.size(),
-                 handle_ptr->get_stream());
+      if (n_dense_columns > 0) {
+        device_AD.copy(AD, handle_ptr->get_stream());
+        // AD differs from A once dense columns are dropped, so form_adat needs its own snapshot
+        // of the unscaled values to restore from.
+        d_original_A_values.resize(device_AD.x.size(), handle_ptr->get_stream());
+        raft::copy(d_original_A_values.data(),
+                   device_AD.x.data(),
+                   device_AD.x.size(),
+                   handle_ptr->get_stream());
+        device_AD.to_compressed_row(device_A, handle_ptr->get_stream());
+      } else {
+        // AD == A, so device_AD is seeded straight from device_A_csc_, which also doubles as
+        // form_adat's restore source, and device_AT_csc_ (already CSR(A)) serves as the SpGEMM's
+        // left operand -- neither needs a second copy. Both stay read-only for the whole solve.
+        device_AD.copy(device_A_csc_, handle_ptr->get_stream());
+      }
       // For efficient scaling of AD col we form the col index array
       device_AD.form_col_index(handle_ptr->get_stream());
-      device_AD.to_compressed_row(device_A, handle_ptr->get_stream());
       RAFT_CHECK_CUDA(handle_ptr->get_stream().get());
     }
 
@@ -895,7 +939,7 @@ class iteration_data_t {
   }
 
   // Attach this solve's settings and rewind iterate-dependent state so barrier can
-  // start with the new c. A and Q are unchanged; the previous solve
+  // start with the new c / b. A and Q are unchanged; the previous solve
   // left D and the KKT values at its last iterate. Reuse is QP-only (no cones),
   // so form_*(false) updates values in the existing CSR; no symbolic rebuild.
   bool reset_iterate_state(const simplex_solver_settings_t<i_t, f_t>& settings)
@@ -1066,7 +1110,6 @@ class iteration_data_t {
                                       stream_view_);
 
       settings_.log.debug("augmented nz %d (gpu build)\n", total_nnz);
-      cuopt_assert(A.col_start[n] == AT.col_start[m], "A nz != AT nz");
       handle_ptr->sync_stream();
 
 #ifdef CHECK_SYMMETRY
@@ -1154,10 +1197,11 @@ class iteration_data_t {
 
     {
       raft::common::nvtx::range scope("Barrier: Form ADAT: restore A");
-      raft::copy(device_AD.x.data(),
-                 d_original_A_values.data(),
-                 d_original_A_values.size(),
-                 handle_ptr->get_stream());
+      // device_A_csc_ holds A's unscaled values and is never written, so when AD == A it is the
+      // snapshot; with dense columns removed AD differs and carries its own.
+      const f_t* original_values =
+        n_dense_columns > 0 ? d_original_A_values.data() : device_A_csc_.x.data();
+      raft::copy(device_AD.x.data(), original_values, device_AD.x.size(), handle_ptr->get_stream());
     }
     {
       raft::common::nvtx::range scope("Barrier: Form ADAT: inv_diag prime");
@@ -1199,12 +1243,29 @@ class iteration_data_t {
     if (settings_.concurrent_halt != nullptr && *settings_.concurrent_halt == 1) { return; }
     if (first_call) {
       raft::common::nvtx::range scope("Barrier: Form ADAT: cusparse init");
+      // With no dense columns AD == A, so CSC(A^T) is the CSR(A) the SpGEMM needs and
+      // device_AT_csc_ is used directly instead of keeping a second copy in device_A.
+      const bool own_csr = n_dense_columns > 0;
+      const i_t A_rows   = own_csr ? device_A.m : device_AT_csc_.n;
+      const i_t A_cols   = own_csr ? device_A.n : device_AT_csc_.m;
+      const i_t A_nnz    = own_csr ? device_A.nz_max : device_AT_csc_.nz_max;
+      i_t* A_offsets     = own_csr ? device_A.row_start.data() : device_AT_csc_.col_start.data();
+      i_t* A_indices     = own_csr ? device_A.j.data() : device_AT_csc_.i.data();
+      f_t* A_values      = own_csr ? device_A.x.data() : device_AT_csc_.x.data();
       try {
         if (!cusparse_info_) {
           cusparse_info_ = std::make_unique<cusparse_info_t<i_t, f_t>>(handle_ptr);
         }
-        initialize_cusparse_data<i_t, f_t>(
-          handle_ptr, device_A, device_AD, device_ADAT, spgemm_info());
+        initialize_cusparse_data<i_t, f_t>(handle_ptr,
+                                           A_rows,
+                                           A_cols,
+                                           A_nnz,
+                                           A_offsets,
+                                           A_indices,
+                                           A_values,
+                                           device_AD,
+                                           device_ADAT,
+                                           spgemm_info());
       } catch (const raft::cuda_error& e) {
         settings_.log.printf("Error in initialize_cusparse_data: %s\n", e.what());
         return;
@@ -1214,7 +1275,7 @@ class iteration_data_t {
 
     {
       raft::common::nvtx::range scope("Barrier: Form ADAT: ADAT multiply");
-      multiply_kernels<i_t, f_t>(handle_ptr, device_A, device_AD, device_ADAT, spgemm_info());
+      multiply_kernels<i_t, f_t>(handle_ptr, device_ADAT, spgemm_info());
       handle_ptr->sync_stream();
     }
 
@@ -2210,7 +2271,6 @@ class iteration_data_t {
   rmm::device_uvector<f_t> d_original_A_values;
 
   csc_matrix_t<i_t, f_t> AD;
-  csc_matrix_t<i_t, f_t> AT;
   csc_matrix_t<i_t, f_t> ADAT;
   // csc_matrix_t<i_t, f_t> augmented;
   device_csr_matrix_t<i_t, f_t> device_augmented;
@@ -2431,10 +2491,18 @@ void cholesky_debug_check(const iteration_data_t<i_t, f_t>& data,
 }
 
 template <typename i_t, typename f_t>
-barrier_solver_t<i_t, f_t>::barrier_solver_t(const lp_problem_t<i_t, f_t>& lp,
-                                             const simplex::presolve_info_t<i_t, f_t>& presolve,
-                                             const simplex_solver_settings_t<i_t, f_t>& settings)
-  : lp(lp), settings(settings), presolve_info(presolve), stream_view_(lp.handle_ptr->get_stream())
+barrier_solver_t<i_t, f_t>::barrier_solver_t(
+  const lp_problem_t<i_t, f_t>& lp,
+  const simplex::presolve_info_t<i_t, f_t>& presolve,
+  const simplex_solver_settings_t<i_t, f_t>& settings,
+  std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_A,
+  std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_Q)
+  : lp(lp),
+    settings(settings),
+    presolve_info(presolve),
+    stream_view_(lp.handle_ptr->get_stream()),
+    device_A_(std::move(device_A)),
+    device_Q_(std::move(device_Q))
 {
 }
 
@@ -4863,8 +4931,13 @@ lp_status_t barrier_solver_t<i_t, f_t>::solve(
       Qin           = xf->barrier_Q.get();
     }
     if (lp.Q.n > 0) { create_Q(lp, *Qin); }
-    owned_data = std::make_unique<iteration_data_t<i_t, f_t>>(
-      lp, num_upper_bounds, presolve_info.direct_free_variables, *Qin, settings);
+    owned_data         = std::make_unique<iteration_data_t<i_t, f_t>>(lp,
+                                                              num_upper_bounds,
+                                                              presolve_info.direct_free_variables,
+                                                              *Qin,
+                                                              settings,
+                                                              std::move(device_A_),
+                                                              std::move(device_Q_));
     lp_status_t status = barrier_advanced_solve(start_time, solution, *owned_data);
     return store_or_clear_cache(cache, owned_data, status);
   } catch (const raft::cuda_error& e) {
@@ -4911,6 +4984,18 @@ void apply_barrier_linear_objective(iteration_data_t<int, double>& data,
     data.d_c_.data(), data.c.data(), static_cast<std::size_t>(n), data.handle_ptr->get_stream());
 }
 
+void apply_barrier_rhs(iteration_data_t<int, double>& data, double const* barrier_b, int m)
+{
+  cuopt_expects(
+    barrier_b != nullptr && static_cast<int>(data.b.size()) == m &&
+      static_cast<int>(data.d_b_.size()) == m,
+    error_type_t::ValidationError,
+    "update_rhs: barrier RHS size does not match cached iteration_data_t.");
+  std::copy(barrier_b, barrier_b + m, data.b.data());
+  raft::copy(
+    data.d_b_.data(), data.b.data(), static_cast<std::size_t>(m), data.handle_ptr->get_stream());
+}
+
 #ifdef DUAL_SIMPLEX_INSTANTIATE_DOUBLE
 template bool validate_barrier_cone_layout<int, double>(
   const lp_problem_t<int, double>& problem, const simplex_solver_settings_t<int, double>& settings);

cpp/src/barrier/barrier.hpp

@@ -20,6 +20,7 @@
 #include <cuda/stream>
 #include <rmm/device_uvector.hpp>
 
+#include <memory>
 #include <utility>
 
 namespace cuopt::mathematical_optimization {
@@ -36,12 +37,21 @@ bool validate_barrier_cone_layout(const simplex::lp_problem_t<i_t, f_t>& problem
 template <typename i_t, typename f_t>
 class iteration_data_t;  // Forward declare
 
+template <typename i_t, typename f_t>
+class device_csc_matrix_t;  // Forward declare
+
 template <typename i_t, typename f_t>
 class barrier_solver_t {
  public:
+  // `device_A` / `device_Q` are the scaled matrices the GPU scaling already left on device, taken
+  // over here so they are neither downloaded there nor uploaded again. Null means the solver
+  // uploads them from `lp`. Only solve() consumes them; solve_with_cache() reuses the cached
+  // iteration_data_t and never looks at them.
   barrier_solver_t(const simplex::lp_problem_t<i_t, f_t>& lp,
                    const simplex::presolve_info_t<i_t, f_t>& presolve,
-                   const simplex::simplex_solver_settings_t<i_t, f_t>& settings);
+                   const simplex::simplex_solver_settings_t<i_t, f_t>& settings,
+                   std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_A = nullptr,
+                   std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_Q = nullptr);
   simplex::lp_status_t solve(f_t start_time,
                              simplex::lp_solution_t<i_t, f_t>& solution,
                              cuopt::mathematical_optimization::barrier_cache_t* cache = nullptr);
@@ -117,6 +127,9 @@ class barrier_solver_t {
   const simplex::simplex_solver_settings_t<i_t, f_t>& settings;
   const simplex::presolve_info_t<i_t, f_t>& presolve_info;
   cuda::stream_ref stream_view_;
+  // Handed over to iteration_data_t by solve(), which empties them.
+  std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_A_;
+  std::shared_ptr<device_csc_matrix_t<i_t, f_t>> device_Q_;
 };
 
 }  // namespace cuopt::mathematical_optimization::barrier