Align all arrays in mjData's buffer to 64 byte boundaries.

The introduction of additional mjData arrays for the implicit integrator caused in a shift in data alignment in memory. In certain combinations of machines and models, this change resulted in ~5% performance regression that appears to be attributable to L2 cache misses.

This change aligns all arrays in mjData to the boundary of a typical L2 cache line size. Benchmarking shows that this resolves the performance regression on a Broadwell machine.

PiperOrigin-RevId: 451652783
Change-Id: I4ed2c2d03453a83a0b4a9b8e224e2fd7c0dadea9

doc/APIreference.rst

@@ -5501,7 +5501,7 @@ mju_malloc
 
    void* mju_malloc(size_t size);
 
-Allocate memory; byte-align on 8; pad size to multiple of 8.
+Allocate memory; byte-align on 64; pad size to multiple of 64.
 
 .. _mju_free:
 

include/mujoco/mujoco.h

@@ -732,7 +732,7 @@ MJAPI void mju_warning_s(const char* msg, const char* text);
 // Clear user error and memory handlers.
 MJAPI void mju_clearHandlers(void);
 
-// Allocate memory; byte-align on 8; pad size to multiple of 8.
+// Allocate memory; byte-align on 64; pad size to multiple of 64.
 MJAPI void* mju_malloc(size_t size);
 
 // Free memory, using free() by default.

introspect/functions.py

@@ -4659,7 +4659,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  type=ValueType(name='size_t'),
              ),
          ),
-         doc='Allocate memory; byte-align on 8; pad size to multiple of 8.',
+         doc='Allocate memory; byte-align on 64; pad size to multiple of 64.',
      )),
     ('mju_free',
      FunctionDecl(

src/engine/engine_forward.c

@@ -626,7 +626,9 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 
   // compute Xfinal
   d->time = time + h;
-  mju_copy(d->qpos, X[0], nq+nv+na);
+  mju_copy(d->qpos, X[0], nq);
+  mju_copy(d->qvel, X[0]+nq, nv);
+  mju_copy(d->act, X[0]+nq+nv, na);
   mj_integratePos(m, d->qpos, dX, h);
   mju_addToScl(d->qvel, dX+nv, h, nv);
   if (na) {

src/engine/engine_io.c

@@ -250,13 +250,9 @@ void mj_defaultStatistic(mjStatistic* stat) {
 static const int ID = 54321;
 
 
-// number of bytes to be skipped to achieve alignment
+// number of bytes to be skipped to achieve 64-byte alignment
-static unsigned int SKIP(intptr_t offset, unsigned int align) {
+static unsigned int SKIP(intptr_t offset) {
-  // replace structure with int size (mjContact starts with int)
+  const unsigned int align = 64;
-  if (align>sizeof(mjtNum)) {
-    align = sizeof(int);
-  }
-
   // compute skipped bytes
   return (align - (offset % align)) % align;
 }
@@ -338,10 +334,10 @@ static void mj_setPtrModel(mjModel* m) {
   MJMODEL_POINTERS_PREAMBLE(m);
 
   // assign pointers with padding
-#define X(type, name, nr, nc)                                           \
+#define X(type, name, nr, nc)                             \
-  m->name = (type*)(ptr + SKIP((intptr_t)ptr, sizeof(type)));           \
+  m->name = (type*)(ptr + SKIP((intptr_t)ptr));           \
-  ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(m->name, ptr));                \
+  ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(m->name, ptr));  \
-  ptr += SKIP((intptr_t)ptr, sizeof(type)) + sizeof(type)*(m->nr)*(nc);
+  ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc);
 
   MJMODEL_POINTERS
 #undef X
@@ -439,9 +435,9 @@ mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int njnt,
 
   // compute buffer size
   m->nbuffer = 0;
-#define X(type, name, nr, nc)                                            \
+#define X(type, name, nr, nc)                              \
-  m->nbuffer += SKIP(offset, sizeof(type)) + sizeof(type)*(m->nr)*(nc);  \
+  m->nbuffer += SKIP(offset) + sizeof(type)*(m->nr)*(nc);  \
-  offset += SKIP(offset, sizeof(type)) + sizeof(type)*(m->nr)*(nc);
+  offset += SKIP(offset) + sizeof(type)*(m->nr)*(nc);
 
   MJMODEL_POINTERS
 #undef X
@@ -763,10 +759,10 @@ static void mj_setPtrData(const mjModel* m, mjData* d) {
   MJDATA_POINTERS_PREAMBLE(m);
 
   // assign pointers with padding
-#define X(type, name, nr, nc)                                           \
+#define X(type, name, nr, nc)                             \
-  d->name = (type*)(ptr + SKIP((intptr_t)ptr, sizeof(type)));           \
+  d->name = (type*)(ptr + SKIP((intptr_t)ptr));           \
-  ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr));                \
+  ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr));  \
-  ptr += SKIP((intptr_t)ptr, sizeof(type)) + sizeof(type)*(m->nr)*(nc);
+  ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc);
 
   MJDATA_POINTERS
 #undef X
@@ -795,9 +791,9 @@ static mjData* _makeData(const mjModel* m) {
 
   // compute buffer size
   d->nbuffer = 0;
-#define X(type, name, nr, nc)                                            \
+#define X(type, name, nr, nc)                              \
-  d->nbuffer += SKIP(offset, sizeof(type)) + sizeof(type)*(m->nr)*(nc);  \
+  d->nbuffer += SKIP(offset) + sizeof(type)*(m->nr)*(nc);  \
-  offset += SKIP(offset, sizeof(type)) + sizeof(type)*(m->nr)*(nc);
+  offset += SKIP(offset) + sizeof(type)*(m->nr)*(nc);
 
   MJDATA_POINTERS
 #undef X

src/engine/engine_util_errmem.c

@@ -187,7 +187,7 @@ void mju_warning_s(const char* msg, const char* text) {
 
 //------------------------------ malloc and free ---------------------------------------------------
 
-// allocate memory; byte-align on 8; pad size to multiple of 8
+// allocate memory; byte-align on 64; pad size to multiple of 64
 void* mju_malloc(size_t size) {
   void* ptr = 0;
 
@@ -198,13 +198,13 @@ void* mju_malloc(size_t size) {
 
   // default allocator
   else {
-    // pad size to multiple of 8
+    // pad size to multiple of 64
-    if ((size%8)) {
+    if ((size%64)) {
-      size += 8 - (size%8);
+      size += 64 - (size%64);
     }
 
     // allocate
-    ptr = mju_alignedMalloc(size, 8);
+    ptr = mju_alignedMalloc(size, 64);
   }
 
   // error if null pointer