From 01f346932591a781eb80252f3ed71938a6e01209 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Wed, 24 Jun 2026 07:31:37 -0700 Subject: [PATCH 01/14] Add support for sparsity exploitation for SOCP of large dimensionality Signed-off-by: yuwenchen95 --- .../cuopt/linear_programming/constants.h | 1 + .../pdlp/solver_settings.hpp | 1 + cpp/src/barrier/barrier.cu | 491 +++++++++++--- cpp/src/barrier/iterative_refinement.hpp | 1 - cpp/src/barrier/second_order_cone_kernels.cuh | 601 ++++++++++++++++-- .../barrier/second_order_cone_reduction.cuh | 116 ++-- .../dual_simplex/simplex_solver_settings.hpp | 2 + cpp/src/math_optimization/solver_settings.cu | 1 + cpp/src/pdlp/solve.cu | 1 + cpp/tests/socp/CMakeLists.txt | 1 + cpp/tests/socp/second_order_cone_kernels.cu | 355 ++++++++++- cpp/tests/socp/solve_barrier_socp.cu | 210 ++++++ cpp/tests/socp/sparse_augmented_kkt_test.cu | 353 ++++++++++ 13 files changed, 1920 insertions(+), 214 deletions(-) create mode 100644 cpp/tests/socp/sparse_augmented_kkt_test.cu diff --git a/cpp/include/cuopt/linear_programming/constants.h b/cpp/include/cuopt/linear_programming/constants.h index 024fae7583..041cc38688 100644 --- a/cpp/include/cuopt/linear_programming/constants.h +++ b/cpp/include/cuopt/linear_programming/constants.h @@ -48,6 +48,7 @@ #define CUOPT_ORDERING "ordering" #define CUOPT_BARRIER_DUAL_INITIAL_POINT "barrier_dual_initial_point" #define CUOPT_BARRIER_ITERATIVE_REFINEMENT "barrier_iterative_refinement" +#define CUOPT_BARRIER_SOC_THRESHOLD "barrier_soc_threshold" #define CUOPT_BARRIER_STEP_SCALE "barrier_step_scale" #define CUOPT_ELIMINATE_DENSE_COLUMNS "eliminate_dense_columns" #define CUOPT_CUDSS_DETERMINISTIC "cudss_deterministic" diff --git a/cpp/include/cuopt/linear_programming/pdlp/solver_settings.hpp b/cpp/include/cuopt/linear_programming/pdlp/solver_settings.hpp index b30286f9ce..dd1979de6c 100644 --- a/cpp/include/cuopt/linear_programming/pdlp/solver_settings.hpp +++ b/cpp/include/cuopt/linear_programming/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{5}; f_t barrier_step_scale{0.9}; bool save_best_primal_so_far{false}; /** diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 1cbeed5cf3..845624b6ef 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -40,6 +40,8 @@ #include #include +#include + #include #include #include @@ -262,6 +264,16 @@ class iteration_data_t { d_augmented_diagonal_indices_(0, lp.handle_ptr->get_stream()), d_cone_csr_indices_(0, lp.handle_ptr->get_stream()), d_cone_Q_values_(0, lp.handle_ptr->get_stream()), + d_dense_cone_block_offsets_(0, lp.handle_ptr->get_stream()), + d_dense_cone_ids_(0, lp.handle_ptr->get_stream()), + d_sparse_hessian_diag_(0, lp.handle_ptr->get_stream()), + d_sparse_hessian_Q_(0, lp.handle_ptr->get_stream()), + d_sparse_exp_v_col_(0, lp.handle_ptr->get_stream()), + d_sparse_exp_u_col_(0, lp.handle_ptr->get_stream()), + d_sparse_exp_v_row_(0, lp.handle_ptr->get_stream()), + d_sparse_exp_u_row_(0, lp.handle_ptr->get_stream()), + d_sparse_expansion_D_(0, lp.handle_ptr->get_stream()), + d_sparse_Hs_diag_(0, lp.handle_ptr->get_stream()), use_augmented(false), has_factorization(false), n_direct_free_linear(0), @@ -406,7 +418,8 @@ class iteration_data_t { std::span(lp.second_order_cone_dims.data(), lp.second_order_cone_dims.size()), raft::device_span{}, raft::device_span{}, - stream_view_); + stream_view_, + settings.barrier_soc_threshold); cuopt_assert(cone_count() > 0, "second-order cone topology must contain at least one cone"); cuopt_assert(cone_entry_count() == total_cone_dim, "second-order cone entry count mismatch"); } @@ -506,7 +519,7 @@ class iteration_data_t { if (use_augmented) { settings.log.printf("Linear system : augmented\n"); - const i_t augmented_size = lp.num_cols + lp.num_rows; + const i_t augmented_size = augmented_system_size(lp.num_cols, lp.num_rows); d_augmented_rhs_.resize(augmented_size, stream_view_); d_augmented_soln_.resize(augmented_size, stream_view_); } else { @@ -599,8 +612,9 @@ class iteration_data_t { if (settings.concurrent_halt != nullptr && *settings.concurrent_halt == 1) { return; } { raft::common::nvtx::range scope("Barrier: LP Data: Cholesky init"); - i_t factorization_size = use_augmented ? lp.num_rows + lp.num_cols : lp.num_rows; - chol = std::make_unique>( + i_t factorization_size = + use_augmented ? augmented_system_size(lp.num_cols, lp.num_rows) : lp.num_rows; + chol = std::make_unique>( handle_ptr, settings, factorization_size); chol->set_positive_definite(false); } @@ -653,16 +667,23 @@ class iteration_data_t { i_t cone_end() const { return cone_start() + cone_entry_count(); } - i_t linear_xz_size(std::size_t full_xz_size) const + i_t augmented_expansion_count() const { - return has_cones() ? cone_start() : static_cast(full_xz_size); + return has_cones() && cones().has_sparse_cones() ? cones().expansion_var_count() : i_t(0); } + i_t augmented_system_size(i_t n, i_t m) const { return n + m + augmented_expansion_count(); } + bool is_cone_variable(i_t variable) const { return has_cones() && variable >= cone_start() && variable < cone_end(); } + i_t linear_xz_size(std::size_t full_xz_size) const + { + return has_cones() ? cone_start() : static_cast(full_xz_size); + } + f_t complementarity_degree(std::size_t num_primal_variables, i_t num_upper_bounds) const { const bool has_soc = has_cones(); @@ -678,52 +699,89 @@ class iteration_data_t { void form_augmented(bool first_call = false) { - i_t n = A.n; - i_t m = A.m; - i_t nnzA = A.col_start[n]; - i_t nnzQ = Q.n > 0 ? Q.col_start[n] : 0; - i_t factorization_size = n + m; + i_t n = A.n; + i_t m = A.m; + i_t nnzA = A.col_start[n]; + i_t nnzQ = Q.n > 0 ? Q.col_start[n] : 0; + + const bool has_soc = has_cones(); + const i_t m_c = cone_entry_count(); + const i_t p = augmented_expansion_count(); + i_t factorization_size = augmented_system_size(n, m); - const bool has_soc = has_cones(); - const i_t m_c = cone_entry_count(); - i_t total_block_nnz = 0; + i_t dense_soc_kkt_nnz = 0; std::vector cone_offsets_host; - std::vector cone_block_offsets_host; + std::vector dense_cone_block_offsets_host; + std::vector dense_cone_ids_host; + std::vector cone_sparse_idx; + std::vector sparse_cone_ids_host; + std::vector sparse_entry_offsets; std::vector local_to_cone; if (first_call) { + const i_t n_sparse_soc = has_soc ? cones().n_sparse_cones : i_t(0); if (has_soc) { raft::common::nvtx::range scope("Barrier: augmented: SOC offsets"); - // Initialize the cone block offsets const i_t n_cones = cone_count(); + // TODO: Once form_augmented() is moved to GPU, we can avoid D2H round-trip. cone_offsets_host.resize(n_cones + 1); - cone_block_offsets_host.resize(n_cones + 1); raft::copy( cone_offsets_host.data(), cones().cone_offsets.data(), n_cones + 1, stream_view_); handle_ptr->sync_stream(); + + cone_sparse_idx.assign(n_cones, i_t(-1)); + if (n_sparse_soc > 0) { + // TODO: Once form_augmented() is moved to GPU, we can avoid D2H round-trip for + // sparse_cone_ids / sparse_cone_dims. + sparse_cone_ids_host = cuopt::host_copy(cones().sparse_cone_ids, stream_view_); + for (i_t s = 0; s < n_sparse_soc; ++s) { + cone_sparse_idx[sparse_cone_ids_host[s]] = s; + } + sparse_entry_offsets.resize(n_sparse_soc + 1, 0); + const auto sparse_dims_host = cuopt::host_copy(cones().sparse_cone_dims, stream_view_); + for (i_t s = 0; s < n_sparse_soc; ++s) { + sparse_entry_offsets[s + 1] = sparse_entry_offsets[s] + sparse_dims_host[s]; + } + } + + dense_cone_ids_host.reserve(n_cones); + dense_cone_block_offsets_host = {0}; + dense_cone_block_offsets_host.reserve(n_cones + 1); for (i_t k = 0; k < n_cones; ++k) { - const auto q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; - cone_block_offsets_host[k + 1] = cone_block_offsets_host[k] + q_k * q_k; + if (cone_sparse_idx[k] < 0) { + dense_cone_ids_host.push_back(k); + const auto q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; + dense_cone_block_offsets_host.push_back(dense_cone_block_offsets_host.back() + + q_k * q_k); + } } - total_block_nnz = static_cast(cone_block_offsets_host[n_cones]); + dense_soc_kkt_nnz = static_cast(dense_cone_block_offsets_host.back()); - // Precompute: for each local cone entry, which cone does it belong to? local_to_cone.resize(m_c); for (i_t k = 0; k < n_cones; ++k) { const i_t lo = cone_offsets_host[k]; const i_t hi = cone_offsets_host[k + 1]; - for (i_t p = lo; p < hi; p++) { - local_to_cone[p] = k; + for (i_t idx = lo; idx < hi; idx++) { + local_to_cone[idx] = k; } } } - i_t new_nnz = 2 * nnzA + n + m + nnzQ + total_block_nnz; // conservative estimate of nnz - csr_matrix_t augmented_CSR(n + m, n + m, new_nnz); - std::vector augmented_diagonal_indices(n + m, -1); - std::vector cone_csr_indices_host(total_block_nnz, -1); - std::vector cone_Q_values_host(total_block_nnz, f_t(0)); + const size_t n_sparse_entries = p > 0 ? cones().n_sparse_cone_entries : size_t{0}; + const i_t sparse_soc_kkt_nnz = static_cast(5 * n_sparse_entries + p); + i_t new_nnz = 2 * nnzA + n + m + p + nnzQ + dense_soc_kkt_nnz + sparse_soc_kkt_nnz; + csr_matrix_t augmented_CSR(factorization_size, factorization_size, new_nnz); + std::vector augmented_diagonal_indices(factorization_size, -1); + std::vector cone_csr_indices_host(dense_soc_kkt_nnz, -1); + std::vector cone_Q_values_host(dense_soc_kkt_nnz, f_t(0)); + std::vector sparse_hessian_diag_host(n_sparse_entries, -1); + std::vector sparse_hessian_Q_host(n_sparse_entries, f_t(0)); + std::vector sparse_exp_v_col_host(n_sparse_entries, -1); + std::vector sparse_exp_u_col_host(n_sparse_entries, -1); + std::vector sparse_exp_v_row_host(n_sparse_entries, -1); + std::vector sparse_exp_u_row_host(n_sparse_entries, -1); + std::vector sparse_expansion_D_host(p, -1); i_t q = 0; i_t off_diag_Qnz = 0; @@ -735,50 +793,101 @@ class iteration_data_t { const bool is_cone_row = is_cone_variable(i); if (is_cone_row) { - // Determine which cone this variable belongs to and its local row i_t local_idx = i - cone_start(); i_t k = local_to_cone[local_idx]; i_t local_r = static_cast(static_cast(local_idx) - cone_offsets_host[k]); i_t q_k = static_cast(cone_offsets_host[k + 1] - cone_offsets_host[k]); i_t cone_col_start = cone_start() + static_cast(cone_offsets_host[k]); - i_t block_base = static_cast(cone_block_offsets_host[k]) + local_r * q_k; + i_t cone_col_end = cone_col_start + q_k; - // Merge-join: Q entries (sorted) with dense cone block columns (contiguous) i_t qp = (nnzQ > 0) ? Q.col_start[i] : 0; i_t q_end = (nnzQ > 0) ? Q.col_start[i + 1] : 0; - // Q entries before cone block - for (; qp < q_end && Q.i[qp] < cone_col_start; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - - // Dense cone block, absorbing any Q entries that fall inside - for (i_t c = 0; c < q_k; c++) { - i_t col = cone_col_start + c; - f_t q_contrib = f_t(0); - f_t initial_val = (c == local_r) ? f_t(-dual_perturb) - : f_t(0); // diagonal entry of the cone block column + // Row of sparse SOC. + if (cone_sparse_idx[k] >= 0) { + const i_t sparse_idx = cone_sparse_idx[k]; + const size_t flat_idx = sparse_entry_offsets[sparse_idx] + local_r; + const i_t exp_v_col = n + m + 2 * sparse_idx; + const i_t exp_u_col = n + m + 2 * sparse_idx + 1; + + // Handle off-diagonal entries of sparse SOC before diagonal entry. + for (; qp < q_end && Q.i[qp] < cone_col_start; qp++) { + augmented_CSR.j[q] = Q.i[qp]; + augmented_CSR.x[q++] = -Q.x[qp]; + off_diag_Qnz++; + } - if (qp < q_end && Q.i[qp] == col) { + // Handle diagonal entry of sparse SOC. + f_t q_contrib = f_t(0); + if (qp < q_end && Q.i[qp] == i) { q_contrib = Q.x[qp]; qp++; } + sparse_hessian_diag_host[flat_idx] = q; + sparse_hessian_Q_host[flat_idx] = q_contrib; + augmented_diagonal_indices[i] = q; + augmented_CSR.j[q] = i; + augmented_CSR.x[q++] = -dual_perturb - q_contrib; + + // Handle expansion columns of sparse SOC. + sparse_exp_v_col_host[flat_idx] = q; + augmented_CSR.j[q] = exp_v_col; + augmented_CSR.x[q++] = f_t(0); + + sparse_exp_u_col_host[flat_idx] = q; + augmented_CSR.j[q] = exp_u_col; + augmented_CSR.x[q++] = f_t(0); + + // Handle off-diagonal entries of dense SOC before diagonal entry. + for (; qp < q_end && Q.i[qp] < cone_col_end; qp++) { + if (Q.i[qp] == i) { continue; } + augmented_CSR.j[q] = Q.i[qp]; + augmented_CSR.x[q++] = -Q.x[qp]; + off_diag_Qnz++; + } + for (; qp < q_end; qp++) { + augmented_CSR.j[q] = Q.i[qp]; + augmented_CSR.x[q++] = -Q.x[qp]; + off_diag_Qnz++; + } + } else { + // Row of dense SOC. + i_t dense_idx = 0; + for (i_t t = 0; t < k; ++t) { + if (cone_sparse_idx[t] < 0) { ++dense_idx; } + } + i_t block_base = + static_cast(dense_cone_block_offsets_host[dense_idx]) + local_r * q_k; - cone_csr_indices_host[block_base + c] = q; - cone_Q_values_host[block_base + c] = q_contrib; - if (col == i) { augmented_diagonal_indices[i] = q; } - augmented_CSR.j[q] = col; - augmented_CSR.x[q++] = initial_val - q_contrib; - } + for (; qp < q_end && Q.i[qp] < cone_col_start; qp++) { + augmented_CSR.j[q] = Q.i[qp]; + augmented_CSR.x[q++] = -Q.x[qp]; + off_diag_Qnz++; + } + + for (i_t c = 0; c < q_k; c++) { + i_t col = cone_col_start + c; + f_t q_contrib = f_t(0); + f_t initial_val = (c == local_r) ? f_t(-dual_perturb) : f_t(0); + + if (qp < q_end && Q.i[qp] == col) { + q_contrib = Q.x[qp]; + qp++; + } + + cone_csr_indices_host[block_base + c] = q; + cone_Q_values_host[block_base + c] = q_contrib; + if (col == i) { augmented_diagonal_indices[i] = q; } + augmented_CSR.j[q] = col; + augmented_CSR.x[q++] = initial_val - q_contrib; + } - // Q entries after cone block - for (; qp < q_end; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; + for (; qp < q_end; qp++) { + augmented_CSR.j[q] = Q.i[qp]; + augmented_CSR.x[q++] = -Q.x[qp]; + off_diag_Qnz++; + } } } else if (nnzQ == 0) { augmented_diagonal_indices[i] = q; @@ -816,27 +925,77 @@ class iteration_data_t { } for (i_t k = n; k < n + m; ++k) { - // A block, we can use AT in csc directly augmented_CSR.row_start[k] = q; const i_t l = k - n; const i_t col_beg = AT.col_start[l]; const i_t col_end = AT.col_start[l + 1]; - for (i_t p = col_beg; p < col_end; ++p) { - augmented_CSR.j[q] = AT.i[p]; - augmented_CSR.x[q++] = AT.x[p]; + for (i_t idx = col_beg; idx < col_end; ++idx) { + augmented_CSR.j[q] = AT.i[idx]; + augmented_CSR.x[q++] = AT.x[idx]; } augmented_diagonal_indices[k] = q; augmented_CSR.j[q] = k; augmented_CSR.x[q++] = primal_perturb; } - augmented_CSR.row_start[n + m] = q; - augmented_CSR.nz_max = q; + + // Handle expansion rows of sparse SOC. + for (i_t s = 0; s < n_sparse_soc; ++s) { + const i_t k = sparse_cone_ids_host[s]; + const i_t q_k = static_cast(cone_offsets_host[k + 1] - cone_offsets_host[k]); + const i_t cone_col_start = cone_start() + static_cast(cone_offsets_host[k]); + const size_t flat_base = sparse_entry_offsets[s]; + const i_t prow_v = n + m + 2 * s; + const i_t prow_u = n + m + 2 * s + 1; + + augmented_CSR.row_start[prow_v] = q; + for (i_t j = 0; j < q_k; ++j) { + sparse_exp_v_row_host[flat_base + j] = q; + augmented_CSR.j[q] = cone_col_start + j; + augmented_CSR.x[q++] = f_t(0); + } + sparse_expansion_D_host[2 * s] = q; + augmented_CSR.j[q] = prow_v; + augmented_CSR.x[q++] = f_t(0); + + augmented_CSR.row_start[prow_u] = q; + for (i_t j = 0; j < q_k; ++j) { + sparse_exp_u_row_host[flat_base + j] = q; + augmented_CSR.j[q] = cone_col_start + j; + augmented_CSR.x[q++] = f_t(0); + } + sparse_expansion_D_host[2 * s + 1] = q; + augmented_CSR.j[q] = prow_u; + augmented_CSR.x[q++] = f_t(0); + } + + augmented_CSR.row_start[factorization_size] = q; + augmented_CSR.nz_max = q; augmented_CSR.j.resize(q); augmented_CSR.x.resize(q); - i_t expected_nnz = 2 * nnzA + (n - m_c) + total_block_nnz + m + off_diag_Qnz; + i_t expected_nnz = + 2 * nnzA + (n - m_c) + dense_soc_kkt_nnz + m + off_diag_Qnz + sparse_soc_kkt_nnz; settings_.log.debug("augmented nz %d predicted %d\n", q, expected_nnz); - cuopt_assert(q == expected_nnz, "augmented nnz != predicted"); + cuopt_assert(q == expected_nnz, "augmented nz != predicted"); cuopt_assert(A.col_start[n] == AT.col_start[m], "A nz != AT nz"); + + if (n_sparse_entries > 0) { + for (size_t e = 0; e < n_sparse_entries; ++e) { + cuopt_assert(sparse_hessian_diag_host[e] >= 0, + "sparse Hessian diagonal CSR index unset"); + cuopt_assert(sparse_exp_v_col_host[e] >= 0, + "sparse expansion v-column CSR index unset"); + cuopt_assert(sparse_exp_u_col_host[e] >= 0, + "sparse expansion u-column CSR index unset"); + cuopt_assert(sparse_exp_v_row_host[e] >= 0, "sparse expansion v-row CSR index unset"); + cuopt_assert(sparse_exp_u_row_host[e] >= 0, "sparse expansion u-row CSR index unset"); + } + for (i_t s = 0; s < n_sparse_soc; ++s) { + cuopt_assert(sparse_expansion_D_host[2 * s] >= 0, + "sparse expansion v diagonal CSR index unset"); + cuopt_assert(sparse_expansion_D_host[2 * s + 1] >= 0, + "sparse expansion u diagonal CSR index unset"); + } + } } { @@ -850,16 +1009,72 @@ class iteration_data_t { handle_ptr->get_stream()); if (has_soc) { - d_cone_csr_indices_.resize(total_block_nnz, handle_ptr->get_stream()); + d_cone_csr_indices_.resize(dense_soc_kkt_nnz, handle_ptr->get_stream()); raft::copy(d_cone_csr_indices_.data(), cone_csr_indices_host.data(), - total_block_nnz, + dense_soc_kkt_nnz, handle_ptr->get_stream()); - d_cone_Q_values_.resize(total_block_nnz, handle_ptr->get_stream()); + d_cone_Q_values_.resize(dense_soc_kkt_nnz, handle_ptr->get_stream()); raft::copy(d_cone_Q_values_.data(), cone_Q_values_host.data(), - total_block_nnz, + dense_soc_kkt_nnz, + handle_ptr->get_stream()); + + d_dense_cone_block_offsets_.resize(dense_cone_block_offsets_host.size(), + handle_ptr->get_stream()); + raft::copy(d_dense_cone_block_offsets_.data(), + dense_cone_block_offsets_host.data(), + dense_cone_block_offsets_host.size(), handle_ptr->get_stream()); + d_dense_cone_ids_.resize(dense_cone_ids_host.size(), handle_ptr->get_stream()); + if (!dense_cone_ids_host.empty()) { + raft::copy(d_dense_cone_ids_.data(), + dense_cone_ids_host.data(), + dense_cone_ids_host.size(), + handle_ptr->get_stream()); + } + + const size_t n_sparse_entries = cones().n_sparse_cone_entries; + d_sparse_hessian_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_hessian_Q_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_exp_v_col_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_exp_u_col_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_exp_v_row_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_exp_u_row_.resize(n_sparse_entries, handle_ptr->get_stream()); + d_sparse_expansion_D_.resize(p, handle_ptr->get_stream()); + d_sparse_Hs_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); + if (n_sparse_entries > 0) { + raft::copy(d_sparse_hessian_diag_.data(), + sparse_hessian_diag_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + raft::copy(d_sparse_hessian_Q_.data(), + sparse_hessian_Q_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + raft::copy(d_sparse_exp_v_col_.data(), + sparse_exp_v_col_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + raft::copy(d_sparse_exp_u_col_.data(), + sparse_exp_u_col_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + raft::copy(d_sparse_exp_v_row_.data(), + sparse_exp_v_row_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + raft::copy(d_sparse_exp_u_row_.data(), + sparse_exp_u_row_host.data(), + n_sparse_entries, + handle_ptr->get_stream()); + } + if (p > 0) { + raft::copy(d_sparse_expansion_D_.data(), + sparse_expansion_D_host.data(), + sparse_expansion_D_host.size(), + handle_ptr->get_stream()); + } } handle_ptr->sync_stream(); @@ -868,11 +1083,11 @@ class iteration_data_t { csc_matrix_t augmented_transpose(1, 1, 1); augmented.transpose(augmented_transpose); settings_.log.printf("Aug nnz %d Aug^T nnz %d\n", - augmented.col_start[m + n], - augmented_transpose.col_start[m + n]); + augmented.col_start[factorization_size], + augmented_transpose.col_start[factorization_size]); augmented.check_matrix(); augmented_transpose.check_matrix(); - csc_matrix_t error(m + n, m + n, 1); + csc_matrix_t error(factorization_size, factorization_size, 1); add(augmented, augmented_transpose, 1.0, -1.0, error); settings_.log.printf("|| Aug - Aug^T ||_1 %e\n", error.norm1()); cuopt_assert(error.norm1() <= 1e-2, "|| Aug - Aug^T ||_1 > 1e-2"); @@ -909,14 +1124,41 @@ class iteration_data_t { RAFT_CHECK_CUDA(handle_ptr->get_stream()); if (has_soc) { - scatter_hessian_into_augmented(cones(), - device_augmented.x, - d_cone_csr_indices_, - d_cone_Q_values_, - handle_ptr->get_stream(), - dual_perturb); - RAFT_CHECK_CUDA(handle_ptr->get_stream()); + if (cones().has_sparse_cones()) { + launch_get_Hs_sparse(cones(), d_sparse_Hs_diag_, handle_ptr->get_stream()); + scatter_sparse_hessian_diag_into_augmented(device_augmented.x, + d_sparse_hessian_diag_, + d_sparse_Hs_diag_, + d_sparse_hessian_Q_, + handle_ptr->get_stream(), + dual_perturb); + update_sparse_expansion_in_augmented(device_augmented.x, + d_sparse_exp_v_col_, + d_sparse_exp_u_col_, + d_sparse_exp_v_row_, + d_sparse_exp_u_row_, + d_sparse_expansion_D_, + cones().sparse_v, + cones().sparse_u, + cones().eta, + cones().sparse_cone_ids, + handle_ptr->get_stream(), + dual_perturb); + RAFT_CHECK_CUDA(handle_ptr->get_stream()); + } + if (cones().n_dense_cones() > 0) { + scatter_dense_hessian_into_augmented(cones(), + device_augmented.x, + d_cone_csr_indices_, + d_cone_Q_values_, + d_dense_cone_block_offsets_, + d_dense_cone_ids_, + handle_ptr->get_stream(), + dual_perturb); + RAFT_CHECK_CUDA(handle_ptr->get_stream()); + } } + handle_ptr->sync_stream(); } } @@ -1400,6 +1642,15 @@ class iteration_data_t { } } + void restore_saved_iterate() + { + x = x_save; + y = y_save; + z = z_save; + v = v_save; + w = w_save; + } + void to_solution(const lp_problem_t& lp, i_t iterations, f_t objective, @@ -1809,18 +2060,28 @@ class iteration_data_t { const i_t m = A.m; const i_t n = A.n; const bool has_soc = has_cones(); + const i_t p = augmented_expansion_count(); + const i_t sys_size = augmented_system_size(n, m); + cuopt_assert(static_cast(x.size()) >= sys_size, "augmented_multiply: x too small"); + cuopt_assert(static_cast(y.size()) >= sys_size, "augmented_multiply: y too small"); rmm::device_uvector d_x1(n, handle_ptr->get_stream()); rmm::device_uvector d_x2(m, handle_ptr->get_stream()); rmm::device_uvector d_y1(n, handle_ptr->get_stream()); rmm::device_uvector d_y2(m, handle_ptr->get_stream()); + rmm::device_uvector d_y_exp(p, handle_ptr->get_stream()); + rmm::device_uvector d_y_exp_orig(p, handle_ptr->get_stream()); raft::copy(d_x1.data(), x.data(), n, handle_ptr->get_stream()); raft::copy(d_x2.data(), x.data() + n, m, handle_ptr->get_stream()); raft::copy(d_y1.data(), y.data(), n, handle_ptr->get_stream()); raft::copy(d_y2.data(), y.data() + n, m, handle_ptr->get_stream()); + if (p > 0) { + raft::copy(d_y_exp_orig.data(), y.data() + n + m, p, handle_ptr->get_stream()); + thrust::fill_n(rmm::exec_policy(stream_view_), d_y_exp.begin(), p, f_t(0)); + } - // y1 <- alpha ( -D * x_1 + A^T x_2) + beta * y1 + // y1 <- alpha ( -(Q + D + H) * x_1 + A^T x_2) + beta * y1 rmm::device_uvector d_r1(n, handle_ptr->get_stream()); thrust::fill_n(rmm::exec_policy(stream_view_), d_r1.begin(), n, f_t(0)); @@ -1835,14 +2096,31 @@ class iteration_data_t { stream_view_); RAFT_CHECK_CUDA(stream_view_); - // r1 <- D * x_1 + H x_1 (cone Hessian block H = S^T S; accumulate_cone_hessian_matvec) + // r1 <- D * x_1 + H x_1 on cone rows + // (dense cones: explicit dense H block; sparse cones: rank-2 expansion, which adds + // Hs_diag .* x_cone to r1 here and writes the expansion rows into d_y_exp) if (has_soc) { const i_t m_c = cone_entry_count(); - accumulate_cone_hessian_matvec(raft::device_span(d_x1.data() + cone_start(), m_c), - cones(), - raft::device_span(d_r1.data() + cone_start(), m_c), - stream_view_); - RAFT_CHECK_CUDA(stream_view_); + if (cones().has_sparse_cones()) { + launch_sparse_augmented_matvec( + raft::device_span(x.data(), x.size()), + raft::device_span(d_r1.data(), d_r1.size()), + raft::device_span(d_y_exp.data(), d_y_exp.size()), + cones(), + raft::device_span(d_sparse_Hs_diag_.data(), d_sparse_Hs_diag_.size()), + cone_start(), + n, + m, + handle_ptr->get_stream()); + RAFT_CHECK_CUDA(stream_view_); + } + if (cones().n_dense_cones() > 0) { + launch_dense_hessian_matvec(raft::device_span(d_x1.data() + cone_start(), m_c), + cones(), + raft::device_span(d_r1.data() + cone_start(), m_c), + stream_view_); + RAFT_CHECK_CUDA(stream_view_); + } } // r1 <- Q x1 + D x1 + H x1 (cone: same H as above) @@ -1861,8 +2139,13 @@ class iteration_data_t { // matrix_vector_multiply(A, alpha, x1, beta, y2); cusparse_view_.spmv(alpha, d_x1, beta, d_y2); + if (p > 0) { + axpy(alpha, d_y_exp.data(), beta, d_y_exp_orig.data(), d_y_exp.data(), p, stream_view_); + } + raft::copy(y.data(), d_y1.data(), n, stream_view_); raft::copy(y.data() + n, d_y2.data(), m, stream_view_); + if (p > 0) { raft::copy(y.data() + n + m, d_y_exp.data(), p, stream_view_); } handle_ptr->sync_stream(); } @@ -1944,6 +2227,16 @@ class iteration_data_t { rmm::device_uvector d_augmented_diagonal_indices_; rmm::device_uvector d_cone_csr_indices_; rmm::device_uvector d_cone_Q_values_; + rmm::device_uvector d_dense_cone_block_offsets_; + rmm::device_uvector d_dense_cone_ids_; + rmm::device_uvector d_sparse_hessian_diag_; + rmm::device_uvector d_sparse_hessian_Q_; + rmm::device_uvector d_sparse_exp_v_col_; + rmm::device_uvector d_sparse_exp_u_col_; + rmm::device_uvector d_sparse_exp_v_row_; + rmm::device_uvector d_sparse_exp_u_row_; + rmm::device_uvector d_sparse_expansion_D_; + rmm::device_uvector d_sparse_Hs_diag_; bool indefinite_Q; cusparse_view_t cusparse_Q_view_; @@ -2064,7 +2357,7 @@ void cholesky_debug_check(const iteration_data_t& data, // return; srand(42); - i_t vec_size = use_augmented ? lp.num_cols + lp.num_rows : lp.num_rows; + i_t vec_size = use_augmented ? data.augmented_system_size(lp.num_cols, lp.num_rows) : lp.num_rows; // 1. Create a random test vector dense_vector_t test_vec(vec_size); for (size_t i = 0; i < test_vec.size(); i++) { @@ -2224,14 +2517,16 @@ int barrier_solver_t::initial_point(iteration_data_t& data) data.inv_diag.pairwise_product(Fu, DinvFu); dense_vector_t q(lp.num_rows); if (use_augmented) { - dense_vector_t rhs(lp.num_cols + lp.num_rows); + const i_t aug_size = data.augmented_system_size(lp.num_cols, lp.num_rows); + dense_vector_t rhs(aug_size); + rhs.set_scalar(0.0); for (i_t k = 0; k < lp.num_cols; k++) { rhs[k] = -Fu[k]; } for (i_t k = 0; k < lp.num_rows; k++) { rhs[lp.num_cols + k] = rhs_x[k]; } - dense_vector_t soln(lp.num_cols + lp.num_rows); + dense_vector_t soln(aug_size); i_t solve_status = data.chol->solve(rhs, soln); struct op_t { op_t(iteration_data_t& data) : data_(data) {} @@ -2374,12 +2669,13 @@ int barrier_solver_t::initial_point(iteration_data_t& data) } } } else if (use_augmented) { - dense_vector_t dual_rhs(lp.num_cols + lp.num_rows); + const i_t aug_size = data.augmented_system_size(lp.num_cols, lp.num_rows); + dense_vector_t dual_rhs(aug_size); dual_rhs.set_scalar(0.0); for (i_t k = 0; k < lp.num_cols; k++) { dual_rhs[k] = data.c[k]; } - dense_vector_t py(lp.num_cols + lp.num_rows); + dense_vector_t py(aug_size); data.chol->solve(dual_rhs, py); for (i_t k = 0; k < lp.num_cols; k++) { data.z[k] = py[k]; @@ -2646,6 +2942,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t(data.d_x_.data() + cone_var_start, m_c); cones.z = raft::device_span(data.d_z_.data() + cone_var_start, m_c); launch_nt_scaling(cones, stream_view_); + if (cones.has_sparse_cones()) { launch_update_scaling_sparse(cones, stream_view_); } } max_residual = 0.0; @@ -2852,6 +3149,14 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t 0) { + const i_t expansion_start = lp.num_cols + lp.num_rows; + thrust::fill_n(rmm::exec_policy(stream_view_), + data.d_augmented_rhs_.begin() + expansion_start, + expansion_count, + f_t(0)); + } data.chol->solve(data.d_augmented_rhs_, data.d_augmented_soln_); struct op_t { op_t(iteration_data_t& data) : data_(data) {} @@ -3956,11 +4261,7 @@ lp_status_t barrier_solver_t::check_for_suboptimal_solution( data.relative_dual_residual_save < settings.barrier_relaxed_optimality_tol && data.relative_complementarity_residual_save < settings.barrier_relaxed_complementarity_tol) { settings.log.printf("Restoring previous solution\n"); - raft::copy(data.x.data(), data.d_x_.data(), data.d_x_.size(), stream_view_); - raft::copy(data.y.data(), data.d_y_.data(), data.d_y_.size(), stream_view_); - raft::copy(data.z.data(), data.d_z_.data(), data.d_z_.size(), stream_view_); - raft::copy(data.v.data(), data.d_v_.data(), data.d_v_.size(), stream_view_); - RAFT_CUDA_TRY(cudaStreamSynchronize(stream_view_)); + data.restore_saved_iterate(); data.to_solution(lp, iter, primal_objective_save, diff --git a/cpp/src/barrier/iterative_refinement.hpp b/cpp/src/barrier/iterative_refinement.hpp index 807d055787..bb87984185 100644 --- a/cpp/src/barrier/iterative_refinement.hpp +++ b/cpp/src/barrier/iterative_refinement.hpp @@ -24,7 +24,6 @@ #include #include -#include #include #include diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index 5a674fef70..1c3e527299 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -202,12 +202,26 @@ struct cone_data_t { rmm::device_uvector w; // [n_cone_entries], unit-J-norm NT direction rmm::device_uvector lambda; // [n_cone_entries], NT point lambda = W^{-T} z + // Sparse SOC rank-2 scaling (cones with dimension > soc_threshold). + i_t soc_threshold; + i_t n_sparse_cones; + size_t n_sparse_cone_entries; // sum of sparse cone dimensions + rmm::device_uvector sparse_cone_ids; // [n_sparse_cones], indices into cone arrays + rmm::device_uvector sparse_cone_dims; // [n_sparse_cones], dimension of each sparse cone + rmm::device_uvector d; // [n_sparse_cones], corner of rank-2 diagonal block + rmm::device_uvector sparse_v; // [n_sparse_cone_entries], rank-2 vector v per sparse cone + rmm::device_uvector sparse_u; // [n_sparse_cone_entries], rank-2 vector u per sparse cone + rmm::device_uvector + sparse_entry_offsets; // [n_sparse_cones + 1], packed prefix offsets of sparse cone entries + rmm::device_uvector cone_is_sparse; // [n_cones], 1 if cone uses sparse KKT expansion + cone_scratch_t scratch; cone_data_t(std::span cone_dimensions_host, raft::device_span x_in, raft::device_span z_in, - rmm::cuda_stream_view stream) + rmm::cuda_stream_view stream, + i_t soc_threshold_in = 5) : n_cones(cone_dimensions_host.size()), n_cone_entries( std::reduce(cone_dimensions_host.begin(), cone_dimensions_host.end(), size_t{0})), @@ -220,8 +234,58 @@ struct cone_data_t { eta(n_cones, stream), w(n_cone_entries, stream), lambda(n_cone_entries, stream), - scratch(n_cones, n_cone_entries, stream) + sparse_cone_ids(0, stream), + sparse_cone_dims(0, stream), + d(0, stream), + sparse_v(0, stream), + sparse_u(0, stream), + sparse_entry_offsets(0, stream), + cone_is_sparse(n_cones, stream), + scratch(n_cones, n_cone_entries, stream), + soc_threshold(soc_threshold_in), + n_sparse_cones(0), + n_sparse_cone_entries(0) { + thrust::fill(rmm::exec_policy(stream), cone_is_sparse.begin(), cone_is_sparse.end(), 0); + + std::vector sparse_cone_ids_host; + std::vector sparse_cone_dims_host; + sparse_cone_ids_host.reserve(n_cones); + sparse_cone_dims_host.reserve(n_cones); + for (i_t cone = 0; cone < n_cones; ++cone) { + if (cone_dimensions_host[cone] > soc_threshold) { + sparse_cone_ids_host.push_back(cone); + sparse_cone_dims_host.push_back(cone_dimensions_host[cone]); + } + } + n_sparse_cones = static_cast(sparse_cone_ids_host.size()); + n_sparse_cone_entries = + std::reduce(sparse_cone_dims_host.begin(), sparse_cone_dims_host.end(), size_t{0}); + + if (n_sparse_cones > 0) { + sparse_cone_ids.resize(n_sparse_cones, stream); + sparse_cone_dims.resize(n_sparse_cones, stream); + d.resize(n_sparse_cones, stream); + sparse_v.resize(n_sparse_cone_entries, stream); + sparse_u.resize(n_sparse_cone_entries, stream); + sparse_entry_offsets.resize(n_sparse_cones + 1, stream); + raft::copy(sparse_cone_ids.data(), sparse_cone_ids_host.data(), n_sparse_cones, stream); + raft::copy(sparse_cone_dims.data(), sparse_cone_dims_host.data(), n_sparse_cones, stream); + std::vector cone_is_sparse_host(n_cones, 0); + for (i_t cone : sparse_cone_ids_host) { + cone_is_sparse_host[cone] = 1; + } + raft::copy(cone_is_sparse.data(), cone_is_sparse_host.data(), n_cones, stream); + + // Packed prefix offsets of sparse cone entries. + std::vector sparse_entry_offsets_host(n_sparse_cones + 1, 0); + for (i_t s = 0; s < n_sparse_cones; ++s) { + sparse_entry_offsets_host[s + 1] = sparse_entry_offsets_host[s] + sparse_cone_dims_host[s]; + } + raft::copy( + sparse_entry_offsets.data(), sparse_entry_offsets_host.data(), n_sparse_cones + 1, stream); + } + raft::copy(cone_dimensions.data(), cone_dimensions_host.data(), n_cones, stream); cone_offsets.set_element_to_zero_async(0, stream); auto policy = rmm::exec_policy(stream); @@ -242,6 +306,16 @@ struct cone_data_t { element_cone_ids.begin()); segmented_sum.template prepare_workspace(stream); } + + // True when at least one cone is large enough (dim > soc_threshold) to use the + // rank-2 sparse KKT expansion instead of a dense q x q block. + bool has_sparse_cones() const { return n_sparse_cones > 0; } + + // Number of cones kept dense (dim <= soc_threshold). + i_t n_dense_cones() const { return n_cones - n_sparse_cones; } + + // Extra augmented-system columns/rows: two expansion variables (v, u) per sparse cone. + i_t expansion_var_count() const { return 2 * n_sparse_cones; } }; template @@ -482,6 +556,94 @@ void launch_nt_scaling(cone_data_t& cones, rmm::cuda_stream_view strea RAFT_CUDA_TRY(cudaPeekAtLastError()); } +// One block per sparse cone. Recompute the rank-2 factors (corner d and the +// vectors v, u, both scaled by eta^2) from the current NT direction w so that +// the implicit block reproduces the dense H = eta^2 (2 w w^T - J). +template +__global__ void update_scaling_sparse_kernel(raft::device_span w, + raft::device_span eta, + raft::device_span d, + raft::device_span sparse_v, + raft::device_span sparse_u, + raft::device_span cone_offsets, + raft::device_span sparse_cone_dims, + raft::device_span sparse_cone_ids, + raft::device_span sparse_entry_offsets, + i_t n_sparse_cones) +{ + const i_t sparse_idx = blockIdx.x; + if (sparse_idx >= n_sparse_cones) { return; } + + __shared__ f_t s_mem[4]; + + const i_t cone_idx = sparse_cone_ids[sparse_idx]; + const size_t cone_off = cone_offsets[cone_idx]; + const i_t cone_dim = sparse_cone_dims[sparse_idx]; + const i_t block_start = sparse_entry_offsets[sparse_idx]; + + if (threadIdx.x == 0) { + const f_t alpha = f_t(2) * w[cone_off]; + const f_t wsq = f_t(2) * w[cone_off] * w[cone_off] - f_t(1); + const f_t wsq_safe = f_t(0.5) * (wsq + sqrt(wsq * wsq + f_t(1))); + const f_t wsqinv = f_t(1) / wsq_safe; + const f_t di = f_t(0.5) * wsqinv; + d[sparse_idx] = di; + const f_t radicand = wsq_safe - di; + const f_t u0 = sqrt(max(radicand, f_t(0))); + const f_t u1 = (u0 > f_t(0)) ? alpha / u0 : f_t(0); + const f_t v0 = f_t(0); + const f_t denom = f_t(2) * wsq_safe - wsqinv; + const f_t v1_arg = (abs(denom) > f_t(1e-12)) ? f_t(2) * (f_t(2) + wsqinv) / denom : f_t(2); + const f_t v1 = sqrt(max(v1_arg, f_t(0))); + const f_t eta_sq = eta[cone_idx] * eta[cone_idx]; + s_mem[0] = eta_sq * u0; + s_mem[1] = eta_sq * u1; + s_mem[2] = eta_sq * v0; + s_mem[3] = eta_sq * v1; + } + __syncthreads(); + + const f_t scaled_u0 = s_mem[0]; + const f_t scaled_u1 = s_mem[1]; + const f_t scaled_v0 = s_mem[2]; + const f_t scaled_v1 = s_mem[3]; + + for (i_t j = threadIdx.x; j < cone_dim; j += blockDim.x) { + if (j == 0) { + sparse_u[block_start + j] = scaled_u0; + sparse_v[block_start + j] = scaled_v0; + } else { + sparse_u[block_start + j] = scaled_u1 * w[cone_off + j]; + sparse_v[block_start + j] = scaled_v1 * w[cone_off + j]; + } + } +} + +/** + * Refresh the rank-2 NT factors (d, v, u) of every sparse cone from the current + * scaling, so the implicit sparse block matches the dense Hessian for this + * iteration. Call after `launch_nt_scaling` has updated w and eta. + */ +template +void launch_update_scaling_sparse(cone_data_t& cones, rmm::cuda_stream_view stream) +{ + if (!cones.has_sparse_cones()) { return; } + + const i_t n_sparse = cones.n_sparse_cones; + update_scaling_sparse_kernel + <<>>(cuopt::make_span(cones.w), + cuopt::make_span(cones.eta), + cuopt::make_span(cones.d), + cuopt::make_span(cones.sparse_v), + cuopt::make_span(cones.sparse_u), + cuopt::make_span(cones.cone_offsets), + cuopt::make_span(cones.sparse_cone_dims), + cuopt::make_span(cones.sparse_cone_ids), + cuopt::make_span(cones.sparse_entry_offsets), + n_sparse); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + template __global__ void __launch_bounds__(soc_block_size) apply_w_inv_write_kernel(raft::device_span v, @@ -546,21 +708,24 @@ __global__ void __launch_bounds__(soc_block_size) template __global__ void __launch_bounds__(soc_block_size) - apply_hessian_write_kernel(raft::device_span v, - raft::device_span out, - raft::device_span w, - raft::device_span eta, - raft::device_span wv_dot, - raft::device_span cone_offsets, - raft::device_span element_cone_ids, - raft::device_span bias, - f_t output_scale, - f_t bias_scale) + apply_hessian_kernel(raft::device_span v, + raft::device_span out, + raft::device_span w, + raft::device_span eta, + raft::device_span wv_dot, + raft::device_span cone_offsets, + raft::device_span element_cone_ids, + raft::device_span cone_is_sparse, + bool dense_cones_only, + raft::device_span bias, + f_t output_scale, + f_t bias_scale) { const size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (idx >= out.size()) { return; } - const i_t cone = element_cone_ids[idx]; + const i_t cone = element_cone_ids[idx]; + if (dense_cones_only && !cone_is_sparse.empty() && cone_is_sparse[cone] != 0) { return; } const size_t cone_off = cone_offsets[cone]; const size_t local_idx = idx - cone_off; @@ -767,7 +932,8 @@ void apply_hessian(raft::device_span v, rmm::cuda_stream_view stream, f_t output_scale = 1, raft::device_span bias = {}, - f_t bias_scale = 0) + f_t bias_scale = 0, + bool dense_cones_only = false) { auto w = cuopt::make_span(cones.w); auto eta = cuopt::make_span(cones.eta); @@ -781,8 +947,19 @@ void apply_hessian(raft::device_span v, cones.segmented_sum(wv_terms, wv_dot, stream); const size_t grid_dim = raft::ceildiv(out.size(), soc_block_size); - apply_hessian_write_kernel<<>>( - v, out, w, eta, wv_dot, cone_offsets, element_cone_ids, bias, output_scale, bias_scale); + apply_hessian_kernel + <<>>(v, + out, + w, + eta, + wv_dot, + cone_offsets, + element_cone_ids, + cuopt::make_span(cones.cone_is_sparse), + dense_cones_only, + bias, + output_scale, + bias_scale); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -804,100 +981,392 @@ void recover_cone_dz_from_target(raft::device_span dx, } /** - * Accumulate the SOC cone-block matvec into an existing output vector. + * Accumulate the dense SOC cone-block matvec into an existing output vector: + * out += H x, where H = S^2, applied to dense cones only. + */ +template +void launch_dense_hessian_matvec(raft::device_span x, + cone_data_t& cones, + raft::device_span out, + rmm::cuda_stream_view stream) +{ + auto out_input = raft::device_span(out.data(), out.size()); + apply_hessian(x, out, cones, stream, 1, out_input, 1, true); +} + +// One block per sparse cone. Write the diagonal of the implicit cone Hessian: +// eta^2 on every entry, with the head entry scaled by the rank-2 corner d. +template +__global__ void get_Hs_sparse_kernel(raft::device_span Hs_diag, + raft::device_span eta, + raft::device_span d, + raft::device_span sparse_cone_ids, + raft::device_span sparse_entry_offsets, + raft::device_span sparse_cone_dims, + i_t n_sparse_cones) +{ + const i_t sparse_idx = blockIdx.x; + if (sparse_idx >= n_sparse_cones) { return; } + + const i_t block_start = sparse_entry_offsets[sparse_idx]; + const i_t block_dim = sparse_cone_dims[sparse_idx]; + const i_t cone_idx = sparse_cone_ids[sparse_idx]; + const f_t eta_sq = eta[cone_idx] * eta[cone_idx]; + + for (i_t j = threadIdx.x; j < block_dim; j += blockDim.x) { + Hs_diag[block_start + j] = eta_sq; + } + __syncthreads(); + + if (threadIdx.x == 0) { Hs_diag[block_start] *= d[sparse_idx]; } +} + +/** + * Compute the packed diagonal `Hs_diag` of the implicit sparse cone Hessians, + * one contiguous block of `sparse_cone_dims[s]` entries per sparse cone. This is + * the diagonal that gets scattered onto the cone rows of the augmented system. + */ +template +void launch_get_Hs_sparse(cone_data_t& cones, + rmm::device_uvector& Hs_diag, + rmm::cuda_stream_view stream) +{ + if (!cones.has_sparse_cones()) { return; } + + const i_t n_sparse = cones.n_sparse_cones; + cuopt_assert(Hs_diag.size() == cones.n_sparse_cone_entries, "Hs_diag size mismatch"); + + get_Hs_sparse_kernel + <<>>(cuopt::make_span(Hs_diag), + cuopt::make_span(cones.eta), + cuopt::make_span(cones.d), + cuopt::make_span(cones.sparse_cone_ids), + cuopt::make_span(cones.sparse_entry_offsets), + cuopt::make_span(cones.sparse_cone_dims), + n_sparse); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +// Scatter one sparse-cone diagonal entry into the augmented value buffer: +// augmented_x[csr_indices[e]] = -Hs_diag[e] - q_values[e] - dual_perturb, +// matching the sign convention of the dense scatter (negated KKT block plus +// any quadratic-objective contribution and diagonal regularization). +template +__global__ void __launch_bounds__(soc_block_size) + scatter_sparse_hessian_diag_kernel(raft::device_span augmented_x, + raft::device_span csr_indices, + raft::device_span Hs_diag, + raft::device_span q_values, + f_t dual_perturb) +{ + const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (e >= csr_indices.size()) { return; } + + augmented_x[csr_indices[e]] = -Hs_diag[e] - q_values[e] - dual_perturb; +} + +/** + * Write the sparse cone Hessian diagonals into the augmented system value buffer. + * + * `csr_indices[e]` gives the value-buffer position of the diagonal for packed + * cone entry `e`; each is set to `-Hs_diag[e] - q_values[e] - dual_perturb`. + */ +template +void scatter_sparse_hessian_diag_into_augmented(rmm::device_uvector& augmented_x, + const rmm::device_uvector& csr_indices, + const rmm::device_uvector& Hs_diag, + const rmm::device_uvector& q_values, + rmm::cuda_stream_view stream, + f_t dual_perturb) +{ + const size_t count = csr_indices.size(); + if (count == 0) { return; } + cuopt_assert(count == Hs_diag.size(), "sparse Hessian index/value size mismatch"); + cuopt_assert(count == q_values.size(), "sparse Hessian Q-value size mismatch"); + + const size_t grid = raft::ceildiv(count, soc_block_size); + scatter_sparse_hessian_diag_kernel + <<>>(cuopt::make_span(augmented_x), + cuopt::make_span(csr_indices), + cuopt::make_span(Hs_diag), + cuopt::make_span(q_values), + dual_perturb); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +/** + * Refresh the expansion (rank-2 coupling) entries of the augmented system. * - * Used by matrix-free products with the primal-reduced KKT block: - * out += H x, where H = S^2. + * For each sparse cone this scatters the rank-2 vectors v, u into their four + * coupling positions -- the v/u columns and the symmetric v/u rows linking the + * cone variables to the two expansion variables -- and writes the expansion + * diagonal +/-(eta^2 + dual_perturb). The Hessian diagonal itself is written + * separately by `scatter_sparse_hessian_diag_into_augmented`. */ template -void accumulate_cone_hessian_matvec(raft::device_span x, +void update_sparse_expansion_in_augmented(rmm::device_uvector& augmented_x, + const rmm::device_uvector& sparse_exp_v_col, + const rmm::device_uvector& sparse_exp_u_col, + const rmm::device_uvector& sparse_exp_v_row, + const rmm::device_uvector& sparse_exp_u_row, + const rmm::device_uvector& sparse_expansion_D, + const rmm::device_uvector& sparse_v, + const rmm::device_uvector& sparse_u, + const rmm::device_uvector& eta, + const rmm::device_uvector& sparse_cone_ids, + rmm::cuda_stream_view stream, + f_t dual_perturb) +{ + const size_t n_coupling = sparse_exp_v_col.size(); + if (n_coupling == 0) { return; } + + cuopt_assert(sparse_exp_u_col.size() == n_coupling, "sparse_exp_u_col size mismatch"); + cuopt_assert(sparse_exp_v_row.size() == n_coupling, "sparse_exp_v_row size mismatch"); + cuopt_assert(sparse_exp_u_row.size() == n_coupling, "sparse_exp_u_row size mismatch"); + cuopt_assert(sparse_v.size() == n_coupling, "sparse_v size mismatch"); + cuopt_assert(sparse_u.size() == n_coupling, "sparse_u size mismatch"); + + thrust::scatter(rmm::exec_policy(stream), + sparse_v.begin(), + sparse_v.end(), + sparse_exp_v_col.begin(), + augmented_x.begin()); + thrust::scatter(rmm::exec_policy(stream), + sparse_u.begin(), + sparse_u.end(), + sparse_exp_u_col.begin(), + augmented_x.begin()); + thrust::scatter(rmm::exec_policy(stream), + sparse_v.begin(), + sparse_v.end(), + sparse_exp_v_row.begin(), + augmented_x.begin()); + thrust::scatter(rmm::exec_policy(stream), + sparse_u.begin(), + sparse_u.end(), + sparse_exp_u_row.begin(), + augmented_x.begin()); + + const i_t n_expansion = static_cast(sparse_expansion_D.size()); + if (n_expansion > 0) { + auto eta_span = cuopt::make_span(eta); + auto sparse_cone_ids_span = cuopt::make_span(sparse_cone_ids); + // Expansion-row diagonal of sparse cone i/2: -(eta^2 + dual_perturb) for the + // v-row (even i), +(eta^2 + dual_perturb) for the u-row (odd i). + auto diag_values = thrust::make_transform_iterator( + thrust::counting_iterator(0), + [eta_span, sparse_cone_ids_span, dual_perturb] HD(i_t i) -> f_t { + const i_t sparse_idx = i / 2; + const f_t eta_val = eta_span[sparse_cone_ids_span[sparse_idx]]; + const f_t sign = (i % 2 == 0) ? f_t(-1) : f_t(1); + return sign * (eta_val * eta_val + dual_perturb); + }); + thrust::scatter(rmm::exec_policy(stream), + diag_values, + diag_values + n_expansion, + sparse_expansion_D.begin(), + augmented_x.begin()); + } +} + +/** + * Accumulate the sparse-SOC expanded KKT block into a matrix-free product. + * + * For each sparse cone this adds to the primal cone rows (before augmented_multiply negates r1): + * r1 += Hs_diag .* x_cone - v * x_exp_v - u * x_exp_u + * so y1 = -r1 matches the explicit CSR coupling (+v, +u on primal rows). + * and to the expansion rows: + * y_exp[2s] += -eta^2 * x_exp_v + v^T x_cone + * y_exp[2s+1] += +eta^2 * x_exp_u + u^T x_cone + * + * `v` and `u` are the rank-2 vectors in cones.sparse_v / cones.sparse_u. + */ +template +__global__ void sparse_augmented_matvec_kernel(raft::device_span x, + raft::device_span r1, + raft::device_span y_exp, + raft::device_span Hs_diag, + raft::device_span sparse_v, + raft::device_span sparse_u, + raft::device_span eta, + raft::device_span sparse_cone_ids, + raft::device_span sparse_cone_dims, + raft::device_span sparse_entry_offsets, + raft::device_span cone_offsets, + i_t cone_var_start, + i_t n_primal, + i_t m_constraints, + i_t n_sparse_cones) +{ + const i_t sparse_idx = blockIdx.x; + if (sparse_idx >= n_sparse_cones) { return; } + + const i_t cone = sparse_cone_ids[sparse_idx]; + const i_t q = sparse_cone_dims[sparse_idx]; + const i_t flat = sparse_entry_offsets[sparse_idx]; + const size_t off = cone_offsets[cone]; + const i_t base = cone_var_start + static_cast(off); + const i_t exp_v = n_primal + m_constraints + 2 * sparse_idx; + const i_t exp_u = exp_v + 1; + const f_t x_exp_v = x[exp_v]; + const f_t x_exp_u = x[exp_u]; + const f_t eta_sq = eta[cone] * eta[cone]; + + f_t partial_dot_v = f_t(0); + f_t partial_dot_u = f_t(0); + for (i_t j = threadIdx.x; j < q; j += blockDim.x) { + const f_t xj = x[base + j]; + const f_t vj = sparse_v[flat + j]; + const f_t uj = sparse_u[flat + j]; + partial_dot_v += vj * xj; + partial_dot_u += uj * xj; + r1[base + j] += Hs_diag[flat + j] * xj - vj * x_exp_v - uj * x_exp_u; + } + + __shared__ f_t s_dot_v[soc_block_size]; + __shared__ f_t s_dot_u[soc_block_size]; + s_dot_v[threadIdx.x] = partial_dot_v; + s_dot_u[threadIdx.x] = partial_dot_u; + __syncthreads(); + + for (i_t stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + s_dot_v[threadIdx.x] += s_dot_v[threadIdx.x + stride]; + s_dot_u[threadIdx.x] += s_dot_u[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + y_exp[2 * sparse_idx] += -eta_sq * x_exp_v + s_dot_v[0]; + y_exp[2 * sparse_idx + 1] += eta_sq * x_exp_u + s_dot_u[0]; + } +} + +/** + * Matrix-free product of the implicit sparse cone blocks (see + * `sparse_augmented_matvec_kernel`). Accumulates the cone-row contribution into + * `r1` and the expansion-row contribution into `y_exp`, one block per sparse cone. + */ +template +void launch_sparse_augmented_matvec(raft::device_span x, + raft::device_span r1, + raft::device_span y_exp, cone_data_t& cones, - raft::device_span out, + raft::device_span Hs_diag, + i_t cone_var_start, + i_t n_primal, + i_t m_constraints, rmm::cuda_stream_view stream) { - auto out_input = raft::device_span(out.data(), out.size()); - apply_hessian(x, out, cones, stream, 1, out_input, 1); + if (!cones.has_sparse_cones()) { return; } + + const i_t n_sparse = cones.n_sparse_cones; + cuopt_assert(Hs_diag.size() == cones.n_sparse_cone_entries, "Hs_diag size mismatch"); + cuopt_assert(y_exp.size() == static_cast(cones.expansion_var_count()), + "expansion output size mismatch"); + + sparse_augmented_matvec_kernel + <<>>(x, + r1, + y_exp, + Hs_diag, + cuopt::make_span(cones.sparse_v), + cuopt::make_span(cones.sparse_u), + cuopt::make_span(cones.eta), + cuopt::make_span(cones.sparse_cone_ids), + cuopt::make_span(cones.sparse_cone_dims), + cuopt::make_span(cones.sparse_entry_offsets), + cuopt::make_span(cones.cone_offsets), + cone_var_start, + n_primal, + m_constraints, + n_sparse); + RAFT_CUDA_TRY(cudaPeekAtLastError()); } +// Scatter one nonzero of a dense cone's q x q NT Hessian block into the +// augmented value buffer. Only cones with dim <= soc_threshold take this path. template __global__ void __launch_bounds__(soc_block_size) - scatter_hessian_into_augmented_kernel(raft::device_span augmented_x, - raft::device_span csr_indices, - raft::device_span q_values, - raft::device_span w, - raft::device_span eta, - raft::device_span cone_offsets, - raft::device_span block_offsets, - i_t n_cones, - f_t dual_perturb_value) + scatter_dense_hessian_into_augmented_kernel(raft::device_span augmented_x, + raft::device_span csr_indices, + raft::device_span q_values, + raft::device_span w, + raft::device_span eta, + raft::device_span cone_offsets, + raft::device_span dense_block_offsets, + raft::device_span dense_cone_ids, + i_t n_dense_cones, + f_t dual_perturb_value) { const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (e >= csr_indices.size()) { return; } i_t lo = 0; - i_t hi = n_cones; + i_t hi = n_dense_cones; while (lo < hi) { const i_t mid = lo + (hi - lo) / 2; - if (block_offsets[mid + 1] <= e) { + if (dense_block_offsets[mid + 1] <= e) { lo = mid + 1; } else { hi = mid; } } - const i_t cone = lo; + const i_t dense_idx = lo; + const i_t cone = dense_cone_ids[dense_idx]; const size_t off = cone_offsets[cone]; const size_t q = cone_offsets[cone + 1] - off; - const size_t blk_off = block_offsets[cone]; + const size_t blk_off = dense_block_offsets[dense_idx]; const size_t local = e - blk_off; const size_t r = local / q; const size_t c = local % q; - const f_t eta_sq = eta[cone] * eta[cone]; - const f_t w0 = w[off]; - const f_t u_r = (r == 0) ? w0 : w[off + r]; - const f_t u_c = (c == 0) ? w0 : w[off + c]; - f_t val = f_t{2} * u_r * eta_sq * u_c; - const f_t diag_correction = (r == 0) ? -eta_sq : eta_sq; - if (r == c) { val += diag_correction; } + const f_t eta_sq = eta[cone] * eta[cone]; + const f_t w0 = w[off]; + const f_t u_r = (r == 0) ? w0 : w[off + r]; + const f_t u_c = (c == 0) ? w0 : w[off + c]; + const f_t val = f_t{2} * u_r * eta_sq * u_c; - augmented_x[csr_indices[e]] = -val - q_values[e]; + f_t entry = -val - q_values[e]; + if (r == c) { + const f_t diag_correction = (r == 0) ? -eta_sq : eta_sq; + entry -= diag_correction + dual_perturb_value; + } + augmented_x[csr_indices[e]] = entry; } +/** + * Write the full q x q NT Hessian blocks of the dense cones (dim <= soc_threshold) + * into the augmented system value buffer at the precomputed CSR positions. + */ template -void scatter_hessian_into_augmented(const cone_data_t& cones, - rmm::device_uvector& augmented_x, - const rmm::device_uvector& csr_indices, - const rmm::device_uvector& q_values, - rmm::cuda_stream_view stream, - f_t dual_perturb_value) +void scatter_dense_hessian_into_augmented(const cone_data_t& cones, + rmm::device_uvector& augmented_x, + const rmm::device_uvector& csr_indices, + const rmm::device_uvector& q_values, + const rmm::device_uvector& dense_block_offsets, + const rmm::device_uvector& dense_cone_ids, + rmm::cuda_stream_view stream, + f_t dual_perturb_value) { const size_t count = csr_indices.size(); if (count == 0) { return; } - cuopt_assert(count == q_values.size(), "cone CSR index and Q-value arrays must match"); - - // TODO: This offset calculation should be done in the barrier layer, - // because it is already done in the barrier layer for the augmented system, see - // cone_block_offsets_host. - rmm::device_uvector block_offsets(cones.n_cones + 1, stream); - block_offsets.set_element_to_zero_async(0, stream); - - auto block_sizes = thrust::make_transform_iterator( - cones.cone_dimensions.begin(), [] HD(i_t q) -> size_t { return static_cast(q) * q; }); - thrust::inclusive_scan( - rmm::exec_policy(stream), block_sizes, block_sizes + cones.n_cones, block_offsets.begin() + 1); + cuopt_assert(count == q_values.size(), "dense cone CSR index and Q-value arrays must match"); - // TODO: use dual_perturb_value for regularization + const i_t n_dense = cones.n_dense_cones(); const size_t grid = raft::ceildiv(count, soc_block_size); - scatter_hessian_into_augmented_kernel + scatter_dense_hessian_into_augmented_kernel <<>>(cuopt::make_span(augmented_x), cuopt::make_span(csr_indices), cuopt::make_span(q_values), cuopt::make_span(cones.w), cuopt::make_span(cones.eta), cuopt::make_span(cones.cone_offsets), - cuopt::make_span(block_offsets), - cones.n_cones, + cuopt::make_span(dense_block_offsets), + cuopt::make_span(dense_cone_ids), + n_dense, dual_perturb_value); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/barrier/second_order_cone_reduction.cuh b/cpp/src/barrier/second_order_cone_reduction.cuh index ada9fcfb6a..6ab380102a 100644 --- a/cpp/src/barrier/second_order_cone_reduction.cuh +++ b/cpp/src/barrier/second_order_cone_reduction.cuh @@ -10,8 +10,8 @@ #include #include +#include #include -#include #include #include @@ -24,7 +24,6 @@ #include #include -#include #include #include @@ -46,14 +45,22 @@ __global__ void __launch_bounds__(warps_per_cta* raft::WarpSize) OutputIt output, value_t init); +template +__global__ void __launch_bounds__(block_dim) + block_per_cone_reduce_kernel(InputIt input, + raft::device_span medium_cone_ids, + raft::device_span cone_offsets, + OutputIt output, + value_t init); + /** * Segmented-sum dispatcher for packed second-order cone vectors. * * Cone dimensions are fixed for a solve, so the constructor partitions cone * ids once by reduction strategy. Each call then reuses those partitions: - * small cones use one warp per cone, medium cones use CUB DeviceSegmentedReduce, + * small cones use one warp per cone, medium cones use one block per cone, * and large cones use CUB DeviceReduce one cone at a time. The object owns the - * CUB workspace for those medium/large paths. Call `prepare_workspace` once + * CUB workspace for the large-cone path. Call `prepare_workspace` once * before using a CUB-backed path. */ template @@ -69,7 +76,7 @@ struct segmented_sum_t { std::vector large_cone_ids; std::vector large_cone_dimensions; - // Maximum CUB temporary storage needed by prepared medium/large reductions. + // Maximum CUB temporary storage needed by prepared large reductions. std::size_t cub_workspace_bytes = 0; rmm::device_buffer cub_workspace; @@ -80,24 +87,6 @@ struct segmented_sum_t { auto input = thrust::make_constant_iterator(value_t{}); auto output = thrust::make_discard_iterator(); - if (!medium_cone_ids.is_empty()) { - const auto medium_begin_offsets = - thrust::make_permutation_iterator(cone_offsets.data(), medium_cone_ids.begin()); - const auto medium_end_offsets = - thrust::make_permutation_iterator(cone_offsets.data() + 1, medium_cone_ids.begin()); - - std::size_t temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceSegmentedReduce::Sum(nullptr, - temp_storage_bytes, - input, - output, - medium_cone_ids.size(), - medium_begin_offsets, - medium_end_offsets, - stream.value())); - cub_workspace_bytes = std::max(cub_workspace_bytes, temp_storage_bytes); - } - for (std::size_t i = 0; i < large_cone_ids.size(); ++i) { std::size_t temp_storage_bytes = 0; RAFT_CUDA_TRY(cub::DeviceReduce::Sum(nullptr, @@ -109,7 +98,7 @@ struct segmented_sum_t { cub_workspace_bytes = std::max(cub_workspace_bytes, temp_storage_bytes); } - if (cub_workspace.size() < cub_workspace_bytes) { + if (cub_workspace_bytes > 0 && cub_workspace.size() < cub_workspace_bytes) { cub_workspace.resize(cub_workspace_bytes, stream); } } @@ -138,31 +127,17 @@ struct segmented_sum_t { } if (!medium_cone_ids.is_empty()) { - cuopt_assert(cub_workspace_bytes > 0 && cub_workspace.size() >= cub_workspace_bytes, - "segmented_sum_t::prepare_workspace must be called before reducing medium or " - "large cones"); - - const auto medium_output = thrust::make_permutation_iterator(output, medium_cone_ids.begin()); - const auto medium_begin_offsets = - thrust::make_permutation_iterator(cone_offsets.data(), medium_cone_ids.begin()); - const auto medium_end_offsets = - thrust::make_permutation_iterator(cone_offsets.data() + 1, medium_cone_ids.begin()); - - std::size_t temp_storage_bytes = cub_workspace_bytes; - RAFT_CUDA_TRY(cub::DeviceSegmentedReduce::Sum(cub_workspace.data(), - temp_storage_bytes, - input, - medium_output, - medium_cone_ids.size(), - medium_begin_offsets, - medium_end_offsets, - stream.value())); + constexpr int medium_block_dim = 256; + const auto n_medium = medium_cone_ids.size(); + block_per_cone_reduce_kernel + <<>>( + input, cuopt::make_span(medium_cone_ids), cone_offsets, output, init); + RAFT_CUDA_TRY(cudaPeekAtLastError()); } if (!large_cone_ids.empty()) { cuopt_assert(cub_workspace_bytes > 0 && cub_workspace.size() >= cub_workspace_bytes, - "segmented_sum_t::prepare_workspace must be called before reducing medium or " - "large cones"); + "segmented_sum_t::prepare_workspace must be called before reducing large cones"); for (std::size_t i = 0; i < large_cone_ids.size(); ++i) { std::size_t temp_storage_bytes = cub_workspace_bytes; @@ -258,4 +233,55 @@ __global__ void __launch_bounds__(warps_per_cta* raft::WarpSize) if (lane_id == 0) { output[cone] = sum; } } +template +__global__ void __launch_bounds__(block_dim) + block_per_cone_reduce_kernel(InputIt input, + raft::device_span medium_cone_ids, + raft::device_span cone_offsets, + OutputIt output, + value_t init) +{ + static_assert(block_dim > 0); + static_assert(block_dim <= 1024); + + constexpr int items_per_thread = 4; + + using block_reduce_t = cub::BlockReduce; + __shared__ typename block_reduce_t::TempStorage temp_storage; + + const auto slot = blockIdx.x; + if (slot >= medium_cone_ids.size()) { return; } + + const auto cone = medium_cone_ids[slot]; + const auto off = cone_offsets[cone]; + const auto dim = cone_offsets[cone + 1] - off; + + value_t acc[items_per_thread]; +#pragma unroll + for (int k = 0; k < items_per_thread; ++k) { + acc[k] = value_t{0}; + } + + // Thread t owns elements t, t+block_dim, ..., t+(items_per_thread-1)*block_dim + // within each tile of `block_dim * items_per_thread` consecutive entries, so + // every warp load stays coalesced. + const std::size_t tile = static_cast(block_dim) * items_per_thread; + for (std::size_t tile_start = 0; tile_start < dim; tile_start += tile) { +#pragma unroll + for (int k = 0; k < items_per_thread; ++k) { + const std::size_t idx = tile_start + threadIdx.x + static_cast(k) * block_dim; + if (idx < dim) { acc[k] = acc[k] + input[off + idx]; } + } + } + + value_t sum = init; +#pragma unroll + for (int k = 0; k < items_per_thread; ++k) { + sum = sum + acc[k]; + } + + sum = block_reduce_t(temp_storage).Sum(sum); + if (threadIdx.x == 0) { output[cone] = sum; } +} + } // namespace cuopt::linear_programming::dual_simplex diff --git a/cpp/src/dual_simplex/simplex_solver_settings.hpp b/cpp/src/dual_simplex/simplex_solver_settings.hpp index 286ac7364f..46a551b9f1 100644 --- a/cpp/src/dual_simplex/simplex_solver_settings.hpp +++ b/cpp/src/dual_simplex/simplex_solver_settings.hpp @@ -69,6 +69,7 @@ struct simplex_solver_settings_t { eliminate_dense_columns(true), barrier_iterative_refinement(true), barrier_step_scale(0.9), + barrier_soc_threshold(5), num_gpus(1), folding(-1), augmented(0), @@ -157,6 +158,7 @@ struct simplex_solver_settings_t { 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 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 diff --git a/cpp/src/math_optimization/solver_settings.cu b/cpp/src/math_optimization/solver_settings.cu index 4a000c6e05..235fc57f30 100644 --- a/cpp/src/math_optimization/solver_settings.cu +++ b/cpp/src/math_optimization/solver_settings.cu @@ -132,6 +132,7 @@ solver_settings_t::solver_settings_t() : pdlp_settings(), mip_settings {CUOPT_DUALIZE, &pdlp_settings.dualize, -1, 1, -1}, {CUOPT_ORDERING, &pdlp_settings.ordering, -1, 1, -1}, {CUOPT_BARRIER_DUAL_INITIAL_POINT, &pdlp_settings.barrier_dual_initial_point, -1, 1, -1}, + {CUOPT_BARRIER_SOC_THRESHOLD, &pdlp_settings.barrier_soc_threshold, 0, 1000000, 5, "SOC dimension above which rank-2 sparse KKT expansion is used"}, {CUOPT_MIP_CUT_PASSES, &mip_settings.max_cut_passes, -1, std::numeric_limits::max(), 10}, {CUOPT_MIP_MIXED_INTEGER_ROUNDING_CUTS, &mip_settings.mir_cuts, -1, 1, -1}, {CUOPT_MIP_MIXED_INTEGER_GOMORY_CUTS, &mip_settings.mixed_integer_gomory_cuts, -1, 1, -1}, diff --git a/cpp/src/pdlp/solve.cu b/cpp/src/pdlp/solve.cu index 5d29e97fa0..3fc1fc7dda 100644 --- a/cpp/src/pdlp/solve.cu +++ b/cpp/src/pdlp/solve.cu @@ -511,6 +511,7 @@ run_barrier(dual_simplex::user_problem_t& user_problem, 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_soc_threshold = settings.barrier_soc_threshold; barrier_settings.barrier_step_scale = settings.barrier_step_scale; barrier_settings.cudss_deterministic = settings.cudss_deterministic; barrier_settings.barrier_relaxed_feasibility_tol = settings.tolerances.relative_primal_tolerance; diff --git a/cpp/tests/socp/CMakeLists.txt b/cpp/tests/socp/CMakeLists.txt index f380984225..7ea1a732d5 100644 --- a/cpp/tests/socp/CMakeLists.txt +++ b/cpp/tests/socp/CMakeLists.txt @@ -5,6 +5,7 @@ ConfigureTest(SOCP_TEST ${CMAKE_CURRENT_SOURCE_DIR}/second_order_cone_kernels.cu + ${CMAKE_CURRENT_SOURCE_DIR}/sparse_augmented_kkt_test.cu ${CMAKE_CURRENT_SOURCE_DIR}/solve_barrier_socp.cu ${CMAKE_CURRENT_SOURCE_DIR}/general_quadratic_test.cu LABELS numopt) diff --git a/cpp/tests/socp/second_order_cone_kernels.cu b/cpp/tests/socp/second_order_cone_kernels.cu index 49126ee335..43c162b4aa 100644 --- a/cpp/tests/socp/second_order_cone_kernels.cu +++ b/cpp/tests/socp/second_order_cone_kernels.cu @@ -202,6 +202,171 @@ TEST(second_order_cone_kernels, nt_scaling_matches_host_reference) } } +TEST(second_order_cone_kernels, nt_scaling_many_small_one_medium_cone) +{ + auto stream = rmm::cuda_stream_default; + + // chainsing-like topology: many small cones plus one dim-1000 medium cone (warp_cone_dim=64). + std::vector cone_dimensions; + cone_dimensions.reserve(21); + for (int i = 0; i < 20; ++i) { + cone_dimensions.push_back(3); + } + cone_dimensions.push_back(1000); + + std::size_t n_cone_entries = 0; + for (const auto dim : cone_dimensions) { + n_cone_entries += static_cast(dim); + } + + std::vector x_host(n_cone_entries); + std::vector z_host(n_cone_entries); + std::size_t offset = 0; + for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { + const auto dim = cone_dimensions[cone]; + x_host[offset] = 100.0 + static_cast(cone); + z_host[offset] = 80.0 + static_cast(cone); + for (int local_idx = 1; local_idx < dim; ++local_idx) { + x_host[offset + local_idx] = 0.001 * static_cast((local_idx % 5) + 1); + z_host[offset + local_idx] = 0.0015 * static_cast((local_idx % 7) + 1); + } + offset += static_cast(dim); + } + + auto x = cuopt::device_copy(x_host, stream); + auto z = cuopt::device_copy(z_host, stream); + cone_data_t cones(cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream); + + EXPECT_EQ(cuopt::host_copy(cones.segmented_sum.medium_cone_ids, stream), (std::vector{20})); + EXPECT_EQ(cones.n_sparse_cones, 1); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + auto eta_host = cuopt::host_copy(cones.eta, stream); + auto w_host = cuopt::host_copy(cones.w, stream); + + std::vector expected_eta(cone_dimensions.size()); + std::vector expected_w(n_cone_entries); + + offset = 0; + for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { + const auto dim = cone_dimensions[cone]; + + double x_tail_sq = 0.0; + double z_tail_sq = 0.0; + for (int local_idx = 1; local_idx < dim; ++local_idx) { + const auto idx = offset + local_idx; + x_tail_sq += x_host[idx] * x_host[idx]; + z_tail_sq += z_host[idx] * z_host[idx]; + } + + const auto x_tail_norm = std::sqrt(x_tail_sq); + const auto z_tail_norm = std::sqrt(z_tail_sq); + const auto x_det = (x_host[offset] - x_tail_norm) * (x_host[offset] + x_tail_norm); + const auto z_det = (z_host[offset] - z_tail_norm) * (z_host[offset] + z_tail_norm); + ASSERT_GT(x_det, 0.0) << "cone " << cone; + ASSERT_GT(z_det, 0.0) << "cone " << cone; + + const auto x_scale = std::sqrt(x_det); + const auto z_scale = std::sqrt(z_det); + + expected_eta[cone] = std::sqrt(z_scale / x_scale); + + double normalized_xz_dot = 0.0; + for (int local_idx = 0; local_idx < dim; ++local_idx) { + const auto idx = offset + local_idx; + normalized_xz_dot += x_host[idx] * z_host[idx] / (x_scale * z_scale); + } + const auto w_det = 2.0 + 2.0 * normalized_xz_dot; + ASSERT_GT(w_det, 0.0) << "cone " << cone; + const auto w_scale = std::sqrt(w_det); + + expected_w[offset] = 0.0; + for (int local_idx = 1; local_idx < dim; ++local_idx) { + const auto idx = offset + local_idx; + expected_w[idx] = (z_host[idx] / z_scale - x_host[idx] / x_scale) / w_scale; + } + + double normalized_tail_sq = 0.0; + for (int local_idx = 1; local_idx < dim; ++local_idx) { + const auto idx = offset + local_idx; + normalized_tail_sq += expected_w[idx] * expected_w[idx]; + } + expected_w[offset] = std::sqrt(1.0 + normalized_tail_sq); + + offset += static_cast(dim); + } + + for (std::size_t i = 0; i < expected_eta.size(); ++i) { + EXPECT_NEAR(eta_host[i], expected_eta[i], 1e-10) << "cone " << i; + } + + for (std::size_t i = 0; i < expected_w.size(); ++i) { + EXPECT_NEAR(w_host[i], expected_w[i], 1e-10) << "entry " << i; + } +} + +TEST(second_order_cone_kernels, cone_step_length_many_small_one_sparse_medium_cone) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions; + for (int i = 0; i < 20; ++i) { + cone_dimensions.push_back(3); + } + cone_dimensions.push_back(1000); + + std::size_t n_cone_entries = 0; + for (const auto dim : cone_dimensions) { + n_cone_entries += static_cast(dim); + } + + std::vector x_host(n_cone_entries); + std::vector z_host(n_cone_entries); + std::vector dx_host(n_cone_entries); + std::vector dz_host(n_cone_entries); + + std::size_t offset = 0; + for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { + const auto dim = cone_dimensions[cone]; + x_host[offset] = 12.0 + static_cast(cone); + z_host[offset] = 14.0 + static_cast(cone); + dx_host[offset] = 0.05; + dz_host[offset] = -0.04; + for (int local_idx = 1; local_idx < dim; ++local_idx) { + const auto idx = offset + local_idx; + x_host[idx] = 0.001 * static_cast((local_idx % 5) + 1); + z_host[idx] = 0.0015 * static_cast((local_idx % 7) + 1); + dx_host[idx] = 0.02 * static_cast((local_idx % 5) - 2); + dz_host[idx] = -0.015 * static_cast((local_idx % 7) - 3); + } + offset += static_cast(dim); + } + + auto x = cuopt::device_copy(x_host, stream); + auto z = cuopt::device_copy(z_host, stream); + auto dx = cuopt::device_copy(dx_host, stream); + auto dz = cuopt::device_copy(dz_host, stream); + + cone_data_t cones(cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + const auto [step_primal, step_dual] = + compute_cone_step_length(cones, + raft::device_span(dx.data(), dx.size()), + raft::device_span(dz.data(), dz.size()), + 1.0, + stream); + + EXPECT_GT(step_primal, 0.0); + EXPECT_GT(step_dual, 0.0); + EXPECT_LE(step_primal, 1.0); + EXPECT_LE(step_dual, 1.0); +} + TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) { auto stream = rmm::cuda_stream_default; @@ -354,7 +519,7 @@ TEST(second_order_cone_kernels, scaling_operators_match_host_reference) raft::device_span(cone_target.data(), cone_target.size()), cuopt::make_span(recovered_dz), stream); - accumulate_cone_hessian_matvec(v_span, cones, cuopt::make_span(accum), stream); + launch_dense_hessian_matvec(v_span, cones, cuopt::make_span(accum), stream); apply_w(v_span, cuopt::make_span(w_tmp), cones, stream); apply_w(raft::device_span(w_tmp.data(), w_tmp.size()), cuopt::make_span(w_squared_v), @@ -374,13 +539,15 @@ TEST(second_order_cone_kernels, scaling_operators_match_host_reference) std::vector expected_w_inv(n_cone_entries); std::vector expected_h(n_cone_entries); std::vector expected_h_unscaled(n_cone_entries); + std::vector expected_accum(n_cone_entries); offset = 0; for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { - const auto dim = cone_dimensions[cone]; - const auto w0 = w_host[offset]; - const auto v0 = v_host[offset]; - const auto eta = eta_host[cone]; + const auto dim = cone_dimensions[cone]; + const auto w0 = w_host[offset]; + const auto v0 = v_host[offset]; + const auto eta = eta_host[cone]; + const bool cone_dense = dim <= cones.soc_threshold; double tail_dot = 0.0; for (int local_idx = 1; local_idx < dim; ++local_idx) { @@ -394,6 +561,8 @@ TEST(second_order_cone_kernels, scaling_operators_match_host_reference) const auto rho = w0 * v0 + tail_dot; expected_h_unscaled[offset] = (eta * eta) * (2.0 * w0 * rho - v0); expected_h[offset] = expected_h_unscaled[offset]; + expected_accum[offset] = + accum_initial_host[offset] + (cone_dense ? expected_h_unscaled[offset] : 0.0); for (int local_idx = 1; local_idx < dim; ++local_idx) { const auto idx = offset + local_idx; @@ -402,6 +571,7 @@ TEST(second_order_cone_kernels, scaling_operators_match_host_reference) expected_w_inv[idx] = (v_host[idx] + (-v0 + tail_dot / (1.0 + w0)) * w_host[idx]) / eta; expected_h_unscaled[idx] = (eta * eta) * (2.0 * w_host[idx] * rho + v_host[idx]); expected_h[idx] = expected_h_unscaled[idx]; + expected_accum[idx] = accum_initial_host[idx] + (cone_dense ? expected_h_unscaled[idx] : 0.0); } offset += static_cast(dim); @@ -414,8 +584,7 @@ TEST(second_order_cone_kernels, scaling_operators_match_host_reference) EXPECT_NEAR(h_out_host[i], w_squared_v_host[i], 1e-9) << "W^2 v vs H v entry " << i; EXPECT_NEAR(recovered_dz_host[i], cone_target_host[i] - expected_h_unscaled[i], 1e-9) << "recovered dz entry " << i; - EXPECT_NEAR(accum_host[i], accum_initial_host[i] + expected_h_unscaled[i], 1e-9) - << "accumulated H entry " << i; + EXPECT_NEAR(accum_host[i], expected_accum[i], 1e-9) << "accumulated H entry " << i; } } @@ -581,4 +750,176 @@ TEST(second_order_cone_kernels, combined_cone_rhs_matches_host_reference) } } +TEST(second_order_cone_kernels, sparse_cone_classification) +{ + auto stream = rmm::cuda_stream_default; + + // threshold=5: cones of dim 3,2 are dense; dim 6,32769 are sparse + std::vector cone_dimensions{3, 6, 2, 32769}; + rmm::device_uvector x(32780, stream); + rmm::device_uvector z(32780, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/5); + + EXPECT_EQ(cones.soc_threshold, 5); + EXPECT_EQ(cones.n_sparse_cones, 2); + EXPECT_EQ(cones.n_dense_cones(), 2); + EXPECT_EQ(cones.expansion_var_count(), 4); + EXPECT_EQ(cones.n_sparse_cone_entries, std::size_t{32775}); + EXPECT_EQ(cones.d.size(), 2u); + EXPECT_EQ(cones.sparse_v.size(), 32775u); + EXPECT_EQ(cones.sparse_u.size(), 32775u); + EXPECT_EQ(cuopt::host_copy(cones.sparse_cone_ids, stream), (std::vector{1, 3})); + EXPECT_EQ(cuopt::host_copy(cones.sparse_cone_dims, stream), (std::vector{6, 32769})); +} + +namespace { + +struct sparse_scaling_head_t { + double d; + double u0; + double u1; + double v0; + double v1; +}; + +// Must match update_scaling_sparse_kernel in second_order_cone_kernels.cuh. +sparse_scaling_head_t sparse_scaling_head_reference(double w0) +{ + const double alpha = 2.0 * w0; + const double wsq = 2.0 * w0 * w0 - 1.0; + const double wsq_safe = 0.5 * (wsq + std::sqrt(wsq * wsq + 1.0)); + const double wsqinv = 1.0 / wsq_safe; + const double d = 0.5 * wsqinv; + const double radicand = wsq_safe - d; + const double u0 = std::sqrt(std::max(radicand, 0.0)); + const double u1 = (u0 > 0.0) ? alpha / u0 : 0.0; + const double v0 = 0.0; + const double denom = 2.0 * wsq_safe - wsqinv; + const double v1_arg = (std::abs(denom) > 1e-12) ? 2.0 * (2.0 + wsqinv) / denom : 2.0; + const double v1 = std::sqrt(std::max(v1_arg, 0.0)); + return {d, u0, u1, v0, v1}; +} + +} // namespace + +TEST(second_order_cone_kernels, update_scaling_sparse_matches_reference) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{3, 6}; + rmm::device_uvector x(9, stream); + rmm::device_uvector z(9, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + ASSERT_EQ(cones.n_sparse_cones, 1); + + // Place interior point in the sparse cone (second cone, dim 6). + std::vector x_host(9, 0.0); + std::vector z_host(9, 0.0); + x_host[3] = 2.0; + z_host[3] = 1.5; + for (int j = 1; j < 6; ++j) { + x_host[3 + j] = 0.1 * j; + z_host[3 + j] = 0.08 * j; + } + x_host[0] = 1.5; + z_host[0] = 1.2; + for (int j = 1; j < 3; ++j) { + x_host[j] = 0.05; + z_host[j] = 0.04; + } + + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + auto d_host = cuopt::host_copy(cones.d, stream); + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto w_host = cuopt::host_copy(cones.w, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + + const int sparse_cone = 1; + const int cone_dim = 6; + const int cone_off = 3; + const double eta = eta_host[sparse_cone]; + + const auto head = sparse_scaling_head_reference(w_host[cone_off]); + const double scale = eta * eta; + + EXPECT_NEAR(d_host[0], head.d, 1e-10); + EXPECT_NEAR(v_host[0], scale * head.v0, 1e-10); + EXPECT_NEAR(u_host[0], scale * head.u0, 1e-10); + for (int j = 1; j < cone_dim; ++j) { + EXPECT_NEAR(v_host[j], scale * head.v1 * w_host[cone_off + j], 1e-10) << "v tail " << j; + EXPECT_NEAR(u_host[j], scale * head.u1 * w_host[cone_off + j], 1e-10) << "u tail " << j; + } +} + +TEST(second_order_cone_kernels, update_scaling_sparse_two_cones) +{ + auto stream = rmm::cuda_stream_default; + + // sparse cones: dim 6 (offset 3) and dim 5 (offset 9) + std::vector cone_dimensions{3, 6, 5}; + rmm::device_uvector x(14, stream); + rmm::device_uvector z(14, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + ASSERT_EQ(cones.n_sparse_cones, 2); + ASSERT_EQ(cones.n_sparse_cone_entries, std::size_t{11}); + + std::vector x_host(14, 0.0); + std::vector z_host(14, 0.0); + for (int cone = 0; cone < 3; ++cone) { + const int off = (cone == 0) ? 0 : (cone == 1 ? 3 : 9); + const int dim = cone_dimensions[cone]; + x_host[off] = 2.0 + 0.1 * cone; + z_host[off] = 1.5 + 0.05 * cone; + for (int j = 1; j < dim; ++j) { + x_host[off + j] = 0.1 * j; + z_host[off + j] = 0.08 * j; + } + } + + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto w_host = cuopt::host_copy(cones.w, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + + const auto check_cone = [&](int sparse_idx, int cone_idx, int cone_off, int cone_dim) { + const double eta = eta_host[cone_idx]; + const auto head = sparse_scaling_head_reference(w_host[cone_off]); + const double scale = eta * eta; + + const int block_start = (sparse_idx == 0) ? 0 : 6; + + EXPECT_NEAR(v_host[block_start], scale * head.v0, 1e-10) << "sparse " << sparse_idx; + EXPECT_NEAR(u_host[block_start], scale * head.u0, 1e-10) << "sparse " << sparse_idx; + for (int j = 1; j < cone_dim; ++j) { + EXPECT_NEAR(v_host[block_start + j], scale * head.v1 * w_host[cone_off + j], 1e-10) + << "v tail sparse " << sparse_idx << " j " << j; + EXPECT_NEAR(u_host[block_start + j], scale * head.u1 * w_host[cone_off + j], 1e-10) + << "u tail sparse " << sparse_idx << " j " << j; + } + }; + + check_cone(0, 1, 3, 6); + check_cone(1, 2, 9, 5); +} + } // namespace cuopt::linear_programming::dual_simplex::test diff --git a/cpp/tests/socp/solve_barrier_socp.cu b/cpp/tests/socp/solve_barrier_socp.cu index f960c3dc3a..1878e43898 100644 --- a/cpp/tests/socp/solve_barrier_socp.cu +++ b/cpp/tests/socp/solve_barrier_socp.cu @@ -824,4 +824,214 @@ TEST(barrier, qp_with_soc_block) EXPECT_NEAR(std::abs(solution.x[3]), 0.0, 1e-4); } +TEST(barrier, sparse_soc_expansion_solves_large_single_cone) +{ + // minimize x_0 + // subject to x_1 = 1 + // (x_0, ..., x_5) in Q^6 + // + // Optimal: x* = (1, 1, 0, 0, 0, 0), obj* = 1 + + raft::handle_t handle{}; + init_handler(&handle); + + using namespace cuopt::linear_programming::dual_simplex; + user_problem_t user_problem(&handle); + + constexpr int m = 1; + constexpr int n = 6; + constexpr int nz = 1; + + user_problem.num_rows = m; + user_problem.num_cols = n; + user_problem.objective.assign(n, 0.0); + user_problem.objective[0] = 1.0; + + user_problem.A.m = m; + user_problem.A.n = n; + user_problem.A.nz_max = nz; + user_problem.A.reallocate(nz); + user_problem.A.col_start = {0, 0, 1, 1, 1, 1, 1}; + user_problem.A.i[0] = 0; + user_problem.A.x[0] = 1.0; + + user_problem.rhs = {1.0}; + user_problem.row_sense = {'E'}; + user_problem.lower.assign(n, 0.0); + user_problem.upper.assign(n, inf); + + user_problem.num_range_rows = 0; + user_problem.problem_name = "sparse_soc_single_large_cone"; + user_problem.cone_var_start = 0; + user_problem.second_order_cone_dims = {6}; + user_problem.var_types.assign(n, variable_type_t::CONTINUOUS); + + simplex_solver_settings_t settings; + settings.barrier = true; + settings.barrier_presolve = true; + settings.dualize = 0; + settings.barrier_soc_threshold = 4; + settings.barrier_iterative_refinement = true; + + lp_solution_t solution(m, n); + auto status = solve_linear_program_with_barrier(user_problem, settings, solution); + EXPECT_EQ(status, lp_status_t::OPTIMAL); + EXPECT_NEAR(solution.objective, 1.0, 1e-4); + EXPECT_NEAR(solution.x[0], 1.0, 1e-4); + EXPECT_NEAR(solution.x[1], 1.0, 1e-4); + for (int j = 2; j < n; ++j) { + EXPECT_NEAR(std::abs(solution.x[j]), 0.0, 1e-4) << "index " << j; + } +} + +TEST(barrier, mixed_dense_and_sparse_soc_blocks) +{ + // Variables ordered as [l1, l2 | dense Q^3 | sparse Q^6]. + // + // minimize t_dense + t_sparse + // subject to l1 - u_dense = 0 + // l2 - u_sparse = 0 + // l1 + l2 >= 3 + // l1 - l2 = 1 + // + // Same optimum as mixed_linear_and_two_soc_blocks_with_inequality: obj* = 3. + raft::handle_t handle{}; + init_handler(&handle); + + using namespace cuopt::linear_programming::dual_simplex; + user_problem_t user_problem(&handle); + + constexpr int m = 4; + constexpr int n = 11; + constexpr int nz = 8; + + user_problem.num_rows = m; + user_problem.num_cols = n; + user_problem.objective.assign(n, 0.0); + user_problem.objective[2] = 1.0; // t_dense (head of dense Q^3 block) + user_problem.objective[5] = 1.0; // t_sparse (head of sparse Q^6 block) + + user_problem.A.m = m; + user_problem.A.n = n; + user_problem.A.nz_max = nz; + user_problem.A.reallocate(nz); + // Columns: l1, l2, t_d, u_d, v_d, t_s, u_s, v_s, w_s, y_s, z_s + user_problem.A.col_start = {0, 3, 6, 6, 7, 7, 7, 8, 8, 8, 8, 8}; + user_problem.A.i[0] = 0; + user_problem.A.x[0] = 1.0; + user_problem.A.i[1] = 2; + user_problem.A.x[1] = 1.0; + user_problem.A.i[2] = 3; + user_problem.A.x[2] = 1.0; + user_problem.A.i[3] = 1; + user_problem.A.x[3] = 1.0; + user_problem.A.i[4] = 2; + user_problem.A.x[4] = 1.0; + user_problem.A.i[5] = 3; + user_problem.A.x[5] = -1.0; + user_problem.A.i[6] = 0; + user_problem.A.x[6] = -1.0; + user_problem.A.i[7] = 1; + user_problem.A.x[7] = -1.0; + + user_problem.rhs = {0.0, 0.0, 3.0, 1.0}; + user_problem.row_sense = {'E', 'E', 'G', 'E'}; + user_problem.lower.assign(n, 0.0); + user_problem.upper.assign(n, inf); + + user_problem.num_range_rows = 0; + user_problem.problem_name = "mixed_dense_and_sparse_soc_blocks"; + user_problem.cone_var_start = 2; + user_problem.second_order_cone_dims = {3, 6}; + user_problem.var_types.assign(n, variable_type_t::CONTINUOUS); + + simplex_solver_settings_t settings; + settings.barrier = true; + settings.barrier_presolve = true; + settings.dualize = 0; + settings.scale_columns = true; + settings.barrier_soc_threshold = 4; + settings.barrier_iterative_refinement = true; + + lp_solution_t solution(m, n); + auto status = solve_linear_program_with_barrier(user_problem, settings, solution); + + EXPECT_EQ(status, lp_status_t::OPTIMAL); + EXPECT_NEAR(solution.objective, 3.0, 1e-4); + EXPECT_NEAR(solution.x[0], 2.0, 1e-4); + EXPECT_NEAR(solution.x[1], 1.0, 1e-4); + EXPECT_NEAR(solution.x[2], 2.0, 1e-4); + EXPECT_NEAR(solution.x[3], 2.0, 1e-4); + EXPECT_NEAR(std::abs(solution.x[4]), 0.0, 1e-4); + EXPECT_NEAR(solution.x[5], 1.0, 1e-4); + EXPECT_NEAR(solution.x[6], 1.0, 1e-4); + for (int j = 7; j < n; ++j) { + EXPECT_NEAR(std::abs(solution.x[j]), 0.0, 1e-4) << "index " << j; + } +} + +TEST(barrier, sparse_soc_expansion_solves_dim_500_cone) +{ + // minimize x_0 + // subject to x_1 = 1 + // (x_0, ..., x_499) in Q^500 + // + // Optimal: x* = (1, 1, 0, ..., 0), obj* = 1 + // Sparse cone (dim 500 > default threshold 5) with barrier IR enabled. + raft::handle_t handle{}; + init_handler(&handle); + + using namespace cuopt::linear_programming::dual_simplex; + user_problem_t user_problem(&handle); + + constexpr int m = 1; + constexpr int n = 500; + constexpr int nz = 1; + + user_problem.num_rows = m; + user_problem.num_cols = n; + user_problem.objective.assign(n, 0.0); + user_problem.objective[0] = 1.0; + + user_problem.A.m = m; + user_problem.A.n = n; + user_problem.A.nz_max = nz; + user_problem.A.reallocate(nz); + // x_1 = 1: nonzero in column 1. + user_problem.A.col_start.assign(n + 1, 1); + user_problem.A.col_start[0] = 0; + user_problem.A.col_start[1] = 0; + user_problem.A.i[0] = 0; + user_problem.A.x[0] = 1.0; + + user_problem.rhs = {1.0}; + user_problem.row_sense = {'E'}; + user_problem.lower.assign(n, 0.0); + user_problem.upper.assign(n, inf); + + user_problem.num_range_rows = 0; + user_problem.problem_name = "sparse_soc_dim_500"; + user_problem.cone_var_start = 0; + user_problem.second_order_cone_dims = {500}; + user_problem.var_types.assign(n, variable_type_t::CONTINUOUS); + + simplex_solver_settings_t settings; + settings.barrier = true; + settings.barrier_presolve = true; + settings.dualize = 0; + settings.barrier_soc_threshold = 5; + settings.barrier_iterative_refinement = true; + + lp_solution_t solution(m, n); + auto status = solve_linear_program_with_barrier(user_problem, settings, solution); + + EXPECT_EQ(status, lp_status_t::OPTIMAL); + EXPECT_NEAR(solution.objective, 1.0, 1e-3); + EXPECT_NEAR(solution.x[0], 1.0, 1e-3); + EXPECT_NEAR(solution.x[1], 1.0, 1e-3); + for (int j = 2; j < n; ++j) { + EXPECT_NEAR(std::abs(solution.x[j]), 0.0, 1e-3) << "index " << j; + } +} + } // namespace cuopt::linear_programming::dual_simplex::test diff --git a/cpp/tests/socp/sparse_augmented_kkt_test.cu b/cpp/tests/socp/sparse_augmented_kkt_test.cu new file mode 100644 index 0000000000..400a66a536 --- /dev/null +++ b/cpp/tests/socp/sparse_augmented_kkt_test.cu @@ -0,0 +1,353 @@ +/* clang-format off */ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +/* clang-format on */ + +#include + +#include + +#include + +#include + +#include +#include + +namespace cuopt::linear_programming::dual_simplex::test { + +TEST(sparse_augmented_kkt, cone_counts_and_expansion_size) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{3, 6, 5}; + rmm::device_uvector x(14, stream); + rmm::device_uvector z(14, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + EXPECT_EQ(cones.n_sparse_cones, 2); + EXPECT_EQ(cones.n_dense_cones(), 1); + EXPECT_EQ(cones.expansion_var_count(), 4); + EXPECT_EQ(cones.n_sparse_cone_entries, 11u); +} + +TEST(sparse_augmented_kkt, launch_get_Hs_sparse) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{3, 6}; + rmm::device_uvector x(9, stream); + rmm::device_uvector z(9, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + ASSERT_EQ(cones.n_sparse_cones, 1); + + std::vector x_host(9, 0.0); + std::vector z_host(9, 0.0); + x_host[3] = 2.0; + z_host[3] = 1.5; + for (int j = 1; j < 6; ++j) { + x_host[3 + j] = 0.1 * j; + z_host[3 + j] = 0.08 * j; + } + x_host[0] = 1.5; + z_host[0] = 1.2; + for (int j = 1; j < 3; ++j) { + x_host[j] = 0.05; + z_host[j] = 0.04; + } + + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + rmm::device_uvector Hs_diag(cones.n_sparse_cone_entries, stream); + launch_get_Hs_sparse(cones, Hs_diag, stream); + + auto d_host = cuopt::host_copy(cones.d, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + auto hs_host = cuopt::host_copy(Hs_diag, stream); + + const int sparse_cone = 1; + const double eta_sq = eta_host[sparse_cone] * eta_host[sparse_cone]; + + EXPECT_NEAR(hs_host[0], eta_sq * d_host[0], 1e-10); + for (int j = 1; j < 6; ++j) { + EXPECT_NEAR(hs_host[j], eta_sq, 1e-10) << "tail index " << j; + } +} + +TEST(sparse_augmented_kkt, scatter_and_update_sparse_expansion) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{6}; + rmm::device_uvector x(6, stream); + rmm::device_uvector z(6, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + ASSERT_EQ(cones.n_sparse_cones, 1); + ASSERT_EQ(cones.expansion_var_count(), 2); + + std::vector x_host{2.0, 0.2, 0.3, 0.4, 0.5, 0.6}; + std::vector z_host{1.5, 0.1, 0.15, 0.2, 0.25, 0.3}; + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + const int nnz = 32; + + rmm::device_uvector augmented_x(nnz, stream); + thrust::fill(rmm::exec_policy(stream), augmented_x.begin(), augmented_x.end(), 0.0); + + // One coupling/Hessian CSR entry per sparse-cone element, matching production where every + // *_col/*_row index array has size n_sparse_cone_entries (== cones.sparse_v.size()). + const int q_dim = 6; + ASSERT_EQ(cones.n_sparse_cone_entries, static_cast(q_dim)); + + std::vector sparse_hessian_diag(q_dim); + std::vector sparse_hessian_Q(q_dim, 0.0); + std::vector sparse_exp_v_col(q_dim); + std::vector sparse_exp_u_col(q_dim); + std::vector sparse_exp_v_row(q_dim); + std::vector sparse_exp_u_row(q_dim); + std::vector sparse_expansion_D{30, 31}; + for (int j = 0; j < q_dim; ++j) { + sparse_hessian_diag[j] = j; // CSR slots 0..5 + sparse_exp_v_col[j] = 6 + j; // 6..11 + sparse_exp_u_col[j] = 12 + j; // 12..17 + sparse_exp_v_row[j] = 18 + j; // 18..23 + sparse_exp_u_row[j] = 24 + j; // 24..29 + } + + auto d_sparse_hessian_diag = cuopt::device_copy(sparse_hessian_diag, stream); + auto d_sparse_hessian_Q = cuopt::device_copy(sparse_hessian_Q, stream); + auto d_sparse_exp_v_col = cuopt::device_copy(sparse_exp_v_col, stream); + auto d_sparse_exp_u_col = cuopt::device_copy(sparse_exp_u_col, stream); + auto d_sparse_exp_v_row = cuopt::device_copy(sparse_exp_v_row, stream); + auto d_sparse_exp_u_row = cuopt::device_copy(sparse_exp_u_row, stream); + auto d_sparse_expansion_D = cuopt::device_copy(sparse_expansion_D, stream); + rmm::device_uvector d_sparse_Hs_diag(cones.n_sparse_cone_entries, stream); + + launch_get_Hs_sparse(cones, d_sparse_Hs_diag, stream); + scatter_sparse_hessian_diag_into_augmented( + augmented_x, d_sparse_hessian_diag, d_sparse_Hs_diag, d_sparse_hessian_Q, stream, 0.0); + update_sparse_expansion_in_augmented(augmented_x, + d_sparse_exp_v_col, + d_sparse_exp_u_col, + d_sparse_exp_v_row, + d_sparse_exp_u_row, + d_sparse_expansion_D, + cones.sparse_v, + cones.sparse_u, + cones.eta, + cones.sparse_cone_ids, + stream, + 0.0); + + auto aug_host = cuopt::host_copy(augmented_x, stream); + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto hs_host = cuopt::host_copy(d_sparse_Hs_diag, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + + for (int j = 0; j < q_dim; ++j) { + EXPECT_NEAR(aug_host[j], -hs_host[j], 1e-10) << "hessian diag " << j; + EXPECT_NEAR(aug_host[6 + j], v_host[j], 1e-10) << "v col " << j; + EXPECT_NEAR(aug_host[12 + j], u_host[j], 1e-10) << "u col " << j; + EXPECT_NEAR(aug_host[18 + j], v_host[j], 1e-10) << "v row " << j; + EXPECT_NEAR(aug_host[24 + j], u_host[j], 1e-10) << "u row " << j; + } + + const double eta_sq = eta_host[0] * eta_host[0]; + EXPECT_NEAR(aug_host[30], -eta_sq, 1e-10); + EXPECT_NEAR(aug_host[31], eta_sq, 1e-10); +} + +TEST(sparse_augmented_kkt, sparse_augmented_matvec) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{6}; + rmm::device_uvector x(6, stream); + rmm::device_uvector z(6, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + ASSERT_EQ(cones.n_sparse_cones, 1); + ASSERT_EQ(cones.expansion_var_count(), 2); + + std::vector x_host{2.0, 0.2, 0.3, 0.4, 0.5, 0.6}; + std::vector z_host{1.5, 0.1, 0.15, 0.2, 0.25, 0.3}; + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + const int n_primal = 6; + const int m_rows = 1; + const int p = cones.expansion_var_count(); + const int sys_size = n_primal + m_rows + p; + + std::vector x_vec(sys_size, 0.0); + x_vec[0] = 1.1; + x_vec[1] = 0.3; + x_vec[2] = 0.4; + x_vec[3] = 0.2; + x_vec[4] = 0.5; + x_vec[5] = 0.6; + x_vec[n_primal + m_rows] = 0.25; // expansion v + x_vec[n_primal + m_rows + 1] = -0.15; // expansion u + + rmm::device_uvector d_x(sys_size, stream); + rmm::device_uvector d_r1(n_primal, stream); + rmm::device_uvector d_y_exp(p, stream); + rmm::device_uvector d_hs(cones.n_sparse_cone_entries, stream); + + raft::copy(d_x.data(), x_vec.data(), sys_size, stream); + thrust::fill(rmm::exec_policy(stream), d_r1.begin(), d_r1.end(), 0.0); + thrust::fill(rmm::exec_policy(stream), d_y_exp.begin(), d_y_exp.end(), 0.0); + + launch_get_Hs_sparse(cones, d_hs, stream); + launch_sparse_augmented_matvec(raft::device_span(d_x.data(), d_x.size()), + raft::device_span(d_r1.data(), d_r1.size()), + raft::device_span(d_y_exp.data(), d_y_exp.size()), + cones, + raft::device_span(d_hs.data(), d_hs.size()), + /*cone_var_start=*/0, + n_primal, + m_rows, + stream); + + auto r1_host = cuopt::host_copy(d_r1, stream); + auto yexp_host = cuopt::host_copy(d_y_exp, stream); + auto hs_host = cuopt::host_copy(d_hs, stream); + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + + const double eta_sq = eta_host[0] * eta_host[0]; + double dot_v = 0.0; + double dot_u = 0.0; + for (int j = 0; j < 6; ++j) { + dot_v += v_host[j] * x_vec[j]; + dot_u += u_host[j] * x_vec[j]; + const double expected = hs_host[j] * x_vec[j] - v_host[j] * x_vec[n_primal + m_rows] - + u_host[j] * x_vec[n_primal + m_rows + 1]; + EXPECT_NEAR(r1_host[j], expected, 1e-10) << "primal row " << j; + } + + EXPECT_NEAR(yexp_host[0], -eta_sq * x_vec[n_primal + m_rows] + dot_v, 1e-10); + EXPECT_NEAR(yexp_host[1], eta_sq * x_vec[n_primal + m_rows + 1] + dot_u, 1e-10); +} + +TEST(sparse_augmented_kkt, update_scaling_sparse_dim_1000) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{1000}; + rmm::device_uvector x(1000, stream); + rmm::device_uvector z(1000, stream); + + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/5); + + ASSERT_EQ(cones.n_sparse_cones, 1); + ASSERT_EQ(cones.n_sparse_cone_entries, 1000u); + ASSERT_EQ(cones.expansion_var_count(), 2); + + std::vector x_host(1000); + std::vector z_host(1000); + x_host[0] = 100.0; + z_host[0] = 80.0; + for (int j = 1; j < 1000; ++j) { + x_host[j] = 0.001 * ((j % 5) + 1); + z_host[j] = 0.0015 * ((j % 7) + 1); + } + + raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); + raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); + + launch_nt_scaling(cones, stream); + launch_update_scaling_sparse(cones, stream); + + auto d_host = cuopt::host_copy(cones.d, stream); + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + + EXPECT_GT(d_host[0], 0.0); + EXPECT_GT(eta_host[0], 0.0); + for (int j = 0; j < 1000; ++j) { + EXPECT_TRUE(std::isfinite(v_host[j])) << "v entry " << j; + EXPECT_TRUE(std::isfinite(u_host[j])) << "u entry " << j; + } + + rmm::device_uvector Hs_diag(cones.n_sparse_cone_entries, stream); + launch_get_Hs_sparse(cones, Hs_diag, stream); + + const int n_primal = 1000; + const int m_rows = 1; + const int p = cones.expansion_var_count(); + const int sys_size = n_primal + m_rows + p; + + std::vector x_vec(sys_size, 0.0); + for (int j = 0; j < 1000; ++j) { + x_vec[j] = x_host[j]; + } + x_vec[n_primal + m_rows] = 0.25; + x_vec[n_primal + m_rows + 1] = -0.15; + + rmm::device_uvector d_x(sys_size, stream); + rmm::device_uvector d_r1(n_primal, stream); + rmm::device_uvector d_y_exp(p, stream); + + raft::copy(d_x.data(), x_vec.data(), sys_size, stream); + thrust::fill(rmm::exec_policy(stream), d_r1.begin(), d_r1.end(), 0.0); + thrust::fill(rmm::exec_policy(stream), d_y_exp.begin(), d_y_exp.end(), 0.0); + + launch_sparse_augmented_matvec(raft::device_span(d_x.data(), d_x.size()), + raft::device_span(d_r1.data(), d_r1.size()), + raft::device_span(d_y_exp.data(), d_y_exp.size()), + cones, + raft::device_span(Hs_diag.data(), Hs_diag.size()), + /*cone_var_start=*/0, + n_primal, + m_rows, + stream); + + auto r1_host = cuopt::host_copy(d_r1, stream); + auto yexp_host = cuopt::host_copy(d_y_exp, stream); + auto hs_host = cuopt::host_copy(Hs_diag, stream); + + const double eta_sq = eta_host[0] * eta_host[0]; + const double x_exp_v = x_vec[n_primal + m_rows]; + const double x_exp_u = x_vec[n_primal + m_rows + 1]; + double dot_v = 0.0; + double dot_u = 0.0; + for (int j = 0; j < 1000; ++j) { + dot_v += v_host[j] * x_vec[j]; + dot_u += u_host[j] * x_vec[j]; + const double expected = hs_host[j] * x_vec[j] - v_host[j] * x_exp_v - u_host[j] * x_exp_u; + EXPECT_NEAR(r1_host[j], expected, 1e-9) << "primal row " << j; + } + + EXPECT_NEAR(yexp_host[0], -eta_sq * x_exp_v + dot_v, 1e-9); + EXPECT_NEAR(yexp_host[1], eta_sq * x_exp_u + dot_u, 1e-9); +} + +} // namespace cuopt::linear_programming::dual_simplex::test From 4ed64ef5f011f44a998e260809bfb6ab8bdabc9a Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Mon, 29 Jun 2026 07:04:28 -0700 Subject: [PATCH 02/14] Stricter IR tolerance for SOCP with sparse socs Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 12 ++++++++---- cpp/src/barrier/iterative_refinement.hpp | 25 +++++++++++++++--------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 0a5377ea01..45b1feab4a 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -683,9 +683,11 @@ class iteration_data_t { i_t cone_end() const { return cone_start() + cone_entry_count(); } + bool has_sparse_cones() const { return has_cones() && cones_->n_sparse_cones > 0; } + i_t augmented_expansion_count() const { - return has_cones() && cones().has_sparse_cones() ? cones().expansion_var_count() : i_t(0); + return has_sparse_cones() ? cones().expansion_var_count() : i_t(0); } i_t augmented_system_size(i_t n, i_t m) const { return n + m + augmented_expansion_count(); } @@ -2563,7 +2565,8 @@ int barrier_solver_t::initial_point(iteration_data_t& data) } op(data); if (settings.barrier_iterative_refinement) { - iterative_refinement(op, rhs, soln); + const f_t ir_tol = data.has_sparse_cones() ? f_t(1e-12) : f_t(1e-8); + iterative_refinement(op, rhs, soln, ir_tol); } for (i_t k = 0; k < lp.num_cols; k++) { @@ -3196,8 +3199,9 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t(op, data.d_augmented_rhs_, data.d_augmented_soln_); + const f_t ir_tol = data.has_sparse_cones() ? f_t(1e-12) : f_t(1e-8); + const f_t solve_err = iterative_refinement( + op, data.d_augmented_rhs_, data.d_augmented_soln_, ir_tol); if (solve_err > 1e-1) { settings.log.printf("|| Aug (dx, dy) - aug_rhs || %e after IR\n", solve_err); } diff --git a/cpp/src/barrier/iterative_refinement.hpp b/cpp/src/barrier/iterative_refinement.hpp index 8ad0bb4c49..d7bc8276a7 100644 --- a/cpp/src/barrier/iterative_refinement.hpp +++ b/cpp/src/barrier/iterative_refinement.hpp @@ -88,7 +88,8 @@ f_t vector_norm2(const rmm::device_uvector& x) template f_t iterative_refinement_simple(T& op, const rmm::device_uvector& b, - rmm::device_uvector& x) + rmm::device_uvector& x, + f_t tol = 1e-8) { rmm::device_uvector x_sav(x, x.stream()); @@ -105,7 +106,7 @@ f_t iterative_refinement_simple(T& op, } rmm::device_uvector delta_x(x.size(), op.data_.handle_ptr->get_stream()); i_t iter = 0; - while (error > 1e-8 && iter < 30) { + while (error > tol && iter < 30) { thrust::fill(op.data_.handle_ptr->get_thrust_policy(), delta_x.data(), delta_x.data() + delta_x.size(), @@ -154,7 +155,8 @@ f_t iterative_refinement_simple(T& op, template f_t iterative_refinement_gmres(T& op, const rmm::device_uvector& b, - rmm::device_uvector& x) + rmm::device_uvector& x, + f_t tol = 1e-8) { // Parameters // Ideally, we do not need to restart here. But having restarts helps as a checkpoint to get @@ -162,7 +164,6 @@ f_t iterative_refinement_gmres(T& op, // are not converging after some point const int max_restarts = 3; const int m = 10; // Krylov space dimension - const f_t tol = 1e-8; rmm::device_uvector r(x.size(), x.stream()); rmm::device_uvector x_sav(x, x.stream()); @@ -188,7 +189,7 @@ f_t iterative_refinement_gmres(T& op, f_t norm_r = vector_norm_inf(r); if (show_info) { CUOPT_LOG_INFO("GMRES IR: initial residual = %e, |b| = %e", norm_r, bnorm); } - if (norm_r <= 1e-8) { return norm_r; } + if (norm_r <= tol) { return norm_r; } f_t residual = norm_r; f_t best_residual = norm_r; @@ -392,13 +393,16 @@ f_t iterative_refinement_gmres(T& op, } template -f_t iterative_refinement(T& op, const dense_vector_t& b, dense_vector_t& x) +f_t iterative_refinement(T& op, + const dense_vector_t& b, + dense_vector_t& x, + f_t tol = 1e-8) { rmm::device_uvector 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 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(op, d_b, d_x); + auto err = iterative_refinement_gmres(op, d_b, d_x, tol); raft::copy(x.data(), d_x.data(), x.size(), op.data_.handle_ptr->get_stream()); @@ -407,9 +411,12 @@ f_t iterative_refinement(T& op, const dense_vector_t& b, dense_vector_ } template -f_t iterative_refinement(T& op, const rmm::device_uvector& b, rmm::device_uvector& x) +f_t iterative_refinement(T& op, + const rmm::device_uvector& b, + rmm::device_uvector& x, + f_t tol = 1e-8) { - return iterative_refinement_gmres(op, b, x); + return iterative_refinement_gmres(op, b, x, tol); } } // namespace cuopt::mathematical_optimization::barrier From 9622be8510d12fc039121e7ca736365d39bbbadd Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Mon, 6 Jul 2026 10:25:34 -0700 Subject: [PATCH 03/14] explicit matrix-vector product in IR Signed-off-by: yuwenchen95 --- .../mathematical_optimization/constants.h | 1 + .../pdlp/solver_settings.hpp | 1 + cpp/src/barrier/barrier.cu | 165 ++++++++++++++---- cpp/src/barrier/cusparse_view.cu | 107 ++++++++++++ cpp/src/barrier/cusparse_view.hpp | 29 +++ cpp/src/barrier/second_order_cone_kernels.cuh | 123 +++++++++++++ .../dual_simplex/simplex_solver_settings.hpp | 2 + cpp/src/math_optimization/solver_settings.cu | 1 + cpp/src/pdlp/solve.cu | 2 + 9 files changed, 396 insertions(+), 35 deletions(-) diff --git a/cpp/include/cuopt/mathematical_optimization/constants.h b/cpp/include/cuopt/mathematical_optimization/constants.h index 041cc38688..215c7b60d2 100644 --- a/cpp/include/cuopt/mathematical_optimization/constants.h +++ b/cpp/include/cuopt/mathematical_optimization/constants.h @@ -48,6 +48,7 @@ #define CUOPT_ORDERING "ordering" #define CUOPT_BARRIER_DUAL_INITIAL_POINT "barrier_dual_initial_point" #define CUOPT_BARRIER_ITERATIVE_REFINEMENT "barrier_iterative_refinement" +#define CUOPT_BARRIER_CSR_IR_MATVEC "barrier_csr_ir_matvec" #define CUOPT_BARRIER_SOC_THRESHOLD "barrier_soc_threshold" #define CUOPT_BARRIER_STEP_SCALE "barrier_step_scale" #define CUOPT_ELIMINATE_DENSE_COLUMNS "eliminate_dense_columns" diff --git a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp index 8122bf92e0..6ba98f983f 100644 --- a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp +++ b/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}; + bool barrier_csr_ir_matvec{false}; i_t barrier_soc_threshold{5}; f_t barrier_step_scale{0.9}; bool save_best_primal_so_far{false}; diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 45b1feab4a..bfd554e647 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -290,6 +290,7 @@ class iteration_data_t { d_sparse_exp_u_row_(0, lp.handle_ptr->get_stream()), d_sparse_expansion_D_(0, lp.handle_ptr->get_stream()), d_sparse_Hs_diag_(0, lp.handle_ptr->get_stream()), + d_dense_cone_diag_csr_indices_(0, lp.handle_ptr->get_stream()), use_augmented(false), has_factorization(false), n_direct_free_linear(0), @@ -312,6 +313,12 @@ class iteration_data_t { d_r1_prime_(lp.num_cols, lp.handle_ptr->get_stream()), d_augmented_rhs_(0, lp.handle_ptr->get_stream()), d_augmented_soln_(0, lp.handle_ptr->get_stream()), + d_aug_x1_(0, lp.handle_ptr->get_stream()), + d_aug_x2_(0, lp.handle_ptr->get_stream()), + d_aug_y1_(0, lp.handle_ptr->get_stream()), + d_aug_y2_(0, lp.handle_ptr->get_stream()), + d_aug_y_exp_(0, lp.handle_ptr->get_stream()), + d_aug_y_exp_orig_(0, lp.handle_ptr->get_stream()), d_c_(lp.num_cols, lp.handle_ptr->get_stream()), d_b_(lp.num_rows, lp.handle_ptr->get_stream()), d_upper_(0, lp.handle_ptr->get_stream()), @@ -538,6 +545,17 @@ class iteration_data_t { const i_t augmented_size = augmented_system_size(lp.num_cols, lp.num_rows); d_augmented_rhs_.resize(augmented_size, stream_view_); d_augmented_soln_.resize(augmented_size, stream_view_); + d_aug_x1_.resize(lp.num_cols, stream_view_); + d_aug_x2_.resize(lp.num_rows, stream_view_); + d_aug_y1_.resize(lp.num_cols, stream_view_); + d_aug_y2_.resize(lp.num_rows, stream_view_); + d_aug_y_exp_.resize(augmented_expansion_count(), stream_view_); + d_aug_y_exp_orig_.resize(augmented_expansion_count(), stream_view_); + if (settings.barrier_csr_ir_matvec) { + settings.log.printf("IR matvec : csr\n"); + } else { + settings.log.printf("IR matvec : matrix-free\n"); + } } else { settings.log.printf("Linear system : ADAT\n"); } @@ -647,6 +665,10 @@ class iteration_data_t { } if (settings.concurrent_halt != nullptr && *settings.concurrent_halt == 1) { return; } symbolic_status = chol->analyze(device_augmented); + if (use_csr_ir_matvec()) { + augmented_cusparse_view_ = + std::make_unique>(handle_ptr, device_augmented); + } } else { { raft::common::nvtx::range form_scope("Barrier: LP Data: form ADAT"); @@ -690,6 +712,8 @@ class iteration_data_t { return has_sparse_cones() ? cones().expansion_var_count() : i_t(0); } + bool use_csr_ir_matvec() const { return settings_.barrier_csr_ir_matvec && use_augmented; } + i_t augmented_system_size(i_t n, i_t m) const { return n + m + augmented_expansion_count(); } bool is_cone_variable(i_t variable) const @@ -793,6 +817,7 @@ class iteration_data_t { std::vector augmented_diagonal_indices(factorization_size, -1); std::vector cone_csr_indices_host(dense_soc_kkt_nnz, -1); std::vector cone_Q_values_host(dense_soc_kkt_nnz, f_t(0)); + std::vector dense_cone_diag_indices_host; std::vector sparse_hessian_diag_host(n_sparse_entries, -1); std::vector sparse_hessian_Q_host(n_sparse_entries, f_t(0)); std::vector sparse_exp_v_col_host(n_sparse_entries, -1); @@ -896,7 +921,10 @@ class iteration_data_t { cone_csr_indices_host[block_base + c] = q; cone_Q_values_host[block_base + c] = q_contrib; - if (col == i) { augmented_diagonal_indices[i] = q; } + if (col == i) { + augmented_diagonal_indices[i] = q; + dense_cone_diag_indices_host.push_back(q); + } augmented_CSR.j[q] = col; augmented_CSR.x[q++] = initial_val - q_contrib; } @@ -1052,6 +1080,15 @@ class iteration_data_t { handle_ptr->get_stream()); } + d_dense_cone_diag_csr_indices_.resize(dense_cone_diag_indices_host.size(), + handle_ptr->get_stream()); + if (!dense_cone_diag_indices_host.empty()) { + raft::copy(d_dense_cone_diag_csr_indices_.data(), + dense_cone_diag_indices_host.data(), + dense_cone_diag_indices_host.size(), + handle_ptr->get_stream()); + } + const size_t n_sparse_entries = cones().n_sparse_cone_entries; d_sparse_hessian_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); d_sparse_hessian_Q_.resize(n_sparse_entries, handle_ptr->get_stream()); @@ -2071,6 +2108,41 @@ class iteration_data_t { __host__ __device__ T operator()(T x, T y) const { return alpha * x + beta * y; } }; + void strip_augmented_perturbation() + { + raft::common::nvtx::range fun_scope("Barrier: strip augmented perturbation"); + const i_t n = A.n; + const i_t m = A.m; + const i_t linear_n = has_cones() ? cone_start() : n; + strip_augmented_perturbation_values(device_augmented.x, + d_augmented_diagonal_indices_, + d_Q_diag_, + d_is_direct_free_linear_, + d_sparse_hessian_diag_, + d_sparse_expansion_D_, + d_dense_cone_diag_csr_indices_, + linear_n, + n, + m, + dual_perturb, + primal_perturb, + stream_view_); + } + + void augmented_csr_multiply(f_t alpha, + const rmm::device_uvector& x, + f_t beta, + rmm::device_uvector& 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"); + cuopt_assert(augmented_cusparse_view_ != nullptr, "augmented cusparse view not initialized"); + const i_t sys_size = augmented_system_size(A.n, A.m); + cuopt_assert(static_cast(x.size()) >= sys_size, "augmented_csr_multiply: x too small"); + cuopt_assert(static_cast(y.size()) >= sys_size, "augmented_csr_multiply: y too small"); + augmented_cusparse_view_->spmv(alpha, x, beta, y); + } + // y <- alpha * Augmented * x + beta * y void augmented_multiply(f_t alpha, const rmm::device_uvector& x, @@ -2085,47 +2157,39 @@ class iteration_data_t { cuopt_assert(static_cast(x.size()) >= sys_size, "augmented_multiply: x too small"); cuopt_assert(static_cast(y.size()) >= sys_size, "augmented_multiply: y too small"); - rmm::device_uvector d_x1(n, handle_ptr->get_stream()); - rmm::device_uvector d_x2(m, handle_ptr->get_stream()); - rmm::device_uvector d_y1(n, handle_ptr->get_stream()); - rmm::device_uvector d_y2(m, handle_ptr->get_stream()); - rmm::device_uvector d_y_exp(p, handle_ptr->get_stream()); - rmm::device_uvector d_y_exp_orig(p, handle_ptr->get_stream()); - - raft::copy(d_x1.data(), x.data(), n, handle_ptr->get_stream()); - raft::copy(d_x2.data(), x.data() + n, m, handle_ptr->get_stream()); - raft::copy(d_y1.data(), y.data(), n, handle_ptr->get_stream()); - raft::copy(d_y2.data(), y.data() + n, m, handle_ptr->get_stream()); + raft::copy(d_aug_x1_.data(), x.data(), n, handle_ptr->get_stream()); + raft::copy(d_aug_x2_.data(), x.data() + n, m, handle_ptr->get_stream()); + raft::copy(d_aug_y1_.data(), y.data(), n, handle_ptr->get_stream()); + raft::copy(d_aug_y2_.data(), y.data() + n, m, handle_ptr->get_stream()); if (p > 0) { - raft::copy(d_y_exp_orig.data(), y.data() + n + m, p, handle_ptr->get_stream()); - thrust::fill_n(rmm::exec_policy(stream_view_), d_y_exp.begin(), p, f_t(0)); + raft::copy(d_aug_y_exp_orig_.data(), y.data() + n + m, p, handle_ptr->get_stream()); + thrust::fill_n(rmm::exec_policy(stream_view_), d_aug_y_exp_.begin(), p, f_t(0)); } // y1 <- alpha ( -(Q + D + H) * x_1 + A^T x_2) + beta * y1 - rmm::device_uvector d_r1(n, handle_ptr->get_stream()); - thrust::fill_n(rmm::exec_policy(stream_view_), d_r1.begin(), n, f_t(0)); + thrust::fill_n(rmm::exec_policy(stream_view_), d_r1_.begin(), n, f_t(0)); // r1 <- D * x_1 on linear indices; barrier D is zero on direct free variables const i_t linear_n = has_soc ? cone_start() : n; - pairwise_multiply_skip_direct_free_linear(d_x1.data(), + pairwise_multiply_skip_direct_free_linear(d_aug_x1_.data(), d_diag_.data(), d_is_direct_free_linear_.data(), - d_r1.data(), + d_r1_.data(), linear_n, stream_view_); RAFT_CHECK_CUDA(stream_view_); // r1 <- D * x_1 + H x_1 on cone rows // (dense cones: explicit dense H block; sparse cones: rank-2 expansion, which adds - // Hs_diag .* x_cone to r1 here and writes the expansion rows into d_y_exp) + // Hs_diag .* x_cone to r1 here and writes the expansion rows into d_aug_y_exp_) if (has_soc) { const i_t m_c = cone_entry_count(); if (cones().has_sparse_cones()) { launch_sparse_augmented_matvec( raft::device_span(x.data(), x.size()), - raft::device_span(d_r1.data(), d_r1.size()), - raft::device_span(d_y_exp.data(), d_y_exp.size()), + raft::device_span(d_r1_.data(), d_r1_.size()), + raft::device_span(d_aug_y_exp_.data(), d_aug_y_exp_.size()), cones(), raft::device_span(d_sparse_Hs_diag_.data(), d_sparse_Hs_diag_.size()), cone_start(), @@ -2135,10 +2199,11 @@ class iteration_data_t { RAFT_CHECK_CUDA(stream_view_); } if (cones().n_dense_cones() > 0) { - launch_dense_hessian_matvec(raft::device_span(d_x1.data() + cone_start(), m_c), - cones(), - raft::device_span(d_r1.data() + cone_start(), m_c), - stream_view_); + launch_dense_hessian_matvec( + raft::device_span(d_aug_x1_.data() + cone_start(), m_c), + cones(), + raft::device_span(d_r1_.data() + cone_start(), m_c), + stream_view_); RAFT_CHECK_CUDA(stream_view_); } } @@ -2146,26 +2211,32 @@ class iteration_data_t { // r1 <- Q x1 + D x1 + H x1 (cone: same H as above) if (Q.n > 0) { // matrix_vector_multiply(Q, 1.0, x1, 1.0, r1); - cusparse_Q_view_.spmv(1.0, d_x1, 1.0, d_r1); + cusparse_Q_view_.spmv(1.0, d_aug_x1_, 1.0, d_r1_); } // y1 <- - alpha * r1 + beta * y1 // flip the sign of r1 = (Q x1 + D x1 + H x1) - axpy(-alpha, d_r1.data(), beta, d_y1.data(), d_y1.data(), n, stream_view_); + axpy(-alpha, d_r1_.data(), beta, d_aug_y1_.data(), d_aug_y1_.data(), n, stream_view_); // matrix_transpose_vector_multiply(A, alpha, x2, 1.0, y1); - cusparse_view_.transpose_spmv(alpha, d_x2, 1.0, d_y1); + cusparse_view_.transpose_spmv(alpha, d_aug_x2_, 1.0, d_aug_y1_); // y2 <- alpha ( A*x) + beta * y2 // matrix_vector_multiply(A, alpha, x1, beta, y2); - cusparse_view_.spmv(alpha, d_x1, beta, d_y2); + cusparse_view_.spmv(alpha, d_aug_x1_, beta, d_aug_y2_); if (p > 0) { - axpy(alpha, d_y_exp.data(), beta, d_y_exp_orig.data(), d_y_exp.data(), p, stream_view_); + axpy(alpha, + d_aug_y_exp_.data(), + beta, + d_aug_y_exp_orig_.data(), + d_aug_y_exp_.data(), + p, + stream_view_); } - raft::copy(y.data(), d_y1.data(), n, stream_view_); - raft::copy(y.data() + n, d_y2.data(), m, stream_view_); - if (p > 0) { raft::copy(y.data() + n + m, d_y_exp.data(), p, stream_view_); } + raft::copy(y.data(), d_aug_y1_.data(), n, stream_view_); + raft::copy(y.data() + n, d_aug_y2_.data(), m, stream_view_); + if (p > 0) { raft::copy(y.data() + n + m, d_aug_y_exp_.data(), p, stream_view_); } handle_ptr->sync_stream(); } @@ -2257,6 +2328,7 @@ class iteration_data_t { rmm::device_uvector d_sparse_exp_u_row_; rmm::device_uvector d_sparse_expansion_D_; rmm::device_uvector d_sparse_Hs_diag_; + rmm::device_uvector d_dense_cone_diag_csr_indices_; bool indefinite_Q; cusparse_view_t cusparse_Q_view_; @@ -2274,6 +2346,7 @@ class iteration_data_t { f_t primal_perturb{1e-8}; std::unique_ptr> chol; + std::unique_ptr> augmented_cusparse_view_; bool has_factorization; bool has_solve_info; @@ -2313,6 +2386,12 @@ class iteration_data_t { rmm::device_uvector d_r1_prime_; rmm::device_uvector d_augmented_rhs_; rmm::device_uvector d_augmented_soln_; + rmm::device_uvector d_aug_x1_; + rmm::device_uvector d_aug_x2_; + rmm::device_uvector d_aug_y1_; + rmm::device_uvector d_aug_y2_; + rmm::device_uvector d_aug_y_exp_; + rmm::device_uvector d_aug_y_exp_orig_; rmm::device_uvector d_c_; rmm::device_uvector d_b_; rmm::device_uvector d_upper_; @@ -2516,6 +2595,10 @@ int barrier_solver_t::initial_point(iteration_data_t& data) #ifdef CHOLESKY_DEBUG_CHECK cholesky_debug_check(data, lp, use_augmented); #endif + if (data.use_csr_ir_matvec()) { + // Strip augmented perturbation for CSR IR matvec path + data.strip_augmented_perturbation(); + } } else { status = data.chol->factorize(data.device_ADAT); } @@ -2556,7 +2639,11 @@ int barrier_solver_t::initial_point(iteration_data_t& data) f_t beta, rmm::device_uvector& 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& b, rmm::device_uvector& x) const { @@ -3082,6 +3169,10 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_tfactorize(data.device_augmented); + if (data.use_csr_ir_matvec()) { + // Strip augmented perturbation for CSR IR matvec path + data.strip_augmented_perturbation(); + } #ifdef CHOLESKY_DEBUG_CHECK cholesky_debug_check(data, lp, use_augmented); @@ -3190,7 +3281,11 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_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& b, rmm::device_uvector& x) const diff --git a/cpp/src/barrier/cusparse_view.cu b/cpp/src/barrier/cusparse_view.cu index 2585084097..77945b81f7 100644 --- a/cpp/src/barrier/cusparse_view.cu +++ b/cpp/src/barrier/cusparse_view.cu @@ -423,4 +423,111 @@ cusparse_view_t::transpose_spmv, std::alloca double beta, std::vector>& y); +template +augmented_cusparse_view_t::augmented_cusparse_view_t( + raft::handle_t const* handle_ptr, device_csr_matrix_t& matrix) + : handle_ptr_(handle_ptr), + spmv_buffer_(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()), + num_rows_(matrix.m) +{ + RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsesetpointermode( + handle_ptr->get_cusparse_handle(), CUSPARSE_POINTER_MODE_DEVICE, handle_ptr->get_stream())); + + const i_t rows = matrix.m; + const i_t cols = matrix.n; + const i_t nnz = matrix.nz_max; + + RAFT_CUSPARSE_TRY(cusparseCreateCsr(&A_, + rows, + cols, + nnz, + matrix.row_start.data(), + matrix.j.data(), + matrix.x.data(), + CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_BASE_ZERO, + CUDA_R_64F)); + + cusparseDnVecDescr_t x; + cusparseDnVecDescr_t y; + rmm::device_uvector d_x(cols, handle_ptr_->get_stream()); + rmm::device_uvector 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(rows), + &buffer_size_spmv, + handle_ptr_->get_stream())); + spmv_buffer_.resize(buffer_size_spmv, handle_ptr_->get_stream()); + RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(x)); + RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y)); +} + +template +augmented_cusparse_view_t::~augmented_cusparse_view_t() +{ + CUOPT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(A_)); +} + +template +pdlp::cusparse_dn_vec_descr_wrapper_t augmented_cusparse_view_t::create_vector( + rmm::device_uvector const& vec) +{ + pdlp::cusparse_dn_vec_descr_wrapper_t descr; + descr.create(vec.size(), const_cast(vec.data())); + return descr; +} + +template +void augmented_cusparse_view_t::spmv(f_t alpha, + const rmm::device_uvector& x, + f_t beta, + rmm::device_uvector& y) +{ + pdlp::cusparse_dn_vec_descr_wrapper_t x_cusparse = create_vector(x); + pdlp::cusparse_dn_vec_descr_wrapper_t y_cusparse = create_vector(y); + spmv(alpha, x_cusparse, beta, y_cusparse); +} + +template +void augmented_cusparse_view_t::spmv(f_t alpha, + pdlp::cusparse_dn_vec_descr_wrapper_t const& x, + f_t beta, + pdlp::cusparse_dn_vec_descr_wrapper_t const& y) +{ + cuopt_assert(alpha == f_t(1) || alpha == f_t(-1), "Only alpha 1 or -1 supported"); + cuopt_assert(beta == f_t(1) || beta == f_t(-1) || beta == f_t(0), + "Only beta 1 or -1 or 0 supported"); + rmm::device_scalar* d_beta = &d_one_; + if (beta == f_t(0)) + d_beta = &d_zero_; + else if (beta == f_t(-1)) + d_beta = &d_minus_one_; + raft::sparse::detail::cusparsespmv(handle_ptr_->get_cusparse_handle(), + CUSPARSE_OPERATION_NON_TRANSPOSE, + (alpha == 1) ? d_one_.data() : d_minus_one_.data(), + A_, + x, + d_beta->data(), + y, + get_spmv_alg(num_rows_), + (f_t*)spmv_buffer_.data(), + handle_ptr_->get_stream()); +} + +template class augmented_cusparse_view_t; + } // namespace cuopt::mathematical_optimization::barrier diff --git a/cpp/src/barrier/cusparse_view.hpp b/cpp/src/barrier/cusparse_view.hpp index 6af4813bc2..6846de420f 100644 --- a/cpp/src/barrier/cusparse_view.hpp +++ b/cpp/src/barrier/cusparse_view.hpp @@ -6,6 +6,8 @@ /* clang-format on */ #pragma once +#include + #include #include @@ -73,4 +75,31 @@ class cusparse_view_t { rmm::device_scalar d_minus_one_; rmm::device_scalar d_zero_; }; + +// cuSparse CSR SpMV view for the augmented KKT matrix stored in device_csr_matrix_t. +// The sparsity pattern is fixed; numerical values are read from matrix.x in place. +template +class augmented_cusparse_view_t { + public: + augmented_cusparse_view_t(raft::handle_t const* handle_ptr, + device_csr_matrix_t& matrix); + ~augmented_cusparse_view_t(); + + pdlp::cusparse_dn_vec_descr_wrapper_t create_vector(rmm::device_uvector const& vec); + + void spmv(f_t alpha, rmm::device_uvector const& x, f_t beta, rmm::device_uvector& y); + void spmv(f_t alpha, + pdlp::cusparse_dn_vec_descr_wrapper_t const& x, + f_t beta, + pdlp::cusparse_dn_vec_descr_wrapper_t const& y); + + private: + raft::handle_t const* handle_ptr_{nullptr}; + cusparseSpMatDescr_t A_{nullptr}; + rmm::device_buffer spmv_buffer_; + rmm::device_scalar d_one_; + rmm::device_scalar d_minus_one_; + rmm::device_scalar d_zero_; + i_t num_rows_{0}; +}; } // namespace cuopt::mathematical_optimization::barrier diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index 5bd787c836..04924a23ff 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -1169,6 +1169,129 @@ void update_sparse_expansion_in_augmented(rmm::device_uvector& augmented_x, } } +// Remove factorization-only regularization from the assembled augmented CSR values so +// iterative refinement can use a single SpMV against the true KKT operator. +template +__global__ void strip_linear_primal_diag_kernel(raft::device_span augmented_x, + raft::device_span diag_indices, + raft::device_span q_diag, + raft::device_span is_direct_free_linear, + f_t dual_perturb, + i_t linear_n) +{ + const i_t j = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (j >= linear_n) { return; } + + const i_t idx = diag_indices[j]; + if (is_direct_free_linear[j]) { + const f_t q_val = q_diag.size() > 0 ? q_diag[j] : f_t(0); + augmented_x[idx] = -q_val; + } else { + augmented_x[idx] += dual_perturb; + } +} + +template +__global__ void strip_dual_diag_kernel(raft::device_span augmented_x, + raft::device_span diag_indices, + f_t primal_perturb, + i_t n, + i_t m) +{ + const i_t l = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (l >= m) { return; } + + const i_t row = n + l; + augmented_x[diag_indices[row]] -= primal_perturb; +} + +template +__global__ void strip_indexed_offset_kernel(raft::device_span augmented_x, + raft::device_span csr_indices, + f_t offset) +{ + const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (e >= csr_indices.size()) { return; } + + augmented_x[csr_indices[e]] += offset; +} + +template +__global__ void strip_expansion_diag_kernel(raft::device_span augmented_x, + raft::device_span expansion_D, + f_t dual_perturb) +{ + const i_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= static_cast(expansion_D.size())) { return; } + + const f_t sign = (i % 2 == 0) ? f_t(-1) : f_t(1); + augmented_x[expansion_D[i]] -= sign * dual_perturb; +} + +template +void strip_augmented_perturbation_values( + rmm::device_uvector& augmented_x, + const rmm::device_uvector& augmented_diagonal_indices, + const rmm::device_uvector& q_diag, + const rmm::device_uvector& is_direct_free_linear, + const rmm::device_uvector& sparse_hessian_diag, + const rmm::device_uvector& sparse_expansion_D, + const rmm::device_uvector& dense_cone_diag_csr_indices, + i_t linear_n, + i_t n, + i_t m, + f_t dual_perturb, + f_t primal_perturb, + rmm::cuda_stream_view stream) +{ + if (linear_n > 0) { + const size_t grid = raft::ceildiv(linear_n, soc_block_size); + strip_linear_primal_diag_kernel + <<>>(cuopt::make_span(augmented_x), + cuopt::make_span(augmented_diagonal_indices), + cuopt::make_span(q_diag), + cuopt::make_span(is_direct_free_linear), + dual_perturb, + linear_n); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (m > 0) { + const size_t grid = raft::ceildiv(m, soc_block_size); + strip_dual_diag_kernel + <<>>(cuopt::make_span(augmented_x), + cuopt::make_span(augmented_diagonal_indices), + primal_perturb, + n, + m); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (sparse_hessian_diag.size() > 0) { + const size_t count = sparse_hessian_diag.size(); + const size_t grid = raft::ceildiv(count, soc_block_size); + strip_indexed_offset_kernel<<>>( + cuopt::make_span(augmented_x), cuopt::make_span(sparse_hessian_diag), dual_perturb); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (sparse_expansion_D.size() > 0) { + const size_t count = sparse_expansion_D.size(); + const size_t grid = raft::ceildiv(count, soc_block_size); + strip_expansion_diag_kernel<<>>( + cuopt::make_span(augmented_x), cuopt::make_span(sparse_expansion_D), dual_perturb); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (dense_cone_diag_csr_indices.size() > 0) { + const size_t count = dense_cone_diag_csr_indices.size(); + const size_t grid = raft::ceildiv(count, soc_block_size); + strip_indexed_offset_kernel<<>>( + cuopt::make_span(augmented_x), cuopt::make_span(dense_cone_diag_csr_indices), dual_perturb); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + /** * Accumulate the sparse-SOC expanded KKT block into a matrix-free product. * diff --git a/cpp/src/dual_simplex/simplex_solver_settings.hpp b/cpp/src/dual_simplex/simplex_solver_settings.hpp index 46499e807c..8726a95612 100644 --- a/cpp/src/dual_simplex/simplex_solver_settings.hpp +++ b/cpp/src/dual_simplex/simplex_solver_settings.hpp @@ -68,6 +68,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(5), num_gpus(1), @@ -157,6 +158,7 @@ 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 + bool barrier_csr_ir_matvec; // true to use CSR SpMV for augmented IR matvec 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) diff --git a/cpp/src/math_optimization/solver_settings.cu b/cpp/src/math_optimization/solver_settings.cu index cfeafcb948..740f0cc006 100644 --- a/cpp/src/math_optimization/solver_settings.cu +++ b/cpp/src/math_optimization/solver_settings.cu @@ -191,6 +191,7 @@ solver_settings_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}, // Diving heuristic hyper-parameters (hidden from default --help: name contains "hyper_") {CUOPT_MIP_HYPER_DIVING_SHOW_TYPE, &mip_settings.diving_params.show_type, false, "log diving heuristic type when it finds a new incumbent"}, diff --git a/cpp/src/pdlp/solve.cu b/cpp/src/pdlp/solve.cu index e6756a5125..e6b373ee73 100644 --- a/cpp/src/pdlp/solve.cu +++ b/cpp/src/pdlp/solve.cu @@ -509,6 +509,7 @@ std::tuple, 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.cudss_deterministic = settings.cudss_deterministic; @@ -693,6 +694,7 @@ static optimization_problem_solution_t 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; From e5b81d7ea79e9c5e0f6a04f06eea5360cbe7d4f8 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Thu, 9 Jul 2026 08:30:53 -0700 Subject: [PATCH 04/14] Remove a slow inner loop in CSR matrix buildup Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index bfd554e647..c73e9f62ac 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -756,6 +756,7 @@ class iteration_data_t { std::vector cone_offsets_host; std::vector dense_cone_block_offsets_host; std::vector dense_cone_ids_host; + std::vector dense_cone_idx_by_cone; std::vector cone_sparse_idx; std::vector sparse_cone_ids_host; std::vector sparse_entry_offsets; @@ -788,14 +789,18 @@ class iteration_data_t { } dense_cone_ids_host.reserve(n_cones); + dense_cone_idx_by_cone.assign(n_cones, i_t(-1)); dense_cone_block_offsets_host = {0}; dense_cone_block_offsets_host.reserve(n_cones + 1); + i_t dense_idx = 0; for (i_t k = 0; k < n_cones; ++k) { if (cone_sparse_idx[k] < 0) { + dense_cone_idx_by_cone[k] = dense_idx; dense_cone_ids_host.push_back(k); const auto q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; dense_cone_block_offsets_host.push_back(dense_cone_block_offsets_host.back() + q_k * q_k); + ++dense_idx; } } dense_soc_kkt_nnz = static_cast(dense_cone_block_offsets_host.back()); @@ -896,10 +901,8 @@ class iteration_data_t { } } else { // Row of dense SOC. - i_t dense_idx = 0; - for (i_t t = 0; t < k; ++t) { - if (cone_sparse_idx[t] < 0) { ++dense_idx; } - } + const i_t dense_idx = dense_cone_idx_by_cone[k]; + cuopt_assert(dense_idx >= 0, "dense cone index unset"); i_t block_base = static_cast(dense_cone_block_offsets_host[dense_idx]) + local_r * q_k; From dc2e7f0b5d8ec937963a6914fb42c708b7954e73 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Fri, 10 Jul 2026 02:51:11 -0700 Subject: [PATCH 05/14] clean up cusparse_view & add more ranges for profiling Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 134 +++++++----- cpp/src/barrier/cusparse_view.cu | 250 +++++++++-------------- cpp/src/barrier/cusparse_view.hpp | 48 ++--- cpp/src/barrier/iterative_refinement.hpp | 34 ++- cpp/src/dual_simplex/presolve.cpp | 4 + cpp/src/dual_simplex/scaling.cpp | 3 + cpp/src/pdlp/translate.hpp | 3 + 7 files changed, 227 insertions(+), 249 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index c73e9f62ac..635f475a31 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -667,7 +667,7 @@ class iteration_data_t { symbolic_status = chol->analyze(device_augmented); if (use_csr_ir_matvec()) { augmented_cusparse_view_ = - std::make_unique>(handle_ptr, device_augmented); + std::make_unique>(handle_ptr, device_augmented); } } else { { @@ -2152,6 +2152,7 @@ class iteration_data_t { f_t beta, rmm::device_uvector& y) { + raft::common::nvtx::range fun_scope("Barrier: augmented_multiply"); const i_t m = A.m; const i_t n = A.n; const bool has_soc = has_cones(); @@ -2175,18 +2176,22 @@ class iteration_data_t { // r1 <- D * x_1 on linear indices; barrier D is zero on direct free variables const i_t linear_n = has_soc ? cone_start() : n; - pairwise_multiply_skip_direct_free_linear(d_aug_x1_.data(), - d_diag_.data(), - d_is_direct_free_linear_.data(), - d_r1_.data(), - linear_n, - stream_view_); - RAFT_CHECK_CUDA(stream_view_); + { + raft::common::nvtx::range scope("Barrier: augmented_multiply: D * x1 (linear)"); + pairwise_multiply_skip_direct_free_linear(d_aug_x1_.data(), + d_diag_.data(), + d_is_direct_free_linear_.data(), + d_r1_.data(), + linear_n, + stream_view_); + RAFT_CHECK_CUDA(stream_view_); + } // r1 <- D * x_1 + H x_1 on cone rows // (dense cones: explicit dense H block; sparse cones: rank-2 expansion, which adds // Hs_diag .* x_cone to r1 here and writes the expansion rows into d_aug_y_exp_) if (has_soc) { + raft::common::nvtx::range scope("Barrier: augmented_multiply: cone Hessian (H * x1)"); const i_t m_c = cone_entry_count(); if (cones().has_sparse_cones()) { launch_sparse_augmented_matvec( @@ -2213,28 +2218,32 @@ class iteration_data_t { // r1 <- Q x1 + D x1 + H x1 (cone: same H as above) if (Q.n > 0) { + raft::common::nvtx::range scope("Barrier: augmented_multiply: Q * x1"); // matrix_vector_multiply(Q, 1.0, x1, 1.0, r1); cusparse_Q_view_.spmv(1.0, d_aug_x1_, 1.0, d_r1_); } - // y1 <- - alpha * r1 + beta * y1 - // flip the sign of r1 = (Q x1 + D x1 + H x1) - axpy(-alpha, d_r1_.data(), beta, d_aug_y1_.data(), d_aug_y1_.data(), n, stream_view_); - - // matrix_transpose_vector_multiply(A, alpha, x2, 1.0, y1); - cusparse_view_.transpose_spmv(alpha, d_aug_x2_, 1.0, d_aug_y1_); - // y2 <- alpha ( A*x) + beta * y2 - // matrix_vector_multiply(A, alpha, x1, beta, y2); - cusparse_view_.spmv(alpha, d_aug_x1_, beta, d_aug_y2_); - - if (p > 0) { - axpy(alpha, - d_aug_y_exp_.data(), - beta, - d_aug_y_exp_orig_.data(), - d_aug_y_exp_.data(), - p, - stream_view_); + { + raft::common::nvtx::range scope("Barrier: augmented_multiply: A products (A^T x2, A x1)"); + // y1 <- - alpha * r1 + beta * y1 + // flip the sign of r1 = (Q x1 + D x1 + H x1) + axpy(-alpha, d_r1_.data(), beta, d_aug_y1_.data(), d_aug_y1_.data(), n, stream_view_); + + // matrix_transpose_vector_multiply(A, alpha, x2, 1.0, y1); + cusparse_view_.transpose_spmv(alpha, d_aug_x2_, 1.0, d_aug_y1_); + // y2 <- alpha ( A*x) + beta * y2 + // matrix_vector_multiply(A, alpha, x1, beta, y2); + cusparse_view_.spmv(alpha, d_aug_x1_, beta, d_aug_y2_); + + if (p > 0) { + axpy(alpha, + d_aug_y_exp_.data(), + beta, + d_aug_y_exp_orig_.data(), + d_aug_y_exp_.data(), + p, + stream_view_); + } } raft::copy(y.data(), d_aug_y1_.data(), n, stream_view_); @@ -2349,7 +2358,7 @@ class iteration_data_t { f_t primal_perturb{1e-8}; std::unique_ptr> chol; - std::unique_ptr> augmented_cusparse_view_; + std::unique_ptr> augmented_cusparse_view_; bool has_factorization; bool has_solve_info; @@ -3051,6 +3060,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t(data.d_x_.data() + cone_var_start, m_c); cones.z = raft::device_span(data.d_z_.data() + cone_var_start, m_c); @@ -3157,22 +3167,27 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_tfactorize(data.device_augmented); + { + raft::common::nvtx::range fun_scope("Barrier: factorize"); + status = data.chol->factorize(data.device_augmented); + } if (data.use_csr_ir_matvec()) { + raft::common::nvtx::range fun_scope("Barrier: strip_augmented_perturbation"); // Strip augmented perturbation for CSR IR matvec path data.strip_augmented_perturbation(); } @@ -3181,16 +3196,21 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_tfactorize(data.device_ADAT); + { + raft::common::nvtx::range fun_scope("Barrier: factorize"); + status = data.chol->factorize(data.device_ADAT); + } } data.has_factorization = true; data.num_factorizations++; @@ -3297,6 +3317,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t( op, data.d_augmented_rhs_, data.d_augmented_soln_, ir_tol); @@ -3330,7 +3351,10 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t::solve(f_t start_time, lp_solution_t::solve(f_t start_time, lp_solution_t +void cusparse_view_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 cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, const simplex::csc_matrix_t& A) @@ -135,6 +169,7 @@ cusparse_view_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()) @@ -148,7 +183,7 @@ cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, if (debug) { printf("A hash: %zu\n", A.hash()); } simplex::csr_matrix_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& offsets = A_csr.row_start; @@ -164,7 +199,7 @@ cusparse_view_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(), @@ -177,7 +212,7 @@ cusparse_view_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(), @@ -191,58 +226,61 @@ cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, cusparseDnVecDescr_t x; cusparseDnVecDescr_t y; rmm::device_uvector d_x(cols, handle_ptr_->get_stream()); - rmm::device_uvector d_y(rows, handle_ptr_->get_stream()); + rmm::device_uvector 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()); + 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); - 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()); + RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(x)); + RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y)); +} - 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()); +template +cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, + device_csr_matrix_t& matrix) + : 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_(matrix.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 = matrix.n; + const i_t nnz = matrix.nz_max; + + RAFT_CUSPARSE_TRY(cusparseCreateCsr(&A_, + rows_, + cols, + nnz, + matrix.row_start.data(), + matrix.j.data(), + matrix.x.data(), + CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_BASE_ZERO, + CUDA_R_64F)); + + cusparseDnVecDescr_t x; + cusparseDnVecDescr_t y; + rmm::device_uvector d_x(cols, handle_ptr_->get_stream()); + rmm::device_uvector 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_); - 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)); } @@ -251,7 +289,7 @@ template cusparse_view_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 @@ -310,7 +348,7 @@ void cusparse_view_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()); } @@ -322,6 +360,7 @@ void cusparse_view_t::transpose_spmv(f_t alpha, f_t beta, std::vector& 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); @@ -334,6 +373,7 @@ void cusparse_view_t::transpose_spmv(f_t alpha, f_t beta, rmm::device_uvector& y) { + cuopt_assert(A_T_ != nullptr, "transpose_spmv requires an A^T descriptor"); pdlp::cusparse_dn_vec_descr_wrapper_t x_cusparse = create_vector(x); pdlp::cusparse_dn_vec_descr_wrapper_t y_cusparse = create_vector(y); transpose_spmv(alpha, x_cusparse, beta, y_cusparse); @@ -345,6 +385,7 @@ void cusparse_view_t::transpose_spmv(f_t alpha, f_t beta, pdlp::cusparse_dn_vec_descr_wrapper_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"); @@ -423,111 +464,4 @@ cusparse_view_t::transpose_spmv, std::alloca double beta, std::vector>& y); -template -augmented_cusparse_view_t::augmented_cusparse_view_t( - raft::handle_t const* handle_ptr, device_csr_matrix_t& matrix) - : handle_ptr_(handle_ptr), - spmv_buffer_(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()), - num_rows_(matrix.m) -{ - RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsesetpointermode( - handle_ptr->get_cusparse_handle(), CUSPARSE_POINTER_MODE_DEVICE, handle_ptr->get_stream())); - - const i_t rows = matrix.m; - const i_t cols = matrix.n; - const i_t nnz = matrix.nz_max; - - RAFT_CUSPARSE_TRY(cusparseCreateCsr(&A_, - rows, - cols, - nnz, - matrix.row_start.data(), - matrix.j.data(), - matrix.x.data(), - CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_BASE_ZERO, - CUDA_R_64F)); - - cusparseDnVecDescr_t x; - cusparseDnVecDescr_t y; - rmm::device_uvector d_x(cols, handle_ptr_->get_stream()); - rmm::device_uvector 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(rows), - &buffer_size_spmv, - handle_ptr_->get_stream())); - spmv_buffer_.resize(buffer_size_spmv, handle_ptr_->get_stream()); - RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(x)); - RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y)); -} - -template -augmented_cusparse_view_t::~augmented_cusparse_view_t() -{ - CUOPT_CUSPARSE_TRY_NO_THROW(cusparseDestroySpMat(A_)); -} - -template -pdlp::cusparse_dn_vec_descr_wrapper_t augmented_cusparse_view_t::create_vector( - rmm::device_uvector const& vec) -{ - pdlp::cusparse_dn_vec_descr_wrapper_t descr; - descr.create(vec.size(), const_cast(vec.data())); - return descr; -} - -template -void augmented_cusparse_view_t::spmv(f_t alpha, - const rmm::device_uvector& x, - f_t beta, - rmm::device_uvector& y) -{ - pdlp::cusparse_dn_vec_descr_wrapper_t x_cusparse = create_vector(x); - pdlp::cusparse_dn_vec_descr_wrapper_t y_cusparse = create_vector(y); - spmv(alpha, x_cusparse, beta, y_cusparse); -} - -template -void augmented_cusparse_view_t::spmv(f_t alpha, - pdlp::cusparse_dn_vec_descr_wrapper_t const& x, - f_t beta, - pdlp::cusparse_dn_vec_descr_wrapper_t const& y) -{ - cuopt_assert(alpha == f_t(1) || alpha == f_t(-1), "Only alpha 1 or -1 supported"); - cuopt_assert(beta == f_t(1) || beta == f_t(-1) || beta == f_t(0), - "Only beta 1 or -1 or 0 supported"); - rmm::device_scalar* d_beta = &d_one_; - if (beta == f_t(0)) - d_beta = &d_zero_; - else if (beta == f_t(-1)) - d_beta = &d_minus_one_; - raft::sparse::detail::cusparsespmv(handle_ptr_->get_cusparse_handle(), - CUSPARSE_OPERATION_NON_TRANSPOSE, - (alpha == 1) ? d_one_.data() : d_minus_one_.data(), - A_, - x, - d_beta->data(), - y, - get_spmv_alg(num_rows_), - (f_t*)spmv_buffer_.data(), - handle_ptr_->get_stream()); -} - -template class augmented_cusparse_view_t; - } // namespace cuopt::mathematical_optimization::barrier diff --git a/cpp/src/barrier/cusparse_view.hpp b/cpp/src/barrier/cusparse_view.hpp index 6846de420f..37ed9acd39 100644 --- a/cpp/src/barrier/cusparse_view.hpp +++ b/cpp/src/barrier/cusparse_view.hpp @@ -19,17 +19,19 @@ #include -// 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 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 simplex::csc_matrix_t& A); + + // Wire cuSparse SpMV over existing device CSR buffers (no copy). Forward SpMV only. + cusparse_view_t(raft::handle_t const* handle_ptr, device_csr_matrix_t& matrix); ~cusparse_view_t(); pdlp::cusparse_dn_vec_descr_wrapper_t create_vector(rmm::device_uvector const& vec); @@ -61,45 +63,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 A_offsets_; rmm::device_uvector A_indices_; rmm::device_uvector A_data_; - cusparseSpMatDescr_t A_; + cusparseSpMatDescr_t A_{nullptr}; rmm::device_uvector A_T_offsets_; rmm::device_uvector A_T_indices_; rmm::device_uvector 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 d_one_; rmm::device_scalar d_minus_one_; rmm::device_scalar d_zero_; -}; - -// cuSparse CSR SpMV view for the augmented KKT matrix stored in device_csr_matrix_t. -// The sparsity pattern is fixed; numerical values are read from matrix.x in place. -template -class augmented_cusparse_view_t { - public: - augmented_cusparse_view_t(raft::handle_t const* handle_ptr, - device_csr_matrix_t& matrix); - ~augmented_cusparse_view_t(); - - pdlp::cusparse_dn_vec_descr_wrapper_t create_vector(rmm::device_uvector const& vec); - - void spmv(f_t alpha, rmm::device_uvector const& x, f_t beta, rmm::device_uvector& y); - void spmv(f_t alpha, - pdlp::cusparse_dn_vec_descr_wrapper_t const& x, - f_t beta, - pdlp::cusparse_dn_vec_descr_wrapper_t const& y); - - private: - raft::handle_t const* handle_ptr_{nullptr}; - cusparseSpMatDescr_t A_{nullptr}; - rmm::device_buffer spmv_buffer_; - rmm::device_scalar d_one_; - rmm::device_scalar d_minus_one_; - rmm::device_scalar d_zero_; - i_t num_rows_{0}; + i_t rows_{0}; }; } // namespace cuopt::mathematical_optimization::barrier diff --git a/cpp/src/barrier/iterative_refinement.hpp b/cpp/src/barrier/iterative_refinement.hpp index d7bc8276a7..dc4a30f09c 100644 --- a/cpp/src/barrier/iterative_refinement.hpp +++ b/cpp/src/barrier/iterative_refinement.hpp @@ -22,6 +22,8 @@ #include +#include + #include #include #include @@ -112,7 +114,10 @@ f_t iterative_refinement_simple(T& op, delta_x.data() + delta_x.size(), 0.0); RAFT_CHECK_CUDA(op.data_.handle_ptr->get_stream()); - op.solve(r, delta_x); + { + raft::common::nvtx::range solve_scope("Barrier: iterative_refinement: linear system solve"); + op.solve(r, delta_x); + } thrust::transform(op.data_.handle_ptr->get_thrust_policy(), x.data(), @@ -158,6 +163,7 @@ f_t iterative_refinement_gmres(T& op, rmm::device_uvector& x, f_t tol = 1e-8) { + raft::common::nvtx::range fun_scope("Barrier: iterative_refinement: gmres"); // Parameters // Ideally, we do not need to restart here. But having restarts helps as a checkpoint to get // better solutions in case of true residual is far from the measured residual and true residuals @@ -183,11 +189,14 @@ f_t iterative_refinement_gmres(T& op, f_t rel_res = 1.0; int outer_iter = 0; - // r = b - A*x - raft::copy(r.data(), b.data(), b.size(), x.stream()); - op.a_multiply(-1.0, x, 1.0, r); - - f_t norm_r = vector_norm_inf(r); + f_t norm_r; + { + raft::common::nvtx::range scope("Barrier: iterative_refinement: initial residual"); + // r = b - A*x + raft::copy(r.data(), b.data(), b.size(), x.stream()); + op.a_multiply(-1.0, x, 1.0, r); + norm_r = vector_norm_inf(r); + } if (show_info) { CUOPT_LOG_INFO("GMRES IR: initial residual = %e, |b| = %e", norm_r, bnorm); } if (norm_r <= tol) { return norm_r; } @@ -196,6 +205,7 @@ f_t iterative_refinement_gmres(T& op, // Main loop while (residual > tol && outer_iter < max_restarts) { + raft::common::nvtx::range restart_scope("Barrier: iterative_refinement: restart"); // For right preconditioning: Apply preconditioner on Krylov directions, not on the residual. // So, start GMRES on r = b - A*x. v0 = r / ||r|| std::vector> V; @@ -221,8 +231,12 @@ f_t iterative_refinement_gmres(T& op, // Hessenberg building int k = 0; for (; k < m; ++k) { + raft::common::nvtx::range arnoldi_scope("Barrier: iterative_refinement: arnoldi step"); // Z[k] = M^{-1} V[k], i.e., apply right preconditioner and store - op.solve(V[k], Z[k]); + { + raft::common::nvtx::range solve_scope("Barrier: iterative_refinement: linear system solve"); + op.solve(V[k], Z[k]); + } // Check if solve produced NaN (indicates cuDSS failure) f_t z_norm = vector_norm_inf(Z[k]); @@ -232,7 +246,10 @@ f_t iterative_refinement_gmres(T& op, } // w = A * Z[k] - op.a_multiply(1.0, Z[k], 0.0, V[k + 1]); + { + raft::common::nvtx::range matvec_scope("Barrier: iterative_refinement: matvec"); + op.a_multiply(1.0, Z[k], 0.0, V[k + 1]); + } // Modified Gram-Schmidt orthogonalization for (int j = 0; j <= k; ++j) { @@ -307,6 +324,7 @@ f_t iterative_refinement_gmres(T& op, } } // end Arnoldi loop + raft::common::nvtx::range update_scope("Barrier: iterative_refinement: solution update"); // Solve least squares H y = e // Back-substitution (H is (k+1)xk upper Hessenberg, cs/sin already applied) std::fill(y.begin(), y.end(), 0.0); diff --git a/cpp/src/dual_simplex/presolve.cpp b/cpp/src/dual_simplex/presolve.cpp index 9c55185c6d..6a5fca8565 100644 --- a/cpp/src/dual_simplex/presolve.cpp +++ b/cpp/src/dual_simplex/presolve.cpp @@ -13,6 +13,8 @@ #include #include +#include + #include #include #include @@ -683,6 +685,7 @@ void convert_user_problem(const user_problem_t& user_problem, std::vector& new_slacks, dualize_info_t& dualize_info) { + raft::common::nvtx::range fun_scope("QCQP: convert_user_problem"); constexpr bool verbose = false; if (verbose) { settings.log.printf("Converting problem with %d rows and %d columns and %d nonzeros\n", @@ -945,6 +948,7 @@ i_t presolve(const lp_problem_t& original, lp_problem_t& problem, presolve_info_t& presolve_info) { + raft::common::nvtx::range fun_scope("QCQP: presolve"); problem = original; const i_t linear_cols = linear_variable_count(problem); const bool has_cones = !problem.second_order_cone_dims.empty(); diff --git a/cpp/src/dual_simplex/scaling.cpp b/cpp/src/dual_simplex/scaling.cpp index ac96496ec6..e7597fff6a 100644 --- a/cpp/src/dual_simplex/scaling.cpp +++ b/cpp/src/dual_simplex/scaling.cpp @@ -8,6 +8,8 @@ #include #include +#include + #include namespace cuopt::mathematical_optimization::simplex { @@ -19,6 +21,7 @@ i_t scaling(const lp_problem_t& unscaled, std::vector& column_scaling, std::vector& row_scaling) { + raft::common::nvtx::range fun_scope("QCQP: scaling"); scaled = unscaled; i_t m = scaled.num_rows; i_t n = scaled.num_cols; diff --git a/cpp/src/pdlp/translate.hpp b/cpp/src/pdlp/translate.hpp index 3c6119e164..1b29ab080a 100644 --- a/cpp/src/pdlp/translate.hpp +++ b/cpp/src/pdlp/translate.hpp @@ -18,6 +18,8 @@ #include +#include + #include #include #include @@ -204,6 +206,7 @@ template static simplex::user_problem_t cuopt_optimization_problem_to_user_problem( raft::handle_t const* handle_ptr, optimization_problem_t& model) { + raft::common::nvtx::range fun_scope("QCQP: cuopt_optimization_problem_to_user_problem"); simplex::user_problem_t user_problem(handle_ptr); i_t const m = model.get_n_constraints(); From 0a37e5d545df8ba3ab802a1f4b46849321ca2125 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Tue, 14 Jul 2026 09:11:07 -0700 Subject: [PATCH 06/14] Update of parameters Signed-off-by: yuwenchen95 --- cpp/include/cuopt/mathematical_optimization/constants.h | 2 -- .../cuopt/mathematical_optimization/pdlp/solver_settings.hpp | 2 +- cpp/src/barrier/second_order_cone_kernels.cuh | 2 +- cpp/src/dual_simplex/simplex_solver_settings.hpp | 2 +- cpp/src/math_optimization/solver_settings.cu | 2 -- 5 files changed, 3 insertions(+), 7 deletions(-) diff --git a/cpp/include/cuopt/mathematical_optimization/constants.h b/cpp/include/cuopt/mathematical_optimization/constants.h index f08133d2f2..5e720a645e 100644 --- a/cpp/include/cuopt/mathematical_optimization/constants.h +++ b/cpp/include/cuopt/mathematical_optimization/constants.h @@ -48,8 +48,6 @@ #define CUOPT_ORDERING "ordering" #define CUOPT_BARRIER_DUAL_INITIAL_POINT "barrier_dual_initial_point" #define CUOPT_BARRIER_ITERATIVE_REFINEMENT "barrier_iterative_refinement" -#define CUOPT_BARRIER_CSR_IR_MATVEC "barrier_csr_ir_matvec" -#define CUOPT_BARRIER_SOC_THRESHOLD "barrier_soc_threshold" #define CUOPT_BARRIER_STEP_SCALE "barrier_step_scale" #define CUOPT_ELIMINATE_DENSE_COLUMNS "eliminate_dense_columns" #define CUOPT_CUDSS_DETERMINISTIC "cudss_deterministic" diff --git a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp index 6ba98f983f..c35fcc2551 100644 --- a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp +++ b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp @@ -284,7 +284,7 @@ class pdlp_solver_settings_t { pdlp_precision_t pdlp_precision{pdlp_precision_t::DefaultPrecision}; bool barrier_iterative_refinement{true}; bool barrier_csr_ir_matvec{false}; - i_t barrier_soc_threshold{5}; + i_t barrier_soc_threshold{100}; f_t barrier_step_scale{0.9}; bool save_best_primal_so_far{false}; /** diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index 04924a23ff..e108bf1f36 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -221,7 +221,7 @@ struct cone_data_t { raft::device_span x_in, raft::device_span z_in, rmm::cuda_stream_view stream, - i_t soc_threshold_in = 5) + i_t soc_threshold_in = 100) : n_cones(cone_dimensions_host.size()), n_cone_entries( std::reduce(cone_dimensions_host.begin(), cone_dimensions_host.end(), size_t{0})), diff --git a/cpp/src/dual_simplex/simplex_solver_settings.hpp b/cpp/src/dual_simplex/simplex_solver_settings.hpp index 18d1952331..0c5a743d56 100644 --- a/cpp/src/dual_simplex/simplex_solver_settings.hpp +++ b/cpp/src/dual_simplex/simplex_solver_settings.hpp @@ -70,7 +70,7 @@ struct simplex_solver_settings_t { barrier_iterative_refinement(true), barrier_csr_ir_matvec(false), barrier_step_scale(0.9), - barrier_soc_threshold(5), + barrier_soc_threshold(100), num_gpus(1), folding(-1), augmented(0), diff --git a/cpp/src/math_optimization/solver_settings.cu b/cpp/src/math_optimization/solver_settings.cu index f874533cc8..8bfd7cd8ba 100644 --- a/cpp/src/math_optimization/solver_settings.cu +++ b/cpp/src/math_optimization/solver_settings.cu @@ -132,7 +132,6 @@ solver_settings_t::solver_settings_t() : pdlp_settings(), mip_settings {CUOPT_DUALIZE, &pdlp_settings.dualize, -1, 1, -1}, {CUOPT_ORDERING, &pdlp_settings.ordering, -1, 1, -1}, {CUOPT_BARRIER_DUAL_INITIAL_POINT, &pdlp_settings.barrier_dual_initial_point, -1, 1, -1}, - {CUOPT_BARRIER_SOC_THRESHOLD, &pdlp_settings.barrier_soc_threshold, 0, 1000000, 5, "SOC dimension above which rank-2 sparse KKT expansion is used"}, {CUOPT_MIP_CUT_PASSES, &mip_settings.max_cut_passes, -1, std::numeric_limits::max(), 10}, {CUOPT_MIP_MIXED_INTEGER_ROUNDING_CUTS, &mip_settings.mir_cuts, -1, 1, -1}, {CUOPT_MIP_MIXED_INTEGER_GOMORY_CUTS, &mip_settings.mixed_integer_gomory_cuts, -1, 1, -1}, @@ -192,7 +191,6 @@ solver_settings_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}, // Diving heuristic hyper-parameters (hidden from default --help: name contains "hyper_") {CUOPT_MIP_HYPER_DIVING_SHOW_TYPE, &mip_settings.diving_params.show_type, false, "log diving heuristic type when it finds a new incumbent"}, From fd86c56f6422c443a9110dc4bf1233abf24fb15a Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Mon, 20 Jul 2026 06:10:13 -0700 Subject: [PATCH 07/14] optimize operations in compute_target_mu; clean up code Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 268 ++++++++++-------- cpp/src/barrier/barrier.hpp | 10 +- cpp/src/barrier/device_sparse_matrix.cuh | 54 ++++ cpp/src/barrier/second_order_cone_kernels.cuh | 34 ++- cpp/tests/socp/second_order_cone_kernels.cu | 25 +- 5 files changed, 233 insertions(+), 158 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 0227824ca3..9ec0e68926 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -129,30 +129,38 @@ template stream.value()); } -// step size computation for nonnegative and free variables +// Step size computation for nonnegative and free variables. Fuses two independent +// same-length reductions (e.g. (w, dw) and (v, dv), or (x, dx) and (z, dz)) into a +// single kernel launch + single host read instead of two. template -static f_t max_nonnegative_step_length_in_range( - transform_reduce_helper_t& transform_reduce_helper, - const rmm::device_uvector& x, - const rmm::device_uvector& dx, +static f2_t max_nonnegative_step_length_pair_in_range( + transform_reduce_pair_helper_t& transform_reduce_pair_helper, + const rmm::device_uvector& x1, + const rmm::device_uvector& dx1, + const rmm::device_uvector& x2, + const rmm::device_uvector& dx2, i_t len, const rmm::device_uvector& is_direct_free_linear, bool apply_direct_free_mask, rmm::cuda_stream_view stream) { - if (len <= 0) { return f_t(1); } - - return transform_reduce_helper.transform_reduce( - thrust::make_zip_iterator(dx.data(), x.data(), is_direct_free_linear.data()), - thrust::minimum{}, - [apply_direct_free_mask] HD(const thrust::tuple& t) { - const f_t dx_val = thrust::get<0>(t); - const f_t x_val = thrust::get<1>(t); - const i_t is_free = thrust::get<2>(t); - return (!(apply_direct_free_mask && is_free) && dx_val < f_t(0.0)) ? -x_val / dx_val - : f_t(1.0); + if (len <= 0) { return f2_t{f_t(1), f_t(1)}; } + + return transform_reduce_pair_helper.transform_reduce( + thrust::make_zip_iterator( + dx1.data(), x1.data(), dx2.data(), x2.data(), is_direct_free_linear.data()), + [apply_direct_free_mask] HD(const thrust::tuple& t) { + const f_t dx1_val = thrust::get<0>(t); + const f_t x1_val = thrust::get<1>(t); + const f_t dx2_val = thrust::get<2>(t); + const f_t x2_val = thrust::get<3>(t); + const i_t is_free = thrust::get<4>(t); + const bool masked = apply_direct_free_mask && is_free; + const f_t a = (!masked && dx1_val < f_t(0.0)) ? -x1_val / dx1_val : f_t(1.0); + const f_t b = (!masked && dx2_val < f_t(0.0)) ? -x2_val / dx2_val : f_t(1.0); + return f2_t{a, b}; }, - f_t(1.0), + f2_t{f_t(1.0), f_t(1.0)}, len, stream); } @@ -320,10 +328,10 @@ class iteration_data_t { d_dv_(0, lp.handle_ptr->get_stream()), d_dw_(0, lp.handle_ptr->get_stream()), d_dw_aff_(0, lp.handle_ptr->get_stream()), - d_dx_aff_(0, lp.handle_ptr->get_stream()), + d_dx_aff_(lp.num_cols, lp.handle_ptr->get_stream()), d_dv_aff_(0, lp.handle_ptr->get_stream()), - d_dz_aff_(0, lp.handle_ptr->get_stream()), - d_dy_aff_(0, lp.handle_ptr->get_stream()), + d_dz_aff_(lp.num_cols, lp.handle_ptr->get_stream()), + d_dy_aff_(lp.num_rows, lp.handle_ptr->get_stream()), d_primal_residual_(0, lp.handle_ptr->get_stream()), d_dual_residual_(lp.num_cols, lp.handle_ptr->get_stream()), d_bound_residual_(0, lp.handle_ptr->get_stream()), @@ -345,6 +353,7 @@ class iteration_data_t { restrict_u_(0), d_restrict_u_(0, lp.handle_ptr->get_stream()), transform_reduce_helper_(lp.handle_ptr->get_stream()), + transform_reduce_pair_helper_(lp.handle_ptr->get_stream()), sum_reduce_helper_(lp.handle_ptr->get_stream()), indefinite_Q(false), Q_diagonal(false), @@ -466,6 +475,9 @@ class iteration_data_t { if (n_upper_bounds > 0) { settings.log.printf("Upper bounds : %d\n", n_upper_bounds); } + + d_dw_aff_.resize(n_upper_bounds, stream_view_); + d_dv_aff_.resize(n_upper_bounds, stream_view_); } std::vector dense_columns_unordered; @@ -2436,6 +2448,7 @@ class iteration_data_t { rmm::device_uvector d_restrict_u_; transform_reduce_helper_t transform_reduce_helper_; + transform_reduce_pair_helper_t transform_reduce_pair_helper_; sum_reduce_helper_t sum_reduce_helper_; bool cone_combined_step_; @@ -2903,13 +2916,10 @@ void barrier_solver_t::gpu_compute_residuals(const rmm::device_uvector cuda::std::minus<>{}, stream_view_.value()); RAFT_CHECK_CUDA(stream_view_); - if (data.Q.n > 0) { - auto descr_dual_residual = data.cusparse_view_.create_vector(data.d_dual_residual_); - data.cusparse_Q_view_.spmv(1.0, cusparse_d_x, 1.0, descr_dual_residual); - } - // Compute dual_residual = c - A'*y - z + E*v - auto cusparse_d_y = data.cusparse_view_.create_vector(d_y); auto descr_dual_residual = data.cusparse_view_.create_vector(data.d_dual_residual_); + if (data.Q.n > 0) { data.cusparse_Q_view_.spmv(1.0, cusparse_d_x, 1.0, descr_dual_residual); } + // Compute dual_residual = c - A'*y - z + E*v + auto cusparse_d_y = data.cusparse_view_.create_vector(d_y); data.cusparse_view_.transpose_spmv(-1.0, cusparse_d_y, 1.0, descr_dual_residual); if (data.n_upper_bounds > 0) { @@ -2981,21 +2991,30 @@ void barrier_solver_t::gpu_compute_residual_norms(const rmm::device_uv } template -f_t barrier_solver_t::compute_nonnegative_step_length(iteration_data_t& data, - const rmm::device_uvector& x, - const rmm::device_uvector& dx) +std::pair barrier_solver_t::compute_nonnegative_step_length_pair( + iteration_data_t& data, + const rmm::device_uvector& x1, + const rmm::device_uvector& dx1, + const rmm::device_uvector& x2, + const rmm::device_uvector& dx2) { - const bool has_soc = data.has_cones() && static_cast(x.size()) >= data.cone_end(); + assert(x1.size() == x2.size()); + + const bool has_soc = data.has_cones() && static_cast(x1.size()) >= data.cone_end(); // SOCP layout is [linear | cone]; stop at cone_start() - const i_t linear_len = has_soc ? data.cone_start() : static_cast(x.size()); - return max_nonnegative_step_length_in_range(data.transform_reduce_helper_, - x, - dx, + const i_t linear_len = has_soc ? data.cone_start() : static_cast(x1.size()); + const f2_t result = + max_nonnegative_step_length_pair_in_range(data.transform_reduce_pair_helper_, + x1, + dx1, + x2, + dx2, linear_len, data.d_is_direct_free_linear_, - static_cast(x.size()) == lp.num_cols, + static_cast(x1.size()) == lp.num_cols, stream_view_); + return {result.a, result.b}; } template @@ -3079,12 +3098,6 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t 0) { cub::DeviceTransform::Transform( @@ -3101,6 +3114,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t 0 && data.Q_diagonal) { cub::DeviceTransform::Transform( cuda::std::make_tuple(data.d_Q_diag_.data(), data.d_diag_.data()), @@ -3109,10 +3123,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t 0 && data.Q_diagonal) { cub::DeviceTransform::Transform( cuda::std::make_tuple( data.d_diag_.data(), data.d_is_direct_free_linear_.data(), data.d_Q_diag_.data()), @@ -3134,10 +3145,8 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t::gpu_compute_search_direction(iteration_data_t::gpu_compute_search_direction(iteration_data_t::compute_target_mu( raft::common::nvtx::range fun_scope("Barrier: compute_target_mu"); const bool has_soc = data.has_cones(); - f_t complementarity_aff_sum = 0.0; - // TMP no copy and data should always be on the GPU - f_t step_primal_aff = std::min(compute_nonnegative_step_length(data, data.d_w_, data.d_dw_aff_), - compute_nonnegative_step_length(data, data.d_x_, data.d_dx_aff_)); - f_t step_dual_aff = std::min(compute_nonnegative_step_length(data, data.d_v_, data.d_dv_aff_), - compute_nonnegative_step_length(data, data.d_z_, data.d_dz_aff_)); + const auto [primal_w, dual_v] = + compute_nonnegative_step_length_pair(data, data.d_w_, data.d_dw_, data.d_v_, data.d_dv_); + const auto [primal_x, dual_z] = + compute_nonnegative_step_length_pair(data, data.d_x_, data.d_dx_, data.d_z_, data.d_dz_); + f_t step_primal_aff = std::min(primal_w, primal_x); + f_t step_dual_aff = std::min(dual_v, dual_z); if (has_soc) { i_t cs = data.cone_start(); i_t mc = data.cone_entry_count(); - auto [cone_p, cone_d] = + const f_t cone_combined = compute_cone_step_length(data.cones(), - raft::device_span(data.d_dx_aff_.data() + cs, mc), - raft::device_span(data.d_dz_aff_.data() + cs, mc), + raft::device_span(data.d_dx_.data() + cs, mc), + raft::device_span(data.d_dz_.data() + cs, mc), f_t(1), stream_view_); - step_primal_aff = std::min(step_primal_aff, cone_p); - step_dual_aff = std::min(step_dual_aff, cone_d); + step_primal_aff = std::min(step_primal_aff, cone_combined); + step_dual_aff = std::min(step_dual_aff, cone_combined); } if (data.Q.n > 0 || has_soc) { step_primal_aff = step_dual_aff = std::min(step_primal_aff, step_dual_aff); } - // Compute complementarity_xz_aff_sum = sum(x_aff * z_aff), - // where x_aff = x + step_primal_aff * dx_aff and z_aff = z + step_dual_aff * dz_aff - // Here the update of x_aff and z_aff are done temporarily and sum of their products is - // computed without storing intermediate results. + // Compute complementarity_xz_aff_sum and complementarity_wv_aff_sum. Save the affine direction as + // a side effect. + raft::device_span x_span(data.d_x_.data(), data.d_x_.size()); + raft::device_span z_span(data.d_z_.data(), data.d_z_.size()); + raft::device_span dx_span(data.d_dx_.data(), data.d_dx_.size()); + raft::device_span dz_span(data.d_dz_.data(), data.d_dz_.size()); + raft::device_span dx_aff_span(data.d_dx_aff_.data(), data.d_dx_aff_.size()); + raft::device_span dz_aff_span(data.d_dz_aff_.data(), data.d_dz_aff_.size()); + f_t complementarity_xz_aff_sum = data.transform_reduce_helper_.transform_reduce( - thrust::make_zip_iterator( - data.d_x_.data(), data.d_z_.data(), data.d_dx_aff_.data(), data.d_dz_aff_.data()), + thrust::make_counting_iterator(0), cuda::std::plus{}, - [step_primal_aff, step_dual_aff] HD(const thrust::tuple t) { - const f_t x = thrust::get<0>(t); - const f_t z = thrust::get<1>(t); - const f_t dx_aff = thrust::get<2>(t); - const f_t dz_aff = thrust::get<3>(t); + [step_primal_aff, step_dual_aff, x_span, z_span, dx_span, dz_span, dx_aff_span, dz_aff_span] HD( + size_t idx) { + const f_t dx = dx_span[idx]; + const f_t dz = dz_span[idx]; + + dx_aff_span[idx] = dx; + dz_aff_span[idx] = dz; - const f_t x_aff = x + step_primal_aff * dx_aff; - const f_t z_aff = z + step_dual_aff * dz_aff; + const f_t x_aff = x_span[idx] + step_primal_aff * dx; + const f_t z_aff = z_span[idx] + step_dual_aff * dz; const f_t complementarity_xz_aff = x_aff * z_aff; @@ -3934,20 +3949,26 @@ void barrier_solver_t::compute_target_mu( data.d_x_.size(), stream_view_); - // Here the update of w_aff and v_aff are done temporarily and sum of their products is - // computed without storing intermediate results. + raft::device_span w_span(data.d_w_.data(), data.d_w_.size()); + raft::device_span v_span(data.d_v_.data(), data.d_v_.size()); + raft::device_span dw_span(data.d_dw_.data(), data.d_dw_.size()); + raft::device_span dv_span(data.d_dv_.data(), data.d_dv_.size()); + raft::device_span dw_aff_span(data.d_dw_aff_.data(), data.d_dw_aff_.size()); + raft::device_span dv_aff_span(data.d_dv_aff_.data(), data.d_dv_aff_.size()); + f_t complementarity_wv_aff_sum = data.transform_reduce_helper_.transform_reduce( - thrust::make_zip_iterator( - data.d_w_.data(), data.d_v_.data(), data.d_dw_aff_.data(), data.d_dv_aff_.data()), + thrust::make_counting_iterator(0), cuda::std::plus{}, - [step_primal_aff, step_dual_aff] HD(const thrust::tuple t) { - const f_t w = thrust::get<0>(t); - const f_t v = thrust::get<1>(t); - const f_t dw_aff = thrust::get<2>(t); - const f_t dv_aff = thrust::get<3>(t); + [step_primal_aff, step_dual_aff, w_span, v_span, dw_span, dv_span, dw_aff_span, dv_aff_span] HD( + size_t idx) { + const f_t dw = dw_span[idx]; + const f_t dv = dv_span[idx]; + + dw_aff_span[idx] = dw; + dv_aff_span[idx] = dv; - const f_t w_aff = w + step_primal_aff * dw_aff; - const f_t v_aff = v + step_dual_aff * dv_aff; + const f_t w_aff = w_span[idx] + step_primal_aff * dw; + const f_t v_aff = v_span[idx] + step_dual_aff * dv; const f_t complementarity_wv_aff = w_aff * v_aff; @@ -3957,11 +3978,14 @@ void barrier_solver_t::compute_target_mu( data.d_w_.size(), stream_view_); - complementarity_aff_sum = complementarity_xz_aff_sum + complementarity_wv_aff_sum; - const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); - mu_aff = complementarity_aff_sum / mu_denom; - sigma = std::max(0.0, std::min(1.0, std::pow(mu_aff / mu, 3.0))); - new_mu = sigma * mu_aff; + // Sum the complementarity terms and save the affine direction. + f_t complementarity_aff_sum = complementarity_xz_aff_sum + complementarity_wv_aff_sum; + raft::copy(data.d_dy_aff_.data(), data.d_dy_.data(), data.d_dy_.size(), stream_view_); + + const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); + mu_aff = complementarity_aff_sum / mu_denom; + sigma = std::max(0.0, std::min(1.0, std::pow(mu_aff / mu, 3.0))); + new_mu = sigma * mu_aff; } template @@ -4002,12 +4026,15 @@ void barrier_solver_t::compute_cc_rhs(iteration_data_t& data stream_view_.value()); RAFT_CHECK_CUDA(stream_view_); // Zero the corrector RHS on device - thrust::fill(rmm::exec_policy(stream_view_), data.d_h_.begin(), data.d_h_.end(), f_t(0)); - thrust::fill( - rmm::exec_policy(stream_view_), data.d_dual_rhs_.begin(), data.d_dual_rhs_.end(), f_t(0)); - thrust::fill( - rmm::exec_policy(stream_view_), data.d_bound_rhs_.begin(), data.d_bound_rhs_.end(), f_t(0)); - thrust::fill(rmm::exec_policy(stream_view_), data.d_dw_.begin(), data.d_dw_.end(), f_t(0)); + RAFT_CUDA_TRY(cudaMemsetAsync(data.d_h_.data(), 0, sizeof(f_t) * data.d_h_.size(), stream_view_)); + RAFT_CUDA_TRY(cudaMemsetAsync( + data.d_dual_rhs_.data(), 0, sizeof(f_t) * data.d_dual_rhs_.size(), stream_view_)); + if (data.n_upper_bounds > 0) { + RAFT_CUDA_TRY(cudaMemsetAsync( + data.d_bound_rhs_.data(), 0, sizeof(f_t) * data.d_bound_rhs_.size(), stream_view_)); + RAFT_CUDA_TRY( + cudaMemsetAsync(data.d_dw_.data(), 0, sizeof(f_t) * data.d_dw_.size(), stream_view_)); + } data.cone_combined_step_ = has_soc; data.cone_sigma_mu_ = has_soc ? new_mu : f_t(0); } @@ -4071,22 +4098,25 @@ void barrier_solver_t::compute_primal_dual_step_length(iteration_data_ f_t max_step_primal = 0.0; f_t max_step_dual = 0.0; - max_step_primal = std::min(compute_nonnegative_step_length(data, data.d_w_, data.d_dw_), - compute_nonnegative_step_length(data, data.d_x_, data.d_dx_)); - max_step_dual = std::min(compute_nonnegative_step_length(data, data.d_v_, data.d_dv_), - compute_nonnegative_step_length(data, data.d_z_, data.d_dz_)); + + const auto [primal_w, dual_v] = + compute_nonnegative_step_length_pair(data, data.d_w_, data.d_dw_, data.d_v_, data.d_dv_); + const auto [primal_x, dual_z] = + compute_nonnegative_step_length_pair(data, data.d_x_, data.d_dx_, data.d_z_, data.d_dz_); + max_step_primal = std::min(primal_w, primal_x); + max_step_dual = std::min(dual_v, dual_z); if (has_soc) { i_t cs = data.cone_start(); i_t mc = data.cone_entry_count(); - auto [cone_primal, cone_dual] = + const f_t cone_combined = compute_cone_step_length(data.cones(), raft::device_span(data.d_dx_.data() + cs, mc), raft::device_span(data.d_dz_.data() + cs, mc), f_t(1), stream_view_); - max_step_primal = std::min(max_step_primal, cone_primal); - max_step_dual = std::min(max_step_dual, cone_dual); + max_step_primal = std::min(max_step_primal, cone_combined); + max_step_dual = std::min(max_step_dual, cone_combined); } step_primal = step_scale * max_step_primal; @@ -4631,20 +4661,6 @@ lp_status_t barrier_solver_t::solve(f_t start_time, lp_solution_t #include + +#include 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& data, - const rmm::device_uvector& x, - const rmm::device_uvector& dx); + std::pair compute_nonnegative_step_length_pair(iteration_data_t& data, + const rmm::device_uvector& x1, + const rmm::device_uvector& dx1, + const rmm::device_uvector& x2, + const rmm::device_uvector& dx2); i_t gpu_compute_search_direction(iteration_data_t& data, f_t& dual_perturb, f_t& primal_perturb, diff --git a/cpp/src/barrier/device_sparse_matrix.cuh b/cpp/src/barrier/device_sparse_matrix.cuh index 90df0c86cf..71e86484a6 100644 --- a/cpp/src/barrier/device_sparse_matrix.cuh +++ b/cpp/src/barrier/device_sparse_matrix.cuh @@ -88,6 +88,60 @@ struct transform_reduce_helper_t { } }; +template +struct f2_t { + f_t a; + f_t b; +}; + +template +struct f2_min_t { + HD f2_t operator()(const f2_t& lhs, const f2_t& rhs) const + { + return f2_t{cuda::std::min(lhs.a, rhs.a), cuda::std::min(lhs.b, rhs.b)}; + } +}; + +template +struct transform_reduce_pair_helper_t { + rmm::device_buffer buffer_data; + rmm::device_scalar> 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{a, b} pair; the two + // components are reduced independently (elementwise min) in a single kernel launch. + template + f2_t transform_reduce(InputIteratorT input, + TransformOpT transform_op, + f2_t init, + i_t size, + rmm::cuda_stream_view stream_view) + { + f2_min_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 struct csc_view_t { raft::device_span col_start; diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index e108bf1f36..bee818618a 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -26,6 +26,7 @@ #include #include #include +#include #include #include @@ -1495,7 +1496,7 @@ void scatter_dense_hessian_into_augmented(const cone_data_t& cones, } /** - * Compute the maximum primal and dual step lengths that keep SOC blocks + * Compute the combined (primal and dual) maximum step length that keeps SOC blocks * feasible: * * x + alpha dx in Q, z + alpha dz in Q, alpha <= alpha_max. @@ -1513,11 +1514,11 @@ void scatter_dense_hessian_into_augmented(const cone_data_t& cones, * crossing, and the final reductions take the global minimum over cones. */ template -std::pair compute_cone_step_length(cone_data_t& cones, - raft::device_span dx, - raft::device_span dz, - f_t alpha_max, - rmm::cuda_stream_view stream) +f_t compute_cone_step_length(cone_data_t& cones, + raft::device_span dx, + raft::device_span dz, + f_t alpha_max, + rmm::cuda_stream_view stream) { auto cone_offsets = cuopt::make_span(cones.cone_offsets); auto element_cone_ids = cuopt::make_span(cones.element_cone_ids); @@ -1565,18 +1566,15 @@ std::pair compute_cone_step_length(cone_data_t& cones, run_pass(cones.x, dx, alpha_primal); run_pass(cones.z, dz, alpha_dual); - const f_t primal = thrust::reduce(rmm::exec_policy(stream), - alpha_primal.begin(), - alpha_primal.end(), - alpha_max, - thrust::minimum()); - const f_t dual = thrust::reduce(rmm::exec_policy(stream), - alpha_dual.begin(), - alpha_dual.end(), - alpha_max, - thrust::minimum()); - - return {primal, dual}; + return thrust::transform_reduce( + rmm::exec_policy(stream), + thrust::make_zip_iterator(alpha_primal.begin(), alpha_dual.begin()), + thrust::make_zip_iterator(alpha_primal.end(), alpha_dual.end()), + [] HD(const thrust::tuple& t) -> f_t { + return cuda::std::min(thrust::get<0>(t), thrust::get<1>(t)); + }, + alpha_max, + thrust::minimum()); } /** diff --git a/cpp/tests/socp/second_order_cone_kernels.cu b/cpp/tests/socp/second_order_cone_kernels.cu index 9cac42d2b9..334c099519 100644 --- a/cpp/tests/socp/second_order_cone_kernels.cu +++ b/cpp/tests/socp/second_order_cone_kernels.cu @@ -354,17 +354,15 @@ TEST(second_order_cone_kernels, cone_step_length_many_small_one_sparse_medium_co launch_nt_scaling(cones, stream); launch_update_scaling_sparse(cones, stream); - const auto [step_primal, step_dual] = + const auto step_combined = compute_cone_step_length(cones, raft::device_span(dx.data(), dx.size()), raft::device_span(dz.data(), dz.size()), 1.0, stream); - EXPECT_GT(step_primal, 0.0); - EXPECT_GT(step_dual, 0.0); - EXPECT_LE(step_primal, 1.0); - EXPECT_LE(step_dual, 1.0); + EXPECT_GT(step_combined, 0.0); + EXPECT_LE(step_combined, 1.0); } TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) @@ -439,7 +437,7 @@ TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) auto dz = cuopt::device_copy(dz_host, stream); cone_data_t cones(cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream); - const auto [step_primal, step_dual] = + const auto step_combined = compute_cone_step_length(cones, raft::device_span(dx.data(), dx.size()), raft::device_span(dz.data(), dz.size()), @@ -448,14 +446,19 @@ TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) const auto primal_per_cone = cuopt::host_copy(cones.scratch.step_alpha_primal, stream); const auto dual_per_cone = cuopt::host_copy(cones.scratch.step_alpha_dual, stream); - EXPECT_NEAR( - step_primal, *std::min_element(primal_per_cone.begin(), primal_per_cone.end()), 1e-12); - EXPECT_NEAR(step_dual, *std::min_element(dual_per_cone.begin(), dual_per_cone.end()), 1e-12); + EXPECT_NEAR(step_combined, + std::min(*std::min_element(primal_per_cone.begin(), primal_per_cone.end()), + *std::min_element(dual_per_cone.begin(), dual_per_cone.end())), + 1e-12); + // Only the combined min(primal, dual) per cone is consumed by callers now: use it for + // both the x/dx and z/dz feasibility checks below, which is still valid since it is + // <= each side's own feasibility boundary. for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { + const double combined_cone = std::min(primal_per_cone[cone], dual_per_cone[cone]); EXPECT_GT(primal_per_cone[cone], 0.0) << "primal cone " << cone; EXPECT_GT(dual_per_cone[cone], 0.0) << "dual cone " << cone; - expect_cone_feasible_after_step(x_host, dx_host, primal_per_cone[cone], cone, "primal"); - expect_cone_feasible_after_step(z_host, dz_host, dual_per_cone[cone], cone, "dual"); + expect_cone_feasible_after_step(x_host, dx_host, combined_cone, cone, "primal"); + expect_cone_feasible_after_step(z_host, dz_host, combined_cone, cone, "dual"); } } From 14d510cc29337dde3786a2123dfbf602e21f7883 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Tue, 21 Jul 2026 01:42:21 -0700 Subject: [PATCH 08/14] Fuse the sparse SOC augmented-system assembly and unify the bucket index device function Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 32 +-- cpp/src/barrier/second_order_cone_kernels.cuh | 271 +++++++----------- cpp/tests/socp/sparse_augmented_kkt_test.cu | 242 ++++++++-------- 3 files changed, 231 insertions(+), 314 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 9ec0e68926..70c182571d 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -1185,25 +1185,18 @@ class iteration_data_t { if (has_soc) { if (cones().has_sparse_cones()) { - launch_get_Hs_sparse(cones(), d_sparse_Hs_diag_, handle_ptr->get_stream()); - scatter_sparse_hessian_diag_into_augmented(device_augmented.x, - d_sparse_hessian_diag_, - d_sparse_Hs_diag_, - d_sparse_hessian_Q_, - handle_ptr->get_stream(), - dual_perturb); - update_sparse_expansion_in_augmented(device_augmented.x, - d_sparse_exp_v_col_, - d_sparse_exp_u_col_, - d_sparse_exp_v_row_, - d_sparse_exp_u_row_, - d_sparse_expansion_D_, - cones().sparse_v, - cones().sparse_u, - cones().eta, - cones().sparse_cone_ids, - handle_ptr->get_stream(), - dual_perturb); + scatter_sparse_hessian_into_augmented(cones(), + device_augmented.x, + d_sparse_Hs_diag_, + d_sparse_hessian_diag_, + d_sparse_hessian_Q_, + d_sparse_exp_v_col_, + d_sparse_exp_u_col_, + d_sparse_exp_v_row_, + d_sparse_exp_u_row_, + d_sparse_expansion_D_, + handle_ptr->get_stream(), + dual_perturb); RAFT_CHECK_CUDA(handle_ptr->get_stream()); } if (cones().n_dense_cones() > 0) { @@ -3164,7 +3157,6 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t x, apply_hessian(x, out, cones, stream, 1, out_input, 1, true); } -// One block per sparse cone. Write the diagonal of the implicit cone Hessian: -// eta^2 on every entry, with the head entry scaled by the rank-2 corner d. -template -__global__ void get_Hs_sparse_kernel(raft::device_span Hs_diag, - raft::device_span eta, - raft::device_span d, - raft::device_span sparse_cone_ids, - raft::device_span sparse_entry_offsets, - raft::device_span sparse_cone_dims, - i_t n_sparse_cones) +// Bucket index owning `entry` via upper-bound search on the exclusive prefix ends in +// offsets[1..n_buckets]. +template +__device__ i_t bucket_index(raft::device_span offsets, + offset_t entry, + i_t n_buckets) { - const i_t sparse_idx = blockIdx.x; - if (sparse_idx >= n_sparse_cones) { return; } - - const i_t block_start = sparse_entry_offsets[sparse_idx]; - const i_t block_dim = sparse_cone_dims[sparse_idx]; - const i_t cone_idx = sparse_cone_ids[sparse_idx]; - const f_t eta_sq = eta[cone_idx] * eta[cone_idx]; - - for (i_t j = threadIdx.x; j < block_dim; j += blockDim.x) { - Hs_diag[block_start + j] = eta_sq; + i_t lo = 0; + i_t hi = n_buckets; + while (lo < hi) { + const i_t mid = lo + (hi - lo) / 2; + if (offsets[mid + 1] <= entry) { + lo = mid + 1; + } else { + hi = mid; + } } - __syncthreads(); - - if (threadIdx.x == 0) { Hs_diag[block_start] *= d[sparse_idx]; } -} - -/** - * Compute the packed diagonal `Hs_diag` of the implicit sparse cone Hessians, - * one contiguous block of `sparse_cone_dims[s]` entries per sparse cone. This is - * the diagonal that gets scattered onto the cone rows of the augmented system. - */ -template -void launch_get_Hs_sparse(cone_data_t& cones, - rmm::device_uvector& Hs_diag, - rmm::cuda_stream_view stream) -{ - if (!cones.has_sparse_cones()) { return; } - - const i_t n_sparse = cones.n_sparse_cones; - cuopt_assert(Hs_diag.size() == cones.n_sparse_cone_entries, "Hs_diag size mismatch"); - - get_Hs_sparse_kernel - <<>>(cuopt::make_span(Hs_diag), - cuopt::make_span(cones.eta), - cuopt::make_span(cones.d), - cuopt::make_span(cones.sparse_cone_ids), - cuopt::make_span(cones.sparse_entry_offsets), - cuopt::make_span(cones.sparse_cone_dims), - n_sparse); - RAFT_CUDA_TRY(cudaPeekAtLastError()); + return lo; } -// Scatter one sparse-cone diagonal entry into the augmented value buffer: -// augmented_x[csr_indices[e]] = -Hs_diag[e] - q_values[e] - dual_perturb, -// matching the sign convention of the dense scatter (negated KKT block plus -// any quadratic-objective contribution and diagonal regularization). template -__global__ void __launch_bounds__(soc_block_size) - scatter_sparse_hessian_diag_kernel(raft::device_span augmented_x, - raft::device_span csr_indices, - raft::device_span Hs_diag, - raft::device_span q_values, - f_t dual_perturb) +__global__ void __launch_bounds__(soc_block_size) scatter_sparse_hessian_into_augmented_kernel( + raft::device_span augmented_x, + raft::device_span Hs_diag, + raft::device_span eta, + raft::device_span d, + raft::device_span sparse_cone_ids, + raft::device_span sparse_entry_offsets, + i_t n_sparse_cones, + raft::device_span hessian_diag_csr_indices, + raft::device_span q_values, + raft::device_span sparse_v, + raft::device_span sparse_u, + raft::device_span exp_v_col, + raft::device_span exp_u_col, + raft::device_span exp_v_row, + raft::device_span exp_u_row, + raft::device_span sparse_expansion_D, + f_t dual_perturb) { const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (e >= csr_indices.size()) { return; } - - augmented_x[csr_indices[e]] = -Hs_diag[e] - q_values[e] - dual_perturb; + if (e >= Hs_diag.size()) { return; } + + const i_t e_i = static_cast(e); // single cast, reused below + const i_t sparse_idx = bucket_index(sparse_entry_offsets, e_i, n_sparse_cones); + const i_t cone_idx = sparse_cone_ids[sparse_idx]; + const bool is_head = e_i == sparse_entry_offsets[sparse_idx]; + const f_t eta_val = eta[cone_idx]; + const f_t eta_sq = eta_val * eta_val; + const f_t hs_val = is_head ? eta_sq * d[sparse_idx] : eta_sq; + Hs_diag[e] = hs_val; + + augmented_x[hessian_diag_csr_indices[e]] = -hs_val - q_values[e] - dual_perturb; + + const f_t v_val = sparse_v[e]; + const f_t u_val = sparse_u[e]; + augmented_x[exp_v_col[e]] = v_val; + augmented_x[exp_u_col[e]] = u_val; + augmented_x[exp_v_row[e]] = v_val; + augmented_x[exp_u_row[e]] = u_val; + + if (is_head) { + const f_t perturbed_eta_sq = eta_sq + dual_perturb; + augmented_x[sparse_expansion_D[2 * sparse_idx]] = -perturbed_eta_sq; + augmented_x[sparse_expansion_D[2 * sparse_idx + 1]] = perturbed_eta_sq; + } } /** - * Write the sparse cone Hessian diagonals into the augmented system value buffer. + * Fused sparse SOC augmented-system assembly: one kernel launch with one thread per packed + * sparse entry (E threads total), writing `Hs_diag`, the Hessian CSR scatter, and the four + * rank-2 couplings, with each cone's unique head thread additionally applying the rank-2 + * corner scale and writing that cone's two expansion diagonals. * - * `csr_indices[e]` gives the value-buffer position of the diagonal for packed - * cone entry `e`; each is set to `-Hs_diag[e] - q_values[e] - dual_perturb`. + * `Hs_diag` is left populated for downstream matrix-free matvec / iterative refinement use. */ template -void scatter_sparse_hessian_diag_into_augmented(rmm::device_uvector& augmented_x, - const rmm::device_uvector& csr_indices, - const rmm::device_uvector& Hs_diag, - const rmm::device_uvector& q_values, - rmm::cuda_stream_view stream, - f_t dual_perturb) +void scatter_sparse_hessian_into_augmented(cone_data_t& cones, + rmm::device_uvector& augmented_x, + rmm::device_uvector& Hs_diag, + const rmm::device_uvector& hessian_diag_csr_indices, + const rmm::device_uvector& q_values, + const rmm::device_uvector& exp_v_col, + const rmm::device_uvector& exp_u_col, + const rmm::device_uvector& exp_v_row, + const rmm::device_uvector& exp_u_row, + const rmm::device_uvector& sparse_expansion_D, + rmm::cuda_stream_view stream, + f_t dual_perturb) { - const size_t count = csr_indices.size(); - if (count == 0) { return; } - cuopt_assert(count == Hs_diag.size(), "sparse Hessian index/value size mismatch"); - cuopt_assert(count == q_values.size(), "sparse Hessian Q-value size mismatch"); + if (!cones.has_sparse_cones()) { return; } - const size_t grid = raft::ceildiv(count, soc_block_size); - scatter_sparse_hessian_diag_kernel - <<>>(cuopt::make_span(augmented_x), - cuopt::make_span(csr_indices), - cuopt::make_span(Hs_diag), - cuopt::make_span(q_values), - dual_perturb); + const i_t n_sparse = cones.n_sparse_cones; + const size_t E = cones.n_sparse_cone_entries; + const size_t entry_grid = raft::ceildiv(E, soc_block_size); + scatter_sparse_hessian_into_augmented_kernel + <<>>( + cuopt::make_span(augmented_x), + cuopt::make_span(Hs_diag), + cuopt::make_span(cones.eta), + cuopt::make_span(cones.d), + cuopt::make_span(cones.sparse_cone_ids), + cuopt::make_span(cones.sparse_entry_offsets), + n_sparse, + cuopt::make_span(hessian_diag_csr_indices), + cuopt::make_span(q_values), + cuopt::make_span(cones.sparse_v), + cuopt::make_span(cones.sparse_u), + cuopt::make_span(exp_v_col), + cuopt::make_span(exp_u_col), + cuopt::make_span(exp_v_row), + cuopt::make_span(exp_u_row), + cuopt::make_span(sparse_expansion_D), + dual_perturb); RAFT_CUDA_TRY(cudaPeekAtLastError()); } -/** - * Refresh the expansion (rank-2 coupling) entries of the augmented system. - * - * For each sparse cone this scatters the rank-2 vectors v, u into their four - * coupling positions -- the v/u columns and the symmetric v/u rows linking the - * cone variables to the two expansion variables -- and writes the expansion - * diagonal +/-(eta^2 + dual_perturb). The Hessian diagonal itself is written - * separately by `scatter_sparse_hessian_diag_into_augmented`. - */ -template -void update_sparse_expansion_in_augmented(rmm::device_uvector& augmented_x, - const rmm::device_uvector& sparse_exp_v_col, - const rmm::device_uvector& sparse_exp_u_col, - const rmm::device_uvector& sparse_exp_v_row, - const rmm::device_uvector& sparse_exp_u_row, - const rmm::device_uvector& sparse_expansion_D, - const rmm::device_uvector& sparse_v, - const rmm::device_uvector& sparse_u, - const rmm::device_uvector& eta, - const rmm::device_uvector& sparse_cone_ids, - rmm::cuda_stream_view stream, - f_t dual_perturb) -{ - const size_t n_coupling = sparse_exp_v_col.size(); - if (n_coupling == 0) { return; } - - cuopt_assert(sparse_exp_u_col.size() == n_coupling, "sparse_exp_u_col size mismatch"); - cuopt_assert(sparse_exp_v_row.size() == n_coupling, "sparse_exp_v_row size mismatch"); - cuopt_assert(sparse_exp_u_row.size() == n_coupling, "sparse_exp_u_row size mismatch"); - cuopt_assert(sparse_v.size() == n_coupling, "sparse_v size mismatch"); - cuopt_assert(sparse_u.size() == n_coupling, "sparse_u size mismatch"); - - thrust::scatter(rmm::exec_policy(stream), - sparse_v.begin(), - sparse_v.end(), - sparse_exp_v_col.begin(), - augmented_x.begin()); - thrust::scatter(rmm::exec_policy(stream), - sparse_u.begin(), - sparse_u.end(), - sparse_exp_u_col.begin(), - augmented_x.begin()); - thrust::scatter(rmm::exec_policy(stream), - sparse_v.begin(), - sparse_v.end(), - sparse_exp_v_row.begin(), - augmented_x.begin()); - thrust::scatter(rmm::exec_policy(stream), - sparse_u.begin(), - sparse_u.end(), - sparse_exp_u_row.begin(), - augmented_x.begin()); - - const i_t n_expansion = static_cast(sparse_expansion_D.size()); - if (n_expansion > 0) { - auto eta_span = cuopt::make_span(eta); - auto sparse_cone_ids_span = cuopt::make_span(sparse_cone_ids); - // Expansion-row diagonal of sparse cone i/2: -(eta^2 + dual_perturb) for the - // v-row (even i), +(eta^2 + dual_perturb) for the u-row (odd i). - auto diag_values = thrust::make_transform_iterator( - thrust::counting_iterator(0), - [eta_span, sparse_cone_ids_span, dual_perturb] HD(i_t i) -> f_t { - const i_t sparse_idx = i / 2; - const f_t eta_val = eta_span[sparse_cone_ids_span[sparse_idx]]; - const f_t sign = (i % 2 == 0) ? f_t(-1) : f_t(1); - return sign * (eta_val * eta_val + dual_perturb); - }); - thrust::scatter(rmm::exec_policy(stream), - diag_values, - diag_values + n_expansion, - sparse_expansion_D.begin(), - augmented_x.begin()); - } -} - -// Remove factorization-only regularization from the assembled augmented CSR values so -// iterative refinement can use a single SpMV against the true KKT operator. template __global__ void strip_linear_primal_diag_kernel(raft::device_span augmented_x, raft::device_span diag_indices, @@ -1427,18 +1367,7 @@ __global__ void __launch_bounds__(soc_block_size) const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (e >= csr_indices.size()) { return; } - i_t lo = 0; - i_t hi = n_dense_cones; - while (lo < hi) { - const i_t mid = lo + (hi - lo) / 2; - if (dense_block_offsets[mid + 1] <= e) { - lo = mid + 1; - } else { - hi = mid; - } - } - - const i_t dense_idx = lo; + const i_t dense_idx = bucket_index(dense_block_offsets, e, n_dense_cones); const i_t cone = dense_cone_ids[dense_idx]; const size_t off = cone_offsets[cone]; const size_t q = cone_offsets[cone + 1] - off; diff --git a/cpp/tests/socp/sparse_augmented_kkt_test.cu b/cpp/tests/socp/sparse_augmented_kkt_test.cu index 9c0b7480ce..af0b9845c6 100644 --- a/cpp/tests/socp/sparse_augmented_kkt_test.cu +++ b/cpp/tests/socp/sparse_augmented_kkt_test.cu @@ -18,13 +18,40 @@ namespace cuopt::mathematical_optimization::barrier::test { +namespace { + +// Packed Hs_diag reference: eta^2 on every entry, head scaled by rank-2 corner d. +std::vector expected_Hs_diag(const cone_data_t& cones, + rmm::cuda_stream_view stream) +{ + const int E = static_cast(cones.n_sparse_cone_entries); + auto d_host = cuopt::host_copy(cones.d, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + auto sparse_cone_ids_host = cuopt::host_copy(cones.sparse_cone_ids, stream); + auto sparse_entry_offsets_host = cuopt::host_copy(cones.sparse_entry_offsets, stream); + + std::vector hs(E); + for (int s = 0; s < cones.n_sparse_cones; ++s) { + const int head = sparse_entry_offsets_host[s]; + const int end = sparse_entry_offsets_host[s + 1]; + const int cone_idx = sparse_cone_ids_host[s]; + const double eta_sq = eta_host[cone_idx] * eta_host[cone_idx]; + for (int e = head; e < end; ++e) { + hs[e] = (e == head) ? eta_sq * d_host[s] : eta_sq; + } + } + return hs; +} + +} // namespace + TEST(sparse_augmented_kkt, cone_counts_and_expansion_size) { auto stream = rmm::cuda_stream_default; - std::vector cone_dimensions{3, 6, 5}; - rmm::device_uvector x(14, stream); - rmm::device_uvector z(14, stream); + std::vector cone_dimensions{3, 6}; + rmm::device_uvector x(9, stream); + rmm::device_uvector z(9, stream); cone_data_t cones( cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); @@ -35,32 +62,36 @@ TEST(sparse_augmented_kkt, cone_counts_and_expansion_size) EXPECT_EQ(cones.n_sparse_cone_entries, 11u); } -TEST(sparse_augmented_kkt, launch_get_Hs_sparse) +TEST(sparse_augmented_kkt, scatter_sparse_hessian_into_augmented) { auto stream = rmm::cuda_stream_default; - std::vector cone_dimensions{3, 6}; - rmm::device_uvector x(9, stream); - rmm::device_uvector z(9, stream); + // Two sparse cones so the fused entry-parallel kernel has more than one sparse-cone + // boundary to get right. + std::vector cone_dimensions{6, 5}; + rmm::device_uvector x(11, stream); + rmm::device_uvector z(11, stream); cone_data_t cones( cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); - ASSERT_EQ(cones.n_sparse_cones, 1); + ASSERT_EQ(cones.n_sparse_cones, 2); + ASSERT_EQ(cones.n_sparse_cone_entries, 11u); + ASSERT_EQ(cones.expansion_var_count(), 4); - std::vector x_host(9, 0.0); - std::vector z_host(9, 0.0); - x_host[3] = 2.0; - z_host[3] = 1.5; + std::vector x_host(11, 0.0); + std::vector z_host(11, 0.0); + x_host[0] = 2.0; + z_host[0] = 1.5; for (int j = 1; j < 6; ++j) { - x_host[3 + j] = 0.1 * j; - z_host[3 + j] = 0.08 * j; + x_host[j] = 0.1 * j; + z_host[j] = 0.08 * j; } - x_host[0] = 1.5; - z_host[0] = 1.2; - for (int j = 1; j < 3; ++j) { - x_host[j] = 0.05; - z_host[j] = 0.04; + x_host[6] = 1.8; + z_host[6] = 1.4; + for (int j = 1; j < 5; ++j) { + x_host[6 + j] = 0.09 * j; + z_host[6 + j] = 0.07 * j; } raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); @@ -69,111 +100,77 @@ TEST(sparse_augmented_kkt, launch_get_Hs_sparse) launch_nt_scaling(cones, stream); launch_update_scaling_sparse(cones, stream); - rmm::device_uvector Hs_diag(cones.n_sparse_cone_entries, stream); - launch_get_Hs_sparse(cones, Hs_diag, stream); - - auto d_host = cuopt::host_copy(cones.d, stream); - auto eta_host = cuopt::host_copy(cones.eta, stream); - auto hs_host = cuopt::host_copy(Hs_diag, stream); - - const int sparse_cone = 1; - const double eta_sq = eta_host[sparse_cone] * eta_host[sparse_cone]; - - EXPECT_NEAR(hs_host[0], eta_sq * d_host[0], 1e-10); - for (int j = 1; j < 6; ++j) { - EXPECT_NEAR(hs_host[j], eta_sq, 1e-10) << "tail index " << j; + const int E = static_cast(cones.n_sparse_cone_entries); + const double dual_perturb = 0.02; + const auto hs_expected = expected_Hs_diag(cones, stream); + + // Distinct augmented-value slots per packed entry: hessian diag, then the four rank-2 + // couplings, then two expansion diagonals per sparse cone. + std::vector hessian_diag_csr(E); + std::vector q_values(E); + std::vector exp_v_col(E); + std::vector exp_u_col(E); + std::vector exp_v_row(E); + std::vector exp_u_row(E); + for (int e = 0; e < E; ++e) { + hessian_diag_csr[e] = e; + q_values[e] = 0.01 * (e + 1); + exp_v_col[e] = 11 + e; + exp_u_col[e] = 22 + e; + exp_v_row[e] = 33 + e; + exp_u_row[e] = 44 + e; } -} - -TEST(sparse_augmented_kkt, scatter_and_update_sparse_expansion) -{ - auto stream = rmm::cuda_stream_default; - - std::vector cone_dimensions{6}; - rmm::device_uvector x(6, stream); - rmm::device_uvector z(6, stream); + std::vector sparse_expansion_D{55, 56, 57, 58}; + const int nnz = 59; - cone_data_t cones( - cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); - - ASSERT_EQ(cones.n_sparse_cones, 1); - ASSERT_EQ(cones.expansion_var_count(), 2); - - std::vector x_host{2.0, 0.2, 0.3, 0.4, 0.5, 0.6}; - std::vector z_host{1.5, 0.1, 0.15, 0.2, 0.25, 0.3}; - raft::copy(cones.x.data(), x_host.data(), x_host.size(), stream); - raft::copy(cones.z.data(), z_host.data(), z_host.size(), stream); - - launch_nt_scaling(cones, stream); - launch_update_scaling_sparse(cones, stream); - - const int nnz = 32; + auto d_hessian_diag_csr = cuopt::device_copy(hessian_diag_csr, stream); + auto d_q_values = cuopt::device_copy(q_values, stream); + auto d_exp_v_col = cuopt::device_copy(exp_v_col, stream); + auto d_exp_u_col = cuopt::device_copy(exp_u_col, stream); + auto d_exp_v_row = cuopt::device_copy(exp_v_row, stream); + auto d_exp_u_row = cuopt::device_copy(exp_u_row, stream); + auto d_sparse_expansion_D = cuopt::device_copy(sparse_expansion_D, stream); rmm::device_uvector augmented_x(nnz, stream); thrust::fill(rmm::exec_policy(stream), augmented_x.begin(), augmented_x.end(), 0.0); - - // One coupling/Hessian CSR entry per sparse-cone element, matching production where every - // *_col/*_row index array has size n_sparse_cone_entries (== cones.sparse_v.size()). - const int q_dim = 6; - ASSERT_EQ(cones.n_sparse_cone_entries, static_cast(q_dim)); - - std::vector sparse_hessian_diag(q_dim); - std::vector sparse_hessian_Q(q_dim, 0.0); - std::vector sparse_exp_v_col(q_dim); - std::vector sparse_exp_u_col(q_dim); - std::vector sparse_exp_v_row(q_dim); - std::vector sparse_exp_u_row(q_dim); - std::vector sparse_expansion_D{30, 31}; - for (int j = 0; j < q_dim; ++j) { - sparse_hessian_diag[j] = j; // CSR slots 0..5 - sparse_exp_v_col[j] = 6 + j; // 6..11 - sparse_exp_u_col[j] = 12 + j; // 12..17 - sparse_exp_v_row[j] = 18 + j; // 18..23 - sparse_exp_u_row[j] = 24 + j; // 24..29 + rmm::device_uvector d_hs_actual(E, stream); + + scatter_sparse_hessian_into_augmented(cones, + augmented_x, + d_hs_actual, + d_hessian_diag_csr, + d_q_values, + d_exp_v_col, + d_exp_u_col, + d_exp_v_row, + d_exp_u_row, + d_sparse_expansion_D, + stream, + dual_perturb); + + auto hs_actual_host = cuopt::host_copy(d_hs_actual, stream); + auto aug_host = cuopt::host_copy(augmented_x, stream); + auto v_host = cuopt::host_copy(cones.sparse_v, stream); + auto u_host = cuopt::host_copy(cones.sparse_u, stream); + auto eta_host = cuopt::host_copy(cones.eta, stream); + auto sparse_cone_ids_host = cuopt::host_copy(cones.sparse_cone_ids, stream); + + for (int e = 0; e < E; ++e) { + EXPECT_NEAR(hs_actual_host[e], hs_expected[e], 1e-10) << "Hs_diag entry " << e; + EXPECT_NEAR(aug_host[e], -hs_actual_host[e] - q_values[e] - dual_perturb, 1e-10) + << "hessian diag " << e; + EXPECT_NEAR(aug_host[11 + e], v_host[e], 1e-10) << "v col " << e; + EXPECT_NEAR(aug_host[22 + e], u_host[e], 1e-10) << "u col " << e; + EXPECT_NEAR(aug_host[33 + e], v_host[e], 1e-10) << "v row " << e; + EXPECT_NEAR(aug_host[44 + e], u_host[e], 1e-10) << "u row " << e; } - auto d_sparse_hessian_diag = cuopt::device_copy(sparse_hessian_diag, stream); - auto d_sparse_hessian_Q = cuopt::device_copy(sparse_hessian_Q, stream); - auto d_sparse_exp_v_col = cuopt::device_copy(sparse_exp_v_col, stream); - auto d_sparse_exp_u_col = cuopt::device_copy(sparse_exp_u_col, stream); - auto d_sparse_exp_v_row = cuopt::device_copy(sparse_exp_v_row, stream); - auto d_sparse_exp_u_row = cuopt::device_copy(sparse_exp_u_row, stream); - auto d_sparse_expansion_D = cuopt::device_copy(sparse_expansion_D, stream); - rmm::device_uvector d_sparse_Hs_diag(cones.n_sparse_cone_entries, stream); - - launch_get_Hs_sparse(cones, d_sparse_Hs_diag, stream); - scatter_sparse_hessian_diag_into_augmented( - augmented_x, d_sparse_hessian_diag, d_sparse_Hs_diag, d_sparse_hessian_Q, stream, 0.0); - update_sparse_expansion_in_augmented(augmented_x, - d_sparse_exp_v_col, - d_sparse_exp_u_col, - d_sparse_exp_v_row, - d_sparse_exp_u_row, - d_sparse_expansion_D, - cones.sparse_v, - cones.sparse_u, - cones.eta, - cones.sparse_cone_ids, - stream, - 0.0); - - auto aug_host = cuopt::host_copy(augmented_x, stream); - auto v_host = cuopt::host_copy(cones.sparse_v, stream); - auto u_host = cuopt::host_copy(cones.sparse_u, stream); - auto hs_host = cuopt::host_copy(d_sparse_Hs_diag, stream); - auto eta_host = cuopt::host_copy(cones.eta, stream); - - for (int j = 0; j < q_dim; ++j) { - EXPECT_NEAR(aug_host[j], -hs_host[j], 1e-10) << "hessian diag " << j; - EXPECT_NEAR(aug_host[6 + j], v_host[j], 1e-10) << "v col " << j; - EXPECT_NEAR(aug_host[12 + j], u_host[j], 1e-10) << "u col " << j; - EXPECT_NEAR(aug_host[18 + j], v_host[j], 1e-10) << "v row " << j; - EXPECT_NEAR(aug_host[24 + j], u_host[j], 1e-10) << "u row " << j; + for (int s = 0; s < cones.n_sparse_cones; ++s) { + const int cone_idx = sparse_cone_ids_host[s]; + const double eta_sq = eta_host[cone_idx] * eta_host[cone_idx]; + EXPECT_NEAR(aug_host[55 + 2 * s], -(eta_sq + dual_perturb), 1e-10) << "expansion v " << s; + EXPECT_NEAR(aug_host[55 + 2 * s + 1], eta_sq + dual_perturb, 1e-10) << "expansion u " << s; } - - const double eta_sq = eta_host[0] * eta_host[0]; - EXPECT_NEAR(aug_host[30], -eta_sq, 1e-10); - EXPECT_NEAR(aug_host[31], eta_sq, 1e-10); } TEST(sparse_augmented_kkt, sparse_augmented_matvec) @@ -213,16 +210,17 @@ TEST(sparse_augmented_kkt, sparse_augmented_matvec) x_vec[n_primal + m_rows] = 0.25; // expansion v x_vec[n_primal + m_rows + 1] = -0.15; // expansion u + const auto hs_host = expected_Hs_diag(cones, stream); + auto d_hs = cuopt::device_copy(hs_host, stream); + rmm::device_uvector d_x(sys_size, stream); rmm::device_uvector d_r1(n_primal, stream); rmm::device_uvector d_y_exp(p, stream); - rmm::device_uvector d_hs(cones.n_sparse_cone_entries, stream); raft::copy(d_x.data(), x_vec.data(), sys_size, stream); thrust::fill(rmm::exec_policy(stream), d_r1.begin(), d_r1.end(), 0.0); thrust::fill(rmm::exec_policy(stream), d_y_exp.begin(), d_y_exp.end(), 0.0); - launch_get_Hs_sparse(cones, d_hs, stream); launch_sparse_augmented_matvec(raft::device_span(d_x.data(), d_x.size()), raft::device_span(d_r1.data(), d_r1.size()), raft::device_span(d_y_exp.data(), d_y_exp.size()), @@ -235,7 +233,6 @@ TEST(sparse_augmented_kkt, sparse_augmented_matvec) auto r1_host = cuopt::host_copy(d_r1, stream); auto yexp_host = cuopt::host_copy(d_y_exp, stream); - auto hs_host = cuopt::host_copy(d_hs, stream); auto v_host = cuopt::host_copy(cones.sparse_v, stream); auto u_host = cuopt::host_copy(cones.sparse_u, stream); auto eta_host = cuopt::host_copy(cones.eta, stream); @@ -297,8 +294,8 @@ TEST(sparse_augmented_kkt, update_scaling_sparse_dim_1000) EXPECT_TRUE(std::isfinite(u_host[j])) << "u entry " << j; } - rmm::device_uvector Hs_diag(cones.n_sparse_cone_entries, stream); - launch_get_Hs_sparse(cones, Hs_diag, stream); + const auto hs_host = expected_Hs_diag(cones, stream); + auto d_hs = cuopt::device_copy(hs_host, stream); const int n_primal = 1000; const int m_rows = 1; @@ -324,7 +321,7 @@ TEST(sparse_augmented_kkt, update_scaling_sparse_dim_1000) raft::device_span(d_r1.data(), d_r1.size()), raft::device_span(d_y_exp.data(), d_y_exp.size()), cones, - raft::device_span(Hs_diag.data(), Hs_diag.size()), + raft::device_span(d_hs.data(), d_hs.size()), /*cone_var_start=*/0, n_primal, m_rows, @@ -332,7 +329,6 @@ TEST(sparse_augmented_kkt, update_scaling_sparse_dim_1000) auto r1_host = cuopt::host_copy(d_r1, stream); auto yexp_host = cuopt::host_copy(d_y_exp, stream); - auto hs_host = cuopt::host_copy(Hs_diag, stream); const double eta_sq = eta_host[0] * eta_host[0]; const double x_exp_v = x_vec[n_primal + m_rows]; From 0d08c0a0ef52a186554095c461c3768d694058a1 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Tue, 21 Jul 2026 06:56:52 -0700 Subject: [PATCH 09/14] test fix Signed-off-by: yuwenchen95 --- cpp/tests/socp/sparse_augmented_kkt_test.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/tests/socp/sparse_augmented_kkt_test.cu b/cpp/tests/socp/sparse_augmented_kkt_test.cu index af0b9845c6..9d9d9eaac9 100644 --- a/cpp/tests/socp/sparse_augmented_kkt_test.cu +++ b/cpp/tests/socp/sparse_augmented_kkt_test.cu @@ -49,9 +49,9 @@ TEST(sparse_augmented_kkt, cone_counts_and_expansion_size) { auto stream = rmm::cuda_stream_default; - std::vector cone_dimensions{3, 6}; - rmm::device_uvector x(9, stream); - rmm::device_uvector z(9, stream); + std::vector cone_dimensions{3, 6, 5}; + rmm::device_uvector x(14, stream); + rmm::device_uvector z(14, stream); cone_data_t cones( cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); From 47b612ab9495bb8dc8928758b8a4a595510f1714 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Tue, 21 Jul 2026 12:11:20 -0700 Subject: [PATCH 10/14] Optimize SOC cone step-length computation Signed-off-by: yuwenchen95 --- cpp/src/barrier/second_order_cone_kernels.cuh | 402 +++++++++++++----- .../barrier/second_order_cone_reduction.cuh | 6 + cpp/tests/socp/second_order_cone_kernels.cu | 59 ++- 3 files changed, 354 insertions(+), 113 deletions(-) diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index 272e73d335..7a25c5da70 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -57,6 +57,21 @@ namespace cuopt::mathematical_optimization::barrier { inline constexpr int soc_block_size = 256; +/** + * Tail aggregates for the cone step-length reduction (CUB value type). + */ +template +struct step_tail_sums_t { + f_t du_tail_sq{}; + f_t u_tail_du_tail{}; + f_t u_tail_sq{}; + + __host__ __device__ step_tail_sums_t operator+(step_tail_sums_t o) const + { + return {du_tail_sq + o.du_tail_sq, u_tail_du_tail + o.u_tail_du_tail, u_tail_sq + o.u_tail_sq}; + } +}; + /** * Reusable device workspace for second-order cone kernels. * @@ -75,15 +90,22 @@ struct cone_scratch_t { rmm::device_uvector step_alpha_primal; // [n_cones] rmm::device_uvector step_alpha_dual; // [n_cones] + // Large-cone step-length CUB outputs, one packed slot per large cone. + rmm::device_uvector> step_large_tail_sums; // [n_large] + // TODO: Consider moving this out to the barrier layer when we wire it in rmm::device_uvector temp_cone; // [n_cone_entries] - cone_scratch_t(i_t n_cones_in, size_t n_cone_entries_in, rmm::cuda_stream_view stream) + cone_scratch_t(i_t n_cones_in, + size_t n_cone_entries_in, + size_t n_large, + rmm::cuda_stream_view stream) : n_cones(n_cones_in), n_cone_entries(n_cone_entries_in), slots(0, stream), step_alpha_primal(0, stream), step_alpha_dual(0, stream), + step_large_tail_sums(0, stream), temp_cone(0, stream) { const size_t n_cones_size = static_cast(n_cones); @@ -91,6 +113,7 @@ struct cone_scratch_t { slots.resize(n_cones_size * static_cast(n_slots), stream); step_alpha_primal.resize(n_cones_size, stream); step_alpha_dual.resize(n_cones_size, stream); + if (n_large > 0) { step_large_tail_sums.resize(n_large, stream); } temp_cone.resize(n_cone_entries, stream); } @@ -120,57 +143,6 @@ struct to_size_t_t { } }; -template -HD f_t cone_step_length_from_scalars( - f_t u0, f_t du0, f_t du_tail_sq, f_t u_tail_du_tail, f_t u_tail_sq, f_t alpha_max) -{ - const f_t a = du0 * du0 - du_tail_sq; - const f_t b = u0 * du0 - u_tail_du_tail; - const f_t c_raw = u0 * u0 - u_tail_sq; - const f_t c = c_raw > 0 ? c_raw : 0; - const f_t disc = b * b - a * c; - f_t alpha = alpha_max; - - if (u0 >= 0 && du0 < 0) { alpha = cuda::std::min(alpha, -u0 / du0); } - - if ((a > 0 && b > 0) || disc < 0) { return alpha; } - - if (a == 0) { - return alpha; - } else if (c == 0) { - alpha = a >= 0 ? alpha : 0; - } else { - const f_t t = -(b + copysign(sqrt(disc), b)); - f_t r1 = c / t; - f_t r2 = t / a; - if (r1 < 0) { r1 = alpha; } - if (r2 < 0) { r2 = alpha; } - alpha = cuda::std::min(alpha, cuda::std::min(r1, r2)); - } - - return alpha; -} - -template -__global__ void __launch_bounds__(soc_block_size) - step_length_single_kernel(raft::device_span u, - raft::device_span du, - raft::device_span alpha, - raft::device_span du_tail_sq, - raft::device_span u_tail_du_tail, - raft::device_span u_tail_sq, - raft::device_span cone_offsets, - f_t alpha_max, - i_t n_cones) -{ - const i_t cone = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - if (cone >= n_cones) { return; } - - const size_t off = cone_offsets[cone]; - alpha[cone] = cone_step_length_from_scalars( - u[off], du[off], du_tail_sq[cone], u_tail_du_tail[cone], u_tail_sq[cone], alpha_max); -} - /** * Device storage for second-order cone topology, NT scaling, and iterate views. * @@ -242,7 +214,7 @@ struct cone_data_t { sparse_u(0, stream), sparse_entry_offsets(0, stream), cone_is_sparse(n_cones, stream), - scratch(n_cones, n_cone_entries, stream), + scratch(n_cones, n_cone_entries, segmented_sum.large_cone_ids.size(), stream), soc_threshold(soc_threshold_in), n_sparse_cones(0), n_sparse_cone_entries(0) @@ -305,7 +277,7 @@ struct cone_data_t { thrust::make_counting_iterator(0), thrust::make_counting_iterator(n_cone_entries), element_cone_ids.begin()); - segmented_sum.template prepare_workspace(stream); + segmented_sum.template prepare_workspace>(stream); } // True when at least one cone is large enough (dim > soc_threshold) to use the @@ -1064,8 +1036,7 @@ __global__ void __launch_bounds__(soc_block_size) scatter_sparse_hessian_into_au } /** - * Fused sparse SOC augmented-system assembly: one kernel launch with one thread per packed - * sparse entry (E threads total), writing `Hs_diag`, the Hessian CSR scatter, and the four + * Sparse SOC augmented-system assembly: writing `Hs_diag`, the Hessian CSR scatter, and the four * rank-2 couplings, with each cone's unique head thread additionally applying the rank-2 * corner scale and writing that cone's two expansion diagonals. * @@ -1424,23 +1395,267 @@ void scatter_dense_hessian_into_augmented(const cone_data_t& cones, RAFT_CUDA_TRY(cudaPeekAtLastError()); } +// ============================================================================= +// Cone step length +// +// Max alpha keeping u + alpha du in the SOC for each packed cone. Size-aware +// path: one pass over each cone's tail accumulates +// (||du_tail||^2, , ||u_tail||^2), then solves the quadratic +// boundary. Small/medium/large partitions come from segmented_sum_t. +// ============================================================================= + +template +HD f_t cone_step_length_from_scalars( + f_t u0, f_t du0, f_t du_tail_sq, f_t u_tail_du_tail, f_t u_tail_sq, f_t alpha_max) +{ + const f_t a = du0 * du0 - du_tail_sq; + const f_t b = u0 * du0 - u_tail_du_tail; + const f_t c_raw = u0 * u0 - u_tail_sq; + const f_t c = c_raw > 0 ? c_raw : 0; + const f_t disc = b * b - a * c; + f_t alpha = alpha_max; + + if (u0 >= 0 && du0 < 0) { alpha = cuda::std::min(alpha, -u0 / du0); } + + if ((a > 0 && b > 0) || disc < 0) { return alpha; } + + if (a == 0) { + return alpha; + } else if (c == 0) { + alpha = a >= 0 ? alpha : 0; + } else { + const f_t t = -(b + copysign(sqrt(disc), b)); + f_t r1 = c / t; + f_t r2 = t / a; + if (r1 < 0) { r1 = alpha; } + if (r2 < 0) { r2 = alpha; } + alpha = cuda::std::min(alpha, cuda::std::min(r1, r2)); + } + + return alpha; +} + /** - * Compute the combined (primal and dual) maximum step length that keeps SOC blocks - * feasible: + * One warp per small cone: accumulate the three tail scalars, then solve for + * alpha[cone]. Tail-only (local index 0 is the SOC head). + */ +template +__global__ void __launch_bounds__(warps_per_cta* raft::WarpSize) + step_length_small_kernel(raft::device_span u, + raft::device_span du, + raft::device_span alpha, + raft::device_span small_cone_ids, + raft::device_span cone_offsets, + f_t alpha_max) +{ + static_assert(warps_per_cta > 0); + static_assert(warps_per_cta * raft::WarpSize <= 1024); + + using warp_reduce_t = cub::WarpReduce, raft::WarpSize>; + __shared__ typename warp_reduce_t::TempStorage temp_storage[warps_per_cta]; + + const auto lane_id = raft::laneId(); + const auto warp_idx = threadIdx.x / raft::WarpSize; + const auto slot = blockIdx.x * warps_per_cta + warp_idx; + if (slot >= small_cone_ids.size()) { return; } + + const i_t cone = small_cone_ids[slot]; + const size_t off = cone_offsets[cone]; + const size_t dim = cone_offsets[cone + 1] - off; + + step_tail_sums_t tail_sums{}; + for (size_t i = lane_id + 1; i < dim; i += raft::WarpSize) { + const f_t ui = u[off + i]; + const f_t dui = du[off + i]; + tail_sums.du_tail_sq += dui * dui; + tail_sums.u_tail_du_tail += ui * dui; + tail_sums.u_tail_sq += ui * ui; + } + + tail_sums = warp_reduce_t(temp_storage[warp_idx]).Sum(tail_sums); + + if (lane_id == 0) { + alpha[cone] = cone_step_length_from_scalars(u[off], + du[off], + tail_sums.du_tail_sq, + tail_sums.u_tail_du_tail, + tail_sums.u_tail_sq, + alpha_max); + } +} + +/** + * One block per medium cone: three-scalar tail reduction + step solve. + */ +template +__global__ void __launch_bounds__(block_dim) + step_length_medium_kernel(raft::device_span u, + raft::device_span du, + raft::device_span alpha, + raft::device_span medium_cone_ids, + raft::device_span cone_offsets, + f_t alpha_max) +{ + static_assert(block_dim > 0); + static_assert(block_dim <= 1024); + + constexpr int items_per_thread = 4; + + using block_reduce_t = cub::BlockReduce; + __shared__ typename block_reduce_t::TempStorage temp_storage; + + const auto slot = blockIdx.x; + if (slot >= medium_cone_ids.size()) { return; } + + const i_t cone = medium_cone_ids[slot]; + const size_t off = cone_offsets[cone]; + const size_t dim = cone_offsets[cone + 1] - off; + + f_t acc_du_sq[items_per_thread]{}; + f_t acc_u_du[items_per_thread]{}; + f_t acc_u_sq[items_per_thread]{}; + + const size_t tile = static_cast(block_dim) * items_per_thread; + for (size_t tile_start = 0; tile_start < dim; tile_start += tile) { +#pragma unroll + for (int k = 0; k < items_per_thread; ++k) { + const size_t idx = tile_start + threadIdx.x + static_cast(k) * block_dim; + if (idx > 0 && idx < dim) { + const f_t ui = u[off + idx]; + const f_t dui = du[off + idx]; + acc_du_sq[k] += dui * dui; + acc_u_du[k] += ui * dui; + acc_u_sq[k] += ui * ui; + } + } + } + + f_t du_sq = f_t{0}; + f_t u_du = f_t{0}; + f_t u_sq = f_t{0}; +#pragma unroll + for (int k = 0; k < items_per_thread; ++k) { + du_sq += acc_du_sq[k]; + u_du += acc_u_du[k]; + u_sq += acc_u_sq[k]; + } + + du_sq = block_reduce_t(temp_storage).Sum(du_sq); + __syncthreads(); + u_du = block_reduce_t(temp_storage).Sum(u_du); + __syncthreads(); + u_sq = block_reduce_t(temp_storage).Sum(u_sq); + + if (threadIdx.x == 0) { + alpha[cone] = cone_step_length_from_scalars(u[off], du[off], du_sq, u_du, u_sq, alpha_max); + } +} + +template +__global__ void step_length_large_solve_kernel(raft::device_span u, + raft::device_span du, + raft::device_span alpha, + raft::device_span> sums, + raft::device_span large_cone_ids, + raft::device_span cone_offsets, + f_t alpha_max) +{ + const size_t slot = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (slot >= large_cone_ids.size()) { return; } + + const i_t cone = large_cone_ids[slot]; + const size_t off = cone_offsets[cone]; + const auto& s = sums[slot]; + alpha[cone] = cone_step_length_from_scalars( + u[off], du[off], s.du_tail_sq, s.u_tail_du_tail, s.u_tail_sq, alpha_max); +} + +/** + * Size-aware step length: one pass over each cone's tail, then write + * alpha[cone]. + */ +template +void launch_cone_step_length(segmented_sum_t& partitions, + raft::device_span u, + raft::device_span du, + raft::device_span alpha, + raft::device_span> large_sums, + f_t alpha_max, + rmm::cuda_stream_view stream) +{ + constexpr int warps_per_cta = 8; + if (!partitions.small_cone_ids.is_empty()) { + const auto n_small = partitions.small_cone_ids.size(); + const auto grid = (n_small + warps_per_cta - 1) / warps_per_cta; + step_length_small_kernel + <<>>( + u, + du, + alpha, + cuopt::make_span(partitions.small_cone_ids), + partitions.cone_offsets, + alpha_max); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (!partitions.medium_cone_ids.is_empty()) { + constexpr int medium_block_dim = 256; + step_length_medium_kernel + <<>>( + u, + du, + alpha, + cuopt::make_span(partitions.medium_cone_ids), + partitions.cone_offsets, + alpha_max); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + if (!partitions.large_cone_ids.empty()) { + const auto n_large = partitions.large_cone_ids.size(); + + for (std::size_t i = 0; i < n_large; ++i) { + const size_t off = partitions.large_cone_offsets[i]; + const i_t dim = partitions.large_cone_dimensions[i]; + std::size_t temp_bytes = partitions.cub_workspace_bytes; + + auto input = + thrust::make_transform_iterator(thrust::make_counting_iterator(0), + [u, du, off] HD(size_t local) -> step_tail_sums_t { + if (local == 0) { return {}; } + const f_t ui = u[off + local]; + const f_t dui = du[off + local]; + return {dui * dui, ui * dui, ui * ui}; + }); + + RAFT_CUDA_TRY(cub::DeviceReduce::Sum(partitions.cub_workspace.data(), + temp_bytes, + input, + large_sums.data() + i, + dim, + stream.value())); + } + + raft::device_span> large_sums_c(large_sums.data(), + large_sums.size()); + constexpr int large_solve_block_dim = 256; + const auto grid = raft::ceildiv(n_large, static_cast(large_solve_block_dim)); + step_length_large_solve_kernel<<>>( + u, + du, + alpha, + large_sums_c, + cuopt::make_span(partitions.large_cone_ids_device), + partitions.cone_offsets, + alpha_max); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +/** + * Combined (primal and dual) maximum step length keeping SOC blocks feasible: * * x + alpha dx in Q, z + alpha dz in Q, alpha <= alpha_max. - * - * For one cone u + alpha du, feasibility is - * - * u_0 + alpha du_0 >= ||u_tail + alpha du_tail||. - * - * Squaring gives the quadratic - * - * c + 2 b alpha + a alpha^2 >= 0, - * - * where c = det_J(u), b = u_0 du_0 - , and - * a = det_J(du). The per-cone kernel below solves for the first boundary - * crossing, and the final reductions take the global minimum over cones. */ template f_t compute_cone_step_length(cone_data_t& cones, @@ -1449,51 +1664,14 @@ f_t compute_cone_step_length(cone_data_t& cones, f_t alpha_max, rmm::cuda_stream_view stream) { - auto cone_offsets = cuopt::make_span(cones.cone_offsets); - auto element_cone_ids = cuopt::make_span(cones.element_cone_ids); - auto slot_0 = cones.scratch.template get_slot<0>(); - auto slot_1 = cones.scratch.template get_slot<1>(); - auto slot_2 = cones.scratch.template get_slot<2>(); - - auto run_pass = [&](raft::device_span u, - raft::device_span du, - raft::device_span alpha) { - auto du_tail_sq_terms = - thrust::make_transform_iterator(thrust::make_counting_iterator(0), - [du, cone_offsets, element_cone_ids] HD(size_t idx) -> f_t { - const i_t cone = element_cone_ids[idx]; - return idx == cone_offsets[cone] ? 0 : du[idx] * du[idx]; - }); - cones.segmented_sum(du_tail_sq_terms, slot_0, stream); - - auto u_tail_du_tail_terms = thrust::make_transform_iterator( - thrust::make_counting_iterator(0), - [u, du, cone_offsets, element_cone_ids] HD(size_t idx) -> f_t { - const i_t cone = element_cone_ids[idx]; - return idx == cone_offsets[cone] ? 0 : u[idx] * du[idx]; - }); - cones.segmented_sum(u_tail_du_tail_terms, slot_1, stream); - - auto u_tail_sq_terms = - thrust::make_transform_iterator(thrust::make_counting_iterator(0), - [u, cone_offsets, element_cone_ids] HD(size_t idx) -> f_t { - const i_t cone = element_cone_ids[idx]; - return idx == cone_offsets[cone] ? 0 : u[idx] * u[idx]; - }); - cones.segmented_sum(u_tail_sq_terms, slot_2, stream); - - const size_t grid_dim = - raft::ceildiv(static_cast(cones.n_cones), soc_block_size); - step_length_single_kernel<<>>( - u, du, alpha, slot_0, slot_1, slot_2, cone_offsets, alpha_max, cones.n_cones); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - }; - auto alpha_primal = cuopt::make_span(cones.scratch.step_alpha_primal); auto alpha_dual = cuopt::make_span(cones.scratch.step_alpha_dual); + raft::device_span x(cones.x.data(), cones.x.size()); + raft::device_span z(cones.z.data(), cones.z.size()); - run_pass(cones.x, dx, alpha_primal); - run_pass(cones.z, dz, alpha_dual); + auto large_sums = cuopt::make_span(cones.scratch.step_large_tail_sums); + launch_cone_step_length(cones.segmented_sum, x, dx, alpha_primal, large_sums, alpha_max, stream); + launch_cone_step_length(cones.segmented_sum, z, dz, alpha_dual, large_sums, alpha_max, stream); return thrust::transform_reduce( rmm::exec_policy(stream), diff --git a/cpp/src/barrier/second_order_cone_reduction.cuh b/cpp/src/barrier/second_order_cone_reduction.cuh index a9629a0d23..bed06572a9 100644 --- a/cpp/src/barrier/second_order_cone_reduction.cuh +++ b/cpp/src/barrier/second_order_cone_reduction.cuh @@ -75,6 +75,7 @@ struct segmented_sum_t { std::vector large_cone_offsets; std::vector large_cone_ids; std::vector large_cone_dimensions; + rmm::device_uvector large_cone_ids_device; // device copy for batched kernels // Maximum CUB temporary storage needed by prepared large reductions. std::size_t cub_workspace_bytes = 0; @@ -163,6 +164,7 @@ struct segmented_sum_t { : cone_offsets(cone_offsets_in), small_cone_ids(0, stream), medium_cone_ids(0, stream), + large_cone_ids_device(0, stream), cub_workspace(0, stream) { std::vector small_cone_ids_host; @@ -193,6 +195,10 @@ struct segmented_sum_t { cuopt::device_copy(medium_cone_ids, medium_cone_ids_host, stream); need_sync = true; } + if (!large_cone_ids.empty()) { + cuopt::device_copy(large_cone_ids_device, large_cone_ids, stream); + need_sync = true; + } if (need_sync) { stream.synchronize(); } } }; diff --git a/cpp/tests/socp/second_order_cone_kernels.cu b/cpp/tests/socp/second_order_cone_kernels.cu index 334c099519..4193838272 100644 --- a/cpp/tests/socp/second_order_cone_kernels.cu +++ b/cpp/tests/socp/second_order_cone_kernels.cu @@ -369,7 +369,7 @@ TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) { auto stream = rmm::cuda_stream_default; - std::vector cone_dimensions{3, 65, 32769}; + std::vector cone_dimensions{3, 65, 32769, 40000}; std::size_t n_cone_entries = 0; for (const auto dim : cone_dimensions) { n_cone_entries += static_cast(dim); @@ -431,6 +431,55 @@ TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) EXPECT_GE(u0 * u0 + cone_tol, tail_sq) << label << " cone " << cone_idx; }; + const auto host_cone_step_length_from_scalars = [](double u0, + double du0, + double du_tail_sq, + double u_tail_du_tail, + double u_tail_sq, + double alpha_max_in) { + const double a = du0 * du0 - du_tail_sq; + const double b = u0 * du0 - u_tail_du_tail; + const double c_raw = u0 * u0 - u_tail_sq; + const double c = c_raw > 0 ? c_raw : 0; + const double disc = b * b - a * c; + double alpha = alpha_max_in; + + if (u0 >= 0 && du0 < 0) { alpha = std::min(alpha, -u0 / du0); } + + if ((a > 0 && b > 0) || disc < 0) { return alpha; } + + if (a == 0) { + return alpha; + } else if (c == 0) { + alpha = a >= 0 ? alpha : 0; + } else { + const double t = -(b + std::copysign(std::sqrt(disc), b)); + double r1 = c / t; + double r2 = t / a; + if (r1 < 0) { r1 = alpha; } + if (r2 < 0) { r2 = alpha; } + alpha = std::min(alpha, std::min(r1, r2)); + } + + return alpha; + }; + + const auto host_cone_step_length_one = + [&](std::vector const& u, std::vector const& du, std::size_t off, int dim) { + double du_tail_sq = 0.0; + double u_tail_du_tail = 0.0; + double u_tail_sq = 0.0; + for (int j = 1; j < dim; ++j) { + const double ui = u[off + static_cast(j)]; + const double dui = du[off + static_cast(j)]; + du_tail_sq += dui * dui; + u_tail_du_tail += ui * dui; + u_tail_sq += ui * ui; + } + return host_cone_step_length_from_scalars( + u[off], du[off], du_tail_sq, u_tail_du_tail, u_tail_sq, alpha_max); + }; + auto x = cuopt::device_copy(x_host, stream); auto z = cuopt::device_copy(z_host, stream); auto dx = cuopt::device_copy(dx_host, stream); @@ -454,6 +503,14 @@ TEST(second_order_cone_kernels, cone_step_length_keeps_iterate_in_cone) // both the x/dx and z/dz feasibility checks below, which is still valid since it is // <= each side's own feasibility boundary. for (std::size_t cone = 0; cone < cone_dimensions.size(); ++cone) { + const std::size_t off = cone_block_offset(cone); + const auto dim = cone_dimensions[cone]; + const double host_primal = host_cone_step_length_one(x_host, dx_host, off, dim); + const double host_dual = host_cone_step_length_one(z_host, dz_host, off, dim); + // Reduction order may differ slightly from the GPU path. + EXPECT_NEAR(primal_per_cone[cone], host_primal, 1e-10) << "primal host ref cone " << cone; + EXPECT_NEAR(dual_per_cone[cone], host_dual, 1e-10) << "dual host ref cone " << cone; + const double combined_cone = std::min(primal_per_cone[cone], dual_per_cone[cone]); EXPECT_GT(primal_per_cone[cone], 0.0) << "primal cone " << cone; EXPECT_GT(dual_per_cone[cone], 0.0) << "dual cone " << cone; From 2b18b6d2fa6b648c13e333ccc6e73793501a1a3e Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Wed, 22 Jul 2026 10:11:11 -0700 Subject: [PATCH 11/14] Addressing notation and comment issues Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 70c182571d..6118adce54 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -40,8 +40,6 @@ #include #include -#include - #include #include #include @@ -52,6 +50,8 @@ #include #include +#include + namespace cuopt::mathematical_optimization::barrier { using simplex::compute_user_objective; @@ -799,7 +799,7 @@ class iteration_data_t { if (cone_sparse_idx[k] < 0) { dense_cone_idx_by_cone[k] = dense_idx; dense_cone_ids_host.push_back(k); - const auto q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; + const i_t q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; dense_cone_block_offsets_host.push_back(dense_cone_block_offsets_host.back() + q_k * q_k); ++dense_idx; @@ -980,9 +980,9 @@ class iteration_data_t { const i_t l = k - n; const i_t col_beg = AT.col_start[l]; const i_t col_end = AT.col_start[l + 1]; - for (i_t idx = col_beg; idx < col_end; ++idx) { - augmented_CSR.j[q] = AT.i[idx]; - augmented_CSR.x[q++] = AT.x[idx]; + for (i_t p = col_beg; p < col_end; ++p) { + augmented_CSR.j[q] = AT.i[p]; + augmented_CSR.x[q++] = AT.x[p]; } augmented_diagonal_indices[k] = q; augmented_CSR.j[q] = k; @@ -999,9 +999,9 @@ class iteration_data_t { const i_t prow_u = n + m + 2 * s + 1; augmented_CSR.row_start[prow_v] = q; - for (i_t j = 0; j < q_k; ++j) { - sparse_exp_v_row_host[flat_base + j] = q; - augmented_CSR.j[q] = cone_col_start + j; + for (i_t t = 0; t < q_k; ++t) { + sparse_exp_v_row_host[flat_base + t] = q; + augmented_CSR.j[q] = cone_col_start + t; augmented_CSR.x[q++] = f_t(0); } sparse_expansion_D_host[2 * s] = q; @@ -1009,9 +1009,9 @@ class iteration_data_t { augmented_CSR.x[q++] = f_t(0); augmented_CSR.row_start[prow_u] = q; - for (i_t j = 0; j < q_k; ++j) { - sparse_exp_u_row_host[flat_base + j] = q; - augmented_CSR.j[q] = cone_col_start + j; + for (i_t t = 0; t < q_k; ++t) { + sparse_exp_u_row_host[flat_base + t] = q; + augmented_CSR.j[q] = cone_col_start + t; augmented_CSR.x[q++] = f_t(0); } sparse_expansion_D_host[2 * s + 1] = q; @@ -1026,7 +1026,7 @@ class iteration_data_t { i_t expected_nnz = 2 * nnzA + (n - m_c) + dense_soc_kkt_nnz + m + off_diag_Qnz + sparse_soc_kkt_nnz; settings_.log.debug("augmented nz %d predicted %d\n", q, expected_nnz); - cuopt_assert(q == expected_nnz, "augmented nz != predicted"); + cuopt_assert(q == expected_nnz, "augmented nnz != predicted"); cuopt_assert(A.col_start[n] == AT.col_start[m], "A nz != AT nz"); if (n_sparse_entries > 0) { @@ -3910,8 +3910,10 @@ void barrier_solver_t::compute_target_mu( step_primal_aff = step_dual_aff = std::min(step_primal_aff, step_dual_aff); } - // Compute complementarity_xz_aff_sum and complementarity_wv_aff_sum. Save the affine direction as - // a side effect. + // Compute complementarity_xz_aff_sum = sum(x_aff * z_aff), + // where x_aff = x + step_primal_aff * dx_aff and z_aff = z + step_dual_aff * dz_aff + // Here the update of x_aff and z_aff are done temporarily and sum of their products is + // computed without storing intermediate results. raft::device_span x_span(data.d_x_.data(), data.d_x_.size()); raft::device_span z_span(data.d_z_.data(), data.d_z_.size()); raft::device_span dx_span(data.d_dx_.data(), data.d_dx_.size()); @@ -3941,6 +3943,8 @@ void barrier_solver_t::compute_target_mu( data.d_x_.size(), stream_view_); + // Here the update of w_aff and v_aff are done temporarily and sum of their products is + // computed without storing intermediate results. raft::device_span w_span(data.d_w_.data(), data.d_w_.size()); raft::device_span v_span(data.d_v_.data(), data.d_v_.size()); raft::device_span dw_span(data.d_dw_.data(), data.d_dw_.size()); From 088bd7f66c748cba9b852cf15f2414bff76bc318 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Wed, 22 Jul 2026 10:16:32 -0700 Subject: [PATCH 12/14] Remove redundant log Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 5 ----- 1 file changed, 5 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 6118adce54..e456564ccd 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -553,11 +553,6 @@ class iteration_data_t { d_aug_y2_.resize(lp.num_rows, stream_view_); d_aug_y_exp_.resize(augmented_expansion_count(), stream_view_); d_aug_y_exp_orig_.resize(augmented_expansion_count(), stream_view_); - if (settings.barrier_csr_ir_matvec) { - settings.log.printf("IR matvec : csr\n"); - } else { - settings.log.printf("IR matvec : matrix-free\n"); - } } else { settings.log.printf("Linear system : ADAT\n"); } From a4c1ed1b7a020b05d2ed6e97523c2ee3d807ca59 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Thu, 23 Jul 2026 14:27:09 -0700 Subject: [PATCH 13/14] Remove barrier_csr_ir_matvec CSR IR matvec path Signed-off-by: yuwenchen95 --- .../pdlp/solver_settings.hpp | 1 - cpp/src/barrier/barrier.cu | 80 +----------- cpp/src/barrier/cusparse_view.cu | 48 ------- cpp/src/barrier/cusparse_view.hpp | 3 - cpp/src/barrier/second_order_cone_kernels.cuh | 121 ------------------ .../dual_simplex/simplex_solver_settings.hpp | 2 - cpp/src/pdlp/solve.cu | 2 - 7 files changed, 3 insertions(+), 254 deletions(-) diff --git a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp index c35fcc2551..2384d03693 100644 --- a/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp +++ b/cpp/include/cuopt/mathematical_optimization/pdlp/solver_settings.hpp @@ -283,7 +283,6 @@ class pdlp_solver_settings_t { bool eliminate_dense_columns{true}; pdlp_precision_t pdlp_precision{pdlp_precision_t::DefaultPrecision}; bool barrier_iterative_refinement{true}; - 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}; diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index e456564ccd..06e3f06304 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -288,7 +288,6 @@ class iteration_data_t { d_sparse_exp_u_row_(0, lp.handle_ptr->get_stream()), d_sparse_expansion_D_(0, lp.handle_ptr->get_stream()), d_sparse_Hs_diag_(0, lp.handle_ptr->get_stream()), - d_dense_cone_diag_csr_indices_(0, lp.handle_ptr->get_stream()), use_augmented(false), has_factorization(false), n_direct_free_linear(0), @@ -662,10 +661,6 @@ class iteration_data_t { } if (settings.concurrent_halt != nullptr && *settings.concurrent_halt == 1) { return; } symbolic_status = chol->analyze(device_augmented); - if (use_csr_ir_matvec()) { - augmented_cusparse_view_ = - std::make_unique>(handle_ptr, device_augmented); - } } else { { raft::common::nvtx::range form_scope("Barrier: LP Data: form ADAT"); @@ -709,8 +704,6 @@ class iteration_data_t { return has_sparse_cones() ? cones().expansion_var_count() : i_t(0); } - bool use_csr_ir_matvec() const { return settings_.barrier_csr_ir_matvec && use_augmented; } - i_t augmented_system_size(i_t n, i_t m) const { return n + m + augmented_expansion_count(); } bool is_cone_variable(i_t variable) const @@ -819,7 +812,6 @@ class iteration_data_t { std::vector augmented_diagonal_indices(factorization_size, -1); std::vector cone_csr_indices_host(dense_soc_kkt_nnz, -1); std::vector cone_Q_values_host(dense_soc_kkt_nnz, f_t(0)); - std::vector dense_cone_diag_indices_host; std::vector sparse_hessian_diag_host(n_sparse_entries, -1); std::vector sparse_hessian_Q_host(n_sparse_entries, f_t(0)); std::vector sparse_exp_v_col_host(n_sparse_entries, -1); @@ -921,10 +913,7 @@ class iteration_data_t { cone_csr_indices_host[block_base + c] = q; cone_Q_values_host[block_base + c] = q_contrib; - if (col == i) { - augmented_diagonal_indices[i] = q; - dense_cone_diag_indices_host.push_back(q); - } + if (col == i) { augmented_diagonal_indices[i] = q; } augmented_CSR.j[q] = col; augmented_CSR.x[q++] = initial_val - q_contrib; } @@ -1080,15 +1069,6 @@ class iteration_data_t { handle_ptr->get_stream()); } - d_dense_cone_diag_csr_indices_.resize(dense_cone_diag_indices_host.size(), - handle_ptr->get_stream()); - if (!dense_cone_diag_indices_host.empty()) { - raft::copy(d_dense_cone_diag_csr_indices_.data(), - dense_cone_diag_indices_host.data(), - dense_cone_diag_indices_host.size(), - handle_ptr->get_stream()); - } - const size_t n_sparse_entries = cones().n_sparse_cone_entries; d_sparse_hessian_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); d_sparse_hessian_Q_.resize(n_sparse_entries, handle_ptr->get_stream()); @@ -2099,41 +2079,6 @@ class iteration_data_t { __host__ __device__ T operator()(T x, T y) const { return alpha * x + beta * y; } }; - void strip_augmented_perturbation() - { - raft::common::nvtx::range fun_scope("Barrier: strip augmented perturbation"); - const i_t n = A.n; - const i_t m = A.m; - const i_t linear_n = has_cones() ? cone_start() : n; - strip_augmented_perturbation_values(device_augmented.x, - d_augmented_diagonal_indices_, - d_Q_diag_, - d_is_direct_free_linear_, - d_sparse_hessian_diag_, - d_sparse_expansion_D_, - d_dense_cone_diag_csr_indices_, - linear_n, - n, - m, - dual_perturb, - primal_perturb, - stream_view_); - } - - void augmented_csr_multiply(f_t alpha, - const rmm::device_uvector& x, - f_t beta, - rmm::device_uvector& 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"); - cuopt_assert(augmented_cusparse_view_ != nullptr, "augmented cusparse view not initialized"); - const i_t sys_size = augmented_system_size(A.n, A.m); - cuopt_assert(static_cast(x.size()) >= sys_size, "augmented_csr_multiply: x too small"); - cuopt_assert(static_cast(y.size()) >= sys_size, "augmented_csr_multiply: y too small"); - augmented_cusparse_view_->spmv(alpha, x, beta, y); - } - // y <- alpha * Augmented * x + beta * y void augmented_multiply(f_t alpha, const rmm::device_uvector& x, @@ -2328,7 +2273,6 @@ class iteration_data_t { rmm::device_uvector d_sparse_exp_u_row_; rmm::device_uvector d_sparse_expansion_D_; rmm::device_uvector d_sparse_Hs_diag_; - rmm::device_uvector d_dense_cone_diag_csr_indices_; bool indefinite_Q; cusparse_view_t cusparse_Q_view_; @@ -2346,7 +2290,6 @@ class iteration_data_t { f_t primal_perturb{1e-8}; std::unique_ptr> chol; - std::unique_ptr> augmented_cusparse_view_; bool has_factorization; bool has_solve_info; @@ -2596,10 +2539,6 @@ int barrier_solver_t::initial_point(iteration_data_t& data) #ifdef CHOLESKY_DEBUG_CHECK cholesky_debug_check(data, lp, use_augmented); #endif - if (data.use_csr_ir_matvec()) { - // Strip augmented perturbation for CSR IR matvec path - data.strip_augmented_perturbation(); - } } else { status = data.chol->factorize(data.device_ADAT); } @@ -2640,11 +2579,7 @@ int barrier_solver_t::initial_point(iteration_data_t& data) f_t beta, rmm::device_uvector& y) const { - if (data_.use_csr_ir_matvec()) { - data_.augmented_csr_multiply(alpha, x, beta, y); - } else { - data_.augmented_multiply(alpha, x, beta, y); - } + data_.augmented_multiply(alpha, x, beta, y); } void solve(rmm::device_uvector& b, rmm::device_uvector& x) const { @@ -3168,11 +3103,6 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_tfactorize(data.device_augmented); } - if (data.use_csr_ir_matvec()) { - raft::common::nvtx::range fun_scope("Barrier: strip_augmented_perturbation"); - // Strip augmented perturbation for CSR IR matvec path - data.strip_augmented_perturbation(); - } #ifdef CHOLESKY_DEBUG_CHECK cholesky_debug_check(data, lp, use_augmented); @@ -3274,11 +3204,7 @@ i_t barrier_solver_t::gpu_compute_search_direction(iteration_data_t& y) { - if (data_.use_csr_ir_matvec()) { - data_.augmented_csr_multiply(alpha, x, beta, y); - } else { - data_.augmented_multiply(alpha, x, beta, y); - } + data_.augmented_multiply(alpha, x, beta, y); } void solve(rmm::device_uvector& b, rmm::device_uvector& x) const diff --git a/cpp/src/barrier/cusparse_view.cu b/cpp/src/barrier/cusparse_view.cu index 7c2443f2bc..f787bed8f2 100644 --- a/cpp/src/barrier/cusparse_view.cu +++ b/cpp/src/barrier/cusparse_view.cu @@ -238,54 +238,6 @@ cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, RAFT_CUSPARSE_TRY(cusparseDestroyDnVec(y)); } -template -cusparse_view_t::cusparse_view_t(raft::handle_t const* handle_ptr, - device_csr_matrix_t& matrix) - : 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_(matrix.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 = matrix.n; - const i_t nnz = matrix.nz_max; - - RAFT_CUSPARSE_TRY(cusparseCreateCsr(&A_, - rows_, - cols, - nnz, - matrix.row_start.data(), - matrix.j.data(), - matrix.x.data(), - CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_BASE_ZERO, - CUDA_R_64F)); - - cusparseDnVecDescr_t x; - cusparseDnVecDescr_t y; - rmm::device_uvector d_x(cols, handle_ptr_->get_stream()); - rmm::device_uvector 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 cusparse_view_t::~cusparse_view_t() { diff --git a/cpp/src/barrier/cusparse_view.hpp b/cpp/src/barrier/cusparse_view.hpp index c2f35bd4a6..ea6bf363b9 100644 --- a/cpp/src/barrier/cusparse_view.hpp +++ b/cpp/src/barrier/cusparse_view.hpp @@ -29,9 +29,6 @@ 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& A); - - // Wire cuSparse SpMV over existing device CSR buffers (no copy). Forward SpMV only. - cusparse_view_t(raft::handle_t const* handle_ptr, device_csr_matrix_t& matrix); ~cusparse_view_t(); pdlp::cusparse_dn_vec_descr_wrapper_t create_vector(rmm::device_uvector const& vec); diff --git a/cpp/src/barrier/second_order_cone_kernels.cuh b/cpp/src/barrier/second_order_cone_kernels.cuh index 7a25c5da70..b605b04192 100644 --- a/cpp/src/barrier/second_order_cone_kernels.cuh +++ b/cpp/src/barrier/second_order_cone_kernels.cuh @@ -1083,127 +1083,6 @@ void scatter_sparse_hessian_into_augmented(cone_data_t& cones, RAFT_CUDA_TRY(cudaPeekAtLastError()); } -template -__global__ void strip_linear_primal_diag_kernel(raft::device_span augmented_x, - raft::device_span diag_indices, - raft::device_span q_diag, - raft::device_span is_direct_free_linear, - f_t dual_perturb, - i_t linear_n) -{ - const i_t j = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (j >= linear_n) { return; } - - const i_t idx = diag_indices[j]; - if (is_direct_free_linear[j]) { - const f_t q_val = q_diag.size() > 0 ? q_diag[j] : f_t(0); - augmented_x[idx] = -q_val; - } else { - augmented_x[idx] += dual_perturb; - } -} - -template -__global__ void strip_dual_diag_kernel(raft::device_span augmented_x, - raft::device_span diag_indices, - f_t primal_perturb, - i_t n, - i_t m) -{ - const i_t l = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (l >= m) { return; } - - const i_t row = n + l; - augmented_x[diag_indices[row]] -= primal_perturb; -} - -template -__global__ void strip_indexed_offset_kernel(raft::device_span augmented_x, - raft::device_span csr_indices, - f_t offset) -{ - const size_t e = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (e >= csr_indices.size()) { return; } - - augmented_x[csr_indices[e]] += offset; -} - -template -__global__ void strip_expansion_diag_kernel(raft::device_span augmented_x, - raft::device_span expansion_D, - f_t dual_perturb) -{ - const i_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= static_cast(expansion_D.size())) { return; } - - const f_t sign = (i % 2 == 0) ? f_t(-1) : f_t(1); - augmented_x[expansion_D[i]] -= sign * dual_perturb; -} - -template -void strip_augmented_perturbation_values( - rmm::device_uvector& augmented_x, - const rmm::device_uvector& augmented_diagonal_indices, - const rmm::device_uvector& q_diag, - const rmm::device_uvector& is_direct_free_linear, - const rmm::device_uvector& sparse_hessian_diag, - const rmm::device_uvector& sparse_expansion_D, - const rmm::device_uvector& dense_cone_diag_csr_indices, - i_t linear_n, - i_t n, - i_t m, - f_t dual_perturb, - f_t primal_perturb, - rmm::cuda_stream_view stream) -{ - if (linear_n > 0) { - const size_t grid = raft::ceildiv(linear_n, soc_block_size); - strip_linear_primal_diag_kernel - <<>>(cuopt::make_span(augmented_x), - cuopt::make_span(augmented_diagonal_indices), - cuopt::make_span(q_diag), - cuopt::make_span(is_direct_free_linear), - dual_perturb, - linear_n); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - if (m > 0) { - const size_t grid = raft::ceildiv(m, soc_block_size); - strip_dual_diag_kernel - <<>>(cuopt::make_span(augmented_x), - cuopt::make_span(augmented_diagonal_indices), - primal_perturb, - n, - m); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - if (sparse_hessian_diag.size() > 0) { - const size_t count = sparse_hessian_diag.size(); - const size_t grid = raft::ceildiv(count, soc_block_size); - strip_indexed_offset_kernel<<>>( - cuopt::make_span(augmented_x), cuopt::make_span(sparse_hessian_diag), dual_perturb); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - if (sparse_expansion_D.size() > 0) { - const size_t count = sparse_expansion_D.size(); - const size_t grid = raft::ceildiv(count, soc_block_size); - strip_expansion_diag_kernel<<>>( - cuopt::make_span(augmented_x), cuopt::make_span(sparse_expansion_D), dual_perturb); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - if (dense_cone_diag_csr_indices.size() > 0) { - const size_t count = dense_cone_diag_csr_indices.size(); - const size_t grid = raft::ceildiv(count, soc_block_size); - strip_indexed_offset_kernel<<>>( - cuopt::make_span(augmented_x), cuopt::make_span(dense_cone_diag_csr_indices), dual_perturb); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } -} - /** * Accumulate the sparse-SOC expanded KKT block into a matrix-free product. * diff --git a/cpp/src/dual_simplex/simplex_solver_settings.hpp b/cpp/src/dual_simplex/simplex_solver_settings.hpp index e809b4dded..d0a4457d7f 100644 --- a/cpp/src/dual_simplex/simplex_solver_settings.hpp +++ b/cpp/src/dual_simplex/simplex_solver_settings.hpp @@ -70,7 +70,6 @@ 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), @@ -163,7 +162,6 @@ 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 - bool barrier_csr_ir_matvec; // true to use CSR SpMV for augmented IR matvec 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) diff --git a/cpp/src/pdlp/solve.cu b/cpp/src/pdlp/solve.cu index 712e711099..0993afc6ce 100644 --- a/cpp/src/pdlp/solve.cu +++ b/cpp/src/pdlp/solve.cu @@ -509,7 +509,6 @@ std::tuple, 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.cudss_deterministic = settings.cudss_deterministic; @@ -694,7 +693,6 @@ static optimization_problem_solution_t 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; From d02f67d1220f6617430d25a7c8274772d4e77171 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Fri, 24 Jul 2026 07:53:37 -0700 Subject: [PATCH 14/14] Move CSR KKT buildup to GPU Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 459 +++---------- cpp/src/barrier/csr_kkt_build.cuh | 725 ++++++++++++++++++++ cpp/src/barrier/device_sparse_matrix.cuh | 3 +- cpp/tests/socp/sparse_augmented_kkt_test.cu | 217 ++++++ 4 files changed, 1031 insertions(+), 373 deletions(-) create mode 100644 cpp/src/barrier/csr_kkt_build.cuh diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 06e3f06304..f72df6a3e6 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -268,6 +269,7 @@ class iteration_data_t { 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()), + augmented_csr_metadata_(lp.handle_ptr->get_stream()), d_original_A_values(0, lp.handle_ptr->get_stream()), device_A_x_values(0, lp.handle_ptr->get_stream()), d_inv_diag_prime(0, lp.handle_ptr->get_stream()), @@ -278,8 +280,7 @@ class iteration_data_t { d_augmented_diagonal_indices_(0, lp.handle_ptr->get_stream()), d_cone_csr_indices_(0, lp.handle_ptr->get_stream()), d_cone_Q_values_(0, lp.handle_ptr->get_stream()), - d_dense_cone_block_offsets_(0, lp.handle_ptr->get_stream()), - d_dense_cone_ids_(0, lp.handle_ptr->get_stream()), + d_dense_cone_diag_csr_indices_(0, lp.handle_ptr->get_stream()), d_sparse_hessian_diag_(0, lp.handle_ptr->get_stream()), d_sparse_hessian_Q_(0, lp.handle_ptr->get_stream()), d_sparse_exp_v_col_(0, lp.handle_ptr->get_stream()), @@ -581,6 +582,7 @@ class iteration_data_t { if (n_upper_bounds > 0 || (has_Q && !use_augmented)) { diag.inverse(inv_diag); } // TMP diag and inv_diag should directly created and filled on the GPU raft::copy(d_inv_diag.data(), inv_diag.data(), inv_diag.size(), stream_view_); + raft::copy(d_diag_.data(), diag.data(), diag.size(), stream_view_); inv_sqrt_diag.set_scalar(1.0); if (n_upper_bounds > 0 || (has_Q && !use_augmented)) { inv_diag.sqrt(inv_sqrt_diag); } } @@ -621,6 +623,22 @@ class iteration_data_t { AD.transpose(AT); } + if (use_augmented) { + raft::common::nvtx::range scope("Barrier: augmented: device CSC upload"); + device_A_csc_.emplace(A, handle_ptr->get_stream()); + device_AT_csc_.emplace(AT, handle_ptr->get_stream()); + if (Q.n > 0 && Q.col_start[Q.n] > 0) { + device_Q_csc_.emplace(Q, handle_ptr->get_stream()); + } else { + // Keep an empty but correctly shaped Q so device views are never zero-sized/uninitialized. + device_Q_csc_.emplace(A.n, A.n, 0, handle_ptr->get_stream()); + thrust::fill(rmm::exec_policy(handle_ptr->get_stream()), + device_Q_csc_->col_start.begin(), + device_Q_csc_->col_start.end(), + i_t(0)); + } + } + // 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"); @@ -733,7 +751,6 @@ class iteration_data_t { { i_t n = A.n; i_t m = A.m; - i_t nnzA = A.col_start[n]; i_t nnzQ = Q.n > 0 ? Q.col_start[n] : 0; const bool has_soc = has_cones(); @@ -741,379 +758,74 @@ class iteration_data_t { const i_t p = augmented_expansion_count(); i_t factorization_size = augmented_system_size(n, m); - i_t dense_soc_kkt_nnz = 0; + if (first_call) { + raft::common::nvtx::range scope("Barrier: augmented: device CSR build"); + cuopt_assert( + device_A_csc_.has_value() && device_AT_csc_.has_value() && device_Q_csc_.has_value(), + "augmented CSR build requires device CSC matrices"); - std::vector cone_offsets_host; - std::vector dense_cone_block_offsets_host; - std::vector dense_cone_ids_host; - std::vector dense_cone_idx_by_cone; - std::vector cone_sparse_idx; - std::vector sparse_cone_ids_host; - std::vector sparse_entry_offsets; - std::vector local_to_cone; + const size_t n_sparse_entries = has_soc && p > 0 ? cones().n_sparse_cone_entries : size_t{0}; - if (first_call) { - const i_t n_sparse_soc = has_soc ? cones().n_sparse_cones : i_t(0); if (has_soc) { - raft::common::nvtx::range scope("Barrier: augmented: SOC offsets"); - const i_t n_cones = cone_count(); - // TODO: Once form_augmented() is moved to GPU, we can avoid D2H round-trip. - cone_offsets_host.resize(n_cones + 1); - raft::copy( - cone_offsets_host.data(), cones().cone_offsets.data(), n_cones + 1, stream_view_); - handle_ptr->sync_stream(); - - cone_sparse_idx.assign(n_cones, i_t(-1)); - if (n_sparse_soc > 0) { - // TODO: Once form_augmented() is moved to GPU, we can avoid D2H round-trip for - // sparse_cone_ids / sparse_cone_dims. - sparse_cone_ids_host = cuopt::host_copy(cones().sparse_cone_ids, stream_view_); - for (i_t s = 0; s < n_sparse_soc; ++s) { - cone_sparse_idx[sparse_cone_ids_host[s]] = s; - } - sparse_entry_offsets.resize(n_sparse_soc + 1, 0); - const auto sparse_dims_host = cuopt::host_copy(cones().sparse_cone_dims, stream_view_); - for (i_t s = 0; s < n_sparse_soc; ++s) { - sparse_entry_offsets[s + 1] = sparse_entry_offsets[s] + sparse_dims_host[s]; - } - } - - dense_cone_ids_host.reserve(n_cones); - dense_cone_idx_by_cone.assign(n_cones, i_t(-1)); - dense_cone_block_offsets_host = {0}; - dense_cone_block_offsets_host.reserve(n_cones + 1); - i_t dense_idx = 0; - for (i_t k = 0; k < n_cones; ++k) { - if (cone_sparse_idx[k] < 0) { - dense_cone_idx_by_cone[k] = dense_idx; - dense_cone_ids_host.push_back(k); - const i_t q_k = cone_offsets_host[k + 1] - cone_offsets_host[k]; - dense_cone_block_offsets_host.push_back(dense_cone_block_offsets_host.back() + - q_k * q_k); - ++dense_idx; - } - } - dense_soc_kkt_nnz = static_cast(dense_cone_block_offsets_host.back()); - - local_to_cone.resize(m_c); - for (i_t k = 0; k < n_cones; ++k) { - const i_t lo = cone_offsets_host[k]; - const i_t hi = cone_offsets_host[k + 1]; - for (i_t idx = lo; idx < hi; idx++) { - local_to_cone[idx] = k; - } - } + build_augmented_csr_metadata(cones(), augmented_csr_metadata_, stream_view_); + } else { + augmented_csr_metadata_.dense_soc_kkt_nnz = 0; } - const size_t n_sparse_entries = p > 0 ? cones().n_sparse_cone_entries : size_t{0}; - const i_t sparse_soc_kkt_nnz = static_cast(5 * n_sparse_entries + p); - i_t new_nnz = 2 * nnzA + n + m + p + nnzQ + dense_soc_kkt_nnz + sparse_soc_kkt_nnz; - csr_matrix_t augmented_CSR(factorization_size, factorization_size, new_nnz); - std::vector augmented_diagonal_indices(factorization_size, -1); - std::vector cone_csr_indices_host(dense_soc_kkt_nnz, -1); - std::vector cone_Q_values_host(dense_soc_kkt_nnz, f_t(0)); - std::vector sparse_hessian_diag_host(n_sparse_entries, -1); - std::vector sparse_hessian_Q_host(n_sparse_entries, f_t(0)); - std::vector sparse_exp_v_col_host(n_sparse_entries, -1); - std::vector sparse_exp_u_col_host(n_sparse_entries, -1); - std::vector sparse_exp_v_row_host(n_sparse_entries, -1); - std::vector sparse_exp_u_row_host(n_sparse_entries, -1); - std::vector sparse_expansion_D_host(p, -1); - i_t q = 0; - i_t off_diag_Qnz = 0; - - { - raft::common::nvtx::range scope("Barrier: augmented: host CSR build"); - for (i_t i = 0; i < n; i++) { - augmented_CSR.row_start[i] = q; - - const bool is_cone_row = is_cone_variable(i); - - if (is_cone_row) { - i_t local_idx = i - cone_start(); - i_t k = local_to_cone[local_idx]; - i_t local_r = - static_cast(static_cast(local_idx) - cone_offsets_host[k]); - i_t q_k = static_cast(cone_offsets_host[k + 1] - cone_offsets_host[k]); - i_t cone_col_start = cone_start() + static_cast(cone_offsets_host[k]); - i_t cone_col_end = cone_col_start + q_k; - - i_t qp = (nnzQ > 0) ? Q.col_start[i] : 0; - i_t q_end = (nnzQ > 0) ? Q.col_start[i + 1] : 0; - - // Row of sparse SOC. - if (cone_sparse_idx[k] >= 0) { - const i_t sparse_idx = cone_sparse_idx[k]; - const size_t flat_idx = sparse_entry_offsets[sparse_idx] + local_r; - const i_t exp_v_col = n + m + 2 * sparse_idx; - const i_t exp_u_col = n + m + 2 * sparse_idx + 1; - - // Handle off-diagonal entries of sparse SOC before diagonal entry. - for (; qp < q_end && Q.i[qp] < cone_col_start; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - - // Handle diagonal entry of sparse SOC. - f_t q_contrib = f_t(0); - if (qp < q_end && Q.i[qp] == i) { - q_contrib = Q.x[qp]; - qp++; - } - sparse_hessian_diag_host[flat_idx] = q; - sparse_hessian_Q_host[flat_idx] = q_contrib; - augmented_diagonal_indices[i] = q; - augmented_CSR.j[q] = i; - augmented_CSR.x[q++] = -dual_perturb - q_contrib; - - // Handle expansion columns of sparse SOC. - sparse_exp_v_col_host[flat_idx] = q; - augmented_CSR.j[q] = exp_v_col; - augmented_CSR.x[q++] = f_t(0); - - sparse_exp_u_col_host[flat_idx] = q; - augmented_CSR.j[q] = exp_u_col; - augmented_CSR.x[q++] = f_t(0); - - // Handle off-diagonal entries of dense SOC before diagonal entry. - for (; qp < q_end && Q.i[qp] < cone_col_end; qp++) { - if (Q.i[qp] == i) { continue; } - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - for (; qp < q_end; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - } else { - // Row of dense SOC. - const i_t dense_idx = dense_cone_idx_by_cone[k]; - cuopt_assert(dense_idx >= 0, "dense cone index unset"); - i_t block_base = - static_cast(dense_cone_block_offsets_host[dense_idx]) + local_r * q_k; - - for (; qp < q_end && Q.i[qp] < cone_col_start; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - - for (i_t c = 0; c < q_k; c++) { - i_t col = cone_col_start + c; - f_t q_contrib = f_t(0); - f_t initial_val = (c == local_r) ? f_t(-dual_perturb) : f_t(0); - - if (qp < q_end && Q.i[qp] == col) { - q_contrib = Q.x[qp]; - qp++; - } - - cone_csr_indices_host[block_base + c] = q; - cone_Q_values_host[block_base + c] = q_contrib; - if (col == i) { augmented_diagonal_indices[i] = q; } - augmented_CSR.j[q] = col; - augmented_CSR.x[q++] = initial_val - q_contrib; - } - - for (; qp < q_end; qp++) { - augmented_CSR.j[q] = Q.i[qp]; - augmented_CSR.x[q++] = -Q.x[qp]; - off_diag_Qnz++; - } - } - } else if (nnzQ == 0) { - augmented_diagonal_indices[i] = q; - augmented_CSR.j[q] = i; - augmented_CSR.x[q++] = -diag[i] - dual_perturb; - } else { - // Q is symmetric - const i_t q_col_beg = Q.col_start[i]; - const i_t q_col_end = Q.col_start[i + 1]; - bool has_diagonal = false; - for (i_t p = q_col_beg; p < q_col_end; ++p) { - augmented_CSR.j[q] = Q.i[p]; - if (Q.i[p] == i) { - has_diagonal = true; - augmented_diagonal_indices[i] = q; - augmented_CSR.x[q++] = -Q.x[p] - diag[i] - dual_perturb; - } else { - off_diag_Qnz++; - augmented_CSR.x[q++] = -Q.x[p]; - } - } - if (!has_diagonal) { - augmented_diagonal_indices[i] = q; - augmented_CSR.j[q] = i; - augmented_CSR.x[q++] = -diag[i] - dual_perturb; - } - } - // AT block, we can use A in csc directly - const i_t col_beg = A.col_start[i]; - const i_t col_end = A.col_start[i + 1]; - for (i_t p = col_beg; p < col_end; ++p) { - augmented_CSR.j[q] = A.i[p] + n; - augmented_CSR.x[q++] = A.x[p]; - } - } - - for (i_t k = n; k < n + m; ++k) { - augmented_CSR.row_start[k] = q; - const i_t l = k - n; - const i_t col_beg = AT.col_start[l]; - const i_t col_end = AT.col_start[l + 1]; - for (i_t p = col_beg; p < col_end; ++p) { - augmented_CSR.j[q] = AT.i[p]; - augmented_CSR.x[q++] = AT.x[p]; - } - augmented_diagonal_indices[k] = q; - augmented_CSR.j[q] = k; - augmented_CSR.x[q++] = primal_perturb; - } - - // Handle expansion rows of sparse SOC. - for (i_t s = 0; s < n_sparse_soc; ++s) { - const i_t k = sparse_cone_ids_host[s]; - const i_t q_k = static_cast(cone_offsets_host[k + 1] - cone_offsets_host[k]); - const i_t cone_col_start = cone_start() + static_cast(cone_offsets_host[k]); - const size_t flat_base = sparse_entry_offsets[s]; - const i_t prow_v = n + m + 2 * s; - const i_t prow_u = n + m + 2 * s + 1; - - augmented_CSR.row_start[prow_v] = q; - for (i_t t = 0; t < q_k; ++t) { - sparse_exp_v_row_host[flat_base + t] = q; - augmented_CSR.j[q] = cone_col_start + t; - augmented_CSR.x[q++] = f_t(0); - } - sparse_expansion_D_host[2 * s] = q; - augmented_CSR.j[q] = prow_v; - augmented_CSR.x[q++] = f_t(0); - - augmented_CSR.row_start[prow_u] = q; - for (i_t t = 0; t < q_k; ++t) { - sparse_exp_u_row_host[flat_base + t] = q; - augmented_CSR.j[q] = cone_col_start + t; - augmented_CSR.x[q++] = f_t(0); - } - sparse_expansion_D_host[2 * s + 1] = q; - augmented_CSR.j[q] = prow_u; - augmented_CSR.x[q++] = f_t(0); - } - - augmented_CSR.row_start[factorization_size] = q; - augmented_CSR.nz_max = q; - augmented_CSR.j.resize(q); - augmented_CSR.x.resize(q); - i_t expected_nnz = - 2 * nnzA + (n - m_c) + dense_soc_kkt_nnz + m + off_diag_Qnz + sparse_soc_kkt_nnz; - settings_.log.debug("augmented nz %d predicted %d\n", q, expected_nnz); - cuopt_assert(q == expected_nnz, "augmented nnz != predicted"); - cuopt_assert(A.col_start[n] == AT.col_start[m], "A nz != AT nz"); - - if (n_sparse_entries > 0) { - for (size_t e = 0; e < n_sparse_entries; ++e) { - cuopt_assert(sparse_hessian_diag_host[e] >= 0, - "sparse Hessian diagonal CSR index unset"); - cuopt_assert(sparse_exp_v_col_host[e] >= 0, - "sparse expansion v-column CSR index unset"); - cuopt_assert(sparse_exp_u_col_host[e] >= 0, - "sparse expansion u-column CSR index unset"); - cuopt_assert(sparse_exp_v_row_host[e] >= 0, "sparse expansion v-row CSR index unset"); - cuopt_assert(sparse_exp_u_row_host[e] >= 0, "sparse expansion u-row CSR index unset"); - } - for (i_t s = 0; s < n_sparse_soc; ++s) { - cuopt_assert(sparse_expansion_D_host[2 * s] >= 0, - "sparse expansion v diagonal CSR index unset"); - cuopt_assert(sparse_expansion_D_host[2 * s + 1] >= 0, - "sparse expansion u diagonal CSR index unset"); - } - } + augmented_csr_side_buffers_t side{d_augmented_diagonal_indices_, + d_cone_csr_indices_, + d_cone_Q_values_, + d_dense_cone_diag_csr_indices_, + d_sparse_hessian_diag_, + d_sparse_hessian_Q_, + d_sparse_exp_v_col_, + d_sparse_exp_u_col_, + d_sparse_exp_v_row_, + d_sparse_exp_u_row_, + d_sparse_expansion_D_}; + + raft::device_span element_cone_ids_span{}; + raft::device_span cone_offsets_span{}; + raft::device_span sparse_cone_ids_span{}; + raft::device_span sparse_entry_offsets_span{}; + size_t n_sparse_cone_entries = 0; + if (has_soc) { + element_cone_ids_span = cuopt::make_span(cones().element_cone_ids); + cone_offsets_span = cuopt::make_span(cones().cone_offsets); + sparse_cone_ids_span = cuopt::make_span(cones().sparse_cone_ids); + sparse_entry_offsets_span = cuopt::make_span(cones().sparse_entry_offsets); + n_sparse_cone_entries = cones().n_sparse_cone_entries; } - { - raft::common::nvtx::range scope("Barrier: augmented: device upload"); - device_augmented.copy(augmented_CSR, handle_ptr->get_stream()); - d_augmented_diagonal_indices_.resize(augmented_diagonal_indices.size(), - handle_ptr->get_stream()); - raft::copy(d_augmented_diagonal_indices_.data(), - augmented_diagonal_indices.data(), - augmented_diagonal_indices.size(), - handle_ptr->get_stream()); - - if (has_soc) { - d_cone_csr_indices_.resize(dense_soc_kkt_nnz, handle_ptr->get_stream()); - raft::copy(d_cone_csr_indices_.data(), - cone_csr_indices_host.data(), - dense_soc_kkt_nnz, - handle_ptr->get_stream()); - d_cone_Q_values_.resize(dense_soc_kkt_nnz, handle_ptr->get_stream()); - raft::copy(d_cone_Q_values_.data(), - cone_Q_values_host.data(), - dense_soc_kkt_nnz, - handle_ptr->get_stream()); - - d_dense_cone_block_offsets_.resize(dense_cone_block_offsets_host.size(), - handle_ptr->get_stream()); - raft::copy(d_dense_cone_block_offsets_.data(), - dense_cone_block_offsets_host.data(), - dense_cone_block_offsets_host.size(), - handle_ptr->get_stream()); - d_dense_cone_ids_.resize(dense_cone_ids_host.size(), handle_ptr->get_stream()); - if (!dense_cone_ids_host.empty()) { - raft::copy(d_dense_cone_ids_.data(), - dense_cone_ids_host.data(), - dense_cone_ids_host.size(), - handle_ptr->get_stream()); - } - - const size_t n_sparse_entries = cones().n_sparse_cone_entries; - d_sparse_hessian_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_hessian_Q_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_exp_v_col_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_exp_u_col_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_exp_v_row_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_exp_u_row_.resize(n_sparse_entries, handle_ptr->get_stream()); - d_sparse_expansion_D_.resize(p, handle_ptr->get_stream()); - d_sparse_Hs_diag_.resize(n_sparse_entries, handle_ptr->get_stream()); - if (n_sparse_entries > 0) { - raft::copy(d_sparse_hessian_diag_.data(), - sparse_hessian_diag_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - raft::copy(d_sparse_hessian_Q_.data(), - sparse_hessian_Q_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - raft::copy(d_sparse_exp_v_col_.data(), - sparse_exp_v_col_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - raft::copy(d_sparse_exp_u_col_.data(), - sparse_exp_u_col_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - raft::copy(d_sparse_exp_v_row_.data(), - sparse_exp_v_row_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - raft::copy(d_sparse_exp_u_row_.data(), - sparse_exp_u_row_host.data(), - n_sparse_entries, - handle_ptr->get_stream()); - } - if (p > 0) { - raft::copy(d_sparse_expansion_D_.data(), - sparse_expansion_D_host.data(), - sparse_expansion_D_host.size(), - handle_ptr->get_stream()); - } - } + if (n_sparse_entries > 0) { d_sparse_Hs_diag_.resize(n_sparse_entries, stream_view_); } + + const i_t total_nnz = + build_augmented_csr_on_device(n, + m, + p, + cone_start(), + m_c, + nnzQ, + dual_perturb, + primal_perturb, + *device_A_csc_, + *device_Q_csc_, + *device_AT_csc_, + raft::device_span{d_diag_.data(), d_diag_.size()}, + element_cone_ids_span, + cone_offsets_span, + sparse_cone_ids_span, + sparse_entry_offsets_span, + n_sparse_cone_entries, + augmented_csr_metadata_, + device_augmented, + side, + 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(); - handle_ptr->sync_stream(); - } #ifdef CHECK_SYMMETRY csc_matrix_t augmented_transpose(1, 1, 1); augmented.transpose(augmented_transpose); @@ -1179,8 +891,8 @@ class iteration_data_t { device_augmented.x, d_cone_csr_indices_, d_cone_Q_values_, - d_dense_cone_block_offsets_, - d_dense_cone_ids_, + augmented_csr_metadata_.dense_block_offsets, + augmented_csr_metadata_.dense_cone_ids, handle_ptr->get_stream(), dual_perturb); RAFT_CHECK_CUDA(handle_ptr->get_stream()); @@ -2234,6 +1946,10 @@ class iteration_data_t { csc_matrix_t ADAT; // csc_matrix_t augmented; device_csr_matrix_t device_augmented; + std::optional> device_A_csc_; + std::optional> device_Q_csc_; + std::optional> device_AT_csc_; + augmented_csr_metadata_t augmented_csr_metadata_; device_csr_matrix_t device_ADAT; device_csr_matrix_t device_A; @@ -2263,8 +1979,7 @@ class iteration_data_t { rmm::device_uvector d_augmented_diagonal_indices_; rmm::device_uvector d_cone_csr_indices_; rmm::device_uvector d_cone_Q_values_; - rmm::device_uvector d_dense_cone_block_offsets_; - rmm::device_uvector d_dense_cone_ids_; + rmm::device_uvector d_dense_cone_diag_csr_indices_; rmm::device_uvector d_sparse_hessian_diag_; rmm::device_uvector d_sparse_hessian_Q_; rmm::device_uvector d_sparse_exp_v_col_; diff --git a/cpp/src/barrier/csr_kkt_build.cuh b/cpp/src/barrier/csr_kkt_build.cuh new file mode 100644 index 0000000000..7efbbd466e --- /dev/null +++ b/cpp/src/barrier/csr_kkt_build.cuh @@ -0,0 +1,725 @@ +/* clang-format off */ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +/* clang-format on */ +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +#include + +namespace cuopt::mathematical_optimization::barrier { + +constexpr int augmented_csr_block_size = 256; + +template +struct augmented_csr_metadata_t { + explicit augmented_csr_metadata_t(rmm::cuda_stream_view stream) + : sparse_ids_by_cone(0, stream), + dense_ids_by_cone(0, stream), + dense_cone_entry_rank(0, stream), + dense_block_offsets(0, stream), + dense_cone_ids(0, stream) + { + } + + rmm::device_uvector sparse_ids_by_cone; + rmm::device_uvector dense_ids_by_cone; + rmm::device_uvector dense_cone_entry_rank; + rmm::device_uvector dense_block_offsets; + rmm::device_uvector dense_cone_ids; + i_t dense_soc_kkt_nnz{0}; +}; + +// Host-side handle referencing the persistent side-index device buffers. Used to resize them on +// the host before the fill kernel runs. rmm::device_uvector is a host-only type, so these +// references must NOT be dereferenced inside a kernel. +template +struct augmented_csr_side_buffers_t { + rmm::device_uvector& augmented_diagonal_indices; + rmm::device_uvector& cone_csr_indices; + rmm::device_uvector& cone_Q_values; + rmm::device_uvector& dense_cone_diag_csr_indices; + rmm::device_uvector& sparse_hessian_diag; + rmm::device_uvector& sparse_hessian_Q; + rmm::device_uvector& sparse_exp_v_col; + rmm::device_uvector& sparse_exp_u_col; + rmm::device_uvector& sparse_exp_v_row; + rmm::device_uvector& sparse_exp_u_row; + rmm::device_uvector& sparse_expansion_D; +}; + +// Device-usable view of the side-index buffers (plain spans, safe to pass to kernels). +template +struct augmented_csr_side_views_t { + raft::device_span augmented_diagonal_indices; + raft::device_span cone_csr_indices; + raft::device_span cone_Q_values; + raft::device_span dense_cone_diag_csr_indices; + raft::device_span sparse_hessian_diag; + raft::device_span sparse_hessian_Q; + raft::device_span sparse_exp_v_col; + raft::device_span sparse_exp_u_col; + raft::device_span sparse_exp_v_row; + raft::device_span sparse_exp_u_row; + raft::device_span sparse_expansion_D; +}; + +template +augmented_csr_side_views_t make_side_views(augmented_csr_side_buffers_t& b) +{ + return augmented_csr_side_views_t{cuopt::make_span(b.augmented_diagonal_indices), + cuopt::make_span(b.cone_csr_indices), + cuopt::make_span(b.cone_Q_values), + cuopt::make_span(b.dense_cone_diag_csr_indices), + cuopt::make_span(b.sparse_hessian_diag), + cuopt::make_span(b.sparse_hessian_Q), + cuopt::make_span(b.sparse_exp_v_col), + cuopt::make_span(b.sparse_exp_u_col), + cuopt::make_span(b.sparse_exp_v_row), + cuopt::make_span(b.sparse_exp_u_row), + cuopt::make_span(b.sparse_expansion_D)}; +} + +template +__global__ void scatter_sparse_ids_by_cone_kernel(raft::device_span sparse_ids_by_cone, + raft::device_span sparse_cone_ids, + i_t n_sparse) +{ + const i_t s = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (s >= n_sparse) { return; } + sparse_ids_by_cone[sparse_cone_ids[s]] = s; +} + +template +__global__ void build_dense_ids_by_cone_kernel(raft::device_span dense_ids_by_cone, + raft::device_span cone_is_sparse, + raft::device_span dense_prefix, + i_t n_cones) +{ + const i_t k = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (k >= n_cones) { return; } + dense_ids_by_cone[k] = cone_is_sparse[k] ? i_t(-1) : dense_prefix[k]; +} + +template +__global__ void compact_dense_cone_ids_kernel(raft::device_span dense_cone_ids, + raft::device_span dense_prefix, + raft::device_span cone_is_sparse, + i_t n_cones) +{ + const i_t k = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (k >= n_cones || cone_is_sparse[k]) { return; } + dense_cone_ids[dense_prefix[k]] = k; +} + +template +__global__ void build_dense_block_sizes_kernel(raft::device_span dense_block_sizes, + raft::device_span dense_cone_ids, + raft::device_span cone_offsets, + i_t n_dense) +{ + const i_t d = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (d >= n_dense) { return; } + const i_t k = dense_cone_ids[d]; + const size_t q = cone_offsets[k + 1] - cone_offsets[k]; + dense_block_sizes[d] = q * q; +} + +template +__global__ void build_dense_cone_entry_rank_kernel(raft::device_span dense_cone_entry_rank, + raft::device_span element_cone_ids, + raft::device_span cone_is_sparse, + raft::device_span dense_entry_prefix, + i_t n_cone_entries) +{ + const i_t idx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= n_cone_entries) { return; } + const i_t cone = element_cone_ids[idx]; + dense_cone_entry_rank[idx] = cone_is_sparse[cone] ? i_t(-1) : dense_entry_prefix[idx]; +} + +// Count nnz of a cone primal row's Hessian part. The Q entries before and after +// the cone column block [cone_col_start, cone_col_end) are counted identically +// for sparse and dense cones; only the in-block contribution differs: +// - dense : a full q_k-wide dense block (in-block Q entries are subsumed by it). +// - sparse: the 3 structural entries (diagonal + two expansion columns v, u), +// plus the in-block off-diagonal Q entries. +template +__device__ i_t count_q_cone(const csc_view_t& Q, + i_t row, + i_t cone_col_start, + i_t cone_col_end, + i_t q_k, + bool is_sparse, + i_t nnzQ) +{ + i_t count = is_sparse ? i_t(3) : q_k; + if (nnzQ == 0) { return count; } + i_t qp = Q.col_start[row]; + const i_t q_end = Q.col_start[row + 1]; + // Q entries left of the cone block. + for (; qp < q_end && Q.i[qp] < cone_col_start; ++qp) { + ++count; + } + // In-block Q entries: dense subsumes them in the q_k columns, so they add + // nothing; sparse counts the off-diagonals and skips the diagonal (col == row). + for (; qp < q_end && Q.i[qp] < cone_col_end; ++qp) { + if (is_sparse && Q.i[qp] != row) { ++count; } + } + // Q entries right of the cone block. + count += q_end - qp; + + return count; +} + +template +__device__ i_t count_primal_row_nnz_device(i_t row, + i_t cone_start, + i_t m_c, + i_t nnzQ, + const csc_view_t& A, + const csc_view_t& Q, + raft::device_span element_cone_ids, + raft::device_span cone_offsets, + raft::device_span sparse_ids_by_cone) +{ + i_t count = 0; + const bool is_cone_row = (row >= cone_start) && (row < cone_start + m_c); + if (is_cone_row) { + const i_t local_idx = row - cone_start; + const i_t k = element_cone_ids[local_idx]; + const i_t q_k = static_cast(cone_offsets[k + 1] - cone_offsets[k]); + const i_t cone_col_start = cone_start + static_cast(cone_offsets[k]); + const i_t cone_col_end = cone_col_start + q_k; + const bool is_sparse = sparse_ids_by_cone[k] >= 0; + count += count_q_cone(Q, row, cone_col_start, cone_col_end, q_k, is_sparse, nnzQ); + } else if (nnzQ == 0) { + ++count; + } else { + const i_t q_col_beg = Q.col_start[row]; + const i_t q_col_end = Q.col_start[row + 1]; + count += q_col_end - q_col_beg; + bool has_diagonal = false; + for (i_t qp = q_col_beg; qp < q_col_end; ++qp) { + if (Q.i[qp] == row) { + has_diagonal = true; + break; + } + } + if (!has_diagonal) { ++count; } + } + count += A.col_start[row + 1] - A.col_start[row]; + return count; +} + +template +__global__ void count_augmented_row_nnz_kernel(i_t factorization_size, + i_t n, + i_t m, + i_t cone_start, + i_t m_c, + i_t nnzQ, + csc_view_t A, + csc_view_t Q, + csc_view_t AT, + raft::device_span element_cone_ids, + raft::device_span cone_offsets, + raft::device_span sparse_cone_ids, + raft::device_span sparse_ids_by_cone, + raft::device_span row_nnz) +{ + const i_t row = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (row >= factorization_size) { return; } + + if (row < n) { + row_nnz[row] = count_primal_row_nnz_device( + row, cone_start, m_c, nnzQ, A, Q, element_cone_ids, cone_offsets, sparse_ids_by_cone); + return; + } + + if (row < n + m) { + const i_t l = row - n; + row_nnz[row] = AT.col_start[l + 1] - AT.col_start[l] + 1; + return; + } + + const i_t exp_row = row - (n + m); + const i_t s = exp_row / 2; + const i_t k = sparse_cone_ids[s]; + const i_t q_k = static_cast(cone_offsets[k + 1] - cone_offsets[k]); + row_nnz[row] = q_k + 1; +} + +template +__global__ void fill_augmented_csr_row_kernel(i_t factorization_size, + i_t n, + i_t m, + i_t p, + i_t cone_start, + i_t m_c, + i_t nnzQ, + f_t dual_perturb, + f_t primal_perturb, + csc_view_t A, + csc_view_t Q, + csc_view_t AT, + raft::device_span diag, + raft::device_span row_start, + raft::device_span j, + raft::device_span x, + raft::device_span element_cone_ids, + raft::device_span cone_offsets, + raft::device_span sparse_cone_ids, + raft::device_span sparse_ids_by_cone, + raft::device_span dense_ids_by_cone, + raft::device_span dense_block_offsets, + raft::device_span dense_cone_entry_rank, + raft::device_span sparse_entry_offsets, + augmented_csr_side_views_t side) +{ + const i_t row = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (row >= factorization_size) { return; } + + i_t q = row_start[row]; + + if (row < n) { + const bool is_cone_row = (row >= cone_start) && (row < cone_start + m_c); + + // Sparse-cone expansion columns are the largest column indices in the row + // (>= n + m), so to keep the row column-sorted they are emitted after the + // A^T block below. exp_v_col >= 0 marks extra columns to emit. + size_t exp_flat_idx = 0; + i_t exp_v_col = -1; + i_t exp_u_col = 0; + + if (is_cone_row) { + const i_t local_idx = row - cone_start; + const i_t k = element_cone_ids[local_idx]; + const i_t local_r = static_cast(static_cast(local_idx) - cone_offsets[k]); + const i_t q_k = static_cast(cone_offsets[k + 1] - cone_offsets[k]); + const i_t cone_col_start = cone_start + static_cast(cone_offsets[k]); + const i_t sparse_idx = sparse_ids_by_cone[k]; + + if (sparse_idx >= 0) { + // Column-sorted Hessian region: Q entries with column < row, then the + // diagonal (folding Q's diagonal if present), then Q entries with + // column > row. + i_t qp = Q.col_start[row]; + const i_t q_end = Q.col_start[row + 1]; + + for (; qp < q_end && Q.i[qp] < row; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + + f_t q_contrib = f_t(0); + if (qp < q_end && Q.i[qp] == row) { + q_contrib = Q.x[qp]; + ++qp; + } + const size_t flat_idx = sparse_entry_offsets[sparse_idx] + local_r; + side.sparse_hessian_diag[flat_idx] = q; + side.sparse_hessian_Q[flat_idx] = q_contrib; + side.augmented_diagonal_indices[row] = q; + j[q] = row; + x[q++] = -dual_perturb - q_contrib; + + for (; qp < q_end; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + + // Defer the expansion columns (largest columns) until after the A^T block. + exp_flat_idx = flat_idx; + exp_v_col = n + m + 2 * sparse_idx; + exp_u_col = n + m + 2 * sparse_idx + 1; + } else { + const i_t dense_idx = dense_ids_by_cone[k]; + const i_t block_base = static_cast(dense_block_offsets[dense_idx]) + local_r * q_k; + + if (nnzQ > 0) { + i_t qp = Q.col_start[row]; + const i_t q_end = Q.col_start[row + 1]; + for (; qp < q_end && Q.i[qp] < cone_col_start; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + for (i_t c = 0; c < q_k; ++c) { + const i_t col = cone_col_start + c; + f_t q_contrib = f_t(0); + const f_t initial_val = (c == local_r) ? f_t(-dual_perturb) : f_t(0); + if (qp < q_end && Q.i[qp] == col) { + q_contrib = Q.x[qp]; + ++qp; + } + side.cone_csr_indices[block_base + c] = q; + side.cone_Q_values[block_base + c] = q_contrib; + if (col == row) { + side.augmented_diagonal_indices[row] = q; + const i_t dense_rank = dense_cone_entry_rank[local_idx]; + if (dense_rank >= 0) { side.dense_cone_diag_csr_indices[dense_rank] = q; } + } + j[q] = col; + x[q++] = initial_val - q_contrib; + } + for (; qp < q_end; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + } else { + for (i_t c = 0; c < q_k; ++c) { + const i_t col = cone_col_start + c; + const f_t initial_val = (c == local_r) ? f_t(-dual_perturb) : f_t(0); + side.cone_csr_indices[block_base + c] = q; + side.cone_Q_values[block_base + c] = f_t(0); + if (col == row) { + side.augmented_diagonal_indices[row] = q; + const i_t dense_rank = dense_cone_entry_rank[local_idx]; + if (dense_rank >= 0) { side.dense_cone_diag_csr_indices[dense_rank] = q; } + } + j[q] = col; + x[q++] = initial_val; + } + } + } + } else if (nnzQ == 0) { + side.augmented_diagonal_indices[row] = q; + j[q] = row; + x[q++] = -diag[row] - dual_perturb; + } else { + // Column-sorted: Q entries with column < row, the diagonal in place + // (folding Q's diagonal if present), then Q entries with column > row. + i_t qp = Q.col_start[row]; + const i_t q_end = Q.col_start[row + 1]; + for (; qp < q_end && Q.i[qp] < row; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + f_t q_diag = f_t(0); + if (qp < q_end && Q.i[qp] == row) { + q_diag = Q.x[qp]; + ++qp; + } + side.augmented_diagonal_indices[row] = q; + j[q] = row; + x[q++] = -q_diag - diag[row] - dual_perturb; + for (; qp < q_end; ++qp) { + j[q] = Q.i[qp]; + x[q++] = -Q.x[qp]; + } + } + + // A^T block columns in [n, n + m). + const i_t col_beg = A.col_start[row]; + const i_t col_end = A.col_start[row + 1]; + for (i_t p_idx = col_beg; p_idx < col_end; ++p_idx) { + j[q] = A.i[p_idx] + n; + x[q++] = A.x[p_idx]; + } + + // Sparse-cone expansion columns (columns >= n + m). + if (exp_v_col >= 0) { + side.sparse_exp_v_col[exp_flat_idx] = q; + j[q] = exp_v_col; + x[q++] = f_t(0); + side.sparse_exp_u_col[exp_flat_idx] = q; + j[q] = exp_u_col; + x[q++] = f_t(0); + } + return; + } + + // Fill A row and the corresponding augmented diagonal entry. + if (row < n + m) { + const i_t l = row - n; + const i_t col_beg = AT.col_start[l]; + const i_t col_end = AT.col_start[l + 1]; + for (i_t idx = col_beg; idx < col_end; ++idx) { + j[q] = AT.i[idx]; + x[q++] = AT.x[idx]; + } + side.augmented_diagonal_indices[row] = q; + j[q] = row; + x[q++] = primal_perturb; + return; + } + + // Fill the expansion column and D slot. + const i_t exp_row = row - (n + m); + const i_t s = exp_row / 2; + const i_t k = sparse_cone_ids[s]; + const i_t q_k = static_cast(cone_offsets[k + 1] - cone_offsets[k]); + const i_t cone_col_start = cone_start + static_cast(cone_offsets[k]); + const size_t flat_base = sparse_entry_offsets[s]; + + raft::device_span exp_row_idx = + (exp_row % 2 == 0) ? side.sparse_exp_v_row : side.sparse_exp_u_row; + for (i_t jj = 0; jj < q_k; ++jj) { + exp_row_idx[flat_base + jj] = q; + j[q] = cone_col_start + jj; + x[q++] = f_t(0); + } + side.sparse_expansion_D[exp_row] = q; + j[q] = n + m + exp_row; + x[q++] = f_t(0); +} + +template +void build_augmented_csr_metadata(const cone_data_t& cones, + augmented_csr_metadata_t& metadata, + rmm::cuda_stream_view stream) +{ + raft::common::nvtx::range scope("Barrier: augmented: device CSR metadata"); + const i_t n_cones = cones.n_cones; + const i_t n_dense = cones.n_dense_cones(); + const i_t n_sparse = cones.n_sparse_cones; + const size_t m_c = cones.n_cone_entries; + + metadata.sparse_ids_by_cone.resize(n_cones, stream); + thrust::fill(rmm::exec_policy(stream), + metadata.sparse_ids_by_cone.begin(), + metadata.sparse_ids_by_cone.end(), + i_t(-1)); + if (n_sparse > 0) { + const size_t grid = raft::ceildiv(n_sparse, augmented_csr_block_size); + scatter_sparse_ids_by_cone_kernel<<>>( + cuopt::make_span(metadata.sparse_ids_by_cone), + cuopt::make_span(cones.sparse_cone_ids), + n_sparse); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + metadata.dense_ids_by_cone.resize(n_cones, stream); + metadata.dense_cone_ids.resize(n_dense, stream); + metadata.dense_block_offsets.resize(static_cast(n_dense) + 1, stream); + + if (n_dense > 0) { + rmm::device_uvector is_dense_cone(n_cones, stream); + thrust::transform(rmm::exec_policy(stream), + cones.cone_is_sparse.begin(), + cones.cone_is_sparse.end(), + is_dense_cone.begin(), + [] __device__(i_t is_sparse) { return is_sparse ? i_t(0) : i_t(1); }); + + rmm::device_uvector dense_prefix(n_cones, stream); + thrust::exclusive_scan( + rmm::exec_policy(stream), is_dense_cone.begin(), is_dense_cone.end(), dense_prefix.begin()); + + const size_t grid = raft::ceildiv(n_cones, augmented_csr_block_size); + build_dense_ids_by_cone_kernel<<>>( + cuopt::make_span(metadata.dense_ids_by_cone), + cuopt::make_span(cones.cone_is_sparse), + cuopt::make_span(dense_prefix), + n_cones); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + compact_dense_cone_ids_kernel<<>>( + cuopt::make_span(metadata.dense_cone_ids), + cuopt::make_span(dense_prefix), + cuopt::make_span(cones.cone_is_sparse), + n_cones); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + rmm::device_uvector dense_block_sizes(n_dense, stream); + const size_t dense_grid = raft::ceildiv(n_dense, augmented_csr_block_size); + build_dense_block_sizes_kernel + <<>>( + cuopt::make_span(dense_block_sizes), + cuopt::make_span(metadata.dense_cone_ids), + cuopt::make_span(cones.cone_offsets), + n_dense); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + thrust::exclusive_scan(rmm::exec_policy(stream), + dense_block_sizes.begin(), + dense_block_sizes.end(), + metadata.dense_block_offsets.begin()); + // exclusive_scan writes only n_dense entries; the final offset (total block nnz) is not + // produced by the scan, so compute and store it explicitly before reading it back. + const size_t total_block_nnz = metadata.dense_block_offsets.element(n_dense - 1, stream) + + dense_block_sizes.element(n_dense - 1, stream); + metadata.dense_block_offsets.set_element_async(n_dense, total_block_nnz, stream); + metadata.dense_soc_kkt_nnz = static_cast(total_block_nnz); + } else { + metadata.dense_soc_kkt_nnz = 0; + metadata.dense_block_offsets.set_element_to_zero_async(0, stream); + } + + metadata.dense_cone_entry_rank.resize(m_c, stream); + if (m_c > 0) { + rmm::device_uvector is_dense_entry(m_c, stream); + // Each cone entry is dense iff its owning cone is dense. Map entry -> cone via + // element_cone_ids, then look up cone_is_sparse[cone]. NOTE: this must be a unary + // transform over element_cone_ids (length m_c); zipping directly against + // cone_is_sparse.begin() would read cone_is_sparse (length n_cones < m_c) out of bounds + // and index it by entry position instead of by cone. + const i_t* cone_is_sparse_ptr = cones.cone_is_sparse.data(); + thrust::transform(rmm::exec_policy(stream), + cones.element_cone_ids.begin(), + cones.element_cone_ids.end(), + is_dense_entry.begin(), + [cone_is_sparse_ptr] __device__(i_t cone) { + return cone_is_sparse_ptr[cone] ? i_t(0) : i_t(1); + }); + + rmm::device_uvector dense_entry_prefix(m_c, stream); + thrust::exclusive_scan(rmm::exec_policy(stream), + is_dense_entry.begin(), + is_dense_entry.end(), + dense_entry_prefix.begin()); + + const size_t entry_grid = raft::ceildiv(m_c, augmented_csr_block_size); + build_dense_cone_entry_rank_kernel + <<>>( + cuopt::make_span(metadata.dense_cone_entry_rank), + cuopt::make_span(cones.element_cone_ids), + cuopt::make_span(cones.cone_is_sparse), + cuopt::make_span(dense_entry_prefix), + static_cast(m_c)); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +template +i_t build_augmented_csr_on_device(i_t n, + i_t m, + i_t p, + i_t cone_start, + i_t m_c, + i_t nnzQ, + f_t dual_perturb, + f_t primal_perturb, + device_csc_matrix_t& A, + device_csc_matrix_t& Q, + device_csc_matrix_t& AT, + raft::device_span diag, + raft::device_span element_cone_ids, + raft::device_span cone_offsets, + raft::device_span sparse_cone_ids, + raft::device_span sparse_entry_offsets, + size_t n_sparse_cone_entries, + augmented_csr_metadata_t& metadata, + device_csr_matrix_t& device_augmented, + augmented_csr_side_buffers_t side, + rmm::cuda_stream_view stream) +{ + const i_t factorization_size = n + m + p; + const csc_view_t A_view = A.view(); + const csc_view_t Q_view = Q.view(); + const csc_view_t AT_view = AT.view(); + + rmm::device_uvector row_nnz(factorization_size, stream); + { + raft::common::nvtx::range scope("Barrier: augmented: device CSR count"); + const size_t grid = raft::ceildiv(factorization_size, augmented_csr_block_size); + count_augmented_row_nnz_kernel<<>>( + factorization_size, + n, + m, + cone_start, + m_c, + nnzQ, + A_view, + Q_view, + AT_view, + element_cone_ids, + cone_offsets, + sparse_cone_ids, + cuopt::make_span(metadata.sparse_ids_by_cone), + cuopt::make_span(row_nnz)); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + i_t total_nnz = 0; + { + raft::common::nvtx::range scope("Barrier: augmented: device CSR scan"); + device_augmented.m = factorization_size; + device_augmented.n = factorization_size; + device_augmented.row_start.resize(static_cast(factorization_size) + 1, stream); + + // Inclusive scan of the per-row counts into row_start[1..], with row_start[0] = 0. + device_augmented.row_start.set_element_to_zero_async(0, stream); + thrust::inclusive_scan(rmm::exec_policy(stream), + row_nnz.begin(), + row_nnz.end(), + device_augmented.row_start.begin() + 1); + + total_nnz = device_augmented.row_start.element(factorization_size, stream); + device_augmented.nz_max = total_nnz; + device_augmented.j.resize(total_nnz, stream); + device_augmented.x.resize(total_nnz, stream); + } + + // Resize a side buffer and initialize every element to a sentinel in one shot. + // A size of 0 makes both the resize and the fill no-ops. + auto resize_and_fill = [&](auto& buf, size_t size, auto value) { + buf.resize(size, stream); + thrust::fill(rmm::exec_policy(stream), buf.begin(), buf.end(), value); + }; + + resize_and_fill(side.augmented_diagonal_indices, factorization_size, i_t(-1)); + resize_and_fill(side.cone_csr_indices, metadata.dense_soc_kkt_nnz, i_t(-1)); + resize_and_fill(side.cone_Q_values, metadata.dense_soc_kkt_nnz, f_t(0)); + + const size_t n_sparse_entries = n_sparse_cone_entries; + resize_and_fill(side.sparse_hessian_diag, n_sparse_entries, i_t(-1)); + resize_and_fill(side.sparse_hessian_Q, n_sparse_entries, f_t(0)); + resize_and_fill(side.sparse_exp_v_col, n_sparse_entries, i_t(-1)); + resize_and_fill(side.sparse_exp_u_col, n_sparse_entries, i_t(-1)); + resize_and_fill(side.sparse_exp_v_row, n_sparse_entries, i_t(-1)); + resize_and_fill(side.sparse_exp_u_row, n_sparse_entries, i_t(-1)); + resize_and_fill(side.sparse_expansion_D, p, i_t(-1)); + + const i_t n_dense_entries = static_cast(m_c) - static_cast(n_sparse_cone_entries); + resize_and_fill(side.dense_cone_diag_csr_indices, std::max(0, n_dense_entries), i_t(-1)); + + { + raft::common::nvtx::range scope("Barrier: augmented: device CSR fill"); + auto side_views = make_side_views(side); + const size_t grid = raft::ceildiv(factorization_size, augmented_csr_block_size); + fill_augmented_csr_row_kernel<<>>( + factorization_size, + n, + m, + p, + cone_start, + m_c, + nnzQ, + dual_perturb, + primal_perturb, + A_view, + Q_view, + AT_view, + diag, + cuopt::make_span(device_augmented.row_start), + cuopt::make_span(device_augmented.j), + cuopt::make_span(device_augmented.x), + element_cone_ids, + cone_offsets, + sparse_cone_ids, + cuopt::make_span(metadata.sparse_ids_by_cone), + cuopt::make_span(metadata.dense_ids_by_cone), + cuopt::make_span(metadata.dense_block_offsets), + cuopt::make_span(metadata.dense_cone_entry_rank), + sparse_entry_offsets, + side_views); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + return total_nnz; +} + +} // namespace cuopt::mathematical_optimization::barrier diff --git a/cpp/src/barrier/device_sparse_matrix.cuh b/cpp/src/barrier/device_sparse_matrix.cuh index 71e86484a6..e74df80785 100644 --- a/cpp/src/barrier/device_sparse_matrix.cuh +++ b/cpp/src/barrier/device_sparse_matrix.cuh @@ -185,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); diff --git a/cpp/tests/socp/sparse_augmented_kkt_test.cu b/cpp/tests/socp/sparse_augmented_kkt_test.cu index 9d9d9eaac9..14195dae35 100644 --- a/cpp/tests/socp/sparse_augmented_kkt_test.cu +++ b/cpp/tests/socp/sparse_augmented_kkt_test.cu @@ -5,6 +5,8 @@ */ /* clang-format on */ +#include +#include #include #include @@ -346,4 +348,219 @@ TEST(sparse_augmented_kkt, update_scaling_sparse_dim_1000) EXPECT_NEAR(yexp_host[1], eta_sq * x_exp_u + dot_u, 1e-9); } +TEST(sparse_augmented_kkt, gpu_augmented_csr_metadata_matches_host) +{ + auto stream = rmm::cuda_stream_default; + + std::vector cone_dimensions{3, 6, 5}; + rmm::device_uvector x(14, stream); + rmm::device_uvector z(14, stream); + cone_data_t cones( + cone_dimensions, cuopt::make_span(x), cuopt::make_span(z), stream, /*soc_threshold=*/4); + + augmented_csr_metadata_t metadata(stream); + build_augmented_csr_metadata(cones, metadata, stream); + + auto sparse_idx_host = cuopt::host_copy(metadata.sparse_ids_by_cone, stream); + auto dense_idx_host = cuopt::host_copy(metadata.dense_ids_by_cone, stream); + auto dense_ids_host = cuopt::host_copy(metadata.dense_cone_ids, stream); + auto block_offsets_host = cuopt::host_copy(metadata.dense_block_offsets, stream); + + std::vector expected_sparse_idx(cones.n_cones, -1); + auto sparse_cone_ids_host = cuopt::host_copy(cones.sparse_cone_ids, stream); + for (int s = 0; s < cones.n_sparse_cones; ++s) { + expected_sparse_idx[sparse_cone_ids_host[s]] = s; + } + ASSERT_EQ(expected_sparse_idx.size(), sparse_idx_host.size()); + for (size_t e = 0; e < expected_sparse_idx.size(); ++e) { + EXPECT_EQ(expected_sparse_idx[e], sparse_idx_host[e]) << "sparse_ids_by_cone index " << e; + } + + int dense_count = 0; + auto cone_is_sparse_host = cuopt::host_copy(cones.cone_is_sparse, stream); + for (int k = 0; k < cones.n_cones; ++k) { + if (cone_is_sparse_host[k] == 0) { + EXPECT_EQ(dense_idx_host[k], dense_count); + ASSERT_LT(dense_count, static_cast(dense_ids_host.size())); + EXPECT_EQ(dense_ids_host[dense_count], k); + ++dense_count; + } else { + EXPECT_EQ(dense_idx_host[k], -1); + } + } + EXPECT_EQ(dense_count, cones.n_dense_cones()); + EXPECT_EQ(metadata.dense_soc_kkt_nnz, block_offsets_host[dense_count]); +} + +// Verifies the GPU-built augmented KKT CSR structure (row_start + column +// indices) and values for a concrete mixed SOCP: one linear variable, one +// dense Q^3 cone, one sparse Q^4 cone, quadratic cost, and two constraints. +// The augmented CSR is emitted column-sorted per row. +TEST(sparse_augmented_kkt, augmented_csr_indices_mixed_dense_sparse_qp) +{ + // Augmented KKT for the QP-SOCP: + // + // minimize (1/2) x^T Q x + // subject to x_0 + x_4 = b_0 (constraint 0) + // x_1 + x_6 = b_1 (constraint 1) + // (x_1, x_2, x_3) in Q^3 (dense cone) + // (x_4, x_5, x_6, x_7) in Q^4 (sparse cone) + // + // Q = diag(2, 3, ..., 9) with symmetric off-diagonals Q[0,4]=0.5, Q[1,7]=0.25. + // This checks only the built KKT sparsity/values, so the objective linear term + // and the right-hand side b are irrelevant and left unset. + using i_t = int; + using f_t = double; + auto stream = rmm::cuda_stream_default; + + // Layout: 1 linear var, dense Q^3 cone (cols [1,4)), sparse Q^4 cone (cols + // [4,8)), 2 constraints. Factorization size = n + m + p = 8 + 2 + 2 = 12. + constexpr i_t n_linear = 1; + constexpr i_t n = 8; // 1 linear + 3 (dense) + 4 (sparse) + constexpr i_t m = 2; + constexpr i_t cone_start = n_linear; + constexpr i_t m_c = 7; // 3 + 4 cone entries + const f_t dual_perturb = 0.01; + const f_t primal_perturb = 0.001; + + // Constraint matrix A (m x n, CSC by variable column): var0 -> c0, var1 -> c1, + // var4 -> c0, var6 -> c1. + csc_matrix_t A(m, n, 4); + A.col_start = {0, 1, 2, 2, 2, 3, 3, 4, 4}; + A.i = {0, 1, 0, 1}; + A.x = {1.0, 1.0, 1.0, 1.0}; + + csc_matrix_t AT(n, m, 4); + A.transpose(AT); + + // Quadratic cost Q (n x n, symmetric CSC): diagonal plus off-diagonal pairs + // (0,4) and (1,7) to exercise sorted insertion of the diagonal and of Q + // entries outside the cone block. + csc_matrix_t Q(n, n, 12); + Q.col_start = {0, 2, 4, 5, 6, 8, 9, 10, 12}; + Q.i = {0, 4, 1, 7, 2, 3, 0, 4, 5, 6, 1, 7}; + Q.x = {2.0, 0.5, 3.0, 0.25, 4.0, 5.0, 0.5, 6.0, 7.0, 8.0, 0.25, 9.0}; + const i_t nnzQ = Q.col_start[n]; + + std::vector diag(n, 0.1); + + device_csc_matrix_t d_A(A, stream); + device_csc_matrix_t d_AT(AT, stream); + device_csc_matrix_t d_Q(Q, stream); + auto d_diag = cuopt::device_copy(diag, stream); + + rmm::device_uvector cone_x(m_c, stream); + rmm::device_uvector cone_z(m_c, stream); + std::vector cone_dimensions{3, 4}; + cone_data_t cones(cone_dimensions, + cuopt::make_span(cone_x), + cuopt::make_span(cone_z), + stream, + /*soc_threshold=*/3); + const i_t p = static_cast(cones.expansion_var_count()); + ASSERT_EQ(cones.n_sparse_cones, 1); + ASSERT_EQ(p, 2); + ASSERT_EQ(cones.n_sparse_cone_entries, 4u); + + augmented_csr_metadata_t metadata(stream); + build_augmented_csr_metadata(cones, metadata, stream); + + device_csr_matrix_t device_augmented(stream); + rmm::device_uvector d_augmented_diagonal_indices(0, stream); + rmm::device_uvector d_cone_csr_indices(0, stream); + rmm::device_uvector d_cone_Q_values(0, stream); + rmm::device_uvector d_dense_cone_diag_csr_indices(0, stream); + rmm::device_uvector d_sparse_hessian_diag(0, stream); + rmm::device_uvector d_sparse_hessian_Q(0, stream); + rmm::device_uvector d_sparse_exp_v_col(0, stream); + rmm::device_uvector d_sparse_exp_u_col(0, stream); + rmm::device_uvector d_sparse_exp_v_row(0, stream); + rmm::device_uvector d_sparse_exp_u_row(0, stream); + rmm::device_uvector d_sparse_expansion_D(0, stream); + + augmented_csr_side_buffers_t side{d_augmented_diagonal_indices, + d_cone_csr_indices, + d_cone_Q_values, + d_dense_cone_diag_csr_indices, + d_sparse_hessian_diag, + d_sparse_hessian_Q, + d_sparse_exp_v_col, + d_sparse_exp_u_col, + d_sparse_exp_v_row, + d_sparse_exp_u_row, + d_sparse_expansion_D}; + + const i_t total_nnz = build_augmented_csr_on_device( + n, + m, + p, + cone_start, + m_c, + nnzQ, + dual_perturb, + primal_perturb, + d_A, + d_Q, + d_AT, + raft::device_span{d_diag.data(), d_diag.size()}, + raft::device_span{cones.element_cone_ids.data(), cones.element_cone_ids.size()}, + raft::device_span{cones.cone_offsets.data(), cones.cone_offsets.size()}, + raft::device_span{cones.sparse_cone_ids.data(), cones.sparse_cone_ids.size()}, + raft::device_span{cones.sparse_entry_offsets.data(), + cones.sparse_entry_offsets.size()}, + cones.n_sparse_cone_entries, + metadata, + device_augmented, + side, + stream); + + // Expected column-sorted CSR. Row blocks: primal [0,8), dual [8,10), + // expansion [10,12). Column blocks: primal [0,8), dual [8,10), expansion cols + // {10 (v), 11 (u)}. + const std::vector expected_row_start{0, 3, 8, 11, 14, 19, 22, 26, 30, 33, 36, 41, 46}; + const std::vector expected_j{ + 0, 4, 8, // row0 linear var0: diag, Q(0,4), A^T(c0) + 1, 2, 3, 7, 9, // row1 dense r0: block[1..3], Q(1,7), A^T(c1) + 1, 2, 3, // row2 dense r1 + 1, 2, 3, // row3 dense r2 + 0, 4, 8, 10, 11, // row4 sparse r0: Q(4,0), diag, A^T(c0), exp v, exp u + 5, 10, 11, // row5 sparse r1 + 6, 9, 10, 11, // row6 sparse r2: diag, A^T(c1), exp v, exp u + 1, 7, 10, 11, // row7 sparse r3: Q(7,1), diag, exp v, exp u + 0, 4, 8, // row8 dual c0: A^T rows {0,4}, diag + 1, 6, 9, // row9 dual c1: A^T rows {1,6}, diag + 4, 5, 6, 7, 10, // row10 expansion v: cone cols, diag + 4, 5, 6, 7, 11}; // row11 expansion u: cone cols, diag + const std::vector expected_x{-2.11, -0.5, 1.0, // row0 + -3.01, 0.0, 0.0, -0.25, 1.0, // row1 + 0.0, -4.01, 0.0, // row2 + 0.0, 0.0, -5.01, // row3 + -0.5, -6.01, 1.0, 0.0, 0.0, // row4 + -7.01, 0.0, 0.0, // row5 + -8.01, 1.0, 0.0, 0.0, // row6 + -0.25, -9.01, 0.0, 0.0, // row7 + 1.0, 1.0, 0.001, // row8 + 1.0, 1.0, 0.001, // row9 + 0.0, 0.0, 0.0, 0.0, 0.0, // row10 + 0.0, 0.0, 0.0, 0.0, 0.0}; // row11 + + EXPECT_EQ(total_nnz, 46); + + auto row_start_host = cuopt::host_copy(device_augmented.row_start, stream); + auto j_host = cuopt::host_copy(device_augmented.j, stream); + auto x_host = cuopt::host_copy(device_augmented.x, stream); + + ASSERT_EQ(row_start_host.size(), expected_row_start.size()); + for (size_t r = 0; r < expected_row_start.size(); ++r) { + EXPECT_EQ(row_start_host[r], expected_row_start[r]) << "row_start[" << r << "]"; + } + + ASSERT_EQ(j_host.size(), expected_j.size()); + ASSERT_EQ(x_host.size(), expected_x.size()); + for (size_t e = 0; e < expected_j.size(); ++e) { + EXPECT_EQ(j_host[e], expected_j[e]) << "j[" << e << "]"; + EXPECT_NEAR(x_host[e], expected_x[e], 1e-12) << "x[" << e << "]"; + } +} + } // namespace cuopt::mathematical_optimization::barrier::test