diff --git a/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu b/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu index ba141b413f..e5d0ab7b33 100644 --- a/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu +++ b/cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu @@ -99,7 +99,7 @@ void multi_gpu_engine_t::distributed_bound_objective_rescaling(f_t c_s // -------- Refresh halo of cumulative scalings ----------------------------- // Refreshes the halo copies of the cumulative variable + constraint scalings on -// every shard. Called before and after each matrix-scaling pass in ruiz and pock-chambolle +// every shard. Called around the matrix-scaling passes (Curtis-Reid, Ruiz, Pock-Chambolle) template void multi_gpu_engine_t::refresh_halo_cummulative_scalings() { @@ -111,6 +111,43 @@ void multi_gpu_engine_t::refresh_halo_cummulative_scalings() }); } +// -------- Distributed Curtis-Reid scaling --------------------------------- +// Owned rows of A and owned columns of A_T are complete, so each log-mean is +// local. The other axis's log-scale is read at halo indices, so it is exchanged +// between the row pass and the column pass. +template +void multi_gpu_engine_t::distributed_curtis_reid_scaling(int num_iter, i_t n_global_vars) +{ + raft::common::nvtx::range scope("distributed_curtis_reid_scaling"); + + for_each_shard( + [](auto& shard) { shard.sub_pdlp->get_initial_scaling_strategy().curtis_reid_init(); }); + + for (int it = 0; it < num_iter; ++it) { + for_each_shard([](auto& shard) { + shard.sub_pdlp->get_initial_scaling_strategy().curtis_reid_row_iteration(); + }); + halo_exchange_cstr([](pdlp_solver_t& p) -> auto& { + return p.get_initial_scaling_strategy().get_iteration_constraint_matrix_scaling(); + }); + + for_each_shard([](auto& shard) { + shard.sub_pdlp->get_initial_scaling_strategy().curtis_reid_col_iteration(); + }); + halo_exchange_var([](pdlp_solver_t& p) -> auto& { + return p.get_initial_scaling_strategy().get_iteration_variable_scaling(); + }); + } + + for_each_shard( + [](auto& shard) { shard.sub_pdlp->get_initial_scaling_strategy().curtis_reid_folding(); }); + + // Downstream passes read halo copies of the cumulative scaling. + refresh_halo_cummulative_scalings(); + + synchronize_shards(); +} + // -------- Distributed Ruiz inf-scaling ------------------------------------ // Each shard owns its rows AND its columns and stores both complete (h_A = // owned rows, h_A_t = owned columns) @@ -171,6 +208,10 @@ void multi_gpu_engine_t::distributed_scaling(pdlp_hyper_params_t const // 1) Matrix scaling passes populate the cumulative row/col scalings on // every shard. Each pass keeps the halo copies refreshed internally. + // Curtis-Reid is a prescale and is skipped inside MIP, matching single-GPU. + if (hyper_params.do_curtis_reid_scaling && !inside_mip) { + distributed_curtis_reid_scaling(hyper_params.number_of_curtis_reid_iterations, n_global_vars); + } if (hyper_params.do_ruiz_scaling) { distributed_ruiz_inf_scaling(hyper_params.default_l_inf_ruiz_iterations, n_global_vars); } @@ -411,6 +452,7 @@ void multi_gpu_engine_t::distributed_compute_initial_primal_weight( template void multi_gpu_engine_t::gather_potential_next_solutions_to_master(); \ template void multi_gpu_engine_t::refresh_halo_cummulative_scalings(); \ template void multi_gpu_engine_t::distributed_bound_objective_rescaling(F_TYPE); \ + template void multi_gpu_engine_t::distributed_curtis_reid_scaling(int, int); \ template void multi_gpu_engine_t::distributed_ruiz_inf_scaling(int, int); \ template void multi_gpu_engine_t::distributed_pock_chambolle_scaling(F_TYPE, int); \ template void multi_gpu_engine_t::distributed_scaling( \ diff --git a/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp index da22ee556e..59ccefdcf6 100644 --- a/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp +++ b/cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp @@ -458,7 +458,7 @@ struct multi_gpu_engine_t { // -------- High-level algorithms (defined in distributed_algorithms.cu) --- // Refreshes the halo copies of the cumulative variable + constraint scalings on - // every shard. Used by the matrix-scaling passes (Ruiz, Pock-Chambolle) + // every shard. Used by the matrix-scaling passes (Curtis-Reid, Ruiz, Pock-Chambolle) void refresh_halo_cummulative_scalings(); // Global bound/objective rescaling: allreduce the owned partial squared norms @@ -466,6 +466,12 @@ struct multi_gpu_engine_t { // scalar on every shard. void distributed_bound_objective_rescaling(f_t c_scaling_weight); + // Distributed Curtis-Reid prescaling. Each iteration is a shard-local row log-mean, + // a constraint-halo exchange of that log-scale, a shard-local column log-mean, and a + // variable-halo exchange. After the last iteration every shard folds + // cumulative *= exp(clamp(log_scale)) and the cumulative halo is refreshed. + void distributed_curtis_reid_scaling(int num_iter, i_t n_global_vars); + // Distributed Ruiz inf-scaling (num_iter passes). Each shard computes both its // owned-row and owned-column inf-norms locally then broadcasts the cumulative scalings to all // shards. @@ -477,7 +483,8 @@ struct multi_gpu_engine_t { // Full distributed scaling entry point. Mirrors what scale_problem() does in // single-GPU by orchestrating: - // - Ruiz inf-scaling -> populates cumulative row/col scalings + // - Curtis-Reid prescaling (skipped inside MIP) -> populates cumulative row/col scalings + // - Ruiz inf-scaling -> same // - Pock-Chambolle scaling -> same // - per-shard apply_cummulative_scaling_to_problem() // - global bound/objective rescaling via distributed_bound_objective_rescaling diff --git a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu index 68e485694f..433d1c18f7 100644 --- a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu +++ b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu @@ -479,6 +479,78 @@ __global__ void curtis_reid_col_kernel(i_t n_variables, } } +template +void pdlp_initial_scaling_strategy_t::curtis_reid_init() +{ + thrust::fill(handle_ptr_->get_thrust_policy(), + iteration_constraint_matrix_scaling_.begin(), + iteration_constraint_matrix_scaling_.end(), + f_t(0)); + thrust::fill(handle_ptr_->get_thrust_policy(), + iteration_variable_scaling_.begin(), + iteration_variable_scaling_.end(), + f_t(0)); +} + +template +void pdlp_initial_scaling_strategy_t::curtis_reid_row_iteration() +{ + constexpr i_t number_of_threads = 128; + if (dual_size_h_ <= 0) return; + curtis_reid_row_kernel + <<>>( + op_problem_scaled_.view(), + cummulative_constraint_matrix_scaling_.data(), + cummulative_variable_scaling_.data(), + iteration_variable_scaling_.data(), + iteration_constraint_matrix_scaling_.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +template +void pdlp_initial_scaling_strategy_t::curtis_reid_col_iteration() +{ + constexpr i_t number_of_threads = 128; + if (primal_size_h_ <= 0) return; + curtis_reid_col_kernel + <<>>( + primal_size_h_, + A_T_.data(), + A_T_offsets_.data(), + A_T_indices_.data(), + cummulative_constraint_matrix_scaling_.data(), + cummulative_variable_scaling_.data(), + iteration_constraint_matrix_scaling_.data(), + iteration_variable_scaling_.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +template +void pdlp_initial_scaling_strategy_t::curtis_reid_folding() +{ + // Fold the converged log-domain fit into the cumulative scale (exp + clamp, see + // a_times_exp_clamped_b): cummulative *= exp(clamp(log_scale)). Unlike Ruiz/ + // Pock-Chambolle's a_divides_sqrt_b_bounded fold (which incorporates *this iteration's* + // norm into a running cumulative), Curtis-Reid already produces the final multiplicative + // scale factor directly, so a straight multiply is correct here. + // + // clamp_bound = 30 (exp(+-30) ~ [9.4e-14, 1.07e13]) bounds the scale-factor range a + // pathological log-domain fit could produce. + constexpr f_t clamp_bound = f_t(30); + raft::linalg::binaryOp(cummulative_constraint_matrix_scaling_.data(), + cummulative_constraint_matrix_scaling_.data(), + iteration_constraint_matrix_scaling_.data(), + dual_size_h_, + a_times_exp_clamped_b(clamp_bound), + stream_view_.get()); + raft::linalg::binaryOp(cummulative_variable_scaling_.data(), + cummulative_variable_scaling_.data(), + iteration_variable_scaling_.data(), + primal_size_h_, + a_times_exp_clamped_b(clamp_bound), + stream_view_.get()); +} + // Curtis-Reid prescaling (A. R. Curtis, J. K. Reid, "On the Automatic Scaling of Matrices // for Gaussian Elimination", IMA J. Applied Mathematics, 1972; also IIASA Collaborative // Paper CP-81-037, https://pure.iiasa.ac.at/id/eprint/1766/7/CP-81-037.pdf): a log-domain @@ -486,6 +558,7 @@ __global__ void curtis_reid_col_kernel(i_t n_variables, // sum((log|a_ij| - row_log_scale[i] - col_log_scale[j])^2) via alternating per-row/ // per-column log-mean fixed-point iteration. This port's sequence and defaults are // inspired by the HPR-LP-C codebase (https://github.com/PolyU-IOR/HPR-LP-C). +// Single-GPU entry point. Distributed PDLP calls the init/row/col/fold pieces directly. template void pdlp_initial_scaling_strategy_t::curtis_reid_scaling( i_t number_of_curtis_reid_iterations) @@ -500,60 +573,16 @@ void pdlp_initial_scaling_strategy_t::curtis_reid_scaling( // as-if-already-scaled by the current cummulative_* factors (like Ruiz/Pock-Chambolle's // own kernels do); Curtis-Reid always runs first in compute_scaling_vectors(), so // cummulative_* is still all-1 here in practice. - auto& row_log_scale = iteration_constraint_matrix_scaling_; - auto& col_log_scale = iteration_variable_scaling_; - RAFT_CUDA_TRY( - cudaMemsetAsync(row_log_scale.data(), 0, sizeof(f_t) * dual_size_h_, stream_view_.get())); - RAFT_CUDA_TRY( - cudaMemsetAsync(col_log_scale.data(), 0, sizeof(f_t) * primal_size_h_, stream_view_.get())); + curtis_reid_init(); - constexpr i_t number_of_threads = 128; for (i_t iter = 0; iter < number_of_curtis_reid_iterations; ++iter) { - curtis_reid_row_kernel - <<>>( - op_problem_scaled_.view(), - cummulative_constraint_matrix_scaling_.data(), - cummulative_variable_scaling_.data(), - col_log_scale.data(), - row_log_scale.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - - curtis_reid_col_kernel - <<>>( - primal_size_h_, - A_T_.data(), - A_T_offsets_.data(), - A_T_indices_.data(), - cummulative_constraint_matrix_scaling_.data(), - cummulative_variable_scaling_.data(), - row_log_scale.data(), - col_log_scale.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); + curtis_reid_row_iteration(); + curtis_reid_col_iteration(); } if (running_mip_) { reset_integer_variables(); } - // Fold the converged log-domain fit into the cumulative scale (exp + clamp, see - // a_times_exp_clamped_b): cummulative *= exp(clamp(log_scale)). Unlike Ruiz/ - // Pock-Chambolle's a_divides_sqrt_b_bounded fold (which incorporates *this iteration's* - // norm into a running cumulative), Curtis-Reid already produces the final multiplicative - // scale factor directly, so a straight multiply is correct here. - // - // clamp_bound = 30 (exp(+-30) ~ [9.4e-14, 1.07e13]) bounds the scale-factor range a - // pathological log-domain fit could produce. - constexpr f_t clamp_bound = f_t(30); - raft::linalg::binaryOp(cummulative_constraint_matrix_scaling_.data(), - cummulative_constraint_matrix_scaling_.data(), - row_log_scale.data(), - dual_size_h_, - a_times_exp_clamped_b(clamp_bound), - stream_view_.get()); - raft::linalg::binaryOp(cummulative_variable_scaling_.data(), - cummulative_variable_scaling_.data(), - col_log_scale.data(), - primal_size_h_, - a_times_exp_clamped_b(clamp_bound), - stream_view_.get()); + curtis_reid_folding(); } template diff --git a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh index 1f8e24ad26..90a3f76306 100644 --- a/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh +++ b/cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cuh @@ -129,10 +129,22 @@ class pdlp_initial_scaling_strategy_t { void ruiz_iter_local(); // Shard-local end-to-end Pock-Chambolle pass. Exposed for distributed PDLP: void pock_chambolle_scaling(f_t alpha); - // Curtis-Reid prescaling pass -- see the implementation in initial_scaling.cu for - // details and references. Not exposed to distributed PDLP yet (no cross-shard-coherent - // version written). + // Curtis-Reid prescaling. Single-GPU orchestrator: curtis_reid_init, then alternate + // curtis_reid_row_iteration / curtis_reid_col_iteration, then curtis_reid_folding. + // Distributed PDLP calls the pieces itself so a halo exchange can sit between + // the row and column passes. See initial_scaling.cu for the algorithm and references. void curtis_reid_scaling(i_t number_of_curtis_reid_iterations); + // Zero both log-scale vectors, halo included. The first row pass reads column log-scales. + void curtis_reid_init(); + // One row log-mean pass. Writes iteration_constraint_matrix_scaling_ from the current + // column log-scales in iteration_variable_scaling_. + void curtis_reid_row_iteration(); + // One column log-mean pass. Writes iteration_variable_scaling_ from the current + // row log-scales in iteration_constraint_matrix_scaling_. + void curtis_reid_col_iteration(); + // Fold the converged log-scales into the cumulative scaling: + // cumulative *= exp(clamp(log_scale)). + void curtis_reid_folding(); // Iteration_* scratch buffers used by ruiz_iter_local / // pock_chambolle_scaling. Exposed mutably so distributed PDLP can grow // them back to full size after the ctor's release (see distributed_scaling). diff --git a/cpp/src/pdlp/pdlp.cu b/cpp/src/pdlp/pdlp.cu index fdb8b64d7d..28351bbb72 100644 --- a/cpp/src/pdlp/pdlp.cu +++ b/cpp/src/pdlp/pdlp.cu @@ -529,9 +529,10 @@ pdlp_solver_t::pdlp_solver_t( // ----- 5. Per-shard settings ----- pdlp_solver_settings_t sub_pdlp_settings = settings; sub_pdlp_settings.num_gpus = 1; - // Disable automatic ruiz and pock-chambolle in the initial_scaling ctor: the - // distributed pipeline computes them via distributed_scaling using the - // GLOBAL problem. + // Disable automatic matrix scaling in the initial_scaling ctor: the + // distributed pipeline computes Curtis-Reid, Ruiz, and Pock-Chambolle via + // distributed_scaling using the global problem. + sub_pdlp_settings.hyper_params.do_curtis_reid_scaling = false; sub_pdlp_settings.hyper_params.do_ruiz_scaling = false; sub_pdlp_settings.hyper_params.do_pock_chambolle_scaling = false; diff --git a/cpp/tests/linear_programming/pdlp_distributed_test.cu b/cpp/tests/linear_programming/pdlp_distributed_test.cu index 0d0509c8d1..b3831efcfc 100644 --- a/cpp/tests/linear_programming/pdlp_distributed_test.cu +++ b/cpp/tests/linear_programming/pdlp_distributed_test.cu @@ -47,8 +47,6 @@ static void expect_distributed_matches_base(raft::handle_t const& handle, pdlp_solver_settings_t base_settings{}; base_settings.method = method_t::PDLP; - // Curtis-Reid scaling is not supported yet for multi-GPU. - base_settings.hyper_params.do_curtis_reid_scaling = false; auto base_op = mps_data_model_to_optimization_problem(&handle, problem); auto base = solve_lp(base_op, base_settings);