Add mju_gather and mju_scatter, private engine functions.

PiperOrigin-RevId: 750245294
Change-Id: I3a14ccdd55a324d3fe206e388fb2a513772f5f5e

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++) {
+  // gather LD <- M (legacy to CSR) and factorize in-place
-    d->qLD[i] = d->qM[d->mapM2M[i]];
+  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);
 }
 

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++) {
+      mju_gather(d->qH, d->qM, d->mapM2M, nM);
-        d->qH[i] = d->qM[d->mapM2M[i]];
-      }
       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
+      // gather qLU <- qM (lower to full)
-      for (int i=0; i < nD; i++) {
+      mju_gather(d->qLU, d->qM, d->mapM2D, nD);
-        d->qLU[i] = d->qM[d->mapM2D[i]];
-      }
 
       // 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++) {
+      mju_gather(MhB, d->qDeriv, d->mapD2M, nM);
-        MhB[i] = d->qDeriv[d->mapD2M[i]];
-      }
 
       // set MhB = M - dt*qDeriv
       mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
-      // copy into qH
+      // gather qH <- MhB (legacy to CSR)
-      for (int i=0; i < nM; i++) {
+      mju_gather(d->qH, MhB, d->mapM2M, nM);
-        d->qH[i] = MhB[d->mapM2M[i]];
-      }
 
       // factorize in-place
       mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);

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
+    // gather qLU <- qM (lower to full)
-    for (int i=0; i < nD; i++) {
+    mju_gather(d->qLU, d->qM, d->mapM2D, nD);
-      d->qLU[i] = d->qM[d->mapM2D[i]];
-    }
 
     // set qLU = qM - dt*qDeriv
     mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);

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
+    // gather C <- qM (legacy to CSR)
-    for (int i=0; i < m->nC; i++) {
+    mju_gather(ctx->C, d->qM, d->mapM2C, m->nC);
-      ctx->C[i] = d->qM[d->mapM2C[i]];
-    }
 
     // initialize Hessian rowadr, rownnz
     mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv,

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++) {
+    mju_gather(C, d->qM, d->mapM2C, nC);
-      C[i] = d->qM[d->mapM2C[i]];
-    }
 
-    mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C,
+    // add to dst
-                  d->C_rownnz, d->C_rowadr, d->C_colind);
+    mj_addMSparse(m, d, dst, rownnz, rowadr, colind, C, d->C_rownnz, d->C_rowadr, d->C_colind);
     mj_freeStack(d);
   }
 

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++) {

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);
 

test/benchmark/factorI_benchmark_test.cc

@@ -19,6 +19,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
 #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++) {
+  mju_gather(M, d->qM, d->mapM2M, m->nM);
-    M[i] = d->qM[d->mapM2M[i]];
-  }
 
   // LDlegacy: legacy LD matrix (size nM)
   mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);

test/benchmark/inertia_benchmark_test.cc

@@ -20,6 +20,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
 #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++) {
+  mju_gather(M, d->qM, d->mapM2M, m->nM);
-    M[i] = d->qM[d->mapM2M[i]];
-  }
 
   // LDlegacy: legacy LD matrix (size nM)
   mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);

test/benchmark/solveLD_benchmark_test.cc

@@ -19,6 +19,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
 #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++) {
+  mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nM);
-    LDlegacy[d->mapM2M[i]] = d->qLD[i];
-  }
 
   // benchmark
   while (state.KeepRunningBatch(kNumBenchmarkSteps)) {

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 <algorithm>
@@ -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<mjtNum> LDlegacy(nM);
-  for (int i=0; i < nM; i++) {
+  mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM);
-    LDlegacy[d->mapM2M[i]] = d->qLD[i];
-  }
 
   // compare LD and LDs densified matrices
   vector<mjtNum> 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<mjtNum> LDlegacy(nM);
-  for (int i=0; i < nM; i++) {
+  mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM);
-    LDlegacy[d->mapM2M[i]] = d->qLD[i];
-  }
 
   // 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<mjtNum> qLD(nM);
-  for (int i=0; i < nM; i++) {
+  mju_gather(qLD.data(), d->qM, d->mapM2M, nM);
-    qLD[i] = d->qM[d->mapM2M[i]];  // mj_factorI is in-place
-  }
 
   vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
   vector<mjtNum> qLDiagInv(nv, 0);