diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 71efa04a..754ab5cc 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -34,6 +34,7 @@ #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" #include "engine/engine_vfs.h" +#include "thread/thread_pool.h" #ifdef ADDRESS_SANITIZER #include @@ -1215,15 +1216,29 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { } +static void maybe_lock_alloc_mutex(mjData* d) { + if (d->threadpool != 0) { + mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool); + } +} + +static void maybe_unlock_alloc_mutex(mjData* d) { + if (d->threadpool != 0) { + mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool); + } +} + // allocate memory from the mjData arena void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) { + maybe_lock_alloc_mutex(d); size_t misalignment = fastmod(d->parena, alignment); size_t padding = misalignment ? alignment - misalignment : 0; // check size size_t bytes_available = d->narena - d->pstack; if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) { + maybe_unlock_alloc_mutex(d); return NULL; } @@ -1240,56 +1255,46 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) { __msan_allocated_memory(result, bytes); #endif + maybe_unlock_alloc_mutex(d); return result; } - -// internal: allocate size bytes on the mjData stack +// internal: allocate size bytes on the provided stack shard // declared inline so that modular arithmetic with specific alignments can be optimized out -static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { +static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size, size_t alignment) { // return NULL if empty if (mjUNLIKELY(!size)) { return NULL; } - // size of entire arena/stack in bytes - size_t stack_size_bytes = d->narena; - - // end of the arena - uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + stack_size_bytes; - - // current top of the stack - uintptr_t end_ptr = end_of_arena_ptr - d->pstack; - // start of the memory to be allocated to the buffer - uintptr_t start_ptr = end_ptr - (size + mjREDZONE); + uintptr_t start_ptr = stack_info->top - (size + mjREDZONE); // align the pointer start_ptr -= fastmod(start_ptr, alignment); // new top of the stack - uintptr_t new_pstack_ptr = start_ptr - mjREDZONE; - size_t new_pstack = end_of_arena_ptr - new_pstack_ptr; + uintptr_t new_top_ptr = start_ptr - mjREDZONE; // exclude red zone from stack usage statistics - size_t current_alloc_usage = end_ptr - new_pstack_ptr - 2 * mjREDZONE; - size_t usage = current_alloc_usage + d->pstack; + size_t current_alloc_usage = stack_info->top - new_top_ptr - 2 * mjREDZONE; + size_t usage = current_alloc_usage + (stack_info->bottom - stack_info->top); // check size - size_t stack_available_bytes = end_ptr - ((uintptr_t)d->arena + d->parena); - size_t stack_required_bytes = end_ptr - new_pstack_ptr; + size_t stack_available_bytes = stack_info->top - stack_info->limit; + size_t stack_required_bytes = stack_info->top - new_top_ptr; if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) { mju_error("mj_stackAlloc: insufficient memory: max = %zu, available = %zu, requested = %zu " "(ne = %d, nf = %d, nefc = %d, ncon = %d)", - stack_size_bytes, stack_available_bytes, stack_required_bytes, + stack_info->bottom - stack_info->limit, stack_available_bytes, stack_required_bytes, d->ne, d->nf, d->nefc, d->ncon); } #ifdef ADDRESS_SANITIZER // actual stack usage (without red zone bytes) is stored in the red zone - if (d->pstack) { - char* prev_pstack_ptr = (char*)(end_of_arena_ptr - d->pstack); + if (stack_info->top != stack_info->bottom) { + char* prev_pstack_ptr = (char*)(stack_info->top); size_t prev_misalign = (uintptr_t)prev_pstack_ptr % _Alignof(size_t); size_t* prev_usage_ptr = (size_t*)(prev_pstack_ptr + @@ -1300,9 +1305,9 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { } // store new stack usage in the red zone - size_t misalign = new_pstack_ptr % _Alignof(size_t); + size_t misalign = new_top_ptr % _Alignof(size_t); size_t* usage_ptr = - (size_t*)(new_pstack_ptr + (misalign ? _Alignof(size_t) - misalign : 0)); + (size_t*)(new_top_ptr + (misalign ? _Alignof(size_t) - misalign : 0)); ASAN_UNPOISON_MEMORY_REGION(usage_ptr, sizeof(size_t)); *usage_ptr = usage; ASAN_POISON_MEMORY_REGION(usage_ptr, sizeof(size_t)); @@ -1311,45 +1316,95 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size); #endif - // update pstack and max usage statistics - d->pstack = new_pstack; - d->maxuse_stack = mjMAX(d->maxuse_stack, usage); - d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena); + // update max usage statistics + stack_info->top = new_top_ptr; + if (!d->threadpool) { + d->maxuse_stack = mjMAX(d->maxuse_stack, usage); + d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena); + } else { + size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); + d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage); + } return (void*)start_ptr; } +static inline mjStackInfo get_stack_info_from_data(mjData* d) { + mjStackInfo stack_info; + stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena; + stack_info.top = stack_info.bottom - d->pstack; + stack_info.limit = (uintptr_t)d->arena + (uintptr_t)d->parena; + stack_info.stack_base = d->pbase; + + return stack_info; +} + + +// internal: allocate size bytes in mjData +// declared inline so that modular arithmetic with specific alignments can be optimized out +static inline void* stackalloc(mjData* d, size_t size, size_t alignment) { + if (!d->threadpool) { + mjStackInfo stack_info = get_stack_info_from_data(d); + + void* result = stackallocinternal(d, &stack_info, size, alignment); + + d->pstack = stack_info.bottom - stack_info.top; + + return result; + } + + size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); + mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); + return stackallocinternal(d, stack_info, size, alignment); +} + + +// mjStackInfo mark stack frame, inline so ASAN errors point to correct code unit +#ifdef ADDRESS_SANITIZER +__attribute__((always_inline)) +#endif +static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) { + size_t top_old = stack_info->top; + mjStackFrame* s = + (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame)); + s->pbase = stack_info->stack_base; + s->pstack = top_old; +#ifdef ADDRESS_SANITIZER + // store the program counter to the caller so that we can compare against mj_freeStack later + s->pc = __sanitizer_return_address(); +#endif + stack_info->stack_base = (uintptr_t) s; +} + // mjData mark stack frame #ifdef ADDRESS_SANITIZER __attribute__((noinline)) #endif void mj_markStack(mjData* d) { - size_t pstack_old = d->pstack; - mjStackFrame* s = - (mjStackFrame*) stackalloc(d, sizeof(mjStackFrame), _Alignof(mjStackFrame)); - s->pbase = d->pbase; - s->pstack = pstack_old; -#ifdef ADDRESS_SANITIZER - // store the program counter to the caller so that we can compare against mj_freeStack later - s->pc = __sanitizer_return_address(); -#endif - d->pbase = d->pstack - mjREDZONE; -} - - - -// mjData free stack frame -#ifdef ADDRESS_SANITIZER -__attribute__((noinline)) -#endif -void mj_freeStack(mjData* d) { - if (mjUNLIKELY(!d->pbase)) { + if (!d->threadpool) { + mjStackInfo stack_info = get_stack_info_from_data(d); + markstackinternal(d, &stack_info); + d->pstack = stack_info.bottom - stack_info.top; + d->pbase = stack_info.stack_base; return; } - mjStackFrame* s = (mjStackFrame*) ((char*)d->arena + d->narena - d->pbase); + size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); + mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); + markstackinternal(d, stack_info); +} + +#ifdef ADDRESS_SANITIZER +__attribute__((always_inline)) +#endif +static inline void freestackinternal(mjStackInfo* stack_info) { + if (mjUNLIKELY(!stack_info->stack_base)) { + return; + } + + mjStackFrame* s = (mjStackFrame*) stack_info->stack_base; #ifdef ADDRESS_SANITIZER // raise an error if caller function name doesn't match the most recent caller of mj_markStack if (!_mj_comparePcFuncName(s->pc, __sanitizer_return_address())) { @@ -1365,15 +1420,34 @@ void mj_freeStack(mjData* d) { #endif // restore pbase and pstack - d->pbase = s->pbase; - d->pstack = s->pstack; + stack_info->stack_base = s->pbase; + stack_info->top = s->pstack; // if running under asan, poison the newly freed memory region #ifdef ADDRESS_SANITIZER - ASAN_POISON_MEMORY_REGION((char*)d->arena + d->parena, d->narena - d->pstack - d->parena); + ASAN_POISON_MEMORY_REGION((char*)stack_info->limit, stack_info->top - stack_info->limit); #endif } + +// mjData free stack frame +#ifdef ADDRESS_SANITIZER +__attribute__((noinline)) +#endif +void mj_freeStack(mjData* d) { + if (!d->threadpool) { + mjStackInfo stack_info = get_stack_info_from_data(d); + freestackinternal(&stack_info); + d->pstack = stack_info.bottom - stack_info.top; + d->pbase = stack_info.stack_base; + return; + } + + size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool); + mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id); + freestackinternal(stack_info); +} + void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) { return stackalloc(d, bytes, alignment); } diff --git a/src/thread/thread_pool.cc b/src/thread/thread_pool.cc index 1b1723ae..c36097fa 100644 --- a/src/thread/thread_pool.cc +++ b/src/thread/thread_pool.cc @@ -14,15 +14,19 @@ #include "thread/thread_pool.h" +#include + #include #include #include #include +#include #include #include #include #include +#include #include "engine/engine_crossplatform.h" #include "engine/engine_util_errmem.h" #include "thread/thread_queue.h" @@ -32,6 +36,11 @@ namespace mujoco { namespace { constexpr size_t kThreadPoolQueueSize = 640; +// Each thread being run will be assigned a worker_id. +// 0: main thread +// 1->n: workers +thread_local size_t worker_id = 0; + struct WorkerThread { // Shutdown function passed to running threads to ensure clean shutdown. static void* ShutdownFunction(void* args) { @@ -58,11 +67,15 @@ class ThreadPoolImpl : public mjThreadPool { // initialize worker threads for (int i = 0; i < std::min(num_worker, mjMAXTHREADS); ++i) { WorkerThread worker{ - std::make_unique(ThreadPoolWorker, this)}; + std::make_unique(ThreadPoolWorker, this, i)}; workers_.push_back(std::move(worker)); } } + size_t NumberOfThreads() { + return workers_.size(); + } + // start a task in the threadpool void Enqueue(mjTask* task) { if (mjUNLIKELY(GetAtomicTaskStatus(task).exchange(mjTASK_QUEUED) != @@ -89,11 +102,39 @@ class ThreadPoolImpl : public mjThreadPool { } } + // registers a worker ID for a given thread + void RegisterWorker(const size_t input_worker_id) { + worker_id = input_worker_id; + } + + // gets the worker id of the current thread + size_t GetWorkerId() { + return worker_id; + } + + void LockAlloc() { + alloc_mutex_.lock(); + } + + void UnlockAlloc() { + alloc_mutex_.unlock(); + } + + bool IsThreadPoolBound() { + return thread_pool_bound_; + } + + void BindThreadPool() { + thread_pool_bound_ = true; + } + ~ThreadPoolImpl() { Shutdown(); } private: // method executed by running threads - static void ThreadPoolWorker(ThreadPoolImpl* thread_pool) { + static void ThreadPoolWorker( + ThreadPoolImpl* thread_pool, const size_t thread_index) { + worker_id = thread_index + 1; while (!thread_pool->shutdown_) { auto task = static_cast(thread_pool->lockless_queue_.pop()); task->args = task->func(task->args); @@ -109,11 +150,128 @@ class ThreadPoolImpl : public mjThreadPool { // queue of tasks to execute mujoco::LocklessQueue lockless_queue_; + + // Mutex to protect arena allocations. + std::mutex alloc_mutex_; + + // Whether or not a ThreadPool was bound using mju_bindThreadPool. + bool thread_pool_bound_ = false; }; // create a thread pool mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) { - return new ThreadPoolImpl(number_of_threads); + return reinterpret_cast(new ThreadPoolImpl(number_of_threads)); +} + +// gets the number of shards the stack is currently broken into +static size_t GetNumberOfShards(mjData* d) { + if (!d->threadpool) { + return 1; + } + return mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool) + 1; +} + +// returns the stack information for the specified thread's shard +mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id) { + auto thread_pool = (ThreadPoolImpl*)d->threadpool; + if (!thread_pool || !thread_pool->IsThreadPoolBound()) { + mju_error("Thread Pool not bound, use mju_bindThreadPool to add an mjThreadPool to mjData"); + } + + // number of threads running in the threadpool plus the main thread + size_t number_of_shards = GetNumberOfShards(d); + + // size of entire arena/stack in bytes + size_t total_arena_size_bytes = d->narena; + + // set the shard cursor to the end of the arena + uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + total_arena_size_bytes; + + // each thread including the main one will get an equal shard of the stack + size_t bytes_per_shard = total_arena_size_bytes / (2 * (number_of_shards)); + + // ensure the shard is larger than the cache line + size_t misalignment = bytes_per_shard % mju_getDestructiveInterferenceSize(); + + if (misalignment != 0) { + bytes_per_shard += mju_getDestructiveInterferenceSize() - misalignment; + } + + if (bytes_per_shard * number_of_shards > total_arena_size_bytes) { + mju_error("Arena is not large enough for %zu shards", number_of_shards); + } + + uintptr_t result = (end_of_arena_ptr - (thread_id + 1) * bytes_per_shard); + + // align the end of the shard to be mjStackInfo. + misalignment = result % alignof(mjStackInfo); + result -= misalignment; + + return (mjStackInfo*) result; +} + +// shards the stack for each thread +static void ConfigureMultiThreadedStack(mjData* d) { + if (!d->threadpool) { + mju_error("No thread pool specified for multithreaded operation"); + } + + size_t number_of_shards = GetNumberOfShards(d); + + // current top of the stack + uintptr_t current_limit = (uintptr_t)d->arena + d->narena - d->pstack; + + // set the shard cursor to the end of the arena + uintptr_t begin_shard_cursor_ptr = (uintptr_t)d->arena + d->narena; + + for (size_t shard_index = 0; shard_index < number_of_shards; ++shard_index) { + mjStackInfo* end_shard_cursor_ptr = mju_getStackInfoForThread(d, shard_index); +#ifdef ADDRESS_SANITIZER + // unpoison stack info + ASAN_UNPOISON_MEMORY_REGION((void*)end_shard_cursor_ptr, sizeof(mjStackInfo)); +#endif + // handle the main thread's stack which may already have data in it + if (shard_index == 0) { + // abort if the current stack is already larger than the portion of the stack + // that would be reserved for the main thread + if ((uintptr_t)end_shard_cursor_ptr > current_limit) { + mju_error("mj_bindThreadPool: sharding stack - existing stack larger than shard size: current_size = %zu, " + "max_size = %zu", current_limit, (uintptr_t) end_shard_cursor_ptr); + } + end_shard_cursor_ptr->top = current_limit; + end_shard_cursor_ptr->stack_base = d->pbase; + } else { + // all other stacks are empty because threads have not been used yet + end_shard_cursor_ptr->top = begin_shard_cursor_ptr; + end_shard_cursor_ptr->stack_base = 0; + } + + end_shard_cursor_ptr->bottom = begin_shard_cursor_ptr; + end_shard_cursor_ptr->limit = (uintptr_t)end_shard_cursor_ptr + sizeof(mjStackInfo); + begin_shard_cursor_ptr = (uintptr_t)end_shard_cursor_ptr - 1; + } +} + +// adds a thread pool to mjData and configures it for multi-threaded use. +void mju_bindThreadPool(mjData* d, mjThreadPool* thread_pool) { + if (d->threadpool) { + mju_error("Thread Pool already bound to mjData"); + } + + d->threadpool = (uintptr_t) thread_pool; + ((ThreadPoolImpl*)thread_pool)->BindThreadPool(); + ConfigureMultiThreadedStack(d); +} + +// gets the number of running threads in the thread pool. +size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool) { + auto thread_pool_impl = static_cast(thread_pool); + return thread_pool_impl->NumberOfThreads(); +} + +size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) { + auto thread_pool_impl = static_cast(thread_pool); + return thread_pool_impl->GetWorkerId(); } // start a task in the threadpool @@ -128,4 +286,23 @@ void mju_threadPoolDestroy(mjThreadPool* thread_pool) { thread_pool_impl->Shutdown(); delete thread_pool_impl; } + +// locks the allocation mutex to protect Stack and Arena allocations +void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool) { + auto thread_pool_impl = static_cast(thread_pool); + thread_pool_impl->LockAlloc(); +} + +// unlocks the allocation mutex to protect Stack and Arena allocations +void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool) { + auto thread_pool_impl = static_cast(thread_pool); + thread_pool_impl->UnlockAlloc(); +} + +// Get the destructive interference size for the architecture. +size_t mju_getDestructiveInterferenceSize(void) { + // return std::hardware_destructive_interference_size; + return 128; +} + } // namespace mujoco diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h index 1d539437..fb6920de 100644 --- a/src/thread/thread_pool.h +++ b/src/thread/thread_pool.h @@ -19,21 +19,64 @@ #include #include +#include #ifdef __cplusplus namespace mujoco { extern "C" { #endif +// MultiThreaded Stack will be an approximately 50/50 split of the entire buffer, with a little +// wiggle for alignment and caching concerns. The basic layout is to reuse the existing single +// threaded markers, and then create shards for each thread to use as its stack. +// Not to scale. +// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------| +// |Used Arena|Free Arena |Shard1 |Shard1 |Shard1 |Shard0 |Shard0 |Shard0 | +// |%%%%%%%%%%| |StackInfo |Free Stack |Used Stack |StackInfo |Free Stack |Used Stack | +// |%%%%%%%%%%| | | |%%%%%%%%%%%| | |%%%%%%%%%%%| +// |%%%%%%%%%%| | | |%%%%%%%%%%%| | |%%%%%%%%%%%| +// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------| +// d->arena d->parena d->pstack shard1->stack_info shard1->bottom_of_stack shard1->bottom_of_stack +// shard1->stack_info shard0->stack_info shard0->current_stack +// shard1->top_of_stack shard1->top_of_stack +// shard1->current_stack +typedef struct { + uintptr_t bottom; // First memory address available to the stack + uintptr_t top; // Current memory address used by the stack + uintptr_t limit; // Top limit of the stack (note this is smaller than bottom, stack grows down) + uintptr_t stack_base; // Current stack base for mark and free stack +} mjStackInfo; + // Create a thread pool with the specified number of threads running. MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads); +// Returns the stack information for the specified thread's shard. +mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id); + +// Adds a thread pool to mjData and configures it for multi-threaded use. +MJAPI void mju_bindThreadPool(mjData* d, mjThreadPool* thread_pool); + +// Gets the number of running threads in the thread pool. +MJAPI size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool); + +// Gets the ID of the current thread being executed +MJAPI size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool); + // Enqueue a task in a thread pool. MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task); +// Locks the allocation mutex to protect Arena allocations. +MJAPI void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool); + +// Unlocks the allocation mutex to protect Arena allocations. +MJAPI void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool); + // Destroy a thread pool. MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool); +// Get the destructive interference size for the architecture. +MJAPI size_t mju_getDestructiveInterferenceSize(void); + #ifdef __cplusplus } // extern "C" } // namespace mujoco diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index 4c6443a4..668ffda0 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -31,6 +31,7 @@ #include #include #include "src/engine/engine_util_errmem.h" +#include "src/thread/thread_pool.h" #include "test/fixture.h" namespace mujoco { @@ -764,6 +765,82 @@ TEST_F(EngineIoTest, CanMarkAndFreeStack) { mj_deleteModel(model); } +struct TestFunctionArgs_ { + mjData* d; + int input; + int stack_output; + int arena_output; + size_t output_thread_worker; +}; +typedef TestFunctionArgs_ TestFunctionArgs; + +void* TestFunction(void* args) { + TestFunctionArgs* test_args = static_cast(args); + test_args->output_thread_worker = + mju_threadPoolCurrentWorkerId((mjThreadPool*)test_args->d->threadpool); + mj_markStack(test_args->d); + int* test_ints = mj_stackAllocInt(test_args->d, 10); + test_ints[0] = test_args->input; + test_args->stack_output = test_ints[0]; + + int* test_arena_ints = + (int*)mj_arenaAllocByte(test_args->d, sizeof(int) * 10, _Alignof(int)); + test_arena_ints[0] = test_args->input; + test_args->arena_output = test_arena_ints[0]; + + + mj_freeStack(test_args->d); + return nullptr; +} + +TEST_F(EngineIoTest, TestStackShardingForThreads) { + constexpr char xml[] = R"( + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data(); + + mjData* data = mj_makeData(model); + ASSERT_THAT(data, NotNull()); + mjThreadPool* thread_pool = mju_threadPoolCreate(10); + mju_bindThreadPool(data, thread_pool); + + constexpr int kTasks = 1000; + TestFunctionArgs test_function_args[kTasks]; + mjTask tasks[kTasks]; + for (int i = 0; i < kTasks; ++i) { + test_function_args[i].d = data; + test_function_args[i].input = i; + mju_defaultTask(&tasks[i]); + tasks[i].func = TestFunction; + tasks[i].args = &test_function_args[i]; + mju_threadPoolEnqueue(thread_pool, &tasks[i]); + } + + mj_markStack(data); + int* test_ints = mj_stackAllocInt(data, 10); + test_ints[0] = 1; + mj_freeStack(data); + + for (int i = 0; i < kTasks; ++i) { + mju_taskJoin(&tasks[i]); + } + + for (int i = 0; i < kTasks; ++i) { + EXPECT_EQ(test_function_args[i].input, test_function_args[i].stack_output); + EXPECT_EQ(test_function_args[i].input, test_function_args[i].arena_output); + } + + mj_deleteData(data); + mj_deleteModel(model); + mju_threadPoolDestroy(thread_pool); +} + #ifdef ADDRESS_SANITIZER void MarkFreeStack(mjData* d, bool free) { mj_markStack(d);