diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index eeb25a50..c867ec95 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -397,13 +397,7 @@ static int mju_compare(const int* vec1, const int* vec2, int n) { #endif // scalar part - for (; i #include #include #include @@ -25,6 +26,8 @@ namespace mujoco { namespace { +using FuncPtr = decltype(&mju_combineSparse); + // number of steps to roll out before benhmarking static const int kNumWarmupSteps = 500; @@ -35,6 +38,72 @@ std::vector AsVector(const mjtNum* array, int n) { // ----------------------------- old functions -------------------------------- +int compare_baseline(const int* vec1, + const int* vec2, + int n) { + int i = 0; + + for (; i < n; i++) { + if (vec1[i] != vec2[i]) { + return 0; + } + } + + return 1; +} + +void addToSclScl(mjtNum* res, + const mjtNum* vec, + mjtNum scl1, + mjtNum scl2, + int n) { + int i = 0; + + for (; i < n; i++) { + res[i] = res[i]*scl1 + vec[i]*scl2; + } +} + +int compare_memcmp(const int* vec1, + const int* vec2, + int n) { + return !memcmp(vec1, vec2, n*sizeof(int)); +} + +int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst, + const mjtNum* src, int n, + mjtNum a, mjtNum b, + int dst_nnz, int src_nnz, + int* dst_ind, + const int* src_ind, + mjtNum* buf, int* buf_ind) { + // check for identical pattern + if (compare_baseline(dst_ind, src_ind, dst_nnz)) { + // combine mjtNum data directly + addToSclScl(dst, src, a, b, dst_nnz); + return dst_nnz; + } else { + return 0; + } +} + +int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst, + const mjtNum* src, int n, + mjtNum a, mjtNum b, + int dst_nnz, int src_nnz, + int* dst_ind, + const int* src_ind, + mjtNum* buf, int* buf_ind) { + // check for identical pattern + if (compare_memcmp(dst_ind, src_ind, dst_nnz)) { + // combine mjtNum data directly + addToSclScl(dst, src, a, b, dst_nnz); + return dst_nnz; + } else { + return 0; + } +} + mjtNum ABSL_ATTRIBUTE_NOINLINE dotSparse_1(const mjtNum* vec1, const mjtNum* vec2, const int nnz1, @@ -217,5 +286,78 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1( } BENCHMARK(BM_MatVecSparse_1); +static void BM_combineSparse(benchmark::State& state, FuncPtr func) { + static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml"); + mjData* d = mj_makeData(m); + + // warm-up rollout to get a typical state + for (int i=0; i < kNumWarmupSteps; i++) { + mj_step(m, d); + } + + // allocate + mjMARKSTACK; + mjtNum* H = mj_stackAlloc(d, m->nv*m->nv); + int* rownnz = (int*)mj_stackAlloc(d, m->nv); + int* rowadr = (int*)mj_stackAlloc(d, m->nv); + int* colind = (int*)mj_stackAlloc(d, m->nv*m->nv); + + // compute D corresponding to quad states + mjtNum* D = mj_stackAlloc(d, d->nefc); + for (int i = 0; i < d->nefc; i++) { + if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) { + D[i] = d->efc_D[i]; + } else { + D[i] = 0; + } + } + + // compute H = J'*D*J, uncompressed layout + 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); + + // compute H = M + J'*D*J + mj_addM(m, d, H, rownnz, rowadr, colind); + + // time benchmark + for (auto s : state) { + for (int r = m->nv-1; r >= 0; r--) { + for (int i = 0; i < rownnz[r]-1; i++) { + int adr = rowadr[r]; + int c = colind[adr+i]; + // true arguments should be i+1 and colind+rowadr[r] + // but instead we repeat rownnz[c] and colind+rowadr[c] + // in order to trigger all if's in combineSparse + func(H+rowadr[c], H+rowadr[r], c+1, 1, -H[adr+i], + rownnz[c], rownnz[c], + colind+rowadr[c], colind+rowadr[c], NULL, NULL); + } + } + } + + // finalize + mjFREESTACK; + mj_deleteData(d); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_new( + benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_combineSparse(state, &combineSparse_new); +} +BENCHMARK(BM_combineSparse_new); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old( + benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_combineSparse(state, &combineSparse_baseline); +} +BENCHMARK(BM_combineSparse_old); + } // namespace } // namespace mujoco