chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enable multithreading of Newton solver.

Browse Source 
PiperOrigin-RevId: 760723081
Change-Id: Ib7d3c9cec2fdf7dd0e46a570f37c0a6fc5424c7b
This commit is contained in:
Yuval Tassa
2025-05-19 12:17:22 -07:00
committed by Copybara-Service
parent 7edbdd0ad6
commit 94803d772c
1 changed files with 25 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_forward.c
+25 -23
View File
@@ -659,32 +659,35 @@ struct mjSolIslandArgs_ {
}; };
typedef struct mjSolIslandArgs_ mjSolIslandArgs; typedef struct mjSolIslandArgs_ mjSolIslandArgs;
// extract arguments, pass to solver // extract arguments, pass to CG solver
void* mj_solCG_island_wrapper(void* args) { static void* CG_wrapper(void* args) {
mjSolIslandArgs* solargs = (mjSolIslandArgs*) args; mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations); mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
return NULL; return NULL;
} }
// extract arguments, pass to Newton solver
static void* Newton_wrapper(void* args) {
mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
mj_solNewton_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
return NULL;
}
// CG solver, multi-threaded over islands // CG solver, multi-threaded over islands
void mj_solCG_island_multithreaded(const mjModel* m, mjData* d) { static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
mj_markStack(d); mj_markStack(d);
// allocate array of arguments to be passed to threads // allocate array of arguments to be passed to threads
mjSolIslandArgs* sol_cg_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs); mjSolIslandArgs* sol_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs);
mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask); mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask);
for (int island = 0; island < d->nisland; ++island) { for (int island = 0; island < d->nisland; ++island) {
sol_cg_island_args[island].m = m; sol_island_args[island].m = m;
sol_cg_island_args[island].d = d; sol_island_args[island].d = d;
sol_cg_island_args[island].island = island; sol_island_args[island].island = island;
mju_defaultTask(&tasks[island]); mju_defaultTask(&tasks[island]);
tasks[island].func = mj_solCG_island_wrapper; tasks[island].func = flg_Newton ? Newton_wrapper : CG_wrapper;
tasks[island].args = &sol_cg_island_args[island]; tasks[island].args = &sol_island_args[island];
mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]); mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]);
} }
@@ -740,23 +743,22 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc); mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc);
mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc); mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc);
// solve per island // solve per island, with or without threads
if (m->opt.solver == mjSOL_CG) { if (!d->threadpool) {
if (!d->threadpool) { // no threadpool, loop over islands
// no threadpool, loop over islands for (int island=0; island < nisland; island++) {
for (int island=0; island < nisland; island++) { if (m->opt.solver == mjSOL_NEWTON) {
mj_solNewton_island(m, d, island, m->opt.iterations);
} else {
mj_solCG_island(m, d, island, m->opt.iterations); mj_solCG_island(m, d, island, m->opt.iterations);
} }
} else {
// have threadpool, solve using threads
mj_solCG_island_multithreaded(m, d);
} }
} else { } else {
for (int island=0; island < nisland; island++) { // have threadpool, solve using threads
mj_solNewton_island(m, d, island, m->opt.iterations); solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
}
} }
// copy back solver outputs (scatter dofs since ni <= nv) // copy back solver outputs (scatter dofs since ni <= nv)
mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof); mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof);
mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof); mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);