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++;
test/benchmark/chol_benchmark_test.cc
+26 -11
View File
@@ -68,13 +68,15 @@ struct HessianData {
// D diagonal
std::vector<mjtNum> D;
int nefc;
void Setup(const mjModel* m, mjData* d) {
// initialize simulation state
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
nv = m->nv;
int nefc = d->nefc;
nefc = d->nefc;
// compute D corresponding to quad states
D.resize(nefc);
@@ -205,14 +207,27 @@ mjModel* GetModel() {
return m;
}
template <Size S>
HessianData& GetHessianData() {
static HessianData data;
static bool initialized = false;
if (!initialized) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
data.Setup(m, d);
mj_deleteData(d);
initialized = true;
}
return data;
}
// old implementation benchmark
template <Size S>
static void BM_chol_old(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<mjtNum> L_work(hd.nL);
std::vector<int> L_colind_work(hd.nL);
@@ -239,8 +254,7 @@ static void BM_chol_symbolic(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<int> L_colind_work(hd.nL);
std::vector<int> LT_rownnz_work(hd.nv);
@@ -266,8 +280,7 @@ static void BM_chol_numeric(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
std::vector<mjtNum> L_work(hd.nL);
std::vector<int> L_colind_work(hd.nL);
@@ -371,9 +384,10 @@ template <Size S>
static void BM_update_old(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
int nv = hd.nv;
@@ -433,9 +447,10 @@ template <Size S>
static void BM_update_new(benchmark::State& state) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
mj_resetDataKeyframe(m, d, 0);
mj_forward(m, d);
HessianData hd;
hd.Setup(m, d);
HessianData& hd = GetHessianData<S>();
int nv = hd.nv;
test/benchmark/engine_util_sparse_benchmark_test.cc
+48 -209
View File
@@ -38,116 +38,7 @@ static const int kNumWarmupSteps = 500;
// ----------------------------- old functions --------------------------------
void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag,
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr, const int* colind,
const int* rowsuper, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, mjData* d, int* unused) {
mj_markStack(d);
int* chain = mj_stackAllocInt(d, 2 * nc);
mjtNum* buffer = mj_stackAllocNum(d, nc);
for (int r = 0; r < nc; r++) {
res_rowadr[r] = r * nc;
}
for (int r = 0; r < nc; r++) {
if (rowsuperT && r > 0 && rowsuperT[r - 1] > 0) {
res_rownnz[r] = res_rownnz[r - 1];
memcpy(res_colind + res_rowadr[r], res_colind + res_rowadr[r - 1],
res_rownnz[r] * sizeof(int));
if (rownnzT[r]) {
res_colind[res_rowadr[r] + res_rownnz[r]] = r;
res_rownnz[r]++;
}
} else {
int nchain = 0;
int inew = 0, iold = nc;
int lastadded = -1;
for (int i = 0; i < rownnzT[r]; i++) {
int c = colindT[rowadrT[r] + i];
if (rowsuper && lastadded >= 0 &&
(c - lastadded) <= rowsuper[lastadded]) {
continue;
} else {
lastadded = c;
}
int adr = inew;
inew = iold;
iold = adr;
int nnewchain = 0;
adr = 0;
int end = rowadr[c] + rownnz[c];
for (int adr1 = rowadr[c]; adr1 < end; adr1++) {
int col_mat = colind[adr1];
while (adr < nchain && chain[iold + adr] < col_mat &&
chain[iold + adr] <= r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
if (col_mat > r) {
break;
}
if (adr < nchain && chain[iold + adr] == col_mat) {
adr++;
}
chain[inew + nnewchain++] = col_mat;
}
while (adr < nchain && chain[iold + adr] <= r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
nchain = nnewchain;
}
res_rownnz[r] = nchain;
if (nchain) {
memcpy(res_colind + res_rowadr[r], chain + inew, nchain * sizeof(int));
}
}
}
for (int r = 0; r < nc; r++) {
int adr = res_rowadr[r];
for (int i = 0; i < res_rownnz[r]; i++) {
buffer[res_colind[adr + i]] = 0;
}
for (int i = 0; i < rownnzT[r]; i++) {
int c = colindT[rowadrT[r] + i];
mjtNum matTrc = matT[rowadrT[r] + i];
if (diag) {
matTrc *= diag[c];
}
int end = rowadr[c] + rownnz[c];
for (int adr = rowadr[c]; adr < end; adr++) {
int adr1;
if ((adr1 = colind[adr]) > r) {
break;
}
buffer[adr1] += matTrc * mat[adr];
}
}
adr = res_rowadr[r];
for (int i = 0; i < res_rownnz[r]; i++) {
res[adr + i] = buffer[res_colind[adr + i]];
}
}
for (int r = 1; r < nc; r++) {
int end = res_rowadr[r] + res_rownnz[r] - 1;
for (int adr = res_rowadr[r]; adr < end; adr++) {
int adr1 = res_rowadr[res_colind[adr]] + res_rownnz[res_colind[adr]]++;
res[adr1] = res[adr];
res_colind[adr1] = r;
}
}
mj_freeStack(d);
}
// transpose sparse matrix (uncompressed)
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
@@ -506,15 +397,27 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
}
BENCHMARK(BM_combineSparse_old);
enum class Size { H2_100, H100 };
template <Size S>
const char* ModelPath() {
if constexpr (S == Size::H2_100) {
return "../test/benchmark/testdata/2humanoid100_chol.xml";
} else {
return "../test/benchmark/testdata/100_humanoids_chol.xml";
}
}
enum class Supernode {
None,
PostProcess,
Inline
};
template <Size S>
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
Supernode super) {
static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml");
static mjModel* m = LoadModelFromPath(ModelPath<S>());
// force use of sparse matrices
m->opt.jacobian = mjJAC_SPARSE;
@@ -553,131 +456,67 @@ static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func,
}
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_old(benchmark::State& state) {
BM_transposeSparse_2H100_old(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &transposeSparse_baseline, Supernode::None);
BM_transposeSparse<Size::H2_100>(state, &transposeSparse_baseline,
Supernode::None);
}
BENCHMARK(BM_transposeSparse_old);
BENCHMARK(BM_transposeSparse_2H100_old);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_new(benchmark::State& state) {
BM_transposeSparse_2H100_new(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse, Supernode::None);
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
Supernode::None);
}
BENCHMARK(BM_transposeSparse_new);
BENCHMARK(BM_transposeSparse_2H100_new);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_superpost(benchmark::State& state) {
BM_transposeSparse_2H100_superpost(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse, Supernode::PostProcess);
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
Supernode::PostProcess);
}
BENCHMARK(BM_transposeSparse_superpost);
BENCHMARK(BM_transposeSparse_2H100_superpost);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_superinline(benchmark::State& state) {
BM_transposeSparse_2H100_superinline(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse(state, &mju_transposeSparse, Supernode::Inline);
}
BENCHMARK(BM_transposeSparse_superinline);
static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
static mjModel* m =
LoadModelFromPath("../test/benchmark/testdata/2humanoid100.xml");
// force use of sparse matrices, Newton solver, no islands
m->opt.jacobian = mjJAC_SPARSE;
m->opt.solver = mjSOL_NEWTON;
m->opt.disableflags |= mjDSBL_ISLAND;
mjData* d = mj_makeData(m);
// warm-up rollout to get a typical state
while (d->time < 2) {
mj_step(m, d);
}
// allocate
mj_markStack(d);
mjtNum* H = mj_stackAllocNum(d, m->nv * m->nv);
int* rownnz = mj_stackAllocInt(d, m->nv);
int* rowadr = mj_stackAllocInt(d, m->nv);
int* colind = mj_stackAllocInt(d, m->nv * m->nv);
int* diagind = mj_stackAllocInt(d, m->nv);
// compute D corresponding to quad states
mjtNum* D = mj_stackAllocNum(d, d->nefc);
for (int i = 0; i < d->nefc; i++) {
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
D[i] = d->efc_D[i];
} else {
D[i] = 0;
}
}
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
int* JT_colind = mj_stackAllocInt(d, d->nJ);
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
// time benchmark
if (func) {
mju_sqrMatTDSparseCount(rownnz, rowadr, m->nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
JT_rownnz, JT_rowadr,
JT_colind, nullptr, d, 1);
for (auto s : state) {
// compute H = J'*D*J, compressed layout
func(H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, diagind);
}
} else {
for (auto s : state) {
// baseline depends on efc_J_rowsuper
mju_superSparse(d->nefc, d->efc_J_rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
// compute H = J'*D*J, uncompressed layout
mju_sqrMatTDSparse_baseline(
H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, /*unused=*/nullptr);
}
}
// finalize
mj_freeStack(d);
mj_deleteData(d);
state.SetItemsProcessed(state.iterations());
BM_transposeSparse<Size::H2_100>(state, &mju_transposeSparse,
Supernode::Inline);
}
BENCHMARK(BM_transposeSparse_2H100_superinline);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_sqrMatTDSparse_col(benchmark::State& state) {
BM_transposeSparse_100H_old(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse);
BM_transposeSparse<Size::H100>(state, &transposeSparse_baseline,
Supernode::None);
}
BENCHMARK(BM_sqrMatTDSparse_col);
BENCHMARK(BM_transposeSparse_100H_old);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_sqrMatTDSparse_row(benchmark::State& state) {
BM_transposeSparse_100H_new(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTDSparse(state, &mju_sqrMatTDSparse_row);
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse, Supernode::None);
}
BENCHMARK(BM_sqrMatTDSparse_row);
BENCHMARK(BM_transposeSparse_100H_new);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_sqrMatTDSparse_uncompressed(benchmark::State& state) {
BM_transposeSparse_100H_superpost(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTDSparse(state, nullptr);
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse,
Supernode::PostProcess);
}
BENCHMARK(BM_sqrMatTDSparse_uncompressed);
BENCHMARK(BM_transposeSparse_100H_superpost);
void ABSL_ATTRIBUTE_NO_TAIL_CALL
BM_transposeSparse_100H_superinline(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_transposeSparse<Size::H100>(state, &mju_transposeSparse,
Supernode::Inline);
}
BENCHMARK(BM_transposeSparse_100H_superinline);
} // namespace
} // namespace mujoco
test/benchmark/sqrmat_benchmark_test.cc
+423
View File
@@ -0,0 +1,423 @@
// Copyright 2026 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// Benchmarks for sparse matrix operations.
#include <cstring>
#include <vector>
#include "benchmark/benchmark.h"
#include <absl/base/attributes.h>
#include <mujoco/mjdata.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_util_sparse.h"
#include "test/fixture.h"
namespace mujoco {
namespace {
// ================================ Test Data ==================================
// Stores pre-computed sparse matrix inputs extracted from MuJoCo simulations.
// Each benchmark computes its own outputs (H, L, etc.) from these inputs.
struct SparseTestData {
// Dimensions
int nv; // number of DoFs
int nefc; // number of constraint rows
int nJ; // nnz in J
// J (Jacobian) - nefc x nv sparse
std::vector<mjtNum> J;
std::vector<int> J_rownnz, J_rowadr, J_colind, J_rowsuper;
// J' (transpose)
std::vector<mjtNum> JT;
std::vector<int> JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper;
// D (diagonal weights for constraints)
std::vector<mjtNum> D;
// M structure (mass matrix, lower triangle)
std::vector<int> M_rownnz, M_rowadr, M_colind;
void Setup(const mjModel* m, mjData* d) {
// initialize simulation state
mj_resetDataKeyframe(m, d, 0);
mj_step(m, d);
mj_forward(m, d);
nv = m->nv;
nefc = d->nefc;
nJ = d->nJ;
// copy J
J.assign(d->efc_J, d->efc_J + nJ);
J_rownnz.assign(d->efc_J_rownnz, d->efc_J_rownnz + nefc);
J_rowadr.assign(d->efc_J_rowadr, d->efc_J_rowadr + nefc);
J_colind.assign(d->efc_J_colind, d->efc_J_colind + nJ);
J_rowsuper.assign(d->efc_J_rowsuper, d->efc_J_rowsuper + nefc);
// transpose J
JT.assign(nJ, 0);
JT_rownnz.assign(nv, 0);
JT_rowadr.assign(nv, 0);
JT_colind.assign(nJ, 0);
JT_rowsuper.assign(nv, 0);
mju_transposeSparse(JT.data(), J.data(), nefc, nv, JT_rownnz.data(),
JT_rowadr.data(), JT_colind.data(), JT_rowsuper.data(),
J_rownnz.data(), J_rowadr.data(), J_colind.data());
// compute D corresponding to quadratic constraint states
D.resize(nefc);
for (int i = 0; i < nefc; i++) {
if (d->efc_state[i] == mjCNSTRSTATE_QUADRATIC) {
D[i] = d->efc_D[i];
} else {
D[i] = 0;
}
}
// copy M structure
M_rownnz.assign(m->M_rownnz, m->M_rownnz + nv);
M_rowadr.assign(m->M_rowadr, m->M_rowadr + nv);
int nM = M_rowadr[nv - 1] + M_rownnz[nv - 1];
M_colind.assign(m->M_colind, m->M_colind + nM);
}
};
// ================================ Model Sizes ================================
enum class Size { H2_100, H100 };
template <Size S>
const char* ModelPath() {
if constexpr (S == Size::H2_100) {
return "../test/benchmark/testdata/2humanoid100_chol.xml";
} else {
return "../test/benchmark/testdata/100_humanoids_chol.xml";
}
}
template <Size S>
mjModel* GetModel() {
static mjModel* m = LoadModelFromPath(ModelPath<S>());
m->opt.jacobian = mjJAC_SPARSE;
m->opt.solver = mjSOL_NEWTON;
m->opt.disableflags |= mjDSBL_ISLAND;
return m;
}
template <Size S>
SparseTestData& GetData() {
static SparseTestData data;
static bool initialized = false;
if (!initialized) {
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
data.Setup(m, d);
mj_deleteData(d);
initialized = true;
}
return data;
}
// ========================== Baseline Implementations =========================
// Baseline sqrMatTD (uncompressed layout, from old implementation)
void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
mjtNum* res, const mjtNum* mat, const mjtNum* matT, const mjtNum* diag,
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr, const int* colind,
const int* rowsuper, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, mjData* d) {
mj_markStack(d);
int* chain = mj_stackAllocInt(d, 2 * nc);
mjtNum* buffer = mj_stackAllocNum(d, nc);
for (int r = 0; r < nc; r++) {
res_rowadr[r] = r * nc;
}
for (int r = 0; r < nc; r++) {
if (rowsuperT && r > 0 && rowsuperT[r - 1] > 0) {
res_rownnz[r] = res_rownnz[r - 1];
memcpy(res_colind + res_rowadr[r], res_colind + res_rowadr[r - 1],
res_rownnz[r] * sizeof(int));
if (rownnzT[r]) {
res_colind[res_rowadr[r] + res_rownnz[r]] = r;
res_rownnz[r]++;
}
} else {
int nchain = 0;
int inew = 0, iold = nc;
int lastadded = -1;
for (int i = 0; i < rownnzT[r]; i++) {
int c = colindT[rowadrT[r] + i];
if (rowsuper && lastadded >= 0 &&
(c - lastadded) <= rowsuper[lastadded]) {
continue;
} else {
lastadded = c;
}
int adr = inew;
inew = iold;
iold = adr;
int nnewchain = 0;
adr = 0;
int end = rowadr[c] + rownnz[c];
for (int adr1 = rowadr[c]; adr1 < end; adr1++) {
int col_mat = colind[adr1];
while (adr < nchain && chain[iold + adr] < col_mat &&
chain[iold + adr] <= r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
if (col_mat > r) {
break;
}
if (adr < nchain && chain[iold + adr] == col_mat) {
adr++;
}
chain[inew + nnewchain++] = col_mat;
}
while (adr < nchain && chain[iold + adr] <= r) {
chain[inew + nnewchain++] = chain[iold + adr++];
}
nchain = nnewchain;
}
res_rownnz[r] = nchain;
if (nchain) {
memcpy(res_colind + res_rowadr[r], chain + inew, nchain * sizeof(int));
}
}
}
for (int r = 0; r < nc; r++) {
int adr = res_rowadr[r];
for (int i = 0; i < res_rownnz[r]; i++) {
buffer[res_colind[adr + i]] = 0;
}
for (int i = 0; i < rownnzT[r]; i++) {
int c = colindT[rowadrT[r] + i];
mjtNum matTrc = matT[rowadrT[r] + i];
if (diag) {
matTrc *= diag[c];
}
int end = rowadr[c] + rownnz[c];
for (int adr2 = rowadr[c]; adr2 < end; adr2++) {
int adr1;
if ((adr1 = colind[adr2]) > r) {
break;
}
buffer[adr1] += matTrc * mat[adr2];
}
}
adr = res_rowadr[r];
for (int i = 0; i < res_rownnz[r]; i++) {
res[adr + i] = buffer[res_colind[adr + i]];
}
}
for (int r = 1; r < nc; r++) {
int end = res_rowadr[r] + res_rownnz[r] - 1;
for (int adr = res_rowadr[r]; adr < end; adr++) {
int adr1 = res_rowadr[res_colind[adr]] + res_rownnz[res_colind[adr]]++;
res[adr1] = res[adr];
res_colind[adr1] = r;
}
}
mj_freeStack(d);
}
// ========================== SqrMatTD Benchmarks ==============================
enum class SqrMatTDVariant {
kBaseline,
kRow,
kCol,
kSplitCol
};
template <Size S>
static void BM_sqrMatTD_impl(benchmark::State& state, SqrMatTDVariant variant) {
SparseTestData& data = GetData<S>();
mjModel* m = GetModel<S>();
mjData* d = mj_makeData(m);
int nv = data.nv;
// nothing to benchmark if no constraints
if (data.nefc == 0) {
for (auto s : state) {}
mj_deleteData(d);
return;
}
// allocate H output (uncompressed for baseline, compressed for others)
int max_nnz = (variant == SqrMatTDVariant::kBaseline) ? nv * nv : 0;
std::vector<mjtNum> H;
std::vector<int> H_rownnz(nv);
std::vector<int> H_rowadr(nv);
std::vector<int> H_colind;
std::vector<int> diagind(nv);
if (variant == SqrMatTDVariant::kBaseline) {
H.resize(max_nnz);
H_colind.resize(max_nnz);
} else if (variant == SqrMatTDVariant::kSplitCol ||
variant == SqrMatTDVariant::kCol) {
// use symbolic to count nnz
int nH = mju_sqrMatTDSparseSymbolic(
H_rownnz.data(), H_rowadr.data(), nullptr, nullptr,
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
data.JT_colind.data(), data.JT_rowsuper.data(), d);
H.resize(nH);
H_colind.resize(nH);
} else {
// row: use Count (lower triangle only)
mju_sqrMatTDSparseCount(
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
int nH = H_rowadr[nv - 1] + H_rownnz[nv - 1];
H.resize(nH);
H_colind.resize(nH);
}
for (auto s : state) {
switch (variant) {
case SqrMatTDVariant::kBaseline:
mju_superSparse(data.nefc, data.J_rowsuper.data(), data.J_rownnz.data(),
data.J_rowadr.data(), data.J_colind.data());
mju_sqrMatTDSparse_baseline(
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
data.J_rowsuper.data(), data.JT_rownnz.data(),
data.JT_rowadr.data(), data.JT_colind.data(),
data.JT_rowsuper.data(), d);
break;
case SqrMatTDVariant::kRow:
mju_sqrMatTDSparseCount(
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
mju_sqrMatTDSparse_row(
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
nullptr, data.JT_rownnz.data(), data.JT_rowadr.data(),
data.JT_colind.data(), data.JT_rowsuper.data(), d, nullptr);
break;
case SqrMatTDVariant::kCol:
mju_sqrMatTDSparseCount(
H_rownnz.data(), H_rowadr.data(), nv, data.J_rownnz.data(),
data.J_rowadr.data(), data.J_colind.data(), data.JT_rownnz.data(),
data.JT_rowadr.data(), data.JT_colind.data(), nullptr, d, 0);
mju_sqrMatTDSparse(
H.data(), data.J.data(), data.JT.data(), data.D.data(), data.nefc,
nv, H_rownnz.data(), H_rowadr.data(), H_colind.data(),
data.J_rownnz.data(), data.J_rowadr.data(), data.J_colind.data(),
nullptr, data.JT_rownnz.data(), data.JT_rowadr.data(),
data.JT_colind.data(), data.JT_rowsuper.data(), d, nullptr);
break;
case SqrMatTDVariant::kSplitCol:
mju_sqrMatTDSparseSymbolic(
H_rownnz.data(), H_rowadr.data(), nullptr, nullptr,
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
data.JT_colind.data(), data.JT_rowsuper.data(), d);
mju_sqrMatTDSparseSymbolic(
H_rownnz.data(), H_rowadr.data(), H_colind.data(), nullptr,
data.nefc, nv, data.J_rownnz.data(), data.J_rowadr.data(),
data.J_colind.data(), data.JT_rownnz.data(), data.JT_rowadr.data(),
data.JT_colind.data(), data.JT_rowsuper.data(), d);
mju_sqrMatTDSparseNumeric(
H.data(), nv, H_rownnz.data(), H_rowadr.data(),
H_colind.data(), nullptr, data.J.data(), data.J_rownnz.data(),
data.J_rowadr.data(), data.J_colind.data(), data.JT.data(),
data.JT_rownnz.data(), data.JT_rowadr.data(), data.JT_colind.data(),
data.JT_rowsuper.data(), data.D.data(), d);
break;
}
}
mj_deleteData(d);
state.SetItemsProcessed(state.iterations());
}
void BM_sqrMatTD_2H100_baseline(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kBaseline);
}
BENCHMARK(BM_sqrMatTD_2H100_baseline);
void BM_sqrMatTD_2H100_row(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kRow);
}
BENCHMARK(BM_sqrMatTD_2H100_row);
void BM_sqrMatTD_2H100_col(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kCol);
}
BENCHMARK(BM_sqrMatTD_2H100_col);
void BM_sqrMatTD_2H100_splitCol(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H2_100>(state, SqrMatTDVariant::kSplitCol);
}
BENCHMARK(BM_sqrMatTD_2H100_splitCol);
void BM_sqrMatTD_100H_baseline(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kBaseline);
}
BENCHMARK(BM_sqrMatTD_100H_baseline);
void BM_sqrMatTD_100H_row(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kRow);
}
BENCHMARK(BM_sqrMatTD_100H_row);
void BM_sqrMatTD_100H_col(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kCol);
}
BENCHMARK(BM_sqrMatTD_100H_col);
void BM_sqrMatTD_100H_splitCol(benchmark::State& state) {
MujocoErrorTestGuard guard;
BM_sqrMatTD_impl<Size::H100>(state, SqrMatTDVariant::kSplitCol);
}
BENCHMARK(BM_sqrMatTD_100H_splitCol);
} // namespace
} // namespace mujoco
int main(int argc, char** argv) {
benchmark::Initialize(&argc, argv);
benchmark::RunSpecifiedBenchmarks();
return 0;
}
test/engine/engine_util_sparse_test.cc
+428 -372
View File
@@ -14,10 +14,11 @@
// Tests for engine/engine_util_sparse.c
#include <array>
#include "src/engine/engine_util_sparse.h"
#include <array>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mujoco.h>
@@ -360,6 +361,45 @@ TEST_F(EngineUtilSparseTest, MjuCompressSparse) {
EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected_minval1, 6));
}
// helper: run split-col approach and return dense result
static void SqrMatTDSplitCol(
std::vector<mjtNum>& dense_result, int nr, int nc,
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,
int* out_diagind, mjData* d) {
// count mode
std::vector<int> H_rownnz(nc, 0);
std::vector<int> H_rowadr(nc, 0);
int nnz = mju_sqrMatTDSparseSymbolic(
H_rownnz.data(), H_rowadr.data(), nullptr,
out_diagind, nr, nc, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, d);
// fill mode
std::vector<int> H_colind(nnz);
mju_sqrMatTDSparseSymbolic(
H_rownnz.data(), H_rowadr.data(), H_colind.data(),
out_diagind, nr, nc, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, d);
// numeric phase
std::vector<mjtNum> H(nnz, 0);
mju_sqrMatTDSparseNumeric(
H.data(), nc, H_rownnz.data(), H_rowadr.data(),
H_colind.data(), out_diagind, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, diag, d);
// densify
dense_result.assign(nc * nc, 0);
for (int r = 0; r < nc; r++) {
for (int j = 0; j < H_rownnz[r]; j++) {
int c = H_colind[H_rowadr[r] + j];
dense_result[r*nc + c] = H[H_rowadr[r] + j];
}
}
}
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
// 0 0 0
// M = 0 0 0
@@ -378,29 +418,13 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
int rownnzT[] = {3, 3, 3};
int rowadrT[] = {0, 3, 6};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, nullptr,
diagindH, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0));
mj_deleteData(data);
mj_deleteModel(model);
@@ -424,27 +448,12 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparseLower) {
int rownnzT[] = {3, 3, 3};
int rowadrT[] = {0, 3, 6};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, nullptr,
nullptr, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 0);
EXPECT_THAT(rownnzH, ElementsAre(1, 2, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 1, 3));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, nullptr);
EXPECT_THAT(matH, ElementsAre(12, 0, 0, 0, 6, 0, 12, 3, 14));
EXPECT_THAT(colindH, ElementsAre(0, 0, 0, 0, 1, 0, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(1, 2, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(12, 0, 0, 0, 6, 0, 12, 3, 14));
mj_deleteData(data);
mj_deleteModel(model);
@@ -468,31 +477,13 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
int rownnzT[] = {3, 3, 3};
int rowadrT[] = {0, 3, 6};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, nullptr,
diagindH, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(12, 0, 12, 0, 6, 3, 12, 3, 14));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(diagindH, ElementsAre(0, 4, 8));
EXPECT_THAT(dense, ElementsAre(12, 0, 12, 0, 6, 3, 12, 3, 14));
mj_deleteData(data);
mj_deleteModel(model);
@@ -516,31 +507,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
int rownnzT[] = {2, 2, 0};
int rowadrT[] = {0, 2, 4};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0));
EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 4, 35, 0, 0, 0, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 1, 0, 2, 0, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 1));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(66, 4, 0, 4, 35, 0, 0, 0, 0));
mj_deleteData(data);
mj_deleteModel(model);
@@ -564,32 +539,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3b) {
int rownnzT[] = {2, 2, 1};
int rowadrT[] = {0, 2, 4};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(2, 3, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 2, 5));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(26, 2, 0, 2, 13, 12, 12, 16, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 1, 2, 1, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 3, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(diagindH, ElementsAre(0, 4, 7));
EXPECT_THAT(dense, ElementsAre(26, 2, 0, 2, 13, 12, 0, 12, 16));
mj_deleteData(data);
mj_deleteModel(model);
@@ -613,32 +571,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
int rownnzT[] = {2, 0, 2};
int rowadrT[] = {0, 2, 2};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4};
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 0, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 2, 2));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 0, 0, 0, 4, 35, 0));
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 1, 0, 0, 0, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(66, 0, 4, 0, 0, 0, 4, 0, 35));
mj_deleteData(data);
mj_deleteModel(model);
@@ -662,30 +603,13 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
int rownnzT[] = {2, 1, 1};
int rowadrT[] = {0, 2, 3};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, nullptr,
diagindH, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(5, 6, 4, 6, 9, 0, 4, 16, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(5, 6, 4, 6, 9, 0, 4, 0, 16));
mj_deleteData(data);
mj_deleteModel(model);
@@ -709,30 +633,13 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
int rownnzT[] = {1, 1, 2};
int rowadrT[] = {0, 1, 2};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, nullptr,
diagindH, data);
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 2, 3));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(1, 2, 0, 4, 0, 0, 2, 13, 0));
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 1, 0, 0, 0, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(diagindH, ElementsAre(0, 3, 7));
EXPECT_THAT(dense, ElementsAre(1, 0, 2, 0, 4, 0, 2, 0, 13));
mj_deleteData(data);
mj_deleteModel(model);
@@ -756,31 +663,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
int rownnzT[] = {2, 2};
int rowadrT[] = {0, 2};
mjtNum matH[] = {0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0};
int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0};
int diagindH[] = {0, 0};
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 2, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
int diagindH[2];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 2, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 2));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 2);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 2, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 4, 35));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 1));
EXPECT_THAT(rownnzH, ElementsAre(2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 2));
EXPECT_THAT(dense, ElementsAre(66, 4, 4, 35));
mj_deleteData(data);
mj_deleteModel(model);
@@ -803,31 +694,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
int rownnzT[] = {2, 1, 1};
int rowadrT[] = {0, 2, 3};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 2, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 2, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(14, 18, 8, 18, 27, 0, 8, 32, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(14, 18, 8, 18, 27, 0, 8, 0, 32));
mj_deleteData(data);
mj_deleteModel(model);
@@ -851,31 +726,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
int rownnzT[] = {3, 3, 3};
int rowadrT[] = {0, 3, 6};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, nullptr, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(69, 77, 80, 77, 99, 108, 80, 108, 120));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(69, 77, 80, 77, 99, 108, 80, 108, 120));
mj_deleteData(data);
mj_deleteModel(model);
@@ -900,31 +759,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
int rowadrT[] = {0, 3, 6};
int rowsuperT[] = {2, 1, 0};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {1, 2, 1};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(18, 17, 14, 17, 23, 19, 14, 19, 18));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(18, 17, 14, 17, 23, 19, 14, 19, 18));
mj_deleteData(data);
mj_deleteModel(model);
@@ -949,31 +792,15 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
int rowadrT[] = {0, 1, 3};
int rowsuperT[] = {0, 1, 0};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
int diagindH[3];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 3, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(dense, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10));
mj_deleteData(data);
mj_deleteModel(model);
@@ -998,33 +825,16 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
int rowadrT[] = {0, 1, 2, 4};
int rowsuperT[] = {1, 0, 1, 0};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0};
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 4, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
int diagindH[4];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 4, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4));
EXPECT_THAT(rowadrH, ElementsAre(0, 4, 8, 12));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 4);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 4, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, diagindH);
EXPECT_THAT(matH,
EXPECT_THAT(dense,
ElementsAre(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 10, 1, 1, 10, 10));
EXPECT_THAT(colindH,
ElementsAre(0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3));
EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4));
EXPECT_THAT(rowadrH, ElementsAre(0, 4, 8, 12));
mj_deleteData(data);
mj_deleteModel(model);
@@ -1049,35 +859,16 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
int rowadrT[] = {0, 3, 6, 6, 6};
int rowsuperT[] = {1, 0, 2, 1, 0};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0, 0};
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 5, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
int diagindH[5];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 3, 5, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, diag,
diagindH, data);
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0, 0, 0));
EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4, 4, 4));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 5);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 5, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 0,
3, 0, 0, 0, 0, 4, 0, 0, 0, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 1, 1, 1));
EXPECT_THAT(rowadrH, ElementsAre(0, 5, 10, 15, 20));
EXPECT_THAT(dense, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0));
mj_deleteData(data);
mj_deleteModel(model);
@@ -1100,40 +891,305 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
int rowadrT[] = {0, 1, 2, 3, 4, 5, 6};
int rowsuperT[] = {3, 2, 1, 0, 2, 1, 0};
mjtNum matH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0, 0, 0, 0};
// test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 7, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7));
EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42));
// test computation
mju_sqrMatTDUncompressedInit(rowadrH, 7);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 1, 7, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, diagindH);
int diagindH[7];
std::vector<mjtNum> dense;
SqrMatTDSplitCol(dense, 1, 7, mat, rownnz, rowadr, colind,
matT, rownnzT, rowadrT, colindT, rowsuperT, nullptr,
diagindH, data);
EXPECT_THAT(
matH, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2,
2, 2, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 2, 2, 2,
2, 4, 4, 4, 2, 2, 2, 2, 4, 4, 4));
EXPECT_THAT(colindH,
ElementsAre(0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3,
4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 0,
1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6));
dense, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1,
1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 4, 4,
4, 2, 2, 2, 2, 4, 4, 4, 2, 2, 2, 2, 4, 4, 4));
EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7));
EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42));
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparseSymbolic) {
// Simple dense 2x2 matrix:
// 1 2
// M = 3 4
//
// M'M (lower triangle) should have 3 elements: (0,0), (1,0), (1,1)
mjModel* model = LoadModelFromString("<mujoco/>");
mjData* data = mj_makeData(model);
// M in CSR: row 0 has cols 0,1; row 1 has cols 0,1
int colind[] = {0, 1, 0, 1};
int rownnz[] = {2, 2};
int rowadr[] = {0, 2};
// compute transpose using mju_transposeSparse
mjtNum mat[] = {1, 2, 3, 4};
mjtNum matT[4];
int colindT[4];
int rownnzT[2];
int rowadrT[2];
mju_transposeSparse(matT, mat, 2, 2, rownnzT, rowadrT, colindT, nullptr,
rownnz, rowadr, colind);
// use old function as ground truth
int rownnzH_expected[] = {0, 0};
int rowadrH_expected[] = {0, 0};
int nnz_expected = mju_sqrMatTDSparseCount(
rownnzH_expected, rowadrH_expected, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, nullptr, data, /*flg_upper=*/0);
// verify: lower triangle should have 3 elements: (0,0), (1,0), (1,1)
EXPECT_EQ(nnz_expected, 3);
EXPECT_THAT(rownnzH_expected, ElementsAre(1, 2));
EXPECT_THAT(rowadrH_expected, ElementsAre(0, 1));
// test count mode of new function
int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0};
int nnz = mju_sqrMatTDSparseSymbolic(rownnzH, rowadrH, nullptr, nullptr,
2, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, nullptr, data);
EXPECT_EQ(nnz, nnz_expected);
EXPECT_THAT(rownnzH, ElementsAre(rownnzH_expected[0], rownnzH_expected[1]));
EXPECT_THAT(rowadrH, ElementsAre(rowadrH_expected[0], rowadrH_expected[1]));
// test fill mode
std::vector<int> colindH(nnz, -1);
mju_sqrMatTDSparseSymbolic(rownnzH, rowadrH, colindH.data(), nullptr, 2, 2,
rownnz, rowadr, colind, rownnzT, rowadrT,
colindT, nullptr, data);
// verify: row 0 should have {0}, row 1 should have {0, 1}
EXPECT_THAT(colindH, ElementsAre(0, 0, 1));
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparseSymbolicUpper) {
// Test flg_upper=1: count both lower and upper triangle
// Same matrix as previous test
mjModel* model = LoadModelFromString("<mujoco/>");
mjData* data = mj_makeData(model);
int colind[] = {0, 1, 0, 1};
int rownnz[] = {2, 2};
int rowadr[] = {0, 2};
mjtNum mat[] = {1, 2, 3, 4};
mjtNum matT[4];
int colindT[4];
int rownnzT[2];
int rowadrT[2];
mju_transposeSparse(matT, mat, 2, 2, rownnzT, rowadrT, colindT, nullptr,
rownnz, rowadr, colind);
// use old function as ground truth with flg_upper=1
int rownnzH_expected[] = {0, 0};
int rowadrH_expected[] = {0, 0};
int nnz_expected = mju_sqrMatTDSparseCount(
rownnzH_expected, rowadrH_expected, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, nullptr, data, /*flg_upper=*/1);
// test new function with diagind (upper triangle)
int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0};
int diagindH[] = {0, 0};
int nnz = mju_sqrMatTDSparseSymbolic(rownnzH, rowadrH, nullptr, diagindH,
2, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, nullptr, data);
EXPECT_EQ(nnz, nnz_expected);
EXPECT_THAT(rownnzH, ElementsAre(rownnzH_expected[0], rownnzH_expected[1]));
EXPECT_THAT(rowadrH, ElementsAre(rowadrH_expected[0], rowadrH_expected[1]));
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparseSymbolicSupernode) {
// Test supernode exploitation with a matrix that has supernodes
// M has two rows with identical sparsity pattern
mjModel* model = LoadModelFromString("<mujoco/>");
mjData* data = mj_makeData(model);
// 3x2 matrix where rows 1 and 2 have same pattern
// 1 0
// M = 2 3
// 4 5
int colind[] = {0, 0, 1, 0, 1};
int rownnz[] = {1, 2, 2};
int rowadr[] = {0, 1, 3};
mjtNum mat[] = {1, 2, 3, 4, 5};
mjtNum matT[5];
int colindT[5];
int rownnzT[2];
int rowadrT[2];
mju_transposeSparse(matT, mat, 3, 2, rownnzT, rowadrT, colindT, nullptr,
rownnz, rowadr, colind);
// compute rowsuperT
int rowsuperT[2];
mju_superSparse(2, rowsuperT, rownnzT, rowadrT, colindT);
// use old function as ground truth
int rownnzH_expected[] = {0, 0};
int rowadrH_expected[] = {0, 0};
int nnz_expected = mju_sqrMatTDSparseCount(
rownnzH_expected, rowadrH_expected, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, rowsuperT, data, /*flg_upper=*/0);
// test new function with supernodes
int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0};
int nnz = mju_sqrMatTDSparseSymbolic(rownnzH, rowadrH, nullptr, nullptr,
3, 2, rownnz, rowadr, colind, rownnzT,
rowadrT, colindT, rowsuperT, data);
EXPECT_EQ(nnz, nnz_expected);
EXPECT_THAT(rownnzH, ElementsAre(rownnzH_expected[0], rownnzH_expected[1]));
EXPECT_THAT(rowadrH, ElementsAre(rowadrH_expected[0], rowadrH_expected[1]));
// test fill mode with supernodes
std::vector<int> colindH(nnz, -1);
mju_sqrMatTDSparseSymbolic(rownnzH, rowadrH, colindH.data(), nullptr, 3, 2,
rownnz, rowadr, colind, rownnzT, rowadrT,
colindT, rowsuperT, data);
// verify all filled
for (int i = 0; i < nnz; i++) {
EXPECT_GE(colindH[i], 0) << "colindH[" << i << "] not filled";
}
// verify numeric phase with supernodes
std::vector<mjtNum> resH(nnz);
mjtNum diag[] = {1, 1, 1, 1, 1}; // dummy diagonal
mju_sqrMatTDSparseNumeric(resH.data(), 2, rownnzH, rowadrH, colindH.data(),
nullptr, mat, rownnz, rowadr, colind, matT, rownnzT,
rowadrT, colindT, rowsuperT, diag, data);
// ground truth numeric
std::vector<mjtNum> res_expected(4);
std::vector<int> colindH_expected(4);
int rownnzH_exp[] = {0, 0};
int rowadrH_exp[] = {0, 2};
mju_sqrMatTDSparse(res_expected.data(), mat, matT, diag, 3, 2, rownnzH_exp,
rowadrH_exp, colindH_expected.data(), rownnz, rowadr,
colind, nullptr, rownnzT, rowadrT, colindT, rowsuperT,
data, nullptr);
// compare values (sparse result vs sparse ground truth)
for (int r = 0; r < 2; r++) {
for (int i = 0; i < rownnzH[r]; i++) {
// find matching col in ground truth
int c = colindH[rowadrH[r] + i];
mjtNum val = resH[rowadrH[r] + i];
bool found = false;
for (int j = 0; j < rownnzH_exp[r]; j++) {
if (colindH_expected[rowadrH_exp[r] + j] == c) {
EXPECT_NEAR(val, res_expected[rowadrH_exp[r] + j], 1e-14);
found = true;
break;
}
}
EXPECT_TRUE(found) << "Column " << c
<< " not found in ground truth for row " << r;
}
}
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparseNumeric) {
// Test numeric phase using symbolic phase + existing function as ground truth
// 1 2
// M = 3 4
mjModel* model = LoadModelFromString("<mujoco/>");
mjData* data = mj_makeData(model);
int colind[] = {0, 1, 0, 1};
int rownnz[] = {2, 2};
int rowadr[] = {0, 2};
mjtNum mat[] = {1, 2, 3, 4};
// compute transpose
mjtNum matT[4];
int colindT[4];
int rownnzT[2];
int rowadrT[2];
mju_transposeSparse(matT, mat, 2, 2, rownnzT, rowadrT, colindT, nullptr,
rownnz, rowadr, colind);
// compute supernodes
int rowsuperT[2];
mju_superSparse(2, rowsuperT, rownnzT, rowadrT, colindT);
mjtNum diag[] = {2, 3}; // diagonal weighting matrix
// test both diagind cases: lower-only (diagind=NULL) and both triangles
// (diagind!=NULL)
for (int use_diagind = 0; use_diagind <= 1; use_diagind++) {
// compute sparsity pattern using symbolic phase
int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0};
int diagindH[] = {0, 0};
int nnz = mju_sqrMatTDSparseSymbolic(
rownnzH, rowadrH, nullptr, use_diagind ? diagindH : nullptr, 2, 2,
rownnz, rowadr, colind, rownnzT, rowadrT, colindT, nullptr, data);
std::vector<int> colindH(nnz);
mju_sqrMatTDSparseSymbolic(
rownnzH, rowadrH, colindH.data(), use_diagind ? diagindH : nullptr, 2,
2, rownnz, rowadr, colind, rownnzT, rowadrT, colindT, nullptr, data);
// compute values using numeric phase
std::vector<mjtNum> resH(nnz);
mju_sqrMatTDSparseNumeric(resH.data(), 2, rownnzH, rowadrH,
colindH.data(), use_diagind ? diagindH : nullptr,
mat, rownnz, rowadr, colind, matT, rownnzT,
rowadrT, colindT, rowsuperT, diag, data);
// compute ground truth using existing mju_sqrMatTDSparse
// use uncompressed storage to give the old function enough room
std::vector<mjtNum> res_expected(4); // 2x2 uncompressed
std::vector<int> colindH_expected(4);
int rownnzH_exp[] = {0, 0};
int rowadrH_exp[] = {0, 2};
int diagind_exp[] = {0, 0};
mju_sqrMatTDSparse(res_expected.data(), mat, matT, diag, 2, 2, rownnzH_exp,
rowadrH_exp, colindH_expected.data(), rownnz, rowadr,
colind, nullptr, rownnzT, rowadrT, colindT, nullptr,
data, use_diagind ? diagind_exp : nullptr);
// check that rownnz matches (nnz may differ due to compressed vs
// uncompressed storage)
EXPECT_EQ(rownnzH[0], rownnzH_exp[0])
<< "rownnz[0] mismatch for use_diagind=" << use_diagind;
EXPECT_EQ(rownnzH[1], rownnzH_exp[1])
<< "rownnz[1] mismatch for use_diagind=" << use_diagind;
// compare column indices and values for each row
for (int r = 0; r < 2; r++) {
for (int j = 0; j < rownnzH[r]; j++) {
int idx = rowadrH[r] + j;
int idx_exp = rowadrH_exp[r] + j;
EXPECT_EQ(colindH[idx], colindH_expected[idx_exp])
<< "colind mismatch at row " << r << " pos " << j
<< " for use_diagind=" << use_diagind;
EXPECT_NEAR(resH[idx], res_expected[idx_exp], 1e-10)
<< "value mismatch at row " << r << " pos " << j
<< " for use_diagind=" << use_diagind;
}
}
}
mj_deleteData(data);
mj_deleteModel(model);