From e42370c982ea7a4fe37be20f31c7fb32be283f38 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 30 Dec 2024 04:18:11 -0800 Subject: [PATCH] Avoid allocation and copying in implicit solver's `addJTBJSparse` PiperOrigin-RevId: 710679558 Change-Id: Ie75fdcd1127ff619c2668f568ab6a8a11e079dda --- src/engine/engine_derivative.c | 20 +++++--------------- src/engine/engine_util_sparse.c | 10 +++++----- src/engine/engine_util_sparse.h | 8 ++++---- 3 files changed, 14 insertions(+), 24 deletions(-) diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index d4c764cb..deb2c0c9 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -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); } diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 27b26576..895179e2 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -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 -void mju_addToSparseInc(mjtNum* dst, const mjtNum* src, - int nnzdst, const int* inddst, - int nnzsrc, const int* indsrc) { +// dst += scl*src, only at common non-zero indices +void mju_addToSclSparseInc(mjtNum* dst, const mjtNum* src, + int nnzdst, const int* inddst, + int nnzsrc, const int* indsrc, mjtNum scl) { if (!nnzdst || !nnzsrc) { return; } @@ -368,7 +368,7 @@ void mju_addToSparseInc(mjtNum* dst, const mjtNum* src, // common non-zero index if (inds == indd) { // add - dst[adrd] += src[adrs]; + dst[adrd] += scl * src[adrs]; // advance src if (++adrs < nnzsrc) { diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 319af57b..48b41979 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -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, int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind); -// dst += src, only at common non-zero indices -void mju_addToSparseInc(mjtNum* dst, const mjtNum* src, - int nnzdst, const int* inddst, - int nnzsrc, const int* indsrc); +// dst += scl * src, only at common non-zero indices +void mju_addToSclSparseInc(mjtNum* dst, const mjtNum* src, + int nnzdst, const int* inddst, + int nnzsrc, const int* indsrc, mjtNum scl); // 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,