From e6d77650f727f026da5922b3ca1857b4b9ad5b40 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 20 Apr 2026 03:09:31 -0700 Subject: [PATCH] Refactor `mju_combineSparse` to eliminate temporary buffers. Combine sparse vectors in-place by first counting total `nnz` and then working backwards from the end. This removes the need for temporary buffers in `mju_combineSparse` and its callers and speeds up the function by ~10%. PiperOrigin-RevId: 902530210 Change-Id: I4f48c327103552ab968d3915399c6067367bec9f --- src/engine/engine_core_constraint.c | 9 +-- src/engine/engine_solver.c | 19 ++---- src/engine/engine_support.c | 12 +--- src/engine/engine_util_solve.c | 9 +-- src/engine/engine_util_sparse.c | 11 ++- src/engine/engine_util_sparse.h | 67 ++++++++++--------- test/benchmark/chol_benchmark_test.cc | 8 +-- .../engine_util_sparse_benchmark_test.cc | 8 +-- 8 files changed, 57 insertions(+), 86 deletions(-) diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 5eab4da6..7849836e 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -496,12 +496,12 @@ static void mj_equalityAnchors(const mjModel* m, const mjData* d, int eq_id, // equality constraints void mj_instantiateEquality(const mjModel* m, mjData* d) { int issparse = mj_isSparse(m), nv = m->nv; - int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL, *buf_ind = NULL; + int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL; int flex_edgeadr, flex_edgenum; int flex_vertadr, flex_vertnum; mjtNum cpos[6], pos[2][3], ref[2], dif, deriv; mjtNum quat[4], quat1[4], quat2[4], quat3[4], axis[3]; - mjtNum *jac[2], *jacdif, *data, *sparse_buf = NULL; + mjtNum *jac[2], *jacdif, *data; // disabled or no equality constraints: return if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax == 0) { @@ -520,8 +520,6 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { if (issparse) { chain = mjSTACKALLOC(d, nv, int); chain2 = mjSTACKALLOC(d, nv, int); - buf_ind = mjSTACKALLOC(d, nv, int); - sparse_buf = mjSTACKALLOC(d, nv, mjtNum); } // find active equality constraints @@ -689,8 +687,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // compute Jacobian: sparse or dense if (issparse) { - NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain, - chain2, sparse_buf, buf_ind); + NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain, chain2); } else { mju_addToScl(jac[0], jac[1], -deriv, nv); } diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index cc9e01ce..1644f841 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -829,8 +829,6 @@ typedef struct { int* L_rowadr; // Hessian factor row addresses (nv x 1) int* LT_rownnz; // Hessian factor transpose row nonzeros (nv x 1) int* LT_rowadr; // Hessian factor transpose row addresses (nv x 1) - int* buf_ind; // index buffer for sparse addition (nv x 1) - mjtNum* buf_val; // value buffer for sparse addition (nv x 1) // Newton arrays, computed-size (MakeHessian) int nH; // number of nonzeros in Hessian H @@ -979,9 +977,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) { if (flg_Newton) { nNum += nefc + nv; // D, cholupd if (is_elliptic) nNum += 6*nv; // LTJ - if (is_sparse) { - nNum += nv; // buf_val - } else { + if (!is_sparse) { nNum += nv*nv; // L (dense) if (is_elliptic) nNum += nv*nv; // Lcone (dense) } @@ -993,7 +989,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) { size_t nInt = nefc; // oldstate if (is_sparse) { nInt += 3*nv + nJ; // JT sparse - if (flg_Newton) nInt += 9*nv; // Newton sparse + if (flg_Newton) nInt += 8*nv; // Newton sparse } // allocate mjtNum and int blocks @@ -1018,9 +1014,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) { if (is_elliptic) { ctx->LTJ = numblock; numblock += 6*nv; } - if (is_sparse) { - ctx->buf_val = numblock; numblock += nv; - } else { + if (!is_sparse) { ctx->nL = nv*nv; ctx->L = numblock; numblock += ctx->nL; ctx->Lcone = is_elliptic ? numblock : NULL; @@ -1049,7 +1043,6 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) { ctx->L_rowadr = intblock; intblock += nv; ctx->LT_rownnz = intblock; intblock += nv; ctx->LT_rowadr = intblock; intblock += nv; - ctx->buf_ind = intblock; intblock += nv; } // sparse: compute Jacobian transpose @@ -1616,8 +1609,7 @@ static void MakeHessian(mjData* d, mjPrimalContext* ctx) { // add mass matrix: H = J'*D*J + M mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, - ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, - ctx->buf_val, ctx->buf_ind); + ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); // compute H' sparse structure (upper triangle, required for symbolic Cholesky) mju_transposeSparse(NULL, NULL, nv, nv, ctx->HT_rownnz, ctx->HT_rowadr, ctx->HT_colind, NULL, @@ -1691,8 +1683,7 @@ static void FactorizeHessian(mjData* d, mjPrimalContext* ctx, int flg_recompute) // add mass matrix: H = J'*D*J + C mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, - ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, - ctx->buf_val, ctx->buf_ind); + ctx->M, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); } // numeric sparse factorization: L = chol(H) using pre-computed sparsity pattern diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 27537d39..a8725bf7 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -420,17 +420,11 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind) { int nv = m->nv; + // sparse if (rownnz && rowadr && colind) { - mj_markStack(d); - mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum); - int* buf_ind = mjSTACKALLOC(d, nv, int); - - mju_addToMatSparse(dst, rownnz, rowadr, colind, nv, - d->M, m->M_rownnz, m->M_rowadr, m->M_colind, - buf_val, buf_ind); - - mj_freeStack(d); + mju_addToMatSparse(dst, rownnz, rowadr, colind, nv, d->M, + m->M_rownnz, m->M_rowadr, m->M_colind); } // dense diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 259f7c59..1f8b839b 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -144,10 +144,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, int* rownnz, const int* rowadr, int* colind, mjData* d) { int rank = n; - - mj_markStack(d); - mjtNum* buf = mjSTACKALLOC(d, n, mjtNum); - int* buf_ind = mjSTACKALLOC(d, n, int); + (void) d; // backpass over rows for (int r=n-1; r >= 0; r--) { @@ -175,15 +172,13 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, // mat(c,0:c) = mat(c,0:c) - mat(r,c) * mat(r,0:c) int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], 1, -mat[adr+i], - rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r], - buf, buf_ind); + rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r]); // assign new nnz to row c rownnz[c] = nnz_c; } } - mj_freeStack(d); return rank; } diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 05960567..cfed023b 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -198,15 +198,14 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int } -// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +// add sparse matrix M to sparse destination matrix void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, const mjtNum* M, const int* M_rownnz, const int* M_rowadr, - const int* M_colind, - mjtNum* buf_val, int* buf_ind) { + const int* M_colind) { for (int i=0; i < nr; i++) { rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1, rownnz[i], M_rownnz[i], colind + rowadr[i], - M_colind + M_rowadr[i], buf_val, buf_ind); + M_colind + M_rowadr[i]); } } @@ -256,8 +255,8 @@ void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat, for (int k=diag-1; k >= 0; k--) { int j = ind[k]; mjtNum val = row[k]; - res[i] += val * vec[j]; // strict lower - res[j] += val * vec[i]; // strict upper + res[i] += val * vec[j]; // strict lower + res[j] += val * vec[i]; // strict upper } } } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 1350ccff..68bda070 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -62,11 +62,10 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); -// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +// add sparse matrix M to sparse destination matrix MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, const mjtNum* M, const int* M_rownnz, const int* M_rowadr, - const int* M_colind, - mjtNum* buf_val, int* buf_ind); + const int* M_colind); // add symmetric matrix (only lower triangle represented) to dense matrix MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n, @@ -294,8 +293,7 @@ void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, i // combine two sparse vectors: dst = a*dst + b*src, return nnz of result static inline int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b, - int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind, - mjtNum* buf, int* buf_ind) { + int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind) { // check for identical pattern if (dst_nnz == src_nnz) { if (mju_compare(dst_ind, src_ind, dst_nnz)) { @@ -305,49 +303,54 @@ int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b, } } - // copy dst into buf - if (dst_nnz) { - memcpy(buf, dst, dst_nnz * sizeof(mjtNum)); - memcpy(buf_ind, dst_ind, dst_nnz * sizeof(int)); - } + // compute total nnz of result + int nnz = mju_combineSparseCount(dst_nnz, src_nnz, dst_ind, src_ind); - // prepare to merge buf and src into dst - int bi = 0, si = 0, nnz = 0; - int buf_nnz = dst_nnz; + // set up read/write pointers at end of arrays + int bi = dst_nnz - 1, si = src_nnz - 1, w = nnz - 1; - // merge vectors - while (bi < buf_nnz && si < src_nnz) { - int badr = buf_ind[bi]; + // merge backwards + while (bi >= 0 && si >= 0) { + int badr = dst_ind[bi]; int sadr = src_ind[si]; if (badr == sadr) { - dst[nnz] = a*buf[bi++] + b*src[si++]; - dst_ind[nnz++] = badr; + dst[w] = a*dst[bi] + b*src[si]; + dst_ind[w] = badr; + bi--; + si--; } - // buf only - else if (badr < sadr) { - dst[nnz] = a*buf[bi++]; - dst_ind[nnz++] = badr; + // dst only + else if (badr > sadr) { + dst[w] = a*dst[bi]; + dst_ind[w] = badr; + bi--; } // src only else { - dst[nnz] = b*src[si++]; - dst_ind[nnz++] = sadr; + dst[w] = b*src[si]; + dst_ind[w] = sadr; + si--; } + w--; } - // the rest of src only - while (si < src_nnz) { - dst[nnz] = b*src[si]; - dst_ind[nnz++] = src_ind[si++]; + // remaining src elements + while (si >= 0) { + dst[w] = b*src[si]; + dst_ind[w] = src_ind[si]; + si--; + w--; } - // the rest of buf only - while (bi < buf_nnz) { - dst[nnz] = a*buf[bi]; - dst_ind[nnz++] = buf_ind[bi++]; + // remaining dst elements: already in place, scale by a + if (a != 1) { + while (bi >= 0) { + dst[bi] *= a; + bi--; + } } return nnz; diff --git a/test/benchmark/chol_benchmark_test.cc b/test/benchmark/chol_benchmark_test.cc index 004e0c2a..ba8d94bf 100644 --- a/test/benchmark/chol_benchmark_test.cc +++ b/test/benchmark/chol_benchmark_test.cc @@ -22,7 +22,6 @@ #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" @@ -352,10 +351,6 @@ constexpr int kNumUpdateVectors = 25; 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]]; @@ -372,10 +367,9 @@ int ABSL_ATTRIBUTE_NOINLINE mju_cholUpdateSparse_old( 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); + colind + adr); i = i - 1 + (new_x_nnz - i); } - mj_freeStack(d); return rank; } diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index a82be98e..3a28a291 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -99,8 +99,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst, mjtNum a, mjtNum b, int dst_nnz, int src_nnz, int* dst_ind, - const int* src_ind, - mjtNum* buf, int* buf_ind) { + const int* src_ind) { // check for identical pattern if (compare_baseline(dst_ind, src_ind, dst_nnz)) { // combine mjtNum data directly @@ -116,8 +115,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst, mjtNum a, mjtNum b, int dst_nnz, int src_nnz, int* dst_ind, - const int* src_ind, - mjtNum* buf, int* buf_ind) { + const int* src_ind) { // check for identical pattern if (compare_memcmp(dst_ind, src_ind, dst_nnz)) { // combine mjtNum data directly @@ -372,7 +370,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) { // in order to trigger all if's in combineSparse func(H+rowadr[c], H+rowadr[r], 1, -H[adr+i], rownnz[c], rownnz[c], - colind+rowadr[c], colind+rowadr[c], NULL, NULL); + colind+rowadr[c], colind+rowadr[c]); } } }