chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Dynamically allocate contact and efc_ arrays on a new memory arena.

Browse Source 
- Add private function `mj_arenaAlloc`. This is used internally to allocate memory from the arena.

- Add private function `mj_nefc` to count constraints. This function returns a tight upper bound on `d->nefc`. The number of counted constraints can be slightly bigger than exact `d->nefc` in the case of constraints with empty Jacobian, as when placing a frictional tendon between two world sites.

- Add new `memory` attribute to the `size` XML element for specification of arena memory size. This attribute is mutually exclusive with `nstack` and `njmax` specifications, which are now deprecated (but left around for the time being for legacy compatibility).

- Move `d->stack` to the end of the new arena space. The stack now grows in reverse from the end.

PiperOrigin-RevId: 479341539
Change-Id: Ie019c202e0908577ffc6f833a37920858116f667
This commit is contained in:
Saran Tunyasuvunakool
2022-10-06 10:02:35 -07:00
committed by Copybara-Service
parent 4d85a464cc
commit 58fd72f53d
29 changed files with 1281 additions and 433 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_collision_driver.c
+2 -1
View File
@@ -19,6 +19,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjxmacro.h>
#include "engine/engine_callback.h"
#include "engine/engine_collision_convex.h"
#include "engine/engine_collision_primitive.h"
@@ -55,7 +56,7 @@ void mj_collision(const mjModel* m, mjData* d) {
int *broadphasepair = 0;
mjMARKSTACK;
// clear size
// reset the size of the contact array
d->ncon = 0;
// return if disabled
src/engine/engine_core_constraint.c
+220 -19
View File
@@ -13,13 +13,15 @@
// limitations under the License.
#include "engine/engine_core_constraint.h"
#include <stdio.h>
#include <stddef.h>
#include <string.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include "engine/engine_collision_driver.h"
#include <mujoco/mjxmacro.h>
#include "engine/engine_array_safety.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
@@ -101,14 +103,27 @@ mjtNum mj_assignMargin(const mjModel* m, mjtNum source) {
// add contact to d->contact list; return 0 if success; 1 if buffer full
int mj_addContact(const mjModel* m, mjData* d, const mjContact* con) {
// if out of space, warn and return error
if (d->ncon >= m->nconmax) {
mj_warning(d, mjWARN_CONTACTFULL, m->nconmax);
// if nconmax is specified and ncon >= nconmax, warn and return error
if (m->nconmax != -1 && d->ncon >= m->nconmax) {
mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
return 1;
}
// move arena pointer back to the end of the existing contact array and invalidate efc_ arrays
d->parena = d->ncon * sizeof(mjContact);
d->nefc = 0;
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS
#undef X
d->contact = d->arena;
// copy contact
d->contact[d->ncon] = *con;
mjContact* dst = mj_arenaAlloc(d, sizeof(mjContact), _Alignof(mjContact));
if (!dst) {
mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
return 1;
}
*dst = *con;
// increase counter, return success
d->ncon++;
@@ -127,12 +142,6 @@ int mj_addConstraint(const mjModel* m, mjData* d,
int *nnz = d->efc_J_rownnz, *adr = d->efc_J_rowadr, *ind = d->efc_J_colind;
mjtNum *J = d->efc_J;
// if out of space, warn and return error
if (nefc+size > m->njmax) {
mj_warning(d, mjWARN_CNSTRFULL, m->njmax);
return 1;
}
// init empty guard for constraints other than contact
if (type==mjCNSTR_CONTACT_FRICTIONLESS ||
type==mjCNSTR_CONTACT_PYRAMIDAL ||
@@ -348,7 +357,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
mjtNum *jac[2], *jacdif, *data, *sparse_buf = NULL;
mjMARKSTACK;
// disabled or no equality contraints: return
// disabled or no equality constraints: return
if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax==0) {
return;
}
@@ -774,12 +783,6 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
b1 = m->geom_bodyid[con->geom1];
b2 = m->geom_bodyid[con->geom2];
// check size here, because pyramid rows are added incrementally
if (d->nefc + (dim==1 ? 1 : (ispyramid ? 2*(dim-1) : dim)) > m->njmax) {
mj_warning(d, mjWARN_CNSTRFULL, m->njmax);
break;
}
// save efc_address
con->efc_address = d->nefc;
@@ -1238,6 +1241,164 @@ void mj_makeImpedance(const mjModel* m, mjData* d) {
//------------------------------------- constraint counting ----------------------------------------
// count equality constraints
static inline int mj_ne(const mjModel* m, const mjData* d) {
// disabled or no equality constraints: return
if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax==0) {
return 0;
}
int ne = 0;
for (int i=0; i<m->neq; i++) {
if (!m->eq_active[i]) {
continue;
}
// process according to type
switch (m->eq_type[i]) {
case mjEQ_CONNECT:
ne += 3;
break;
case mjEQ_WELD:
ne += 6;
break;
case mjEQ_JOINT:
case mjEQ_TENDON:
ne++;
break;
default: // SHOULD NOT OCCUR
mju_error_i("Invalid equality constraint type %d", m->eq_type[i]);
}
}
return ne;
}
// count frictional constraints
static inline int mj_nf(const mjModel* m, const mjData* d) {
// disabled: return
if (mjDISABLED(mjDSBL_FRICTIONLOSS)) {
return 0;
}
int nf = 0;
const int nv = m->nv;
const int ntendon = m->ntendon;
// count frictional dofs
for (int i=0; i<nv; i++) {
nf += (m->dof_frictionloss[i] > 0);
}
// count frictional tendons
for (int i=0; i<ntendon; i++) {
nf += (m->tendon_frictionloss[i] > 0);
}
return nf;
}
// count limit constraints
static inline int mj_nl(const mjModel* m, const mjData* d) {
// disabled: return
if (mjDISABLED(mjDSBL_LIMIT)) {
return 0;
}
int nl = 0;
const int njnt = m->njnt;
const int ntendon = m->ntendon;
// count limited joints
for (int i=0; i<njnt; i++) {
if (!m->jnt_limited[i]) {
continue;
}
// slides and hinges can have active limits on two sides, check both
if (m->jnt_type[i]==mjJNT_SLIDE || m->jnt_type[i]==mjJNT_HINGE) {
// get margin
mjtNum margin = m->jnt_margin[i];
// get joint value
mjtNum value = d->qpos[m->jnt_qposadr[i]];
// check lower and upper limits
for (int side=-1; side<=1; side+=2) {
// compute distance (negative: penetration)
mjtNum dist = side * (m->jnt_range[2*i+(side+1)/2] - value);
// detect joint limit
if (dist<margin) {
nl++;
}
}
} else {
nl++;
}
}
// count limited tendons
for (int i=0; i<ntendon; i++) {
nl += m->tendon_limited[i];
}
return nl;
}
// count contact constraints
static inline int mj_nc(const mjModel* m, const mjData* d) {
// disabled or no contacts: return
int ncon = d->ncon;
if (mjDISABLED(mjDSBL_CONTACT) || ncon==0) {
return 0;
}
int nc = 0;
int ispyramid = mj_isPyramidal(m);
// find contacts to be counted
for (int i=0; i<ncon; i++) {
mjContact* con = d->contact + i;
if (con->exclude) {
continue;
}
int dim = con->dim;
// dim 1: single constraint
if (dim==1) {
nc++;
}
// dim > 1: depends on cone type
else {
nc += (ispyramid ? 2*(dim-1) : dim);
}
}
return nc;
}
// count all constraints
static inline int mj_nefc(const mjModel* m, const mjData* d) {
return mj_ne(m, d) + mj_nf(m, d) + mj_nl(m, d) + mj_nc(m, d);
}
//---------------------------- top-level API for constraint construction ---------------------------
// driver: call all functions above
@@ -1246,16 +1407,56 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
d->ne = d->nf = d->nefc = 0;
// disabled or Jacobian not allocated: return
if (mjDISABLED(mjDSBL_CONSTRAINT) || m->njmax==0) {
if (mjDISABLED(mjDSBL_CONSTRAINT)) {
return;
}
int nefc_allocated = mj_nefc(m, d);
d->nefc = nefc_allocated;
#undef MJ_M
#define MJ_M(n) m->n
#undef MJ_D
#define MJ_D(n) d->n
// move arena pointer to end of contact array
d->parena = d->ncon * sizeof(mjContact);
#define X(type, name, nr, nc) \
d->name = mj_arenaAlloc(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \
if (!d->name) { \
mj_warning(d, mjWARN_CNSTRFULL, d->nstack * sizeof(mjtNum)); \
d->nefc = 0; \
return; \
}
MJDATA_ARENA_POINTERS_PRIMAL
if (mj_isDual(m)) {
MJDATA_ARENA_POINTERS_DUAL
}
#undef X
#undef MJ_M
#define MJ_M(n) n
#undef MJ_D
#define MJ_D(n) n
d->nefc = 0;
// instantiate all elements of Jacobian
mj_instantiateEquality(m, d);
mj_instantiateFriction(m, d);
mj_instantiateLimit(m, d);
mj_instantiateContact(m, d);
if (d->nefc > nefc_allocated) {
char msg[1024];
mjSNPRINTF(
msg, "nefc under-allocation: found nefc=%d but allocated only %d", d->nefc, nefc_allocated);
mju_error(msg);
}
// collect memory use statistics
d->maxuse_con = mjMAX(d->maxuse_con, d->ncon);
d->maxuse_efc = mjMAX(d->maxuse_efc, d->nefc);
src/engine/engine_io.c
+77 -45
View File
@@ -24,37 +24,17 @@
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjxmacro.h>
#include "engine/engine_array_safety.h"
#include "engine/engine_macro.h"
#include "engine/engine_plugin.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_vfs.h"
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
#elif defined(_MSC_VER)
#define ASAN_POISON_MEMORY_REGION(addr, size)
#define ASAN_UNPOISON_MEMORY_REGION(addr, size)
#else
#define ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#define ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#endif
#ifdef MEMORY_SANITIZER
#include <sanitizer/msan_interface.h>
#endif
#ifdef _MSC_VER
#pragma warning (disable: 4305) // disable MSVC warning: truncation from 'double' to 'float'
#endif
#ifndef __has_builtin
#define __has_builtin(x) 0
#endif
#define PTRDIFF(x, y) ((void*)(x) - (void*)(y))
//------------------------------ mjLROpt -----------------------------------------------------------
// set default options for length range computation
@@ -260,14 +240,6 @@ void mj_defaultStatistic(mjStatistic* stat) {
static const int ID = 54321;
// number of bytes to be skipped to achieve 64-byte alignment
static unsigned int SKIP(intptr_t offset) {
const unsigned int align = 64;
// compute skipped bytes
return (align - (offset % align)) % align;
}
// count ints in mjModel
static int getnint(void) {
@@ -840,6 +812,13 @@ static void mj_setPtrData(const mjModel* m, mjData* d) {
if (d->nbuffer != sz) {
mju_error("mjData buffer size mismatch");
}
// zero-initialize arena pointers
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS
#undef X
d->contact = d->arena;
}
@@ -859,7 +838,7 @@ static mjData* _makeData(const mjModel* m) {
// compute buffer size
d->nbuffer = 0;
d->buffer = d->stack = NULL;
d->buffer = d->arena = NULL;
#define X(type, name, nr, nc) \
if (!safeAddToBufferSize(&offset, &d->nbuffer, sizeof(type), m->nr, nc)) { \
mju_free(d); \
@@ -880,12 +859,12 @@ static mjData* _makeData(const mjModel* m) {
mju_error("Could not allocate mjData buffer");
}
// allocate stack
d->stack = (mjtNum*) mju_malloc(d->nstack * sizeof(mjtNum));
if (!d->stack) {
// allocate arena
d->arena = mju_malloc(d->nstack * sizeof(mjtNum));
if (!d->arena) {
mju_free(d->buffer);
mju_free(d);
mju_error("Could not allocate mjData stack");
mju_error("Could not allocate mjData arena");
}
// set pointers into buffer, reset data
@@ -917,7 +896,7 @@ mjData* mj_makeData(const mjModel* m) {
// copy mjData, if dest==NULL create new data
mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
void* save_buffer;
mjtNum* save_stack;
void* save_arena;
// allocate new data if needed
if (!dest) {
@@ -939,10 +918,10 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
// save pointers, copy everything, restore pointers
save_buffer = dest->buffer;
save_stack = dest->stack;
save_arena = dest->arena;
*dest = *src;
dest->buffer = save_buffer;
dest->stack = save_stack;
dest->arena = save_arena;
mj_setPtrData(m, dest);
// save plugin_data, since the X macro copying block below will override it
@@ -965,6 +944,16 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
#undef X
}
// copy arena memory
memcpy(dest->arena, src->arena, src->nstack * sizeof(mjtNum));
#define X(type, name, nr, nc) \
dest->name = src->name ? (type*)((char*)dest->arena + PTRDIFF(src->name, src->arena)) : NULL;
MJDATA_ARENA_POINTERS
#undef X
// restore contact pointer
dest->contact = dest->arena;
// restore plugin_data
if (plugin_data_size) {
memcpy(dest->plugin_data, save_plugin_data, plugin_data_size);
@@ -986,25 +975,60 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
// allocate memory from the mjData arena
void* mj_arenaAlloc(mjData* d, int bytes, int alignment) {
int misalignment = d->parena % alignment;
int padding = misalignment ? alignment - misalignment : 0;
// check size
size_t bytes_available = (d->nstack - d->pstack) * sizeof(mjtNum);
if (d->parena + padding + bytes > bytes_available) {
return NULL;
}
// allocate, update max, return pointer to buffer
void* result = (char*)d->arena + d->parena + padding;
d->parena += padding + bytes;
d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
return result;
}
// allocate size mjtNums on the mjData stack
mjtNum* mj_stackAlloc(mjData* d, int size) {
mjtNum* result;
// return NULL if empty
if (!size) {
return 0;
}
// check size
if (d->pstack + size > d->nstack) {
mju_error("Stack overflow");
size_t stack_available_bytes = d->nstack * sizeof(mjtNum) - d->parena;
size_t stack_required_bytes = (d->pstack + size) * sizeof(mjtNum);
if (stack_required_bytes > stack_available_bytes) {
char err[256];
mjSNPRINTF(err, "stack overflow: max = %zu, available = %zu, requested = %zu "
"(ne = %d, nf = %d, nefc = %d, ncon = %d)",
d->nstack * sizeof(mjtNum), stack_available_bytes, stack_required_bytes,
d->ne, d->nf, d->nefc, d->ncon);
mju_error(err);
}
// allocate, update max, return pointer to buffer
result = (mjtNum*)d->stack + d->pstack;
// allocate at end of arena
char* end_ptr = (char*)d->arena + d->nstack * sizeof(mjtNum);
char* result = end_ptr - (d->pstack + size + 1) * sizeof(mjtNum);
#ifdef ADDRESS_SANITIZER
if ((uintptr_t)result % sizeof(mjtNum)) {
mju_error("mj_stackAlloc fails to align to sizeof(mjtNum)");
}
#endif
// update max, return pointer to buffer
d->pstack += size;
d->maxuse_stack = mjMAX(d->maxuse_stack, d->pstack);
return result;
d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack*sizeof(mjtNum) + d->parena);
return (mjtNum*)result;
}
@@ -1022,8 +1046,16 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
// clear stack pointer
d->pstack = 0;
// clear arena pointers
d->parena = 0;
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS
#undef X
d->contact = d->arena;
// clear memory utilization stats
d->maxuse_stack = 0;
d->maxuse_arena = 0;
d->maxuse_con = 0;
d->maxuse_efc = 0;
@@ -1162,7 +1194,7 @@ void mj_deleteData(mjData* d) {
}
}
mju_free(d->buffer);
mju_free(d->stack);
mju_free(d->arena);
mju_free(d);
}
}
src/engine/engine_io.h
+3
View File
@@ -96,6 +96,9 @@ 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
void* mj_arenaAlloc(mjData* d, int bytes, int alignment);
// mjData stack allocate
MJAPI mjtNum* mj_stackAlloc(mjData* d, int size);
src/engine/engine_macro.h
+33
View File
@@ -15,6 +15,8 @@
#ifndef MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_
#define MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_
#include <stdint.h>
#include "engine/engine_callback.h" // IWYU pragma: export
//-------------------------------- utility macros --------------------------------------------------
@@ -40,4 +42,35 @@
#define TM_START1 mjtNum _tm1 = (mjcb_time ? mjcb_time() : 0);
#define TM_END1(i) {d->timer[i].duration += ((mjcb_time ? mjcb_time() : 0) - _tm1); d->timer[i].number++;}
//-------------------------- sanitizer macros ------------------------------------------------------
#ifdef ADDRESS_SANITIZER
#include <sanitizer/asan_interface.h>
#elif defined(_MSC_VER)
#define ASAN_POISON_MEMORY_REGION(addr, size)
#define ASAN_UNPOISON_MEMORY_REGION(addr, size)
#else
#define ASAN_POISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#define ASAN_UNPOISON_MEMORY_REGION(addr, size) ((void)(addr), (void)(size))
#endif
#ifdef MEMORY_SANITIZER
#include <sanitizer/msan_interface.h>
#endif
#ifndef __has_builtin
#define __has_builtin(x) 0
#endif
//-------------------------- pointer arithmetic ----------------------------------------------------
#define PTRDIFF(x, y) ((char*)(x) - (char*)(y))
// number of bytes to be skipped to achieve 64-byte alignment
static inline unsigned int SKIP(intptr_t offset) {
const unsigned int align = 64;
// compute skipped bytes
return (align - (offset % align)) % align;
}
#endif // MUJOCO_SRC_ENGINE_ENGINE_MACRO_H_
src/engine/engine_print.c
+18 -1
View File
@@ -37,6 +37,7 @@
#define FLOAT_FORMAT "% -9.2g"
#define FLOAT_FORMAT_MAX_LEN 20
#define INT_FORMAT " %d"
#define SIZE_T_FORMAT " %zu"
#define NAME_FORMAT "%-21s"
@@ -46,6 +47,9 @@
// print 2D array of mjtNum into file
static void printArray(const char* str, int nr, int nc, const mjtNum* data, FILE* fp,
const char* float_format) {
if (!data) {
return;
}
if (nr && nc) {
fprintf(fp, "%s\n ", str);
for (int r=0; r<nr; r++) {
@@ -62,6 +66,9 @@ static void printArray(const char* str, int nr, int nc, const mjtNum* data, FILE
// print 2D array of int into file
static void printArrayInt(const char* str, int nr, int nc, const int* data, FILE* fp) {
if (!data) {
return;
}
if (nr && nc) {
fprintf(fp, "%s\n ", str);
for (int r=0; r<nr; r++) {
@@ -80,6 +87,9 @@ static void printArrayInt(const char* str, int nr, int nc, const int* data, FILE
static void printSparse(const char* str, const mjtNum* mat, int nr,
const int* rownnz, const int* rowadr,
const int* colind, FILE* fp, const char* float_format) {
if (!mat) {
return;
}
fprintf(fp, "%s\n ", str);
for (int r=0; r<nr; r++) {
@@ -98,6 +108,9 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
// print vector
static void printVector(const char* str, const mjtNum* data, int n, FILE* fp,
const char* float_format) {
if (!data) {
return;
}
// print str
fprintf(fp, "%s", str);
@@ -754,7 +767,11 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
fprintf(fp, "SIZES\n");
#define X(type, name) \
{ \
const char* format = _Generic(d->name, int : INT_FORMAT, default : NULL); \
const char* format = _Generic( \
d->name, \
int : INT_FORMAT, \
size_t : SIZE_T_FORMAT, \
default : NULL); \
if (format) { \
fprintf(fp, " "); \
fprintf(fp, NAME_FORMAT, #name); \
src/engine/engine_util_misc.c
+7 -2
View File
@@ -956,11 +956,16 @@ const char* mju_warningText(int warning, int info) {
break;
case mjWARN_CONTACTFULL:
mjSNPRINTF(str, "Pre-allocated contact buffer is full. Increase nconmax above %d.", info);
mjSNPRINTF(str,
"Too many contacts. Either the arena memory is full, or nconmax is specified and is "
"exceeded. Increase arena memory allocation, or increase/remove nconmax. "
"(ncon = %d)", info);
break;
case mjWARN_CNSTRFULL:
mjSNPRINTF(str, "Pre-allocated constraint buffer is full. Increase njmax above %d.", info);
mjSNPRINTF(str,
"Insufficient arena memory for the number of constraints generated. "
"Increase arena memory allocation above %d bytes.", info);
break;
case mjWARN_VGEOMFULL: