diff --git a/CMakeLists.txt b/CMakeLists.txt index 1caaa05c..bc6c5126 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -59,6 +59,7 @@ set(MUJOCO_HEADERS include/mujoco/mjmodel.h include/mujoco/mjplugin.h include/mujoco/mjrender.h + include/mujoco/mjthread.h include/mujoco/mjtnum.h include/mujoco/mjui.h include/mujoco/mjvisualize.h @@ -88,6 +89,7 @@ add_subdirectory(src/engine) add_subdirectory(src/user) add_subdirectory(src/xml) add_subdirectory(src/render) +add_subdirectory(src/thread) add_subdirectory(src/ui) target_compile_definitions(mujoco PRIVATE _GNU_SOURCE CCD_STATIC_DEFINE MUJOCO_DLL_EXPORTS -DMC_IMPLEM_ENABLE) diff --git a/doc/APIreference/APIfunctions.rst b/doc/APIreference/APIfunctions.rst index fca904c7..ce8309ea 100644 --- a/doc/APIreference/APIfunctions.rst +++ b/doc/APIreference/APIfunctions.rst @@ -32,6 +32,7 @@ API function can be classified as: - :ref:`Derivatives`. - :ref:`Plugin` related functions. - :ref:`Macros`. +- :ref:`Thread` related functions. .. TODO(b/273075045): Better category-label namespacing. diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index 81a5ff17..d8ff40ad 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -712,7 +712,26 @@ Options for configuring the automatic :ref:`actuator length-range computation`__ Defines data structures required by :ref:`engine plugins`. +`mjthread.h `__ + Defines data structures and functions required by :ref:`thread`. .. _inVersion: diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 51522ca3..408e62f8 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -402,6 +402,9 @@ struct mjData_ { mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) mjtNum* efc_force; // constraint force in constraint space (nefc x 1) int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + + // ThreadPool for multithreaded operations + uintptr_t threadpool; }; typedef struct mjData_ mjData; diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h new file mode 100644 index 00000000..1538e954 --- /dev/null +++ b/include/mujoco/mjthread.h @@ -0,0 +1,47 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_INCLUDE_MJTHREAD_H_ +#define MUJOCO_INCLUDE_MJTHREAD_H_ + +// C API for MuJoCo threading +#ifdef __cplusplus +extern "C" { +#endif + +#include + +#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[48]; +}; +typedef struct mjTask_ mjTask; + +struct mjThreadPool_ { + char buffer[6208]; +}; +typedef struct mjThreadPool_ mjThreadPool; + +typedef void*(*mjStartRoutine_)(void*); +typedef mjStartRoutine_ mjStartRoutine; + +#ifdef __cplusplus +} +#endif + + +#endif // MUJOCO_INCLUDE_MJTHREAD_H_ diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index bccfcde2..2ef1c7ee 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -644,7 +644,8 @@ X( int, nnzJ ) \ X( int, ncon ) \ X( int, nisland ) \ - X( mjtNum, time ) + X( mjtNum, time ) \ + X( uintptr_t, threadpool ) // vector fields of mjData diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index ce813376..c1c6eb41 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -36,6 +36,7 @@ extern "C" { #include #include #include +#include #include #include #include @@ -1304,6 +1305,21 @@ MJAPI const mjpResourceProvider* mjp_getResourceProvider(const char* resource_na // If invalid slot number, return NULL. MJAPI const mjpResourceProvider* mjp_getResourceProviderAtSlot(int slot); +//---------------------- Thread ------------------------------------------------------------------- + +// Creates a thread pool with the specified number of threads running. +MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); + +// Enqueues a task in a thread pool. +MJAPI void mju_threadPoolEnqueue( + mjThreadPool* thread_pool, mjTask* task, void*(start_routine)(void*), + void* args); + +// Waits for a task to complete. +MJAPI void mju_taskJoin(mjTask* task); + +// Destroys a thread pool. +MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool); #if defined(__cplusplus) } diff --git a/introspect/ast_nodes.py b/introspect/ast_nodes.py index 926cff42..9aa58c00 100644 --- a/introspect/ast_nodes.py +++ b/introspect/ast_nodes.py @@ -66,6 +66,7 @@ class ValueType: def __init__(self, name: str, is_const: bool = False, is_volatile: bool = False): is_valid_type_name = ( + name == 'void *(*)(void *)' or VALID_TYPE_NAME_PATTERN.fullmatch(name) or _is_valid_integral_type(name)) and name not in C_INVALID_TYPE_NAMES if not is_valid_type_name: diff --git a/introspect/functions.py b/introspect/functions.py index 984f4eda..cf2417d2 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -8257,4 +8257,76 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Look up a resource provider by slot number returned by mjp_registerResourceProvider. If invalid slot number, return NULL.', # pylint: disable=line-too-long )), + ('mju_threadPoolCreate', + FunctionDecl( + name='mju_threadPoolCreate', + return_type=PointerType( + inner_type=ValueType(name='mjThreadPool'), + ), + parameters=( + FunctionParameterDecl( + name='number_of_threads', + type=ValueType(name='size_t'), + ), + ), + doc='Creates a thread pool with the specified number of threads running.', # 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'), + ), + ), + 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.', + )), + ('mju_threadPoolDestroy', + FunctionDecl( + name='mju_threadPoolDestroy', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='thread_pool', + type=PointerType( + inner_type=ValueType(name='mjThreadPool'), + ), + ), + ), + doc='Destroys a thread pool.', + )), ]) diff --git a/introspect/structs.py b/introspect/structs.py index 1573fb0c..4bdc211c 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4475,6 +4475,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='constraint state (mjtConstraintState) (nefc x 1)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='threadpool', + type=ValueType(name='uintptr_t'), + doc='ThreadPool for multithreaded operations', + ), ), )), ('mjvPerturb', @@ -6978,6 +6983,36 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), ), )), + ('mjTask', + StructDecl( + name='mjTask', + declname='struct mjTask_', + fields=( + StructFieldDecl( + name='buffer', + type=ArrayType( + inner_type=ValueType(name='char'), + extents=(48,), + ), + doc='', + ), + ), + )), + ('mjThreadPool', + StructDecl( + name='mjThreadPool', + declname='struct mjThreadPool_', + fields=( + StructFieldDecl( + name='buffer', + type=ArrayType( + inner_type=ValueType(name='char'), + extents=(6208,), + ), + doc='', + ), + ), + )), ('mjuiState', StructDecl( name='mjuiState', diff --git a/introspect/type_parsing.py b/introspect/type_parsing.py index 80bf9668..dfc2ab53 100644 --- a/introspect/type_parsing.py +++ b/introspect/type_parsing.py @@ -68,6 +68,8 @@ def _parse_maybe_pointer( ast_nodes.PointerType]] ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]: """Internal-only helper that parses a type that may be a pointer type.""" + if type_name == 'void *(*)(void *)': + return ast_nodes.ValueType(name=type_name) p = type_name.rfind('*') if p != -1: leftover, is_qualifier = _parse_qualifiers( @@ -107,6 +109,9 @@ def _peel_nested_parens(input_str: str) -> MutableSequence[str]: A sequence of substrings enclosed with in respective parentheses. See the description above for the precise detail of the output. """ + if input_str == 'void *(*)(void *)': + return ['void *(*)(void *)'] + start = input_str.find('(') end = input_str.rfind(')') @@ -146,4 +151,4 @@ def parse_type( def parse_function_return_type( type_name: str ) -> Union[ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]: - return parse_type(type_name[:type_name.rfind('(')]) + return parse_type(type_name[:type_name.find('(')]) diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index f713c718..5df0f59e 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1114,6 +1114,8 @@ static mjData* _makeData(const mjModel* m) { } } + d->threadpool = 0; + return d; } @@ -1202,6 +1204,8 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { } } + dest->threadpool = src->threadpool; + return dest; } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index d9455c4d..6e215a9c 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -778,7 +778,8 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, #define X(type, name) \ if (strcmp(#name, "pstack") != 0 && \ strcmp(#name, "pbase") != 0 && \ - strcmp(#name, "parena") != 0) { \ + strcmp(#name, "parena") != 0 && \ + strcmp(#name, "threadpool") != 0) { \ const char* format = _Generic( \ d->name, \ int : INT_FORMAT, \ @@ -794,6 +795,16 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, MJDATA_SCALAR #undef X + + int threadpool = 0; + if (d->threadpool) { + threadpool = 1; + } + fprintf(fp, " "); + fprintf(fp, NAME_FORMAT, "threadpool"); + fprintf(fp, INT_FORMAT, threadpool); + fprintf(fp, "\n"); + fprintf(fp, "\n"); // WARNING diff --git a/src/thread/CMakeLists.txt b/src/thread/CMakeLists.txt new file mode 100644 index 00000000..948cd2e7 --- /dev/null +++ b/src/thread/CMakeLists.txt @@ -0,0 +1,23 @@ +# 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 + lockless_queue.h + task.cc + task.h + thread_pool.cc + thread_pool.h +) + +target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS}) diff --git a/src/thread/lockless_queue.h b/src/thread/lockless_queue.h new file mode 100644 index 00000000..fbe21b96 --- /dev/null +++ b/src/thread/lockless_queue.h @@ -0,0 +1,154 @@ +// 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_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/task.cc b/src/thread/task.cc new file mode 100644 index 00000000..acb1858b --- /dev/null +++ b/src/thread/task.cc @@ -0,0 +1,24 @@ +// 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/task.h" + +#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(); +} diff --git a/src/thread/task.h b/src/thread/task.h new file mode 100644 index 00000000..390a4571 --- /dev/null +++ b/src/thread/task.h @@ -0,0 +1,69 @@ +// 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: + enum Status { + QUEUED, + COMPLETE, + }; + + static void Initialize( + Task* task, + std::function 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: + std::function 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 new file mode 100644 index 00000000..dd060737 --- /dev/null +++ b/src/thread/thread_pool.cc @@ -0,0 +1,52 @@ +// 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 "thread/task.h" + +static constexpr size_t kMaxThreads = 128; + +// 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)); +} + +// 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); +} + +// 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; +} + diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h new file mode 100644 index 00000000..d8526d54 --- /dev/null +++ b/src/thread/thread_pool.h @@ -0,0 +1,103 @@ +// 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_THREAD_POOL_H_ +#define MUJOCO_SRC_THREAD_THREAD_POOL_H_ + +#ifdef __cplusplus + +#include +#include +#include +#include + +#include "thread/lockless_queue.h" +#include "thread/task.h" + +namespace mujoco { + +static constexpr size_t kThreadPoolQueueSize = 640; + +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)); + } + } + + // start a task in the threadpool + void Enqueue( + Task* task, std::function 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[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_; +}; + +} // namespace mujoco + +#endif // __cplusplus + +#endif // MUJOCO_SRC_THREAD_THREAD_POOL_H_ diff --git a/test/thread/lockless_queue_test.cc b/test/thread/lockless_queue_test.cc new file mode 100644 index 00000000..774f0ac4 --- /dev/null +++ b/test/thread/lockless_queue_test.cc @@ -0,0 +1,46 @@ +// 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/lockless_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/test/thread/mjthread_test.cc b/test/thread/mjthread_test.cc new file mode 100644 index 00000000..ba3ee7e9 --- /dev/null +++ b/test/thread/mjthread_test.cc @@ -0,0 +1,90 @@ +// 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 new file mode 100644 index 00000000..7bd53274 --- /dev/null +++ b/test/thread/task_test.cc @@ -0,0 +1,46 @@ +// 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 new file mode 100644 index 00000000..a488e21f --- /dev/null +++ b/test/thread/thread_pool_test.cc @@ -0,0 +1,134 @@ +// 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_pool.h" + +#include +#include +#include +#include +#include + +#include +#include "src/thread/task.h" + +namespace mujoco { +namespace { + +struct TestFunctionArgs { + int input; + // make this atomic to avoid red-herring tsan failures. + std::atomic output; +}; + +void* test_function(void* args) { + TestFunctionArgs* test_function_args = static_cast(args); + test_function_args->output = test_function_args->input; + return nullptr; +} + +TEST(TestMjThreadPool, TestMjThreadPool10Threads) { + ThreadPool<10> thread_pool(10); + + constexpr int kTasks = 1000; + TestFunctionArgs test_function_args[kTasks]; + Task 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])); + } + + for (int i = 0; i < kTasks; ++i) { + tasks[i].Join(); + } + + for (int i = 0; i < kTasks; ++i) { + EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); + } + + thread_pool.Shutdown(); +} + +TEST(TestMjThreadPool, TestMjThreadPool100Threads) { + ThreadPool<100> thread_pool(100); + + constexpr int kTasks = 1000; + TestFunctionArgs test_function_args[kTasks]; + Task 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])); + } + + for (int i = 0; i < kTasks; ++i) { + tasks[i].Join(); + } + + for (int i = 0; i < kTasks; ++i) { + EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); + } + + thread_pool.Shutdown(); +} + +TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) { + ThreadPool<10> thread_pool(10); + + constexpr int kTasks = 20; + TestFunctionArgs test_function_args[kTasks]; + Task 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; + 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 + thread_pool.Enqueue( + &tasks[i], test_function, static_cast(&test_function_args[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) { + tasks[i].Join(); + } + + for (int i = 0; i < kTasks; ++i) { + EXPECT_EQ(test_function_args[i].input, test_function_args[i].output); + } + + thread_pool.Shutdown(); +} + +} // namespace +} // namespace mujoco diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index e886c3c5..37a1278a 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -56,6 +56,7 @@ public const bool mjEXTERNC = true; public const bool THIRD_PARTY_MUJOCO_MJRENDER_H_ = true; public const int mjNAUX = 10; public const int mjMAXTEXTURE = 1000; +public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true; public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTNUM_H_ = true; public const bool mjUSEDOUBLE = true; public const double mjMINVAL = 1e-15; @@ -1736,6 +1737,7 @@ public unsafe struct mjData_ { public double* efc_b; public double* efc_force; public int* efc_state; + public UIntPtr threadpool; } [StructLayout(LayoutKind.Sequential)] @@ -2354,6 +2356,16 @@ public unsafe struct mjrContext_ { public int readPixelFormat; } +[StructLayout(LayoutKind.Sequential)] +public unsafe struct mjTask_ { + public fixed sbyte buffer[48]; +} + +[StructLayout(LayoutKind.Sequential)] +public unsafe struct mjThreadPool_ { + public fixed sbyte buffer[6208]; +} + [StructLayout(LayoutKind.Sequential)] public unsafe struct mjuiState_ { public int nrect; @@ -3930,5 +3942,14 @@ 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); } }