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
@@ -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 <sanitizer/asan_interface.h>
@@ -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);
 }