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

Add mj_stackAllocInt to get correct size for allocating ints on mjData stack. Reducing stack memory usage by 10% - 15%.

Browse Source 
PiperOrigin-RevId: 517672195
Change-Id: Ie14e6c9c99e8b9e1c5d3686b77631f2df0651c67
This commit is contained in:
Kyle Bayes
2023-03-18 13:03:11 -07:00
committed by Copybara-Service
parent fae5e822c0
commit 8696247f9d
17 changed files with 103 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_collision_driver.c
+3 -2
View File
@@ -388,8 +388,9 @@ void mj_collision(const mjModel* m, mjData* d) {
// dynamic only or merge; apply exclude
else {
// call broadphase collision detector
broadphasepair = (int*)mj_stackAlloc(d, (m->nbody*(m->nbody-1))/2);
nbodypair = mj_broadphase(m, d, broadphasepair, (m->nbody*(m->nbody-1))/2);
int npairs = (m->nbody*(m->nbody - 1))/2;
broadphasepair = mj_stackAllocInt(d, npairs);
nbodypair = mj_broadphase(m, d, broadphasepair, npairs);
unsigned int last_signature = -1;
// loop over body pairs (broadphase or all)
src/engine/engine_core_constraint.c
+13 -13
View File
@@ -374,9 +374,9 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
jac[1] = mj_stackAlloc(d, 6*nv);
jacdif = mj_stackAlloc(d, 6*nv);
if (issparse) {
chain = (int*)mj_stackAlloc(d, nv);
chain2 = (int*)mj_stackAlloc(d, nv);
buf_ind = (int*)mj_stackAlloc(d, nv);
chain = mj_stackAllocInt(d, nv);
chain2 = mj_stackAllocInt(d, nv);
buf_ind = mj_stackAllocInt(d, nv);
sparse_buf = mj_stackAlloc(d, nv);
}
@@ -775,7 +775,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
jac1r = mj_stackAlloc(d, 3*NV);
jac2r = mj_stackAlloc(d, 3*NV);
if (issparse) {
chain = (int*)mj_stackAlloc(d, NV);
chain = mj_stackAllocInt(d, NV);
}
// find contacts to be included
@@ -1321,8 +1321,8 @@ static inline int mj_ne(const mjModel* m, mjData* d, int* nnz) {
mjMARKSTACK;
if (nnz) {
chain = (int*)mj_stackAlloc(d, nv);
chain2 = (int*)mj_stackAlloc(d, nv);
chain = mj_stackAllocInt(d, nv);
chain2 = mj_stackAllocInt(d, nv);
}
// find active equality constraints
@@ -1692,13 +1692,13 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// sparse
if (mj_isSparse(m)) {
// space for JM2 and JM2T indices
int* rownnz = (int*)mj_stackAlloc(d, nefc);
int* rowadr = (int*)mj_stackAlloc(d, nefc);
int* colind = (int*)mj_stackAlloc(d, nefc*nv);
int* rowsuper = (int*)mj_stackAlloc(d, nefc);
int* rownnzT = (int*)mj_stackAlloc(d, nv);
int* rowadrT = (int*)mj_stackAlloc(d, nv);
int* colindT = (int*)mj_stackAlloc(d, nv*nefc);
int* rownnz = mj_stackAllocInt(d, nefc);
int* rowadr = mj_stackAllocInt(d, nefc);
int* colind = mj_stackAllocInt(d, nefc*nv);
int* rowsuper = mj_stackAllocInt(d, nefc);
int* rownnzT = mj_stackAllocInt(d, nv);
int* rowadrT = mj_stackAllocInt(d, nv);
int* colindT = mj_stackAllocInt(d, nv*nefc);
// construct JM2 = backsubM2(J')' by rows
for (int r=0; r<nefc; r++) {
src/engine/engine_core_smooth.c
+3 -3
View File
@@ -399,8 +399,8 @@ void mj_tendon(const mjModel* m, mjData* d) {
jacdif = mj_stackAlloc(d, 3*nv);
tmp = mj_stackAlloc(d, nv);
if (issparse) {
chain = (int*)mj_stackAlloc(d, nv);
buf_ind = (int*)mj_stackAlloc(d, nv);
chain = mj_stackAllocInt(d, nv);
buf_ind = mj_stackAllocInt(d, nv);
sparse_buf = mj_stackAlloc(d, nv);
}
@@ -869,7 +869,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
jacdifp = mj_stackAlloc(d, 3*nv);
jac1p = mj_stackAlloc(d, 3*nv);
jac2p = mj_stackAlloc(d, 3*nv);
chain = issparse ? (int*)mj_stackAlloc(d, nv) : NULL;
chain = issparse ? mj_stackAllocInt(d, nv) : NULL;
}
// clear efc_force and moment_exclude
src/engine/engine_derivative.c
+3 -3
View File
@@ -1122,8 +1122,8 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) {
int rownnz[6], rowadr[6];
mjtNum* J = mj_stackAlloc(d, 6*nv);
mjtNum* tmp = mj_stackAlloc(d, 3*nv);
int* colind = (int*) mj_stackAlloc(d, 6*nv);
int* colind_compressed = (int*) mj_stackAlloc(d, 6*nv);
int* colind = mj_stackAllocInt(d, 6*nv);
int* colind_compressed = mj_stackAllocInt(d, 6*nv);
mjtNum lvel[6], wind[6], lwind[6];
mjtNum geom_interaction_coef, magnus_lift_coef, kutta_lift_coef;
@@ -1239,7 +1239,7 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, int i)
int rownnz[6], rowadr[6];
mjtNum* J = mj_stackAlloc(d, 6*nv);
mjtNum* tmp = mj_stackAlloc(d, 3*nv);
int* colind = (int*) mj_stackAlloc(d, 6*nv);
int* colind = mj_stackAllocInt(d, 6*nv);
mjtNum lvel[6], wind[6], lwind[6], box[3], B;
mjtNum* inertia = m->body_inertia + 3*i;
src/engine/engine_derivative_fd.c
+1 -1
View File
@@ -185,7 +185,7 @@ void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps) {
mjMARKSTACK;
mjtNum* qfrc_passive = mj_stackAlloc(d, nv);
mjtNum* fd = mj_stackAlloc(d, nv);
int* cnt = (int*)mj_stackAlloc(d, nv);
int* cnt = mj_stackAllocInt(d, nv);
// clear row counters
memset(cnt, 0, nv*sizeof(int));
src/engine/engine_io.c
+15 -2
View File
@@ -819,7 +819,7 @@ static void makeDSparse(const mjModel* m, mjData* d) {
int* colind = d->D_colind;
mjMARKSTACK;
int* remaining = (int*)mj_stackAlloc(d, nv);
int* remaining = mj_stackAllocInt(d, nv);
// compute rownnz
memset(rownnz, 0, nv * sizeof(int));
@@ -912,7 +912,7 @@ static void makeBSparse(const mjModel* m, mjData* d) {
// allocate and clear incremental row counts
mjMARKSTACK;
int* cnt = (int*)mj_stackAlloc(d, nbody);
int* cnt = mj_stackAllocInt(d, nbody);
memset(cnt, 0, sizeof(int) * nbody);
// add subtree dofs to colind
@@ -1232,6 +1232,19 @@ mjtNum* mj_stackAlloc(mjData* d, int size) {
int* mj_stackAllocInt(mjData* d, int size) {
// optimize for mjtNum being twice the size of int
if (2*sizeof(int) == sizeof(mjtNum)) {
return (int*)mj_stackAlloc(d, (size + 1) >> 1);
}
// arbitrary bytes sizes
int new_size = (sizeof(int)*size + sizeof(mjtNum) - 1) / sizeof(mjtNum);
return (int*)mj_stackAlloc(d, new_size);
}
// clear data, set defaults
static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
//------------------------------ save plugin state and data
src/engine/engine_io.h
+4 -1
View File
@@ -102,9 +102,12 @@ 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
// mjData stack allocate for array of mjtNums
MJAPI mjtNum* mj_stackAlloc(mjData* d, int size);
// mjData stack allocate for array of ints
MJAPI int* mj_stackAllocInt(mjData* d, int size);
// de-allocate data
MJAPI void mj_deleteData(mjData* d);
src/engine/engine_solver.c
+6 -6
View File
@@ -309,7 +309,7 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
mjContact* con;
mjMARKSTACK;
mjtNum* ARinv = mj_stackAlloc(d, nefc);
int* oldstate = (int*)mj_stackAlloc(d, nefc);
int* oldstate = mj_stackAllocInt(d, nefc);
// precompute inverse diagonal of AR
ARdiaginv(m, d, ARinv, 0);
@@ -511,7 +511,7 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
mjContact* con;
mjMARKSTACK;
mjtNum* ARinv = mj_stackAlloc(d, nefc);
int* oldstate = (int*)mj_stackAlloc(d, nefc);
int* oldstate = mj_stackAllocInt(d, nefc);
// precompute inverse diagonal of A
ARdiaginv(m, d, ARinv, 1);
@@ -768,9 +768,9 @@ static void CGallocate(const mjModel* m, mjData* d,
if (flg_Newton) {
ctx->H = mj_stackAlloc(d, nv*nv);
ctx->Hcone = mj_stackAlloc(d, nv*nv);
ctx->rownnz = (int*)mj_stackAlloc(d, nv);
ctx->rowadr = (int*)mj_stackAlloc(d, nv);
ctx->colind = (int*)mj_stackAlloc(d, nv*nv);
ctx->rownnz = mj_stackAllocInt(d, nv);
ctx->rowadr = mj_stackAllocInt(d, nv);
ctx->colind = mj_stackAllocInt(d, nv*nv);
}
}
@@ -1507,7 +1507,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int maxiter, int flg_New
Mgradold = mj_stackAlloc(d, nv);
Mgraddif = mj_stackAlloc(d, nv);
}
int* oldstate = (int*)mj_stackAlloc(d, nefc);
int* oldstate = mj_stackAllocInt(d, nefc);
// initialize matrix-vector products
mj_mulM(m, d, ctx.Ma, d->qacc);
src/engine/engine_support.c
+1 -1
View File
@@ -984,7 +984,7 @@ void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* sr
mjMARKSTACK;
// init remaining
int* remaining = (int*)mj_stackAlloc(d, nv);
int* remaining = mj_stackAllocInt(d, nv);
memcpy(remaining, d->D_rownnz, nv * sizeof(int));
// copy data
src/engine/engine_util_sparse.c
+3 -3
View File
@@ -418,8 +418,8 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
mjData* d) {
// allocate space for accumulation buffer and matT
mjMARKSTACK;
@@ -434,7 +434,7 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
// these mark the currently set columns in the dense row buffer,
// used for when creating the resulting sparse row
int* markers = (int*) mj_stackAlloc(d, nc);
int* markers = mj_stackAllocInt(d, nc);
for (int i=0; i<nc; i++) {
int* cols = res_colind+res_rowadr[i];