chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor mju_combineSparse to eliminate temporary buffers.

Browse Source 
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
This commit is contained in:
Yuval Tassa
2026-04-20 03:09:31 -07:00
committed by Copybara-Service
parent 3230cf99f9
commit e6d77650f7
8 changed files with 57 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+3 -6
View File
@@ -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);
}
src/engine/engine_solver.c
+5 -14
View File
@@ -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
src/engine/engine_support.c
+3 -9
View File
@@ -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
src/engine/engine_util_solve.c
+2 -7
View File
@@ -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;
}
src/engine/engine_util_sparse.c
+5 -6
View File
@@ -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
}
}
}
src/engine/engine_util_sparse.h
+35 -32
View File
@@ -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;
test/benchmark/chol_benchmark_test.cc
+1 -7
View File
@@ -22,7 +22,6 @@
#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"
@@ -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;
}
test/benchmark/engine_util_sparse_benchmark_test.cc
+3 -5
View File
@@ -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]);
}
}
}