Distributed curtis reid

cpp/src/pdlp/distributed_pdlp/distributed_algorithms.cu

@@ -99,7 +99,7 @@ void multi_gpu_engine_t<i_t, f_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 <typename i_t, typename f_t>
 void multi_gpu_engine_t<i_t, f_t>::refresh_halo_cummulative_scalings()
 {
@@ -111,6 +111,43 @@ void multi_gpu_engine_t<i_t, f_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 <typename i_t, typename f_t>
+void multi_gpu_engine_t<i_t, f_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<i_t, f_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<i_t, f_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<i_t, f_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<i_t, f_t>::distributed_compute_initial_primal_weight(
   template void multi_gpu_engine_t<int, F_TYPE>::gather_potential_next_solutions_to_master();     \
   template void multi_gpu_engine_t<int, F_TYPE>::refresh_halo_cummulative_scalings();             \
   template void multi_gpu_engine_t<int, F_TYPE>::distributed_bound_objective_rescaling(F_TYPE);   \
+  template void multi_gpu_engine_t<int, F_TYPE>::distributed_curtis_reid_scaling(int, int);       \
   template void multi_gpu_engine_t<int, F_TYPE>::distributed_ruiz_inf_scaling(int, int);          \
   template void multi_gpu_engine_t<int, F_TYPE>::distributed_pock_chambolle_scaling(F_TYPE, int); \
   template void multi_gpu_engine_t<int, F_TYPE>::distributed_scaling(                             \

cpp/src/pdlp/distributed_pdlp/multi_gpu_engine.hpp

@@ -458,14 +458,20 @@ 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
   // of the constraint bounds and (weighted) objective, then apply the identical
   // 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

cpp/src/pdlp/initial_scaling_strategy/initial_scaling.cu

@@ -479,13 +479,86 @@ __global__ void curtis_reid_col_kernel(i_t n_variables,
   }
 }
 
+template <typename i_t, typename f_t>
+void pdlp_initial_scaling_strategy_t<i_t, f_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 <typename i_t, typename f_t>
+void pdlp_initial_scaling_strategy_t<i_t, f_t>::curtis_reid_row_iteration()
+{
+  constexpr i_t number_of_threads = 128;
+  if (dual_size_h_ <= 0) return;
+  curtis_reid_row_kernel<i_t, f_t, number_of_threads>
+    <<<dual_size_h_, number_of_threads, 0, stream_view_.get()>>>(
+      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 <typename i_t, typename f_t>
+void pdlp_initial_scaling_strategy_t<i_t, f_t>::curtis_reid_col_iteration()
+{
+  constexpr i_t number_of_threads = 128;
+  if (primal_size_h_ <= 0) return;
+  curtis_reid_col_kernel<i_t, f_t, number_of_threads>
+    <<<primal_size_h_, number_of_threads, 0, stream_view_.get()>>>(
+      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 <typename i_t, typename f_t>
+void pdlp_initial_scaling_strategy_t<i_t, f_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<f_t>(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<f_t>(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
 // least-squares fit run *before* Ruiz/Pock-Chambolle, minimizing
 // 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 <typename i_t, typename f_t>
 void pdlp_initial_scaling_strategy_t<i_t, f_t>::curtis_reid_scaling(
   i_t number_of_curtis_reid_iterations)
@@ -500,60 +573,16 @@ void pdlp_initial_scaling_strategy_t<i_t, f_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<i_t, f_t, number_of_threads>
-      <<<dual_size_h_, number_of_threads, 0, stream_view_.get()>>>(
-        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<i_t, f_t, number_of_threads>
-      <<<primal_size_h_, number_of_threads, 0, stream_view_.get()>>>(
-        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<f_t>(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<f_t>(clamp_bound),
-                         stream_view_.get());
+  curtis_reid_folding();
 }
 
 template <typename i_t, typename f_t>

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).

cpp/src/pdlp/pdlp.cu

@@ -529,9 +529,10 @@ pdlp_solver_t<i_t, f_t>::pdlp_solver_t(
   // ----- 5. Per-shard settings -----
   pdlp_solver_settings_t<i_t, f_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;
 

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<int, double> 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<int, double>(&handle, problem);
   auto base    = solve_lp(base_op, base_settings);