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

Multithread parallelizable operations in mj_fwdPosition

Browse Source 
PiperOrigin-RevId: 572576240
Change-Id: I70eb518385366a0027c5dad44ac15e3904d1b4a3
This commit is contained in:
Matthew Bennice
2023-10-11 08:00:57 -07:00
committed by Copybara-Service
parent ba03ebdadd
commit 2842245ca9
6 changed files with 207 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+6
View File
@@ -31,6 +31,12 @@ New features
3. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
pipeline.
If engine-level threading is enabled, the following operations will be multi-threaded:
- Island constraint resolution, if island discovery is :ref:`enable flag<option-flag-island>` and the :ref:`CG<option-solver>` is selected.
- Inertia-related computations and collision detection will happen in parallel.
.. youtube:: ra2bTiZHGlw
:align: right
:width: 240px
sample/testspeed.cc
+1 -1
View File
@@ -132,7 +132,7 @@ int main(int argc, char** argv) {
" nstep 10000 number of steps per rollout\n"
" nthread 1 number of threads for which to run parallel rollouts\n"
" ctrlnoise 0.01 scale of pseudo-random noise injected into actuators\n"
" npoolthread 1 number of threads in engine-internal threadpool\n"
" npoolthread 0 number of threads in engine-internal threadpool\n"
"\n"
"Note: If the model has a keyframe named \"test\", it will be loaded prior to simulation\n");
}
src/engine/engine_forward.c
+50 -3
View File
@@ -95,6 +95,30 @@ void mj_checkAcc(const mjModel* m, mjData* d) {
//-------------------------- solver components -----------------------------------------------------
// args for internal functions in mj_fwdPosition
struct mjFwdPositionArgs_ {
const mjModel* m;
mjData* d;
};
typedef struct mjFwdPositionArgs_ mjFwdPositionArgs;
// wrapper for mj_crb and mj_factorM
void* mj_inertialThreaded(void* args) {
mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args;
mj_crb(forward_args->m, forward_args->d); // timed internally (POS_INERTIA)
mj_factorM(forward_args->m, forward_args->d); // timed internally (POS_INERTIA)
return NULL;
}
// wrapper for mj_collision
void* mj_collisionThreaded(void* args) {
mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args;
mj_collision(forward_args->m, forward_args->d); // timed internally (POS_COLLISION)
return NULL;
}
// position-dependent computations
void mj_fwdPosition(const mjModel* m, mjData* d) {
TM_START1;
@@ -106,10 +130,33 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
mj_tendon(m, d);
TM_END(mjTIMER_POS_KINEMATICS);
mj_crb(m, d); // timed internally (POS_INERTIA)
mj_factorM(m, d); // timed internally (POS_INERTIA)
// no threadpool: inertia and collision on main thread
if (!d->threadpool) {
mj_crb(m, d); // timed internally (POS_INERTIA)
mj_factorM(m, d); // timed internally (POS_INERTIA)
mj_collision(m, d); // timed internally (POS_COLLISION)
}
mj_collision(m, d); // timed internally (POS_COLLISION)
// have threadpool: inertia and collision on seperate threads
else {
mjTask tasks[2];
mjFwdPositionArgs forward_args;
forward_args.m = m;
forward_args.d = d;
mju_defaultTask(&tasks[0]);
tasks[0].func = mj_inertialThreaded;
tasks[0].args = &forward_args;
mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[0]);
mju_defaultTask(&tasks[1]);
tasks[1].func = mj_collisionThreaded;
tasks[1].args = &forward_args;
mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[1]);
mju_taskJoin(&tasks[0]);
mju_taskJoin(&tasks[1]);
}
TM_RESTART;
mj_makeConstraint(m, d);
test/engine/CMakeLists.txt
+3
View File
@@ -78,6 +78,9 @@ target_link_libraries(engine_solver_test fixture gmock)
mujoco_test(engine_support_test)
target_link_libraries(engine_support_test fixture gmock)
mujoco_test(engine_thread_test)
target_link_libraries(engine_thread_test fixture gmock)
mujoco_test(engine_util_blas_test)
target_link_libraries(engine_util_blas_test fixture gmock)
test/engine/engine_island_test.cc
-56
View File
@@ -387,61 +387,5 @@ TEST_F(IslandTest, IslandEfcElliptic) {
mj_deleteModel(model);
}
TEST_F(IslandTest, IslandSingleAndMultiThreadedMatch) {
auto model_path = GetTestDataFilePath("benchmark/testdata/22_humanoids.xml");
std::array<char, 1024> error;
mjModel* model =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mjModel* model_threaded =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model_threaded->opt.solver = mjSOL_CG; // use CG solver
model_threaded->opt.enableflags |= mjENBL_ISLAND; // enable islands
mjData* data = mj_makeData(model);
mjData* data_threaded = mj_makeData(model_threaded);
// warm-up rollout to get a steady state
static constexpr int kNumWarmupSteps = 100;
for (int i = 0; i < kNumWarmupSteps; i++) {
mj_step(model, data);
}
// sync the models and data
int spec = mjSTATE_INTEGRATION;
int size = mj_stateSize(model, spec);
std::vector<mjtNum> initial_state(size);
mj_getState(model, data, initial_state.data(), spec);
mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
// bind a threadpool to the data_threaded
mjThreadPool* threadpool = mju_threadPoolCreate(10);
mju_bindThreadPool(data_threaded, threadpool);
for (int i = 0; i < 10; ++i) {
mj_step(model, data);
mj_step(model_threaded, data_threaded);
}
// compare the mjData's.
{
MJDATA_POINTERS_PREAMBLE((model))
#define X(type, name, nr, nc) \
EXPECT_EQ(std::memcmp(data->name, data_threaded->name, \
sizeof(type)*(model->nr)*(nc)), \
0) << "mjData::" #name " differs";
MJDATA_POINTERS
#undef X
}
mj_deleteData(data);
mj_deleteModel(model);
mj_deleteData(data_threaded);
mj_deleteModel(model_threaded);
mju_threadPoolDestroy(threadpool);
}
} // namespace
} // namespace mujoco
test/engine/engine_thread_test.cc
+147
View File
@@ -0,0 +1,147 @@
// Copyright 2023 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.
// Tests for engine/engine_island.c.
#include <array>
#include <cstring>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjthread.h>
#include <mujoco/mjxmacro.h>
#include <mujoco/mujoco.h>
#include "src/thread/thread_pool.h"
#include "test/fixture.h"
namespace mujoco {
namespace {
using ThreadTest = MujocoTest;
TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
auto model_path = GetTestDataFilePath("benchmark/testdata/22_humanoids.xml");
std::array<char, 1024> error;
mjModel* model =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model->opt.solver = mjSOL_CG; // use CG solver
mjModel* model_threaded =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model_threaded->opt.solver = mjSOL_CG; // use CG solver
mjData* data = mj_makeData(model);
mjData* data_threaded = mj_makeData(model_threaded);
// warm-up rollout to get a steady state
static constexpr int kNumWarmupSteps = 100;
for (int i = 0; i < kNumWarmupSteps; i++) {
mj_step(model, data);
}
// sync the models and data
int spec = mjSTATE_INTEGRATION;
int size = mj_stateSize(model, spec);
std::vector<mjtNum> initial_state(size);
mj_getState(model, data, initial_state.data(), spec);
mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
// bind a threadpool to the data_threaded
mjThreadPool* threadpool = mju_threadPoolCreate(10);
mju_bindThreadPool(data_threaded, threadpool);
for (int i = 0; i < 10; ++i) {
mj_step(model, data);
mj_step(model_threaded, data_threaded);
}
// compare the mjData's.
{
MJDATA_POINTERS_PREAMBLE((model))
#define X(type, name, nr, nc) \
EXPECT_EQ(std::memcmp(data->name, data_threaded->name, \
sizeof(type)*(model->nr)*(nc)), \
0) << "mjData::" #name " differs";
MJDATA_POINTERS
#undef X
}
mj_deleteData(data);
mj_deleteModel(model);
mj_deleteData(data_threaded);
mj_deleteModel(model_threaded);
mju_threadPoolDestroy(threadpool);
}
TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
auto model_path = GetTestDataFilePath("benchmark/testdata/22_humanoids.xml");
std::array<char, 1024> error;
mjModel* model =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model->opt.solver = mjSOL_CG; // use CG solver
model->opt.enableflags |= mjENBL_ISLAND; // enable islands
mjModel* model_threaded =
mj_loadXML(model_path.c_str(), nullptr, error.data(), error.size());
model_threaded->opt.solver = mjSOL_CG; // use CG solver
model_threaded->opt.enableflags |= mjENBL_ISLAND; // enable islands
mjData* data = mj_makeData(model);
mjData* data_threaded = mj_makeData(model_threaded);
// warm-up rollout to get a steady state
static constexpr int kNumWarmupSteps = 100;
for (int i = 0; i < kNumWarmupSteps; i++) {
mj_step(model, data);
}
// sync the models and data
int spec = mjSTATE_INTEGRATION;
int size = mj_stateSize(model, spec);
std::vector<mjtNum> initial_state(size);
mj_getState(model, data, initial_state.data(), spec);
mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
// bind a threadpool to the data_threaded
mjThreadPool* threadpool = mju_threadPoolCreate(10);
mju_bindThreadPool(data_threaded, threadpool);
for (int i = 0; i < 10; ++i) {
mj_step(model, data);
mj_step(model_threaded, data_threaded);
}
// compare the mjData's.
{
MJDATA_POINTERS_PREAMBLE((model))
#define X(type, name, nr, nc) \
EXPECT_EQ(std::memcmp(data->name, data_threaded->name, \
sizeof(type)*(model->nr)*(nc)), \
0) << "mjData::" #name " differs";
MJDATA_POINTERS
#undef X
}
mj_deleteData(data);
mj_deleteModel(model);
mj_deleteData(data_threaded);
mj_deleteModel(model_threaded);
mju_threadPoolDestroy(threadpool);
}
} // namespace
} // namespace mujoco