diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 1a60b6e6..3e625f49 100644 --- a/src/engine/engine_core_constraint.c +++ b/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 diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 37b83024..635fa6af 100644 --- a/src/engine/engine_solver.c +++ b/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 diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 44dddb92..c9cf3bab 100644 --- a/src/engine/engine_util_sparse.c +++ b/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]]++; - res_colind[c] = r; + int c = colind[i]; + 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]; + } + } + } } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 33f65513..f6c2f7c2 100644 --- a/src/engine/engine_util_sparse.h +++ b/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 diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index 6a1b3835..82a04a3f 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/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 = diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index edbd2558..36cbda57 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/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));