Use CSR representation of augmented matrix for multiplication in iterative refinement

cpp/include/cuopt/mathematical_optimization/constants.h

@@ -50,6 +50,7 @@
 #define CUOPT_POSTSOLVE_INFO                        "postsolve_info"
 #define CUOPT_BARRIER_PRESOLVE_BOUND_FREE_VARIABLES "barrier_presolve_bound_free_variables"
 #define CUOPT_BARRIER_ITERATIVE_REFINEMENT          "barrier_iterative_refinement"
+#define CUOPT_BARRIER_CSR_IR_MATVEC                 "barrier_csr_ir_matvec"
 #define CUOPT_BARRIER_STEP_SCALE                    "barrier_step_scale"
 #define CUOPT_ELIMINATE_DENSE_COLUMNS               "eliminate_dense_columns"
 #define CUOPT_CUDSS_DETERMINISTIC                   "cudss_deterministic"
@@ -241,6 +242,9 @@
 #define CUOPT_BARRIER_ITERATIVE_REFINEMENT_OFF 0
 #define CUOPT_BARRIER_ITERATIVE_REFINEMENT_ON  1
 
+#define CUOPT_BARRIER_CSR_IR_MATVEC_OFF 0
+#define CUOPT_BARRIER_CSR_IR_MATVEC_ON  1
+
 /* @brief Scalar problem attribute selectors
  * Passed as cuopt_int_t; the valid set depends on the accessor's value type. */
 #define CUOPT_ATTR_NUM_VARIABLES             0

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

@@ -304,6 +304,9 @@ class pdlp_solver_settings_t {
   bool eliminate_dense_columns{true};
   pdlp_precision_t pdlp_precision{pdlp_precision_t::DefaultPrecision};
   bool barrier_iterative_refinement{true};
+  // true to use a single cuSPARSE SpMV over the unperturbed augmented CSR for the barrier
+  // solver's IR matvec, instead of the matrix-free path. Experimental; default off.
+  bool barrier_csr_ir_matvec{false};
   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

@@ -709,6 +709,8 @@ class iteration_data_t {
 
   i_t augmented_system_size(i_t n, i_t m) const { return n + m + augmented_expansion_count(); }
 
+  bool use_csr_ir_matvec() const { return settings_.barrier_csr_ir_matvec && use_augmented; }
+
   bool is_cone_variable(i_t variable) const
   {
     return has_cones() && variable >= cone_start() && variable < cone_end();
@@ -1898,6 +1900,55 @@ class iteration_data_t {
     handle_ptr->sync_stream();
   }
 
+  // Undo the dual_perturb/primal_perturb regularization baked into device_augmented.x by
+  // form_augmented, in place. Must be called after every chol->factorize(device_augmented) and
+  // before augmented_csr_multiply is used, since IR's matvec needs the true unperturbed KKT
+  // operator while the factorization itself must stay regularized for stability. No-op unless
+  // use_csr_ir_matvec().
+  void strip_augmented_perturbation()
+  {
+    if (!use_csr_ir_matvec()) { return; }
+    raft::common::nvtx::range fun_scope("Barrier: strip_augmented_perturbation");
+    cuopt::mathematical_optimization::barrier::strip_augmented_perturbation<i_t, f_t>(
+      A.n,
+      A.m,
+      augmented_expansion_count(),
+      dual_perturb,
+      primal_perturb,
+      d_augmented_diagonal_indices_,
+      cone_kkt_data_,
+      device_augmented,
+      stream_view_);
+    handle_ptr->sync_stream();
+  }
+
+  // Lazily wire a no-copy cuSparse view over device_augmented's buffers. Built once (the
+  // augmented CSR's sparsity pattern is only constructed on first_call, so device_augmented.x's
+  // pointer is stable thereafter); rebuilt defensively if that pointer ever changes.
+  void ensure_augmented_csr_view()
+  {
+    if (cusparse_augmented_view_ != nullptr &&
+        cusparse_augmented_view_data_ptr_ == device_augmented.x.data()) {
+      return;
+    }
+    cusparse_augmented_view_ =
+      std::make_unique<cusparse_view_t<i_t, f_t>>(handle_ptr, device_augmented);
+    cusparse_augmented_view_data_ptr_ = device_augmented.x.data();
+  }
+
+  // Drop-in alternative to augmented_multiply(): a single cuSPARSE SpMV over the already-
+  // factorized, perturbation-stripped device_augmented CSR buffer.
+  void augmented_csr_multiply(f_t alpha,
+                              const rmm::device_uvector<f_t>& x,
+                              f_t beta,
+                              rmm::device_uvector<f_t>& y)
+  {
+    raft::common::nvtx::range fun_scope("Barrier: augmented_csr_multiply");
+    cuopt_assert(use_csr_ir_matvec(), "augmented_csr_multiply requires CSR IR matvec path");
+    ensure_augmented_csr_view();
+    cusparse_augmented_view_->spmv(alpha, x, beta, y);
+  }
+
   raft::handle_t const* handle_ptr;
   i_t n_upper_bounds;
   dense_vector_t<i_t, i_t> upper_bounds;
@@ -1982,6 +2033,11 @@ class iteration_data_t {
 
   std::unique_ptr<sparse_cholesky_base_t<i_t, f_t>> chol;
 
+  // No-copy cuSparse SpMV view over device_augmented, used by augmented_csr_multiply() when
+  // use_csr_ir_matvec() is enabled. Built lazily by ensure_augmented_csr_view().
+  std::unique_ptr<cusparse_view_t<i_t, f_t>> cusparse_augmented_view_;
+  const f_t* cusparse_augmented_view_data_ptr_{nullptr};
+
   bool has_factorization;
   bool has_solve_info;
   i_t num_factorizations;
@@ -2226,6 +2282,7 @@ int barrier_solver_t<i_t, f_t>::initial_point(iteration_data_t<i_t, f_t>& data)
   i_t status;
   if (use_augmented) {
     status = data.chol->factorize(data.device_augmented);
+    data.strip_augmented_perturbation();
 
 #ifdef CHOLESKY_DEBUG_CHECK
     cholesky_debug_check(data, lp, use_augmented);
@@ -2270,7 +2327,11 @@ int barrier_solver_t<i_t, f_t>::initial_point(iteration_data_t<i_t, f_t>& data)
                       f_t beta,
                       rmm::device_uvector<f_t>& y) const
       {
-        data_.augmented_multiply(alpha, x, beta, y);
+        if (data_.use_csr_ir_matvec()) {
+          data_.augmented_csr_multiply(alpha, x, beta, y);
+        } else {
+          data_.augmented_multiply(alpha, x, beta, y);
+        }
       }
       void solve(rmm::device_uvector<f_t>& b, rmm::device_uvector<f_t>& x) const
       {
@@ -2793,6 +2854,7 @@ i_t barrier_solver_t<i_t, f_t>::gpu_compute_search_direction(iteration_data_t<i_
       {
         raft::common::nvtx::range fun_scope("Barrier: factorize");
         status = data.chol->factorize(data.device_augmented);
+        data.strip_augmented_perturbation();
       }
 
 #ifdef CHOLESKY_DEBUG_CHECK
@@ -2895,7 +2957,11 @@ i_t barrier_solver_t<i_t, f_t>::gpu_compute_search_direction(iteration_data_t<i_
                       f_t beta,
                       rmm::device_uvector<f_t>& y)
       {
-        data_.augmented_multiply(alpha, x, beta, y);
+        if (data_.use_csr_ir_matvec()) {
+          data_.augmented_csr_multiply(alpha, x, beta, y);
+        } else {
+          data_.augmented_multiply(alpha, x, beta, y);
+        }
       }
 
       void solve(rmm::device_uvector<f_t>& b, rmm::device_uvector<f_t>& x) const

cpp/src/barrier/csr_kkt_build.cuh

@@ -744,4 +744,89 @@ i_t build_augmented_csr_on_device(i_t n,
   return total_nnz;
 }
 
+// Undo the dual_perturb/primal_perturb regularization that fill_augmented_csr_row_kernel bakes
+// into device_augmented.x, in place. The augmented matrix must stay regularized for
+// Cholesky factorization (stability), but the CSR-based IR matvec (augmented_csr_multiply) needs
+// the true unperturbed KKT operator -- so this is applied to the already-factorized buffer,
+// after chol->factorize() and before any CSR-based SpMV against it. Not idempotent: must run
+// exactly once per refresh of device_augmented.x.
+template <std::integral i_t, std::floating_point f_t>
+__global__ void strip_primal_block_diag_kernel(raft::device_span<f_t> augmented_x,
+                                               raft::device_span<const i_t> diag_indices,
+                                               f_t dual_perturb,
+                                               i_t n)
+{
+  const i_t row = static_cast<i_t>(blockIdx.x * blockDim.x + threadIdx.x);
+  if (row >= n) { return; }
+  augmented_x[diag_indices[row]] += dual_perturb;
+}
+
+template <std::integral i_t, std::floating_point f_t>
+__global__ void strip_constraint_block_diag_kernel(raft::device_span<f_t> augmented_x,
+                                                   raft::device_span<const i_t> diag_indices,
+                                                   f_t primal_perturb,
+                                                   i_t n,
+                                                   i_t m)
+{
+  const i_t l = static_cast<i_t>(blockIdx.x * blockDim.x + threadIdx.x);
+  if (l >= m) { return; }
+  augmented_x[diag_indices[n + l]] -= primal_perturb;
+}
+
+template <std::integral i_t, std::floating_point f_t>
+__global__ void strip_sparse_expansion_D_kernel(raft::device_span<f_t> augmented_x,
+                                                raft::device_span<const i_t> sparse_expansion_D,
+                                                f_t dual_perturb)
+{
+  const i_t e = static_cast<i_t>(blockIdx.x * blockDim.x + threadIdx.x);
+  if (e >= static_cast<i_t>(sparse_expansion_D.size())) { return; }
+  const i_t idx = sparse_expansion_D[e];
+  if (idx < 0) { return; }
+  const f_t sign = (e % 2 == 0) ? f_t(1) : f_t(-1);
+  augmented_x[idx] += sign * dual_perturb;
+}
+
+template <std::integral i_t, std::floating_point f_t>
+void strip_augmented_perturbation(i_t n,
+                                  i_t m,
+                                  i_t p,
+                                  f_t dual_perturb,
+                                  f_t primal_perturb,
+                                  rmm::device_uvector<i_t>& augmented_diagonal_indices,
+                                  cone_kkt_data_t<i_t, f_t>& cone_data,
+                                  device_csr_matrix_t<i_t, f_t>& device_augmented,
+                                  rmm::cuda_stream_view stream)
+{
+  raft::common::nvtx::range scope("Barrier: strip augmented perturbation");
+  if (n > 0) {
+    const size_t grid = raft::ceildiv<size_t>(n, augmented_csr_block_size);
+    strip_primal_block_diag_kernel<i_t, f_t><<<grid, augmented_csr_block_size, 0, stream.value()>>>(
+      cuopt::make_span(device_augmented.x),
+      cuopt::make_span(augmented_diagonal_indices),
+      dual_perturb,
+      n);
+    RAFT_CUDA_TRY(cudaPeekAtLastError());
+  }
+  if (m > 0) {
+    const size_t grid = raft::ceildiv<size_t>(m, augmented_csr_block_size);
+    strip_constraint_block_diag_kernel<i_t, f_t>
+      <<<grid, augmented_csr_block_size, 0, stream.value()>>>(
+        cuopt::make_span(device_augmented.x),
+        cuopt::make_span(augmented_diagonal_indices),
+        primal_perturb,
+        n,
+        m);
+    RAFT_CUDA_TRY(cudaPeekAtLastError());
+  }
+  if (p > 0) {
+    const size_t grid = raft::ceildiv<size_t>(static_cast<size_t>(p), augmented_csr_block_size);
+    strip_sparse_expansion_D_kernel<i_t, f_t>
+      <<<grid, augmented_csr_block_size, 0, stream.value()>>>(
+        cuopt::make_span(device_augmented.x),
+        cuopt::make_span(cone_data.sparse_expansion_D),
+        dual_perturb);
+    RAFT_CUDA_TRY(cudaPeekAtLastError());
+  }
+}
+
 }  // namespace cuopt::mathematical_optimization::barrier

cpp/src/barrier/cusparse_view.cu

@@ -238,6 +238,54 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
   RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y));
 }
 
+template <typename i_t, typename f_t>
+cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
+                                           device_csr_matrix_t<i_t, f_t>& csr)
+  : handle_ptr_(handle_ptr),
+    A_offsets_(0, handle_ptr->get_stream()),
+    A_indices_(0, handle_ptr->get_stream()),
+    A_data_(0, handle_ptr->get_stream()),
+    A_T_offsets_(0, handle_ptr->get_stream()),
+    A_T_indices_(0, handle_ptr->get_stream()),
+    A_T_data_(0, handle_ptr->get_stream()),
+    spmv_buffer_(0, handle_ptr->get_stream()),
+    spmv_buffer_transpose_(0, handle_ptr->get_stream()),
+    d_one_(f_t(1), handle_ptr->get_stream()),
+    d_minus_one_(f_t(-1), handle_ptr->get_stream()),
+    d_zero_(f_t(0), handle_ptr->get_stream()),
+    rows_(csr.m)
+{
+  RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsesetpointermode(
+    handle_ptr->get_cusparse_handle(), CUSPARSE_POINTER_MODE_DEVICE, handle_ptr->get_stream()));
+
+  const i_t cols = csr.n;
+  const i_t nnz  = csr.nz_max;
+
+  cusparseCreateCsr(&A_,
+                    rows_,
+                    cols,
+                    nnz,
+                    csr.row_start.data(),
+                    csr.j.data(),
+                    csr.x.data(),
+                    CUSPARSE_INDEX_32I,
+                    CUSPARSE_INDEX_32I,
+                    CUSPARSE_INDEX_BASE_ZERO,
+                    CUDA_R_64F);
+
+  cusparseDnVecDescr_t x;
+  cusparseDnVecDescr_t y;
+  rmm::device_uvector<f_t> d_x(cols, handle_ptr_->get_stream());
+  rmm::device_uvector<f_t> d_y(rows_, handle_ptr_->get_stream());
+  RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsecreatednvec(&x, d_x.size(), d_x.data()));
+  RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsecreatednvec(&y, d_y.size(), d_y.data()));
+
+  init_spmv_buffer_and_preprocess(A_, x, y, spmv_buffer_, rows_);
+
+  RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(x));
+  RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y));
+}
+
 template <typename i_t, typename f_t>
 cusparse_view_t<i_t, f_t>::~cusparse_view_t()
 {

cpp/src/barrier/cusparse_view.hpp

@@ -29,6 +29,11 @@ class cusparse_view_t {
   // Copy CSC -> owned CSR + CSC-transpose, with preprocess. Supports forward and transpose SpMV.
   // TMP matrix data should already be on the GPU and in CSR not CSC
   cusparse_view_t(raft::handle_t const* handle_ptr, const csc_matrix_t<i_t, f_t>& A);
+
+  // Wire cuSparse SpMV over existing device CSR buffers (no copy). Forward SpMV only (A_T_ stays
+  // null). The caller must keep `csr` alive and its row_start/j arrays unresized for the life of
+  // this view; only the contents of csr.x may change between spmv() calls.
+  cusparse_view_t(raft::handle_t const* handle_ptr, device_csr_matrix_t<i_t, f_t>& csr);
   ~cusparse_view_t();
 
   pdlp::cusparse_dn_vec_descr_wrapper_t<f_t> create_vector(rmm::device_uvector<f_t> const& vec);

cpp/src/dual_simplex/simplex_solver_settings.hpp

@@ -70,6 +70,7 @@ struct simplex_solver_settings_t {
       barrier(false),
       eliminate_dense_columns(true),
       barrier_iterative_refinement(true),
+      barrier_csr_ir_matvec(false),
       barrier_step_scale(0.9),
       barrier_soc_threshold(100),
       num_gpus(1),
@@ -165,8 +166,11 @@ struct simplex_solver_settings_t {
   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
+    barrier_csr_ir_matvec;  // true to use a single cuSPARSE SpMV over the unperturbed augmented
+                            // CSR for the IR matvec, instead of the matrix-free augmented_multiply
+  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/math_optimization/solver_settings.cu

@@ -207,6 +207,7 @@ solver_settings_t<i_t, f_t>::solver_settings_t() : pdlp_settings(), mip_settings
     {CUOPT_CUDSS_DETERMINISTIC, &pdlp_settings.cudss_deterministic, false},
     {CUOPT_DUAL_POSTSOLVE, &pdlp_settings.dual_postsolve, true},
     {CUOPT_BARRIER_ITERATIVE_REFINEMENT, &pdlp_settings.barrier_iterative_refinement, true},
+    {CUOPT_BARRIER_CSR_IR_MATVEC, &pdlp_settings.barrier_csr_ir_matvec, false},
     {CUOPT_MIP_PROBING, &mip_settings.probing, true},
     {CUOPT_USE_DISTRIBUTED_PDLP, &pdlp_settings.use_distributed_pdlp, false},
     // Diving heuristic hyper-parameters (hidden from default --help: name contains "hyper_")

cpp/src/pdlp/solve.cu

@@ -515,6 +515,7 @@ std::tuple<simplex::lp_solution_t<i_t, f_t>, simplex::lp_status_t, f_t, f_t, f_t
   barrier_settings.crossover                       = settings.crossover;
   barrier_settings.eliminate_dense_columns         = settings.eliminate_dense_columns;
   barrier_settings.barrier_iterative_refinement    = settings.barrier_iterative_refinement;
+  barrier_settings.barrier_csr_ir_matvec           = settings.barrier_csr_ir_matvec;
   barrier_settings.barrier_soc_threshold           = settings.barrier_soc_threshold;
   barrier_settings.barrier_step_scale              = settings.barrier_step_scale;
   barrier_settings.qcqp_ruiz_equilibration         = settings.qcqp_ruiz_equilibration;
@@ -703,6 +704,7 @@ static optimization_problem_solution_t<i_t, double> run_pdlp_solver_in_fp32(
   fs.all_primal_feasible          = settings.all_primal_feasible;
   fs.eliminate_dense_columns      = settings.eliminate_dense_columns;
   fs.barrier_iterative_refinement = settings.barrier_iterative_refinement;
+  fs.barrier_csr_ir_matvec        = settings.barrier_csr_ir_matvec;
   fs.barrier_step_scale           = settings.barrier_step_scale;
   fs.pdlp_precision               = pdlp_precision_t::DefaultPrecision;
   fs.method                       = method_t::PDLP;