Speed up hfield JIT time by reducing XLA code generation overhead.

PiperOrigin-RevId: 653766462
Change-Id: I9d7e7fa3a4d33ae0fa9d00babd5c5478eb4eff88

mjx/mujoco/mjx/_src/collision_convex.py

@@ -1079,9 +1079,11 @@ def _hfield_collision(
   bmask = jp.array([True, True, False])
 
   # process all prisms in sub-grid
-  prisms = []
+  rs = jp.repeat(jp.arange(subgrid_size[1]), subgrid_size[0])
-  for r in range(subgrid_size[1]):
+  cs = jp.tile(jp.arange(subgrid_size[0]), subgrid_size[1])
-    for c in range(subgrid_size[0]):
+
+  @jax.vmap
+  def make_prisms(r, c):
     ri, ci = rmin + r, cmin + c
 
     # ensure ri, ci are in the bounds of the hfield
@@ -1106,7 +1108,7 @@ def _hfield_collision(
     top = jp.array([p1, p2, p3])
     bottom = jp.array([p1, p3, p2]) * bmask + bvert
     vert = jp.concatenate([bottom, top])
-      prisms.append(mesh.hfield_prism(vert))
+    prism1 = mesh.hfield_prism(vert)
 
     p3 = p2
     p2 = [
@@ -1120,10 +1122,13 @@ def _hfield_collision(
     # NB: If the order of verts is updated above, the corresponding
     # hfield_prism function must be updated to ensure that all faces have the
     # correct winding order.
-      prisms.append(mesh.hfield_prism(vert))
+    prism2 = mesh.hfield_prism(vert)
 
-  n_prisms = len(prisms)
+    return prism1, prism2
-  prisms = jax.tree_util.tree_map(lambda *x: jp.stack(x), *prisms)
+
+  prism1, prism2 = make_prisms(rs, cs)
+  n_prisms = 2 * rs.shape[0]
+  prisms = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), prism1, prism2)
   dist, pos, n = jax.vmap(collider_fn, in_axes=[None, 0])(
       obj.replace(pos=obj_pos, mat=obj_mat), prisms
   )