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

Add per-island PGS solver dispatch.

Browse Source 
Total testspeed runtime for `2humanoids100.xml` reduced by 27.6% (63.4 -> 49.5s) due to early termination on small islands

PiperOrigin-RevId: 906910915
Change-Id: If55ad468c3680ef44eda7000455a77f8003b3122
This commit is contained in:
Yuval Tassa
2026-04-28 05:13:22 -07:00
committed by Copybara-Service
parent 00fff8c780
commit 647af382c1
7 changed files with 146 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_forward.c
+57 -36
View File
@@ -953,7 +953,7 @@ static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
// compute efc_b, efc_force, qfrc_constraint; update qacc
void mj_fwdConstraint(const mjModel* m, mjData* d) {
TM_START;
int nv = m->nv, nefc = d->nefc, nisland = d->nisland;
int nv = m->nv, nefc = d->nefc, nisland = d->nisland, nidof;
// always clear qfrc_constraint
mju_zero(d->qfrc_constraint, nv);
@@ -970,50 +970,69 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
mj_mulJacVec(m, d, d->efc_b, d->qacc_smooth);
mju_subFrom(d->efc_b, d->efc_aref, nefc);
// check for invalid solver type
if (m->opt.solver != mjSOL_PGS && m->opt.solver != mjSOL_CG && m->opt.solver != mjSOL_NEWTON) {
mjERROR("unknown solver type %d", m->opt.solver);
}
// warmstart solver
warmstart(m, d);
mju_zeroInt(d->solver_niter, mjNISLAND);
// check if islands are supported
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) &&
nisland > 0 &&
m->opt.noslip_iterations == 0 &&
(m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON);
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
// run solver over constraint islands
if (islands_supported) {
int nidof = d->nidof;
// copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof);
mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof);
mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof);
mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc);
mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc);
// solve per island, with or without threads
if (!d->threadpool) {
// no threadpool, loop over islands
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS:
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_solPGS_island(m, d, island, m->opt.iterations);
}
} else {
// have threadpool, solve using threads
solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
break;
case mjSOL_CG:
case mjSOL_NEWTON:
// copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
nidof = d->nidof;
mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof);
mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof);
mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof);
mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc);
mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc);
// solve per island, with or without threads
if (!d->threadpool) {
// no threadpool, loop over islands
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);
}
}
} else {
// have threadpool, solve using threads
solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
}
// copy back solver outputs (scatter dofs since ni <= nv)
mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof);
mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
break;
}
// copy back solver outputs (scatter dofs since ni <= nv)
mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof);
mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
// run noslip solver per island if enabled
if (m->opt.noslip_iterations > 0) {
for (int island=0; island < nisland; island++) {
mj_solNoSlip_island(m, d, island, m->opt.noslip_iterations);
}
}
}
// run solver over all constraints
// run solver over all constraints (monolithic)
else {
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS: // PGS
@@ -1027,15 +1046,17 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
case mjSOL_NEWTON: // Newton
mj_solNewton(m, d, m->opt.iterations);
break;
}
default:
mjERROR("unknown solver type %d", m->opt.solver);
// run noslip solver if enabled
if (m->opt.noslip_iterations > 0) {
mj_solNoSlip(m, d, m->opt.noslip_iterations);
}
}
// run noslip solver if enabled
if (m->opt.noslip_iterations > 0) {
mj_solNoSlip(m, d, m->opt.noslip_iterations);
// dual solvers: map efc_force to joint space (always monolithic)
if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
mj_dualFinish(m, d);
}
TM_END(mjTIMER_CONSTRAINT);
src/engine/engine_solver.c
+37 -14
View File
@@ -66,7 +66,6 @@ static void saveStats(const mjModel* m, mjData* d, int island, int iter,
// finalize dual solver: map to joint space
// TODO: b/295296178 - add island support to Dual solvers
static void dualFinish(const mjModel* m, mjData* d) {
// map constraint force to joint space
mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force);
@@ -77,10 +76,15 @@ static void dualFinish(const mjModel* m, mjData* d) {
}
// PGS: map efc_force to joint space
void mj_dualFinish(const mjModel* m, mjData* d) {
dualFinish(m, d);
}
// compute 1/diag(AR)
// res[c] = 1 / AR[efclist[c], efclist[c]] for c = 0..nefc-1
// efclist is NULL for monolithic (sequential) iteration
// TODO: b/295296178 - add island support to Dual solvers
static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res,
int nefc, const int* efclist, int flg_subR) {
const mjtNum *AR = d->efc_AR;
@@ -118,7 +122,6 @@ static void ARdiaginv(const mjModel* m, const mjData* d, mjtNum* res,
// extract diagonal block from AR, clamp diag to 1e-10 if flg_subR
// TODO: b/295296178 - add island support to Dual solvers
static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
int start, int n, int flg_subR) {
int nefc = d->nefc;
@@ -178,7 +181,6 @@ static void extractBlock(const mjModel* m, const mjData* d, mjtNum* Ac,
// compute residual for one block
// TODO: b/295296178 - add island support to Dual solvers
static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int dim, int flg_subR) {
int nefc = d->nefc;
@@ -208,7 +210,6 @@ static void residual(const mjModel* m, const mjData* d, mjtNum* res, int i, int
// compute cost change
// TODO: b/295296178 - add island support to Dual solvers
static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
const mjtNum* res, int dim) {
mjtNum change;
@@ -235,7 +236,6 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
// set efc_state to dual constraint state; return nactive
// iterates over efclist (or sequentially if NULL), classifies by ne/nf ranges
// TODO: b/295296178 - add island support to Dual solvers
static int dualState(const mjData* d, int* state,
int ne, int nf, int nefc, const int* efclist) {
const mjtNum* force = d->efc_force;
@@ -384,7 +384,6 @@ static void solveQCQP(mjtNum* force, int i, int dim,
// island: island index for stats (use -1 for monolithic, mapped to 0)
// ne, nf, nefc: constraint type counts
// efclist: maps list position c to monolithic efc index (NULL for sequential)
// TODO: b/295296178 - add island support to Dual solvers
static void solPGS(const mjModel* m, mjData* d, int island,
int ne, int nf, int nefc,
const int* efclist, int maxiter) {
@@ -572,10 +571,20 @@ static void solPGS(const mjModel* m, mjData* d, int island,
}
// PGS entry point (monolithic)
// PGS entry point (monolithic, no dualFinish — caller handles it)
void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
solPGS(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
dualFinish(m, d);
}
// PGS entry point (one island)
void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter) {
int ne = d->island_ne[island];
int nf = d->island_nf[island];
int nefc = d->island_nefc[island];
int iefcadr = d->island_iefcadr[island];
solPGS(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
}
@@ -754,8 +763,10 @@ static void solNoSlip(const mjModel* m, mjData* d, int island,
improvement *= scale;
// save noslip stats after all the entries from regular solver
int stats_iter = iter + d->solver_niter[island_stat];
saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
if (island_stat < mjNISLAND) {
int stats_iter = iter + d->solver_niter[island_stat];
saveStats(m, d, island_stat, stats_iter, improvement, 0, 0, nactive, nchange, 0, 0);
}
// increment iteration count
iter++;
@@ -767,16 +778,28 @@ static void solNoSlip(const mjModel* m, mjData* d, int island,
}
// update solver iterations
d->solver_niter[island_stat] += iter;
if (island_stat < mjNISLAND) {
d->solver_niter[island_stat] += iter;
}
mj_freeStack(d);
}
// NoSlip entry point (monolithic)
// NoSlip entry point (monolithic, no dualFinish — caller handles it)
void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
solNoSlip(m, d, /*island=*/-1, d->ne, d->nf, d->nefc, /*efclist=*/NULL, maxiter);
dualFinish(m, d);
}
// NoSlip entry point (one island)
void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter) {
int ne = d->island_ne[island];
int nf = d->island_nf[island];
int nefc = d->island_nefc[island];
int iefcadr = d->island_iefcadr[island];
solNoSlip(m, d, island, ne, nf, nefc, d->map_iefc2efc + iefcadr, maxiter);
}
src/engine/engine_solver.h
+9
View File
@@ -35,10 +35,19 @@ void mj_solNewton(const mjModel* m, mjData* d, int maxiter);
//------------------------------ per-island solvers ------------------------------------------------
// PGS solver (one island, no dualFinish — caller handles it)
void mj_solPGS_island(const mjModel* m, mjData* d, int island, int maxiter);
// NoSlip solver (one island, no dualFinish — caller handles it)
void mj_solNoSlip_island(const mjModel* m, mjData* d, int island, int maxiter);
// CG solver
void mj_solCG_island(const mjModel* m, mjData* d, int island, int maxiter);
// Newton entry point
void mj_solNewton_island(const mjModel* m, mjData* d, int island, int maxiter);
// map efc_force to joint space (used after dual island dispatch)
void mj_dualFinish(const mjModel* m, mjData* d);
#endif // MUJOCO_SRC_ENGINE_ENGINE_SOLVER_H_