From 436b5a8e1f35b6acc5c64944b58a4834e8d38f0c Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 12 May 2025 16:18:51 -0700 Subject: [PATCH] Add `mju_addToSymSparse` - Internal engine function to add a symmetric sparse matrix to a dense matrix. - Also minor refactors to related functions. PiperOrigin-RevId: 757954336 Change-Id: I8d8a48bcbbd5d6c618ae097a87df2bc1e6f5ef1d --- src/engine/engine_support.c | 16 +++++------ src/engine/engine_support.h | 4 +-- src/engine/engine_util_blas.c | 22 ++++++++++----- src/engine/engine_util_blas.h | 4 +++ src/engine/engine_util_sparse.c | 37 +++++++++++++++++++++++++- src/engine/engine_util_sparse.h | 11 ++++++++ test/engine/engine_util_sparse_test.cc | 34 +++++++++++++++++++++++ 7 files changed, 109 insertions(+), 19 deletions(-) diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index b7036659..5eb0c9fd 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1089,20 +1089,18 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, // add inertia matrix to sparse destination matrix void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst, - int* rownnz, int* rowadr, int* colind, mjtNum* M, - int* M_rownnz, int* M_rowadr, int* M_colind) { + int* rownnz, int* rowadr, int* colind, const mjtNum* M, + const int* M_rownnz, const int* M_rowadr, const int* M_colind) { int nv = m->nv; mj_markStack(d); + mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum); int* buf_ind = mjSTACKALLOC(d, nv, int); - mjtNum* sparse_buf = mjSTACKALLOC(d, nv, mjtNum); - // add to destination - for (int i=0; i < nv; i++) { - rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1, - rownnz[i], M_rownnz[i], colind + rowadr[i], - M_colind + M_rowadr[i], sparse_buf, buf_ind); - } + mju_addToMatSparse(dst, rownnz, rowadr, colind, nv, + M, M_rownnz, M_rowadr, M_colind, + buf_val, buf_ind); + mj_freeStack(d); } diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index 818e9622..5933f6a7 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -137,8 +137,8 @@ MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, // add inertia matrix to sparse destination matrix MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst, - int* rownnz, int* rowadr, int* colind, mjtNum* M, - int* M_rownnz, int* M_rowadr, int* M_colind); + int* rownnz, int* rowadr, int* colind, const mjtNum* M, + const int* M_rownnz, const int* M_rowadr, const int* M_colind); // add inertia matrix to dense destination matrix MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst); diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c index c870c6fb..95055d72 100644 --- a/src/engine/engine_util_blas.c +++ b/src/engine/engine_util_blas.c @@ -843,9 +843,9 @@ void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, } - -// compute M'*diag*M (diag=NULL: compute M'*M) -void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) { +// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional +void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, + int nr, int nc, int flg_upper) { mjtNum tmp; // half of MatMat routine: only lower triangle @@ -870,15 +870,23 @@ void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, in } } - // make symmetric - for (int i=0; i < nc; i++) { - for (int j=i+1; j < nc; j++) { - res[i*nc+j] = res[j*nc+i]; + // flg_upper is set: make symmetric + if (flg_upper) { + for (int i=0; i < nc; i++) { + for (int j=i+1; j < nc; j++) { + res[i*nc+j] = res[j*nc+i]; + } } } } +// compute M'*diag*M (diag=NULL: compute M'*M) +void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) { + mju_sqrMatTD_impl(res, mat, diag, nr, nc, /*flg_upper=*/ 1); +} + + // multiply matrices, first argument transposed void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, diff --git a/src/engine/engine_util_blas.h b/src/engine/engine_util_blas.h index 619a3225..4b2519c9 100644 --- a/src/engine/engine_util_blas.h +++ b/src/engine/engine_util_blas.h @@ -223,6 +223,10 @@ MJAPI void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, MJAPI void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2, int r1, int c1, int c2); +// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional +void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc, + int flg_upper); + // compute M'*diag*M (diag=NULL: compute M'*M) MJAPI void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc); diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 250020e4..8d364c5e 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -186,6 +186,41 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int } +// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, + const mjtNum* M, const int* M_rownnz, const int* M_rowadr, + const int* M_colind, + mjtNum* buf_val, int* buf_ind) { + for (int i=0; i < nr; i++) { + rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1, + rownnz[i], M_rownnz[i], colind + rowadr[i], + M_colind + M_rowadr[i], buf_val, buf_ind); + } +} + + +// add symmetric matrix (lower triangle) to dense matrix, upper triangle optional +void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind, int flg_upper) { + for (int i=0; i < n; i++) { + int start = rowadr[i]; + int end = start + rownnz[i]; + for (int adr=start; adr < end; adr++) { + mjtNum val = mat[adr]; + int j = colind[adr]; + + // lower + diagonal + res[i*n + j] += val; + + // strict upper + if (flg_upper && j < i) { + res[j*n + i] += val; + } + } + } +} + + // multiply symmetric matrix (only lower triangle represented) by vector: // res = (mat + strict_upper(mat')) * vec @@ -214,7 +249,7 @@ void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat, // off-diagonals const int* ind = colind + adr; - for (int k=0; k < diag; k++) { + for (int k=diag-1; k >= 0; k--) { int j = ind[k]; mjtNum val = row[k]; res[i] += val * vec[j]; // strict lower diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index ea26014b..616b4d21 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -50,6 +50,17 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); +// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers +MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr, + const mjtNum* M, const int* M_rownnz, const int* M_rowadr, + const int* M_colind, + mjtNum* buf_val, int* buf_ind); + +// add symmetric matrix (only lower triangle represented) to dense matrix +MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind, + int flg_upper); + // multiply symmetric matrix (only lower triangle represented) by vector: // res = (mat + strict_upper(mat')) * vec MJAPI void mju_mulSymVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n, diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 2f1e087f..10875a67 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -1036,6 +1036,40 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) { EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24)); } +TEST_F(EngineUtilSparseTest, MjuAddToSymSparse) { + // 1 2 4 + // M = 2 3 0 + // 4 0 5 + + // only lower triangle represented + mjtNum mat[] = {1, 2, 3, 4, 5}; + int colind[] = {0, 0, 1, 0, 2}; + int rownnz[] = {1, 2, 2}; + int rowadr[] = {0, 1, 3}; + + // 0 0 0 + // A = 5 4 2 + // 4 3 2 + mjtNum A[] = {0, 0, 0, + 5, 4, 2, + 4, 3, 2}; + + mju_addToSymSparse(A, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/1); + EXPECT_THAT(AsVector(A, 9), ElementsAre(1, 2, 4, + 7, 7, 2, + 8, 3, 7)); + + // same as A + mjtNum B[] = {0, 0, 0, + 5, 4, 2, + 4, 3, 2}; + + mju_addToSymSparse(B, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/0); + EXPECT_THAT(AsVector(B, 9), ElementsAre(1, 0, 0, + 7, 7, 2, + 8, 3, 7)); +} + TEST_F(EngineUtilSparseTest, MjuMulSymVecSparse) { constexpr int n = 4; constexpr int nnz = 9;