Use diagonal index to speed up mju_solveLUSparse.

PiperOrigin-RevId: 704303282
Change-Id: I72c3fbcd31a5564fc792b73efe9d12a151b4d60f

src/engine/engine_forward.c

@@ -967,7 +967,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
-    mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_colind);
+    mju_solveLUSparse(qacc, d->qLU, qfrc, nv, d->D_rownnz, d->D_rowadr, d->D_diag, d->D_colind);
   }
 
   // IMPLICITFAST

src/engine/engine_util_solve.c

@@ -651,41 +651,31 @@ void mju_factorLUSparse(mjtNum* LU, int n, int* scratch,
 
 // solve mat*res=vec given LU factorization of mat
 void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n,
-                       const int* rownnz, const int* rowadr, const int* colind) {
-  //------------------ solve (U+I)*res = vec
+                       const int* rownnz, const int* rowadr, const int* diag, const int* colind) {
+  // solve (U+I)*res = vec
   for (int i=n-1; i >= 0; i--) {
     // init: diagonal of (U+I) is 1
     res[i] = vec[i];
 
-    // res[i] -= sum_k>i res[k]*LU(i,k)
-    int j = rownnz[i] - 1;
-    while (colind[rowadr[i]+j] > i) {
-      res[i] -= res[colind[rowadr[i]+j]] * LU[rowadr[i]+j];
-      j--;
-    }
-
-    // make sure j points to diagonal
-    if (colind[rowadr[i]+j] != i) {
-      mjERROR("diagonal of U not reached");
+    int d1 = diag[i]+1;
+    int nnz = rownnz[i] - d1;
+    if (nnz > 0) {
+      int adr = rowadr[i] + d1;
+      res[i] -= mju_dotSparse(LU+adr, res, nnz, colind+adr, /*flg_unc1=*/0);
     }
   }
 
   //------------------ solve L*res(new) = res
   for (int i=0; i < n; i++) {
     // res[i] -= sum_k<i res[k]*LU(i,k)
-    int j = 0;
-    while (colind[rowadr[i]+j] < i) {
-      res[i] -= res[colind[rowadr[i]+j]] * LU[rowadr[i]+j];
-      j++;
+    int d = diag[i];
+    int adr = rowadr[i];
+    if (d > 0) {
+      res[i] -= mju_dotSparse(LU+adr, res, d, colind+adr, /*flg_unc1=*/0);
     }
 
     // divide by diagonal element of L
-    res[i] /= LU[rowadr[i]+j];
-
-    // make sure j points to diagonal
-    if (colind[rowadr[i]+j] != i) {
-      mjERROR("diagonal of L not reached");
-    }
+    res[i] /= LU[adr + d];
   }
 }
 

src/engine/engine_util_solve.h

@@ -82,7 +82,7 @@ void mju_factorLUSparse(mjtNum *LU, int n, int* scratch,
 
 // solve mat*res=vec given LU factorization of mat
 void mju_solveLUSparse(mjtNum *res, const mjtNum *LU, const mjtNum* vec, int n,
-                       const int *rownnz, const int *rowadr, const int *colind);
+                       const int *rownnz, const int *rowadr, const int* diag, const int *colind);
 
 // eigenvalue decomposition of symmetric 3x3 matrix
 MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);