chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 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
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);