From 26918875006dec8d924dcb779fb2d8cab64a3b4d Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 11 Dec 2024 09:03:39 -0800
Subject: [PATCH] Add engine-internal convenience macro for allocating typed
 arrays, improve error message.

PiperOrigin-RevId: 705126655
Change-Id: I2bd8fada6d33a919d2fb82297f93ac57958355a4
---
 python/mujoco/bindings_test.py       |  2 +-
 src/engine/engine_collision_driver.c | 36 ++++++--------
 src/engine/engine_collision_gjk.c    | 16 +++----
 src/engine/engine_collision_sdf.c    |  4 +-
 src/engine/engine_core_constraint.c  | 60 ++++++++++++------------
 src/engine/engine_core_smooth.c      | 54 ++++++++++-----------
 src/engine/engine_derivative.c       | 40 ++++++++--------
 src/engine/engine_derivative_fd.c    | 54 ++++++++++-----------
 src/engine/engine_forward.c          | 34 +++++++-------
 src/engine/engine_inverse.c          | 18 +++----
 src/engine/engine_io.c               | 56 ++++++++++++++--------
 src/engine/engine_io.h               | 10 +++-
 src/engine/engine_island.c           | 12 ++---
 src/engine/engine_passive.c          |  2 +-
 src/engine/engine_print.c            |  2 +-
 src/engine/engine_ray.c              |  4 +-
 src/engine/engine_sensor.c           |  4 +-
 src/engine/engine_setconst.c         | 14 +++---
 src/engine/engine_solver.c           | 70 ++++++++++++++--------------
 src/engine/engine_support.c          | 36 +++++++-------
 src/engine/engine_util_container.c   |  3 +-
 src/engine/engine_util_solve.c       |  8 ++--
 src/engine/engine_util_sparse.c      | 10 ++--
 src/engine/engine_vis_interact.c     |  4 +-
 src/engine/engine_vis_visualize.c    |  4 +-
 25 files changed, 283 insertions(+), 274 deletions(-)

diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index daed5d2c..6bfebb90 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -1060,7 +1060,7 @@ Euler integrator, semi-implicit in velocity.
   def test_can_raise_error(self):
     self.data.pstack = self.data.narena
     with self.assertRaisesRegex(
-        mujoco.FatalError, r'\Amj_stackAlloc: insufficient memory:'
+        mujoco.FatalError, r'\Amj_stackAlloc: out of memory, stack overflow'
     ):
       mujoco.mj_forward(self.model, self.data)
 
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index ef5710a6..f46d8157 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -289,7 +289,7 @@ void mj_collision(const mjModel* m, mjData* d) {
   // broadphase collision detector
   TM_START;
   int nmaxpairs = (nbodyflex*(nbodyflex - 1))/2;
-  int* broadphasepair = mj_stackAllocInt(d, nmaxpairs);
+  int* broadphasepair = mjSTACKALLOC(d, nmaxpairs, int);
   int nbfpair = mj_broadphase(m, d, broadphasepair, nmaxpairs);
   unsigned int last_signature = -1;
   TM_END(mjTIMER_COL_BROAD);
@@ -368,8 +368,7 @@ void mj_collision(const mjModel* m, mjData* d) {
       int n = ncon_after - ncon_before;
       if (n > 1) {
         mj_markStack(d);
-        mjContact* buf = (mjContact*)mj_stackAllocByte(d, n * sizeof(mjContact),
-                                                       _Alignof(mjContact));
+        mjContact* buf = mjSTACKALLOC(d, n, mjContact);
         contactSort(d->contact + ncon_before, buf, n, (void*)m);
         mj_freeStack(d);
       }
@@ -496,15 +495,6 @@ struct mjCollisionTree_ {
 typedef struct mjCollisionTree_ mjCollisionTree;
 
 
-
-// collision tree allocation
-static mjCollisionTree* mj_stackAllocTree(mjData* d, int max_stack) {
-  return (mjCollisionTree*) mj_stackAllocByte(
-    d, max_stack * sizeof(mjCollisionTree), _Alignof(mjCollisionTree));
-}
-
-
-
 // checks if the proposed collision pair is already present in pair_geom and calls narrow phase
 void mj_collideGeomPair(const mjModel* m, mjData* d, int g1, int g2, int merged,
                         int startadr, int pairadr) {
@@ -667,7 +657,7 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
   // TODO(b/273737633): Store bvh max depths to make this bound tighter.
   const int max_stack = (isbody1 ? m->body_bvhnum[bf1] : m->flex_bvhnum[f1]) +
                         (isbody2 ? m->body_bvhnum[bf2] : m->flex_bvhnum[f2]);
-  mjCollisionTree* stack = mj_stackAllocTree(d, max_stack);
+  mjCollisionTree* stack = mjSTACKALLOC(d, max_stack, mjCollisionTree);
 
   int nstack = 1;
   stack[0].node1 = stack[0].node2 = 0;
@@ -1037,8 +1027,8 @@ static int mj_SAP(mjData* d, const mjtNum* aamm, int n, int axis, int* pair, int
   }
 
   // allocate sort buffer
-  mjtSAP* sortbuf = (mjtSAP*) mj_stackAllocByte(d, 2*n*sizeof(mjtSAP), _Alignof(mjtSAP));
-  mjtSAP* activebuf = (mjtSAP*) mj_stackAllocByte(d, 2*n*sizeof(mjtSAP), _Alignof(mjtSAP));
+  mjtSAP* sortbuf = mjSTACKALLOC(d, 2*n, mjtSAP);
+  mjtSAP* activebuf = mjSTACKALLOC(d, 2*n, mjtSAP);
 
   // init sortbuf with specified axis
   for (int i=0; i < n; i++) {
@@ -1049,7 +1039,7 @@ static int mj_SAP(mjData* d, const mjtNum* aamm, int n, int axis, int* pair, int
   }
 
   // sort along specified axis
-  mjtSAP* buf = (mjtSAP*) mj_stackAllocByte(d, 2*n*sizeof(mjtSAP), _Alignof(mjtSAP));
+  mjtSAP* buf = mjSTACKALLOC(d, 2*n, mjtSAP);
   SAPsort(sortbuf, buf, 2*n, NULL);
 
   // define the other two axes
@@ -1235,7 +1225,7 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
 
   // allocate collidable bodyflex ids, construct list
   mj_markStack(d);
-  int* bfid = mj_stackAllocInt(d, nbodyflex);
+  int* bfid = mjSTACKALLOC(d, nbodyflex, int);
   int ncollide = 0;
   for (int i=1; i < nbodyflex; i++) {
     if (canCollide(m, i)) {
@@ -1245,14 +1235,14 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
 
   if (ncollide > 1) {
     // allocate and construct AAMMs for collidable only
-    mjtNum* aamm = mj_stackAllocNum(d, 6*ncollide);
+    mjtNum* aamm = mjSTACKALLOC(d, 6*ncollide, mjtNum);
     for (int i=0; i < ncollide; i++) {
       makeAAMM(m, d, aamm+6*i, bfid[i], frame);
     }
 
     // call SAP
     int maxsappair = ncollide*(ncollide-1)/2;
-    int* sappair = mj_stackAllocInt(d, maxsappair);
+    int* sappair = mjSTACKALLOC(d, maxsappair, int);
     int nsappair = mj_SAP(d, aamm, ncollide, 0, sappair, maxsappair);
     if (nsappair < 0) {
       mjERROR("SAP failed");
@@ -1283,7 +1273,7 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
 
   // sort bodyflex pairs by signature
   if (npair > 1) {
-    int* buf = mj_stackAllocInt(d, npair);
+    int* buf = mjSTACKALLOC(d, npair, int);
     bfsort(bfpair, buf, npair, NULL);
   }
 
@@ -1800,7 +1790,7 @@ void mj_collideFlexSAP(const mjModel* m, mjData* d, int f) {
   mj_markStack(d);
 
   // allocate and construct active element ids
-  int* elid = mj_stackAllocInt(d, m->flex_elemnum[f]);
+  int* elid = mjSTACKALLOC(d, m->flex_elemnum[f], int);
   int nactive = 0;
   int flex_elemnum = m->flex_elemnum[f];
   for (int i=0; i < flex_elemnum; i++) {
@@ -1816,7 +1806,7 @@ void mj_collideFlexSAP(const mjModel* m, mjData* d, int f) {
   }
 
   // allocate and construct AAMMs for active elements
-  mjtNum* aamm = mj_stackAllocNum(d, 6*nactive);
+  mjtNum* aamm = mjSTACKALLOC(d, 6*nactive, mjtNum);
   const mjtNum* elemaabb = d->flexelem_aabb + 6*m->flex_elemadr[f];
   for (int i=0; i < nactive; i++) {
     mju_sub3(aamm+6*i+0, elemaabb+6*elid[i], elemaabb+6*elid[i]+3);
@@ -1829,7 +1819,7 @@ void mj_collideFlexSAP(const mjModel* m, mjData* d, int f) {
 
   // call SAP; hard limit on number of pairs to avoid out-of-memory
   int maxsappair = mjMIN(nactive*(nactive-1)/2, 1000000);
-  int* sappair = mj_stackAllocInt(d, maxsappair);
+  int* sappair = mjSTACKALLOC(d, maxsappair, int);
   int nsappair = mj_SAP(d, aamm, nactive, axis, sappair, maxsappair);
   if (nsappair < 0) {
     mjERROR("SAP failed");
diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c
index 1be326de..e8c4653f 100644
--- a/src/engine/engine_collision_gjk.c
+++ b/src/engine/engine_collision_gjk.c
@@ -1291,8 +1291,8 @@ static mjtNum epa(mjCCDStatus* status, Polytope* pt, mjCCDObj* obj1, mjCCDObj* o
   // initialize horizon
   Horizon h;
   mj_markStack(d);
-  h.indices = mj_stackAllocInt(d, 6 + status->max_iterations);
-  h.edges = mj_stackAllocInt(d, 6 + status->max_iterations);
+  h.indices = mjSTACKALLOC(d, 6 + status->max_iterations, int);
+  h.edges = mjSTACKALLOC(d, 6 + status->max_iterations, int);
   h.nedges = 0;
   h.pt = pt;
 
@@ -1483,9 +1483,9 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
     pt.nfaces = pt.nmap = pt.nverts = 0;
 
     // allocate memory for vertices
-    pt.verts  = mj_stackAllocNum(d, 3*(5 + N));
-    pt.verts1 = mj_stackAllocNum(d, 3*(5 + N));
-    pt.verts2 = mj_stackAllocNum(d, 3*(5 + N));
+    pt.verts  = mjSTACKALLOC(d, 3*(5 + N), mjtNum);
+    pt.verts1 = mjSTACKALLOC(d, 3*(5 + N), mjtNum);
+    pt.verts2 = mjSTACKALLOC(d, 3*(5 + N), mjtNum);
 
     // allocate memory for faces
     pt.maxfaces = (6*N > 1000) ? 6*N : 1000;  // use 1000 faces as lower bound
@@ -1497,11 +1497,9 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
     size_t max_size = mj_stackBytesAvailable(d) - 12*(N * sizeof(int));
     if (size1 + size2 > max_size) {
       pt.maxfaces = max_size / (sizeof(Face) + sizeof(Face*));
-      size1 = sizeof(Face) * pt.maxfaces;
-      size2 = sizeof(Face*) * pt.maxfaces;
     }
-    pt.faces = mj_stackAllocByte(d, size1, _Alignof(Face));
-    pt.map = mj_stackAllocByte(d, size2, _Alignof(Face*));
+    pt.faces = mjSTACKALLOC(d, pt.maxfaces, Face);
+    pt.map = mjSTACKALLOC(d, pt.maxfaces, Face*);
 
     int ret;
     if (status->nsimplex == 2) {
diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c
index e01de86f..c6fe8eba 100644
--- a/src/engine/engine_collision_sdf.c
+++ b/src/engine/engine_collision_sdf.c
@@ -516,9 +516,7 @@ static void collideBVH(const mjModel* m, mjData* d, int g,
     int node;
   };
   typedef struct CollideTreeArgs_ CollideTreeArgs;
-  CollideTreeArgs* stack = (CollideTreeArgs*) mj_stackAllocByte(
-    d, max_stack * sizeof(CollideTreeArgs), _Alignof(CollideTreeArgs));
-
+  CollideTreeArgs* stack = mjSTACKALLOC(d, max_stack, CollideTreeArgs);
   int nstack = 0;
   stack[nstack].node = 0;
   nstack++;
diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
index 2c62248a..8d380e9e 100644
--- a/src/engine/engine_core_constraint.c
+++ b/src/engine/engine_core_constraint.c
@@ -485,14 +485,14 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // allocate space
-  jac[0] = mj_stackAllocNum(d, 6*nv);
-  jac[1] = mj_stackAllocNum(d, 6*nv);
-  jacdif = mj_stackAllocNum(d, 6*nv);
+  jac[0] = mjSTACKALLOC(d, 6*nv, mjtNum);
+  jac[1] = mjSTACKALLOC(d, 6*nv, mjtNum);
+  jacdif = mjSTACKALLOC(d, 6*nv, mjtNum);
   if (issparse) {
-    chain = mj_stackAllocInt(d, nv);
-    chain2 = mj_stackAllocInt(d, nv);
-    buf_ind = mj_stackAllocInt(d, nv);
-    sparse_buf = mj_stackAllocNum(d, nv);
+    chain = mjSTACKALLOC(d, nv, int);
+    chain2 = mjSTACKALLOC(d, nv, int);
+    buf_ind = mjSTACKALLOC(d, nv, int);
+    sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
   }
 
   // find active equality constraints
@@ -756,7 +756,7 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // allocate Jacobian
-  jac = mj_stackAllocNum(d, nv);
+  jac = mjSTACKALLOC(d, nv, mjtNum);
 
   // find frictional dofs
   for (int i=0; i < nv; i++) {
@@ -813,7 +813,7 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // allocate Jacobian
-  jac = mj_stackAllocNum(d, nv);
+  jac = mjSTACKALLOC(d, nv, mjtNum);
 
   // find joint limits
   for (int i=0; i < m->njnt; i++) {
@@ -953,16 +953,16 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // allocate Jacobian
-  jac = mj_stackAllocNum(d, 6*nv);
-  jacdif = mj_stackAllocNum(d, 6*nv);
+  jac = mjSTACKALLOC(d, 6*nv, mjtNum);
+  jacdif = mjSTACKALLOC(d, 6*nv, mjtNum);
   jacdifp = jacdif;
   jacdifr = jacdif + 3*nv;
-  jac1p = mj_stackAllocNum(d, 3*nv);
-  jac2p = mj_stackAllocNum(d, 3*nv);
-  jac1r = mj_stackAllocNum(d, 3*nv);
-  jac2r = mj_stackAllocNum(d, 3*nv);
+  jac1p = mjSTACKALLOC(d, 3*nv, mjtNum);
+  jac2p = mjSTACKALLOC(d, 3*nv, mjtNum);
+  jac1r = mjSTACKALLOC(d, 3*nv, mjtNum);
+  jac2r = mjSTACKALLOC(d, 3*nv, mjtNum);
   if (issparse) {
-    chain = mj_stackAllocInt(d, nv);
+    chain = mjSTACKALLOC(d, nv, int);
   }
 
   // find contacts to be included
@@ -1589,8 +1589,8 @@ static int mj_jacSumCount(const mjModel* m, mjData* d, int* chain,
   int nv = m->nv, NV;
 
   mj_markStack(d);
-  int* bodychain = mj_stackAllocInt(d, nv);
-  int* tempchain = mj_stackAllocInt(d, nv);
+  int* bodychain = mjSTACKALLOC(d, nv, int);
+  int* tempchain = mjSTACKALLOC(d, nv, int);
 
   // set first
   NV = mj_bodyChain(m, body[0], chain);
@@ -1643,8 +1643,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
   mj_markStack(d);
 
   if (nnz) {
-    chain = mj_stackAllocInt(d, nv);
-    chain2 = mj_stackAllocInt(d, nv);
+    chain = mjSTACKALLOC(d, nv, int);
+    chain2 = mjSTACKALLOC(d, nv, int);
   }
 
   // find active equality constraints
@@ -1870,7 +1870,7 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
   }
 
   mj_markStack(d);
-  int *chain = mj_stackAllocInt(d, m->nv);
+  int *chain = mjSTACKALLOC(d, m->nv, int);
 
   for (int i=0; i < ncon; i++) {
     mjContact* con = d->contact + i;
@@ -2068,19 +2068,19 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // space for backsubM2(J')' and its traspose
-  mjtNum* JM2 = mj_stackAllocNum(d, nefc*nv);
-  mjtNum* JM2T = mj_stackAllocNum(d, nv*nefc);
+  mjtNum* JM2 = mjSTACKALLOC(d, nefc*nv, mjtNum);
+  mjtNum* JM2T = mjSTACKALLOC(d, nv*nefc, mjtNum);
 
   // sparse
   if (mj_isSparse(m)) {
     // space for JM2 and JM2T indices
-    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);
+    int* rownnz = mjSTACKALLOC(d, nefc, int);
+    int* rowadr = mjSTACKALLOC(d, nefc, int);
+    int* colind = mjSTACKALLOC(d, nefc*nv, int);
+    int* rowsuper = mjSTACKALLOC(d, nefc, int);
+    int* rownnzT = mjSTACKALLOC(d, nv, int);
+    int* rowadrT = mjSTACKALLOC(d, nv, int);
+    int* colindT = mjSTACKALLOC(d, nv*nefc, int);
 
     // construct JM2 = backsubM2(J')' by rows
     for (int r=0; r < nefc; r++) {
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index a11895b5..ca16cb30 100644
--- a/src/engine/engine_core_smooth.c
+++ b/src/engine/engine_core_smooth.c
@@ -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);
diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c
index 15d48149..d4c764cb 100644
--- a/src/engine/engine_derivative.c
+++ b/src/engine/engine_derivative.c
@@ -395,11 +395,11 @@ void mjd_rne_vel_dense(const mjModel* m, mjData* d) {
   mjtNum mat[36], mat1[36], mat2[36], dmul[36], tmp[6];
 
   mj_markStack(d);
-  mjtNum* Dcvel = mj_stackAllocNum(d, nbody*6*nv);
-  mjtNum* Dcdofdot = mj_stackAllocNum(d, nv*6*nv);
-  mjtNum* Dcacc = mj_stackAllocNum(d, nbody*6*nv);
-  mjtNum* Dcfrcbody = mj_stackAllocNum(d, nbody*6*nv);
-  mjtNum* row = mj_stackAllocNum(d, nv);
+  mjtNum* Dcvel = mjSTACKALLOC(d, nbody*6*nv, mjtNum);
+  mjtNum* Dcdofdot = mjSTACKALLOC(d, nv*6*nv, mjtNum);
+  mjtNum* Dcacc = mjSTACKALLOC(d, nbody*6*nv, mjtNum);
+  mjtNum* Dcfrcbody = mjSTACKALLOC(d, nbody*6*nv, mjtNum);
+  mjtNum* row = mjSTACKALLOC(d, nv, mjtNum);
 
   // compute Dcvel and Dcdofdot
   mjd_comVel_vel_dense(m, d, Dcvel, Dcdofdot);
@@ -610,11 +610,11 @@ static void mjd_rne_vel(const mjModel* m, mjData* d) {
   mjtNum mat[36], mat1[36], mat2[36], dmul[36], tmp[6];
 
   mj_markStack(d);
-  mjtNum* Dcdofdot = mj_stackAllocNum(d, 6*m->nD);
-  mjtNum* Dcvel = mj_stackAllocNum(d, 6*m->nB);
-  mjtNum* Dcacc = mj_stackAllocNum(d, 6*m->nB);
-  mjtNum* Dcfrcbody = mj_stackAllocNum(d, 6*m->nB);
-  mjtNum* row = mj_stackAllocNum(d, nv);
+  mjtNum* Dcdofdot = mjSTACKALLOC(d, 6*m->nD, mjtNum);
+  mjtNum* Dcvel = mjSTACKALLOC(d, 6*m->nB, mjtNum);
+  mjtNum* Dcacc = mjSTACKALLOC(d, 6*m->nB, mjtNum);
+  mjtNum* Dcfrcbody = mjSTACKALLOC(d, 6*m->nB, mjtNum);
+  mjtNum* row = mjSTACKALLOC(d, nv, mjtNum);
 
   // clear
   mju_zero(Dcdofdot, 6*m->nD);
@@ -695,7 +695,7 @@ static void addJTBJ(const mjModel* m, mjData* d, const mjtNum* J, const mjtNum*
 
   // allocate dense row
   mj_markStack(d);
-  mjtNum* row = mj_stackAllocNum(d, nv);
+  mjtNum* row = mjSTACKALLOC(d, nv, mjtNum);
 
   // process non-zero elements of B
   for (int i=0; i < n; i++) {
@@ -734,7 +734,7 @@ static void addJTBJSparse(
 
   // allocate row
   mj_markStack(d);
-  mjtNum* row = mj_stackAllocNum(d, nv);
+  mjtNum* row = mjSTACKALLOC(d, nv, mjtNum);
 
   // compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
   for (int i = 0; i < n; i++) {
@@ -829,7 +829,7 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
 
   // allocate dense actuator_moment row
   mj_markStack(d);
-  mjtNum* moment = mj_stackAllocNum(d, nv);
+  mjtNum* moment = mjSTACKALLOC(d, nv, mjtNum);
 
   // process actuators
   for (int i=0; i < nu; i++) {
@@ -1181,10 +1181,10 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) {
   int nv = m->nv;
   int nnz = nv;
   int rownnz[6], rowadr[6];
-  mjtNum* J = mj_stackAllocNum(d, 6*nv);
-  mjtNum* tmp = mj_stackAllocNum(d, 3*nv);
-  int* colind = mj_stackAllocInt(d, 6*nv);
-  int* colind_compressed = mj_stackAllocInt(d, 6*nv);
+  mjtNum* J = mjSTACKALLOC(d, 6*nv, mjtNum);
+  mjtNum* tmp = mjSTACKALLOC(d, 3*nv, mjtNum);
+  int* colind = mjSTACKALLOC(d, 6*nv, int);
+  int* colind_compressed = mjSTACKALLOC(d, 6*nv, int);
 
   mjtNum lvel[6], wind[6], lwind[6];
   mjtNum geom_interaction_coef, magnus_lift_coef, kutta_lift_coef;
@@ -1287,9 +1287,9 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, int i) {
 
   int nv = m->nv;
   int rownnz[6], rowadr[6];
-  mjtNum* J = mj_stackAllocNum(d, 6*nv);
-  mjtNum* tmp = mj_stackAllocNum(d, 3*nv);
-  int* colind = mj_stackAllocInt(d, 6*nv);
+  mjtNum* J = mjSTACKALLOC(d, 6*nv, mjtNum);
+  mjtNum* tmp = mjSTACKALLOC(d, 3*nv, mjtNum);
+  int* colind = mjSTACKALLOC(d, 6*nv, int);
 
   mjtNum lvel[6], wind[6], lwind[6], box[3], B;
   mjtNum* inertia = m->body_inertia + 3*i;
diff --git a/src/engine/engine_derivative_fd.c b/src/engine/engine_derivative_fd.c
index 239779af..b5741232 100644
--- a/src/engine/engine_derivative_fd.c
+++ b/src/engine/engine_derivative_fd.c
@@ -176,9 +176,9 @@ void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps) {
   int nv = m->nv;
 
   mj_markStack(d);
-  mjtNum* qfrc_passive = mj_stackAllocNum(d, nv);
-  mjtNum* fd = mj_stackAllocNum(d, nv);
-  int* cnt = mj_stackAllocInt(d, nv);
+  mjtNum* qfrc_passive = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* fd = mjSTACKALLOC(d, nv, mjtNum);
+  int* cnt = mjSTACKALLOC(d, nv, int);
 
   // clear row counters
   mju_zeroInt(cnt, nv);
@@ -227,10 +227,10 @@ void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) {
   int nv = m->nv;
 
   mj_markStack(d);
-  mjtNum* plus = mj_stackAllocNum(d, nv);
-  mjtNum* minus = mj_stackAllocNum(d, nv);
-  mjtNum* fd = mj_stackAllocNum(d, nv);
-  int* cnt = mj_stackAllocInt(d, nv);
+  mjtNum* plus = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* minus = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* fd = mjSTACKALLOC(d, nv, mjtNum);
+  int* cnt = mjSTACKALLOC(d, nv, int);
 
   // clear row counters
   mju_zeroInt(cnt, nv);
@@ -314,20 +314,20 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
   unsigned int restore_spec = mjSTATE_FULLPHYSICS | mjSTATE_CTRL;
   restore_spec |= mjDISABLED(mjDSBL_WARMSTART) ? 0 : mjSTATE_WARMSTART;
 
-  mjtNum *fullstate  = mj_stackAllocNum(d, mj_stateSize(m, restore_spec));
-  mjtNum *state      = mj_stackAllocNum(d, nq+nv+na);  // current state
-  mjtNum *next       = mj_stackAllocNum(d, nq+nv+na);  // next state
-  mjtNum *next_plus  = mj_stackAllocNum(d, nq+nv+na);  // forward-nudged next state
-  mjtNum *next_minus = mj_stackAllocNum(d, nq+nv+na);  // backward-nudged next state
+  mjtNum *fullstate  = mjSTACKALLOC(d, mj_stateSize(m, restore_spec), mjtNum);
+  mjtNum *state      = mjSTACKALLOC(d, nq+nv+na, mjtNum);  // current state
+  mjtNum *next       = mjSTACKALLOC(d, nq+nv+na, mjtNum);  // next state
+  mjtNum *next_plus  = mjSTACKALLOC(d, nq+nv+na, mjtNum);  // forward-nudged next state
+  mjtNum *next_minus = mjSTACKALLOC(d, nq+nv+na, mjtNum);  // backward-nudged next state
 
   // sensors
   int skipsensor = !DsDq && !DsDv && !DsDa && !DsDu;
-  mjtNum *sensor       = skipsensor ? NULL : mj_stackAllocNum(d, ns);  // sensor values
-  mjtNum *sensor_plus  = skipsensor ? NULL : mj_stackAllocNum(d, ns);  // forward-nudged sensors
-  mjtNum *sensor_minus = skipsensor ? NULL : mj_stackAllocNum(d, ns);  // backward-nudged sensors
+  mjtNum *sensor       = skipsensor ? NULL : mjSTACKALLOC(d, ns, mjtNum);  // sensor values
+  mjtNum *sensor_plus  = skipsensor ? NULL : mjSTACKALLOC(d, ns, mjtNum);  // forward-nudged
+  mjtNum *sensor_minus = skipsensor ? NULL : mjSTACKALLOC(d, ns, mjtNum);  // backward-nudged
 
   // controls
-  mjtNum *ctrl = mj_stackAllocNum(d, nu);
+  mjtNum *ctrl = mjSTACKALLOC(d, nu, mjtNum);
 
   // save current inputs
   mj_getState(m, d, fullstate, restore_spec);
@@ -485,7 +485,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
 
   // finite-difference positions: skip=mjSTAGE_NONE
   if (DyDq || DsDq) {
-    mjtNum *dpos  = mj_stackAllocNum(d, nv);  // allocate position perturbation
+    mjtNum *dpos  = mjSTACKALLOC(d, nv, mjtNum);  // allocate position perturbation
     for (int i=0; i < nv; i++) {
       // nudge forward
       mju_zero(dpos, nv);
@@ -563,10 +563,10 @@ void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_cente
   mj_markStack(d);
 
   // allocate transposed matrices
-  mjtNum *AT = A ? mj_stackAllocNum(d, ndx*ndx) : NULL;  // state-transition matrix   (transposed)
-  mjtNum *BT = B ? mj_stackAllocNum(d, nu*ndx) : NULL;   // control-transition matrix (transposed)
-  mjtNum *CT = C ? mj_stackAllocNum(d, ndx*ns) : NULL;   // state-observation matrix   (transposed)
-  mjtNum *DT = D ? mj_stackAllocNum(d, nu*ns) : NULL;    // control-observation matrix (transposed)
+  mjtNum *AT = A ? mjSTACKALLOC(d, ndx*ndx, mjtNum) : NULL;  // state-transition     (transposed)
+  mjtNum *BT = B ? mjSTACKALLOC(d, nu*ndx, mjtNum) : NULL;   // control-transition   (transposed)
+  mjtNum *CT = C ? mjSTACKALLOC(d, ndx*ns, mjtNum) : NULL;   // state-observation    (transposed)
+  mjtNum *DT = D ? mjSTACKALLOC(d, nu*ns, mjtNum) : NULL;    // control-observation  (transposed)
 
   // set offset pointers
   if (A) {
@@ -629,11 +629,11 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
 
   // local vectors
   mj_markStack(d);
-  mjtNum *pos        = mj_stackAllocNum(d, nq);                      // position
-  mjtNum *force      = mj_stackAllocNum(d, nv);                      // force
-  mjtNum *force_plus = mj_stackAllocNum(d, nv);                      // nudged force
-  mjtNum *sensor     = skipsensor ? NULL : mj_stackAllocNum(d, ns);  // sensor values
-  mjtNum *mass       = DmDq ? mj_stackAllocNum(d, nM) : NULL;        // mass matrix
+  mjtNum *pos        = mjSTACKALLOC(d, nq, mjtNum);                      // position
+  mjtNum *force      = mjSTACKALLOC(d, nv, mjtNum);                      // force
+  mjtNum *force_plus = mjSTACKALLOC(d, nv, mjtNum);                      // nudged force
+  mjtNum *sensor     = skipsensor ? NULL : mjSTACKALLOC(d, ns, mjtNum);  // sensor values
+  mjtNum *mass       = DmDq ? mjSTACKALLOC(d, nM, mjtNum) : NULL;        // mass matrix
 
   // save current positions
   mju_copy(pos, d->qpos, nq);
@@ -687,7 +687,7 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
 
   // position: skip = mjSTAGE_NONE
   if (DfDq || DsDq || DmDq) {
-    mjtNum *dpos  = mj_stackAllocNum(d, nv);  // allocate position perturbation
+    mjtNum *dpos  = mjSTACKALLOC(d, nv, mjtNum);  // allocate position perturbation
     for (int i=0; i < nv; i++) {
       // nudge
       mju_zero(dpos, nv);
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index 181db4d3..b90ba5d8 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -286,7 +286,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
 
   // local, clamped copy of ctrl
   mj_markStack(d);
-  mjtNum *ctrl = mj_stackAllocNum(d, nu);
+  mjtNum *ctrl = mjSTACKALLOC(d, nu, mjtNum);
   mju_copy(ctrl, d->ctrl, nu);
   if (!mjDISABLED(mjDSBL_CLAMPCTRL)) {
     clampVec(ctrl, m->actuator_ctrlrange, m->actuator_ctrllimited, nu, NULL);
@@ -531,7 +531,7 @@ static void warmstart(const mjModel* m, mjData* d) {
   // warmstart with best of (qacc_warmstart, qacc_smooth)
   if (!mjDISABLED(mjDSBL_WARMSTART)) {
     mj_markStack(d);
-    mjtNum* jar = mj_stackAllocNum(d, nefc);
+    mjtNum* jar = mjSTACKALLOC(d, nefc, mjtNum);
 
     // start with qacc = qacc_warmstart
     mju_copy(d->qacc, d->qacc_warmstart, nv);
@@ -548,7 +548,7 @@ static void warmstart(const mjModel* m, mjData* d) {
     if (m->opt.solver == mjSOL_PGS) {
       // cost(force_warmstart)
       mjtNum PGS_warmstart = mju_dot(d->efc_force, d->efc_b, nefc);
-      mjtNum* ARf = mj_stackAllocNum(d, nefc);
+      mjtNum* ARf = mjSTACKALLOC(d, nefc, mjtNum);
       if (mj_isSparse(m))
         mju_mulMatVecSparse(ARf, d->efc_AR, d->efc_force, nefc,
                             d->efc_AR_rownnz, d->efc_AR_rowadr,
@@ -568,7 +568,7 @@ static void warmstart(const mjModel* m, mjData* d) {
     // non-PGS
     else {
       // add Gauss to cost(qacc_warmstart)
-      mjtNum* Ma = mj_stackAllocNum(d, nv);
+      mjtNum* Ma = mjSTACKALLOC(d, nv, mjtNum);
       mj_mulM(m, d, Ma, d->qacc_warmstart);
       for (int i=0; i < nv; i++) {
         cost_warmstart += 0.5*(Ma[i]-d->qfrc_smooth[i])*(d->qacc_warmstart[i]-d->qacc_smooth[i]);
@@ -618,9 +618,9 @@ void* mj_solCG_island_wrapper(void* args) {
 void mj_solCG_island_multithreaded(const mjModel* m, mjData* d) {
   mj_markStack(d);
   // allocate array of arguments to be passed to threads
-  mjSolIslandArgs* sol_cg_island_args =
-    mj_stackAllocByte(d, sizeof(mjSolIslandArgs) * d->nisland, _Alignof(mjSolIslandArgs));
-  mjTask* tasks = mj_stackAllocByte(d, sizeof(mjTask)  * d->nisland, _Alignof(mjTask));
+  mjSolIslandArgs* sol_cg_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs);
+  mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask);
+
 
   for (int island = 0; island < d->nisland; ++island) {
     sol_cg_island_args[island].m = m;
@@ -772,8 +772,8 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
   int nv = m->nv, nM = m->nM;
   mj_markStack(d);
-  mjtNum* qfrc = mj_stackAllocNum(d, nv);
-  mjtNum* qacc = mj_stackAllocNum(d, nv);
+  mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
 
   // check for dof damping if disable flag is not set
   int dof_damping = 0;
@@ -794,7 +794,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   // damping: integrate implicitly
   else {
     if (!skipfactor) {
-      mjtNum* MhB = mj_stackAllocNum(d, nM);
+      mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum);
 
       // MhB = M + h*diag(B)
       mju_copy(MhB, d->qM, nM);
@@ -857,10 +857,10 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 
   // allocate space for intermediate solutions
   mj_markStack(d);
-  dX = mj_stackAllocNum(d, 2*nv+na);
+  dX = mjSTACKALLOC(d, 2*nv+na, mjtNum);
   for (int i=0; i < N; i++) {
-    X[i] = mj_stackAllocNum(d, nq+nv+na);
-    F[i] = mj_stackAllocNum(d, nv+na);
+    X[i] = mjSTACKALLOC(d, nq+nv+na, mjtNum);
+    F[i] = mjSTACKALLOC(d, nv+na, mjtNum);
   }
 
   // precompute C and T;  C,T,A have size (N-1)
@@ -941,8 +941,8 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   int nv = m->nv, nM = m->nM, nD = m->nD;
 
   mj_markStack(d);
-  mjtNum* qfrc = mj_stackAllocNum(d, nv);
-  mjtNum* qacc = mj_stackAllocNum(d, nv);
+  mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
 
   // set qfrc = qfrc_smooth + qfrc_constraint
   mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
@@ -962,7 +962,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
 
       // factorize qLU
-      int* scratch = mj_stackAllocInt(d, nv);
+      int* scratch = mjSTACKALLOC(d, nv, int);
       mju_factorLUSparse(d->qLU, nv, scratch, d->D_rownnz, d->D_rowadr, d->D_colind);
     }
 
@@ -977,7 +977,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mjd_smooth_vel(m, d, /* flg_bias = */ 0);
 
       // modified mass matrix MhB = qDeriv[Lower]
-      mjtNum* MhB = mj_stackAllocNum(d, nM);
+      mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum);
       for (int i=0; i < nM; i++) {
         MhB[i] = d->qDeriv[d->mapD2M[i]];
       }
diff --git a/src/engine/engine_inverse.c b/src/engine/engine_inverse.c
index f582728e..f372947a 100644
--- a/src/engine/engine_inverse.c
+++ b/src/engine/engine_inverse.c
@@ -76,7 +76,7 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
   mjtNum *qacc = d->qacc;
 
   mj_markStack(d);
-  mjtNum* qfrc = mj_stackAllocNum(d, nv);
+  mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
 
   // use selected integrator
   switch ((mjtIntegrator) m->opt.integrator) {
@@ -132,11 +132,11 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
     mjd_smooth_vel(m, d, /* flg_bias = */ 0);
 
     // save mass matrix
-    mjtNum* qMsave = mj_stackAllocNum(d, m->nM);
+    mjtNum* qMsave = mjSTACKALLOC(d, m->nM, mjtNum);
     mju_copy(qMsave, d->qM, m->nM);
 
     // set M = M - dt*qDeriv (reduced to M nonzeros)
-    mjtNum* qDerivReduced = mj_stackAllocNum(d, m->nM);
+    mjtNum* qDerivReduced = mjSTACKALLOC(d, m->nM, mjtNum);
     for (int i=0; i < nM; i++) {
       qDerivReduced[i] = d->qDeriv[d->mapD2M[i]];
     }
@@ -171,7 +171,7 @@ void mj_invConstraint(const mjModel* m, mjData* d) {
   }
 
   mj_markStack(d);
-  mjtNum* jar = mj_stackAllocNum(d, nefc);
+  mjtNum* jar = mjSTACKALLOC(d, nefc, mjtNum);
 
   // compute jar = Jac*qacc - aref
   mj_mulJacVec(m, d, jar, d->qacc);
@@ -218,7 +218,7 @@ void mj_inverseSkip(const mjModel* m, mjData* d,
 
   if (mjENABLED(mjENBL_INVDISCRETE)) {
     // save current qacc
-    qacc = mj_stackAllocNum(d, nv);
+    qacc = mjSTACKALLOC(d, nv, mjtNum);
     mju_copy(qacc, d->qacc, nv);
 
     // modify qacc in-place
@@ -271,10 +271,10 @@ void mj_compareFwdInv(const mjModel* m, mjData* d) {
 
   // allocate
   mj_markStack(d);
-  qforce = mj_stackAllocNum(d, nv);
-  dif = mj_stackAllocNum(d, nv);
-  save_qfrc_constraint = mj_stackAllocNum(d, nv);
-  save_efc_force = mj_stackAllocNum(d, nefc);
+  qforce = mjSTACKALLOC(d, nv, mjtNum);
+  dif = mjSTACKALLOC(d, nv, mjtNum);
+  save_qfrc_constraint = mjSTACKALLOC(d, nv, mjtNum);
+  save_efc_force = mjSTACKALLOC(d, nefc, mjtNum);
 
   // qforce = qfrc_applied + J'*xfrc_applied + qfrc_actuator
   //  should equal result of inverse dynamics
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 59704140..34cc8dbb 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -931,7 +931,7 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
   }
 
   mj_markStack(d);
-  int* remaining = mj_stackAllocInt(d, nv);
+  int* remaining = mjSTACKALLOC(d, nv, int);
 
   // compute rownnz
   mju_zeroInt(rownnz, nv);
@@ -1044,7 +1044,7 @@ static void makeBSparse(const mjModel* m, mjData* d) {
 
   // allocate and clear incremental row counts
   mj_markStack(d);
-  int* cnt = mj_stackAllocInt(d, nbody);
+  int* cnt = mjSTACKALLOC(d, nbody, int);
   mju_zeroInt(cnt, nbody);
 
   // add subtree dofs to colind
@@ -1134,7 +1134,7 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
   mj_markStack(d);
 
   // init remaining
-  int* remaining = mj_stackAllocInt(d, nv);
+  int* remaining = mjSTACKALLOC(d, nv, int);
   mju_copyInt(remaining, rownnz, nv);
 
   // copy data
@@ -1202,7 +1202,7 @@ static void makeDmap(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // make mapM2D
-  int* M = mj_stackAllocInt(d, nM);
+  int* M = mjSTACKALLOC(d, nM, int);
   for (int i=0; i < nM; i++) M[i] = i;
   for (int i=0; i < nD; i++) d->mapM2D[i] = -1;
   copyM2Sparse(m, d, d->mapM2D, M, /*reduced=*/0);
@@ -1215,7 +1215,7 @@ static void makeDmap(const mjModel* m, mjData* d) {
   }
 
   // make mapD2M
-  int* D = mj_stackAllocInt(d, nD);
+  int* D = mjSTACKALLOC(d, nD, int);
   for (int i=0; i < nD; i++) D[i] = i;
   for (int i=0; i < nM; i++) d->mapD2M[i] = -1;
   copyD2MSparse(m, d, d->mapD2M, D);
@@ -1591,7 +1591,8 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
 
 // internal: allocate size bytes on the provided stack shard
 // declared inline so that modular arithmetic with specific alignments can be optimized out
-static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size, size_t alignment) {
+static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_t size,
+    size_t alignment, const char* caller, int line) {
   // return NULL if empty
   if (mjUNLIKELY(!size)) {
     return NULL;
@@ -1614,10 +1615,19 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
   size_t stack_available_bytes = stack_info->top - stack_info->limit;
   size_t stack_required_bytes = stack_info->top - new_top_ptr;
   if (mjUNLIKELY(stack_required_bytes > stack_available_bytes)) {
-    mju_error("mj_stackAlloc: insufficient memory: max = %zu, available = %zu, requested = %zu "
-              "(ne = %d, nf = %d, nefc = %d, ncon = %d)",
+    char info[1024];
+    if (caller) {
+      snprintf(info, sizeof(info), " at %s, line %d", caller, line);
+    } else {
+      info[0] = '\0';
+    }
+    mju_error("mj_stackAlloc: out of memory, stack overflow%s\n"
+              "  max = %zu, available = %zu, requested = %zu\n"
+              "  nefc = %d, ncon = %d",
+              info,
               stack_info->bottom - stack_info->limit, stack_available_bytes, stack_required_bytes,
-              d->ne, d->nf, d->nefc, d->ncon);
+              d->nefc, d->ncon);
+
   }
 
 #ifdef ADDRESS_SANITIZER
@@ -1652,20 +1662,20 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
 
 // internal: allocate size bytes in mjData
 // declared inline so that modular arithmetic with specific alignments can be optimized out
-static inline void* stackalloc(mjData* d, size_t size, size_t alignment) {
+static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
+                               const char* caller, int line) {
+  // single threaded allocation
   if (!d->threadpool) {
     mjStackInfo stack_info = get_stack_info_from_data(d);
-
-    void* result = stackallocinternal(d, &stack_info, size, alignment);
-
+    void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
     d->pstack = stack_info.bottom - stack_info.top;
-
     return result;
   }
 
+  // multi threaded allocation
   size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
   mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  return stackallocinternal(d, stack_info, size, alignment);
+  return stackallocinternal(d, stack_info, size, alignment, caller, line);
 }
 
 
@@ -1677,7 +1687,7 @@ __attribute__((always_inline))
 static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
   size_t top_old = stack_info->top;
   mjStackFrame* s =
-    (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame));
+    (mjStackFrame*) stackallocinternal(d, stack_info, sizeof(mjStackFrame), _Alignof(mjStackFrame), NULL, 0);
   s->pbase = stack_info->stack_base;
   s->pstack = top_old;
 #ifdef ADDRESS_SANITIZER
@@ -1779,7 +1789,15 @@ size_t mj_stackBytesAvailable(mjData* d) {
 
 // allocate bytes on the stack
 void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment) {
-  return stackalloc(d, bytes, alignment);
+  return stackalloc(d, bytes, alignment, NULL, 0);
+}
+
+
+
+// allocate bytes on the stack, with caller information
+void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
+                        const char* caller, int line) {
+  return stackalloc(d, bytes, alignment, caller, line);
 }
 
 
@@ -1789,7 +1807,7 @@ mjtNum* mj_stackAllocNum(mjData* d, size_t size) {
   if (mjUNLIKELY(size >= SIZE_MAX / sizeof(mjtNum))) {
     mjERROR("requested size is too large (more than 2^64 bytes).");
   }
-  return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum));
+  return (mjtNum*) stackalloc(d, size * sizeof(mjtNum), _Alignof(mjtNum), NULL, 0);
 }
 
 
@@ -1799,7 +1817,7 @@ int* mj_stackAllocInt(mjData* d, size_t size) {
   if (mjUNLIKELY(size >= SIZE_MAX / sizeof(int))) {
     mjERROR("requested size is too large (more than 2^64 bytes).");
   }
-  return (int*) stackalloc(d, size * sizeof(int), _Alignof(int));
+  return (int*) stackalloc(d, size * sizeof(int), _Alignof(int), NULL, 0);
 }
 
 
diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h
index a6fac94e..298452fa 100644
--- a/src/engine/engine_io.h
+++ b/src/engine/engine_io.h
@@ -134,9 +134,17 @@ void mj__freeStack(mjData* d) __attribute__((noinline));
 // returns the number of bytes available on the stack
 MJAPI size_t mj_stackBytesAvailable(mjData* d);
 
-// mjData stack allocate
+// allocate bytes on the stack
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
 
+// allocate bytes on the stack, with added caller information
+MJAPI void* mj_stackAllocInfo(mjData* d, size_t bytes, size_t alignment,
+                              const char* caller, int line);
+
+// macro to allocate a stack array of given type, adds caller information
+#define mjSTACKALLOC(d, num, type) \
+(type*) mj_stackAllocInfo(d, (num) * sizeof(type), _Alignof(type), __func__, __LINE__)
+
 // mjData stack allocate for array of mjtNums
 MJAPI mjtNum* mj_stackAllocNum(mjData* d, size_t size);
 
diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c
index 26d67213..9a8be763 100644
--- a/src/engine/engine_island.c
+++ b/src/engine/engine_island.c
@@ -426,14 +426,14 @@ void mj_island(const mjModel* m, mjData* d) {
 
   // allocate edge array
   int nedge_max = countMaxEdge(m, d);
-  int* edge = mj_stackAllocInt(d, 2*nedge_max);
+  int* edge = mjSTACKALLOC(d, 2*nedge_max, int);
 
   // get tree-tree edges and rownnz counts from efc arrays
-  int* rownnz = mj_stackAllocInt(d, ntree);  // number of edges per tree
+  int* rownnz = mjSTACKALLOC(d, ntree, int);  // number of edges per tree
   int nedge = findEdges(m, d, rownnz, edge, nedge_max);
 
   // compute starting address of tree's column indices while resetting rownnz
-  int* rowadr = mj_stackAllocInt(d, ntree);
+  int* rowadr = mjSTACKALLOC(d, ntree, int);
   rowadr[0] = 0;
   for (int r=1; r < ntree; r++) {
     rowadr[r] = rowadr[r-1] + rownnz[r-1];
@@ -442,7 +442,7 @@ void mj_island(const mjModel* m, mjData* d) {
   rownnz[ntree-1] = 0;
 
   // copy column indices: list each tree's neighbors
-  int* colind = mj_stackAllocInt(d, nedge);
+  int* colind = mjSTACKALLOC(d, nedge, int);
   for (int e=0; e < nedge; e++) {
     int row = edge[2*e];
     int col = edge[2*e + 1];
@@ -450,8 +450,8 @@ void mj_island(const mjModel* m, mjData* d) {
   }
 
   // discover islands
-  int* tree_island = mj_stackAllocInt(d, ntree);  // id of island assigned to tree
-  int* stack = mj_stackAllocInt(d, nedge);
+  int* tree_island = mjSTACKALLOC(d, ntree, int);  // id of island assigned to tree
+  int* stack = mjSTACKALLOC(d, nedge, int);
   d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack);
 
   // allocate island arrays on arena
diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c
index 8dcf46b8..3365ac73 100644
--- a/src/engine/engine_passive.c
+++ b/src/engine/engine_passive.c
@@ -134,7 +134,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
     mjtNum kD = m->flex_damping[f] / m->opt.timestep;
 
     mj_markStack(d);
-    mjtNum* qfrc = mj_stackAllocNum(d, 3*m->flex_vertnum[f]);
+    mjtNum* qfrc = mjSTACKALLOC(d, 3*m->flex_vertnum[f], mjtNum);
     mju_zero(qfrc, 3*m->flex_vertnum[f]);
 
     // compute force element-by-element
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 10ad5668..4de41fca 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -892,7 +892,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
 
   // allocate full inertia if it's small
   if (m->nv <= 200) {
-    M = mj_stackAllocNum(d, m->nv*m->nv);
+    M = mjSTACKALLOC(d, m->nv*m->nv, mjtNum);
   }
 
 #ifdef MEMORY_SANITIZER
diff --git a/src/engine/engine_ray.c b/src/engine/engine_ray.c
index d8824bf4..db58e69c 100644
--- a/src/engine/engine_ray.c
+++ b/src/engine/engine_ray.c
@@ -1351,8 +1351,8 @@ void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum*
   mj_markStack(d);
 
   // allocate source
-  mjtNum* geom_ba = mj_stackAllocNum(d, 4*m->ngeom);
-  int* geom_eliminate = mj_stackAllocInt(d, m->ngeom);
+  mjtNum* geom_ba = mjSTACKALLOC(d, 4*m->ngeom, mjtNum);
+  int* geom_eliminate = mjSTACKALLOC(d, m->ngeom, int);
 
   // initialize source
   mju_multiRayPrepare(m, d, pnt, NULL, geomgroup, flg_static, bodyexclude,
diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c
index a82fb303..56966941 100644
--- a/src/engine/engine_sensor.c
+++ b/src/engine/engine_sensor.c
@@ -996,15 +996,13 @@ void mj_energyPos(const mjModel* m, mjData* d) {
 
 // velocity-dependent energy (kinetic)
 void mj_energyVel(const mjModel* m, mjData* d) {
-  mjtNum *vec;
-
   // return if disabled (already cleared in potential)
   if (!mjENABLED(mjENBL_ENERGY)) {
     return;
   }
 
   mj_markStack(d);
-  vec = mj_stackAllocNum(d, m->nv);
+  mjtNum *vec = mjSTACKALLOC(d, m->nv, mjtNum);
 
   // kinetic energy:  0.5 * qvel' * M * qvel
   mj_mulM(m, d, vec, d->qvel);
diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c
index 04fcec99..046ba9d3 100644
--- a/src/engine/engine_setconst.c
+++ b/src/engine/engine_setconst.c
@@ -65,22 +65,22 @@ static void set0(mjModel* m, mjData* d) {
   int nv = m->nv;
   mjtNum A[36] = {0}, pos[3], quat[4];
   mj_markStack(d);
-  mjtNum* jac = mj_stackAllocNum(d, 6*nv);
-  mjtNum* tmp = mj_stackAllocNum(d, 6*nv);
-  mjtNum* moment = mj_stackAllocNum(d, nv);
+  mjtNum* jac = mjSTACKALLOC(d, 6*nv, mjtNum);
+  mjtNum* tmp = mjSTACKALLOC(d, 6*nv, mjtNum);
+  mjtNum* moment = mjSTACKALLOC(d, nv, mjtNum);
   int* cammode = 0;
   int* lightmode = 0;
 
   // save camera and light mode, set to fixed
   if (m->ncam) {
-    cammode = mj_stackAllocInt(d, m->ncam);
+    cammode = mjSTACKALLOC(d, m->ncam, int);
     for (int i=0; i < m->ncam; i++) {
       cammode[i] = m->cam_mode[i];
       m->cam_mode[i] = mjCAMLIGHT_FIXED;
     }
   }
   if (m->nlight) {
-    lightmode = mj_stackAllocInt(d, m->nlight);
+    lightmode = mjSTACKALLOC(d, m->nlight, int);
     for (int i=0; i < m->nlight; i++) {
       lightmode[i] = m->light_mode[i];
       m->light_mode[i] = mjCAMLIGHT_FIXED;
@@ -427,7 +427,7 @@ static void setStat(mjModel* m, mjData* d) {
   mjtNum xmax[3] = {-1E+10, -1E+10, -1E+10};
   mjtNum rbound;
   mj_markStack(d);
-  mjtNum* body = mj_stackAllocNum(d, m->nbody);
+  mjtNum* body = mjSTACKALLOC(d, m->nbody, mjtNum);
 
   // compute bounding box of bodies, joint centers, geoms and sites
   for (int i=1; i < m->nbody; i++) {
@@ -595,7 +595,7 @@ static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
 
   // dense actuator_moment row
   mj_markStack(d);
-  mjtNum* moment = mj_stackAllocNum(d, nv);
+  mjtNum* moment = mjSTACKALLOC(d, nv, mjtNum);
   mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + index,
                    d->moment_rowadr + index, d->moment_colind);
 
diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c
index cb4553ea..ac506fa4 100644
--- a/src/engine/engine_solver.c
+++ b/src/engine/engine_solver.c
@@ -330,8 +330,8 @@ void mj_solPGS(const mjModel* m, mjData* d, int maxiter) {
   const mjtNum *floss = d->efc_frictionloss;
   mjtNum *force = d->efc_force;
   mj_markStack(d);
-  mjtNum* ARinv = mj_stackAllocNum(d, nefc);
-  int* oldstate = mj_stackAllocInt(d, nefc);
+  mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum);
+  int* oldstate = mjSTACKALLOC(d, nefc, int);
 
   // TODO: b/295296178 - Use island index (currently hardcoded to 0)
   int island = 0;
@@ -555,8 +555,8 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
   mjtNum v[5], Ac[25], bc[5], res[5], oldforce[5], delta[5], mid, y, K0, K1;
   mjContact* con;
   mj_markStack(d);
-  mjtNum* ARinv = mj_stackAllocNum(d, nefc);
-  int* oldstate = mj_stackAllocInt(d, nefc);
+  mjtNum* ARinv = mjSTACKALLOC(d, nefc, mjtNum);
+  int* oldstate = mjSTACKALLOC(d, nefc, int);
 
   // TODO: b/295296178 - Use island index (currently hardcoded to 0)
   int island = 0;
@@ -837,28 +837,28 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
   ctx->efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island];
 
   // common arrays
-  ctx->Jaref  = mj_stackAllocNum(d, nefc);
-  ctx->Jv     = mj_stackAllocNum(d, nefc);
-  ctx->Ma     = mj_stackAllocNum(d, nv);
-  ctx->Mv     = mj_stackAllocNum(d, nv);
-  ctx->grad   = mj_stackAllocNum(d, nv);
-  ctx->Mgrad  = mj_stackAllocNum(d, nv);
-  ctx->search = mj_stackAllocNum(d, nv);
-  ctx->quad   = mj_stackAllocNum(d, nefc*3);
+  ctx->Jaref  = mjSTACKALLOC(d, nefc, mjtNum);
+  ctx->Jv     = mjSTACKALLOC(d, nefc, mjtNum);
+  ctx->Ma     = mjSTACKALLOC(d, nv, mjtNum);
+  ctx->Mv     = mjSTACKALLOC(d, nv, mjtNum);
+  ctx->grad   = mjSTACKALLOC(d, nv, mjtNum);
+  ctx->Mgrad  = mjSTACKALLOC(d, nv, mjtNum);
+  ctx->search = mjSTACKALLOC(d, nv, mjtNum);
+  ctx->quad   = mjSTACKALLOC(d, nefc*3, mjtNum);
 
   // Newton only, known-size arrays
   ctx->flg_Newton = flg_Newton;
   if (flg_Newton) {
-    ctx->D    = mj_stackAllocNum(d, nefc);
+    ctx->D    = mjSTACKALLOC(d, nefc, mjtNum);
 
     // sparse Newton only
     if (mj_isSparse(m)) {
-      ctx->C          = mj_stackAllocNum(d, m->nC);
-      ctx->H_rowadr   = mj_stackAllocInt(d, nv);
-      ctx->H_rownnz   = mj_stackAllocInt(d, nv);
-      ctx->H_lowernnz = mj_stackAllocInt(d, nv);
-      ctx->L_rownnz   = mj_stackAllocInt(d, nv);
-      ctx->L_rowadr   = mj_stackAllocInt(d, nv);
+      ctx->C          = mjSTACKALLOC(d, m->nC, mjtNum);
+      ctx->H_rowadr   = mjSTACKALLOC(d, nv, int);
+      ctx->H_rownnz   = mjSTACKALLOC(d, nv, int);
+      ctx->H_lowernnz = mjSTACKALLOC(d, nv, int);
+      ctx->L_rownnz   = mjSTACKALLOC(d, nv, int);
+      ctx->L_rowadr   = mjSTACKALLOC(d, nv, int);
     }
   }
 }
@@ -1416,8 +1416,8 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
     }
 
     // allocate H_colind and H
-    ctx->H_colind = mj_stackAllocInt(d, ctx->nH);
-    ctx->H = mj_stackAllocNum(d, ctx->nH);
+    ctx->H_colind = mjSTACKALLOC(d, ctx->nH, int);
+    ctx->H = mjSTACKALLOC(d, ctx->nH, mjtNum);
 
     // compute H = J'*D*J
     mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, ctx->D, nefc, nv,
@@ -1440,10 +1440,10 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
     }
 
     // allocate L_colind, L, Lcone
-    ctx->L_colind = mj_stackAllocInt(d, ctx->nL);
-    ctx->L = mj_stackAllocNum(d, ctx->nL);
+    ctx->L_colind = mjSTACKALLOC(d, ctx->nL, int);
+    ctx->L = mjSTACKALLOC(d, ctx->nL, mjtNum);
     if (m->opt.cone == mjCONE_ELLIPTIC) {
-      ctx->Lcone = mj_stackAllocNum(d, ctx->nL);
+      ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum);
     }
 
     // count nonzeros in rows of H lower triangle
@@ -1471,9 +1471,9 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
   else {
     // allocate L, Lcone
     ctx->nL = nv*nv;
-    ctx->L = mj_stackAllocNum(d, ctx->nL);
+    ctx->L = mjSTACKALLOC(d, ctx->nL, mjtNum);
     if (m->opt.cone == mjCONE_ELLIPTIC) {
-      ctx->Lcone = mj_stackAllocNum(d, ctx->nL);
+      ctx->Lcone = mjSTACKALLOC(d, ctx->nL, mjtNum);
     }
 
     // compute H = M + J'*D*J
@@ -1572,9 +1572,9 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
   mj_markStack(d);
 
   // storage for L'*J
-  mjtNum* LTJ = mj_stackAllocNum(d, 6*nv);
-  mjtNum* LTJ_row = mj_stackAllocNum(d, nv);
-  int* LTJ_ind = mj_stackAllocInt(d, nv);
+  mjtNum* LTJ = mjSTACKALLOC(d, 6*nv, mjtNum);
+  mjtNum* LTJ_row = mjSTACKALLOC(d, nv, mjtNum);
+  int* LTJ_ind = mjSTACKALLOC(d, nv, int);
 
   // add contributions
   for (int i=0; i < nefc; i++) {
@@ -1646,8 +1646,8 @@ static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, co
   mj_markStack(d);
 
   // local space
-  mjtNum* vec = mj_stackAllocNum(d, nv);
-  int* vec_ind = mj_stackAllocInt(d, nv);
+  mjtNum* vec = mjSTACKALLOC(d, nv, mjtNum);
+  int* vec_ind = mjSTACKALLOC(d, nv, int);
 
   // clear update counter
   ctx->nupdate = 0;
@@ -1727,11 +1727,11 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
 
   // allocate local storage
   if (!flg_Newton) {
-    gradold     = mj_stackAllocNum(d, nv);
-    Mgradold    = mj_stackAllocNum(d, nv);
-    Mgraddif    = mj_stackAllocNum(d, nv);
+    gradold     = mjSTACKALLOC(d, nv, mjtNum);
+    Mgradold    = mjSTACKALLOC(d, nv, mjtNum);
+    Mgraddif    = mjSTACKALLOC(d, nv, mjtNum);
   }
-  int* oldstate = mj_stackAllocInt(d, nefc);
+  int* oldstate = mjSTACKALLOC(d, nefc, int);
 
   // initialize matrix-vector products
   int flg_vecunc = 1;  // d->qacc is uncompressed
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index ff8c1d4c..e77520e2 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -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);
 
diff --git a/src/engine/engine_util_container.c b/src/engine/engine_util_container.c
index 21d190c1..98fa11ca 100644
--- a/src/engine/engine_util_container.c
+++ b/src/engine/engine_util_container.c
@@ -24,8 +24,7 @@
 
 // stack allocate and initialize new mjArrayList
 mjArrayList* mju_arrayListCreate(mjData* d, size_t element_size, size_t initial_capacity) {
-  mjArrayList* array_list = (mjArrayList*) mj_stackAllocByte(
-    d, sizeof(mjArrayList), _Alignof(mjArrayList));
+  mjArrayList* array_list = mjSTACKALLOC(d, 1, mjArrayList);
   initial_capacity = mjMAX(1, initial_capacity);
   array_list->d = d;
   array_list->element_size = element_size;
diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c
index 4a455359..30fdc651 100644
--- a/src/engine/engine_util_solve.c
+++ b/src/engine/engine_util_solve.c
@@ -148,8 +148,8 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
   int rank = n;
 
   mj_markStack(d);
-  mjtNum* buf = mj_stackAllocNum(d, n);
-  int* buf_ind = mj_stackAllocInt(d, n);
+  mjtNum* buf = mjSTACKALLOC(d, n, mjtNum);
+  int* buf_ind = mjSTACKALLOC(d, n, int);
 
   // backpass over rows
   for (int r=n-1; r >= 0; r--) {
@@ -241,8 +241,8 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
                          const int* rownnz, const int* rowadr, int* colind, int x_nnz, int* x_ind,
                          mjData* d) {
   mj_markStack(d);
-  int* buf_ind = mj_stackAllocInt(d, n);
-  mjtNum* sparse_buf = mj_stackAllocNum(d, n);
+  int* buf_ind = mjSTACKALLOC(d, n, int);
+  mjtNum* sparse_buf = mjSTACKALLOC(d, n, mjtNum);
 
   // backpass over rows corresponding to non-zero x(r)
   int rank = n, i = x_nnz - 1;
diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c
index 434a0ce4..87580444 100644
--- a/src/engine/engine_util_sparse.c
+++ b/src/engine/engine_util_sparse.c
@@ -678,7 +678,7 @@ void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
                             const int* rownnzT, const int* rowadrT, const int* colindT,
                             const int* rowsuperT, mjData* d) {
   mj_markStack(d);
-  int* chain = mj_stackAllocInt(d, 2*nr);
+  int* chain = mjSTACKALLOC(d, 2*nr, int);
   int nchain = 0;
   int* res_colind = NULL;
 
@@ -784,11 +784,11 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
   mj_markStack(d);
 
   // a dense row buffer that stores the current row in the resulting matrix
-  mjtNum* buffer = mj_stackAllocNum(d, nc);
+  mjtNum* buffer = mjSTACKALLOC(d, nc, mjtNum);
 
   // these mark the currently set columns in the dense row buffer,
   // used for when creating the resulting sparse row
-  int* markers = mj_stackAllocInt(d, nc);
+  int* markers = mjSTACKALLOC(d, nc, int);
 
   for (int i=0; i < nc; i++) {
     int* cols = res_colind+res_rowadr[i];
@@ -899,8 +899,8 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
 int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
                       int n, mjData* d) {
   mj_markStack(d);
-  int* parent = mj_stackAllocInt(d, n);
-  int* flag = mj_stackAllocInt(d, n);
+  int* parent = mjSTACKALLOC(d, n, int);
+  int* flag = mjSTACKALLOC(d, n, int);
 
   // loop over rows in reverse order
   for (int r = n - 1; r >= 0; r--) {
diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c
index cea65167..eb3530b1 100644
--- a/src/engine/engine_vis_interact.c
+++ b/src/engine/engine_vis_interact.c
@@ -539,8 +539,8 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
   int sel = pert->select;
   mjtNum headpos[3], forward[3], dif[3];
 
-  mjtNum* jac = mj_stackAllocNum(d, 3*nv);
-  mjtNum* jacM2 = mj_stackAllocNum(d, 3*nv);
+  mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mjtNum* jacM2 = mjSTACKALLOC(d, 3*nv, mjtNum);
 
   // invalid selected body: return
   if (sel <= 0 || sel >= m->nbody) {
diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c
index 5f619fdd..8f5be4e0 100644
--- a/src/engine/engine_vis_visualize.c
+++ b/src/engine/engine_vis_visualize.c
@@ -1817,7 +1817,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
 
           // allocate catenary
           mj_markStack(d);
-          mjtNum* catenary = mj_stackAllocNum(d, 3*ncatenary);
+          mjtNum* catenary = mjSTACKALLOC(d, 3*ncatenary, mjtNum);
 
           // points along catenary path
           int npoints = mjv_catenary(x0, x1, m->opt.gravity, length, catenary, ncatenary);
@@ -2520,7 +2520,7 @@ void mjv_updateActiveFlex(const mjModel* m, mjData* d, mjvScene* scn, const mjvO
     else {
       // allocate and clear vertex normals for smoothing
       mj_markStack(d);
-      mjtNum* vertnorm = mj_stackAllocNum(d, 3*m->flex_vertnum[f]);
+      mjtNum* vertnorm = mjSTACKALLOC(d, 3*m->flex_vertnum[f], mjtNum);
       mju_zero(vertnorm, 3*m->flex_vertnum[f]);
 
       // add vertex normals: top element sides in 2D, shell fragments in 3D