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

CSR implementation of mj_solveLD.

Browse Source 
PiperOrigin-RevId: 704316200
Change-Id: Ibaff0284e40b3ebbe43bb489b6211ce739270e27
This commit is contained in:
Yuval Tassa
2024-12-09 09:33:20 -08:00
committed by Copybara-Service
parent 537fc2ff45
commit 1c4c7b012c
4 changed files with 152 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+32
View File
@@ -1575,6 +1575,38 @@ void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
}
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// like mj_solveLD, but using the CSR representation of L
void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
const int* rownnz, const int* rowadr, const int* diag, const int* colind) {
// x <- L^-T x
for (int i=nv-2; i >= 0; i--) {
int d1 = diag[i] + 1;
int nnz = rownnz[i] - d1;
if (nnz > 0) {
int adr = rowadr[i] + d1;
x[i] -= mju_dotSparse(qLDs+adr, x, nnz, colind+adr, /*flg_unc1=*/0);
}
}
// x(i) /= D(i,i)
for (int i=0; i < nv; i++) {
x[i] *= qLDiagInv[i];
}
// x <- L^-1 x
for (int i=1; i < nv; i++) {
int d = diag[i];
if (d > 0) {
int adr = rowadr[i];
x[i] -= mju_dotSparse(qLDs+adr, x, d, colind+adr, /*flg_unc1=*/0);
}
}
}
// sparse backsubstitution: x = inv(L'*D*L)*y
// use factorization in d
void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
src/engine/engine_core_smooth.h
+6 -1
View File
@@ -55,10 +55,15 @@ MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD,
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
MJAPI void mj_factorM(const mjModel* m, mjData* d);
// sparse backsubstitution: x = inv(L'*D*L)*y
// sparse backsubstitution: x = inv(L'*D*L)*x
MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
const mjtNum* qLD, const mjtNum* qLDiagInv);
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// like mj_solveLD, but using the CSR representation of L
MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
const int* rownnz, const int* rowadr, const int* diag, const int* colind);
// sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d
MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
test/benchmark/engine_core_smooth_benchmark_test.cc
+37 -85
View File
@@ -19,85 +19,28 @@
#include <absl/base/attributes.h>
#include <mujoco/mjdata.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_core_smooth.h"
#include "test/fixture.h"
namespace mujoco {
namespace {
// number of steps to roll out before benhmarking
// number of steps to roll out before benchmarking
static const int kNumWarmupSteps = 200;
// number of steps to benchmark
static const int kNumBenchmarkSteps = 50;
// ----------------------------- old functions --------------------------------
void ABSL_ATTRIBUTE_NOINLINE solveLD_baseline(const mjModel* m, mjtNum* x,
const mjtNum* y,
const mjtNum* qLD,
const mjtNum* qLDiagInv) {
mjtNum tmp;
// local copies of key variables
int* dof_Madr = m->dof_Madr;
int* dof_parentid = m->dof_parentid;
int nv = m->nv;
// x = y
if (x != y) {
mju_copy(x, y, nv);
}
// x <- inv(L') * x; skip simple, exploit sparsity of input vector
for (int i=nv-1; i >= 0; i--) {
if (!m->dof_simplenum[i] && (tmp = x[i])) {
// init
int Madr_ij = dof_Madr[i]+1;
int j = dof_parentid[i];
// traverse ancestors backwards
while (j >= 0) {
x[j] -= qLD[Madr_ij++]*tmp; // x(j) -= L(i,j) * x(i)
// advance to parent
j = dof_parentid[j];
}
}
}
// x <- inv(D) * x
for (int i=0; i < nv; i++) {
x[i] *= qLDiagInv[i]; // x(i) /= L(i,i)
}
// x <- inv(L) * x; skip simple
for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) {
// init
int Madr_ij = dof_Madr[i]+1;
int j = dof_parentid[i];
// traverse ancestors backwards
tmp = x[i];
while (j>= 0) {
tmp -= qLD[Madr_ij++]*x[j]; // x(i) -= L(i,j) * x(j)
// advance to parent
j = dof_parentid[j];
}
x[i] = tmp;
}
}
}
void solveM_baseline(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y) {
solveLD_baseline(m, x, y, d->qLD, d->qLDiagInv);
}
// ----------------------------- benchmark ------------------------------------
static void BM_solveLD(benchmark::State& state, bool new_function) {
static mjModel* m = LoadModelFromPath("plugin/elasticity/coil.xml");
static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
static mjModel* m;
if (coil) {
m = LoadModelFromPath("plugin/elasticity/coil.xml");
} else {
m = LoadModelFromPath("humanoid/humanoid100.xml");
}
mjData* d = mj_makeData(m);
// warm-up rollout to get a typical state
@@ -117,24 +60,21 @@ static void BM_solveLD(benchmark::State& state, bool new_function) {
grad[i] = Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
}
// save state
std::vector<mjtNum> qpos = AsVector(d->qpos, m->nq);
std::vector<mjtNum> qvel = AsVector(d->qvel, m->nv);
std::vector<mjtNum> act = AsVector(d->act, m->na);
std::vector<mjtNum> warmstart = AsVector(d->qacc_warmstart, m->nv);
// CSR matrix
mjtNum* LDs = mj_stackAllocNum(d, m->nC);
for (int i=0; i < m->nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
}
// reset state, benchmark subsequent kNumBenchmarkSteps steps
while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
mju_copy(d->qpos, qpos.data(), m->nq);
mju_copy(d->qvel, qvel.data(), m->nv);
mju_copy(d->act, act.data(), m->na);
mju_copy(d->qacc_warmstart, warmstart.data(), m->nv);
for (int i=0; i < kNumBenchmarkSteps; i++) {
if (new_function) {
if (featherstone) {
mj_solveM(m, d, res, grad, 1);
} else {
solveM_baseline(m, d, res, grad);
mju_copy(res, grad, m->nv);
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind);
}
}
}
@@ -145,17 +85,29 @@ static void BM_solveLD(benchmark::State& state, bool new_function) {
state.SetItemsProcessed(state.iterations());
}
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_new(benchmark::State& state) {
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_FS(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_solveLD(state, true);
BM_solveLD(state, /*featherstone=*/true, /*coil=*/true);
}
BENCHMARK(BM_solveLD_new);
BENCHMARK(BM_solveLD_COIL_FS);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_old(benchmark::State& state) {
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_CSR(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_solveLD(state, false);
BM_solveLD(state, /*featherstone=*/false, /*coil=*/true);
}
BENCHMARK(BM_solveLD_old);
BENCHMARK(BM_solveLD_COIL_CSR);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_H100_FS(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_solveLD(state, /*featherstone=*/true, /*coil=*/false);
}
BENCHMARK(BM_solveLD_H100_FS);
void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_H100_CSR(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_solveLD(state, /*featherstone=*/false, /*coil=*/false);
}
BENCHMARK(BM_solveLD_H100_CSR);
} // namespace
} // namespace mujoco
test/engine/engine_core_smooth_test.cc
+77 -1
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_sparse.h"
#include <string>
#include <string_view>
@@ -31,6 +32,7 @@
namespace mujoco {
namespace {
using ::std::vector;
using ::testing::Each;
using ::testing::ElementsAre;
using ::testing::Eq;
@@ -261,7 +263,7 @@ TEST_F(CoreSmoothTest, EqualityBodySite) {
while (data->time < 0.1) {
mj_step(model, data);
}
std::vector<mjtNum> sdata = AsVector(data->sensordata, model->nsensordata);
vector<mjtNum> sdata = AsVector(data->sensordata, model->nsensordata);
// reset
mj_resetData(model, data);
@@ -405,5 +407,79 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
mj_deleteModel(model);
}
TEST_F(CoreSmoothTest, SolveLD2) {
static constexpr char xml[] = R"(
<mujoco>
<default>
<geom type="capsule" size="0.1"/>
<joint axis="0 1 0"/>
</default>
<worldbody>
<body>
<geom fromto="0 0 0 0 0 1"/>
<joint/>
<body pos="0 0 1">
<geom fromto="0 0 0 1 0 1"/>
<joint/>
</body>
<body pos="0 0 1">
<geom fromto="0 0 0 -1 0 1"/>
<joint/>
<body pos="-1 0 1">
<geom fromto="0 0 0 1 0 1"/>
<joint/>
</body>
<body pos="-1 0 1">
<geom fromto="0 0 0 -1 0 1"/>
<joint/>
</body>
</body>
</body>
</worldbody>
</mujoco>
)";
mjModel* m = LoadModelFromString(xml);
mjData* d = mj_makeData(m);
mj_forward(m, d);
int nv = m->nv;
int nC = m->nC;
// copy LD into LDs: CSR format
vector<mjtNum> LDs(nC);
for (int i=0; i < nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
}
// compare LD and LDs densified matrices
vector<mjtNum> LDdense(nv*nv);
mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv,
d->C_rownnz, d->C_rowadr, d->C_colind);
vector<mjtNum> LDdense2(nv*nv);
mj_fullM(m, LDdense2.data(), d->qLD);
// expect dense matrices to match exactly
for (int i=0; i < nv*nv; i++) EXPECT_EQ(LDdense[i], LDdense2[i]);
// compare LD and LDs vector solve
vector<mjtNum> vec(nv);
vector<mjtNum> vec2(nv);
for (int i=0; i < nv; i++) vec[i] = vec2[i] = 20 + 30*i;
for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0;
mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv,
d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv; i++) {
EXPECT_FLOAT_EQ(vec[i], vec2[i]);
}
mj_deleteData(d);
mj_deleteModel(m);
}
} // namespace
} // namespace mujoco