Further speedup to CSR back-substitution using dof_simplenum.

PiperOrigin-RevId: 711440268 Change-Id: I81cd9a6a8b8ec78d08cfbc34f60a9216ea753833
2025-01-02 09:03:54 -08:00
parent 69c9ac074a
commit 4510c6d290
5 changed files with 85 additions and 15 deletions
@@ -0,0 +1,105 @@
+// 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 <absl/base/attributes.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mujoco.h>
+#include "src/engine/engine_core_smooth.h"
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+// number of steps to benchmark
+static const int kNumBenchmarkSteps = 50;
+
+// ----------------------------- benchmark ------------------------------------
+
+static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
+  static mjModel* m;
+  if (coil) {
+    m = LoadModelFromPath("plugin/elasticity/coil.xml");
+  } else {
+    m = LoadModelFromPath("humanoid/humanoid100.xml");
+  }
+
+  mjData* d = mj_makeData(m);
+  mj_forward(m, d);
+
+  // allocate input and output vectors
+  mj_markStack(d);
+  mjtNum *vec = mj_stackAllocNum(d, m->nv);
+  mjtNum *res = mj_stackAllocNum(d, m->nv);
+
+  // arbitrary input vector
+  for (int i=0; i < m->nv; i++) {
+    vec[i] = 0.2 + 0.3*i;
+  }
+
+  // make CSR matrix
+  mjtNum* LDs = mj_stackAllocNum(d, m->nC);
+  for (int i=0; i < m->nC; i++) {
+    LDs[i] = d->qLD[d->mapM2C[i]];
+  }
+
+  // benchmark
+  while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
+    for (int i=0; i < kNumBenchmarkSteps; i++) {
+      if (featherstone) {
+        mj_solveM(m, d, res, vec, 1);
+      } else {
+        mju_copy(res, vec, m->nv);
+        mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
+                    d->C_rownnz, d->C_rowadr, d->C_diag, m->dof_simplenum,
+                    d->C_colind);
+      }
+    }
+  }
+
+  // finalize
+  mj_freeStack(d);
+  mj_deleteData(d);
+  mj_deleteModel(m);
+  state.SetItemsProcessed(state.iterations());
+}
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_FS(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, /*featherstone=*/true, /*coil=*/true);
+}
+BENCHMARK(BM_solveLD_COIL_FS);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_CSR(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, /*featherstone=*/false, /*coil=*/true);
+}
+BENCHMARK(BM_solveLD_COIL_CSR);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_H100_FS(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, /*featherstone=*/true, /*coil=*/false);
+}
+BENCHMARK(BM_solveLD_H100_FS);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_H100_CSR(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, /*featherstone=*/false, /*coil=*/false);
+}
+BENCHMARK(BM_solveLD_H100_CSR);
+
+}  // namespace
+}  // namespace mujoco