CSR implementation of mj_solveLD.

PiperOrigin-RevId: 704316200
Change-Id: Ibaff0284e40b3ebbe43bb489b6211ce739270e27

src/engine/engine_core_smooth.c

@@ -1575,6 +1575,38 @@ void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
 }
 
 
+
+// in-place sparse backsubstitution:  x = inv(L'*D*L)*x
+//  like mj_solveLD, but using the CSR representation of L
+void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
+                 const int* rownnz, const int* rowadr, const int* diag, const int* colind) {
+  // x <- L^-T x
+  for (int i=nv-2; i >= 0; i--) {
+    int d1 = diag[i] + 1;
+    int nnz = rownnz[i] - d1;
+    if (nnz > 0) {
+      int adr = rowadr[i] + d1;
+      x[i] -= mju_dotSparse(qLDs+adr, x, nnz, colind+adr, /*flg_unc1=*/0);
+    }
+  }
+
+  // x(i) /= D(i,i)
+  for (int i=0; i < nv; i++) {
+    x[i] *= qLDiagInv[i];
+  }
+
+  // x <- L^-1 x
+  for (int i=1; i < nv; i++) {
+    int d = diag[i];
+    if (d > 0) {
+      int adr = rowadr[i];
+      x[i] -= mju_dotSparse(qLDs+adr, x, d, colind+adr, /*flg_unc1=*/0);
+    }
+  }
+}
+
+
+
 // 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) {

src/engine/engine_core_smooth.h

@@ -55,10 +55,15 @@ MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD,
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
 
-// sparse backsubstitution:  x = inv(L'*D*L)*y
+// sparse backsubstitution:  x = inv(L'*D*L)*x
 MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
                       const mjtNum* qLD, const mjtNum* qLDiagInv);
 
+// in-place sparse backsubstitution:  x = inv(L'*D*L)*x
+//  like mj_solveLD, but using the CSR representation of L
+MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
+                       const int* rownnz, const int* rowadr, const int* diag, const int* colind);
+
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 

test/benchmark/engine_core_smooth_benchmark_test.cc

@@ -19,85 +19,28 @@
 #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 roll out before benhmarking
+// number of steps to roll out before benchmarking
 static const int kNumWarmupSteps = 200;
 
 // number of steps to benchmark
 static const int kNumBenchmarkSteps = 50;
 
-// ----------------------------- 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 void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
-  static mjModel* m = LoadModelFromPath("plugin/elasticity/coil.xml");
+  static mjModel* m;
+  if (coil) {
+    m = LoadModelFromPath("plugin/elasticity/coil.xml");
+  } else {
+    m = LoadModelFromPath("humanoid/humanoid100.xml");
+  }
+
   mjData* d = mj_makeData(m);
 
   // warm-up rollout to get a typical state
@@ -117,24 +60,21 @@ static void BM_solveLD(benchmark::State& state, bool new_function) {
     grad[i] = Ma[i] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
   }
 
-  // save state
+  // CSR matrix
-  std::vector<mjtNum> qpos = AsVector(d->qpos, m->nq);
+  mjtNum* LDs = mj_stackAllocNum(d, m->nC);
-  std::vector<mjtNum> qvel = AsVector(d->qvel, m->nv);
+  for (int i=0; i < m->nC; i++) {
-  std::vector<mjtNum> act = AsVector(d->act, m->na);
+    LDs[i] = d->qLD[d->mapM2C[i]];
-  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) {
+      if (featherstone) {
         mj_solveM(m, d, res, grad, 1);
       } else {
-        solveM_baseline(m, d, res, grad);
+        mju_copy(res, grad, m->nv);
+        mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
+                    d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind);
       }
     }
   }
@@ -145,17 +85,29 @@ static void BM_solveLD(benchmark::State& state, bool new_function) {
   state.SetItemsProcessed(state.iterations());
 }
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_new(benchmark::State& state) {
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_FS(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  BM_solveLD(state, true);
+  BM_solveLD(state, /*featherstone=*/true, /*coil=*/true);
 }
-BENCHMARK(BM_solveLD_new);
+BENCHMARK(BM_solveLD_COIL_FS);
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_old(benchmark::State& state) {
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_solveLD_COIL_CSR(benchmark::State& state) {
   MujocoErrorTestGuard guard;
-  BM_solveLD(state, false);
+  BM_solveLD(state, /*featherstone=*/false, /*coil=*/true);
 }
-BENCHMARK(BM_solveLD_old);
+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

test/engine/engine_core_smooth_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_core_smooth.c.
 
 #include "src/engine/engine_core_smooth.h"
+#include "src/engine/engine_util_sparse.h"
 
 #include <string>
 #include <string_view>
@@ -31,6 +32,7 @@
 namespace mujoco {
 namespace {
 
+using ::std::vector;
 using ::testing::Each;
 using ::testing::ElementsAre;
 using ::testing::Eq;
@@ -261,7 +263,7 @@ TEST_F(CoreSmoothTest, EqualityBodySite) {
   while (data->time < 0.1) {
     mj_step(model, data);
   }
-  std::vector<mjtNum> sdata = AsVector(data->sensordata, model->nsensordata);
+  vector<mjtNum> sdata = AsVector(data->sensordata, model->nsensordata);
 
   // reset
   mj_resetData(model, data);
@@ -405,5 +407,79 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
   mj_deleteModel(model);
 }
 
+TEST_F(CoreSmoothTest, SolveLD2) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <default>
+      <geom type="capsule" size="0.1"/>
+      <joint axis="0 1 0"/>
+    </default>
+
+    <worldbody>
+      <body>
+        <geom fromto="0 0 0 0 0 1"/>
+        <joint/>
+        <body pos="0 0 1">
+          <geom fromto="0 0 0 1 0 1"/>
+          <joint/>
+        </body>
+        <body pos="0 0 1">
+          <geom fromto="0 0 0 -1 0 1"/>
+          <joint/>
+          <body pos="-1 0 1">
+            <geom fromto="0 0 0 1 0 1"/>
+            <joint/>
+          </body>
+          <body pos="-1 0 1">
+            <geom fromto="0 0 0 -1 0 1"/>
+            <joint/>
+          </body>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  mjModel* m = LoadModelFromString(xml);
+  mjData* d = mj_makeData(m);
+  mj_forward(m, d);
+
+  int nv = m->nv;
+  int nC = m->nC;
+
+  // copy LD into LDs: CSR format
+  vector<mjtNum> LDs(nC);
+  for (int i=0; i < nC; i++) {
+    LDs[i] = d->qLD[d->mapM2C[i]];
+  }
+
+  // compare LD and LDs densified matrices
+  vector<mjtNum> LDdense(nv*nv);
+  mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv,
+                   d->C_rownnz, d->C_rowadr, d->C_colind);
+  vector<mjtNum> LDdense2(nv*nv);
+  mj_fullM(m, LDdense2.data(), d->qLD);
+
+  // expect dense matrices to match exactly
+  for (int i=0; i < nv*nv; i++) EXPECT_EQ(LDdense[i], LDdense2[i]);
+
+  // compare LD and LDs vector solve
+  vector<mjtNum> vec(nv);
+  vector<mjtNum> vec2(nv);
+  for (int i=0; i < nv; i++) vec[i] = vec2[i] = 20 + 30*i;
+  for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0;
+
+  mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv);
+  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv,
+              d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind);
+
+  // expect vectors to match up to floating point precision
+  for (int i=0; i < nv; i++) {
+    EXPECT_FLOAT_EQ(vec[i], vec2[i]);
+  }
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 }  // namespace
 }  // namespace mujoco