chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Add mju_gather and mju_scatter, private engine functions.

Browse Source 
PiperOrigin-RevId: 750245294
Change-Id: I3a14ccdd55a324d3fe206e388fb2a513772f5f5e
This commit is contained in:
Yuval Tassa
2025-04-22 10:32:42 -07:00
committed by Copybara-Service
parent 99490163df
commit c2ac0d724e
11 changed files with 58 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+4 -4
View File
@@ -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
+7 -15
View File
@@ -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
+3 -4
View File
@@ -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
+2 -4
View File
@@ -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
+4 -6
View File
@@ -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
+18
View File
@@ -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
+6
View File
@@ -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
+2 -3
View File
@@ -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++) {
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);
test/benchmark/inertia_benchmark_test.cc
+2 -3
View File
@@ -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++) {
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);
test/benchmark/solveLD_benchmark_test.cc
+3 -4
View File
@@ -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++) {
LDlegacy[d->mapM2M[i]] = d->qLD[i];
}
mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nM);
// benchmark
while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
test/engine/engine_core_smooth_test.cc
+7 -12
View File
@@ -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++) {
LDlegacy[d->mapM2M[i]] = d->qLD[i];
}
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM);
// 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++) {
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<mjtNum> 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<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
vector<mjtNum> qLDiagInv(nv, 0);