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

Add nnz argument for size of allocated memory to mju_dense2sparse and check if this number is too small for number of non-zeros.

Browse Source 
PiperOrigin-RevId: 690545825
Change-Id: I0e31cb907a1151cb2e4766d5f0cdfecb5629a2a6
This commit is contained in:
Taylor Howell
2024-10-28 03:46:00 -07:00
committed by Copybara-Service
parent 1c424644dd
commit a36f2cccb6
4 changed files with 63 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+13 -2
View File
@@ -148,8 +148,13 @@ mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const in
// convert matrix from dense to sparse
void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind) {
// nnz is size of res and colind, return 1 if too small, 0 otherwise
int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind, int nnz) {
if (nnz <= 0) {
return 1;
}
int adr = 0;
// find non-zeros and construct sparse
@@ -161,6 +166,11 @@ void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
// find non-zeros
for (int c=0; c < nc; c++) {
if (mat[r*nc+c]) {
// check for out of bounds
if (adr >= nnz) {
return 1;
}
// record index and count
colind[adr] = c;
rownnz[r]++;
@@ -170,6 +180,7 @@ void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
}
}
}
return 0;
}
src/engine/engine_util_sparse.h
+3 -2
View File
@@ -34,8 +34,9 @@ MJAPI mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, co
int nnz2, const int* ind2, int flg_unc2);
// convert matrix from dense to sparse
MJAPI void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind);
// nnz is size of res and colind, return 1 if too small, 0 otherwise
MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind, int nnz);
// convert matrix from sparse to dense
MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
test/engine/engine_island_test.cc
+5 -5
View File
@@ -55,7 +55,7 @@ TEST_F(IslandTest, FloodFillSingleton) {
int rowadr[nr];
int colind[nnz];
mjtNum res[nnz]; // unused
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind, nnz);
// outputs / scratch
int island[nr];
@@ -82,7 +82,7 @@ TEST_F(IslandTest, FloodFill1) {
int rowadr[nr];
int colind[nnz];
mjtNum res[nnz]; // unused
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind, nnz);
// outputs / stack
int island[nr];
@@ -112,7 +112,7 @@ TEST_F(IslandTest, FloodFill2) {
int rowadr[nr];
int colind[nnz];
mjtNum res[nnz]; // unused
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind, nnz);
// outputs / stack
int island[nr];
@@ -140,7 +140,7 @@ TEST_F(IslandTest, FloodFill3a) {
int rowadr[nr];
int colind[nnz];
mjtNum res[nnz]; // unused
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind, nnz);
// outputs / stack
int island[nr];
@@ -174,7 +174,7 @@ TEST_F(IslandTest, FloodFill3b) {
int rowadr[nr];
int colind[nnz];
mjtNum res[nnz]; // unused
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind);
mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind, nnz);
// outputs / stack
int island[nr];
test/engine/engine_util_sparse_test.cc
+42 -6
View File
@@ -984,7 +984,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrA[2];
int colindA[4];
int rownnzA_factor[2];
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA);
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA, 4);
int nnzA = mju_cholFactorNNZ(rownnzA_factor,
rownnzA, rowadrA, colindA, nA, d);
@@ -1000,7 +1000,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrB[3];
int colindB[9];
int rownnzB_factor[3];
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB);
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB, 9);
int nnzB = mju_cholFactorNNZ(rownnzB_factor,
rownnzB, rowadrB, colindB, nB, d);
@@ -1016,7 +1016,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrC[3];
int colindC[9];
int rownnzC_factor[3];
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC);
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC, 9);
int nnzC = mju_cholFactorNNZ(rownnzC_factor,
rownnzC, rowadrC, colindC, nC, d);
@@ -1033,7 +1033,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrD[4];
int colindD[16];
int rownnzD_factor[4];
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD);
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD, 16);
int nnzD = mju_cholFactorNNZ(rownnzD_factor,
rownnzD, rowadrD, colindD, nD, d);
@@ -1050,11 +1050,11 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
mjtNum mat[] = {1, 2, 0,
0, 3, 4};
mjtNum mat_sparse[6];
mjtNum mat_sparse[4];
int rownnz[2];
int rowadr[2];
int colind[4];
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind);
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 4);
// multiply: res = mat' * vec
mjtNum vec[] = {5, 6};
@@ -1064,5 +1064,41 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
}
TEST_F(EngineUtilSparseTest, MjuDenseToSparse) {
int nr = 2;
int nc = 2;
mjtNum mat[] = {1, 2,
0, 3};
mjtNum mat_sparse[4];
int rownnz[2];
int rowadr[2];
int colind[4];
// nnz == number of non-zeros
int status3 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 3);
EXPECT_EQ(status3, 0);
// nnz > number of non-zeros
int status4 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 4);
EXPECT_EQ(status4, 0);
// nnz < number of non-zeros
int status2 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 2);
EXPECT_EQ(status2, 1);
// nnz == 0
int status0 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 0);
EXPECT_EQ(status0, 1);
}
} // namespace
} // namespace mujoco