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

Implement a more performant mju_transposeSparse that doesn't require dense memory allocation.

Browse Source 
PiperOrigin-RevId: 505044194
Change-Id: Ibe0e39e3ad711b2b5bdad45180ca4f83b5c8fc82
This commit is contained in:
Kyle Bayes
2023-01-27 00:19:14 -08:00
committed by Copybara-Service
parent 4b4ee27f64
commit c741dfce7d
5 changed files with 255 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/engine_util_sparse_benchmark_test.cc
+73 -3
View File
@@ -26,9 +26,10 @@
namespace mujoco {
namespace {
using FuncPtr = decltype(&mju_combineSparse);
using CombineFuncPtr = decltype(&mju_combineSparse);
using TransposeFuncPtr = decltype(&mju_transposeSparse);
// number of steps to roll out before benhmarking
// number of steps to roll out before benchmarking
static const int kNumWarmupSteps = 500;
// copy array into vector
@@ -38,6 +39,27 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
// ----------------------------- old functions --------------------------------
// 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) {
memset(res_rownnz, 0, nc * sizeof(int));
for (int rt = 0; rt < nc; rt++) {
res_rowadr[rt] = rt * nr;
}
for (int r = 0; r < nr; r++) {
for (int ci = 0; ci < rownnz[r]; ci++) {
int rt = colind[rowadr[r] + ci];
res_colind[rt * nr + res_rownnz[rt]] = r;
res[rt * nr + res_rownnz[rt]] = mat[rowadr[r] + ci];
res_rownnz[rt]++;
}
}
mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind);
}
int compare_baseline(const int* vec1,
const int* vec2,
int n) {
@@ -286,8 +308,9 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1(
}
BENCHMARK(BM_MatVecSparse_1);
static void BM_combineSparse(benchmark::State& state, FuncPtr func) {
static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
mjData* d = mj_makeData(m);
// warm-up rollout to get a typical state
@@ -359,5 +382,52 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
}
BENCHMARK(BM_combineSparse_old);
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
static mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml");
// force use of sparse matrices
m->opt.jacobian = mjJAC_SPARSE;
mjData* d = mj_makeData(m);
// warm-up rollout to get a typical state
while (d-> time < 2) {
mj_step(m, d);
}
mjMARKSTACK;
// need uncompressed layout
mjtNum* res = mj_stackAlloc(d, m->nv * d->nefc);
int* res_rownnz = (int*)mj_stackAlloc(d, m->nv);
int* res_rowadr = (int*)mj_stackAlloc(d, m->nv);
int* res_colind = (int*)mj_stackAlloc(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,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
}
mjFREESTACK;
mj_deleteData(d);
state.SetItemsProcessed(state.iterations());
}
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_new(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse);
}
BENCHMARK(BM_transposeSparse_new);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_old(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &transposeSparse_baseline);
}
BENCHMARK(BM_transposeSparse_old);
} // namespace
} // namespace mujoco