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

Refactor mj_mulM to use CSR structure

Browse Source 
PiperOrigin-RevId: 796450207
Change-Id: I2a9d567732614674547c7ae02ff0d3d4bbf1b654
This commit is contained in:
Yuval Tassa
2025-08-18 09:08:45 -07:00
committed by Copybara-Service
parent c2ff60e630
commit 424226b640
4 changed files with 3 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+1 -4
View File
@@ -474,10 +474,7 @@ pedagogical examples.
.. mujoco-include:: mj_mulM
This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format
that the user should not attempt to manipulate directly. Alternatively one can convert qM to a dense matrix with
mj_fullM and then user regular matrix-vector multiplication, but this is slower because it no longer benefits from
sparsity.
This function multiplies the joint-space inertia matrix stored in ``mjData.M`` by a vector.
.. _mj_mulM2:
doc/APIreference/functions_override.rst
+1 -4
View File
@@ -276,10 +276,7 @@ pedagogical examples.
.. _mj_mulM:
This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format
that the user should not attempt to manipulate directly. Alternatively one can convert qM to a dense matrix with
mj_fullM and then user regular matrix-vector multiplication, but this is slower because it no longer benefits from
sparsity.
This function multiplies the joint-space inertia matrix stored in ``mjData.M`` by a vector.
.. _mj_applyFT:
src/engine/engine_support.c
+1 -55
View File
@@ -983,63 +983,9 @@ void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) {
// multiply vector by inertia matrix (implementation)
void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M,
const int* Madr, const int* parentid, const int* simplenum) {
mju_zero(res, nv);
for (int i=0; i < nv; i++) {
#ifdef mjUSEAVX
// simple: diagonal multiplication, AVX
if (simplenum[i] >= 4) {
// init
__m256d result, val1, val2;
// parallel computation
val1 = _mm256_loadu_pd(vec+i);
val2 = _mm256_set_pd(M[Madr[i+3]],
M[Madr[i+2]],
M[Madr[i+1]],
M[Madr[i+0]]);
result = _mm256_mul_pd(val1, val2);
// store result
_mm256_storeu_pd(res+i, result);
// skip rest of block
i += 3;
continue;
}
#endif
// address in M
int adr = Madr[i];
// compute diagonal
res[i] = M[adr]*vec[i];
// simple dof: continue
if (simplenum[i]) {
continue;
}
// compute off-diagonals
int j = parentid[i];
while (j >= 0) {
adr++;
res[i] += M[adr]*vec[j];
res[j] += M[adr]*vec[i];
// advance to parent
j = parentid[j];
}
}
}
// multiply vector by inertia matrix
void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
mj_mulM_impl(res, vec, m->nv, d->qM, m->dof_Madr, m->dof_parentid, m->dof_simplenum);
mju_mulSymVecSparse(res, d->M, vec, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind);
}
src/engine/engine_support.h
-4
View File
@@ -123,10 +123,6 @@ MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
// convert sparse inertia matrix M into full matrix
MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
// multiply vector by inertia matrix (implementation)
MJAPI void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M,
const int* Madr, const int* parentid, const int* simplenum);
// multiply vector by inertia matrix
MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);