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

Separate mj_addM into dense and sparse functions.

Browse Source 
PiperOrigin-RevId: 539066339
Change-Id: Ide3f9644fee46a5fb1a2e79a25be162b929e2421
This commit is contained in:
Kyle Bayes
2023-06-09 06:43:34 -07:00
committed by Copybara-Service
parent 466c9aca82
commit 81d74c56e1
4 changed files with 133 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+2 -2
View File
@@ -1375,7 +1375,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J
mj_addM(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
// factorize H, uncompressed layout
int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
@@ -1404,7 +1404,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
else {
// compute H = M + J'*D*J
mju_sqrMatTD(ctx->H, d->efc_J, D, nefc, nv);
mj_addM(m, d, ctx->H, NULL, NULL, NULL);
mj_addMDense(m, d, ctx->H);
// factorize H
mju_cholFactor(ctx->H, nv, mjMINVAL);
src/engine/engine_support.c
+122 -106
View File
@@ -831,151 +831,167 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
// destination can be sparse uncompressed, or dense when all int* are NULL
void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind) {
int adr, adr1, nv = m->nv;
// sparse
if (rownnz && rowadr && colind) {
// special processing of simple dofs
int simplecnt = 0;
for (int i=0; i < nv; i++) {
if (m->dof_simplenum[i]) {
// count simple
simplecnt++;
mj_addMSparse(m, d, dst, rownnz, rowadr, colind);
}
// empty row: create entry
if (!rownnz[i]) {
colind[rowadr[i]] = i;
dst[rowadr[i]] = d->qM[m->dof_Madr[i]];
rownnz[i] = 1;
}
// dense
else {
mj_addMDense(m, d, dst);
}
}
// non-empty row: assume dof is in dst (J'*D*J satisfies this)
else {
// find dof in row, add
adr = rowadr[i];
int end = adr + rownnz[i];
while (adr < end)
if (colind[adr] == i) {
dst[adr] += d->qM[m->dof_Madr[i]];
break;
} else {
adr++;
}
// not found: error
if (adr >= end) {
mju_error("mj_addM sparse: dst row expected to be empty");
// add inertia matrix to sparse uncompressed destination matrix
void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind) {
int adr, adr1, nv = m->nv;
// special processing of simple dofs
int simplecnt = 0;
for (int i=0; i < nv; i++) {
if (m->dof_simplenum[i]) {
// count simple
simplecnt++;
// empty row: create entry
if (!rownnz[i]) {
colind[rowadr[i]] = i;
dst[rowadr[i]] = d->qM[m->dof_Madr[i]];
rownnz[i] = 1;
}
// non-empty row: assume dof is in dst (J'*D*J satisfies this)
else {
// find dof in row, add
adr = rowadr[i];
int end = adr + rownnz[i];
while (adr < end)
if (colind[adr] == i) {
dst[adr] += d->qM[m->dof_Madr[i]];
break;
} else {
adr++;
}
// not found: error
if (adr >= end) {
mju_error("mj_addM sparse: dst row expected to be empty");
}
}
}
}
// done if all simple
if (simplecnt == nv) {
return;
}
// done if all simple
if (simplecnt == nv) {
return;
}
// allocate space for sparse M
mjMARKSTACK;
mjtNum* M = mj_stackAlloc(d, nv*nv);
int* M_rownnz = (int*) mj_stackAlloc(d, nv);
int* M_rowadr = (int*) mj_stackAlloc(d, nv);
int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
int* buf_ind = (int*) mj_stackAlloc(d, nv);
mjtNum* sparse_buf = mj_stackAlloc(d, nv);
// allocate space for sparse M
mjMARKSTACK;
mjtNum* M = mj_stackAlloc(d, nv*nv);
int* M_rownnz = (int*) mj_stackAlloc(d, nv);
int* M_rowadr = (int*) mj_stackAlloc(d, nv);
int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
int* buf_ind = (int*) mj_stackAlloc(d, nv);
mjtNum* sparse_buf = mj_stackAlloc(d, nv);
// convert M into sparse format, lower-triangular
for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) {
// backward pass over dofs: construct M_row(i) in reverse order
adr = m->dof_Madr[i];
int j = i;
adr1 = 0;
while (j >= 0) {
// assign
M[i*nv+adr1] = d->qM[adr];
M_colind[i*nv+adr1] = j;
// convert M into sparse format, lower-triangular
for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) {
// backward pass over dofs: construct M_row(i) in reverse order
adr = m->dof_Madr[i];
int j = i;
adr1 = 0;
while (j >= 0) {
// assign
M[i*nv+adr1] = d->qM[adr];
M_colind[i*nv+adr1] = j;
// count columns
adr1++;
// count columns
adr1++;
// advance
adr++;
j = m->dof_parentid[j];
}
// advance
adr++;
j = m->dof_parentid[j];
}
// assign row descriptors
M_rownnz[i] = adr1;
M_rowadr[i] = i*nv;
// assign row descriptors
M_rownnz[i] = adr1;
M_rowadr[i] = i*nv;
// reverse order
for (int k=0; k < adr1/2; k++) {
mjtNum tmp = M[i*nv+k];
M[i*nv+k] = M[i*nv+adr1-1-k];
M[i*nv+adr1-1-k] = tmp;
// reverse order
for (int k=0; k < adr1/2; k++) {
mjtNum tmp = M[i*nv+k];
M[i*nv+k] = M[i*nv+adr1-1-k];
M[i*nv+adr1-1-k] = tmp;
int tmpi = M_colind[i*nv+k];
M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
M_colind[i*nv+adr1-1-k] = tmpi;
}
int tmpi = M_colind[i*nv+k];
M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
M_colind[i*nv+adr1-1-k] = tmpi;
}
}
}
// make symmetric
for (int i=1; i < nv; i++) {
if (!m->dof_simplenum[i]) {
for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
// add to row given by column index
adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
M[adr1] = M[k];
M_colind[adr1] = i;
}
// make symmetric
for (int i=1; i < nv; i++) {
if (!m->dof_simplenum[i]) {
for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
// add to row given by column index
adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
M[adr1] = M[k];
M_colind[adr1] = i;
}
}
}
// add to destination
for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) {
int new_nnz =
// add to destination
for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) {
int new_nnz =
mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
rownnz[i], M_rownnz[i],
colind + rowadr[i], M_colind + M_rowadr[i],
sparse_buf, buf_ind);
rownnz[i] = new_nnz;
}
rownnz[i] = new_nnz;
}
mjFREESTACK;
}
// dense
else {
for (int i=0; i < nv; i++) {
adr = m->dof_Madr[i];
int j = i;
while (j >= 0) {
// add
dst[i*nv+j] += d->qM[adr];
if (j < i) {
dst[j*nv+i] += d->qM[adr];
}
mjFREESTACK;
}
// only diagonal if simplenum
if (m->dof_simplenum[i]) {
break;
}
// advance
j = m->dof_parentid[j];
adr++;
// add inertia matrix to dense destination matrix
void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) {
int nv = m->nv;
for (int i = 0; i < nv; i++) {
int adr = m->dof_Madr[i];
int j = i;
while (j >= 0) {
// add
dst[i*nv+j] += d->qM[adr];
if (j < i) {
dst[j*nv+i] += d->qM[adr];
}
// only diagonal if simplenum
if (m->dof_simplenum[i]) {
break;
}
// advance
j = m->dof_parentid[j];
adr++;
}
}
}
//-------------------------- sparse system matrix conversion ---------------------------------------
// dst[D] = src[M], handle different sparsity representations
src/engine/engine_support.h
+7
View File
@@ -104,6 +104,13 @@ MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind);
// add inertia matrix to sparse uncompressed destination matrix
MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind);
// add inertia matrix to dense destination matrix
MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
//-------------------------- sparse system matrix conversion ---------------------------------------
test/benchmark/engine_util_sparse_benchmark_test.cc
+2 -1
View File
@@ -20,6 +20,7 @@
#include <absl/base/attributes.h>
#include <mujoco/mjdata.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_support.h"
#include "src/engine/engine_util_sparse.h"
#include "test/fixture.h"
@@ -458,7 +459,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J
mj_addM(m, d, H, rownnz, rowadr, colind);
mj_addMSparse(m, d, H, rownnz, rowadr, colind);
// time benchmark
for (auto s : state) {