Small change to testspeed file.

PiperOrigin-RevId: 618323795
Change-Id: I3a6307b8ae6d9af0d17c0ae85eae673d60112ddf

mjx/mujoco/mjx/_src/support.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 """Engine support functions."""
-
 from typing import Optional, Tuple, Union
 
 import jax

mjx/mujoco/mjx/_src/test_util.py

@@ -14,14 +14,96 @@
 # ==============================================================================
 """Utilities for testing."""
 
+import os
 import sys
+import time
 from typing import Dict, Tuple
 from xml.etree import ElementTree as ET
 
 from etils import epath
+import jax
 import mujoco
+# pylint: disable=g-importing-member
+from mujoco.mjx._src import forward
+from mujoco.mjx._src import io
+# pylint: enable=g-importing-member
 import numpy as np
 
+
+def _measure(fn, *args) -> Tuple[float, float]:
+  """Reports jit time and op time for a function."""
+
+  beg = time.perf_counter()
+  compiled_fn = fn.lower(*args).compile()
+  end = time.perf_counter()
+  jit_time = end - beg
+
+  beg = time.perf_counter()
+  result = compiled_fn(*args)
+  jax.block_until_ready(result)
+  end = time.perf_counter()
+  run_time = end - beg
+
+  return jit_time, run_time
+
+
+def benchmark(
+    m: mujoco.MjModel,
+    nstep: int = 1000,
+    batch_size: int = 1024,
+    unroll_steps: int = 1,
+    solver: str = 'cg',
+    iterations: int = 1,
+    ls_iterations: int = 4,
+) -> Tuple[float, float, int]:
+  """Benchmark a model."""
+
+  xla_flags = os.environ.get('XLA_FLAGS', '')
+  xla_flags += ' --xla_gpu_triton_gemm_any=True'
+  os.environ['XLA_FLAGS'] = xla_flags
+
+  m.opt.solver = {
+      'cg': mujoco.mjtSolver.mjSOL_CG,
+      'newton': mujoco.mjtSolver.mjSOL_NEWTON,
+  }[solver.lower()]
+  m.opt.iterations = iterations
+  m.opt.ls_iterations = ls_iterations
+  m = io.put_model(m)
+
+  @jax.pmap
+  def init(key):
+    key = jax.random.split(key, batch_size // jax.device_count())
+
+    @jax.vmap
+    def random_init(key):
+      d = io.make_data(m)
+      qvel = 0.01 * jax.random.normal(key, shape=(m.nv,))
+      d = d.replace(qvel=qvel)
+      return d
+
+    return random_init(key)
+
+  key = jax.random.split(jax.random.key(0), jax.device_count())
+  d = init(key)
+  jax.block_until_ready(d)
+
+  @jax.pmap
+  def unroll(d):
+    @jax.vmap
+    def step(d, _):
+      d = forward.step(m, d)
+      return d, None
+
+    d, _ = jax.lax.scan(step, d, None, length=nstep, unroll=unroll_steps)
+
+    return d
+
+  jit_time, run_time = _measure(unroll, d)
+  steps = nstep * batch_size
+
+  return jit_time, run_time, steps
+
+
 _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
 _DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact']
 _DYN_PRMS = ['0.189', '2.1']