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

Use two-step sparse Cholesky in Newton solver.

Browse Source 
17% overall speedup for `100_humanoids.xml` as measured by `testspeed`

PiperOrigin-RevId: 846723342
Change-Id: Ibc41487bfc576c992f14853640d0b6000f5b47b2
This commit is contained in:
Yuval Tassa
2025-12-19 07:19:01 -08:00
committed by Copybara-Service
parent 45b0153067
commit 76e0a78b32
1 changed files with 31 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+31 -48
View File
@@ -815,18 +815,24 @@ struct _mjCGContext {
mjtNum* D; // constraint inertia (nefc x 1) mjtNum* D; // constraint inertia (nefc x 1)
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* H_lowernnz; // Hessian lower triangle row nonzeros (nv x 1) int* HT_rownnz; // Hessian transpose row nonzeros (nv x 1)
int* HT_rowadr; // Hessian transpose row addresses (nv x 1)
int* L_rownnz; // Hessian factor row nonzeros (nv x 1) int* L_rownnz; // Hessian factor row nonzeros (nv x 1)
int* L_rowadr; // Hessian factor row addresses (nv x 1) int* L_rowadr; // Hessian factor row addresses (nv x 1)
int* LT_rownnz; // Hessian factor transpose row nonzeros (nv x 1)
int* LT_rowadr; // Hessian factor transpose row addresses (nv x 1)
int* buf_ind; // index buffer for sparse addition (nv x 1) int* buf_ind; // index buffer for sparse addition (nv x 1)
mjtNum* buf_val; // value buffer for sparse addition (nv x 1) mjtNum* buf_val; // value buffer for sparse addition (nv x 1)
// Newton arrays, computed-size (MakeHessian) // Newton arrays, computed-size (MakeHessian)
int nH; // number of nonzeros in Hessian H int nH; // number of nonzeros in Hessian H
int* H_colind; // Hessian column indices (nH x 1) int* H_colind; // Hessian column indices (nH x 1)
int* HT_colind; // Hessian transpose column indices (nH x 1)
mjtNum* H; // Hessian (nH x 1) mjtNum* H; // Hessian (nH x 1)
int nL; // number of nonzeros in Cholesky factor L int nL; // number of nonzeros in Cholesky factor L
int* L_colind; // Cholesky factor column indices (nL x 1) int* L_colind; // Cholesky factor column indices (nL x 1)
int* LT_colind; // Cholesky factor transpose column indices (nL x 1)
int* LT_map; // CSC-to-CSR index mapping (nL x 1)
mjtNum* L; // Cholesky factor (nL x 1) mjtNum* L; // Cholesky factor (nL x 1)
mjtNum* Lcone; // Cholesky factor with cone contributions (nL x 1) mjtNum* Lcone; // Cholesky factor with cone contributions (nL x 1)
@@ -989,9 +995,12 @@ static void CGallocate(mjData* d, mjCGContext* ctx, int flg_Newton) {
if (ctx->is_sparse) { if (ctx->is_sparse) {
ctx->H_rowadr = mjSTACKALLOC(d, nv, int); ctx->H_rowadr = mjSTACKALLOC(d, nv, int);
ctx->H_rownnz = mjSTACKALLOC(d, nv, int); ctx->H_rownnz = mjSTACKALLOC(d, nv, int);
ctx->H_lowernnz = mjSTACKALLOC(d, nv, int); ctx->HT_rownnz = mjSTACKALLOC(d, nv, int);
ctx->HT_rowadr = mjSTACKALLOC(d, nv, int);
ctx->L_rownnz = mjSTACKALLOC(d, nv, int); ctx->L_rownnz = mjSTACKALLOC(d, nv, int);
ctx->L_rowadr = mjSTACKALLOC(d, nv, int); ctx->L_rowadr = mjSTACKALLOC(d, nv, int);
ctx->LT_rownnz = mjSTACKALLOC(d, nv, int);
ctx->LT_rowadr = mjSTACKALLOC(d, nv, int);
ctx->buf_val = mjSTACKALLOC(d, nv, mjtNum); ctx->buf_val = mjSTACKALLOC(d, nv, mjtNum);
ctx->buf_ind = mjSTACKALLOC(d, nv, int); ctx->buf_ind = mjSTACKALLOC(d, nv, int);
} }
@@ -1549,22 +1558,16 @@ static void MakeHessian(mjData* d, mjCGContext* ctx) {
ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
ctx->buf_val, ctx->buf_ind); ctx->buf_val, ctx->buf_ind);
// transiently compute H'; mju_cholFactorNNZ is memory-contiguous in upper triangle layout // compute H' (upper triangle, required for symbolic Cholesky)
mj_markStack(d); ctx->HT_colind = mjSTACKALLOC(d, ctx->nH, int);
int* HT_rownnz = mjSTACKALLOC(d, nv, int); mju_transposeSparse(NULL, NULL, nv, nv, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, NULL,
int* HT_rowadr = mjSTACKALLOC(d, nv, int);
int* HT_colind = mjSTACKALLOC(d, ctx->nH, int);
mju_transposeSparse(NULL, NULL, nv, nv,
HT_rownnz, HT_rowadr, HT_colind, NULL,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind); ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind);
// count total and row non-zeros of reverse-Cholesky factors L and LT // count total and row non-zeros of reverse-Cholesky factors L and LT
int* LT_rownnz_temp = mjSTACKALLOC(d, nv, int);
int* LT_rowadr_temp = mjSTACKALLOC(d, nv, int);
ctx->nL = mju_cholFactorSymbolic(NULL, ctx->L_rownnz, ctx->L_rowadr, NULL, ctx->nL = mju_cholFactorSymbolic(NULL, ctx->L_rownnz, ctx->L_rowadr, NULL,
LT_rownnz_temp, LT_rowadr_temp, NULL, ctx->LT_rownnz, ctx->LT_rowadr, NULL,
HT_rownnz, HT_rowadr, HT_colind, nv, d); ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind,
mj_freeStack(d); nv, d);
// allocate L_colind, L, Lcone // allocate L_colind, L, Lcone
ctx->L_colind = mjSTACKALLOC(d, ctx->nL, int); ctx->L_colind = mjSTACKALLOC(d, ctx->nL, int);
@@ -1573,25 +1576,15 @@ static void MakeHessian(mjData* d, mjCGContext* ctx) {
ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum); ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum);
} }
// count nonzeros in rows of H lower triangle // allocate LT (CSC representation of L)
for (int r = 0; r < nv; r++) { ctx->LT_colind = mjSTACKALLOC(d, ctx->nL, int);
const int* colind = ctx->H_colind + ctx->H_rowadr[r]; ctx->LT_map = mjSTACKALLOC(d, ctx->nL, int);
int rownnz = ctx->H_rownnz[r];
// count nonzeros up to diagonal (inclusive) for row r // symbolic Cholesky: populate L_colind and LT structures
int nnz = 1; mju_cholFactorSymbolic(ctx->L_colind, ctx->L_rownnz, ctx->L_rowadr,
while (nnz < rownnz && colind[nnz - 1] < r) { ctx->LT_colind, ctx->LT_rownnz, ctx->LT_rowadr, ctx->LT_map,
nnz++; ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind,
} nv, d);
// last row element is not the diagonal; SHOULD NOT OCCUR
if (colind[nnz - 1] != r) {
mjERROR("Newton solver Hessian has zero diagonal on row %d", r);
}
// save row nonzeros
ctx->H_lowernnz[r] = nnz;
}
} }
// dense // dense
@@ -1643,27 +1636,17 @@ static void FactorizeHessian(mjData* d, mjCGContext* ctx, int flg_recompute) {
ctx->buf_val, ctx->buf_ind); ctx->buf_val, ctx->buf_ind);
} }
// copy H lower-triangle into L, fill-in already accounted for // numeric sparse factorization: L = chol(H) using pre-computed sparsity pattern
for (int r = 0; r < nv; r++) { int rank = mju_cholFactorNumeric(
int nnz = ctx->H_lowernnz[r]; ctx->L, nv, mjMINVAL,
mju_copy(ctx->L + ctx->L_rowadr[r], ctx->H + ctx->H_rowadr[r], nnz); ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind,
mju_copyInt(ctx->L_colind + ctx->L_rowadr[r], ctx->H_colind + ctx->H_rowadr[r], nnz); ctx->LT_rownnz, ctx->LT_rowadr, ctx->LT_colind, ctx->LT_map,
ctx->L_rownnz[r] = nnz; ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, d);
}
// in-place sparse factorization: L = chol(H)
int rank = mju_cholFactorSparse(ctx->L, nv, mjMINVAL,
ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, d);
// rank-deficient; SHOULD NOT OCCUR // rank-deficient; SHOULD NOT OCCUR
if (rank != nv) { if (rank != nv) {
mjERROR("rank-deficient sparse Hessian"); mjERROR("rank-deficient sparse Hessian");
} }
// pre-counted nL does not match post-factorization nL; SHOULD NOT OCCUR
if (ctx->nL != ctx->L_rowadr[nv-1] + ctx->L_rownnz[nv-1]) {
mjERROR("mismatch between pre-counted and post-factorization L nonzeros");
}
} }
// dense // dense