Optimize and refactor combineSparse and combineSparseCount.

PiperOrigin-RevId: 558126391
Change-Id: Iead14c31136e789560f3c8124b84bd4e383158f4

src/engine/engine_core_constraint.c

@@ -546,9 +546,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
 
           // compute Jacobian: sparse or dense
           if (issparse) {
-            NV = mju_combineSparse(jac[0], jac[1], nv, 1, -deriv,
+            NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain,
-                                   NV, NV2, chain, chain2,
+                                   chain2, sparse_buf, buf_ind);
-                                   sparse_buf, buf_ind);
           } else {
             mju_addToScl(jac[0], jac[1], -deriv, nv);
           }
@@ -1310,26 +1309,27 @@ void mj_makeImpedance(const mjModel* m, mjData* d) {
 
 // count the number of non-zeros in the sum of two sparse vectors
 int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
-  int a = 0;
+  int a = 0, b = 0, c_nnz = 0;
-  int b = 0;
-  int nnz = 0;
 
-  // while there are elements remaining in both a_ind and b_ind
+  // count c_nnz: nonzero indices common to both a and b
   while (a < a_nnz && b < b_nnz) {
-    // add the smaller element of either a_ind[a] or b_ind[b] to the combined nnz
+    // common index, increment everything
-    ++nnz;
+    if (a_ind[a] == b_ind[b]) {
+      c_nnz++;
+      a++;
+      b++;
+    }
 
-    // if a_ind[a] == b_ind[b], increment both a and b so that we don't double count
+    // update smallest index
-    // otherwise, increment the index pointing to the smaller element
+    else if (a_ind[a] < b_ind[b]) {
-    int aa = a;
+      a++;
-    int bb = b;
+    } else {
-    if (a_ind[aa] <= b_ind[bb]) ++a;
+      b++;
-    if (a_ind[aa] >= b_ind[bb]) ++b;
+    }
   }
 
-  // count remaining elements from the vector with larger nnz
+  // union minus the intersection
-  nnz += (a_nnz - a) + (b_nnz - b);
+  return a_nnz + b_nnz - c_nnz;
-  return nnz;
 }
 
 

src/engine/engine_core_smooth.c

@@ -566,9 +566,9 @@ void mj_tendon(const mjModel* m, mjData* d) {
             mju_mulMatTVec(tmp, jacdif, dif, 3, NV);
 
             // add to existing
-            rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, nv, 1, 1/divisor,
+            rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor,
-                                          rownnz[i], NV, colind+rowadr[i], chain,
+                                          rownnz[i], NV, colind+rowadr[i],
-                                          sparse_buf, buf_ind);
+                                          chain, sparse_buf, buf_ind);
           }
 
           // dense

src/engine/engine_support.c

@@ -1053,7 +1053,7 @@ void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
 
   // add to destination
   for (int i=0; i < nv; i++) {
-    rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
+    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], sparse_buf, buf_ind);
   }

src/engine/engine_util_solve.c

@@ -190,7 +190,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
       int c = colind[adr+i];
 
       // 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], c + 1, 1, -mat[adr+i],
+      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],
                                     sparse_buf, buf_ind);
 
@@ -278,7 +278,7 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
     mat[adr+nnz-1] = r;
 
     // update row:  mat(r,1:r-1) = (mat(r,1:r-1) + s*x(1:r-1)) / c
-    int new_nnz = mju_combineSparse(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
+    int new_nnz = mju_combineSparse(mat + adr, x, 1 / c, (flg_plus ? s / c : -s / c),
                                     nnz-1, i, colind + adr, x_ind,
                                     sparse_buf, buf_ind);
 
@@ -288,9 +288,8 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
     }
 
     // update x:  x(1:r-1) = c*x(1:r-1) - s*mat(r,1:r-1)
-    int new_x_nnz = mju_combineSparse(x, mat+adr, n, c, -s,
+    int new_x_nnz = mju_combineSparse(x, mat+adr, c, -s, i, nnz-1, x_ind,
-                                      i, nnz-1, x_ind, colind+adr,
+                                      colind+adr, sparse_buf, buf_ind);
-                                      sparse_buf, buf_ind);
 
     // update i, correct for changing x
     i = i - 1 + (new_x_nnz - i);

src/engine/engine_util_sparse.c

@@ -208,7 +208,7 @@ static int mju_compare(const int* vec1, const int* vec2, int n) {
 
 
 // combine two sparse vectors: dst = a*dst + b*src, return nnz of result
-int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+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,
                       mjtNum* buf, int* buf_ind) {
   // check for identical pattern
@@ -226,40 +226,45 @@ int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
     memcpy(buf_ind, dst_ind, dst_nnz*sizeof(int));
   }
 
-  // prepare to merge buf and scr into dst
+  // prepare to merge buf and src into dst
   int bi = 0, si = 0, nnz = 0;
   int buf_nnz = dst_nnz;
-  int badr = bi < buf_nnz ? buf_ind[bi] : n+1;
-  int sadr = si < src_nnz ? src_ind[si] : n+1;
 
   // merge vectors
-  while (bi < buf_nnz || si < src_nnz) {
+  while (bi < buf_nnz && si < src_nnz) {
-    // both
+    int badr = buf_ind[bi];
+    int sadr = src_ind[si];
+
     if (badr == sadr) {
       dst[nnz] = a*buf[bi++] + b*src[si++];
       dst_ind[nnz++] = badr;
-
-      badr = bi < buf_nnz ? buf_ind[bi] : n+1;
-      sadr = si < src_nnz ? src_ind[si] : n+1;
     }
 
-    // dst only
+    // buf only
     else if (badr < sadr) {
       dst[nnz] = a*buf[bi++];
       dst_ind[nnz++] = badr;
-
-      badr = bi < buf_nnz ? buf_ind[bi] : n+1;
     }
 
     // src only
     else {
       dst[nnz] = b*src[si++];
       dst_ind[nnz++] = sadr;
-
-      sadr = si < src_nnz ? src_ind[si] : n+1;
     }
   }
 
+  // the rest of src only
+  while (si < src_nnz) {
+    dst[nnz] = b*src[si];
+    dst_ind[nnz++] = src_ind[si++];
+  }
+
+  // the rest of buf only
+  while (bi < buf_nnz) {
+    dst[nnz] = a*buf[bi];
+    dst_ind[nnz++] = buf_ind[bi++];
+  }
+
   return nnz;
 }
 

src/engine/engine_util_sparse.h

@@ -52,7 +52,7 @@ MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc,
                               int* rownnz, int* rowadr, int* colind);
 
 // combine two sparse vectors: dst = a*dst + b*src, return nnz of result
-int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+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,
                       mjtNum* buf, int* buf_ind);
 

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -206,7 +206,7 @@ int compare_memcmp(const int* vec1,
 }
 
 int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
-                                                   const mjtNum* src, int n,
+                                                   const mjtNum* src,
                                                    mjtNum a, mjtNum b,
                                                    int dst_nnz, int src_nnz,
                                                    int* dst_ind,
@@ -223,7 +223,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
 }
 
 int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst,
-                                              const mjtNum* src, int n,
+                                              const mjtNum* src,
                                               mjtNum a, mjtNum b,
                                               int dst_nnz, int src_nnz,
                                               int* dst_ind,
@@ -470,7 +470,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
         // true arguments should be i+1 and colind+rowadr[r]
         // but instead we repeat rownnz[c] and colind+rowadr[c]
         // in order to trigger all if's in combineSparse
-         func(H+rowadr[c], H+rowadr[r], c+1, 1, -H[adr+i],
+         func(H+rowadr[c], H+rowadr[r], 1, -H[adr+i],
               rownnz[c], rownnz[c],
               colind+rowadr[c], colind+rowadr[c], NULL, NULL);
       }