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

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