Convert qLD to CSR format.

PiperOrigin-RevId: 723955038
Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b

doc/APIreference/functions.rst

@@ -423,7 +423,8 @@ Get name of object with the specified mjtObj type and id, returns NULL if name n
 
 .. mujoco-include:: mj_fullM
 
-Convert sparse inertia matrix M into full (i.e. dense) matrix.
+Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
+|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``.
 
 .. _mj_mulM:
 

doc/APIreference/functions_override.rst

@@ -233,6 +233,11 @@ found, the function will return ``distmax`` and ``fromto``, if given, will be se
    In order to determine whether a geom pair uses ``mjc_Convex``, inspect the table at the top of
    `engine_collision_driver.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__.
 
+.. _mj_fullM:
+
+Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
+|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``.
+
 .. _mj_mulM:
 
 This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format

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)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC 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               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData

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)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC 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               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData

include/mujoco/mjxmacro.h

@@ -637,7 +637,7 @@
     X   ( mjtNum,    actuator_moment,   nJmom,       1           ) \
     X   ( mjtNum,    crb,               nbody,       10          ) \
     X   ( mjtNum,    qM,                nM,          1           ) \
-    X   ( mjtNum,    qLD,               nM,          1           ) \
+    X   ( mjtNum,    qLD,               nC,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
     XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
@@ -654,7 +654,7 @@
     X   ( mjtNum,    qfrc_passive,      nv,          1           ) \
     X   ( mjtNum,    subtree_linvel,    nbody,       3           ) \
     X   ( mjtNum,    subtree_angmom,    nbody,       3           ) \
-    X   ( mjtNum,    qH,                nM,          1           ) \
+    X   ( mjtNum,    qH,                nC,          1           ) \
     X   ( mjtNum,    qHDiagInv,         nv,          1           ) \
     X   ( int,       B_rownnz,          nbody,       1           ) \
     X   ( int,       B_rowadr,          nbody,       1           ) \

introspect/structs.py

@@ -5261,7 +5261,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of M (sparse)",
-                 array_extent=('nM',),
+                 array_extent=('nC',),
              ),
              StructFieldDecl(
                  name='qLDiagInv',
@@ -5397,7 +5397,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of modified M",
-                 array_extent=('nM',),
+                 array_extent=('nC',),
              ),
              StructFieldDecl(
                  name='qHDiagInv',

mjx/mujoco/mjx/_src/io.py

@@ -368,7 +368,7 @@ def make_data(
         'efc_aref': (nefc, float),
         'efc_force': (nefc, float),
         '_qM_sparse': (m.nM, float),
-        '_qLD_sparse': (m.nM, float),
+        '_qLD_sparse': (m.nC, float),
         '_qLDiagInv_sparse': (m.nv, float),
     }
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -91,7 +91,10 @@ class SmoothTest(absltest.TestCase):
     _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'qLD')
+    qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
+    for i in range(m.nC):
+      qLDLegacy[d.mapM2C[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')
     _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')

mjx/mujoco/mjx/_src/types.py

@@ -1348,7 +1348,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                   (nM,)
+    _qLD_sparse: qLD in sparse representation                   (nC,)
     _qLDiagInv_sparse: qLDiagInv in sparse representation       (nv,)
   """  # fmt: skip
   # constant sizes:

src/engine/engine_core_constraint.c

@@ -2131,7 +2131,8 @@ 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++) {
-    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
 
   // sparse
@@ -2238,13 +2239,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 
     // === in-place sparse back-substitution:  B <- B * M^-1/2
 
-    // make qLD
-    int nC = m->nC;
-    mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum);
-    for (int i=0; i < nC; i++) {
-      qLD[i] = d->qLD[d->mapM2C[i]];
-    }
-
     // sparse backsubM2 (half of LD back-substitution)
     for (int r=0; r < nefc; r++) {
       int nnzB = B_rownnz[r];
@@ -2258,7 +2252,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
         }
         int j = B_colind[i];
         int adrC = d->C_rowadr[j];
-        mju_addToSclSparseInc(B + adrB, qLD + adrC,
+        mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
                               nnzB, B_colind + adrB,
                               d->C_rownnz[j]-1, d->C_colind + adrC, -b);
       }

src/engine/engine_core_smooth.c

@@ -1465,7 +1465,11 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 void mj_factorM(const mjModel* m, mjData* d) {
   TM_START;
-  mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
+  int nC = m->nC;
+  for (int i=0; i < nC; i++) {
+    d->qLD[i] = d->qM[d->mapM2C[i]];
+  }
+  mj_factorIs(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
   TM_ADD(mjTIMER_POS_INERTIA);
 }
 
@@ -1709,18 +1713,20 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
   if (x != y) {
     mju_copy(x, y, n*m->nv);
   }
-  mj_solveLD(m, x, n, d->qLD, d->qLDiagInv);
+  mj_solveLDs(x, d->qLD, d->qLDiagInv, m->nv, n,
+              d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 }
 
 
 // in-place sparse backsubstitution for one island:  x = inv(L'*D*L)*x
 //  L is in lower triangle of qLD; D is on diagonal of qLD
-void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) {
+void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
   // if no islands, call mj_solveLD
   const mjtNum* qLD = d->qLD;
   const mjtNum* qLDiagInv = d->qLDiagInv;
   if (island < 0) {
-    mj_solveLD(m, x, 1, qLD, qLDiagInv);
+    mj_solveLDs(x, qLD, qLDiagInv, m->nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     return;
   }
 
@@ -1730,14 +1736,6 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
   const int* colind = d->C_colind;
   const int* diagnum = m->dof_simplenum;
 
-  // temporary: make local CSR version of qLD
-  int nC = m->nC;
-  mj_markStack(d);
-  mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum);
-  for (int i=0; i < nC; i++) {
-    qLDs[i] = d->qLD[d->mapM2C[i]];
-  }
-
   // local constants: island specific
   int ndof = d->island_dofnum[island];
   const int* dofind = d->island_dofind + d->island_dofadr[island];
@@ -1751,7 +1749,7 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
       int start = rowadr[i];
       int end = start + rownnz[i] - 1;
       for (int adr=end-1; adr >= start; adr--) {
-        x[islandind[colind[adr]]] -= qLDs[adr] * x_k;
+        x[islandind[colind[adr]]] -= qLD[adr] * x_k;
       }
     }
   }
@@ -1773,11 +1771,9 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
     int start = rowadr[i];
     int end = start + rownnz[i] - 1;
     for (int adr=end-1; adr >= start; adr--) {
-      x[k] -= x[islandind[colind[adr]]] * qLDs[adr];
+      x[k] -= x[islandind[colind[adr]]] * qLD[adr];
     }
   }
-
-  mj_freeStack(d);
 }
 
 
@@ -1785,23 +1781,18 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
                 const mjtNum* sqrtInvD, int n) {
+  int nv = m->nv;
+
   // local copies of key variables
-  int nv = m->nv, nC = m->nC;
   const int* rownnz = d->C_rownnz;
   const int* rowadr = d->C_rowadr;
   const int* colind = d->C_colind;
   const int* diagnum = m->dof_simplenum;
+  const mjtNum* qLD = d->qLD;
 
   // x = y
   mju_copy(x, y, n * nv);
 
-  // temporary: make local CSR version of qLD
-  mj_markStack(d);
-  mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum);
-  for (int i=0; i < nC; i++) {
-    qLD[i] = d->qLD[d->mapM2C[i]];
-  }
-
   // x <- L^-T x
   for (int i=nv-1; i > 0; i--) {
     // skip diagonal rows
@@ -1831,8 +1822,6 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
       x[i+offset] *= invD_i;
     }
   }
-
-  mj_freeStack(d);
 }
 
 

src/engine/engine_core_smooth.h

@@ -71,9 +71,8 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 
-// TODO(tassa): Restore mjData const-ness.
 // sparse backsubstitution for one island:  x = inv(L'*D*L)*x, use factorization in d
-MJAPI void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* x, int island);
+MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,

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, nM = m->nM;
+  int nv = m->nv, nC = m->nC;
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
   mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
@@ -794,22 +794,23 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   // damping: integrate implicitly
   else {
     if (!skipfactor) {
-      mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum);
+      // qH = M + h*diag(B)
-
+      for (int i=0; i < nC; i++) {
-      // MhB = M + h*diag(B)
+        d->qH[i] = d->qM[d->mapM2C[i]];
-      mju_copy(MhB, d->qM, nM);
+      }
       for (int i=0; i < nv; i++) {
-        MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i];
+        d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
       }
 
-      // factor
+      // factorize in-place
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
+      mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     }
 
     // solve
     mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
     mju_copy(qacc, qfrc, m->nv);
-    mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
+    mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
   }
 
   // advance state and time
@@ -938,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;
+  int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC;
 
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -985,13 +986,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       // set MhB = M - dt*qDeriv
       mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
-      // factorize
+      // copy into qH
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
+      for (int i=0; i < nC; i++) {
+        d->qH[i] = MhB[d->mapM2C[i]];
+      }
+
+      // factorize in-place
+      mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
     mju_copy(qacc, qfrc, nv);
-    mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
+    mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+
   } else {
     mjERROR("integrator must be implicit or implicitfast");
   }

src/engine/engine_print.c

@@ -1127,7 +1127,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
 
   printInertia("QM", d->qM, m, fp, float_format);
 
-  printInertia("QLD", d->qLD, m, fp, float_format);
+  printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
+              d->C_rowadr, d->C_colind, fp, float_format);
   printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
 
   // B sparse structure

src/engine/engine_support.c

@@ -1087,38 +1087,25 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
 void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
   int  nv = m->nv;
   const mjtNum* qLD = d->qLD;
-  const int* dofMadr = m->dof_Madr;
 
   mju_zero(res, nv);
 
   // res = L * vec
   for (int i=0; i < nv; i++) {
-    // simple: diagonal
+    // diagonal
-    if (m->dof_simplenum[i]) {
+    res[i] = vec[i];
-      res[i] = vec[i];
-    }
 
-    // regular: full multiplication
+    // non-simple: add off-diagonals
-    else {
+    if (!m->dof_simplenum[i]) {
-      // diagonal
+      int adr = d->C_rowadr[i];
-      res[i] += vec[i];
+      res[i] += mju_dotSparse(qLD+adr, vec, d->C_rownnz[i] - 1, d->C_colind+adr, /*flg_unc1=*/0);
-
-      // off-diagonal
-      int j = m->dof_parentid[i];
-      int adr = dofMadr[i] + 1;
-      while (j >= 0) {
-        res[i] += qLD[adr]*vec[j];
-
-        // advance to next element
-        j = m->dof_parentid[j];
-        adr++;
-      }
     }
   }
 
   // res *= sqrt(D)
   for (int i=0; i < nv; i++) {
-    res[i] *= mju_sqrt(qLD[dofMadr[i]]);
+    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    res[i] *= mju_sqrt(qLD[diag]);
   }
 }
 

src/engine/engine_vis_interact.c

@@ -556,7 +556,8 @@ 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++) {
-    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
   mj_jac(m, d, jac, NULL, selpos, sel);
   mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3);

test/benchmark/factorI_benchmark_test.cc

@@ -43,21 +43,23 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
   // allocate inputs and outputs
   mj_markStack(d);
 
-  // CSR matrices
+  // M: mass matrix in CSR format
-  mjtNum* Ms = mj_stackAllocNum(d, m->nC);
+  mjtNum* M = mj_stackAllocNum(d, m->nC);
-  mjtNum* LDs = mj_stackAllocNum(d, m->nC);
   for (int i=0; i < m->nC; i++) {
-    Ms[i] = d->qM[d->mapM2C[i]];
+    M[i] = d->qM[d->mapM2C[i]];
   }
 
+  // LDlegacy: legacy LD matrix (size nM)
+  mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
+
   // benchmark
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
     for (int i=0; i < kNumBenchmarkSteps; i++) {
       if (legacy) {
-        mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
+        mj_factorI(m, d, d->qM, LDlegacy, d->qLDiagInv);
       } else {
-        mju_copy(LDs, Ms, m->nC);
+        mju_copy(d->qLD, M, m->nC);
-        mj_factorIs(LDs, d->qLDiagInv, m->nv,
+        mj_factorIs(d->qLD, d->qLDiagInv, m->nv,
                     d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
       }
     }

test/benchmark/inertia_benchmark_test.cc

@@ -45,13 +45,15 @@ static void BM_solve(benchmark::State& state, SolveType type) {
   // allocate input and output vectors
   mj_markStack(d);
 
-  // make CSR matrix
+  // M: mass matrix in CSR format
-  mjtNum* Ms = mj_stackAllocNum(d, m->nC);
+  mjtNum* M = mj_stackAllocNum(d, m->nC);
-  mjtNum* LDs = mj_stackAllocNum(d, m->nC);
   for (int i=0; i < m->nC; i++) {
-    Ms[i] = d->qM[d->mapM2C[i]];
+    M[i] = d->qM[d->mapM2C[i]];
   }
 
+  // LDlegacy: legacy LD matrix (size nM)
+  mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
+
   // arbitrary input vector
   mjtNum *res = mj_stackAllocNum(d, m->nv);
   mjtNum *vec = mj_stackAllocNum(d, m->nv);
@@ -62,17 +64,18 @@ static void BM_solve(benchmark::State& state, SolveType type) {
   // benchmark
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
     for (int i=0; i < kNumBenchmarkSteps; i++) {
+      mju_copy(res, vec, m->nv);
       switch (type) {
         case SolveType::kLegacy:
-          mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
+          mj_factorI(m, d, d->qM, LDlegacy, d->qLDiagInv);
+          mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv);
           mj_solveM(m, d, res, vec, 1);
           break;
         case SolveType::kCsr:
-          mju_copy(LDs, Ms, m->nC);
+          mju_copy(d->qLD, M, m->nC);
-          mj_factorIs(LDs, d->qLDiagInv, m->nv,
+          mj_factorIs(d->qLD, d->qLDiagInv, m->nv,
                       d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
-          mju_copy(res, vec, m->nv);
+          mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1,
-          mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
                       d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
       }
     }

test/benchmark/solveLD_benchmark_test.cc

@@ -50,20 +50,21 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
     vec[i] = 0.2 + 0.3*i;
   }
 
-  // make CSR matrix
+  // make legacy matrix
-  mjtNum* LDs = mj_stackAllocNum(d, m->nC);
+  mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
+  mju_zero(LDlegacy, m->nM);
   for (int i=0; i < m->nC; i++) {
-    LDs[i] = d->qLD[d->mapM2C[i]];
+    LDlegacy[d->mapM2C[i]] = d->qLD[i];
   }
 
   // benchmark
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
     for (int i=0; i < kNumBenchmarkSteps; i++) {
+      mju_copy(res, vec, m->nv);
       if (featherstone) {
-        mj_solveM(m, d, res, vec, 1);
+        mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv);
       } else {
-        mju_copy(res, vec, m->nv);
+        mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1,
-        mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
                     d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
       }
     }

test/engine/engine_core_smooth_test.cc

@@ -449,9 +449,8 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
 
       // expect corresponding values to match
       for (int j=0; j < dofnum; j++) {
-        EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14));
+        EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12));
       }
-
       mju_free(res_i);
     }
 
@@ -475,21 +474,21 @@ TEST_F(CoreSmoothTest, FactorI) {
 
   // dense L matrix
   int nv = model->nv;
-  vector<mjtNum> Ldense(nv*nv);
+  vector<mjtNum> Ldense(nv*nv, 0);
-  mj_fullM(model, Ldense.data(), data->qLD);
+  mju_sparse2dense(Ldense.data(), data->qLD, nv, nv,
-  // clear upper triangle, set diagonal to 1
+                   data->C_rownnz, data->C_rowadr, data->C_colind);
   for (int i=0; i < nv; i++) {
-    for (int j=i; j < nv; j++) {
+    // set diagonal to 1
-      Ldense[i*nv+j] = i == j ? 1 : 0;
+    Ldense[i*nv+i] = 1;
-    }
   }
 
   // dense D matrix
   vector<mjtNum> Ddense(nv*nv);
-  mj_fullM(model, Ddense.data(), data->qLD);
+  mju_sparse2dense(Ddense.data(), data->qLD, nv, nv,
-  // clear everything but the diagonal
+                   data->C_rownnz, data->C_rowadr, data->C_colind);
   for (int i=0; i < nv; i++) {
     for (int j=0; j < nv; j++) {
+      // zero everything except the diagonal
       if (i != j) Ddense[i*nv+j] = 0;
     }
   }
@@ -521,20 +520,21 @@ TEST_F(CoreSmoothTest, SolveLDs) {
   mj_forward(m, d);
 
   int nv = m->nv;
+  int nM = m->nM;
   int nC = m->nC;
 
-  // copy LD into LDs: CSR format
+  // copy M into LD: Legacy format
-  vector<mjtNum> LDs(nC);
+  vector<mjtNum> LDlegacy(nM, 0);
   for (int i=0; i < nC; i++) {
-    LDs[i] = d->qLD[d->mapM2C[i]];
+    LDlegacy[d->mapM2C[i]] = d->qLD[i];
   }
 
   // compare LD and LDs densified matrices
   vector<mjtNum> LDdense(nv*nv);
-  mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv,
+  mju_sparse2dense(LDdense.data(), d->qLD, nv, nv,
                    d->C_rownnz, d->C_rowadr, d->C_colind);
   vector<mjtNum> LDdense2(nv*nv);
-  mj_fullM(m, LDdense2.data(), d->qLD);
+  mj_fullM(m, LDdense2.data(), LDlegacy.data());
 
   // expect lower triangles to match exactly
   for (int i=0; i < nv; i++) {
@@ -543,14 +543,14 @@ TEST_F(CoreSmoothTest, SolveLDs) {
     }
   }
 
-  // compare LD and LDs vector solve
+  // compare legacy and CSR LD vector solve
   vector<mjtNum> vec(nv);
   vector<mjtNum> vec2(nv);
   for (int i=0; i < nv; i++) vec[i] = vec2[i] = 20 + 30*i;
   for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0;
 
-  mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv);
+  mj_solveLD(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
-  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, 1,
+  mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
   // expect vectors to match up to floating point precision
@@ -572,12 +572,13 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
   mj_forward(m, d);
 
   int nv = m->nv;
+  int nM = m->nM;
   int nC = m->nC;
 
-  // copy LD into LDs: CSR format
+  // copy LD into LDlegacy: Legacy format
-  vector<mjtNum> LDs(nC);
+  vector<mjtNum> LDlegacy(nM, 0);
   for (int i=0; i < nC; i++) {
-    LDs[i] = d->qLD[d->mapM2C[i]];
+    LDlegacy[d->mapM2C[i]] = d->qLD[i];
   }
 
   // compare n LD and LDs vector solve
@@ -587,8 +588,8 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
   for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i;
   for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0;
 
-  mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv);
+  mj_solveLD(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
-  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
+  mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
   // expect vectors to match up to floating point precision
@@ -609,19 +610,12 @@ TEST_F(CoreSmoothTest, SolveM2) {
   mjData* d = mj_makeData(m);
   mj_forward(m, d);
 
-  int nv = m->nv;
-  int nC = m->nC;
-
-  // copy LD into LDs: CSR format
-  vector<mjtNum> LDs(nC);
-  for (int i=0; i < nC; i++) {
-    LDs[i] = d->qLD[d->mapM2C[i]];
-  }
-
   // inverse square root of D from inertia LDL decomposition
+  int nv = m->nv;
   vector<mjtNum> sqrtInvD(nv);
   for (int i=0; i < nv; i++) {
-    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
 
   // compare full solve and half solve
@@ -633,7 +627,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
   vector<mjtNum> res(nv*n);
 
   mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
-  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
+  mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
   // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
@@ -655,25 +649,32 @@ TEST_F(CoreSmoothTest, FactorIs) {
   mjData* d = mj_makeData(m);
   mj_forward(m, d);
 
-  int nC = m->nC, nv = m->nv;
+  int nC = m->nC, nM = m->nM, nv = m->nv;
 
-  // copy qM into LDs, qLD into qLDexpected: CSR format
+  // copy qM into into qLDlegacy and factorize
-  vector<mjtNum> qLDsExpected(nC);
+  vector<mjtNum> qLDlegacy(nM);
-  vector<mjtNum> qLDs(nC);
+  mj_factorI(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv);
+
+  // copy qLDlegacy into qLDexpected: CSR format
+  vector<mjtNum> qLDexpected(nC);
   for (int i=0; i < nC; i++) {
-    int index = d->mapM2C[i];
+    qLDexpected[i] = qLDlegacy[d->mapM2C[i]];  // mj_factorIs is in-place
-    qLDs[i] = d->qM[index];  // mj_factorIs is in-place
+  }
-    qLDsExpected[i] = d->qLD[index];
+
+  // copy qM into qLD: CSR format
+  vector<mjtNum> qLD(nC);
+  for (int i=0; i < nC; i++) {
+    qLD[i] = d->qM[d->mapM2C[i]];  // mj_factorIs is in-place
   }
 
   vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
   vector<mjtNum> qLDiagInv(nv, 0);
 
-  mj_factorIs(qLDs.data(), qLDiagInv.data(), nv,
+  mj_factorIs(qLD.data(), qLDiagInv.data(), nv,
               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
   // expect outputs to match to floating point precision
-  EXPECT_THAT(qLDs, Pointwise(DoubleNear(1e-12), qLDsExpected));
+  EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));
   EXPECT_THAT(qLDiagInv, Pointwise(DoubleNear(1e-12), qLDiagInvExpected));
 
   /* uncomment for debugging

test/engine/engine_derivative_test.cc

@@ -436,7 +436,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
       Ac[i*nv + i]  = -m->jnt_stiffness[i];
       Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
     }
-    mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv);
+    mj_solveLDs(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
     // A = [dt*Ac; Ac]
     mju_transpose(A, Ac, 2*nv, nv);
@@ -463,7 +464,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
     mjtNum *BcT = mj_stackAllocNum(d, nv*nu);
     mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
                      d->moment_rowadr, d->moment_colind);
-    mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv);
+    mj_solveLDs(Bc, d->qH, d->qHDiagInv, nv, nu,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     mju_transpose(BcT, Bc, nu, nv);
     mju_scl(B, BcT, dt*dt, nu*nv);
     mju_scl(B+nu*nv, BcT, dt, nu*nv);

test/pipeline_test.cc

@@ -61,7 +61,8 @@ TEST_F(PipelineTest, SparseDenseEquivalent) {
     std::vector<mjtNum> qacc_sparse = AsVector(data->qacc, model->nv);
 
     // expect accelerations to be insignificantly different
-    EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse));
+    EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse))
+        << "failed equivalence for solver=" << solver;
   }
 
   mj_deleteData(data);