Rework sparse addM function to support compressed sparse matrices.

Note:

- M is still a nv x nv uncompressed sparse matrix.
- NNZ precounting still needs to be implemented using the combineSparseCount helper function.

PiperOrigin-RevId: 554452710
Change-Id: I14adf58fc3acd4ed82d864d7c2e9e96ee6b0377d

src/engine/engine_solver.c

@@ -1364,7 +1364,15 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
 
   // sparse
   if (mj_isSparse(m)) {
+    // create sparse inertia matrix M (uncompressed)
+    mjtNum* M = mj_stackAlloc(d, nv*nv);
+    int* M_rownnz = (int*) mj_stackAlloc(d, nv);
+    int* M_rowadr = (int*) mj_stackAlloc(d, nv);
+    int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
+    mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
+
     // compute H = J'*D*J
+
     // TODO(b/266802572): remove uncompressed layout
     mju_sqrMatTDUncompressedInit(ctx->rowadr, nv);
     mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
@@ -1375,7 +1383,8 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
                        d->efc_JT_colind, d->efc_JT_rowsuper, d);
 
     // compute H = M + J'*D*J
-    mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
+    mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind,
+                  M, M_rownnz, M_rowadr, M_colind);
 
     // factorize H, uncompressed layout
     int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,

src/engine/engine_support.c

@@ -948,7 +948,17 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
              int* rownnz, int* rowadr, int* colind) {
   // sparse
   if (rownnz && rowadr && colind) {
-    mj_addMSparse(m, d, dst, rownnz, rowadr, colind);
+    int nv = m->nv;
+    mjMARKSTACK;
+    // create sparse inertia matrix M (uncompressed)
+    mjtNum* M = mj_stackAlloc(d, nv*nv);
+    int* M_rownnz = (int*) mj_stackAlloc(d, nv);
+    int* M_rowadr = (int*) mj_stackAlloc(d, nv);
+    int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
+    mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
+    mj_addMSparse(m, d, dst, rownnz, rowadr, colind, M,
+                  M_rownnz, M_rowadr, M_colind);
+    mjFREESTACK;
   }
 
   // dense
@@ -959,122 +969,87 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
 
 
 
-// add inertia matrix to sparse uncompressed destination matrix
+// create inertia matrix M (uncompressed)
-void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
+void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
-                   int* rownnz, int* rowadr, int* colind) {
+                      int* M_rownnz, int* M_rowadr, int* M_colind) {
   int adr, adr1, nv = m->nv;
 
-  // special processing of simple dofs
+  // build M into sparse format, lower-triangular
-  int simplecnt = 0;
   for (int i=0; i < nv; i++) {
+    M_rowadr[i] = i*nv;
+    adr = m->dof_Madr[i];
+
     if (m->dof_simplenum[i]) {
-      // count simple
+      M[i*nv] = d->qM[adr];
-      simplecnt++;
+      M_colind[i*nv] = i;
-
+      M_rownnz[i] = 1;
-      // empty row: create entry
+      continue;
-      if (!rownnz[i]) {
-        colind[rowadr[i]] = i;
-        dst[rowadr[i]] = d->qM[m->dof_Madr[i]];
-        rownnz[i] = 1;
-      }
-
-      // non-empty row: assume dof is in dst (J'*D*J satisfies this)
-      else {
-        // find dof in row, add
-        adr = rowadr[i];
-        int end = adr + rownnz[i];
-        while (adr < end)
-          if (colind[adr] == i) {
-            dst[adr] += d->qM[m->dof_Madr[i]];
-            break;
-          } else {
-            adr++;
-          }
-
-        // not found: error
-        if (adr >= end) {
-          mjERROR("dst row expected to be empty");
-        }
-      }
     }
-  }
 
-  // done if all simple
+    // backward pass over dofs: construct M_row(i) in reverse order
-  if (simplecnt == nv) {
+    int j = i;
-    return;
+    adr1 = 0;
-  }
+    while (j >= 0) {
+      // assign
+      M[i*nv+adr1] = d->qM[adr];
+      M_colind[i*nv+adr1] = j;
 
-  // allocate space for sparse M
+      // count columns
-  mjMARKSTACK;
+      adr1++;
-  mjtNum* M = mj_stackAlloc(d, nv*nv);
-  int* M_rownnz = (int*) mj_stackAlloc(d, nv);
-  int* M_rowadr = (int*) mj_stackAlloc(d, nv);
-  int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
-  int* buf_ind = (int*) mj_stackAlloc(d, nv);
-  mjtNum* sparse_buf = mj_stackAlloc(d, nv);
 
-  // convert M into sparse format, lower-triangular
+      // advance
-  for (int i=0; i < nv; i++) {
+      adr++;
-    if (!m->dof_simplenum[i]) {
+      j = m->dof_parentid[j];
-      // backward pass over dofs: construct M_row(i) in reverse order
+    }
-      adr = m->dof_Madr[i];
-      int j = i;
-      adr1 = 0;
-      while (j >= 0) {
-        // assign
-        M[i*nv+adr1] = d->qM[adr];
-        M_colind[i*nv+adr1] = j;
 
-        // count columns
+    // assign row descriptors
-        adr1++;
+    M_rownnz[i] = adr1;
 
-        // advance
+    // reverse order
-        adr++;
+    for (int k=0; k < adr1/2; k++) {
-        j = m->dof_parentid[j];
+      mjtNum tmp = M[i*nv+k];
-      }
+      M[i*nv+k] = M[i*nv+adr1-1-k];
+      M[i*nv+adr1-1-k] = tmp;
 
-      // assign row descriptors
+      int tmpi = M_colind[i*nv+k];
-      M_rownnz[i] = adr1;
+      M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
-      M_rowadr[i] = i*nv;
+      M_colind[i*nv+adr1-1-k] = tmpi;
-
-      // reverse order
-      for (int k=0; k < adr1/2; k++) {
-        mjtNum tmp = M[i*nv+k];
-        M[i*nv+k] = M[i*nv+adr1-1-k];
-        M[i*nv+adr1-1-k] = tmp;
-
-        int tmpi = M_colind[i*nv+k];
-        M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
-        M_colind[i*nv+adr1-1-k] = tmpi;
-      }
     }
   }
 
   // make symmetric
   for (int i=1; i < nv; i++) {
-    if (!m->dof_simplenum[i]) {
+    if (m->dof_simplenum[i]) {
-      for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
+      continue;
-        // add to row given by column index
+    }
-        adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
+
-        M[adr1] = M[k];
+    for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
-        M_colind[adr1] = i;
+      // add to row given by column index
-      }
+      adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
+      M[adr1] = M[k];
+      M_colind[adr1] = i;
     }
   }
+}
+
+
+
+// add inertia matrix to sparse destination matrix
+void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
+                   int* rownnz, int* rowadr, int* colind, mjtNum* M,
+                   int* M_rownnz, int* M_rowadr, int* M_colind) {
+  int nv = m->nv;
+
+  mjMARKSTACK;
+  int* buf_ind = (int*) mj_stackAlloc(d, nv);
+  mjtNum* sparse_buf = mj_stackAlloc(d, nv);
 
   // add to destination
   for (int i=0; i < nv; i++) {
-    if (!m->dof_simplenum[i]) {
+    rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
-      int new_nnz =
+                                  rownnz[i], M_rownnz[i], colind + rowadr[i],
-          mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
+                                  M_colind + M_rowadr[i], sparse_buf, buf_ind);
-                            rownnz[i], M_rownnz[i],
-                            colind + rowadr[i], M_colind + M_rowadr[i],
-                            sparse_buf, buf_ind);
-
-      rownnz[i] = new_nnz;
-    }
   }
-
   mjFREESTACK;
 }
 

src/engine/engine_support.h

@@ -116,9 +116,14 @@ MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
 MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
                    int* rownnz, int* rowadr, int* colind);
 
-// add inertia matrix to sparse uncompressed destination matrix
+// create inertia matrix M (uncompressed)
+MJAPI void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
+                            int* M_rownnz, int* M_rowadr, int* M_colind);
+
+// add inertia matrix to sparse destination matrix
 MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
-                         int* rownnz, int* rowadr, int* colind);
+                         int* rownnz, int* rowadr, int* colind, mjtNum* M,
+                         int* M_rownnz, int* M_rowadr, int* M_colind);
 
 // add inertia matrix to dense destination matrix
 MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -452,14 +452,14 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
   // compute H = J'*D*J, uncompressed layout
   mju_sqrMatTDUncompressedInit(rowadr, m->nv);
   mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv,
-                      rownnz, rowadr, colind,
+                     rownnz, rowadr, colind,
-                      d->efc_J_rownnz, d->efc_J_rowadr,
+                     d->efc_J_rownnz, d->efc_J_rowadr,
-                      d->efc_J_colind, d->efc_J_rowsuper,
+                     d->efc_J_colind, d->efc_J_rowsuper,
-                      d->efc_JT_rownnz, d->efc_JT_rowadr,
+                     d->efc_JT_rownnz, d->efc_JT_rowadr,
-                      d->efc_JT_colind, d->efc_JT_rowsuper, d);
+                     d->efc_JT_colind, d->efc_JT_rowsuper, d);
 
   // compute H = M + J'*D*J
-  mj_addMSparse(m, d, H, rownnz, rowadr, colind);
+  mj_addM(m, d, H, rownnz, rowadr, colind);
 
   // time benchmark
   for (auto s : state) {