chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Speed up sparse supernode detection by combining it with transposition.

Browse Source 
PiperOrigin-RevId: 760673008
Change-Id: I6d66580e675fd86e5b974859383f93f87482ca16
This commit is contained in:
Yuval Tassa
2025-05-19 10:13:20 -07:00
committed by Copybara-Service
parent ced630181d
commit 45fc15b844
6 changed files with 115 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+2 -6
View File
@@ -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
+1 -1
View File
@@ -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
+33 -7
View File
@@ -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];
}
}
}
}
src/engine/engine_util_sparse.h
+2 -2
View File
@@ -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
+36 -8
View File
@@ -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
+41 -12
View File
@@ -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));