Solve mathematical programs in parallel

bindings/pydrake/solvers/solvers_py_mathematicalprogram.cc

@@ -1682,7 +1682,100 @@ void BindFreeFunctions(py::module m) {
               const std::optional<SolverOptions>&>(&solvers::Solve),
           py::arg("prog"), py::arg("initial_guess") = py::none(),
           py::arg("solver_options") = py::none(), doc.Solve.doc_3args)
-      .def("GetProgramType", &solvers::GetProgramType, doc.GetProgramType.doc);
+      .def("GetProgramType", &solvers::GetProgramType, doc.GetProgramType.doc)
+      //      .def("SolveInParallel",
+      //          py::overload_cast<const std::vector<const
+      //          MathematicalProgram*>&,
+      //              const std::vector<const Eigen::VectorXd*>*,
+      //              const std::vector<const SolverOptions*>*,
+      //              const std::vector<std::optional<SolverId>>*, const
+      //              Parallelism, bool>(&solvers::SolveInParallel),
+      //          py::arg("progs"), py::arg("initial_guesses"),
+      //          py::arg("solver_options"), py::arg("solver_ids"),
+      //          py::arg("parallelism") = Parallelism::Max(),
+      //          py::arg("dynamic_schedule") = false,
+      //          doc.SolveInParallel
+      //              .doc_6args_progs_initial_guesses_solver_options_solver_ids_parallelism_dynamic_schedule)
+      //      .def("SolveInParallel",
+      //          py::overload_cast<const std::vector<const
+      //          MathematicalProgram*>&,
+      //              const std::vector<const Eigen::VectorXd*>*,
+      //              const std::optional<SolverOptions>&,
+      //              const std::optional<SolverId>&, const Parallelism, bool>(
+      //              &solvers::SolveInParallel),
+      //          py::arg("progs"), py::arg("initial_guesses") = nullptr,
+      //          py::arg("solver_options") = py::none(),
+      //          py::arg("solver_id") = py::none(),
+      //          py::arg("parallelism") = Parallelism::Max(),
+      //          py::arg("dynamic_schedule") = false,
+      //          doc.SolveInParallel
+      //              .doc_6args_progs_initial_guesses_solver_options_solver_id_parallelism_dynamic_schedule)
+      // Claude's attempt
+      .def(
+          "SolveInParallel",
+          [](const std::vector<const MathematicalProgram*>& progs,
+              const std::optional<std::vector<Eigen::VectorXd>>&
+                  initial_guesses,
+              const std::optional<std::vector<SolverOptions>>& solver_options,
+              const std::optional<std::vector<std::optional<SolverId>>>&
+                  solver_ids,
+              const Parallelism& parallelism, bool dynamic_schedule) {
+            std::vector<const Eigen::VectorXd*> initial_guesses_ptrs;
+            if (initial_guesses.has_value()) {
+              initial_guesses_ptrs.reserve(initial_guesses->size());
+              for (const auto& guess : *initial_guesses) {
+                initial_guesses_ptrs.push_back(&guess);
+              }
+            }
+
+            std::vector<const SolverOptions*> solver_options_ptrs;
+            if (solver_options.has_value()) {
+              solver_options_ptrs.reserve(solver_options->size());
+              for (const auto& options : *solver_options) {
+                solver_options_ptrs.push_back(&options);
+              }
+            }
+
+            return solvers::SolveInParallel(progs,
+                initial_guesses.has_value() ? &initial_guesses_ptrs : nullptr,
+                solver_options.has_value() ? &solver_options_ptrs : nullptr,
+                solver_ids.has_value() ? &solver_ids.value() : nullptr,
+                parallelism, dynamic_schedule);
+          },
+          py::arg("progs"), py::arg("initial_guesses") = std::nullopt,
+          py::arg("solver_options") = std::nullopt,
+          py::arg("solver_ids") = std::nullopt,
+          py::arg("parallelism") = Parallelism::Max(),
+          py::arg("dynamic_schedule") = false,
+          doc.SolveInParallel
+              .doc_6args_progs_initial_guesses_solver_options_solver_ids_parallelism_dynamic_schedule)
+      .def(
+          "SolveInParallel",
+          [](const std::vector<const MathematicalProgram*>& progs,
+              const std::optional<std::vector<Eigen::VectorXd>>&
+                  initial_guesses,
+              const std::optional<SolverOptions>& solver_options,
+              const std::optional<SolverId>& solver_id,
+              const Parallelism& parallelism, bool dynamic_schedule) {
+            std::vector<const Eigen::VectorXd*> initial_guesses_ptrs;
+            if (initial_guesses.has_value()) {
+              initial_guesses_ptrs.reserve(initial_guesses->size());
+              for (const auto& guess : *initial_guesses) {
+                initial_guesses_ptrs.push_back(&guess);
+              }
+            }
+
+            return solvers::SolveInParallel(progs,
+                initial_guesses.has_value() ? &initial_guesses_ptrs : nullptr,
+                solver_options, solver_id, parallelism, dynamic_schedule);
+          },
+          py::arg("progs"), py::arg("initial_guesses") = std::nullopt,
+          py::arg("solver_options") = std::nullopt,
+          py::arg("solver_id") = std::nullopt,
+          py::arg("parallelism") = Parallelism::Max(),
+          py::arg("dynamic_schedule") = false,
+          doc.SolveInParallel
+              .doc_6args_progs_initial_guesses_solver_options_solver_id_parallelism_dynamic_schedule);
 }
 
 }  // namespace

bindings/pydrake/solvers/test/mathematicalprogram_test.py

@@ -7,7 +7,7 @@
 import scipy.sparse
 
 from pydrake.autodiffutils import AutoDiffXd
-from pydrake.common import kDrakeAssertIsArmed
+from pydrake.common import kDrakeAssertIsArmed, Parallelism
 from pydrake.common.test_utilities import numpy_compare
 from pydrake.forwarddiff import jacobian
 from pydrake.math import ge
@@ -1527,6 +1527,81 @@ def test_mathematical_program_result(self):
         self.assertEqual(result.get_solution_result(),
                          mp.SolutionResult.kSolutionResultNotSet)
 
+    def test_solve_in_parallel(self):
+        prog = mp.MathematicalProgram()
+        x = prog.NewContinuousVariables(2)
+        prog.AddLinearConstraint(x[0] + x[1] == 2)
+        prog.AddQuadraticCost(x[0] ** 2, is_convex=True)
+
+        num_progs = 4
+        progs = [prog for _ in range(num_progs)]
+        initial_guesses = [np.zeros(2) for _ in range(num_progs)]
+        solver_ids = [ScsSolver().solver_id() for _ in range(num_progs)]
+        solver_options = [SolverOptions() for _ in range(num_progs)]
+
+        # results = mp.SolveInParallel(progs)
+
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=initial_guesses,
+                                     solver_options=solver_options,
+                                     solver_ids=solver_ids,
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=None,
+                                     solver_options=solver_options,
+                                     solver_ids=solver_ids,
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=initial_guesses,
+                                     solver_options=None,
+                                     solver_ids=solver_ids,
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=initial_guesses,
+                                     solver_options=solver_options,
+                                     solver_ids=None,
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        # Now we test the overload
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=None,
+                                     solver_options=SolverOptions(),
+                                     solver_id=ScsSolver().solver_id(),
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        # Ensure that all options being None does not cause ambiguity.
+        results = mp.SolveInParallel(progs=progs,
+                                     initial_guesses=None,
+                                     solver_options=None,
+                                     solver_id=None,
+                                     parallelism=Parallelism.Max(),
+                                     dynamic_schedule=False)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
+        # Ensure default arguments do not cause ambiguity.
+        results = mp.SolveInParallel(progs=progs)
+        self.assertEqual(len(results), len(progs))
+        self.assertTrue(all([r.is_success() for r in results]))
+
 
 class DummySolverInterface(SolverInterface):
 

solvers/BUILD.bazel

@@ -726,10 +726,12 @@ drake_cc_library(
         ":mathematical_program",
         ":mathematical_program_result",
         ":solver_base",
+        "//common:parallelism",
     ],
     implementation_deps = [
         ":choose_best_solver",
         "//common:nice_type_name",
+        "@common_robotics_utilities",
     ],
 )
 
@@ -2003,6 +2005,9 @@ drake_cc_googletest(
         "//common/test_utilities:eigen_matrix_compare",
         "//common/test_utilities:expect_throws_message",
     ],
+    # Running with multiple threads is an essential part of testing
+    # SolveInParallel.
+    num_threads = 2
 )
 
 drake_cc_googletest(

solvers/solve.cc

@@ -1,6 +1,11 @@
 #include "drake/solvers/solve.h"
 
 #include <memory>
+#include <unordered_map>
+#include <utility>
+#include <variant>
+
+#include <common_robotics_utilities/parallelism.hpp>
 
 #include "drake/common/nice_type_name.h"
 #include "drake/common/text_logging.h"
@@ -9,6 +14,11 @@
 
 namespace drake {
 namespace solvers {
+using common_robotics_utilities::parallelism::DegreeOfParallelism;
+using common_robotics_utilities::parallelism::DynamicParallelForIndexLoop;
+using common_robotics_utilities::parallelism::ParallelForBackend;
+using common_robotics_utilities::parallelism::StaticParallelForIndexLoop;
+
 MathematicalProgramResult Solve(
     const MathematicalProgram& prog,
     const std::optional<Eigen::VectorXd>& initial_guess,
@@ -31,5 +41,154 @@ MathematicalProgramResult Solve(
 MathematicalProgramResult Solve(const MathematicalProgram& prog) {
   return Solve(prog, {}, {});
 }
+
+namespace {
+std::vector<MathematicalProgramResult> SolveInParallelImpl(
+    const std::vector<const MathematicalProgram*>& progs,
+    const std::vector<const Eigen::VectorXd*>* initial_guesses,
+    const std::vector<const SolverOptions*>* solver_options,
+    const std::vector<std::optional<SolverId>>* solver_ids,
+    const Parallelism parallelism, bool dynamic_schedule) {
+  DRAKE_THROW_UNLESS(std::all_of(progs.begin(), progs.end(), [](auto prog) {
+    return prog != nullptr;
+  }));
+  if (initial_guesses != nullptr) {
+    DRAKE_THROW_UNLESS(initial_guesses->size() == progs.size());
+  }
+  if (solver_options != nullptr) {
+    DRAKE_THROW_UNLESS(solver_options->size() == progs.size());
+  }
+  if (solver_ids != nullptr) {
+    DRAKE_THROW_UNLESS(solver_ids->size() == progs.size());
+  }
+
+  std::vector<MathematicalProgramResult> results{progs.size()};
+  std::vector<uint8_t> result_is_populated(progs.size(), 0);
+
+  // As unique types of solvers are encountered by the threads, we will cache
+  // the solvers into this data structure so that we don't have to create and
+  // destroy them at every call to Solve.
+  std::vector<std::unordered_map<SolverId, std::unique_ptr<SolverInterface>>>
+      solvers(parallelism.num_threads());
+
+  // Extracts the options from the variant and sets the maximum thread number to
+  // 1 if set_max_threads_to_one is true.
+  auto GetOptions = [&solver_options](const int64_t i,
+                                      bool set_max_threads_to_one) {
+    // Copy the solver options and make sure that they solve with only 1
+    // thread. This is achieved by setting the
+    // CommonSolverOption.kMaxThreads to 1. If the solver options
+    // explicitly specify to use more threads using a solver specific
+    // option, then more threads than 1 may be used.
+    bool option_available =
+        solver_options != nullptr && solver_options->at(i) != nullptr;
+    SolverOptions options =
+        option_available ? *solver_options->at(i) : SolverOptions();
+    if (set_max_threads_to_one) {
+      options.SetOption(CommonSolverOption::kMaxThreads, 1);
+    }
+    return options;
+  };
+
+  auto DoSolveParallel = [&](const int thread_num, const int64_t i) {
+    unused(thread_num);
+
+    if (!progs[i]->IsThreadSafe()) {
+      // If this program is not thread safe, exit after identifying the
+      // necessary solver and solve it later.
+      return;
+    }
+
+    // Access the required solver.
+    const SolverId solver_id =
+        solver_ids == nullptr
+            ? ChooseBestSolver(*(progs.at(i)))
+            : solver_ids->at(i).value_or(ChooseBestSolver(*(progs.at(i))));
+    if (!solvers.at(thread_num).contains(solver_id)) {
+      // If this thread has not solved this type of program yet, save the solver
+      // for use later.
+      solvers.at(thread_num).insert({solver_id, MakeSolver(solver_id)});
+    }
+    const SolverInterface* solver = solvers.at(thread_num).at(solver_id).get();
+    DRAKE_THROW_UNLESS(solver->AreProgramAttributesSatisfied(*(progs.at(i))));
+
+    // Convert the initial guess into the requisite optional.
+    std::optional<Eigen::VectorXd> initial_guess{std::nullopt};
+    if (initial_guesses != nullptr && initial_guesses->at(i) != nullptr) {
+      initial_guess = *(initial_guesses->at(i));
+    }
+
+    const SolverOptions options = GetOptions(i, true /* solving_in_parallel */);
+
+    // Solve the program.
+    solver->Solve(*(progs.at(i)), initial_guess, options, &(results.at(i)));
+    result_is_populated.at(i) = true;
+  };
+
+  if (dynamic_schedule) {
+    DynamicParallelForIndexLoop(DegreeOfParallelism(parallelism.num_threads()),
+                                0, ssize(progs), DoSolveParallel,
+                                ParallelForBackend::BEST_AVAILABLE);
+  } else {
+    StaticParallelForIndexLoop(DegreeOfParallelism(parallelism.num_threads()),
+                               0, ssize(progs), DoSolveParallel,
+                               ParallelForBackend::BEST_AVAILABLE);
+  }
+
+  // Now finish solving the programs which cannot be solved in parallel and
+  // write the final results to the results vector.
+  for (int i = 0; i < ssize(progs); ++i) {
+    if (result_is_populated.at(i)) {
+      continue;
+    }
+    // Convert the initial guess into the requisite optional.
+    std::optional<Eigen::VectorXd> initial_guess{std::nullopt};
+    if (initial_guesses != nullptr && initial_guesses->at(i) != nullptr) {
+      initial_guess = *(initial_guesses->at(i));
+    }
+
+    const SolverOptions options =
+        GetOptions(i, false /* solving_in_parallel */);
+
+    // We will use just the solvers at the 0th index as cached solvers.
+    const SolverId solver_id =
+        solver_ids->at(i).value_or(ChooseBestSolver(*(progs.at(i))));
+    if (!solvers.at(0).contains(solver_id)) {
+      // If this thread has not solved this type of program yet, save the
+      // solver for use later.
+      solvers.at(0).insert({solver_id, MakeSolver(solver_id)});
+    }
+    solvers.at(0).at(solver_id)->Solve(*(progs.at(i)), initial_guess, options,
+                                       &(results.at(i)));
+  }
+  return results;
+}
+
+}  // namespace
+
+std::vector<MathematicalProgramResult> SolveInParallel(
+    const std::vector<const MathematicalProgram*>& progs,
+    const std::vector<const Eigen::VectorXd*>* initial_guesses,
+    const std::vector<const SolverOptions*>* solver_options,
+    const std::vector<std::optional<SolverId>>* solver_ids,
+    const Parallelism parallelism, bool dynamic_schedule) {
+  return SolveInParallelImpl(progs, initial_guesses, solver_options, solver_ids,
+                             parallelism, dynamic_schedule);
+}
+
+std::vector<MathematicalProgramResult> SolveInParallel(
+    const std::vector<const MathematicalProgram*>& progs,
+    const std::vector<const Eigen::VectorXd*>* initial_guesses,
+    const std::optional<SolverOptions>& solver_options,
+    const std::optional<SolverId>& solver_id, const Parallelism parallelism,
+    bool dynamic_schedule) {
+  std::vector<const SolverOptions*> solver_options_vec{
+      progs.size(),
+      solver_options.has_value() ? &(solver_options.value()) : nullptr};
+  std::vector<std::optional<SolverId>> solver_ids{progs.size(), solver_id};
+  return SolveInParallelImpl(progs, initial_guesses, &solver_options_vec,
+                             &solver_ids, parallelism, dynamic_schedule);
+}
+
 }  // namespace solvers
 }  // namespace drake