Switch mjData.{qH,qLD} from reduced ("C") to full ("M") inertia matrix structure. No performance impact of extra zeros because of existing "simple dof" skipping mechanism.

yuvaltassa · copybara-github · commit 0f563ecf31bf · 2025-03-04T17:34:35.000-08:00
PiperOrigin-RevId: 733523931
Change-Id: Ic8d8a152dda5532331c239cb6b4ce7d8d09b7fff
diff --git a/doc/includes/references.h b/doc/includes/references.h
@@ -272,7 +272,7 @@ struct mjData_ {
   mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -305,7 +305,7 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
@@ -300,7 +300,7 @@ struct mjData_ {
   mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -333,7 +333,7 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
@@ -650,7 +650,7 @@
     X   ( mjtNum,    actuator_moment,   nJmom,       1           ) \
     X   ( mjtNum,    crb,               nbody,       10          ) \
     X   ( mjtNum,    qM,                nM,          1           ) \
-    X   ( mjtNum,    qLD,               nC,          1           ) \
+    X   ( mjtNum,    qLD,               nM,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
     XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
@@ -667,7 +667,7 @@
     X   ( mjtNum,    qfrc_passive,      nv,          1           ) \
     X   ( mjtNum,    subtree_linvel,    nbody,       3           ) \
     X   ( mjtNum,    subtree_angmom,    nbody,       3           ) \
-    X   ( mjtNum,    qH,                nC,          1           ) \
+    X   ( mjtNum,    qH,                nM,          1           ) \
     X   ( mjtNum,    qHDiagInv,         nv,          1           ) \
     X   ( int,       B_rownnz,          nbody,       1           ) \
     X   ( int,       B_rowadr,          nbody,       1           ) \
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -372,7 +372,7 @@ def make_data(
         'efc_aref': (nefc, float),
         'efc_force': (nefc, float),
         '_qM_sparse': (m.nM, float),
-        '_qLD_sparse': (m.nC, float),
+        '_qLD_sparse': (m.nM, float),
         '_qLDiagInv_sparse': (m.nv, float),
     }
 
@@ -511,7 +511,7 @@ def get_data_into(
         value = value[dof_i, dof_j]
       elif field.name == 'qLD' and not support.is_sparse(m):
         # TODO(erikfrey): provide correct qLDs
-        value = np.zeros(m.nC)
+        value = np.zeros(m.nM)
       elif field.name == 'qLDiagInv' and not support.is_sparse(m):
         value = np.ones(m.nv)
 
diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -92,8 +92,8 @@ def test_smooth(self):
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
     qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
-    for i in range(m.nC):
-      qLDLegacy[d.mapM2C[i]] = d.qLD[i]
+    for i in range(m.nM):
+      qLDLegacy[d.mapM2M[i]] = d.qLD[i]
     _assert_eq(qLDLegacy, dx.qLD, 'qLD')
     _assert_attr_eq(d, dx, 'qLDiagInv')
     _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
@@ -1352,7 +1352,7 @@ class Data(PyTreeNode):
     efc_aref: reference pseudo-acceleration                     (nefc,)
     efc_force: constraint force in constraint space             (nefc,)
     _qM_sparse: qM in sparse representation                     (nM,)
-    _qLD_sparse: qLD in sparse representation                   (nC,)
+    _qLD_sparse: qLD in sparse representation                   (nM,)
     _qLDiagInv_sparse: qLDiagInv in sparse representation       (nv,)
   """  # fmt: skip
   # constant sizes:
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
@@ -5356,7 +5356,7 @@
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of M (sparse)",
-                 array_extent=('nC',),
+                 array_extent=('nM',),
              ),
              StructFieldDecl(
                  name='qLDiagInv',
@@ -5492,7 +5492,7 @@
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of modified M",
-                 array_extent=('nC',),
+                 array_extent=('nM',),
              ),
              StructFieldDecl(
                  name='qHDiagInv',
diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
@@ -2131,7 +2131,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
   // inverse square root of D from inertia LDL decomposition
   mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
   for (int i=0; i < nv; i++) {
-    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
     sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
 
@@ -2167,11 +2167,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
           continue;
         }
 
-        // traverse row j of C, marking new unique nonzeros
-        int nnzC = d->C_rownnz[j];
-        int adrC = d->C_rowadr[j];
-        for (int k=0; k < nnzC; k++) {
-          int c = d->C_colind[adrC + k];
+        // traverse row j of M, marking new unique nonzeros
+        int nnzM = d->M_rownnz[j];
+        int adrM = d->M_rowadr[j];
+        for (int k=0; k < nnzM; k++) {
+          int c = d->M_colind[adrM + k];
           if (marker[c] != r) {
             marker[c] = r;
             nnz++;
@@ -2251,10 +2251,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
           continue;
         }
         int j = B_colind[i];
-        int adrC = d->C_rowadr[j];
-        mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
+        int adrM = d->M_rowadr[j];
+        mju_addToSclSparseInc(B + adrB, d->qLD + adrM,
                               nnzB, B_colind + adrB,
-                              d->C_rownnz[j]-1, d->C_colind + adrC, -b);
+                              d->M_rownnz[j]-1, d->M_colind + adrM, -b);
       }
 
       // B(r,:) <- sqrt(inv(D)) * B(r,:)
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
@@ -1469,11 +1469,11 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 void mj_factorM(const mjModel* m, mjData* d) {
   TM_START;
-  int nC = m->nC;
-  for (int i=0; i < nC; i++) {
-    d->qLD[i] = d->qM[d->mapM2C[i]];
+  int nM = m->nM;
+  for (int i=0; i < nM; i++) {
+    d->qLD[i] = d->qM[d->mapM2M[i]];
   }
-  mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+  mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
   TM_ADD(mjTIMER_POS_INERTIA);
 }
 
@@ -1715,7 +1715,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
     mju_copy(x, y, n*m->nv);
   }
   mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
-             d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
 }
 
 
@@ -1727,14 +1727,14 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
   const mjtNum* qLDiagInv = d->qLDiagInv;
   if (island < 0) {
     mj_solveLD(x, qLD, qLDiagInv, m->nv, 1,
-               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
     return;
   }
 
   // local copies of key variables
-  const int* rownnz = d->C_rownnz;
-  const int* rowadr = d->C_rowadr;
-  const int* colind = d->C_colind;
+  const int* rownnz = d->M_rownnz;
+  const int* rowadr = d->M_rowadr;
+  const int* colind = d->M_colind;
   const int* diagnum = m->dof_simplenum;
 
   // local constants: island specific
@@ -1785,9 +1785,9 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
   int nv = m->nv;
 
   // local copies of key variables
-  const int* rownnz = d->C_rownnz;
-  const int* rowadr = d->C_rowadr;
-  const int* colind = d->C_colind;
+  const int* rownnz = d->M_rownnz;
+  const int* rowadr = d->M_rowadr;
+  const int* colind = d->M_colind;
   const int* diagnum = m->dof_simplenum;
   const mjtNum* qLD = d->qLD;
 
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
@@ -770,7 +770,7 @@ static void mj_advance(const mjModel* m, mjData* d,
 // Euler integrator, semi-implicit in velocity, possibly skipping factorisation
 void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nv = m->nv, nC = m->nC;
+  int nv = m->nv, nM = m->nM;
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
   mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
@@ -795,22 +795,22 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   else {
     if (!skipfactor) {
       // qH = M + h*diag(B)
-      for (int i=0; i < nC; i++) {
-        d->qH[i] = d->qM[d->mapM2C[i]];
+      for (int i=0; i < nM; i++) {
+        d->qH[i] = d->qM[d->mapM2M[i]];
       }
       for (int i=0; i < nv; i++) {
-        d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
+        d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
       }
 
       // factorize in-place
-      mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
     }
 
     // solve
     mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
     mju_copy(qacc, qfrc, m->nv);
     mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
-               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
   }
 
   // advance state and time
@@ -939,7 +939,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 // fully implicit in velocity, possibly skipping factorization
 void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC;
+  int nv = m->nv, nM = m->nM, nD = m->nD;
 
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -987,18 +987,18 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
       // copy into qH
-      for (int i=0; i < nC; i++) {
-        d->qH[i] = MhB[d->mapM2C[i]];
+      for (int i=0; i < nM; i++) {
+        d->qH[i] = MhB[d->mapM2M[i]];
       }
 
       // factorize in-place
-      mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
     mju_copy(qacc, qfrc, nv);
     mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
-               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
 
   } else {
     mjERROR("integrator must be implicit or implicitfast");
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
@@ -1127,12 +1127,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
 
   printInertia("QM", d->qM, m, fp, float_format);
 
-  printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
-              d->C_rowadr, d->C_colind, fp, float_format);
+  printSparse("QLD", d->qLD, m->nv, d->M_rownnz,
+              d->M_rowadr, d->M_colind, fp, float_format);
   printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
 
   if (!mju_isZero(d->qHDiagInv, m->nv)) {
-    printSparse("QH", d->qH, m->nv, d->C_rownnz, d->C_rowadr, d->C_colind, fp, float_format);
+    printSparse("QH", d->qH, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind, fp, float_format);
     printArray("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format);
   }
 
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
@@ -1097,14 +1097,14 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
 
     // non-simple: add off-diagonals
     if (!m->dof_simplenum[i]) {
-      int adr = d->C_rowadr[i];
-      res[i] += mju_dotSparse(qLD+adr, vec, d->C_rownnz[i] - 1, d->C_colind+adr, /*flg_unc1=*/0);
+      int adr = d->M_rowadr[i];
+      res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr, /*flg_unc1=*/0);
     }
   }
 
   // res *= sqrt(D)
   for (int i=0; i < nv; i++) {
-    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
     res[i] *= mju_sqrt(qLD[diag]);
   }
 }
diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c
@@ -556,7 +556,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
 
   // compute average spatial inertia at selection point
   for (int i=0; i < nv; i++) {
-    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
     sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
   mj_jac(m, d, jac, NULL, selpos, sel);
diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc
@@ -44,9 +44,9 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
   mj_markStack(d);
 
   // M: mass matrix in CSR format
-  mjtNum* M = mj_stackAllocNum(d, m->nC);
-  for (int i=0; i < m->nC; i++) {
-    M[i] = d->qM[d->mapM2C[i]];
+  mjtNum* M = mj_stackAllocNum(d, m->nM);
+  for (int i=0; i < m->nM; i++) {
+    M[i] = d->qM[d->mapM2M[i]];
   }
 
   // LDlegacy: legacy LD matrix (size nM)
@@ -60,7 +60,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
       } else {
         mju_copy(d->qLD, M, m->nC);
         mj_factorI(d->qLD, d->qLDiagInv, m->nv,
-                   d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+                   d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
       }
     }
   }
diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc
@@ -46,9 +46,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
   mj_markStack(d);
 
   // M: mass matrix in CSR format
-  mjtNum* M = mj_stackAllocNum(d, m->nC);
-  for (int i=0; i < m->nC; i++) {
-    M[i] = d->qM[d->mapM2C[i]];
+  mjtNum* M = mj_stackAllocNum(d, m->nM);
+  for (int i=0; i < m->nM; i++) {
+    M[i] = d->qM[d->mapM2M[i]];
   }
 
   // LDlegacy: legacy LD matrix (size nM)
@@ -74,9 +74,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
         case SolveType::kCsr:
           mju_copy(d->qLD, M, m->nC);
           mj_factorI(d->qLD, d->qLDiagInv, m->nv,
-                     d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+                     d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
           mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
-                     d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+                     d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
       }
     }
   }
diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc
@@ -52,9 +52,8 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
 
   // make legacy matrix
   mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
-  mju_zero(LDlegacy, m->nM);
-  for (int i=0; i < m->nC; i++) {
-    LDlegacy[d->mapM2C[i]] = d->qLD[i];
+  for (int i=0; i < m->nM; i++) {
+    LDlegacy[d->mapM2M[i]] = d->qLD[i];
   }
 
   // benchmark
@@ -65,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
         mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv);
       } else {
         mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
-                   d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+                   d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
       }
     }
   }
diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc
diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc
diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc
diff --git a/test/fixture.h b/test/fixture.h
