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

Refactor Newton solver: move Hessian from arena back to the stack.

Browse Source 
The changes in 9a0dc20821, moving the Hessian from the stack to the arena, should be rolled back: they prevent future threading over islands. Unlike the stack, arena allocations are not thread-friendly. The reason for the original move was to save memory, but the savings are small: `O(ctx->nH)` and only linear in `nv`. The significant reduction of the Cholesky factor, from quadratic in `nv` to quadratic in the largest dof island — the reduction afforded by 2dd518734f — remains in place.

Also refactor and improve readability.

PiperOrigin-RevId: 685717356
Change-Id: Ia9ec3e44a62d459a3b9cffb578cf6479d5aa1d7f
This commit is contained in:
Yuval Tassa
2024-10-14 08:32:44 -07:00
committed by Copybara-Service
parent ac91a7639d
commit 7ec94f46d4
4 changed files with 260 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+11 -6
View File
@@ -738,7 +738,7 @@ void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) {
// res_rowadr is required to be precomputed
void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
int* res_rownnz, int* res_rowadr, int* res_colind,
int* res_rownnz, const int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
@@ -859,13 +859,17 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
// compute row non-zeros of reverse-Cholesky factor L, return total
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
const int* rowadr, const int* colind, int n) {
int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
int n, mjData* d) {
mj_markStack(d);
int* parent = mj_stackAllocInt(d, n);
int* flag = mj_stackAllocInt(d, n);
// loop over rows in reverse order
for (int r = n - 1; r >= 0; r--) {
parent[r] = -1;
flag[r] = r;
L_rownnz[r] = 0;
L_rownnz[r] = 1; // start with 1 for diagonal
int start = rowadr[r];
int end = start + rownnz[r];
// loop over non-zero columns
@@ -886,10 +890,11 @@ int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
}
}
// add 1 for diagonal, accumulate sum
mj_freeStack(d);
// accumulate sum
int sum = 0;
for (int r = 0; r < n; r++) {
L_rownnz[r]++;
sum += L_rownnz[r];
}