From 52ddcbc81a5e4686157e924dd58da810a1b35aa0 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Sat, 11 Jul 2026 09:48:54 +0530 Subject: [PATCH] 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) {