chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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_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;