Skip to content

Commit

Permalink
rewrite python bindings
Browse files Browse the repository at this point in the history
  • Loading branch information
AlexandreAmice committed Dec 4, 2024
1 parent 77a388b commit 100e435
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 144 deletions.
94 changes: 67 additions & 27 deletions bindings/pydrake/solvers/solvers_py_mathematicalprogram.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1637,7 +1637,7 @@ void BindFreeFunctions(py::module m) {
// Inside of the lambda we'll demote nullopts back to nullptrs. Note
// that SolverOptions is not necessarily cheap to copy, so we still
// carefully accept it by-pointer. The VectorXd is always necessarily
// copied when going form numpy to Eigen so we still pass it by-value.
// copied when going from numpy to Eigen so we still pass it by-value.
// TODO(Alexandre.Amice) change to generator form.
[](std::vector<const MathematicalProgram*> progs,
std::optional<std::vector<std::optional<Eigen::VectorXd>>>
Expand All @@ -1646,25 +1646,41 @@ void BindFreeFunctions(py::module m) {
solver_options,
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 ? &(*guess) : nullptr);
const std::function<const MathematicalProgram*(int64_t, int64_t)>
progs_generator =
[&progs](int64_t, int64_t i) -> const MathematicalProgram* {
return progs.at(i);
};
auto initial_guess_generator =
[&initial_guesses](
int64_t, int64_t i) -> std::optional<Eigen::VectorXd> {
if (initial_guesses.has_value()) {
return initial_guesses->at(i);
} else {
return std::nullopt;
}
}
std::vector<const SolverOptions*> solver_options_ptrs;
if (solver_options.has_value()) {
solver_options_ptrs.reserve(solver_options->size());
for (const auto& option : *solver_options) {
solver_options_ptrs.push_back(option ? *option : nullptr);
};
auto solver_options_generator =
[&solver_options](
int64_t, int64_t i) -> std::optional<SolverOptions> {
if (solver_options.has_value() &&
*(solver_options->at(i)) != nullptr) {
return **(solver_options->at(i));
} else {
return std::nullopt;
}
}
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) : nullptr, parallelism,
dynamic_schedule);
};
auto solver_ids_generator =
[&solver_ids](int64_t, int64_t i) -> std::optional<SolverId> {
if (solver_ids.has_value()) {
return solver_ids->at(i);
} else {
return std::nullopt;
}
};
return solvers::SolveInParallel(progs_generator, 0, ssize(progs),
initial_guess_generator, solver_options_generator,
solver_ids_generator, parallelism, dynamic_schedule);
},
py::arg("progs"), py::arg("initial_guesses") = std::nullopt,
py::arg("solver_options") = std::nullopt,
Expand All @@ -1683,16 +1699,37 @@ void BindFreeFunctions(py::module m) {
const std::optional<SolverId>& solver_id,
const Parallelism& parallelism, bool dynamic_schedule) {
// TODO(Alexandre.Amice) change to generator form.
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 ? &(*guess) : nullptr);
// 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 ? &(*guess) :
// nullptr);
// }
// }
// return solvers::SolveInParallel(progs,
// initial_guesses.has_value() ?
// &initial_guesses_ptrs : nullptr, solver_options,
// solver_id, parallelism, dynamic_schedule);

const std::function<const MathematicalProgram*(int64_t, int64_t)>
progs_generator =
[&progs](int64_t, int64_t i) -> const MathematicalProgram* {
return progs.at(i);
};
auto initial_guess_generator =
[&initial_guesses](
int64_t, int64_t i) -> std::optional<Eigen::VectorXd> {
if (initial_guesses.has_value()) {
return initial_guesses->at(i);
} else {
return std::nullopt;
}
}
return solvers::SolveInParallel(progs,
initial_guesses.has_value() ? &initial_guesses_ptrs : nullptr,
solver_options, solver_id, parallelism, dynamic_schedule);
};
return solvers::SolveInParallel(progs_generator, 0, ssize(progs),
initial_guess_generator, solver_options, solver_id, parallelism,
dynamic_schedule);
},
py::arg("progs"), py::arg("initial_guesses") = std::nullopt,
py::arg("solver_options") = std::nullopt,
Expand All @@ -1702,6 +1739,9 @@ void BindFreeFunctions(py::module m) {
py::call_guard<py::gil_scoped_release>(),
doc.SolveInParallel
.doc_6args_conststdvector_conststdvector_constdrakesolversSolverOptions_conststdoptional_drakeParallelism_bool);
// N.B. Do NOT bind the generator forms of SolveInParallel directly. Callbacks
// to Python state from parallel C++ is difficult and likely to lead to
// deadlocks.
}

} // namespace
Expand Down
149 changes: 42 additions & 107 deletions geometry/optimization/convex_set.cc
Original file line number Diff line number Diff line change
Expand Up @@ -122,132 +122,67 @@ bool IsBoundedParallel(const ConvexSet& s, Parallelism parallelism) {
ConstructEmptyBoundednessProgram(&(progs[i]), s);
}

// Pre-allocate empty MathematicalProgramResults for each thread.
std::vector<MathematicalProgramResult> results(parallelism.num_threads());

// Pre-allocate solver instances.
const solvers::SolverId solver_id = solvers::ChooseBestSolver(progs[0]);
std::vector<std::unique_ptr<solvers::SolverInterface>> solver_interfaces(
parallelism.num_threads());

// Pre-allocate the solver options.
std::vector<solvers::SolverOptions> options(parallelism.num_threads());

for (int i = 0; i < parallelism.num_threads(); ++i) {
solver_interfaces[i] = solvers::MakeSolver(solver_id);
options[i].SetOption(solvers::CommonSolverOption::kMaxThreads, 1);
}

std::atomic<bool> certified_unbounded = false;
auto index_to_dimension = [&](const int64_t i) -> int {
return i / 2;
};

// This worker lambda maps the index i to a dimension and direction to check
// boundedness. If unboundedness has already been verified, it just exits
// early. If unboundedness is verified in this iteration, certified_unbounded
// will be updated to reflect that fact. For a given index i, the dimension is
// int(i / 2), and if i % 2 == 0, then we maximize, otherwise, we minimize.
auto solve_ith = [&](const int thread_num, const int64_t i) {
// boundedness. If unboundedness has already been verified, it returns
// nullptr. If unboundedness is verified in this iteration,
// certified_unbounded will be updated to reflect that fact. For a given index
// i, the dimension is int(i / 2), and if i % 2 == 0, then we maximize,
// otherwise, we minimize.
auto make_ith = [&](int64_t thread_num, int64_t i) -> MathematicalProgram* {
if (certified_unbounded.load()) {
return;
// Exit early if unboundedness has already been verified.
return nullptr;
}

const int dimension = i / 2;
const int dimension = index_to_dimension(i);
bool maximize = i % 2 == 0;

// Update the objective vector. By construction, each MathematicalProgram
// has one cost (the linear cost).
progs[thread_num].linear_costs()[0].evaluator()->update_coefficient_entry(
dimension, maximize ? -1 : 1);
DRAKE_ASSERT(progs[thread_num].IsThreadSafe());
solver_interfaces[thread_num]->Solve(progs[thread_num], std::nullopt,
options[thread_num],
&(results[thread_num]));
if (ProgramResultImpliesUnbounded(results[thread_num])) {
certified_unbounded.store(true);
}

// Reset the objective vector.
progs[thread_num].linear_costs()[0].evaluator()->update_coefficient_entry(
dimension, 0);
return &(progs[thread_num]);
};

std::function<void(MathematicalProgram**, const MathematicalProgramResult&,
int64_t, int64_t)>
teardown_ith = [&](MathematicalProgram**,
const MathematicalProgramResult& result_i,
int64_t thread_num, int64_t i) {
const int dimension = index_to_dimension(i);
progs[thread_num]
.linear_costs()[0]
.evaluator()
->update_coefficient_entry(dimension, 0);
if (ProgramResultImpliesUnbounded(result_i)) {
certified_unbounded.store(true);
}
};

// We run the loop from 0 to 2 * s.ambient_dimension(), since two programs are
// solved for each dimension (maximizing and minimizing). All programs are the
// same size, so static scheduling is appropriate.
StaticParallelForIndexLoop(DegreeOfParallelism(parallelism.num_threads()), 0,
2 * s.ambient_dimension(), solve_ith,
ParallelForBackend::BEST_AVAILABLE);

std::vector<MathematicalProgramResult> results =
SolveInParallel<MathematicalProgram*>(
make_ith, // prog_generator
0, // range_start
2 * s.ambient_dimension(), // range_end
[](int64_t, int64_t) -> std::optional<Eigen::VectorXd> {
return std::nullopt;
}, // initial_guesses_generator
nullptr, // solver_options
std::nullopt, // solver_id
Parallelism::Max(), // parallelism
false, // dynamic_schedule
&teardown_ith // prog_teardown
); // NOLINT
return !certified_unbounded.load();
}
//bool IsBoundedParallel2(const ConvexSet& s, Parallelism parallelism) {
// // Pre-allocate programs (which will be updated and solved within the parallel
// // loop).
// std::vector<MathematicalProgram> progs(parallelism.num_threads());
// for (int i = 0; i < ssize(progs); ++i) {
// ConstructEmptyBoundednessProgram(&(progs[i]), s);
// }
//
// std::atomic<bool> certified_unbounded = false;
// auto index_to_dimension = [&](const int64_t i) -> int {
// return i / 2;
// };
//
// // This worker lambda maps the index i to a dimension and direction to check
// // boundedness. If unboundedness has already been verified, it returns
// // nullptr. If unboundedness is verified in this iteration,
// // certified_unbounded will be updated to reflect that fact. For a given index
// // i, the dimension is int(i / 2), and if i % 2 == 0, then we maximize,
// // otherwise, we minimize.
// auto make_ith = [&](int64_t thread_num, int64_t i) -> MathematicalProgram* {
// if (certified_unbounded.load()) {
// // Exit early if unboundedness has already been verified.
// return nullptr;
// }
//
// const int dimension = index_to_dimension(i);
// bool maximize = i % 2 == 0;
//
// // Update the objective vector. By construction, each MathematicalProgram
// // has one cost (the linear cost).
// progs[thread_num].linear_costs()[0].evaluator()->update_coefficient_entry(
// dimension, maximize ? -1 : 1);
// return &(progs[thread_num]);
// };
//
// std::function<void(MathematicalProgram**, const MathematicalProgramResult&,
// int64_t, int64_t)>
// teardown_ith = [&](MathematicalProgram**,
// const MathematicalProgramResult& result_i,
// int64_t thread_num, int64_t i) {
// const int dimension = index_to_dimension(i);
// progs[thread_num]
// .linear_costs()[0]
// .evaluator()
// ->update_coefficient_entry(dimension, 0);
// if (ProgramResultImpliesUnbounded(result_i)) {
// certified_unbounded.store(true);
// }
// };
//
// // We run the loop from 0 to 2 * s.ambient_dimension(), since two programs are
// // solved for each dimension (maximizing and minimizing). All programs are the
// // same size, so static scheduling is appropriate.
// std::vector<MathematicalProgramResult> results =
// SolveInParallel<MathematicalProgram*>(
// make_ith, // prog_generator
// 0, // range_start
// 2 * s.ambient_dimension(), // range_end
// [](int64_t, int64_t) -> std::optional<Eigen::VectorXd> {
// return std::nullopt;
// }, // initial_guesses_generator
// nullptr, // solver_options
// std::nullopt, // solver_id
// Parallelism::Max(), // parallelism
// false, // dynamic_schedule
// &teardown_ith // prog_teardown
// ); // NOLINT
// return !certified_unbounded.load();
//}
} // namespace

bool ConvexSet::GenericDoIsBounded(Parallelism parallelism) const {
Expand Down
6 changes: 1 addition & 5 deletions geometry/optimization/graph_of_convex_sets.cc
Original file line number Diff line number Diff line change
Expand Up @@ -826,13 +826,9 @@ std::set<EdgeId> GraphOfConvexSets::PreprocessShortestPath(
// remainder of the function.
int nE = edge_id_list.size();

// TODO(cohnt): Rewrite with a parallel for loop where each thread creates and
// solves the preprocessing program, to avoid having to use batching together
// with SolveInParallel.

// TODO(cohnt): Rewrite using generators to avoid batching.
// Given an edge (u,v) check if a path from source to u and another from v to
// target exist without sharing edges.

int preprocessing_parallel_batch_size = 1000;
int num_batches = 1 + (nE / preprocessing_parallel_batch_size);
for (int batch_idx = 0; batch_idx < num_batches; ++batch_idx) {
Expand Down
18 changes: 13 additions & 5 deletions solvers/solve.cc
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include "drake/solvers/choose_best_solver.h"
#include "drake/solvers/ipopt_solver.h"
#include "drake/solvers/solver_interface.h"

namespace drake {
namespace solvers {
using common_robotics_utilities::parallelism::DegreeOfParallelism;
Expand Down Expand Up @@ -187,12 +188,11 @@ std::vector<MathematicalProgramResult> SolveInParallel(
bool dynamic_schedule,
const std::function<void(T*, const MathematicalProgramResult&, int64_t,
int64_t)>* prog_teardown) {
const std::optional<SolverOptions> returned_options =
solver_options == nullptr ? std::nullopt
: std::optional<SolverOptions>{*solver_options};
auto solver_options_generator =
[&solver_options](int64_t, int64_t) -> std::optional<SolverOptions> {
static const std::optional<SolverOptions> returned_options =
solver_options == nullptr
? std::nullopt
: std::optional<SolverOptions>{*solver_options};
[&returned_options](int64_t, int64_t) -> std::optional<SolverOptions> {
return returned_options;
};
auto solver_id_generator = [&solver_id](int64_t,
Expand Down Expand Up @@ -289,6 +289,14 @@ std::vector<MathematicalProgramResult> SolveInParallel(
const SolverOptions* solver_options,
const std::optional<SolverId>& solver_id, Parallelism parallelism,
bool dynamic_schedule) {
DRAKE_THROW_UNLESS(
std::all_of(progs.begin(), progs.end(), [](const auto& prog) {
return prog != nullptr;
}));
if (initial_guesses != nullptr) {
DRAKE_THROW_UNLESS(initial_guesses->size() == progs.size());
}

auto prog_generator =
[&progs](const int, const int64_t i) -> const MathematicalProgram* {
return GetRawPointer(progs[i]);
Expand Down

0 comments on commit 100e435

Please sign in to comment.