diff --git a/CMakeLists.txt b/CMakeLists.txt index 0423d3f4..9a9a34b4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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}" ) diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 9890735b..2f6a4f69 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -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 `_. .. _glNumericSizes: diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index ba1a6135..fbc2e6d9 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -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: diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 4daaaf44..7a409cc2 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -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. - -.. _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. +Create a thread pool with nthread worker threads. .. _Standardmath: diff --git a/doc/changelog.rst b/doc/changelog.rst index 699d0dd0..c5b5377d 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,16 @@ General - Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`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` and the diff --git a/doc/includes/references.h b/doc/includes/references.h index 60b8d699..1ce969d7 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -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_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); +void mju_threadpool(mjData* d, int nthread); mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child, const char* prefix, const char* suffix); mjsBody* mjs_addBody(mjsBody* body, const mjsDefault* def); diff --git a/doc/programming/index.rst b/doc/programming/index.rst index 24c2402f..25aa9ee9 100644 --- a/doc/programming/index.rst +++ b/doc/programming/index.rst @@ -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` interface - to process tasks asynchronously. To enable use in MuJoCo, create a thread pool and assign it to the - ``mjData.threadpool`` field. + 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, call ``mju_threadpool``. 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 `. `mjplugin.h `__ Defines data structures required by :ref:`engine plugins`. -`mjthread.h `__ - Defines data structures and functions required by :ref:`thread`. `mjmacro.h `__ Defines C macros that are useful in user code. `mjxmacro.h `__ diff --git a/doc/programming/samples.rst b/doc/programming/samples.rst index 8a964c86..ee6c7feb 100644 --- a/doc/programming/samples.rst +++ b/doc/programming/samples.rst @@ -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. diff --git a/doc/programming/simulation.rst b/doc/programming/simulation.rst index 702a9768..2f636221 100644 --- a/doc/programming/simulation.rst +++ b/doc/programming/simulation.rst @@ -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 -``mjData.threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across -:ref:`islands` — can be distributed across worker threads. Note that within-step threading currently has -significant memory overhead and is still a work in progress. +MuJoCo has support for within-step multi-threading. When a thread pool is initialized via +``mju_threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across +:ref:`islands` — can be distributed across worker threads. 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 diff --git a/include/mujoco/mjassert.h b/include/mujoco/mjassert.h index 169cbb30..36651ac5 100644 --- a/include/mujoco/mjassert.h +++ b/include/mujoco/mjassert.h @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -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); diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 486c382a..2ee67f95 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -20,7 +20,7 @@ #include #include -#include + @@ -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 }; diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h deleted file mode 100644 index 153ff0cc..00000000 --- a/include/mujoco/mjthread.h +++ /dev/null @@ -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_ diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index adb7686a..04539f5e 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.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( uintptr_t, threadpool ) + X( mjtNum, time ) // 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 ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 4a1384a7..1af40764 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -31,7 +31,6 @@ #include #include #include -#include #include #include #include @@ -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. -MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); - -// 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); +// Create a thread pool with nthread worker threads. +MJAPI void mju_threadpool(mjData* d, int nthread); //---------------------------------- Attachment ---------------------------------------------------- diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index e509bce3..05dfa362 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -725,6 +725,9 @@ PYBIND11_MODULE(_functions, pymodule) { Def(pymodule); Def(pymodule); + // Thread pool + Def(pymodule); + // Ray collision Def( pymodule, diff --git a/python/mujoco/introspect/enums.py b/python/mujoco/introspect/enums.py index 2247a805..3611253e 100644 --- a/python/mujoco/introspect/enums.py +++ b/python/mujoco/introspect/enums.py @@ -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', diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py index d2323d9e..5deefe3c 100644 --- a/python/mujoco/introspect/functions.py +++ b/python/mujoco/introspect/functions.py @@ -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', - 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', + name='mju_threadpool', return_type=ValueType(name='void'), parameters=( FunctionParameterDecl( @@ -10011,75 +9997,11 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), FunctionParameterDecl( - name='thread_pool', - type=PointerType( - inner_type=ValueType(name='void'), - ), + name='nthread', + type=ValueType(name='int'), ), ), - doc='Adds a thread pool to mjData and configures it for multi-threaded use.', # pylint: disable=line-too-long - )), - ('mju_threadPoolEnqueue', - FunctionDecl( - name='mju_threadPoolEnqueue', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='thread_pool', - type=PointerType( - inner_type=ValueType(name='mjThreadPool'), - ), - ), - FunctionParameterDecl( - name='task', - type=PointerType( - inner_type=ValueType(name='mjTask'), - ), - ), - ), - 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.', + doc='Create a thread pool with nthread worker threads.', )), ('mjs_attach', FunctionDecl( diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 119219ee..8fa54694 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -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'), diff --git a/sample/testspeed.cc b/sample/testspeed.cc index f14e6572..0cedb767 100644 --- a/sample/testspeed.cc +++ b/sample/testspeed.cc @@ -218,8 +218,7 @@ int main(int argc, char** argv) { // make and bind threadpool if (npoolthread > 1) { - mjThreadPool* threadpool = mju_threadPoolCreate(npoolthread); - mju_bindThreadPool(d[id], threadpool); + mju_threadpool(d[id], npoolthread); } } @@ -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 diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 4249d7a9..61940816 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -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: // Load key + case 6: // Adjust 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[5].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.range[0] = 0; + this->ui0.sect[SECT_SIMULATION].item[6].slider.range[1] = mjMAX(0, 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[11].slider.divisions = nhistory_; + this->ui0.sect[SECT_SIMULATION].item[12].slider.range[0] = 1 - nhistory_; + this->ui0.sect[SECT_SIMULATION].item[12].slider.divisions = nhistory_; // detect image sensors for visualization DetectImageSensors(this, this->m_); diff --git a/simulate/simulate.h b/simulate/simulate.h index 27727ee7..367d95a3 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -153,6 +153,7 @@ class Simulate { std::optional save_mjb; std::optional print_model; std::optional 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"}, diff --git a/src/engine/CMakeLists.txt b/src/engine/CMakeLists.txt index e61c1dc4..baa23887 100644 --- a/src/engine/CMakeLists.txt +++ b/src/engine/CMakeLists.txt @@ -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 diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index 66d0cf0f..dd70d99e 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -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 - mj_markStack(d); - mjc_setCCDBuffer(mj_stackAllocByte(d, ccd_size, sizeof(mjtNum))); + mjContactArg arg; + arg.ccd_size = ccd_size; + 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++) { - int g1 = pairbuffer[4*i + 0]; - int g2 = pairbuffer[4*i + 1]; - int ipair = pairbuffer[4*i + 2]; - int idx = pairbuffer[4*i + 3]; - mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]]; - 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); - } + // dispatch narrowphase to threads with local stack allocation for EPA + { + mj_markStack(d); + arg.epabuffer = mj_stackAllocByte(d, ccd_size * nthread, sizeof(mjtNum)); + mju_dispatch(m, d, collisionTask, &arg, nchunk); + mj_freeStack(d); } - // 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; diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 921fbae4..59dcd459 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -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 "thread/thread_task.h" +#include "engine/engine_thread.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 - if (!d->threadpool) { - // inertia, timed internally (POS_INERTIA) - mj_makeM(m, d); - mj_factorM(m, d); + // inertia, timed internally (POS_INERTIA) + mj_makeM(m, d); + mj_factorM(m, d); - // collision, timed internally (POS_COLLISION) - 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]); - } + // collision, timed internally (POS_COLLISION) + mj_collision(m, d); 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 -struct mjSolIslandArgs_ { - const mjModel* m; - mjData* d; - int island; -}; -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]); +// mju_dispatch callback: solve one island +static void solveIslandTask(const mjModel* m, mjData* d, void* arg, int thread_id, int 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); } - - 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 - 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); - } + mju_dispatch(m, d, solveIslandTask, NULL, nisland); // copy back solver outputs (scatter dofs since ni <= nv) mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof); diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 6a1e10fe..ff23d33a 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -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); } diff --git a/src/engine/engine_memory.c b/src/engine/engine_memory.c index 741575b7..9ecdea1e 100644 --- a/src/engine/engine_memory.c +++ b/src/engine/engine_memory.c @@ -16,6 +16,7 @@ #include // IWYU pragma: keep #include +#include #include #include #include @@ -26,7 +27,7 @@ #include // IWYU pragma: keep #include "engine/engine_crossplatform.h" #include "engine/engine_util_errmem.h" -#include "thread/thread_pool.h" + #ifdef ADDRESS_SANITIZER #include @@ -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_ptr; - 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); + mjStackInfo stack_info = get_stack_info_from_data(d); + stack_usage = stack_usage_redzone(&stack_info); #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_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); - } + d->maxuse_stack = mjMAX(d->maxuse_stack, usage); + d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena); 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 - if (!d->threadpool) { - 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; + // size zero: no-op + if (!size) { + return NULL; } - // multi threaded allocation - size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); - mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); - return stackallocinternal(d, stack_info, size, alignment, caller, line); + // call in mju_dispatch: atomically reserve space on the stack + if (d->threadlock) { + size_t alloc_size = size + alignment - 1 + 2 * mjREDZONE; + 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) { - mjStackInfo stack_info = get_stack_info_from_data(d); - markstackinternal(d, &stack_info); - d->pstack = stack_info.bottom - stack_info.top; - d->pbase = stack_info.stack_base; + // no-op if called from mju_dispatch + if (d->threadlock) { return; } - size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); - mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); - markstackinternal(d, stack_info); + mjStackInfo stack_info = get_stack_info_from_data(d); + 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) { - 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; + if (d->threadlock) { return; } - size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); - mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); - freestackinternal(stack_info); -} - - -// 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; - } + 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; } diff --git a/src/engine/engine_memory.h b/src/engine/engine_memory.h index bbfd836b..f7969eb4 100644 --- a/src/engine/engine_memory.h +++ b/src/engine/engine_memory.h @@ -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); diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index 83db1d49..c8309e77 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -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 "thread/thread_task.h" +#include "engine/engine_thread.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 - if (d->threadpool && ncon >= kTactileParallelThreshold) { - int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool); - int batch_size = (ncon + nthreads - 1) / nthreads; - int ntasks = (ncon + batch_size - 1) / batch_size; + // parallel path: use mj_batch to process taxel batches + int nthread = mju_numThread(d); + if (nthread > 0 && ncon >= kTactileParallelThreshold) { + int batch_size = (ncon + nthread - 1) / nthread; + int ntask = (ncon + batch_size - 1) / batch_size; - mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask); - mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs); + mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntask, mjTactileTaskArgs); - for (int t = 0; t < ntasks; t++) { - task_args[t].m = m; - task_args[t].d = d; + for (int t = 0; t < ntask; t++) { 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_taskJoin(&tasks[t]); - } + mju_dispatch(m, d, tactileTask, task_args, ntask); } // 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 diff --git a/src/engine/engine_thread.cc b/src/engine/engine_thread.cc new file mode 100644 index 00000000..3f521d21 --- /dev/null +++ b/src/engine/engine_thread.cc @@ -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 +#include +#include +#include + +#include +#include +#include +#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 next_{0}; + + // atomic counter for number of workers who completed their tasks + alignas(64) std::atomic ndone_{0}; + + // alternating signal from -1, 1 to start / halt the worker threads, + // set to 0 to force all workers to exit + std::atomic signal_{1}; + + std::vector threads_; +}; + + + +// create a thread pool with nthread threads +void mju_threadpool(mjData* d, int nthread) { + if (d->threadpool) { + ThreadPoolContext* ctx = + reinterpret_cast(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(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(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(d->threadpool); + return ctx ? ctx->ThreadCount() + 1 : 1; +} diff --git a/src/engine/engine_thread.h b/src/engine/engine_thread.h new file mode 100644 index 00000000..52f8ba64 --- /dev/null +++ b/src/engine/engine_thread.h @@ -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 +#include +#include + +#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_ diff --git a/src/thread/CMakeLists.txt b/src/thread/CMakeLists.txt deleted file mode 100644 index e320c20b..00000000 --- a/src/thread/CMakeLists.txt +++ /dev/null @@ -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}) diff --git a/src/thread/thread_pool.cc b/src/thread/thread_pool.cc deleted file mode 100644 index af117f6a..00000000 --- a/src/thread/thread_pool.cc +++ /dev/null @@ -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 - -#include -#include -#include -#include -#include -#include -#include -#include - -#include // IWYU pragma: keep -#include -#include -#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 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(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 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(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 shutdown_ = false; - - // OS threads that are running in this pool - std::vector workers_; - - // queue of tasks to execute - mujoco::LocklessQueue 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(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(thread_pool); - return thread_pool_impl->NumberOfThreads(); -} - -size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) { - auto thread_pool_impl = static_cast(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(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(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(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(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 diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h deleted file mode 100644 index 54d852d9..00000000 --- a/src/thread/thread_pool.h +++ /dev/null @@ -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 - -#include -#include -#include - -#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_ diff --git a/src/thread/thread_queue.h b/src/thread/thread_queue.h deleted file mode 100644 index 3178f8ec..00000000 --- a/src/thread/thread_queue.h +++ /dev/null @@ -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 -#include -#include -#include - -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 -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 read_cursor_ = 0; - std::atomic write_cursor_ = 0; - std::atomic maximum_read_cursor_ = 0; - - std::atomic buffer_[(buffer_capacity + 1)]; -}; - -} // namespace mujoco - -#endif // MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_ diff --git a/src/thread/thread_task.cc b/src/thread/thread_task.cc deleted file mode 100644 index 77ce2bfd..00000000 --- a/src/thread/thread_task.cc +++ /dev/null @@ -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 - -#include - -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 diff --git a/src/thread/thread_task.h b/src/thread/thread_task.h deleted file mode 100644 index 0ed6f468..00000000 --- a/src/thread/thread_task.h +++ /dev/null @@ -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 -#include - -#ifdef __cplusplus -#include -#include -#include -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; -inline std::atomic& GetAtomicTaskStatus(mjTask* task) { - static_assert(sizeof(std::atomic) == sizeof(TaskStatus)); - static_assert(alignof(std::atomic) == alignof(TaskStatus)); - static_assert(std::atomic::is_always_lock_free); - return *std::launder(reinterpret_cast*>( - const_cast(&task->status))); -} -} // namespace mujoco -#endif // __cplusplus - -#endif // MUJOCO_SRC_THREAD_THREAD_TASK_H_ diff --git a/test/benchmark/ccd_benchmark_test.cc b/test/benchmark/ccd_benchmark_test.cc index a7768dce..9f1ad4bf 100644 --- a/test/benchmark/ccd_benchmark_test.cc +++ b/test/benchmark/ccd_benchmark_test.cc @@ -15,8 +15,8 @@ // A benchmark for parsing and compiling models from XML. #include -#include #include +#include #include #include @@ -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 - BM_BoxMesh_NativeCCD(benchmark::State& state) { - static TestHarness harness(kBoxMeshPath, "boxmesh.xml (nativeccd)"); +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxMesh(benchmark::State& state) { + int nthread = state.range(0); + static TestHarness harness(kBoxMeshPath, "boxmesh.xml"); + harness.SetThreads(nthread); harness.RunBenchmark(state); } -BENCHMARK(BM_BoxMesh_NativeCCD); - -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); +BENCHMARK(BM_BoxMesh)->Arg(0)->Arg(kNumThreads); 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 - BM_Ellipsoid_NativeCCD(benchmark::State& state) { - static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (nativeccd)"); +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Ellipsoid(benchmark::State& state) { + int nthread = state.range(0); + 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 - BM_Ellipsoid_LibCCD(benchmark::State& state) { - static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (libccd)", - mjDSBL_NATIVECCD); +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed(benchmark::State& state) { + int nthread = state.range(0); + static TestHarness harness(kMixedPath, "mixed.xml"); + harness.SetThreads(nthread); harness.RunBenchmark(state); } -BENCHMARK(BM_Ellipsoid_LibCCD); - -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); +BENCHMARK(BM_Mixed)->Arg(0)->Arg(kNumThreads); } // namespace } // namespace mujoco diff --git a/test/benchmark/thread_performance_test.cc b/test/benchmark/thread_performance_test.cc index aa4c04a8..0bb3df47 100644 --- a/test/benchmark/thread_performance_test.cc +++ b/test/benchmark/thread_performance_test.cc @@ -17,7 +17,6 @@ #include #include -#include #include #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 > 1) { - threadpool = mju_threadPoolCreate(nthread); - mju_bindThreadPool(data, threadpool); + if (nthread) { + mju_threadpool(data, nthread); } // 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 > 1) { - threadpool = mju_threadPoolCreate(nthread); - mju_bindThreadPool(data, threadpool); + if (nthread) { + mju_threadpool(data, nthread); } // 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_Step22Humanoids)->Arg(1)->Arg(6); +BENCHMARK(BM_StepHumanoid200)->Arg(0)->Arg(kNumThreads); +BENCHMARK(BM_Step22Humanoids)->Arg(0)->Arg(kNumThreads); } // namespace } // namespace mujoco diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index 0100ef4c..c83ff769 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -32,7 +32,7 @@ #include #include #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_ { - mjData* d; - int input; - int stack_output; - int arena_output; - size_t output_thread_worker; +struct TestFunctionArgs { + int stack_output[1000]; + int arena_output[1000]; }; -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(args); - test_args->output_thread_worker = - mju_threadPoolCurrentWorkerId((mjThreadPool*)test_args->d->threadpool); - mj_markStack(test_args->d); - int* test_ints = mj_stackAllocInt(test_args->d, 10); - test_ints[0] = test_args->input; - 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; + mj_markStack(d); + int* test_ints = mj_stackAllocInt(d, 10); + test_ints[0] = j; + test_args->stack_output[j] = test_ints[0]; + mj_freeStack(d); } 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_bindThreadPool(data, thread_pool); + mju_threadpool(data, 10); constexpr int kTasks = 1000; - TestFunctionArgs test_function_args[kTasks]; - mjTask tasks[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); + TestFunctionArgs test_function_args; + mju_dispatch(model, data, TestFunction, &test_function_args, kTasks); 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].stack_output); - EXPECT_EQ(test_function_args[i].input, test_function_args[i].arena_output); + EXPECT_EQ(i, test_function_args.stack_output[i]); } mj_deleteData(data); mj_deleteModel(model); - mju_threadPoolDestroy(thread_pool); } #ifdef ADDRESS_SANITIZER diff --git a/test/engine/engine_thread_test.cc b/test/engine/engine_thread_test.cc index fd4ceae3..384f8ede 100644 --- a/test/engine/engine_thread_test.cc +++ b/test/engine/engine_thread_test.cc @@ -22,10 +22,9 @@ #include #include #include -#include #include #include -#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_bindThreadPool(data_threaded, threadpool); + mju_threadpool(data_threaded, 10); 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_bindThreadPool(data_threaded, threadpool); + mju_threadpool(data_threaded, 10); 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 diff --git a/test/thread/CMakeLists.txt b/test/thread/CMakeLists.txt deleted file mode 100644 index ada7b7ef..00000000 --- a/test/thread/CMakeLists.txt +++ /dev/null @@ -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) diff --git a/test/thread/thread_pool_test.cc b/test/thread/thread_pool_test.cc deleted file mode 100644 index 2532ec65..00000000 --- a/test/thread/thread_pool_test.cc +++ /dev/null @@ -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 -#include -#include -#include -#include - -#include -#include - -namespace mujoco { -namespace { - -struct TestFunctionArgs_ { - int input; - // make this atomic to avoid red-herring tsan failures. - std::atomic output; -}; -typedef struct TestFunctionArgs_ TestFunctionArgs; - -void* test_function(void* args) { - TestFunctionArgs* test_function_args = static_cast(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 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([&, i] { - // synchronize all threads adding to the thread_pool at the same time - { - std::unique_lock 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 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 diff --git a/test/thread/thread_queue_test.cc b/test/thread/thread_queue_test.cc deleted file mode 100644 index 7b68b6f9..00000000 --- a/test/thread/thread_queue_test.cc +++ /dev/null @@ -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 - -#include - -namespace mujoco { -namespace { - -constexpr size_t kBufferCapacity = 640; - -TEST(TestMujocoLocklessQueue, TestMujocoLocklessQueue) { - LocklessQueue 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 diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 0a6ba9b4..e25c0986 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -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); } } diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 80d0675a..735c4b1e 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -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(ptr_->maxuse_stack); } void set_maxuse_stack(int value) { ptr_->maxuse_stack = static_cast(value); } - emscripten::val maxuse_threadstack() const { - return emscripten::val(emscripten::typed_memory_view(128, ptr_->maxuse_threadstack)); - } int maxuse_arena() const { return static_cast(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") - .value("mjTASK_NEW", mjTASK_NEW) - .value("mjTASK_QUEUED", mjTASK_QUEUED) - .value("mjTASK_COMPLETED", mjTASK_COMPLETED); enum_("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); diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py index fedb130f..086746b3 100644 --- a/wasm/codegen/generators/constants.py +++ b/wasm/codegen/generators/constants.py @@ -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", diff --git a/wasm/tests/enums_test.ts b/wasm/tests/enums_test.ts index d80fb48e..6f16e9f4 100644 --- a/wasm/tests/enums_test.ts +++ b/wasm/tests/enums_test.ts @@ -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(); });