From 7b0fbc63f899d7d6ae125dd663fa1b176f139e9a Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 14 Nov 2022 07:30:51 -0800 Subject: [PATCH] Speed improvement with no AVX. 40% on Linux Intel Xeon and 60% on ARM Mac. PiperOrigin-RevId: 488360660 Change-Id: I423269cd362fdcc2434393b691638ac6ed778ef2 --- doc/changelog.rst | 9 +- src/engine/engine_util_sparse.c | 16 ++ src/engine/engine_util_sparse.h | 11 +- test/benchmark/CMakeLists.txt | 16 ++ .../engine_core_smooth_benchmark_test.cc | 1 - .../engine_util_sparse_benchmark_test.cc | 221 ++++++++++++++++++ test/engine/engine_util_sparse_test.cc | 44 ++++ 7 files changed, 310 insertions(+), 8 deletions(-) create mode 100644 test/benchmark/engine_util_sparse_benchmark_test.cc create mode 100644 test/engine/engine_util_sparse_test.cc diff --git a/doc/changelog.rst b/doc/changelog.rst index 200a9548..e56af43b 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -43,8 +43,13 @@ General - Improved particle :ref:`composite` type, which now permits a user-specified geometry and multiple joints. See the two new examples: `particle_free.xml `_ and - `particle_free2d.xml `_ -- Better mj_solveLD performance. This gives an improvement of 14% on a Intel Skylake Xeon with 36 cores. + `particle_free2d.xml `_. +- Performance improvements for non-AVX configurations: + + - 14% faster ``mj_solveLD`` using `restrict `_. See `engine_core_smooth_benchmark_test + `_. + - 50% faster ``mju_dotSparse`` using manual loop unroll. See `engine_util_sparse_benchmark_test + `_. Version 2.3.0 (October 18, 2022) -------------------------------- diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 71a6d981..eeb25a50 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -73,6 +73,22 @@ mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, high64 = _mm_unpackhi_pd(vlow, vlow); res = _mm_cvtsd_f64(_mm_add_sd(vlow, high64)); } + +#else + int n_4 = nnz1 - 4; + + mjtNum res0 = 0; + mjtNum res1 = 0; + mjtNum res2 = 0; + mjtNum res3 = 0; + + for (; i<=n_4; i+=4) { + res0 += vec1[i+0] * vec2[ind1[i+0]]; + res1 += vec1[i+1] * vec2[ind1[i+1]]; + res2 += vec1[i+2] * vec2[ind1[i+2]]; + res3 += vec1[i+3] * vec2[ind1[i+3]]; + } + res = (res0 + res2) + (res1 + res3); #endif // scalar part diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 8db98f9f..d02f01af 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -16,6 +16,7 @@ #define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_ #include +#include #include #ifdef __cplusplus @@ -25,8 +26,8 @@ extern "C" { //------------------------------ sparse operations ------------------------------------------------- // dot-product, first vector is sparse -mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, - const int nnz1, const int* ind1); +MJAPI mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2, + const int nnz1, const int* ind1); // dot-product, both vectors are sparse mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, @@ -42,9 +43,9 @@ void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); // multiply sparse matrix and dense vector: res = mat * vec -void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, - int nr, const int* rownnz, const int* rowadr, - const int* colind, const int* rowsuper); +MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, + int nr, const int* rownnz, const int* rowadr, + const int* colind, const int* rowsuper); // compress layout of sparse matrix void mju_compressSparse(mjtNum* mat, int nr, int nc, diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index dddd0bbb..9c132c33 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -63,3 +63,19 @@ target_link_libraries( gmock benchmark::benchmark ) + +mujoco_benchmark_test(engine_core_smooth_benchmark_test) +target_link_libraries( + engine_core_smooth_benchmark_test + fixture + gmock + benchmark::benchmark +) + +mujoco_benchmark_test(engine_util_sparse_benchmark_test) +target_link_libraries( + engine_util_sparse_benchmark_test + fixture + gmock + benchmark::benchmark +) diff --git a/test/benchmark/engine_core_smooth_benchmark_test.cc b/test/benchmark/engine_core_smooth_benchmark_test.cc index 68713711..fc8f72f0 100644 --- a/test/benchmark/engine_core_smooth_benchmark_test.cc +++ b/test/benchmark/engine_core_smooth_benchmark_test.cc @@ -107,7 +107,6 @@ static void BM_solveLD(benchmark::State& state, bool new_function) { // warm-up rollout to get a typcal state for (int i=0; i < kNumWarmupSteps; i++) { - AddCtrlNoise(m, d, i); mj_step(m, d); } diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc new file mode 100644 index 00000000..a78dc6be --- /dev/null +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -0,0 +1,221 @@ +// Copyright 2021 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// A benchmark for comparing different implementations of mj_solveLD. + +#include +#include +#include +#include +#include +#include "src/engine/engine_util_sparse.h" +#include "test/fixture.h" + +namespace mujoco { +namespace { + +// number of steps to roll out before benhmarking +static const int kNumWarmupSteps = 500; + +// copy array into vector +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} + +// ----------------------------- old functions -------------------------------- + +mjtNum ABSL_ATTRIBUTE_NOINLINE dotSparse_1(const mjtNum* vec1, + const mjtNum* vec2, + const int nnz1, + const int* ind1) { + int i = 0; + mjtNum res = 0; + + // scalar part + for (; i < nnz1; i++) { + res += vec1[i] * vec2[ind1[i]]; + } + + return res; +} + +mjtNum ABSL_ATTRIBUTE_NOINLINE dotSparse_8(const mjtNum* vec1, + const mjtNum* vec2, + const int nnz1, + const int* ind1) { + int i = 0; + mjtNum res = 0; + + int n_8 = nnz1 - 8; + + mjtNum res0 = 0; + mjtNum res1 = 0; + mjtNum res2 = 0; + mjtNum res3 = 0; + mjtNum res4 = 0; + mjtNum res5 = 0; + mjtNum res6 = 0; + mjtNum res7 = 0; + + for (; i <= n_8; i+=8) { + res0 += vec1[i+0] * vec2[ind1[i+0]]; + res1 += vec1[i+1] * vec2[ind1[i+1]]; + res2 += vec1[i+2] * vec2[ind1[i+2]]; + res3 += vec1[i+3] * vec2[ind1[i+3]]; + res4 += vec1[i+4] * vec2[ind1[i+4]]; + res5 += vec1[i+5] * vec2[ind1[i+5]]; + res6 += vec1[i+6] * vec2[ind1[i+6]]; + res7 += vec1[i+7] * vec2[ind1[i+7]]; + } + res = ((res0 + res2) + (res1 + res3)) + ((res4 + res6) + (res5 + res7)); + + // process remaining + int n_i = nnz1 - i; + if (n_i == 7) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]] + + vec1[i+2]*vec2[ind1[i+2]] + vec1[i+3]*vec2[ind1[i+3]] + + vec1[i+4]*vec2[ind1[i+4]] + vec1[i+5]*vec2[ind1[i+5]] + + vec1[i+6]*vec2[ind1[i+6]]; + } else if (n_i == 6) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]] + + vec1[i+2]*vec2[ind1[i+2]] + vec1[i+3]*vec2[ind1[i+3]] + + vec1[i+4]*vec2[ind1[i+4]] + vec1[i+5]*vec2[ind1[i+5]]; + } else if (n_i == 5) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]] + + vec1[i+2]*vec2[ind1[i+2]] + vec1[i+3]*vec2[ind1[i+3]] + + vec1[i+4]*vec2[ind1[i+4]]; + } else if (n_i == 4) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]] + + vec1[i+2]*vec2[ind1[i+2]] + vec1[i+1]*vec2[ind1[i+3]]; + } else if (n_i == 3) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]] + + vec1[i+2]*vec2[ind1[i+2]]; + } else if (n_i == 2) { + res += vec1[i+0]*vec2[ind1[i+0]] + vec1[i+1]*vec2[ind1[i+1]]; + } else if (n_i == 1) { + res += vec1[i+0]*vec2[ind1[i+0]]; + } + + return res; +} + +void ABSL_ATTRIBUTE_NOINLINE mulMatVecSparse_1(mjtNum* res, + const mjtNum* mat, + const mjtNum* vec, + int nr, + const int* rownnz, + const int* rowadr, + const int* colind, + const int* rowsuper) { + for (int r=0; r < nr; r++) { + res[r] = dotSparse_1( + mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]); + } +} + +void ABSL_ATTRIBUTE_NOINLINE mulMatVecSparse_8(mjtNum* res, + const mjtNum* mat, + const mjtNum* vec, + int nr, + const int* rownnz, + const int* rowadr, + const int* colind, + const int* rowsuper) { + for (int r=0; r < nr; r++) { + res[r] = dotSparse_8( + mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]); + } +} + +// ----------------------------- benchmark ------------------------------------ + +static void BM_MatVecSparse(benchmark::State& state, int unroll) { + static mjModel* m = LoadModelFromPath("composite/cloth.xml"); + mjData* d = mj_makeData(m); + + // warm-up rollout to get a typical state + for (int i=0; i < kNumWarmupSteps; i++) { + mj_step(m, d); + } + + // allocate gradient + mjMARKSTACK; + mjtNum *Ma = mj_stackAlloc(d, m->nv); + mjtNum *vec = mj_stackAlloc(d, m->nv); + mjtNum *res = mj_stackAlloc(d, m->nv); + mjtNum *grad = mj_stackAlloc(d, m->nv); + mjtNum *Mgrad = mj_stackAlloc(d, m->nv); + + // compute gradient + mj_mulM(m, d, Ma, d->qacc); + for (int i=0; i < m->nv; i++) { + grad[i] = Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i]; + } + + // compute search direction + mj_solveM(m, d, Mgrad, grad, 1); + mju_scl(vec, Mgrad, -1, m->nv); + + // save state + std::vector qpos = AsVector(d->qpos, m->nq); + std::vector qvel = AsVector(d->qvel, m->nv); + std::vector act = AsVector(d->act, m->na); + std::vector warmstart = AsVector(d->qacc_warmstart, m->nv); + + // time benchmark + for (auto s : state) { + if (unroll == 4) { + mju_mulMatVecSparse(res, d->efc_J, vec, d->nefc, + d->efc_J_rownnz, d->efc_J_rowadr, + d->efc_J_colind, d->efc_J_rowsuper); + } else if (unroll == 1) { + mulMatVecSparse_1(res, d->efc_J, vec, d->nefc, + d->efc_J_rownnz, d->efc_J_rowadr, + d->efc_J_colind, d->efc_J_rowsuper); + } else if (unroll == 8) { + mulMatVecSparse_8(res, d->efc_J, vec, d->nefc, + d->efc_J_rownnz, d->efc_J_rowadr, + d->efc_J_colind, d->efc_J_rowsuper); + } + } + + // finalize + mjFREESTACK; + mj_deleteData(d); + state.SetItemsProcessed(state.iterations()); +} + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_8( + benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_MatVecSparse(state, 8); +} +BENCHMARK(BM_MatVecSparse_8); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_4( + benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_MatVecSparse(state, 4); +} +BENCHMARK(BM_MatVecSparse_4); + +void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1( + benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_MatVecSparse(state, 1); +} +BENCHMARK(BM_MatVecSparse_1); + +} // namespace +} // namespace mujoco diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc new file mode 100644 index 00000000..21f4229a --- /dev/null +++ b/test/engine/engine_util_sparse_test.cc @@ -0,0 +1,44 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Tests for engine/engine_util_sparse.c + +#include "src/engine/engine_util_sparse.h" + +#include +#include "test/fixture.h" + +namespace mujoco { +namespace { + +using EngineUtilSparseTest = MujocoTest; + +TEST_F(EngineUtilSparseTest, MjuDot) { + mjtNum a[] = {1, 2, 3, 4, 5, 6, 7}; + mjtNum b[] = {7, 0, 6, 0, 0, 5, 0, 0, 0, 4, 0, 0, 0, 3, 0, 0, 2, 0, 1}; + int i[] = {0, 2, 5, 9, 13, 16, 18}; + + // test various vector lengths as mju_dotSparse adds numbers in groups of four + EXPECT_EQ(mju_dotSparse(a, b, 0, i), 0); + EXPECT_EQ(mju_dotSparse(a, b, 1, i), 7); + EXPECT_EQ(mju_dotSparse(a, b, 2, i), 7 + 2*6); + EXPECT_EQ(mju_dotSparse(a, b, 3, i), 7 + 2*6 + 3*5); + EXPECT_EQ(mju_dotSparse(a, b, 4, i), 7 + 2*6 + 3*5 + 4*4); + EXPECT_EQ(mju_dotSparse(a, b, 5, i), 7 + 2*6 + 3*5 + 4*4 + 5*3); + EXPECT_EQ(mju_dotSparse(a, b, 6, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2); + EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7); +} + +} // namespace +} // namespace mujoco