Make iterative refinement methods flexible

Signed-off-by: yuwenchen95 <yuwchen@nvidia.com>

Author: yuwenchen95

cpp/include/cuopt/mathematical_optimization/constants.h

@@ -49,6 +49,7 @@
 #define CUOPT_BARRIER_DUAL_INITIAL_POINT           "barrier_dual_initial_point"
 #define CUOPT_POSTSOLVE_INFO                       "postsolve_info"
 #define CUOPT_BARRIER_ITERATIVE_REFINEMENT         "barrier_iterative_refinement"
+#define CUOPT_BARRIER_ITERATIVE_REFINEMENT_METHOD  "barrier_iterative_refinement_method"
 #define CUOPT_BARRIER_STEP_SCALE                   "barrier_step_scale"
 #define CUOPT_ELIMINATE_DENSE_COLUMNS              "eliminate_dense_columns"
 #define CUOPT_CUDSS_DETERMINISTIC                  "cudss_deterministic"

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

@@ -307,6 +307,7 @@ class pdlp_solver_settings_t {
   bool eliminate_dense_columns{true};
   pdlp_precision_t pdlp_precision{pdlp_precision_t::DefaultPrecision};
   bool barrier_iterative_refinement{true};
+  i_t barrier_iterative_refinement_method{1};  // 0: fixed_point, 1: gmres (default)
   i_t barrier_soc_threshold{100};
   f_t barrier_step_scale{0.9};
   bool save_best_primal_so_far{false};

cpp/src/barrier/barrier.cu

@@ -2279,7 +2279,8 @@ int barrier_solver_t<i_t, f_t>::initial_point(iteration_data_t<i_t, f_t>& data)
 
     if (settings.barrier_iterative_refinement) {
       const f_t ir_tol = data.has_sparse_cones() ? f_t(1e-12) : f_t(1e-8);
-      iterative_refinement<i_t, f_t, op_t>(op, rhs, soln, ir_tol);
+      iterative_refinement<i_t, f_t, op_t>(
+        op, rhs, soln, ir_tol, settings.barrier_iterative_refinement_method);
     }
 
     for (i_t k = 0; k < lp.num_cols; k++) {
@@ -2904,9 +2905,13 @@ i_t barrier_solver_t<i_t, f_t>::gpu_compute_search_direction(iteration_data_t<i_
     } op(data);
     if (settings.barrier_iterative_refinement) {
       raft::common::nvtx::range fun_scope("Barrier: iterative_refinement");
-      const f_t ir_tol    = data.has_sparse_cones() ? f_t(1e-12) : f_t(1e-8);
-      const f_t solve_err = iterative_refinement<i_t, f_t, op_t>(
-        op, data.d_augmented_rhs_, data.d_augmented_soln_, ir_tol);
+      const f_t ir_tol = data.has_sparse_cones() ? f_t(1e-12) : f_t(1e-8);
+      const f_t solve_err =
+        iterative_refinement<i_t, f_t, op_t>(op,
+                                             data.d_augmented_rhs_,
+                                             data.d_augmented_soln_,
+                                             ir_tol,
+                                             settings.barrier_iterative_refinement_method);
       if (solve_err > 1e-1) {
         settings.log.printf("|| Aug (dx, dy) - aug_rhs || %e after IR\n", solve_err);
       }
@@ -2992,8 +2997,8 @@ i_t barrier_solver_t<i_t, f_t>::gpu_compute_search_direction(iteration_data_t<i_
             data_.gpu_solve_adat(b, x);
           }
         } adat_op(data);
-        const f_t adat_solve_err =
-          iterative_refinement<i_t, f_t, adat_op_t>(adat_op, data.d_h_, data.d_dy_);
+        const f_t adat_solve_err = iterative_refinement<i_t, f_t, adat_op_t>(
+          adat_op, data.d_h_, data.d_dy_, f_t(1e-8), settings.barrier_iterative_refinement_method);
         if (adat_solve_err > 1e-1) {
           settings.log.printf("||ADAT*dy - h|| %e after IR\n", adat_solve_err);
         }

cpp/src/barrier/iterative_refinement.hpp

@@ -55,10 +55,10 @@ struct subtract_scaled_op {
 };
 
 template <typename i_t, typename f_t, typename T>
-f_t iterative_refinement_simple(T& op,
-                                const rmm::device_uvector<f_t>& b,
-                                rmm::device_uvector<f_t>& x,
-                                f_t tol = 1e-8)
+f_t iterative_refinement_fixed_point(T& op,
+                                     const rmm::device_uvector<f_t>& b,
+                                     rmm::device_uvector<f_t>& x,
+                                     f_t tol = 1e-8)
 {
   rmm::device_uvector<f_t> x_sav(x, x.stream());
 
@@ -365,13 +365,15 @@ template <typename i_t, typename f_t, typename T>
 f_t iterative_refinement(T& op,
                          const dense_vector_t<i_t, f_t>& b,
                          dense_vector_t<i_t, f_t>& x,
-                         f_t tol = 1e-8)
+                         f_t tol    = 1e-8,
+                         i_t method = 1)
 {
   rmm::device_uvector<f_t> d_b(b.size(), op.data_.handle_ptr->get_stream());
   raft::copy(d_b.data(), b.data(), b.size(), op.data_.handle_ptr->get_stream());
   rmm::device_uvector<f_t> d_x(x.size(), op.data_.handle_ptr->get_stream());
   raft::copy(d_x.data(), x.data(), x.size(), op.data_.handle_ptr->get_stream());
-  auto err = iterative_refinement_gmres<i_t, f_t, T>(op, d_b, d_x, tol);
+  auto err = (method == 0) ? iterative_refinement_fixed_point<i_t, f_t, T>(op, d_b, d_x, tol)
+                           : iterative_refinement_gmres<i_t, f_t, T>(op, d_b, d_x, tol);
 
   raft::copy(x.data(), d_x.data(), x.size(), op.data_.handle_ptr->get_stream());
 
@@ -383,9 +385,11 @@ template <typename i_t, typename f_t, typename T>
 f_t iterative_refinement(T& op,
                          const rmm::device_uvector<f_t>& b,
                          rmm::device_uvector<f_t>& x,
-                         f_t tol = 1e-8)
+                         f_t tol    = 1e-8,
+                         i_t method = 1)
 {
-  return iterative_refinement_gmres<i_t, f_t, T>(op, b, x, tol);
+  return (method == 0) ? iterative_refinement_fixed_point<i_t, f_t, T>(op, b, x, tol)
+                       : iterative_refinement_gmres<i_t, f_t, T>(op, b, x, tol);
 }
 
 }  // namespace cuopt::mathematical_optimization::barrier

cpp/src/dual_simplex/simplex_solver_settings.hpp

@@ -74,6 +74,7 @@ struct simplex_solver_settings_t {
       barrier(false),
       eliminate_dense_columns(true),
       barrier_iterative_refinement(true),
+      barrier_iterative_refinement_method(1),
       barrier_step_scale(0.9),
       barrier_soc_threshold(100),
       num_gpus(1),
@@ -172,10 +173,11 @@ struct simplex_solver_settings_t {
                                    // (only takes effect when a cuDSS threading layer is loaded)
   bool barrier;                    // true to use barrier method, false to use dual simplex method
   bool deterministic;  // true to use B&B deterministic mode, false to use non-deterministic mode
-  bool eliminate_dense_columns;       // true to eliminate dense columns from A*D*A^T
-  bool barrier_iterative_refinement;  // true to use iterative refinement for barrier method
-  f_t barrier_step_scale;             // step scale for barrier method
-  i_t barrier_soc_threshold;          // SOC dimension above which rank-2 sparse scaling is used
+  bool eliminate_dense_columns;             // true to eliminate dense columns from A*D*A^T
+  bool barrier_iterative_refinement;        // true to use iterative refinement for barrier method
+  i_t barrier_iterative_refinement_method;  // 0: fixed_point, 1: gmres
+  f_t barrier_step_scale;                   // step scale for barrier method
+  i_t barrier_soc_threshold;  // SOC dimension above which rank-2 sparse scaling is used
   int num_gpus;   // Number of GPUs to use (maximum of 2 gpus are supported at the moment)
   i_t folding;    // -1 automatic, 0 don't fold, 1 fold
   i_t augmented;  // -1 automatic, 0 to solve with ADAT, 1 to solve with augmented system

cpp/src/grpc/codegen/field_registry.yaml

@@ -573,6 +573,10 @@ pdlp_settings:
       field_num: 37
       type: int32
       optional: true
+  - barrier_iterative_refinement_method:
+      field_num: 38
+      type: int32
+      optional: true
   - save_best_primal_so_far:
       field_num: 28
       type: bool

cpp/src/grpc/codegen/generated/cuopt_remote_data.proto

@@ -198,6 +198,7 @@ message PDLPSolverSettings {
   bool cudss_hybrid_mode = 35;
   bool cudss_hybrid_execute_mode = 36;
   optional int32 cudss_host_nthreads = 37;
+  optional int32 barrier_iterative_refinement_method = 38;
   PDLPWarmStartData warm_start_data = 50;
 }
 

cpp/src/grpc/codegen/generated/generated_pdlp_settings_to_proto.inc

@@ -40,6 +40,7 @@
   pb_settings->set_cudss_hybrid_mode(settings.cudss_hybrid_mode);
   pb_settings->set_cudss_hybrid_execute_mode(settings.cudss_hybrid_execute_mode);
   pb_settings->set_cudss_host_nthreads(settings.cudss_host_nthreads);
+  pb_settings->set_barrier_iterative_refinement_method(settings.barrier_iterative_refinement_method);
   pb_settings->set_save_best_primal_so_far(settings.save_best_primal_so_far);
   pb_settings->set_first_primal_feasible(settings.first_primal_feasible);
   pb_settings->set_pdlp_precision(static_cast<int32_t>(settings.pdlp_precision));

cpp/src/grpc/codegen/generated/generated_proto_to_pdlp_settings.inc

@@ -90,6 +90,9 @@
   if (pb_settings.has_cudss_host_nthreads()) {
     settings.cudss_host_nthreads = pb_settings.cudss_host_nthreads();
   }
+  if (pb_settings.has_barrier_iterative_refinement_method()) {
+    settings.barrier_iterative_refinement_method = pb_settings.barrier_iterative_refinement_method();
+  }
   settings.save_best_primal_so_far = pb_settings.save_best_primal_so_far();
   settings.first_primal_feasible = pb_settings.first_primal_feasible();
   if (pb_settings.has_pdlp_precision()) {

cpp/src/math_optimization/solver_settings.cu

@@ -138,6 +138,7 @@ solver_settings_t<i_t, f_t>::solver_settings_t() : pdlp_settings(), mip_settings
     {CUOPT_ORDERING, &pdlp_settings.ordering, -1, 1, -1},
     {CUOPT_CUDSS_ND_NLEVELS, &pdlp_settings.cudss_nd_nlevels, -1, std::numeric_limits<i_t>::max(), -1},
     {CUOPT_CUDSS_HOST_NTHREADS, &pdlp_settings.cudss_host_nthreads, -1, std::numeric_limits<i_t>::max(), -1},
+    {CUOPT_BARRIER_ITERATIVE_REFINEMENT_METHOD, &pdlp_settings.barrier_iterative_refinement_method, 0, 1, 1},
     {CUOPT_BARRIER_DUAL_INITIAL_POINT, &pdlp_settings.barrier_dual_initial_point, -1, 1, -1},
     {CUOPT_POSTSOLVE_INFO, &pdlp_settings.postsolve_info, -1, 1, -1},
     {CUOPT_MIP_CUT_PASSES, &mip_settings.max_cut_passes, -1, std::numeric_limits<i_t>::max(), 10},