diff --git a/cpp/src/barrier/barrier.cu b/cpp/src/barrier/barrier.cu index 8c6988649c..f92744f14a 100644 --- a/cpp/src/barrier/barrier.cu +++ b/cpp/src/barrier/barrier.cu @@ -230,6 +230,175 @@ 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) + { + 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); + } + + 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. + void sync(rmm::cuda_stream_view stream_view) + { + raft::copy(h_results_.data(), d_results_.data(), static_cast(kCount), stream_view); + stream_view.synchronize(); + } + + // 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_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]; } + f_t xTQx() const { return h_results_[kXTQx]; } + + private: + enum Slot : i_t { + kPrimalResidual = 0, + kBoundResidual, + kDualResidual, + kComplXzLinear, + kComplWv, + kComplCone, + kMuXzSum, + kMuWvSum, + kCTx, + kBTy, + kUTv, + kXTQx, + 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); + } + + 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_; +}; + template class iteration_data_t { public: @@ -356,6 +525,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()), + reduce_helper_(lp.handle_ptr->get_stream()), indefinite_Q(false), Q_diagonal(false), symbolic_status(0), @@ -2079,6 +2249,8 @@ class iteration_data_t { transform_reduce_pair_helper_t transform_reduce_pair_helper_; sum_reduce_helper_t sum_reduce_helper_; + barrier_reduce_helper_t reduce_helper_; + bool cone_combined_step_; f_t cone_sigma_mu_; @@ -2570,46 +2742,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, @@ -3828,91 +3960,70 @@ 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; -} + auto& rh = data.reduce_helper_; -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_); + 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); - 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, - d_cx.data(), - 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, - d_by.data(), - 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, - d_uv.data(), - stream_view_)); - f_t quad_objective = 0.0; + // 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. + 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_); + rh.cone_complementarity_residual_async(cone_dot, stream_view_); + } + rh.mu_terms_async( + data.d_complementarity_xz_residual_, data.d_complementarity_wv_residual_, 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); - 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(), - stream_view_)); - quad_objective = 0.5 * d_xQx.value(stream_view_); + rh.xTQx_async(data.d_Qx_, data.d_x_, cublas_handle, 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; + rh.sync(stream_view_); + + 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_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; + dual_objective = rh.bTy() - rh.uTv() - quad_objective; #ifdef CHECK_OBJECTIVE_GAP rmm::device_scalar d_xz(stream_view_); @@ -4214,29 +4325,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); @@ -4245,14 +4352,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 / @@ -4389,12 +4488,13 @@ lp_status_t barrier_solver_t::solve(f_t start_time, lp_solution_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, f_t& step_primal, f_t& step_dual); - void compute_residual_norms(iteration_data_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 +71,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,