From 6245edae2836933bd50a6aada14ad2b76ef35c03 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 16 Aug 2023 04:54:53 -0700 Subject: [PATCH] Improve island discovery performance PiperOrigin-RevId: 557444755 Change-Id: I22a2093b563e8e144de9534917955ac403fa5599 --- src/engine/engine_island.c | 267 ++++++++++++------------------ src/engine/engine_island.h | 5 - test/engine/engine_island_test.cc | 57 ------- 3 files changed, 110 insertions(+), 219 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index dccf809c..ca78a119 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -32,91 +32,6 @@ #include #endif -// clear island-related arena pointers in mjData -static void clearIsland(mjData* d, size_t parena) { -#define X(type, name, nr, nc) d->name = NULL; - MJDATA_ARENA_POINTERS_ISLAND -#undef X - d->nefc = 0; - d->nisland = 0; - d->parena = parena; - - // poison remaining memory -#ifdef ADDRESS_SANITIZER - ASAN_POISON_MEMORY_REGION( - (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena); -#endif -} - -// comparison function for lexicographic edge sorting -quicksortfunc(edgecompare, context, edge0, edge1) { - int* e0 = (int*)edge0; - int* e1 = (int*)edge1; - int v00 = e0[0]; - int v10 = e1[0]; - - if (v00 < v10) { - return -1; - } - - if (v00 == v10) { - int v01 = e0[1]; - int v11 = e1[1]; - - if (v01 < v11) { - return -1; - } - - if (v01 == v11) { - return 0; - } - } - - return 1; -} - - - -// construct sparse matrix from non-unique, unsorted edge array, return number of nonzeros -int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr) { - if (!ne) { - return 0; - } - - // sort edges - mjQUICKSORT(edge, ne, 2*sizeof(int), edgecompare, NULL); - - // construct sparse - int nnz = 0; // number of nonzeros - int e = 0; // current edge - for (int r=0; r < nr; r++) { - // init row - rownnz[r] = 0; - rowadr[r] = nnz; - - // copy values while making unique and checking indices - while (e < ne && edge[2*e] == r) { - int v0 = edge[2*e]; - int v1 = edge[2*e + 1]; - - // skip if duplicate - if (rownnz[r] && v0 == edge[2*e - 2] && v1 == edge[2*e - 1]) { - e++; - continue; - } - - // check for invalid indices - if (v0 < 0 || v0 >= nr) mjERROR("invalid row index %d in edge %d", v0, e); - if (v1 < 0 || v1 >= nr) mjERROR("invalid column index %d in edge %d", v1, e); - - // copy column index, increment nnz, e, rownnz - colind[nnz++] = edge[2*(e++) + 1]; - rownnz[r]++; - } - } - - return nnz; -} // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix @@ -128,6 +43,7 @@ int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int // colind (nnz) - matrix column indices // stack (nnz) - stack space // returns number of islands +// note: column indices are not required to be unique or sorted int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, int* stack) { // initialize island count, set ids to -1 @@ -172,6 +88,24 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons +// clear island-related arena pointers in mjData +static void clearIsland(mjData* d, size_t parena) { +#define X(type, name, nr, nc) d->name = NULL; + MJDATA_ARENA_POINTERS_ISLAND +#undef X + d->nefc = 0; + d->nisland = 0; + d->parena = parena; + + // poison remaining memory +#ifdef ADDRESS_SANITIZER + ASAN_POISON_MEMORY_REGION( + (char*)d->arena + d->parena, (d->nstack - d->pstack) * sizeof(mjtNum) - d->parena); +#endif +} + + + // return upper bound on number of tree-tree edges static int countMaxEdge(const mjModel* m, const mjData* d) { int nedge_max = 0; @@ -194,61 +128,6 @@ static int countMaxEdge(const mjModel* m, const mjData* d) { -// add tree-tree edge array: check size, add flipped edge if non-self -static int addEdge(int* edge, int nedge, int tree1, int tree2, int nedge_max) { - // handle the static tree - if (tree1 == -1 && tree2 == -1) { - mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR - return 0; - } - if (tree1 == -1) tree1 = tree2; - if (tree2 == -1) tree2 = tree1; - - // previous edge - int p1 = nedge ? edge[2*nedge - 2] : -1; - int p2 = nedge ? edge[2*nedge - 1] : -1; - - // === self edge - if (tree1 == tree2) { - // same as previous edge, return - if (nedge && tree1 == p1 && tree1 == p2) { - return nedge; - } - - // check size - if (nedge >= nedge_max) { - mjERROR("edge array too small"); - return 0; - } - - // add tree1-tree1 self-edge - edge[2*nedge + 0] = tree1; - edge[2*nedge + 1] = tree1; - return nedge + 1; - } - - // === non-self edge - if (nedge && ((tree1 == p1 && tree2 == p2) || (tree1 == p2 && tree2 == p1))) { - // same as previous edge, return - return nedge; - } - - // check size - if (nedge + 2 > nedge_max) { - mjERROR("edge array too small"); - return 0; - } - - // add tree1-tree2 and tree2-tree1 - edge[2*nedge + 0] = tree1; - edge[2*nedge + 1] = tree2; - edge[2*nedge + 2] = tree2; - edge[2*nedge + 3] = tree1; - return nedge + 2; -} - - - // return id of next tree in Jacobian row i that is different from tree, -1 if not found // write the index of the found tree to *index if given // start search from *index if given, otherwise 0 @@ -299,13 +178,75 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind -// find tree-tree edges -static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max) { +// add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge +// return current number of edges +static int addEdge(int* treenedge, int* edge, int nedge, int tree1, int tree2, int nedge_max) { + // handle the static tree + if (tree1 == -1 && tree2 == -1) { + mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR + return 0; + } + if (tree1 == -1) tree1 = tree2; + if (tree2 == -1) tree2 = tree1; + + // previous edge + int p1 = nedge ? edge[2*nedge - 2] : -1; + int p2 = nedge ? edge[2*nedge - 1] : -1; + + // === self edge + if (tree1 == tree2) { + // same as previous edge, return + if (nedge && tree1 == p1 && tree1 == p2) { + return nedge; + } + + // check size + if (nedge >= nedge_max) { + mjERROR("edge array too small"); + return 0; + } + + // add tree1-tree1 self-edge + edge[2*nedge + 0] = tree1; + edge[2*nedge + 1] = tree1; + treenedge[tree1]++; + return nedge + 1; + } + + // === non-self edge + if (nedge && ((tree1 == p1 && tree2 == p2) || (tree1 == p2 && tree2 == p1))) { + // same as previous edge, return + return nedge; + } + + // check size + if (nedge + 2 > nedge_max) { + mjERROR("edge array too small"); + return 0; + } + + // add tree1-tree2 and tree2-tree1 + edge[2*nedge + 0] = tree1; + edge[2*nedge + 1] = tree2; + edge[2*nedge + 2] = tree2; + edge[2*nedge + 3] = tree1; + treenedge[tree1]++; + treenedge[tree2]++; + return nedge + 2; +} + + + +// find tree-tree edges, increment treenedge counters, return total number of edges +static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edge, int nedge_max) { int nefc = d->nefc; int efc_type = -1; int efc_id = -1; int tree1, tree2; + // clear treenedge + memset(treenedge, 0, m->ntree*sizeof(int)); + int nedge = 0; for (int i=0; i < nefc; i++) { // row i is still in the same constraint: skip @@ -320,14 +261,14 @@ static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max // joint friction if (efc_type == mjCNSTR_FRICTION_DOF) { tree1 = m->dof_treeid[efc_id]; - nedge = addEdge(edge, nedge, tree1, tree1, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); continue; } // joint limit if (efc_type == mjCNSTR_LIMIT_JOINT) { tree1 = m->dof_treeid[m->jnt_dofadr[efc_id]]; - nedge = addEdge(edge, nedge, tree1, tree1, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); continue; } @@ -337,7 +278,7 @@ static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max efc_type == mjCNSTR_CONTACT_ELLIPTIC) { tree1 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]]; tree2 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]]; - nedge = addEdge(edge, nedge, tree1, tree2, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); continue; } @@ -347,7 +288,7 @@ static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) { tree1 = m->body_treeid[m->eq_obj1id[efc_id]]; tree2 = m->body_treeid[m->eq_obj2id[efc_id]]; - nedge = addEdge(edge, nedge, tree1, tree2, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); continue; } } @@ -359,15 +300,15 @@ static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max if (tree2 == -1) { // 1 tree found: add self-edge - nedge = addEdge(edge, nedge, tree1, tree1, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); } else { // 2 trees found: add edge, keep scanning and adding until no more trees - nedge = addEdge(edge, nedge, tree1, tree2, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); int tree3 = treeNext(m, d, tree2, i, &index); while (tree3 > -1 && tree3 != tree2) { tree1 = tree2; tree2 = tree3; - nedge = addEdge(edge, nedge, tree1, tree2, nedge_max); + nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); tree3 = treeNext(m, d, tree2, i, &index); } } @@ -376,6 +317,8 @@ static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max return nedge; } + + // discover islands: // nisland, island_dofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext void mj_island(const mjModel* m, mjData* d) { @@ -393,20 +336,30 @@ void mj_island(const mjModel* m, mjData* d) { int nedge_max = countMaxEdge(m, d); int* edge = mj_stackAllocInt(d, 2*nedge_max); - // find tree-tree edges - int nedge = findEdges(m, d, edge, nedge_max); + // get tree-tree edges and rownnz counts from efc arrays + int* rownnz = mj_stackAllocInt(d, ntree); // number of edges per tree + int nedge = findEdges(m, d, rownnz, edge, nedge_max); - // TODO: b/295296178 - don't add flipped edges in findEdges, symmetrize in mj_edge2sparse instead - - // construct adjacency matrix from edges - int* rownnz = mj_stackAllocInt(d, ntree); + // compute starting address of tree's column indices while resetting rownnz int* rowadr = mj_stackAllocInt(d, ntree); + rowadr[0] = 0; + for (int r=1; r < ntree; r++) { + rowadr[r] = rowadr[r-1] + rownnz[r-1]; + rownnz[r-1] = 0; + } + rownnz[ntree-1] = 0; + + // copy column indices: list each tree's neighbors int* colind = mj_stackAllocInt(d, nedge); - int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, nedge, ntree); + for (int e=0; e < nedge; e++) { + int row = edge[2*e]; + int col = edge[2*e + 1]; + colind[rowadr[row] + rownnz[row]++] = col; + } // discover islands int* tree_island = mj_stackAllocInt(d, ntree); // id of island assigned to tree - int* stack = mj_stackAllocInt(d, nnz); + int* stack = mj_stackAllocInt(d, nedge); d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack); // ========== begin arena allocation of MJDATA_ARENA_POINTERS_ISLAND diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 20c0b303..a724ac33 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -24,11 +24,6 @@ extern "C" { #endif -//-------------------------- utility functions ----------------------------------------------------- - -// construct sparse matrix from non-unique, unsorted edge array, return number of nonzeros -MJAPI int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr); - // 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, int* scratch); diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 72dacece..20d235e7 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -29,69 +29,12 @@ namespace mujoco { namespace { using ::testing::ElementsAre; -using ::testing::ElementsAreArray; using IslandTest = MujocoTest; std::vector AsVector(const int* array, int n) { return std::vector(array, array + n); } -TEST_F(IslandTest, EdgeToSparse2) { - // unsorted edges, with duplication - constexpr int ne = 4; - constexpr int nr = 5; - int edge[2*ne] = { - 3, 4, - 1, 1, - 3, 4, - 1, 1 - }; - - int rownnz[nr]; - int rowadr[nr]; - int colind[ne]; - - int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr); - - constexpr int expected_nnz = 2; - - EXPECT_EQ(nnz, expected_nnz); - - EXPECT_THAT(rownnz, ElementsAre(0, 1, 0, 1, 0)); - EXPECT_THAT(rowadr, ElementsAre(0, 0, 1, 1, 2)); - - int expected_colind[expected_nnz] = {1, 4}; - EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz)); -} - -TEST_F(IslandTest, EdgeToSparse3) { - // unsorted edges, with duplication - constexpr int ne = 3; - constexpr int nr = 1; - int edge[2*ne] = { - 0, 0, - 0, 0, - 0, 0 - }; - - int rownnz[nr]; - int rowadr[nr]; - int colind[ne]; - - int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr); - - constexpr int expected_nnz = 1; - - EXPECT_EQ(nnz, expected_nnz); - - EXPECT_THAT(rownnz, ElementsAre(1)); - EXPECT_THAT(rowadr, ElementsAre(0)); - - int expected_colind[expected_nnz] = {0}; - EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz)); -} - - TEST_F(IslandTest, FloodFillSingleton) { // adjacency matrix for the graph 0 1 2 // U U