diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index affe394a..7e3784ba 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -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); diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 0ff17891..86544881 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -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 diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index 35e0b64c..3b3bb90b 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -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 --------------------------------------- diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index aa3d5ca1..ca8489a4 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -20,6 +20,7 @@ #include #include #include +#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) {