From 2d7d5319f7a01fabd29f38ed64c9447ae100fd17 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Mon, 7 Aug 2023 06:06:33 -0700 Subject: [PATCH] Rework sparse addM function to support compressed sparse matrices. 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 --- src/engine/engine_solver.c | 11 +- src/engine/engine_support.c | 165 ++++++++---------- src/engine/engine_support.h | 9 +- .../engine_util_sparse_benchmark_test.cc | 12 +- 4 files changed, 93 insertions(+), 104 deletions(-) diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index f8fd516f..6985dfb6 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -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, diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 46fdcb5e..00367254 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -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; } diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index dd1c5c7b..d80cf079 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -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); diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index 06687399..45a13df5 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -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) {