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

Refactor mju_cholUpdateSparse to use a dense accumulator, 1.5-2x function speedup.

Browse Source 
PiperOrigin-RevId: 846754379
Change-Id: Ie599658db2907a26095072d97d07b99a693050dc
This commit is contained in:
Yuval Tassa
2025-12-19 08:56:43 -08:00
committed by Copybara-Service
parent 2935af23f0
commit 849a650143
4 changed files with 236 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+6 -14
View File
@@ -1685,8 +1685,6 @@ static void HessianCone(mjData* d, mjCGContext* ctx) {
// storage for L'*J
mjtNum* LTJ = mjSTACKALLOC(d, 6*nv, mjtNum);
mjtNum* LTJ_row = mjSTACKALLOC(d, nv, mjtNum);
int* LTJ_ind = mjSTACKALLOC(d, nv, int);
// add contributions
for (int i=0; i < nefc; i++) {
@@ -1713,13 +1711,9 @@ static void HessianCone(mjData* d, mjCGContext* ctx) {
// update
for (int r=0; r < dim; r++) {
// copy data for this row
mju_copy(LTJ_row, LTJ+r*nnz, nnz);
mju_copyInt(LTJ_ind, ctx->J_colind+ctx->J_rowadr[i+r], nnz);
// update
mju_cholUpdateSparse(ctx->Lcone, LTJ_row, nv, 1,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, LTJ_ind, d);
mju_cholUpdateSparse(ctx->Lcone, LTJ+r*nnz, nv, 1,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz,
ctx->J_colind+ctx->J_rowadr[i+r], d);
}
}
@@ -1758,7 +1752,6 @@ static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate)
// local space
mjtNum* vec = mjSTACKALLOC(d, nv, mjtNum);
int* vec_ind = mjSTACKALLOC(d, nv, int);
// clear update counter
ctx->nupdate = 0;
@@ -1784,14 +1777,13 @@ static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate)
// get nnz and adr of row i
const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i];
// scale vec, copy colind
// scale vec
mju_scl(vec, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz);
mju_copyInt(vec_ind, ctx->J_colind+adr, nnz);
// sparse update or downdate
rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, vec_ind,
d);
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz,
ctx->J_colind+adr, d);
} else {
mju_scl(vec, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv);
rank = mju_cholUpdate(ctx->L, vec, nv, flag_update);