Avoid allocation and copying in implicit solver's addJTBJSparse

PiperOrigin-RevId: 710679558
Change-Id: Ie75fdcd1127ff619c2668f568ab6a8a11e079dda

Author: yuvaltassa

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);
 }
 
 

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) {

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,