From 17aa1cb3deed0ba842b5f6788ddbc949ebe30437 Mon Sep 17 00:00:00 2001 From: YUWEN Chen Date: Tue, 25 Aug 2026 03:59:51 -0700 Subject: [PATCH 1/7] Reorder computation for residuals, mu and objectives at the begin of barrier method and the end of each barrier iteration, which reduce the number of synchronization required Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 342 ++++++++++--------------- cpp/src/barrier/barrier.hpp | 25 +- cpp/src/linear_algebra/vector_math.cuh | 58 +++++ 3 files changed, 198 insertions(+), 227 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index c164296a25..0eabf6bac1 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -347,6 +347,9 @@ class iteration_data_t { 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()), + d_scalar_batch_(kNumScalarBatchSlots, lp.handle_ptr->get_stream()), + h_scalar_batch_(kNumScalarBatchSlots), + d_reduce_tmp_(0, lp.handle_ptr->get_stream()), indefinite_Q(false), Q_diagonal(false), symbolic_status(0), @@ -2073,6 +2076,14 @@ class iteration_data_t { transform_reduce_pair_helper_t transform_reduce_pair_helper_; sum_reduce_helper_t sum_reduce_helper_; + // Staging area for compute_residual_norms_mu_and_objective: several independent GPU + // reductions/dot-products write into slots of d_scalar_batch_, then a single copy into + // h_scalar_batch_ + one stream sync reads them all back at once instead of one sync each. + static constexpr i_t kNumScalarBatchSlots = 12; + rmm::device_uvector d_scalar_batch_; + pinned_dense_vector_t h_scalar_batch_; + rmm::device_buffer d_reduce_tmp_; + bool cone_combined_step_; f_t cone_sigma_mu_; @@ -2564,46 +2575,6 @@ void barrier_solver_t::gpu_compute_residuals(const rmm::device_uvector RAFT_CHECK_CUDA(stream_view_); } -template -void barrier_solver_t::gpu_compute_residual_norms(const rmm::device_uvector& d_w, - const rmm::device_uvector& d_x, - const rmm::device_uvector& d_y, - const rmm::device_uvector& d_v, - const rmm::device_uvector& d_z, - iteration_data_t& data, - f_t& primal_residual_norm, - f_t& dual_residual_norm, - f_t& complementarity_residual_norm) -{ - raft::common::nvtx::range fun_scope("Barrier: GPU compute_residual_norms"); - - gpu_compute_residuals(d_w, d_x, d_y, d_v, d_z, data); - primal_residual_norm = - std::max(device_vector_norm_inf(data.d_primal_residual_, stream_view_), - device_vector_norm_inf(data.d_bound_residual_, stream_view_)); - dual_residual_norm = device_vector_norm_inf(data.d_dual_residual_, stream_view_); - const bool has_soc = data.has_cones(); - const i_t linear_xz_size = data.linear_xz_size(data.d_complementarity_xz_residual_.size()); - auto linear_xz_span = - raft::device_span(data.d_complementarity_xz_residual_.data(), linear_xz_size); - complementarity_residual_norm = - std::max(device_vector_norm_inf(linear_xz_span, stream_view_), - device_vector_norm_inf(data.d_complementarity_wv_residual_, stream_view_)); - if (has_soc) { - f_t cone_complementarity_norm = f_t(0); - raft::device_span cone_dot = data.cones().scratch.template get_slot<0>(); - data.cones().segmented_sum( - data.d_complementarity_xz_residual_.data() + data.cone_start(), cone_dot, stream_view_); - cone_complementarity_norm = thrust::reduce(rmm::exec_policy(stream_view_), - cone_dot.begin(), - cone_dot.end(), - f_t(0), - thrust::maximum()); - complementarity_residual_norm = - std::max(complementarity_residual_norm, cone_complementarity_norm); - } -} - template std::pair barrier_solver_t::compute_nonnegative_step_length_pair( iteration_data_t& data, @@ -3822,47 +3793,86 @@ void barrier_solver_t::compute_next_iterate(iteration_data_t } template -void barrier_solver_t::compute_residual_norms(iteration_data_t& data, - f_t& primal_residual_norm, - f_t& dual_residual_norm, - f_t& complementarity_residual_norm) +void barrier_solver_t::compute_residual_norms_mu_and_objective( + iteration_data_t& data, + f_t& primal_residual_norm, + f_t& dual_residual_norm, + f_t& complementarity_residual_norm, + f_t& mu, + f_t& primal_objective, + f_t& dual_objective) { - raft::common::nvtx::range fun_scope("Barrier: compute_residual_norms"); - gpu_compute_residual_norms(data.d_w_, - data.d_x_, - data.d_y_, - data.d_v_, - data.d_z_, - data, - primal_residual_norm, - dual_residual_norm, - complementarity_residual_norm); -} + raft::common::nvtx::range fun_scope("Barrier: compute_residual_norms_mu_and_objective"); -template -void barrier_solver_t::compute_mu(iteration_data_t& data, f_t& mu) -{ - raft::common::nvtx::range fun_scope("Barrier: compute_mu"); + gpu_compute_residuals(data.d_w_, data.d_x_, data.d_y_, data.d_v_, data.d_z_, data); - const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); - mu = (data.sum_reduce_helper_.sum(data.d_complementarity_xz_residual_.begin(), - data.d_complementarity_xz_residual_.size(), - stream_view_) + - data.sum_reduce_helper_.sum(data.d_complementarity_wv_residual_.begin(), - data.d_complementarity_wv_residual_.size(), - stream_view_)) / - mu_denom; -} + constexpr i_t kSlotPrimalResidual = 0; + constexpr i_t kSlotBoundResidual = 1; + constexpr i_t kSlotDualResidual = 2; + constexpr i_t kSlotComplXzLinear = 3; + constexpr i_t kSlotComplWv = 4; + constexpr i_t kSlotComplCone = 5; + constexpr i_t kSlotMuXzSum = 6; + constexpr i_t kSlotMuWvSum = 7; + constexpr i_t kSlotCx = 8; + constexpr i_t kSlotBy = 9; + constexpr i_t kSlotUv = 10; + constexpr i_t kSlotXQx = 11; -template -void barrier_solver_t::compute_primal_dual_objective(iteration_data_t& data, - f_t& primal_objective, - f_t& dual_objective) -{ - raft::common::nvtx::range fun_scope("Barrier: compute_primal_dual_objective"); - rmm::device_scalar d_cx(stream_view_); - rmm::device_scalar d_by(stream_view_); - rmm::device_scalar d_uv(stream_view_); + f_t* d_batch = data.d_scalar_batch_.data(); + + const bool has_soc = data.has_cones(); + const i_t linear_xz_size = data.linear_xz_size(data.d_complementarity_xz_residual_.size()); + auto linear_xz_span = + raft::device_span(data.d_complementarity_xz_residual_.data(), linear_xz_size); + + // All enqueue calls below must stay on stream_view_: correctness relies on strict + // single-stream FIFO ordering, so that the single sync at the bottom is enough for every + // result to be ready on the host. + enqueue_norm_inf_into(data.d_primal_residual_.data(), + data.d_primal_residual_.size(), + d_batch + kSlotPrimalResidual, + data.d_reduce_tmp_, + stream_view_); + enqueue_norm_inf_into(data.d_bound_residual_.data(), + data.d_bound_residual_.size(), + d_batch + kSlotBoundResidual, + data.d_reduce_tmp_, + stream_view_); + enqueue_norm_inf_into(data.d_dual_residual_.data(), + data.d_dual_residual_.size(), + d_batch + kSlotDualResidual, + data.d_reduce_tmp_, + stream_view_); + enqueue_norm_inf_into(linear_xz_span.data(), + linear_xz_span.size(), + d_batch + kSlotComplXzLinear, + data.d_reduce_tmp_, + stream_view_); + enqueue_norm_inf_into(data.d_complementarity_wv_residual_.data(), + data.d_complementarity_wv_residual_.size(), + d_batch + kSlotComplWv, + data.d_reduce_tmp_, + stream_view_); + + if (has_soc) { + raft::device_span cone_dot = data.cones().scratch.template get_slot<0>(); + data.cones().segmented_sum( + data.d_complementarity_xz_residual_.data() + data.cone_start(), cone_dot, stream_view_); + enqueue_max_into( + cone_dot.data(), cone_dot.size(), d_batch + kSlotComplCone, data.d_reduce_tmp_, stream_view_); + } + + enqueue_sum_into(data.d_complementarity_xz_residual_.data(), + data.d_complementarity_xz_residual_.size(), + d_batch + kSlotMuXzSum, + data.d_reduce_tmp_, + stream_view_); + enqueue_sum_into(data.d_complementarity_wv_residual_.data(), + data.d_complementarity_wv_residual_.size(), + d_batch + kSlotMuWvSum, + data.d_reduce_tmp_, + stream_view_); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_c_.size(), @@ -3870,7 +3880,7 @@ void barrier_solver_t::compute_primal_dual_objective(iteration_data_t< 1, data.d_x_.data(), 1, - d_cx.data(), + d_batch + kSlotCx, stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_b_.size(), @@ -3878,7 +3888,7 @@ void barrier_solver_t::compute_primal_dual_objective(iteration_data_t< 1, data.d_y_.data(), 1, - d_by.data(), + d_batch + kSlotBy, stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_restrict_u_.size(), @@ -3886,118 +3896,43 @@ void barrier_solver_t::compute_primal_dual_objective(iteration_data_t< 1, data.d_v_.data(), 1, - d_uv.data(), + d_batch + kSlotUv, stream_view_)); - f_t quad_objective = 0.0; if (data.Q.n > 0) { auto cusparse_d_x = data.cusparse_view_.create_vector(data.d_x_); auto cusparse_Qx = data.cusparse_view_.create_vector(data.d_Qx_); data.cusparse_Q_view_.spmv(1.0, cusparse_d_x, 0.0, cusparse_Qx); - rmm::device_scalar d_xQx(stream_view_); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_Qx_.size(), data.d_Qx_.data(), 1, data.d_x_.data(), 1, - d_xQx.data(), + d_batch + kSlotXQx, stream_view_)); - quad_objective = 0.5 * d_xQx.value(stream_view_); } - primal_objective = d_cx.value(stream_view_) + quad_objective; - dual_objective = d_by.value(stream_view_) - d_uv.value(stream_view_) - quad_objective; - -#ifdef CHECK_OBJECTIVE_GAP - rmm::device_scalar d_xz(stream_view_); - rmm::device_scalar d_wv(stream_view_); - rmm::device_scalar d_rdx(stream_view_); - rmm::device_scalar d_rpy(stream_view_); - rmm::device_scalar d_rwv(stream_view_); - rmm::device_scalar d_p(stream_view_); - rmm::device_scalar d_y(stream_view_); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_x_.size(), - data.d_x_.data(), - 1, - data.d_z_.data(), - 1, - d_xz.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_w_.size(), - data.d_w_.data(), - 1, - data.d_v_.data(), - 1, - d_wv.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_x_.size(), - data.d_x_.data(), - 1, - data.d_dual_residual_.data(), - 1, - d_rdx.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_y_.size(), - data.d_y_.data(), - 1, - data.d_primal_residual_.data(), - 1, - d_rpy.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_bound_residual_.size(), - data.d_bound_residual_.data(), - 1, - data.d_v_.data(), - 1, - d_rwv.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_primal_residual_.size(), - data.d_primal_residual_.data(), - 1, - data.d_primal_residual_.data(), - 1, - d_p.data(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_y_.size(), - data.d_y_.data(), - 1, - data.d_y_.data(), - 1, - d_y.data(), - stream_view_)); - f_t xz = d_xz.value(stream_view_); - f_t wv = d_wv.value(stream_view_); - f_t rdx = d_rdx.value(stream_view_); - f_t rpy = d_rpy.value(stream_view_); - f_t rwv = d_rwv.value(stream_view_); - f_t p = d_p.value(stream_view_); - f_t y = d_y.value(stream_view_); - + raft::copy(data.h_scalar_batch_.data(), + data.d_scalar_batch_.data(), + data.kNumScalarBatchSlots, + stream_view_); stream_view_.synchronize(); - f_t objective_gap_1 = primal_objective - dual_objective; - f_t objective_gap_2 = xz + wv + rdx - rpy + rwv; - - settings.log.printf("Objective gap 1: %.2e, Objective gap 2: %.2e Diff: %.2e\n", - objective_gap_1, - objective_gap_2, - std::abs(objective_gap_1 - objective_gap_2)); - settings.log.printf( - "Objective - Complementarity: %.2e, rdx: %.2e, rpy: %.2e, rwv: %.2e, p: %.2e, y: %.2e\n", - std::abs(objective_gap_1 - (xz + wv)), - rdx, - rpy, - rwv, - p, - y); -#endif + const f_t* h = data.h_scalar_batch_.data(); + + primal_residual_norm = std::max(h[kSlotPrimalResidual], h[kSlotBoundResidual]); + dual_residual_norm = h[kSlotDualResidual]; + complementarity_residual_norm = std::max(h[kSlotComplXzLinear], h[kSlotComplWv]); + if (has_soc) { + complementarity_residual_norm = std::max(complementarity_residual_norm, h[kSlotComplCone]); + } + + const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); + mu = (h[kSlotMuXzSum] + h[kSlotMuWvSum]) / mu_denom; + + const f_t quad_objective = (data.Q.n > 0) ? 0.5 * h[kSlotXQx] : f_t(0); + primal_objective = h[kSlotCx] + quad_objective; + dual_objective = h[kSlotBy] - h[kSlotUv] - quad_objective; } template @@ -4198,29 +4133,25 @@ lp_status_t barrier_solver_t::solve(f_t start_time, lp_solution_t(data.b, stream_view_); f_t norm_c = vector_norm_inf(data.c, stream_view_); - f_t quad_objective = 0.0; - if (data.Q.n > 0) { - dense_vector_t Qx(data.Q.n); - matrix_vector_multiply(data.Q, 1.0, data.x, 0.0, Qx); - quad_objective = 0.5 * data.x.inner_product(Qx); - } - f_t primal_objective = data.c.inner_product(data.x) + quad_objective; + dense_vector_t upper(lp.upper); + data.gather_upper_bounds(upper, data.restrict_u_); + data.d_restrict_u_.resize(data.restrict_u_.size(), stream_view_); + raft::copy( + data.d_restrict_u_.data(), data.restrict_u_.data(), data.restrict_u_.size(), stream_view_); + + f_t primal_residual_norm, dual_residual_norm, complementarity_residual_norm; + f_t mu; + f_t primal_objective, dual_objective; + compute_residual_norms_mu_and_objective(data, + primal_residual_norm, + dual_residual_norm, + complementarity_residual_norm, + mu, + primal_objective, + dual_objective); f_t relative_primal_residual = primal_residual_norm / (1.0 + norm_b); f_t relative_dual_residual = dual_residual_norm / (1.0 + norm_c); @@ -4229,14 +4160,6 @@ lp_status_t barrier_solver_t::solve(f_t start_time, lp_solution_t upper(lp.upper); - data.gather_upper_bounds(upper, data.restrict_u_); - data.d_restrict_u_.resize(data.restrict_u_.size(), stream_view_); - raft::copy( - data.d_restrict_u_.data(), data.restrict_u_.data(), data.restrict_u_.size(), stream_view_); - f_t dual_objective = - data.b.inner_product(data.y) - data.restrict_u_.inner_product(data.v) - quad_objective; - f_t objective_gap_abs = std::abs(primal_objective - dual_objective); f_t objective_gap_rel = objective_gap_abs / @@ -4373,12 +4296,13 @@ lp_status_t barrier_solver_t::solve(f_t start_time, lp_solution_t& data, - f_t& primal_residual_norm, - f_t& dual_residual_norm, - f_t& complementarity_residual_norm); - void compute_mu(iteration_data_t& data, f_t& mu); - void compute_primal_dual_objective(iteration_data_t& data, - f_t& primal_objective, - f_t& dual_objective); + void compute_residual_norms_mu_and_objective(iteration_data_t& data, + f_t& primal_residual_norm, + f_t& dual_residual_norm, + f_t& complementarity_residual_norm, + f_t& mu, + f_t& primal_objective, + f_t& dual_objective); // To be able to directly pass lambdas to transform functions public: @@ -81,16 +80,6 @@ class barrier_solver_t { rmm::device_uvector const& d_v, rmm::device_uvector const& d_z, iteration_data_t& data); - void gpu_compute_residual_norms(const rmm::device_uvector& d_w, - const rmm::device_uvector& d_x, - const rmm::device_uvector& d_y, - const rmm::device_uvector& d_v, - const rmm::device_uvector& d_z, - iteration_data_t& data, - f_t& primal_residual_norm, - f_t& dual_residual_norm, - f_t& complementarity_residual_norm); - std::pair compute_nonnegative_step_length_pair(iteration_data_t& data, const rmm::device_uvector& x1, const rmm::device_uvector& dx1, diff --git a/cpp/src/linear_algebra/vector_math.cuh b/cpp/src/linear_algebra/vector_math.cuh index ac9d24001b..8554523d6e 100644 --- a/cpp/src/linear_algebra/vector_math.cuh +++ b/cpp/src/linear_algebra/vector_math.cuh @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -68,6 +69,63 @@ f_t device_custom_vector_norm_inf(InputIteratorT in, i_t size, rmm::cuda_stream_ return d_out.value(stream_view); } +// Same reduction as device_custom_vector_norm_inf, but writes into a caller-supplied device +// pointer (and reuses a caller-supplied temp-storage buffer) instead of allocating a private +// rmm::device_scalar and blocking on .value(). Lets callers batch several reductions and defer +// the host readback to a single copy + sync. +template +void enqueue_norm_inf_into( + InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) +{ + if (size == 0) { + RAFT_CUDA_TRY(cudaMemsetAsync(out, 0, sizeof(f_t), stream_view.value())); + return; + } + size_t temp_storage_bytes = 0; + f_t init = 0; + auto custom_op = norm_inf_max{}; + cub::DeviceReduce::Reduce( + nullptr, temp_storage_bytes, in, out, size, custom_op, init, stream_view); + + tmp.resize(temp_storage_bytes, stream_view); + + cub::DeviceReduce::Reduce( + tmp.data(), temp_storage_bytes, in, out, size, custom_op, init, stream_view); +} + +// Sum reduction into a caller-supplied device pointer/temp-storage buffer, deferring the host +// readback (see enqueue_norm_inf_into). +template +void enqueue_sum_into( + InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) +{ + size_t temp_storage_bytes = 0; + cub::DeviceReduce::Sum(nullptr, temp_storage_bytes, in, out, size, stream_view); + + tmp.resize(temp_storage_bytes, stream_view); + + cub::DeviceReduce::Sum(tmp.data(), temp_storage_bytes, in, out, size, stream_view); +} + +// Max reduction (with a floor of 0, matching this codebase's existing +// thrust::reduce(..., f_t(0), thrust::maximum()) usage) into a caller-supplied device +// pointer/temp-storage buffer, deferring the host readback (see enqueue_norm_inf_into). +template +void enqueue_max_into( + InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) +{ + size_t temp_storage_bytes = 0; + f_t init = 0; + auto custom_op = thrust::maximum{}; + cub::DeviceReduce::Reduce( + nullptr, temp_storage_bytes, in, out, size, custom_op, init, stream_view); + + tmp.resize(temp_storage_bytes, stream_view); + + cub::DeviceReduce::Reduce( + tmp.data(), temp_storage_bytes, in, out, size, custom_op, init, stream_view); +} + template f_t device_vector_norm_inf(const rmm::device_uvector& in, rmm::cuda_stream_view stream_view) { From a6f49a6a72f6531facb14a619c73407d275a1dff Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Wed, 26 Aug 2026 08:13:41 -0700 Subject: [PATCH 2/7] Rename new vectors Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 45 ++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 0eabf6bac1..227f253a03 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -347,9 +347,9 @@ class iteration_data_t { 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()), - d_scalar_batch_(kNumScalarBatchSlots, lp.handle_ptr->get_stream()), - h_scalar_batch_(kNumScalarBatchSlots), - d_reduce_tmp_(0, lp.handle_ptr->get_stream()), + d_reduction_results_(kNumScalarBatchSlots, lp.handle_ptr->get_stream()), + h_reduction_results_(kNumScalarBatchSlots), + d_reduce_temp_storage_(0, lp.handle_ptr->get_stream()), indefinite_Q(false), Q_diagonal(false), symbolic_status(0), @@ -2077,12 +2077,12 @@ class iteration_data_t { sum_reduce_helper_t sum_reduce_helper_; // Staging area for compute_residual_norms_mu_and_objective: several independent GPU - // reductions/dot-products write into slots of d_scalar_batch_, then a single copy into - // h_scalar_batch_ + one stream sync reads them all back at once instead of one sync each. + // reductions/dot-products write into slots of d_reduction_results_, then a single copy into + // h_reduction_results_ + one stream sync reads them all back at once instead of one sync each. static constexpr i_t kNumScalarBatchSlots = 12; - rmm::device_uvector d_scalar_batch_; - pinned_dense_vector_t h_scalar_batch_; - rmm::device_buffer d_reduce_tmp_; + rmm::device_uvector d_reduction_results_; + pinned_dense_vector_t h_reduction_results_; + rmm::device_buffer d_reduce_temp_storage_; bool cone_combined_step_; f_t cone_sigma_mu_; @@ -3819,7 +3819,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( constexpr i_t kSlotUv = 10; constexpr i_t kSlotXQx = 11; - f_t* d_batch = data.d_scalar_batch_.data(); + f_t* d_batch = data.d_reduction_results_.data(); const bool has_soc = data.has_cones(); const i_t linear_xz_size = data.linear_xz_size(data.d_complementarity_xz_residual_.size()); @@ -3832,46 +3832,49 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( enqueue_norm_inf_into(data.d_primal_residual_.data(), data.d_primal_residual_.size(), d_batch + kSlotPrimalResidual, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); enqueue_norm_inf_into(data.d_bound_residual_.data(), data.d_bound_residual_.size(), d_batch + kSlotBoundResidual, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); enqueue_norm_inf_into(data.d_dual_residual_.data(), data.d_dual_residual_.size(), d_batch + kSlotDualResidual, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); enqueue_norm_inf_into(linear_xz_span.data(), linear_xz_span.size(), d_batch + kSlotComplXzLinear, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); enqueue_norm_inf_into(data.d_complementarity_wv_residual_.data(), data.d_complementarity_wv_residual_.size(), d_batch + kSlotComplWv, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); if (has_soc) { raft::device_span cone_dot = data.cones().scratch.template get_slot<0>(); data.cones().segmented_sum( data.d_complementarity_xz_residual_.data() + data.cone_start(), cone_dot, stream_view_); - enqueue_max_into( - cone_dot.data(), cone_dot.size(), d_batch + kSlotComplCone, data.d_reduce_tmp_, stream_view_); + enqueue_max_into(cone_dot.data(), + cone_dot.size(), + d_batch + kSlotComplCone, + data.d_reduce_temp_storage_, + stream_view_); } enqueue_sum_into(data.d_complementarity_xz_residual_.data(), data.d_complementarity_xz_residual_.size(), d_batch + kSlotMuXzSum, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); enqueue_sum_into(data.d_complementarity_wv_residual_.data(), data.d_complementarity_wv_residual_.size(), d_batch + kSlotMuWvSum, - data.d_reduce_tmp_, + data.d_reduce_temp_storage_, stream_view_); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), @@ -3912,13 +3915,13 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( stream_view_)); } - raft::copy(data.h_scalar_batch_.data(), - data.d_scalar_batch_.data(), + raft::copy(data.h_reduction_results_.data(), + data.d_reduction_results_.data(), data.kNumScalarBatchSlots, stream_view_); stream_view_.synchronize(); - const f_t* h = data.h_scalar_batch_.data(); + const f_t* h = data.h_reduction_results_.data(); primal_residual_norm = std::max(h[kSlotPrimalResidual], h[kSlotBoundResidual]); dual_residual_norm = h[kSlotDualResidual]; From ea4f38a680e0ec53d107f763a1a5bcac7170ff6b Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Thu, 27 Aug 2026 02:26:41 -0700 Subject: [PATCH 3/7] Code clean up: rename functions and create a new struct barrier_reduce_helper_t addressing operations related to termination check Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 251 ++++++++++++++++--------- cpp/src/barrier/barrier.hpp | 9 - cpp/src/linear_algebra/vector_math.cuh | 58 ------ 3 files changed, 162 insertions(+), 156 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 227f253a03..3c6dbe8886 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -221,6 +221,144 @@ static void fill_linear_cc_rhs(raft::device_span out, RAFT_CHECK_CUDA(stream); } +// Batches the independent GPU reductions/dot-products needed by +// compute_residual_norms_mu_and_objective (primal/dual/complementarity residual norms, mu, and +// primal/dual objectives) into one on-device results buffer and one host readback + stream sync. +template +class barrier_reduce_helper_t { + public: + explicit barrier_reduce_helper_t(rmm::cuda_stream_view stream_view) + : d_results_(kCount, stream_view), h_results_(kCount), d_temp_storage_(0, stream_view) + { + } + + void primal_residual_norm_async(const rmm::device_uvector& d_primal_residual, + const rmm::device_uvector& d_bound_residual, + rmm::cuda_stream_view stream_view) + { + norm_inf_async( + kPrimalResidual, d_primal_residual.data(), d_primal_residual.size(), stream_view); + norm_inf_async(kBoundResidual, d_bound_residual.data(), d_bound_residual.size(), stream_view); + } + + void dual_residual_norm_async(const rmm::device_uvector& d_dual_residual, + rmm::cuda_stream_view stream_view) + { + norm_inf_async(kDualResidual, d_dual_residual.data(), d_dual_residual.size(), stream_view); + } + + void complementarity_residual_norm_async(raft::device_span linear_xz, + const rmm::device_uvector& d_wv, + rmm::cuda_stream_view stream_view) + { + norm_inf_async(kComplXzLinear, linear_xz.data(), linear_xz.size(), stream_view); + norm_inf_async(kComplWv, d_wv.data(), d_wv.size(), stream_view); + } + + void cone_complementarity_residual_async(raft::device_span cone_dot, + rmm::cuda_stream_view stream_view) + { + has_soc_ = true; + max_async(kComplCone, cone_dot.data(), cone_dot.size(), stream_view); + } + + void mu_terms_async(const rmm::device_uvector& d_xz, + const rmm::device_uvector& d_wv, + rmm::cuda_stream_view stream_view) + { + sum_async(kMuXzSum, d_xz.data(), d_xz.size(), stream_view); + sum_async(kMuWvSum, d_wv.data(), d_wv.size(), stream_view); + } + + // Raw device slots for the caller's own cublasdot() calls. + f_t* cx_slot() { return d_results_.data() + kCx; } + f_t* by_slot() { return d_results_.data() + kBy; } + f_t* uv_slot() { return d_results_.data() + kUv; } + f_t* xqx_slot() { return d_results_.data() + kXQx; } + + // Single batched device-to-host copy + the one stream synchronize needed before any accessor + // below can be read. + void sync(rmm::cuda_stream_view stream_view) + { + raft::copy(h_results_.data(), d_results_.data(), static_cast(kCount), stream_view); + stream_view.synchronize(); + } + + f_t primal_residual_norm() const + { + return std::max(h_results_[kPrimalResidual], h_results_[kBoundResidual]); + } + f_t dual_residual_norm() const { return h_results_[kDualResidual]; } + f_t complementarity_residual_norm() const + { + f_t result = std::max(h_results_[kComplXzLinear], h_results_[kComplWv]); + if (has_soc_) { result = std::max(result, h_results_[kComplCone]); } + return result; + } + f_t mu(f_t mu_denom) const { return (h_results_[kMuXzSum] + h_results_[kMuWvSum]) / mu_denom; } + f_t cx() const { return h_results_[kCx]; } + f_t by() const { return h_results_[kBy]; } + f_t uv() const { return h_results_[kUv]; } + f_t xqx() const { return h_results_[kXQx]; } + + private: + enum Slot : i_t { + kPrimalResidual = 0, + kBoundResidual, + kDualResidual, + kComplXzLinear, + kComplWv, + kComplCone, + kMuXzSum, + kMuWvSum, + kCx, + kBy, + kUv, + kXQx, + kCount + }; + + template + void reduce_async( + Slot slot, const f_t* in, i_t size, ReduceOpT op, f_t init, rmm::cuda_stream_view stream_view) + { + f_t* out = d_results_.data() + slot; + if (size == 0) { + RAFT_CUDA_TRY(cudaMemsetAsync(out, 0, sizeof(f_t), stream_view.value())); + return; + } + size_t temp_storage_bytes = 0; + cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, in, out, size, op, init, stream_view); + d_temp_storage_.resize(temp_storage_bytes, stream_view); + cub::DeviceReduce::Reduce( + d_temp_storage_.data(), temp_storage_bytes, in, out, size, op, init, stream_view); + } + + void norm_inf_async(Slot slot, const f_t* in, i_t size, rmm::cuda_stream_view stream_view) + { + reduce_async(slot, in, size, norm_inf_max{}, f_t(0), stream_view); + } + + void max_async(Slot slot, const f_t* in, i_t size, rmm::cuda_stream_view stream_view) + { + reduce_async(slot, in, size, thrust::maximum{}, f_t(0), stream_view); + } + + void sum_async(Slot slot, const f_t* in, i_t size, rmm::cuda_stream_view stream_view) + { + f_t* out = d_results_.data() + slot; + size_t temp_storage_bytes = 0; + cub::DeviceReduce::Sum(nullptr, temp_storage_bytes, in, out, size, stream_view); + d_temp_storage_.resize(temp_storage_bytes, stream_view); + cub::DeviceReduce::Sum(d_temp_storage_.data(), temp_storage_bytes, in, out, size, stream_view); + } + + rmm::device_uvector d_results_; + pinned_dense_vector_t h_results_; + rmm::device_buffer d_temp_storage_; + bool has_soc_ = false; +}; + template class iteration_data_t { public: @@ -347,9 +485,7 @@ class iteration_data_t { 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()), - d_reduction_results_(kNumScalarBatchSlots, lp.handle_ptr->get_stream()), - h_reduction_results_(kNumScalarBatchSlots), - d_reduce_temp_storage_(0, lp.handle_ptr->get_stream()), + reduce_helper_(lp.handle_ptr->get_stream()), indefinite_Q(false), Q_diagonal(false), symbolic_status(0), @@ -2076,13 +2212,7 @@ class iteration_data_t { transform_reduce_pair_helper_t transform_reduce_pair_helper_; sum_reduce_helper_t sum_reduce_helper_; - // Staging area for compute_residual_norms_mu_and_objective: several independent GPU - // reductions/dot-products write into slots of d_reduction_results_, then a single copy into - // h_reduction_results_ + one stream sync reads them all back at once instead of one sync each. - static constexpr i_t kNumScalarBatchSlots = 12; - rmm::device_uvector d_reduction_results_; - pinned_dense_vector_t h_reduction_results_; - rmm::device_buffer d_reduce_temp_storage_; + barrier_reduce_helper_t reduce_helper_; bool cone_combined_step_; f_t cone_sigma_mu_; @@ -3806,76 +3936,28 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( gpu_compute_residuals(data.d_w_, data.d_x_, data.d_y_, data.d_v_, data.d_z_, data); - constexpr i_t kSlotPrimalResidual = 0; - constexpr i_t kSlotBoundResidual = 1; - constexpr i_t kSlotDualResidual = 2; - constexpr i_t kSlotComplXzLinear = 3; - constexpr i_t kSlotComplWv = 4; - constexpr i_t kSlotComplCone = 5; - constexpr i_t kSlotMuXzSum = 6; - constexpr i_t kSlotMuWvSum = 7; - constexpr i_t kSlotCx = 8; - constexpr i_t kSlotBy = 9; - constexpr i_t kSlotUv = 10; - constexpr i_t kSlotXQx = 11; - - f_t* d_batch = data.d_reduction_results_.data(); + auto& rh = data.reduce_helper_; const bool has_soc = data.has_cones(); const i_t linear_xz_size = data.linear_xz_size(data.d_complementarity_xz_residual_.size()); auto linear_xz_span = raft::device_span(data.d_complementarity_xz_residual_.data(), linear_xz_size); - // All enqueue calls below must stay on stream_view_: correctness relies on strict - // single-stream FIFO ordering, so that the single sync at the bottom is enough for every + // All *_async calls below must stay on stream_view_: correctness relies on strict + // single-stream FIFO ordering, so that the single rh.sync() at the bottom is enough for every // result to be ready on the host. - enqueue_norm_inf_into(data.d_primal_residual_.data(), - data.d_primal_residual_.size(), - d_batch + kSlotPrimalResidual, - data.d_reduce_temp_storage_, - stream_view_); - enqueue_norm_inf_into(data.d_bound_residual_.data(), - data.d_bound_residual_.size(), - d_batch + kSlotBoundResidual, - data.d_reduce_temp_storage_, - stream_view_); - enqueue_norm_inf_into(data.d_dual_residual_.data(), - data.d_dual_residual_.size(), - d_batch + kSlotDualResidual, - data.d_reduce_temp_storage_, - stream_view_); - enqueue_norm_inf_into(linear_xz_span.data(), - linear_xz_span.size(), - d_batch + kSlotComplXzLinear, - data.d_reduce_temp_storage_, - stream_view_); - enqueue_norm_inf_into(data.d_complementarity_wv_residual_.data(), - data.d_complementarity_wv_residual_.size(), - d_batch + kSlotComplWv, - data.d_reduce_temp_storage_, - stream_view_); - + rh.primal_residual_norm_async(data.d_primal_residual_, data.d_bound_residual_, stream_view_); + rh.dual_residual_norm_async(data.d_dual_residual_, stream_view_); + rh.complementarity_residual_norm_async( + linear_xz_span, data.d_complementarity_wv_residual_, stream_view_); if (has_soc) { raft::device_span cone_dot = data.cones().scratch.template get_slot<0>(); data.cones().segmented_sum( data.d_complementarity_xz_residual_.data() + data.cone_start(), cone_dot, stream_view_); - enqueue_max_into(cone_dot.data(), - cone_dot.size(), - d_batch + kSlotComplCone, - data.d_reduce_temp_storage_, - stream_view_); + rh.cone_complementarity_residual_async(cone_dot, stream_view_); } - - enqueue_sum_into(data.d_complementarity_xz_residual_.data(), - data.d_complementarity_xz_residual_.size(), - d_batch + kSlotMuXzSum, - data.d_reduce_temp_storage_, - stream_view_); - enqueue_sum_into(data.d_complementarity_wv_residual_.data(), - data.d_complementarity_wv_residual_.size(), - d_batch + kSlotMuWvSum, - data.d_reduce_temp_storage_, - stream_view_); + rh.mu_terms_async( + data.d_complementarity_xz_residual_, data.d_complementarity_wv_residual_, stream_view_); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_c_.size(), @@ -3883,7 +3965,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_x_.data(), 1, - d_batch + kSlotCx, + rh.cx_slot(), stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_b_.size(), @@ -3891,7 +3973,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_y_.data(), 1, - d_batch + kSlotBy, + rh.by_slot(), stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_restrict_u_.size(), @@ -3899,7 +3981,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_v_.data(), 1, - d_batch + kSlotUv, + rh.uv_slot(), stream_view_)); if (data.Q.n > 0) { auto cusparse_d_x = data.cusparse_view_.create_vector(data.d_x_); @@ -3911,31 +3993,22 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_x_.data(), 1, - d_batch + kSlotXQx, + rh.xqx_slot(), stream_view_)); } - raft::copy(data.h_reduction_results_.data(), - data.d_reduction_results_.data(), - data.kNumScalarBatchSlots, - stream_view_); - stream_view_.synchronize(); - - const f_t* h = data.h_reduction_results_.data(); + rh.sync(stream_view_); - primal_residual_norm = std::max(h[kSlotPrimalResidual], h[kSlotBoundResidual]); - dual_residual_norm = h[kSlotDualResidual]; - complementarity_residual_norm = std::max(h[kSlotComplXzLinear], h[kSlotComplWv]); - if (has_soc) { - complementarity_residual_norm = std::max(complementarity_residual_norm, h[kSlotComplCone]); - } + primal_residual_norm = rh.primal_residual_norm(); + dual_residual_norm = rh.dual_residual_norm(); + complementarity_residual_norm = rh.complementarity_residual_norm(); const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); - mu = (h[kSlotMuXzSum] + h[kSlotMuWvSum]) / mu_denom; + mu = rh.mu(mu_denom); - const f_t quad_objective = (data.Q.n > 0) ? 0.5 * h[kSlotXQx] : f_t(0); - primal_objective = h[kSlotCx] + quad_objective; - dual_objective = h[kSlotBy] - h[kSlotUv] - quad_objective; + const f_t quad_objective = (data.Q.n > 0) ? 0.5 * rh.xqx() : f_t(0); + primal_objective = rh.cx() + quad_objective; + dual_objective = rh.by() - rh.uv() - quad_objective; } template diff --git a/cpp/src/barrier/barrier.hpp b/cpp/src/barrier/barrier.hpp index b1e898e188..8312df701e 100644 --- a/cpp/src/barrier/barrier.hpp +++ b/cpp/src/barrier/barrier.hpp @@ -40,15 +40,6 @@ class barrier_solver_t { void my_pop_range(bool debug) const; void create_Q(const simplex::lp_problem_t& lp, csc_matrix_t& Q); int initial_point(iteration_data_t& data); - void compute_residual_norms(const dense_vector_t& w, - const dense_vector_t& x, - const dense_vector_t& y, - const dense_vector_t& v, - const dense_vector_t& z, - iteration_data_t& data, - f_t& primal_residual_norm, - f_t& dual_residual_norm, - f_t& complementarity_residual_norm); void compute_primal_dual_step_length(iteration_data_t& data, f_t step_scale, diff --git a/cpp/src/linear_algebra/vector_math.cuh b/cpp/src/linear_algebra/vector_math.cuh index 8554523d6e..ac9d24001b 100644 --- a/cpp/src/linear_algebra/vector_math.cuh +++ b/cpp/src/linear_algebra/vector_math.cuh @@ -14,7 +14,6 @@ #include #include -#include #include #include #include @@ -69,63 +68,6 @@ f_t device_custom_vector_norm_inf(InputIteratorT in, i_t size, rmm::cuda_stream_ return d_out.value(stream_view); } -// Same reduction as device_custom_vector_norm_inf, but writes into a caller-supplied device -// pointer (and reuses a caller-supplied temp-storage buffer) instead of allocating a private -// rmm::device_scalar and blocking on .value(). Lets callers batch several reductions and defer -// the host readback to a single copy + sync. -template -void enqueue_norm_inf_into( - InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) -{ - if (size == 0) { - RAFT_CUDA_TRY(cudaMemsetAsync(out, 0, sizeof(f_t), stream_view.value())); - return; - } - size_t temp_storage_bytes = 0; - f_t init = 0; - auto custom_op = norm_inf_max{}; - cub::DeviceReduce::Reduce( - nullptr, temp_storage_bytes, in, out, size, custom_op, init, stream_view); - - tmp.resize(temp_storage_bytes, stream_view); - - cub::DeviceReduce::Reduce( - tmp.data(), temp_storage_bytes, in, out, size, custom_op, init, stream_view); -} - -// Sum reduction into a caller-supplied device pointer/temp-storage buffer, deferring the host -// readback (see enqueue_norm_inf_into). -template -void enqueue_sum_into( - InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) -{ - size_t temp_storage_bytes = 0; - cub::DeviceReduce::Sum(nullptr, temp_storage_bytes, in, out, size, stream_view); - - tmp.resize(temp_storage_bytes, stream_view); - - cub::DeviceReduce::Sum(tmp.data(), temp_storage_bytes, in, out, size, stream_view); -} - -// Max reduction (with a floor of 0, matching this codebase's existing -// thrust::reduce(..., f_t(0), thrust::maximum()) usage) into a caller-supplied device -// pointer/temp-storage buffer, deferring the host readback (see enqueue_norm_inf_into). -template -void enqueue_max_into( - InputIteratorT in, i_t size, f_t* out, rmm::device_buffer& tmp, rmm::cuda_stream_view stream_view) -{ - size_t temp_storage_bytes = 0; - f_t init = 0; - auto custom_op = thrust::maximum{}; - cub::DeviceReduce::Reduce( - nullptr, temp_storage_bytes, in, out, size, custom_op, init, stream_view); - - tmp.resize(temp_storage_bytes, stream_view); - - cub::DeviceReduce::Reduce( - tmp.data(), temp_storage_bytes, in, out, size, custom_op, init, stream_view); -} - template f_t device_vector_norm_inf(const rmm::device_uvector& in, rmm::cuda_stream_view stream_view) { From 520e99516513d4ffee96d098e17b4c9c05fa10d8 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Thu, 27 Aug 2026 14:06:58 -0700 Subject: [PATCH 4/7] restore 'CHECK_OBJECTIVE_GAP' back Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 91 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 70851ab458..6fdbcdf691 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -4015,6 +4015,97 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( const f_t quad_objective = (data.Q.n > 0) ? 0.5 * rh.xqx() : f_t(0); primal_objective = rh.cx() + quad_objective; dual_objective = rh.by() - rh.uv() - quad_objective; + +#ifdef CHECK_OBJECTIVE_GAP + rmm::device_scalar d_xz(stream_view_); + rmm::device_scalar d_wv(stream_view_); + rmm::device_scalar d_rdx(stream_view_); + rmm::device_scalar d_rpy(stream_view_); + rmm::device_scalar d_rwv(stream_view_); + rmm::device_scalar d_p(stream_view_); + rmm::device_scalar d_y(stream_view_); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_x_.size(), + data.d_x_.data(), + 1, + data.d_z_.data(), + 1, + d_xz.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_w_.size(), + data.d_w_.data(), + 1, + data.d_v_.data(), + 1, + d_wv.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_x_.size(), + data.d_x_.data(), + 1, + data.d_dual_residual_.data(), + 1, + d_rdx.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_y_.size(), + data.d_y_.data(), + 1, + data.d_primal_residual_.data(), + 1, + d_rpy.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_bound_residual_.size(), + data.d_bound_residual_.data(), + 1, + data.d_v_.data(), + 1, + d_rwv.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_primal_residual_.size(), + data.d_primal_residual_.data(), + 1, + data.d_primal_residual_.data(), + 1, + d_p.data(), + stream_view_)); + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), + data.d_y_.size(), + data.d_y_.data(), + 1, + data.d_y_.data(), + 1, + d_y.data(), + stream_view_)); + f_t xz = d_xz.value(stream_view_); + f_t wv = d_wv.value(stream_view_); + f_t rdx = d_rdx.value(stream_view_); + f_t rpy = d_rpy.value(stream_view_); + f_t rwv = d_rwv.value(stream_view_); + f_t p = d_p.value(stream_view_); + f_t y = d_y.value(stream_view_); + + stream_view_.synchronize(); + + f_t objective_gap_1 = primal_objective - dual_objective; + f_t objective_gap_2 = xz + wv + rdx - rpy + rwv; + + settings.log.printf("Objective gap 1: %.2e, Objective gap 2: %.2e Diff: %.2e\n", + objective_gap_1, + objective_gap_2, + std::abs(objective_gap_1 - objective_gap_2)); + settings.log.printf( + "Objective - Complementarity: %.2e, rdx: %.2e, rpy: %.2e, rwv: %.2e, p: %.2e, y: %.2e\n", + std::abs(objective_gap_1 - (xz + wv)), + rdx, + rpy, + rwv, + p, + y); +#endif } template From df960e4577e28655ccf2d14cf0aefea0234cf9d3 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Fri, 28 Aug 2026 03:39:17 -0700 Subject: [PATCH 5/7] rename scalars Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 6fdbcdf691..222ae818c7 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -280,10 +280,10 @@ class barrier_reduce_helper_t { } // Raw device slots for the caller's own cublasdot() calls. - f_t* cx_slot() { return d_results_.data() + kCx; } - f_t* by_slot() { return d_results_.data() + kBy; } - f_t* uv_slot() { return d_results_.data() + kUv; } - f_t* xqx_slot() { return d_results_.data() + kXQx; } + f_t* cTx_slot() { return d_results_.data() + kCTx; } + f_t* bTy_slot() { return d_results_.data() + kBTy; } + f_t* uTv_slot() { return d_results_.data() + kUTv; } + f_t* xTQx_slot() { return d_results_.data() + kXTQx; } // Single batched device-to-host copy + the one stream synchronize needed before any accessor // below can be read. @@ -305,10 +305,10 @@ class barrier_reduce_helper_t { return result; } f_t mu(f_t mu_denom) const { return (h_results_[kMuXzSum] + h_results_[kMuWvSum]) / mu_denom; } - f_t cx() const { return h_results_[kCx]; } - f_t by() const { return h_results_[kBy]; } - f_t uv() const { return h_results_[kUv]; } - f_t xqx() const { return h_results_[kXQx]; } + f_t cTx() const { return h_results_[kCTx]; } + f_t bTy() const { return h_results_[kBTy]; } + f_t uTv() const { return h_results_[kUTv]; } + f_t xTQx() const { return h_results_[kXTQx]; } private: enum Slot : i_t { @@ -320,10 +320,10 @@ class barrier_reduce_helper_t { kComplCone, kMuXzSum, kMuWvSum, - kCx, - kBy, - kUv, - kXQx, + kCTx, + kBTy, + kUTv, + kXTQx, kCount }; @@ -3971,7 +3971,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_x_.data(), 1, - rh.cx_slot(), + rh.cTx_slot(), stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_b_.size(), @@ -3979,7 +3979,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_y_.data(), 1, - rh.by_slot(), + rh.bTy_slot(), stream_view_)); RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), data.d_restrict_u_.size(), @@ -3987,7 +3987,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_v_.data(), 1, - rh.uv_slot(), + rh.uTv_slot(), stream_view_)); if (data.Q.n > 0) { auto cusparse_d_x = data.cusparse_view_.create_vector(data.d_x_); @@ -3999,7 +3999,7 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( 1, data.d_x_.data(), 1, - rh.xqx_slot(), + rh.xTQx_slot(), stream_view_)); } @@ -4012,9 +4012,9 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); mu = rh.mu(mu_denom); - const f_t quad_objective = (data.Q.n > 0) ? 0.5 * rh.xqx() : f_t(0); - primal_objective = rh.cx() + quad_objective; - dual_objective = rh.by() - rh.uv() - quad_objective; + const f_t quad_objective = (data.Q.n > 0) ? 0.5 * rh.xTQx() : f_t(0); + primal_objective = rh.cTx() + quad_objective; + dual_objective = rh.bTy() - rh.uTv() - quad_objective; #ifdef CHECK_OBJECTIVE_GAP rmm::device_scalar d_xz(stream_view_); From 5263121d82ef1ffb84d05674e21d65d98ef34374 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Fri, 28 Aug 2026 03:52:13 -0700 Subject: [PATCH 6/7] move math of barrier operations out of new class Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 222ae818c7..26536eb358 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -267,7 +267,6 @@ class barrier_reduce_helper_t { void cone_complementarity_residual_async(raft::device_span cone_dot, rmm::cuda_stream_view stream_view) { - has_soc_ = true; max_async(kComplCone, cone_dot.data(), cone_dot.size(), stream_view); } @@ -293,18 +292,15 @@ class barrier_reduce_helper_t { stream_view.synchronize(); } - f_t primal_residual_norm() const - { - return std::max(h_results_[kPrimalResidual], h_results_[kBoundResidual]); - } + // Raw reduced values; the caller combines these into residual norms, mu, and objectives. + f_t primal_residual() const { return h_results_[kPrimalResidual]; } + f_t bound_residual() const { return h_results_[kBoundResidual]; } f_t dual_residual_norm() const { return h_results_[kDualResidual]; } - f_t complementarity_residual_norm() const - { - f_t result = std::max(h_results_[kComplXzLinear], h_results_[kComplWv]); - if (has_soc_) { result = std::max(result, h_results_[kComplCone]); } - return result; - } - f_t mu(f_t mu_denom) const { return (h_results_[kMuXzSum] + h_results_[kMuWvSum]) / mu_denom; } + f_t complementarity_xz_linear() const { return h_results_[kComplXzLinear]; } + f_t complementarity_wv() const { return h_results_[kComplWv]; } + f_t complementarity_cone() const { return h_results_[kComplCone]; } + f_t mu_xz_sum() const { return h_results_[kMuXzSum]; } + f_t mu_wv_sum() const { return h_results_[kMuWvSum]; } f_t cTx() const { return h_results_[kCTx]; } f_t bTy() const { return h_results_[kBTy]; } f_t uTv() const { return h_results_[kUTv]; } @@ -365,7 +361,6 @@ class barrier_reduce_helper_t { rmm::device_uvector d_results_; pinned_dense_vector_t h_results_; rmm::device_buffer d_temp_storage_; - bool has_soc_ = false; }; template @@ -4005,12 +4000,17 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( rh.sync(stream_view_); - primal_residual_norm = rh.primal_residual_norm(); - dual_residual_norm = rh.dual_residual_norm(); - complementarity_residual_norm = rh.complementarity_residual_norm(); + primal_residual_norm = std::max(rh.primal_residual(), rh.bound_residual()); + dual_residual_norm = rh.dual_residual_norm(); + + complementarity_residual_norm = std::max(rh.complementarity_xz_linear(), rh.complementarity_wv()); + if (has_soc) { + complementarity_residual_norm = + std::max(complementarity_residual_norm, rh.complementarity_cone()); + } const f_t mu_denom = data.complementarity_degree(data.x.size(), data.n_upper_bounds); - mu = rh.mu(mu_denom); + mu = (rh.mu_xz_sum() + rh.mu_wv_sum()) / mu_denom; const f_t quad_objective = (data.Q.n > 0) ? 0.5 * rh.xTQx() : f_t(0); primal_objective = rh.cTx() + quad_objective; From 5fac90bd36fc28efbb7c07e140142f4761963b44 Mon Sep 17 00:00:00 2001 From: yuwenchen95 Date: Fri, 28 Aug 2026 04:12:14 -0700 Subject: [PATCH 7/7] Abstract scalar computation into helper class 'barrier_reduce_helper_t' Signed-off-by: yuwenchen95 --- cpp/src/barrier/barrier.cu | 83 +++++++++++++++++++++----------------- 1 file changed, 46 insertions(+), 37 deletions(-) diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 26536eb358..f92744f14a 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -278,11 +278,37 @@ class barrier_reduce_helper_t { sum_async(kMuWvSum, d_wv.data(), d_wv.size(), stream_view); } - // Raw device slots for the caller's own cublasdot() calls. - f_t* cTx_slot() { return d_results_.data() + kCTx; } - f_t* bTy_slot() { return d_results_.data() + kBTy; } - f_t* uTv_slot() { return d_results_.data() + kUTv; } - f_t* xTQx_slot() { return d_results_.data() + kXTQx; } + void cTx_async(const rmm::device_uvector& d_c, + const rmm::device_uvector& d_x, + cublasHandle_t cublas_handle, + rmm::cuda_stream_view stream_view) + { + dot_async(kCTx, d_c, d_x, cublas_handle, stream_view); + } + + void bTy_async(const rmm::device_uvector& d_b, + const rmm::device_uvector& d_y, + cublasHandle_t cublas_handle, + rmm::cuda_stream_view stream_view) + { + dot_async(kBTy, d_b, d_y, cublas_handle, stream_view); + } + + void uTv_async(const rmm::device_uvector& d_u, + const rmm::device_uvector& d_v, + cublasHandle_t cublas_handle, + rmm::cuda_stream_view stream_view) + { + dot_async(kUTv, d_u, d_v, cublas_handle, stream_view); + } + + void xTQx_async(const rmm::device_uvector& d_Qx, + const rmm::device_uvector& d_x, + cublasHandle_t cublas_handle, + rmm::cuda_stream_view stream_view) + { + dot_async(kXTQx, d_Qx, d_x, cublas_handle, stream_view); + } // Single batched device-to-host copy + the one stream synchronize needed before any accessor // below can be read. @@ -358,6 +384,16 @@ class barrier_reduce_helper_t { cub::DeviceReduce::Sum(d_temp_storage_.data(), temp_storage_bytes, in, out, size, stream_view); } + void dot_async(Slot slot, + const rmm::device_uvector& a, + const rmm::device_uvector& b, + cublasHandle_t cublas_handle, + rmm::cuda_stream_view stream_view) + { + RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot( + cublas_handle, a.size(), a.data(), 1, b.data(), 1, d_results_.data() + slot, stream_view)); + } + rmm::device_uvector d_results_; pinned_dense_vector_t h_results_; rmm::device_buffer d_temp_storage_; @@ -3960,42 +3996,15 @@ void barrier_solver_t::compute_residual_norms_mu_and_objective( rh.mu_terms_async( data.d_complementarity_xz_residual_, data.d_complementarity_wv_residual_, stream_view_); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_c_.size(), - data.d_c_.data(), - 1, - data.d_x_.data(), - 1, - rh.cTx_slot(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_b_.size(), - data.d_b_.data(), - 1, - data.d_y_.data(), - 1, - rh.bTy_slot(), - stream_view_)); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_restrict_u_.size(), - data.d_restrict_u_.data(), - 1, - data.d_v_.data(), - 1, - rh.uTv_slot(), - stream_view_)); + cublasHandle_t cublas_handle = lp.handle_ptr->get_cublas_handle(); + rh.cTx_async(data.d_c_, data.d_x_, cublas_handle, stream_view_); + rh.bTy_async(data.d_b_, data.d_y_, cublas_handle, stream_view_); + rh.uTv_async(data.d_restrict_u_, data.d_v_, cublas_handle, stream_view_); if (data.Q.n > 0) { auto cusparse_d_x = data.cusparse_view_.create_vector(data.d_x_); auto cusparse_Qx = data.cusparse_view_.create_vector(data.d_Qx_); data.cusparse_Q_view_.spmv(1.0, cusparse_d_x, 0.0, cusparse_Qx); - RAFT_CUBLAS_TRY(raft::linalg::detail::cublasdot(lp.handle_ptr->get_cublas_handle(), - data.d_Qx_.size(), - data.d_Qx_.data(), - 1, - data.d_x_.data(), - 1, - rh.xTQx_slot(), - stream_view_)); + rh.xTQx_async(data.d_Qx_, data.d_x_, cublas_handle, stream_view_); } rh.sync(stream_view_);