Separate mj_addM into dense and sparse functions.

PiperOrigin-RevId: 539066339
Change-Id: Ide3f9644fee46a5fb1a2e79a25be162b929e2421

src/engine/engine_solver.c

@@ -1375,7 +1375,7 @@ 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_addM(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
+    mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
 
     // factorize H, uncompressed layout
     int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
@@ -1404,7 +1404,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
   else {
     // compute H = M + J'*D*J
     mju_sqrMatTD(ctx->H, d->efc_J, D, nefc, nv);
-    mj_addM(m, d, ctx->H, NULL, NULL, NULL);
+    mj_addMDense(m, d, ctx->H);
 
     // factorize H
     mju_cholFactor(ctx->H, nv, mjMINVAL);

src/engine/engine_support.c

@@ -831,10 +831,24 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
 //  destination can be sparse uncompressed, or dense when all int* are NULL
 void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
              int* rownnz, int* rowadr, int* colind) {
-  int adr, adr1, nv = m->nv;
-
   // sparse
   if (rownnz && rowadr && colind) {
+    mj_addMSparse(m, d, dst, rownnz, rowadr, colind);
+  }
+
+  // dense
+  else {
+    mj_addMDense(m, d, dst);
+  }
+}
+
+
+
+// add inertia matrix to sparse uncompressed destination matrix
+void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
+                   int* rownnz, int* rowadr, int* colind) {
+  int adr, adr1, nv = m->nv;
+
   // special processing of simple dofs
   int simplecnt = 0;
   for (int i=0; i < nv; i++) {
@@ -947,12 +961,16 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
   }
 
   mjFREESTACK;
-  }
+}
 
-  // dense
+
-  else {
+
-    for (int i=0; i < nv; i++) {
+// add inertia matrix to dense destination matrix
-      adr = m->dof_Madr[i];
+void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) {
+  int nv = m->nv;
+
+  for (int i = 0; i < nv; i++) {
+    int adr = m->dof_Madr[i];
     int j = i;
     while (j >= 0) {
       // add
@@ -971,11 +989,9 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
       adr++;
     }
   }
-  }
 }
 
 
-
 //-------------------------- sparse system matrix conversion ---------------------------------------
 
 // dst[D] = src[M], handle different sparsity representations

src/engine/engine_support.h

@@ -104,6 +104,13 @@ 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
+MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
+                         int* rownnz, int* rowadr, int* colind);
+
+// add inertia matrix to dense destination matrix
+MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
+
 
 //-------------------------- sparse system matrix conversion ---------------------------------------
 

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -20,6 +20,7 @@
 #include <absl/base/attributes.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
+#include "src/engine/engine_support.h"
 #include "src/engine/engine_util_sparse.h"
 #include "test/fixture.h"
 
@@ -458,7 +459,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
                       d->efc_JT_colind, d->efc_JT_rowsuper, d);
 
   // compute H = M + J'*D*J
-  mj_addM(m, d, H, rownnz, rowadr, colind);
+  mj_addMSparse(m, d, H, rownnz, rowadr, colind);
 
   // time benchmark
   for (auto s : state) {