Skip to content

Commit ef258b4

Browse files
AlexandreAmiceAlexandreAmice
committed
first draft
1 parent c1856c2 commit ef258b4

File tree

2 files changed

+159
-49
lines changed

2 files changed

+159
-49
lines changed

‎solvers/solve.cc

Lines changed: 109 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -61,37 +61,125 @@ std::vector<MathematicalProgramResult> SolveInParallel(
6161
DRAKE_THROW_UNLESS(solver_ids->size() == progs.size());
6262
}
6363

64+
// TODO(Alexandre.Amice) is there a way around this clone?
65+
auto prog_generator = [&progs](const int thread_num, const int64_t i) {
66+
unused(thread_num);
67+
return progs[i]->Clone();
68+
};
69+
70+
auto initial_guess_generator =
71+
[&initial_guesses](int64_t thread_num,
72+
int64_t i) -> std::optional<Eigen::VectorXd> {
73+
unused(thread_num);
74+
if (initial_guesses != nullptr && initial_guesses->at(i) != nullptr) {
75+
return *(initial_guesses->at(i));
76+
} else {
77+
return std::nullopt;
78+
}
79+
};
80+
81+
auto solver_options_generator =
82+
[&solver_options](int64_t thread_num,
83+
int64_t i) -> std::optional<SolverOptions> {
84+
unused(thread_num);
85+
if (solver_options != nullptr) {
86+
return *(solver_options->at(i));
87+
} else {
88+
return std::nullopt;
89+
}
90+
};
91+
92+
auto solver_ids_generator =
93+
[&solver_ids](int64_t thread_num, int64_t i) -> std::optional<SolverId> {
94+
unused(thread_num);
95+
if (solver_ids != nullptr && solver_ids->at(i).has_value()) {
96+
return solver_ids->at(i);
97+
} else {
98+
return std::nullopt;
99+
}
100+
};
101+
102+
return SolveInParallel(
103+
prog_generator, initial_guess_generator, solver_options_generator,
104+
solver_ids_generator, static_cast<int64_t>(0),
105+
static_cast<int64_t>(progs.size()), parallelism, dynamic_schedule);
106+
}
107+
108+
std::vector<MathematicalProgramResult> SolveInParallel(
109+
const std::vector<const MathematicalProgram*>& progs,
110+
const std::vector<const Eigen::VectorXd*>* initial_guesses,
111+
const SolverOptions* solver_options,
112+
const std::optional<SolverId>& solver_id, const Parallelism parallelism,
113+
const bool dynamic_schedule) {
114+
// Broadcast the option and id arguments into vectors (if given).
115+
std::optional<std::vector<const SolverOptions*>> broadcast_options;
116+
std::optional<std::vector<std::optional<SolverId>>> broadcast_ids;
117+
if (solver_options != nullptr) {
118+
broadcast_options.emplace(progs.size(), solver_options);
119+
}
120+
if (solver_id.has_value()) {
121+
broadcast_ids.emplace(progs.size(), solver_id);
122+
}
123+
// Delegate to the primary overload.
124+
return SolveInParallel(
125+
progs, initial_guesses,
126+
broadcast_options.has_value() ? &(*broadcast_options) : nullptr,
127+
broadcast_ids.has_value() ? &(*broadcast_ids) : nullptr, // BR
128+
parallelism, dynamic_schedule);
129+
}
130+
131+
std::vector<MathematicalProgramResult> SolveInParallel(
132+
const std::function<std::unique_ptr<MathematicalProgram>(int64_t, int64_t)>&
133+
prog_generator,
134+
const std::function<std::optional<Eigen::VectorXd>(int64_t, int64_t)>&
135+
initial_guesses_generator,
136+
const std::function<std::optional<SolverOptions>(int64_t, int64_t)>&
137+
solver_options_generator,
138+
const std::function<std::optional<SolverId>(int64_t, int64_t)>&
139+
solver_ids_generator,
140+
const int64_t range_start, const int64_t range_end, Parallelism parallelism,
141+
bool dynamic_schedule) {
64142
// Pre-allocate the results (i.e., don't use push_back) so that we can safely
65143
// write to the vector on multiple threads concurrently.
66-
std::vector<MathematicalProgramResult> results{progs.size()};
144+
std::vector<MathematicalProgramResult> results{
145+
static_cast<std::size_t>(range_end)};
67146

68147
// Track which results are set during the par-for loop. Some programs are not
69148
// thread safe so will be skipped inside the loop, and we'll need to circle
70149
// back and solve them serially later. (N.B. we cannot use vector<bool> here
71150
// because it's not thread safe.)
72-
std::vector<uint8_t> result_is_populated(progs.size(), 0);
151+
std::vector<uint8_t> result_is_populated(range_end, 0);
73152

74153
// As unique types of solvers are encountered by the threads, we will cache
75154
// the solvers into this data structure so that we don't have to create and
76155
// destroy them at every call to Solve.
77156
std::vector<std::unordered_map<SolverId, std::unique_ptr<SolverInterface>>>
78157
solvers(parallelism.num_threads());
79158

159+
// This caches programs which are skipped in the parallel for loop because
160+
// they are not threadsafe. This avoids having to create the program twice.
161+
std::vector<std::unique_ptr<MathematicalProgram>> cached_progs(
162+
static_cast<std::size_t>(range_end));
163+
80164
// The worker lambda behaves slightly differently depending on whether we are
81165
// in the par-for loop or in the single-threaded cleanup pass later on.
82166
// This is the worker callback for the i'th program.
83167
auto solve_ith = [&](const bool in_parallel, const int thread_num,
84168
const int64_t i) {
85-
// If this program is not thread safe, then skip it and save it for later.
86-
if (in_parallel && !progs[i]->IsThreadSafe()) {
169+
std::unique_ptr<MathematicalProgram> prog =
170+
cached_progs[i] == nullptr ? prog_generator(thread_num, i)
171+
: std::move(cached_progs[i]);
172+
// If this program is not thread safe, then skip it and save it for
173+
// later.
174+
if (in_parallel && !prog->IsThreadSafe()) {
87175
return;
88176
}
89177

90178
// Access (or choose) the required solver.
91-
const SolverId solver_id =
92-
((solver_ids != nullptr) && (*solver_ids)[i].has_value())
93-
? *((*solver_ids)[i])
94-
: ChooseBestSolver(*progs[i]);
179+
std::optional<SolverId> solver_id = solver_ids_generator(thread_num, i);
180+
if (!solver_id.has_value()) {
181+
solver_id = ChooseBestSolver(*prog);
182+
}
95183

96184
// IPOPT's MUMPs linear solver is not thread-safe. Therefore, if we are
97185
// operating in parallel we need to skip this program.
@@ -104,12 +192,12 @@ std::vector<MathematicalProgramResult> SolveInParallel(
104192
}
105193

106194
// Find (or create) the specified solver.
107-
auto solver_iter = solvers[thread_num].find(solver_id);
195+
auto solver_iter = solvers[thread_num].find(solver_id.value());
108196
if (solver_iter == solvers[thread_num].end()) {
109197
// If this thread has not solved this type of program yet, save the solver
110198
// for use later.
111-
solver_iter = solvers[thread_num].emplace_hint(solver_iter, solver_id,
112-
MakeSolver(solver_id));
199+
solver_iter = solvers[thread_num].emplace_hint(
200+
solver_iter, solver_id.value(), MakeSolver(solver_id.value()));
113201
}
114202
const SolverInterface& solver = *(solver_iter->second);
115203

@@ -118,12 +206,9 @@ std::vector<MathematicalProgramResult> SolveInParallel(
118206
// (during the serial cleanup pass) it is allowed `parallelism`. Note that
119207
// if the user set a solver-specific options to use more threads, that could
120208
// cause it to exceed its limit.
121-
std::optional<SolverOptions> new_options;
122-
if ((solver_options != nullptr) && ((*solver_options)[i] != nullptr)) {
123-
// TODO(jwnimmer-tri) The SolverBase should offer an overload to take a
124-
// list of options, so that we can avoid this copying.
125-
new_options.emplace(*((*solver_options)[i]));
126-
} else {
209+
std::optional<SolverOptions> new_options =
210+
solver_options_generator(thread_num, i);
211+
if (!new_options.has_value()) {
127212
new_options.emplace();
128213
}
129214
new_options->SetOption(CommonSolverOption::kMaxThreads,
@@ -132,13 +217,11 @@ std::vector<MathematicalProgramResult> SolveInParallel(
132217
// Convert the initial guess from nullable to optional.
133218
// TODO(jwnimmer-tri) The SolverBase should offer a nullable overload so
134219
// that we can avoid this copying.
135-
std::optional<Eigen::VectorXd> initial_guess;
136-
if ((initial_guesses != nullptr) && ((*initial_guesses)[i] != nullptr)) {
137-
initial_guess = *((*initial_guesses)[i]);
138-
}
220+
std::optional<Eigen::VectorXd> initial_guess =
221+
initial_guesses_generator(thread_num, i);
139222

140223
// Solve the program.
141-
solver.Solve(*(progs[i]), initial_guess, new_options, &(results[i]));
224+
solver.Solve(*prog, initial_guess, new_options, &(results[i]));
142225
result_is_populated[i] = true;
143226
};
144227
const auto solve_ith_parallel = [&](const int thread_num, const int64_t i) {
@@ -154,11 +237,11 @@ std::vector<MathematicalProgramResult> SolveInParallel(
154237
if (parallelism.num_threads() > 1) {
155238
if (dynamic_schedule) {
156239
DynamicParallelForIndexLoop(
157-
DegreeOfParallelism(parallelism.num_threads()), 0, ssize(progs),
158-
solve_ith_parallel, ParallelForBackend::BEST_AVAILABLE);
240+
DegreeOfParallelism(parallelism.num_threads()), range_start,
241+
range_end, solve_ith_parallel, ParallelForBackend::BEST_AVAILABLE);
159242
} else {
160243
StaticParallelForIndexLoop(DegreeOfParallelism(parallelism.num_threads()),
161-
0, ssize(progs), solve_ith_parallel,
244+
range_start, range_end, solve_ith_parallel,
162245
ParallelForBackend::BEST_AVAILABLE);
163246
}
164247
}
@@ -169,7 +252,7 @@ std::vector<MathematicalProgramResult> SolveInParallel(
169252
solvers.resize(1);
170253

171254
// Finish solving the programs that couldn't be solved in parallel.
172-
for (size_t i = 0; i < progs.size(); ++i) {
255+
for (int64_t i = range_start; i < range_end; ++i) {
173256
if (!result_is_populated[i]) {
174257
solve_ith_serial(i);
175258
}
@@ -178,28 +261,5 @@ std::vector<MathematicalProgramResult> SolveInParallel(
178261
return results;
179262
}
180263

181-
std::vector<MathematicalProgramResult> SolveInParallel(
182-
const std::vector<const MathematicalProgram*>& progs,
183-
const std::vector<const Eigen::VectorXd*>* initial_guesses,
184-
const SolverOptions* solver_options,
185-
const std::optional<SolverId>& solver_id, const Parallelism parallelism,
186-
const bool dynamic_schedule) {
187-
// Broadcast the option and id arguments into vectors (if given).
188-
std::optional<std::vector<const SolverOptions*>> broadcast_options;
189-
std::optional<std::vector<std::optional<SolverId>>> broadcast_ids;
190-
if (solver_options != nullptr) {
191-
broadcast_options.emplace(progs.size(), solver_options);
192-
}
193-
if (solver_id.has_value()) {
194-
broadcast_ids.emplace(progs.size(), solver_id);
195-
}
196-
// Delegate to the primary overload.
197-
return SolveInParallel(
198-
progs, initial_guesses,
199-
broadcast_options.has_value() ? &(*broadcast_options) : nullptr,
200-
broadcast_ids.has_value() ? &(*broadcast_ids) : nullptr, // BR
201-
parallelism, dynamic_schedule);
202-
}
203-
204264
} // namespace solvers
205265
} // namespace drake

‎solvers/solve.h

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#pragma once
22

3+
#include <functional>
4+
#include <memory>
35
#include <optional>
46
#include <string>
57
#include <vector>
@@ -104,5 +106,53 @@ std::vector<MathematicalProgramResult> SolveInParallel(
104106
Parallelism parallelism = Parallelism::Max(),
105107
bool dynamic_schedule = false);
106108

109+
/**
110+
* Provides the same functionality as SolveInParallel, but the programs, initial
111+
* guesses, solver options and solver ids are created using a generator.
112+
*
113+
* The input to the generator is the current thread number and an index i and
114+
* the output is the ith program, guess, option and solver id respectively. The
115+
* index i is iterated from range_start to range_end.
116+
*
117+
* The output of prog_generator cannot be a nullptr.
118+
*
119+
* The output of the initial_guesses_generator, solver_options_generator and
120+
* solver_ids can be nullptr or std::optional.
121+
*
122+
* @note Please ensure that all generators are thread safe.
123+
* @return
124+
*/
125+
std::vector<MathematicalProgramResult> SolveInParallel(
126+
const std::function<std::unique_ptr<MathematicalProgram>(int64_t, int64_t)>&
127+
prog_generator,
128+
const std::function<std::optional<Eigen::VectorXd>(int64_t, int64_t)>&
129+
initial_guesses_generator,
130+
const std::function<std::optional<SolverOptions>(int64_t, int64_t)>&
131+
solver_options_generator,
132+
const std::function<std::optional<SolverId>(int64_t, int64_t)>&
133+
solver_ids_generator,
134+
const int64_t range_start, const int64_t range_end,
135+
Parallelism parallelism = Parallelism::Max(),
136+
bool dynamic_schedule = false);
137+
138+
/**
139+
* Provides the same functionality as SolveInParallel, but the programs and
140+
* initial guesses are created using a generator. The same solver option, and
141+
* solver ids are used when solving all the programs.
142+
*
143+
* The input to the generator is an integer i and the output is the ith program.
144+
*
145+
* The output of prog_generator cannot be a nullptr.
146+
* @return
147+
*/
148+
std::vector<MathematicalProgramResult> SolveInParallel(
149+
const std::function<std::unique_ptr<MathematicalProgram>(int64_t, int64_t)>&
150+
prog_generator,
151+
const const std::function<std::optional<Eigen::VectorXd>(int64_t, int64_t)>&
152+
initial_guesses_generatorconst SolverOptions* solver_options = nullptr,
153+
const std::optional<SolverId>& solver_id = std::nullopt,
154+
Parallelism parallelism = Parallelism::Max(),
155+
bool dynamic_schedule = false);
156+
107157
} // namespace solvers
108158
} // namespace drake

0 commit comments

Comments
 (0)