diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 4be2c28a..38ca5bf2 100644 --- a/src/engine/engine_solver.c +++ b/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); diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 94fdcb3d..be43b1c8 100644 --- a/src/engine/engine_util_solve.c +++ b/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); diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 50c31c15..523bc04a 100644 --- a/src/engine/engine_util_solve.h +++ b/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 diff --git a/test/benchmark/chol_benchmark_test.cc b/test/benchmark/chol_benchmark_test.cc index 68ec8f51..769f990d 100644 --- a/test/benchmark/chol_benchmark_test.cc +++ b/test/benchmark/chol_benchmark_test.cc @@ -19,8 +19,10 @@ #include #include +#include #include #include +#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 +static void BM_update_old(benchmark::State& state) { + mjModel* m = GetModel(); + 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 L_work(hd.nL); + std::vector L_colind_work(hd.nL); + std::vector 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> update_vecs(kNumUpdateVectors); + std::vector> 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 x_work(nv); + std::vector 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 +static void BM_update_new(benchmark::State& state) { + mjModel* m = GetModel(); + mjData* d = mj_makeData(m); + + HessianData hd; + hd.Setup(m, d); + + int nv = hd.nv; + + // factorize L using new method + std::vector L_work(hd.nL); + std::vector L_colind_work(hd.nL); + std::vector LT_rownnz_work(nv); + std::vector LT_rowadr_work(nv); + std::vector LT_colind_work(hd.nL); + std::vector 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> update_vecs(kNumUpdateVectors); + std::vector> 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(state); +} +BENCHMARK(BM_update_old_L); + +void BM_update_new_L(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_update_new(state); +} +BENCHMARK(BM_update_new_L); + +void BM_update_old_XL(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_update_old(state); +} +BENCHMARK(BM_update_old_XL); + +void BM_update_new_XL(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_update_new(state); +} +BENCHMARK(BM_update_new_XL); + } // namespace } // namespace mujoco