Merge pull request #44 from erikfrey/testspeed_fixes

Testspeed fixes

Author: erikfrey

mujoco/mjx/_src/io.py

@@ -353,11 +353,11 @@ def put_data(
   # TODO(team): move to Model?
   if nconmax == -1:
     # TODO(team): heuristic for nconmax
-    nconmax = 512
+    nconmax = max(512, mjd.ncon * nworld)
   d.nconmax = nconmax
   if njmax == -1:
     # TODO(team): heuristic for njmax
-    njmax = 512
+    njmax = max(512, mjd.nefc * nworld)
   d.njmax = njmax
 
   if nworld * mjd.nefc > njmax:

mujoco/mjx/_src/test_util.py

@@ -46,41 +46,41 @@ def fixture(fname: str, keyframe: int = -1, sparse: bool = True):
 
 def benchmark(
   fn: Callable[[types.Model, types.Data], None],
-  m: mujoco.MjModel,
+  mjm: mujoco.MjModel,
+  mjd: mujoco.MjData,
   nstep: int = 1000,
   batch_size: int = 1024,
   unroll_steps: int = 1,
   solver: str = "cg",
   iterations: int = 1,
   ls_iterations: int = 4,
-  nefc_total: int = 0,
+  nconmax: int = -1,
+  njmax: int = -1,
 ) -> Tuple[float, float, int]:
   """Benchmark a model."""
 
   if solver == "cg":
-    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
+    mjm.opt.solver = mujoco.mjtSolver.mjSOL_CG
   elif solver == "newton":
-    m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
+    mjm.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
 
-  m.opt.iterations = iterations
-  m.opt.ls_iterations = ls_iterations
+  mjm.opt.iterations = iterations
+  mjm.opt.ls_iterations = ls_iterations
 
-  mx = io.put_model(m)
-  dx = io.make_data(m, nworld=batch_size, njmax=nefc_total)
-  dx.nefc_total = wp.array([nefc_total], dtype=wp.int32, ndim=1)
+  m = io.put_model(mjm)
+  d = io.put_data(mjm, mjd, nworld=batch_size, nconmax=nconmax, njmax=njmax)
 
-  wp.clear_kernel_cache()
   jit_beg = time.perf_counter()
-  fn(mx, dx)
+  fn(m, d)
   # double warmup to work around issues with compilation during graph capture:
-  fn(mx, dx)
+  fn(m, d)
   jit_end = time.perf_counter()
   jit_duration = jit_end - jit_beg
   wp.synchronize()
 
   # capture the whole smooth.kinematic() function as a CUDA graph
   with wp.ScopedCapture() as capture:
-    fn(mx, dx)
+    fn(m, d)
   graph = capture.graph
 
   run_beg = time.perf_counter()

mujoco/mjx/testspeed.py

@@ -49,12 +49,18 @@
 _IS_SPARSE = flags.DEFINE_bool(
   "is_sparse", True, "if model should create sparse mass matrices"
 )
-_NEFC_TOTAL = flags.DEFINE_integer(
-  "nefc_total", 0, "total number of efc for batch of worlds"
+_NCONMAX = flags.DEFINE_integer(
+  "nconmax", -1, "Maximum number of contacts in a batch physics step."
+)
+_NJMAX = flags.DEFINE_integer(
+  "njmax", -1, "Maximum number of constraints in a batch physics step."
 )
 _OUTPUT = flags.DEFINE_enum(
   "output", "text", ["text", "tsv"], "format to print results"
 )
+_CLEAR_KERNEL_CACHE = flags.DEFINE_bool(
+  "clear_kernel_cache", False, "Clear kernel cache (to calculate full JIT time)"
+)
 
 
 def _main(argv: Sequence[str]):
@@ -74,20 +80,32 @@ def _main(argv: Sequence[str]):
   else:
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
 
+  d = mujoco.MjData(m)
+  if m.nkey > 0:
+    mujoco.mj_resetDataKeyframe(m, d, 0)
+  # populate some constraints
+  mujoco.mj_forward(m, d)
+
+  if _CLEAR_KERNEL_CACHE.value:
+    wp.clear_kernel_cache()
+
   print(
     f"Model nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} is_sparse: {_IS_SPARSE.value}"
   )
+  print(f"Data ncon: {d.ncon} nefc: {d.nefc}")
   print(f"Rolling out {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...")
   jit_time, run_time, steps = mjx.benchmark(
     mjx.__dict__[_FUNCTION.value],
     m,
+    d,
     _NSTEP.value,
     _BATCH_SIZE.value,
     _UNROLL.value,
     _SOLVER.value,
     _ITERATIONS.value,
     _LS_ITERATIONS.value,
-    _NEFC_TOTAL.value,
+    _NCONMAX.value,
+    _NJMAX.value,
   )
 
   name = argv[0]
@@ -99,7 +117,7 @@ def _main(argv: Sequence[str]):
  Total simulation time: {run_time:.2f} s
  Total steps per second: {steps / run_time:,.0f}
  Total realtime factor: {steps * m.opt.timestep / run_time:,.2f} x
- Total time per step: {1e6 * run_time / steps:.2f} µs""")
+ Total time per step: {1e9 * run_time / steps:.2f} ns""")
   elif _OUTPUT.value == "tsv":
     name = name.split("/")[-1].replace("testspeed_", "")
     print(f"{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}")