Use memcmp for mju_compare if no AVX. 15% faster on Humanoid.

PiperOrigin-RevId: 493324356
Change-Id: I5f71dc1a8ae40c18864ea80f00a5f86ff952f72c

src/engine/engine_util_sparse.c

@@ -397,13 +397,7 @@ static int mju_compare(const int* vec1, const int* vec2, int n) {
 #endif
 
   // scalar part
-  for (; i<n; i++) {
+  return !memcmp(vec1+i, vec2+i, (n-i)*sizeof(int));
-    if (vec1[i]!=vec2[i]) {
-      return 0;
-    }
-  }
-
-  return 1;
 }
 
 // combine two sparse vectors: dst = a*dst + b*src, return nnz of result

src/engine/engine_util_sparse.h

@@ -70,14 +70,14 @@ void mju_superSparse(int nr, int* rowsuper,
                      const int* rownnz, const int* rowadr, const int* colind);
 
 // compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
-void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
+MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
-                        const mjtNum* diag, int nr, int nc,
+                              const mjtNum* diag, int nr, int nc,
-                        int* res_rownnz, int* res_rowadr, int* res_colind,
+                              int* res_rownnz, int* res_rowadr, int* res_colind,
-                        const int* rownnz, const int* rowadr,
+                              const int* rownnz, const int* rowadr,
-                        const int* colind, const int* rowsuper,
+                              const int* colind, const int* rowsuper,
-                        const int* rownnzT, const int* rowadrT,
+                              const int* rownnzT, const int* rowadrT,
-                        const int* colindT, const int* rowsuperT,
+                              const int* colindT, const int* rowsuperT,
-                        mjData* d);
+                              mjData* d);
 
 
 #ifdef __cplusplus

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -14,6 +14,7 @@
 
 // A benchmark for comparing different implementations of mj_solveLD.
 
+#include <cstddef>
 #include <benchmark/benchmark.h>
 #include <gtest/gtest.h>
 #include <absl/base/attributes.h>
@@ -25,6 +26,8 @@
 namespace mujoco {
 namespace {
 
+using FuncPtr = decltype(&mju_combineSparse);
+
 // number of steps to roll out before benhmarking
 static const int kNumWarmupSteps = 500;
 
@@ -35,6 +38,72 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
 
 // ----------------------------- old functions --------------------------------
 
+int compare_baseline(const int* vec1,
+                     const int* vec2,
+                     int n) {
+  int i = 0;
+
+  for (; i < n; i++) {
+    if (vec1[i] != vec2[i]) {
+      return 0;
+    }
+  }
+
+  return 1;
+}
+
+void addToSclScl(mjtNum* res,
+                 const mjtNum* vec,
+                 mjtNum scl1,
+                 mjtNum scl2,
+                 int n) {
+  int i = 0;
+
+  for (; i < n; i++) {
+    res[i] = res[i]*scl1 + vec[i]*scl2;
+  }
+}
+
+int compare_memcmp(const int* vec1,
+                const int* vec2,
+                int n) {
+  return !memcmp(vec1, vec2, n*sizeof(int));
+}
+
+int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
+                                                   const mjtNum* src, int n,
+                                                   mjtNum a, mjtNum b,
+                                                   int dst_nnz, int src_nnz,
+                                                   int* dst_ind,
+                                                   const int* src_ind,
+                                                   mjtNum* buf, int* buf_ind) {
+  // check for identical pattern
+  if (compare_baseline(dst_ind, src_ind, dst_nnz)) {
+    // combine mjtNum data directly
+    addToSclScl(dst, src, a, b, dst_nnz);
+    return dst_nnz;
+  } else {
+    return 0;
+  }
+}
+
+int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst,
+                                              const mjtNum* src, int n,
+                                              mjtNum a, mjtNum b,
+                                              int dst_nnz, int src_nnz,
+                                              int* dst_ind,
+                                              const int* src_ind,
+                                              mjtNum* buf, int* buf_ind) {
+  // check for identical pattern
+  if (compare_memcmp(dst_ind, src_ind, dst_nnz)) {
+    // combine mjtNum data directly
+    addToSclScl(dst, src, a, b, dst_nnz);
+    return dst_nnz;
+  } else {
+    return 0;
+  }
+}
+
 mjtNum ABSL_ATTRIBUTE_NOINLINE dotSparse_1(const mjtNum* vec1,
                                            const mjtNum* vec2,
                                            const int nnz1,
@@ -217,5 +286,78 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1(
 }
 BENCHMARK(BM_MatVecSparse_1);
 
+static void BM_combineSparse(benchmark::State& state, FuncPtr func) {
+  static mjModel* m = LoadModelFromPath("humanoid/humanoid.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
+  mjMARKSTACK;
+  mjtNum* H = mj_stackAlloc(d, m->nv*m->nv);
+  int* rownnz = (int*)mj_stackAlloc(d, m->nv);
+  int* rowadr = (int*)mj_stackAlloc(d, m->nv);
+  int* colind = (int*)mj_stackAlloc(d, m->nv*m->nv);
+
+  // compute D corresponding to quad states
+  mjtNum* D = mj_stackAlloc(d, d->nefc);
+  for (int i = 0; i < d->nefc; i++) {
+    if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
+      D[i] = d->efc_D[i];
+    } else {
+      D[i] = 0;
+    }
+  }
+
+  // compute H = J'*D*J, uncompressed layout
+  mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv,
+                      rownnz, rowadr, colind,
+                      d->efc_J_rownnz, d->efc_J_rowadr,
+                      d->efc_J_colind, d->efc_J_rowsuper,
+                      d->efc_JT_rownnz, d->efc_JT_rowadr,
+                      d->efc_JT_colind, d->efc_JT_rowsuper, d);
+
+  // compute H = M + J'*D*J
+  mj_addM(m, d, H, rownnz, rowadr, colind);
+
+  // time benchmark
+  for (auto s : state) {
+    for (int r = m->nv-1; r >= 0; r--) {
+      for (int i = 0; i < rownnz[r]-1; i++) {
+        int adr = rowadr[r];
+        int c = colind[adr+i];
+        // true arguments should be i+1 and colind+rowadr[r]
+        // but instead we repeat rownnz[c] and colind+rowadr[c]
+        // in order to trigger all if's in combineSparse
+         func(H+rowadr[c], H+rowadr[r], c+1, 1, -H[adr+i],
+              rownnz[c], rownnz[c],
+              colind+rowadr[c], colind+rowadr[c], NULL, NULL);
+      }
+    }
+  }
+
+  // finalize
+  mjFREESTACK;
+  mj_deleteData(d);
+  state.SetItemsProcessed(state.iterations());
+}
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_new(
+  benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_combineSparse(state, &combineSparse_new);
+}
+BENCHMARK(BM_combineSparse_new);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
+  benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_combineSparse(state, &combineSparse_baseline);
+}
+BENCHMARK(BM_combineSparse_old);
+
 }  // namespace
 }  // namespace mujoco