From 537fc2ff45df0303b74381488007e52430216046 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Mon, 9 Dec 2024 08:52:31 -0800
Subject: [PATCH] Use diagonal index to speed up mju_solveLUSparse.

PiperOrigin-RevId: 704303282
Change-Id: I72c3fbcd31a5564fc792b73efe9d12a151b4d60f
---
 src/engine/engine_forward.c    |  2 +-
 src/engine/engine_util_solve.c | 34 ++++++++++++----------------------
 src/engine/engine_util_solve.h |  2 +-
 3 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index e77639e5..181db4d3 100644
--- a/src/engine/engine_forward.c
+++ b/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
diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c
index 10b50157..4a455359 100644
--- a/src/engine/engine_util_solve.c
+++ b/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];
   }
 }
 
diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h
index 308e2d77..91ea13cc 100644
--- a/src/engine/engine_util_solve.h
+++ b/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]);