Add engine-internal convenience macro for allocating typed arrays, improve error message.

PiperOrigin-RevId: 705126655
Change-Id: I2bd8fada6d33a919d2fb82297f93ac57958355a4
This commit is contained in:
Yuval Tassa
2024-12-11 09:03:39 -08:00
committed by Copybara-Service
parent 3ca97248a3
commit 2691887500
25 changed files with 283 additions and 274 deletions
+27 -27
View File
@@ -184,7 +184,7 @@ void mj_comPos(const mjModel* m, mjData* d) {
int nbody = m->nbody, njnt = m->njnt;
mjtNum offset[3], axis[3];
mj_markStack(d);
mjtNum* mass_subtree = mj_stackAllocNum(d, m->nbody);
mjtNum* mass_subtree = mjSTACKALLOC(d, m->nbody, mjtNum);
// clear subtree
mju_zero(mass_subtree, m->nbody);
@@ -393,7 +393,7 @@ void mj_camlight(const mjModel* m, mjData* d) {
// update dynamic BVH; leaf aabbs must be updated before call
void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) {
mj_markStack(d);
int* modified = mj_stackAllocInt(d, bvhnum);
int* modified = mjSTACKALLOC(d, bvhnum, int);
mju_zeroInt(modified, bvhnum);
// mark leafs as modified
@@ -526,10 +526,10 @@ void mj_flex(const mjModel* m, mjData* d) {
// allocate space
mj_markStack(d);
mjtNum* jac1 = mj_stackAllocNum(d, 3*nv);
mjtNum* jac2 = mj_stackAllocNum(d, 3*nv);
mjtNum* jacdif = mj_stackAllocNum(d, 3*nv);
int* chain = issparse ? mj_stackAllocInt(d, nv) : NULL;
mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL;
// clear Jacobian: sparse or dense
if (issparse) {
@@ -631,14 +631,14 @@ void mj_tendon(const mjModel* m, mjData* d) {
// allocate space
mj_markStack(d);
jac1 = mj_stackAllocNum(d, 3*nv);
jac2 = mj_stackAllocNum(d, 3*nv);
jacdif = mj_stackAllocNum(d, 3*nv);
tmp = mj_stackAllocNum(d, nv);
jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
tmp = mjSTACKALLOC(d, nv, mjtNum);
if (issparse) {
chain = mj_stackAllocInt(d, nv);
buf_ind = mj_stackAllocInt(d, nv);
sparse_buf = mj_stackAllocNum(d, nv);
chain = mjSTACKALLOC(d, nv, int);
buf_ind = mjSTACKALLOC(d, nv, int);
sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
}
// clear results
@@ -863,9 +863,9 @@ void mj_transmission(const mjModel* m, mjData* d) {
// allocate Jacbians
mj_markStack(d);
mjtNum* jac = mj_stackAllocNum(d, 3*nv);
mjtNum* jacA = mj_stackAllocNum(d, 3*nv);
mjtNum* jacS = mj_stackAllocNum(d, 3*nv);
mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacA = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacS = mjSTACKALLOC(d, 3*nv, mjtNum);
// define stack variables required for body transmission, don't allocate
int issparse = mj_isSparse(m);
@@ -1088,7 +1088,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
// reference site defined
else {
int refid = m->actuator_trnid[2*i+1];
if (!jacref) jacref = mj_stackAllocNum(d, 3*nv);
if (!jacref) jacref = mjSTACKALLOC(d, 3*nv, mjtNum);
// initialize last dof address for each body
int b0 = m->body_weldid[m->site_bodyid[id]];
@@ -1190,7 +1190,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
mju_mulMatVec3(wrench, d->site_xmat+9*refid, gear+3);
// moment_tmp: global Jacobian projected on wrench, add to moment
if (!moment_tmp) moment_tmp = mj_stackAllocNum(d, nv);
if (!moment_tmp) moment_tmp = mjSTACKALLOC(d, nv, mjtNum);
mju_mulMatTVec(moment_tmp, jacS, wrench, 3, nv);
mju_addTo(moment+adr, moment_tmp, nv);
}
@@ -1220,12 +1220,12 @@ void mj_transmission(const mjModel* m, mjData* d) {
{
// allocate stack variables for the first mjTRN_BODY
if (!efc_force) {
efc_force = mj_stackAllocNum(d, d->nefc);
moment_exclude = mj_stackAllocNum(d, nv);
jacdifp = mj_stackAllocNum(d, 3*nv);
jac1p = mj_stackAllocNum(d, 3*nv);
jac2p = mj_stackAllocNum(d, 3*nv);
chain = issparse ? mj_stackAllocInt(d, nv) : NULL;
efc_force = mjSTACKALLOC(d, d->nefc, mjtNum);
moment_exclude = mjSTACKALLOC(d, nv, mjtNum);
jacdifp = mjSTACKALLOC(d, 3*nv, mjtNum);
jac1p = mjSTACKALLOC(d, 3*nv, mjtNum);
jac2p = mjSTACKALLOC(d, 3*nv, mjtNum);
chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL;
}
// clear efc_force and moment_exclude
@@ -1804,7 +1804,7 @@ void mj_subtreeVel(const mjModel* m, mjData* d) {
int nbody = m->nbody;
mjtNum dx[3], dv[3], dp[3], dL[3];
mj_markStack(d);
mjtNum* body_vel = mj_stackAllocNum(d, 6*m->nbody);
mjtNum* body_vel = mjSTACKALLOC(d, 6*m->nbody, mjtNum);
// bodywise quantities
for (int i=0; i < nbody; i++) {
@@ -1871,8 +1871,8 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) {
int nbody = m->nbody, nv = m->nv;
mjtNum tmp[6], tmp1[6];
mj_markStack(d);
mjtNum* loc_cacc = mj_stackAllocNum(d, m->nbody*6);
mjtNum* loc_cfrc_body = mj_stackAllocNum(d, m->nbody*6);
mjtNum* loc_cacc = mjSTACKALLOC(d, m->nbody*6, mjtNum);
mjtNum* loc_cfrc_body = mjSTACKALLOC(d, m->nbody*6, mjtNum);
// set world acceleration to -gravity
mju_zero(loc_cacc, 6);