From c2ac0d724e866045e89eee954cfaa22e5c52b4a5 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 22 Apr 2025 10:32:42 -0700 Subject: [PATCH] Add `mju_gather` and `mju_scatter`, private engine functions. PiperOrigin-RevId: 750245294 Change-Id: I3a14ccdd55a324d3fe206e388fb2a513772f5f5e --- src/engine/engine_core_smooth.c | 8 ++++---- src/engine/engine_forward.c | 22 +++++++--------------- src/engine/engine_inverse.c | 7 +++---- src/engine/engine_solver.c | 6 ++---- src/engine/engine_support.c | 10 ++++------ src/engine/engine_util_misc.c | 18 ++++++++++++++++++ src/engine/engine_util_misc.h | 6 ++++++ test/benchmark/factorI_benchmark_test.cc | 5 ++--- test/benchmark/inertia_benchmark_test.cc | 5 ++--- test/benchmark/solveLD_benchmark_test.cc | 7 +++---- test/engine/engine_core_smooth_test.cc | 19 +++++++------------ 11 files changed, 58 insertions(+), 55 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 0bb97dc7..09ea01c2 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1646,11 +1646,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 nM = m->nM; - for (int i=0; i < nM; i++) { - d->qLD[i] = d->qM[d->mapM2M[i]]; - } + + // gather LD <- M (legacy to CSR) and factorize in-place + mju_gather(d->qLD, d->qM, d->mapM2M, m->nM); 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); } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 5771db32..42520046 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -840,9 +840,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { else { if (!skipfactor) { // qH = M + h*diag(B) - for (int i=0; i < nM; i++) { - d->qH[i] = d->qM[d->mapM2M[i]]; - } + mju_gather(d->qH, d->qM, d->mapM2M, nM); for (int i=0; i < nv; i++) { d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i]; } @@ -999,10 +997,8 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // compute analytical derivative qDeriv mjd_smooth_vel(m, d, /* flg_bias = */ 1); - // set qLU = qM - for (int i=0; i < nD; i++) { - d->qLU[i] = d->qM[d->mapM2D[i]]; - } + // gather qLU <- qM (lower to full) + mju_gather(d->qLU, d->qM, d->mapM2D, nD); // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); @@ -1022,19 +1018,15 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // compute analytical derivative qDeriv; skip rne derivative mjd_smooth_vel(m, d, /* flg_bias = */ 0); - // modified mass matrix MhB = qDeriv[Lower] + // modified mass matrix: gather MhB <- qDeriv (full to lower) mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum); - for (int i=0; i < nM; i++) { - MhB[i] = d->qDeriv[d->mapD2M[i]]; - } + mju_gather(MhB, d->qDeriv, d->mapD2M, nM); // set MhB = M - dt*qDeriv mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); - // copy into qH - for (int i=0; i < nM; i++) { - d->qH[i] = MhB[d->mapM2M[i]]; - } + // gather qH <- MhB (legacy to CSR) + mju_gather(d->qH, MhB, d->mapM2M, nM); // factorize in-place mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); diff --git a/src/engine/engine_inverse.c b/src/engine/engine_inverse.c index c4042401..57005f1e 100644 --- a/src/engine/engine_inverse.c +++ b/src/engine/engine_inverse.c @@ -31,6 +31,7 @@ #include "engine/engine_support.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" +#include "engine/engine_util_misc.h" #include "engine/engine_util_sparse.h" // position-dependent computations @@ -115,10 +116,8 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) { // compute qDeriv mjd_smooth_vel(m, d, /* flg_bias = */ 1); - // set qLU = qM - for (int i=0; i < nD; i++) { - d->qLU[i] = d->qM[d->mapM2D[i]]; - } + // gather qLU <- qM (lower to full) + mju_gather(d->qLU, d->qM, d->mapM2D, nD); // set qLU = qM - dt*qDeriv mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index f23d8bf5..8e83ace5 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -1394,10 +1394,8 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { // sparse if (mj_isSparse(m)) { - // copy values of reduced sparse inertia matrix C - for (int i=0; i < m->nC; i++) { - ctx->C[i] = d->qM[d->mapM2C[i]]; - } + // gather C <- qM (legacy to CSR) + mju_gather(ctx->C, d->qM, d->mapM2C, m->nC); // initialize Hessian rowadr, rownnz mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv, diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index b7fc4d7c..b7e2699d 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1130,14 +1130,12 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int nC = m->nC; mj_markStack(d); - // create reduced sparse inertia matrix C + // gather C <- qM (legacy to CSR) mjtNum* C = mjSTACKALLOC(d, nC, mjtNum); - for (int i=0; i < nC; i++) { - C[i] = d->qM[d->mapM2C[i]]; - } + mju_gather(C, d->qM, d->mapM2C, nC); - mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, - d->C_rownnz, d->C_rowadr, d->C_colind); + // add to dst + mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, d->C_rownnz, d->C_rowadr, d->C_colind); mj_freeStack(d); } diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index 528b8f21..bec28aa3 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -1390,6 +1390,24 @@ void mju_n2d(double* res, const mjtNum* vec, int n) { +// gather +void mju_gather(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[i] = vec[ind[i]]; + } +} + + + +// scatter +void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + for (int i=0; i < n; i++) { + res[ind[i]] = vec[i]; + } +} + + + // insertion sort, increasing order void mju_insertionSort(mjtNum* list, int n) { for (int i=1; i < n; i++) { diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index e266f814..de10f907 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -156,6 +156,12 @@ MJAPI void mju_d2n(mjtNum* res, const double* vec, int n); // convert from mjtNum to double MJAPI void mju_n2d(double* res, const mjtNum* vec, int n); +// gather +MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n); + +// scatter +MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n); + // insertion sort, increasing order MJAPI void mju_insertionSort(mjtNum* list, int n); diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc index dd73d3da..f2feea1e 100644 --- a/test/benchmark/factorI_benchmark_test.cc +++ b/test/benchmark/factorI_benchmark_test.cc @@ -19,6 +19,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -45,9 +46,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { // M: mass matrix in CSR format mjtNum* M = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - M[i] = d->qM[d->mapM2M[i]]; - } + mju_gather(M, d->qM, d->mapM2M, m->nM); // LDlegacy: legacy LD matrix (size nM) mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index e787fa24..3bdeddfd 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -20,6 +20,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -47,9 +48,7 @@ static void BM_solve(benchmark::State& state, SolveType type) { // M: mass matrix in CSR format mjtNum* M = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - M[i] = d->qM[d->mapM2M[i]]; - } + mju_gather(M, d->qM, d->mapM2M, m->nM); // LDlegacy: legacy LD matrix (size nM) mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 64276ba9..6204425f 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -19,6 +19,7 @@ #include #include #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "test/fixture.h" namespace mujoco { @@ -50,11 +51,9 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { vec[i] = 0.2 + 0.3*i; } - // make legacy matrix + // scatter into legacy matrix mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); - for (int i=0; i < m->nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nM); // benchmark while (state.KeepRunningBatch(kNumBenchmarkSteps)) { diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 1d383922..1c71475a 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -15,6 +15,7 @@ // Tests for engine/engine_core_smooth.c. #include "src/engine/engine_core_smooth.h" +#include "src/engine/engine_util_misc.h" #include "src/engine/engine_util_sparse.h" #include @@ -754,11 +755,9 @@ TEST_F(CoreSmoothTest, SolveLDs) { int nv = m->nv; int nM = m->nM; - // copy M into LD: Legacy format + // scatter M into LD: Legacy format vector LDlegacy(nM); - for (int i=0; i < nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM); // compare LD and LDs densified matrices vector LDdense(nv*nv); @@ -805,11 +804,9 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { int nv = m->nv; int nM = m->nM; - // copy LD into LDlegacy: Legacy format + // scatter LD into LDlegacy: Legacy format vector LDlegacy(nM); - for (int i=0; i < nM; i++) { - LDlegacy[d->mapM2M[i]] = d->qLD[i]; - } + mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM); // compare n LD and LDs vector solve int n = 3; @@ -891,11 +888,9 @@ TEST_F(CoreSmoothTest, FactorIs) { qLDexpected[i] = qLDlegacy[d->mapM2M[i]]; } - // copy qM into qLD: CSR format + // gather qM into qLD: CSR format vector qLD(nM); - for (int i=0; i < nM; i++) { - qLD[i] = d->qM[d->mapM2M[i]]; // mj_factorI is in-place - } + mju_gather(qLD.data(), d->qM, d->mapM2M, nM); vector qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); vector qLDiagInv(nv, 0);