Reuse sparse Dof-Dof matrix to make compressed sparse M in mj_addM.

PiperOrigin-RevId: 556823353
Change-Id: I650ac83767d60c69859aae6e38aa331f43814378

src/engine/engine_solver.c

@@ -1364,12 +1364,12 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
 
   // sparse
   if (mj_isSparse(m)) {
-    // create sparse inertia matrix M (uncompressed)
+    // create sparse inertia matrix M
-    mjtNum* M = mj_stackAlloc(d, nv*nv);
+    int nnz = m->nD;  // use sparse dof-dof matrix
-    int* M_rownnz = (int*) mj_stackAlloc(d, nv);
+    int* M_rownnz = (int*) mj_stackAlloc(d, nv);  // actual nnz count
-    int* M_rowadr = (int*) mj_stackAlloc(d, nv);
+    int* M_colind = (int*) mj_stackAlloc(d, nnz);
-    int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
+    mjtNum* M = mj_stackAlloc(d, nnz);
-    mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
+    mj_makeMSparse(m, d, M, M_rownnz, NULL, M_colind);
 
     // compute H = J'*D*J
 
@@ -1384,7 +1384,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
 
     // compute H = M + J'*D*J
     mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind,
-                  M, M_rownnz, M_rowadr, M_colind);
+                  M, M_rownnz, NULL, M_colind);
 
     // factorize H, uncompressed layout
     int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,

src/engine/engine_support.c

@@ -950,14 +950,15 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
   if (rownnz && rowadr && colind) {
     int nv = m->nv;
     mjMARKSTACK;
-    // create sparse inertia matrix M (uncompressed)
+    // create sparse inertia matrix M
-    mjtNum* M = mj_stackAlloc(d, nv*nv);
+    int nnz = m->nD;  // use sparse dof-dof matrix
-    int* M_rownnz = (int*) mj_stackAlloc(d, nv);
+    int* M_rownnz = (int*) mj_stackAlloc(d, nv);  // actual nnz count
-    int* M_rowadr = (int*) mj_stackAlloc(d, nv);
+    int* M_colind = (int*) mj_stackAlloc(d, nnz);
-    int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
+    mjtNum* M = mj_stackAlloc(d, nnz);
-    mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
+
+    mj_makeMSparse(m, d, M, M_rownnz, NULL, M_colind);
     mj_addMSparse(m, d, dst, rownnz, rowadr, colind, M,
-                  M_rownnz, M_rowadr, M_colind);
+                  M_rownnz, NULL, M_colind);
     mjFREESTACK;
   }
 
@@ -969,65 +970,64 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
 
 
 
-// create inertia matrix M (uncompressed)
+// make inertia matrix M
-void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
+void mj_makeMSparse(const mjModel* m, mjData* d, mjtNum* M,
-                      int* M_rownnz, int* M_rowadr, int* M_colind) {
+                    int* M_rownnz, int* M_rowadr, int* M_colind) {
-  int adr, adr1, nv = m->nv;
+  int nv = m->nv;
+  // currently the sparse dof-dof matrix D row addresses are used, since D has
+  // the same predetermined sparsity structure as M, however with simple bodies
+  // M has less non-zeros and can be precounted for further memory reduction
+  if (M_rowadr == NULL) {
+    M_rowadr = d->D_rowadr;
+  }
 
-  // build M into sparse format, lower-triangular
+  // build M into sparse format, lower triangle
-  for (int i=0; i < nv; i++) {
+  for (int i = 0; i < nv; i++) {
-    M_rowadr[i] = i*nv;
+    int Madr = m->dof_Madr[i];
-    adr = m->dof_Madr[i];
 
+    // simple, fill diagonal only
     if (m->dof_simplenum[i]) {
-      M[i*nv] = d->qM[adr];
-      M_colind[i*nv] = i;
       M_rownnz[i] = 1;
+      M[M_rowadr[i]] = d->qM[Madr];
+      M_colind[M_rowadr[i]] = i;
       continue;
     }
 
     // backward pass over dofs: construct M_row(i) in reverse order
-    int j = i;
+    int col = M_rowadr[i]; // current column in row i
-    adr1 = 0;
+    for (int j = i; j >= 0; j = m->dof_parentid[j]) {
-    while (j >= 0) {
+      M[col] = d->qM[Madr++];
-      // assign
+      M_colind[col++] = j;
-      M[i*nv+adr1] = d->qM[adr];
-      M_colind[i*nv+adr1] = j;
-
-      // count columns
-      adr1++;
-
-      // advance
-      adr++;
-      j = m->dof_parentid[j];
     }
 
-    // assign row descriptors
+    // track nnz of lower triangle for row i
-    M_rownnz[i] = adr1;
+    int nnz = M_rownnz[i] = col - M_rowadr[i];
 
     // reverse order
-    for (int k=0; k < adr1/2; k++) {
+    int end = nnz >> 1;
-      mjtNum tmp = M[i*nv+k];
+    for (int j = 0; j < end; j++) {
-      M[i*nv+k] = M[i*nv+adr1-1-k];
+      int a1 = M_rowadr[i] + j;              // address 1
-      M[i*nv+adr1-1-k] = tmp;
+      int a2 = (M_rowadr[i] + nnz - 1) - j;  // address 2
 
-      int tmpi = M_colind[i*nv+k];
+      // swap M data on row i
-      M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
+      mjtNum val = M[a1];
-      M_colind[i*nv+adr1-1-k] = tmpi;
+      M[a1] = M[a2];
+      M[a2] = val;
+
+      // swap M column indices on row i
+      int ind = M_colind[a1];
+      M_colind[a1] = M_colind[a2];
+      M_colind[a2] = ind;
     }
   }
 
-  // make symmetric
+  // fill upper triangle
-  for (int i=1; i < nv; i++) {
+  for (int i = 1; i < nv; i++) {
-    if (m->dof_simplenum[i]) {
+    int end = M_rowadr[i] + M_rownnz[i] - 1;
-      continue;
+    for (int j = M_rowadr[i]; j < end; j++) {
-    }
+      int a = M_rowadr[M_colind[j]] + M_rownnz[M_colind[j]]++;
-
+      M[a] = M[j];
-    for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
+      M_colind[a] = 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;
     }
   }
 }
@@ -1039,6 +1039,12 @@ 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;
+  // currently the sparse dof-dof matrix D row addresses are used, since D has
+  // the same predetermined sparsity structure as M, however with simple bodies
+  // M has less non-zeros and can be precounted for further memory reduction
+  if (M_rowadr == NULL) {
+    M_rowadr = d->D_rowadr;
+  }
 
   mjMARKSTACK;
   int* buf_ind = (int*) mj_stackAlloc(d, nv);

src/engine/engine_support.h

@@ -116,9 +116,9 @@ 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);
 
-// create inertia matrix M (uncompressed)
+// make inertia matrix M
-MJAPI void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
+MJAPI void mj_makeMSparse(const mjModel* m, mjData* d, mjtNum* M,
-                            int* M_rownnz, int* M_rowadr, int* M_colind);
+                          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,

test/engine/engine_support_test.cc

@@ -16,10 +16,10 @@
 
 #include <random>
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
 
@@ -34,6 +34,7 @@ using ::testing::DoubleNear;
 using ::testing::ContainsRegex;
 using ::testing::MatchesRegex;
 using ::testing::Pointwise;
+using ::testing::ElementsAreArray;
 using JacobianTest = MujocoTest;
 static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
 
@@ -405,5 +406,52 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
   mj_deleteModel(model);
 }
 
+using AddMTest = MujocoTest;
+
+TEST_F(AddMTest, DenseSameAsSparse) {
+  mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml");
+  mjData* d = mj_makeData(m);
+  int nv = m->nv;
+
+  // force use of sparse matrices
+  m->opt.jacobian = mjJAC_SPARSE;
+
+  // warm-up rollout to get a typical state
+  while (d->time < 2) {
+    mj_step(m, d);
+  }
+
+  // dense zero matrix
+  std::vector<mjtNum> dst_sparse = std::vector(nv * nv, 0.0);
+
+  // sparse zero matrix
+  std::vector<mjtNum> dst_dense = std::vector(nv * nv, 0.0);
+  std::vector<int> rownnz = std::vector(nv, nv);
+  std::vector<int> rowadr = std::vector(nv, 0);
+  std::vector<int> colind = std::vector(nv * nv, 0);
+
+  // set sparse structure
+  for (int i = 0; i < nv; i++) {
+    rowadr[i] = i * nv;
+    for (int j = 0; j < nv; j++) {
+      colind[rowadr[i] + j] = j;
+    }
+  }
+
+  // sparse addM
+  mj_addM(m, d, dst_sparse.data(), rownnz.data(),
+          rowadr.data(), colind.data());
+
+  // dense addM
+  mj_addM(m, d, dst_dense.data(), NULL, NULL, NULL);
+
+  // dense comparison, should be same matrix
+  EXPECT_THAT(dst_dense, ElementsAreArray(dst_sparse));
+
+  // clean up
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco