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

Clean up mju_sqrMatTDSparse

Browse Source 
PiperOrigin-RevId: 711486739
Change-Id: I5f9145dcd8522a9cb372cbdf60611d7d05cb524e
This commit is contained in:
Yuval Tassa
2025-01-02 12:00:53 -08:00
committed by Copybara-Service
parent 4510c6d290
commit 6c880eb0fa
1 changed files with 34 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+34 -20
View File
@@ -742,7 +742,8 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
int* markers = mjSTACKALLOC(d, nc, int); int* markers = mjSTACKALLOC(d, nc, int);
for (int i=0; i < nc; i++) { for (int i=0; i < nc; i++) {
int* cols = res_colind+res_rowadr[i]; int rowadr_i = res_rowadr[i];
int* cols = res_colind + rowadr_i;
res_rownnz[i] = 0; res_rownnz[i] = 0;
buffer[i] = 0; buffer[i] = 0;
@@ -755,18 +756,26 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
} }
// iterate through each row of M' // iterate through each row of M'
int end = rowadrT[i] + rownnzT[i]; int adrT = rowadrT[i];
for (int r = rowadrT[i]; r < end; r++) { int end_r = adrT + rownnzT[i];
for (int r = adrT; r < end_r; r++) {
int t = colindT[r]; int t = colindT[r];
mjtNum v = diag ? matT[r] * diag[t] : matT[r]; int adr = rowadr[t];
for (int c=rowadr[t]; c < rowadr[t]+rownnz[t]; c++) { int end_c = adr + rownnz[t];
for (int c=adr; c < end_c; c++) {
int cc = colind[c]; int cc = colind[c];
// ignore upper triangle // ignore upper triangle
if (cc > i) { if (cc > i) {
break; break;
} }
buffer[cc] += v*mat[c]; // add value to buffer
if (diag) {
buffer[cc] += matT[r] * diag[t] * mat[c];
} else {
buffer[cc] += matT[r] * mat[c];
}
// only need to insert nnz if not marked // only need to insert nnz if not marked
if (!markers[cc]) { if (!markers[cc]) {
@@ -810,22 +819,26 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
} }
} }
end = res_rownnz[i]; end_r = res_rownnz[i];
// rowsuperT: reuse sparsity, copy into res // rowsuperT: reuse sparsity, copy into res
if (rowsuperT && rowsuperT[i]) { if (rowsuperT && rowsuperT[i]) {
for (int r=0; r < end; r++) { for (int r=0; r < end_r; r++) {
res[res_rowadr[i] + r] = buffer[cols[r]]; int c = cols[r];
buffer[cols[r]] = 0; res[rowadr_i + r] = buffer[c];
buffer[c] = 0;
} }
} else { }
// clear out buffers since sparsity cannot be reused
for (int r=0; r < end; r++) { // clear out buffers, sparsity cannot be reused
int cc = cols[r]; else {
res[res_rowadr[i] + r] = buffer[cc]; for (int r=0; r < end_r; r++) {
res_colind[res_rowadr[i] + r] = cc; int c = cols[r];
buffer[cc] = 0; int adr = rowadr_i + r;
markers[cc] = 0; res[adr] = buffer[c];
res_colind[adr] = c;
buffer[c] = 0;
markers[c] = 0;
} }
} }
} }
@@ -833,8 +846,9 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
// fill upper triangle // fill upper triangle
for (int i=0; i < nc; i++) { for (int i=0; i < nc; i++) {
int end = res_rowadr[i] + res_rownnz[i] - 1; int start = res_rowadr[i];
for (int j=res_rowadr[i]; j < end; j++) { int end = start + res_rownnz[i] - 1;
for (int j=start; j < end; j++) {
int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++; int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
res[adr] = res[j]; res[adr] = res[j];
res_colind[adr] = i; res_colind[adr] = i;