Modify memory allocation in MuJoCo to be thread safe:

1) protect mj_arenaAllocBytes with mutexes 2) create shards for each thread in the stack and update mj_stackAllocBytes to allocate memory within each shard for a given thread PiperOrigin-RevId: 568315726 Change-Id: I0dee6694f2a5200fa4df22ade0e68dfaebf637fc
2023-09-25 13:52:33 -07:00
parent dff0bc2683
commit ff4158efff
4 changed files with 426 additions and 55 deletions
@@ -14,15 +14,19 @@

 #include "thread/thread_pool.h"

+#include <stdint.h>
+
 #include <algorithm>
 #include <atomic>
 #include <cstddef>
 #include <memory>
+#include <mutex>
 #include <thread>
 #include <utility>
 #include <vector>

 #include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>
 #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<std::thread>(ThreadPoolWorker, this)};
+        std::make_unique<std::thread>(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<mjTask*>(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<void*, kThreadPoolQueueSize> 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<mjThreadPool*>(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<ThreadPoolImpl*>(thread_pool);
+  return thread_pool_impl->NumberOfThreads();
+}
+
+size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) {
+  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(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<ThreadPoolImpl*>(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<ThreadPoolImpl*>(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
@@ -19,21 +19,64 @@

 #include <mujoco/mjexport.h>
 #include <mujoco/mjthread.h>
+#include <mujoco/mujoco.h>

 #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