diff --git a/doc/changelog.rst b/doc/changelog.rst index 19601f9e..16a3d108 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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` and the :ref:`CG` is selected. + - Inertia-related computations and collision detection will happen in parallel. + + .. youtube:: ra2bTiZHGlw :align: right :width: 240px diff --git a/sample/testspeed.cc b/sample/testspeed.cc index 09d80a3d..b13c9238 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -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"); } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 5746d865..9433a81b 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -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); diff --git a/test/engine/CMakeLists.txt b/test/engine/CMakeLists.txt index c810a055..1e083b36 100644 --- a/test/engine/CMakeLists.txt +++ b/test/engine/CMakeLists.txt @@ -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) diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 758dcc89..c0f5afb8 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -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 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 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 diff --git a/test/engine/engine_thread_test.cc b/test/engine/engine_thread_test.cc new file mode 100644 index 00000000..1575d640 --- /dev/null +++ b/test/engine/engine_thread_test.cc @@ -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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#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 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 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 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 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