Improve island discovery performance

PiperOrigin-RevId: 557444755
Change-Id: I22a2093b563e8e144de9534917955ac403fa5599

src/engine/engine_island.c

@@ -32,91 +32,6 @@
   #include <sanitizer/msan_interface.h>
 #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
+// add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge
-static int findEdges(const mjModel* m, const mjData* d, int* edge, int nedge_max) {
+//   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
+  // get tree-tree edges and rownnz counts from efc arrays
-  int nedge = findEdges(m, d, edge, nedge_max);
+  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
+  // compute starting address of tree's column indices while resetting rownnz
-
-  // construct adjacency matrix from edges
-  int* rownnz = mj_stackAllocInt(d, ntree);
   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

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);

test/engine/engine_island_test.cc

@@ -29,69 +29,12 @@ namespace mujoco {
 namespace {
 
 using ::testing::ElementsAre;
-using ::testing::ElementsAreArray;
 using IslandTest = MujocoTest;
 
 std::vector<int> AsVector(const int* array, int n) {
   return std::vector<int>(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