diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 6cc5a21f..a11895b5 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -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) { diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 35a78b43..a1fab753 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -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); diff --git a/test/benchmark/engine_core_smooth_benchmark_test.cc b/test/benchmark/engine_core_smooth_benchmark_test.cc index fc6c6a89..9a80ae8e 100644 --- a/test/benchmark/engine_core_smooth_benchmark_test.cc +++ b/test/benchmark/engine_core_smooth_benchmark_test.cc @@ -19,85 +19,28 @@ #include #include #include +#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 qpos = AsVector(d->qpos, m->nq); - std::vector qvel = AsVector(d->qvel, m->nv); - std::vector act = AsVector(d->act, m->na); - std::vector 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 diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 8b2b4f2a..f27d1670 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/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_sparse.h" #include #include @@ -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 sdata = AsVector(data->sensordata, model->nsensordata); + vector 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"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + )"; + 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 LDs(nC); + for (int i=0; i < nC; i++) { + LDs[i] = d->qLD[d->mapM2C[i]]; + } + + // compare LD and LDs densified matrices + vector LDdense(nv*nv); + mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv, + d->C_rownnz, d->C_rowadr, d->C_colind); + vector 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 vec(nv); + vector 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