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

Avoid allocation and copying in implicit solver's addJTBJSparse

Browse Source 
PiperOrigin-RevId: 710679558
Change-Id: Ie75fdcd1127ff619c2668f568ab6a8a11e079dda
This commit is contained in:
Yuval Tassa
2024-12-30 04:18:11 -08:00
committed by Copybara-Service
parent 8cb253b6ca
commit e42370c982
3 changed files with 14 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_derivative.c
+5 -15
View File
@@ -730,11 +730,6 @@ static void addJTBJSparse(
const mjModel* m, mjData* d, const mjtNum* J,
const mjtNum* B, int n, int offset,
const int* J_rownnz, const int* J_rowadr, const int* J_colind) {
int nv = m->nv;
// allocate row
mj_markStack(d);
mjtNum* row = mjSTACKALLOC(d, nv, mjtNum);
// compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
for (int i = 0; i < n; i++) {
@@ -749,19 +744,14 @@ static void addJTBJSparse(
int ik = J_rowadr[offset_i] + k;
int colik = J_colind[ik];
// row = J(i,k)*B(i,j)*J(j,:)
mju_scl(row, J + J_rowadr[offset_j], J[ik]*B[i*n+j], J_rownnz[offset_j]);
// qDeriv(k,:) += row
mju_addToSparseInc(d->qDeriv + d->D_rowadr[colik], row,
d->D_rownnz[colik], d->D_colind + d->D_rowadr[colik],
J_rownnz[offset_j], J_colind + J_rowadr[offset_j]);
// qDeriv(k,:) += J(j,:) * J(i,k)*B(i,j)
mju_addToSclSparseInc(d->qDeriv + d->D_rowadr[colik], J + J_rowadr[offset_j],
d->D_rownnz[colik], d->D_colind + d->D_rowadr[colik],
J_rownnz[offset_j], J_colind + J_rowadr[offset_j],
J[ik]*B[i*n+j]);
}
}
}
// free space
mj_freeStack(d);
}