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

Use diagonal index to speed up mju_solveLUSparse.

Browse Source 
PiperOrigin-RevId: 704303282
Change-Id: I72c3fbcd31a5564fc792b73efe9d12a151b4d60f
This commit is contained in:
Yuval Tassa
2024-12-09 08:52:31 -08:00
committed by Copybara-Service
parent bd1d94c282
commit 537fc2ff45
3 changed files with 14 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_forward.c
+1 -1
View File
@@ -967,7 +967,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
} }
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_colind); mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_diag, d->D_colind);
} }
// IMPLICITFAST // IMPLICITFAST
src/engine/engine_util_solve.c
+12 -22
View File
@@ -651,41 +651,31 @@ void mju_factorLUSparse(mjtNum* LU, int n, int* scratch,
// solve mat*res=vec given LU factorization of mat // solve mat*res=vec given LU factorization of mat
void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n, void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n,
const int* rownnz, const int* rowadr, const int* colind) { const int* rownnz, const int* rowadr, const int* diag, const int* colind) {
//------------------ solve (U+I)*res = vec // solve (U+I)*res = vec
for (int i=n-1; i >= 0; i--) { for (int i=n-1; i >= 0; i--) {
// init: diagonal of (U+I) is 1 // init: diagonal of (U+I) is 1
res[i] = vec[i]; res[i] = vec[i];
// res[i] -= sum_k>i res[k]*LU(i,k) int d1 = diag[i]+1;
int j = rownnz[i] - 1; int nnz = rownnz[i] - d1;
while (colind[rowadr[i]+j] > i) { if (nnz > 0) {
res[i] -= res[colind[rowadr[i]+j]] * LU[rowadr[i]+j]; int adr = rowadr[i] + d1;
j--; res[i] -= mju_dotSparse(LU+adr, res, nnz, colind+adr, /*flg_unc1=*/0);
}
// make sure j points to diagonal
if (colind[rowadr[i]+j] != i) {
mjERROR("diagonal of U not reached");
} }
} }
//------------------ solve L*res(new) = res //------------------ solve L*res(new) = res
for (int i=0; i < n; i++) { for (int i=0; i < n; i++) {
// res[i] -= sum_k<i res[k]*LU(i,k) // res[i] -= sum_k<i res[k]*LU(i,k)
int j = 0; int d = diag[i];
while (colind[rowadr[i]+j] < i) { int adr = rowadr[i];
res[i] -= res[colind[rowadr[i]+j]] * LU[rowadr[i]+j]; if (d > 0) {
j++; res[i] -= mju_dotSparse(LU+adr, res, d, colind+adr, /*flg_unc1=*/0);
} }
// divide by diagonal element of L // divide by diagonal element of L
res[i] /= LU[rowadr[i]+j]; res[i] /= LU[adr + d];
// make sure j points to diagonal
if (colind[rowadr[i]+j] != i) {
mjERROR("diagonal of L not reached");
}
} }
} }
src/engine/engine_util_solve.h
+1 -1
View File
@@ -82,7 +82,7 @@ void mju_factorLUSparse(mjtNum *LU, int n, int* scratch,
// solve mat*res=vec given LU factorization of mat // solve mat*res=vec given LU factorization of mat
void mju_solveLUSparse(mjtNum *res, const mjtNum *LU, const mjtNum* vec, int n, void mju_solveLUSparse(mjtNum *res, const mjtNum *LU, const mjtNum* vec, int n,
const int *rownnz, const int *rowadr, const int *colind); const int *rownnz, const int *rowadr, const int* diag, const int *colind);
// eigenvalue decomposition of symmetric 3x3 matrix // eigenvalue decomposition of symmetric 3x3 matrix
MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]); MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);