Cache benchmark data in engine_util_sparse_benchmark_test
PiperOrigin-RevId: 902651349 Change-Id: I760e3f696b37699366c40c3c93a74e2b1b35678e
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
2d12dee025
commit
a8a5afc8dc
@@ -14,7 +14,6 @@
|
||||
|
||||
// A benchmark for comparing different implementations of mj_solveLD.
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
@@ -31,14 +30,179 @@ namespace {
|
||||
|
||||
using CombineFuncPtr = decltype(&mju_combineSparse);
|
||||
using TransposeFuncPtr = decltype(&mju_transposeSparse);
|
||||
using SqrMatTDFuncPtr = decltype(&mju_sqrMatTDSparse);
|
||||
|
||||
// number of steps to roll out before benchmarking
|
||||
static const int kNumWarmupSteps = 500;
|
||||
// ================================ Cached Data ================================
|
||||
|
||||
// ----------------------------- old functions --------------------------------
|
||||
// ---- MatVecSparse data ----
|
||||
struct MatVecData {
|
||||
int nv;
|
||||
int nefc;
|
||||
int nJ;
|
||||
std::vector<mjtNum> efc_J;
|
||||
std::vector<int> efc_J_rownnz, efc_J_rowadr, efc_J_colind, efc_J_rowsuper;
|
||||
std::vector<mjtNum> vec;
|
||||
};
|
||||
|
||||
MatVecData& GetMatVecData() {
|
||||
static MatVecData data = [] {
|
||||
MatVecData d;
|
||||
mjModel* m = LoadModelFromPath("flex/flag.xml");
|
||||
mjData* dat = mj_makeData(m);
|
||||
|
||||
for (int i = 0; i < 500; i++) {
|
||||
mj_step(m, dat);
|
||||
}
|
||||
|
||||
d.nv = m->nv;
|
||||
d.nefc = dat->nefc;
|
||||
d.nJ = dat->nJ;
|
||||
d.efc_J.assign(dat->efc_J, dat->efc_J + d.nJ);
|
||||
d.efc_J_rownnz.assign(dat->efc_J_rownnz, dat->efc_J_rownnz + d.nefc);
|
||||
d.efc_J_rowadr.assign(dat->efc_J_rowadr, dat->efc_J_rowadr + d.nefc);
|
||||
d.efc_J_colind.assign(dat->efc_J_colind, dat->efc_J_colind + d.nJ);
|
||||
d.efc_J_rowsuper.assign(dat->efc_J_rowsuper, dat->efc_J_rowsuper + d.nefc);
|
||||
|
||||
// compute direction: vec = -M^{-1} * (Ma - qfrc_smooth - qfrc_constraint)
|
||||
mj_markStack(dat);
|
||||
mjtNum* Ma = mj_stackAllocNum(dat, m->nv);
|
||||
mjtNum* grad = mj_stackAllocNum(dat, m->nv);
|
||||
mjtNum* Mgrad = mj_stackAllocNum(dat, m->nv);
|
||||
mj_mulM(m, dat, Ma, dat->qacc);
|
||||
for (int i = 0; i < m->nv; i++) {
|
||||
grad[i] = Ma[i] - dat->qfrc_smooth[i] - dat->qfrc_constraint[i];
|
||||
}
|
||||
mj_solveM(m, dat, Mgrad, grad, 1);
|
||||
d.vec.resize(m->nv);
|
||||
mju_scl(d.vec.data(), Mgrad, -1, m->nv);
|
||||
mj_freeStack(dat);
|
||||
|
||||
mj_deleteData(dat);
|
||||
mj_deleteModel(m);
|
||||
return d;
|
||||
}();
|
||||
return data;
|
||||
}
|
||||
|
||||
// ---- CombineSparse data ----
|
||||
struct CombineData {
|
||||
int nv;
|
||||
std::vector<mjtNum> H;
|
||||
std::vector<int> rownnz, rowadr, colind;
|
||||
};
|
||||
|
||||
CombineData& GetCombineData() {
|
||||
static CombineData data = [] {
|
||||
CombineData cd;
|
||||
mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
for (int i = 0; i < 500; i++) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
cd.nv = m->nv;
|
||||
mj_markStack(d);
|
||||
mjtNum* H = mj_stackAllocNum(d, m->nv*m->nv);
|
||||
int* rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* colind = mj_stackAllocInt(d, m->nv*m->nv);
|
||||
int* diagind = mj_stackAllocInt(d, m->nv);
|
||||
|
||||
mjtNum* D = mj_stackAllocNum(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;
|
||||
}
|
||||
}
|
||||
|
||||
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_colind = mj_stackAllocInt(d, d->nJ);
|
||||
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
|
||||
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
|
||||
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// compute H = J'*D*J, uncompressed layout
|
||||
mju_sqrMatTDUncompressedInit(rowadr, m->nv);
|
||||
mju_sqrMatTDSparse(H, d->efc_J, 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,
|
||||
JT_rownnz, JT_rowadr,
|
||||
JT_colind, JT_rowsuper, d,
|
||||
diagind);
|
||||
|
||||
// compute H = M + J'*D*J
|
||||
mj_addM(m, d, H, rownnz, rowadr, colind);
|
||||
|
||||
// copy to persistent storage
|
||||
int nH = rowadr[m->nv-1] + m->nv; // uncompressed: rowadr[r] = r*nv
|
||||
cd.H.assign(H, H + nH);
|
||||
cd.rownnz.assign(rownnz, rownnz + m->nv);
|
||||
cd.rowadr.assign(rowadr, rowadr + m->nv);
|
||||
cd.colind.assign(colind, colind + nH);
|
||||
|
||||
mj_freeStack(d);
|
||||
mj_deleteData(d);
|
||||
mj_deleteModel(m);
|
||||
return cd;
|
||||
}();
|
||||
return data;
|
||||
}
|
||||
|
||||
// ---- TransposeSparse data ----
|
||||
struct TransposeData {
|
||||
int nv;
|
||||
int nefc;
|
||||
int nJ;
|
||||
std::vector<mjtNum> efc_J;
|
||||
std::vector<int> efc_J_rownnz, efc_J_rowadr, efc_J_colind;
|
||||
};
|
||||
|
||||
enum class Size { H2_100, H100 };
|
||||
|
||||
template <Size S>
|
||||
const char* ModelPath() {
|
||||
if constexpr (S == Size::H2_100) {
|
||||
return "../test/benchmark/testdata/2humanoid100_chol.xml";
|
||||
} else {
|
||||
return "../test/benchmark/testdata/100_humanoids_chol.xml";
|
||||
}
|
||||
}
|
||||
|
||||
template <Size S>
|
||||
TransposeData& GetTransposeData() {
|
||||
static TransposeData data = [] {
|
||||
TransposeData td;
|
||||
mjModel* m = LoadModelFromPath(ModelPath<S>());
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
while (d->time < 2) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
td.nv = m->nv;
|
||||
td.nefc = d->nefc;
|
||||
td.nJ = d->nJ;
|
||||
td.efc_J.assign(d->efc_J, d->efc_J + d->nJ);
|
||||
td.efc_J_rownnz.assign(d->efc_J_rownnz, d->efc_J_rownnz + d->nefc);
|
||||
td.efc_J_rowadr.assign(d->efc_J_rowadr, d->efc_J_rowadr + d->nefc);
|
||||
td.efc_J_colind.assign(d->efc_J_colind, d->efc_J_colind + d->nJ);
|
||||
|
||||
mj_deleteData(d);
|
||||
mj_deleteModel(m);
|
||||
return td;
|
||||
}();
|
||||
return data;
|
||||
}
|
||||
|
||||
// ================================ old functions ==============================
|
||||
|
||||
// transpose sparse matrix (uncompressed)
|
||||
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
|
||||
@@ -229,61 +393,31 @@ void ABSL_ATTRIBUTE_NOINLINE mulMatVecSparse_8(mjtNum* res,
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------- benchmark ------------------------------------
|
||||
// ----------------------------- benchmark -------------------------------------
|
||||
|
||||
static void BM_MatVecSparse(benchmark::State& state, int unroll) {
|
||||
static mjModel* m = LoadModelFromPath("flex/flag.xml");
|
||||
mjData* d = mj_makeData(m);
|
||||
MatVecData& data = GetMatVecData();
|
||||
std::vector<mjtNum> res(data.nefc);
|
||||
|
||||
// warm-up rollout to get a typical state
|
||||
for (int i=0; i < kNumWarmupSteps; i++) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
// allocate gradient
|
||||
mj_markStack(d);
|
||||
mjtNum *Ma = mj_stackAllocNum(d, m->nv);
|
||||
mjtNum *vec = mj_stackAllocNum(d, m->nv);
|
||||
mjtNum *res = mj_stackAllocNum(d, d->nefc);
|
||||
mjtNum *grad = mj_stackAllocNum(d, m->nv);
|
||||
mjtNum *Mgrad = mj_stackAllocNum(d, m->nv);
|
||||
|
||||
// compute gradient
|
||||
mj_mulM(m, d, Ma, d->qacc);
|
||||
for (int i=0; i < m->nv; i++) {
|
||||
grad[i] = Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
|
||||
}
|
||||
|
||||
// compute search direction
|
||||
mj_solveM(m, d, Mgrad, grad, 1);
|
||||
mju_scl(vec, Mgrad, -1, m->nv);
|
||||
|
||||
// save state
|
||||
std::vector<mjtNum> qpos = AsVector(d->qpos, m->nq);
|
||||
std::vector<mjtNum> qvel = AsVector(d->qvel, m->nv);
|
||||
std::vector<mjtNum> act = AsVector(d->act, m->na);
|
||||
std::vector<mjtNum> warmstart = AsVector(d->qacc_warmstart, m->nv);
|
||||
|
||||
// time benchmark
|
||||
for (auto s : state) {
|
||||
if (unroll == 4) {
|
||||
mju_mulMatVecSparse(res, d->efc_J, vec, d->nefc,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr,
|
||||
d->efc_J_colind, d->efc_J_rowsuper);
|
||||
mju_mulMatVecSparse(res.data(), data.efc_J.data(), data.vec.data(),
|
||||
data.nefc, data.efc_J_rownnz.data(),
|
||||
data.efc_J_rowadr.data(), data.efc_J_colind.data(),
|
||||
data.efc_J_rowsuper.data());
|
||||
} else if (unroll == 1) {
|
||||
mulMatVecSparse_1(res, d->efc_J, vec, d->nefc,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr,
|
||||
d->efc_J_colind, d->efc_J_rowsuper);
|
||||
mulMatVecSparse_1(res.data(), data.efc_J.data(), data.vec.data(),
|
||||
data.nefc, data.efc_J_rownnz.data(),
|
||||
data.efc_J_rowadr.data(), data.efc_J_colind.data(),
|
||||
data.efc_J_rowsuper.data());
|
||||
} else if (unroll == 8) {
|
||||
mulMatVecSparse_8(res, d->efc_J, vec, d->nefc,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr,
|
||||
d->efc_J_colind, d->efc_J_rowsuper);
|
||||
mulMatVecSparse_8(res.data(), data.efc_J.data(), data.vec.data(),
|
||||
data.nefc, data.efc_J_rownnz.data(),
|
||||
data.efc_J_rowadr.data(), data.efc_J_colind.data(),
|
||||
data.efc_J_rowsuper.data());
|
||||
}
|
||||
}
|
||||
|
||||
// finalize
|
||||
mj_freeStack(d);
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
@@ -309,75 +443,30 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1(
|
||||
BENCHMARK(BM_MatVecSparse_1);
|
||||
|
||||
static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
|
||||
static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
CombineData& data = GetCombineData();
|
||||
|
||||
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
|
||||
mj_markStack(d);
|
||||
mjtNum* H = mj_stackAllocNum(d, m->nv*m->nv);
|
||||
int* rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* colind = mj_stackAllocInt(d, m->nv*m->nv);
|
||||
int* diagind = mj_stackAllocInt(d, m->nv);
|
||||
|
||||
// compute D corresponding to quad states
|
||||
mjtNum* D = mj_stackAllocNum(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;
|
||||
}
|
||||
}
|
||||
|
||||
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
|
||||
int* JT_colind = mj_stackAllocInt(d, d->nJ);
|
||||
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
|
||||
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
|
||||
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
|
||||
// compute H = J'*D*J, uncompressed layout
|
||||
mju_sqrMatTDUncompressedInit(rowadr, m->nv);
|
||||
mju_sqrMatTDSparse(H, d->efc_J, 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,
|
||||
JT_rownnz, JT_rowadr,
|
||||
JT_colind, JT_rowsuper, d,
|
||||
diagind);
|
||||
|
||||
// compute H = M + J'*D*J
|
||||
mj_addM(m, d, H, rownnz, rowadr, colind);
|
||||
// make working copies that get modified each iteration
|
||||
std::vector<mjtNum> H = data.H;
|
||||
std::vector<int> rownnz = data.rownnz;
|
||||
std::vector<int> rowadr = data.rowadr;
|
||||
std::vector<int> colind = data.colind;
|
||||
|
||||
// time benchmark
|
||||
for (auto s : state) {
|
||||
for (int r = m->nv-1; r >= 0; r--) {
|
||||
for (int r = data.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], 1, -H[adr+i],
|
||||
func(H.data()+rowadr[c], H.data()+rowadr[r], 1, -H[adr+i],
|
||||
rownnz[c], rownnz[c],
|
||||
colind+rowadr[c], colind+rowadr[c]);
|
||||
colind.data()+rowadr[c], colind.data()+rowadr[c]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// finalize
|
||||
mj_freeStack(d);
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
@@ -395,17 +484,6 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
|
||||
}
|
||||
BENCHMARK(BM_combineSparse_old);
|
||||
|
||||
enum class Size { H2_100, H100 };
|
||||
|
||||
template <Size S>
|
||||
const char* ModelPath() {
|
||||
if constexpr (S == Size::H2_100) {
|
||||
return "../test/benchmark/testdata/2humanoid100_chol.xml";
|
||||
} else {
|
||||
return "../test/benchmark/testdata/100_humanoids_chol.xml";
|
||||
}
|
||||
}
|
||||
|
||||
enum class Supernode {
|
||||
None,
|
||||
PostProcess,
|
||||
@@ -415,44 +493,33 @@ enum class Supernode {
|
||||
template <Size S>
|
||||
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
|
||||
Supernode super) {
|
||||
static mjModel* m = LoadModelFromPath(ModelPath<S>());
|
||||
TransposeData& data = GetTransposeData<S>();
|
||||
|
||||
// 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);
|
||||
}
|
||||
|
||||
mj_markStack(d);
|
||||
|
||||
// need uncompressed layout
|
||||
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);
|
||||
// allocate output buffers (uncompressed layout)
|
||||
std::vector<mjtNum> res(data.nv * data.nefc);
|
||||
std::vector<int> res_rownnz(data.nv);
|
||||
std::vector<int> res_rowadr(data.nv);
|
||||
std::vector<int> res_rowsuper(data.nv);
|
||||
std::vector<int> res_colind(data.nv * data.nefc);
|
||||
|
||||
// time benchmark
|
||||
for (auto s : state) {
|
||||
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);
|
||||
int* rowsuper =
|
||||
(super == Supernode::Inline) ? res_rowsuper.data() : nullptr;
|
||||
func(res.data(), data.efc_J.data(), data.nefc, data.nv,
|
||||
res_rownnz.data(), res_rowadr.data(), res_colind.data(), rowsuper,
|
||||
data.efc_J_rownnz.data(), data.efc_J_rowadr.data(),
|
||||
data.efc_J_colind.data());
|
||||
if (super == Supernode::PostProcess) {
|
||||
mju_superSparse(m->nv, res_rowsuper,
|
||||
res_rownnz, res_rowadr, res_colind);
|
||||
mju_superSparse(data.nv, res_rowsuper.data(),
|
||||
res_rownnz.data(), res_rowadr.data(), res_colind.data());
|
||||
}
|
||||
}
|
||||
|
||||
mj_freeStack(d);
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_2H100_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
|
||||
Reference in New Issue
Block a user