Refactor and speed up mj_crb and add a benchmark.

This is a no-op "prefactor" of `mj_crb` to reduce the number of lines changed in the upcoming sleeping CL. The main change here is that `mj_crb` now avoids accessing the model and data pointer repeatedly, but instead has the input and output pointers explicitly declared as local variables. A benchmark test found a healthy **14% perf bump** due to two changes:
- Adding `restrict` to `mju_mulInertVec`
- The function-local pointers.

Adding `restrict` to the local pointers had no effect. Note that `mj_crb` is not a particularly expensive function so these speed bumps are not significant per se, but rather indicative of possible future gains with these techniques.

```
Benchmark                Time(ns)        CPU(ns)     Iterations
---------------------------------------------------------------

ORGINAL BASELINE
BM_CRB_BASELINE_mean         4325           4348         993200  230.052k items/s

BASELINE + RESTRICT
BM_CRB_BASELINE_mean         4057           4090        1186150  244.573k items/s

LOCAL POINTERS + RESTRICT
BM_CRB_mean                  3772           3800        1001750  263.209k items/s
```

PiperOrigin-RevId: 815798225
Change-Id: Iffdf57b544e8e10562807c617f73ca4ccd1c614f

src/engine/engine_core_smooth.c

@@ -1513,30 +1513,44 @@ void mj_tendonArmature(const mjModel* m, mjData* d) {
 
 // composite rigid body inertia algorithm
 void mj_crb(const mjModel* m, mjData* d) {
-  int nv = m->nv;
-  mjtNum buf[6];
+  int nv = m->nv, nbody = m->nbody;
+
+  // outputs
   mjtNum* crb = d->crb;
+  mjtNum* M   = d->M;
+
+  // inputs
+  const mjtNum* cinert       = d->cinert;
+  const mjtNum* cdof         = d->cdof;
+  const mjtNum* dof_M0       = m->dof_M0;
+  const mjtNum* dof_armature = m->dof_armature;
+  const int* rownnz          = m->M_rownnz;
+  const int* rowadr          = m->M_rowadr;
+  const int* body_parentid   = m->body_parentid;
+  const int* dof_parentid    = m->dof_parentid;
+  const int* dof_simplenum   = m->dof_simplenum;
+  const int* dof_bodyid      = m->dof_bodyid;
 
   // crb = cinert
-  mju_copy(crb, d->cinert, 10*m->nbody);
+  mju_copy(crb, cinert, 10*nbody);
 
   // backward pass over bodies, accumulate composite inertias
-  for (int i=m->nbody - 1; i > 0; i--) {
-    if (m->body_parentid[i] > 0) {
-      mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10);
+  for (int i=nbody - 1; i > 0; i--) {
+    if (body_parentid[i]) {
+      mju_addTo(crb + 10*body_parentid[i], crb + 10*i, 10);
     }
   }
 
   // clear M
-  mju_zero(d->M, m->nC);
+  mju_zero(M, m->nC);
 
   // dense forward pass over dofs
   for (int i=0; i < nv; i++) {
     // process block of diagonals (simple bodies)
-    if (m->dof_simplenum[i]) {
-      int n = i + m->dof_simplenum[i];
+    if (dof_simplenum[i]) {
+      int n = i + dof_simplenum[i];
       for (; i < n; i++) {
-        d->M[m->M_rowadr[i]] = m->dof_M0[i];
+        M[rowadr[i]] = dof_M0[i];
       }
 
       // finish or else fall through with next row
@@ -1546,16 +1560,17 @@ void mj_crb(const mjModel* m, mjData* d) {
     }
 
     // init M(i,i) with armature inertia
-    int Madr_ij = m->M_rowadr[i] + m->M_rownnz[i] - 1;
-    d->M[Madr_ij] = m->dof_armature[i];
+    int Madr_ij = rowadr[i] + rownnz[i] - 1;
+    M[Madr_ij] = dof_armature[i];
 
     // precompute buf = crb_body_i * cdof_i
-    mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i);
+    mjtNum buf[6];
+    mju_mulInertVec(buf, crb+10*dof_bodyid[i], cdof+6*i);
 
     // sparse backward pass over ancestors
-    for (int j=i; j >= 0; j = m->dof_parentid[j]) {
+    for (int j=i; j >= 0; j = dof_parentid[j]) {
       // M(i,j) += cdof_j * (crb_body_i * cdof_i)
-      d->M[Madr_ij--] += mju_dot(d->cdof+6*j, buf, 6);
+      M[Madr_ij--] += mju_dot(cdof+6*j, buf, 6);
     }
   }
 }

src/engine/engine_util_spatial.c

@@ -431,7 +431,7 @@ void mju_inertCom(mjtNum res[10], const mjtNum inert[3], const mjtNum mat[9],
 
 
 // multiply 6D vector (rotation, translation) by 6D inertia matrix
-void mju_mulInertVec(mjtNum res[6], const mjtNum i[10], const mjtNum v[6]) {
+void mju_mulInertVec(mjtNum* restrict res, const mjtNum i[10], const mjtNum v[6]) {
   res[0] = i[0]*v[0] + i[3]*v[1] + i[4]*v[2] - i[8]*v[4] + i[7]*v[5];
   res[1] = i[3]*v[0] + i[1]*v[1] + i[5]*v[2] + i[8]*v[3] - i[6]*v[5];
   res[2] = i[4]*v[0] + i[5]*v[1] + i[2]*v[2] - i[7]*v[3] + i[6]*v[4];

src/engine/engine_util_spatial.h

@@ -102,7 +102,7 @@ void mju_inertCom(mjtNum res[10], const mjtNum inert[3], const mjtNum mat[9],
 void mju_dofCom(mjtNum res[6], const mjtNum axis[3], const mjtNum offset[3]);
 
 // multiply 6D vector (rotation, translation) by 6D inertia matrix
-void mju_mulInertVec(mjtNum res[6], const mjtNum inert[10], const mjtNum vec[6]);
+MJAPI void mju_mulInertVec(mjtNum res[6], const mjtNum inert[10], const mjtNum vec[6]);
 
 // multiply dof matrix by vector
 void mju_mulDofVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);

test/benchmark/CMakeLists.txt

@@ -67,6 +67,12 @@ mujoco_test(
   ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
 )
 
+mujoco_test(
+  crb_benchmark_test
+  MAIN_TARGET benchmark::benchmark_main
+  ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers
+)
+
 mujoco_test(
   engine_util_sparse_benchmark_test
   MAIN_TARGET benchmark::benchmark_main

test/benchmark/crb_benchmark_test.cc

@@ -0,0 +1,124 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// A benchmark for comparing different implementations of mj_crb.
+
+#include <benchmark/benchmark.h>
+#include <absl/base/attributes.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mujoco.h>
+#include "src/engine/engine_core_smooth.h"
+#include "src/engine/engine_util_spatial.h"
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+// number of steps to benchmark
+static const int kNumBenchmarkSteps = 50;
+
+// ----------------------------- previous implementation -----------------------
+
+void ABSL_ATTRIBUTE_NOINLINE mj_crb_baseline(const mjModel* m, mjData* d) {
+  int nv = m->nv;
+  mjtNum buf[6];
+  mjtNum* crb = d->crb;
+
+  // crb = cinert
+  mju_copy(crb, d->cinert, 10*m->nbody);
+
+  // backward pass over bodies, accumulate composite inertias
+  for (int i=m->nbody - 1; i > 0; i--) {
+    if (m->body_parentid[i] > 0) {
+      mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10);
+    }
+  }
+
+  // clear M
+  mju_zero(d->M, m->nC);
+
+  // dense forward pass over dofs
+  for (int i=0; i < nv; i++) {
+    // process block of diagonals (simple bodies)
+    if (m->dof_simplenum[i]) {
+      int n = i + m->dof_simplenum[i];
+      for (; i < n; i++) {
+        d->M[m->M_rowadr[i]] = m->dof_M0[i];
+      }
+
+      // finish or else fall through with next row
+      if (i == nv) {
+        break;
+      }
+    }
+
+    // init M(i,i) with armature inertia
+    int Madr_ij = m->M_rowadr[i] + m->M_rownnz[i] - 1;
+    d->M[Madr_ij] = m->dof_armature[i];
+
+    // precompute buf = crb_body_i * cdof_i
+    mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i);
+
+    // sparse backward pass over ancestors
+    for (int j=i; j >= 0; j = m->dof_parentid[j]) {
+      // M(i,j) += cdof_j * (crb_body_i * cdof_i)
+      d->M[Madr_ij--] += mju_dot(d->cdof+6*j, buf, 6);
+    }
+  }
+}
+
+
+// ----------------------------- benchmark ------------------------------------
+
+static void BM_crb(benchmark::State& state, bool baseline) {
+  static mjModel* m = nullptr;
+  static mjData* d = nullptr;
+
+  if (!m) {
+    m = LoadModelFromPath("../test/benchmark/testdata/inertia.xml");
+    d = mj_makeData(m);
+  }
+  mj_forward(m, d);
+
+  // benchmark
+  while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
+    if (baseline) {
+      for (int i=0; i < kNumBenchmarkSteps; i++) {
+        mj_crb_baseline(m, d);
+      }
+    } else {
+      for (int i=0; i < kNumBenchmarkSteps; i++) {
+        mj_crb(m, d);
+      }
+    }
+  }
+
+  // finalize
+  state.SetItemsProcessed(state.iterations());
+}
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_CRB(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_crb(state, /*baseline=*/false);
+}
+BENCHMARK(BM_CRB);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_CRB_BASELINE(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_crb(state, /*baseline=*/true);
+}
+BENCHMARK(BM_CRB_BASELINE);
+
+}  // namespace
+}  // namespace mujoco