diff --git a/cmake/MujocoOptions.cmake b/cmake/MujocoOptions.cmake index 3e3f080f..be68de3d 100644 --- a/cmake/MujocoOptions.cmake +++ b/cmake/MujocoOptions.cmake @@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__. set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5 - -Wno-maybe-uninitialized + -Wno-maybe-uninitialized -Wno-c++20-extensions ) endif() + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator) + endif() endif() include(MujocoHarden) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 6df1f7fa..1f33d81b 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -3444,7 +3444,7 @@ mju_threadPoolCreate .. mujoco-include:: mju_threadPoolCreate -Creates a thread pool with the specified number of threads running. +Create a thread pool with the specified number of threads running. .. _mju_threadPoolEnqueue: @@ -3453,16 +3453,7 @@ mju_threadPoolEnqueue .. mujoco-include:: mju_threadPoolEnqueue -Enqueues a task in a thread pool. - -.. _mju_taskJoin: - -mju_taskJoin -~~~~~~~~~~~~ - -.. mujoco-include:: mju_taskJoin - -Waits for a task to complete. +Enqueue a task in a thread pool. .. _mju_threadPoolDestroy: @@ -3471,5 +3462,23 @@ mju_threadPoolDestroy .. mujoco-include:: mju_threadPoolDestroy -Destroys a thread pool. +Destroy a thread pool. + +.. _mju_defaultTask: + +mju_defaultTask +~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_defaultTask + +Initialize an mjTask. + +.. _mju_taskJoin: + +mju_taskJoin +~~~~~~~~~~~~ + +.. mujoco-include:: mju_taskJoin + +Wait for a task to complete. diff --git a/doc/includes/references.h b/doc/includes/references.h index 7fa7dca4..7b4ec3e9 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -1443,14 +1443,21 @@ struct mjrContext_ { // custom OpenGL context int readPixelFormat; // default color pixel format for mjr_readPixels }; typedef struct mjrContext_ mjrContext; -struct mjTask_ { - char buffer[24]; -}; -typedef struct mjTask_ mjTask; +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_ { - char buffer[6208]; + 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 mjtButton_ { // mouse button mjBUTTON_NONE = 0, // no button mjBUTTON_LEFT, // left button @@ -2588,9 +2595,8 @@ int mjp_resourceProviderCount(void); const mjpResourceProvider* mjp_getResourceProvider(const char* resource_name); const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot); mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); -void mju_threadPoolEnqueue( -ThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*), -id* args); -void mju_taskJoin(mjTask* task); +void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task); void mju_threadPoolDestroy(mjThreadPool* thread_pool); +void mju_defaultTask(mjTask* task); +void mju_taskJoin(mjTask* task); // NOLINTEND diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h index 66290370..c716e70f 100644 --- a/include/mujoco/mjthread.h +++ b/include/mujoco/mjthread.h @@ -15,33 +15,26 @@ #ifndef MUJOCO_INCLUDE_MJTHREAD_H_ #define MUJOCO_INCLUDE_MJTHREAD_H_ -// C API for MuJoCo threading -#ifdef __cplusplus -extern "C" { -#endif +typedef enum mjtTaskStatus_ { // status values for mjTask + mjTASK_NEW = 0, // newly created + mjTASK_QUEUED, // enqueued in a thread pool + mjTASK_COMPLETED // completed execution +} mjtTaskStatus; -#include - -#include - -// These types are implemented in C++, they're just used as opaque pointers in C -// to provide type safety for functions. -struct mjTask_ { - char buffer[24]; -}; -typedef struct mjTask_ mjTask; +// function pointer type for mjTask +typedef void* (*mjfTask)(void*); +// An opaque type representing a thread pool. struct mjThreadPool_ { - char buffer[6208]; + int nworker; // number of workers in the pool }; typedef struct mjThreadPool_ mjThreadPool; -typedef void*(*mjStartRoutine_)(void*); -typedef mjStartRoutine_ mjStartRoutine; - -#ifdef __cplusplus -} -#endif - +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/mujoco.h b/include/mujoco/mujoco.h index 8abb1e1f..c7964c73 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1307,19 +1307,21 @@ MJAPI const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot); //---------------------- Thread ------------------------------------------------------------------- -// Creates a thread pool with the specified number of threads running. +// Create a thread pool with the specified number of threads running. MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); -// Enqueues a task in a thread pool. -MJAPI void mju_threadPoolEnqueue( - mjThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*), - void* args); +// Enqueue a task in a thread pool. +MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task); -// Waits for a task to complete. +// 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); -// Destroys a thread pool. -MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool); #if defined(__cplusplus) } diff --git a/introspect/codegen/generate_enums.py b/introspect/codegen/generate_enums.py index 68e9592f..59d7ff3b 100644 --- a/introspect/codegen/generate_enums.py +++ b/introspect/codegen/generate_enums.py @@ -56,7 +56,10 @@ class MjEnumVisitor: child_kind = child.get('kind') if child_kind == 'EnumConstantDecl': next_idx = values[-1][1] + 1 if values else 0 - value = int(child['inner'][0].get('value', next_idx)) + if 'inner' in child: + value = int(child['inner'][0].get('value', next_idx)) + else: + value = next_idx values.append((child['name'], value)) return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values)) diff --git a/introspect/enums.py b/introspect/enums.py index 6be0abe7..c769d225 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -642,6 +642,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjFONT_BIG', 2), ]), )), + ('mjtTaskStatus', + EnumDecl( + name='mjtTaskStatus', + declname='enum mjtTaskStatus_', + values=dict([ + ('mjTASK_NEW', 0), + ('mjTASK_QUEUED', 1), + ('mjTASK_COMPLETED', 2), + ]), + )), ('mjtButton', EnumDecl( name='mjtButton', diff --git a/introspect/functions.py b/introspect/functions.py index e5323e4c..a27e95db 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -8269,7 +8269,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ type=ValueType(name='size_t'), ), ), - doc='Creates a thread pool with the specified number of threads running.', # pylint: disable=line-too-long + doc='Create a thread pool with the specified number of threads running.', # pylint: disable=line-too-long )), ('mju_threadPoolEnqueue', FunctionDecl( @@ -8288,32 +8288,8 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ inner_type=ValueType(name='mjTask'), ), ), - FunctionParameterDecl( - name='start_routine', - type=ValueType(name='void *(*)(void *)'), - ), - FunctionParameterDecl( - name='args', - type=PointerType( - inner_type=ValueType(name='void'), - ), - ), ), - doc='Enqueues a task in a thread pool.', - )), - ('mju_taskJoin', - FunctionDecl( - name='mju_taskJoin', - return_type=ValueType(name='void'), - parameters=( - FunctionParameterDecl( - name='task', - type=PointerType( - inner_type=ValueType(name='mjTask'), - ), - ), - ), - doc='Waits for a task to complete.', + doc='Enqueue a task in a thread pool.', )), ('mju_threadPoolDestroy', FunctionDecl( @@ -8327,6 +8303,34 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc='Destroys a thread pool.', + 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.', )), ]) diff --git a/introspect/structs.py b/introspect/structs.py index b9877656..4f6a1a6c 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -7007,33 +7007,39 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), ), )), - ('mjTask', - StructDecl( - name='mjTask', - declname='struct mjTask_', - fields=( - StructFieldDecl( - name='buffer', - type=ArrayType( - inner_type=ValueType(name='char'), - extents=(24,), - ), - doc='', - ), - ), - )), ('mjThreadPool', StructDecl( name='mjThreadPool', declname='struct mjThreadPool_', fields=( StructFieldDecl( - name='buffer', - type=ArrayType( - inner_type=ValueType(name='char'), - extents=(6208,), + 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='', + doc='arguments to func', + ), + StructFieldDecl( + name='status', + type=ValueType(name='int', is_volatile=True), + doc='status of the task', ), ), )), diff --git a/sample/cmake/SampleOptions.cmake b/sample/cmake/SampleOptions.cmake index 3e3f080f..be68de3d 100644 --- a/sample/cmake/SampleOptions.cmake +++ b/sample/cmake/SampleOptions.cmake @@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__. set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5 - -Wno-maybe-uninitialized + -Wno-maybe-uninitialized -Wno-c++20-extensions ) endif() + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator) + endif() endif() include(MujocoHarden) diff --git a/simulate/cmake/SimulateOptions.cmake b/simulate/cmake/SimulateOptions.cmake index 3e3f080f..be68de3d 100644 --- a/simulate/cmake/SimulateOptions.cmake +++ b/simulate/cmake/SimulateOptions.cmake @@ -99,9 +99,12 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR (CMAKE_CXX_COMPILER_ID MATCHES "Clang if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") # Set -Wimplicit-fallthrough=5 to only allow fallthrough annotation via __attribute__. set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wimplicit-fallthrough=5 - -Wno-maybe-uninitialized + -Wno-maybe-uninitialized -Wno-c++20-extensions ) endif() + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + set(EXTRA_COMPILE_OPTIONS ${EXTRA_COMPILE_OPTIONS} -Wno-c++20-designator) + endif() endif() include(MujocoHarden) diff --git a/src/thread/CMakeLists.txt b/src/thread/CMakeLists.txt index 948cd2e7..e320c20b 100644 --- a/src/thread/CMakeLists.txt +++ b/src/thread/CMakeLists.txt @@ -13,11 +13,11 @@ # limitations under the License. set(MUJOCO_THREAD_SRCS - lockless_queue.h - task.cc - task.h 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/task.h b/src/thread/task.h deleted file mode 100644 index 1a4f8c28..00000000 --- a/src/thread/task.h +++ /dev/null @@ -1,70 +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. -// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h" -// IWYU pragma: friend "third_party/(py/)?mujoco/.*" - -#ifndef MUJOCO_SRC_THREAD_TASK_H_ -#define MUJOCO_SRC_THREAD_TASK_H_ - -#ifdef __cplusplus - -#include -#include -#include - -namespace mujoco { - -class Task { - public: - using FunctionPtr = void* (*)(void*); - enum Status { - QUEUED, - COMPLETE, - }; - - static void Initialize( - Task* task, - FunctionPtr start_routine, - void* args) { - // instantiate a task at the pointer passed in - new(task) Task(); - task->start_routine_ = start_routine; - task->args_ = args; - task->status_ = Status::QUEUED; - } - - void Execute() { - args_ = start_routine_(args_); - status_ = Status::COMPLETE; - } - - void Join() { - while (status_ != Status::COMPLETE) { - std::this_thread::yield(); - } - } - - private: - FunctionPtr start_routine_; - - void* args_; - - std::atomic status_ = Status::QUEUED; -}; - -} // namespace mujoco - -#endif // __cplusplus - -#endif // MUJOCO_SRC_THREAD_TASK_H_ diff --git a/src/thread/thread_pool.cc b/src/thread/thread_pool.cc index dd060737..8842a242 100644 --- a/src/thread/thread_pool.cc +++ b/src/thread/thread_pool.cc @@ -14,39 +14,114 @@ #include "thread/thread_pool.h" +#include #include +#include +#include +#include #include -#include -#include "thread/task.h" +#include "engine/engine_crossplatform.h" +#include "engine/engine_util_errmem.h" +#include "thread/thread_queue.h" +#include "thread/thread_task.h" -static constexpr size_t kMaxThreads = 128; +namespace mujoco { +namespace { +constexpr size_t kThreadPoolQueueSize = 640; + +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_ { + .func = &ShutdownFunction, + .args = nullptr, + }; +}; +} // 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{.nworker = num_worker} { + // initialize worker threads + for (int i = 0; i < num_worker; ++i) { + workers_.push_back( + {std::make_unique(ThreadPoolWorker, this)}); + } + } + + // 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(); + } + } + + ~ThreadPoolImpl() { Shutdown(); } + + private: + // method executed by running threads + static void ThreadPoolWorker(ThreadPoolImpl* thread_pool) { + 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_; +}; // create a thread pool mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) { - mujoco::ThreadPool* thread_pool = - new mujoco::ThreadPool(number_of_threads); - return static_cast(static_cast(thread_pool)); + return new ThreadPoolImpl(number_of_threads); } // start a task in the threadpool -void mju_threadPoolEnqueue( - mjThreadPool* thread_pool, mjTask* task, mjStartRoutine start_routine, - void* args) { - mujoco::ThreadPool* thread_pool_ptr = - static_cast*>( - static_cast(thread_pool)); - thread_pool_ptr->Enqueue( - static_cast(static_cast(task)), start_routine, - args); +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) { - mujoco::ThreadPool* thread_pool_ptr = - static_cast*>( - static_cast(thread_pool)); - thread_pool_ptr->Shutdown(); - delete thread_pool_ptr; + auto thread_pool_impl = static_cast(thread_pool); + thread_pool_impl->Shutdown(); + delete thread_pool_impl; } - +} // namespace mujoco diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h index 11b51145..1d539437 100644 --- a/src/thread/thread_pool.h +++ b/src/thread/thread_pool.h @@ -11,92 +11,32 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h" -// IWYU pragma: friend "third_party/(py/)?mujoco/.*" #ifndef MUJOCO_SRC_THREAD_THREAD_POOL_H_ #define MUJOCO_SRC_THREAD_THREAD_POOL_H_ +#include + +#include +#include + #ifdef __cplusplus - -#include -#include -#include - -#include "thread/lockless_queue.h" -#include "thread/task.h" - namespace mujoco { +extern "C" { +#endif -static constexpr size_t kThreadPoolQueueSize = 640; +// Create a thread pool with the specified number of threads running. +MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); -template -class ThreadPool { - public: - ThreadPool(size_t number_of_threads) - : number_of_threads_(number_of_threads) { - for (int i = 0; i < number_of_threads_; ++i) { - threads_[i] = std::thread(ThreadPoolWorker, static_cast(this)); - } - } +// Enqueue a task in a thread pool. +MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task); - // start a task in the threadpool - void Enqueue( - Task* task, Task::FunctionPtr start_routine, void* args) { - Task::Initialize(task, start_routine, args); - lockless_queue_.push(static_cast(task)); - } - - // shutdown the threadpool - void Shutdown() { - if (shutdown_) { - return; - } - - shutdown_ = true; - Task shutdown_tasks[max_number_of_threads]; - for (int i = 0; i < number_of_threads_; ++i) { - Enqueue(&shutdown_tasks[i], ShutdownFunction, nullptr); - } - - for (int i = 0; i < number_of_threads_; ++i) { - threads_[i].join(); - } - } - - ~ThreadPool() { Shutdown(); } - - private: - // method executed by running threads - static void ThreadPoolWorker(void* arg) { - ThreadPool* thread_pool = - static_cast*>(arg); - while (!thread_pool->shutdown_) { - Task* task = static_cast(thread_pool->lockless_queue_.pop()); - task->Execute(); - } - } - - // shutdown function passed to running threads to ensure cleans shutdown - static void* ShutdownFunction(void* args) { - return NULL; - } - - // is the thread pool is being shut down - std::atomic shutdown_ = false; - - // actual number of running threads in the threadpool - const size_t number_of_threads_; - - // OS threads that are running in this pool - std::thread threads_[max_number_of_threads]; - - // queue of tasks to execute - LocklessQueue lockless_queue_; -}; +// Destroy a thread pool. +MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool); +#ifdef __cplusplus +} // extern "C" } // namespace mujoco - #endif // __cplusplus #endif // MUJOCO_SRC_THREAD_THREAD_POOL_H_ diff --git a/src/thread/lockless_queue.h b/src/thread/thread_queue.h similarity index 97% rename from src/thread/lockless_queue.h rename to src/thread/thread_queue.h index fbe21b96..3178f8ec 100644 --- a/src/thread/lockless_queue.h +++ b/src/thread/thread_queue.h @@ -11,8 +11,6 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -// IWYU pragma: private, include "third_party/mujoco/include/mujoco.h" -// IWYU pragma: friend "third_party/(py/)?mujoco/.*" #ifndef MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_ #define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_ diff --git a/src/thread/task.cc b/src/thread/thread_task.cc similarity index 68% rename from src/thread/task.cc rename to src/thread/thread_task.cc index acb1858b..77ce2bfd 100644 --- a/src/thread/task.cc +++ b/src/thread/thread_task.cc @@ -12,13 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "thread/task.h" +#include "thread/thread_task.h" + +#include #include -#include -// waits for a task to complete -void mju_taskJoin(mjTask* task) { - mujoco::Task* task_ptr = static_cast(static_cast(task)); - task_ptr->Join(); +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 new file mode 100644 index 00000000..81c63dd7 --- /dev/null +++ b/src/thread/thread_task.h @@ -0,0 +1,50 @@ +// 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 +#include + +#include +#include + +#ifdef __cplusplus +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/thread/CMakeLists.txt b/test/thread/CMakeLists.txt new file mode 100644 index 00000000..920bea0d --- /dev/null +++ b/test/thread/CMakeLists.txt @@ -0,0 +1,19 @@ +# 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) +target_link_libraries(thread_pool_test fixture gmock) + +mujoco_test(thread_queue_test) +target_link_libraries(thread_queue_test fixture gmock) diff --git a/test/thread/mjthread_test.cc b/test/thread/mjthread_test.cc deleted file mode 100644 index ba3ee7e9..00000000 --- a/test/thread/mjthread_test.cc +++ /dev/null @@ -1,90 +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 "src/thread/task.h" -#include "src/thread/thread_pool.h" - -namespace { - -struct TestFunctionArgs_ { - int input; - // make this atomic to avoid red-herring tsan failures. - std::atomic 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, EnsureStructClassSizeMatch) { - EXPECT_EQ(sizeof(mjTask), sizeof(mujoco::Task)); - EXPECT_EQ(sizeof(mjThreadPool), sizeof(mujoco::ThreadPool<128>)); -} - -TEST(TestMjThreadPool, TestMjThreadPool10Threads) { - mjThreadPool* thread_pool = mju_threadPoolCreate(10); - - TestFunctionArgs test_function_args[1000]; - mjTask tasks[1000]; - for (int i = 0; i < 1000; ++i) { - test_function_args[i].input = i; - mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function, - (void*)&test_function_args[i]); - } - - for (int i = 0; i < 1000; ++i) { - mju_taskJoin(&tasks[i]); - } - - for (int i = 0; i < 1000; ++i) { - EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); - } - mju_threadPoolDestroy(thread_pool); -} - -TEST(TestMjThreadPool, TestMjThreadPool100Threads) { - mjThreadPool* thread_pool = mju_threadPoolCreate(100); - - TestFunctionArgs test_function_args[1000]; - mjTask tasks[1000]; - for (int i = 0; i < 1000; ++i) { - test_function_args[i].input = i; - mju_threadPoolEnqueue(thread_pool, &tasks[i], test_function, - (void*)&test_function_args[i]); - } - - for (int i = 0; i < 1000; ++i) { - mju_taskJoin(&tasks[i]); - } - - for (int i = 0; i < 1000; ++i) { - EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); - } - - mju_threadPoolDestroy(thread_pool); -} - -} // namespace diff --git a/test/thread/task_test.cc b/test/thread/task_test.cc deleted file mode 100644 index 7bd53274..00000000 --- a/test/thread/task_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/task.h" - -#include - -namespace mujoco { -namespace { - -struct TestFunctionArgs { - int input; - int output; -}; - -void* test_function(void* args) { - TestFunctionArgs* test_function_args = (TestFunctionArgs*)args; - test_function_args->output = test_function_args->input; - return nullptr; -} - -TEST(TestMjThread, TestMjThread) { - TestFunctionArgs test_function_args; - test_function_args.input = 1; - test_function_args.output = 2; - Task task; - Task::Initialize( - &task, test_function, static_cast(&test_function_args)); - task.Execute(); - task.Join(); - EXPECT_EQ(test_function_args.input, test_function_args.output); -} - -} // namespace -} // namespace mujoco diff --git a/test/thread/thread_pool_test.cc b/test/thread/thread_pool_test.cc index a488e21f..2532ec65 100644 --- a/test/thread/thread_pool_test.cc +++ b/test/thread/thread_pool_test.cc @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "src/thread/thread_pool.h" - #include #include #include @@ -21,75 +19,82 @@ #include #include -#include "src/thread/task.h" +#include namespace mujoco { namespace { -struct TestFunctionArgs { +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) { - ThreadPool<10> thread_pool(10); + mjThreadPool* thread_pool = mju_threadPoolCreate(10); constexpr int kTasks = 1000; TestFunctionArgs test_function_args[kTasks]; - Task tasks[kTasks]; + mjTask tasks[kTasks]; for (int i = 0; i < kTasks; ++i) { test_function_args[i].input = i; - thread_pool.Enqueue( - &tasks[i], test_function, static_cast(&test_function_args[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) { - tasks[i].Join(); + mju_taskJoin(&tasks[i]); } for (int i = 0; i < kTasks; ++i) { EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); } - - thread_pool.Shutdown(); + mju_threadPoolDestroy(thread_pool); } TEST(TestMjThreadPool, TestMjThreadPool100Threads) { - ThreadPool<100> thread_pool(100); + mjThreadPool* thread_pool = mju_threadPoolCreate(100); constexpr int kTasks = 1000; TestFunctionArgs test_function_args[kTasks]; - Task tasks[kTasks]; + mjTask tasks[kTasks]; for (int i = 0; i < kTasks; ++i) { test_function_args[i].input = i; - thread_pool.Enqueue( - &tasks[i], test_function, static_cast(&test_function_args[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) { - tasks[i].Join(); + mju_taskJoin(&tasks[i]); } for (int i = 0; i < kTasks; ++i) { EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); } - thread_pool.Shutdown(); + mju_threadPoolDestroy(thread_pool); } TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) { - ThreadPool<10> thread_pool(10); + mjThreadPool* thread_pool = mju_threadPoolCreate(10); constexpr int kTasks = 20; TestFunctionArgs test_function_args[kTasks]; - Task tasks[kTasks]; + mjTask tasks[kTasks]; std::unique_ptr enqueue_threads[kTasks]; // add tasks to the thread pool from many threads @@ -98,6 +103,10 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) { 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 { @@ -105,8 +114,7 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) { start_cv.wait(lock, [&] { return start; }); } // enqueue outside the lock, to get some concurrency - thread_pool.Enqueue( - &tasks[i], test_function, static_cast(&test_function_args[i])); + mju_threadPoolEnqueue(thread_pool, &tasks[i]); }); } { @@ -120,14 +128,14 @@ TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) { } for (int i = 0; i < kTasks; ++i) { - tasks[i].Join(); + mju_taskJoin(&tasks[i]); } for (int i = 0; i < kTasks; ++i) { EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); } - thread_pool.Shutdown(); + mju_threadPoolDestroy(thread_pool); } } // namespace diff --git a/test/thread/lockless_queue_test.cc b/test/thread/thread_queue_test.cc similarity index 97% rename from test/thread/lockless_queue_test.cc rename to test/thread/thread_queue_test.cc index 774f0ac4..7b68b6f9 100644 --- a/test/thread/lockless_queue_test.cc +++ b/test/thread/thread_queue_test.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "src/thread/lockless_queue.h" +#include "src/thread/thread_queue.h" #include diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 685233a1..aae87bf5 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -403,6 +403,11 @@ public enum mjtFont : int{ mjFONT_SHADOW = 1, mjFONT_BIG = 2, } +public enum mjtTaskStatus : int{ + mjTASK_NEW = 0, + mjTASK_QUEUED = 1, + mjTASK_COMPLETED = 2, +} public enum mjtButton : int{ mjBUTTON_NONE = 0, mjBUTTON_LEFT = 1, @@ -2360,16 +2365,6 @@ public unsafe struct mjrContext_ { public int readPixelFormat; } -[StructLayout(LayoutKind.Sequential)] -public unsafe struct mjTask_ { - public fixed sbyte buffer[24]; -} - -[StructLayout(LayoutKind.Sequential)] -public unsafe struct mjThreadPool_ { - public fixed sbyte buffer[6208]; -} - [StructLayout(LayoutKind.Sequential)] public unsafe struct mjuiState_ { public int nrect; @@ -3950,14 +3945,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 mjThreadPool_* mju_threadPoolCreate(UIntPtr number_of_threads); - -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mju_taskJoin(mjTask_* task); - -[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mju_threadPoolDestroy(mjThreadPool_* thread_pool); } }