Allow CSR back-substitution to handle multiple vectors.

PiperOrigin-RevId: 713229673
Change-Id: I7a5b43fe966cf9e482bd41e2eea6c30dd3ffa1d4

Author: yuvaltassa

src/engine/engine_core_smooth.c

@@ -1607,39 +1607,88 @@ 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* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
+void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
                  const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
-  // x <- L^-T x
-  for (int i=nv-1; i > 0; i--) {
-    // skip diagonal (simple) rows, exploit sparsity of input vector
-    if (diagnum[i] || x[i] == 0) {
-      continue;
+  // single vector
+  if (n == 1) {
+    // x <- L^-T x
+    for (int i=nv-1; i > 0; i--) {
+      // skip diagonal rows, zero elements in input vector
+      mjtNum x_i = x[i];
+      if (x_i == 0 || diagnum[i]) {
+        continue;
+      }
+
+      int start = rowadr[i];
+      int end = start + rownnz[i] - 1;
+      for (int adr=start; adr < end; adr++) {
+        x[colind[adr]] -= qLDs[adr] * x_i;
+      }
     }
 
-    int d = rownnz[i] - 1;
-    int adr_i = rowadr[i];
-    mjtNum x_i = x[i];
-    for (int j=0; j < d; j++) {
-      int adr = adr_i + j;
-      x[colind[adr]] -= qLDs[adr] * x_i;
+    // x <- D^-1 x
+    for (int i=0; i < nv; i++) {
+      x[i] *= qLDiagInv[i];
     }
-  }
 
-  // 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++) {
+      // skip diagonal rows
+      if (diagnum[i]) {
+        i += diagnum[i] - 1;  // iterating forward: skip ahead, adjust i
+        continue;
+      }
+
+      int adr = rowadr[i];
+      x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0);
+    }
   }
 
-  // x <- L^-1 x
-  for (int i=1; i < nv; i++) {
-    // skip diagonal (simple) rows
-    if (diagnum[i]) {
-      i += diagnum[i] - 1;  // when iterating forward we can skip ahead
-      continue;
+  // multiple vectors
+  else {
+    // x <- L^-T x
+    for (int i=nv-1; i > 0; i--) {
+      // skip diagonal rows
+      if (diagnum[i]) {
+        continue;
+      }
+
+      int start = rowadr[i];
+      int end = start + rownnz[i] - 1;
+      for (int adr=start; adr < end; adr++) {
+        int j = colind[adr];
+        mjtNum val = qLDs[adr];
+        for (int offset=0; offset < n*nv; offset+=nv) {
+          mjtNum x_i;
+          if ((x_i = x[i+offset])) {
+            x[j+offset] -= val * x_i;
+          }
+        }
+      }
+    }
+
+    // x <- D^-1 x
+    for (int i=0; i < nv; i++) {
+      mjtNum invD_i = qLDiagInv[i];
+      for (int offset=0; offset < n*nv; offset+=nv) {
+        x[i+offset] *= invD_i;
+      }
     }
 
-    int adr = rowadr[i];
-    x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0);
+    // x <- L^-1 x
+    for (int i=1; i < nv; i++) {
+      // skip diagonal rows
+      if (diagnum[i]) {
+        i += diagnum[i] - 1;  // iterating forward: skip ahead, adjust i
+        continue;
+      }
+
+      int adr = rowadr[i];
+      int d = rownnz[i] - 1;
+      for (int offset=0; offset < n*nv; offset+=nv) {
+        x[i+offset] -= mju_dotSparse(qLDs+adr, x+offset, d, colind+adr, /*flg_unc1=*/0);
+      }
+    }
   }
 }
 

src/engine/engine_core_smooth.h

@@ -64,8 +64,8 @@ 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,
+//  handle n vectors at once
+MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
                        const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
 
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d

test/benchmark/inertia_benchmark_test.cc

@@ -72,7 +72,7 @@ static void BM_solve(benchmark::State& state, SolveType type) {
           mj_factorIs(LDs, d->qLDiagInv, m->nv,
                       d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
           mju_copy(res, vec, m->nv);
-          mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
+          mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
                       d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
       }
     }

test/benchmark/solveLD_benchmark_test.cc

@@ -63,7 +63,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
         mj_solveM(m, d, res, vec, 1);
       } else {
         mju_copy(res, vec, m->nv);
-        mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
+        mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
                     d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
       }
     }

test/engine/engine_core_smooth_test.cc

@@ -495,7 +495,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
   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,
+  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, 1,
               d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 
   // expect vectors to match up to floating point precision
@@ -507,6 +507,44 @@ TEST_F(CoreSmoothTest, SolveLDs) {
   mj_deleteModel(m);
 }
 
+TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
+  const std::string xml_path = GetTestDataFilePath(kInertiaPath);
+  char error[1024];
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
+
+  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 n LD and LDs vector solve
+  int n = 3;
+  vector<mjtNum> vec(nv*n);
+  vector<mjtNum> vec2(nv*n);
+  for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i;
+  for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0;
+
+  mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv);
+  mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
+              d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+
+  // expect vectors to match up to floating point precision
+  for (int i=0; i < nv*n; i++) {
+    EXPECT_FLOAT_EQ(vec[i], vec2[i]);
+  }
+
+  mj_deleteData(d);
+  mj_deleteModel(m);
+}
+
 TEST_F(CoreSmoothTest, FactorIs) {
   const std::string xml_path = GetTestDataFilePath(kInertiaPath);
   char error[1024];