Add stage_in and stage_out for MJX-Warp.

PiperOrigin-RevId: 868180029
Change-Id: I0d1b17b6023d32498766b5d1d9e5624bd31cda3d

mjx/mujoco/mjx/warp/collision_driver.py

@@ -233,8 +233,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       'collision_pair': d._impl.collision_pair.shape,
       'collision_pairid': d._impl.collision_pairid.shape,
       'collision_worldid': d._impl.collision_worldid.shape,
-      'geom_xmat': d.geom_xmat.shape,
-      'geom_xpos': d.geom_xpos.shape,
       'nacon': d._impl.nacon.shape,
       'ncollision': d._impl.ncollision.shape,
       'contact__dim': d._impl.contact__dim.shape,
@@ -253,15 +251,13 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _collision_shim,
-      num_outputs=20,
+      num_outputs=18,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
           'collision_pair',
           'collision_pairid',
           'collision_worldid',
-          'geom_xmat',
-          'geom_xpos',
           'nacon',
           'ncollision',
           'contact__dim',
@@ -278,6 +274,27 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'contact__type',
           'contact__worldid',
       },
+      stage_in_argnames={
+          'geom_aabb',
+          'geom_friction',
+          'geom_gap',
+          'geom_margin',
+          'geom_rbound',
+          'geom_size',
+          'geom_solimp',
+          'geom_solmix',
+          'geom_solref',
+          'geom_xmat',
+          'geom_xpos',
+          'hfield_data',
+          'pair_friction',
+          'pair_gap',
+          'pair_margin',
+          'pair_solimp',
+          'pair_solref',
+          'pair_solreffriction',
+      },
+      stage_out_argnames={},
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
@@ -373,23 +390,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       '_impl.collision_pair': out[0],
       '_impl.collision_pairid': out[1],
       '_impl.collision_worldid': out[2],
-      'geom_xmat': out[3],
+      '_impl.nacon': out[3],
-      'geom_xpos': out[4],
+      '_impl.ncollision': out[4],
-      '_impl.nacon': out[5],
+      '_impl.contact__dim': out[5],
-      '_impl.ncollision': out[6],
+      '_impl.contact__dist': out[6],
-      '_impl.contact__dim': out[7],
+      '_impl.contact__frame': out[7],
-      '_impl.contact__dist': out[8],
+      '_impl.contact__friction': out[8],
-      '_impl.contact__frame': out[9],
+      '_impl.contact__geom': out[9],
-      '_impl.contact__friction': out[10],
+      '_impl.contact__geomcollisionid': out[10],
-      '_impl.contact__geom': out[11],
+      '_impl.contact__includemargin': out[11],
-      '_impl.contact__geomcollisionid': out[12],
+      '_impl.contact__pos': out[12],
-      '_impl.contact__includemargin': out[13],
+      '_impl.contact__solimp': out[13],
-      '_impl.contact__pos': out[14],
+      '_impl.contact__solref': out[14],
-      '_impl.contact__solimp': out[15],
+      '_impl.contact__solreffriction': out[15],
-      '_impl.contact__solref': out[16],
+      '_impl.contact__type': out[16],
-      '_impl.contact__solreffriction': out[17],
+      '_impl.contact__worldid': out[17],
-      '_impl.contact__type': out[18],
-      '_impl.contact__worldid': out[19],
   })
   return d
 

mjx/mujoco/mjx/warp/ffi.py

@@ -101,6 +101,8 @@ def jax_callable_variadic_tuple(
     vmap_method: Optional[str] = None,
     output_dims: Optional[dict[str, tuple[int, ...]]] = None,
     in_out_argnames: Optional[Sequence[str]] = None,
+    stage_in_argnames: Optional[Sequence[str]] = None,
+    stage_out_argnames: Optional[Sequence[str]] = None,
 ):
   """Wraps a JAX callable to support variadic tuples and dataclasses."""
 
@@ -130,6 +132,8 @@ def jax_callable_variadic_tuple(
         vmap_method=vmap_method,
         output_dims=output_dims,
         in_out_argnames=in_out_argnames,
+        stage_in_argnames=stage_in_argnames,
+        stage_out_argnames=stage_out_argnames,
     )
 
     flat_args, in_tree = jax.tree.flatten(args)

mjx/mujoco/mjx/warp/forward_test.py

@@ -218,7 +218,7 @@ class StepTest(parameterized.TestCase):
       batch_size=(1, 7),
       # NOTE: GraphMode.JAX is incompatible with MuJoCo Warp at the moment,
       # even when setting graph_conditional=False.
-      graph_mode=('WARP',),
+      graph_mode=('WARP', 'WARP_STAGED'),
   )
   def test_step(self, xml: str, batch_size: int, graph_mode: str):
     if not _FORCE_TEST:

mjx/mujoco/mjx/warp/smooth.py

@@ -136,9 +136,6 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
-      'mocap_pos': d.mocap_pos.shape,
-      'mocap_quat': d.mocap_quat.shape,
-      'qpos': d.qpos.shape,
       'site_xmat': d.site_xmat.shape,
       'site_xpos': d.site_xpos.shape,
       'xanchor': d.xanchor.shape,
@@ -151,15 +148,52 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _kinematics_shim,
-      num_outputs=14,
+      num_outputs=11,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
           'geom_xmat',
           'geom_xpos',
+          'site_xmat',
+          'site_xpos',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      },
+      stage_in_argnames={
+          'body_ipos',
+          'body_iquat',
+          'body_pos',
+          'body_quat',
+          'geom_pos',
+          'geom_quat',
+          'geom_xmat',
+          'geom_xpos',
+          'jnt_axis',
+          'jnt_pos',
           'mocap_pos',
           'mocap_quat',
           'qpos',
+          'qpos0',
+          'site_pos',
+          'site_quat',
+          'site_xmat',
+          'site_xpos',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      },
+      stage_out_argnames={
+          'geom_xmat',
+          'geom_xpos',
           'site_xmat',
           'site_xpos',
           'xanchor',
@@ -216,18 +250,15 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
   d = d.tree_replace({
       'geom_xmat': out[0],
       'geom_xpos': out[1],
-      'mocap_pos': out[2],
+      'site_xmat': out[2],
-      'mocap_quat': out[3],
+      'site_xpos': out[3],
-      'qpos': out[4],
+      'xanchor': out[4],
-      'site_xmat': out[5],
+      'xaxis': out[5],
-      'site_xpos': out[6],
+      'ximat': out[6],
-      'xanchor': out[7],
+      'xipos': out[7],
-      'xaxis': out[8],
+      'xmat': out[8],
-      'ximat': out[9],
+      'xpos': out[9],
-      'xipos': out[10],
+      'xquat': out[10],
-      'xmat': out[11],
-      'xpos': out[12],
-      'xquat': out[13],
   })
   return d
 
@@ -243,25 +274,6 @@ def kinematics_vmap(unused_axis_size, is_batched, m, d):
   return d, is_batched[1]
 
 
-_m = mjwarp.Model(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
-)
-_d = mjwarp.Data(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init}
-)
-_o = mjwarp.Option(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init}
-)
-_s = mjwarp.Statistic(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init}
-)
-_c = mjwarp.Contact(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init}
-)
-_e = mjwarp.Constraint(
-    **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
-)
-
 @ffi.format_args_for_warp
 def _tendon_shim(
     # Model
@@ -348,12 +360,6 @@ def _tendon_shim(
 
 def _tendon_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'cdof': d.cdof.shape,
-      'geom_xmat': d.geom_xmat.shape,
-      'geom_xpos': d.geom_xpos.shape,
-      'qpos': d.qpos.shape,
-      'site_xpos': d.site_xpos.shape,
-      'subtree_com': d.subtree_com.shape,
       'ten_J': d._impl.ten_J.shape,
       'ten_length': d.ten_length.shape,
       'ten_wrapadr': d._impl.ten_wrapadr.shape,
@@ -363,16 +369,10 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _tendon_shim,
-      num_outputs=12,
+      num_outputs=6,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
-          'cdof',
-          'geom_xmat',
-          'geom_xpos',
-          'qpos',
-          'site_xpos',
-          'subtree_com',
           'ten_J',
           'ten_length',
           'ten_wrapadr',
@@ -380,6 +380,17 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
           'wrap_obj',
           'wrap_xpos',
       },
+      stage_in_argnames={
+          'cdof',
+          'geom_size',
+          'geom_xmat',
+          'geom_xpos',
+          'qpos',
+          'site_xpos',
+          'subtree_com',
+          'ten_length',
+      },
+      stage_out_argnames={'ten_length'},
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
@@ -421,18 +432,12 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
       d._impl.wrap_xpos,
   )
   d = d.tree_replace({
-      'cdof': out[0],
+      '_impl.ten_J': out[0],
-      'geom_xmat': out[1],
+      'ten_length': out[1],
-      'geom_xpos': out[2],
+      '_impl.ten_wrapadr': out[2],
-      'qpos': out[3],
+      '_impl.ten_wrapnum': out[3],
-      'site_xpos': out[4],
+      '_impl.wrap_obj': out[4],
-      'subtree_com': out[5],
+      '_impl.wrap_xpos': out[5],
-      '_impl.ten_J': out[6],
-      'ten_length': out[7],
-      '_impl.ten_wrapadr': out[8],
-      '_impl.ten_wrapnum': out[9],
-      '_impl.wrap_obj': out[10],
-      '_impl.wrap_xpos': out[11],
   })
   return d
 

mjx/mujoco/mjx/warp/testspeed.py

@@ -22,6 +22,7 @@ from typing import Any, Callable, Sequence, Tuple
 from absl import app
 from absl import flags
 import jax
+import jax.numpy as jnp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
@@ -29,6 +30,7 @@ from mujoco.mjx.warp import collision_driver as wp_collision
 from mujoco.mjx.warp import forward as wp_forward
 from mujoco.mjx.warp import smooth as wp_smooth
 import mujoco.mjx.third_party.mujoco_warp as mjwarp
+import numpy as np
 import warp as wp
 from mujoco.mjx.third_party.warp._src.jax_experimental import ffi as warp_ffi
 
@@ -50,6 +52,18 @@ _WP_KERNEL_CACHE_DIR = flags.DEFINE_string(
     None,
     'Path to the Warp kernel cache directory.',
 )
+_GRAPH_MODE = flags.DEFINE_enum(
+    'graph_mode',
+    'WARP',
+    ['NONE', 'WARP', 'WARP_STAGED', 'WARP_STAGED_EX'],
+    'Graph capture mode for JAX WARP FFI benchmark.',
+)
+_BENCHMARK = flags.DEFINE_enum(
+    'benchmark',
+    'jax_warp',
+    ['jax_warp', 'jax', 'warp'],
+    'Which benchmark to run.',
+)
 _COMPILER_OPTIONS = {'xla_gpu_graph_min_graph_size': 1}
 jax_jit = functools.partial(jax.jit, compiler_options=_COMPILER_OPTIONS)
 
@@ -95,10 +109,14 @@ def benchmark(
     d = mjx.make_data(
         m, impl=mx.impl, naconmax=_NACONMAX.value, njmax=_NJMAX.value
     )
+    # Initialize from first keyframe if available
+    if m.nkey > 0:
+      d = d.replace(qpos=m.key_qpos[0], ctrl=m.key_ctrl[0])
     return d
 
   key = jax.random.split(jax.random.key(0), nenv)
   d = jax_jit(init)(key)
+
   jax.block_until_ready(d)
 
   @jax_jit
@@ -115,177 +133,6 @@ def benchmark(
   return jit_time, run_time, steps
 
 
-def benchmark_raw_jax_warp(
-    m: mujoco.MjModel,
-    nstep: int = 1000,
-    nenv: int = 8192,
-    unroll_steps: int = 4,
-    function: str = 'kinematics',
-):
-  if function not in ('kinematics', 'forward', 'step', 'collision'):
-    raise NotImplementedError(
-        f'{function} is not implemented for raw warp speed test.'
-    )
-
-  def warp_fn(
-      qpos_in: wp.array2d(dtype=wp.float32),
-      qpos_out: wp.array2d(dtype=wp.float32),
-      time_: wp.array(dtype=wp.float32),
-      xpos: wp.array2d(dtype=wp.vec3),
-      xquat: wp.array2d(dtype=wp.quat),
-      xmat: wp.array2d(dtype=wp.mat33),
-      xipos: wp.array2d(dtype=wp.vec3),
-      ximat: wp.array2d(dtype=wp.mat33),
-      xanchor: wp.array2d(dtype=wp.vec3),
-      xaxis: wp.array2d(dtype=wp.vec3),
-      geom_xpos: wp.array2d(dtype=wp.vec3),
-      geom_xmat: wp.array2d(dtype=wp.mat33),
-  ):
-    wp.copy(d.qpos, qpos_in)
-    if function == 'kinematics':
-      mjwarp.kinematics(mw, d)
-    elif function == 'forward':
-      mjwarp.forward(mw, d)
-    elif function == 'step':
-      mjwarp.step(mw, d)
-    elif function == 'collision':
-      mjwarp.collision(mw, d)
-    else:
-      raise NotImplementedError(f'{function} not implemented in speed test.')
-    wp.copy(qpos_out, d.qpos)
-    wp.copy(time_, d.time)
-    wp.copy(xpos, d.xpos)
-    wp.copy(xquat, d.xquat)
-    wp.copy(xmat, d.xmat)
-    wp.copy(xipos, d.xipos)
-    wp.copy(ximat, d.ximat)
-    wp.copy(xanchor, d.xanchor)
-    wp.copy(xaxis, d.xaxis)
-    wp.copy(geom_xpos, d.geom_xpos)
-    wp.copy(geom_xmat, d.geom_xmat)
-
-  def unroll(
-      qpos,
-      qpos_out,
-      time_,
-      xpos,
-      xquat,
-      xmat,
-      xipos,
-      ximat,
-      xanchor,
-      xaxis,
-      geom_xpos,
-      geom_xmat,
-  ):
-    def step(carry, _):
-      qpos_, *_ = carry
-      out = warp_fn_jax(qpos_ + 0.0 * qpos_)
-      out = tuple(out)
-      return (out[0],) + out, None
-
-    (
-        qpos,
-        qpos_out,
-        time_,
-        xpos,
-        xquat,
-        xmat,
-        xipos,
-        ximat,
-        xanchor,
-        xaxis,
-        geom_xpos,
-        geom_xmat,
-    ), _ = jax.lax.scan(
-        step,
-        (
-            qpos,
-            qpos_out,
-            time_,
-            xpos,
-            xquat,
-            xmat,
-            xipos,
-            ximat,
-            xanchor,
-            xaxis,
-            geom_xpos,
-            geom_xmat,
-        ),
-        length=nstep,
-        unroll=unroll_steps,
-    )
-
-    return (
-        qpos,
-        qpos_out,
-        time_,
-        xpos,
-        xquat,
-        xmat,
-        xipos,
-        ximat,
-        xanchor,
-        xaxis,
-        geom_xpos,
-        geom_xmat,
-    )
-
-  output_dims = {
-      'time_': (nenv,),
-      'qpos_out': (nenv, m.nq),
-      'xpos': (nenv, m.nbody, 3),
-      'xquat': (nenv, m.nbody, 4),
-      'xmat': (nenv, m.nbody, 3, 3),
-      'xipos': (nenv, m.nbody, 3),
-      'ximat': (nenv, m.nbody, 3, 3),
-      'xanchor': (nenv, m.njnt, 3),
-      'xaxis': (nenv, m.njnt, 3),
-      'geom_xpos': (nenv, m.ngeom, 3),
-      'geom_xmat': (nenv, m.ngeom, 3, 3),
-  }
-  warp_fn_jax = warp_ffi.jax_callable(
-      warp_fn,
-      num_outputs=11,
-      output_dims=output_dims,
-      graph_mode=warp_ffi.GraphMode.JAX,
-  )
-
-  @jax.vmap
-  def init(key):
-    d = mjx.make_data(m, impl='jax')
-    return d
-
-  key = jax.random.split(jax.random.key(0), nenv)
-  dx = jax_jit(init)(key)
-  d_ = mujoco.MjData(m)
-  mw = mjwarp.put_model(m)
-  d = mjwarp.put_data(
-      m, d_, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value
-  )
-
-  jax_unroll_fn = jax_jit(unroll)
-  jit_time, run_time = _measure(
-      jax_unroll_fn,
-      dx.qpos,
-      dx.qpos,
-      dx.time,
-      dx.xpos,
-      dx.xquat,
-      dx.xmat,
-      dx.xipos,
-      dx.ximat,
-      dx.xanchor,
-      dx.xaxis,
-      dx.geom_xpos,
-      dx.geom_xmat,
-  )
-  steps = nstep * nenv
-
-  return jit_time, run_time, steps
-
-
 def _compile_fn(fn, m, d):
   fn(m, d)
   fn(m, d)
@@ -310,9 +157,16 @@ def benchmark_raw_warp(
 
   mw = mjwarp.put_model(m)
   dw = mjwarp.make_data(
-      m, nworld=nenv, nconmax=_NACONMAX.value, njmax=_NJMAX.value
+      m, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value
   )
 
+  # Initialize from first keyframe if available
+  if m.nkey > 0:
+    qpos_init = np.tile(m.key_qpos[0], (nenv, 1)).astype(np.float32)
+    ctrl_init = np.tile(m.key_ctrl[0], (nenv, 1)).astype(np.float32)
+    wp.copy(dw.qpos, wp.from_numpy(qpos_init))
+    wp.copy(dw.ctrl, wp.from_numpy(ctrl_init))
+
   if function == 'kinematics':
     fn = mjwarp.kinematics
   elif function == 'forward':
@@ -353,8 +207,16 @@ def _main(_: Sequence[str]):
   except Exception as _:
     m = mujoco.MjModel.from_xml_path(modelfile)
 
-  mx = mjx.put_model(m, impl='jax')
+  benchmark_type = _BENCHMARK.value
-  mw = mjx.put_model(m, impl='warp')
+  graph_mode = getattr(warp_ffi.GraphMode, _GRAPH_MODE.value)
+
+  # Only allocate the model needed for the specific benchmark
+  mx = None
+  mw = None
+  if benchmark_type == 'jax_warp':
+    mw = mjx.put_model(m, impl='warp', graph_mode=graph_mode)
+  elif benchmark_type == 'jax':
+    mx = mjx.put_model(m, impl='jax')
 
   if function_ == 'kinematics':
     func_warp = jax.vmap(wp_smooth.kinematics, in_axes=(None, 0))
@@ -377,38 +239,24 @@ def _main(_: Sequence[str]):
   print(f' nenv                 : {nenv}')
   print(f' nstep                : {nstep}')
   print(f' timestep             : {m.opt.timestep}')
-  print(f' unroll               : {unroll}\n')
+  print(f' unroll               : {unroll}')
+  print(f' benchmark            : {benchmark_type}')
+  print(f' graph_mode           : {_GRAPH_MODE.value}\n')
 
-  for name, mx_, op in (
+  if benchmark_type == 'jax_warp':
-      ('JAX WARP FFI', mw, func_warp),
+    jit_time, run_time, steps = benchmark(m, mw, func_warp, nstep, nenv, unroll)
-      ('Pure JAX', mx, func_jax),
+    print(f' JAX WARP FFI (GraphMode: {_GRAPH_MODE.value}):')
-  ):
+  elif benchmark_type == 'jax':
-    if op is not None:
+    jit_time, run_time, steps = benchmark(m, mx, func_jax, nstep, nenv, unroll)
-      jit_time, run_time, steps = benchmark(m, mx_, op, nstep, nenv, unroll)
+    print(' Pure JAX:')
+  elif benchmark_type == 'warp':
+    jit_time, run_time, steps = benchmark_raw_warp(
+        m, nstep, nenv, unroll, function=function_
+    )
+    print(' Pure WARP:')
+  else:
+    raise ValueError(f'Unknown benchmark type: {benchmark_type}')
 
-      print(f' {name}:')
-      print(f' JIT time             : {jit_time:.2f} s')
-      print(f' simulation time      : {run_time:.2f} s')
-      print(f' steps per second     : {steps / run_time:,.0f}')
-      print(
-          f' realtime factor      : {steps * m.opt.timestep / run_time:.2f} x'
-      )
-      print(f' time per step        : {1e6 * run_time / steps:.2f} µs\n')
-
-  jit_time, run_time, steps = benchmark_raw_jax_warp(
-      m, nstep, nenv, unroll, function=function_
-  )
-  print(' Pure JAX-WARP:')
-  print(f' JIT time             : {jit_time:.2f} s')
-  print(f' simulation time      : {run_time:.2f} s')
-  print(f' steps per second     : {steps / run_time:,.0f}')
-  print(f' realtime factor      : {steps * m.opt.timestep / run_time:.2f} x')
-  print(f' time per step        : {1e6 * run_time / steps:.2f} µs\n')
-
-  jit_time, run_time, steps = benchmark_raw_warp(
-      m, nstep, nenv, unroll, function=function_
-  )
-  print(' Pure WARP:')
   print(f' JIT time             : {jit_time:.2f} s')
   print(f' simulation time      : {run_time:.2f} s')
   print(f' steps per second     : {steps / run_time:,.0f}')

mjx/mujoco/mjx/warp/types.py

@@ -23,7 +23,6 @@ from jax import tree_util
 from jax.interpreters import batching
 from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
-
 if typing.TYPE_CHECKING:
   GraphMode = int
 else:
@@ -44,6 +43,7 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
+
   adr: np.ndarray
   size: int