Multithreaded mesh compilation.

PiperOrigin-RevId: 651358120 Change-Id: Id6272aac84ec6ac1b709d2f4a44b7593307bd158
2024-07-11 04:20:59 -07:00
parent e7301edd20
commit 3d1d1d0718
7 changed files with 173 additions and 38 deletions
@@ -3664,7 +3664,7 @@ Associate this body with an :ref:`engine plugin<exPlugin>`. Either :at:`plugin`
 The :el:`attach` element is used to insert a sub-tree of bodies from another model into this model's kinematic tree.
 Unlike :ref:`include<include>`, which is implemented in the parser and is equivalent to copying and pasting XML from
 one file into another, :el:`attach` is implemented in the model compiler. In order to use this element, the sub-model
-must first be defined as an :ref:`asset<model-asset>`. When creating an attachment, the top body of the attached subtree
+must first be defined as an :ref:`asset<asset-model>`. When creating an attachment, the top body of the attached subtree
 is specified, and all referencing elements outside the kinematic tree (e.g., sensors and actuators), are
 also copied into the top-level model. Additionally, any elements referenced from within the attached subtree (e.g.
 defaults and assets) will be copied in to the top-level model. :el:`attach` is a :ref:`meta-element`, so upon saving
@@ -57,12 +57,13 @@ General
 14. Quaternions in ``mjData->qpos`` and ``mjData->mocap_quat`` are no longer normalized in-place by
    :ref:`mj_kinematics`. Instead they are normalized when they are used. After the first step, quaternions in
    ``mjData->qpos`` will be normalized.
+15. Mesh loading in the compiler, which is usually the slowest part of the loading process, is now multi-threaded.

 MJX
 ~~~
-15. Added support for :ref:`elliptic friction cones<option-cone>`.
-16. Fixed a bug that resulted in less-optimal linesearch solutions for some difficult constraint settings.
-17. Fixed a bug in the Newton solver that sometimes resulted in less-optimal gradients.
+16. Added support for :ref:`elliptic friction cones<option-cone>`.
+17. Fixed a bug that resulted in less-optimal linesearch solutions for some difficult constraint settings.
+18. Fixed a bug in the Newton solver that sometimes resulted in less-optimal gradients.


 .. youtube:: P83tKA1iz2Y
@@ -71,18 +72,17 @@ MJX

 Simulate
 ^^^^^^^^
-18. Added improved tutorial video.
-19. Improved the Brownian noise generator.
-
-|br| |br| |br| |br|
+19. Added improved tutorial video.
+20. Improved the Brownian noise generator.
+21. Now displaying model load times if they are longer than 0.25 seconds.

 Python bindings
 ^^^^^^^^^^^^^^^
-20. Fixed a memory leak when using ``copy.deepcopy()`` on a ``mujoco.MjData`` instance (:github:issue:`1572`).
+22. Fixed a memory leak when using ``copy.deepcopy()`` on a ``mujoco.MjData`` instance (:github:issue:`1572`).

 Bug fixes
 ^^^^^^^^^
-21. Fix an issue where ``mj_copyData`` (or ``copy.copy()`` in the Python bindings) was not copying contact information
+23. Fix an issue where ``mj_copyData`` (or ``copy.copy()`` in the Python bindings) was not copying contact information
    correctly (:github:issue:`1710`).

 Version 3.1.6 (Jun 3, 2024)
@@ -235,9 +235,9 @@ mjModel* LoadModel(const char* file, mj::Simulate& sim) {
  auto load_interval = mj::Simulate::Clock::now() - load_start;
  double load_seconds = Seconds(load_interval).count();

-  // if no error and load took more than 1/2 seconds, report load time
-  if (!loadError[0] && load_seconds > 0.5) {
-    mju::sprintf_arr(loadError, "Model loaded in %.1g seconds", load_seconds);
+  // if no error and load took more than 1/4 seconds, report load time
+  if (!loadError[0] && load_seconds > 0.25) {
+    mju::sprintf_arr(loadError, "Model loaded in %.2g seconds", load_seconds);
  }

  mju::strcpy_arr(sim.load_error, loadError);
@@ -1376,7 +1376,7 @@ void mjCMesh::Process() {

    // perform computation with convex mesh if volume is negative
    if (GetVolumeRef(type) <= 0 && exactmeshinertia) {
-      mju_warning("Malformed mesh %s, computing mesh inertia from convex hull", name.c_str());
+      mju_warning("Malformed mesh '%s', computing mesh inertia from convex hull", name.c_str());
      exactmeshinertia = false;
      ComputeVolume(CoM, type, facecen, exactmeshinertia);
    }
@@ -15,6 +15,7 @@
 #include "user/user_model.h"

 #include <algorithm>
+#include <atomic>
 #include <cmath>
 #include <array>
 #include <csetjmp>
@@ -22,7 +23,9 @@
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
+#include <exception>
 #include <map>
+#include <mutex>
 #include <string>
 #include <string_view>
 #include <thread>
@@ -51,6 +54,7 @@ namespace {
 namespace mju = ::mujoco::util;
 using std::string;
 using std::vector;
+constexpr int kMaxCompilerThreads = 16;
 }  // namespace

 //---------------------------------- CONSTRUCTOR AND DESTRUCTOR ------------------------------------
@@ -133,7 +137,9 @@ mjCModel& mjCModel::operator=(const mjCModel& other) {
    mjCDef* subtree = new mjCDef(*other.defaults_[0]);
    *this += *subtree;
    for (const auto& [name, def] : other.def_map) {
-      std::size_t index = std::find(other.defaults_.begin(), other.defaults_.end(), def) - other.defaults_.begin();
+      std::size_t index =
+          std::find(other.defaults_.begin(), other.defaults_.end(), def) -
+          other.defaults_.begin();
      def_map[name] = defaults_[index];
    }

@@ -1552,8 +1558,9 @@ void mjCModel::LengthRange(mjModel* m, mjData* data) {
    m->opt.timestep = LRopt.timestep;
  }

-  // number of threads available, max 16
-  const int nthread = mjMIN(16, std::thread::hardware_concurrency()/2);
+  // number of threads available
+  int hardware_threads = std::thread::hardware_concurrency();
+  const int nthread = mjMAX(1, mjMIN(kMaxCompilerThreads, hardware_threads/2));

  // count actuators that need computation
  int cnt = 0;
@@ -1592,8 +1599,8 @@ void mjCModel::LengthRange(mjModel* m, mjData* data) {
  // multiple threads
  else {
    // allocate mjData for each thread
-    char err[16][200];
-    mjData* pdata[16] = {data};
+    char err[kMaxCompilerThreads][200];
+    mjData* pdata[kMaxCompilerThreads] = {data};
    for (int i=1; i<nthread; i++) {
      pdata[i] = mj_makeData(m);
    }
@@ -1605,7 +1612,7 @@ void mjCModel::LengthRange(mjModel* m, mjData* data) {
    }

    // prepare thread function arguments, clear errors
-    LRThreadArg arg[16];
+    LRThreadArg arg[kMaxCompilerThreads];
    for (int i=0; i<nthread; i++) {
      LRThreadArg temp = {m, pdata[i], i*num, num, &LRopt, err[i], 200};
      arg[i] = temp;
@@ -1613,7 +1620,7 @@ void mjCModel::LengthRange(mjModel* m, mjData* data) {
    }

    // launch threads
-    std::thread th[16];
+    std::thread th[kMaxCompilerThreads];
    for (int i=0; i<nthread; i++) {
      th[i] = std::thread(LRfunc, arg+i);
    }
@@ -3184,8 +3191,9 @@ void mjCModel::ProcessLists(bool checkrepeat) {


 // error handler for low-level engine
+constexpr int kErrorBufferSize = 500;
 static thread_local std::jmp_buf error_jmp_buf;
-static thread_local char errortext[500] = "";
+static thread_local char errortext[kErrorBufferSize] = "";
 static void errorhandler(const char* msg) {
  mju::strcpy_arr(errortext, msg);
  std::longjmp(error_jmp_buf, 1);
@@ -3193,9 +3201,14 @@ static void errorhandler(const char* msg) {


 // warning handler for low-level engine
-static thread_local char warningtext[500] = "";
+static thread_local char warningtext[kErrorBufferSize] = "";       // top-level warning buffer
+static thread_local std::string* local_warningtext_ptr = nullptr;  // sub-thread warning buffer
 static void warninghandler(const char* msg) {
-  mju::strcpy_arr(warningtext, msg);
+  if (local_warningtext_ptr) {
+    *local_warningtext_ptr = msg;
+  } else {
+    mju::strcpy_arr(warningtext, msg);
+  }
 }


@@ -3279,6 +3292,86 @@ mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
 }


+
+// mesh compilation function to be used in threads
+void CompileMesh(mjCMesh* mesh, const mjVFS* vfs, std::exception_ptr& exception,
+                 std::mutex& exception_mutex, std::string* warningtext) {
+  // set warning text buffer to this local thread
+  local_warningtext_ptr = warningtext;
+  auto previous_handler = _mjPRIVATE__get_tls_warning_fn();
+  _mjPRIVATE__set_tls_warning_fn(warninghandler);
+
+  // compile the mesh, catch exception to be rethrown later
+  try {
+    mesh->Compile(vfs);
+  } catch (...) {
+    std::lock_guard<std::mutex> lock(exception_mutex);
+    if (!exception) {
+      exception = std::current_exception();
+    }
+  }
+
+  // restore warning handler to top-level
+  _mjPRIVATE__set_tls_warning_fn(previous_handler);
+  local_warningtext_ptr = nullptr;
+}
+
+
+
+// multi-threaded mesh compilation
+void mjCModel::CompileMeshes(const mjVFS* vfs) {
+  std::vector<std::thread> threads;
+  int nmesh = meshes_.size();
+  int hardware_threads = std::thread::hardware_concurrency();
+  int nthread = std::max(1, std::min(kMaxCompilerThreads, hardware_threads / 2));
+
+  threads.reserve(nthread);
+
+  // holds an exception thrown by a worker thread
+  std::exception_ptr exception;
+  std::mutex except_mutex;
+
+  std::atomic_int next_mesh = 0;
+  std::vector<std::string> mesh_warningtext(nmesh);
+  for (int i = 0; i < nthread; ++i) {
+    threads.emplace_back([&] {
+      for (int meshid = next_mesh++; meshid < nmesh; meshid = next_mesh++) {
+        auto& mesh = meshes_[meshid];
+        CompileMesh(mesh, vfs, exception, except_mutex,
+                    &mesh_warningtext[meshid]);
+      }
+    });
+  }
+
+  // join threads
+  for (auto& thread : threads) {
+    if (thread.joinable()) {
+      thread.join();
+    }
+  }
+
+  // concatenate all warnings from threads, copy into warningtext
+  std::string concatenated_warnings;
+  bool has_warning = false;
+  for (int i = 0; i < nmesh; i++) {
+    if (!mesh_warningtext[i].empty()) {
+      if (has_warning) {
+        concatenated_warnings += "\n";
+      }
+      concatenated_warnings += mesh_warningtext[i];
+      has_warning = true;
+    }
+  }
+  mju::strcpy_arr(warningtext, concatenated_warnings.c_str());
+
+  // if exception was caught, rethrow it
+  if (exception) {
+    std::rethrow_exception(exception);
+  }
+}
+
+
+
 void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
  // check if nan test works
  double test = mjNAN;
@@ -3352,8 +3445,14 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
  SetNuser();

  // compile meshes (needed for geom compilation)
-  for (int i=0; i<meshes_.size(); i++) {
-    meshes_[i]->Compile(vfs);
+  if (usethread) {
+    // multi-threaded mesh compile
+    CompileMeshes(vfs);
+  } else {
+    // single-threaded mesh compile
+    for (int i=0; i < meshes_.size(); i++) {
+      meshes_[i]->Compile(vfs);
+    }
  }

  // compile objects in kinematic tree
@@ -261,6 +261,9 @@ class mjCModel : public mjCModel_, private mjSpec {
  // clear objects allocated by Compile
  void Clear();

+  // multi-threaded mesh compilation
+  void CompileMeshes(const mjVFS* vfs);
+
  // if asset name is missing, set to filename
  template<class T> void SetDefaultNames(std::vector<T*>& assets);

@@ -24,6 +24,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <absl/strings/str_format.h>
+#include <absl/strings/str_replace.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mujoco.h>
@@ -72,6 +73,7 @@ using ::testing::ElementsAre;
 using ::testing::HasSubstr;
 using ::testing::IsNull;
 using ::testing::NotNull;
+using ::testing::StartsWith;

 // ------------- test invalid filenames ----------------------------------------

@@ -745,25 +747,56 @@ TEST_F(MjCMeshTest, VolumeNegativeDefaultsLegacy) {
  <mujoco>
    <compiler exactmeshinertia="true"/>
    <asset>
-      <mesh name="example_mesh"
-        vertex="0 0 0  1 0 0  0 1 0  0 0 1"
-        face="3 0 2  0 3 1  1 3 2  0 1 2" />
+      MESH_DEFINITIONS
    </asset>
    <worldbody>
      <body>
-        <geom type="mesh" mesh="example_mesh"/>
+        GEOM_DEFINITIONS
      </body>
    </worldbody>
  </mujoco>
  )";
-  std::array<char, 1024> error;
-  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
-  EXPECT_THAT(model, NotNull());
-  EXPECT_LE(mju_abs(model->geom_size[0]), 1);
-  EXPECT_LE(mju_abs(model->geom_size[1]), 1);
-  EXPECT_LE(mju_abs(model->geom_size[2]), 1);
-  EXPECT_THAT(error.data(), HasSubstr("Malformed"));
-  mj_deleteModel(model);
+
+  static constexpr char bad_mesh[] = R"(
+      <mesh name="bad_mesh%d"
+        vertex="0 0 0  1 0 0  0 1 0  0 0 1"
+        face="3 0 2  0 3 1  1 3 2  0 1 2"/>\n"
+  )";
+  static constexpr char geom[] = R"(<geom type="mesh" mesh="bad_mesh%d"/>\n)";
+
+  for (int nmesh : {3, 16, 17, 50}) {
+    std::string mesh_definitions = "";
+    for (int i = 1; i < nmesh+1; i++) {
+      mesh_definitions += absl::StrFormat(bad_mesh, i);
+    }
+
+    std::string geom_definitions = "";
+    for (int i = 1; i < nmesh+1; i++) {
+      geom_definitions += absl::StrFormat(geom, i);
+    }
+
+    std::string xml_str = xml;
+    absl::StrReplaceAll({{"MESH_DEFINITIONS", mesh_definitions},
+                         {"GEOM_DEFINITIONS", geom_definitions}}, &xml_str);
+
+    std::array<char, 1024> error;
+    mjModel* model = LoadModelFromString(xml_str.c_str(),
+                                         error.data(), error.size());
+    EXPECT_THAT(model, NotNull()) << error.data();
+    EXPECT_LE(mju_abs(model->geom_size[0]), 1);
+    EXPECT_LE(mju_abs(model->geom_size[1]), 1);
+    EXPECT_LE(mju_abs(model->geom_size[2]), 1);
+
+    EXPECT_THAT(error.data(), StartsWith("Malformed mesh 'bad_mesh1'"));
+
+    // first 7 warnings fit in the length-500 warning buffer
+    for (int i = 2; i < mjMIN(nmesh+1, 8); ++i) {
+      std::string msg = "Malformed mesh 'bad_mesh" + std::to_string(i);
+      EXPECT_THAT(error.data(), HasSubstr(msg));
+    }
+
+    mj_deleteModel(model);
+  }
 }

 TEST_F(MjCMeshTest, VolumeTooSmallAllowedWorld) {