Add support for sparsity exploitation for SOCs of large dimensionality

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

@@ -283,6 +283,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_soc_threshold{100};
   f_t barrier_step_scale{0.9};
   bool save_best_primal_so_far{false};
   /**

cpp/src/barrier/barrier.hpp

@@ -16,6 +16,8 @@
 #include <math_optimization/tic_toc.hpp>
 
 #include <rmm/device_uvector.hpp>
+
+#include <utility>
 namespace cuopt::mathematical_optimization::barrier {
 
 /** Validates SOC layout on an simplex::lp_problem_t before barrier presolve/solve. */
@@ -89,9 +91,11 @@ class barrier_solver_t {
                                   f_t& dual_residual_norm,
                                   f_t& complementarity_residual_norm);
 
-  f_t compute_nonnegative_step_length(iteration_data_t<i_t, f_t>& data,
-                                      const rmm::device_uvector<f_t>& x,
-                                      const rmm::device_uvector<f_t>& dx);
+  std::pair<f_t, f_t> compute_nonnegative_step_length_pair(iteration_data_t<i_t, f_t>& data,
+                                                           const rmm::device_uvector<f_t>& x1,
+                                                           const rmm::device_uvector<f_t>& dx1,
+                                                           const rmm::device_uvector<f_t>& x2,
+                                                           const rmm::device_uvector<f_t>& dx2);
   i_t gpu_compute_search_direction(iteration_data_t<i_t, f_t>& data,
                                    f_t& dual_perturb,
                                    f_t& primal_perturb,

cpp/src/barrier/cusparse_view.cu

@@ -125,6 +125,40 @@ static cusparseSpMVAlg_t get_spmv_alg([[maybe_unused]] int num_rows)
   return CUSPARSE_SPMV_CSR_ALG2;
 }
 
+template <typename i_t, typename f_t>
+void cusparse_view_t<i_t, f_t>::init_spmv_buffer_and_preprocess(cusparseSpMatDescr_t mat,
+                                                                cusparseDnVecDescr_t x,
+                                                                cusparseDnVecDescr_t y,
+                                                                rmm::device_buffer& buffer,
+                                                                i_t rows)
+{
+  const auto spmv_alg     = get_spmv_alg(rows);
+  size_t buffer_size_spmv = 0;
+  RAFT_CUSPARSE_TRY(
+    raft::sparse::detail::cusparsespmv_buffersize(handle_ptr_->get_cusparse_handle(),
+                                                  CUSPARSE_OPERATION_NON_TRANSPOSE,
+                                                  d_one_.data(),
+                                                  mat,
+                                                  x,
+                                                  d_one_.data(),
+                                                  y,
+                                                  spmv_alg,
+                                                  &buffer_size_spmv,
+                                                  handle_ptr_->get_stream()));
+  buffer.resize(buffer_size_spmv, handle_ptr_->get_stream());
+
+  my_cusparsespmv_preprocess(handle_ptr_->get_cusparse_handle(),
+                             CUSPARSE_OPERATION_NON_TRANSPOSE,
+                             d_one_.data(),
+                             mat,
+                             x,
+                             d_one_.data(),
+                             y,
+                             spmv_alg,
+                             buffer.data(),
+                             handle_ptr_->get_stream());
+}
+
 template <typename i_t, typename f_t>
 cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
                                            const csc_matrix_t<i_t, f_t>& A)
@@ -136,6 +170,7 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
     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())
@@ -149,7 +184,7 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
   if (debug) { printf("A hash: %zu\n", A.hash()); }
   csr_matrix_t<i_t, f_t> A_csr(A.m, A.n, 1);
   A.to_compressed_row(A_csr);
-  i_t rows                        = A_csr.m;
+  rows_                           = A_csr.m;
   i_t cols                        = A_csr.n;
   i_t nnz                         = A_csr.x.size();
   const std::vector<i_t>& offsets = A_csr.row_start;
@@ -165,7 +200,7 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
   A_T_data_    = device_copy(A.x, handle_ptr->get_stream());
 
   cusparseCreateCsr(&A_,
-                    rows,
+                    rows_,
                     cols,
                     nnz,
                     A_offsets_.data(),
@@ -178,7 +213,7 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
 
   cusparseCreateCsr(&A_T_,
                     cols,
-                    rows,
+                    rows_,
                     nnz,
                     A_T_offsets_.data(),
                     A_T_indices_.data(),
@@ -192,58 +227,13 @@ cusparse_view_t<i_t, f_t>::cusparse_view_t(raft::handle_t const* handle_ptr,
   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());
+  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()));
 
-  size_t buffer_size_spmv = 0;
-  RAFT_CUSPARSE_TRY(
-    raft::sparse::detail::cusparsespmv_buffersize(handle_ptr_->get_cusparse_handle(),
-                                                  CUSPARSE_OPERATION_NON_TRANSPOSE,
-                                                  d_one_.data(),
-                                                  A_,
-                                                  x,
-                                                  d_one_.data(),
-                                                  y,
-                                                  get_spmv_alg(A_offsets_.size() - 1),
-                                                  &buffer_size_spmv,
-                                                  handle_ptr_->get_stream()));
-  spmv_buffer_.resize(buffer_size_spmv, handle_ptr_->get_stream());
-
-  my_cusparsespmv_preprocess(handle_ptr_->get_cusparse_handle(),
-                             CUSPARSE_OPERATION_NON_TRANSPOSE,
-                             d_one_.data(),
-                             A_,
-                             x,
-                             d_one_.data(),
-                             y,
-                             get_spmv_alg(A_offsets_.size() - 1),
-                             spmv_buffer_.data(),
-                             handle_ptr->get_stream());
+  init_spmv_buffer_and_preprocess(A_, x, y, spmv_buffer_, rows_);
+  init_spmv_buffer_and_preprocess(A_T_, y, x, spmv_buffer_transpose_, A_T_offsets_.size() - 1);
 
-  RAFT_CUSPARSE_TRY(
-    raft::sparse::detail::cusparsespmv_buffersize(handle_ptr_->get_cusparse_handle(),
-                                                  CUSPARSE_OPERATION_NON_TRANSPOSE,
-                                                  d_one_.data(),
-                                                  A_T_,
-                                                  y,
-                                                  d_one_.data(),
-                                                  x,
-                                                  get_spmv_alg(A_T_offsets_.size() - 1),
-                                                  &buffer_size_spmv,
-                                                  handle_ptr_->get_stream()));
-  spmv_buffer_transpose_.resize(buffer_size_spmv, handle_ptr_->get_stream());
-
-  my_cusparsespmv_preprocess(handle_ptr_->get_cusparse_handle(),
-                             CUSPARSE_OPERATION_NON_TRANSPOSE,
-                             d_one_.data(),
-                             A_T_,
-                             y,
-                             d_one_.data(),
-                             x,
-                             get_spmv_alg(A_T_offsets_.size() - 1),
-                             spmv_buffer_transpose_.data(),
-                             handle_ptr->get_stream());
   RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(x));
   RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y));
 }
@@ -252,7 +242,7 @@ template <typename i_t, typename f_t>
 cusparse_view_t<i_t, f_t>::~cusparse_view_t()
 {
   CUOPT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(A_));
-  CUOPT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(A_T_));
+  if (A_T_ != nullptr) { CUOPT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(A_T_)); }
 }
 
 template <typename i_t, typename f_t>
@@ -311,7 +301,7 @@ void cusparse_view_t<i_t, f_t>::spmv(f_t alpha,
                                      x,
                                      d_beta->data(),
                                      y,
-                                     get_spmv_alg(A_offsets_.size() - 1),
+                                     get_spmv_alg(rows_),
                                      (f_t*)spmv_buffer_.data(),
                                      handle_ptr_->get_stream());
 }
@@ -323,6 +313,7 @@ void cusparse_view_t<i_t, f_t>::transpose_spmv(f_t alpha,
                                                f_t beta,
                                                std::vector<f_t, AllocatorB>& y)
 {
+  cuopt_assert(A_T_ != nullptr, "transpose_spmv requires an A^T descriptor");
   auto d_x = device_copy(x, handle_ptr_->get_stream());
   auto d_y = device_copy(y, handle_ptr_->get_stream());
   transpose_spmv(alpha, d_x, beta, d_y);
@@ -335,6 +326,7 @@ void cusparse_view_t<i_t, f_t>::transpose_spmv(f_t alpha,
                                                f_t beta,
                                                rmm::device_uvector<f_t>& y)
 {
+  cuopt_assert(A_T_ != nullptr, "transpose_spmv requires an A^T descriptor");
   pdlp::cusparse_dn_vec_descr_wrapper_t<f_t> x_cusparse = create_vector(x);
   pdlp::cusparse_dn_vec_descr_wrapper_t<f_t> y_cusparse = create_vector(y);
   transpose_spmv(alpha, x_cusparse, beta, y_cusparse);
@@ -346,6 +338,7 @@ void cusparse_view_t<i_t, f_t>::transpose_spmv(f_t alpha,
                                                f_t beta,
                                                pdlp::cusparse_dn_vec_descr_wrapper_t<f_t> const& y)
 {
+  cuopt_assert(A_T_ != nullptr, "transpose_spmv requires an A^T descriptor");
   // Would be simpler if we could pass host data direct;y but other cusparse calls with the same
   // handler depend on device data
   cuopt_assert(alpha == f_t(1) || alpha == f_t(-1), "Only alpha 1 or -1 supported");

cpp/src/barrier/cusparse_view.hpp

@@ -6,6 +6,8 @@
 /* clang-format on */
 #pragma once
 
+#include <barrier/device_sparse_matrix.cuh>
+
 #include <linear_algebra/sparse_matrix.hpp>
 
 #include <pdlp/cusparse_view.hpp>
@@ -17,15 +19,14 @@
 
 #include <raft/core/handle.hpp>
 
-// Lightweight cuSparse view
-// Only owns data linked to the associated matrix
-// Associated dense vector should be owned by the calling object
-// This allows handling many different X Y vector along with one common matrix
+// Lightweight cuSparse view over a sparse matrix descriptor. The dense vectors
+// are owned by the caller, which allows many x/y pairs to share one matrix view.
 namespace cuopt::mathematical_optimization::barrier {
 
 template <typename i_t, typename f_t>
 class cusparse_view_t {
  public:
+  // 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);
   ~cusparse_view_t();
@@ -59,18 +60,25 @@ class cusparse_view_t {
   raft::handle_t const* handle_ptr_{nullptr};
 
  private:
+  void init_spmv_buffer_and_preprocess(cusparseSpMatDescr_t mat,
+                                       cusparseDnVecDescr_t x,
+                                       cusparseDnVecDescr_t y,
+                                       rmm::device_buffer& buffer,
+                                       i_t rows);
+
   rmm::device_uvector<i_t> A_offsets_;
   rmm::device_uvector<i_t> A_indices_;
   rmm::device_uvector<f_t> A_data_;
-  cusparseSpMatDescr_t A_;
+  cusparseSpMatDescr_t A_{nullptr};
   rmm::device_uvector<i_t> A_T_offsets_;
   rmm::device_uvector<i_t> A_T_indices_;
   rmm::device_uvector<f_t> A_T_data_;
-  cusparseSpMatDescr_t A_T_;
+  cusparseSpMatDescr_t A_T_{nullptr};
   rmm::device_buffer spmv_buffer_;
   rmm::device_buffer spmv_buffer_transpose_;
   rmm::device_scalar<f_t> d_one_;
   rmm::device_scalar<f_t> d_minus_one_;
   rmm::device_scalar<f_t> d_zero_;
+  i_t rows_{0};
 };
 }  // namespace cuopt::mathematical_optimization::barrier

cpp/src/barrier/device_sparse_matrix.cuh

@@ -88,6 +88,60 @@ struct transform_reduce_helper_t {
   }
 };
 
+template <typename f_t>
+struct f2_t {
+  f_t a;
+  f_t b;
+};
+
+template <typename f_t>
+struct f2_min_t {
+  HD f2_t<f_t> operator()(const f2_t<f_t>& lhs, const f2_t<f_t>& rhs) const
+  {
+    return f2_t<f_t>{cuda::std::min(lhs.a, rhs.a), cuda::std::min(lhs.b, rhs.b)};
+  }
+};
+
+template <typename f_t>
+struct transform_reduce_pair_helper_t {
+  rmm::device_buffer buffer_data;
+  rmm::device_scalar<f2_t<f_t>> out;
+  size_t buffer_size;
+
+  transform_reduce_pair_helper_t(rmm::cuda_stream_view stream_view)
+    : buffer_data(0, stream_view), out(stream_view)
+  {
+  }
+
+  // TransformOpT must map each input element to an f2_t<f_t>{a, b} pair; the two
+  // components are reduced independently (elementwise min) in a single kernel launch.
+  template <typename InputIteratorT, typename TransformOpT, typename i_t>
+  f2_t<f_t> transform_reduce(InputIteratorT input,
+                             TransformOpT transform_op,
+                             f2_t<f_t> init,
+                             i_t size,
+                             rmm::cuda_stream_view stream_view)
+  {
+    f2_min_t<f_t> reduce_op{};
+    cub::DeviceReduce::TransformReduce(
+      nullptr, buffer_size, input, out.data(), size, reduce_op, transform_op, init, stream_view);
+
+    buffer_data.resize(buffer_size, stream_view);
+
+    cub::DeviceReduce::TransformReduce(buffer_data.data(),
+                                       buffer_size,
+                                       input,
+                                       out.data(),
+                                       size,
+                                       reduce_op,
+                                       transform_op,
+                                       init,
+                                       stream_view);
+
+    return out.value(stream_view);
+  }
+};
+
 template <typename i_t, typename f_t>
 struct csc_view_t {
   raft::device_span<i_t> col_start;
@@ -131,7 +185,8 @@ class device_csc_matrix_t {
       nz_max(A.col_start[A.n]),
       col_start(A.col_start.size(), stream),
       i(A.i.size(), stream),
-      x(A.x.size(), stream)
+      x(A.x.size(), stream),
+      col_index(0, stream)
   {
     col_start = cuopt::device_copy(A.col_start, stream);
     i         = cuopt::device_copy(A.i, stream);