diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 9bbece11..170d1f3c 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -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++) { diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index b03ad913..b02302fc 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -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, diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 9d78e409..05960567 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -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 #include @@ -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; diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 4168cb8b..1350ccff 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -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++; diff --git a/test/benchmark/chol_benchmark_test.cc b/test/benchmark/chol_benchmark_test.cc index 769f990d..004e0c2a 100644 --- a/test/benchmark/chol_benchmark_test.cc +++ b/test/benchmark/chol_benchmark_test.cc @@ -68,13 +68,15 @@ struct HessianData { // D diagonal std::vector 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 +HessianData& GetHessianData() { + static HessianData data; + static bool initialized = false; + if (!initialized) { + mjModel* m = GetModel(); + mjData* d = mj_makeData(m); + data.Setup(m, d); + mj_deleteData(d); + initialized = true; + } + return data; +} + // old implementation benchmark template static void BM_chol_old(benchmark::State& state) { mjModel* m = GetModel(); mjData* d = mj_makeData(m); - HessianData hd; - hd.Setup(m, d); + HessianData& hd = GetHessianData(); std::vector L_work(hd.nL); std::vector L_colind_work(hd.nL); @@ -239,8 +254,7 @@ static void BM_chol_symbolic(benchmark::State& state) { mjModel* m = GetModel(); mjData* d = mj_makeData(m); - HessianData hd; - hd.Setup(m, d); + HessianData& hd = GetHessianData(); std::vector L_colind_work(hd.nL); std::vector LT_rownnz_work(hd.nv); @@ -266,8 +280,7 @@ static void BM_chol_numeric(benchmark::State& state) { mjModel* m = GetModel(); mjData* d = mj_makeData(m); - HessianData hd; - hd.Setup(m, d); + HessianData& hd = GetHessianData(); std::vector L_work(hd.nL); std::vector L_colind_work(hd.nL); @@ -371,9 +384,10 @@ template static void BM_update_old(benchmark::State& state) { mjModel* m = GetModel(); mjData* d = mj_makeData(m); + mj_resetDataKeyframe(m, d, 0); + mj_forward(m, d); - HessianData hd; - hd.Setup(m, d); + HessianData& hd = GetHessianData(); int nv = hd.nv; @@ -433,9 +447,10 @@ template static void BM_update_new(benchmark::State& state) { mjModel* m = GetModel(); mjData* d = mj_makeData(m); + mj_resetDataKeyframe(m, d, 0); + mj_forward(m, d); - HessianData hd; - hd.Setup(m, d); + HessianData& hd = GetHessianData(); int nv = hd.nv; diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index 1537c398..a82be98e 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -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 +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 static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func, Supernode super) { - static mjModel* m = LoadModelFromPath("humanoid/humanoid100.xml"); + static mjModel* m = LoadModelFromPath(ModelPath()); // 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(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(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(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(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(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(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(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(state, &mju_transposeSparse, + Supernode::Inline); +} +BENCHMARK(BM_transposeSparse_100H_superinline); } // namespace } // namespace mujoco diff --git a/test/benchmark/sqrmat_benchmark_test.cc b/test/benchmark/sqrmat_benchmark_test.cc new file mode 100644 index 00000000..9e6896d4 --- /dev/null +++ b/test/benchmark/sqrmat_benchmark_test.cc @@ -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 +#include + +#include "benchmark/benchmark.h" +#include +#include +#include +#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 J; + std::vector J_rownnz, J_rowadr, J_colind, J_rowsuper; + + // J' (transpose) + std::vector JT; + std::vector JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper; + + // D (diagonal weights for constraints) + std::vector D; + + // M structure (mass matrix, lower triangle) + std::vector 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 +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 +mjModel* GetModel() { + static mjModel* m = LoadModelFromPath(ModelPath()); + m->opt.jacobian = mjJAC_SPARSE; + m->opt.solver = mjSOL_NEWTON; + m->opt.disableflags |= mjDSBL_ISLAND; + return m; +} + +template +SparseTestData& GetData() { + static SparseTestData data; + static bool initialized = false; + if (!initialized) { + mjModel* m = GetModel(); + 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 +static void BM_sqrMatTD_impl(benchmark::State& state, SqrMatTDVariant variant) { + SparseTestData& data = GetData(); + mjModel* m = GetModel(); + 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 H; + std::vector H_rownnz(nv); + std::vector H_rowadr(nv); + std::vector H_colind; + std::vector 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(state, SqrMatTDVariant::kBaseline); +} +BENCHMARK(BM_sqrMatTD_2H100_baseline); + +void BM_sqrMatTD_2H100_row(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kRow); +} +BENCHMARK(BM_sqrMatTD_2H100_row); + +void BM_sqrMatTD_2H100_col(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kCol); +} +BENCHMARK(BM_sqrMatTD_2H100_col); + +void BM_sqrMatTD_2H100_splitCol(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kSplitCol); +} +BENCHMARK(BM_sqrMatTD_2H100_splitCol); + +void BM_sqrMatTD_100H_baseline(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kBaseline); +} +BENCHMARK(BM_sqrMatTD_100H_baseline); + +void BM_sqrMatTD_100H_row(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kRow); +} +BENCHMARK(BM_sqrMatTD_100H_row); + +void BM_sqrMatTD_100H_col(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(state, SqrMatTDVariant::kCol); +} +BENCHMARK(BM_sqrMatTD_100H_col); + +void BM_sqrMatTD_100H_splitCol(benchmark::State& state) { + MujocoErrorTestGuard guard; + BM_sqrMatTD_impl(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; +} diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index be74f9dd..67eadbbc 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -14,10 +14,11 @@ // Tests for engine/engine_util_sparse.c -#include - #include "src/engine/engine_util_sparse.h" +#include +#include + #include #include #include @@ -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& 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 H_rownnz(nc, 0); + std::vector 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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(""); + 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 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(""); + 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(""); + 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 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 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 res_expected(4); + std::vector 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(""); + 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 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 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 res_expected(4); // 2x2 uncompressed + std::vector 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);