From 215433f5f7b8ffa0e11ddfde0a3f0f74836f4a71 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Wed, 30 Aug 2023 04:28:17 -0700
Subject: [PATCH] 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 | 3 ++-
 include/mujoco/mujoco.h        | 3 ++-
 introspect/functions.py        | 2 +-
 src/engine/engine_io.c         | 9 +++++++--
 src/engine/engine_io.h         | 4 ++--
 5 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index 963bc409..c80c8358 100644
--- a/doc/APIreference/functions.rst
+++ b/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:
 
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 706def25..19c37f67 100644
--- a/include/mujoco/mujoco.h
+++ b/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.
diff --git a/introspect/functions.py b/introspect/functions.py
index 87b1bde2..a5bbf038 100644
--- a/introspect/functions.py
+++ b/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(
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index d62044e7..5b13811f 100644
--- a/src/engine/engine_io.c
+++ b/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;
diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h
index 0916d769..137a00f4 100644
--- a/src/engine/engine_io.h
+++ b/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