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
+18 -18
View File
@@ -459,7 +459,7 @@ void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr
void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
int nv = m->nv;
mj_markStack(d);
mjtNum* jacp_b = mj_stackAllocNum(d, 3*nv);
mjtNum* jacp_b = mjSTACKALLOC(d, 3*nv, mjtNum);
// clear output
mju_zero(jacp, 3*nv);
@@ -505,8 +505,8 @@ void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacA
// get full Jacobian of point
mj_markStack(d);
mjtNum* jacp = (jacPoint ? jacPoint : mj_stackAllocNum(d, 3*nv));
mjtNum* jacr = mj_stackAllocNum(d, 3*nv);
mjtNum* jacp = (jacPoint ? jacPoint : mjSTACKALLOC(d, 3*nv, mjtNum));
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mj_jac(m, d, jacp, jacr, point, body);
// jacAxis_col = cross(jacr_col, axis)
@@ -741,15 +741,15 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
mj_markStack(d);
mjtNum* jtmp = mj_stackAllocNum(d, flg_rot ? 6*nv : 3*nv);
mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
mjtNum* jp = jtmp;
mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
// sparse
if (mj_isSparse(m)) {
mjtNum* buf = mj_stackAllocNum(d, flg_rot ? 6*nv : 3*nv);
int* buf_ind = mj_stackAllocInt(d, nv);
int* bodychain = mj_stackAllocInt(d, nv);
mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
int* buf_ind = mjSTACKALLOC(d, nv, int);
int* bodychain = mjSTACKALLOC(d, nv, int);
// set first
NV = mj_bodyChain(m, body[0], chain);
@@ -878,10 +878,10 @@ void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
mj_markStack(d);
// stack allocations
mjtNum* jacp = mj_stackAllocNum(d, 3*nv);
mjtNum* jacr = mj_stackAllocNum(d, 3*nv);
mjtNum* term1 = mj_stackAllocNum(d, 3*nv);
mjtNum* term2 = mj_stackAllocNum(d, 3*nv);
mjtNum* jacp = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacr = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* term2 = mjSTACKALLOC(d, 3*nv, mjtNum);
// clear output
mju_zero(mat, 3*nv);
@@ -1153,7 +1153,7 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
mj_markStack(d);
// create reduced sparse inertia matrix C
mjtNum* C = mj_stackAllocNum(d, nC);
mjtNum* C = mjSTACKALLOC(d, nC, mjtNum);
for (int i=0; i < nC; i++) {
C[i] = d->qM[d->mapM2C[i]];
}
@@ -1178,8 +1178,8 @@ void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int nv = m->nv;
mj_markStack(d);
int* buf_ind = mj_stackAllocInt(d, nv);
mjtNum* sparse_buf = mj_stackAllocNum(d, nv);
int* buf_ind = mjSTACKALLOC(d, nv, int);
mjtNum* sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
// add to destination
for (int i=0; i < nv; i++) {
@@ -1230,9 +1230,9 @@ void mj_applyFT(const mjModel* m, mjData* d,
// allocate local variables
mj_markStack(d);
mjtNum* jacp = force ? mj_stackAllocNum(d, 3*nv) : NULL;
mjtNum* jacr = torque ? mj_stackAllocNum(d, 3*nv) : NULL;
mjtNum* qforce = mj_stackAllocNum(d, nv);
mjtNum* jacp = force ? mjSTACKALLOC(d, 3*nv, mjtNum) : NULL;
mjtNum* jacr = torque ? mjSTACKALLOC(d, 3*nv, mjtNum) : NULL;
mjtNum* qforce = mjSTACKALLOC(d, nv, mjtNum);
// make sure body is in range
if (body < 0 || body >= m->nbody) {
@@ -1242,7 +1242,7 @@ void mj_applyFT(const mjModel* m, mjData* d,
// sparse case
if (mj_isSparse(m)) {
// construct chain and sparse Jacobians
int* chain = mj_stackAllocInt(d, nv);
int* chain = mjSTACKALLOC(d, nv, int);
int NV = mj_bodyChain(m, body, chain);
mj_jacSparse(m, d, jacp, jacr, point, body, NV, chain);