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

Speed up sparse supernode detection by combining it with transposition.

Browse Source 
PiperOrigin-RevId: 760673008
Change-Id: I6d66580e675fd86e5b974859383f93f87482ca16
This commit is contained in:
Yuval Tassa
2025-05-19 10:13:20 -07:00
committed by Copybara-Service
parent ced630181d
commit 45fc15b844
6 changed files with 115 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/engine_util_sparse_benchmark_test.cc
+36 -8
View File
@@ -152,8 +152,8 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
// transpose sparse matrix (uncompressed)
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz,
int* res_rowadr, int* res_colind, const int* rownnz, const int* rowadr,
const int* colind) {
int* res_rowadr, int* res_colind, int* res_rowsuper,
const int* rownnz, const int* rowadr, const int* colind) {
memset(res_rownnz, 0, nc * sizeof(int));
for (int rt = 0; rt < nc; rt++) {
res_rowadr[rt] = rt * nr;
@@ -497,7 +497,14 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
}
BENCHMARK(BM_combineSparse_old);
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
enum class Supernode {
None,
PostProcess,
Inline
};
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
Supernode super) {
static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
// force use of sparse matrices
@@ -516,12 +523,19 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
mjtNum* res = mj_stackAllocNum(d, m->nv * d->nefc);
int* res_rownnz = mj_stackAllocInt(d, m->nv);
int* res_rowadr = mj_stackAllocInt(d, m->nv);
int* res_rowsuper = mj_stackAllocInt(d, m->nv);
int* res_colind = mj_stackAllocInt(d, m->nv * d->nefc);
// time benchmark
for (auto s : state) {
func(res, d->efc_J, d->nefc, m->nv, res_rownnz, res_rowadr, res_colind,
int* rowsuper = (super == Supernode::Inline) ? res_rowsuper : nullptr;
func(res, d->efc_J, d->nefc, m->nv,
res_rownnz, res_rowadr, res_colind, rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
if (super == Supernode::PostProcess) {
mju_superSparse(m->nv, res_rowsuper,
res_rownnz, res_rowadr, res_colind);
}
}
mj_freeStack(d);
@@ -529,19 +543,33 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
state.SetItemsProcessed(state.iterations());
}
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_old(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &transposeSparse_baseline, Supernode::None);
}
BENCHMARK(BM_transposeSparse_old);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_new(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse);
BM_transposeSparse(state, &mju_transposeSparse, Supernode::None);
}
BENCHMARK(BM_transposeSparse_new);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_old(benchmark::State& state) {
BM_transposeSparse_superpost(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &transposeSparse_baseline);
BM_transposeSparse(state, &mju_transposeSparse, Supernode::PostProcess);
}
BENCHMARK(BM_transposeSparse_old);
BENCHMARK(BM_transposeSparse_superpost);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_superinline(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse, Supernode::Inline);
}
BENCHMARK(BM_transposeSparse_superinline);
static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
static mjModel* m =