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
This commit is contained in:
Yuval Tassa
2025-10-06 10:53:09 -07:00
committed by Copybara-Service
parent b75af81c75
commit ab9575d8a6
5 changed files with 162 additions and 17 deletions
+30 -15
View File
@@ -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);
}
}
}
+1 -1
View File
@@ -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];
+1 -1
View File
@@ -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);