Add threading primitives.

PiperOrigin-RevId: 562649644
Change-Id: I9c35b270e4b3b50cc7eb5a152f24d71def8e2fcd

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