diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 65fec4e0..3ed6f063 100644 --- a/src/engine/engine_core_smooth.c +++ b/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); } } } diff --git a/src/engine/engine_util_spatial.c b/src/engine/engine_util_spatial.c index 1e433467..a20f0b7e 100644 --- a/src/engine/engine_util_spatial.c +++ b/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]; diff --git a/src/engine/engine_util_spatial.h b/src/engine/engine_util_spatial.h index e27dfa57..c90cb27f 100644 --- a/src/engine/engine_util_spatial.h +++ b/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); diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index dc56b611..3f0be08e 100644 --- a/test/benchmark/CMakeLists.txt +++ b/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 diff --git a/test/benchmark/crb_benchmark_test.cc b/test/benchmark/crb_benchmark_test.cc new file mode 100644 index 00000000..bd643fc8 --- /dev/null +++ b/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 +#include +#include +#include +#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