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

Use memcmp for mju_compare if no AVX. 15% faster on Humanoid.

Browse Source 
PiperOrigin-RevId: 493324356
Change-Id: I5f71dc1a8ae40c18864ea80f00a5f86ff952f72c
This commit is contained in:
Alessio Quaglino
2022-12-06 09:13:50 -08:00
committed by Copybara-Service
parent 7f459ab8ca
commit 2d304d65e5
3 changed files with 151 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+1 -7
View File
@@ -397,13 +397,7 @@ static int mju_compare(const int* vec1, const int* vec2, int n) {
#endif
// scalar part
for (; i<n; i++) {
if (vec1[i]!=vec2[i]) {
return 0;
}
}
return 1;
return !memcmp(vec1+i, vec2+i, (n-i)*sizeof(int));
}
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
src/engine/engine_util_sparse.h
+8 -8
View File
@@ -70,14 +70,14 @@ void mju_superSparse(int nr, int* rowsuper,
const int* rownnz, const int* rowadr, const int* colind);
// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
mjData* d);
MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
mjData* d);
#ifdef __cplusplus
test/benchmark/engine_util_sparse_benchmark_test.cc
+142
View File
@@ -14,6 +14,7 @@
// A benchmark for comparing different implementations of mj_solveLD.
#include <cstddef>
#include <benchmark/benchmark.h>
#include <gtest/gtest.h>
#include <absl/base/attributes.h>
@@ -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<mjtNum> 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