Remove thread_local EPA data in favor of using mjData stack.

PiperOrigin-RevId: 900201369
Change-Id: I82fbe9bf0ef9ea117c2124353d876dd42a57fb3a

src/engine/engine_collision_convex.c

@@ -34,17 +34,6 @@
 
 #define mjMINVAL2 (mjMINVAL * mjMINVAL)
 
-// allocate callback for EPA in nativeccd
-static void* ccd_allocate(void* data, size_t nbytes) {
-  mj_markStack((mjData*)data);
-  return mj_stackAllocByte((mjData*)data, nbytes, sizeof(mjtNum));
-}
-
-// free callback for EPA in nativeccd
-static void ccd_free(void* data, void* buffer) {
-  mj_freeStack((mjData*)data);
-}
-
 // ccd prism first dir
 static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
   ccdVec3Set(vec, 0, 0, 1);
@@ -98,15 +87,15 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
   mjtNum dist;
 
   // set config
+  mj_markStack(d);
   config.max_iterations = m->opt.ccd_iterations;
   config.tolerance = m->opt.ccd_tolerance;
   config.max_contacts = ncon;
   config.dist_cutoff = 0;  // no geom distances needed
-  config.context = (void*)d;
+  config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
-  config.alloc = ccd_allocate;
-  config.free = ccd_free;
 
   if ((dist = mjc_ccd(&config, &status, obj1, obj2)) < 0) {
+    mj_freeStack(d);
     int nwitness = status.nx;
     for (int i = 0; i < nwitness; i++, con++) {
       con->dist = margin + dist;
@@ -119,6 +108,7 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
     }
     return nwitness;
   }
+  mj_freeStack(d);
   return 0;
 }
 

src/engine/engine_collision_driver.c

@@ -516,6 +516,7 @@ void mj_collision(const mjModel* m, mjData* d) {
       }
     }
   }
+  mj_freeStack(d);
 
   // finish merging predefined geom pairs
   for (; pairadr < npair; pairadr++) {
@@ -572,8 +573,6 @@ void mj_collision(const mjModel* m, mjData* d) {
 
   // end narrowphase and midphase timer
   TM_END(mjTIMER_COL_NARROW);
-
-  mj_freeStack(d);
   TM_END1(mjTIMER_POS_COLLISION);
 }
 

src/engine/engine_collision_gjk.c

@@ -22,7 +22,6 @@
 #include <mujoco/mjtnum.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_convex.h"
-#include "engine/engine_macro.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 
@@ -2208,16 +2207,17 @@ static inline void inflate(mjCCDStatus* status, mjtNum margin1, mjtNum margin2)
 }
 
 
+// return size in bytes of the buffer needed for mjc_ccd for a given number of iterations
+size_t mjc_ccdSize(int iterations) {
+  return (sizeof(Face) * 6 * iterations)           // faces in polytope
+         + (sizeof(Face*) * 6 * iterations)        // map in polytope
+         + (sizeof(Vertex) * (5 + iterations))     // vertices in polytope
+         + 2 * (24 * sizeof(int));                 // horizon data
+}
+
+
 // general convex collision detection
 mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2) {
-  // pre-allocate static memory for low iterations
-  void* buffer = NULL;
-  static mjTHREADLOCAL Vertex vert_data[5 + mjMAX_EPA_ITERATIONS];
-  static mjTHREADLOCAL Face face_data[6 * mjMAX_EPA_ITERATIONS];
-  static mjTHREADLOCAL Face* map_data[6 * mjMAX_EPA_ITERATIONS];
-  static mjTHREADLOCAL int index_data[6 + mjMAX_EPA_ITERATIONS];
-  static mjTHREADLOCAL int edge_data[6 + mjMAX_EPA_ITERATIONS];
-
   // setup
   obj1->center(status->x1, obj1);
   obj2->center(status->x2, obj2);
@@ -2295,42 +2295,24 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
     return status->dist;
   }
 
-  if (status->dist <= config->tolerance && status->nsimplex > 1) {
+  if (status->dist <= config->tolerance && status->nsimplex > 1 && config->buffer) {
     status->dist = 0;  // assume touching
     Polytope pt;
     pt.nfaces = pt.nmap = pt.nverts = pt.horizon.nedges = 0;
 
-    // allocate memory via static thread-local storage
+    // allocate memory for polytope
     int N = config->max_iterations;
-    if (N <= mjMAX_EPA_ITERATIONS) {
+    pt.maxfaces = 6 * N;
-      pt.maxfaces = 6 * mjMAX_EPA_ITERATIONS;
+    uint8_t* buffer = config->buffer;
-      pt.verts = vert_data;
+    pt.verts = (Vertex*)buffer;
-      pt.faces = face_data;
+    buffer += sizeof(Vertex) * (5 + N);
-      pt.map = map_data;
+    pt.faces = (Face*)buffer;
-      pt.horizon.indices = index_data;
+    buffer += sizeof(Face) * (6 * N);
-      pt.horizon.edges = edge_data;
+    pt.map = (Face**)buffer;
-    }
+    buffer += sizeof(Face*) * (6 * N);
-
+    pt.horizon.indices = (int*)buffer;
-    // static storage insufficient, allocate with callback
+    buffer += sizeof(int) * 24;
-    else {
+    pt.horizon.edges = (int*)buffer;
-      size_t nbytes = (sizeof(Face) * 6 * N)      // faces in polytope
-                    + (sizeof(Face*) * 6 * N)     // map in polytope
-                    + (sizeof(Vertex) * (5 + N))  // vertices in polytope
-                    + 2*(sizeof(int) * (6 + N));  // horizon data
-
-      pt.maxfaces = 6 * N;
-      buffer = config->alloc(config->context, nbytes);
-      uint8_t* bbuffer = (uint8_t*)buffer;
-      pt.verts = (Vertex*)bbuffer;
-      bbuffer += sizeof(Vertex) * (5 + N);
-      pt.faces = (Face*)bbuffer;
-      bbuffer += sizeof(Face) * (6 * N);
-      pt.map = (Face**)bbuffer;
-      bbuffer += sizeof(Face*) * (6 * N);
-      pt.horizon.indices = (int*)bbuffer;
-      bbuffer += sizeof(int) * (6 + N);
-      pt.horizon.edges = (int*)bbuffer;
-    }
 
     int ret;
     if (status->nsimplex == 2) {
@@ -2350,9 +2332,6 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
       }
     }
   }
-  if (buffer) {
-    config->free(config->context, buffer);
-  }
   return status->dist;
 }
 

src/engine/engine_collision_gjk.h

@@ -35,9 +35,6 @@ extern "C" {
   #define mjMAX_LIMIT FLT_MAX
 #endif
 
-// max number of EPA iterations
-#define mjMAX_EPA_ITERATIONS 170
-
 // tolerance for normal alignment of two faces (cosine of 1.6e-3)
 #define mjFACE_TOL 0.99999872
 
@@ -77,13 +74,7 @@ typedef struct {
   mjtNum tolerance;    // tolerance used by GJK and EPA
   int max_contacts;    // set to max number of contact points to recover
   mjtNum dist_cutoff;  // set to max geom distance to recover
-  void* context;       // opaque data pointer passed to callbacks
+  void* buffer;        // buffer memory for polytope (should be sized given by mjc_ccdSize)
-
-  // callback to allocate memory for polytope (only needed for penetration recovery)
-  void*(*alloc)(void* context, size_t nbytes);
-
-  // callback to free memory from alloc callback
-  void(*free)(void* context, void* buffer);
 } mjCCDConfig;
 
 // data produced from running GJK and EPA
@@ -108,6 +99,9 @@ typedef struct {
   int nsimplex;
 } mjCCDStatus;
 
+// return size in bytes of the buffer needed for mjc_ccd for a given number of iterations
+MJAPI size_t mjc_ccdSize(int iterations);
+
 // run general convex collision detection, returns positive for distance, negative for penetration
 MJAPI mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, mjCCDObj* obj2);
 #ifdef __cplusplus

src/engine/engine_support.c

@@ -526,6 +526,7 @@ void mj_xfrcAccumulate(const mjModel* m, mjData* d, mjtNum* qfrc) {
 // returns the smallest distance between two geoms (using nativeccd)
 static mjtNum mj_geomDistanceCCD(const mjModel* m, mjData* d, int g1, int g2,
                                  mjtNum distmax, mjtNum fromto[6]) {
+  mj_markStack(d);
   mjCCDConfig config;
   mjCCDStatus status;
 
@@ -534,12 +535,14 @@ static mjtNum mj_geomDistanceCCD(const mjModel* m, mjData* d, int g1, int g2,
   config.tolerance = m->opt.ccd_tolerance;
   config.max_contacts = 1;        // want contacts
   config.dist_cutoff = distmax;   // want geom distances
+  config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
 
   mjCCDObj obj1, obj2;
   mjc_initCCDObj(&obj1, m, d, g1, 0);
   mjc_initCCDObj(&obj2, m, d, g2, 0);
 
   mjtNum dist = mjc_ccd(&config, &status, &obj1, &obj2);
+  mj_freeStack(d);
 
   // witness points are only computed if dist <= distmax
   if (fromto && status.nx > 0) {

test/engine/engine_collision_gjk_test.cc

@@ -61,14 +61,6 @@ constexpr char kEllipsoidXml[] = R"(
   </keyframe>
 </mujoco>)";
 
-void* CCDAllocate(void* data, std::size_t nbytes) {
-  return new std::byte[nbytes];
-}
-
-void CCDFree(void* data, void* buffer) {
-  delete [] (std::byte*)buffer;
-}
-
 mjtNum GeomDist(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3],
                 mjtNum x2[3], mjtNum cutoff = mjMAX_LIMIT) {
   mjCCDConfig config;
@@ -79,6 +71,7 @@ mjtNum GeomDist(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3],
   config.tolerance = kTolerance,
   config.max_contacts = 0;   // no geom contacts needed
   config.dist_cutoff = cutoff;
+  config.buffer = nullptr;
 
   mjCCDObj obj1, obj2;
   mjc_initCCDObj(&obj1, m, d, g1, 0);
@@ -129,14 +122,13 @@ int Penetration(mjCCDStatus& status, mjtNum& depth, std::vector<mjtNum>& dir,
   mjCCDConfig config;
 
   // set config
+  auto buffer = std::vector<std::byte>(mjc_ccdSize(kMaxIterations));
   config.max_iterations = kMaxIterations;
   config.tolerance = kTolerance;
   config.max_contacts = max_contacts;
   config.dist_cutoff = 0;  // no geom distances needed
   config.max_contacts = max_contacts;
-  config.context = nullptr;
+  config.buffer = buffer.data();
-  config.alloc = CCDAllocate;
-  config.free = CCDFree;
 
   mjtNum dist = mjc_ccd(&config, &status, &obj1, &obj2);
   if (dist < 0) {