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

Rework sparse addM function to support compressed sparse matrices.

Browse Source 
Note:

- M is still a nv x nv uncompressed sparse matrix.
- NNZ precounting still needs to be implemented using the combineSparseCount helper function.

PiperOrigin-RevId: 554452710
Change-Id: I14adf58fc3acd4ed82d864d7c2e9e96ee6b0377d
This commit is contained in:
Kyle Bayes
2023-08-07 06:06:33 -07:00
committed by Copybara-Service
parent 6d7b49ccf5
commit 2d7d5319f7
4 changed files with 93 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+10 -1
View File
@@ -1364,7 +1364,15 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// sparse
if (mj_isSparse(m)) {
// create sparse inertia matrix M (uncompressed)
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);
mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
// compute H = J'*D*J
// TODO(b/266802572): remove uncompressed layout
mju_sqrMatTDUncompressedInit(ctx->rowadr, nv);
mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
@@ -1375,7 +1383,8 @@ 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_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind);
mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind,
M, M_rownnz, M_rowadr, M_colind);
// factorize H, uncompressed layout
int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
src/engine/engine_support.c
+70 -95
View File
@@ -948,7 +948,17 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind) {
// sparse
if (rownnz && rowadr && colind) {
mj_addMSparse(m, d, dst, rownnz, rowadr, colind);
int nv = m->nv;
mjMARKSTACK;
// create sparse inertia matrix M (uncompressed)
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);
mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
mj_addMSparse(m, d, dst, rownnz, rowadr, colind, M,
M_rownnz, M_rowadr, M_colind);
mjFREESTACK;
}
// dense
@@ -959,122 +969,87 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
// add inertia matrix to sparse uncompressed destination matrix
void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind) {
// create inertia matrix M (uncompressed)
void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind) {
int adr, adr1, nv = m->nv;
// special processing of simple dofs
int simplecnt = 0;
// build M into sparse format, lower-triangular
for (int i=0; i < nv; i++) {
M_rowadr[i] = i*nv;
adr = m->dof_Madr[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) {
mjERROR("dst row expected to be empty");
}
}
M[i*nv] = d->qM[adr];
M_colind[i*nv] = i;
M_rownnz[i] = 1;
continue;
}
}
// done if all simple
if (simplecnt == nv) {
return;
}
// backward pass over dofs: construct M_row(i) in reverse order
int j = i;
adr1 = 0;
while (j >= 0) {
// assign
M[i*nv+adr1] = d->qM[adr];
M_colind[i*nv+adr1] = j;
// 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);
// count columns
adr1++;
// 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;
// advance
adr++;
j = m->dof_parentid[j];
}
// count columns
adr1++;
// assign row descriptors
M_rownnz[i] = adr1;
// advance
adr++;
j = m->dof_parentid[j];
}
// 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;
// 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;
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;
}
if (m->dof_simplenum[i]) {
continue;
}
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 inertia matrix to sparse destination matrix
void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind) {
int nv = m->nv;
mjMARKSTACK;
int* buf_ind = (int*) mj_stackAlloc(d, nv);
mjtNum* sparse_buf = mj_stackAlloc(d, nv);
// 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] = 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);
}
mjFREESTACK;
}
src/engine/engine_support.h
+7 -2
View File
@@ -116,9 +116,14 @@ 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
// create inertia matrix M (uncompressed)
MJAPI void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind);
// add inertia matrix to sparse destination matrix
MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind);
int* rownnz, int* rowadr, int* colind, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind);
// add inertia matrix to dense destination matrix
MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
test/benchmark/engine_util_sparse_benchmark_test.cc
+6 -6
View File
@@ -452,14 +452,14 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
// compute H = J'*D*J, uncompressed layout
mju_sqrMatTDUncompressedInit(rowadr, m->nv);
mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv,
rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr,
d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d);
rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr,
d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J
mj_addMSparse(m, d, H, rownnz, rowadr, colind);
mj_addM(m, d, H, rownnz, rowadr, colind);
// time benchmark
for (auto s : state) {