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

Refactor islands to be memory contiguous.

Browse Source 
PiperOrigin-RevId: 755803476
Change-Id: I41972b07e0d5ef5d0117c94f565b93367b87458b
This commit is contained in:
Yuval Tassa
2025-05-07 05:05:34 -07:00
committed by Copybara-Service
parent 449de73430
commit ecb769fc3a
30 changed files with 1742 additions and 1116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/engine/engine_island_test.cc
+185 -28
View File
@@ -208,17 +208,19 @@ TEST_F(IslandTest, Abacus) {
int nv = model->nv;
int nefc = data->nefc;
int nisland = data->nisland;
int nidof = data->nidof;
// 4 dofs, 12 constraints, 2 islands
EXPECT_EQ(nv, 4);
EXPECT_EQ(nidof, 3);
EXPECT_EQ(nefc, 12); // 3 pyramidal contacts
EXPECT_EQ(nisland, 2);
// the islands begin at dofs 0 and 1
EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1));
EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1));
// number of dofs in the 2 islands
EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2));
EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2));
// dof 0 in island 0
// dof 1 in no island
@@ -228,19 +230,19 @@ TEST_F(IslandTest, Abacus) {
// dof 0 constitutes first island
// dofs 2, 3 are the second island
// last index is unassigned since dof 1 is unconstrained
EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 2, 3, -1));
EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 2, 3, 1));
// dof 0 constitutes first island
// dofs 1 is unassigned
// dofs 2, 3 are second island
EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, -1, 0, 1));
EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 3, 1, 2));
// island 0 starts at constraint 0
// island 1 starts at constraint 4
EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4));
EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4));
// number of constraints in the 2 islands
EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 8));
EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 8));
// first contact (4 constraints) is in island 0
// second contact (8 constraints) is in island 1
@@ -248,7 +250,7 @@ TEST_F(IslandTest, Abacus) {
ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1));
// index lists for islands 0 and 1
EXPECT_THAT(AsVector(data->island_efcind, nefc),
EXPECT_THAT(AsVector(data->map_iefc2efc, nefc),
ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11));
// reset, push 0 to the left, 3 to the right, 1,2 to the middle
@@ -266,18 +268,20 @@ TEST_F(IslandTest, Abacus) {
// local variables
nefc = data->nefc;
nisland = data->nisland;
nidof = data->nidof;
EXPECT_EQ(nisland, 3);
EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 1, 3));
EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2, 1));
EXPECT_EQ(nidof, 4);
EXPECT_THAT(AsVector(data->island_idofadr, nisland), ElementsAre(0, 1, 3));
EXPECT_THAT(AsVector(data->island_nv, nisland), ElementsAre(1, 2, 1));
EXPECT_THAT(AsVector(data->dof_island, nv), ElementsAre(0, 1, 1, 2));
EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 1, 2, 3));
EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, 0, 1, 0));
EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4, 8));
EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 4, 4));
EXPECT_THAT(AsVector(data->map_idof2dof, nv), ElementsAre(0, 1, 2, 3));
EXPECT_THAT(AsVector(data->map_dof2idof, nv), ElementsAre(0, 1, 2, 3));
EXPECT_THAT(AsVector(data->island_iefcadr, nisland), ElementsAre(0, 4, 8));
EXPECT_THAT(AsVector(data->island_nefc, nisland), ElementsAre(4, 4, 4));
EXPECT_THAT(AsVector(data->efc_island, nefc),
ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2));
EXPECT_THAT(AsVector(data->island_efcind, nefc),
EXPECT_THAT(AsVector(data->map_iefc2efc, nefc),
ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11));
mj_deleteData(data);
@@ -311,27 +315,30 @@ TEST_F(IslandTest, DenseSparse) {
int nisland = data1->nisland;
// expect sparse and dense to be identical
EXPECT_EQ(data1->nidof, data2->nidof);
EXPECT_EQ(data1->nefc, data2->nefc);
EXPECT_EQ(data1->nisland, data2->nisland);
EXPECT_EQ(data1->nefc, data2->nefc);
EXPECT_EQ(AsVector(data1->island_dofadr, nisland),
AsVector(data2->island_dofadr, nisland));
EXPECT_EQ(AsVector(data1->island_dofnum, nisland),
AsVector(data2->island_dofnum, nisland));
EXPECT_EQ(AsVector(data1->island_idofadr, nisland),
AsVector(data2->island_idofadr, nisland));
EXPECT_EQ(AsVector(data1->island_nv, nisland),
AsVector(data2->island_nv, nisland));
EXPECT_EQ(AsVector(data1->dof_island, nv),
AsVector(data2->dof_island, nv));
EXPECT_EQ(AsVector(data1->island_dofind, nv),
AsVector(data2->island_dofind, nv));
EXPECT_EQ(AsVector(data1->dof_islandind, nv),
AsVector(data2->dof_islandind, nv));
EXPECT_EQ(AsVector(data1->island_efcadr, nisland),
AsVector(data2->island_efcadr, nisland));
EXPECT_EQ(AsVector(data1->island_efcnum, nisland),
AsVector(data2->island_efcnum, nisland));
EXPECT_EQ(AsVector(data1->map_idof2dof, nv),
AsVector(data2->map_idof2dof, nv));
EXPECT_EQ(AsVector(data1->map_dof2idof, nv),
AsVector(data2->map_dof2idof, nv));
EXPECT_EQ(AsVector(data1->island_iefcadr, nisland),
AsVector(data2->island_iefcadr, nisland));
EXPECT_EQ(AsVector(data1->island_nefc, nisland),
AsVector(data2->island_nefc, nisland));
EXPECT_EQ(AsVector(data1->efc_island, nefc),
AsVector(data2->efc_island, nefc));
EXPECT_EQ(AsVector(data1->island_efcind, nefc),
AsVector(data2->island_efcind, nefc));
EXPECT_EQ(AsVector(data1->map_iefc2efc, nefc),
AsVector(data2->map_iefc2efc, nefc));
EXPECT_EQ(AsVector(data1->map_efc2iefc, nefc),
AsVector(data2->map_efc2iefc, nefc));
mj_deleteData(data2);
mj_deleteData(data1);
@@ -361,6 +368,156 @@ TEST_F(IslandTest, IslandEfc) {
mj_deleteModel(model);
}
static const char* const k2H100Path = "engine/testdata/island/2humanoid100.xml";
TEST_F(IslandTest, IslandJacobian) {
for (const char* local_path : {kIlslandEfcPath, k2H100Path}) {
const std::string xml_path = GetTestDataFilePath(local_path);
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
int jac0 = m->opt.jacobian;
mjData* d = mj_makeData(m);
for (mjtNum t_stop : {0.0, 0.2, 2.0}) {
while (d->time < t_stop) {
mj_step(m, d);
}
for (mjtJacobian jac : {mjJAC_DENSE, mjJAC_SPARSE}) {
m->opt.jacobian = jac;
mj_forward(m, d);
int nv = m->nv;
int nefc = d->nefc;
int nisland = d->nisland;
int nidof = d->nidof;
mjtNum* J = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nv);
mjtNum* iJ = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nidof);
// get local dense Jacobian
if (jac == mjJAC_DENSE) {
mju_copy(J, d->efc_J, nefc * nv);
mju_copy(iJ, d->iefc_J, nefc * nidof);
} else {
mju_sparse2dense(J, d->efc_J, nefc, nv, d->efc_J_rownnz,
d->efc_J_rowadr, d->efc_J_colind);
}
// compare random access in efc_J to contiguous memory in iefc_J
for (int island=0; island < nisland; island++) {
int idof = d->island_idofadr[island];
int iefc = d->island_iefcadr[island];
int nefc_island = d->island_nefc[island];
int nv_island = d->island_nv[island];
// === test J
// get pointer to J_island, dense (nefc_island x nv_island) submatrix
mjtNum* J_island;
if (jac == mjJAC_DENSE) {
// point to starting address of island in efc_J
J_island = iJ + iefc * nidof;
} else {
// dense copy of island in iJ (here used as scratch)
mju_sparse2dense(iJ, d->iefc_J, nefc_island, nv_island,
d->iefc_J_rownnz + iefc,
d->iefc_J_rowadr + iefc,
d->iefc_J_colind);
J_island = iJ;
}
// sequential memory in J_island equals random access memory in J
for (int i=0; i < nefc_island; i++) {
for (int j=0; j < nv_island; j++) {
int efc = d->map_iefc2efc[iefc + i];
int dof = d->map_idof2dof[idof + j];
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);
mju_free(J);
}
// reset opt.jacobian to initial value
m->opt.jacobian = jac0;
}
mj_deleteData(d);
mj_deleteModel(m);
}
}
TEST_F(IslandTest, IslandInertia) {
for (const char* local_path : {kIlslandEfcPath, k2H100Path}) {
const std::string xml_path = GetTestDataFilePath(local_path);
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
int nv = m->nv;
mjData* d = mj_makeData(m);
mjtNum* M = (mjtNum*)mju_malloc(sizeof(mjtNum) * nv * nv);
for (mjtNum t_stop : {0.0, 0.2, 2.0}) {
while (d->time < t_stop) {
mj_step(m, d);
}
mj_forward(m, d);
int nisland = d->nisland;
// get dense inertia (lower only)
mj_fullM(m, M, d->qM);
// compare iM sub-matrix to full M
for (int island=0; island < nisland; island++) {
int nvi = d->island_nv[island];
mjtNum* Mi = (mjtNum*)mju_malloc(sizeof(mjtNum) * nvi * nvi);
int adr = d->island_idofadr[island];
mju_sparse2dense(Mi, d->iM, nvi, nvi,
d->iM_rownnz + adr,
d->iM_rowadr + adr,
d->iM_colind);
// compare Mi to M (lower triangle only)
for (int i=0; i < nvi; i++) {
for (int j=0; j <= i; j++) {
int dofi = d->map_idof2dof[adr + j];
int dofj = d->map_idof2dof[adr + i];
EXPECT_EQ(Mi[i * nvi + j], M[dofi * nv + dofj]);
}
}
mju_free(Mi);
}
}
mju_free(M);
mj_deleteData(d);
mj_deleteModel(m);
}
}
TEST_F(IslandTest, IslandEfcElliptic) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);