Merge pull request #81 from erikfrey/no_alloc

Remove inline allocs from smooth and xfrc_accumulate.

Author: erikfrey

mujoco/mjx/_src/forward.py

@@ -635,27 +635,18 @@ def qfrc_actuator_kernel(
 def fwd_acceleration(m: Model, d: Data):
   """Add up all non-constraint forces, compute qacc_smooth."""
 
-  qfrc_applied = d.qfrc_applied
-  qfrc_accumulated = xfrc_accumulate(m, d)
-
   @kernel
-  def _qfrc_smooth(
-    d: Data,
-    qfrc_applied: wp.array(ndim=2, dtype=wp.float32),
-    qfrc_accumulated: wp.array(ndim=2, dtype=wp.float32),
-  ):
+  def _qfrc_smooth(d: Data):
     worldid, dofid = wp.tid()
     d.qfrc_smooth[worldid, dofid] = (
       d.qfrc_passive[worldid, dofid]
       - d.qfrc_bias[worldid, dofid]
       + d.qfrc_actuator[worldid, dofid]
-      + qfrc_applied[worldid, dofid]
-      + qfrc_accumulated[worldid, dofid]
+      + d.qfrc_applied[worldid, dofid]
     )
 
-  wp.launch(
-    _qfrc_smooth, dim=(d.nworld, m.nv), inputs=[d, qfrc_applied, qfrc_accumulated]
-  )
+  wp.launch(_qfrc_smooth, dim=(d.nworld, m.nv), inputs=[d])
+  xfrc_accumulate(m, d, d.qfrc_smooth)
 
   smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
 

mujoco/mjx/_src/io.py

@@ -441,6 +441,9 @@ def make_data(
   d.qfrc_constraint = wp.zeros((nworld, mjm.nv), dtype=wp.float32)
   d.qacc_smooth = wp.zeros((nworld, mjm.nv), dtype=wp.float32)
 
+  d.rne_cacc = wp.zeros(shape=(d.nworld, mjm.nbody), dtype=wp.spatial_vector)
+  d.rne_cfrc = wp.zeros(shape=(d.nworld, mjm.nbody), dtype=wp.spatial_vector)
+
   d.xfrc_applied = wp.zeros((nworld, mjm.nbody), dtype=wp.spatial_vector)
 
   # internal tmp arrays
@@ -661,6 +664,9 @@ def tile(x):
   d.contact.efc_address = wp.array(con_efc_address, dtype=wp.int32, ndim=1)
   d.contact.worldid = wp.array(con_worldid, dtype=wp.int32, ndim=1)
 
+  d.rne_cacc = wp.zeros(shape=(d.nworld, mjm.nbody), dtype=wp.spatial_vector)
+  d.rne_cfrc = wp.zeros(shape=(d.nworld, mjm.nbody), dtype=wp.spatial_vector)
+
   d.efc = _constraint(mjm.nv, d.nworld, d.njmax)
   d.efc.J = wp.array(efc_J_fill, dtype=wp.float32, ndim=2)
   d.efc.D = wp.array(efc_D_fill, dtype=wp.float32, ndim=1)

mujoco/mjx/_src/smooth.py

@@ -403,74 +403,68 @@ def factor_m(m: Model, d: Data):
 def rne(m: Model, d: Data):
   """Computes inverse dynamics using Newton-Euler algorithm."""
 
-  cacc = wp.zeros(shape=(d.nworld, m.nbody), dtype=wp.spatial_vector)
-  cfrc = wp.zeros(shape=(d.nworld, m.nbody), dtype=wp.spatial_vector)
-
   @kernel
-  def cacc_gravity(m: Model, cacc: wp.array(dtype=wp.spatial_vector, ndim=2)):
+  def cacc_gravity(m: Model, d: Data):
     worldid = wp.tid()
-    cacc[worldid, 0] = wp.spatial_vector(wp.vec3(0.0), -m.opt.gravity)
+    d.rne_cacc[worldid, 0] = wp.spatial_vector(wp.vec3(0.0), -m.opt.gravity)
 
   @kernel
   def cacc_level(
     m: Model,
     d: Data,
-    cacc: wp.array(dtype=wp.spatial_vector, ndim=2),
     leveladr: int,
   ):
     worldid, nodeid = wp.tid()
     bodyid = m.body_tree[leveladr + nodeid]
     dofnum = m.body_dofnum[bodyid]
     pid = m.body_parentid[bodyid]
     dofadr = m.body_dofadr[bodyid]
-    local_cacc = cacc[worldid, pid]
+    local_cacc = d.rne_cacc[worldid, pid]
     for i in range(dofnum):
       local_cacc += d.cdof_dot[worldid, dofadr + i] * d.qvel[worldid, dofadr + i]
-    cacc[worldid, bodyid] = local_cacc
+    d.rne_cacc[worldid, bodyid] = local_cacc
 
   @kernel
-  def frc_fn(
-    d: Data,
-    cfrc: wp.array(dtype=wp.spatial_vector, ndim=2),
-    cacc: wp.array(dtype=wp.spatial_vector, ndim=2),
-  ):
+  def frc_fn(d: Data):
     worldid, bodyid = wp.tid()
-    frc = math.inert_vec(d.cinert[worldid, bodyid], cacc[worldid, bodyid])
+    frc = math.inert_vec(d.cinert[worldid, bodyid], d.rne_cacc[worldid, bodyid])
     frc += math.motion_cross_force(
       d.cvel[worldid, bodyid],
       math.inert_vec(d.cinert[worldid, bodyid], d.cvel[worldid, bodyid]),
     )
-    cfrc[worldid, bodyid] += frc
+    d.rne_cfrc[worldid, bodyid] = frc
 
   @kernel
-  def cfrc_fn(m: Model, cfrc: wp.array(dtype=wp.spatial_vector, ndim=2), leveladr: int):
+  def cfrc_fn(m: Model, d: Data, leveladr: int):
     worldid, nodeid = wp.tid()
     bodyid = m.body_tree[leveladr + nodeid]
     pid = m.body_parentid[bodyid]
-    wp.atomic_add(cfrc[worldid], pid, cfrc[worldid, bodyid])
+    wp.atomic_add(d.rne_cfrc[worldid], pid, d.rne_cfrc[worldid, bodyid])
 
   @kernel
-  def qfrc_bias(m: Model, d: Data, cfrc: wp.array(dtype=wp.spatial_vector, ndim=2)):
+  def qfrc_bias(m: Model, d: Data):
     worldid, dofid = wp.tid()
     bodyid = m.dof_bodyid[dofid]
-    d.qfrc_bias[worldid, dofid] = wp.dot(d.cdof[worldid, dofid], cfrc[worldid, bodyid])
+    d.qfrc_bias[worldid, dofid] = wp.dot(
+      d.cdof[worldid, dofid], d.rne_cfrc[worldid, bodyid]
+    )
 
-  wp.launch(cacc_gravity, dim=[d.nworld], inputs=[m, cacc])
+  wp.launch(cacc_gravity, dim=[d.nworld], inputs=[m, d])
 
   body_treeadr = m.body_treeadr.numpy()
   for i in range(len(body_treeadr)):
     beg = body_treeadr[i]
     end = m.nbody if i == len(body_treeadr) - 1 else body_treeadr[i + 1]
-    wp.launch(cacc_level, dim=(d.nworld, end - beg), inputs=[m, d, cacc, beg])
+    wp.launch(cacc_level, dim=(d.nworld, end - beg), inputs=[m, d, beg])
 
-  wp.launch(frc_fn, dim=[d.nworld, m.nbody], inputs=[d, cfrc, cacc])
+  wp.launch(frc_fn, dim=[d.nworld, m.nbody], inputs=[d])
 
   for i in reversed(range(len(body_treeadr))):
     beg = body_treeadr[i]
     end = m.nbody if i == len(body_treeadr) - 1 else body_treeadr[i + 1]
-    wp.launch(cfrc_fn, dim=[d.nworld, end - beg], inputs=[m, cfrc, beg])
+    wp.launch(cfrc_fn, dim=[d.nworld, end - beg], inputs=[m, d, beg])
 
-  wp.launch(qfrc_bias, dim=[d.nworld, m.nv], inputs=[m, d, cfrc])
+  wp.launch(qfrc_bias, dim=[d.nworld, m.nv], inputs=[m, d])
 
 
 @event_scope

mujoco/mjx/_src/support.py

@@ -113,40 +113,25 @@ def _mul_m_sparse_ij(
     )
 
 
-@wp.kernel
-def process_level(
-  body_tree: wp.array(ndim=1, dtype=int),
-  body_parentid: wp.array(ndim=1, dtype=int),
-  dof_bodyid: wp.array(ndim=1, dtype=int),
-  mask: wp.array2d(dtype=wp.bool),
-  beg: int,
-):
-  dofid, tid_y = wp.tid()
-  j = beg + tid_y
-  el = body_tree[j]
-  parent_id = body_parentid[el]
-  parent_val = mask[dofid, parent_id]
-  mask[dofid, el] = parent_val or (dof_bodyid[dofid] == el)
-
-
-@wp.kernel
-def compute_qfrc(
-  d: Data,
-  m: Model,
-  mask: wp.array2d(dtype=wp.bool),
-  qfrc_total: array2df,
-):
-  worldid, dofid = wp.tid()
-  accumul = float(0.0)
-  cdof_vec = d.cdof[worldid, dofid]
-  rotational_cdof = wp.vec3(cdof_vec[0], cdof_vec[1], cdof_vec[2])
-
-  jac = wp.spatial_vector(
-    cdof_vec[3], cdof_vec[4], cdof_vec[5], cdof_vec[0], cdof_vec[1], cdof_vec[2]
-  )
-
-  for bodyid in range(m.nbody):
-    if mask[dofid, bodyid]:
+@event_scope
+def xfrc_accumulate(m: Model, d: Data, qfrc: array2df):
+  @wp.kernel
+  def _accumulate(m: Model, d: Data, qfrc: array2df):
+    worldid, dofid = wp.tid()
+    cdof = d.cdof[worldid, dofid]
+    rotational_cdof = wp.vec3(cdof[0], cdof[1], cdof[2])
+    jac = wp.spatial_vector(cdof[3], cdof[4], cdof[5], cdof[0], cdof[1], cdof[2])
+
+    dof_bodyid = m.dof_bodyid[dofid]
+    accumul = float(0.0)
+
+    for bodyid in range(dof_bodyid, m.nbody):
+      # any body that is in the subtree of dof_bodyid is part of the jacobian
+      parentid = bodyid
+      while parentid != 0 and parentid != dof_bodyid:
+        parentid = m.body_parentid[parentid]
+      if parentid == 0:
+        continue  # body is not part of the subtree
       offset = d.xipos[worldid, bodyid] - d.subtree_com[worldid, m.body_rootid[bodyid]]
       cross_term = wp.cross(rotational_cdof, offset)
       accumul += wp.dot(jac, d.xfrc_applied[worldid, bodyid]) + wp.dot(
@@ -158,30 +143,9 @@ def compute_qfrc(
         ),
       )
 
-  qfrc_total[worldid, dofid] = accumul
-
-
-@event_scope
-def xfrc_accumulate(m: Model, d: Data) -> array2df:
-  body_treeadr_np = m.body_treeadr.numpy()
-  mask = wp.zeros((m.nv, m.nbody), dtype=wp.bool)
-
-  for i in range(len(body_treeadr_np)):
-    beg = body_treeadr_np[i]
-    end = m.nbody if i == len(body_treeadr_np) - 1 else body_treeadr_np[i + 1]
-
-    if end > beg:
-      wp.launch(
-        kernel=process_level,
-        dim=[m.nv, (end - beg)],
-        inputs=[m.body_tree, m.body_parentid, m.dof_bodyid, mask, beg],
-      )
-
-  qfrc_total = wp.zeros((d.nworld, m.nv), dtype=float)
-
-  wp.launch(kernel=compute_qfrc, dim=(d.nworld, m.nv), inputs=[d, m, mask, qfrc_total])
+    qfrc[worldid, dofid] += accumul
 
-  return qfrc_total
+  wp.launch(kernel=_accumulate, dim=(d.nworld, m.nv), inputs=[m, d, qfrc])
 
 
 @wp.func

mujoco/mjx/_src/support_test.py

@@ -60,20 +60,14 @@ def test_xfrc_accumulated(self):
     mjm, mjd, m, d = test_util.fixture("pendula.xml")
     xfrc = np.random.randn(*d.xfrc_applied.numpy().shape)
     d.xfrc_applied = wp.from_numpy(xfrc, dtype=wp.spatial_vector)
-    qfrc = xfrc_accumulate(m, d)
+    qfrc = wp.zeros((1, mjm.nv), dtype=wp.float32)
+    xfrc_accumulate(m, d, qfrc)
 
     qfrc_expected = np.zeros(m.nv)
     xfrc = xfrc[0]
-    mjd.xfrc_applied[:] = xfrc
     for i in range(1, m.nbody):
       mujoco.mj_applyFT(
-        mjm,
-        mjd,
-        mjd.xfrc_applied[i, :3],
-        mjd.xfrc_applied[i, 3:],
-        mjd.xipos[i],
-        i,
-        qfrc_expected,
+        mjm, mjd, xfrc[i, :3], xfrc[i, 3:], mjd.xipos[i], i, qfrc_expected
       )
     np.testing.assert_almost_equal(qfrc.numpy()[0], qfrc_expected, 6)
 

mujoco/mjx/_src/types.py

@@ -464,6 +464,10 @@ class Data:
   contact: Contact
   efc: Constraint
 
+  # arrays used for smooth.rne
+  rne_cacc: wp.array(dtype=wp.spatial_vector, ndim=2)
+  rne_cfrc: wp.array(dtype=wp.spatial_vector, ndim=2)
+
   # temp arrays
   qfrc_integration: wp.array(dtype=wp.float32, ndim=2)
   qacc_integration: wp.array(dtype=wp.float32, ndim=2)