Add new mju_threadpool API function, and delete old threading API.

PiperOrigin-RevId: 922838541 Change-Id: Id9f7e0fb298ffde61fcc49a802dc78971858ce51
2026-05-28 10:09:07 -07:00
parent a22fc2423a
commit b935d4153c
47 changed files with 576 additions and 1755 deletions
@@ -118,7 +118,6 @@ set(MUJOCO_HEADERS
    include/mujoco/mjsan.h
    include/mujoco/mjspec.h
    include/mujoco/mjspecmacro.h
    include/mujoco/mjthread.h
    include/mujoco/mjtype.h
    include/mujoco/mjui.h
    include/mujoco/mjvisualize.h
@@ -135,6 +134,7 @@ endif()
 # Emscripten does not support SHARED libs
 if(NOT EMSCRIPTEN)
  find_package(Threads REQUIRED)
  add_library(mujoco SHARED ${MUJOCO_RESOURCE_FILES})
 else()
  add_library(mujoco STATIC ${MUJOCO_RESOURCE_FILES})
@@ -160,7 +160,6 @@ add_subdirectory(src/engine)
 add_subdirectory(src/user)
 add_subdirectory(src/xml)
 add_subdirectory(src/xml/mjz)
 add_subdirectory(src/thread)
 if(NOT EMSCRIPTEN AND NOT MUJOCO_USE_FILAMENT_MJR_COMPAT)
  add_subdirectory(src/render/classic)
  add_subdirectory(src/ui)
@@ -204,6 +203,10 @@ target_link_libraries(
          tinyxml2
 )
 if(NOT EMSCRIPTEN)
  target_link_libraries(mujoco PRIVATE Threads::Threads)
 endif()
 set_target_properties(
  mujoco PROPERTIES VERSION "${mujoco_VERSION}" PUBLIC_HEADER "${MUJOCO_HEADERS}"
 )
@@ -452,10 +452,6 @@ indicated otherwise.
   * - ``mjMINAWAKE``
     - 10
     - The minimum number of timesteps that must pass after a tree is awoken, before it is allowed to go back to sleep.
   * - ``mjMAXTHREAD``
     - 128
     - Maximum number of OS threads that can be used in a thread pool.
       Defined in `mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`_.
 .. _glNumericSizes:
@@ -1013,36 +1013,6 @@ Asset cache used by the compiler to avoid repeated slow recompilation. See :ref:
 .. mujoco-include:: mjCache
 .. _mjtTaskStatus:
 mjtTaskStatus
 ~~~~~~~~~~~~~
 Status values for :ref:`mjTask`.
 .. mujoco-include:: mjtTaskStatus
 .. _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.
 The ``status`` field uses values from :ref:`mjtTaskStatus`.
 .. 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:
@@ -3375,59 +3375,14 @@ If no match, return NULL.
 Threads
 ^^^^^^^
-.. _mju_threadPoolCreate:
+.. _mju_threadpool:
-`mju_threadPoolCreate <#mju_threadPoolCreate>`__
+`mju_threadpool <#mju_threadpool>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. mujoco-include:: mju_threadPoolCreate
+.. mujoco-include:: mju_threadpool
-Create a thread pool with the specified number of threads running.
+Create a thread pool with nthread worker threads.
 .. _mju_bindThreadPool:
 `mju_bindThreadPool <#mju_bindThreadPool>`__
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. mujoco-include:: mju_bindThreadPool
 Adds a thread pool to mjData and configures it for multi-threaded use.
 .. _mju_threadPoolEnqueue:
 `mju_threadPoolEnqueue <#mju_threadPoolEnqueue>`__
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. mujoco-include:: mju_threadPoolEnqueue
 Enqueue a task in a thread pool.
 .. _mju_threadPoolDestroy:
 `mju_threadPoolDestroy <#mju_threadPoolDestroy>`__
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. mujoco-include:: mju_threadPoolDestroy
 Destroy a thread pool.
 .. _mju_defaultTask:
 `mju_defaultTask <#mju_defaultTask>`__
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. mujoco-include:: mju_defaultTask
 Initialize an mjTask.
 .. _mju_taskJoin:
 `mju_taskJoin <#mju_taskJoin>`__
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. mujoco-include:: mju_taskJoin
 Wait for a task to complete.
 .. _Standardmath:
@@ -10,6 +10,16 @@ General
 - Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp<body-flexcomp>` element for
  programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as
  ``body.make_flex()`` in Python.
 - Added :ref:`mju_threadpool`, a new function for creating a thread pool on an ``mjData`` instance. When a thread pool
  is initialized, parts of the simulation pipeline, such as collision detection and constraint solving across islands,
  are parallelized. The thread pool is automatically destroyed when the ``mjData`` is freed.
  .. admonition:: Breaking API changes
     :class: attention
     - The header file ``mjthread.h`` was removed along with the engine threading API.
     **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
 Bug fixes
 ^^^^^^^^^
@@ -1708,7 +1718,7 @@ New features
   If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
   see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
-5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
+5. Added ``mjThreadPool`` and ``mjTask`` which allow for multi-threaded operations within the MuJoCo engine
   pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
   - Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
@@ -89,9 +89,12 @@ struct mjData_ {
  // arena pointer
  size_t  parena;            // first available byte in arena
  // threading
  uintptr_t threadpool;      // thread pool pointer
  mjtBool threadlock;        // disable stack freeing during threaded execution
  // memory utilization statistics
  mjtSize maxuse_stack;                       // maximum stack allocation in bytes (mutable)
  mjtSize maxuse_threadstack[mjMAXTHREAD];    // maximum stack allocation per thread in bytes
  mjtSize maxuse_arena;                       // maximum arena allocation in bytes
  int     maxuse_con;                         // maximum number of contacts
  int     maxuse_efc;                         // maximum number of scalar constraints
@@ -394,9 +397,6 @@ struct mjData_ {
  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
  mjtNum* ifrc_constraint;   // constraint force                                 (nidof x 1)
  // thread pool pointer
  uintptr_t threadpool;
  // compilation signature
  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
@@ -2146,21 +2146,6 @@ typedef struct mjsDefault_ {       // default specification
  mjsTendon* tendon;               // tendon defaults
  mjsActuator* actuator;           // actuator defaults
 } mjsDefault;
 typedef enum mjtTaskStatus_ {  // status values for mjTask
  mjTASK_NEW = 0,              // newly created
  mjTASK_QUEUED,               // enqueued in a thread pool
  mjTASK_COMPLETED             // completed execution
 } mjtTaskStatus;
 struct mjThreadPool_ {
  int nworker;  // number of workers in the pool
 };
 typedef struct mjThreadPool_ mjThreadPool;
 struct mjTask_ {        // a task that can be executed by a thread pool.
  mjfTask func;         // pointer to the function that implements the task
  void* args;           // arguments to func
  volatile int status;  // status of the task
 };
 typedef struct mjTask_ mjTask;
 typedef enum mjtDisableBit_ {     // disable default feature bitflags
  mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
  mjDSBL_EQUALITY     = 1<<1,     // equality constraints
@@ -3658,12 +3643,7 @@ void mju_getResourceDir(mjResource* resource, const char** dir, int* ndir);
 int mju_isModifiedResource(const mjResource* resource, const char* timestamp);
 mjSpec* mju_decodeResource(mjResource* resource, const char* content_type,
                           const mjVFS* vfs);
-mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
+void mju_threadpool(mjData* d, int nthread);
 void mju_bindThreadPool(mjData* d, void* thread_pool);
 void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
 void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 void mju_defaultTask(mjTask* task);
 void mju_taskJoin(mjTask* task);
 mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
                       const char* prefix, const char* suffix);
 mjsBody* mjs_addBody(mjsBody* body, const mjsDefault* def);
@@ -31,9 +31,8 @@ OpenGL renderer
   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 is written in C++ and exposed in C. It provides a :ref:`mjThreadPool<mjThreadPool>` interface
+   The threading framework is written in C++ and exposed in C. It provides a thread pool interface
-   to process tasks asynchronously. To enable use in MuJoCo, create a thread pool and assign it to the
+   to process tasks asynchronously. To enable use in MuJoCo, call ``mju_threadpool``.
   ``mjData.threadpool`` field.
 UI framework
   The UI framework 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
@@ -171,8 +170,6 @@ links below, to make this documentation self-contained.
   Defines enums and structs used for :doc:`procedural model editing <modeledit>`.
 `mjplugin.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjplugin.h>`__
   Defines data structures required by :ref:`engine plugins<exPlugin>`.
 `mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`__
   Defines data structures and functions required by :ref:`thread<Thread>`.
 `mjmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmacro.h>`__
   Defines C macros that are useful in user code.
 `mjxmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`__
@@ -60,7 +60,7 @@ Where the command line arguments are
  keyframe named "test" is present in the model, it is used as the initial state.
 - The ``ctrlnoise`` argument prevents models from settling into a static state where, due to warmstarts, one can
  measure artificially faster simulation.
- When ``npoolthread > 1`` is specified, an engine-internal :ref:`mjThreadPool` is created with the specified number of
+- When ``npoolthread > 1`` is specified, an engine-internal thread pool is created with the specified number of
  threads, to speed up simulation of large scenes. Note that while it is possible to use both ``nthread`` and
  ``npoolthread``, the scenarios for which one would want these different types of multithreading are usually mutually
  exclusive.
@@ -562,10 +562,9 @@ external force computed by inverse dynamics.
 Multi-threading
 ~~~~~~~~~~~~~~~
-MuJoCo has experimental support for within-step multi-threading. When a :ref:`mjThreadPool` is assigned to
+MuJoCo has support for within-step multi-threading. When a thread pool is initialized via
-``mjData.threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
+``mju_threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
-:ref:`islands<siSleep>` — can be distributed across worker threads. Note that within-step threading currently has
+:ref:`islands<siSleep>` — can be distributed across worker threads.
 significant memory overhead and is still a work in progress.
 The more common and well-supported use of multi-threading is to speed up sampling operations that are
 common in more advanced applications. Simulation is inherently serial over time (the output of one mj_step is the
@@ -20,7 +20,6 @@
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjrender.h>
 #include <mujoco/mjspec.h>
 #include <mujoco/mjthread.h>
 #include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
@@ -123,9 +122,6 @@ MJ_ASSERT_SIZE(mjtEvent, 4);
 MJ_ASSERT_SIZE(mjtItem, 4);
 MJ_ASSERT_SIZE(mjtSection, 4);
 // mjthread.h
 MJ_ASSERT_SIZE(mjtTaskStatus, 4);
 // mjplugin.h
 MJ_ASSERT_SIZE(mjtPluginCapabilityBit, 4);
@@ -20,7 +20,7 @@
 #include <mujoco/mjtype.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjthread.h>
+
@@ -112,9 +112,12 @@ struct mjData_ {
  // arena pointer
  size_t  parena;            // first available byte in arena
  // threading
  uintptr_t threadpool;      // thread pool pointer
  mjtBool threadlock;        // disable stack freeing during threaded execution
  // memory utilization statistics
  mjtSize maxuse_stack;                       // maximum stack allocation in bytes (mutable)
  mjtSize maxuse_threadstack[mjMAXTHREAD];    // maximum stack allocation per thread in bytes
  mjtSize maxuse_arena;                       // maximum arena allocation in bytes
  int     maxuse_con;                         // maximum number of contacts
  int     maxuse_efc;                         // maximum number of scalar constraints
@@ -417,9 +420,6 @@ struct mjData_ {
  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
  mjtNum* ifrc_constraint;   // constraint force                                 (nidof x 1)
  // thread pool pointer
  uintptr_t threadpool;
  // compilation signature
  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
@@ -1,42 +0,0 @@
 // 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_
 #define mjMAXTHREAD 128        // maximum number of threads in a thread pool
 typedef enum mjtTaskStatus_ {  // status values for mjTask
  mjTASK_NEW = 0,              // newly created
  mjTASK_QUEUED,               // enqueued in a thread pool
  mjTASK_COMPLETED             // completed execution
 } mjtTaskStatus;
 // function pointer type for mjTask
 typedef void* (*mjfTask)(void*);
 // An opaque type representing a thread pool.
 struct mjThreadPool_ {
  int nworker;  // number of workers in the pool
 };
 typedef struct mjThreadPool_ mjThreadPool;
 struct mjTask_ {        // a task that can be executed by a thread pool.
  mjfTask func;         // pointer to the function that implements the task
  void* args;           // arguments to func
  volatile int status;  // status of the task
 };
 typedef struct mjTask_ mjTask;
 #endif  // MUJOCO_INCLUDE_MJTHREAD_H_
@@ -1014,6 +1014,7 @@
    X( size_t,    pstack             ) \
    X( size_t,    pbase              ) \
    X( size_t,    parena             ) \
    X( uintptr_t, threadpool         ) \
    X( size_t,    maxuse_stack       ) \
    X( size_t,    maxuse_arena       ) \
    X( int,       maxuse_con         ) \
@@ -1036,13 +1037,11 @@
    X( mjtBool,   flg_energyvel      ) \
    X( mjtBool,   flg_subtreevel     ) \
    X( mjtBool,   flg_rnepost        ) \
-    X( mjtNum,    time               ) \
+    X( mjtNum,    time               )
    X( uintptr_t, threadpool         )
 // vector fields of mjData
 #define MJDATA_VECTOR                                                \
    X( size_t,         maxuse_threadstack, mjMAXTHREAD,  1         ) \
    X( mjSolverStat,   solver,             mjNISLAND,    mjNSOLVER ) \
    X( int,            solver_niter,       mjNISLAND,    1         ) \
    X( int,            solver_nnz,         mjNISLAND,    1         ) \
@@ -31,7 +31,6 @@
 #include <mujoco/mjrender.h>
 #include <mujoco/mjsan.h>
 #include <mujoco/mjspec.h>
 #include <mujoco/mjthread.h>
 #include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
@@ -1593,23 +1592,8 @@ MJAPI mjSpec* mju_decodeResource(mjResource* resource, const char* content_type,
 //---------------------------------- Threads -------------------------------------------------------
-// Create a thread pool with the specified number of threads running.
+// Create a thread pool with nthread worker threads.
-MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
+MJAPI void mju_threadpool(mjData* d, int nthread);
 // Adds a thread pool to mjData and configures it for multi-threaded use.
 MJAPI void mju_bindThreadPool(mjData* d, void* thread_pool);
 // Enqueue a task in a thread pool.
 MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
 // Destroy a thread pool.
 MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 // Initialize an mjTask.
 MJAPI void mju_defaultTask(mjTask* task);
 // Wait for a task to complete.
 MJAPI void mju_taskJoin(mjTask* task);
 //---------------------------------- Attachment ----------------------------------------------------
@@ -725,6 +725,9 @@ PYBIND11_MODULE(_functions, pymodule) {
  Def<traits::mj_version>(pymodule);
  Def<traits::mj_versionString>(pymodule);
  // Thread pool
  Def<traits::mju_threadpool>(pymodule);
  // Ray collision
  Def<traits::mj_multiRay>(
      pymodule,
@@ -609,16 +609,6 @@ ENUMS: Mapping[str, EnumDecl] = dict([
             ('mjS_AWAKE', 1),
         ]),
     )),
    ('mjtTaskStatus',
     EnumDecl(
         name='mjtTaskStatus',
         declname='enum mjtTaskStatus_',
         values=dict([
             ('mjTASK_NEW', 0),
             ('mjTASK_QUEUED', 1),
             ('mjTASK_COMPLETED', 2),
         ]),
     )),
    ('mjtGeomInertia',
     EnumDecl(
         name='mjtGeomInertia',
@@ -9985,23 +9985,9 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
         ),
         doc='Find the decoder for a resource and return the decoded spec. The caller takes ownership of the spec and is responsible for cleaning it up.',  # pylint: disable=line-too-long
     )),
-    ('mju_threadPoolCreate',
+    ('mju_threadpool',
     FunctionDecl(
-         name='mju_threadPoolCreate',
+         name='mju_threadpool',
         return_type=PointerType(
             inner_type=ValueType(name='mjThreadPool'),
         ),
         parameters=(
             FunctionParameterDecl(
                 name='number_of_threads',
                 type=ValueType(name='size_t'),
             ),
         ),
         doc='Create a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
     )),
    ('mju_bindThreadPool',
     FunctionDecl(
         name='mju_bindThreadPool',
         return_type=ValueType(name='void'),
         parameters=(
             FunctionParameterDecl(
@@ -10011,75 +9997,11 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                 ),
             ),
             FunctionParameterDecl(
-                 name='thread_pool',
+                 name='nthread',
-                 type=PointerType(
+                 type=ValueType(name='int'),
                     inner_type=ValueType(name='void'),
                 ),
             ),
         ),
-         doc='Adds a thread pool to mjData and configures it for multi-threaded use.',  # pylint: disable=line-too-long
+         doc='Create a thread pool with nthread worker threads.',
     )),
    ('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'),
                 ),
             ),
         ),
         doc='Enqueue a task in a thread pool.',
     )),
    ('mju_threadPoolDestroy',
     FunctionDecl(
         name='mju_threadPoolDestroy',
         return_type=ValueType(name='void'),
         parameters=(
             FunctionParameterDecl(
                 name='thread_pool',
                 type=PointerType(
                     inner_type=ValueType(name='mjThreadPool'),
                 ),
             ),
         ),
         doc='Destroy a thread pool.',
     )),
    ('mju_defaultTask',
     FunctionDecl(
         name='mju_defaultTask',
         return_type=ValueType(name='void'),
         parameters=(
             FunctionParameterDecl(
                 name='task',
                 type=PointerType(
                     inner_type=ValueType(name='mjTask'),
                 ),
             ),
         ),
         doc='Initialize an mjTask.',
     )),
    ('mju_taskJoin',
     FunctionDecl(
         name='mju_taskJoin',
         return_type=ValueType(name='void'),
         parameters=(
             FunctionParameterDecl(
                 name='task',
                 type=PointerType(
                     inner_type=ValueType(name='mjTask'),
                 ),
             ),
         ),
         doc='Wait for a task to complete.',
     )),
    ('mjs_attach',
     FunctionDecl(
@@ -5104,42 +5104,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
             ),
         ),
     )),
    ('mjThreadPool',
     StructDecl(
         name='mjThreadPool',
         declname='struct mjThreadPool_',
         fields=(
             StructFieldDecl(
                 name='nworker',
                 type=ValueType(name='int'),
                 doc='number of workers in the pool',
             ),
         ),
     )),
    ('mjTask',
     StructDecl(
         name='mjTask',
         declname='struct mjTask_',
         fields=(
             StructFieldDecl(
                 name='func',
                 type=ValueType(name='mjfTask'),
                 doc='pointer to the function that implements the task',
             ),
             StructFieldDecl(
                 name='args',
                 type=PointerType(
                     inner_type=ValueType(name='void'),
                 ),
                 doc='arguments to func',
             ),
             StructFieldDecl(
                 name='status',
                 type=ValueType(name='int', is_volatile=True),
                 doc='status of the task',
             ),
         ),
     )),
    ('mjPreContact',
     StructDecl(
         name='mjPreContact',
@@ -5422,19 +5386,21 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                 type=ValueType(name='size_t'),
                 doc='first available byte in arena',
             ),
             StructFieldDecl(
                 name='threadpool',
                 type=ValueType(name='uintptr_t'),
                 doc='thread pool pointer',
             ),
             StructFieldDecl(
                 name='threadlock',
                 type=ValueType(name='mjtBool'),
                 doc='disable stack freeing during threaded execution',
             ),
             StructFieldDecl(
                 name='maxuse_stack',
                 type=ValueType(name='mjtSize'),
                 doc='maximum stack allocation in bytes (mutable)',
             ),
             StructFieldDecl(
                 name='maxuse_threadstack',
                 type=ArrayType(
                     inner_type=ValueType(name='mjtSize'),
                     extents=(128,),
                 ),
                 doc='maximum stack allocation per thread in bytes',
             ),
             StructFieldDecl(
                 name='maxuse_arena',
                 type=ValueType(name='mjtSize'),
@@ -6903,11 +6869,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                 doc='constraint force',
                 array_extent=('nidof',),
             ),
             StructFieldDecl(
                 name='threadpool',
                 type=ValueType(name='uintptr_t'),
                 doc='thread pool pointer',
             ),
             StructFieldDecl(
                 name='signature',
                 type=ValueType(name='uint64_t'),
@@ -218,8 +218,7 @@ int main(int argc, char** argv) {
    // make and bind threadpool
    if (npoolthread > 1) {
-      mjThreadPool* threadpool = mju_threadPoolCreate(npoolthread);
+      mju_threadpool(d[id], npoolthread);
      mju_bindThreadPool(d[id], threadpool);
    }
  }
@@ -348,11 +347,7 @@ int main(int argc, char** argv) {
  // free per-thread data
  for (int id=0; id < nthread; id++) {
    mjThreadPool* threadpool = (mjThreadPool*) d[id]->threadpool;
    mj_deleteData(d[id]);
    if (threadpool) {
      mju_threadPoolDestroy(threadpool);
    }
  }
  // finalize
@@ -1642,33 +1642,36 @@ void UiEvent(mjuiState* state) {
    // simulation section
    else if (it && it->sectionid==SECT_SIMULATION) {
      switch (it->itemid) {
-      case 1:             // Reset
+      case 1:             // Threadpool
        sim->pending_.update_threadpool = true;
        break;
      case 2:             // Reset
        sim->pending_.reset = true;
        break;
-      case 2:             // Reload
+      case 3:             // Reload
        sim->uiloadrequest.fetch_add(1);
        break;
-      case 3:             // Align
+      case 4:             // Align
        sim->pending_.align = true;
        break;
-      case 4:             // Copy key
+      case 5:             // Copy key
        sim->pending_.copy_key = true;
        sim->pending_.copy_key_full_precision = sim->platform_ui->IsShiftKeyPressed();
        break;
-      case 5:             // Adjust key
+      case 6:             // Adjust key
-      case 6:             // Load key
+      case 7:             // Load key
        sim->pending_.load_key = true;
        break;
-      case 7:             // Save key
+      case 8:             // Save key
        sim->pending_.save_key = true;
        break;
-      case 11:            // History scrubber
+      case 12:            // History scrubber
        sim->run = 0;
        sim->pending_.load_from_history = true;
        mjui0_update_section(sim, SECT_SIMULATION);
@@ -2174,6 +2177,11 @@ void Simulate::Sync(bool state_only) {
    pending_.print_data = std::nullopt;
  }
  if (pending_.update_threadpool) {
    mju_threadpool(d_, nthread);
    pending_.update_threadpool = false;
  }
  if (pending_.reset) {
    mj_resetData(m_, d_);
    mj_forward(m_, d_);
@@ -2561,13 +2569,13 @@ void Simulate::LoadOnRenderThread() {
  }
  // set keyframe range and divisions
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.range[0] = 0;
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.range[0] = 0;
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.range[1] = mjMAX(0, this->m_->nkey - 1);
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.range[1] = mjMAX(0, this->m_->nkey - 1);
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.divisions = mjMAX(1, this->m_->nkey - 1);
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.divisions = mjMAX(1, this->m_->nkey - 1);
  // set scrubber range and divisions
-  this->ui0.sect[SECT_SIMULATION].item[11].slider.range[0] = 1 - nhistory_;
+  this->ui0.sect[SECT_SIMULATION].item[12].slider.range[0] = 1 - nhistory_;
-  this->ui0.sect[SECT_SIMULATION].item[11].slider.divisions = nhistory_;
+  this->ui0.sect[SECT_SIMULATION].item[12].slider.divisions = nhistory_;
  // detect image sensors for visualization
  DetectImageSensors(this, this->m_);
@@ -153,6 +153,7 @@ class Simulate {
    std::optional<std::string> save_mjb;
    std::optional<std::string> print_model;
    std::optional<std::string> print_data;
    bool update_threadpool;
    bool reset;
    bool align;
    bool copy_key;
@@ -205,6 +206,10 @@ class Simulate {
  // simulation
  int run = 1;
  // number of workers in threadpool
  int nthread = 0;
  // atomics for cross-thread messages
  std::atomic_int exitrequest = 0;
  std::atomic_int droploadrequest = 0;
@@ -320,9 +325,10 @@ class Simulate {
  // simulation section of UI
-  const mjuiDef def_simulation[14] = {
+  const mjuiDef def_simulation[15] = {
    {mjITEM_SECTION,   "Simulation",    mjPRESERVE, nullptr,     "AS"},
    {mjITEM_RADIO,     "",              5, &this->run,           "Pause\nRun"},
    {mjITEM_EDITINT,   "Num threads",   5, &this->nthread,       "1"},
    {mjITEM_BUTTON,    "Reset",         2, nullptr,              " #259"},
    {mjITEM_BUTTON,    "Reload",        5, nullptr,              "CL"},
    {mjITEM_BUTTON,    "Align",         2, nullptr,              "CA"},
@@ -75,6 +75,8 @@ set(MUJOCO_ENGINE_SRCS
    engine_sort.h
    engine_support.c
    engine_support.h
    engine_thread.cc
    engine_thread.h
    engine_util_blas.c
    engine_util_blas.h
    engine_util_blas_avx.h
@@ -26,6 +26,7 @@
 #include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_thread.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_util.h"
 #include "engine/engine_inline.h"
@@ -178,10 +179,8 @@ static inline mjtNum getGap(const mjModel* m, int g1, int g2, int ipair) {
 static inline void resetArena(mjData* d) {
  d->parena = d->ncon * sizeof(mjContact);
 #ifdef ADDRESS_SANITIZER
  if (!d->threadpool) {
    ASAN_POISON_MEMORY_REGION(
      (char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
  }
 #endif
 }
@@ -937,6 +936,7 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 }
 // binary search between two bodyflex trees
 void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                    int merged, int startadr, int pairadr) {
  int nbody = m->nbody, nbvhstatic = m->nbvhstatic;
@@ -1814,13 +1814,67 @@ static void mj_makeCapsule(const mjModel* m, mjData* d, int f, const int vid[2],
 }
 // struct for collision task
 typedef struct {
  mjPreContact* conbuffer;  // pre-contact buffer returned by collision functions
  int* nconbuffer;          // contact count for each collision pair
  char* epabuffer;          // buffer for nativeccd
  int ccd_size;             // size of nativeccd buffer
  const int* pairbuffer;    // collision pairs (g1, g2, ipair, index into conbuffer)
  int npair;                // number of collision pairs
  int chunksize;            // number of pairs to process per task
  int maxcon;               // maximum number of contacts (size of conbuffer)
 } mjContactArg;
 static void collisionTask(const mjModel* m, mjData* d, void* arg, int thread_id, int idx) {
  mjContactArg* conargs = (mjContactArg*)arg;
  mjPreContact* conbuffer = conargs->conbuffer;
  char* epabuffer = conargs->epabuffer;
  int chunksize = conargs->chunksize;
  int globalidx = chunksize * idx;
  const int* pair = conargs->pairbuffer + 4 * globalidx;
  int* ncon = conargs->nconbuffer + chunksize * idx;
  int npair = conargs->npair;
  int n = mjMIN(chunksize, npair - globalidx);
  mjc_setCCDBuffer(epabuffer + thread_id * conargs->ccd_size);
  for (int i = 0; i < n; i++) {
    int g1 = pair[4*i + 0];
    int g2 = pair[4*i + 1];
    int ipair = pair[4*i + 2];
    int conpos = pair[4*i + 3];
    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
    mjtNum margin = getMargin(m, g1, g2, ipair);
    mjtNum gap = getGap(m, g1, g2, ipair);
    ncon[i] = collision_func(m, d, conbuffer + conpos, g1, g2, margin + gap);
    // SHOULD NOT OCCUR
    int expected_max = (globalidx + i + 1 < npair ? pair[4*(i+1) + 3] : conargs->maxcon) - conpos;
    if (ncon[i] > expected_max) {
      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
              "expected at most %d from mj_maxContact", ncon[i], g1, g2, expected_max);
    }
  }
  mjc_setCCDBuffer(NULL);
 }
 // compute contacts for a batch of collision pairs contained in a buffer of
 // stride 3 ints (g1, g2, ipair)
 // if buffer is NULL, results are read from arena starting at parena
 void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena) {
  int nthread = mju_numThread(d);
  int ccd_size = mjc_ccdSize(m->opt.ccd_iterations);
  mjtNum margin, gap;
  // try to balance load of 5 chunks per thread (chunksize should be divisible by 16)
  int chunksize = npair / mjMAX(1, 5 * nthread);
  chunksize = mjMAX(16, (chunksize + 15) & ~15);  // round up to next 16
  int nchunk = (npair + chunksize - 1) / chunksize;
  // set buffer and arena pointer
  if (!buffer) {
    buffer = (const int*) ((char*) d->arena + parena);
@@ -1830,9 +1884,6 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
  mj_markStack(d);
  // buffer store how many contacts are generated for each pair
  int* nconbuffer = mj_stackAllocInt(d, npair);
  // buffer for pair data (g1, g2, ipair, index into conbuffer)
  int* pairbuffer = mj_stackAllocInt(d, 4 * npair);
  int maxcon = 0;
@@ -1850,42 +1901,30 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
    maxcon += mj_maxContact(m, g1, g2, margin + gap > 0);
  }
  // buffer for precontact data
  mjPreContact* conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
  // buffer data has been copied to metadata on the stack;
  // reclaim arena space so contacts can overwrite the buffer region
  d->parena = parena;
-  // set buffer for nativeccd
+  mjContactArg arg;
-  mj_markStack(d);
+  arg.ccd_size = ccd_size;
-  mjc_setCCDBuffer(mj_stackAllocByte(d, ccd_size, sizeof(mjtNum)));
+  arg.pairbuffer = pairbuffer;
  arg.nconbuffer = mjSTACKALLOC(d, npair, int);
  arg.conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
  arg.npair = npair;
  arg.chunksize = chunksize;
  arg.maxcon = maxcon;
-  for (int i = 0; i < npair; i++) {
+  // dispatch narrowphase to threads with local stack allocation for EPA
-    int g1 = pairbuffer[4*i + 0];
+  {
-    int g2 = pairbuffer[4*i + 1];
+    mj_markStack(d);
-    int ipair = pairbuffer[4*i + 2];
+    arg.epabuffer = mj_stackAllocByte(d, ccd_size * nthread, sizeof(mjtNum));
-    int idx = pairbuffer[4*i + 3];
+    mju_dispatch(m, d, collisionTask, &arg, nchunk);
-    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
+    mj_freeStack(d);
    margin = getMargin(m, g1, g2, ipair);
    gap = getGap(m, g1, g2, ipair);
    nconbuffer[i] = collision_func(m, d, conbuffer + idx, g1, g2, margin + gap);
    // SHOULD NOT OCCUR
    int expected_max = (i + 1 < npair ? pairbuffer[4*(i+1) + 3] : maxcon) - idx;
    if (nconbuffer[i] > expected_max) {
      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
              "expected at most %d from mj_maxContact", nconbuffer[i], g1, g2, expected_max);
    }
  }
  // set nativeccd buffer back to NULL
  mjc_setCCDBuffer(NULL);
  mj_freeStack(d);
  int ncon = 0;
  for (int i = 0; i < npair; i++) {
-    ncon += nconbuffer[i];
+    ncon += arg.nconbuffer[i];
  }
  if (ncon == 0) {
@@ -1906,7 +1945,7 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
  // fill in contact data
  int conpos = 0;
  for (int i = 0; i < npair; i++) {
-    if (!(ncon = nconbuffer[i]))
+    if (!(ncon = arg.nconbuffer[i]))
      continue;
    int condim;
@@ -1928,7 +1967,7 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
      mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
    }
-    mjPreContact* bc = conbuffer + pairbuffer[4*i + 3];
+    mjPreContact* bc = arg.conbuffer + pairbuffer[4*i+3];
    margin = getMargin(m, g1, g2, ipair);
    for (int j=0; j < ncon; j++) {
      mjContact* c = con + conpos + j;
@@ -44,8 +44,7 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_solve.h"
 #include "engine/engine_util_sparse.h"
-#include "thread/thread_pool.h"
+#include "engine/engine_thread.h"
 #include "thread/thread_task.h"
@@ -116,28 +115,6 @@ 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_makeM(forward_args->m, forward_args->d);
  mj_factorM(forward_args->m, forward_args->d);
  return NULL;
 }
 // wrapper for mj_collision
 void* mj_collisionThreaded(void* args) {
  mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args;
  mj_collision(forward_args->m, forward_args->d);
  return NULL;
 }
 // kinematics-related computations
 void mj_fwdKinematics(const mjModel* m, mjData* d) {
  mj_kinematics(m, d);
@@ -162,36 +139,12 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
  TM_END(mjTIMER_POS_KINEMATICS);
-  // no threadpool: inertia and collision on main thread
+  // inertia, timed internally (POS_INERTIA)
-  if (!d->threadpool) {
+  mj_makeM(m, d);
-    // inertia, timed internally (POS_INERTIA)
+  mj_factorM(m, d);
    mj_makeM(m, d);
    mj_factorM(m, d);
-    // collision, timed internally (POS_COLLISION)
+  // collision, timed internally (POS_COLLISION)
-    mj_collision(m, d);
+  mj_collision(m, d);
  }
  // have threadpool: inertia and collision on separate 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]);
  }
  if (mj_wakeCollision(m, d)) {
    mj_updateSleep(m, d);
@@ -909,51 +862,13 @@ static void warmstart(const mjModel* m, mjData* d) {
 }
-// struct encapsulating arguments to thread task
+// mju_dispatch callback: solve one island
-struct mjSolIslandArgs_ {
+static void solveIslandTask(const mjModel* m, mjData* d, void* arg, int thread_id, int island) {
-  const mjModel* m;
+  if (m->opt.solver == mjSOL_NEWTON) {
-  mjData* d;
+    mj_solNewton_island(m, d, island, m->opt.iterations);
-  int island;
+  } else {
-};
+    mj_solCG_island(m, d, island, m->opt.iterations);
 typedef struct mjSolIslandArgs_ mjSolIslandArgs;
 // extract arguments, pass to CG solver
 static void* CG_wrapper(void* args) {
  mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
  mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
  return NULL;
 }
 // extract arguments, pass to Newton solver
 static void* Newton_wrapper(void* args) {
  mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
  mj_solNewton_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
  return NULL;
 }
 // CG solver, multi-threaded over islands
 static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
  mj_markStack(d);
  // allocate array of arguments to be passed to threads
  mjSolIslandArgs* sol_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs);
  mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask);
  for (int island = 0; island < d->nisland; ++island) {
    sol_island_args[island].m = m;
    sol_island_args[island].d = d;
    sol_island_args[island].island = island;
    mju_defaultTask(&tasks[island]);
    tasks[island].func = flg_Newton ? Newton_wrapper : CG_wrapper;
    tasks[island].args = &sol_island_args[island];
    mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]);
  }
  for (int island = 0; island < d->nisland; ++island) {
    mju_taskJoin(&tasks[island]);
  }
  mj_freeStack(d);
 }
@@ -1009,20 +924,7 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
      mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
      mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
-      // solve per island, with or without threads
+      mju_dispatch(m, d, solveIslandTask, NULL, nisland);
      if (!d->threadpool) {
        // no threadpool, loop over islands
        for (int island=0; island < nisland; island++) {
          if (m->opt.solver == mjSOL_NEWTON) {
            mj_solNewton_island(m, d, island, m->opt.iterations);
          } else {
            mj_solCG_island(m, d, island, m->opt.iterations);
          }
        }
      } else {
        // have threadpool, solve using threads
        solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
      }
      // copy back solver outputs (scatter dofs since ni <= nv)
      mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
@@ -33,6 +33,7 @@
 #include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_sleep.h"
 #include "engine/engine_thread.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -1081,6 +1082,7 @@ void mj_makeRawData(mjData** dest, const mjModel* m) {
  // clear threadpool
  d->threadpool = 0;
  d->threadlock = 0;
  // clear nplugin (overwritten by _initPlugin)
  d->nplugin = 0;
@@ -1140,6 +1142,7 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
  *dest = *src;
  dest->buffer = save_buffer;
  dest->arena = save_arena;
  dest->threadpool = 0;
  mj_setPtrData(m, dest);
  // save plugin_data, since the X macro copying block below will override it
@@ -1239,8 +1242,6 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
    }
  }
  dest->threadpool = src->threadpool;
  return dest;
 }
@@ -1300,7 +1301,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
  // clear memory utilization stats
  d->maxuse_stack = 0;
  memset(d->maxuse_threadstack, 0, mjMAXTHREAD*sizeof(mjtSize));
  d->maxuse_arena = 0;
  d->maxuse_con = 0;
  d->maxuse_efc = 0;
@@ -1572,6 +1572,7 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) {
 // de-allocate mjData
 void mj_deleteData(mjData* d) {
  if (d) {
    mju_threadpool(d, 0);
    freeDataBuffers(d);
    mju_free(d);
  }
@@ -16,6 +16,7 @@
 #include <inttypes.h>  // IWYU pragma: keep
 #include <limits.h>
 #include <stdatomic.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
@@ -26,7 +27,7 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_util_errmem.h"
-#include "thread/thread_pool.h"
+
 #ifdef ADDRESS_SANITIZER
  #include <sanitizer/asan_interface.h>
@@ -57,25 +58,20 @@ static inline size_t fastmod(size_t a, size_t b) {
  return a % b;
 }
 typedef struct {
  uintptr_t bottom;      // first memory address available to the stack
  uintptr_t top;         // current memory address used by the stack
  uintptr_t limit;       // top limit of the stack (stack grows down)
  uintptr_t stack_base;  // current stack base for mark and free stack
 } mjStackInfo;
 typedef struct {
  size_t pbase;   // value of d->pbase immediately before mj_markStack
  size_t pstack;  // value of d->pstack immediately before mj_markStack
  void* pc;       // program counter of the call site of mj_markStack (only set when under asan)
 } mjStackFrame;
 static void maybe_lock_alloc_mutex(mjData* d) {
  if (d->threadpool != 0) {
    mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
  }
 }
 static void maybe_unlock_alloc_mutex(mjData* d) {
  if (d->threadpool != 0) {
    mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
  }
 }
 static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
  mjStackInfo stack_info;
  stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
@@ -110,14 +106,12 @@ static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
 // allocate memory from the mjData arena
 void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
  maybe_lock_alloc_mutex(d);
  size_t misalignment = fastmod(d->parena, alignment);
  size_t padding = misalignment ? alignment - misalignment : 0;
  // check size
  size_t bytes_available = d->narena - d->pstack;
  if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
    maybe_unlock_alloc_mutex(d);
    return NULL;
  }
@@ -125,16 +119,8 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
  // under ASAN, get stack usage from red zone
 #ifdef ADDRESS_SANITIZER
-  mjStackInfo stack_info;
+  mjStackInfo stack_info = get_stack_info_from_data(d);
-  mjStackInfo* stack_info_ptr;
+  stack_usage = stack_usage_redzone(&stack_info);
  if (!d->threadpool) {
    stack_info = get_stack_info_from_data(d);
    stack_info_ptr = &stack_info;
  } else {
    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
    stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
  }
  stack_usage = stack_usage_redzone(stack_info_ptr);
 #endif
  // allocate, update max, return pointer to buffer
@@ -150,7 +136,6 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
  __msan_allocated_memory(result, bytes);
 #endif
  maybe_unlock_alloc_mutex(d);
  return result;
 }
@@ -212,13 +197,8 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
  // update max usage statistics
  stack_info->top = new_top_ptr;
-  if (!d->threadpool) {
+  d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
-    d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
    d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
  } else {
    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
    d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
  }
  return (void*)start_ptr;
 }
@@ -228,18 +208,46 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
 // declared inline so that modular arithmetic with specific alignments can be optimized out
 static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
                               const char* caller, int line) {
-  // single threaded allocation
+  // size zero: no-op
-  if (!d->threadpool) {
+  if (!size) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
+    return NULL;
    void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
    d->pstack = stack_info.bottom - stack_info.top;
    return result;
  }
-  // multi threaded allocation
+  // call in mju_dispatch: atomically reserve space on the stack
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  if (d->threadlock) {
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+    size_t alloc_size = size + alignment - 1 + 2 * mjREDZONE;
-  return stackallocinternal(d, stack_info, size, alignment, caller, line);
+    size_t old_pstack = atomic_fetch_add_explicit(
        (_Atomic size_t*)&d->pstack, alloc_size, memory_order_relaxed);
    // check for stack overflow
    size_t stack_available_bytes = (size_t)d->narena - d->parena;
    if (mjUNLIKELY(old_pstack + alloc_size > stack_available_bytes)) {
      char info[1024];
      if (caller) {
        snprintf(info, sizeof(info), " at %s, line %d", caller, line);
      } else {
        info[0] = '\0';
      }
      mju_error(
          "mj_stackAlloc: out of memory, stack overflow%s (threadlock)\n"
          "  max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
          "\n nefc = %d, ncon = %d",
          info, (uintptr_t)stack_available_bytes,
          (uintptr_t)(stack_available_bytes - old_pstack),
          (uintptr_t)alloc_size, d->nefc, d->ncon);
    }
    uintptr_t bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
    uintptr_t start_ptr = bottom - old_pstack - size - mjREDZONE;
    start_ptr -= fastmod(start_ptr, alignment);
    ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
    return (void*)start_ptr;
  }
  mjStackInfo stack_info = get_stack_info_from_data(d);
  void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
  d->pstack = stack_info.bottom - stack_info.top;
  return result;
 }
@@ -268,17 +276,15 @@ void mj_markStack(mjData* d)
 void mj__markStack(mjData* d)
 #endif
 {
-  if (!d->threadpool) {
+  // no-op if called from mju_dispatch
-    mjStackInfo stack_info = get_stack_info_from_data(d);
+  if (d->threadlock) {
    markstackinternal(d, &stack_info);
    d->pstack = stack_info.bottom - stack_info.top;
    d->pbase = stack_info.stack_base;
    return;
  }
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo stack_info = get_stack_info_from_data(d);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  markstackinternal(d, &stack_info);
-  markstackinternal(d, stack_info);
+  d->pstack = stack_info.bottom - stack_info.top;
  d->pbase = stack_info.stack_base;
 }
@@ -319,30 +325,14 @@ void mj_freeStack(mjData* d)
 void mj__freeStack(mjData* d)
 #endif
 {
-  if (!d->threadpool) {
+  if (d->threadlock) {
    mjStackInfo stack_info = get_stack_info_from_data(d);
    freestackinternal(&stack_info);
    d->pstack = stack_info.bottom - stack_info.top;
    d->pbase = stack_info.stack_base;
    return;
  }
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
+  mjStackInfo stack_info = get_stack_info_from_data(d);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
+  freestackinternal(&stack_info);
-  freestackinternal(stack_info);
+  d->pstack = stack_info.bottom - stack_info.top;
-}
+  d->pbase = stack_info.stack_base;
 // returns the number of bytes available on the stack
 size_t mj_stackBytesAvailable(mjData* d) {
  if (!d->threadpool) {
    mjStackInfo stack_info = get_stack_info_from_data(d);
    return stack_info.top - stack_info.limit;
  } else {
    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
    mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
    return stack_info->top - stack_info->limit;
  }
 }
@@ -49,9 +49,6 @@ void mj__freeStack(mjData* d) __attribute__((noinline));
 #endif  // ADDRESS_SANITIZER
 // returns the number of bytes available on the stack
 MJAPI size_t mj_stackBytesAvailable(mjData* d);
 // allocate bytes on the stack
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
@@ -35,8 +35,7 @@
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_spatial.h"
-#include "thread/thread_pool.h"
+#include "engine/engine_thread.h"
 #include "thread/thread_task.h"
@@ -64,8 +63,6 @@ mjPARTIAL_SORT(ContactSelect, ContactInfo, ContactInfoCompare);
 // arguments for parallel tactile sensor computation
 typedef struct mjTactileTaskArgs_ {
  const mjModel* m;
  mjData* d;
  int sensor_id;
  int mesh_id;
  int geom_id;
@@ -80,10 +77,8 @@ typedef struct mjTactileTaskArgs_ {
 // worker function for parallel tactile computation over taxel batches
-static void* tactile_taxel_batch(void* args) {
+static void* tactile_taxel_batch(const mjModel* m, mjData* d, void* args) {
  mjTactileTaskArgs* t = (mjTactileTaskArgs*)args;
  const mjModel* m = t->m;
  mjData* d = t->d;
  int mesh_id = t->mesh_id;
  int geom_id = t->geom_id;
  int parent_weld = t->parent_weld;
@@ -193,6 +188,12 @@ static void* tactile_taxel_batch(void* args) {
 }
 static void tactileTask(const mjModel* m, mjData* d, void* arg, int thread_id, int task_id) {
  mjTactileTaskArgs* args_array = (mjTactileTaskArgs*)arg;
  tactile_taxel_batch(m, d, &args_array[task_id]);
 }
 // apply cutoff to sensor i, clamping values in data buffer
 static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
  mjtNum cutoff = m->sensor_cutoff[i];
@@ -1261,18 +1262,15 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
      // threshold for parallelization (taxel count below which sequential is faster)
      const int kTactileParallelThreshold = 1000;
-      // parallel path: use threadpool to process taxel batches
+      // parallel path: use mj_batch to process taxel batches
-      if (d->threadpool && ncon >= kTactileParallelThreshold) {
+      int nthread = mju_numThread(d);
-        int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool);
+      if (nthread > 0 && ncon >= kTactileParallelThreshold) {
-        int batch_size = (ncon + nthreads - 1) / nthreads;
+        int batch_size = (ncon + nthread - 1) / nthread;
-        int ntasks = (ncon + batch_size - 1) / batch_size;
+        int ntask = (ncon + batch_size - 1) / batch_size;
-        mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask);
+        mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntask, mjTactileTaskArgs);
        mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs);
-        for (int t = 0; t < ntasks; t++) {
+        for (int t = 0; t < ntask; t++) {
          task_args[t].m = m;
          task_args[t].d = d;
          task_args[t].sensor_id = i;
          task_args[t].mesh_id = mesh_id;
          task_args[t].geom_id = geom_id;
@@ -1283,22 +1281,13 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
          task_args[t].start_taxel = t * batch_size;
          task_args[t].end_taxel = mju_min((t+1) * batch_size, ncon);
          task_args[t].forcesT = forcesT;
          mju_defaultTask(&tasks[t]);
          tasks[t].func = tactile_taxel_batch;
          tasks[t].args = &task_args[t];
          mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[t]);
        }
-        for (int t = 0; t < ntasks; t++) {
+        mju_dispatch(m, d, tactileTask, task_args, ntask);
          mju_taskJoin(&tasks[t]);
        }
      }
      // sequential path: call tactile_taxel_batch with full range
      else {
        mjTactileTaskArgs args;
        args.m = m;
        args.d = d;
        args.sensor_id = i;
        args.mesh_id = mesh_id;
        args.geom_id = geom_id;
@@ -1309,7 +1298,7 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
        args.start_taxel = 0;
        args.end_taxel = ncon;
        args.forcesT = forcesT;
-        tactile_taxel_batch(&args);
+        tactile_taxel_batch(m, d, &args);
      }
      // compute sensor output
@@ -0,0 +1,190 @@
 // Copyright 2026 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 "engine/engine_thread.h"
 #include <atomic>
 #include <cstdint>
 #include <thread>
 #include <vector>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_memory.h"
 // context for thread pool stored on mjData
 class ThreadPoolContext {
 public:
  explicit ThreadPoolContext(int nthread) : threads_(nthread) {
    for (int i = 0; i < nthread; i++) {
      threads_[i] = std::thread(&ThreadPoolContext::Worker, this, i + 1);
    }
  }
  // non-copyable, non-movable
  ThreadPoolContext(const ThreadPoolContext&) = delete;
  ThreadPoolContext& operator=(const ThreadPoolContext&) = delete;
  ~ThreadPoolContext() {
    signal_.store(0, std::memory_order_release);
    signal_.notify_all();
    for (auto& thread : threads_) {
      if (thread.joinable()) {
        thread.join();
      }
    }
  }
  // dispatch tasks to the thread pool and work on them on the main thread
  void Dispatch(const mjModel* model, mjData* data, mjTaskFunc func, void* arg,
                int ntask) {
    func_ = func;
    model_ = model;
    data_ = data;
    arg_ = arg;
    ntask_ = ntask;
    next_.store(0, std::memory_order_relaxed);
    ndone_.store(0, std::memory_order_relaxed);
    signal_.store(-signal_.load(std::memory_order_relaxed),
                 std::memory_order_release);
    signal_.notify_all();
    // process tasks on main thread
    while (true) {
      int taskId = next_.fetch_add(1, std::memory_order_relaxed);
      if (taskId >= ntask_) {
        break;
      }
      func_(model_, data_, arg_, 0, taskId);
    }
    // busy wait for rest of workers to finish
    int nthread = threads_.size();
    while (ndone_.load(std::memory_order_acquire) < nthread) {
    }
  }
  int ThreadCount() const { return threads_.size(); }
 private:
  // worker loop for each worker thread
  void Worker(int threadId) {
    int status = 1;
    // main loop waiting for next batch of tasks
    while (true) {
      // wait until signal atomic is notified and sign flips
      signal_.wait(status, std::memory_order_acquire);
      // if signal was set to zero, halt
      status = signal_.load(std::memory_order_acquire);
      if (status == 0) {
        return;
      }
      // subloop to process tasks for the current batch
      while (true) {
        int taskId = next_.fetch_add(1, std::memory_order_relaxed);
        if (taskId >= ntask_) {
          break;
        }
        func_(model_, data_, arg_, threadId, taskId);
      }
      // let main thread know this worker is done
      ndone_.fetch_add(1, std::memory_order_release);
    }
  }
  // arguments for the current batch set by Dispatch
  const mjModel* model_;
  mjData* data_;
  mjTaskFunc func_;
  void* arg_;
  int ntask_;  // total number of tasks for workers to do
  // atomic for each worker to grab the next task
  std::atomic<int> next_{0};
  // atomic counter for number of workers who completed their tasks
  alignas(64) std::atomic<int> ndone_{0};
  // alternating signal from -1, 1 to start / halt the worker threads,
  // set to 0 to force all workers to exit
  std::atomic<int> signal_{1};
  std::vector<std::thread> threads_;
 };
 // create a thread pool with nthread threads
 void mju_threadpool(mjData* d, int nthread) {
  if (d->threadpool) {
    ThreadPoolContext* ctx =
        reinterpret_cast<ThreadPoolContext*>(d->threadpool);
    // same size, nothing to do
    if (nthread == ctx->ThreadCount()) {
        return;
    }
    delete ctx;
    d->threadpool = 0;  // null out in case nthread == 0
  }
  if (nthread >= 1) {
    d->threadpool = reinterpret_cast<uintptr_t>(new ThreadPoolContext(nthread));
  }
 }
 // dispatch ntask tasks to the thread pool; passes arg into func along with
 // thread_id and task_id
 void mju_dispatch(const mjModel* m, mjData* d, mjTaskFunc func, void* arg,
                  int ntask) {
  // no thread pool or trivial number of tasks: run on main thread
  if (!d->threadpool || ntask < 2) {
    for (int i = 0; i < ntask; i++) {
      func(m, d, arg, 0, i);
    }
    return;
  }
  ThreadPoolContext& ctx = *reinterpret_cast<ThreadPoolContext*>(d->threadpool);
  // lock mjData and mark stack frame, memory will be freed after thread completion
  if (!d->threadlock) {
    mj_markStack(d);
    d->threadlock = true;
  }
  ctx.Dispatch(m, d, func, arg, ntask);
  if (d->threadlock) {
    // update max usage statistics
    d->maxuse_stack = mjMAX(d->maxuse_stack, d->pstack);
    d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack + d->parena);
    // unlock mjData and free stack used during worker execution
    d->threadlock = false;
    mj_freeStack(d);
  }
 }
 // return total number of threads in the pool (including main thread)
 int mju_numThread(const mjData* d) {
  ThreadPoolContext* ctx = reinterpret_cast<ThreadPoolContext*>(d->threadpool);
  return ctx ? ctx->ThreadCount() + 1 : 1;
 }
@@ -0,0 +1,41 @@
 // Copyright 2026 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_SRC_ENGINE_ENGINE_THREAD_H_
 #define MUJOCO_SRC_ENGINE_ENGINE_THREAD_H_
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
 #ifdef __cplusplus
 extern "C" {
 #endif
 // dispatch function for mju_dispatch
 typedef void (*mjTaskFunc)(const mjModel* m, mjData* d, void* arg, int thread_id, int task_id);
 // create a thread pool with nthread worker threads.
 MJAPI void mju_threadpool(mjData* d, int nthread);
 // return total number of threads in the pool (including main thread)
 MJAPI int mju_numThread(const mjData* d);
 // dispatch ntask tasks to the thread pool; passes arg into func along with thread_id and task_id
 MJAPI void mju_dispatch(const mjModel* m, mjData* d, mjTaskFunc func, void* arg, int ntask);
 #ifdef __cplusplus
 }
 #endif
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_THREAD_H_
@@ -1,23 +0,0 @@
 # 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
    thread_pool.cc
    thread_pool.h
    thread_queue.h
    thread_task.cc
    thread_task.h
 )
 target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS})
@@ -1,312 +0,0 @@
 // 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 <stdint.h>
 #include <algorithm>
 #include <atomic>
 #include <cstddef>
 #include <memory>
 #include <mutex>
 #include <thread>
 #include <utility>
 #include <vector>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjthread.h>
 #include <mujoco/mujoco.h>
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_util_errmem.h"
 #include "thread/thread_queue.h"
 #include "thread/thread_task.h"
 namespace mujoco {
 namespace {
 constexpr size_t kThreadPoolQueueSize = 640;
 // Each thread being run will be assigned a worker_id.
 //   0: main thread
 //   1->n: workers
 thread_local size_t worker_id = 0;
 struct WorkerThread {
  // Shutdown function passed to running threads to ensure clean shutdown.
  static void* ShutdownFunction(void* args) {
    return nullptr;
  }
  // Thread for the worker.
  std::unique_ptr<std::thread> thread_;
  // An mjTask for shutting down this worker.
  mjTask shutdown_task_ {
    &ShutdownFunction,
    nullptr,
    mjTASK_NEW
  };
 };
 }  // namespace
 // Concrete C++ class definition for mjThreadPool.
 // (The public mjThreadPool C struct is an opaque one.)
 class ThreadPoolImpl : public mjThreadPool {
 public:
  ThreadPoolImpl(int num_worker) : mjThreadPool{num_worker} {
    // initialize worker threads
    for (int i = 0; i < std::min(num_worker, mjMAXTHREAD); ++i) {
      WorkerThread worker{
        std::make_unique<std::thread>(ThreadPoolWorker, this, i)};
      workers_.push_back(std::move(worker));
    }
  }
  size_t NumberOfThreads() {
    return workers_.size();
  }
  // start a task in the threadpool
  void Enqueue(mjTask* task) {
    if (mjUNLIKELY(GetAtomicTaskStatus(task).exchange(mjTASK_QUEUED) !=
                   mjTASK_NEW)) {
      mjERROR("task->status is not mjTASK_NEW");
    }
    lockless_queue_.push(task);
  }
  // shutdown the threadpool
  void Shutdown() {
    if (shutdown_) {
      return;
    }
    shutdown_ = true;
    std::vector<mjTask> shutdown_tasks(workers_.size());
    for (auto& worker : workers_) {
      Enqueue(&worker.shutdown_task_);
    }
    for (auto& worker : workers_) {
      worker.thread_->join();
    }
  }
  // registers a worker ID for a given thread
  void RegisterWorker(const size_t input_worker_id) {
    worker_id = input_worker_id;
  }
  // gets the worker id of the current thread
  size_t GetWorkerId() {
    return worker_id;
  }
  void LockAlloc() {
    alloc_mutex_.lock();
  }
  void UnlockAlloc() {
    alloc_mutex_.unlock();
  }
  bool IsThreadPoolBound() {
    return thread_pool_bound_;
  }
  void BindThreadPool() {
    thread_pool_bound_ = true;
  }
  ~ThreadPoolImpl() { Shutdown(); }
 private:
  // method executed by running threads
  static void ThreadPoolWorker(
      ThreadPoolImpl* thread_pool, const size_t thread_index) {
    worker_id = thread_index + 1;
    while (!thread_pool->shutdown_) {
      auto task = static_cast<mjTask*>(thread_pool->lockless_queue_.pop());
      task->args = task->func(task->args);
      GetAtomicTaskStatus(task).store(mjTASK_COMPLETED);
    }
  }
  // indicates whether the thread pool is being shut down
  std::atomic<bool> shutdown_ = false;
  // OS threads that are running in this pool
  std::vector<WorkerThread> workers_;
  // queue of tasks to execute
  mujoco::LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
  // Mutex to protect arena allocations.
  std::mutex alloc_mutex_;
  // Whether or not a ThreadPool was bound using mju_bindThreadPool.
  bool thread_pool_bound_ = false;
 };
 // create a thread pool
 mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
  return reinterpret_cast<mjThreadPool*>(new ThreadPoolImpl(number_of_threads));
 }
 // gets the number of shards the stack is currently broken into
 static size_t GetNumberOfShards(mjData* d) {
  if (!d->threadpool) {
    return 1;
  }
  return mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool) + 1;
 }
 // returns the stack information for the specified thread's shard
 mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id) {
  auto thread_pool = (ThreadPoolImpl*)d->threadpool;
  if (!thread_pool || !thread_pool->IsThreadPoolBound()) {
    mju_error("Thread Pool not bound, use mju_bindThreadPool to add an mjThreadPool to mjData");
  }
  // number of threads running in the threadpool plus the main thread
  size_t number_of_shards = GetNumberOfShards(d);
  // size of entire arena/stack in bytes
  size_t total_arena_size_bytes = d->narena;
  // set the shard cursor to the end of the arena
  uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + total_arena_size_bytes;
  // each thread including the main one will get an equal shard of the stack
  size_t bytes_per_shard = total_arena_size_bytes / (2 * (number_of_shards));
  // ensure the shard is larger than the cache line
  size_t misalignment = bytes_per_shard % mju_getDestructiveInterferenceSize();
  if (misalignment != 0) {
    bytes_per_shard += mju_getDestructiveInterferenceSize() - misalignment;
  }
  if (bytes_per_shard * number_of_shards > total_arena_size_bytes) {
    mju_error("Arena is not large enough for %zu shards", number_of_shards);
  }
  uintptr_t result = (end_of_arena_ptr - (thread_id + 1) * bytes_per_shard);
  // align the end of the shard to be mjStackInfo.
  misalignment = result % alignof(mjStackInfo);
  result -= misalignment;
 #ifdef ADDRESS_SANITIZER
  // Ensure StackInfo is always accessible
  ASAN_UNPOISON_MEMORY_REGION((void*)result, sizeof(mjStackInfo));
 #endif
  return (mjStackInfo*) result;
 }
 // shards the stack for each thread
 static void  ConfigureMultiThreadedStack(mjData* d) {
  if (!d->threadpool) {
    mju_error("No thread pool specified for multithreaded operation");
  }
  size_t number_of_shards = GetNumberOfShards(d);
  // current top of the stack
  uintptr_t current_limit = (uintptr_t)d->arena + d->narena - d->pstack;
  // set the shard cursor to the end of the arena
  uintptr_t begin_shard_cursor_ptr = (uintptr_t)d->arena + d->narena;
  for (size_t shard_index = 0; shard_index < number_of_shards; ++shard_index) {
    mjStackInfo* end_shard_cursor_ptr = mju_getStackInfoForThread(d, shard_index);
 #ifdef ADDRESS_SANITIZER
  // unpoison stack info
  ASAN_UNPOISON_MEMORY_REGION((void*)end_shard_cursor_ptr, sizeof(mjStackInfo));
 #endif
    // handle the main thread's stack which may already have data in it
    if (shard_index == 0) {
      // abort if the current stack is already larger than the portion of the stack
      // that would be reserved for the main thread
      if ((uintptr_t)end_shard_cursor_ptr > current_limit) {
        mju_error("mj_bindThreadPool: sharding stack - existing stack larger than shard size: current_size = %zu, "
                  "max_size = %zu", current_limit, (uintptr_t) end_shard_cursor_ptr);
      }
      end_shard_cursor_ptr->top = current_limit;
      end_shard_cursor_ptr->stack_base = d->pbase;
    } else {
      // all other stacks are empty because threads have not been used yet
      end_shard_cursor_ptr->top = begin_shard_cursor_ptr;
      end_shard_cursor_ptr->stack_base = 0;
    }
    end_shard_cursor_ptr->bottom = begin_shard_cursor_ptr;
    end_shard_cursor_ptr->limit = (uintptr_t)end_shard_cursor_ptr + sizeof(mjStackInfo);
    begin_shard_cursor_ptr = (uintptr_t)end_shard_cursor_ptr - 1;
  }
 }
 // adds a thread pool to mjData and configures it for multi-threaded use.
 void mju_bindThreadPool(mjData* d, void* thread_pool) {
  if (d->threadpool) {
    mju_error("Thread Pool already bound to mjData");
  }
  d->threadpool = (uintptr_t) thread_pool;
  ((ThreadPoolImpl*)thread_pool)->BindThreadPool();
  ConfigureMultiThreadedStack(d);
 }
 // gets the number of running threads in the thread pool.
 size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  return thread_pool_impl->NumberOfThreads();
 }
 size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  return thread_pool_impl->GetWorkerId();
 }
 // start a task in the threadpool
 void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  thread_pool_impl->Enqueue(task);
 }
 // shutdown the threadpool and free the memory
 void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  thread_pool_impl->Shutdown();
  delete thread_pool_impl;
 }
 // locks the allocation mutex to protect Stack and Arena allocations
 void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  thread_pool_impl->LockAlloc();
 }
 // unlocks the allocation mutex to protect Stack and Arena allocations
 void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool) {
  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
  thread_pool_impl->UnlockAlloc();
 }
 // Get the destructive interference size for the architecture.
 size_t mju_getDestructiveInterferenceSize(void) {
  // return std::hardware_destructive_interference_size;
  return 128;
 }
 }  // namespace mujoco
@@ -1,85 +0,0 @@
 // 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_SRC_THREAD_THREAD_POOL_H_
 #define MUJOCO_SRC_THREAD_THREAD_POOL_H_
 #include <stddef.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjthread.h>
 #include <mujoco/mujoco.h>
 #ifdef __cplusplus
 namespace mujoco {
 extern "C" {
 #endif
 // MultiThreaded Stack will be an approximately 50/50 split of the entire buffer, with a little
 // wiggle for alignment and caching concerns. The basic layout is to reuse the existing single
 // threaded markers, and then create shards for each thread to use as its stack.
 // Not to scale.
 // |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
 // |Used Arena|Free Arena |Shard1     |Shard1     |Shard1     |Shard0    |Shard0     |Shard0     |
 // |%%%%%%%%%%|           |StackInfo  |Free Stack |Used Stack |StackInfo |Free Stack |Used Stack |
 // |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
 // |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
 // |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
 // d->arena   d->parena   d->pstack   shard1->stack_info      shard1->bottom_of_stack            shard1->bottom_of_stack
 //                        shard1->stack_info                  shard0->stack_info     shard0->current_stack
 //                                    shard1->top_of_stack               shard1->top_of_stack
 //                                                shard1->current_stack
 typedef struct {
  uintptr_t bottom;          // First memory address available to the stack
  uintptr_t top;             // Current memory address used by the stack
  uintptr_t limit;             // Top limit of the stack (note this is smaller than bottom, stack grows down)
  uintptr_t stack_base;               // Current stack base for mark and free stack
 } mjStackInfo;
 // Create a thread pool with the specified number of threads running.
 MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
 // Returns the stack information for the specified thread's shard.
 mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id);
 // Adds a thread pool to mjData and configures it for multi-threaded use.
 MJAPI void mju_bindThreadPool(mjData* d, void* thread_pool);
 // Gets the number of running threads in the thread pool.
 MJAPI size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool);
 // Gets the ID of the current thread being executed
 MJAPI size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool);
 // Enqueue a task in a thread pool.
 MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
 // Locks the allocation mutex to protect Arena allocations.
 MJAPI void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool);
 // Unlocks the allocation mutex to protect Arena allocations.
 MJAPI void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool);
 // Destroy a thread pool.
 MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 // Get the destructive interference size for the architecture.
 MJAPI size_t mju_getDestructiveInterferenceSize(void);
 #ifdef __cplusplus
 }  // extern "C"
 }  // namespace mujoco
 #endif  // __cplusplus
 #endif  // MUJOCO_SRC_THREAD_THREAD_POOL_H_
@@ -1,152 +0,0 @@
 // 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_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_
@@ -1,33 +0,0 @@
 // 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_task.h"
 #include <thread>
 #include <mujoco/mjthread.h>
 namespace mujoco {
 void mju_defaultTask(mjTask* task) {
  task->func = nullptr;
  task->args = nullptr;
  task->status = mjTASK_NEW;
 }
 void mju_taskJoin(mjTask* task) {
  while (GetAtomicTaskStatus(task) != mjTASK_COMPLETED) {
    std::this_thread::yield();
  }
 }
 }  // namespace mujoco
@@ -1,49 +0,0 @@
 // 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_SRC_THREAD_THREAD_TASK_H_
 #define MUJOCO_SRC_THREAD_THREAD_TASK_H_
 #include <mujoco/mjexport.h>
 #include <mujoco/mjthread.h>
 #ifdef __cplusplus
 #include <atomic>
 #include <new>
 #include <type_traits>
 namespace mujoco {
 extern "C" {
 #endif
 // Initialize an mjTask.
 MJAPI void mju_defaultTask(mjTask* task);
 // Wait for a task to complete.
 MJAPI void mju_taskJoin(mjTask* task);
 #ifdef __cplusplus
 }  // extern "C"
 using TaskStatus = std::remove_volatile_t<decltype(mjTask::status)>;
 inline std::atomic<TaskStatus>& GetAtomicTaskStatus(mjTask* task) {
  static_assert(sizeof(std::atomic<TaskStatus>) == sizeof(TaskStatus));
  static_assert(alignof(std::atomic<TaskStatus>) == alignof(TaskStatus));
  static_assert(std::atomic<TaskStatus>::is_always_lock_free);
  return *std::launder(reinterpret_cast<std::atomic<TaskStatus>*>(
      const_cast<TaskStatus*>(&task->status)));
 }
 }  // namespace mujoco
 #endif  // __cplusplus
 #endif  // MUJOCO_SRC_THREAD_THREAD_TASK_H_
@@ -15,8 +15,8 @@
 // A benchmark for parsing and compiling models from XML.
 #include <cstddef>
 #include <vector>
 #include <string>
 #include <vector>
 #include <benchmark/benchmark.h>
 #include <absl/base/attributes.h>
@@ -36,6 +36,9 @@ static const int kNumWarmupSteps = 1000;
 // number of steps to benchmark (before resetting state)
 static const int kBatchSize = 50;
 // number of threads to test
 static const int kNumThreads = 9;
 static const char kBoxMeshPath[] =
    "../test/engine/testdata/collision_convex/perf/boxmesh.xml";
 static const char kBoxBoxPath[] =
@@ -47,11 +50,10 @@ static const char kMixedPath[] =
 class TestHarness {
 public:
-  TestHarness(const char* xml_path, std::string label, int disable_flags = 0) {
+  TestHarness(const char* xml_path, std::string label) {
    // Fail test if there are any mujoco errors
    MujocoErrorTestGuard guard;
    model_ = LoadModelFromPath(xml_path);
    model_->opt.disableflags |= disable_flags;
    data_ = mj_makeData(model_);
    for (int i=0; i < kNumWarmupSteps; i++) {
      mj_step(model_, data_);
@@ -62,13 +64,18 @@ class TestHarness {
    int size = mj_stateSize(model_, spec_);
    initial_state_.resize(size);
    mj_getState(model_, data_, initial_state_.data(), spec_);
-    label_ = label;
+    name_ = label;
  }
  void Reset() {
    mj_setState(model_, data_, initial_state_.data(), spec_);
  }
  void SetThreads(int nthread) {
    mju_threadpool(data_, nthread);
    nthread_ = nthread;
  }
  void RunBenchmark(benchmark::State& state) {
    std::size_t ncon = 0;
    while (state.KeepRunningBatch(kBatchSize)) {
@@ -81,7 +88,9 @@ class TestHarness {
      }
    }
-    state.SetLabel(label_);
+    std::string label = name_ + " " + std::to_string(nthread_ + 1) +
                        " thread(s)";
    state.SetLabel(label);
    state.SetItemsProcessed(ncon);  // report number of contacts per second
  }
@@ -91,7 +100,8 @@ class TestHarness {
  }
 private:
-  std::string label_;
+  std::string name_;
  int nthread_ = 0;
  int spec_;
  mjModel* model_;
  mjData* data_;
@@ -103,62 +113,47 @@ class TestHarness {
 // separately in CPU profiles (and don't get replaced with raw calls to
 // run_parse_benchmark).
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxMesh(benchmark::State& state) {
-    BM_BoxMesh_NativeCCD(benchmark::State& state) {
+  int nthread = state.range(0);
-  static TestHarness harness(kBoxMeshPath, "boxmesh.xml (nativeccd)");
+  static TestHarness harness(kBoxMeshPath, "boxmesh.xml");
  harness.SetThreads(nthread);
  harness.RunBenchmark(state);
 }
-BENCHMARK(BM_BoxMesh_NativeCCD);
+BENCHMARK(BM_BoxMesh)->Arg(0)->Arg(kNumThreads);
 void ABSL_ATTRIBUTE_NO_TAIL_CALL
    BM_BoxMesh_LibCCD(benchmark::State& state) {
  static TestHarness harness(kBoxMeshPath, "boxmesh.xml (libccd)",
                             mjDSBL_NATIVECCD);
  harness.RunBenchmark(state);
 }
 BENCHMARK(BM_BoxMesh_LibCCD);
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox(benchmark::State& state) {
-  static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)", mjDSBL_NATIVECCD);
+  int nthread = state.range(0);
  static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)");
  harness.SetThreads(nthread);
  harness.RunBenchmark(state);
 }
-BENCHMARK(BM_BoxBox);
+BENCHMARK(BM_BoxBox)->Arg(0)->Arg(kNumThreads);
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox_NativeCCD(benchmark::State& state) {
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBoxConvex(benchmark::State& state) {
  int nthread = state.range(0);
  mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_Convex;
  static TestHarness harness(kBoxBoxPath, "box.xml (NativeCCD)");
  harness.SetThreads(nthread);
  harness.RunBenchmark(state);
  mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_BoxBox;
 }
-BENCHMARK(BM_BoxBox_NativeCCD);
+BENCHMARK(BM_BoxBoxConvex)->Arg(0)->Arg(kNumThreads);
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Ellipsoid(benchmark::State& state) {
-    BM_Ellipsoid_NativeCCD(benchmark::State& state) {
+  int nthread = state.range(0);
-  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (nativeccd)");
+  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml");
  harness.SetThreads(nthread);
  harness.RunBenchmark(state);
 }
-BENCHMARK(BM_Ellipsoid_NativeCCD);
+BENCHMARK(BM_Ellipsoid)->Arg(0)->Arg(kNumThreads);
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed(benchmark::State& state) {
-    BM_Ellipsoid_LibCCD(benchmark::State& state) {
+  int nthread = state.range(0);
-  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (libccd)",
+  static TestHarness harness(kMixedPath, "mixed.xml");
-                             mjDSBL_NATIVECCD);
+  harness.SetThreads(nthread);
  harness.RunBenchmark(state);
 }
-BENCHMARK(BM_Ellipsoid_LibCCD);
+BENCHMARK(BM_Mixed)->Arg(0)->Arg(kNumThreads);
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_NativeCCD(benchmark::State& state) {
  static TestHarness harness(kMixedPath, "mixed.xml (nativeccd)");
  harness.RunBenchmark(state);
 }
 BENCHMARK(BM_Mixed_NativeCCD);
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_LibCCD(benchmark::State& state) {
  static TestHarness harness(kMixedPath, "mixed.xml (libccd)",
                             mjDSBL_NATIVECCD);
  harness.RunBenchmark(state);
 }
 BENCHMARK(BM_Mixed_LibCCD);
 }  // namespace
 }  // namespace mujoco
@@ -17,7 +17,6 @@
 #include <benchmark/benchmark.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjthread.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
@@ -30,6 +29,9 @@ static const int kNumWarmupSteps = 500;
 // number of steps to benchmark (before resetting state)
 static const int kBatchSize = 50;
 // number of threads to test
 static const int kNumThreads = 6;
 void BM_StepHumanoid200(benchmark::State& state) {
  int nthread = state.range(0);
  std::string label = std::to_string(nthread) + " thread(s)";
@@ -43,10 +45,8 @@ void BM_StepHumanoid200(benchmark::State& state) {
  model->opt.disableflags &= ~mjDSBL_ISLAND;  // enable islands
  mjData* data = mj_makeData(model);
-  mjThreadPool* threadpool = nullptr;
+  if (nthread) {
-  if (nthread > 1) {
+    mju_threadpool(data, nthread);
    threadpool = mju_threadPoolCreate(nthread);
    mju_bindThreadPool(data, threadpool);
  }
  // warm-up rollout to get a steady state
@@ -73,9 +73,6 @@ void BM_StepHumanoid200(benchmark::State& state) {
  state.SetLabel(label);
  state.SetItemsProcessed(state.iterations());
  mj_deleteData(data);
  if (threadpool) {
    mju_threadPoolDestroy(threadpool);
  }
 }
 void BM_Step22Humanoids(benchmark::State& state) {
@@ -90,10 +87,8 @@ void BM_Step22Humanoids(benchmark::State& state) {
  model->opt.disableflags &= ~mjDSBL_ISLAND;  // enable islands
  mjData* data = mj_makeData(model);
-  mjThreadPool* threadpool = nullptr;
+  if (nthread) {
-  if (nthread > 1) {
+    mju_threadpool(data, nthread);
    threadpool = mju_threadPoolCreate(nthread);
    mju_bindThreadPool(data, threadpool);
  }
  // warm-up rollout to get a steady state
@@ -123,12 +118,9 @@ void BM_Step22Humanoids(benchmark::State& state) {
  state.SetLabel(label);
  state.SetItemsProcessed(state.iterations());
  mj_deleteData(data);
  if (threadpool) {
    mju_threadPoolDestroy(threadpool);
  }
 }
-BENCHMARK(BM_StepHumanoid200)->Arg(1)->Arg(6);
+BENCHMARK(BM_StepHumanoid200)->Arg(0)->Arg(kNumThreads);
-BENCHMARK(BM_Step22Humanoids)->Arg(1)->Arg(6);
+BENCHMARK(BM_Step22Humanoids)->Arg(0)->Arg(kNumThreads);
 }  // namespace
 }  // namespace mujoco
@@ -32,7 +32,7 @@
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_util_errmem.h"
-#include "src/thread/thread_pool.h"
+#include "src/engine/engine_thread.h"
 #include "test/fixture.h"
 namespace mujoco {
@@ -804,32 +804,18 @@ TEST_F(EngineIoTest, VeryLargeMemory) {
  }
 }
-struct TestFunctionArgs_ {
+struct TestFunctionArgs {
-  mjData* d;
+  int stack_output[1000];
-  int input;
+  int arena_output[1000];
  int stack_output;
  int arena_output;
  size_t output_thread_worker;
 };
 typedef TestFunctionArgs_ TestFunctionArgs;
-void* TestFunction(void* args) {
+void TestFunction(const mjModel* m, mjData* d, void* args, int i, int j) {
  TestFunctionArgs* test_args = static_cast<TestFunctionArgs*>(args);
-  test_args->output_thread_worker =
+  mj_markStack(d);
-      mju_threadPoolCurrentWorkerId((mjThreadPool*)test_args->d->threadpool);
+  int* test_ints = mj_stackAllocInt(d, 10);
-  mj_markStack(test_args->d);
+  test_ints[0] = j;
-  int* test_ints = mj_stackAllocInt(test_args->d, 10);
+  test_args->stack_output[j] = test_ints[0];
-  test_ints[0] = test_args->input;
+  mj_freeStack(d);
  test_args->stack_output = test_ints[0];
  int* test_arena_ints =
      (int*)mj_arenaAllocByte(test_args->d, sizeof(int) * 10, alignof(int));
  test_arena_ints[0] = test_args->input;
  test_args->arena_output = test_arena_ints[0];
  mj_freeStack(test_args->d);
  return nullptr;
 }
 TEST_F(EngineIoTest, TestStackShardingForThreads) {
@@ -846,38 +832,18 @@ TEST_F(EngineIoTest, TestStackShardingForThreads) {
  mjData* data = mj_makeData(model);
  ASSERT_THAT(data, NotNull());
-  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
+  mju_threadpool(data, 10);
  mju_bindThreadPool(data, thread_pool);
  constexpr int kTasks = 1000;
-  TestFunctionArgs test_function_args[kTasks];
+  TestFunctionArgs test_function_args;
-  mjTask tasks[kTasks];
+  mju_dispatch(model, data, TestFunction, &test_function_args, kTasks);
  for (int i = 0; i < kTasks; ++i) {
    test_function_args[i].d = data;
    test_function_args[i].input = i;
    mju_defaultTask(&tasks[i]);
    tasks[i].func = TestFunction;
    tasks[i].args = &test_function_args[i];
    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
  }
  mj_markStack(data);
  int* test_ints = mj_stackAllocInt(data, 10);
  test_ints[0] = 1;
  mj_freeStack(data);
  for (int i = 0; i < kTasks; ++i) {
-    mju_taskJoin(&tasks[i]);
+    EXPECT_EQ(i, test_function_args.stack_output[i]);
  }
  for (int i = 0; i < kTasks; ++i) {
    EXPECT_EQ(test_function_args[i].input, test_function_args[i].stack_output);
    EXPECT_EQ(test_function_args[i].input, test_function_args[i].arena_output);
  }
  mj_deleteData(data);
  mj_deleteModel(model);
  mju_threadPoolDestroy(thread_pool);
 }
 #ifdef ADDRESS_SANITIZER
@@ -22,10 +22,9 @@
 #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 "src/engine/engine_thread.h"
 #include "test/fixture.h"
 namespace mujoco {
@@ -61,8 +60,7 @@ TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
  mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
  // bind a threadpool to the data_threaded
-  mjThreadPool* threadpool = mju_threadPoolCreate(10);
+  mju_threadpool(data_threaded, 10);
  mju_bindThreadPool(data_threaded, threadpool);
  for (int i = 0; i < 10; ++i) {
    mj_step(model, data);
@@ -83,7 +81,6 @@ TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
  mj_deleteModel(model);
  mj_deleteData(data_threaded);
  mj_deleteModel(model_threaded);
  mju_threadPoolDestroy(threadpool);
 }
 TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
@@ -116,8 +113,7 @@ TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
  mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
  // bind a threadpool to the data_threaded
-  mjThreadPool* threadpool = mju_threadPoolCreate(10);
+  mju_threadpool(data_threaded, 10);
  mju_bindThreadPool(data_threaded, threadpool);
  for (int i = 0; i < 10; ++i) {
    mj_step(model, data);
@@ -138,7 +134,6 @@ TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
  mj_deleteModel(model);
  mj_deleteData(data_threaded);
  mj_deleteModel(model_threaded);
  mju_threadPoolDestroy(threadpool);
 }
 }  // namespace
@@ -1,17 +0,0 @@
 # 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.
 mujoco_test(thread_pool_test)
 mujoco_test(thread_queue_test)
@@ -1,142 +0,0 @@
 // 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 <atomic>
 #include <condition_variable>
 #include <memory>
 #include <mutex>
 #include <thread>
 #include <gtest/gtest.h>
 #include <mujoco/mujoco.h>
 namespace mujoco {
 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, TestMjThreadPool10Threads) {
  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
  constexpr int kTasks = 1000;
  TestFunctionArgs test_function_args[kTasks];
  mjTask tasks[kTasks];
  for (int i = 0; i < kTasks; ++i) {
    test_function_args[i].input = i;
    mju_defaultTask(&tasks[i]);
    tasks[i].func = test_function;
    tasks[i].args = &test_function_args[i];
    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
  }
  for (int i = 0; i < kTasks; ++i) {
    mju_taskJoin(&tasks[i]);
  }
  for (int i = 0; i < kTasks; ++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);
  constexpr int kTasks = 1000;
  TestFunctionArgs test_function_args[kTasks];
  mjTask tasks[kTasks];
  for (int i = 0; i < kTasks; ++i) {
    test_function_args[i].input = i;
    mju_defaultTask(&tasks[i]);
    tasks[i].func = test_function;
    tasks[i].args = &test_function_args[i];
    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
  }
  for (int i = 0; i < kTasks; ++i) {
    mju_taskJoin(&tasks[i]);
  }
  for (int i = 0; i < kTasks; ++i) {
    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
  }
  mju_threadPoolDestroy(thread_pool);
 }
 TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
  constexpr int kTasks = 20;
  TestFunctionArgs test_function_args[kTasks];
  mjTask 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;
    mju_defaultTask(&tasks[i]);
    tasks[i].func = &test_function;
    tasks[i].args = &test_function_args[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
      mju_threadPoolEnqueue(thread_pool, &tasks[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) {
    mju_taskJoin(&tasks[i]);
  }
  for (int i = 0; i < kTasks; ++i) {
    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
  }
  mju_threadPoolDestroy(thread_pool);
 }
 }  // namespace
 }  // namespace mujoco
@@ -1,46 +0,0 @@
 // 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_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
@@ -66,8 +66,6 @@ public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSAN_H_ = true;
 public const bool ADDRESS_SANITIZER = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSPEC_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MJSPECMACRO_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true;
 public const int mjMAXTHREAD = 128;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTYPE_H_ = true;
 public const double mjMINVAL = 1e-15;
 public const bool THIRD_PARTY_MUJOCO_MJUI_H_ = true;
@@ -518,11 +516,6 @@ public enum mjtSleepState : int{
  mjS_ASLEEP = 0,
  mjS_AWAKE = 1,
 }
 public enum mjtTaskStatus : int{
  mjTASK_NEW = 0,
  mjTASK_QUEUED = 1,
  mjTASK_COMPLETED = 2,
 }
 public enum mjtGeomInertia : int{
  mjINERTIA_VOLUME = 0,
  mjINERTIA_SHELL = 1,
@@ -1586,135 +1579,9 @@ public unsafe struct mjData_ {
  public UIntPtr pstack;
  public UIntPtr pbase;
  public UIntPtr parena;
  public UIntPtr threadpool;
  public byte threadlock;
  public UInt64 maxuse_stack;
  public UInt64 maxuse_threadstack0;
  public UInt64 maxuse_threadstack1;
  public UInt64 maxuse_threadstack2;
  public UInt64 maxuse_threadstack3;
  public UInt64 maxuse_threadstack4;
  public UInt64 maxuse_threadstack5;
  public UInt64 maxuse_threadstack6;
  public UInt64 maxuse_threadstack7;
  public UInt64 maxuse_threadstack8;
  public UInt64 maxuse_threadstack9;
  public UInt64 maxuse_threadstack10;
  public UInt64 maxuse_threadstack11;
  public UInt64 maxuse_threadstack12;
  public UInt64 maxuse_threadstack13;
  public UInt64 maxuse_threadstack14;
  public UInt64 maxuse_threadstack15;
  public UInt64 maxuse_threadstack16;
  public UInt64 maxuse_threadstack17;
  public UInt64 maxuse_threadstack18;
  public UInt64 maxuse_threadstack19;
  public UInt64 maxuse_threadstack20;
  public UInt64 maxuse_threadstack21;
  public UInt64 maxuse_threadstack22;
  public UInt64 maxuse_threadstack23;
  public UInt64 maxuse_threadstack24;
  public UInt64 maxuse_threadstack25;
  public UInt64 maxuse_threadstack26;
  public UInt64 maxuse_threadstack27;
  public UInt64 maxuse_threadstack28;
  public UInt64 maxuse_threadstack29;
  public UInt64 maxuse_threadstack30;
  public UInt64 maxuse_threadstack31;
  public UInt64 maxuse_threadstack32;
  public UInt64 maxuse_threadstack33;
  public UInt64 maxuse_threadstack34;
  public UInt64 maxuse_threadstack35;
  public UInt64 maxuse_threadstack36;
  public UInt64 maxuse_threadstack37;
  public UInt64 maxuse_threadstack38;
  public UInt64 maxuse_threadstack39;
  public UInt64 maxuse_threadstack40;
  public UInt64 maxuse_threadstack41;
  public UInt64 maxuse_threadstack42;
  public UInt64 maxuse_threadstack43;
  public UInt64 maxuse_threadstack44;
  public UInt64 maxuse_threadstack45;
  public UInt64 maxuse_threadstack46;
  public UInt64 maxuse_threadstack47;
  public UInt64 maxuse_threadstack48;
  public UInt64 maxuse_threadstack49;
  public UInt64 maxuse_threadstack50;
  public UInt64 maxuse_threadstack51;
  public UInt64 maxuse_threadstack52;
  public UInt64 maxuse_threadstack53;
  public UInt64 maxuse_threadstack54;
  public UInt64 maxuse_threadstack55;
  public UInt64 maxuse_threadstack56;
  public UInt64 maxuse_threadstack57;
  public UInt64 maxuse_threadstack58;
  public UInt64 maxuse_threadstack59;
  public UInt64 maxuse_threadstack60;
  public UInt64 maxuse_threadstack61;
  public UInt64 maxuse_threadstack62;
  public UInt64 maxuse_threadstack63;
  public UInt64 maxuse_threadstack64;
  public UInt64 maxuse_threadstack65;
  public UInt64 maxuse_threadstack66;
  public UInt64 maxuse_threadstack67;
  public UInt64 maxuse_threadstack68;
  public UInt64 maxuse_threadstack69;
  public UInt64 maxuse_threadstack70;
  public UInt64 maxuse_threadstack71;
  public UInt64 maxuse_threadstack72;
  public UInt64 maxuse_threadstack73;
  public UInt64 maxuse_threadstack74;
  public UInt64 maxuse_threadstack75;
  public UInt64 maxuse_threadstack76;
  public UInt64 maxuse_threadstack77;
  public UInt64 maxuse_threadstack78;
  public UInt64 maxuse_threadstack79;
  public UInt64 maxuse_threadstack80;
  public UInt64 maxuse_threadstack81;
  public UInt64 maxuse_threadstack82;
  public UInt64 maxuse_threadstack83;
  public UInt64 maxuse_threadstack84;
  public UInt64 maxuse_threadstack85;
  public UInt64 maxuse_threadstack86;
  public UInt64 maxuse_threadstack87;
  public UInt64 maxuse_threadstack88;
  public UInt64 maxuse_threadstack89;
  public UInt64 maxuse_threadstack90;
  public UInt64 maxuse_threadstack91;
  public UInt64 maxuse_threadstack92;
  public UInt64 maxuse_threadstack93;
  public UInt64 maxuse_threadstack94;
  public UInt64 maxuse_threadstack95;
  public UInt64 maxuse_threadstack96;
  public UInt64 maxuse_threadstack97;
  public UInt64 maxuse_threadstack98;
  public UInt64 maxuse_threadstack99;
  public UInt64 maxuse_threadstack100;
  public UInt64 maxuse_threadstack101;
  public UInt64 maxuse_threadstack102;
  public UInt64 maxuse_threadstack103;
  public UInt64 maxuse_threadstack104;
  public UInt64 maxuse_threadstack105;
  public UInt64 maxuse_threadstack106;
  public UInt64 maxuse_threadstack107;
  public UInt64 maxuse_threadstack108;
  public UInt64 maxuse_threadstack109;
  public UInt64 maxuse_threadstack110;
  public UInt64 maxuse_threadstack111;
  public UInt64 maxuse_threadstack112;
  public UInt64 maxuse_threadstack113;
  public UInt64 maxuse_threadstack114;
  public UInt64 maxuse_threadstack115;
  public UInt64 maxuse_threadstack116;
  public UInt64 maxuse_threadstack117;
  public UInt64 maxuse_threadstack118;
  public UInt64 maxuse_threadstack119;
  public UInt64 maxuse_threadstack120;
  public UInt64 maxuse_threadstack121;
  public UInt64 maxuse_threadstack122;
  public UInt64 maxuse_threadstack123;
  public UInt64 maxuse_threadstack124;
  public UInt64 maxuse_threadstack125;
  public UInt64 maxuse_threadstack126;
  public UInt64 maxuse_threadstack127;
  public UInt64 maxuse_arena;
  public int maxuse_con;
  public int maxuse_efc;
@@ -5926,7 +5793,6 @@ public unsafe struct mjData_ {
  public int* efc_state;
  public double* efc_force;
  public double* ifrc_constraint;
  public UIntPtr threadpool;
  public UInt64 signature;
 }
@@ -7725,8 +7591,5 @@ 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 void mju_bindThreadPool(mjData_* d, void* thread_pool);
 }
 }
@@ -6455,15 +6455,24 @@ struct MjData {
  void set_parena(size_t value) {
    ptr_->parena = value;
  }
  uintptr_t threadpool() const {
    return ptr_->threadpool;
  }
  void set_threadpool(uintptr_t value) {
    ptr_->threadpool = value;
  }
  mjtBool threadlock() const {
    return ptr_->threadlock;
  }
  void set_threadlock(mjtBool value) {
    ptr_->threadlock = value;
  }
  int maxuse_stack() const {
    return static_cast<int>(ptr_->maxuse_stack);
  }
  void set_maxuse_stack(int value) {
    ptr_->maxuse_stack = static_cast<mjtSize>(value);
  }
  emscripten::val maxuse_threadstack() const {
    return emscripten::val(emscripten::typed_memory_view(128, ptr_->maxuse_threadstack));
  }
  int maxuse_arena() const {
    return static_cast<int>(ptr_->maxuse_arena);
  }
@@ -7095,12 +7104,6 @@ struct MjData {
  emscripten::val ifrc_constraint() const {
    return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->ifrc_constraint));
  }
  uintptr_t threadpool() const {
    return ptr_->threadpool;
  }
  void set_threadpool(uintptr_t value) {
    ptr_->threadpool = value;
  }
  uint64_t signature() const {
    return ptr_->signature;
  }
@@ -10766,6 +10769,10 @@ void mju_symmetrize_wrapper(const val& res, const NumberArray& mat, int n) {
  mju_symmetrize(res_.data(), mat_.data(), n);
 }
 void mju_threadpool_wrapper(MjData& d, int nthread) {
  mju_threadpool(d.get(), nthread);
 }
 void mju_transformSpatial_wrapper(const val& res, const NumberArray& vec, int flg_force, const NumberArray& newpos, const NumberArray& oldpos, const NumberArray& rotnew2old) {
  UNPACK_VALUE(mjtNum, res);
  UNPACK_ARRAY(mjtNum, vec);
@@ -11467,10 +11474,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
    .value("mjSTEREO_NONE", mjSTEREO_NONE)
    .value("mjSTEREO_QUADBUFFERED", mjSTEREO_QUADBUFFERED)
    .value("mjSTEREO_SIDEBYSIDE", mjSTEREO_SIDEBYSIDE);
  enum_<mjtTaskStatus>("mjtTaskStatus")
    .value("mjTASK_NEW", mjTASK_NEW)
    .value("mjTASK_QUEUED", mjTASK_QUEUED)
    .value("mjTASK_COMPLETED", mjTASK_COMPLETED);
  enum_<mjtTexture>("mjtTexture")
    .value("mjTEXTURE_2D", mjTEXTURE_2D)
    .value("mjTEXTURE_CUBE", mjTEXTURE_CUBE)
@@ -11734,7 +11737,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
    .property("maxuse_con", &MjData::maxuse_con, &MjData::set_maxuse_con, reference())
    .property("maxuse_efc", &MjData::maxuse_efc, &MjData::set_maxuse_efc, reference())
    .property("maxuse_stack", &MjData::maxuse_stack, &MjData::set_maxuse_stack, reference())
    .property("maxuse_threadstack", &MjData::maxuse_threadstack)
    .property("mocap_pos", &MjData::mocap_pos)
    .property("mocap_quat", &MjData::mocap_quat)
    .property("moment_colind", &MjData::moment_colind)
@@ -11804,6 +11806,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
    .property("ten_wrapadr", &MjData::ten_wrapadr)
    .property("ten_wrapnum", &MjData::ten_wrapnum)
    .property("tendon_efcadr", &MjData::tendon_efcadr)
    .property("threadlock", &MjData::threadlock, &MjData::set_threadlock, reference())
    .property("threadpool", &MjData::threadpool, &MjData::set_threadpool, reference())
    .property("time", &MjData::time, &MjData::set_time, reference())
    .property("timer", &MjData::timer, reference())
@@ -13610,6 +13613,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
  function("mju_sum", &mju_sum_wrapper);
  function("mju_sym2dense", &mju_sym2dense_wrapper);
  function("mju_symmetrize", &mju_symmetrize_wrapper);
  function("mju_threadpool", &mju_threadpool_wrapper);
  function("mju_transformSpatial", &mju_transformSpatial_wrapper);
  function("mju_transpose", &mju_transpose_wrapper);
  function("mju_trnVecPose", &mju_trnVecPose_wrapper);
@@ -98,17 +98,6 @@ _SKIPPED_WRITABLE_ERROR: tuple[str, ...] = (
    "mj_printSchema",
 )
 # Omitted thread management functions
 _SKIPPED_THREAD_FUNCTIONS: tuple[str, ...] = (
    # go/keep-sorted start
    "mju_bindThreadPool",
    "mju_defaultTask",
    "mju_taskJoin",
    "mju_threadPoolCreate",
    "mju_threadPoolDestroy",
    "mju_threadPoolEnqueue",
    # go/keep-sorted end
 )
 # Omitted asset cache functions
 _SKIPPED_ASSET_CACHE_FUNCTIONS: tuple[str, ...] = (
@@ -207,7 +196,6 @@ _SKIPPED_UTILITY_FUNCTIONS: tuple[str, ...] = (
 # List of functions that should be skipped during the code generation process.
 SKIPPED_FUNCTIONS: tuple[str, ...] = (
    _SKIPPED_CLASS_METHODS
    + _SKIPPED_THREAD_FUNCTIONS
    + _SKIPPED_MEMORY_FUNCTIONS
    + _SKIPPED_PLUGIN_FUNCTIONS
    + _SKIPPED_GETTERS_AND_SETTERS
@@ -223,8 +211,6 @@ SKIPPED_STRUCTS: tuple[str, ...] = (
    # go/keep-sorted start
    "mjCache",
    "mjSDF",
    "mjTask",
    "mjThreadPool",
    "mjUI",
    "mjVFS",
    "mjrContext",
@@ -152,10 +152,6 @@ describe('Enums', () => {
    expect(mujoco.mjtSDFType).toBeDefined();
  });
  it('mjtTaskStatus should exist', () => {
    expect(mujoco.mjtTaskStatus).toBeDefined();
  });
  it('mjtState should exist', () => {
    expect(mujoco.mjtState).toBeDefined();
  });