chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean up mj_mulM.

Browse Source 
PiperOrigin-RevId: 561404770
Change-Id: Ie802f5d7f1df6bf32d4549d3db6809d0f507d1e1
This commit is contained in:
Yuval Tassa
2023-08-30 11:46:10 -07:00
committed by Copybara-Service
parent f371ca7a4e
commit ac150f751d
1 changed files with 25 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_support.c
+25 -26
View File
@@ -827,25 +827,27 @@ void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) {
// multiply vector by inertia matrix // multiply vector by inertia matrix
void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
int adr, nv = m->nv; int nv = m->nv;
const mjtNum* M = d->qM; const mjtNum* M = d->qM;
const int* dofMadr = m->dof_Madr; const int* Madr = m->dof_Madr;
const int* parentid = m->dof_parentid;
const int* simplenum = m->dof_simplenum;
mju_zero(res, nv); mju_zero(res, nv);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
#ifdef mjUSEAVX #ifdef mjUSEAVX
// simple: diagonal division, AVX // simple: diagonal multiplication, AVX
if (m->dof_simplenum[i] >= 4) { if (simplenum[i] >= 4) {
// init // init
__m256d result, val1, val2; __m256d result, val1, val2;
// parallel computation // parallel computation
val1 = _mm256_loadu_pd(vec+i); val1 = _mm256_loadu_pd(vec+i);
val2 = _mm256_set_pd(M[dofMadr[i+3]], val2 = _mm256_set_pd(M[Madr[i+3]],
M[dofMadr[i+2]], M[Madr[i+2]],
M[dofMadr[i+1]], M[Madr[i+1]],
M[dofMadr[i+0]]); M[Madr[i+0]]);
result = _mm256_mul_pd(val1, val2); result = _mm256_mul_pd(val1, val2);
// store result // store result
@@ -856,29 +858,26 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
continue; continue;
} }
#endif #endif
// address in M
int adr = Madr[i];
// simple: diagonal multiplication // compute diagonal
if (m->dof_simplenum[i]) { res[i] = M[adr]*vec[i];
res[i] = M[dofMadr[i]]*vec[i];
// simple dof: continue
if (simplenum[i]) {
continue;
} }
// regular: full multiplication // compute off-diagonals
else { int j = parentid[i];
// diagonal while (j >= 0) {
adr = dofMadr[i];
res[i] += M[adr]*vec[i];
// off-diagonal
int j = m->dof_parentid[i];
adr++; adr++;
while (j >= 0) { res[i] += M[adr]*vec[j];
res[i] += M[adr]*vec[j]; res[j] += M[adr]*vec[i];
res[j] += M[adr]*vec[i];
// advance to next element // advance to parent
j = m->dof_parentid[j]; j = parentid[j];
adr++;
}
} }
} }
} }