Speed up sparse supernode detection by combining it with transposition.

PiperOrigin-RevId: 760673008
Change-Id: I6d66580e675fd86e5b974859383f93f87482ca16

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -152,8 +152,8 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
 // transpose sparse matrix (uncompressed)
 void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
     mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz,
-    int* res_rowadr, int* res_colind, const int* rownnz, const int* rowadr,
-    const int* colind) {
+    int* res_rowadr, int* res_colind, int* res_rowsuper,
+    const int* rownnz, const int* rowadr, const int* colind) {
   memset(res_rownnz, 0, nc * sizeof(int));
   for (int rt = 0; rt < nc; rt++) {
     res_rowadr[rt] = rt * nr;
@@ -497,7 +497,14 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
 }
 BENCHMARK(BM_combineSparse_old);
 
-static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
+enum class Supernode {
+  None,
+  PostProcess,
+  Inline
+};
+
+static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
+                               Supernode super) {
   static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
 
   // force use of sparse matrices
@@ -516,12 +523,19 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
   mjtNum* res = mj_stackAllocNum(d, m->nv * d->nefc);
   int* res_rownnz = mj_stackAllocInt(d, m->nv);
   int* res_rowadr = mj_stackAllocInt(d, m->nv);
+  int* res_rowsuper = mj_stackAllocInt(d, m->nv);
   int* res_colind = mj_stackAllocInt(d, m->nv * d->nefc);
 
   // time benchmark
   for (auto s : state) {
-    func(res, d->efc_J, d->nefc, m->nv, res_rownnz, res_rowadr, res_colind,
+    int* rowsuper = (super == Supernode::Inline) ? res_rowsuper : nullptr;
+    func(res, d->efc_J, d->nefc, m->nv,
+         res_rownnz, res_rowadr, res_colind, rowsuper,
          d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
+    if (super == Supernode::PostProcess) {
+      mju_superSparse(m->nv, res_rowsuper,
+                      res_rownnz, res_rowadr, res_colind);
+    }
   }
 
   mj_freeStack(d);
@@ -529,19 +543,33 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
   state.SetItemsProcessed(state.iterations());
 }
 
+void ABSL_ATTRIBUTE_NO_TAIL_CALL
+BM_transposeSparse_old(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_transposeSparse(state, &transposeSparse_baseline, Supernode::None);
+}
+BENCHMARK(BM_transposeSparse_old);
+
 void ABSL_ATTRIBUTE_NO_TAIL_CALL
 BM_transposeSparse_new(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  BM_transposeSparse(state, &mju_transposeSparse);
+  BM_transposeSparse(state, &mju_transposeSparse, Supernode::None);
 }
 BENCHMARK(BM_transposeSparse_new);
 
 void ABSL_ATTRIBUTE_NO_TAIL_CALL
-BM_transposeSparse_old(benchmark::State& state) {
+BM_transposeSparse_superpost(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  BM_transposeSparse(state, &transposeSparse_baseline);
+  BM_transposeSparse(state, &mju_transposeSparse, Supernode::PostProcess);
 }
-BENCHMARK(BM_transposeSparse_old);
+BENCHMARK(BM_transposeSparse_superpost);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL
+BM_transposeSparse_superinline(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_transposeSparse(state, &mju_transposeSparse, Supernode::Inline);
+}
+BENCHMARK(BM_transposeSparse_superinline);
 
 static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
   static mjModel* m =

test/engine/engine_util_sparse_test.cc

@@ -173,8 +173,8 @@ TEST_F(EngineUtilSparseTest, MjuTranspose3by3) {
   int rownnzT[] = {0, 0, 0};
   int rowadrT[] = {0, 0, 0};
 
-  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, nullptr,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(matT, ElementsAre(1, 2, 1, 3));
   EXPECT_THAT(colindT, ElementsAre(0, 0, 1, 2));
@@ -197,8 +197,8 @@ TEST_F(EngineUtilSparseTest, MjuTranspose1by3) {
   int rownnzT[] = {0, 0, 0};
   int rowadrT[] = {0, 0, 0};
 
-  mju_transposeSparse(matT, mat, 1, 3, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 1, 3, rownnzT, rowadrT, colindT, nullptr,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(matT, ElementsAre(1, 3));
   EXPECT_THAT(colindT, ElementsAre(0, 0));
@@ -221,8 +221,8 @@ TEST_F(EngineUtilSparseTest, MjuTranspose3by1) {
   int rownnzT[] = {0};
   int rowadrT[] = {0};
 
-  mju_transposeSparse(matT, mat, 3, 1, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 3, 1, rownnzT, rowadrT, colindT, nullptr,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(matT, ElementsAre(1, 3));
   EXPECT_THAT(colindT, ElementsAre(0, 2));
@@ -244,14 +244,43 @@ TEST_F(EngineUtilSparseTest, MjuTransposeDense) {
   int colindT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
   int rownnzT[] = {0, 0, 0};
   int rowadrT[] = {0, 0, 0};
+  int rowsuperT[] = {0, 0, 0};
 
-  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rowsuperT,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(matT, ElementsAre(1, 4, 7, 2, 5, 8, 3, 6, 9));
   EXPECT_THAT(colindT, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
   EXPECT_THAT(rownnzT, ElementsAre(3, 3, 3));
   EXPECT_THAT(rowadrT, ElementsAre(0, 3, 6));
+  EXPECT_THAT(rowsuperT, ElementsAre(2, 1, 0));
+}
+
+TEST_F(EngineUtilSparseTest, MjuTransposeSuper) {
+  //  mat:   0,  1,  2,  0,  3,  4,  5,  0,  0,  0,  0,  0,  0
+  //         0,  0,  0,  6,  7,  8,  9, 10, 11, 12,  0,  0,  0
+  //
+  //  super: 0,  1,  0,  0,  2,  1,  0,  2,  1,  0,  2,  1,  0
+
+  mjtNum mat[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+  int colind[] = {1, 2, 4, 5, 6, 3, 4, 5, 6, 7,  8,  9};
+  int rownnz[] = {5, 7};
+  int rowadr[] = {0, 5};
+
+  mjtNum matT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  int colindT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  int rownnzT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  int rowadrT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  int rowsuperT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+
+  mju_transposeSparse(matT, mat, 2, 13, rownnzT, rowadrT, colindT, rowsuperT,
+                      rownnz, rowadr, colind);
+
+  EXPECT_THAT(matT,    ElementsAre(1, 2, 6, 3, 7, 4, 8, 5, 9, 10, 11, 12));
+  EXPECT_THAT(colindT, ElementsAre(0, 0, 1, 0, 1, 0, 1, 0, 1, 1,  1,  1));
+  EXPECT_THAT(rownnzT, ElementsAre(0, 1, 1, 1, 2, 2, 2, 1, 1,  1,  0,  0,  0));
+  EXPECT_THAT(rowadrT, ElementsAre(0, 0, 1, 2, 3, 5, 7, 9, 10, 11, 12, 12, 12));
+  EXPECT_THAT(rowsuperT, ElementsAre(0, 1, 0, 0, 2, 1, 0, 2, 1, 0, 2,  1,  0));
 }
 
 TEST_F(EngineUtilSparseTest, MjuTranspose1by1) {
@@ -267,8 +296,8 @@ TEST_F(EngineUtilSparseTest, MjuTranspose1by1) {
   int rownnzT[] = {0};
   int rowadrT[] = {0};
 
-  mju_transposeSparse(matT, mat, 1, 1, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 1, 1, rownnzT, rowadrT, colindT, nullptr,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(matT, ElementsAre(1));
   EXPECT_THAT(colindT, ElementsAre(0));
@@ -289,8 +318,8 @@ TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) {
   int rownnzT[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
   int rowadrT[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
 
-  mju_transposeSparse(matT, mat, 10, 10, rownnzT, rowadrT, colindT, rownnz,
-                      rowadr, colind);
+  mju_transposeSparse(matT, mat, 10, 10, rownnzT, rowadrT, colindT, nullptr,
+                      rownnz, rowadr, colind);
 
   EXPECT_THAT(rownnzT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0));
   EXPECT_THAT(rowadrT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0));