From 5ee065483021db684c9e8c3ffff75762c82823e0 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 7 Feb 2025 10:20:41 -0800 Subject: [PATCH] Roll back recent change to mjData.qLD until some issues are resolved. PiperOrigin-RevId: 724389161 Change-Id: I1b30ab0950ab5c5e611b1c6fc3bf3f42f964939c --- doc/APIreference/functions.rst | 3 +- doc/APIreference/functions_override.rst | 5 -- doc/includes/references.h | 4 +- include/mujoco/mjdata.h | 4 +- include/mujoco/mjxmacro.h | 4 +- introspect/structs.py | 4 +- mjx/mujoco/mjx/_src/io.py | 2 +- mjx/mujoco/mjx/_src/smooth_test.py | 5 +- mjx/mujoco/mjx/_src/types.py | 2 +- src/engine/engine_core_constraint.c | 12 +++- src/engine/engine_core_smooth.c | 41 +++++++---- src/engine/engine_core_smooth.h | 3 +- src/engine/engine_forward.c | 34 ++++----- src/engine/engine_print.c | 3 +- src/engine/engine_support.c | 29 +++++--- src/engine/engine_vis_interact.c | 3 +- test/benchmark/factorI_benchmark_test.cc | 16 ++--- test/benchmark/inertia_benchmark_test.cc | 21 +++--- test/benchmark/solveLD_benchmark_test.cc | 13 ++-- test/engine/engine_core_smooth_test.cc | 87 ++++++++++++------------ test/engine/engine_derivative_test.cc | 6 +- test/pipeline_test.cc | 3 +- 22 files changed, 153 insertions(+), 151 deletions(-) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 64c50073..d7b8db51 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -423,8 +423,7 @@ 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. -|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``. +Convert sparse inertia matrix M into full (i.e. dense) matrix. .. _mj_mulM: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 03c1e7fb..9d425623 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -233,11 +233,6 @@ 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 `__. -.. _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 diff --git a/doc/includes/references.h b/doc/includes/references.h index 693a2a99..dcb3f22d 100644 --- 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 index c573a2c7..3c4d889e 100644 --- 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 index 88b802b8..50f0f8e2 100644 --- a/include/mujoco/mjxmacro.h +++ b/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, 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 ) \ @@ -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, 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/introspect/structs.py b/introspect/structs.py index 605fcd4b..abe63628 100644 --- a/introspect/structs.py +++ b/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=('nC',), + array_extent=('nM',), ), 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=('nC',), + array_extent=('nM',), ), StructFieldDecl( name='qHDiagInv', diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 24704b18..6acef205 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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.nC, float), + '_qLD_sparse': (m.nM, float), '_qLDiagInv_sparse': (m.nv, float), } diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index bf9176ba..11cf85fd 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -91,10 +91,7 @@ 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)) - 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, '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') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index a7fd3617..3a7f3f9b 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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 (nC,) + _qLD_sparse: qLD in sparse representation (nM,) _qLDiagInv_sparse: qLDiagInv in sparse representation (nv,) """ # fmt: skip # constant sizes: diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index af773edd..df37ab30 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2131,8 +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; - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); } // sparse @@ -2239,6 +2238,13 @@ 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]; @@ -2252,7 +2258,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { } int j = B_colind[i]; int adrC = d->C_rowadr[j]; - mju_addToSclSparseInc(B + adrB, d->qLD + adrC, + mju_addToSclSparseInc(B + adrB, qLD + adrC, nnzB, B_colind + adrB, d->C_rownnz[j]-1, d->C_colind + adrC, -b); } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 9f73b092..c7bffa00 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1465,11 +1465,7 @@ 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; - 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); + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); TM_ADD(mjTIMER_POS_INERTIA); } @@ -1713,20 +1709,18 @@ 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_solveLDs(x, d->qLD, d->qLDiagInv, m->nv, n, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mj_solveLD(m, x, n, d->qLD, d->qLDiagInv); } // 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, const mjData* d, mjtNum* restrict x, int island) { +void mj_solveM_island(const mjModel* m, 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_solveLDs(x, qLD, qLDiagInv, m->nv, 1, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mj_solveLD(m, x, 1, qLD, qLDiagInv); return; } @@ -1736,6 +1730,14 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int 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]; @@ -1749,7 +1751,7 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=end-1; adr >= start; adr--) { - x[islandind[colind[adr]]] -= qLD[adr] * x_k; + x[islandind[colind[adr]]] -= qLDs[adr] * x_k; } } } @@ -1771,9 +1773,11 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=end-1; adr >= start; adr--) { - x[k] -= x[islandind[colind[adr]]] * qLD[adr]; + x[k] -= x[islandind[colind[adr]]] * qLDs[adr]; } } + + mj_freeStack(d); } @@ -1781,18 +1785,23 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int // 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 @@ -1822,6 +1831,8 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, x[i+offset] *= invD_i; } } + + mj_freeStack(d); } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index a710416c..ec8d168d 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -71,8 +71,9 @@ 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, const mjData* d, mjtNum* x, int island); +MJAPI void mj_solveM_island(const mjModel* m, 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, diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 34fd4da9..0ece01ac 100644 --- 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); @@ -794,23 +794,22 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { // damping: integrate implicitly else { if (!skipfactor) { - // qH = M + h*diag(B) - for (int i=0; i < nC; i++) { - d->qH[i] = d->qM[d->mapM2C[i]]; - } + mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum); + + // MhB = M + h*diag(B) + mju_copy(MhB, d->qM, nM); 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]; + MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i]; } - // factorize in-place - mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + // factor + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); } // solve mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); mju_copy(qacc, qfrc, m->nv); - mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); } // advance state and time @@ -939,7 +938,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); @@ -986,20 +985,13 @@ 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); - // copy into qH - 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); + // factorize + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); } // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc mju_copy(qacc, qfrc, nv); - mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); - + mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); } else { mjERROR("integrator must be implicit or implicitfast"); } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 729a2e54..377234b1 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1127,8 +1127,7 @@ 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); + printInertia("QLD", d->qLD, m, fp, float_format); printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); // B sparse structure diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 32068a77..e5237a61 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1087,25 +1087,38 @@ 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++) { - // diagonal - res[i] = vec[i]; + // simple: diagonal + if (m->dof_simplenum[i]) { + res[i] = vec[i]; + } - // 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); + // regular: full multiplication + else { + // diagonal + res[i] += vec[i]; + + // 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++) { - int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; - res[i] *= mju_sqrt(qLD[diag]); + res[i] *= mju_sqrt(qLD[dofMadr[i]]); } } diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c index 3366e731..6d1acf19 100644 --- a/src/engine/engine_vis_interact.c +++ b/src/engine/engine_vis_interact.c @@ -556,8 +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; - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); } mj_jac(m, d, jac, NULL, selpos, sel); mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3); diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc index c1dedbd9..55118a99 100644 --- a/test/benchmark/factorI_benchmark_test.cc +++ b/test/benchmark/factorI_benchmark_test.cc @@ -43,23 +43,21 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { // allocate inputs and outputs mj_markStack(d); - // M: mass matrix in CSR format - mjtNum* M = mj_stackAllocNum(d, m->nC); + // CSR matrices + mjtNum* Ms = mj_stackAllocNum(d, m->nC); + mjtNum* LDs = mj_stackAllocNum(d, m->nC); for (int i=0; i < m->nC; i++) { - M[i] = d->qM[d->mapM2C[i]]; + Ms[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, LDlegacy, d->qLDiagInv); + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); } else { - mju_copy(d->qLD, M, m->nC); - mj_factorIs(d->qLD, d->qLDiagInv, m->nv, + mju_copy(LDs, Ms, m->nC); + mj_factorIs(LDs, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index df230fc8..7ddc3468 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -45,15 +45,13 @@ static void BM_solve(benchmark::State& state, SolveType type) { // allocate input and output vectors mj_markStack(d); - // M: mass matrix in CSR format - mjtNum* M = mj_stackAllocNum(d, m->nC); + // make CSR matrix + mjtNum* Ms = mj_stackAllocNum(d, m->nC); + mjtNum* LDs = mj_stackAllocNum(d, m->nC); for (int i=0; i < m->nC; i++) { - M[i] = d->qM[d->mapM2C[i]]; + Ms[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); @@ -64,18 +62,17 @@ 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, LDlegacy, d->qLDiagInv); - mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv); + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); mj_solveM(m, d, res, vec, 1); break; case SolveType::kCsr: - mju_copy(d->qLD, M, m->nC); - mj_factorIs(d->qLD, d->qLDiagInv, m->nv, + mju_copy(LDs, Ms, m->nC); + mj_factorIs(LDs, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); - mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1, + mju_copy(res, vec, m->nv); + mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 434d5d2e..18ebe69b 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -50,21 +50,20 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { vec[i] = 0.2 + 0.3*i; } - // make legacy matrix - mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); - mju_zero(LDlegacy, m->nM); + // make CSR matrix + mjtNum* LDs = mj_stackAllocNum(d, m->nC); for (int i=0; i < m->nC; i++) { - LDlegacy[d->mapM2C[i]] = d->qLD[i]; + LDs[i] = d->qLD[d->mapM2C[i]]; } // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { for (int i=0; i < kNumBenchmarkSteps; i++) { - mju_copy(res, vec, m->nv); if (featherstone) { - mj_solveLD(m, res, 1, LDlegacy, d->qLDiagInv); + mj_solveM(m, d, res, vec, 1); } else { - mj_solveLDs(res, d->qLD, d->qLDiagInv, m->nv, 1, + mju_copy(res, vec, m->nv); + mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 05e9348e..47e1ff49 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -449,8 +449,9 @@ 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-12)); + EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14)); } + mju_free(res_i); } @@ -474,21 +475,21 @@ TEST_F(CoreSmoothTest, FactorI) { // dense L matrix int nv = model->nv; - vector Ldense(nv*nv, 0); - mju_sparse2dense(Ldense.data(), data->qLD, nv, nv, - data->C_rownnz, data->C_rowadr, data->C_colind); + vector Ldense(nv*nv); + mj_fullM(model, Ldense.data(), data->qLD); + // clear upper triangle, set diagonal to 1 for (int i=0; i < nv; i++) { - // set diagonal to 1 - Ldense[i*nv+i] = 1; + for (int j=i; j < nv; j++) { + Ldense[i*nv+j] = i == j ? 1 : 0; + } } // dense D matrix vector Ddense(nv*nv); - mju_sparse2dense(Ddense.data(), data->qLD, nv, nv, - data->C_rownnz, data->C_rowadr, data->C_colind); + mj_fullM(model, Ddense.data(), data->qLD); + // clear everything but the diagonal 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; } } @@ -520,21 +521,20 @@ TEST_F(CoreSmoothTest, SolveLDs) { mj_forward(m, d); int nv = m->nv; - int nM = m->nM; int nC = m->nC; - // copy M into LD: Legacy format - vector LDlegacy(nM, 0); + // copy LD into LDs: CSR format + vector LDs(nC); for (int i=0; i < nC; i++) { - LDlegacy[d->mapM2C[i]] = d->qLD[i]; + LDs[i] = d->qLD[d->mapM2C[i]]; } // compare LD and LDs densified matrices vector LDdense(nv*nv); - mju_sparse2dense(LDdense.data(), d->qLD, nv, nv, + mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv, d->C_rownnz, d->C_rowadr, d->C_colind); vector LDdense2(nv*nv); - mj_fullM(m, LDdense2.data(), LDlegacy.data()); + mj_fullM(m, LDdense2.data(), d->qLD); // expect lower triangles to match exactly for (int i=0; i < nv; i++) { @@ -543,14 +543,14 @@ TEST_F(CoreSmoothTest, SolveLDs) { } } - // compare legacy and CSR LD vector solve + // compare LD and LDs vector solve vector vec(nv); vector 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, LDlegacy.data(), d->qLDiagInv); - mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, 1, + mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv); + mj_solveLDs(vec2.data(), LDs.data(), 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,13 +572,12 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { mj_forward(m, d); int nv = m->nv; - int nM = m->nM; int nC = m->nC; - // copy LD into LDlegacy: Legacy format - vector LDlegacy(nM, 0); + // copy LD into LDs: CSR format + vector LDs(nC); for (int i=0; i < nC; i++) { - LDlegacy[d->mapM2C[i]] = d->qLD[i]; + LDs[i] = d->qLD[d->mapM2C[i]]; } // compare n LD and LDs vector solve @@ -588,8 +587,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, LDlegacy.data(), d->qLDiagInv); - mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n, + mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv); + mj_solveLDs(vec2.data(), LDs.data(), 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 @@ -610,12 +609,19 @@ TEST_F(CoreSmoothTest, SolveM2) { mjData* d = mj_makeData(m); mj_forward(m, d); - // inverse square root of D from inertia LDL decomposition int nv = m->nv; + int nC = m->nC; + + // copy LD into LDs: CSR format + vector 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 vector sqrtInvD(nv); for (int i=0; i < nv; i++) { - int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1; - sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); } // compare full solve and half solve @@ -627,7 +633,7 @@ TEST_F(CoreSmoothTest, SolveM2) { vector res(nv*n); mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n); - mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n, + mj_solveLDs(vec2.data(), LDs.data(), 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) @@ -649,32 +655,25 @@ TEST_F(CoreSmoothTest, FactorIs) { mjData* d = mj_makeData(m); mj_forward(m, d); - int nC = m->nC, nM = m->nM, nv = m->nv; + int nC = m->nC, nv = m->nv; - // copy qM into into qLDlegacy and factorize - vector qLDlegacy(nM); - mj_factorI(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv); - - // copy qLDlegacy into qLDexpected: CSR format - vector qLDexpected(nC); + // copy qM into LDs, qLD into qLDexpected: CSR format + vector qLDsExpected(nC); + vector qLDs(nC); for (int i=0; i < nC; i++) { - qLDexpected[i] = qLDlegacy[d->mapM2C[i]]; // mj_factorIs is in-place - } - - // copy qM into qLD: CSR format - vector qLD(nC); - for (int i=0; i < nC; i++) { - qLD[i] = d->qM[d->mapM2C[i]]; // mj_factorIs is in-place + int index = d->mapM2C[i]; + qLDs[i] = d->qM[index]; // mj_factorIs is in-place + qLDsExpected[i] = d->qLD[index]; } vector qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); vector qLDiagInv(nv, 0); - mj_factorIs(qLD.data(), qLDiagInv.data(), nv, + mj_factorIs(qLDs.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(qLD, Pointwise(DoubleNear(1e-12), qLDexpected)); + EXPECT_THAT(qLDs, Pointwise(DoubleNear(1e-12), qLDsExpected)); EXPECT_THAT(qLDiagInv, Pointwise(DoubleNear(1e-12), qLDiagInvExpected)); /* uncomment for debugging diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 7aaf6ad3..1db7731a 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -436,8 +436,7 @@ 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_solveLDs(Ac, d->qH, d->qHDiagInv, nv, 2*nv, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv); // A = [dt*Ac; Ac] mju_transpose(A, Ac, 2*nv, nv); @@ -464,8 +463,7 @@ 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_solveLDs(Bc, d->qH, d->qHDiagInv, nv, nu, - d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv); mju_transpose(BcT, Bc, nu, nv); mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv); diff --git a/test/pipeline_test.cc b/test/pipeline_test.cc index 6f1eb737..f9c40713 100644 --- a/test/pipeline_test.cc +++ b/test/pipeline_test.cc @@ -61,8 +61,7 @@ TEST_F(PipelineTest, SparseDenseEquivalent) { std::vector qacc_sparse = AsVector(data->qacc, model->nv); // expect accelerations to be insignificantly different - EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse)) - << "failed equivalence for solver=" << solver; + EXPECT_THAT(qacc_dense, Pointwise(DoubleNear(tol), qacc_sparse)); } mj_deleteData(data);