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

Compute diagonal indices in mj_sqrMatTDSparse

Browse Source 
PiperOrigin-RevId: 713714087
Change-Id: Icc8eae74e6e47ba1d12a6e9aa0d774ab006cfe38
This commit is contained in:
Yuval Tassa
2025-01-09 10:16:37 -08:00
committed by Copybara-Service
parent daba00ca9d
commit 00010f5848
6 changed files with 53 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/engine_util_sparse_benchmark_test.cc
+6 -4
View File
@@ -43,7 +43,7 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
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, int unused) {
const int* colindT, const int* rowsuperT, mjData* d, int* unused) {
mj_markStack(d);
int* chain = mj_stackAllocInt(d, 2 * nc);
mjtNum* buffer = mj_stackAllocNum(d, nc);
@@ -435,6 +435,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
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);
@@ -454,7 +455,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
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,
/*flg_upper=*/1);
diagind);
// compute H = M + J'*D*J
mj_addM(m, d, H, rownnz, rowadr, colind);
@@ -559,6 +560,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
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);
@@ -579,7 +581,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
func(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, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, /*flg_upper=*/1);
d->efc_JT_rowsuper, d, diagind);
}
} else {
for (auto s : state) {
@@ -592,7 +594,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
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, /*unused=*/0);
d->efc_JT_rowsuper, d, /*unused=*/nullptr);
}
}