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

Remove transposed Jacobian fields from mjData

Browse Source 
No longer computed unconditionally, only ad-hoc where required.

PiperOrigin-RevId: 797753297
Change-Id: I2030fc342c98ff33575b0526dd72e2110c4fcb74
This commit is contained in:
Yuval Tassa
2025-08-21 07:20:29 -07:00
committed by Copybara-Service
parent 5fe4224ae7
commit 2f28473bc1
10 changed files with 33 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/benchmark/engine_util_sparse_benchmark_test.cc
+29 -11
View File
@@ -448,14 +448,23 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
}
}
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, d->efc_JT, D, d->nefc, 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,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d,
JT_rownnz, JT_rowadr,
JT_colind, JT_rowsuper, d,
diagind);
// compute H = M + J'*D*J
@@ -605,19 +614,28 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
}
}
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);
// time benchmark
if (func) {
mju_sqrMatTDSparseCount(rownnz, rowadr, m->nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, nullptr, d, 1);
JT_rownnz, JT_rowadr,
JT_colind, nullptr, d, 1);
for (auto s : state) {
// compute H = J'*D*J, compressed layout
func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
func(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, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, diagind);
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, diagind);
}
} else {
for (auto s : state) {
@@ -627,10 +645,10 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
// compute H = J'*D*J, uncompressed layout
mju_sqrMatTDSparse_baseline(
H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
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,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, /*unused=*/nullptr);
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, /*unused=*/nullptr);
}
}
test/engine/engine_island_test.cc
-21
View File
@@ -467,27 +467,6 @@ TEST_F(IslandTest, IslandJacobian) {
EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]);
}
}
// === test JT (if sparse)
// get pointer to J_island, dense (nefc_island x nv_island) submatrix
if (jac == mjJAC_SPARSE) {
// dense copy of island in iJ (here used as scratch)
mju_sparse2dense(iJ, d->iefc_JT, nv_island, nefc_island,
d->iefc_JT_rownnz + idof,
d->iefc_JT_rowadr + idof,
d->iefc_JT_colind);
J_island = iJ;
// sequential memory in J_island equals random access memory in J
for (int i=0; i < nv_island; i++) {
for (int j=0; j < nefc_island; j++) {
int dof = d->map_idof2dof[idof + i];
int efc = d->map_iefc2efc[iefc + j];
EXPECT_EQ(J_island[i * nefc_island + j], J[efc * nv + dof]);
}
}
}
}
mju_free(iJ);