Add threading primitives.

PiperOrigin-RevId: 562649644
Change-Id: I9c35b270e4b3b50cc7eb5a152f24d71def8e2fcd

CMakeLists.txt

@@ -59,6 +59,7 @@ set(MUJOCO_HEADERS
     include/mujoco/mjmodel.h
     include/mujoco/mjplugin.h
     include/mujoco/mjrender.h
+    include/mujoco/mjthread.h
     include/mujoco/mjtnum.h
     include/mujoco/mjui.h
     include/mujoco/mjvisualize.h
@@ -88,6 +89,7 @@ add_subdirectory(src/engine)
 add_subdirectory(src/user)
 add_subdirectory(src/xml)
 add_subdirectory(src/render)
+add_subdirectory(src/thread)
 add_subdirectory(src/ui)
 
 target_compile_definitions(mujoco PRIVATE _GNU_SOURCE CCD_STATIC_DEFINE MUJOCO_DLL_EXPORTS -DMC_IMPLEM_ENABLE)

doc/APIreference/APIfunctions.rst

@@ -32,6 +32,7 @@ API function can be classified as:
 - :ref:`Derivatives<Derivatives-api>`.
 - :ref:`Plugin<Plugins-api>` related functions.
 - :ref:`Macros<Macros>`.
+- :ref:`Thread<Thread>` related functions.
 
 .. TODO(b/273075045): Better category-label namespacing.
 

doc/APIreference/APItypes.rst

@@ -712,7 +712,26 @@ Options for configuring the automatic :ref:`actuator length-range computation<CL
 
 .. mujoco-include:: mjLROpt
 
+.. _mjTask:
 
+mjTask
+~~~~~~
+
+This is a representation of a task to be run asynchronously inside of an :ref:`mjThreadPool` . It is created in the
+:ref:`mju_threadPoolEnqueue` method of the :ref:`mjThreadPool`  and is used to join the task at completion.
+
+.. mujoco-include:: mjTask
+
+.. _mjThreadPool:
+
+mjThreadPool
+~~~~~~~~~~~~
+
+This is the data structure of the threadpool. It can only be constructed programmatically, and does not
+have an analog in MJCF. In order to enable multi-threaded calculations, a pointer to an existing :ref:`mjThreadPool`
+should be assigned to the ``mjData.threadpool``.
+
+.. mujoco-include:: mjThreadPool
 
 .. _tyStatStructure:
 

doc/APIreference/functions.rst

@@ -3433,3 +3433,43 @@ mjp_getResourceProviderAtSlot
 Look up a resource provider by slot number returned by mjp_registerResourceProvider.
 If invalid slot number, return NULL.
 
+.. _Thread:
+
+Thread
+^^^^^^
+.. _mju_threadPoolCreate:
+
+mju_threadPoolCreate
+~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mju_threadPoolCreate
+
+Creates a thread pool with the specified number of threads running.
+
+.. _mju_threadPoolEnqueue:
+
+mju_threadPoolEnqueue
+~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mju_threadPoolEnqueue
+
+Enqueues a task in a thread pool.
+
+.. _mju_taskJoin:
+
+mju_taskJoin
+~~~~~~~~~~~~
+
+.. mujoco-include:: mju_taskJoin
+
+Waits for a task to complete.
+
+.. _mju_threadPoolDestroy:
+
+mju_threadPoolDestroy
+~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mju_threadPoolDestroy
+
+Destroys a thread pool.
+

doc/changelog.rst

@@ -45,6 +45,7 @@ General
    ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s.
 #. Added new functions ``mj_markStack`` and ``mj_freeStack`` that manages ``mjData`` stack frames in a fully
    encapsulated way (i.e. without having to introduce a local variable at the call site).
+#. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within MuJoCo engine pipeline.
 
 Python bindings
 ^^^^^^^^^^^^^^^

doc/includes/references.h

@@ -375,6 +375,9 @@ struct mjData_ {
   mjtNum* efc_b;            // linear cost term: J*qacc_smooth - aref            (nefc x 1)
   mjtNum* efc_force;        // constraint force in constraint space              (nefc x 1)
   int*    efc_state;        // constraint state (mjtConstraintState)             (nefc x 1)
+
+  // ThreadPool for multithreaded operations
+  uintptr_t threadpool;
 };
 typedef struct mjData_ mjData;
 typedef enum mjtDisableBit_ {     // disable default feature bitflags
@@ -1436,6 +1439,14 @@ struct mjrContext_ {              // custom OpenGL context
   int     readPixelFormat;        // default color pixel format for mjr_readPixels
 };
 typedef struct mjrContext_ mjrContext;
+struct mjTask_ {
+  char buffer[48];
+};
+typedef struct mjTask_ mjTask;
+struct mjThreadPool_ {
+  char buffer[6208];
+};
+typedef struct mjThreadPool_ mjThreadPool;
 typedef enum mjtButton_ {         // mouse button
   mjBUTTON_NONE = 0,              // no button
   mjBUTTON_LEFT,                  // left button
@@ -2572,4 +2583,10 @@ int mjp_registerResourceProvider(const mjpResourceProvider* provider);
 int mjp_resourceProviderCount();
 const mjpResourceProvider* mjp_getResourceProvider(const char* resource_name);
 const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot);
+mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
+void mju_threadPoolEnqueue(
+ThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*),
+id* args);
+void mju_taskJoin(mjTask* task);
+void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 // NOLINTEND

doc/programming/index.rst

@@ -30,6 +30,10 @@ OpenGL renderer
    The renderer is written in C and is based on fixed-function OpenGL. It does not have all the features of
    state-of-the-art rendering engines (and can be replaced with such an engine if desired) but nevertheless it provides
    efficient and informative 3D rendering.
+Thread
+   The Threading framework (new in MuJoCo 3.0) is written in C++ and exposed in C. It provides a ThreadPool interface
+   to process Tasks asynchronously. To enable use in MuJoCo, create a ThreadPool and assign it to the thread_pool field
+   in mjData.
 UI framework
    The UI framework (new in MuJoCo 2.0) is written in C. UI elements are rendered in OpenGL. It has its own event
    mechanism and abstract hooks for keyboard and mouse input. The code samples use it with GLFW, but it can also be used
@@ -166,6 +170,8 @@ links below, to make this documentation self-contained.
    code.
 `mjplugin.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjplugin.h>`__
    Defines data structures required by :ref:`engine plugins<exPlugin>`.
+`mjthread.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`__
+   Defines data structures and functions required by :ref:`thread<Thread>`.
 
 .. _inVersion:
 

include/mujoco/mjdata.h

@@ -402,6 +402,9 @@ struct mjData_ {
   mjtNum* efc_b;            // linear cost term: J*qacc_smooth - aref            (nefc x 1)
   mjtNum* efc_force;        // constraint force in constraint space              (nefc x 1)
   int*    efc_state;        // constraint state (mjtConstraintState)             (nefc x 1)
+
+  // ThreadPool for multithreaded operations
+  uintptr_t threadpool;
 };
 typedef struct mjData_ mjData;
 

include/mujoco/mjthread.h

@@ -0,0 +1,47 @@
+// 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.
+
+#ifndef MUJOCO_INCLUDE_MJTHREAD_H_
+#define MUJOCO_INCLUDE_MJTHREAD_H_
+
+// C API for MuJoCo threading
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stddef.h>
+
+#include <mujoco/mjexport.h>
+
+// These types are implemented in C++, they're just used as opaque pointers in C
+// to provide type safety for functions.
+struct mjTask_ {
+  char buffer[48];
+};
+typedef struct mjTask_ mjTask;
+
+struct mjThreadPool_ {
+  char buffer[6208];
+};
+typedef struct mjThreadPool_ mjThreadPool;
+
+typedef void*(*mjStartRoutine_)(void*);
+typedef mjStartRoutine_ mjStartRoutine;
+
+#ifdef __cplusplus
+}
+#endif
+
+
+#endif  // MUJOCO_INCLUDE_MJTHREAD_H_

include/mujoco/mjxmacro.h

@@ -644,7 +644,8 @@
     X( int,       nnzJ               ) \
     X( int,       ncon               ) \
     X( int,       nisland            ) \
-    X( mjtNum,    time               )
+    X( mjtNum,    time               ) \
+    X( uintptr_t, threadpool         )
 
 
 // vector fields of mjData

include/mujoco/mujoco.h

@@ -36,6 +36,7 @@ extern "C" {
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjrender.h>
+#include <mujoco/mjthread.h>
 #include <mujoco/mjtnum.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
@@ -1304,6 +1305,21 @@ MJAPI const mjpResourceProvider* mjp_getResourceProvider(const char* resource_na
 // If invalid slot number, return NULL.
 MJAPI const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot);
 
+//---------------------- Thread -------------------------------------------------------------------
+
+// Creates a thread pool with the specified number of threads running.
+MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
+
+// Enqueues a task in a thread pool.
+MJAPI void mju_threadPoolEnqueue(
+    mjThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*),
+    void* args);
+
+// Waits for a task to complete.
+MJAPI void mju_taskJoin(mjTask* task);
+
+// Destroys a thread pool.
+MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 
 #if defined(__cplusplus)
 }

introspect/ast_nodes.py

@@ -66,6 +66,7 @@ class ValueType:
   def __init__(self, name: str, is_const: bool = False,
                is_volatile: bool = False):
     is_valid_type_name = (
+        name == 'void *(*)(void *)' or
         VALID_TYPE_NAME_PATTERN.fullmatch(name) or
         _is_valid_integral_type(name)) and name not in C_INVALID_TYPE_NAMES
     if not is_valid_type_name:

introspect/functions.py

@@ -8257,4 +8257,76 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Look up a resource provider by slot number returned by mjp_registerResourceProvider. If invalid slot number, return NULL.',  # pylint: disable=line-too-long
      )),
+    ('mju_threadPoolCreate',
+     FunctionDecl(
+         name='mju_threadPoolCreate',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjThreadPool'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='number_of_threads',
+                 type=ValueType(name='size_t'),
+             ),
+         ),
+         doc='Creates a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
+     )),
+    ('mju_threadPoolEnqueue',
+     FunctionDecl(
+         name='mju_threadPoolEnqueue',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='thread_pool',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjThreadPool'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='start_routine',
+                 type=ValueType(name='void *(*)(void *)'),
+             ),
+             FunctionParameterDecl(
+                 name='args',
+                 type=PointerType(
+                     inner_type=ValueType(name='void'),
+                 ),
+             ),
+         ),
+         doc='Enqueues a task in a thread pool.',
+     )),
+    ('mju_taskJoin',
+     FunctionDecl(
+         name='mju_taskJoin',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='task',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjTask'),
+                 ),
+             ),
+         ),
+         doc='Waits for a task to complete.',
+     )),
+    ('mju_threadPoolDestroy',
+     FunctionDecl(
+         name='mju_threadPoolDestroy',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='thread_pool',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjThreadPool'),
+                 ),
+             ),
+         ),
+         doc='Destroys a thread pool.',
+     )),
 ])

introspect/structs.py

@@ -4475,6 +4475,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='constraint state (mjtConstraintState)             (nefc x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='threadpool',
+                 type=ValueType(name='uintptr_t'),
+                 doc='ThreadPool for multithreaded operations',
+             ),
          ),
      )),
     ('mjvPerturb',
@@ -6978,6 +6983,36 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
+    ('mjTask',
+     StructDecl(
+         name='mjTask',
+         declname='struct mjTask_',
+         fields=(
+             StructFieldDecl(
+                 name='buffer',
+                 type=ArrayType(
+                     inner_type=ValueType(name='char'),
+                     extents=(48,),
+                 ),
+                 doc='',
+             ),
+         ),
+     )),
+    ('mjThreadPool',
+     StructDecl(
+         name='mjThreadPool',
+         declname='struct mjThreadPool_',
+         fields=(
+             StructFieldDecl(
+                 name='buffer',
+                 type=ArrayType(
+                     inner_type=ValueType(name='char'),
+                     extents=(6208,),
+                 ),
+                 doc='',
+             ),
+         ),
+     )),
     ('mjuiState',
      StructDecl(
          name='mjuiState',

introspect/type_parsing.py

@@ -68,6 +68,8 @@ def _parse_maybe_pointer(
                                                    ast_nodes.PointerType]]
 ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
   """Internal-only helper that parses a type that may be a pointer type."""
+  if type_name == 'void *(*)(void *)':
+    return ast_nodes.ValueType(name=type_name)
   p = type_name.rfind('*')
   if p != -1:
     leftover, is_qualifier = _parse_qualifiers(
@@ -107,6 +109,9 @@ def _peel_nested_parens(input_str: str) -> MutableSequence[str]:
     A sequence of substrings enclosed with in respective parentheses. See the
     description above for the precise detail of the output.
   """
+  if input_str == 'void *(*)(void *)':
+    return ['void *(*)(void *)']
+
   start = input_str.find('(')
   end = input_str.rfind(')')
 
@@ -146,4 +151,4 @@ def parse_type(
 def parse_function_return_type(
     type_name: str
 ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
-  return parse_type(type_name[:type_name.rfind('(')])
+  return parse_type(type_name[:type_name.find('(')])

src/engine/engine_io.c

@@ -1114,6 +1114,8 @@ static mjData* _makeData(const mjModel* m) {
     }
   }
 
+  d->threadpool = 0;
+
   return d;
 }
 
@@ -1202,6 +1204,8 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
     }
   }
 
+  dest->threadpool = src->threadpool;
+
   return dest;
 }
 

src/engine/engine_print.c

@@ -778,7 +778,8 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
 #define X(type, name)                                                         \
   if (strcmp(#name, "pstack") != 0 &&                                         \
       strcmp(#name, "pbase") != 0 &&                                          \
-      strcmp(#name, "parena") != 0) {                                         \
+      strcmp(#name, "parena") != 0 &&                                         \
+      strcmp(#name, "threadpool") != 0) {                                     \
     const char* format = _Generic(                                            \
         d->name,                                                              \
         int : INT_FORMAT,                                                     \
@@ -794,6 +795,16 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
 
   MJDATA_SCALAR
 #undef X
+
+  int threadpool = 0;
+  if (d->threadpool) {
+    threadpool = 1;
+  }
+  fprintf(fp, "  ");
+  fprintf(fp, NAME_FORMAT, "threadpool");
+  fprintf(fp, INT_FORMAT, threadpool);
+  fprintf(fp, "\n");
+
   fprintf(fp, "\n");
 
   // WARNING

src/thread/CMakeLists.txt

@@ -0,0 +1,23 @@
+# 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
+#
+#     https://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.
+
+set(MUJOCO_THREAD_SRCS
+    lockless_queue.h
+    task.cc
+    task.h
+    thread_pool.cc
+    thread_pool.h
+)
+
+target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS})

src/thread/lockless_queue.h

@@ -0,0 +1,154 @@
+// 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
+#define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
+
+#include <atomic>
+#include <climits>
+#include <cstddef>
+#include <thread>
+
+namespace mujoco {
+
+// A Lockless Queue allows for sending information quickly between different
+// threads. This is a Multi-Producer Multi-Consumer Lockless Queue allowing for
+// multiple threads to be adding items to the queue while multiple threads are
+// consuming items from the queue. Internally it uses a Ring Buffer for storage
+// so it will not grow as items are added. Push will block if the Queue is full
+// and Pop will block if it is empty.
+//
+// For a basic overview of this category of structures:
+// https://www.linuxjournal.com/content/lock-free-multi-producer-multi-consumer-queue-ring-buffer
+template  <typename T, size_t buffer_capacity>
+class LocklessQueue {
+ public:
+  bool full() const {
+    return full_internal(
+        convert_to_index(read_cursor_), convert_to_index(write_cursor_));
+  }
+
+  bool empty() const {
+    return maximum_read_cursor_ == read_cursor_;
+  }
+
+  // Push an element into the queue.
+  void push(const T& input) {
+    // Reserve a slot in the queue
+    size_t current_write_cursor;
+    size_t dummy_current_write_cursor;
+    size_t next_write_cursor;
+    size_t current_write_index;
+    size_t current_read_index;
+    do {
+      // Check if the queue is full.
+      do {
+        current_write_cursor = write_cursor_.load();
+        current_write_index = convert_to_index(current_write_cursor);
+        next_write_cursor = get_next_cursor(current_write_cursor);
+
+        current_read_index = convert_to_index(read_cursor_.load());
+      } while (full_internal(current_read_index, current_write_index));
+
+      // Once it's not full, attempt to grab a slot to write.
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!write_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+
+    // Write the entry.
+    buffer_[current_write_index].store(input);
+
+    // Increment maximum read cursor. Note here it has to wait if the compare
+    // and exchange fails as another thread might not have completed its write.
+    do {
+      dummy_current_write_cursor = current_write_cursor;
+    } while (!maximum_read_cursor_.compare_exchange_weak(
+                 dummy_current_write_cursor, next_write_cursor));
+  }
+
+  // Pop an element from the queue.
+  T pop() {
+    size_t current_read_cursor;
+    size_t dummy_current_read_cursor;
+    size_t current_read_index;
+    size_t next_read_cursor;
+    size_t current_maximum_read_cursor;
+    size_t current_maximum_read_index;
+    bool empty = false;
+    T result;
+    do {
+      // Wait until the queue has an element
+      do {
+        if (empty) {
+          std::this_thread::yield();
+        }
+        current_read_cursor = read_cursor_.load();
+        current_maximum_read_cursor = maximum_read_cursor_.load();
+
+        current_read_index = convert_to_index(current_read_cursor);
+        current_maximum_read_index = convert_to_index(
+            current_maximum_read_cursor);
+
+        empty = empty_internal(
+            current_read_index, current_maximum_read_index);
+      } while (empty);
+
+      next_read_cursor = get_next_cursor(current_read_cursor);
+
+      // Attempt to grab the element, if unsuccessful then wait for the next
+      // element to arrive.
+      result = buffer_[current_read_index].load();
+      dummy_current_read_cursor = current_read_cursor;
+    } while (!read_cursor_.compare_exchange_weak(
+                 dummy_current_read_cursor, next_read_cursor));
+
+    return result;
+  }
+
+ private:
+  size_t convert_to_index(size_t input) const {
+    return input % internal_buffer_capacity_;
+  }
+
+  size_t get_next_cursor(size_t input) const {
+    return (input + 1) % cursor_max_;
+  }
+
+  size_t get_next_index(size_t input) const {
+    return convert_to_index(get_next_cursor(input));
+  }
+
+  bool full_internal(size_t read_index, size_t write_index) const {
+    return get_next_index(write_index) == read_index;
+  }
+
+  bool empty_internal(size_t read_index, size_t write_index) const {
+    return read_index == write_index;
+  }
+
+  const size_t internal_buffer_capacity_ = buffer_capacity + 1;
+  const size_t cursor_max_ = UINT_MAX - (UINT_MAX % internal_buffer_capacity_);
+
+  std::atomic<size_t> read_cursor_ = 0;
+  std::atomic<size_t> write_cursor_ = 0;
+  std::atomic<size_t> maximum_read_cursor_ = 0;
+
+  std::atomic<T> buffer_[(buffer_capacity + 1)];
+};
+
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_

src/thread/task.cc

@@ -0,0 +1,24 @@
+// 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.
+
+#include "thread/task.h"
+
+#include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
+
+// waits for a task to complete
+void mju_taskJoin(mjTask* task) {
+  mujoco::Task* task_ptr = static_cast<mujoco::Task*>(static_cast<void*>(task));
+  task_ptr->Join();
+}

src/thread/task.h

@@ -0,0 +1,69 @@
+// 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_TASK_H_
+#define MUJOCO_SRC_THREAD_TASK_H_
+
+#ifdef __cplusplus
+
+#include <atomic>
+#include <functional>
+#include <thread>
+
+namespace mujoco {
+
+class Task {
+ public:
+  enum Status {
+    QUEUED,
+    COMPLETE,
+  };
+
+  static void Initialize(
+      Task* task,
+      std::function<void*(void*)> start_routine,
+      void* args) {
+    // instantiate a task at the pointer passed in
+    new(task) Task();
+    task->start_routine_ = start_routine;
+    task->args_ = args;
+    task->status_ = Status::QUEUED;
+  }
+
+  void Execute() {
+    args_ = start_routine_(args_);
+    status_ = Status::COMPLETE;
+  }
+
+  void Join() {
+    while (status_ != Status::COMPLETE) {
+      std::this_thread::yield();
+    }
+  }
+
+ private:
+  std::function<void*(void*)> start_routine_;
+
+  void* args_;
+
+  std::atomic<Status> status_ = Status::QUEUED;
+};
+
+}  // namespace mujoco
+
+#endif  // __cplusplus
+
+#endif  // MUJOCO_SRC_THREAD_TASK_H_

src/thread/thread_pool.cc

@@ -0,0 +1,52 @@
+// 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.
+
+#include "thread/thread_pool.h"
+
+#include <cstddef>
+
+#include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
+#include "thread/task.h"
+
+static constexpr size_t kMaxThreads = 128;
+
+// create a thread pool
+mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool =
+      new mujoco::ThreadPool<kMaxThreads>(number_of_threads);
+  return static_cast<mjThreadPool*>(static_cast<void*>(thread_pool));
+}
+
+// start a task in the threadpool
+void mju_threadPoolEnqueue(
+    mjThreadPool* thread_pool, mjTask* task, mjStartRoutine start_routine,
+    void* args) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
+      static_cast<mujoco::ThreadPool<kMaxThreads>*>(
+          static_cast<void*>(thread_pool));
+  thread_pool_ptr->Enqueue(
+      static_cast<mujoco::Task*>(static_cast<void*>(task)), start_routine,
+      args);
+}
+
+// shutdown the threadpool and free the memory
+void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
+  mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
+      static_cast<mujoco::ThreadPool<kMaxThreads>*>(
+          static_cast<void*>(thread_pool));
+  thread_pool_ptr->Shutdown();
+  delete thread_pool_ptr;
+}
+

src/thread/thread_pool.h

@@ -0,0 +1,103 @@
+// 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.
+// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h"
+// IWYU pragma: friend "third_party/(py/)?mujoco/.*"
+
+#ifndef MUJOCO_SRC_THREAD_THREAD_POOL_H_
+#define MUJOCO_SRC_THREAD_THREAD_POOL_H_
+
+#ifdef __cplusplus
+
+#include <atomic>
+#include <cstddef>
+#include <functional>
+#include <thread>
+
+#include "thread/lockless_queue.h"
+#include "thread/task.h"
+
+namespace mujoco {
+
+static constexpr size_t kThreadPoolQueueSize = 640;
+
+template <size_t max_number_of_threads>
+class ThreadPool {
+ public:
+  ThreadPool(size_t number_of_threads)
+      : number_of_threads_(number_of_threads) {
+    for (int i = 0; i < number_of_threads_; ++i) {
+      threads_[i] = std::thread(ThreadPoolWorker, static_cast<void*>(this));
+    }
+  }
+
+  // start a task in the threadpool
+  void Enqueue(
+      Task* task, std::function<void*(void*)> start_routine, void* args) {
+    Task::Initialize(task, start_routine, args);
+    lockless_queue_.push(static_cast<void*>(task));
+  }
+
+  // shutdown the threadpool
+  void Shutdown() {
+    if (shutdown_) {
+      return;
+    }
+
+    shutdown_ = true;
+    Task shutdown_tasks[number_of_threads_];
+    for (int i = 0; i < number_of_threads_; ++i) {
+      Enqueue(&shutdown_tasks[i], ShutdownFunction, nullptr);
+    }
+
+    for (int i = 0; i < number_of_threads_; ++i) {
+      threads_[i].join();
+    }
+  }
+
+  ~ThreadPool() { Shutdown(); }
+
+ private:
+  // method executed by running threads
+  static void ThreadPoolWorker(void* arg) {
+    ThreadPool<max_number_of_threads>* thread_pool =
+        static_cast<ThreadPool<max_number_of_threads>*>(arg);
+    while (!thread_pool->shutdown_) {
+      Task* task = static_cast<Task*>(thread_pool->lockless_queue_.pop());
+      task->Execute();
+    }
+  }
+
+  // shutdown function passed to running threads to ensure cleans shutdown
+  static void* ShutdownFunction(void* args) {
+    return NULL;
+  }
+
+  // is the thread pool is being shut down
+  std::atomic<bool> shutdown_ = false;
+
+  // actual number of running threads in the threadpool
+  const size_t number_of_threads_;
+
+  // OS threads that are running in this pool
+  std::thread threads_[max_number_of_threads];
+
+  // queue of tasks to execute
+  LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
+};
+
+}  // namespace mujoco
+
+#endif  // __cplusplus
+
+#endif  // MUJOCO_SRC_THREAD_THREAD_POOL_H_

test/thread/lockless_queue_test.cc

@@ -0,0 +1,46 @@
+// 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.
+
+#include "src/thread/lockless_queue.h"
+
+#include <cstddef>
+
+#include <gtest/gtest.h>
+
+namespace mujoco {
+namespace {
+
+constexpr size_t kBufferCapacity = 640;
+
+TEST(TestMujocoLocklessQueue, TestMujocoLocklessQueue) {
+  LocklessQueue<void*, 640> queue;
+  EXPECT_TRUE(queue.empty());
+  int test_integers[kBufferCapacity];
+  for (int h = 0; h < 10; ++h) {
+    for (int i = 0; i < kBufferCapacity; ++i) {
+      test_integers[i] = i;
+      queue.push(&test_integers[i]);
+    }
+    EXPECT_TRUE(queue.full());
+
+    for (int i = 0; i < kBufferCapacity; ++i) {
+      void* output_ptr = queue.pop();
+      ASSERT_EQ(output_ptr, &test_integers[i]);
+    }
+    EXPECT_TRUE(queue.empty());
+  }
+}
+
+}  // namespace
+}  // namespace mujoco

test/thread/mjthread_test.cc

@@ -0,0 +1,90 @@
+// 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.
+
+#include <mujoco/mjthread.h>
+
+#include <atomic>
+
+#include <gtest/gtest.h>
+#include <mujoco/mujoco.h>
+#include "src/thread/task.h"
+#include "src/thread/thread_pool.h"
+
+namespace {
+
+struct TestFunctionArgs_ {
+  int input;
+  // make this atomic to avoid red-herring tsan failures.
+  std::atomic<int> output;
+};
+typedef struct TestFunctionArgs_ TestFunctionArgs;
+
+void* test_function(void* args) {
+  TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
+  if (!test_function_args) {
+    return nullptr;
+  }
+  test_function_args->output = test_function_args->input;
+  return nullptr;
+}
+
+TEST(TestMjThreadPool, EnsureStructClassSizeMatch) {
+  EXPECT_EQ(sizeof(mjTask), sizeof(mujoco::Task));
+  EXPECT_EQ(sizeof(mjThreadPool), sizeof(mujoco::ThreadPool<128>));
+}
+
+TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
+  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
+
+  TestFunctionArgs test_function_args[1000];
+  mjTask tasks[1000];
+  for (int i = 0; i < 1000; ++i) {
+    test_function_args[i].input = i;
+    mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function,
+                          (void*)&test_function_args[i]);
+  }
+
+  for (int i = 0; i < 1000; ++i) {
+    mju_taskJoin(&tasks[i]);
+  }
+
+  for (int i = 0; i < 1000; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
+  }
+  mju_threadPoolDestroy(thread_pool);
+}
+
+TEST(TestMjThreadPool, TestMjThreadPool100Threads) {
+  mjThreadPool* thread_pool = mju_threadPoolCreate(100);
+
+  TestFunctionArgs test_function_args[1000];
+  mjTask tasks[1000];
+  for (int i = 0; i < 1000; ++i) {
+    test_function_args[i].input = i;
+    mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function,
+                          (void*)&test_function_args[i]);
+  }
+
+  for (int i = 0; i < 1000; ++i) {
+    mju_taskJoin(&tasks[i]);
+  }
+
+  for (int i = 0; i < 1000; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
+  }
+
+  mju_threadPoolDestroy(thread_pool);
+}
+
+}  // namespace

test/thread/task_test.cc

@@ -0,0 +1,46 @@
+// 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.
+
+#include "src/thread/task.h"
+
+#include <gtest/gtest.h>
+
+namespace mujoco {
+namespace {
+
+struct TestFunctionArgs {
+  int input;
+  int output;
+};
+
+void* test_function(void* args) {
+  TestFunctionArgs* test_function_args = (TestFunctionArgs*)args;
+  test_function_args->output = test_function_args->input;
+  return nullptr;
+}
+
+TEST(TestMjThread, TestMjThread) {
+  TestFunctionArgs test_function_args;
+  test_function_args.input = 1;
+  test_function_args.output = 2;
+  Task task;
+  Task::Initialize(
+      &task, test_function, static_cast<void*>(&test_function_args));
+  task.Execute();
+  task.Join();
+  EXPECT_EQ(test_function_args.input, test_function_args.output);
+}
+
+}  // namespace
+}  // namespace mujoco

test/thread/thread_pool_test.cc

@@ -0,0 +1,134 @@
+// 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.
+
+#include "src/thread/thread_pool.h"
+
+#include <atomic>
+#include <condition_variable>
+#include <memory>
+#include <mutex>
+#include <thread>
+
+#include <gtest/gtest.h>
+#include "src/thread/task.h"
+
+namespace mujoco {
+namespace {
+
+struct TestFunctionArgs {
+  int input;
+  // make this atomic to avoid red-herring tsan failures.
+  std::atomic<int> output;
+};
+
+void* test_function(void* args) {
+  TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
+  test_function_args->output = test_function_args->input;
+  return nullptr;
+}
+
+TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
+  ThreadPool<10> thread_pool(10);
+
+  constexpr int kTasks = 1000;
+  TestFunctionArgs test_function_args[kTasks];
+  Task tasks[kTasks];
+  for (int i = 0; i < kTasks; ++i) {
+    test_function_args[i].input = i;
+    thread_pool.Enqueue(
+        &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    tasks[i].Join();
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
+  }
+
+  thread_pool.Shutdown();
+}
+
+TEST(TestMjThreadPool, TestMjThreadPool100Threads) {
+  ThreadPool<100> thread_pool(100);
+
+  constexpr int kTasks = 1000;
+  TestFunctionArgs test_function_args[kTasks];
+  Task tasks[kTasks];
+  for (int i = 0; i < kTasks; ++i) {
+    test_function_args[i].input = i;
+    thread_pool.Enqueue(
+        &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    tasks[i].Join();
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
+  }
+
+  thread_pool.Shutdown();
+}
+
+TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
+  ThreadPool<10> thread_pool(10);
+
+  constexpr int kTasks = 20;
+  TestFunctionArgs test_function_args[kTasks];
+  Task tasks[kTasks];
+  std::unique_ptr<std::thread> enqueue_threads[kTasks];
+
+  // add tasks to the thread pool from many threads
+  std::condition_variable start_cv;
+  std::mutex start_mutex;
+  bool start = false;
+  for (int i = 0; i < kTasks; ++i) {
+    test_function_args[i].input = i;
+    enqueue_threads[i] = std::make_unique<std::thread>([&, i] {
+      // synchronize all threads adding to the thread_pool at the same time
+      {
+        std::unique_lock<std::mutex> lock(start_mutex);
+        start_cv.wait(lock, [&] { return start; });
+      }
+      // enqueue outside the lock, to get some concurrency
+      thread_pool.Enqueue(
+          &tasks[i], test_function, static_cast<void*>(&test_function_args[i]));
+    });
+  }
+  {
+    std::unique_lock<std::mutex> lock(start_mutex);
+    start = true;
+  }
+  start_cv.notify_all();
+
+  for (int i = 0; i < kTasks; ++i) {
+    enqueue_threads[i]->join();
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    tasks[i].Join();
+  }
+
+  for (int i = 0; i < kTasks; ++i) {
+    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
+  }
+
+  thread_pool.Shutdown();
+}
+
+}  // namespace
+}  // namespace mujoco

unity/Runtime/Bindings/MjBindings.cs

@@ -56,6 +56,7 @@ public const bool mjEXTERNC = true;
 public const bool THIRD_PARTY_MUJOCO_MJRENDER_H_ = true;
 public const int mjNAUX = 10;
 public const int mjMAXTEXTURE = 1000;
+public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTNUM_H_ = true;
 public const bool mjUSEDOUBLE = true;
 public const double mjMINVAL = 1e-15;
@@ -1736,6 +1737,7 @@ public unsafe struct mjData_ {
   public double* efc_b;
   public double* efc_force;
   public int* efc_state;
+  public UIntPtr threadpool;
 }
 
 [StructLayout(LayoutKind.Sequential)]
@@ -2354,6 +2356,16 @@ public unsafe struct mjrContext_ {
   public int readPixelFormat;
 }
 
+[StructLayout(LayoutKind.Sequential)]
+public unsafe struct mjTask_ {
+  public fixed sbyte buffer[48];
+}
+
+[StructLayout(LayoutKind.Sequential)]
+public unsafe struct mjThreadPool_ {
+  public fixed sbyte buffer[6208];
+}
+
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct mjuiState_ {
   public int nrect;
@@ -3930,5 +3942,14 @@ public static unsafe extern void mjd_subQuat(double* qa, double* qb, double* Da,
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mjd_quatIntegrate(double* vel, double scale, double* Dquat, double* Dvel, double* Dscale);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern mjThreadPool_* mju_threadPoolCreate(UIntPtr number_of_threads);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mju_taskJoin(mjTask_* task);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mju_threadPoolDestroy(mjThreadPool_* thread_pool);
 }
 }