Update APIs - SOCP support

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

@@ -30,6 +30,7 @@ void apply_barrier_rhs(iteration_data_t<int, double>& data, double const* barrie
 namespace cuopt {
 namespace CUOPT_EXPORT mathematical_optimization {
 
+template <typename i_t, typename f_t>
 struct barrier_transform_t;
 
 /**
@@ -63,9 +64,10 @@ class barrier_cache_t {
    */
   barrier::iteration_data_t<int, double>* release_iteration_data();
 
-  void store_transform(std::unique_ptr<barrier_transform_t> transform);
-  [[nodiscard]] barrier_transform_t* transform();
-  [[nodiscard]] barrier_transform_t const* transform() const;
+  template <typename i_t, typename f_t>
+  void store_transform(std::unique_ptr<barrier_transform_t<i_t, f_t>> transform);
+  [[nodiscard]] barrier_transform_t<int, double>* transform();
+  [[nodiscard]] barrier_transform_t<int, double> const* transform() const;
   /** True when an update API has staged new data that the next solve should reuse. */
   [[nodiscard]] bool dirty() const;
   void mark_clean();
@@ -77,13 +79,15 @@ class barrier_cache_t {
    * 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);
+  template <typename i_t, typename f_t>
+  void update_linear_objective(f_t const* c, i_t 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);
+  template <typename i_t, typename f_t>
+  void update_rhs(f_t const* b, i_t m);
 
  private:
   barrier_cache_t(std::unique_ptr<rmm::cuda_stream> stream, std::unique_ptr<raft::handle_t> handle);

cpp/src/barrier/barrier.cu

@@ -932,7 +932,7 @@ class iteration_data_t {
         {
           raft::common::nvtx::range form_scope("Barrier: LP Data: form augmented");
           // Build the sparsity pattern of the augmented system
-          form_augmented(true);
+          form_augmented(augmented_form_t::build);
         }
         if (settings.concurrent_halt != nullptr && *settings.concurrent_halt == 1) { return; }
         symbolic_status = chol->analyze(device_augmented);
@@ -948,9 +948,10 @@ class iteration_data_t {
   }
 
   // Attach this solve's settings and rewind iterate-dependent state so barrier can
-  // 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.
+  // start with the new c / b. A and Q are unchanged. The previous solve left D
+  // and the KKT values at its last iterate. reset_for_reuse (or form_adat(false)) rewrites
+  // values in the existing CSR. The cone block is put back to the initial diagonal a cold
+  // start factorizes; Nesterov-Todd scaling is recomputed from the new point.
   bool reset_iterate_state(const simplex_solver_settings_t<i_t, f_t>& settings)
   {
     if (chol == nullptr || symbolic_status != 0) { return false; }
@@ -987,7 +988,7 @@ class iteration_data_t {
     }
 
     if (use_augmented) {
-      form_augmented(false);
+      form_augmented(augmented_form_t::reset_for_reuse);
     } else {
       form_adat(false);
     }
@@ -1063,7 +1064,12 @@ class iteration_data_t {
     return degree;
   }
 
-  void form_augmented(bool first_call = false)
+  // build: first call, device CSR and metadata.
+  // reset_for_reuse: values only; cone block back to the cold-start matrix.
+  // update: values only; cone block from the current Nesterov-Todd scaling.
+  enum class augmented_form_t { build, reset_for_reuse, update };
+
+  void form_augmented(augmented_form_t mode = augmented_form_t::update)
   {
     i_t n    = A.n;
     i_t m    = A.m;
@@ -1074,7 +1080,7 @@ class iteration_data_t {
     const i_t p            = augmented_expansion_count();
     i_t factorization_size = augmented_system_size(n, m);
 
-    if (first_call) {
+    if (mode == augmented_form_t::build) {
       raft::common::nvtx::range scope("Barrier: augmented: device CSR build");
 
       const size_t n_sparse_cone_entries =
@@ -1165,7 +1171,40 @@ class iteration_data_t {
                          });
       RAFT_CHECK_CUDA(handle_ptr->get_stream().get());
 
-      if (has_soc) {
+      if (has_soc && mode == augmented_form_t::reset_for_reuse) {
+        // Cold initial_point factorizes this diagonal, then the first Newton step
+        // rebuilds the Nesterov-Todd Hessian. Zero w and eta so a dense block
+        // scatter and the matrix-free product both see that same initial matrix.
+        auto stream = handle_ptr->get_stream();
+        thrust::fill(rmm::exec_policy(stream), cones().w.begin(), cones().w.end(), f_t(0));
+        thrust::fill(rmm::exec_policy(stream), cones().eta.begin(), cones().eta.end(), f_t(0));
+        if (cones().has_sparse_cones()) {
+          restore_initial_sparse_cone_block(cones(),
+                                            device_augmented.x,
+                                            cone_kkt_data_.sparse_Hs_diag,
+                                            cone_kkt_data_.sparse_hessian_diag,
+                                            cone_kkt_data_.sparse_hessian_Q,
+                                            cone_kkt_data_.sparse_exp_v_col,
+                                            cone_kkt_data_.sparse_exp_u_col,
+                                            cone_kkt_data_.sparse_exp_v_row,
+                                            cone_kkt_data_.sparse_exp_u_row,
+                                            cone_kkt_data_.sparse_expansion_D,
+                                            stream,
+                                            dual_perturb);
+          RAFT_CHECK_CUDA(stream.get());
+        }
+        if (cones().n_dense_cones() > 0) {
+          scatter_dense_hessian_into_augmented(cones(),
+                                               device_augmented.x,
+                                               cone_kkt_data_.cone_csr_indices,
+                                               cone_kkt_data_.cone_Q_values,
+                                               cone_kkt_data_.dense_block_offsets,
+                                               cone_kkt_data_.dense_cone_ids,
+                                               stream,
+                                               dual_perturb);
+          RAFT_CHECK_CUDA(stream.get());
+        }
+      } else if (has_soc) {
         if (cones().has_sparse_cones()) {
           scatter_sparse_hessian_into_augmented(cones(),
                                                 device_augmented.x,

cpp/src/barrier/barrier_cache.cu

@@ -23,8 +23,9 @@ using barrier_iteration_data_t = barrier::iteration_data_t<int, double>;
 using barrier_iteration_data_ptr =
   std::unique_ptr<barrier_iteration_data_t, void (*)(barrier_iteration_data_t*)>;
 
-static void require_cache(barrier_transform_t const* transform,
-                          barrier_iteration_data_t const* data,
+template <typename i_t, typename f_t>
+static void require_cache(barrier_transform_t<i_t, f_t> const* transform,
+                          barrier::iteration_data_t<i_t, f_t> const* data,
                           char const* api)
 {
   cuopt_expects(transform != nullptr,
@@ -38,7 +39,8 @@ static void require_cache(barrier_transform_t const* transform,
 }
 
 // Re-adds the first solve's barrier-minus-crush shift so the update lands in the presolved model.
-static void add_shift(std::vector<double>& crushed, std::vector<double> const& shift)
+template <typename f_t>
+static void add_shift(std::vector<f_t>& crushed, std::vector<f_t> const& shift)
 {
   cuopt_expects(shift.size() == crushed.size(),
                 error_type_t::ValidationError,
@@ -59,7 +61,7 @@ struct barrier_cache_t::impl {
   std::unique_ptr<rmm::cuda_stream> stream;
   std::unique_ptr<raft::handle_t> handle;
   // Destroy iteration_data before transform: it may const-ref A/Q stored on the transform.
-  std::unique_ptr<barrier_transform_t> transform;
+  std::unique_ptr<barrier_transform_t<int, double>> transform;
   barrier_iteration_data_ptr iteration_data;
   bool linear_objective_dirty{false};
   bool rhs_dirty{false};
@@ -107,14 +109,18 @@ barrier_iteration_data_t* barrier_cache_t::release_iteration_data()
   return impl_->iteration_data.release();
 }
 
-void barrier_cache_t::store_transform(std::unique_ptr<barrier_transform_t> transform)
+template <typename i_t, typename f_t>
+void barrier_cache_t::store_transform(std::unique_ptr<barrier_transform_t<i_t, f_t>> transform)
 {
   impl_->transform = std::move(transform);
 }
 
-barrier_transform_t* barrier_cache_t::transform() { return impl_->transform.get(); }
+barrier_transform_t<int, double>* barrier_cache_t::transform() { return impl_->transform.get(); }
 
-barrier_transform_t const* barrier_cache_t::transform() const { return impl_->transform.get(); }
+barrier_transform_t<int, double> const* barrier_cache_t::transform() const
+{
+  return impl_->transform.get();
+}
 
 bool barrier_cache_t::dirty() const
 {
@@ -131,19 +137,20 @@ void barrier_cache_t::mark_clean()
 
 bool barrier_cache_t::rhs_infeasible() const { return impl_->rhs_infeasible; }
 
-void barrier_cache_t::update_linear_objective(double const* c, int n)
+template <typename i_t, typename f_t>
+void barrier_cache_t::update_linear_objective(f_t const* c, i_t n)
 {
   require_cache(impl_->transform.get(), impl_->iteration_data.get(), "update_linear_objective");
   // Cached Q and c are in minimization space.
-  std::vector<double> user_objective;
+  std::vector<f_t> user_objective;
   if (impl_->transform->maximize && c != nullptr && n > 0) {
     user_objective.assign(c, c + n);
-    for (double& value : user_objective) {
+    for (f_t& value : user_objective) {
       value = -value;
     }
     c = user_objective.data();
   }
-  std::vector<double> crushed;
+  std::vector<f_t> crushed;
   try {
     crushed = crush_user_linear_objective(*impl_->transform, c, n);
   } catch (std::invalid_argument const& e) {
@@ -154,35 +161,36 @@ void barrier_cache_t::update_linear_objective(double const* c, int n)
   auto const& linear_obj_shift = impl_->transform->linear_obj_shift;
   auto const& column_scales    = impl_->transform->column_scales;
   auto const& translated_lower = impl_->transform->presolve_info.removed_lower_bounds;
-  simplex::lp_problem_t<int, double>& barrier_lp = *impl_->transform->barrier_lp;
+  simplex::lp_problem_t<i_t, f_t>& barrier_lp = *impl_->transform->barrier_lp;
   if (!translated_lower.empty() && linear_obj_shift.size() == crushed.size() &&
       column_scales.size() == crushed.size() && barrier_lp.objective.size() == crushed.size()) {
-    double obj_constant_delta = 0.0;
+    f_t obj_constant_delta    = 0.0;
     std::size_t const n_lower = std::min(translated_lower.size(), crushed.size());
     for (std::size_t j = 0; j < n_lower; ++j) {
-      double const crushed_before = barrier_lp.objective[j] - linear_obj_shift[j];
+      f_t const crushed_before = barrier_lp.objective[j] - linear_obj_shift[j];
       obj_constant_delta += (crushed[j] - crushed_before) * column_scales[j] * translated_lower[j];
     }
     barrier_lp.obj_constant += obj_constant_delta;
   }
   add_shift(crushed, linear_obj_shift);
   // The next solve builds its solver from barrier_lp, so keep its objective and the cached
   // iteration workspace on the same c.
-  std::vector<double>& barrier_objective = barrier_lp.objective;
+  std::vector<f_t>& barrier_objective = barrier_lp.objective;
   cuopt_expects(barrier_objective.size() == crushed.size(),
                 error_type_t::ValidationError,
                 "update_linear_objective: crushed objective size does not match the cached "
                 "barrier LP.");
   barrier_objective = crushed;
   barrier::apply_barrier_linear_objective(
-    *impl_->iteration_data, crushed.data(), static_cast<int>(crushed.size()));
+    *impl_->iteration_data, crushed.data(), static_cast<i_t>(crushed.size()));
   impl_->linear_objective_dirty = true;
 }
 
-void barrier_cache_t::update_rhs(double const* b, int m)
+template <typename i_t, typename f_t>
+void barrier_cache_t::update_rhs(f_t const* b, i_t m)
 {
   require_cache(impl_->transform.get(), impl_->iteration_data.get(), "update_rhs");
-  std::vector<double> crushed;
+  std::vector<f_t> crushed;
   std::string error;
   crush_rhs_status_t const status = crush_user_rhs(*impl_->transform, b, m, crushed, error);
   if (status == crush_rhs_status_t::infeasible) {
@@ -198,14 +206,21 @@ void barrier_cache_t::update_rhs(double const* b, int m)
   add_shift(crushed, impl_->transform->rhs_shift);
   // barrier_lp->rhs also seeds the next solve's Mehrotra start, so keep it and the cached
   // workspace on the same b.
-  std::vector<double>& barrier_rhs = impl_->transform->barrier_lp->rhs;
+  std::vector<f_t>& barrier_rhs = impl_->transform->barrier_lp->rhs;
   cuopt_expects(barrier_rhs.size() == crushed.size(),
                 error_type_t::ValidationError,
                 "update_rhs: crushed RHS size does not match the cached barrier LP.");
   barrier_rhs = crushed;
   barrier::apply_barrier_rhs(
-    *impl_->iteration_data, crushed.data(), static_cast<int>(crushed.size()));
+    *impl_->iteration_data, crushed.data(), static_cast<i_t>(crushed.size()));
   impl_->rhs_dirty = true;
 }
 
+#ifdef DUAL_SIMPLEX_INSTANTIATE_DOUBLE
+template void barrier_cache_t::store_transform<int, double>(
+  std::unique_ptr<barrier_transform_t<int, double>>);
+template void barrier_cache_t::update_linear_objective<int, double>(double const*, int);
+template void barrier_cache_t::update_rhs<int, double>(double const*, int);
+#endif
+
 }  // namespace cuopt::mathematical_optimization