chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor thread pool implementation.

Browse Source 
- Make mjTask non-opaque and remove C++ Task class.
- Make C++ thread pool a subclass of a skeletal mjThreadPool C struct.
- Make the mjTask status enum more consistent with the rest of MuJoCo.
- Change mju_threadPoolEnqueue to take just the mjTask. Users must now prepare the mjTask by assigning the function pointer and argument into the struct.
- Rename files in thread/ to be more consistent with the rest of MuJoCo.
- Run threading tests in CMake.
- Allow use of C++20 designated initializers.

Otherwise the functionality remains identical.

PiperOrigin-RevId: 564374675
Change-Id: I37c9894566bc39faf217e5aa97a4e2713a70e467
This commit is contained in:
Saran Tunyasuvunakool
2023-09-11 07:19:50 -07:00
committed by Copybara-Service
parent 0ba10cc4f0
commit 78183e60e1
24 changed files with 379 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/MujocoOptions.cmake
+4 -1
View File
@@ -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)
doc/APIreference/functions.rst
+21 -12
View File
@@ -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.
doc/includes/references.h
+15 -9
View File
@@ -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
include/mujoco/mjthread.h
+15 -22
View File
@@ -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 <stddef.h>
#include <mujoco/mjexport.h>
// These types are implemented in C++, they're just used as opaque pointers in C
// to provide type safety for functions.
struct mjTask_ {
char buffer[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_
include/mujoco/mujoco.h
+10 -8
View File
@@ -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)
}
introspect/codegen/generate_enums.py
+4 -1
View File
@@ -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))
introspect/enums.py
+10
View File
@@ -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',
introspect/functions.py
+31 -27
View File
@@ -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.',
)),
])
introspect/structs.py
+26 -20
View File
@@ -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',
),
),
)),
sample/cmake/SampleOptions.cmake
+4 -1
View File
@@ -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)
simulate/cmake/SimulateOptions.cmake
+4 -1
View File
@@ -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)
src/thread/CMakeLists.txt
+3 -3
View File
@@ -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})
src/thread/task.h
-70
View File
@@ -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 <atomic>
#include <new>
#include <thread>
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_ = Status::QUEUED;
};
} // namespace mujoco
#endif // __cplusplus
#endif // MUJOCO_SRC_THREAD_TASK_H_
src/thread/thread_pool.cc
+96 -21
View File
@@ -14,39 +14,114 @@
#include "thread/thread_pool.h"
#include <atomic>
#include <cstddef>
#include <memory>
#include <thread>
#include <vector>
#include <mujoco/mjthread.h>
#include <mujoco/mujoco.h>
#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<std::thread> 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<std::thread>(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<mjTask> 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<mjTask*>(thread_pool->lockless_queue_.pop());
task->args = task->func(task->args);
GetAtomicTaskStatus(task).store(mjTASK_COMPLETED);
}
}
// indicates whether the thread pool is being shut down
std::atomic<bool> shutdown_ = false;
// OS threads that are running in this pool
std::vector<WorkerThread> workers_;
// queue of tasks to execute
mujoco::LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
};
// create a thread pool
mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
mujoco::ThreadPool<kMaxThreads>* thread_pool =
new mujoco::ThreadPool<kMaxThreads>(number_of_threads);
return static_cast<mjThreadPool*>(static_cast<void*>(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<kMaxThreads>* thread_pool_ptr =
static_cast<mujoco::ThreadPool<kMaxThreads>*>(
static_cast<void*>(thread_pool));
thread_pool_ptr->Enqueue(
static_cast<mujoco::Task*>(static_cast<void*>(task)), start_routine,
args);
void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task) {
auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
thread_pool_impl->Enqueue(task);
}
// shutdown the threadpool and free the memory
void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
mujoco::ThreadPool<kMaxThreads>* thread_pool_ptr =
static_cast<mujoco::ThreadPool<kMaxThreads>*>(
static_cast<void*>(thread_pool));
thread_pool_ptr->Shutdown();
delete thread_pool_ptr;
auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
thread_pool_impl->Shutdown();
delete thread_pool_impl;
}
} // namespace mujoco
src/thread/thread_pool.h
+15 -75
View File
@@ -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 <stddef.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjthread.h>
#ifdef __cplusplus
#include <atomic>
#include <cstddef>
#include <thread>
#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 <size_t max_number_of_threads>
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<void*>(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<void*>(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<max_number_of_threads>* thread_pool =
static_cast<ThreadPool<max_number_of_threads>*>(arg);
while (!thread_pool->shutdown_) {
Task* task = static_cast<Task*>(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<bool> 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<void*, kThreadPoolQueueSize> 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_
src/thread/lockless_queue.h → src/thread/thread_queue.h
-2
View File
@@ -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_
src/thread/task.cc → src/thread/thread_task.cc
+15 -6
View File
@@ -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 <thread>
#include <mujoco/mjthread.h>
#include <mujoco/mujoco.h>
// waits for a task to complete
void mju_taskJoin(mjTask* task) {
mujoco::Task* task_ptr = static_cast<mujoco::Task*>(static_cast<void*>(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
src/thread/thread_task.h
+50
View File
@@ -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 <atomic>
#include <new>
#include <type_traits>
#include <mujoco/mjexport.h>
#include <mujoco/mjthread.h>
#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<decltype(mjTask::status)>;
inline std::atomic<TaskStatus>& GetAtomicTaskStatus(mjTask* task) {
static_assert(sizeof(std::atomic<TaskStatus>) == sizeof(TaskStatus));
static_assert(alignof(std::atomic<TaskStatus>) == alignof(TaskStatus));
static_assert(std::atomic<TaskStatus>::is_always_lock_free);
return *std::launder(reinterpret_cast<std::atomic<TaskStatus>*>(
const_cast<TaskStatus*>(&task->status)));
}
} // namespace mujoco
#endif // __cplusplus
#endif // MUJOCO_SRC_THREAD_THREAD_TASK_H_
test/thread/CMakeLists.txt
+19
View File
@@ -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)
test/thread/mjthread_test.cc
-90
View File
@@ -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 <mujoco/mjthread.h>
#include <atomic>
#include <gtest/gtest.h>
#include <mujoco/mujoco.h>
#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<int> output;
};
typedef struct TestFunctionArgs_ TestFunctionArgs;
void* test_function(void* args) {
TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
if (!test_function_args) {
return nullptr;
}
test_function_args->output = test_function_args->input;
return nullptr;
}
TEST(TestMjThreadPool, 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
test/thread/task_test.cc
-46
View File
@@ -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 <gtest/gtest.h>
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<void*>(&test_function_args));
task.Execute();
task.Join();
EXPECT_EQ(test_function_args.input, test_function_args.output);
}
} // namespace
} // namespace mujoco
test/thread/thread_pool_test.cc
+31 -23
View File
@@ -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 <atomic>
#include <condition_variable>
#include <memory>
@@ -21,75 +19,82 @@
#include <thread>
#include <gtest/gtest.h>
#include "src/thread/task.h"
#include <mujoco/mujoco.h>
namespace mujoco {
namespace {
struct TestFunctionArgs {
struct TestFunctionArgs_ {
int input;
// make this atomic to avoid red-herring tsan failures.
std::atomic<int> output;
};
typedef struct TestFunctionArgs_ TestFunctionArgs;
void* test_function(void* args) {
TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
if (!test_function_args) {
return nullptr;
}
test_function_args->output = test_function_args->input;
return nullptr;
}
TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
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<void*>(&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<void*>(&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<std::thread> 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<std::thread>([&, 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<void*>(&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
test/thread/lockless_queue_test.cc → test/thread/thread_queue_test.cc
+1 -1
View File
@@ -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 <cstddef>
unity/Runtime/Bindings/MjBindings.cs
+5 -19
View File
@@ -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);
}
}