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);
src/engine/engine_util_solve.c
+44 -25
View File
@@ -415,43 +415,62 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// sparse reverse-order Cholesky rank-one update: L'*L +/- x*x'; return rank
// x is sparse, change in sparsity pattern of mat is not allowed
int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
const int* rownnz, const int* rowadr, const int* colind,
int x_nnz, int* x_ind,
int mju_cholUpdateSparse(mjtNum* restrict mat, const mjtNum* restrict x, int n, int flg_plus,
const int* restrict rownnz, const int* restrict rowadr,
const int* restrict colind, int x_nnz, const int* restrict x_ind,
mjData* d) {
// early return if x is empty
if (x_nnz == 0) {
return n;
}
// get starting row: last non-zero entry in x
int start = x_ind[x_nnz - 1];
// allocate dense accumulator for x
mj_markStack(d);
int* buf_ind = mjSTACKALLOC(d, n, int);
mjtNum* sparse_buf = mjSTACKALLOC(d, n, mjtNum);
mjtNum* restrict dense = mjSTACKALLOC(d, start + 1, mjtNum);
mju_zero(dense, start + 1);
// backpass over rows corresponding to non-zero x(r)
int rank = n, i = x_nnz - 1;
while (i >= 0) {
// get rownnz and rowadr for this row
int nnz = rownnz[x_ind[i]], adr = rowadr[x_ind[i]];
// scatter x into dense
mju_scatter(dense, x, x_ind, x_nnz);
// compute quantities
mjtNum tmp = mat[adr+nnz-1]*mat[adr+nnz-1] + (flg_plus ? x[i]*x[i] : -x[i]*x[i]);
// backpass over rows from start down to 0
int rank = n;
for (int row = start; row >= 0; row--) {
// skip if zero
if (dense[row] == 0) continue;
// get rownnz (excluding diagonal), rowadr
int nnz = rownnz[row] - 1;
int adr = rowadr[row];
// update diagonal, handle rank-deficient case
mjtNum diag = mat[adr + nnz];
mjtNum x_row = dense[row];
mjtNum tmp = diag*diag + (flg_plus ? x_row*x_row : -x_row*x_row);
if (tmp < mjMINVAL) {
tmp = mjMINVAL;
rank--;
}
mjtNum r = mju_sqrt(tmp);
mjtNum c = r / mat[adr+nnz-1];
mjtNum s = x[i] / mat[adr+nnz-1];
mat[adr + nnz] = r;
// update diagonal
mat[adr+nnz-1] = r;
// compute Givens rotation parameters https://en.wikipedia.org/wiki/Givens_rotation
mjtNum c = diag / r;
mjtNum s = -x_row / r;
mjtNum s_signed = flg_plus ? -s : s;
// update row: mat(r,1:r-1) = (mat(r,1:r-1) + s*x(1:r-1)) / c
mju_combineSparseInc(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
nnz-1, i, colind + adr, x_ind);
// update row
for (int i = 0; i < nnz; i++) {
int j = colind[adr + i];
mjtNum dense_j = dense[j];
mjtNum mat_val = mat[adr + i];
// update x: x(1:r-1) = c*x(1:r-1) - s*mat(r,1:r-1)
int new_x_nnz = mju_combineSparse(x, mat+adr, c, -s, i, nnz-1, x_ind,
colind+adr, sparse_buf, buf_ind);
// update i, correct for changing x
i = i - 1 + (new_x_nnz - i);
// update mat and dense using the Givens rotation
mat[adr + i] = c*mat_val + s_signed*dense_j;
dense[j] = s*mat_val + c*dense_j;
}
}
mj_freeStack(d);
src/engine/engine_util_solve.h
+2 -2
View File
@@ -62,9 +62,9 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank
// x is sparse, change in sparsity pattern of mat is not allowed
MJAPI int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
MJAPI int mju_cholUpdateSparse(mjtNum* mat, const mjtNum* x, int n, int flg_plus,
const int* rownnz, const int* rowadr, const int* colind,
int x_nnz, int* x_ind, mjData* d);
int x_nnz, const int* x_ind, mjData* d);
// band-dense Cholesky decomposition
// returns minimum value in the factorized diagonal, or 0 if rank-deficient