Multithread parallelizable operations in mj_fwdPosition

PiperOrigin-RevId: 572576240
Change-Id: I70eb518385366a0027c5dad44ac15e3904d1b4a3

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<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

@@ -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

@@ -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)
+  // no threadpool: inertia and collision on main thread
-  mj_factorM(m, d);    // timed internally (POS_INERTIA)
+  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

@@ -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

@@ -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

@@ -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