Switch mj_solveM_island to use CSR representation

PiperOrigin-RevId: 717684720
Change-Id: I7d74a5d4aef4aa5b0c0acddf88c5b8591d63e88d

src/engine/engine_core_smooth.c

@@ -1701,7 +1701,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
 
 // in-place sparse backsubstitution for one island:  x = inv(L'*D*L)*x
 //  L is in lower triangle of qLD; D is on diagonal of qLD
-void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
+void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) {
   // if no islands, call mj_solveLD
   const mjtNum* qLD = d->qLD;
   const mjtNum* qLDiagInv = d->qLDiagInv;
@@ -1710,10 +1710,19 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
     return;
   }
 
-  // local constants: general
+  // local copies of key variables
-  const int* Madr = m->dof_Madr;
+  const int* rownnz = d->C_rownnz;
-  const int* parentid = m->dof_parentid;
+  const int* rowadr = d->C_rowadr;
-  const int* simplenum = m->dof_simplenum;
+  const int* colind = d->C_colind;
+  const int* diagnum = m->dof_simplenum;
+
+  // temporary: make local CSR version of qLD
+  int nC = m->nC;
+  mj_markStack(d);
+  mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum);
+  for (int i=0; i < nC; i++) {
+    qLDs[i] = d->qLD[d->mapM2C[i]];
+  }
 
   // local constants: island specific
   int ndof = d->island_dofnum[island];
@@ -1723,18 +1732,12 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
   // x <- inv(L') * x; skip simple, exploit sparsity of input vector
   for (int k=ndof-1; k >= 0; k--) {
     int i = dofind[k];
-    if (!simplenum[i] && x[k]) {
+    mjtNum x_k;
-      // init
+    if (!diagnum[i] && (x_k = x[k])) {
-      int Madr_ij = Madr[i]+1;
+      int start = rowadr[i];
-      int j = parentid[i];
+      int end = start + rownnz[i] - 1;
-
+      for (int adr=end-1; adr >= start; adr--) {
-      // traverse ancestors backwards
+        x[islandind[colind[adr]]] -= qLDs[adr] * x_k;
-      // read directly from x[l] since j cannot be a parent of itself
-      while (j >= 0) {
-        x[islandind[j]] -= qLD[Madr_ij++]*x[k];         // x(j) -= L(i,j) * x(i)
-
-        // advance to parent
-        j = parentid[j];
       }
     }
   }
@@ -1747,21 +1750,20 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
   // x <- inv(L) * x; skip simple
   for (int k=0; k < ndof; k++) {
     int i = dofind[k];
-    if (!simplenum[i]) {
-      // init
-      int Madr_ij = Madr[i]+1;
-      int j = parentid[i];
 
-      // traverse ancestors backwards
+    // skip diagonal rows
-      // write directly in x[i] since i cannot be a parent of itself
+    if (diagnum[i]) {
-      while (j >= 0) {
+      continue;
-        x[k] -= qLD[Madr_ij++]*x[islandind[j]];             // x(i) -= L(i,j) * x(j)
+    }
 
-        // advance to parent
+    int start = rowadr[i];
-        j = parentid[j];
+    int end = start + rownnz[i] - 1;
-      }
+    for (int adr=end-1; adr >= start; adr--) {
+      x[k] -= x[islandind[colind[adr]]] * qLDs[adr];
     }
   }
+
+  mj_freeStack(d);
 }
 
 

src/engine/engine_core_smooth.h

@@ -71,8 +71,9 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i
 // 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);
 
+// TODO(tassa): Restore mjData const-ness.
 // sparse backsubstitution for one island:  x = inv(L'*D*L)*x, use factorization in d
-MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
+MJAPI void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* x, int island);
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,

src/engine/engine_solver.c

@@ -895,9 +895,9 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
 }
 
 
-
+// TODO(tassa): Restore mjData const-ness.
 // update grad, Mgrad
-static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx) {
+static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
   int nv = ctx->nv;
   const int* dofind = ctx->dofind;
 

test/engine/engine_core_smooth_test.cc

@@ -353,58 +353,63 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) {
 
 static const char* const kIlslandEfcPath =
     "engine/testdata/island/island_efc.xml";
+static const char* const kModelPath =
+    "testdata/model.xml";
 
 TEST_F(CoreSmoothTest, SolveMIsland) {
-  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  for (auto model_path : {kModelPath, kIlslandEfcPath}) {
-  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+    const std::string xml_path = GetTestDataFilePath(model_path);
-  mjData* data = mj_makeData(model);
+    mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
-  int nv = model->nv;
+    mjData* data = mj_makeData(model);
+    int nv = model->nv;
 
-  // allocate vec, fill with arbitrary values, copy to sol
+    // allocate vec, fill with arbitrary values, copy to sol
-  mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
+    mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
-  mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
+    mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
-  for (int i=0; i < nv; i++) {
+    for (int i=0; i < nv; i++) {
-    vec[i] = 0.2 + 0.3*i;
+      vec[i] = 0.2 + 0.3*i;
-  }
+    }
-  mju_copy(res, vec, nv);
+    mju_copy(res, vec, nv);
 
-  // simulate for 0.2 seconds
+    if (model->nkey > 0) mj_resetDataKeyframe(model, data, 0);
-  mj_resetData(model, data);
-  while (data->time < 0.2) {
-    mj_step(model, data);
-  }
-  mj_forward(model, data);
 
-  // divide by mass matrix: sol = M^-1 * vec
+    for (int i=0; i < 6; i++) {
-  mj_solveM(model, data, res, res, 1);
+      mj_step(model, data);
-
-  // iterate over islands
-  for (int i=0; i < data->nisland; i++) {
-    // allocate dof vectors for island
-    int dofnum = data->island_dofnum[i];
-    mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
-
-    // copy values into sol_i
-    int* dofind = data->island_dofind + data->island_dofadr[i];
-    for (int j=0; j < dofnum; j++) {
-      res_i[j] = vec[dofind[j]];
     }
 
-    // divide by mass matrix, for this island
+    mj_forward(model, data);
-    mj_solveM_island(model, data, res_i, i);
 
-    // expect corresponding values to match
+    // divide by mass matrix: sol = M^-1 * vec
-    for (int j=0; j < dofnum; j++) {
+    mj_solveM(model, data, res, res, 1);
-      EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12));
+
+    // iterate over islands
+    for (int i=0; i < data->nisland; i++) {
+      // allocate dof vectors for island
+      int dofnum = data->island_dofnum[i];
+      mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
+
+      // copy values into sol_i
+      int* dofind = data->island_dofind + data->island_dofadr[i];
+      for (int j=0; j < dofnum; j++) {
+        res_i[j] = vec[dofind[j]];
+      }
+
+      // divide by mass matrix, for this island
+      mj_solveM_island(model, data, res_i, i);
+
+      // expect corresponding values to match
+      for (int j=0; j < dofnum; j++) {
+        EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14));
+      }
+
+      mju_free(res_i);
     }
 
-    mju_free(res_i);
+    mju_free(res);
+    mju_free(vec);
+    mj_deleteData(data);
+    mj_deleteModel(model);
   }
-
-  mju_free(res);
-  mju_free(vec);
-  mj_deleteData(data);
-  mj_deleteModel(model);
 }
 
 static const char* const kInertiaPath = "engine/testdata/inertia.xml";