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

2-3x speedup of sparse matrix squaring.

Browse Source 
Split symbolic and numeric phases for sparse `M'*diag*M` computation. Microseconds per call for the monolithic vs the split approach for the 100_humanoids and 2humanoid100 models:

```
+-------+------+----------+------------+---------+
| Model | Arch | Col (µs) | Split (µs) | Speedup |
+-------+------+----------+------------+---------+
| 2H100 | x86  | 238.3    | 74.5       | 3.2x    |
+-------+------+----------+------------+---------+
|       | ARM  | 111.6    | 53.2       | 2.1x    |
+-------+------+----------+------------+---------+
| 100H  | x86  | 1325.3   | 656.2      | 2.0x    |
+-------+------+----------+------------+---------+
|       | ARM  | 594.8    | 306.6      | 1.9x    |
+-------+------+----------+------------+---------+
```

PiperOrigin-RevId: 900154308
Change-Id: Ia6e9b8e196e2ed37b723a0faf60e9731303a9619
This commit is contained in:
Yuval Tassa
2026-04-15 07:19:03 -07:00
committed by Copybara-Service
parent 62cffb1536
commit a2d0e33c0f
8 changed files with 1301 additions and 626 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+16 -10
View File
@@ -2965,10 +2965,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
return;
}
// pre-count A nonzeros (compute AR_rownnz, AR_rowadr)
d->nA = mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc,
BT_rownnz, BT_rowadr, BT_colind,
B_rownnz, B_rowadr, B_colind, B_rowsuper, d, /*flg_upper=*/1);
int* diagind = mjSTACKALLOC(d, nefc, int);
d->nA = mju_sqrMatTDSparseSymbolic(
d->efc_AR_rownnz, d->efc_AR_rowadr, NULL, diagind,
nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
// allocate A values and column indices on arena
d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
@@ -2981,12 +2982,17 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
return;
}
// A = B * B'
int* diagind = mjSTACKALLOC(d, nefc, int);
mju_sqrMatTDSparse(d->efc_AR, BT, B, NULL, nv, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
BT_rownnz, BT_rowadr, BT_colind, NULL,
B_rownnz, B_rowadr, B_colind, B_rowsuper, d, diagind);
// A = B * B': symbolic phase
mju_sqrMatTDSparseSymbolic(
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, diagind,
nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
// A = B * B': numeric phase
mju_sqrMatTDSparseNumeric(
d->efc_AR, nefc, d->efc_AR_rownnz, d->efc_AR_rowadr,
d->efc_AR_colind, diagind, BT, BT_rownnz, BT_rowadr,
BT_colind, B, B_rownnz, B_rowadr, B_colind, B_rowsuper, NULL, d);
// AR = A + diag(R)
for (int i=0; i < nefc; i++) {
src/engine/engine_solver.c
+28 -16
View File
@@ -1530,10 +1530,10 @@ static void MakeHessian(mjData* d, mjPrimalContext* ctx) {
// sparse
if (ctx->is_sparse) {
// initialize Hessian rowadr, rownnz; get total nonzeros
ctx->nH = mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv,
ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind,
ctx->JT_rowsuper, d, /*flg_upper=*/0);
ctx->nH = mju_sqrMatTDSparseSymbolic(
ctx->H_rownnz, ctx->H_rowadr, NULL, NULL,
nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d);
// add M nonzeros to Hessian total (unavoidable overcounting since H_colind is still unknown)
ctx->nH += ctx->M_rowadr[nv - 1] + ctx->M_rownnz[nv - 1];
@@ -1549,12 +1549,18 @@ static void MakeHessian(mjData* d, mjPrimalContext* ctx) {
ctx->H_colind = mjSTACKALLOC(d, ctx->nH, int);
ctx->H = mjSTACKALLOC(d, ctx->nH, mjtNum);
// compute H = J'*D*J
mju_sqrMatTDSparse(ctx->H, ctx->J, ctx->JT, ctx->D, nefc, nv,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, NULL,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper,
d, /*diagind=*/NULL);
// compute H = J'*D*J: symbolic phase
mju_sqrMatTDSparseSymbolic(
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL,
nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d);
// compute H = J'*D*J: numeric phase
mju_sqrMatTDSparseNumeric(
ctx->H, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
NULL, ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind,
ctx->JT_rowsuper, ctx->D, d);
// add mass matrix: H = J'*D*J + C
mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv,
@@ -1626,12 +1632,18 @@ static void FactorizeHessian(mjData* d, mjPrimalContext* ctx, int flg_recompute)
if (ctx->is_sparse) {
// maybe compute H = M + J'*D*J
if (flg_recompute) {
// compute H = J'*D*J
mju_sqrMatTDSparse(ctx->H, ctx->J, ctx->JT, ctx->D, nefc, nv,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, NULL,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper,
d, /*diagind=*/NULL);
// compute H = J'*D*J: symbolic phase
mju_sqrMatTDSparseSymbolic(
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, NULL,
nefc, nv, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper, d);
// compute H = J'*D*J: numeric phase
mju_sqrMatTDSparseNumeric(
ctx->H, nv, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
NULL, ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
ctx->JT, ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind,
ctx->JT_rowsuper, ctx->D, d);
// add mass matrix: H = J'*D*J + C
mju_addToMatSparse(ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv,
src/engine/engine_util_sparse.c
+311 -5
View File
@@ -13,8 +13,6 @@
// limitations under the License.
#include "engine/engine_util_sparse.h"
#include "engine/engine_util_sparse_avx.h" // IWYU pragma: keep
#include <mujoco/mjdata.h>
#include <mujoco/mjmacro.h>
@@ -23,7 +21,7 @@
#include "engine/engine_memory.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_sparse_avx.h" // IWYU pragma: keep
//------------------------------ sparse operations -------------------------------------------------
@@ -723,9 +721,317 @@ void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) {
}
// max number of supernodes handled
// max number of supernodes handled by column-based matrix squaring functions
#define mjMAXSUPER 8
// column-based symbolic phase for sparse matrix squaring: compute sparsity pattern of M'*M
// if res_colind is NULL: count mode, fill res_rownnz/res_rowadr, return nnz
// if res_colind is not NULL: fill mode, write sorted column indices
// if res_diagind is not NULL: also fill upper triangle and output diagonal indices
int mju_sqrMatTDSparseSymbolic(
int* restrict res_rownnz, int* restrict res_rowadr,
int* restrict res_colind, int* restrict res_diagind, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind,
const int* rownnzT, const int* rowadrT, const int* colindT, const int* rowsuperT, mjData* d) {
mj_markStack(d);
// reinterpret M^T as CSC
const int* colnnz = rownnzT;
const int* coladr = rowadrT;
const int* rowind = colindT;
const int* colsuper = rowsuperT;
// reinterpret M as CSC
const int* colnnzT = rownnz;
const int* coladrT = rowadr;
const int* rowindT = colind;
// marker[j] = 1 if row j has been visited in current column batch
int* marker = mjSTACKALLOC(d, nc, int);
mju_zeroInt(marker, nc);
// buffer_idx: list of row indices with nonzeros in current column batch
int* buffer_idx = mjSTACKALLOC(d, nc, int);
// rowstart[r]: first index in row r of M where column > current result column
int* rowstart = mjSTACKALLOC(d, nr, int);
mju_zeroInt(rowstart, nr);
// clear res_rownnz (used for both counting and filling)
mju_zeroInt(res_rownnz, nc);
// process result columns c = 0, 1, ..., nc-1
int ns; // set in the loop
for (int c = 0; c < nc; c += ns) {
int buffer_nnz = 0;
// column c of M^T
int nnz_c = colnnz[c];
int adr_c = coladr[c];
const int* ind_c = rowind + adr_c;
// supernode size: how many consecutive columns share the same sparsity pattern
ns = 1;
int cs;
if (colsuper && (cs = colsuper[c])) {
ns += mjMIN(cs, mjMAXSUPER - 1);
}
// for each row r where M^T[r, c] != 0, look at row r of M
for (int i = 0; i < nnz_c; i++) {
int r = ind_c[i];
int adrT = coladrT[r];
int nnzT = colnnzT[r];
const int* indT = rowindT + adrT;
// scan row r of M, starting from rowstart[r]
for (int k = rowstart[r]; k < nnzT; k++) {
int j = indT[k];
// skip if j <= c: only fill the strict lower triangle
if (j <= c) {
rowstart[r]++;
continue;
}
// new nonzero in row j of result
if (!marker[j]) {
marker[j] = 1;
buffer_idx[buffer_nnz++] = j;
}
}
}
// scatter: update result rows j > c that have nonzeros in this column batch
// fill mode: write column indices, clear markers
if (res_colind) {
for (int i = 0; i < buffer_nnz; i++) {
int j = buffer_idx[i];
marker[j] = 0;
int nm = mjMIN(ns, j - c);
int adr_j = res_rowadr[j] + res_rownnz[j];
for (int s = 0; s < nm; s++) {
res_colind[adr_j + s] = c + s;
}
res_rownnz[j] += nm;
}
// write diagonal entries
for (int s = 0; s < ns; s++) {
int col = c + s;
if (colnnz[col]) {
res_colind[res_rowadr[col] + res_rownnz[col]] = col;
res_rownnz[col]++;
}
}
}
// count mode: just count and clear markers
else {
for (int i = 0; i < buffer_nnz; i++) {
int j = buffer_idx[i];
marker[j] = 0;
int nm = mjMIN(ns, j - c);
res_rownnz[j] += nm;
if (res_diagind) {
for (int s = 0; s < nm; s++) {
res_rownnz[c + s]++;
}
}
}
// count diagonal entries
for (int s = 0; s < ns; s++) {
int col = c + s;
if (colnnz[col]) {
res_rownnz[col]++;
}
}
}
}
// count mode: compute res_rowadr from res_rownnz
if (!res_colind) {
res_rowadr[0] = 0;
for (int r = 1; r < nc; r++) {
res_rowadr[r] = res_rowadr[r - 1] + res_rownnz[r - 1];
}
}
// fill mode with upper triangle: record diagonal positions and mirror from lower
if (res_colind && res_diagind) {
// save current counts (lower + diagonal)
int* lower_nnz = mjSTACKALLOC(d, nc, int);
mju_copyInt(lower_nnz, res_rownnz, nc);
// save diagonal indices
for (int r = 0; r < nc; r++) {
res_diagind[r] = res_rowadr[r] + lower_nnz[r] - 1;
}
// fill upper triangle: for each (r, c) with c < r, write to (c, r)
for (int r = 0; r < nc; r++) {
int adr = res_rowadr[r];
int nnz = lower_nnz[r];
for (int j = 0; j < nnz; j++) {
int col = res_colind[adr + j];
if (col < r) {
res_colind[res_rowadr[col] + res_rownnz[col]++] = r;
}
}
}
}
mj_freeStack(d);
return res_rowadr[nc - 1] + res_rownnz[nc - 1];
}
// numeric phase for sparse matrix squaring: compute values given pre-computed sparsity
// diagind can be NULL, otherwise fills upper triangle and saves diagonal indices
void mju_sqrMatTDSparseNumeric(
mjtNum* restrict res, int nc,
const int* res_rownnz, const int* res_rowadr, const int* res_colind, const int* res_diagind,
const mjtNum* mat, const int* rownnz, const int* rowadr, const int* colind,
const mjtNum* matT, const int* rownnzT, const int* rowadrT, const int* colindT,
const int* rowsuperT, const mjtNum* diag, mjData* d) {
mj_markStack(d);
// dense accumulator for current result row (or batch of rows)
mjtNum* restrict buffer = mjSTACKALLOC(d, nc * mjMAXSUPER, mjtNum);
mju_zero(buffer, nc * mjMAXSUPER);
// process result rows
int ns; // set in the loop
for (int r = 0; r < nc; r += ns) {
// determine supernode size
ns = 1;
if (rowsuperT) {
ns = rowsuperT[r] + 1;
if (ns > mjMAXSUPER) ns = mjMAXSUPER;
}
// single row
if (ns == 1) {
int nnzT_r = rownnzT[r];
int adr_r = rowadrT[r];
// accumulate: res[r, :] = sum over k in M'[r, :] of diag[k] * M'[r, k] * M[k, :]
for (int i = 0; i < nnzT_r; i++) {
int k = colindT[adr_r + i];
mjtNum valT = matT[adr_r + i];
mjtNum scale = diag ? diag[k] * valT : valT;
if (scale == 0) continue;
int adr_k = rowadr[k];
int nnz_k = rownnz[k];
const int* ind_k = colind + adr_k;
const mjtNum* val_k = mat + adr_k;
for (int j = 0; j < nnz_k; j++) {
int c = ind_k[j];
if (c > r) break;
buffer[c] += scale * val_k[j];
}
}
// scatter from dense buffer to sparse result
int res_adr = res_rowadr[r];
int res_nnz = res_rownnz[r];
const int* res_ind = res_colind + res_adr;
mjtNum* res_val = res + res_adr;
for (int j = 0; j < res_nnz; j++) {
int c = res_ind[j];
res_val[j] = buffer[c];
buffer[c] = 0;
}
}
// supernode: ns > 1 rows share the same sparsity pattern
else {
int nnzT_r = rownnzT[r];
int adr_r = rowadrT[r];
// accumulate for ns rows
for (int i = 0; i < nnzT_r; i++) {
int k = colindT[adr_r + i];
// compute scale for all rows
mjtNum scale[mjMAXSUPER];
if (diag) {
mjtNum dk = diag[k];
if (dk == 0) continue;
for (int s = 0; s < ns; s++) {
scale[s] = dk * matT[rowadrT[r + s] + i];
}
} else {
for (int s = 0; s < ns; s++) {
scale[s] = matT[rowadrT[r + s] + i];
}
}
int adr_k = rowadr[k];
int nnz_k = rownnz[k];
const int* ind_k = colind + adr_k;
const mjtNum* val_k = mat + adr_k;
for (int j = 0; j < nnz_k; j++) {
int c = ind_k[j];
if (c > r + ns - 1) break; // skip if beyond block
mjtNum v = val_k[j];
for (int s = 0; s < ns; s++) {
if (c <= r + s) {
buffer[s * nc + c] += scale[s] * v;
}
}
}
}
// scatter
for (int s = 0; s < ns; s++) {
int row = r + s;
int res_adr = res_rowadr[row];
int res_nnz = res_rownnz[row];
const int* res_ind = res_colind + res_adr;
mjtNum* res_val = res + res_adr;
for (int j = 0; j < res_nnz; j++) {
int c = res_ind[j];
res_val[j] = buffer[s*nc + c];
buffer[s*nc + c] = 0;
}
}
}
}
// fill upper triangle: mirror values from lower triangle
if (res_diagind) {
// initialize write positions after diagonal
int* upper_pos = mjSTACKALLOC(d, nc, int);
for (int r = 0; r < nc; r++) {
upper_pos[r] = res_diagind[r] + 1;
}
// for each (r, c) with c < r, write r to row c
for (int r = 0; r < nc; r++) {
int adr = res_rowadr[r];
int lower_nnz = res_diagind[r] - adr + 1;
for (int j = 0; j < lower_nnz; j++) {
int c = res_colind[adr + j];
if (c < r) {
res[upper_pos[c]++] = res[adr + j];
}
}
}
}
mj_freeStack(d);
}
// compute sparse M'*diag*M (diag=NULL: compute M'*M), res_rowadr must be precomputed
void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc,
@@ -1160,7 +1466,7 @@ void mju_blockDiagSparse(mjtNum* restrict res, int* restrict res_rownnz,
}
// end of block reached: update block counter, column offset, next row
if (r + 1 >= row_next && block + 1 < nb ) {
if (r + 1 >= row_next && block + 1 < nb) {
block++;
col_offset = block_c[block];
row_next = block + 1 < nb ? block_r[block + 1] : nr;
src/engine/engine_util_sparse.h
+21 -3
View File
@@ -82,7 +82,7 @@ MJAPI void mju_mulSymVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
MJAPI int mju_compressSparse(mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind, mjtNum minval);
// count the number of non-zeros in the sum of two sparse vectors
// count the number of nonzeros in the sum of two sparse vectors
MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind);
// incomplete combine sparse: dst = a*dst + b*src at common indices
@@ -138,6 +138,24 @@ MJAPI int mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr,
const int* rownnzT, const int* rowadrT, const int* colindT,
const int* rowsuperT, mjData* d, int flg_upper);
// symbolic phase for mju_sqrMatTDSparse: compute sparsity pattern of M'*M
// if res_colind is NULL: count mode, fill res_rownnz/res_rowadr, return nnz
// if res_colind is not NULL: fill mode, write sorted column indices
// if res_diagind is not NULL: also fill upper triangle and output diagonal indices
MJAPI int mju_sqrMatTDSparseSymbolic(
int* res_rownnz, int* res_rowadr, int* res_colind, int* res_diagind, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind,
const int* rownnzT, const int* rowadrT, const int* colindT, const int* rowsuperT, mjData* d);
// numeric phase for mju_sqrMatTDSparse: compute values given pre-computed sparsity
// res_colind, res_rownnz, res_rowadr must be pre-computed by mju_sqrMatTDSparseSymbolic
MJAPI void mju_sqrMatTDSparseNumeric(
mjtNum* res, int nc,
const int* res_rownnz, const int* res_rowadr, const int* res_colind, const int* res_diagind,
const mjtNum* mat, const int* rownnz, const int* rowadr, const int* colind,
const mjtNum* matT, const int* rownnzT, const int* rowadrT, const int* colindT,
const int* rowsuperT, const mjtNum* diag, mjData* d);
// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);
@@ -146,7 +164,7 @@ MJAPI void mju_blockDiag(mjtNum* res, const mjtNum* mat,
int nc_mat, int nc_res, int nb,
const int* perm_r, const int* perm_c,
const int* block_nr, const int* block_nc,
const int* blockadr_r, const int* blockadr_c);
const int* block_r, const int* block_c);
// block-diagonalize a sparse matrix
MJAPI void mju_blockDiagSparse(
@@ -238,7 +256,7 @@ int mj_mergeSorted(int* merge, const int* chain1, int n1, const int* chain2, int
} else if (c1 > c2) {
merge[k++] = c2;
j++;
} else { // c1 == c2
} else { // c1 == c2
merge[k++] = c1;
i++;
j++;