Rename mj_stackAllocByte to mj_stackAlloc and make it public.

PiperOrigin-RevId: 561133040
Change-Id: I4f5ffe51c1287aa0c82af68bfe17ed6dd2ad09da

doc/APIreference/functions.rst

@@ -1232,6 +1232,15 @@ mj_resetDataKeyframe
 
 Reset data, set fields from specified keyframe.
 
+.. _mj_stackAlloc:
+
+mj_stackAlloc
+~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_stackAlloc
+
+Allocate a specific number of bytes on :ref:`mjData` stack. Call mju_error on stack overflow.
+
 .. _mj_stackAllocNum:
 
 mj_stackAllocNum

doc/changelog.rst

@@ -38,20 +38,22 @@ General
    Euler integrator. See the :ref:`Numerical Integration<geIntegration>` section for more details.
 #. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
    :ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
-#. Renamed the function ``mj_stackAlloc`` to ``mj_stackAllocNum``.
+#. Changed the function ``mj_stackAlloc`` to allocate an arbitrary number of bytes, rather than in multiples of
+   ``sizeof(mjtNum)``, and add an additional argument for specifying the alignment of the returned pointer. The existing
+   functionality of allocating ``mjtNum`` arrays is still available through the new function ``mj_stackAllocNum``.
 #. Renamed the ``nstack`` field in ``mjModel`` and ``mjData`` to ``narena``. Changed ``narena``, ``pstack``, and
    ``maxuse_stack`` to count number of bytes rather than number of ``mjtNum``s.
 
 Python bindings
 ^^^^^^^^^^^^^^^
 
-9. Fixed `#870 <https://github.com/deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
-   camera name used the default camera.
+10. Fixed `#870 <https://github.com/deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
+    camera name used the default camera.
 
 Bug fixes
 ^^^^^^^^^
 
-10. Fixed a bug that was causing the geom margins to be ignored during the midphase.
+11. Fixed a bug that was causing the geom margins to be ignored during the midphase.
 
 
 Version 2.3.7 (July 20, 2023)

doc/includes/references.h

@@ -2213,6 +2213,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
 void mj_resetData(const mjModel* m, mjData* d);
 void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
 void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
+void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment);
 mjtNum* mj_stackAllocNum(mjData* d, int size);
 int* mj_stackAllocInt(mjData* d, int size);
 void mj_deleteData(mjData* d);

doc/programming/simulation.rst

@@ -722,9 +722,8 @@ the scope. If not, saving and restoring the stack pointer should be done manuall
 The function :ref:`mj_stackAllocNum` checks if there is enough space, and if so it advances the stack pointer,
 otherwise it triggers an error. It also keeps track of the maximum stack allocation;
 see :ref:`diagnostics <siDiagnostics>` below. Note that :ref:`mj_stackAllocNum` is only used for allocating
-``mjtNum`` arrays, the most common type of array. :ref:`mj_stackAllocInt` is provided for integer array allocation.
-Allocators for other types are also possible, as in
-`engine_collision_driver.c <https://github.com/deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__.
+``mjtNum`` arrays, the most common type of array. :ref:`mj_stackAllocInt` is provided for integer array allocation,
+and :ref:`mj_stackAlloc` is provided for allocation of arbitrary number of bytes and alignment.
 
 .. _siError:
 

include/mujoco/mujoco.h

@@ -188,6 +188,9 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 // Reset data, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
+// Allocate a specific number of bytes on mjData stack. Call mju_error on stack overflow.
+MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment);
+
 // Allocate array of mjtNums on mjData stack. Call mju_error on stack overflow.
 MJAPI mjtNum* mj_stackAllocNum(mjData* d, int size);
 

introspect/functions.py

@@ -677,6 +677,30 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Reset data, set fields from specified keyframe.',
      )),
+    ('mj_stackAlloc',
+     FunctionDecl(
+         name='mj_stackAlloc',
+         return_type=PointerType(
+             inner_type=ValueType(name='void'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='bytes',
+                 type=ValueType(name='size_t'),
+             ),
+             FunctionParameterDecl(
+                 name='alignment',
+                 type=ValueType(name='size_t'),
+             ),
+         ),
+         doc='Allocate a specific number of bytes on mjData stack. Call mju_error on stack overflow.',  # pylint: disable=line-too-long
+     )),
     ('mj_stackAllocNum',
      FunctionDecl(
          name='mj_stackAllocNum',

python/mujoco/bindings_test.py

@@ -948,7 +948,7 @@ Euler integrator, semi-implicit in velocity.
   def test_can_raise_error(self):
     self.data.pstack = self.data.narena
     with self.assertRaisesRegex(mujoco.FatalError,
-                                r'\Amj_stackAllocByte: stack overflow'):
+                                r'\AmjData stack overflow'):
       mujoco.mj_forward(self.model, self.data)
 
   def test_mjcb_time(self):

src/engine/engine_collision_driver.c

@@ -208,7 +208,8 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 }
 
 static mjCollisionTree* mj_stackAllocTree(mjData* d, int max_stack) {
-  return (mjCollisionTree*)mj_stackAllocByte(d, max_stack * sizeof(mjCollisionTree*));
+  return (mjCollisionTree*) mj_stackAlloc(
+      d, max_stack * sizeof(mjCollisionTree), _Alignof(mjCollisionTree));
 }
 
 // binary search between two body trees
@@ -751,8 +752,10 @@ int mj_broadphase(const mjModel* m, mjData* d, int* pair, int maxpair) {
   }
 
   // allocate sort buffer
-  sortbuf = (mjtBroadphase*)mj_stackAllocByte(d, 2 * bufcnt * sizeof(mjtBroadphase));
-  activebuf = (mjtBroadphase*)mj_stackAllocByte(d, 2 *bufcnt * sizeof(mjtBroadphase));
+  sortbuf = (mjtBroadphase*) mj_stackAlloc(
+      d, 2 * bufcnt * sizeof(mjtBroadphase), _Alignof(mjtBroadphase));
+  activebuf = (mjtBroadphase*) mj_stackAlloc(
+      d, 2 * bufcnt * sizeof(mjtBroadphase), _Alignof(mjtBroadphase));
 
   // init sortbuf with axis0
   int k = 0;

src/engine/engine_collision_driver.h

@@ -21,19 +21,10 @@
 
 #ifdef __cplusplus
 extern "C" {
-#elif !defined(__STDC_VERSION__) || __STDC_VERSION__ < 201112L
-// No C11 support in Visual Studio 2019 update 7 and earlier.
-// However, MSVC allows C++ alignas to be used in C code, so
-// we can just skip the #include <stdalign.h>.
-#ifndef _MSC_VER
-#error "Compiler does not support C11."
-#endif
-#else
-#include <stdalign.h>
 #endif
 
 struct mjCollisionTree_ {
-  alignas(mjtNum) int node1;
+  int node1;
   int node2;
 };
 typedef struct mjCollisionTree_ mjCollisionTree;

src/engine/engine_collision_sdf.c

@@ -456,8 +456,8 @@ static void collideBVH(const mjModel* m, mjData* d, int g,
     int node;
   };
   typedef struct CollideTreeArgs_ CollideTreeArgs;
-  CollideTreeArgs* stack = (CollideTreeArgs*) mj_stackAllocByte(
-      d, max_stack * sizeof(CollideTreeArgs));
+  CollideTreeArgs* stack = (CollideTreeArgs*) mj_stackAlloc(
+      d, max_stack * sizeof(CollideTreeArgs), _Alignof(CollideTreeArgs));
 
   int nstack = 0;
   stack[nstack].node = 0;

src/engine/engine_crossplatform.h

@@ -64,7 +64,9 @@
 #endif
 
 #if defined(_MSC_VER) && !defined(__clang__) && !defined(__cplusplus)
-  typedef long double max_align_t;
+  typedef long double mjtMaxAlign;
+#else
+  typedef max_align_t mjtMaxAlign;
 #endif
 
 #endif  // MUJOCO_SRC_ENGINE_ENGINE_CROSSPLATFORM_H_

src/engine/engine_io.c

@@ -27,7 +27,6 @@
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_array_safety.h"  // IWYU pragma: keep
-#include "engine/engine_crossplatform.h"  // IWYU pragma: keep
 #include "engine/engine_resource.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_plugin.h"
@@ -1200,9 +1199,9 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
 
 
 // allocate memory from the mjData arena
-void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
-  int misalignment = d->parena % alignment;
-  int padding = misalignment ? alignment - misalignment : 0;
+void* mj_arenaAlloc(mjData* d, size_t bytes, size_t alignment) {
+  size_t misalignment = d->parena % alignment;
+  size_t padding = misalignment ? alignment - misalignment : 0;
 
   // check size
   size_t bytes_available = d->narena - d->pstack;
@@ -1228,8 +1227,9 @@ void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
 
 
 
-// allocate size bytes on the mjData stack
-void* mj_stackAllocByte(mjData* d, size_t size) {
+// internal: allocate size bytes on the mjData stack
+// 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) {
   // return NULL if empty
   if (!size) {
     return NULL;
@@ -1254,8 +1254,8 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
   // start of the memory to be allocated to the buffer
   uintptr_t start_ptr = end_ptr - (size + mjREDZONE);
 
-  // move start_ptr back to align to max_align_t
-  start_ptr -= start_ptr % _Alignof(max_align_t);  // NOLINT
+  // align the pointer
+  start_ptr -= start_ptr % alignment;
 
   // new top of the stack
   uintptr_t new_pstack_ptr = start_ptr - mjREDZONE;
@@ -1269,7 +1269,7 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
   size_t stack_available_bytes = end_ptr - ((uintptr_t)d->arena + d->parena);
   size_t stack_required_bytes = end_ptr - new_pstack_ptr;
   if (stack_required_bytes > stack_available_bytes) {
-    mjERROR("stack overflow: max = %zu, available = %zu, requested = %zu "
+    mju_error("mjData stack overflow: max = %zu, available = %zu, requested = %zu "
               "(ne = %d, nf = %d, nefc = %d, ncon = %d)",
               stack_size_bytes, stack_available_bytes, stack_required_bytes,
               d->ne, d->nf, d->nefc, d->ncon);
@@ -1303,12 +1303,16 @@ void* mj_stackAllocByte(mjData* d, size_t size) {
   return (void*)start_ptr;
 }
 
+void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment) {
+  return stackalloc(d, bytes, alignment);
+}
+
 mjtNum* mj_stackAllocNum(mjData* d, int size) {
-  return (mjtNum*)mj_stackAllocByte(d, size * sizeof(mjtNum));
+  return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum));
 }
 
 int* mj_stackAllocInt(mjData* d, int size) {
-  return (int*)mj_stackAllocByte(d, size * sizeof(int));
+  return (int*) stackalloc(d, size * sizeof(int), _Alignof(int));
 }
 
 

src/engine/engine_io.h

@@ -20,7 +20,10 @@
 #include <mujoco/mjmodel.h>
 
 #ifdef __cplusplus
+#include <cstddef>
 extern "C" {
+#else
+#include <stddef.h>
 #endif
 
 // internal hash map size factor (2 corresponds to a load factor of 0.5)
@@ -100,7 +103,10 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
 // mjData arena allocate
-MJAPI void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
+MJAPI void* mj_arenaAlloc(mjData* d, size_t bytes, size_t alignment);
+
+// mjData stack allocate
+MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment);
 
 // mjData stack allocate for array of mjtNums
 MJAPI mjtNum* mj_stackAllocNum(mjData* d, int size);
@@ -108,9 +114,6 @@ MJAPI mjtNum* mj_stackAllocNum(mjData* d, int size);
 // mjData stack allocate for array of ints
 MJAPI int* mj_stackAllocInt(mjData* d, int size);
 
-// mjData stack allocate for a specific size of bytes
-MJAPI void* mj_stackAllocByte(mjData* d, size_t size);
-
 // de-allocate data
 MJAPI void mj_deleteData(mjData* d);
 

src/engine/engine_util_container.c

@@ -13,16 +13,19 @@
 // limitations under the License.
 
 #include "engine/engine_util_container.h"
+
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
 #include <mujoco/mujoco.h>
+#include "engine/engine_crossplatform.h"
 #include "engine/engine_io.h"
 
 // stack allocate and initialize new mjArrayList
 mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity) {
-  mjArrayList* array_list = (mjArrayList*) mj_stackAllocByte(d, sizeof(mjArrayList));
+  mjArrayList* array_list = (mjArrayList*) mj_stackAlloc(
+      d, sizeof(mjArrayList), _Alignof(mjArrayList));
   initial_capacity = mjMAX(1, initial_capacity);
   array_list->d = d;
   array_list->element_size = element_size;
@@ -31,7 +34,8 @@ mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_
   array_list->next_segment = NULL;
 
   // allocate array list buffer
-  array_list->buffer = (void*) mj_stackAllocByte(d, element_size * initial_capacity);
+  array_list->buffer = (void*) mj_stackAlloc(
+      d, element_size * initial_capacity, _Alignof(mjtMaxAlign));
   return array_list;
 }
 

test/engine/engine_util_container_test.cc

@@ -33,12 +33,11 @@ constexpr int GetExpectedStackUsageBytes() {
     return 0;
   } else {
     constexpr auto RoundUpToAlignment =
-        [](int x) {
-          constexpr auto kAlignment = alignof(std::max_align_t);
-          return kAlignment * (x / kAlignment + ((x % kAlignment) ? 1 : 0));
+        [](int x, int alignment) {
+          return alignment * (x / alignment + ((x % alignment) ? 1 : 0));
         };
-    return RoundUpToAlignment(sizeof(mjArrayList)) +
-           RoundUpToAlignment(Capacity * sizeof(T)) +
+    return RoundUpToAlignment(sizeof(mjArrayList), alignof(mjArrayList)) +
+           RoundUpToAlignment(Capacity * sizeof(T), alignof(std::max_align_t)) +
            GetExpectedStackUsageBytes<T, N - Capacity, 2 * Capacity>();
   }
 }

unity/Runtime/Bindings/MjBindings.cs

@@ -3087,6 +3087,9 @@ public static unsafe extern void mj_resetDataDebug(mjModel_* m, mjData_* d, byte
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_resetDataKeyframe(mjModel_* m, mjData_* d, int key);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void* mj_stackAlloc(mjData_* d, UIntPtr bytes, UIntPtr alignment);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern double* mj_stackAllocNum(mjData_* d, int size);