From 9fc9823a8c8dcc5a2d0b54ca92ec0ae96418cea5 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 8 Nov 2022 08:38:04 -0800 Subject: [PATCH] Speed improvements in mj_solveLD (14% on Intel Skylake Xeon 36 cores). PiperOrigin-RevId: 486960670 Change-Id: I8fce96427158ec36069dea893500059e93bd40db --- doc/changelog.rst | 1 + src/engine/engine_core_smooth.c | 26 ++- src/engine/engine_core_smooth.h | 2 +- src/engine/engine_forward.c | 3 +- .../engine_core_smooth_benchmark_test.cc | 167 ++++++++++++++++++ test/benchmark/step_benchmark_test.cc | 40 +---- test/engine/engine_derivative_test.cc | 5 +- test/fixture.cc | 30 ++++ test/fixture.h | 8 + 9 files changed, 228 insertions(+), 54 deletions(-) create mode 100644 test/benchmark/engine_core_smooth_benchmark_test.cc diff --git a/doc/changelog.rst b/doc/changelog.rst index 88c84b9d..0ef9fb0c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -38,6 +38,7 @@ General joints. See the two new examples: `particle_free.xml `_ and `particle_free2d.xml `_ +- Better mj_solveLD performance. This gives an improvement of 14% on a Intel Skylake Xeon with 36 cores. Version 2.3.0 (October 18, 2022) -------------------------------- diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index b3380de8..a23540f2 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1088,32 +1088,26 @@ void mj_factorM(const mjModel* m, mjData* d) { // sparse backsubstitution: x = inv(L'*D*L)*y // L is in lower triangle of qLD; D is on diagonal of qLD // handle n vectors at once -void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n, +void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n, 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, n*nv); - } - // single vector if (n==1) { // 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])) { + if (!m->dof_simplenum[i] && x[i]) { // init int Madr_ij = dof_Madr[i]+1; int j = dof_parentid[i]; // traverse ancestors backwards + // read directly from x[i] since i cannot be a parent of itself while (j>=0) { - x[j] -= qLD[Madr_ij++]*tmp; // x(j) -= L(i,j) * x(i) + x[j] -= qLD[Madr_ij++]*x[i]; // x(j) -= L(i,j) * x(i) // advance to parent j = dof_parentid[j]; @@ -1134,14 +1128,13 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n, int j = dof_parentid[i]; // traverse ancestors backwards - tmp = x[i]; + // write directly in x[i] since i cannot be a parent of itself while (j>=0) { - tmp -= qLD[Madr_ij++]*x[j]; // x(i) -= L(i,j) * x(j) + x[i] -= qLD[Madr_ij++]*x[j]; // x(i) -= L(i,j) * x(j) // advance to parent j = dof_parentid[j]; } - x[i] = tmp; } } } @@ -1149,6 +1142,7 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n, // multiple vectors else { int offset; + mjtNum tmp; // x <- inv(L') * x; skip simple for (int i=nv-1; i>=0; i--) { @@ -1204,11 +1198,13 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n, } - // 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) { - mj_solveLD(m, x, y, n, d->qLD, d->qLDiagInv); + if (x != y) { + mju_copy(x, y, n*m->nv); + } + mj_solveLD(m, x, n, d->qLD, d->qLDiagInv); } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 86525c1a..08e5cc44 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -56,7 +56,7 @@ MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, MJAPI void mj_factorM(const mjModel* m, mjData* d); // sparse backsubstitution: x = inv(L'*D*L)*y -MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n, +MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n, const mjtNum* qLD, const mjtNum* qLDiagInv); // sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 7fe75f76..40288d40 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -576,7 +576,8 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { // solve mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); - mj_solveLD(m, qacc, qfrc, 1, d->qH, d->qHDiagInv); + mju_copy(qacc, qfrc, m->nv); + mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv); } // advance state and time diff --git a/test/benchmark/engine_core_smooth_benchmark_test.cc b/test/benchmark/engine_core_smooth_benchmark_test.cc new file mode 100644 index 00000000..68713711 --- /dev/null +++ b/test/benchmark/engine_core_smooth_benchmark_test.cc @@ -0,0 +1,167 @@ +// Copyright 2021 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_solveLD. + +#include +#include +#include +#include +#include +#include "test/fixture.h" + +namespace mujoco { +namespace { + +// number of steps to roll out before benhmarking +static const int kNumWarmupSteps = 500; + +// number of steps to benchmark +static const int kNumBenchmarkSteps = 50; + +// copy array into vector +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} + +// ----------------------------- 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("composite/cloth.xml"); + mjData* d = mj_makeData(m); + + // warm-up rollout to get a typcal state + for (int i=0; i < kNumWarmupSteps; i++) { + AddCtrlNoise(m, d, i); + mj_step(m, d); + } + + // allocate gadient + mjMARKSTACK; + mjtNum *grad = mj_stackAlloc(d, m->nv); + mjtNum *Ma = mj_stackAlloc(d, m->nv); + mjtNum *res = mj_stackAlloc(d, m->nv); + + // compute gradient + mj_mulM(m, d, Ma, d->qacc); + for (int i=0; i < m->nv; i++) { + 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); + + // 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) { + mj_solveM(m, d, res, grad, 1); + } else { + solveM_baseline(m, d, res, grad); + } + } + } + + // finalize + mjFREESTACK; + mj_deleteData(d); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_new(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_solveLD(state, true); +} +BENCHMARK(BM_solveLD_new); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_old(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_solveLD(state, false); +} +BENCHMARK(BM_solveLD_old); + +} // namespace +} // namespace mujoco diff --git a/test/benchmark/step_benchmark_test.cc b/test/benchmark/step_benchmark_test.cc index fc2a7964..c4704603 100644 --- a/test/benchmark/step_benchmark_test.cc +++ b/test/benchmark/step_benchmark_test.cc @@ -28,53 +28,23 @@ namespace mujoco { namespace { -using ::testing::NotNull; - // number of steps to roll out before benhmarking static const int kNumWarmupSteps = 500; // number of steps to benchmark static const int kNumBenchmarkSteps = 50; - // copy array into vector std::vector AsVector(const mjtNum* array, int n) { return std::vector(array, array + n); } -static void add_ctrl_noise(const mjModel* m, mjData* d, int step) { - for (int i = 0; i < m->nu; i++) { - mjtNum center = 0.0; - mjtNum radius = 1.0; - mjtNum* range = m->actuator_ctrlrange + 2 * i; - if (m->actuator_ctrllimited[i]) { - center = (range[1] + range[0]) / 2; - radius = (range[1] - range[0]) / 2; - } - radius *= 0.01; - d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1); - } -} - -static void assert_model_not_null(mjModel* model, - const std::array& error) { - ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); -} - -static mjModel* load_model(const char* model_path) { - const std::string xml_path = GetModelPath(model_path); - std::array error; - mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error.data(), error.size()); - assert_model_not_null(model, error); - return model; -} - static void run_step_benchmark(const mjModel* model, benchmark::State& state) { mjData* data = mj_makeData(model); // warm-up rollout to get a typcal state for (int i=0; i < kNumWarmupSteps; i++) { - add_ctrl_noise(model, data, i); + AddCtrlNoise(model, data, i); mj_step(model, data); } // save state @@ -91,7 +61,7 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) { mju_copy(data->qacc_warmstart, warmstart.data(), model->nv); for (int i=0; i < kNumBenchmarkSteps; i++) { - add_ctrl_noise(model, data, i+kNumWarmupSteps); + AddCtrlNoise(model, data, i+kNumWarmupSteps); mj_step(model, data); } } @@ -107,21 +77,21 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) { void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepCloth(benchmark::State& state) { MujocoErrorTestGuard guard; - static mjModel* model = load_model("composite/cloth.xml"); + static mjModel* model = LoadModelFromPath("composite/cloth.xml"); run_step_benchmark(model, state); } BENCHMARK(BM_StepCloth); void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid(benchmark::State& state) { MujocoErrorTestGuard guard; - static mjModel* model = load_model("humanoid/humanoid.xml"); + static mjModel* model = LoadModelFromPath("humanoid/humanoid.xml"); run_step_benchmark(model, state); } BENCHMARK(BM_StepHumanoid); void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid100(benchmark::State& state) { MujocoErrorTestGuard guard; - static mjModel* model = load_model("humanoid100/humanoid100.xml"); + static mjModel* model = LoadModelFromPath("humanoid100/humanoid100.xml"); run_step_benchmark(model, state); } BENCHMARK(BM_StepHumanoid100); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 47cd03be..574a3523 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -289,7 +289,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { Ac[i*nv + i] = -m->jnt_stiffness[i]; Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; } - mj_solveLD(m, Ac, Ac, 2*nv, d->qH, d->qHDiagInv); + mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv); // A = [dt*Ac; Ac] mju_transpose(A, Ac, 2*nv, nv); @@ -314,7 +314,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { if (B) { mjtNum *Bc = mj_stackAlloc(d, nu*nv); mjtNum *BcT = mj_stackAlloc(d, nv*nu); - mj_solveLD(m, Bc, d->actuator_moment, nu, d->qH, d->qHDiagInv); + mju_copy(Bc, d->actuator_moment, nv*nu); + mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv); mju_transpose(BcT, Bc, nu, nv); mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv); diff --git a/test/fixture.cc b/test/fixture.cc index 347f4f54..937adc10 100644 --- a/test/fixture.cc +++ b/test/fixture.cc @@ -30,6 +30,8 @@ namespace mujoco { namespace { +using ::testing::NotNull; + ABSL_CONST_INIT static absl::Mutex handlers_mutex(absl::kConstInit); static int guard_count ABSL_GUARDED_BY(handlers_mutex) = 0; @@ -80,6 +82,20 @@ mjModel* LoadModelFromString(std::string_view xml, char* error, return m; } +static void AssertModelNotNull(mjModel* model, + const std::array& error) { + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); +} + +mjModel* LoadModelFromPath(const char* model_path) { + const std::string xml_path = GetModelPath(model_path); + std::array error; + mjModel* model = mj_loadXML( + xml_path.c_str(), nullptr, error.data(), error.size()); + AssertModelNotNull(model, error); + return model; +} + const std::string GetFileContents(const char* path) { std::ifstream ifs; ifs.open(path, std::ifstream::in); @@ -118,4 +134,18 @@ const std::string SaveAndReadXml(const mjModel* model) { return contents; } +void AddCtrlNoise(const mjModel* m, mjData* d, int step) { + for (int i = 0; i < m->nu; i++) { + mjtNum center = 0.0; + mjtNum radius = 1.0; + mjtNum* range = m->actuator_ctrlrange + 2 * i; + if (m->actuator_ctrllimited[i]) { + center = (range[1] + range[0]) / 2; + radius = (range[1] - range[0]) / 2; + } + radius *= 0.01; + d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1); + } +} + } // namespace mujoco diff --git a/test/fixture.h b/test/fixture.h index 4addb5dd..c32c200c 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -17,6 +17,7 @@ #include #include +#include #include namespace mujoco { @@ -53,8 +54,15 @@ const std::string GetModelPath(absl::string_view path); mjModel* LoadModelFromString(absl::string_view xml, char* error = nullptr, int error_size = 0); +// Returns a newly-allocated mjModel, loaded from the contents in model_path. +// On failure it asserts that model is null. +mjModel* LoadModelFromPath(const char* model_path); + // Returns a string loaded from first saving the model given an input. const std::string SaveAndReadXml(const mjModel* model); +// Adds control noise. +void AddCtrlNoise(const mjModel* m, mjData* d, int step); + } // namespace mujoco #endif // MUJOCO_TEST_FIXTURE_H_