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

Refactor utility functions for block extraction.

Browse Source 
PiperOrigin-RevId: 908744392
Change-Id: I8d1323946870c90d5e37b1853697fb35d5291e11
This commit is contained in:
Yuval Tassa
2026-05-01 08:24:06 -07:00
committed by Copybara-Service
parent 3367b64a8e
commit f7d31e06f0
2 changed files with 89 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+75 -47
View File
@@ -1376,11 +1376,29 @@ void mju_sqrMatTDSparse_row(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
}
// block-diagonalize a dense matrix
// extract a single block of a dense matrix
// res output matrix
// mat input matrix
// nc_mat number of columns in mat
// nc_res number of columns in res
// nr number of rows in res
// perm_r reverse permutation of rows (res -> mat)
// perm_c reverse permutation of columns (res -> mat)
void mju_block(mjtNum* restrict res, const mjtNum* restrict mat,
int nc_mat, int nc_res, int nr,
const int* restrict perm_r, const int* restrict perm_c) {
for (int r = 0; r < nr; r++) {
mjtNum* res_r = res + r * nc_res;
const mjtNum* mat_r = mat + perm_r[r] * nc_mat;
mju_gather(res_r, mat_r, perm_c, nc_res);
}
}
// block-diagonalize a dense matrix
// res target matrix
// mat source full matrix
// nc_mat number of columns in source matrix
// nc_res number of columns in the block-diagonal target matrix
// nb number of blocks
// perm_r reverse permutation of rows (res -> mat)
// perm_c reverse permutation of columns (res -> mat)
@@ -1393,16 +1411,54 @@ void mju_blockDiag(mjtNum* restrict res, const mjtNum* restrict mat,
const int* restrict perm_r, const int* restrict perm_c,
const int* restrict block_nr, const int* restrict block_nc,
const int* restrict block_r, const int* restrict block_c) {
for (int b=0; b < nb; b++) {
int bnr = block_nr[b];
int bnc = block_nc[b];
const int* adr_r = perm_r + block_r[b];
const int* adr_c = perm_c + block_c[b];
for (int b = 0; b < nb; b++) {
int adr = nc_res * block_r[b];
for (int r = 0; r < bnr; r++) {
for (int c = 0; c < bnc; c++) {
res[adr++] = mat[nc_mat * adr_r[r] + adr_c[c]];
}
mju_block(res + adr, mat, nc_mat, block_nc[b], block_nr[b],
perm_r + block_r[b], perm_c + block_c[b]);
}
}
// extract a single block of a sparse matrix
// res, res2 target matrix values
// res_rownnz non-zeros in each row of target matrix
// res_rowadr row address of each initial row in target matrix
// res_colind column indices for each extracted value (relative)
// mat, mat2 source matrix values
// rownnz non-zeros in each row of source matrix
// rowadr addresses within the source matrix values
// colind source matrix column indices
// nr number of rows to extract
// perm_r row permutation (maps local row to source row)
// perm_c column permutation (maps source col to local col)
// col_offset subtrahend to shift mapped absolute columns into relative block space
// res_offset rowadr starting offset for the extracted submatrix
void mju_blockSparse(mjtNum* restrict res, int* restrict res_rownnz,
int* restrict res_rowadr, int* restrict res_colind,
const mjtNum* restrict mat, const int* restrict rownnz,
const int* restrict rowadr, const int* restrict colind,
int nr,
const int* restrict perm_r, const int* restrict perm_c,
int col_offset, int res_offset,
mjtNum* restrict res2, const mjtNum* restrict mat2) {
for (int r = 0; r < nr; r++) {
int k = perm_r[r];
int nnz = rownnz[k];
res_rownnz[r] = nnz;
int res_adr = (r == 0) ? res_offset : (res_rowadr[r-1] + res_rownnz[r-1]);
res_rowadr[r] = res_adr;
int* res_colind_r = res_colind + (res_adr - res_offset);
int mat_adr = rowadr[k];
const int* colind_k = colind + mat_adr;
for (int j = 0; j < nnz; j++) {
res_colind_r[j] = perm_c[colind_k[j]] - col_offset;
}
mju_copy(res + (res_adr - res_offset), mat + mat_adr, nnz);
if (mat2 && res2) {
mju_copy(res2 + (res_adr - res_offset), mat2 + mat_adr, nnz);
}
}
}
@@ -1433,43 +1489,15 @@ void mju_blockDiagSparse(mjtNum* restrict res, int* restrict res_rownnz,
const int* restrict perm_r, const int* restrict perm_c,
const int* restrict block_r, const int* restrict block_c,
mjtNum* restrict res2, const mjtNum* restrict mat2) {
int block = 0;
int col_offset = block_c[block];
int row_next = block + 1 < nb ? block_r[block + 1] : nr;
for (int r=0; r < nr; r++) {
// row k in mat goes to row r in res
int k = perm_r[r];
for (int b = 0; b < nb; b++) {
int nr_block = (b + 1 < nb ? block_r[b + 1] : nr) - block_r[b];
int res_adr = (block_r[b] == 0) ? 0 : (res_rowadr[block_r[b] - 1] + res_rownnz[block_r[b] - 1]);
// rownnz
int nnz = rownnz[k];
res_rownnz[r] = nnz;
// rowadr
int res_adr = (r == 0) ? 0 : res_rowadr[r-1] + res_rownnz[r-1];
res_rowadr[r] = res_adr;
// colind
int* res_colind_r = res_colind + res_adr;
mjtNum* res_r = res + res_adr;
int mat_adr = rowadr[k];
const int* colind_k = colind + mat_adr;
const mjtNum* mat_k = mat + mat_adr;
for (int j=0; j < nnz; j++) {
res_colind_r[j] = perm_c[colind_k[j]] - col_offset;
}
// values (dense copy: partial order within block is guaranteed)
mju_copy(res_r, mat_k, nnz);
if (mat2 && res2) {
mju_copy(res2 + res_adr, mat2 + mat_adr, nnz);
}
// end of block reached: update block counter, column offset, next row
if (r + 1 >= row_next && block + 1 < nb) {
block++;
col_offset = block_c[block];
row_next = block + 1 < nb ? block_r[block + 1] : nr;
}
mju_blockSparse(res + res_adr, res_rownnz + block_r[b],
res_rowadr + block_r[b], res_colind + res_adr,
mat, rownnz, rowadr, colind,
nr_block, perm_r + block_r[b], perm_c,
block_c[b], res_adr,
res2 ? res2 + res_adr : NULL, mat2);
}
}
src/engine/engine_util_sparse.h
+14 -2
View File
@@ -158,13 +158,25 @@ MJAPI void mju_sqrMatTDSparseNumeric(
// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);
// extract a single block of a dense matrix
void mju_block(mjtNum* res, const mjtNum* mat, int nc_mat, int nc_res, int nr,
const int* perm_r, const int* perm_c);
// block-diagonalize a dense matrix
MJAPI void mju_blockDiag(mjtNum* res, const mjtNum* mat,
int nc_mat, int nc_res, int nb,
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* block_r, const int* block_c);
// extract a single block of a sparse matrix
void mju_blockSparse(
mjtNum* res, int* res_rownnz, int* res_rowadr, int* res_colind,
const mjtNum* mat, const int* rownnz, const int* rowadr, const int* colind,
int nr,
const int* perm_r, const int* perm_c,
int col_offset, int res_offset,
mjtNum* res2, const mjtNum* mat2);
// block-diagonalize a sparse matrix
MJAPI void mju_blockDiagSparse(
mjtNum* res, int* res_rownnz, int* res_rowadr, int* res_colind,