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

Move island-specific sparse matrices from arena to stack.

Browse Source 
PiperOrigin-RevId: 923850345
Change-Id: I9683d7554b15b7cd8c45a8dce7814640aa266452
This commit is contained in:
Yuval Tassa
2026-05-30 03:07:31 -07:00
committed by Copybara-Service
parent 5d782a2bb8
commit 96bf8aea81
11 changed files with 122 additions and 402 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_island.c
-35
View File
@@ -526,15 +526,6 @@ void mj_island(const mjModel* m, mjData* d) {
d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]];
}
// inertia: block-diagonalize both iLD <- qLD and iM <- M
mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind,
nidof, nisland,
d->map_idof2dof, d->map_dof2idof,
d->island_idofadr, d->island_idofadr,
d->iM, d->M);
mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof);
// ------------------------------------- constraints ---------------------------------------------
@@ -578,32 +569,6 @@ void mj_island(const mjModel* m, mjData* d) {
// SHOULD NOT OCCUR
if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount");
// dense: block-diagonalize Jacobian
if (!mj_isSparse(m)) {
mju_blockDiag(d->iefc_J, d->efc_J,
nv, nidof, nisland,
d->map_iefc2efc, d->map_idof2dof,
d->island_nefc, d->island_nv,
d->island_iefcadr, d->island_idofadr);
}
// sparse
else {
// block-diagonalize Jacobian
mju_blockDiagSparse(d->iefc_J, d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind,
d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
nefc, nisland,
d->map_iefc2efc, d->map_dof2idof,
d->island_iefcadr, d->island_idofadr, NULL, NULL);
// recompute rowsuper per island
for (int island=0; island < nisland; island++) {
int adr = d->island_iefcadr[island];
mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr,
d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind);
}
}
// copy position-dependent efc vectors required by solver
mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc);
mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc);
src/engine/engine_print.c
+25 -22
View File
@@ -190,11 +190,12 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
// print block-diagonal dense matrix, embedded in a larger matrix
static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
static void printBlockArray(const char* str, const mjtNum* data, int nc,
int nisland, const int* island_nr, const int* island_nc,
const int* island_r, const int* island_c,
const int* map_r, const int* map_c,
FILE* fp, const char* float_format) {
if (!data || !nr || !nc) {
if (!data || !nisland) {
return;
}
@@ -209,7 +210,6 @@ static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
int bnc = island_nc[b];
int r_start = island_r[b];
int c_start = island_c[b];
const mjtNum* data_ptr = data + r_start * nc;
// print rows for this block
for (int r_block = 0; r_block < bnr; r_block++) {
@@ -220,10 +220,13 @@ static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
fprintf(fp, " ");
}
int row = map_r[r_start + r_block];
// block data
for (int c = 0; c < bnc; c++) {
int col = map_c[c_start + c];
fprintf(fp, " ");
fprintf(fp, float_format, *data_ptr++);
fprintf(fp, float_format, data[row * nc + col]);
}
// trailing dots
@@ -318,6 +321,7 @@ void mj_printBlockSparsity(const char* str, int nr, int nc, int nisland,
const int* island_col_offset,
const int* entity_island,
const int* map_row_to_entity,
const int* map_col_to_entity,
const int* rownnz, const int* rowadr, const int* colind,
const int* rowsuper, FILE* fp) {
// if no rows / columns, or too many columns to be visually useful, return
@@ -343,27 +347,25 @@ void mj_printBlockSparsity(const char* str, int nr, int nc, int nisland,
int c_start = island_col_offset[island];
int bnc = island_block_ncols[island];
int current_nnz = 0;
int adr = rowadr[r];
int adr = rowadr[entity_r];
int nnz = rownnz[entity_r];
char nz_char = (island < 10) ? ('0' + island) : 'x';
for (int c = 0; c < nc; c++) { // c is the global column index
for (int c = 0; c < nc; c++) { // c is the block-space column index
bool nonzero = false;
if (c >= c_start && c < c_start + bnc) {
int c_block = c - c_start; // c_block is the island-local column index
// search for c_block in colind for the current row r
while (current_nnz < rownnz[r] && colind[adr + current_nnz] < c_block) {
current_nnz++;
}
if (current_nnz < rownnz[r] && colind[adr + current_nnz] == c_block) {
nonzero = true;
int target_col = map_col_to_entity[c];
for (int i = 0; i < nnz; i++) {
if (colind[adr + i] == target_col) {
nonzero = true;
break;
}
}
}
fprintf(fp, "%c", nonzero ? nz_char : ' ');
}
fprintf(fp, " |");
if (rowsuper && rowsuper[r] > 0) fprintf(fp, " %d", rowsuper[r]);
if (rowsuper && rowsuper[entity_r] > 0) fprintf(fp, " %d", rowsuper[entity_r]);
fprintf(fp, "\n");
}
for (int c = 0; c < nc + 2; c++) fprintf(fp, "-");
@@ -1477,8 +1479,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
mj_printBlockSparsity("iM: block-diagonal inertia (nnzs are island ids)",
d->nidof, d->nidof, d->nisland,
d->island_nv, d->island_idofadr,
d->dof_island, d->map_idof2dof,
d->iM_rownnz, d->iM_rowadr, d->iM_colind, NULL, fp);
d->dof_island, d->map_idof2dof, d->map_idof2dof,
m->M_rownnz, m->M_rowadr, m->M_colind, NULL, fp);
}
if (!mju_isZero(d->qHDiagInv, m->nv)) {
@@ -1555,9 +1557,10 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
if (!mj_isSparse(m)) {
printArray2d("EFC_J", d->nefc, m->nv, d->efc_J, fp, float_format);
if (d->nisland) {
printBlockArray("IEFC_J", d->iefc_J, d->nefc, d->nidof,
printBlockArray("IEFC_J", d->efc_J, m->nv,
d->nisland, d->island_nefc, d->island_nv,
d->island_iefcadr, d->island_idofadr,
d->map_iefc2efc, d->map_idof2dof,
fp, float_format);
}
printArray2d("EFC_AR", d->nefc, d->nefc, d->efc_AR, fp, float_format);
@@ -1577,9 +1580,9 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
mj_printBlockSparsity("IEFC_J: block-diagonalized constraint Jacobian (nnzs are island ids)",
d->nefc, d->nidof, d->nisland,
d->island_nv, d->island_idofadr,
d->efc_island, d->map_iefc2efc,
d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind,
d->iefc_J_rowsuper, fp);
d->efc_island, d->map_iefc2efc, d->map_idof2dof,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_J_rowsuper, fp);
}
if (mj_isDual(m)) {
src/engine/engine_solver.c
+92 -35
View File
@@ -877,12 +877,12 @@ typedef struct {
mjtNum* qacc;
// inertia
const int* M_rownnz;
const int* M_rowadr;
const int* M_colind;
const mjtNum* M;
const mjtNum* qLD;
const mjtNum* qLDiagInv;
int* M_rownnz;
int* M_rowadr;
int* M_colind;
mjtNum* M;
mjtNum* qLD;
mjtNum* qLDiagInv;
// efc arrays
const mjtNum* efc_D;
@@ -895,11 +895,11 @@ typedef struct {
int* efc_state;
// Jacobians
const int* J_rownnz;
const int* J_rowadr;
const int* J_rowsuper;
const int* J_colind;
const mjtNum* J;
int* J_rownnz;
int* J_rowadr;
int* J_rowsuper;
int* J_colind;
mjtNum* J;
int* JT_rownnz;
int* JT_rowadr;
int* JT_rowsuper;
@@ -1031,14 +1031,6 @@ static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* c
ctx->qacc_smooth = d->iacc_smooth + idofadr;
ctx->qacc = d->iacc + idofadr;
// inertia
ctx->M_rownnz = d->iM_rownnz + idofadr;
ctx->M_rowadr = d->iM_rowadr + idofadr;
ctx->M_colind = d->iM_colind;
ctx->M = d->iM;
ctx->qLD = d->iLD;
ctx->qLDiagInv = d->iLDiagInv + idofadr;
// efc arrays
int iefcadr = d->island_iefcadr[island];
ctx->efc_D = d->iefc_D + iefcadr;
@@ -1049,32 +1041,44 @@ static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* c
ctx->efc_type = d->iefc_type + iefcadr;
ctx->efc_force = d->iefc_force + iefcadr;
ctx->efc_state = d->iefc_state + iefcadr;
// Jacobians
if (!ctx->is_sparse) {
ctx->J = d->iefc_J + d->nidof * iefcadr;
} else {
ctx->J_rownnz = d->iefc_J_rownnz + iefcadr;
ctx->J_rowadr = d->iefc_J_rowadr + iefcadr;
ctx->J_rowsuper = d->iefc_J_rowsuper + iefcadr;
ctx->J_colind = d->iefc_J_colind;
ctx->J = d->iefc_J;
ctx->nJ = ctx->J_rowadr[ctx->nefc-1] + ctx->J_rownnz[ctx->nefc-1]
- ctx->J_rowadr[0];
}
}
}
// allocate fixed-size arrays in mjPrimalContext
// mj_{mark/free}Stack in calling function!
static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
static void PrimalAllocate(const mjModel* m, mjData* d, mjPrimalContext* ctx, int flg_Newton) {
// local sizes and flags
int nv = ctx->nv;
int nefc = ctx->nefc;
int nJ = ctx->is_sparse ? ctx->nJ : 0;
int is_sparse = ctx->is_sparse;
int is_elliptic = ctx->is_elliptic;
int nJ = is_sparse ? d->nJ : 0;
// compute island matrix sizes if needed
int nC = 0;
if (ctx->island >= 0) {
// count nC: number of nonzeros in M block of island (always sparse)
int island = ctx->island;
int idofadr = d->island_idofadr[island];
for (int i = 0; i < nv; i++) {
int dof = d->map_idof2dof[idofadr + i];
nC += m->M_rownnz[dof];
}
// count nJ: number of nonzeros in J block of island (sparse or dense)
if (is_sparse) {
nJ = 0;
int iefcadr = d->island_iefcadr[island];
for (int i = 0; i < nefc; i++) {
int efc = d->map_iefc2efc[iefcadr + i];
nJ += d->efc_J_rownnz[efc];
}
} else {
nJ = nefc * nv;
}
ctx->nJ = nJ;
}
// compute mjtNum block size
size_t nNum = 5*nefc + 5*nv; // common arrays
@@ -1090,6 +1094,11 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
nNum += 3*nv; // CG arrays
}
// add island matrix sizes
if (ctx->island >= 0) {
nNum += 2 * nC + nv + nJ; // iM, iLD, iLDiagInv, iefc_J
}
// compute int block size
size_t nInt = nefc; // oldstate
if (is_sparse) {
@@ -1097,10 +1106,58 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
if (flg_Newton) nInt += 8*nv; // Newton sparse
}
// add island matrix sizes
if (ctx->island >= 0) {
nInt += 2 * nv + nC; // iM_{rownnz, rowadr, colind}
if (is_sparse) {
nInt += 3 * nefc + nJ; // iefc_J_{rownnz, rowadr, rowsuper, colind}
}
}
// allocate mjtNum and int blocks
mjtNum* numblock = mjSTACKALLOC(d, nNum, mjtNum);
int* intblock = mjSTACKALLOC(d, nInt, int);
// populate island matrices if needed
if (ctx->island >= 0) {
int island = ctx->island;
int idofadr = d->island_idofadr[island];
int iefcadr = d->island_iefcadr[island];
ctx->M_rownnz = intblock; intblock += nv;
ctx->M_rowadr = intblock; intblock += nv;
ctx->M_colind = intblock; intblock += nC;
ctx->M = numblock; numblock += nC;
ctx->qLD = numblock; numblock += nC;
ctx->qLDiagInv = numblock; numblock += nv;
mju_blockSparse(ctx->qLD, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind,
nv, d->map_idof2dof + idofadr, d->map_dof2idof,
d->island_idofadr[island], 0, ctx->M, d->M);
mju_gather(ctx->qLDiagInv, d->qLDiagInv, d->map_idof2dof + idofadr, nv);
ctx->J = numblock; numblock += nJ;
if (!is_sparse) {
mju_block(ctx->J, d->efc_J, m->nv, nv, nefc,
d->map_iefc2efc + iefcadr, d->map_idof2dof + idofadr);
} else {
ctx->J_rownnz = intblock; intblock += nefc;
ctx->J_rowadr = intblock; intblock += nefc;
ctx->J_rowsuper = intblock; intblock += nefc;
ctx->J_colind = intblock; intblock += nJ;
mju_blockSparse(ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
nefc, d->map_iefc2efc + iefcadr, d->map_dof2idof,
d->island_idofadr[island], 0, NULL, NULL);
mju_superSparse(nefc, ctx->J_rowsuper, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind);
}
}
// carve mjtNum block
ctx->Jaref = numblock; numblock += nefc;
ctx->Jv = numblock; numblock += nefc;
@@ -1962,7 +2019,7 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
// make context
PrimalPointers(m, d, &ctx, island);
PrimalAllocate(d, &ctx, flg_Newton);
PrimalAllocate(m, d, &ctx, flg_Newton);
// local copies
int nv = ctx.nv;