From ac11e5faa6700ab54409c0713e2d80cd688c217a Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 6 Jan 2025 05:47:13 -0800 Subject: [PATCH] Add CSR implementation of `mj_factorI` PiperOrigin-RevId: 712498431 Change-Id: I13b52e53482ed97da8788875d4d95e2beb5ca7c1 --- src/engine/engine_core_smooth.c | 46 +++++- src/engine/engine_core_smooth.h | 8 +- test/benchmark/CMakeLists.txt | 12 ++ test/benchmark/factorI_benchmark_test.cc | 102 ++++++++++++ test/benchmark/inertia_benchmark_test.cc | 101 ++++++++++++ test/benchmark/solveLD_benchmark_test.cc | 3 +- test/benchmark/testdata/inertia.xml | 199 +++++++++++++++++++++++ test/engine/engine_core_smooth_test.cc | 49 +++++- test/fixture.h | 8 +- 9 files changed, 514 insertions(+), 14 deletions(-) create mode 100644 test/benchmark/factorI_benchmark_test.cc create mode 100644 test/benchmark/inertia_benchmark_test.cc create mode 100644 test/benchmark/testdata/inertia.xml diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 641b7bdc..81e276f6 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1440,10 +1440,9 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu } } - // compute 1/diag(D), 1/sqrt(diag(D)) + // compute 1/diag(D) for (int i=0; i < nv; i++) { - mjtNum qLDi = qLD[dof_Madr[i]]; - qLDiagInv[i] = 1.0/qLDi; + qLDiagInv[i] = 1.0 / qLD[dof_Madr[i]]; } } @@ -1458,6 +1457,40 @@ void mj_factorM(const mjModel* m, mjData* d) { +// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd +// like mj_factorI, but using CSR representation +void mj_factorIs(mjtNum* mat, mjtNum* diaginv, int nv, + const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) { + // backward loop over rows + for (int k=nv-1; k >= 0; k--) { + // get row k's address, diagonal index, inverse diagonal value + int rowadr_k = rowadr[k]; + int diag_k = rowadr_k + rownnz[k] - 1; + mjtNum invD = 1 / mat[diag_k]; + if (diaginv) diaginv[k] = invD; + + // skip if simple + if (diagnum[k]) { + continue; + } + + // update triangle above row k, inclusive + for (int adr=diag_k - 1; adr >= rowadr_k; adr--) { + // tmp = L(k, i) / L(k, k) + mjtNum tmp = mat[adr] * invD; + + // update row i < k: L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k) + int i = colind[adr]; + mju_addToScl(mat + rowadr[i], mat + rowadr_k, -tmp, rownnz[i]); + + // update ith element of row k: L(k, i) /= L(k, k) + mat[adr] = tmp; + } + } +} + + + // in-place sparse backsubstitution: x = inv(L'*D*L)*x // L is in lower triangle of qLD; D is on diagonal of qLD // handle n vectors at once @@ -1575,8 +1608,7 @@ 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* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, - const int* rownnz, const int* rowadr, const int* diagind, const int* diagnum, - const int* colind) { + const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) { // x <- L^-T x for (int i=nv-1; i > 0; i--) { // skip diagonal (simple) rows, exploit sparsity of input vector @@ -1584,7 +1616,7 @@ void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv continue; } - int d = diagind[i]; + int d = rownnz[i] - 1; int adr_i = rowadr[i]; mjtNum x_i = x[i]; for (int j=0; j < d; j++) { @@ -1607,7 +1639,7 @@ void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv } int adr = rowadr[i]; - x[i] -= mju_dotSparse(qLDs+adr, x, diagind[i], colind+adr, /*flg_unc1=*/0); + x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0); } } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index db1af84b..f4733b08 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -51,6 +51,11 @@ MJAPI void mj_crb(const mjModel* m, mjData* d); // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv); +// sparse L'*D*L factorizaton of inertia-like matrix +// like mj_factorI, but using CSR representation +MJAPI void mj_factorIs(mjtNum* mat, mjtNum* diaginv, int nv, + const int* rownnz, const int* rowadr, const int* diagnum, const int* colind); + // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd MJAPI void mj_factorM(const mjModel* m, mjData* d); @@ -61,8 +66,7 @@ MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n, // 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* diagind, const int* diagnum, - const int* colind); + const int* rownnz, const int* rowadr, const int* diagnum, 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/CMakeLists.txt b/test/benchmark/CMakeLists.txt index 0407edc0..0a23fafe 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -43,6 +43,18 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) +mujoco_test( + factorI_benchmark_test + MAIN_TARGET benchmark::benchmark_main + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers +) + +mujoco_test( + inertia_benchmark_test + MAIN_TARGET benchmark::benchmark_main + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers +) + mujoco_test( solveLD_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc new file mode 100644 index 00000000..55118a99 --- /dev/null +++ b/test/benchmark/factorI_benchmark_test.cc @@ -0,0 +1,102 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// A benchmark for comparing different implementations of mj_factorI. + +#include +#include +#include +#include +#include "src/engine/engine_core_smooth.h" +#include "test/fixture.h" + +namespace mujoco { +namespace { + +// number of steps to benchmark +static const int kNumBenchmarkSteps = 50; + +// ----------------------------- benchmark ------------------------------------ + +static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { + static mjModel* m; + if (coil) { + m = LoadModelFromPath("plugin/elasticity/coil.xml"); + } else { + m = LoadModelFromPath("humanoid/humanoid100.xml"); + } + + mjData* d = mj_makeData(m); + mj_forward(m, d); + + // allocate inputs and outputs + mj_markStack(d); + + // CSR matrices + mjtNum* Ms = mj_stackAllocNum(d, m->nC); + mjtNum* LDs = mj_stackAllocNum(d, m->nC); + for (int i=0; i < m->nC; i++) { + Ms[i] = d->qM[d->mapM2C[i]]; + } + + // benchmark + while (state.KeepRunningBatch(kNumBenchmarkSteps)) { + for (int i=0; i < kNumBenchmarkSteps; i++) { + if (legacy) { + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); + } else { + mju_copy(LDs, Ms, m->nC); + mj_factorIs(LDs, d->qLDiagInv, m->nv, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + } + } + } + + // finalize + mj_freeStack(d); + mj_deleteData(d); + mj_deleteModel(m); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_factorI_COIL_LEGACY(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_factorI(state, /*legacy=*/true, /*coil=*/true); +} +BENCHMARK(BM_factorI_COIL_LEGACY); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_factorI_COIL_CSR(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_factorI(state, /*legacy=*/false, /*coil=*/true); +} +BENCHMARK(BM_factorI_COIL_CSR); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_factorI_H100_LEGACY(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_factorI(state, /*legacy=*/true, /*coil=*/false); +} +BENCHMARK(BM_factorI_H100_LEGACY); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_factorI_H100_CSR(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_factorI(state, /*legacy=*/false, /*coil=*/false); +} +BENCHMARK(BM_factorI_H100_CSR); + +} // namespace +} // namespace mujoco diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc new file mode 100644 index 00000000..8ba580b7 --- /dev/null +++ b/test/benchmark/inertia_benchmark_test.cc @@ -0,0 +1,101 @@ +// Copyright 2025 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// A benchmark for comparing legacy and two CSR implementations of inertia +// factor and then solve. + +#include +#include +#include +#include +#include "src/engine/engine_core_smooth.h" +#include "test/fixture.h" + +namespace mujoco { +namespace { + +// number of steps to benchmark +static const int kNumBenchmarkSteps = 50; + +// ----------------------------- benchmark ------------------------------------ + +enum class SolveType { + kLegacy = 0, + kCsr, +}; + +static void BM_solve(benchmark::State& state, SolveType type) { + static mjModel* m; + m = LoadModelFromPath("../test/benchmark/testdata/inertia.xml"); + + mjData* d = mj_makeData(m); + mj_forward(m, d); + + // allocate input and output vectors + mj_markStack(d); + + // make CSR matrix + mjtNum* Ms = mj_stackAllocNum(d, m->nC); + mjtNum* LDs = mj_stackAllocNum(d, m->nC); + for (int i=0; i < m->nC; i++) { + Ms[i] = d->qM[d->mapM2C[i]]; + } + + // arbitrary input vector + mjtNum *res = mj_stackAllocNum(d, m->nv); + mjtNum *vec = mj_stackAllocNum(d, m->nv); + for (int i=0; i < m->nv; i++) { + vec[i] = 0.2 + 0.3*i; + } + + // benchmark + while (state.KeepRunningBatch(kNumBenchmarkSteps)) { + for (int i=0; i < kNumBenchmarkSteps; i++) { + switch (type) { + case SolveType::kLegacy: + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); + mj_solveM(m, d, res, vec, 1); + break; + case SolveType::kCsr: + mju_copy(LDs, Ms, m->nC); + mj_factorIs(LDs, d->qLDiagInv, m->nv, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + mju_copy(res, vec, m->nv); + mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + } + } + } + + // finalize + mj_freeStack(d); + mj_deleteData(d); + mj_deleteModel(m); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solve_LEGACY(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_solve(state, SolveType::kLegacy); +} +BENCHMARK(BM_solve_LEGACY); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solve_CSR(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_solve(state, SolveType::kCsr); +} +BENCHMARK(BM_solve_CSR); + +} // namespace +} // namespace mujoco diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 5564fc78..f2dc39e1 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -64,8 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { } else { mju_copy(res, vec, m->nv); mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, - d->C_rownnz, d->C_rowadr, d->C_diag, m->dof_simplenum, - d->C_colind); + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); } } } diff --git a/test/benchmark/testdata/inertia.xml b/test/benchmark/testdata/inertia.xml new file mode 100644 index 00000000..5d08a751 --- /dev/null +++ b/test/benchmark/testdata/inertia.xml @@ -0,0 +1,199 @@ + + diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 0cf12290..717fde98 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -496,8 +496,7 @@ TEST_F(CoreSmoothTest, SolveLDs) { 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, m->dof_simplenum, - d->C_colind); + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); // expect vectors to match up to floating point precision for (int i=0; i < nv; i++) { @@ -508,5 +507,51 @@ TEST_F(CoreSmoothTest, SolveLDs) { mj_deleteModel(m); } +TEST_F(CoreSmoothTest, FactorIs) { + const std::string xml_path = GetTestDataFilePath(kInertiaPath); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error; + + mjData* d = mj_makeData(m); + mj_forward(m, d); + + int nC = m->nC, nv = m->nv; + + // copy qM into LDs, qLD into qLDexpected: CSR format + vector qLDsExpected(nC); + vector qLDs(nC); + for (int i=0; i < nC; i++) { + int index = d->mapM2C[i]; + qLDs[i] = d->qM[index]; // mj_factorIs is in-place + qLDsExpected[i] = d->qLD[index]; + } + + vector qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); + vector qLDiagInv(nv, 0); + + mj_factorIs(qLDs.data(), qLDiagInv.data(), nv, + d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind); + + // expect outputs to match to floating point precision + EXPECT_THAT(qLDs, Pointwise(DoubleNear(1e-12), qLDsExpected)); + EXPECT_THAT(qLDiagInv, Pointwise(DoubleNear(1e-12), qLDiagInvExpected)); + + /* uncomment for debugging + vector LDdense(nv*nv); + + mju_sparse2dense(LDdense.data(), qLDexpected.data(), nv, nv, + d->C_rownnz, d->C_rowadr, d->C_colind); + PrintMatrix(LDdense.data(), nv, nv, 2); + + mju_sparse2dense(LDdense.data(), qLDs.data(), nv, nv, + d->C_rownnz, d->C_rowadr, d->C_colind); + PrintMatrix(LDdense.data(), nv, nv, 2); + */ + + mj_deleteData(d); + mj_deleteModel(m); +} + } // namespace } // namespace mujoco diff --git a/test/fixture.h b/test/fixture.h index 69224c15..d2860d7c 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -118,7 +118,13 @@ inline void PrintMatrix(const mjtNum* mat, int nrow, int ncol, int p = 5) { std::cerr << "\n"; for (int r = 0; r < nrow; r++) { for (int c = 0; c < ncol; c++) { - std::cerr << std::fixed << std::setw(3 + p) << mat[c + r*ncol] << " "; + mjtNum val = mat[c + r*ncol]; + if (val) { + std::cerr << std::fixed << std::setw(5 + p) << val << " "; + } else { + // don't print exact zeros + std::cerr << std::string(6 + p, ' '); + } } std::cerr << "\n"; }