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

Speed improvements in mj_solveLD (14% on Intel Skylake Xeon 36 cores).

Browse Source 
PiperOrigin-RevId: 486960670
Change-Id: I8fce96427158ec36069dea893500059e93bd40db
This commit is contained in:
Alessio Quaglino
2022-11-08 08:38:04 -08:00
committed by Copybara-Service
parent 9ffde761e8
commit 9fc9823a8c
9 changed files with 228 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1
View File
@@ -38,6 +38,7 @@ General
joints. See the two new examples:
`particle_free.xml <https://github.com/deepmind/mujoco/tree/main/model/composite/particle_free.xml>`_ and
`particle_free2d.xml <https://github.com/deepmind/mujoco/tree/main/model/composite/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)
--------------------------------
src/engine/engine_core_smooth.c
+11 -15
View File
@@ -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);
}
src/engine/engine_core_smooth.h
+1 -1
View File
@@ -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
src/engine/engine_forward.c
+2 -1
View File
@@ -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
test/benchmark/engine_core_smooth_benchmark_test.cc
+167
View File
@@ -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 <benchmark/benchmark.h>
#include <gtest/gtest.h>
#include <absl/base/attributes.h>
#include <mujoco/mjdata.h>
#include <mujoco/mujoco.h>
#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<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(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<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);
// 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
test/benchmark/step_benchmark_test.cc
+5 -35
View File
@@ -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<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(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<char, 1024>& 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<char, 1024> 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);
test/engine/engine_derivative_test.cc
+3 -2
View File
@@ -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);
test/fixture.cc
+30
View File
@@ -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<char, 1024>& 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<char, 1024> 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
test/fixture.h
+8
View File
@@ -17,6 +17,7 @@
#include <gtest/gtest.h>
#include <absl/strings/string_view.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
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_