chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 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); d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J // 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 // factorize H, uncompressed layout
int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL, int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
@@ -1404,7 +1404,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
else { else {
// compute H = M + J'*D*J // compute H = M + J'*D*J
mju_sqrMatTD(ctx->H, d->efc_J, D, nefc, nv); 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 // factorize H
mju_cholFactor(ctx->H, nv, mjMINVAL); 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 // destination can be sparse uncompressed, or dense when all int* are NULL
void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind) { int* rownnz, int* rowadr, int* colind) {
int adr, adr1, nv = m->nv;
// sparse // sparse
if (rownnz && rowadr && colind) { if (rownnz && rowadr && colind) {
// special processing of simple dofs mj_addMSparse(m, d, dst, rownnz, rowadr, colind);
int simplecnt = 0; }
for (int i=0; i < nv; i++) {
if (m->dof_simplenum[i]) {
// count simple
simplecnt++;
// empty row: create entry // dense
if (!rownnz[i]) { else {
colind[rowadr[i]] = i; mj_addMDense(m, d, dst);
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) { // add inertia matrix to sparse uncompressed destination matrix
mju_error("mj_addM sparse: dst row expected to be empty"); 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 // done if all simple
if (simplecnt == nv) { if (simplecnt == nv) {
return; return;
} }
// allocate space for sparse M // allocate space for sparse M
mjMARKSTACK; mjMARKSTACK;
mjtNum* M = mj_stackAlloc(d, nv*nv); mjtNum* M = mj_stackAlloc(d, nv*nv);
int* M_rownnz = (int*) mj_stackAlloc(d, nv); int* M_rownnz = (int*) mj_stackAlloc(d, nv);
int* M_rowadr = (int*) mj_stackAlloc(d, nv); int* M_rowadr = (int*) mj_stackAlloc(d, nv);
int* M_colind = (int*) mj_stackAlloc(d, nv*nv); int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
int* buf_ind = (int*) mj_stackAlloc(d, nv); int* buf_ind = (int*) mj_stackAlloc(d, nv);
mjtNum* sparse_buf = mj_stackAlloc(d, nv); mjtNum* sparse_buf = mj_stackAlloc(d, nv);
// convert M into sparse format, lower-triangular // convert M into sparse format, lower-triangular
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) { if (!m->dof_simplenum[i]) {
// backward pass over dofs: construct M_row(i) in reverse order // backward pass over dofs: construct M_row(i) in reverse order
adr = m->dof_Madr[i]; adr = m->dof_Madr[i];
int j = i; int j = i;
adr1 = 0; adr1 = 0;
while (j >= 0) { while (j >= 0) {
// assign // assign
M[i*nv+adr1] = d->qM[adr]; M[i*nv+adr1] = d->qM[adr];
M_colind[i*nv+adr1] = j; M_colind[i*nv+adr1] = j;
// count columns // count columns
adr1++; adr1++;
// advance // advance
adr++; adr++;
j = m->dof_parentid[j]; j = m->dof_parentid[j];
} }
// assign row descriptors // assign row descriptors
M_rownnz[i] = adr1; M_rownnz[i] = adr1;
M_rowadr[i] = i*nv; M_rowadr[i] = i*nv;
// reverse order // reverse order
for (int k=0; k < adr1/2; k++) { for (int k=0; k < adr1/2; k++) {
mjtNum tmp = M[i*nv+k]; mjtNum tmp = M[i*nv+k];
M[i*nv+k] = M[i*nv+adr1-1-k]; M[i*nv+k] = M[i*nv+adr1-1-k];
M[i*nv+adr1-1-k] = tmp; M[i*nv+adr1-1-k] = tmp;
int tmpi = M_colind[i*nv+k]; int tmpi = M_colind[i*nv+k];
M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k]; M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
M_colind[i*nv+adr1-1-k] = tmpi; M_colind[i*nv+adr1-1-k] = tmpi;
}
} }
} }
}
// make symmetric // make symmetric
for (int i=1; i < nv; i++) { for (int i=1; i < nv; i++) {
if (!m->dof_simplenum[i]) { if (!m->dof_simplenum[i]) {
for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) { for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
// add to row given by column index // add to row given by column index
adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++; adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
M[adr1] = M[k]; M[adr1] = M[k];
M_colind[adr1] = i; M_colind[adr1] = i;
}
} }
} }
}
// add to destination // add to destination
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
if (!m->dof_simplenum[i]) { if (!m->dof_simplenum[i]) {
int new_nnz = int new_nnz =
mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1, mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
rownnz[i], M_rownnz[i], rownnz[i], M_rownnz[i],
colind + rowadr[i], M_colind + M_rowadr[i], colind + rowadr[i], M_colind + M_rowadr[i],
sparse_buf, buf_ind); sparse_buf, buf_ind);
rownnz[i] = new_nnz; rownnz[i] = new_nnz;
}
} }
mjFREESTACK;
} }
// dense mjFREESTACK;
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];
}
// only diagonal if simplenum
if (m->dof_simplenum[i]) {
break;
}
// advance
j = m->dof_parentid[j]; // add inertia matrix to dense destination matrix
adr++; 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 --------------------------------------- //-------------------------- sparse system matrix conversion ---------------------------------------
// dst[D] = src[M], handle different sparsity representations // 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, MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind); 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 --------------------------------------- //-------------------------- 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 <absl/base/attributes.h>
#include <mujoco/mjdata.h> #include <mujoco/mjdata.h>
#include <mujoco/mujoco.h> #include <mujoco/mujoco.h>
#include "src/engine/engine_support.h"
#include "src/engine/engine_util_sparse.h" #include "src/engine/engine_util_sparse.h"
#include "test/fixture.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); d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J // compute H = M + J'*D*J
mj_addM(m, d, H, rownnz, rowadr, colind); mj_addMSparse(m, d, H, rownnz, rowadr, colind);
// time benchmark // time benchmark
for (auto s : state) { for (auto s : state) {