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

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

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

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,

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

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