Implement mj_stackAlloc and mj_arenaAlloc alignment using bitwise-and.

This implementation explicitly avoids calling division instructions (which are generally expensive) even when link-time optimization is not active or if the alloc functions aren't inlined for some reason.

Also make mj_stackAlloc cheaper for external callers who may be linking against MuJoCo dynamically (and therefore cannot benefit from inlining).

PiperOrigin-RevId: 561299565
Change-Id: Iea4be5bf6d351ddf0c9d1fd8bd89a0c207b39309

doc/APIreference/functions.rst

@@ -1239,7 +1239,8 @@ mj_stackAlloc
 
 .. mujoco-include:: mj_stackAlloc
 
-Allocate a specific number of bytes on :ref:`mjData` stack. Call mju_error on stack overflow.
+Allocate a number of bytes on :ref:`mjData` stack at a specific alignment which must be a power of 2.
+Call mju_error on stack overflow.
 
 .. _mj_stackAllocNum:
 

include/mujoco/mujoco.h

@@ -188,7 +188,8 @@ 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.
+// Allocate a number of bytes on mjData stack at a specific alignment which must be a power of 2.
+// 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.

introspect/functions.py

@@ -699,7 +699,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  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
+         doc='Allocate a number of bytes on mjData stack at a specific alignment which must be a power of 2. Call mju_error on stack overflow.',  # pylint: disable=line-too-long
      )),
     ('mj_stackAllocNum',
      FunctionDecl(

src/engine/engine_io.c

@@ -45,6 +45,11 @@
 
 static const int MAX_ARRAY_SIZE = INT_MAX / 4;
 
+// compute a % b assuming that the second argument is a power of 2
+static inline size_t modpow2(size_t a, size_t b) {
+  return a & (b - 1);
+}
+
 //------------------------------ mjLROpt -----------------------------------------------------------
 
 // set default options for length range computation
@@ -1200,7 +1205,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
 
 // allocate memory from the mjData arena
 void* mj_arenaAlloc(mjData* d, size_t bytes, size_t alignment) {
-  size_t misalignment = d->parena % alignment;
+  size_t misalignment = modpow2(d->parena, alignment);
   size_t padding = misalignment ? alignment - misalignment : 0;
 
   // check size
@@ -1255,7 +1260,7 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
   uintptr_t start_ptr = end_ptr - (size + mjREDZONE);
 
   // align the pointer
-  start_ptr -= start_ptr % alignment;
+  start_ptr -= modpow2(start_ptr, alignment);
 
   // new top of the stack
   uintptr_t new_pstack_ptr = start_ptr - mjREDZONE;

src/engine/engine_io.h

@@ -102,10 +102,10 @@ 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);
 
-// mjData arena allocate
+// mjData arena allocate (alignment must be a power of 2)
 MJAPI void* mj_arenaAlloc(mjData* d, size_t bytes, size_t alignment);
 
-// mjData stack allocate
+// mjData stack allocate (alignment must be a power of 2)
 MJAPI void* mj_stackAlloc(mjData* d, size_t bytes, size_t alignment);
 
 // mjData stack allocate for array of mjtNums