From 52ddcbc81a5e4686157e924dd58da810a1b35aa0 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 11 Jul 2026 09:48:54 +0530 Subject: [PATCH 01/10] Build native islands directly from constraint incidence Signed-off-by: teerthsharma --- src/engine/engine_island.c | 226 +++++++++++++++--------------- test/engine/engine_island_test.cc | 190 +++++++++++++++++++++++++ 2 files changed, 300 insertions(+), 116 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 47501298..8d124e5e 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -16,7 +16,6 @@ #include #include -#include #include #include @@ -84,6 +83,68 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) { //-------------------------- flood-fill and graph construction ------------------------------------ +// find the canonical root of an active tree and compress its path +static int dsuFind(int* parent, int tree) { + int root = tree; + while (parent[root] != root) { + root = parent[root]; + } + + while (parent[tree] != tree) { + int next = parent[tree]; + parent[tree] = root; + tree = next; + } + + return root; +} + + +// initialize all trees as inactive +static void dsuInit(int* parent, int ntree) { + mju_fillInt(parent, -1, ntree); +} + + +// activate and union two incident trees; -1 denotes a static endpoint +static void dsuUnion(int* parent, int tree1, int tree2) { + if (tree1 == -1 && tree2 == -1) { + mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR + return; + } + + if (tree1 == -1) tree1 = tree2; + if (tree2 == -1) tree2 = tree1; + + if (parent[tree1] == -1) parent[tree1] = tree1; + if (parent[tree2] == -1) parent[tree2] = tree2; + + int root1 = dsuFind(parent, tree1); + int root2 = dsuFind(parent, tree2); + if (root1 < root2) { + parent[root2] = root1; + } else if (root2 < root1) { + parent[root1] = root2; + } +} + + +// assign deterministic island ids in ascending canonical-root order +static int dsuAssign(int* island, int* parent, int ntree) { + int nisland = 0; + for (int tree=0; tree < ntree; tree++) { + if (parent[tree] == -1) { + island[tree] = -1; + continue; + } + + int root = dsuFind(parent, tree); + island[tree] = root == tree ? nisland++ : island[root]; + } + + return nisland; +} + // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: // island (nr) - island index assigned to vertex, -1 if vertex has no edges @@ -280,61 +341,27 @@ static void treeIterInit(const mjModel* m, const mjData* d, int i, mjTreeIter* i } -// add 0, 1 or 2 edges to uncompressed CSR adjacency matrix -// increment rownnz using tree_tree to de-dupe; return number of edges added -static int addEdge(int* rownnz, int* colind, mjtByte* tree_tree, int ntree, int tree1, int tree2) { - if (tree1 == -1 && tree2 == -1) { - mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR - return 0; - } - - // handle static trees (treat as self-edge) - if (tree1 == -1) tree1 = tree2; - if (tree2 == -1) tree2 = tree1; - - // skip if edge already present - if (tree_tree[tree1*ntree + tree2]) { - return 0; - } - - // add edge - tree_tree[tree1*ntree + tree2] = 1; - colind[tree1*ntree + rownnz[tree1]++] = tree2; // uncompressed format, rowadr is known - - // add flipped edge (off-diagonal) - if (tree1 != tree2) { - tree_tree[tree2*ntree + tree1] = 1; - colind[tree2*ntree + rownnz[tree2]++] = tree1; // uncompressed format, rowadr is known - return 2; - } - - return 1; +// return whether repeated scalar rows of this constraint require separate tree scans +static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) { + return efc_type == mjCNSTR_EQUALITY && + (m->eq_type[efc_id] == mjEQ_FLEX || + m->eq_type[efc_id] == mjEQ_FLEXVERT || + m->eq_type[efc_id] == mjEQ_FLEXSTRAIN); } -// find tree-tree edges (column indices), return total number of edges -// efc_tree: first nonegative tree index of each constraint -static int findEdges(const mjModel* m, const mjData* d, - int* rownnz, int* colind, mjtByte* tree_tree, int* efc_tree, int ntree) { +// activate and union all trees with direct incidence in a constraint +static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) { int nefc = d->nefc; - int nnz = 0; int efc_type = -1; int efc_id = -1; - // clear row nonzeros - mju_zeroInt(rownnz, ntree); - - // iterate over constraints, compute tree-tree edges, assign efc_tree + // iterate over constraints and union incident trees for (int i=0; i < nefc; i++) { - // row i is still in the same constraint: skip it, + // row i is still in the same constraint: skip it if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { // unless it is a flex equality, where the tree pattern changes per dof - if (!(efc_type == mjCNSTR_EQUALITY && - (m->eq_type[efc_id] == mjEQ_FLEX || - m->eq_type[efc_id] == mjEQ_FLEXVERT || - m->eq_type[efc_id] == mjEQ_FLEXSTRAIN))) { - // copy tree assignment from previous constraint and continue - efc_tree[i] = efc_tree[i-1]; + if (!isFlexEquality(m, efc_type, efc_id)) { continue; } } @@ -350,18 +377,12 @@ static int findEdges(const mjModel* m, const mjData* d, if (tree1 != -2) { int tree2 = treeNext(m, d, i, &iter); - // assign tree to constraint, one of (tree1, tree2) must be non-negative - efc_tree[i] = tree1 >= 0 ? tree1 : tree2; - if (efc_tree[i] < 0) { - mjERROR("constraint %d is between two static bodies", i); // SHOULD NOT OCCUR - } - - // add one edge or continue to search for more edges + // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, -1); + dsuUnion(parent, tree1, -1); } else { while (tree2 != -2) { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2); + dsuUnion(parent, tree1, tree2); tree1 = tree2; tree2 = treeNext(m, d, i, &iter); } @@ -370,49 +391,37 @@ static int findEdges(const mjModel* m, const mjData* d, mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR } } +} - // flex stiffness couples all vertices (nodes for interpolated flexes) of a flex without any - // constraint row representing the coupling: union the trees of every stiffness-active flex - // (star around the first dynamic tree). This keeps the partition valid when the implicit - // effective metric (mj_flexCG) carries the stiffness inside the constraint solve. Awake - // trees only: sleeping trees must stay out of islands (mj_sleep invariant, matching the - // constraint filter); waking a flex as a unit remains the wake machinery's job. - for (int f=0; f < m->nflex; f++) { - // mirror the stiffness-activity conditions of engine_derivative's flexStiff_active / - // flexInterp_processed: deformable dim>=2 flex with bending or nonzero stiffness - if (m->flex_rigid[f] || m->flex_dim[f] < 2) { +// assign each constraint from its first non-negative incident tree +static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree_island) { + int efc_type = -1; + int efc_id = -1; + + for (int i=0; i < d->nefc; i++) { + // reuse assignment for repeated scalar rows, except flex equality rows + if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i] && + !isFlexEquality(m, efc_type, efc_id)) { + d->efc_island[i] = d->efc_island[i-1]; continue; } - int sadr = m->flex_stiffnessadr[f]; - if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) { - continue; - } - int num, adr; - const int* bodyid; - if (m->flex_interp[f]) { - num = m->flex_nodenum[f]; - adr = m->flex_nodeadr[f]; - bodyid = m->flex_nodebodyid; + efc_type = d->efc_type[i]; + efc_id = d->efc_id[i]; + + mjTreeIter iter; + treeIterInit(m, d, i, &iter); + + int tree; + do { + tree = treeNext(m, d, i, &iter); + } while (tree == -1); + + if (tree == -2) { + mjERROR("no dynamic tree found for constraint %d", i); // SHOULD NOT OCCUR } else { - num = m->flex_vertnum[f]; - adr = m->flex_vertadr[f]; - bodyid = m->flex_vertbodyid; - } - int tree1 = -1; - for (int j=0; j < num; j++) { - int treeid = m->body_treeid[bodyid[adr+j]]; - if (treeid < 0 || treeid == tree1 || !d->tree_awake[treeid]) { - continue; - } - if (tree1 < 0) { - tree1 = treeid; - } else { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, treeid); - } + d->efc_island[i] = tree_island[tree]; } } - - return nnz; } @@ -431,29 +440,12 @@ void mj_island(const mjModel* m, mjData* d) { mj_markStack(d); - // dense tree-tree adjacency matrix - int ntree2 = ntree * ntree; - mjtByte* tree_tree = mjSTACKALLOC(d, ntree2, mjtByte); - memset(tree_tree, 0, ntree2); - - // CSR representation of tree-tree adjacency matrix (uncompressed) - int* colind = mjSTACKALLOC(d, ntree2, int); - int* rownnz = mjSTACKALLOC(d, ntree, int); - int* rowadr = mjSTACKALLOC(d, ntree, int); - for (int r=0; r < ntree; r++) { - rowadr[r] = r * ntree; - } - - // first non-negative tree index of each constraint, used later for computing efc_island - int* efc_tree = mjSTACKALLOC(d, nefc, int); - - // compute tree-tree adjacency matrix: fill rownnz and colind - int nnz = findEdges(m, d, rownnz, colind, tree_tree, efc_tree, ntree); - - // discover islands + // union direct tree incidence and assign deterministic components + int* parent = mjSTACKALLOC(d, ntree, int); + dsuInit(parent, ntree); + unionConstraintTrees(m, d, parent); int* tree_island = mjSTACKALLOC(d, ntree, int); - int* stack = mjSTACKALLOC(d, nnz, int); - d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack); + d->nisland = dsuAssign(tree_island, parent, ntree); // no islands found: quick return if (!d->nisland) { @@ -570,13 +562,15 @@ void mj_island(const mjModel* m, mjData* d) { // ------------------------------------- constraints --------------------------------------------- + // compute efc_island from first non-negative tree of each constraint + assignConstraintIslands(m, d, tree_island); + // compute efc_island, island_{ne,nf,nefc} mju_zeroInt(d->island_ne, nisland); mju_zeroInt(d->island_nf, nisland); mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { - int island = tree_island[efc_tree[i]]; - d->efc_island[i] = island; + int island = d->efc_island[i]; d->island_nefc[island]++; switch (d->efc_type[i]) { case mjCNSTR_EQUALITY: diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index f520c6df..3a4462ab 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -156,6 +156,196 @@ TEST_F(IslandTest, FloodFill3b) { EXPECT_THAT(island, ElementsAre(0, 1, 1, -1, 0, 0, 0)); } +TEST_F(IslandTest, ProductionStaticFirstAndRepeatedRows) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, 3); + ASSERT_EQ(model->nv, 3); + + // The first equality incidence is static first, then dynamic tree 0. + ASSERT_EQ(model->eq_objtype[0], mjOBJ_SITE); + int body1 = model->site_bodyid[model->eq_obj1id[0]]; + int body2 = model->site_bodyid[model->eq_obj2id[0]]; + EXPECT_EQ(model->body_treeid[body1], -1); + EXPECT_EQ(model->body_treeid[body2], 0); + + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 6); + EXPECT_EQ(data->nisland, 2); + EXPECT_EQ(data->nidof, 3); + EXPECT_EQ(data->ne, 6); + EXPECT_EQ(data->nf, 0); + EXPECT_THAT(AsVector(data->efc_type, data->nefc), + ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, + mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY)); + EXPECT_THAT(AsVector(data->efc_id, data->nefc), ElementsAre(0, 0, 0, 1, 1, 1)); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 1, 1)); + EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 1, 1)); + EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_dof2idof, model->nv), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->map_idof2dof, model->nv), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0, 0, 1, 1, 1)); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(3, 3)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(3, 3)); + EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0, 3)); + EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc), ElementsAre(0, 1, 2, 3, 4, 5)); + EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1, 2, 3, 4, 5)); +} + +TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { + static constexpr char xml[] = R"( + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, 3); + ASSERT_EQ(model->nv, 9); + ASSERT_EQ(model->neq, 1); + ASSERT_EQ(model->eq_type[0], mjEQ_FLEX); + ASSERT_TRUE(mj_isSparse(model.get())); + + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 2); + auto row_trees = [&](int row) { + std::vector trees; + for (int j=0; j < data->efc_J_rownnz[row]; j++) { + int dof = data->efc_J_colind[data->efc_J_rowadr[row] + j]; + int tree = model->dof_treeid[dof]; + if (trees.empty() || trees.back() != tree) { + trees.push_back(tree); + } + } + return trees; + }; + + // Rows share one flex equality id but have different tree incidence. + EXPECT_THAT(row_trees(0), ElementsAre(0, 1)); + EXPECT_THAT(row_trees(1), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->efc_type, data->nefc), + ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY)); + EXPECT_THAT(AsVector(data->efc_id, data->nefc), ElementsAre(0, 0)); + EXPECT_EQ(data->nisland, 1); + EXPECT_EQ(data->nidof, 9); + EXPECT_EQ(data->ne, 2); + EXPECT_EQ(data->nf, 0); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, 0)); + EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(3)); + EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), + ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0)); + EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(9)); + EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_dof2idof, model->nv), + ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8)); + EXPECT_THAT(AsVector(data->map_idof2dof, model->nv), + ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0)); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(2)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(2)); + EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); +} + +TEST_F(IslandTest, BoundedArenaSupports1024Trees) { + constexpr int kTreeCount = 1024; + constexpr size_t kArenaBytes = 2 * 1024 * 1024; + std::string xml = R"( + + + + +)"; + xml.reserve(160 * kTreeCount); + for (int i=0; i < kTreeCount; i++) { + std::string name = std::to_string(i); + xml += ""; + xml += ""; + xml += ""; + if (i < 2) { + xml += ""; + } + xml += ""; + } + xml += R"( + + + + + +)"; + + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, kTreeCount); + ASSERT_EQ(model->narena, kArenaBytes); + MjDataPtr data = MakeData(model); + ASSERT_THAT(data.get(), NotNull()); + + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 3); + ASSERT_EQ(data->nisland, 1); + ASSERT_THAT(data->tree_island, NotNull()); + EXPECT_EQ(data->tree_island[0], 0); + EXPECT_EQ(data->tree_island[1], 0); + for (int tree=2; tree < kTreeCount; tree++) { + EXPECT_EQ(data->tree_island[tree], -1); + } + EXPECT_LE(data->maxuse_arena, kArenaBytes); +} + static const char* const kAbacusPath = "engine/testdata/island/abacus.xml"; TEST_F(IslandTest, Abacus) { From 03203e2f32ba8654f8b6c865f1c726c7a249eecc Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 11 Jul 2026 11:50:50 +0530 Subject: [PATCH 02/10] Update engine_island.c Signed-off-by: teerthsharma --- src/engine/engine_island.c | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 8d124e5e..34403b11 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -130,8 +130,9 @@ static void dsuUnion(int* parent, int tree1, int tree2) { // assign deterministic island ids in ascending canonical-root order -static int dsuAssign(int* island, int* parent, int ntree) { +static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { int nisland = 0; + *nidof = 0; for (int tree=0; tree < ntree; tree++) { if (parent[tree] == -1) { island[tree] = -1; @@ -140,6 +141,7 @@ static int dsuAssign(int* island, int* parent, int ntree) { int root = dsuFind(parent, tree); island[tree] = root == tree ? nisland++ : island[root]; + *nidof += tree_dofnum[tree]; } return nisland; @@ -445,7 +447,8 @@ void mj_island(const mjModel* m, mjData* d) { dsuInit(parent, ntree); unionConstraintTrees(m, d, parent); int* tree_island = mjSTACKALLOC(d, ntree, int); - d->nisland = dsuAssign(tree_island, parent, ntree); + int nidof; + d->nisland = dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); // no islands found: quick return if (!d->nisland) { @@ -454,13 +457,6 @@ void mj_island(const mjModel* m, mjData* d) { return; } - // count nidof: total number of dofs in islands - int nidof = 0; - for (int i=0; i < ntree; i++) { - if (tree_island[i] >= 0) { - nidof += m->tree_dofnum[i]; - } - } d->nidof = nidof; // allocate island arrays on arena @@ -599,17 +595,15 @@ void mj_island(const mjModel* m, mjData* d) { int ic = d->island_iefcadr[island] + island_nefc2[island]++; d->map_efc2iefc[c] = ic; d->map_iefc2efc[ic] = c; + d->iefc_type[ic] = d->efc_type[c]; + d->iefc_id[ic] = d->efc_id[c]; + d->iefc_frictionloss[ic] = d->efc_frictionloss[c]; + d->iefc_D[ic] = d->efc_D[c]; + d->iefc_R[ic] = d->efc_R[c]; } // SHOULD NOT OCCUR if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); - // copy position-dependent efc vectors required by solver - mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc); - mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc); - mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc); - mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc); - mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc); - mj_freeStack(d); } From bc32db7f258a343201265d87a13f555f5eb97ecd Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 11 Jul 2026 19:08:15 +0530 Subject: [PATCH 03/10] Cache island topology data Store island topology in `mjData` so repeated island solves can reuse stable connect/weld equality partitions. Add cache invalidation checks for active equality changes and cover the fast path with an island regression test. Signed-off-by: teerthsharma --- include/mujoco/mjdata.h | 5 + include/mujoco/mjxmacro.h | 6 +- src/engine/engine_io.c | 29 +++- src/engine/engine_island.c | 218 ++++++++++++++++++++++++++++++ test/engine/engine_island_test.cc | 63 +++++++++ 5 files changed, 319 insertions(+), 2 deletions(-) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index c2fddb81..800aefa3 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -198,6 +198,11 @@ typedef struct mjData_ { // sleep state int* tree_asleep; // <0: awake; >=0: index cycle of sleeping trees (ntree x 1) + // opaque internal island topology cache; users must not modify + int* island_cache_tree; // tree and island topology (ntree x 28) + int* island_cache_dof; // DOF topology (nv x 8) + int* island_cache_eq; // equality topology (neq x 40) + // plugins int* plugin; // copy of m->plugin, required for deletion (nplugin x 1) uintptr_t* plugin_data; // pointer to plugin-managed data structure (nplugin x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 8f36c166..71f8ed39 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -929,7 +929,10 @@ X ( mjtNum, qfrc_inverse, nv, 1 ) \ X ( mjtNum, cacc, nbody, 6 ) \ X ( mjtNum, cfrc_int, nbody, 6 ) \ - X ( mjtNum, cfrc_ext, nbody, 6 ) + X ( mjtNum, cfrc_ext, nbody, 6 ) \ + XIC ( int, island_cache_tree, ntree, 28 ) \ + XIC ( int, island_cache_dof, nv, 8 ) \ + XIC ( int, island_cache_eq, neq, 40 ) // macro for annotating that an array size in an X macro is a member of mjData @@ -1080,5 +1083,6 @@ // to obtain only X macros for fields that are relevant for mjvScene creation, // redefine XNV to expand to nothing #define XNV X +#define XIC(type, name, nr, nc) #endif // MUJOCO_MJXMACRO_H_ diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 482f5915..285561fb 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -998,7 +998,11 @@ static void mj_setPtrData(const mjModel* m, mjData* d) { ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr)); \ ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc); + #undef XIC + #define XIC X MJDATA_POINTERS + #undef XIC + #define XIC(type, name, nr, nc) #undef X // check size @@ -1085,7 +1089,11 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { return; \ } + #undef XIC + #define XIC X MJDATA_POINTERS + #undef XIC + #define XIC(type, name, nr, nc) #undef X // copy stack size from model @@ -1108,6 +1116,9 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { // set pointers into buffer mj_setPtrData(m, d); + if (m->ntree) { + d->island_cache_tree[0] = 0; + } // clear threadpool d->threadpool = 0; @@ -1173,6 +1184,9 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int dest->arena = save_arena; dest->threadpool = 0; mj_setPtrData(m, dest); + if (m->ntree) { + dest->island_cache_tree[0] = 0; + } // save plugin_data, since the X macro copying block below will override it const size_t plugin_data_size = sizeof(*dest->plugin_data) * dest->nplugin; @@ -1373,19 +1387,32 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { // fill buffer with debug_value (normally 0) #ifdef ADDRESS_SANITIZER { + #undef XIC + #define XIC(type, name, nr, nc) #define X(type, name, nr, nc) memset(d->name, (int)debug_value, sizeof(type)*(m->nr)*(nc)); MJDATA_POINTERS #undef X + #undef XIC + #define XIC(type, name, nr, nc) } #else - memset(d->buffer, (int)debug_value, d->nbuffer); + size_t cache_offset = d->island_cache_tree ? + (size_t)PTRDIFF(d->island_cache_tree, d->buffer) : d->nbuffer; + memset(d->buffer, (int)debug_value, cache_offset); #endif #ifdef MEMORY_SANITIZER // under MSAN, mark the entire buffer as uninitialized __msan_allocated_memory(d->buffer, d->nbuffer); + if (m->ntree) { + d->island_cache_tree[0] = 0; + } #endif + if (debug_value && m->ntree) { + d->island_cache_tree[0] = 0; + } + // zero out user-settable state and input arrays (MSAN: mark as initialized) mju_zero(d->qpos, m->nq); mju_zero(d->qvel, m->nv); diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 34403b11..207ccad8 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -14,8 +14,10 @@ #include "engine/engine_island.h" +#include #include #include +#include #include #include @@ -427,6 +429,211 @@ static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree } +enum { + kCacheMagic = 0x49534C44, + kCacheHeader = 16, +}; + +typedef struct mjIslandCacheView_ { + int* tree_island; + int* island_ntree; + int* island_itreeadr; + int* map_itree2tree; + int* island_nv; + int* island_idofadr; + int* island_dofadr; + int* island_ne; + int* island_nf; + int* island_nefc; + int* island_iefcadr; + int* tree_dofnum; + int* dof_treeid; + int* dof_island; + int* map_dof2idof; + int* map_idof2dof; + int* efc_island; + int* efc_type; + int* efc_id; + int* map_efc2iefc; + int* map_iefc2efc; + int* eq_active; + int* eq_type; + int* eq_tree1; + int* eq_tree2; +} mjIslandCacheView; + +static mjIslandCacheView islandCacheView(const mjModel* m, const mjData* d) { + int* tree = d->island_cache_tree + kCacheHeader; + int* dof = d->island_cache_dof; + int* eq = d->island_cache_eq; + int nisland = d->island_cache_tree[4]; + int nefc = d->island_cache_tree[3]; + mjIslandCacheView view; +#define TAKE(base, name, count) view.name = base; base += (count) + TAKE(tree, tree_island, m->ntree); + TAKE(tree, island_ntree, nisland); + TAKE(tree, island_itreeadr, nisland); + TAKE(tree, map_itree2tree, m->ntree); + TAKE(tree, island_nv, nisland); + TAKE(tree, island_idofadr, nisland); + TAKE(tree, island_dofadr, nisland); + TAKE(tree, island_ne, nisland); + TAKE(tree, island_nf, nisland); + TAKE(tree, island_nefc, nisland); + TAKE(tree, island_iefcadr, nisland); + TAKE(tree, tree_dofnum, m->ntree); + TAKE(dof, dof_treeid, m->nv); + TAKE(dof, dof_island, m->nv); + TAKE(dof, map_dof2idof, m->nv); + TAKE(dof, map_idof2dof, m->nv); + TAKE(eq, efc_island, nefc); + TAKE(eq, efc_type, nefc); + TAKE(eq, efc_id, nefc); + TAKE(eq, map_efc2iefc, nefc); + TAKE(eq, map_iefc2efc, nefc); + TAKE(eq, eq_active, m->neq); + TAKE(eq, eq_type, m->neq); + TAKE(eq, eq_tree1, m->neq); + TAKE(eq, eq_tree2, m->neq); +#undef TAKE + return view; +} + +static int islandCacheMatches(const mjModel* m, const mjData* d) { + if (!m->ntree || !m->neq || d->island_cache_tree[0] != kCacheMagic || + d->island_cache_tree[1] != m->ntree || d->island_cache_tree[2] != m->nv || + d->island_cache_tree[3] != d->nefc || d->island_cache_tree[6] != m->neq || + d->nefc != d->ne || d->nefc < 0 || (int64_t)d->nefc > 6*(int64_t)m->neq || + d->island_cache_tree[4] < 0 || d->island_cache_tree[4] > m->ntree || + d->island_cache_tree[5] < 0 || d->island_cache_tree[5] > m->nv) { + return 0; + } + mjIslandCacheView view = islandCacheView(m, d); + if (memcmp(view.efc_type, d->efc_type, d->nefc*sizeof(int)) || + memcmp(view.efc_id, d->efc_id, d->nefc*sizeof(int)) || + memcmp(view.dof_treeid, m->dof_treeid, m->nv*sizeof(int)) || + memcmp(view.tree_dofnum, m->tree_dofnum, m->ntree*sizeof(int))) { + return 0; + } + for (int i=0; i < m->neq; i++) { + if (view.eq_active[i] != d->eq_active[i] || view.eq_type[i] != m->eq_type[i] || + (m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD)) { + return 0; + } + int obj1 = m->eq_obj1id[i]; + int obj2 = m->eq_obj2id[i]; + if (m->eq_objtype[i] == mjOBJ_SITE) { + obj1 = m->site_bodyid[obj1]; + obj2 = m->site_bodyid[obj2]; + } + if (view.eq_tree1[i] != m->body_treeid[obj1] || + view.eq_tree2[i] != m->body_treeid[obj2]) { + return 0; + } + } + return 1; +} + +static void restoreIslandCache(const mjModel* m, mjData* d) { + mjIslandCacheView view = islandCacheView(m, d); +#define RESTORE(name, count) mju_copyInt(d->name, view.name, (count)) + RESTORE(tree_island, m->ntree); + RESTORE(island_ntree, d->nisland); + RESTORE(island_itreeadr, d->nisland); + RESTORE(map_itree2tree, m->ntree); + RESTORE(dof_island, m->nv); + RESTORE(island_nv, d->nisland); + RESTORE(island_idofadr, d->nisland); + RESTORE(island_dofadr, d->nisland); + RESTORE(map_dof2idof, m->nv); + RESTORE(map_idof2dof, m->nv); + RESTORE(efc_island, d->nefc); + RESTORE(island_ne, d->nisland); + RESTORE(island_nf, d->nisland); + RESTORE(island_nefc, d->nisland); + RESTORE(island_iefcadr, d->nisland); + RESTORE(map_efc2iefc, d->nefc); + RESTORE(map_iefc2efc, d->nefc); +#undef RESTORE +} + +static void saveIslandCache(const mjModel* m, mjData* d) { + if (!m->ntree || !m->neq || d->nefc != d->ne || + (int64_t)d->nefc > 6*(int64_t)m->neq) { + if (m->ntree) d->island_cache_tree[0] = 0; + return; + } + d->island_cache_tree[0] = 0; + d->island_cache_tree[1] = m->ntree; + d->island_cache_tree[2] = m->nv; + d->island_cache_tree[3] = d->nefc; + d->island_cache_tree[4] = d->nisland; + d->island_cache_tree[5] = d->nidof; + d->island_cache_tree[6] = m->neq; + d->island_cache_tree[7] = 1; + mjIslandCacheView view = islandCacheView(m, d); + for (int i=0; i < m->neq; i++) { + view.eq_active[i] = d->eq_active[i]; + view.eq_type[i] = m->eq_type[i]; + if (m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD) return; + int obj1 = m->eq_obj1id[i]; + int obj2 = m->eq_obj2id[i]; + if (m->eq_objtype[i] == mjOBJ_SITE) { + obj1 = m->site_bodyid[obj1]; + obj2 = m->site_bodyid[obj2]; + } + view.eq_tree1[i] = m->body_treeid[obj1]; + view.eq_tree2[i] = m->body_treeid[obj2]; + } +#define SAVE(name, count) mju_copyInt(view.name, d->name, (count)) + SAVE(tree_island, m->ntree); + SAVE(island_ntree, d->nisland); + SAVE(island_itreeadr, d->nisland); + SAVE(map_itree2tree, m->ntree); + SAVE(dof_island, m->nv); + SAVE(island_nv, d->nisland); + SAVE(island_idofadr, d->nisland); + SAVE(island_dofadr, d->nisland); + SAVE(map_dof2idof, m->nv); + SAVE(map_idof2dof, m->nv); + SAVE(efc_island, d->nefc); + SAVE(island_ne, d->nisland); + SAVE(island_nf, d->nisland); + SAVE(island_nefc, d->nisland); + SAVE(island_iefcadr, d->nisland); + SAVE(map_efc2iefc, d->nefc); + SAVE(map_iefc2efc, d->nefc); + mju_copyInt(view.efc_type, d->efc_type, d->nefc); + mju_copyInt(view.efc_id, d->efc_id, d->nefc); + mju_copyInt(view.dof_treeid, m->dof_treeid, m->nv); + mju_copyInt(view.tree_dofnum, m->tree_dofnum, m->ntree); +#undef SAVE + for (int i=0; i < d->nefc; i++) { + if (d->map_iefc2efc[i] != i) { + d->island_cache_tree[7] = 0; + break; + } + } + d->island_cache_tree[0] = kCacheMagic; +} + +static void copyIslandEfcVectors(mjData* d) { + if (d->island_cache_tree[0] == kCacheMagic && d->island_cache_tree[7]) { + d->iefc_type = d->efc_type; + d->iefc_id = d->efc_id; + d->iefc_frictionloss = d->efc_frictionloss; + d->iefc_D = d->efc_D; + d->iefc_R = d->efc_R; + return; + } + mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, d->nefc); + mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, d->nefc); + mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, d->nefc); + mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, d->nefc); + mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, d->nefc); +} + + //-------------------------- main entry-point ----------------------------------------------------- // discover islands: @@ -440,6 +647,16 @@ void mj_island(const mjModel* m, mjData* d) { return; } + // exact fast path for topology-stable connect/weld equality constraints + if (islandCacheMatches(m, d)) { + d->nisland = d->island_cache_tree[4]; + d->nidof = d->island_cache_tree[5]; + if (!arenaAllocIsland(m, d)) return; + restoreIslandCache(m, d); + copyIslandEfcVectors(d); + return; + } + mj_markStack(d); // union direct tree incidence and assign deterministic components @@ -605,5 +822,6 @@ void mj_island(const mjModel* m, mjData* d) { // SHOULD NOT OCCUR if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); + saveIslandCache(m, d); mj_freeStack(d); } diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3a4462ab..00762da2 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -295,6 +295,69 @@ TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); } +TEST_F(IslandTest, EqualityTopologyCacheInvalidatesExactly) { + static constexpr char xml[] = R"( + + + + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + MjDataPtr data = MakeData(model); + + auto expect_all_connected = [&] { + EXPECT_EQ(data->nisland, 1); + EXPECT_EQ(data->nidof, 3); + EXPECT_EQ(data->nefc, 6); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), + ElementsAre(0, 0, 0)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, 0)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), + ElementsAre(0, 0, 0, 0, 0, 0)); + }; + + // Identical topology is stable across repeated evaluation and ordinary reset. + mj_fwdPosition(model.get(), data.get()); + expect_all_connected(); + mj_fwdPosition(model.get(), data.get()); + expect_all_connected(); + + // Copying a warm cache produces independent data with the same result. + MjDataPtr copy(mj_copyData(nullptr, model.get(), data.get())); + ASSERT_THAT(copy.get(), NotNull()); + mj_fwdPosition(model.get(), copy.get()); + EXPECT_THAT(AsVector(copy->tree_island, model->ntree), + ElementsAre(0, 0, 0)); + + mj_resetData(model.get(), data.get()); + mj_fwdPosition(model.get(), data.get()); + expect_all_connected(); + + // Equality activation changes invalidate the cached partition. + data->eq_active[1] = 0; + mj_fwdPosition(model.get(), data.get()); + EXPECT_EQ(data->nisland, 1); + EXPECT_EQ(data->nidof, 2); + EXPECT_EQ(data->nefc, 3); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, -1)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, -1)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0, 0)); + + // The earlier copy remains independent after the source topology changes. + mj_fwdPosition(model.get(), copy.get()); + EXPECT_THAT(AsVector(copy->tree_island, model->ntree), ElementsAre(0, 0, 0)); +} + TEST_F(IslandTest, BoundedArenaSupports1024Trees) { constexpr int kTreeCount = 1024; constexpr size_t kArenaBytes = 2 * 1024 * 1024; From ba2782140f92f7c9f435351322d1a2fb052ec370 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 11 Jul 2026 19:13:30 +0530 Subject: [PATCH 04/10] Revert unvalidated island topology cache Signed-off-by: teerthsharma --- include/mujoco/mjdata.h | 5 - include/mujoco/mjxmacro.h | 6 +- src/engine/engine_io.c | 29 +--- src/engine/engine_island.c | 218 ------------------------------ test/engine/engine_island_test.cc | 63 --------- 5 files changed, 2 insertions(+), 319 deletions(-) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 800aefa3..c2fddb81 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -198,11 +198,6 @@ typedef struct mjData_ { // sleep state int* tree_asleep; // <0: awake; >=0: index cycle of sleeping trees (ntree x 1) - // opaque internal island topology cache; users must not modify - int* island_cache_tree; // tree and island topology (ntree x 28) - int* island_cache_dof; // DOF topology (nv x 8) - int* island_cache_eq; // equality topology (neq x 40) - // plugins int* plugin; // copy of m->plugin, required for deletion (nplugin x 1) uintptr_t* plugin_data; // pointer to plugin-managed data structure (nplugin x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 71f8ed39..8f36c166 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -929,10 +929,7 @@ X ( mjtNum, qfrc_inverse, nv, 1 ) \ X ( mjtNum, cacc, nbody, 6 ) \ X ( mjtNum, cfrc_int, nbody, 6 ) \ - X ( mjtNum, cfrc_ext, nbody, 6 ) \ - XIC ( int, island_cache_tree, ntree, 28 ) \ - XIC ( int, island_cache_dof, nv, 8 ) \ - XIC ( int, island_cache_eq, neq, 40 ) + X ( mjtNum, cfrc_ext, nbody, 6 ) // macro for annotating that an array size in an X macro is a member of mjData @@ -1083,6 +1080,5 @@ // to obtain only X macros for fields that are relevant for mjvScene creation, // redefine XNV to expand to nothing #define XNV X -#define XIC(type, name, nr, nc) #endif // MUJOCO_MJXMACRO_H_ diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 285561fb..482f5915 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -998,11 +998,7 @@ static void mj_setPtrData(const mjModel* m, mjData* d) { ASAN_POISON_MEMORY_REGION(ptr, PTRDIFF(d->name, ptr)); \ ptr += SKIP((intptr_t)ptr) + sizeof(type)*(m->nr)*(nc); - #undef XIC - #define XIC X MJDATA_POINTERS - #undef XIC - #define XIC(type, name, nr, nc) #undef X // check size @@ -1089,11 +1085,7 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { return; \ } - #undef XIC - #define XIC X MJDATA_POINTERS - #undef XIC - #define XIC(type, name, nr, nc) #undef X // copy stack size from model @@ -1116,9 +1108,6 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { // set pointers into buffer mj_setPtrData(m, d); - if (m->ntree) { - d->island_cache_tree[0] = 0; - } // clear threadpool d->threadpool = 0; @@ -1184,9 +1173,6 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int dest->arena = save_arena; dest->threadpool = 0; mj_setPtrData(m, dest); - if (m->ntree) { - dest->island_cache_tree[0] = 0; - } // save plugin_data, since the X macro copying block below will override it const size_t plugin_data_size = sizeof(*dest->plugin_data) * dest->nplugin; @@ -1387,32 +1373,19 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { // fill buffer with debug_value (normally 0) #ifdef ADDRESS_SANITIZER { - #undef XIC - #define XIC(type, name, nr, nc) #define X(type, name, nr, nc) memset(d->name, (int)debug_value, sizeof(type)*(m->nr)*(nc)); MJDATA_POINTERS #undef X - #undef XIC - #define XIC(type, name, nr, nc) } #else - size_t cache_offset = d->island_cache_tree ? - (size_t)PTRDIFF(d->island_cache_tree, d->buffer) : d->nbuffer; - memset(d->buffer, (int)debug_value, cache_offset); + memset(d->buffer, (int)debug_value, d->nbuffer); #endif #ifdef MEMORY_SANITIZER // under MSAN, mark the entire buffer as uninitialized __msan_allocated_memory(d->buffer, d->nbuffer); - if (m->ntree) { - d->island_cache_tree[0] = 0; - } #endif - if (debug_value && m->ntree) { - d->island_cache_tree[0] = 0; - } - // zero out user-settable state and input arrays (MSAN: mark as initialized) mju_zero(d->qpos, m->nq); mju_zero(d->qvel, m->nv); diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 207ccad8..34403b11 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -14,10 +14,8 @@ #include "engine/engine_island.h" -#include #include #include -#include #include #include @@ -429,211 +427,6 @@ static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree } -enum { - kCacheMagic = 0x49534C44, - kCacheHeader = 16, -}; - -typedef struct mjIslandCacheView_ { - int* tree_island; - int* island_ntree; - int* island_itreeadr; - int* map_itree2tree; - int* island_nv; - int* island_idofadr; - int* island_dofadr; - int* island_ne; - int* island_nf; - int* island_nefc; - int* island_iefcadr; - int* tree_dofnum; - int* dof_treeid; - int* dof_island; - int* map_dof2idof; - int* map_idof2dof; - int* efc_island; - int* efc_type; - int* efc_id; - int* map_efc2iefc; - int* map_iefc2efc; - int* eq_active; - int* eq_type; - int* eq_tree1; - int* eq_tree2; -} mjIslandCacheView; - -static mjIslandCacheView islandCacheView(const mjModel* m, const mjData* d) { - int* tree = d->island_cache_tree + kCacheHeader; - int* dof = d->island_cache_dof; - int* eq = d->island_cache_eq; - int nisland = d->island_cache_tree[4]; - int nefc = d->island_cache_tree[3]; - mjIslandCacheView view; -#define TAKE(base, name, count) view.name = base; base += (count) - TAKE(tree, tree_island, m->ntree); - TAKE(tree, island_ntree, nisland); - TAKE(tree, island_itreeadr, nisland); - TAKE(tree, map_itree2tree, m->ntree); - TAKE(tree, island_nv, nisland); - TAKE(tree, island_idofadr, nisland); - TAKE(tree, island_dofadr, nisland); - TAKE(tree, island_ne, nisland); - TAKE(tree, island_nf, nisland); - TAKE(tree, island_nefc, nisland); - TAKE(tree, island_iefcadr, nisland); - TAKE(tree, tree_dofnum, m->ntree); - TAKE(dof, dof_treeid, m->nv); - TAKE(dof, dof_island, m->nv); - TAKE(dof, map_dof2idof, m->nv); - TAKE(dof, map_idof2dof, m->nv); - TAKE(eq, efc_island, nefc); - TAKE(eq, efc_type, nefc); - TAKE(eq, efc_id, nefc); - TAKE(eq, map_efc2iefc, nefc); - TAKE(eq, map_iefc2efc, nefc); - TAKE(eq, eq_active, m->neq); - TAKE(eq, eq_type, m->neq); - TAKE(eq, eq_tree1, m->neq); - TAKE(eq, eq_tree2, m->neq); -#undef TAKE - return view; -} - -static int islandCacheMatches(const mjModel* m, const mjData* d) { - if (!m->ntree || !m->neq || d->island_cache_tree[0] != kCacheMagic || - d->island_cache_tree[1] != m->ntree || d->island_cache_tree[2] != m->nv || - d->island_cache_tree[3] != d->nefc || d->island_cache_tree[6] != m->neq || - d->nefc != d->ne || d->nefc < 0 || (int64_t)d->nefc > 6*(int64_t)m->neq || - d->island_cache_tree[4] < 0 || d->island_cache_tree[4] > m->ntree || - d->island_cache_tree[5] < 0 || d->island_cache_tree[5] > m->nv) { - return 0; - } - mjIslandCacheView view = islandCacheView(m, d); - if (memcmp(view.efc_type, d->efc_type, d->nefc*sizeof(int)) || - memcmp(view.efc_id, d->efc_id, d->nefc*sizeof(int)) || - memcmp(view.dof_treeid, m->dof_treeid, m->nv*sizeof(int)) || - memcmp(view.tree_dofnum, m->tree_dofnum, m->ntree*sizeof(int))) { - return 0; - } - for (int i=0; i < m->neq; i++) { - if (view.eq_active[i] != d->eq_active[i] || view.eq_type[i] != m->eq_type[i] || - (m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD)) { - return 0; - } - int obj1 = m->eq_obj1id[i]; - int obj2 = m->eq_obj2id[i]; - if (m->eq_objtype[i] == mjOBJ_SITE) { - obj1 = m->site_bodyid[obj1]; - obj2 = m->site_bodyid[obj2]; - } - if (view.eq_tree1[i] != m->body_treeid[obj1] || - view.eq_tree2[i] != m->body_treeid[obj2]) { - return 0; - } - } - return 1; -} - -static void restoreIslandCache(const mjModel* m, mjData* d) { - mjIslandCacheView view = islandCacheView(m, d); -#define RESTORE(name, count) mju_copyInt(d->name, view.name, (count)) - RESTORE(tree_island, m->ntree); - RESTORE(island_ntree, d->nisland); - RESTORE(island_itreeadr, d->nisland); - RESTORE(map_itree2tree, m->ntree); - RESTORE(dof_island, m->nv); - RESTORE(island_nv, d->nisland); - RESTORE(island_idofadr, d->nisland); - RESTORE(island_dofadr, d->nisland); - RESTORE(map_dof2idof, m->nv); - RESTORE(map_idof2dof, m->nv); - RESTORE(efc_island, d->nefc); - RESTORE(island_ne, d->nisland); - RESTORE(island_nf, d->nisland); - RESTORE(island_nefc, d->nisland); - RESTORE(island_iefcadr, d->nisland); - RESTORE(map_efc2iefc, d->nefc); - RESTORE(map_iefc2efc, d->nefc); -#undef RESTORE -} - -static void saveIslandCache(const mjModel* m, mjData* d) { - if (!m->ntree || !m->neq || d->nefc != d->ne || - (int64_t)d->nefc > 6*(int64_t)m->neq) { - if (m->ntree) d->island_cache_tree[0] = 0; - return; - } - d->island_cache_tree[0] = 0; - d->island_cache_tree[1] = m->ntree; - d->island_cache_tree[2] = m->nv; - d->island_cache_tree[3] = d->nefc; - d->island_cache_tree[4] = d->nisland; - d->island_cache_tree[5] = d->nidof; - d->island_cache_tree[6] = m->neq; - d->island_cache_tree[7] = 1; - mjIslandCacheView view = islandCacheView(m, d); - for (int i=0; i < m->neq; i++) { - view.eq_active[i] = d->eq_active[i]; - view.eq_type[i] = m->eq_type[i]; - if (m->eq_type[i] != mjEQ_CONNECT && m->eq_type[i] != mjEQ_WELD) return; - int obj1 = m->eq_obj1id[i]; - int obj2 = m->eq_obj2id[i]; - if (m->eq_objtype[i] == mjOBJ_SITE) { - obj1 = m->site_bodyid[obj1]; - obj2 = m->site_bodyid[obj2]; - } - view.eq_tree1[i] = m->body_treeid[obj1]; - view.eq_tree2[i] = m->body_treeid[obj2]; - } -#define SAVE(name, count) mju_copyInt(view.name, d->name, (count)) - SAVE(tree_island, m->ntree); - SAVE(island_ntree, d->nisland); - SAVE(island_itreeadr, d->nisland); - SAVE(map_itree2tree, m->ntree); - SAVE(dof_island, m->nv); - SAVE(island_nv, d->nisland); - SAVE(island_idofadr, d->nisland); - SAVE(island_dofadr, d->nisland); - SAVE(map_dof2idof, m->nv); - SAVE(map_idof2dof, m->nv); - SAVE(efc_island, d->nefc); - SAVE(island_ne, d->nisland); - SAVE(island_nf, d->nisland); - SAVE(island_nefc, d->nisland); - SAVE(island_iefcadr, d->nisland); - SAVE(map_efc2iefc, d->nefc); - SAVE(map_iefc2efc, d->nefc); - mju_copyInt(view.efc_type, d->efc_type, d->nefc); - mju_copyInt(view.efc_id, d->efc_id, d->nefc); - mju_copyInt(view.dof_treeid, m->dof_treeid, m->nv); - mju_copyInt(view.tree_dofnum, m->tree_dofnum, m->ntree); -#undef SAVE - for (int i=0; i < d->nefc; i++) { - if (d->map_iefc2efc[i] != i) { - d->island_cache_tree[7] = 0; - break; - } - } - d->island_cache_tree[0] = kCacheMagic; -} - -static void copyIslandEfcVectors(mjData* d) { - if (d->island_cache_tree[0] == kCacheMagic && d->island_cache_tree[7]) { - d->iefc_type = d->efc_type; - d->iefc_id = d->efc_id; - d->iefc_frictionloss = d->efc_frictionloss; - d->iefc_D = d->efc_D; - d->iefc_R = d->efc_R; - return; - } - mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, d->nefc); - mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, d->nefc); - mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, d->nefc); - mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, d->nefc); - mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, d->nefc); -} - - //-------------------------- main entry-point ----------------------------------------------------- // discover islands: @@ -647,16 +440,6 @@ void mj_island(const mjModel* m, mjData* d) { return; } - // exact fast path for topology-stable connect/weld equality constraints - if (islandCacheMatches(m, d)) { - d->nisland = d->island_cache_tree[4]; - d->nidof = d->island_cache_tree[5]; - if (!arenaAllocIsland(m, d)) return; - restoreIslandCache(m, d); - copyIslandEfcVectors(d); - return; - } - mj_markStack(d); // union direct tree incidence and assign deterministic components @@ -822,6 +605,5 @@ void mj_island(const mjModel* m, mjData* d) { // SHOULD NOT OCCUR if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); - saveIslandCache(m, d); mj_freeStack(d); } diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 00762da2..3a4462ab 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -295,69 +295,6 @@ TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); } -TEST_F(IslandTest, EqualityTopologyCacheInvalidatesExactly) { - static constexpr char xml[] = R"( - - - - - - - - - - - - -)"; - char error[1024] = {}; - MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); - ASSERT_THAT(model.get(), NotNull()) << error; - MjDataPtr data = MakeData(model); - - auto expect_all_connected = [&] { - EXPECT_EQ(data->nisland, 1); - EXPECT_EQ(data->nidof, 3); - EXPECT_EQ(data->nefc, 6); - EXPECT_THAT(AsVector(data->tree_island, model->ntree), - ElementsAre(0, 0, 0)); - EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, 0)); - EXPECT_THAT(AsVector(data->efc_island, data->nefc), - ElementsAre(0, 0, 0, 0, 0, 0)); - }; - - // Identical topology is stable across repeated evaluation and ordinary reset. - mj_fwdPosition(model.get(), data.get()); - expect_all_connected(); - mj_fwdPosition(model.get(), data.get()); - expect_all_connected(); - - // Copying a warm cache produces independent data with the same result. - MjDataPtr copy(mj_copyData(nullptr, model.get(), data.get())); - ASSERT_THAT(copy.get(), NotNull()); - mj_fwdPosition(model.get(), copy.get()); - EXPECT_THAT(AsVector(copy->tree_island, model->ntree), - ElementsAre(0, 0, 0)); - - mj_resetData(model.get(), data.get()); - mj_fwdPosition(model.get(), data.get()); - expect_all_connected(); - - // Equality activation changes invalidate the cached partition. - data->eq_active[1] = 0; - mj_fwdPosition(model.get(), data.get()); - EXPECT_EQ(data->nisland, 1); - EXPECT_EQ(data->nidof, 2); - EXPECT_EQ(data->nefc, 3); - EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, -1)); - EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 0, -1)); - EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0, 0)); - - // The earlier copy remains independent after the source topology changes. - mj_fwdPosition(model.get(), copy.get()); - EXPECT_THAT(AsVector(copy->tree_island, model->ntree), ElementsAre(0, 0, 0)); -} - TEST_F(IslandTest, BoundedArenaSupports1024Trees) { constexpr int kTreeCount = 1024; constexpr size_t kArenaBytes = 2 * 1024 * 1024; From 5d91d878c2b2f04847130c80953a73180c1d93d6 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 18 Jul 2026 16:59:24 +0530 Subject: [PATCH 05/10] Benchmark and expose disjoint-set islands Signed-off-by: teerthsharma --- src/engine/engine_island.c | 109 ++++--- src/engine/engine_island.h | 6 + test/benchmark/CMakeLists.txt | 6 + test/benchmark/island_benchmark_test.cc | 374 ++++++++++++++++++++++++ test/engine/engine_island_test.cc | 313 ++++++++++++++++++++ 5 files changed, 772 insertions(+), 36 deletions(-) create mode 100644 test/benchmark/island_benchmark_test.cc diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 34403b11..4f53d2db 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -84,7 +84,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) { //-------------------------- flood-fill and graph construction ------------------------------------ // find the canonical root of an active tree and compress its path -static int dsuFind(int* parent, int tree) { +static inline int dsuFind(int* parent, int tree) { int root = tree; while (parent[root] != root) { root = parent[root]; @@ -101,13 +101,13 @@ static int dsuFind(int* parent, int tree) { // initialize all trees as inactive -static void dsuInit(int* parent, int ntree) { +static inline void dsuInit(int* parent, int ntree) { mju_fillInt(parent, -1, ntree); } // activate and union two incident trees; -1 denotes a static endpoint -static void dsuUnion(int* parent, int tree1, int tree2) { +static inline void dsuUnion(int* parent, int tree1, int tree2) { if (tree1 == -1 && tree2 == -1) { mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR return; @@ -119,6 +119,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) { if (parent[tree1] == -1) parent[tree1] = tree1; if (parent[tree2] == -1) parent[tree2] = tree2; + if (parent[tree1] == parent[tree2]) return; + int root1 = dsuFind(parent, tree1); int root2 = dsuFind(parent, tree2); if (root1 < root2) { @@ -130,7 +132,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) { // assign deterministic island ids in ascending canonical-root order -static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { +static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, + int* nidof) { int nisland = 0; *nidof = 0; for (int tree=0; tree < ntree; tree++) { @@ -139,14 +142,41 @@ static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree continue; } - int root = dsuFind(parent, tree); - island[tree] = root == tree ? nisland++ : island[root]; + if (parent[tree] == tree) { + island[tree] = nisland++; + } else { + // Union always links the larger root to the smaller root. Since trees are visited in + // ascending order, this predecessor has already been compressed and assigned an island. + parent[tree] = parent[parent[tree]]; + island[tree] = island[parent[tree]]; + } *nidof += tree_dofnum[tree]; } return nisland; } + +// exported private wrappers for direct unit tests and benchmarks +int _mjPRIVATE_dsuFind(int* parent, int tree) { + return dsuFind(parent, tree); +} + + +void _mjPRIVATE_dsuInit(int* parent, int ntree) { + dsuInit(parent, ntree); +} + + +void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2) { + dsuUnion(parent, tree1, tree2); +} + + +int _mjPRIVATE_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { + return dsuAssign(island, parent, tree_dofnum, ntree, nidof); +} + // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: // island (nr) - island index assigned to vertex, -1 if vertex has no edges @@ -353,7 +383,7 @@ static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) { // activate and union all trees with direct incidence in a constraint -static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) { +static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, int* efc_tree) { int nefc = d->nefc; int efc_type = -1; int efc_id = -1; @@ -361,9 +391,10 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) // iterate over constraints and union incident trees for (int i=0; i < nefc; i++) { // row i is still in the same constraint: skip it - if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { + if (i > 0 && efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { // unless it is a flex equality, where the tree pattern changes per dof if (!isFlexEquality(m, efc_type, efc_id)) { + efc_tree[i] = efc_tree[i-1]; continue; } } @@ -378,6 +409,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) int tree1 = treeNext(m, d, i, &iter); if (tree1 != -2) { int tree2 = treeNext(m, d, i, &iter); + efc_tree[i] = tree1 == -1 ? tree2 : tree1; // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { @@ -393,35 +425,41 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR } } -} -// assign each constraint from its first non-negative incident tree -static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree_island) { - int efc_type = -1; - int efc_id = -1; - - for (int i=0; i < d->nefc; i++) { - // reuse assignment for repeated scalar rows, except flex equality rows - if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i] && - !isFlexEquality(m, efc_type, efc_id)) { - d->efc_island[i] = d->efc_island[i-1]; + // Flex stiffness couples all vertices (nodes for interpolated flexes) without a constraint + // row representing the coupling. Union the awake dynamic trees of each stiffness-active flex. + for (int f=0; f < m->nflex; f++) { + if (m->flex_rigid[f] || m->flex_dim[f] < 2) { + continue; + } + int sadr = m->flex_stiffnessadr[f]; + if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) { continue; } - efc_type = d->efc_type[i]; - efc_id = d->efc_id[i]; - mjTreeIter iter; - treeIterInit(m, d, i, &iter); - - int tree; - do { - tree = treeNext(m, d, i, &iter); - } while (tree == -1); - - if (tree == -2) { - mjERROR("no dynamic tree found for constraint %d", i); // SHOULD NOT OCCUR + int num, adr; + const int* bodyid; + if (m->flex_interp[f]) { + num = m->flex_nodenum[f]; + adr = m->flex_nodeadr[f]; + bodyid = m->flex_nodebodyid; } else { - d->efc_island[i] = tree_island[tree]; + num = m->flex_vertnum[f]; + adr = m->flex_vertadr[f]; + bodyid = m->flex_vertbodyid; + } + + int tree1 = -1; + for (int j=0; j < num; j++) { + int tree2 = m->body_treeid[bodyid[adr+j]]; + if (tree2 < 0 || tree2 == tree1 || !d->tree_awake[tree2]) { + continue; + } + if (tree1 < 0) { + tree1 = tree2; + } else { + dsuUnion(parent, tree1, tree2); + } } } } @@ -443,9 +481,10 @@ void mj_island(const mjModel* m, mjData* d) { mj_markStack(d); // union direct tree incidence and assign deterministic components + int* efc_tree = mjSTACKALLOC(d, nefc, int); int* parent = mjSTACKALLOC(d, ntree, int); dsuInit(parent, ntree); - unionConstraintTrees(m, d, parent); + unionConstraintTrees(m, d, parent, efc_tree); int* tree_island = mjSTACKALLOC(d, ntree, int); int nidof; d->nisland = dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); @@ -558,14 +597,12 @@ void mj_island(const mjModel* m, mjData* d) { // ------------------------------------- constraints --------------------------------------------- - // compute efc_island from first non-negative tree of each constraint - assignConstraintIslands(m, d, tree_island); - // compute efc_island, island_{ne,nf,nefc} mju_zeroInt(d->island_ne, nisland); mju_zeroInt(d->island_nf, nisland); mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { + d->efc_island[i] = tree_island[efc_tree[i]]; int island = d->efc_island[i]; d->island_nefc[island]++; switch (d->efc_type[i]) { diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 140bcdf6..7e239b40 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -23,6 +23,12 @@ extern "C" { #endif +MJAPI int _mjPRIVATE_dsuFind(int* parent, int tree); +MJAPI void _mjPRIVATE_dsuInit(int* parent, int ntree); +MJAPI void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2); +MJAPI int _mjPRIVATE_dsuAssign(int* island, int* parent, + const int* tree_dofnum, int ntree, int* nidof); + // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index b9aff09f..6821aa03 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -79,6 +79,12 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) +mujoco_test( + island_benchmark_test + MAIN_TARGET benchmark::benchmark_main + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers +) + mujoco_test( engine_util_sparse_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/island_benchmark_test.cc b/test/benchmark/island_benchmark_test.cc new file mode 100644 index 00000000..21a922b8 --- /dev/null +++ b/test/benchmark/island_benchmark_test.cc @@ -0,0 +1,374 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Benchmarks island discovery on deterministic geodesic Rips graphs. The +// corpus spans the connectivity transition of points sampled on S^2 and adds +// MuJoCo-relevant static and repeated incidences. Corpus construction and +// validation are deliberately outside the timed region. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "src/engine/engine_island.h" + +namespace mujoco { +namespace { + +struct Edge { + int first; + int second; +}; + +struct GraphCase { + std::string name; + int node_count; + std::vector incidences; + std::vector expected_partition; + int active_nodes; + int expected_components; + int pre_bridge_components; + bool bridge_added; + std::uint64_t expected_checksum; +}; + +struct Point { + double x; + double y; + double z; +}; + +std::uint64_t SplitMix64(std::uint64_t& state) { + state += 0x9e3779b97f4a7c15ULL; + std::uint64_t value = state; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +double Uniform01(std::uint64_t& state) { + return static_cast(SplitMix64(state) >> 11) * 0x1.0p-53; +} + +std::vector SampleSphere(int count, std::uint64_t seed) { + constexpr double kTwoPi = 6.283185307179586476925286766559; + std::vector points; + points.reserve(count); + for (int i = 0; i < count; ++i) { + const double z = 2.0 * Uniform01(seed) - 1.0; + const double angle = kTwoPi * Uniform01(seed); + const double radial = std::sqrt(std::max(0.0, 1.0 - z*z)); + points.push_back({radial * std::cos(angle), radial * std::sin(angle), z}); + } + return points; +} + +double Dot(const Point& a, const Point& b) { + return a.x*b.x + a.y*b.y + a.z*b.z; +} + +std::vector RipsEdges(const std::vector& points, double target_degree) { + const double probability = target_degree / (points.size() - 1); + const double radius = 2.0 * std::asin(std::sqrt(probability)); + const double minimum_dot = std::cos(radius); + std::vector edges; + for (int i = 0; i < static_cast(points.size()); ++i) { + for (int j = i + 1; j < static_cast(points.size()); ++j) { + if (Dot(points[i], points[j]) >= minimum_dot) { + edges.push_back({i, j}); + } + } + } + return edges; +} + +std::vector CanonicalPartition(int node_count, const std::vector& edges) { + std::vector> adjacency(node_count); + std::vector active(node_count, false); + for (const Edge& edge : edges) { + active[edge.first] = true; + active[edge.second] = true; + if (edge.first != edge.second) { + adjacency[edge.first].push_back(edge.second); + adjacency[edge.second].push_back(edge.first); + } + } + + std::vector partition(node_count, -1); + std::queue pending; + for (int start = 0; start < node_count; ++start) { + if (!active[start] || partition[start] != -1) { + continue; + } + partition[start] = start; + pending.push(start); + while (!pending.empty()) { + const int node = pending.front(); + pending.pop(); + for (int neighbor : adjacency[node]) { + if (partition[neighbor] == -1) { + partition[neighbor] = start; + pending.push(neighbor); + } + } + } + } + return partition; +} + +int CountComponents(const std::vector& partition) { + int count = 0; + for (int node = 0; node < static_cast(partition.size()); ++node) { + count += partition[node] == node; + } + return count; +} + +std::uint64_t PartitionChecksum(const std::vector& partition) { + std::uint64_t hash = 1469598103934665603ULL; + for (int value : partition) { + hash ^= static_cast(value); + hash *= 1099511628211ULL; + } + return hash; +} + +std::vector CanonicalizeLabels(const std::vector& labels) { + std::vector minimum(labels.size(), std::numeric_limits::max()); + for (int node = 0; node < static_cast(labels.size()); ++node) { + if (labels[node] >= 0) { + minimum[labels[node]] = std::min(minimum[labels[node]], node); + } + } + std::vector canonical(labels.size(), -1); + for (int node = 0; node < static_cast(labels.size()); ++node) { + if (labels[node] >= 0) { + canonical[node] = minimum[labels[node]]; + } + } + return canonical; +} + +void DeterministicShuffle(std::vector& edges, std::uint64_t seed) { + for (std::size_t i = edges.size(); i > 1; --i) { + const std::size_t j = SplitMix64(seed) % i; + std::swap(edges[i - 1], edges[j]); + } +} + +bool AddCriticalBridge(const std::vector& points, std::vector& edges) { + const std::vector partition = CanonicalPartition(points.size(), edges); + if (CountComponents(partition) < 2) { + return false; + } + + double best_dot = -2.0; + Edge bridge{-1, -1}; + for (int i = 0; i < static_cast(points.size()); ++i) { + for (int j = i + 1; j < static_cast(points.size()); ++j) { + if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] && + Dot(points[i], points[j]) > best_dot) { + best_dot = Dot(points[i], points[j]); + bridge = {i, j}; + } + } + } + if (bridge.first >= 0) { + edges.push_back(bridge); + return true; + } + return false; +} + +GraphCase MakeCase(std::string name, int node_count, double target_degree, + std::uint64_t seed, bool critical_bridge, bool static_rows, + bool repeated_rows) { + const std::vector points = SampleSphere(node_count, seed); + std::vector edges = RipsEdges(points, target_degree); + const int pre_bridge_components = + critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1; + bool bridge_added = false; + if (critical_bridge) { + bridge_added = AddCriticalBridge(points, edges); + } + + const std::vector unique_edges = edges; + if (static_rows) { + for (int node = 0; node < node_count; node += 17) { + edges.push_back({node, node}); + } + } + if (repeated_rows) { + for (std::size_t i = 0; i < unique_edges.size(); i += 11) { + edges.push_back(unique_edges[i]); + edges.push_back({unique_edges[i].second, unique_edges[i].first}); + } + } + DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL); + + std::vector expected = CanonicalPartition(node_count, edges); + const int active_nodes = std::count_if(expected.begin(), expected.end(), + [](int component) { return component >= 0; }); + const int components = CountComponents(expected); + const std::uint64_t checksum = PartitionChecksum(expected); + return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes, + components, pre_bridge_components, bridge_added, checksum}; +} + +const std::vector& Corpus() { + static const std::vector corpus = { + MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false), + MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0), + 0x33960002ULL, true, false, false), + MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)), + 0x33960003ULL, false, false, false), + MakeCase("GroundedStaticRepeated_S2Rips_256", 256, + 2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true), + MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL, + false, false, true), + MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)), + 0x33960006ULL, true, false, false), + }; + return corpus; +} + +struct FloodFillWorkspace { + explicit FloodFillWorkspace(int node_count) + : adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count), + colind(node_count * node_count), stack(node_count * node_count + node_count), + island(node_count) {} + + std::vector adjacency; + std::vector rownnz; + std::vector rowadr; + std::vector colind; + std::vector stack; + std::vector island; +}; + +int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) { + const int n = graph.node_count; + std::fill(work.adjacency.begin(), work.adjacency.end(), 0); + std::fill(work.rownnz.begin(), work.rownnz.end(), 0); + for (const Edge& edge : graph.incidences) { + work.adjacency[edge.first*n + edge.second] = 1; + work.adjacency[edge.second*n + edge.first] = 1; + } + + int address = 0; + for (int row = 0; row < n; ++row) { + work.rowadr[row] = address; + for (int column = 0; column < n; ++column) { + if (work.adjacency[row*n + column]) { + work.colind[address++] = column; + ++work.rownnz[row]; + } + } + } + return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(), + work.colind.data(), work.stack.data()); +} + +struct DsuWorkspace { + explicit DsuWorkspace(int node_count) + : parent(node_count), island(node_count), dof_count(node_count, 1) {} + + std::vector parent; + std::vector island; + std::vector dof_count; +}; + +int RunDsu(const GraphCase& graph, DsuWorkspace& work) { + _mjPRIVATE_dsuInit(work.parent.data(), graph.node_count); + for (const Edge& edge : graph.incidences) { + _mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second); + } + int dof_count = 0; + return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(), + graph.node_count, &dof_count); +} + +bool Validate(const GraphCase& graph) { + FloodFillWorkspace flood(graph.node_count); + DsuWorkspace dsu(graph.node_count); + const int flood_components = RunFloodFill(graph, flood); + const int dsu_components = RunDsu(graph, dsu); + const bool bridge_valid = graph.pre_bridge_components < 0 || + (graph.bridge_added && + graph.pre_bridge_components == graph.expected_components + 1); + return bridge_valid && flood_components == graph.expected_components && + dsu_components == graph.expected_components && + CanonicalizeLabels(flood.island) == graph.expected_partition && + CanonicalizeLabels(dsu.island) == graph.expected_partition && + PartitionChecksum(graph.expected_partition) == graph.expected_checksum; +} + +void BM_FloodFill(benchmark::State& state, const GraphCase* graph) { + if (!Validate(*graph)) { + state.SkipWithError("invalid S2-Rips graph fixture"); + return; + } + FloodFillWorkspace work(graph->node_count); + state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + + " active=" + std::to_string(graph->active_nodes) + + " components=" + std::to_string(graph->expected_components) + + " pre_bridge=" + std::to_string(graph->pre_bridge_components) + + " bridge_added=" + std::to_string(graph->bridge_added) + + " checksum=" + std::to_string(graph->expected_checksum)); + for (auto _ : state) { + int components = RunFloodFill(*graph, work); + benchmark::DoNotOptimize(components); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * graph->incidences.size()); +} + +void BM_Dsu(benchmark::State& state, const GraphCase* graph) { + if (!Validate(*graph)) { + state.SkipWithError("invalid S2-Rips graph fixture"); + return; + } + DsuWorkspace work(graph->node_count); + state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + + " active=" + std::to_string(graph->active_nodes) + + " components=" + std::to_string(graph->expected_components) + + " pre_bridge=" + std::to_string(graph->pre_bridge_components) + + " bridge_added=" + std::to_string(graph->bridge_added) + + " checksum=" + std::to_string(graph->expected_checksum)); + for (auto _ : state) { + int components = RunDsu(*graph, work); + benchmark::DoNotOptimize(components); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * graph->incidences.size()); +} + +const bool kRegistered = [] { + for (const GraphCase& graph : Corpus()) { + benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph); + benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph); + } + return true; +}(); + +} // namespace +} // namespace mujoco diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3a4462ab..be1f2832 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -16,6 +16,8 @@ #include "src/engine/engine_island.h" +#include +#include #include #include @@ -35,6 +37,310 @@ using ::testing::NotNull; using ::testing::Pointwise; using IslandTest = MujocoTest; +TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) { + int parent[] = {8, 6, 7, 5}; + + _mjPRIVATE_dsuInit(parent, 0); + EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5)); + + _mjPRIVATE_dsuInit(parent, 4); + EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1)); +} + +TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) { + int parent[] = {0, 0, 1, 2, 3}; + + EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0); + EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0)); +} + +TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) { + int parent[] = {-1, -1, -1, -1, -1, -1}; + + _mjPRIVATE_dsuUnion(parent, -1, 4); + _mjPRIVATE_dsuUnion(parent, 3, -1); + _mjPRIVATE_dsuUnion(parent, 5, 2); + _mjPRIVATE_dsuUnion(parent, 4, 5); + _mjPRIVATE_dsuUnion(parent, 3, 4); + + EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); + for (int tree = 2; tree < 6; ++tree) { + EXPECT_EQ(_mjPRIVATE_dsuFind(parent, tree), 2); + } + EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); +} + +TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) { + int parent[] = {-1, -1, -1, -1}; + _mjPRIVATE_dsuUnion(parent, 3, 1); + _mjPRIVATE_dsuUnion(parent, 2, 1); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); + + _mjPRIVATE_dsuUnion(parent, 1, 3); + _mjPRIVATE_dsuUnion(parent, 3, 1); + _mjPRIVATE_dsuUnion(parent, 2, 2); + _mjPRIVATE_dsuUnion(parent, -1, 2); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); +} + +TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) { + int self_parent[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(self_parent, 1, 1); + EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1)); + + int static_first[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(static_first, -1, 2); + EXPECT_THAT(static_first, ElementsAre(-1, -1, 2)); + + int static_second[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(static_second, 0, -1); + EXPECT_THAT(static_second, ElementsAre(0, -1, -1)); +} + +TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) { + int distinct_parent[] = {0, 0, 2, 2}; + _mjPRIVATE_dsuUnion(distinct_parent, 1, 3); + EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2)); + + int shared_parent[] = {0, 0, 0, 3}; + _mjPRIVATE_dsuUnion(shared_parent, 1, 2); + EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3)); + + int long_paths[] = {0, 0, 1, 3, 3, 4}; + _mjPRIVATE_dsuUnion(long_paths, 2, 5); + EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) { + int parent[] = {-1, -1, -1, -1, -1, -1}; + _mjPRIVATE_dsuUnion(parent, 0, 1); + _mjPRIVATE_dsuUnion(parent, 1, 2); + _mjPRIVATE_dsuUnion(parent, 2, 0); + _mjPRIVATE_dsuUnion(parent, 0, 2); + _mjPRIVATE_dsuUnion(parent, 3, 4); + _mjPRIVATE_dsuUnion(parent, 4, 5); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3)); + + _mjPRIVATE_dsuUnion(parent, 5, 0); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) { + int parent[] = {-1, 1, 1, 3}; + + EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1), + "self-incidence of the static tree"); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3)); +} + +TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) { + int island[] = {71}; + int parent[] = {72}; + const int tree_dofnum[] = {73}; + int nidof = -1; + + EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0); + EXPECT_EQ(nidof, 0); + EXPECT_THAT(island, ElementsAre(71)); + EXPECT_THAT(parent, ElementsAre(72)); + + parent[0] = -1; + EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0); + EXPECT_EQ(nidof, 0); + EXPECT_THAT(island, ElementsAre(-1)); + EXPECT_THAT(parent, ElementsAre(-1)); +} + +TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) { + int parent[] = {-1, 1, 1, 2, 4, 4, 6}; + const int tree_dofnum[] = {1000, 0, 3, 5, 7, 11, 13}; + int island[] = {9, 9, 9, 9, 9, 9, 9}; + int nidof = -1; + + EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3); + EXPECT_EQ(nidof, 39); + EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2)); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6)); +} + +TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) { + int parent[] = {-1, 1, 1, 2, 4, 4, 5, 7, 7, 8}; + const int tree_dofnum[] = {99, 0, 2, 3, 0, 5, 7, 11, 0, 13}; + int island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9}; + int nidof = -1; + + EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3); + EXPECT_EQ(nidof, 41); + EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2)); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 4, 7, 7, 7)); +} + +TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) { + constexpr int kTreeCount = 4096; + std::vector parent(kTreeCount); + std::vector island(kTreeCount, -2); + std::vector tree_dofnum(kTreeCount); + parent[0] = 0; + int expected_nidof = 0; + for (int tree = 1; tree < kTreeCount; ++tree) { + parent[tree] = tree - 1; + tree_dofnum[tree] = tree % 5; + expected_nidof += tree_dofnum[tree]; + } + + int nidof = -1; + EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + kTreeCount, &nidof), + 1); + EXPECT_EQ(nidof, expected_nidof); + for (int tree = 0; tree < kTreeCount; ++tree) { + EXPECT_EQ(island[tree], 0); + EXPECT_EQ(parent[tree], 0); + } +} + +TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { + constexpr int kTreeCount = 4096; + std::vector parent(kTreeCount); + std::vector island(kTreeCount, -2); + std::vector tree_dofnum(kTreeCount); + _mjPRIVATE_dsuInit(parent.data(), kTreeCount); + + int expected_nidof = 0; + for (int tree = kTreeCount - 1; tree > 0; --tree) { + _mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1); + } + for (int tree = 0; tree < kTreeCount; ++tree) { + tree_dofnum[tree] = tree % 7; + expected_nidof += tree_dofnum[tree]; + } + + int nidof = -1; + EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + kTreeCount, &nidof), + 1); + EXPECT_EQ(nidof, expected_nidof); + for (int tree = 0; tree < kTreeCount; ++tree) { + EXPECT_EQ(island[tree], 0); + EXPECT_EQ(parent[tree], 0); + } +} + +TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { + constexpr uint32_t kSeed = 0x5eed3396u; + constexpr int kTrials = 2000; + uint32_t state = kSeed; + auto next = [&state]() { + state = state * 1664525u + 1013904223u; + return state; + }; + + for (int trial = 0; trial < kTrials; ++trial) { + const int ntree = 1 + next() % 64; + const int nedge = next() % 192; + std::vector> edges; + edges.reserve(nedge); + for (int edge = 0; edge < nedge; ++edge) { + int tree1; + int tree2; + switch (next() % 8) { + case 0: + tree1 = -1; + tree2 = next() % ntree; + break; + case 1: + tree1 = next() % ntree; + tree2 = -1; + break; + case 2: + tree1 = next() % ntree; + tree2 = tree1; + break; + case 3: + if (!edges.empty()) { + const auto& previous = edges[next() % edges.size()]; + tree1 = previous[0]; + tree2 = previous[1]; + break; + } + [[fallthrough]]; + case 4: + if (!edges.empty()) { + const auto& previous = edges[next() % edges.size()]; + tree1 = previous[1]; + tree2 = previous[0]; + break; + } + [[fallthrough]]; + default: + tree1 = next() % ntree; + tree2 = next() % ntree; + break; + } + edges.push_back({tree1, tree2}); + } + + std::vector parent(ntree); + _mjPRIVATE_dsuInit(parent.data(), ntree); + for (const auto& edge : edges) { + _mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]); + } + + std::vector active(ntree); + std::vector> adjacent(ntree); + for (const auto& edge : edges) { + if (edge[0] >= 0) active[edge[0]] = 1; + if (edge[1] >= 0) active[edge[1]] = 1; + if (edge[0] >= 0 && edge[1] >= 0) { + adjacent[edge[0]].push_back(edge[1]); + adjacent[edge[1]].push_back(edge[0]); + } + } + + std::vector expected_island(ntree, -1); + std::vector expected_parent(ntree, -1); + int expected_nisland = 0; + for (int start = 0; start < ntree; ++start) { + if (!active[start] || expected_island[start] != -1) continue; + std::vector pending = {start}; + std::vector component; + expected_island[start] = expected_nisland; + while (!pending.empty()) { + const int tree = pending.back(); + pending.pop_back(); + component.push_back(tree); + for (int neighbor : adjacent[tree]) { + if (expected_island[neighbor] == -1) { + expected_island[neighbor] = expected_nisland; + pending.push_back(neighbor); + } + } + } + for (int tree : component) expected_parent[tree] = start; + ++expected_nisland; + } + + std::vector tree_dofnum(ntree); + int expected_nidof = 0; + for (int tree = 0; tree < ntree; ++tree) { + tree_dofnum[tree] = next() % 8; + if (active[tree]) expected_nidof += tree_dofnum[tree]; + } + std::vector island(ntree, -2); + int nidof = -1; + const int nisland = _mjPRIVATE_dsuAssign( + island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof); + + SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial + << " ntree=" << ntree << " nedge=" << nedge); + EXPECT_EQ(nisland, expected_nisland); + EXPECT_EQ(nidof, expected_nidof); + EXPECT_EQ(island, expected_island); + EXPECT_EQ(parent, expected_parent); + } +} + TEST_F(IslandTest, FloodFillSingleton) { // adjacency matrix for the graph 0 1 2 // U U @@ -532,6 +838,13 @@ TEST_F(IslandTest, IslandEfc) { EXPECT_EQ(data->nf, 2); EXPECT_EQ(data->nl, 1); EXPECT_EQ(data->nefc, 30); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(1, 0, 0, 6)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 1, 1, 0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(6, 17, 1, 6)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), + ElementsAre(0, 3, 3, 3, 3, 3, 3, 1, 2, 0, + 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1)); mj_deleteData(data); mj_deleteModel(model); From d9c8bcbc8f862dea6544cbe08586bbd783038add Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 08:43:20 +0530 Subject: [PATCH 06/10] Remove temporary island benchmark Signed-off-by: teerthsharma --- test/benchmark/CMakeLists.txt | 6 - test/benchmark/island_benchmark_test.cc | 374 ------------------------ 2 files changed, 380 deletions(-) delete mode 100644 test/benchmark/island_benchmark_test.cc diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index 6821aa03..b9aff09f 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -79,12 +79,6 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) -mujoco_test( - island_benchmark_test - MAIN_TARGET benchmark::benchmark_main - ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers -) - mujoco_test( engine_util_sparse_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/island_benchmark_test.cc b/test/benchmark/island_benchmark_test.cc deleted file mode 100644 index 21a922b8..00000000 --- a/test/benchmark/island_benchmark_test.cc +++ /dev/null @@ -1,374 +0,0 @@ -// Copyright 2026 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Benchmarks island discovery on deterministic geodesic Rips graphs. The -// corpus spans the connectivity transition of points sampled on S^2 and adds -// MuJoCo-relevant static and repeated incidences. Corpus construction and -// validation are deliberately outside the timed region. - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "src/engine/engine_island.h" - -namespace mujoco { -namespace { - -struct Edge { - int first; - int second; -}; - -struct GraphCase { - std::string name; - int node_count; - std::vector incidences; - std::vector expected_partition; - int active_nodes; - int expected_components; - int pre_bridge_components; - bool bridge_added; - std::uint64_t expected_checksum; -}; - -struct Point { - double x; - double y; - double z; -}; - -std::uint64_t SplitMix64(std::uint64_t& state) { - state += 0x9e3779b97f4a7c15ULL; - std::uint64_t value = state; - value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; - value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; - return value ^ (value >> 31); -} - -double Uniform01(std::uint64_t& state) { - return static_cast(SplitMix64(state) >> 11) * 0x1.0p-53; -} - -std::vector SampleSphere(int count, std::uint64_t seed) { - constexpr double kTwoPi = 6.283185307179586476925286766559; - std::vector points; - points.reserve(count); - for (int i = 0; i < count; ++i) { - const double z = 2.0 * Uniform01(seed) - 1.0; - const double angle = kTwoPi * Uniform01(seed); - const double radial = std::sqrt(std::max(0.0, 1.0 - z*z)); - points.push_back({radial * std::cos(angle), radial * std::sin(angle), z}); - } - return points; -} - -double Dot(const Point& a, const Point& b) { - return a.x*b.x + a.y*b.y + a.z*b.z; -} - -std::vector RipsEdges(const std::vector& points, double target_degree) { - const double probability = target_degree / (points.size() - 1); - const double radius = 2.0 * std::asin(std::sqrt(probability)); - const double minimum_dot = std::cos(radius); - std::vector edges; - for (int i = 0; i < static_cast(points.size()); ++i) { - for (int j = i + 1; j < static_cast(points.size()); ++j) { - if (Dot(points[i], points[j]) >= minimum_dot) { - edges.push_back({i, j}); - } - } - } - return edges; -} - -std::vector CanonicalPartition(int node_count, const std::vector& edges) { - std::vector> adjacency(node_count); - std::vector active(node_count, false); - for (const Edge& edge : edges) { - active[edge.first] = true; - active[edge.second] = true; - if (edge.first != edge.second) { - adjacency[edge.first].push_back(edge.second); - adjacency[edge.second].push_back(edge.first); - } - } - - std::vector partition(node_count, -1); - std::queue pending; - for (int start = 0; start < node_count; ++start) { - if (!active[start] || partition[start] != -1) { - continue; - } - partition[start] = start; - pending.push(start); - while (!pending.empty()) { - const int node = pending.front(); - pending.pop(); - for (int neighbor : adjacency[node]) { - if (partition[neighbor] == -1) { - partition[neighbor] = start; - pending.push(neighbor); - } - } - } - } - return partition; -} - -int CountComponents(const std::vector& partition) { - int count = 0; - for (int node = 0; node < static_cast(partition.size()); ++node) { - count += partition[node] == node; - } - return count; -} - -std::uint64_t PartitionChecksum(const std::vector& partition) { - std::uint64_t hash = 1469598103934665603ULL; - for (int value : partition) { - hash ^= static_cast(value); - hash *= 1099511628211ULL; - } - return hash; -} - -std::vector CanonicalizeLabels(const std::vector& labels) { - std::vector minimum(labels.size(), std::numeric_limits::max()); - for (int node = 0; node < static_cast(labels.size()); ++node) { - if (labels[node] >= 0) { - minimum[labels[node]] = std::min(minimum[labels[node]], node); - } - } - std::vector canonical(labels.size(), -1); - for (int node = 0; node < static_cast(labels.size()); ++node) { - if (labels[node] >= 0) { - canonical[node] = minimum[labels[node]]; - } - } - return canonical; -} - -void DeterministicShuffle(std::vector& edges, std::uint64_t seed) { - for (std::size_t i = edges.size(); i > 1; --i) { - const std::size_t j = SplitMix64(seed) % i; - std::swap(edges[i - 1], edges[j]); - } -} - -bool AddCriticalBridge(const std::vector& points, std::vector& edges) { - const std::vector partition = CanonicalPartition(points.size(), edges); - if (CountComponents(partition) < 2) { - return false; - } - - double best_dot = -2.0; - Edge bridge{-1, -1}; - for (int i = 0; i < static_cast(points.size()); ++i) { - for (int j = i + 1; j < static_cast(points.size()); ++j) { - if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] && - Dot(points[i], points[j]) > best_dot) { - best_dot = Dot(points[i], points[j]); - bridge = {i, j}; - } - } - } - if (bridge.first >= 0) { - edges.push_back(bridge); - return true; - } - return false; -} - -GraphCase MakeCase(std::string name, int node_count, double target_degree, - std::uint64_t seed, bool critical_bridge, bool static_rows, - bool repeated_rows) { - const std::vector points = SampleSphere(node_count, seed); - std::vector edges = RipsEdges(points, target_degree); - const int pre_bridge_components = - critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1; - bool bridge_added = false; - if (critical_bridge) { - bridge_added = AddCriticalBridge(points, edges); - } - - const std::vector unique_edges = edges; - if (static_rows) { - for (int node = 0; node < node_count; node += 17) { - edges.push_back({node, node}); - } - } - if (repeated_rows) { - for (std::size_t i = 0; i < unique_edges.size(); i += 11) { - edges.push_back(unique_edges[i]); - edges.push_back({unique_edges[i].second, unique_edges[i].first}); - } - } - DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL); - - std::vector expected = CanonicalPartition(node_count, edges); - const int active_nodes = std::count_if(expected.begin(), expected.end(), - [](int component) { return component >= 0; }); - const int components = CountComponents(expected); - const std::uint64_t checksum = PartitionChecksum(expected); - return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes, - components, pre_bridge_components, bridge_added, checksum}; -} - -const std::vector& Corpus() { - static const std::vector corpus = { - MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false), - MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0), - 0x33960002ULL, true, false, false), - MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)), - 0x33960003ULL, false, false, false), - MakeCase("GroundedStaticRepeated_S2Rips_256", 256, - 2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true), - MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL, - false, false, true), - MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)), - 0x33960006ULL, true, false, false), - }; - return corpus; -} - -struct FloodFillWorkspace { - explicit FloodFillWorkspace(int node_count) - : adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count), - colind(node_count * node_count), stack(node_count * node_count + node_count), - island(node_count) {} - - std::vector adjacency; - std::vector rownnz; - std::vector rowadr; - std::vector colind; - std::vector stack; - std::vector island; -}; - -int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) { - const int n = graph.node_count; - std::fill(work.adjacency.begin(), work.adjacency.end(), 0); - std::fill(work.rownnz.begin(), work.rownnz.end(), 0); - for (const Edge& edge : graph.incidences) { - work.adjacency[edge.first*n + edge.second] = 1; - work.adjacency[edge.second*n + edge.first] = 1; - } - - int address = 0; - for (int row = 0; row < n; ++row) { - work.rowadr[row] = address; - for (int column = 0; column < n; ++column) { - if (work.adjacency[row*n + column]) { - work.colind[address++] = column; - ++work.rownnz[row]; - } - } - } - return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(), - work.colind.data(), work.stack.data()); -} - -struct DsuWorkspace { - explicit DsuWorkspace(int node_count) - : parent(node_count), island(node_count), dof_count(node_count, 1) {} - - std::vector parent; - std::vector island; - std::vector dof_count; -}; - -int RunDsu(const GraphCase& graph, DsuWorkspace& work) { - _mjPRIVATE_dsuInit(work.parent.data(), graph.node_count); - for (const Edge& edge : graph.incidences) { - _mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second); - } - int dof_count = 0; - return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(), - graph.node_count, &dof_count); -} - -bool Validate(const GraphCase& graph) { - FloodFillWorkspace flood(graph.node_count); - DsuWorkspace dsu(graph.node_count); - const int flood_components = RunFloodFill(graph, flood); - const int dsu_components = RunDsu(graph, dsu); - const bool bridge_valid = graph.pre_bridge_components < 0 || - (graph.bridge_added && - graph.pre_bridge_components == graph.expected_components + 1); - return bridge_valid && flood_components == graph.expected_components && - dsu_components == graph.expected_components && - CanonicalizeLabels(flood.island) == graph.expected_partition && - CanonicalizeLabels(dsu.island) == graph.expected_partition && - PartitionChecksum(graph.expected_partition) == graph.expected_checksum; -} - -void BM_FloodFill(benchmark::State& state, const GraphCase* graph) { - if (!Validate(*graph)) { - state.SkipWithError("invalid S2-Rips graph fixture"); - return; - } - FloodFillWorkspace work(graph->node_count); - state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + - " active=" + std::to_string(graph->active_nodes) + - " components=" + std::to_string(graph->expected_components) + - " pre_bridge=" + std::to_string(graph->pre_bridge_components) + - " bridge_added=" + std::to_string(graph->bridge_added) + - " checksum=" + std::to_string(graph->expected_checksum)); - for (auto _ : state) { - int components = RunFloodFill(*graph, work); - benchmark::DoNotOptimize(components); - benchmark::ClobberMemory(); - } - state.SetItemsProcessed(state.iterations() * graph->incidences.size()); -} - -void BM_Dsu(benchmark::State& state, const GraphCase* graph) { - if (!Validate(*graph)) { - state.SkipWithError("invalid S2-Rips graph fixture"); - return; - } - DsuWorkspace work(graph->node_count); - state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + - " active=" + std::to_string(graph->active_nodes) + - " components=" + std::to_string(graph->expected_components) + - " pre_bridge=" + std::to_string(graph->pre_bridge_components) + - " bridge_added=" + std::to_string(graph->bridge_added) + - " checksum=" + std::to_string(graph->expected_checksum)); - for (auto _ : state) { - int components = RunDsu(*graph, work); - benchmark::DoNotOptimize(components); - benchmark::ClobberMemory(); - } - state.SetItemsProcessed(state.iterations() * graph->incidences.size()); -} - -const bool kRegistered = [] { - for (const GraphCase& graph : Corpus()) { - benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph); - benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph); - } - return true; -}(); - -} // namespace -} // namespace mujoco From 8d9f230514e7c5231a34cc40d5e9269380549f70 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 08:59:54 +0530 Subject: [PATCH 07/10] Export disjoint-set island helpers directly Signed-off-by: teerthsharma --- src/engine/engine_island.c | 46 +++++---------- src/engine/engine_island.h | 11 ++-- test/engine/engine_island_test.cc | 96 +++++++++++++++---------------- 3 files changed, 67 insertions(+), 86 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 4f53d2db..ae986cd6 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -84,7 +84,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) { //-------------------------- flood-fill and graph construction ------------------------------------ // find the canonical root of an active tree and compress its path -static inline int dsuFind(int* parent, int tree) { +int mj_dsuRoot(int* parent, int tree) { int root = tree; while (parent[root] != root) { root = parent[root]; @@ -101,13 +101,13 @@ static inline int dsuFind(int* parent, int tree) { // initialize all trees as inactive -static inline void dsuInit(int* parent, int ntree) { +void mj_dsuInit(int* parent, int ntree) { mju_fillInt(parent, -1, ntree); } // activate and union two incident trees; -1 denotes a static endpoint -static inline void dsuUnion(int* parent, int tree1, int tree2) { +void mj_dsuMerge(int* parent, int tree1, int tree2) { if (tree1 == -1 && tree2 == -1) { mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR return; @@ -121,8 +121,8 @@ static inline void dsuUnion(int* parent, int tree1, int tree2) { if (parent[tree1] == parent[tree2]) return; - int root1 = dsuFind(parent, tree1); - int root2 = dsuFind(parent, tree2); + int root1 = mj_dsuRoot(parent, tree1); + int root2 = mj_dsuRoot(parent, tree2); if (root1 < root2) { parent[root2] = root1; } else if (root2 < root1) { @@ -132,8 +132,8 @@ static inline void dsuUnion(int* parent, int tree1, int tree2) { // assign deterministic island ids in ascending canonical-root order -static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, - int* nidof) { +int mj_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, + int* nidof) { int nisland = 0; *nidof = 0; for (int tree=0; tree < ntree; tree++) { @@ -145,7 +145,7 @@ static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, in if (parent[tree] == tree) { island[tree] = nisland++; } else { - // Union always links the larger root to the smaller root. Since trees are visited in + // union always links the larger root to the smaller root. Since trees are visited in // ascending order, this predecessor has already been compressed and assigned an island. parent[tree] = parent[parent[tree]]; island[tree] = island[parent[tree]]; @@ -157,26 +157,6 @@ static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, in } -// exported private wrappers for direct unit tests and benchmarks -int _mjPRIVATE_dsuFind(int* parent, int tree) { - return dsuFind(parent, tree); -} - - -void _mjPRIVATE_dsuInit(int* parent, int ntree) { - dsuInit(parent, ntree); -} - - -void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2) { - dsuUnion(parent, tree1, tree2); -} - - -int _mjPRIVATE_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { - return dsuAssign(island, parent, tree_dofnum, ntree, nidof); -} - // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: // island (nr) - island index assigned to vertex, -1 if vertex has no edges @@ -413,10 +393,10 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { - dsuUnion(parent, tree1, -1); + mj_dsuMerge(parent, tree1, -1); } else { while (tree2 != -2) { - dsuUnion(parent, tree1, tree2); + mj_dsuMerge(parent, tree1, tree2); tree1 = tree2; tree2 = treeNext(m, d, i, &iter); } @@ -458,7 +438,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, if (tree1 < 0) { tree1 = tree2; } else { - dsuUnion(parent, tree1, tree2); + mj_dsuMerge(parent, tree1, tree2); } } } @@ -483,11 +463,11 @@ void mj_island(const mjModel* m, mjData* d) { // union direct tree incidence and assign deterministic components int* efc_tree = mjSTACKALLOC(d, nefc, int); int* parent = mjSTACKALLOC(d, ntree, int); - dsuInit(parent, ntree); + mj_dsuInit(parent, ntree); unionConstraintTrees(m, d, parent, efc_tree); int* tree_island = mjSTACKALLOC(d, ntree, int); int nidof; - d->nisland = dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); + d->nisland = mj_dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); // no islands found: quick return if (!d->nisland) { diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 7e239b40..ca48ba20 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -23,11 +23,12 @@ extern "C" { #endif -MJAPI int _mjPRIVATE_dsuFind(int* parent, int tree); -MJAPI void _mjPRIVATE_dsuInit(int* parent, int ntree); -MJAPI void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2); -MJAPI int _mjPRIVATE_dsuAssign(int* island, int* parent, - const int* tree_dofnum, int ntree, int* nidof); +// disjoint-set roots are minimum tree indices; mj_dsuRoot requires parent[tree] >= 0 +MJAPI void mj_dsuInit(int* parent, int ntree); +MJAPI void mj_dsuMerge(int* parent, int tree1, int tree2); +MJAPI int mj_dsuRoot(int* parent, int tree); +MJAPI int mj_dsuAssign(int* island, int* parent, + const int* tree_dofnum, int ntree, int* nidof); // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index be1f2832..241a053c 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -40,96 +40,96 @@ using IslandTest = MujocoTest; TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) { int parent[] = {8, 6, 7, 5}; - _mjPRIVATE_dsuInit(parent, 0); + mj_dsuInit(parent, 0); EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5)); - _mjPRIVATE_dsuInit(parent, 4); + mj_dsuInit(parent, 4); EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1)); } -TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) { +TEST_F(IslandTest, DsuRootReturnsCanonicalRootAndCompressesPath) { int parent[] = {0, 0, 1, 2, 3}; - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0); - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0); + EXPECT_EQ(mj_dsuRoot(parent, 0), 0); + EXPECT_EQ(mj_dsuRoot(parent, 4), 0); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0)); } -TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) { +TEST_F(IslandTest, DsuMergeActivatesEndpointsAndUsesMinimumRoot) { int parent[] = {-1, -1, -1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, -1, 4); - _mjPRIVATE_dsuUnion(parent, 3, -1); - _mjPRIVATE_dsuUnion(parent, 5, 2); - _mjPRIVATE_dsuUnion(parent, 4, 5); - _mjPRIVATE_dsuUnion(parent, 3, 4); + mj_dsuMerge(parent, -1, 4); + mj_dsuMerge(parent, 3, -1); + mj_dsuMerge(parent, 5, 2); + mj_dsuMerge(parent, 4, 5); + mj_dsuMerge(parent, 3, 4); EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); for (int tree = 2; tree < 6; ++tree) { - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, tree), 2); + EXPECT_EQ(mj_dsuRoot(parent, tree), 2); } EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); } -TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) { +TEST_F(IslandTest, DsuMergeRedundantAndReversedEdgesAreIdempotent) { int parent[] = {-1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, 3, 1); - _mjPRIVATE_dsuUnion(parent, 2, 1); + mj_dsuMerge(parent, 3, 1); + mj_dsuMerge(parent, 2, 1); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); - _mjPRIVATE_dsuUnion(parent, 1, 3); - _mjPRIVATE_dsuUnion(parent, 3, 1); - _mjPRIVATE_dsuUnion(parent, 2, 2); - _mjPRIVATE_dsuUnion(parent, -1, 2); + mj_dsuMerge(parent, 1, 3); + mj_dsuMerge(parent, 3, 1); + mj_dsuMerge(parent, 2, 2); + mj_dsuMerge(parent, -1, 2); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); } -TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) { +TEST_F(IslandTest, DsuMergeFastPathActivatesBeforeTestingParents) { int self_parent[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(self_parent, 1, 1); + mj_dsuMerge(self_parent, 1, 1); EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1)); int static_first[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(static_first, -1, 2); + mj_dsuMerge(static_first, -1, 2); EXPECT_THAT(static_first, ElementsAre(-1, -1, 2)); int static_second[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(static_second, 0, -1); + mj_dsuMerge(static_second, 0, -1); EXPECT_THAT(static_second, ElementsAre(0, -1, -1)); } -TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) { +TEST_F(IslandTest, DsuMergeFastPathDistinguishesParentsFromRoots) { int distinct_parent[] = {0, 0, 2, 2}; - _mjPRIVATE_dsuUnion(distinct_parent, 1, 3); + mj_dsuMerge(distinct_parent, 1, 3); EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2)); int shared_parent[] = {0, 0, 0, 3}; - _mjPRIVATE_dsuUnion(shared_parent, 1, 2); + mj_dsuMerge(shared_parent, 1, 2); EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3)); int long_paths[] = {0, 0, 1, 3, 3, 4}; - _mjPRIVATE_dsuUnion(long_paths, 2, 5); + mj_dsuMerge(long_paths, 2, 5); EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3)); } -TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) { +TEST_F(IslandTest, DsuMergeFastPathPreservesCyclesDuplicatesAndForest) { int parent[] = {-1, -1, -1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, 0, 1); - _mjPRIVATE_dsuUnion(parent, 1, 2); - _mjPRIVATE_dsuUnion(parent, 2, 0); - _mjPRIVATE_dsuUnion(parent, 0, 2); - _mjPRIVATE_dsuUnion(parent, 3, 4); - _mjPRIVATE_dsuUnion(parent, 4, 5); + mj_dsuMerge(parent, 0, 1); + mj_dsuMerge(parent, 1, 2); + mj_dsuMerge(parent, 2, 0); + mj_dsuMerge(parent, 0, 2); + mj_dsuMerge(parent, 3, 4); + mj_dsuMerge(parent, 4, 5); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3)); - _mjPRIVATE_dsuUnion(parent, 5, 0); + mj_dsuMerge(parent, 5, 0); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3)); } -TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) { +TEST_F(IslandTest, DsuMergeRejectsStaticSelfIncidence) { int parent[] = {-1, 1, 1, 3}; - EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1), + EXPECT_EQ(MjuErrorMessageFrom(mj_dsuMerge)(parent, -1, -1), "self-incidence of the static tree"); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3)); } @@ -140,13 +140,13 @@ TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) { const int tree_dofnum[] = {73}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0); EXPECT_EQ(nidof, 0); EXPECT_THAT(island, ElementsAre(71)); EXPECT_THAT(parent, ElementsAre(72)); parent[0] = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0); EXPECT_EQ(nidof, 0); EXPECT_THAT(island, ElementsAre(-1)); EXPECT_THAT(parent, ElementsAre(-1)); @@ -158,7 +158,7 @@ TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) { int island[] = {9, 9, 9, 9, 9, 9, 9}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3); EXPECT_EQ(nidof, 39); EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6)); @@ -170,7 +170,7 @@ TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) { int island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3); EXPECT_EQ(nidof, 41); EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 4, 7, 7, 7)); @@ -190,7 +190,7 @@ TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) { } int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + EXPECT_EQ(mj_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), kTreeCount, &nidof), 1); EXPECT_EQ(nidof, expected_nidof); @@ -205,11 +205,11 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { std::vector parent(kTreeCount); std::vector island(kTreeCount, -2); std::vector tree_dofnum(kTreeCount); - _mjPRIVATE_dsuInit(parent.data(), kTreeCount); + mj_dsuInit(parent.data(), kTreeCount); int expected_nidof = 0; for (int tree = kTreeCount - 1; tree > 0; --tree) { - _mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1); + mj_dsuMerge(parent.data(), tree, tree - 1); } for (int tree = 0; tree < kTreeCount; ++tree) { tree_dofnum[tree] = tree % 7; @@ -217,7 +217,7 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { } int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + EXPECT_EQ(mj_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), kTreeCount, &nidof), 1); EXPECT_EQ(nidof, expected_nidof); @@ -282,9 +282,9 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { } std::vector parent(ntree); - _mjPRIVATE_dsuInit(parent.data(), ntree); + mj_dsuInit(parent.data(), ntree); for (const auto& edge : edges) { - _mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]); + mj_dsuMerge(parent.data(), edge[0], edge[1]); } std::vector active(ntree); @@ -329,7 +329,7 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { } std::vector island(ntree, -2); int nidof = -1; - const int nisland = _mjPRIVATE_dsuAssign( + const int nisland = mj_dsuAssign( island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof); SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial From cdda84719154836ec91776c236fbc37d01a792ad Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 09:12:58 +0530 Subject: [PATCH 08/10] Restore static-constraint island diagnostic Signed-off-by: teerthsharma --- src/engine/engine_island.c | 6 +++++- test/engine/engine_island_test.cc | 26 ++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index ae986cd6..42444f42 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -389,7 +389,11 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, int tree1 = treeNext(m, d, i, &iter); if (tree1 != -2) { int tree2 = treeNext(m, d, i, &iter); - efc_tree[i] = tree1 == -1 ? tree2 : tree1; + // assign tree to constraint, one of (tree1, tree2) must be non-negative + efc_tree[i] = tree1 >= 0 ? tree1 : tree2; + if (efc_tree[i] < 0) { + mjERROR("constraint %d is between two static bodies", i); // SHOULD NOT OCCUR + } // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 241a053c..6f04533f 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -531,6 +531,32 @@ TEST_F(IslandTest, ProductionStaticFirstAndRepeatedRows) { EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1, 2, 3, 4, 5)); } +TEST_F(IslandTest, ReportsConstraintBetweenTwoStaticBodies) { + static constexpr char xml[] = R"( + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + ASSERT_GT(data->nefc, 0); + ASSERT_EQ(data->efc_type[0], mjCNSTR_FRICTION_DOF); + + model->dof_treeid[data->efc_id[0]] = -1; + + EXPECT_EQ(MjuErrorMessageFrom(mj_island)(model.get(), data.get()), + "constraint 0 is between two static bodies"); +} + TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { static constexpr char xml[] = R"( From 7ae4f22693bb15ff4ad44276af8ab8d462eef60d Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 09:19:33 +0530 Subject: [PATCH 09/10] Restore flex island invariants in comments Signed-off-by: teerthsharma --- src/engine/engine_island.c | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 42444f42..60a87c87 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -410,9 +410,15 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, } } - // Flex stiffness couples all vertices (nodes for interpolated flexes) without a constraint - // row representing the coupling. Union the awake dynamic trees of each stiffness-active flex. + // flex stiffness couples all vertices (nodes for interpolated flexes) of a flex without any + // constraint row representing the coupling: union the trees of every stiffness-active flex + // (star around the first dynamic tree). This keeps the partition valid when the implicit + // effective metric (mj_flexCG) carries the stiffness inside the constraint solve. Awake + // trees only: sleeping trees must stay out of islands (mj_sleep invariant, matching the + // constraint filter); waking a flex as a unit remains the wake machinery's job. for (int f=0; f < m->nflex; f++) { + // mirror the stiffness-activity conditions of engine_derivative's flexStiff_active / + // flexInterp_processed: deformable dim>=2 flex with bending or nonzero stiffness if (m->flex_rigid[f] || m->flex_dim[f] < 2) { continue; } From 57faf5f63a13464a7b5af3726c88fdce0a39371f Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 09:25:27 +0530 Subject: [PATCH 10/10] Credit linear-memory island discovery Signed-off-by: teerthsharma --- doc/changelog.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index a8d3af11..dfd18c56 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -48,6 +48,7 @@ Engine as originally guarded by ``ngravcomp``. Since the engine uses these integers as flags (zero vs. non-zero), the new flags are honest boolean properties, writeable from the Python bindings at runtime. The field ``ngravcomp`` is deprecated and will be removed in a future release. +- Replaced quadratic scratch in native island discovery with a linear-memory disjoint set. Contribution by :github:user:`teerthsharma`. .. admonition:: Breaking API changes :class: attention