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 mjModel* m, mjData* d, const mjtNum* J,
const mjtNum* B, int n, int offset, const mjtNum* B, int n, int offset,
const int* J_rownnz, const int* J_rowadr, const int* J_colind) { 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) ) // compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
for (int i = 0; i < n; i++) { for (int i = 0; i < n; i++) {
@@ -749,19 +744,14 @@ static void addJTBJSparse(
int ik = J_rowadr[offset_i] + k; int ik = J_rowadr[offset_i] + k;
int colik = J_colind[ik]; int colik = J_colind[ik];
// row = J(i,k)*B(i,j)*J(j,:) // qDeriv(k,:) += J(j,:) * J(i,k)*B(i,j)
mju_scl(row, J + J_rowadr[offset_j], J[ik]*B[i*n+j], J_rownnz[offset_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],
// qDeriv(k,:) += row J_rownnz[offset_j], J_colind + J_rowadr[offset_j],
mju_addToSparseInc(d->qDeriv + d->D_rowadr[colik], row, J[ik]*B[i*n+j]);
d->D_rownnz[colik], d->D_colind + d->D_rowadr[colik],
J_rownnz[offset_j], J_colind + J_rowadr[offset_j]);
} }
} }
} }
// free space
mj_freeStack(d);
} }
src/engine/engine_util_sparse.c
+5 -5
View File
@@ -355,10 +355,10 @@ void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNu
// dst += src, only at common non-zero indices // dst += scl*src, only at common non-zero indices
void mju_addToSparseInc(mjtNum* dst, const mjtNum* src, void mju_addToSclSparseInc(mjtNum* dst, const mjtNum* src,
int nnzdst, const int* inddst, int nnzdst, const int* inddst,
int nnzsrc, const int* indsrc) { int nnzsrc, const int* indsrc, mjtNum scl) {
if (!nnzdst || !nnzsrc) { if (!nnzdst || !nnzsrc) {
return; return;
} }
@@ -368,7 +368,7 @@ void mju_addToSparseInc(mjtNum* dst, const mjtNum* src,
// common non-zero index // common non-zero index
if (inds == indd) { if (inds == indd) {
// add // add
dst[adrd] += src[adrs]; dst[adrd] += scl * src[adrs];
// advance src // advance src
if (++adrs < nnzsrc) { if (++adrs < nnzsrc) {
src/engine/engine_util_sparse.h
+4 -4
View File
@@ -63,10 +63,10 @@ int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b, void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind); int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind);
// dst += src, only at common non-zero indices // dst += scl * src, only at common non-zero indices
void mju_addToSparseInc(mjtNum* dst, const mjtNum* src, void mju_addToSclSparseInc(mjtNum* dst, const mjtNum* src,
int nnzdst, const int* inddst, int nnzdst, const int* inddst,
int nnzsrc, const int* indsrc); int nnzsrc, const int* indsrc, mjtNum scl);
// add to sparse matrix: dst = dst + scl*src, return nnz of result // add to sparse matrix: dst = dst + scl*src, return nnz of result
int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum scl, int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum scl,