Add threading primitives.

PiperOrigin-RevId: 562649644
Change-Id: I9c35b270e4b3b50cc7eb5a152f24d71def8e2fcd

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: