Speed improvements in mj_solveLD (14% on Intel Skylake Xeon 36 cores).

PiperOrigin-RevId: 486960670
Change-Id: I8fce96427158ec36069dea893500059e93bd40db

doc/changelog.rst

@@ -38,6 +38,7 @@ General
   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>`_
+- Better mj_solveLD performance. This gives an improvement of 14% on a Intel Skylake Xeon with 36 cores.
 
 Version 2.3.0 (October 18, 2022)
 --------------------------------

src/engine/engine_core_smooth.c

@@ -1088,32 +1088,26 @@ void mj_factorM(const mjModel* m, mjData* d) {
 // sparse backsubstitution:  x = inv(L'*D*L)*y
 //  L is in lower triangle of qLD; D is on diagonal of qLD
 //  handle n vectors at once
-void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n,
+void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
                 const mjtNum* qLD, const mjtNum* qLDiagInv) {
-  mjtNum tmp;
-
   // local copies of key variables
   int* dof_Madr = m->dof_Madr;
   int* dof_parentid = m->dof_parentid;
   int nv = m->nv;
 
-  // x = y
-  if (x != y) {
-    mju_copy(x, y, n*nv);
-  }
-
   // single vector
   if (n==1) {
     // x <- inv(L') * x; skip simple, exploit sparsity of input vector
     for (int i=nv-1; i>=0; i--) {
-      if (!m->dof_simplenum[i] && (tmp = x[i])) {
+      if (!m->dof_simplenum[i] && x[i]) {
         // init
         int Madr_ij = dof_Madr[i]+1;
         int j = dof_parentid[i];
 
         // traverse ancestors backwards
+        // read directly from x[i] since i cannot be a parent of itself
         while (j>=0) {
-          x[j] -= qLD[Madr_ij++]*tmp;         // x(j) -= L(i,j) * x(i)
+          x[j] -= qLD[Madr_ij++]*x[i];         // x(j) -= L(i,j) * x(i)
 
           // advance to parent
           j = dof_parentid[j];
@@ -1134,14 +1128,13 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n,
         int j = dof_parentid[i];
 
         // traverse ancestors backwards
-        tmp = x[i];
+        // write directly in x[i] since i cannot be a parent of itself
         while (j>=0) {
-          tmp -= qLD[Madr_ij++]*x[j];             // x(i) -= L(i,j) * x(j)
+          x[i] -= qLD[Madr_ij++]*x[j];             // x(i) -= L(i,j) * x(j)
 
           // advance to parent
           j = dof_parentid[j];
         }
-        x[i] = tmp;
       }
     }
   }
@@ -1149,6 +1142,7 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n,
   // multiple vectors
   else {
     int offset;
+    mjtNum tmp;
 
     // x <- inv(L') * x; skip simple
     for (int i=nv-1; i>=0; i--) {
@@ -1204,11 +1198,13 @@ void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n,
 }
 
 
-
 // sparse backsubstitution:  x = inv(L'*D*L)*y
 //  use factorization in d
 void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
-  mj_solveLD(m, x, y, n, d->qLD, d->qLDiagInv);
+  if (x != y) {
+    mju_copy(x, y, n*m->nv);
+  }
+  mj_solveLD(m, x, n, d->qLD, d->qLDiagInv);
 }
 
 

src/engine/engine_core_smooth.h

@@ -56,7 +56,7 @@ MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD,
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
 
 // sparse backsubstitution:  x = inv(L'*D*L)*y
-MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, const mjtNum* y, int n,
+MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
                       const mjtNum* qLD, const mjtNum* qLDiagInv);
 
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d

src/engine/engine_forward.c

@@ -576,7 +576,8 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
 
     // solve
     mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
-    mj_solveLD(m, qacc, qfrc, 1, d->qH, d->qHDiagInv);
+    mju_copy(qacc, qfrc, m->nv);
+    mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
   }
 
   // advance state and time

test/benchmark/engine_core_smooth_benchmark_test.cc

@@ -0,0 +1,167 @@
+// 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 "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+// number of steps to roll out before benhmarking
+static const int kNumWarmupSteps = 500;
+
+// number of steps to benchmark
+static const int kNumBenchmarkSteps = 50;
+
+// copy array into vector
+std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
+  return std::vector<mjtNum>(array, array + n);
+}
+
+// ----------------------------- old functions --------------------------------
+
+void ABSL_ATTRIBUTE_NOINLINE solveLD_baseline(const mjModel* m, mjtNum* x,
+                                              const mjtNum* y,
+                                              const mjtNum* qLD,
+                                              const mjtNum* qLDiagInv) {
+  mjtNum tmp;
+
+  // local copies of key variables
+  int* dof_Madr = m->dof_Madr;
+  int* dof_parentid = m->dof_parentid;
+  int nv = m->nv;
+
+  // x = y
+  if (x != y) {
+    mju_copy(x, y, nv);
+  }
+
+  // x <- inv(L') * x; skip simple, exploit sparsity of input vector
+  for (int i=nv-1; i >= 0; i--) {
+    if (!m->dof_simplenum[i] && (tmp = x[i])) {
+      // init
+      int Madr_ij = dof_Madr[i]+1;
+      int j = dof_parentid[i];
+
+      // traverse ancestors backwards
+      while (j >= 0) {
+        x[j] -= qLD[Madr_ij++]*tmp;         // x(j) -= L(i,j) * x(i)
+
+        // advance to parent
+        j = dof_parentid[j];
+      }
+    }
+  }
+
+  // x <- inv(D) * x
+  for (int i=0; i < nv; i++) {
+    x[i] *= qLDiagInv[i];  // x(i) /= L(i,i)
+  }
+
+  // x <- inv(L) * x; skip simple
+  for (int i=0; i < nv; i++) {
+    if (!m->dof_simplenum[i]) {
+      // init
+      int Madr_ij = dof_Madr[i]+1;
+      int j = dof_parentid[i];
+
+      // traverse ancestors backwards
+      tmp = x[i];
+      while (j>= 0) {
+        tmp -= qLD[Madr_ij++]*x[j];             // x(i) -= L(i,j) * x(j)
+
+        // advance to parent
+        j = dof_parentid[j];
+      }
+      x[i] = tmp;
+    }
+  }
+}
+
+void solveM_baseline(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y) {
+  solveLD_baseline(m, x, y, d->qLD, d->qLDiagInv);
+}
+
+// ----------------------------- benchmark ------------------------------------
+
+static void BM_solveLD(benchmark::State& state, bool new_function) {
+  static mjModel* m = LoadModelFromPath("composite/cloth.xml");
+  mjData* d = mj_makeData(m);
+
+  // warm-up rollout to get a typcal state
+  for (int i=0; i < kNumWarmupSteps; i++) {
+    AddCtrlNoise(m, d, i);
+    mj_step(m, d);
+  }
+
+  // allocate gadient
+  mjMARKSTACK;
+  mjtNum *grad = mj_stackAlloc(d, m->nv);
+  mjtNum *Ma = mj_stackAlloc(d, m->nv);
+  mjtNum *res = 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];
+  }
+
+  // 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);
+
+  // reset state, benchmark subsequent kNumBenchmarkSteps steps
+  while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
+    mju_copy(d->qpos, qpos.data(), m->nq);
+    mju_copy(d->qvel, qvel.data(), m->nv);
+    mju_copy(d->act, act.data(), m->na);
+    mju_copy(d->qacc_warmstart, warmstart.data(), m->nv);
+
+    for (int i=0; i < kNumBenchmarkSteps; i++) {
+      if (new_function) {
+        mj_solveM(m, d, res, grad, 1);
+      } else {
+        solveM_baseline(m, d, res, grad);
+      }
+    }
+  }
+
+  // finalize
+  mjFREESTACK;
+  mj_deleteData(d);
+  state.SetItemsProcessed(state.iterations());
+}
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_new(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, true);
+}
+BENCHMARK(BM_solveLD_new);
+
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_old(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_solveLD(state, false);
+}
+BENCHMARK(BM_solveLD_old);
+
+}  // namespace
+}  // namespace mujoco

test/benchmark/step_benchmark_test.cc

@@ -28,53 +28,23 @@
 namespace mujoco {
 namespace {
 
-using ::testing::NotNull;
-
 // number of steps to roll out before benhmarking
 static const int kNumWarmupSteps = 500;
 
 // number of steps to benchmark
 static const int kNumBenchmarkSteps = 50;
 
-
 // copy array into vector
 std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
   return std::vector<mjtNum>(array, array + n);
 }
 
-static void add_ctrl_noise(const mjModel* m, mjData* d, int step) {
-  for (int i = 0; i < m->nu; i++) {
-    mjtNum center = 0.0;
-    mjtNum radius = 1.0;
-    mjtNum* range = m->actuator_ctrlrange + 2 * i;
-    if (m->actuator_ctrllimited[i]) {
-      center = (range[1] + range[0]) / 2;
-      radius = (range[1] - range[0]) / 2;
-    }
-    radius *= 0.01;
-    d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1);
-  }
-}
-
-static void assert_model_not_null(mjModel* model,
-                                  const std::array<char, 1024>& error) {
-  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
-}
-
-static mjModel* load_model(const char* model_path) {
-  const std::string xml_path = GetModelPath(model_path);
-  std::array<char, 1024> error;
-  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error.data(), error.size());
-  assert_model_not_null(model, error);
-  return model;
-}
-
 static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
   mjData* data = mj_makeData(model);
 
   // warm-up rollout to get a typcal state
   for (int i=0; i < kNumWarmupSteps; i++) {
-    add_ctrl_noise(model, data, i);
+    AddCtrlNoise(model, data, i);
     mj_step(model, data);
   }
   // save state
@@ -91,7 +61,7 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
     mju_copy(data->qacc_warmstart, warmstart.data(), model->nv);
 
     for (int i=0; i < kNumBenchmarkSteps; i++) {
-      add_ctrl_noise(model, data, i+kNumWarmupSteps);
+      AddCtrlNoise(model, data, i+kNumWarmupSteps);
       mj_step(model, data);
     }
   }
@@ -107,21 +77,21 @@ static void run_step_benchmark(const mjModel* model, benchmark::State& state) {
 
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepCloth(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  static mjModel* model = load_model("composite/cloth.xml");
+  static mjModel* model = LoadModelFromPath("composite/cloth.xml");
   run_step_benchmark(model, state);
 }
 BENCHMARK(BM_StepCloth);
 
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  static mjModel* model = load_model("humanoid/humanoid.xml");
+  static mjModel* model = LoadModelFromPath("humanoid/humanoid.xml");
   run_step_benchmark(model, state);
 }
 BENCHMARK(BM_StepHumanoid);
 
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_StepHumanoid100(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  static mjModel* model = load_model("humanoid100/humanoid100.xml");
+  static mjModel* model = LoadModelFromPath("humanoid100/humanoid100.xml");
   run_step_benchmark(model, state);
 }
 BENCHMARK(BM_StepHumanoid100);

test/engine/engine_derivative_test.cc

@@ -289,7 +289,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
       Ac[i*nv + i]  = -m->jnt_stiffness[i];
       Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
     }
-    mj_solveLD(m, Ac, Ac, 2*nv, d->qH, d->qHDiagInv);
+    mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv);
 
     // A = [dt*Ac; Ac]
     mju_transpose(A, Ac, 2*nv, nv);
@@ -314,7 +314,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
   if (B) {
     mjtNum *Bc = mj_stackAlloc(d, nu*nv);
     mjtNum *BcT = mj_stackAlloc(d, nv*nu);
-    mj_solveLD(m, Bc, d->actuator_moment, nu, d->qH, d->qHDiagInv);
+    mju_copy(Bc, d->actuator_moment, nv*nu);
+    mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv);
     mju_transpose(BcT, Bc, nu, nv);
     mju_scl(B, BcT, dt*dt, nu*nv);
     mju_scl(B+nu*nv, BcT, dt, nu*nv);

test/fixture.cc

@@ -30,6 +30,8 @@
 namespace mujoco {
 namespace {
 
+using ::testing::NotNull;
+
 ABSL_CONST_INIT static absl::Mutex handlers_mutex(absl::kConstInit);
 static int guard_count ABSL_GUARDED_BY(handlers_mutex) = 0;
 
@@ -80,6 +82,20 @@ mjModel* LoadModelFromString(std::string_view xml, char* error,
   return m;
 }
 
+static void AssertModelNotNull(mjModel* model,
+                                  const std::array<char, 1024>& error) {
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+}
+
+mjModel* LoadModelFromPath(const char* model_path) {
+  const std::string xml_path = GetModelPath(model_path);
+  std::array<char, 1024> error;
+  mjModel* model = mj_loadXML(
+    xml_path.c_str(), nullptr, error.data(), error.size());
+  AssertModelNotNull(model, error);
+  return model;
+}
+
 const std::string GetFileContents(const char* path) {
   std::ifstream ifs;
   ifs.open(path, std::ifstream::in);
@@ -118,4 +134,18 @@ const std::string SaveAndReadXml(const mjModel* model) {
   return contents;
 }
 
+void AddCtrlNoise(const mjModel* m, mjData* d, int step) {
+  for (int i = 0; i < m->nu; i++) {
+    mjtNum center = 0.0;
+    mjtNum radius = 1.0;
+    mjtNum* range = m->actuator_ctrlrange + 2 * i;
+    if (m->actuator_ctrllimited[i]) {
+      center = (range[1] + range[0]) / 2;
+      radius = (range[1] - range[0]) / 2;
+    }
+    radius *= 0.01;
+    d->ctrl[i] = center + radius * (2 * mju_Halton(step, i + 2) - 1);
+  }
+}
+
 }  // namespace mujoco

test/fixture.h

@@ -17,6 +17,7 @@
 
 #include <gtest/gtest.h>
 #include <absl/strings/string_view.h>
+#include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 
 namespace mujoco {
@@ -53,8 +54,15 @@ const std::string GetModelPath(absl::string_view path);
 mjModel* LoadModelFromString(absl::string_view xml, char* error = nullptr,
                              int error_size = 0);
 
+// Returns a newly-allocated mjModel, loaded from the contents in model_path.
+// On failure it asserts that model is null.
+mjModel* LoadModelFromPath(const char* model_path);
+
 // Returns a string loaded from first saving the model given an input.
 const std::string SaveAndReadXml(const mjModel* model);
 
+// Adds control noise.
+void AddCtrlNoise(const mjModel* m, mjData* d, int step);
+
 }  // namespace mujoco
 #endif  // MUJOCO_TEST_FIXTURE_H_