Speed up sparse supernode detection by combining it with transposition.

PiperOrigin-RevId: 760673008
Change-Id: I6d66580e675fd86e5b974859383f93f87482ca16

src/engine/engine_core_constraint.c

@@ -1996,7 +1996,7 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
   if (mj_isSparse(m)) {
     // transpose
     mju_transposeSparse(d->efc_JT, d->efc_J, d->nefc, m->nv,
-                        d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
+                        d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
                         d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
 
 
@@ -2010,10 +2010,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
     __msan_allocated_memory(d->efc_J_rowsuper, d->nefc);
   #endif  // MEMORY_SANITIZER
 #endif  // mjUSEAVX
-
-    // supernodes of JT
-    mju_superSparse(m->nv, d->efc_JT_rowsuper,
-                    d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind);
   }
 
   // compute diagApprox
@@ -2182,7 +2178,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     int* BT_colind = mjSTACKALLOC(d, nB, int);
     mjtNum* BT = mjSTACKALLOC(d, nB, mjtNum);
     mju_transposeSparse(BT, B, nefc, nv,
-                        BT_rownnz, BT_rowadr, BT_colind,
+                        BT_rownnz, BT_rowadr, BT_colind, NULL,
                         B_rownnz, B_rowadr, B_colind);
 
     // allocate AR row nonzeros and addresses on arena

src/engine/engine_solver.c

@@ -1572,7 +1572,7 @@ static void MakeHessian(mjData* d, mjCGContext* ctx) {
     int* HT_rowadr = mjSTACKALLOC(d, nv, int);
     int* HT_colind = mjSTACKALLOC(d, ctx->nH, int);
     mju_transposeSparse(NULL, NULL, nv, nv,
-                        HT_rownnz, HT_rowadr, HT_colind,
+                        HT_rownnz, HT_rowadr, HT_colind, NULL,
                         ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind);
 
     // count total and row non-zeros of reverse-Cholesky factor L

src/engine/engine_util_sparse.c

@@ -531,9 +531,9 @@ int mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, in
 
 
 
-// transpose sparse matrix
+// transpose sparse matrix, optionally compute row supernodes
 void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
-                         int* res_rownnz, int* res_rowadr, int* res_colind,
+                         int* res_rownnz, int* res_rowadr, int* res_colind, int* res_rowsuper,
                          const int* rownnz, const int* rowadr, const int* colind) {
   // clear number of non-zeros for each row of transposed
   mju_zeroInt(res_rownnz, nc);
@@ -547,22 +547,40 @@ void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
     }
   }
 
+  // init res_rowsuper
+  if (res_rowsuper) {
+    for (int i = 0; i < nc - 1; i++) {
+      res_rowsuper[i] = (res_rownnz[i] == res_rownnz[i + 1]);
+    }
+    res_rowsuper[nc - 1] = 0;
+  }
+
   // compute the row addresses for the transposed matrix
   res_rowadr[0] = 0;
   for (int i = 1; i < nc; i++) {
     res_rowadr[i] = res_rowadr[i-1] + res_rownnz[i-1];
   }
 
-  // iterate through each non-zero entry of mat
+  // iterate through each row (column) of mat (res)
   for (int r = 0; r < nr; r++) {
+    int c_prev = -1;
     int start = rowadr[r];
     int end = start + rownnz[r];
     for (int i = start; i < end; i++) {
       // swap rows with columns and increment res_rowadr
-      int c = res_rowadr[colind[i]]++;
+      int c = colind[i];
-      res_colind[c] = r;
+      int adr = res_rowadr[c]++;
+      res_colind[adr] = r;
       if (res) {
-        res[c] = mat[i];
+        res[adr] = mat[i];
+      }
+
+      // mark non-supernodes
+      if (res_rowsuper) {
+        if (c > 0 && c != c_prev + 1 && res_rowsuper[c - 1]) {
+          res_rowsuper[c - 1] = 0;
+        }
+        c_prev = c;
       }
     }
   }
@@ -571,8 +589,16 @@ void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
   for (int i = nc-1; i > 0; i--) {
     res_rowadr[i] = res_rowadr[i-1];
   }
-
   res_rowadr[0] = 0;
+
+  // accumulate supernodes
+  if (res_rowsuper) {
+    for (int i = nc - 2; i >= 0; i--) {
+      if (res_rowsuper[i]) {
+        res_rowsuper[i] += res_rowsuper[i + 1];
+      }
+    }
+  }
 }
 
 

src/engine/engine_util_sparse.h

@@ -91,9 +91,9 @@ int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum s
 int mju_addChains(int* res, int n, int NV1, int NV2,
                   const int* chain1, const int* chain2);
 
-// transpose sparse matrix
+// transpose sparse matrix, optionally compute row supernodes
 MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
-                               int* res_rownnz, int* res_rowadr, int* res_colind,
+                               int* res_rownnz, int* res_rowadr, int* res_colind, int* res_rowsuper,
                                const int* rownnz, const int* rowadr, const int* colind);
 
 // construct row supernodes

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,
+    int* res_rowadr, int* res_colind, int* res_rowsuper,
-    const int* colind) {
+    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,
+  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, nullptr,
-                      rowadr, colind);
+                      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,
+  mju_transposeSparse(matT, mat, 1, 3, rownnzT, rowadrT, colindT, nullptr,
-                      rowadr, colind);
+                      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,
+  mju_transposeSparse(matT, mat, 3, 1, rownnzT, rowadrT, colindT, nullptr,
-                      rowadr, colind);
+                      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,
+  mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rowsuperT,
-                      rowadr, colind);
+                      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,
+  mju_transposeSparse(matT, mat, 1, 1, rownnzT, rowadrT, colindT, nullptr,
-                      rowadr, colind);
+                      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,
+  mju_transposeSparse(matT, mat, 10, 10, rownnzT, rowadrT, colindT, nullptr,
-                      rowadr, colind);
+                      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));