From a36f2cccb6218c5d069e21236f24d3bb21110778 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 28 Oct 2024 03:46:00 -0700 Subject: [PATCH] Add nnz argument for size of allocated memory to `mju_dense2sparse` and check if this number is too small for number of non-zeros. PiperOrigin-RevId: 690545825 Change-Id: I0e31cb907a1151cb2e4766d5f0cdfecb5629a2a6 --- src/engine/engine_util_sparse.c | 15 ++++++-- src/engine/engine_util_sparse.h | 5 +-- test/engine/engine_island_test.cc | 10 +++--- test/engine/engine_util_sparse_test.cc | 48 ++++++++++++++++++++++---- 4 files changed, 63 insertions(+), 15 deletions(-) diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 89403bca..f6d05cfd 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -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; } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index ae14cf70..fa2647c3 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -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, diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index c0f5afb8..3729e3f0 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -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]; diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 92b8fc8d..ade96c0b 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -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