Refactor mju_cholUpdateSparse to use a dense accumulator, 1.5-2x function speedup.

PiperOrigin-RevId: 846754379
Change-Id: Ie599658db2907a26095072d97d07b99a693050dc

src/engine/engine_solver.c

@@ -1685,8 +1685,6 @@ static void HessianCone(mjData* d, mjCGContext* ctx) {
 
   // storage for L'*J
   mjtNum* LTJ = mjSTACKALLOC(d, 6*nv, mjtNum);
-  mjtNum* LTJ_row = mjSTACKALLOC(d, nv, mjtNum);
-  int* LTJ_ind = mjSTACKALLOC(d, nv, int);
 
   // add contributions
   for (int i=0; i < nefc; i++) {
@@ -1713,13 +1711,9 @@ static void HessianCone(mjData* d, mjCGContext* ctx) {
 
         // update
         for (int r=0; r < dim; r++) {
-          // copy data for this row
-          mju_copy(LTJ_row, LTJ+r*nnz, nnz);
-          mju_copyInt(LTJ_ind, ctx->J_colind+ctx->J_rowadr[i+r], nnz);
-
-          // update
-          mju_cholUpdateSparse(ctx->Lcone, LTJ_row, nv, 1,
-                               ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, LTJ_ind, d);
+          mju_cholUpdateSparse(ctx->Lcone, LTJ+r*nnz, nv, 1,
+                               ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz,
+                               ctx->J_colind+ctx->J_rowadr[i+r], d);
         }
       }
 
@@ -1758,7 +1752,6 @@ static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate)
 
   // local space
   mjtNum* vec = mjSTACKALLOC(d, nv, mjtNum);
-  int* vec_ind = mjSTACKALLOC(d, nv, int);
 
   // clear update counter
   ctx->nupdate = 0;
@@ -1784,14 +1777,13 @@ static void HessianIncremental(mjData* d, mjCGContext* ctx, const int* oldstate)
         // get nnz and adr of row i
         const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i];
 
-        // scale vec, copy colind
+        // scale vec
         mju_scl(vec, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz);
-        mju_copyInt(vec_ind, ctx->J_colind+adr, nnz);
 
         // sparse update or downdate
         rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update,
-                                    ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, vec_ind,
-                                    d);
+                                    ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz,
+                                    ctx->J_colind+adr, d);
       } else {
         mju_scl(vec, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv);
         rank = mju_cholUpdate(ctx->L, vec, nv, flag_update);

src/engine/engine_util_solve.c

@@ -415,43 +415,62 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 
 // sparse reverse-order Cholesky rank-one update: L'*L +/- x*x'; return rank
 //  x is sparse, change in sparsity pattern of mat is not allowed
-int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
-                         const int* rownnz, const int* rowadr, const int* colind,
-                         int x_nnz, int* x_ind,
+int mju_cholUpdateSparse(mjtNum* restrict mat, const mjtNum* restrict x, int n, int flg_plus,
+                         const int* restrict rownnz, const int* restrict rowadr,
+                         const int* restrict colind, int x_nnz, const int* restrict x_ind,
                          mjData* d) {
+  // early return if x is empty
+  if (x_nnz == 0) {
+    return n;
+  }
+
+  // get starting row: last non-zero entry in x
+  int start = x_ind[x_nnz - 1];
+
+  // allocate dense accumulator for x
   mj_markStack(d);
-  int* buf_ind = mjSTACKALLOC(d, n, int);
-  mjtNum* sparse_buf = mjSTACKALLOC(d, n, mjtNum);
+  mjtNum* restrict dense = mjSTACKALLOC(d, start + 1, mjtNum);
+  mju_zero(dense, start + 1);
 
-  // backpass over rows corresponding to non-zero x(r)
-  int rank = n, i = x_nnz - 1;
-  while (i >= 0) {
-    // get rownnz and rowadr for this row
-    int nnz = rownnz[x_ind[i]], adr = rowadr[x_ind[i]];
+  // scatter x into dense
+  mju_scatter(dense, x, x_ind, x_nnz);
 
-    // compute quantities
-    mjtNum tmp = mat[adr+nnz-1]*mat[adr+nnz-1] + (flg_plus ? x[i]*x[i] : -x[i]*x[i]);
+  // backpass over rows from start down to 0
+  int rank = n;
+  for (int row = start; row >= 0; row--) {
+    // skip if zero
+    if (dense[row] == 0) continue;
+
+    // get rownnz (excluding diagonal), rowadr
+    int nnz = rownnz[row] - 1;
+    int adr = rowadr[row];
+
+    // update diagonal, handle rank-deficient case
+    mjtNum diag = mat[adr + nnz];
+    mjtNum x_row = dense[row];
+    mjtNum tmp = diag*diag + (flg_plus ? x_row*x_row : -x_row*x_row);
     if (tmp < mjMINVAL) {
       tmp = mjMINVAL;
       rank--;
     }
     mjtNum r = mju_sqrt(tmp);
-    mjtNum c = r / mat[adr+nnz-1];
-    mjtNum s = x[i] / mat[adr+nnz-1];
+    mat[adr + nnz] = r;
 
-    // update diagonal
-    mat[adr+nnz-1] = r;
+    // compute Givens rotation parameters https://en.wikipedia.org/wiki/Givens_rotation
+    mjtNum c = diag / r;
+    mjtNum s = -x_row / r;
+    mjtNum s_signed = flg_plus ? -s : s;
 
-    // update row:  mat(r,1:r-1) = (mat(r,1:r-1) + s*x(1:r-1)) / c
-    mju_combineSparseInc(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
-                         nnz-1, i, colind + adr, x_ind);
+    // update row
+    for (int i = 0; i < nnz; i++) {
+      int j = colind[adr + i];
+      mjtNum dense_j = dense[j];
+      mjtNum mat_val = mat[adr + i];
 
-    // update x:  x(1:r-1) = c*x(1:r-1) - s*mat(r,1:r-1)
-    int new_x_nnz = mju_combineSparse(x, mat+adr, c, -s, i, nnz-1, x_ind,
-                                      colind+adr, sparse_buf, buf_ind);
-
-    // update i, correct for changing x
-    i = i - 1 + (new_x_nnz - i);
+      // update mat and dense using the Givens rotation
+      mat[adr + i] = c*mat_val + s_signed*dense_j;
+      dense[j]     = s*mat_val + c*dense_j;
+    }
   }
 
   mj_freeStack(d);

src/engine/engine_util_solve.h

@@ -62,9 +62,9 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 
 // sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank
 //  x is sparse, change in sparsity pattern of mat is not allowed
-MJAPI int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
+MJAPI int mju_cholUpdateSparse(mjtNum* mat, const mjtNum* x, int n, int flg_plus,
                                const int* rownnz, const int* rowadr, const int* colind,
-                               int x_nnz, int* x_ind, mjData* d);
+                               int x_nnz, const int* x_ind, mjData* d);
 
 // band-dense Cholesky decomposition
 //  returns minimum value in the factorized diagonal, or 0 if rank-deficient

test/benchmark/chol_benchmark_test.cc

@@ -19,8 +19,10 @@
 #include <vector>
 
 #include <benchmark/benchmark.h>
+#include <absl/base/attributes.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mujoco.h>
+#include "src/engine/engine_memory.h"
 #include "src/engine/engine_support.h"
 #include "src/engine/engine_util_solve.h"
 #include "src/engine/engine_util_sparse.h"
@@ -329,5 +331,187 @@ void BM_numeric_XL(benchmark::State& state) {
 }
 BENCHMARK(BM_numeric_XL);
 
+// -------------------- rank-1 update/downdate benchmarks ----------------------
+
+constexpr int kNumUpdateVectors = 25;
+
+// old implementation using sparse merge
+int ABSL_ATTRIBUTE_NOINLINE mju_cholUpdateSparse_old(
+    mjtNum* mat, mjtNum* x, int n, int flg_plus, const int* rownnz,
+    const int* rowadr, const int* colind, int x_nnz, int* x_ind, mjData* d) {
+  mj_markStack(d);
+  int* buf_ind = mjSTACKALLOC(d, n, int);
+  mjtNum* sparse_buf = mjSTACKALLOC(d, n, mjtNum);
+
+  int rank = n, i = x_nnz - 1;
+  while (i >= 0) {
+    int nnz = rownnz[x_ind[i]], adr = rowadr[x_ind[i]];
+    mjtNum tmp = mat[adr + nnz - 1] * mat[adr + nnz - 1] +
+                 (flg_plus ? x[i] * x[i] : -x[i] * x[i]);
+    if (tmp < mjMINVAL) {
+      tmp = mjMINVAL;
+      rank--;
+    }
+    mjtNum r = mju_sqrt(tmp);
+    mjtNum c = r / mat[adr + nnz - 1];
+    mjtNum s = x[i] / mat[adr + nnz - 1];
+    mat[adr + nnz - 1] = r;
+    mju_combineSparseInc(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
+                         nnz - 1, i, colind + adr, x_ind);
+    int new_x_nnz = mju_combineSparse(x, mat + adr, c, -s, i, nnz - 1, x_ind,
+                                      colind + adr, sparse_buf, buf_ind);
+    i = i - 1 + (new_x_nnz - i);
+  }
+  mj_freeStack(d);
+  return rank;
+}
+
+// old update implementation benchmark
+template <Size S>
+static void BM_update_old(benchmark::State& state) {
+  mjModel* m = GetModel<S>();
+  mjData* d = mj_makeData(m);
+
+  HessianData hd;
+  hd.Setup(m, d);
+
+  int nv = hd.nv;
+
+  // factorize L using old method (we'll just keep updating this)
+  std::vector<mjtNum> L_work(hd.nL);
+  std::vector<int> L_colind_work(hd.nL);
+  std::vector<int> L_rownnz_work(hd.nv);
+  std::memcpy(L_work.data(), hd.L_init.data(), hd.nL * sizeof(mjtNum));
+  std::memcpy(L_colind_work.data(), hd.L_colind_init.data(),
+              hd.nL * sizeof(int));
+  std::memcpy(L_rownnz_work.data(), hd.L_rownnz_init.data(),
+              hd.nv * sizeof(int));
+  mju_cholFactorSparse(L_work.data(), nv, mjMINVAL, L_rownnz_work.data(),
+                       hd.L_rowadr.data(), L_colind_work.data(), d);
+
+  // prepare update vectors: pick kNumUpdateVectors rows from J (constraint
+  // rows) Each J row has DoF indices which is correct for updating L (nv x nv)
+  int nefc = d->nefc;
+  std::vector<std::vector<mjtNum>> update_vecs(kNumUpdateVectors);
+  std::vector<std::vector<int>> update_inds(kNumUpdateVectors);
+  for (int k = 0; k < kNumUpdateVectors; k++) {
+    int row = (k * 7) % nefc;
+    int nnz = d->efc_J_rownnz[row];
+    int adr = d->efc_J_rowadr[row];
+    update_vecs[k].resize(nnz);
+    update_inds[k].resize(nnz);
+    for (int i = 0; i < nnz; i++) {
+      update_vecs[k][i] = d->efc_J[adr + i] * 0.01;
+      update_inds[k][i] = d->efc_J_colind[adr + i];
+    }
+  }
+
+  // working copy of update vector (sized to nv since pattern can grow)
+  std::vector<mjtNum> x_work(nv);
+  std::vector<int> x_ind_work(nv);
+
+  int vec_idx = 0;
+  for (auto s : state) {
+    int nnz = update_inds[vec_idx].size();
+    std::memset(x_work.data(), 0, nv * sizeof(mjtNum));
+    std::memcpy(x_work.data(), update_vecs[vec_idx].data(),
+                nnz * sizeof(mjtNum));
+    std::memcpy(x_ind_work.data(), update_inds[vec_idx].data(),
+                nnz * sizeof(int));
+    mju_cholUpdateSparse_old(L_work.data(), x_work.data(), nv, 1,
+                             L_rownnz_work.data(), hd.L_rowadr.data(),
+                             L_colind_work.data(), nnz, x_ind_work.data(), d);
+    vec_idx = (vec_idx + 1) % kNumUpdateVectors;
+  }
+
+  mj_deleteData(d);
+  state.SetItemsProcessed(state.iterations());
+}
+
+// new update implementation benchmark
+template <Size S>
+static void BM_update_new(benchmark::State& state) {
+  mjModel* m = GetModel<S>();
+  mjData* d = mj_makeData(m);
+
+  HessianData hd;
+  hd.Setup(m, d);
+
+  int nv = hd.nv;
+
+  // factorize L using new method
+  std::vector<mjtNum> L_work(hd.nL);
+  std::vector<int> L_colind_work(hd.nL);
+  std::vector<int> LT_rownnz_work(nv);
+  std::vector<int> LT_rowadr_work(nv);
+  std::vector<int> LT_colind_work(hd.nL);
+  std::vector<int> LT_pos_work(hd.nL);
+
+  mju_cholFactorSymbolic(L_colind_work.data(), hd.L_rownnz.data(),
+                         hd.L_rowadr.data(), LT_colind_work.data(),
+                         LT_rownnz_work.data(), LT_rowadr_work.data(),
+                         LT_pos_work.data(), hd.HT_rownnz.data(),
+                         hd.HT_rowadr.data(), hd.HT_colind.data(), nv, d);
+  mju_cholFactorNumeric(
+      L_work.data(), nv, mjMINVAL, hd.L_rownnz.data(), hd.L_rowadr.data(),
+      L_colind_work.data(), LT_rownnz_work.data(), LT_rowadr_work.data(),
+      LT_colind_work.data(), LT_pos_work.data(), hd.H.data(),
+      hd.H_rownnz.data(), hd.H_rowadr.data(), hd.H_colind.data(), d);
+
+  // prepare update vectors: pick kNumUpdateVectors rows from J (constraint
+  // rows) Each J row has DoF indices which is correct for updating L (nv x nv)
+  int nefc = d->nefc;
+  std::vector<std::vector<mjtNum>> update_vecs(kNumUpdateVectors);
+  std::vector<std::vector<int>> update_inds(kNumUpdateVectors);
+  for (int k = 0; k < kNumUpdateVectors; k++) {
+    int row = (k * 7) % nefc;
+    int nnz = d->efc_J_rownnz[row];
+    int adr = d->efc_J_rowadr[row];
+    update_vecs[k].resize(nnz);
+    update_inds[k].resize(nnz);
+    for (int i = 0; i < nnz; i++) {
+      update_vecs[k][i] = d->efc_J[adr + i] * 0.01;
+      update_inds[k][i] = d->efc_J_colind[adr + i];
+    }
+  }
+
+  int vec_idx = 0;
+  for (auto s : state) {
+    int nnz = update_inds[vec_idx].size();
+    mju_cholUpdateSparse(L_work.data(), update_vecs[vec_idx].data(), nv, 1,
+                         hd.L_rownnz.data(), hd.L_rowadr.data(),
+                         L_colind_work.data(), nnz, update_inds[vec_idx].data(),
+                         d);
+    vec_idx = (vec_idx + 1) % kNumUpdateVectors;
+  }
+
+  mj_deleteData(d);
+  state.SetItemsProcessed(state.iterations());
+}
+
+void BM_update_old_L(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_update_old<Size::L>(state);
+}
+BENCHMARK(BM_update_old_L);
+
+void BM_update_new_L(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_update_new<Size::L>(state);
+}
+BENCHMARK(BM_update_new_L);
+
+void BM_update_old_XL(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_update_old<Size::XL>(state);
+}
+BENCHMARK(BM_update_old_XL);
+
+void BM_update_new_XL(benchmark::State& state) {
+  MujocoErrorTestGuard guard;
+  BM_update_new<Size::XL>(state);
+}
+BENCHMARK(BM_update_new_XL);
+
 }  // namespace
 }  // namespace mujoco