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

Hoist solver stack allocations out of iteration loop

Browse Source 
PiperOrigin-RevId: 900085248
Change-Id: I2bde8fafc8a4801a92c39aa9a13af0ea7475843d
This commit is contained in:
Yuval Tassa
2026-04-15 04:22:16 -07:00
committed by Copybara-Service
parent f24f9ef44d
commit 8f3ed662eb
1 changed files with 14 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+14 -19
View File
@@ -813,6 +813,8 @@ typedef struct {
// Newton arrays, known-size (PrimalAllocate) // Newton arrays, known-size (PrimalAllocate)
mjtNum* D; // constraint inertia (nefc x 1) mjtNum* D; // constraint inertia (nefc x 1)
mjtNum* cholupd; // scratch for rank-1 Cholesky updates (nv x 1)
mjtNum* LTJ; // L'*J for cone Cholesky updates (6 x nv)
int* H_rowadr; // Hessian row addresses (nv x 1) int* H_rowadr; // Hessian row addresses (nv x 1)
int* H_rownnz; // Hessian row nonzeros (nv x 1) int* H_rownnz; // Hessian row nonzeros (nv x 1)
int* HT_rownnz; // Hessian transpose row nonzeros (nv x 1) int* HT_rownnz; // Hessian transpose row nonzeros (nv x 1)
@@ -988,6 +990,10 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
// Newton only, known-size arrays // Newton only, known-size arrays
if (flg_Newton) { if (flg_Newton) {
ctx->D = mjSTACKALLOC(d, nefc, mjtNum); ctx->D = mjSTACKALLOC(d, nefc, mjtNum);
ctx->cholupd = mjSTACKALLOC(d, nv, mjtNum);
if (ctx->is_elliptic) {
ctx->LTJ = mjSTACKALLOC(d, 6*nv, mjtNum);
}
// sparse Newton only // sparse Newton only
if (ctx->is_sparse) { if (ctx->is_sparse) {
@@ -1678,10 +1684,7 @@ static void HessianCone(mjData* d, mjPrimalContext* ctx) {
// start with Hcone = H // start with Hcone = H
mju_copy(ctx->Lcone, ctx->L, ctx->nL); mju_copy(ctx->Lcone, ctx->L, ctx->nL);
mj_markStack(d); mjtNum* LTJ = ctx->LTJ;
// storage for L'*J
mjtNum* LTJ = mjSTACKALLOC(d, 6*nv, mjtNum);
// add contributions // add contributions
for (int i=0; i < nefc; i++) { for (int i=0; i < nefc; i++) {
@@ -1737,18 +1740,13 @@ static void HessianCone(mjData* d, mjPrimalContext* ctx) {
i += (dim-1); i += (dim-1);
} }
} }
mj_freeStack(d);
} }
// incremental update to Hessian factor due to changes in efc_state // incremental update to Hessian factor due to changes in efc_state
static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldstate) { static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldstate) {
int rank, nv = ctx->nv, nefc = ctx->nefc; int rank, nv = ctx->nv, nefc = ctx->nefc;
mj_markStack(d); mjtNum* cholupd = ctx->cholupd;
// local space
mjtNum* vec = mjSTACKALLOC(d, nv, mjtNum);
// clear update counter // clear update counter
ctx->nupdate = 0; ctx->nupdate = 0;
@@ -1769,27 +1767,26 @@ static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldst
// perform update if flagged // perform update if flagged
if (flag_update != -1) { if (flag_update != -1) {
// update with vec = J(i,:)*sqrt(D[i])) // update with cholupd = J(i,:)*sqrt(D[i]))
if (ctx->is_sparse) { if (ctx->is_sparse) {
// get nnz and adr of row i // get nnz and adr of row i
const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i]; const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i];
// scale vec // scale cholupd
mju_scl(vec, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz); mju_scl(cholupd, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz);
// sparse update or downdate // sparse update or downdate
rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update, rank = mju_cholUpdateSparse(ctx->L, cholupd, nv, flag_update,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz,
ctx->J_colind+adr, d); ctx->J_colind+adr, d);
} else { } else {
mju_scl(vec, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv); mju_scl(cholupd, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv);
rank = mju_cholUpdate(ctx->L, vec, nv, flag_update); rank = mju_cholUpdate(ctx->L, cholupd, nv, flag_update);
} }
ctx->nupdate++; ctx->nupdate++;
// recompute H directly if accuracy lost // recompute H directly if accuracy lost
if (rank < nv) { if (rank < nv) {
mj_freeStack(d);
FactorizeHessian(d, ctx, /*flg_recompute=*/1); FactorizeHessian(d, ctx, /*flg_recompute=*/1);
// nothing else to do // nothing else to do
@@ -1802,8 +1799,6 @@ static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldst
if (ctx->ncone) { if (ctx->ncone) {
HessianCone(d, ctx); HessianCone(d, ctx);
} }
mj_freeStack(d);
} }