Remove unnecessary diagnum argument in mj_solveLD and mj_factorI

PiperOrigin-RevId: 758669364
Change-Id: Icfefc4a7a1e7d28da373725358ccf9d5f589f689

src/engine/engine_core_smooth.c

@@ -1653,7 +1653,7 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
 void mj_factorM(const mjModel* m, mjData* d) {
   TM_START;
   mju_copy(d->qLD, d->M, m->nC);
-  mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+  mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind);
   TM_ADD(mjTIMER_POS_INERTIA);
 }
 
@@ -1661,32 +1661,28 @@ void mj_factorM(const mjModel* m, mjData* d) {
 
 // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
 void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv,
-                const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
+                const int* rownnz, const int* rowadr, const int* colind) {
   // backward loop over rows
   for (int k=nv-1; k >= 0; k--) {
     // get row k's address, diagonal index, inverse diagonal value
-    int rowadr_k = rowadr[k];
+    int start = rowadr[k];
-    int diag_k = rowadr_k + rownnz[k] - 1;
+    int diag = rownnz[k] - 1;
-    mjtNum invD = 1 / mat[diag_k];
+    int end = start + diag;
+    mjtNum invD = 1 / mat[end];
     if (diaginv) diaginv[k] = invD;
 
-    // skip if simple
+    // update triangle above row k
-    if (diagnum[k]) {
+    for (int adr=end - 1; adr >= start; adr--) {
-      continue;
-    }
-
-    // update triangle above row k, inclusive
-    for (int adr=diag_k - 1; adr >= rowadr_k; adr--) {
       // tmp = L(k, i) / L(k, k)
       mjtNum tmp = mat[adr] * invD;
 
       // update row i < k:  L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k)
       int i = colind[adr];
-      mju_addToScl(mat + rowadr[i], mat + rowadr_k, -tmp, rownnz[i]);
+      mju_addToScl(mat + rowadr[i], mat + start, -tmp, rownnz[i]);
-
-      // update ith element of row k:  L(k, i) /= L(k, k)
-      mat[adr] = tmp;
     }
+
+    // update row k:  L(k, :) /= L(k, k)
+    mju_scl(mat + start, mat + start, invD, diag);
   }
 }
 
@@ -1807,11 +1803,11 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n,
 
 // in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
-                const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
+                const int* rownnz, const int* rowadr, const int* colind) {
   // x <- L^-T x
   for (int i=nv-1; i > 0; i--) {
     // skip diagonal rows
-    if (diagnum[i]) {
+    if (rownnz[i] == 1) {
       continue;
     }
 
@@ -1862,8 +1858,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv,
   // x <- L^-1 x
   for (int i=1; i < nv; i++) {
     // skip diagonal rows
-    if (diagnum[i]) {
+    if (rownnz[i] == 1) {
-      i += diagnum[i] - 1;  // iterating forward: skip ahead, adjust i
       continue;
     }
 
@@ -1895,7 +1890,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
     mju_copy(x, y, n*m->nv);
   }
   mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
-             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+             d->M_rownnz, d->M_rowadr, d->M_colind);
 }
 
 

src/engine/engine_core_smooth.h

@@ -63,7 +63,7 @@ MJAPI void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M,
 
 // sparse L'*D*L factorizaton of inertia-like matrix
 MJAPI void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv,
-                      const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
+                      const int* rownnz, const int* rowadr, const int* colind);
 
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -75,7 +75,7 @@ MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n,
 // in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 //  handle n vectors at once
 MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
-                      const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
+                      const int* rownnz, const int* rowadr, 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);

src/engine/engine_forward.c

@@ -873,14 +873,14 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
       }
 
       // factorize in-place
-      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind);
     }
 
     // solve
     mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
     mju_copy(qacc, qfrc, m->nv);
     mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
-               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+               d->M_rownnz, d->M_rowadr, d->M_colind);
   }
 
   // advance state and time
@@ -1056,13 +1056,13 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mju_gather(d->qH, MhB, d->mapM2M, nC);
 
       // factorize in-place
-      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+      mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
     mju_copy(qacc, qfrc, nv);
     mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
-               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+               d->M_rownnz, d->M_rowadr, d->M_colind);
 
   } else {
     mjERROR("integrator must be implicit or implicitfast");

src/engine/engine_solver.c

@@ -1071,7 +1071,7 @@ static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) {
   else {
     mju_copy(ctx->Mgrad, ctx->grad, nv);
     mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1,
-               ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
+               ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind);
   }
 }
 

test/benchmark/factorI_benchmark_test.cc

@@ -59,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
       } else {
         mju_copy(d->qLD, M, m->nC);
         mj_factorI(d->qLD, d->qLDiagInv, m->nv,
-                   d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+                   d->M_rownnz, d->M_rowadr, d->M_colind);
       }
     }
   }

test/benchmark/inertia_benchmark_test.cc

@@ -73,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
         case SolveType::kCsr:
           mju_copy(d->qLD, M, m->nC);
           mj_factorI(d->qLD, d->qLDiagInv, m->nv,
-                     d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+                     d->M_rownnz, d->M_rowadr, d->M_colind);
           mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
-                     d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+                     d->M_rownnz, d->M_rowadr, d->M_colind);
       }
     }
   }

test/benchmark/solveLD_benchmark_test.cc

@@ -64,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
         mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv);
       } else {
         mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
-                   d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+                   d->M_rownnz, d->M_rowadr, d->M_colind);
       }
     }
   }

test/engine/engine_core_smooth_test.cc

@@ -722,7 +722,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
 
   mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
   mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
-             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+             d->M_rownnz, d->M_rowadr, d->M_colind);
 
   // expect vectors to match up to floating point precision
   for (int i=0; i < nv; i++) {
@@ -757,7 +757,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
 
   mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
   mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
-             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+             d->M_rownnz, d->M_rowadr, d->M_colind);
 
   // expect vectors to match up to floating point precision
   for (int i=0; i < nv*n; i++) {
@@ -795,7 +795,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
 
   mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
   mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
-             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+             d->M_rownnz, d->M_rowadr, d->M_colind);
 
   // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
   for (int i=0; i < n; i++) {
@@ -834,7 +834,7 @@ TEST_F(CoreSmoothTest, FactorIs) {
   vector<mjtNum> qLDiagInv(nv, 0);
 
   mj_factorI(qLD.data(), qLDiagInv.data(), nv,
-             d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+             d->M_rownnz, d->M_rowadr, d->M_colind);
 
   // expect outputs to match to floating point precision
   EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));

test/engine/engine_derivative_test.cc

@@ -436,7 +436,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
       Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
     }
     mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
-               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+               d->M_rownnz, d->M_rowadr, d->M_colind);
 
     // A = [dt*Ac; Ac]
     mju_transpose(A, Ac, 2*nv, nv);
@@ -464,7 +464,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
     mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
                      d->moment_rowadr, d->moment_colind);
     mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu,
-               d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
+               d->M_rownnz, d->M_rowadr, d->M_colind);
     mju_transpose(BcT, Bc, nu, nv);
     mju_scl(B, BcT, dt*dt, nu*nv);
     mju_scl(B+nu*nv, BcT, dt, nu*nv);