From c741dfce7d72c414ad099c77c4671b3d8bf12271 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Fri, 27 Jan 2023 00:19:14 -0800 Subject: [PATCH] Implement a more performant mju_transposeSparse that doesn't require dense memory allocation. PiperOrigin-RevId: 505044194 Change-Id: Ibe0e39e3ad711b2b5bdad45180ca4f83b5c8fc82 --- doc/changelog.rst | 4 + src/engine/engine_util_sparse.c | 55 ++++--- src/engine/engine_util_sparse.h | 10 +- .../engine_util_sparse_benchmark_test.cc | 76 +++++++++- test/engine/engine_util_sparse_test.cc | 140 ++++++++++++++++++ 5 files changed, 255 insertions(+), 30 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index d2d2e09b..616adf72 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -8,6 +8,10 @@ Upcoming version (not yet released) General ^^^^^^^ +- A more performant mju_transposeSparse has been implemented that doesn't require dense memory allocation. + For a constraint Jacobian matrix from the + `humanoid100.xml `_ model, + this function is 35% faster. - The function :ref:`mj_name2id` is now implemented using a hash function instead of a linear search for better performance. diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index c867ec95..ca28e334 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -528,37 +528,48 @@ void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, i void mju_transposeSparse(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) { - // clear counters for transposed + // clear number of non-zeros for each row of transposed memset(res_rownnz, 0, nc*sizeof(int)); - // set uncompressed layout - for (int rt=0; rt= rownnz[r]) r++; + + // row index becomes the column index + res[res_rowadr[colind[i]]] = mat[i]; + + // column index becomes the row index (increment the rowadr) + res_colind[res_rowadr[colind[i]]++] = r; + } + + // shift back row addresses + for (int i = nc-1; i>0; i--) { + res_rowadr[i] = res_rowadr[i-1]; + } + + res_rowadr[0] = 0; } - // construct row supernodes void mju_superSparse(int nr, int* rowsuper, const int* rownnz, const int* rowadr, const int* colind) { diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index cf9885a5..e9493176 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -48,8 +48,8 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec const int* colind, const int* rowsuper); // compress layout of sparse matrix -void mju_compressSparse(mjtNum* mat, int nr, int nc, - int* rownnz, int* rowadr, int* colind); +MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc, + int* rownnz, int* rowadr, int* colind); // combine two sparse vectors: dst = a*dst + b*src, return nnz of result int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b, @@ -61,9 +61,9 @@ void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNu int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind); // transpose sparse matrix -void mju_transposeSparse(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); +MJAPI void mju_transposeSparse(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); // construct row supernodes void mju_superSparse(int nr, int* rowsuper, diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index bc746718..cece3a81 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -26,9 +26,10 @@ namespace mujoco { namespace { -using FuncPtr = decltype(&mju_combineSparse); +using CombineFuncPtr = decltype(&mju_combineSparse); +using TransposeFuncPtr = decltype(&mju_transposeSparse); -// number of steps to roll out before benhmarking +// number of steps to roll out before benchmarking static const int kNumWarmupSteps = 500; // copy array into vector @@ -38,6 +39,27 @@ std::vector AsVector(const mjtNum* array, int n) { // ----------------------------- old functions -------------------------------- +// 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) { + memset(res_rownnz, 0, nc * sizeof(int)); + for (int rt = 0; rt < nc; rt++) { + res_rowadr[rt] = rt * nr; + } + for (int r = 0; r < nr; r++) { + for (int ci = 0; ci < rownnz[r]; ci++) { + int rt = colind[rowadr[r] + ci]; + res_colind[rt * nr + res_rownnz[rt]] = r; + res[rt * nr + res_rownnz[rt]] = mat[rowadr[r] + ci]; + res_rownnz[rt]++; + } + } + + mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind); +} + int compare_baseline(const int* vec1, const int* vec2, int n) { @@ -286,8 +308,9 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1( } BENCHMARK(BM_MatVecSparse_1); -static void BM_combineSparse(benchmark::State& state, FuncPtr func) { +static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) { static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml"); + mjData* d = mj_makeData(m); // warm-up rollout to get a typical state @@ -359,5 +382,52 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old( } BENCHMARK(BM_combineSparse_old); +static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) { + static mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml"); + + // force use of sparse matrices + m->opt.jacobian = mjJAC_SPARSE; + + mjData* d = mj_makeData(m); + + // warm-up rollout to get a typical state + while (d-> time < 2) { + mj_step(m, d); + } + + mjMARKSTACK; + + // need uncompressed layout + mjtNum* res = mj_stackAlloc(d, m->nv * d->nefc); + int* res_rownnz = (int*)mj_stackAlloc(d, m->nv); + int* res_rowadr = (int*)mj_stackAlloc(d, m->nv); + int* res_colind = (int*)mj_stackAlloc(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, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind); + } + + mjFREESTACK; + mj_deleteData(d); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_transposeSparse_new(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_transposeSparse(state, &mju_transposeSparse); +} +BENCHMARK(BM_transposeSparse_new); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL +BM_transposeSparse_old(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_transposeSparse(state, &transposeSparse_baseline); +} + +BENCHMARK(BM_transposeSparse_old); + } // namespace } // namespace mujoco diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 21f4229a..12777e5d 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -16,12 +16,14 @@ #include "src/engine/engine_util_sparse.h" +#include #include #include "test/fixture.h" namespace mujoco { namespace { +using ::testing::ElementsAre; using EngineUtilSparseTest = MujocoTest; TEST_F(EngineUtilSparseTest, MjuDot) { @@ -40,5 +42,143 @@ TEST_F(EngineUtilSparseTest, MjuDot) { EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7); } +TEST_F(EngineUtilSparseTest, MjuTranspose3by3) { + // 1 2 0 1 0 0 + // 0 1 0 --> 2 1 3 + // 0 3 0 0 0 0 + + mjtNum mat[] = {1, 2, 1, 3}; + int colind[] = {0, 1, 1, 1}; + int rownnz[] = {2, 1, 1}; + int rowadr[] = {0, 2, 3}; + + mjtNum matT[] = {0, 0, 0, 0}; + int colindT[] = {0, 0, 0, 0}; + int rownnzT[] = {0, 0, 0}; + int rowadrT[] = {0, 0, 0}; + + mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rownnz, + rowadr, colind); + + EXPECT_THAT(matT, ElementsAre(1, 2, 1, 3)); + EXPECT_THAT(colindT, ElementsAre(0, 0, 1, 2)); + EXPECT_THAT(rownnzT, ElementsAre(1, 3, 0)); + EXPECT_THAT(rowadrT, ElementsAre(0, 1, 4)); +} + +TEST_F(EngineUtilSparseTest, MjuTranspose1by3) { + // 1 0 3 1 + // --> 0 + // 3 + + mjtNum mat[] = {1, 3}; + int colind[] = {0, 2}; + int rownnz[] = {2}; + int rowadr[] = {0}; + + mjtNum matT[] = {0, 0}; + int colindT[] = {0, 0}; + int rownnzT[] = {0, 0, 0}; + int rowadrT[] = {0, 0, 0}; + + mju_transposeSparse(matT, mat, 1, 3, rownnzT, rowadrT, colindT, rownnz, + rowadr, colind); + + EXPECT_THAT(matT, ElementsAre(1, 3)); + EXPECT_THAT(colindT, ElementsAre(0, 0)); + EXPECT_THAT(rownnzT, ElementsAre(1, 0, 1)); + EXPECT_THAT(rowadrT, ElementsAre(0, 1, 1)); +} + +TEST_F(EngineUtilSparseTest, MjuTranspose3by1) { + // 1 1 0 3 + // 0 --> + // 3 + + mjtNum mat[] = {1, 3}; + int colind[] = {0, 0}; + int rownnz[] = {1, 0, 1}; + int rowadr[] = {0, 1, 1}; + + mjtNum matT[] = {0, 0}; + int colindT[] = {0, 0}; + int rownnzT[] = {0}; + int rowadrT[] = {0}; + + mju_transposeSparse(matT, mat, 3, 1, rownnzT, rowadrT, colindT, rownnz, + rowadr, colind); + + EXPECT_THAT(matT, ElementsAre(1, 3)); + EXPECT_THAT(colindT, ElementsAre(0, 2)); + EXPECT_THAT(rownnzT, ElementsAre(2)); + EXPECT_THAT(rowadrT, ElementsAre(0)); +} + +TEST_F(EngineUtilSparseTest, MjuTransposeDense) { + // 1 2 3 1 4 7 + // 4 5 6 --> 2 5 8 + // 7 8 9 3 6 9 + + mjtNum mat[] = {1, 2, 3, 4, 5, 6, 7, 8, 9}; + int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2}; + int rownnz[] = {3, 3, 3}; + int rowadr[] = {0, 3, 6}; + + mjtNum matT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; + int colindT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; + int rownnzT[] = {0, 0, 0}; + int rowadrT[] = {0, 0, 0}; + + mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, 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)); +} + +TEST_F(EngineUtilSparseTest, MjuTranspose1by1) { + // 1 -> 1 + + mjtNum mat[] = {1}; + int colind[] = {0}; + int rownnz[] = {1}; + int rowadr[] = {0}; + + mjtNum matT[] = {0}; + int colindT[] = {0}; + int rownnzT[] = {0}; + int rowadrT[] = {0}; + + mju_transposeSparse(matT, mat, 1, 1, rownnzT, rowadrT, colindT, rownnz, + rowadr, colind); + + EXPECT_THAT(matT, ElementsAre(1)); + EXPECT_THAT(colindT, ElementsAre(0)); + EXPECT_THAT(rownnzT, ElementsAre(1)); + EXPECT_THAT(rowadrT, ElementsAre(0)); +} + +TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) { + // 0 -> 0 + + mjtNum mat[] = {}; + int colind[] = {}; + int rownnz[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + int rowadr[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; + + mjtNum matT[] = {}; + int colindT[] = {}; + 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); + + 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)); +} + } // namespace } // namespace mujoco