Refactor mju_combineSparse to eliminate temporary buffers.

Combine sparse vectors in-place by first counting total `nnz` and then working backwards from the end. This removes the need for temporary buffers in `mju_combineSparse` and its callers and speeds up the function by ~10%.

PiperOrigin-RevId: 902530210
Change-Id: I4f48c327103552ab968d3915399c6067367bec9f

src/engine/engine_core_constraint.c

@@ -496,12 +496,12 @@ static void mj_equalityAnchors(const mjModel* m, const mjData* d, int eq_id,
 // equality constraints
 void mj_instantiateEquality(const mjModel* m, mjData* d) {
   int issparse = mj_isSparse(m), nv = m->nv;
-  int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL, *buf_ind = NULL;
+  int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL;
   int flex_edgeadr, flex_edgenum;
   int flex_vertadr, flex_vertnum;
   mjtNum cpos[6], pos[2][3], ref[2], dif, deriv;
   mjtNum quat[4], quat1[4], quat2[4], quat3[4], axis[3];
-  mjtNum *jac[2], *jacdif, *data, *sparse_buf = NULL;
+  mjtNum *jac[2], *jacdif, *data;
 
   // disabled or no equality constraints: return
   if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax == 0) {
@@ -520,8 +520,6 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
   if (issparse) {
     chain = mjSTACKALLOC(d, nv, int);
     chain2 = mjSTACKALLOC(d, nv, int);
-    buf_ind = mjSTACKALLOC(d, nv, int);
-    sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
   }
 
   // find active equality constraints
@@ -689,8 +687,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
 
         // compute Jacobian: sparse or dense
         if (issparse) {
-          NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain,
+          NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain, chain2);
-                                  chain2, sparse_buf, buf_ind);
         } else {
           mju_addToScl(jac[0], jac[1], -deriv, nv);
         }

src/engine/engine_solver.c

@@ -829,8 +829,6 @@ typedef struct {
   int* L_rowadr;          // Hessian factor row addresses                 (nv x 1)
   int* LT_rownnz;         // Hessian factor transpose row nonzeros        (nv x 1)
   int* LT_rowadr;         // Hessian factor transpose row addresses       (nv x 1)
-  int* buf_ind;           // index buffer for sparse addition             (nv x 1)
-  mjtNum* buf_val;        // value buffer for sparse addition             (nv x 1)
 
   // Newton arrays, computed-size (MakeHessian)
   int nH;                 // number of nonzeros in Hessian H
@@ -979,9 +977,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
   if (flg_Newton) {
     nNum += nefc + nv;                 // D, cholupd
     if (is_elliptic) nNum += 6*nv;     // LTJ
-    if (is_sparse) {
+    if (!is_sparse) {
-      nNum += nv;                      // buf_val
-    } else {
       nNum += nv*nv;                   // L (dense)
       if (is_elliptic) nNum += nv*nv;  // Lcone (dense)
     }
@@ -993,7 +989,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
   size_t nInt = nefc;                  // oldstate
   if (is_sparse) {
     nInt += 3*nv + nJ;                 // JT sparse
-    if (flg_Newton) nInt += 9*nv;      // Newton sparse
+    if (flg_Newton) nInt += 8*nv;      // Newton sparse
   }
 
   // allocate mjtNum and int blocks
@@ -1018,9 +1014,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
     if (is_elliptic) {
       ctx->LTJ = numblock;  numblock += 6*nv;
     }
-    if (is_sparse) {
+    if (!is_sparse) {
-      ctx->buf_val = numblock;  numblock += nv;
-    } else {
       ctx->nL = nv*nv;
       ctx->L     = numblock;  numblock += ctx->nL;
       ctx->Lcone = is_elliptic ? numblock : NULL;
@@ -1049,7 +1043,6 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
     ctx->L_rowadr   = intblock;  intblock += nv;
     ctx->LT_rownnz  = intblock;  intblock += nv;
     ctx->LT_rowadr  = intblock;  intblock += nv;
-    ctx->buf_ind    = intblock;  intblock += nv;
   }
 
   // sparse: compute Jacobian transpose
@@ -1616,8 +1609,7 @@ static void MakeHessian(mjData* d, mjPrimalContext* ctx) {
 
     // add mass matrix: H = J'*D*J + M
     mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv,
-                       ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
+                       ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind);
-                       ctx->buf_val, ctx->buf_ind);
 
     // compute H' sparse structure (upper triangle, required for symbolic Cholesky)
     mju_transposeSparse(NULL, NULL, nv, nv, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, NULL,
@@ -1691,8 +1683,7 @@ static void FactorizeHessian(mjData* d, mjPrimalContext* ctx, int flg_recompute)
 
       // add mass matrix: H = J'*D*J + C
       mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv,
-                         ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
+                         ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind);
-                         ctx->buf_val, ctx->buf_ind);
     }
 
     // numeric sparse factorization: L = chol(H) using pre-computed sparsity pattern

src/engine/engine_support.c

@@ -420,17 +420,11 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
 void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
              int* rownnz, int* rowadr, int* colind) {
   int nv = m->nv;
+
   // sparse
   if (rownnz && rowadr && colind) {
-    mj_markStack(d);
+    mju_addToMatSparse(dst, rownnz, rowadr, colind, nv, d->M,
-    mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum);
+                       m->M_rownnz, m->M_rowadr, m->M_colind);
-    int* buf_ind = mjSTACKALLOC(d, nv, int);
-
-    mju_addToMatSparse(dst, rownnz, rowadr, colind, nv,
-      d->M, m->M_rownnz, m->M_rowadr, m->M_colind,
-      buf_val, buf_ind);
-
-    mj_freeStack(d);
   }
 
   // dense

src/engine/engine_util_solve.c

@@ -144,10 +144,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
                          int* rownnz, const int* rowadr, int* colind,
                          mjData* d) {
   int rank = n;
-
+  (void) d;
-  mj_markStack(d);
-  mjtNum* buf = mjSTACKALLOC(d, n, mjtNum);
-  int* buf_ind = mjSTACKALLOC(d, n, int);
 
   // backpass over rows
   for (int r=n-1; r >= 0; r--) {
@@ -175,15 +172,13 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
 
       // mat(c,0:c) = mat(c,0:c) - mat(r,c) * mat(r,0:c)
       int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], 1, -mat[adr+i],
-                                    rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r],
+                                    rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r]);
-                                    buf, buf_ind);
 
       // assign new nnz to row c
       rownnz[c] = nnz_c;
     }
   }
 
-  mj_freeStack(d);
   return rank;
 }
 

src/engine/engine_util_sparse.c

@@ -198,15 +198,14 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 }
 
 
-// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
+// add sparse matrix M to sparse destination matrix
 void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
                         const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
-                        const int* M_colind,
+                        const int* M_colind) {
-                        mjtNum* buf_val, int* buf_ind) {
   for (int i=0; i < nr; i++) {
     rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
                                   rownnz[i], M_rownnz[i], colind + rowadr[i],
-                                  M_colind + M_rowadr[i], buf_val, buf_ind);
+                                  M_colind + M_rowadr[i]);
   }
 }
 
@@ -256,8 +255,8 @@ void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat,
     for (int k=diag-1; k >= 0; k--) {
       int j = ind[k];
       mjtNum val = row[k];
-      res[i] += val * vec[j]; // strict lower
+      res[i] += val * vec[j];  // strict lower
-      res[j] += val * vec[i]; // strict upper
+      res[j] += val * vec[i];  // strict upper
     }
   }
 }

src/engine/engine_util_sparse.h

@@ -62,11 +62,10 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
 MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc,
                                 const int* rownnz, const int* rowadr, const int* colind);
 
-// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
+// add sparse matrix M to sparse destination matrix
 MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
                               const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
-                              const int* M_colind,
+                              const int* M_colind);
-                              mjtNum* buf_val, int* buf_ind);
 
 // add symmetric matrix (only lower triangle represented) to dense matrix
 MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n,
@@ -294,8 +293,7 @@ void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, i
 // combine two sparse vectors: dst = a*dst + b*src, return nnz of result
 static inline
 int mju_combineSparse(mjtNum* dst, const mjtNum* src, 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) {
-                      mjtNum* buf, int* buf_ind) {
   // check for identical pattern
   if (dst_nnz == src_nnz) {
     if (mju_compare(dst_ind, src_ind, dst_nnz)) {
@@ -305,49 +303,54 @@ int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
     }
   }
 
-  // copy dst into buf
+  // compute total nnz of result
-  if (dst_nnz) {
+  int nnz = mju_combineSparseCount(dst_nnz, src_nnz, dst_ind, src_ind);
-    memcpy(buf, dst, dst_nnz * sizeof(mjtNum));
-    memcpy(buf_ind, dst_ind, dst_nnz * sizeof(int));
-  }
 
-  // prepare to merge buf and src into dst
+  // set up read/write pointers at end of arrays
-  int bi = 0, si = 0, nnz = 0;
+  int bi = dst_nnz - 1, si = src_nnz - 1, w = nnz - 1;
-  int buf_nnz = dst_nnz;
 
-  // merge vectors
+  // merge backwards
-  while (bi < buf_nnz && si < src_nnz) {
+  while (bi >= 0 && si >= 0) {
-    int badr = buf_ind[bi];
+    int badr = dst_ind[bi];
     int sadr = src_ind[si];
 
     if (badr == sadr) {
-      dst[nnz] = a*buf[bi++] + b*src[si++];
+      dst[w] = a*dst[bi] + b*src[si];
-      dst_ind[nnz++] = badr;
+      dst_ind[w] = badr;
+      bi--;
+      si--;
     }
 
-    // buf only
+    // dst only
-    else if (badr < sadr) {
+    else if (badr > sadr) {
-      dst[nnz] = a*buf[bi++];
+      dst[w] = a*dst[bi];
-      dst_ind[nnz++] = badr;
+      dst_ind[w] = badr;
+      bi--;
     }
 
     // src only
     else {
-      dst[nnz] = b*src[si++];
+      dst[w] = b*src[si];
-      dst_ind[nnz++] = sadr;
+      dst_ind[w] = sadr;
+      si--;
     }
+    w--;
   }
 
-  // the rest of src only
+  // remaining src elements
-  while (si < src_nnz) {
+  while (si >= 0) {
-    dst[nnz] = b*src[si];
+    dst[w] = b*src[si];
-    dst_ind[nnz++] = src_ind[si++];
+    dst_ind[w] = src_ind[si];
+    si--;
+    w--;
   }
 
-  // the rest of buf only
+  // remaining dst elements: already in place, scale by a
-  while (bi < buf_nnz) {
+  if (a != 1) {
-    dst[nnz] = a*buf[bi];
+    while (bi >= 0) {
-    dst_ind[nnz++] = buf_ind[bi++];
+      dst[bi] *= a;
+      bi--;
+    }
   }
 
   return nnz;

test/benchmark/chol_benchmark_test.cc

@@ -22,7 +22,6 @@
 #include <absl/base/attributes.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
-#include "src/engine/engine_memory.h"
 #include "src/engine/engine_support.h"
 #include "src/engine/engine_util_solve.h"
 #include "src/engine/engine_util_sparse.h"
@@ -352,10 +351,6 @@ constexpr int kNumUpdateVectors = 25;
 int ABSL_ATTRIBUTE_NOINLINE mju_cholUpdateSparse_old(
     mjtNum* mat, mjtNum* x, int n, int flg_plus, const int* rownnz,
     const int* rowadr, const int* colind, int x_nnz, int* x_ind, mjData* d) {
-  mj_markStack(d);
-  int* buf_ind = mjSTACKALLOC(d, n, int);
-  mjtNum* sparse_buf = mjSTACKALLOC(d, n, mjtNum);
-
   int rank = n, i = x_nnz - 1;
   while (i >= 0) {
     int nnz = rownnz[x_ind[i]], adr = rowadr[x_ind[i]];
@@ -372,10 +367,9 @@ int ABSL_ATTRIBUTE_NOINLINE mju_cholUpdateSparse_old(
     mju_combineSparseInc(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
                          nnz - 1, i, colind + adr, x_ind);
     int new_x_nnz = mju_combineSparse(x, mat + adr, c, -s, i, nnz - 1, x_ind,
-                                      colind + adr, sparse_buf, buf_ind);
+                                      colind + adr);
     i = i - 1 + (new_x_nnz - i);
   }
-  mj_freeStack(d);
   return rank;
 }
 

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -99,8 +99,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
                                                    mjtNum a, mjtNum b,
                                                    int dst_nnz, int src_nnz,
                                                    int* dst_ind,
-                                                   const int* src_ind,
+                                                   const int* src_ind) {
-                                                   mjtNum* buf, int* buf_ind) {
   // check for identical pattern
   if (compare_baseline(dst_ind, src_ind, dst_nnz)) {
     // combine mjtNum data directly
@@ -116,8 +115,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst,
                                               mjtNum a, mjtNum b,
                                               int dst_nnz, int src_nnz,
                                               int* dst_ind,
-                                              const int* src_ind,
+                                              const int* src_ind) {
-                                              mjtNum* buf, int* buf_ind) {
   // check for identical pattern
   if (compare_memcmp(dst_ind, src_ind, dst_nnz)) {
     // combine mjtNum data directly
@@ -372,7 +370,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
         // in order to trigger all if's in combineSparse
          func(H+rowadr[c], H+rowadr[r], 1, -H[adr+i],
               rownnz[c], rownnz[c],
-              colind+rowadr[c], colind+rowadr[c], NULL, NULL);
+              colind+rowadr[c], colind+rowadr[c]);
       }
     }
   }