Cache benchmark data in engine_util_sparse_benchmark_test

PiperOrigin-RevId: 902651349
Change-Id: I760e3f696b37699366c40c3c93a74e2b1b35678e

test/benchmark/engine_util_sparse_benchmark_test.cc

@@ -14,7 +14,6 @@
 
 // A benchmark for comparing different implementations of mj_solveLD.
 
-#include <cstddef>
 #include <cstring>
 #include <vector>
 
@@ -31,14 +30,179 @@ namespace {
 
 using CombineFuncPtr = decltype(&mju_combineSparse);
 using TransposeFuncPtr = decltype(&mju_transposeSparse);
-using SqrMatTDFuncPtr = decltype(&mju_sqrMatTDSparse);
 
-// number of steps to roll out before benchmarking
+// ================================ Cached Data ================================
-static const int kNumWarmupSteps = 500;
 
-// ----------------------------- old functions --------------------------------
+// ---- MatVecSparse data ----
+struct MatVecData {
+  int nv;
+  int nefc;
+  int nJ;
+  std::vector<mjtNum> efc_J;
+  std::vector<int> efc_J_rownnz, efc_J_rowadr, efc_J_colind, efc_J_rowsuper;
+  std::vector<mjtNum> vec;
+};
 
+MatVecData& GetMatVecData() {
+  static MatVecData data = [] {
+    MatVecData d;
+    mjModel* m = LoadModelFromPath("flex/flag.xml");
+    mjData* dat = mj_makeData(m);
 
+    for (int i = 0; i < 500; i++) {
+      mj_step(m, dat);
+    }
+
+    d.nv = m->nv;
+    d.nefc = dat->nefc;
+    d.nJ = dat->nJ;
+    d.efc_J.assign(dat->efc_J, dat->efc_J + d.nJ);
+    d.efc_J_rownnz.assign(dat->efc_J_rownnz, dat->efc_J_rownnz + d.nefc);
+    d.efc_J_rowadr.assign(dat->efc_J_rowadr, dat->efc_J_rowadr + d.nefc);
+    d.efc_J_colind.assign(dat->efc_J_colind, dat->efc_J_colind + d.nJ);
+    d.efc_J_rowsuper.assign(dat->efc_J_rowsuper, dat->efc_J_rowsuper + d.nefc);
+
+    // compute direction: vec = -M^{-1} * (Ma - qfrc_smooth - qfrc_constraint)
+    mj_markStack(dat);
+    mjtNum* Ma = mj_stackAllocNum(dat, m->nv);
+    mjtNum* grad = mj_stackAllocNum(dat, m->nv);
+    mjtNum* Mgrad = mj_stackAllocNum(dat, m->nv);
+    mj_mulM(m, dat, Ma, dat->qacc);
+    for (int i = 0; i < m->nv; i++) {
+      grad[i] = Ma[i] - dat->qfrc_smooth[i] - dat->qfrc_constraint[i];
+    }
+    mj_solveM(m, dat, Mgrad, grad, 1);
+    d.vec.resize(m->nv);
+    mju_scl(d.vec.data(), Mgrad, -1, m->nv);
+    mj_freeStack(dat);
+
+    mj_deleteData(dat);
+    mj_deleteModel(m);
+    return d;
+  }();
+  return data;
+}
+
+// ---- CombineSparse data ----
+struct CombineData {
+  int nv;
+  std::vector<mjtNum> H;
+  std::vector<int> rownnz, rowadr, colind;
+};
+
+CombineData& GetCombineData() {
+  static CombineData data = [] {
+    CombineData cd;
+    mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
+    m->opt.jacobian = mjJAC_SPARSE;
+    mjData* d = mj_makeData(m);
+
+    for (int i = 0; i < 500; i++) {
+      mj_step(m, d);
+    }
+
+    cd.nv = m->nv;
+    mj_markStack(d);
+    mjtNum* H = mj_stackAllocNum(d, m->nv*m->nv);
+    int* rownnz = mj_stackAllocInt(d, m->nv);
+    int* rowadr = mj_stackAllocInt(d, m->nv);
+    int* colind = mj_stackAllocInt(d, m->nv*m->nv);
+    int* diagind = mj_stackAllocInt(d, m->nv);
+
+    mjtNum* D = mj_stackAllocNum(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;
+      }
+    }
+
+    int* JT_rownnz   = mj_stackAllocInt(d, m->nv);
+    int* JT_rowadr   = mj_stackAllocInt(d, m->nv);
+    int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
+    int* JT_colind   = mj_stackAllocInt(d, d->nJ);
+    mjtNum* JT       = mj_stackAllocNum(d, d->nJ);
+    mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
+                        JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
+                        d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
+
+    // compute H = J'*D*J, uncompressed layout
+    mju_sqrMatTDUncompressedInit(rowadr, m->nv);
+    mju_sqrMatTDSparse(H, d->efc_J, 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,
+                       JT_rownnz, JT_rowadr,
+                       JT_colind, JT_rowsuper, d,
+                       diagind);
+
+    // compute H = M + J'*D*J
+    mj_addM(m, d, H, rownnz, rowadr, colind);
+
+    // copy to persistent storage
+    int nH = rowadr[m->nv-1] + m->nv;  // uncompressed: rowadr[r] = r*nv
+    cd.H.assign(H, H + nH);
+    cd.rownnz.assign(rownnz, rownnz + m->nv);
+    cd.rowadr.assign(rowadr, rowadr + m->nv);
+    cd.colind.assign(colind, colind + nH);
+
+    mj_freeStack(d);
+    mj_deleteData(d);
+    mj_deleteModel(m);
+    return cd;
+  }();
+  return data;
+}
+
+// ---- TransposeSparse data ----
+struct TransposeData {
+  int nv;
+  int nefc;
+  int nJ;
+  std::vector<mjtNum> efc_J;
+  std::vector<int> efc_J_rownnz, efc_J_rowadr, efc_J_colind;
+};
+
+enum class Size { H2_100, H100 };
+
+template <Size S>
+const char* ModelPath() {
+  if constexpr (S == Size::H2_100) {
+    return "../test/benchmark/testdata/2humanoid100_chol.xml";
+  } else {
+    return "../test/benchmark/testdata/100_humanoids_chol.xml";
+  }
+}
+
+template <Size S>
+TransposeData& GetTransposeData() {
+  static TransposeData data = [] {
+    TransposeData td;
+    mjModel* m = LoadModelFromPath(ModelPath<S>());
+    m->opt.jacobian = mjJAC_SPARSE;
+    mjData* d = mj_makeData(m);
+
+    while (d->time < 2) {
+      mj_step(m, d);
+    }
+
+    td.nv = m->nv;
+    td.nefc = d->nefc;
+    td.nJ = d->nJ;
+    td.efc_J.assign(d->efc_J, d->efc_J + d->nJ);
+    td.efc_J_rownnz.assign(d->efc_J_rownnz, d->efc_J_rownnz + d->nefc);
+    td.efc_J_rowadr.assign(d->efc_J_rowadr, d->efc_J_rowadr + d->nefc);
+    td.efc_J_colind.assign(d->efc_J_colind, d->efc_J_colind + d->nJ);
+
+    mj_deleteData(d);
+    mj_deleteModel(m);
+    return td;
+  }();
+  return data;
+}
+
+// ================================ old functions ==============================
 
 // transpose sparse matrix (uncompressed)
 void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
@@ -229,61 +393,31 @@ void ABSL_ATTRIBUTE_NOINLINE mulMatVecSparse_8(mjtNum* res,
   }
 }
 
-// ----------------------------- benchmark ------------------------------------
+// ----------------------------- benchmark -------------------------------------
 
 static void BM_MatVecSparse(benchmark::State& state, int unroll) {
-  static mjModel* m = LoadModelFromPath("flex/flag.xml");
+  MatVecData& data = GetMatVecData();
-  mjData* d = mj_makeData(m);
+  std::vector<mjtNum> res(data.nefc);
 
-  // warm-up rollout to get a typical state
-  for (int i=0; i < kNumWarmupSteps; i++) {
-    mj_step(m, d);
-  }
-
-  // allocate gradient
-  mj_markStack(d);
-  mjtNum *Ma = mj_stackAllocNum(d, m->nv);
-  mjtNum *vec = mj_stackAllocNum(d, m->nv);
-  mjtNum *res = mj_stackAllocNum(d, d->nefc);
-  mjtNum *grad = mj_stackAllocNum(d, m->nv);
-  mjtNum *Mgrad  = mj_stackAllocNum(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,
+      mju_mulMatVecSparse(res.data(), data.efc_J.data(), data.vec.data(),
-                          d->efc_J_rownnz, d->efc_J_rowadr,
+                          data.nefc, data.efc_J_rownnz.data(),
-                          d->efc_J_colind, d->efc_J_rowsuper);
+                          data.efc_J_rowadr.data(), data.efc_J_colind.data(),
+                          data.efc_J_rowsuper.data());
     } else if (unroll == 1) {
-      mulMatVecSparse_1(res, d->efc_J, vec, d->nefc,
+      mulMatVecSparse_1(res.data(), data.efc_J.data(), data.vec.data(),
-                        d->efc_J_rownnz, d->efc_J_rowadr,
+                        data.nefc, data.efc_J_rownnz.data(),
-                        d->efc_J_colind, d->efc_J_rowsuper);
+                        data.efc_J_rowadr.data(), data.efc_J_colind.data(),
+                        data.efc_J_rowsuper.data());
     } else if (unroll == 8) {
-      mulMatVecSparse_8(res, d->efc_J, vec, d->nefc,
+      mulMatVecSparse_8(res.data(), data.efc_J.data(), data.vec.data(),
-                        d->efc_J_rownnz, d->efc_J_rowadr,
+                        data.nefc, data.efc_J_rownnz.data(),
-                        d->efc_J_colind, d->efc_J_rowsuper);
+                        data.efc_J_rowadr.data(), data.efc_J_colind.data(),
+                        data.efc_J_rowsuper.data());
     }
   }
 
-  // finalize
-  mj_freeStack(d);
-  mj_deleteData(d);
   state.SetItemsProcessed(state.iterations());
 }
 
@@ -309,75 +443,30 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1(
 BENCHMARK(BM_MatVecSparse_1);
 
 static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
-  static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
+  CombineData& data = GetCombineData();
-  m->opt.jacobian = mjJAC_SPARSE;
 
-  mjData* d = mj_makeData(m);
+  // make working copies that get modified each iteration
-
+  std::vector<mjtNum> H = data.H;
-  // warm-up rollout to get a typical state
+  std::vector<int> rownnz = data.rownnz;
-  for (int i=0; i < kNumWarmupSteps; i++) {
+  std::vector<int> rowadr = data.rowadr;
-    mj_step(m, d);
+  std::vector<int> colind = data.colind;
-  }
-
-  // allocate
-  mj_markStack(d);
-  mjtNum* H = mj_stackAllocNum(d, m->nv*m->nv);
-  int* rownnz = mj_stackAllocInt(d, m->nv);
-  int* rowadr = mj_stackAllocInt(d, m->nv);
-  int* colind = mj_stackAllocInt(d, m->nv*m->nv);
-  int* diagind = mj_stackAllocInt(d, m->nv);
-
-  // compute D corresponding to quad states
-  mjtNum* D = mj_stackAllocNum(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;
-    }
-  }
-
-  int* JT_rownnz   = mj_stackAllocInt(d, m->nv);
-  int* JT_rowadr   = mj_stackAllocInt(d, m->nv);
-  int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
-  int* JT_colind   = mj_stackAllocInt(d, d->nJ);
-  mjtNum* JT       = mj_stackAllocNum(d, d->nJ);
-  mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
-                      JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
-                      d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
-
-  // compute H = J'*D*J, uncompressed layout
-  mju_sqrMatTDUncompressedInit(rowadr, m->nv);
-  mju_sqrMatTDSparse(H, d->efc_J, 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,
-                     JT_rownnz, JT_rowadr,
-                     JT_colind, JT_rowsuper, d,
-                     diagind);
-
-  // 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 r = data.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], 1, -H[adr+i],
+         func(H.data()+rowadr[c], H.data()+rowadr[r], 1, -H[adr+i],
               rownnz[c], rownnz[c],
-              colind+rowadr[c], colind+rowadr[c]);
+              colind.data()+rowadr[c], colind.data()+rowadr[c]);
       }
     }
   }
 
-  // finalize
-  mj_freeStack(d);
-  mj_deleteData(d);
   state.SetItemsProcessed(state.iterations());
 }
 
@@ -395,17 +484,6 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
 }
 BENCHMARK(BM_combineSparse_old);
 
-enum class Size { H2_100, H100 };
-
-template <Size S>
-const char* ModelPath() {
-  if constexpr (S == Size::H2_100) {
-    return "../test/benchmark/testdata/2humanoid100_chol.xml";
-  } else {
-    return "../test/benchmark/testdata/100_humanoids_chol.xml";
-  }
-}
-
 enum class Supernode {
   None,
   PostProcess,
@@ -415,44 +493,33 @@ enum class Supernode {
 template <Size S>
 static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
                                Supernode super) {
-  static mjModel* m = LoadModelFromPath(ModelPath<S>());
+  TransposeData& data = GetTransposeData<S>();
 
-  // force use of sparse matrices
+  // allocate output buffers (uncompressed layout)
-  m->opt.jacobian = mjJAC_SPARSE;
+  std::vector<mjtNum> res(data.nv * data.nefc);
-
+  std::vector<int> res_rownnz(data.nv);
-  mjData* d = mj_makeData(m);
+  std::vector<int> res_rowadr(data.nv);
-
+  std::vector<int> res_rowsuper(data.nv);
-  // warm-up rollout to get a typical state
+  std::vector<int> res_colind(data.nv * data.nefc);
-  while (d->time < 2) {
-    mj_step(m, d);
-  }
-
-  mj_markStack(d);
-
-  // need uncompressed layout
-  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) {
-    int* rowsuper = (super == Supernode::Inline) ? res_rowsuper : nullptr;
+    int* rowsuper =
-    func(res, d->efc_J, d->nefc, m->nv,
+        (super == Supernode::Inline) ? res_rowsuper.data() : nullptr;
-         res_rownnz, res_rowadr, res_colind, rowsuper,
+    func(res.data(), data.efc_J.data(), data.nefc, data.nv,
-         d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
+         res_rownnz.data(), res_rowadr.data(), res_colind.data(), rowsuper,
+         data.efc_J_rownnz.data(), data.efc_J_rowadr.data(),
+         data.efc_J_colind.data());
     if (super == Supernode::PostProcess) {
-      mju_superSparse(m->nv, res_rowsuper,
+      mju_superSparse(data.nv, res_rowsuper.data(),
-                      res_rownnz, res_rowadr, res_colind);
+                      res_rownnz.data(), res_rowadr.data(), res_colind.data());
     }
   }
 
-  mj_freeStack(d);
-  mj_deleteData(d);
   state.SetItemsProcessed(state.iterations());
 }
 
+
 void ABSL_ATTRIBUTE_NO_TAIL_CALL
 BM_transposeSparse_2H100_old(benchmark::State& state) {
   MujocoErrorTestGuard guard;