Speed improvement with no AVX. 40% on Linux Intel Xeon and 60% on ARM Mac.

PiperOrigin-RevId: 488360660
Change-Id: I423269cd362fdcc2434393b691638ac6ed778ef2

doc/changelog.rst

@@ -43,8 +43,13 @@ General
 - Improved particle :ref:`composite<body-composite>` type, which now permits a user-specified geometry and multiple
   joints. See the two new examples:
   `particle_free.xml <https://github.com/deepmind/mujoco/tree/main/model/composite/particle_free.xml>`_ and
-  `particle_free2d.xml <https://github.com/deepmind/mujoco/tree/main/model/composite/particle_free2d.xml>`_
+  `particle_free2d.xml <https://github.com/deepmind/mujoco/tree/main/model/composite/particle_free2d.xml>`_.
-- Better mj_solveLD performance. This gives an improvement of 14% on a Intel Skylake Xeon with 36 cores.
+- Performance improvements for non-AVX configurations:
+
+  - 14% faster ``mj_solveLD`` using `restrict <https://en.wikipedia.org/wiki/Restrict>`_. See `engine_core_smooth_benchmark_test
+    <https://github.com/deepmind/mujoco/tree/main/test/benchmark/engine_core_smooth_benchmark_test.cc>`_.
+  - 50% faster ``mju_dotSparse`` using manual loop unroll. See `engine_util_sparse_benchmark_test
+    <https://github.com/deepmind/mujoco/tree/main/test/benchmark/engine_util_sparse_benchmark_test.cc>`_.
 
 Version 2.3.0 (October 18, 2022)
 --------------------------------

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

src/engine/engine_util_sparse.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_
 
 #include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
 #include <mujoco/mjtnum.h>
 
 #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,
+MJAPI mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
-                     const int nnz1, const int* ind1);
+                           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,
+MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
-                         int nr, const int* rownnz, const int* rowadr,
+                               int nr, const int* rownnz, const int* rowadr,
-                         const int* colind, const int* rowsuper);
+                               const int* colind, const int* rowsuper);
 
 // compress layout of sparse matrix
 void mju_compressSparse(mjtNum* mat, int nr, int nc,

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
+)

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);
   }
 

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 <benchmark/benchmark.h>
+#include <gtest/gtest.h>
+#include <absl/base/attributes.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mujoco.h>
+#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<mjtNum> AsVector(const mjtNum* array, int n) {
+  return std::vector<mjtNum>(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<mjtNum> qpos = AsVector(d->qpos, m->nq);
+  std::vector<mjtNum> qvel = AsVector(d->qvel, m->nv);
+  std::vector<mjtNum> act = AsVector(d->act, m->na);
+  std::vector<mjtNum> 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

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 <gtest/gtest.h>
+#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