Unrolled matmul for small matrices in MJX improves steps/sec by 10-20%.

PiperOrigin-RevId: 605033188
Change-Id: I9da723d01e1f6511abaf4522b4c8f829a8091772

mjx/mujoco/mjx/_src/math.py

@@ -20,6 +20,30 @@ import jax
 from jax import numpy as jp
 
 
+def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array:
+  """Calculates a @ b via explicit cell value operations.
+
+  This is faster than XLA matmul for small matrices (e.g. 3x3, 4x4).
+
+  Args:
+    a: left hand of matmul operand
+    b: right hand of matmul operand
+  Returns:
+    the matrix product of the inputs.
+  """
+  c = []
+  for i in range(a.shape[0]):
+    row = []
+    for j in range(b.shape[1]):
+      s = 0.0
+      for k in range(a.shape[1]):
+        s += a[i, k] * b[k, j]
+      row.append(s)
+    c.append(row)
+
+  return jp.array(c)
+
+
 def norm(
     x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None
 ) -> jax.Array:

mjx/mujoco/mjx/_src/smooth.py

@@ -146,7 +146,8 @@ def com_pos(m: Model, d: Data) -> Data:
   @jax.vmap
   def inert_com(inert, ximat, off, mass):
     h = jp.cross(off, -jp.eye(3))
-    inert = (ximat * inert) @ ximat.T + h @ h.T * mass
+    inert = math.matmul_unroll((ximat * inert), ximat.T)
+    inert += math.matmul_unroll(h, h.T) * mass
     # cinert is triu(inert), mass * off, mass
     inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
     return jp.concatenate([inert, off * mass, mass[None]])