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

Implement a more performant mju_transposeSparse that doesn't require dense memory allocation.

Browse Source 
PiperOrigin-RevId: 505044194
Change-Id: Ibe0e39e3ad711b2b5bdad45180ca4f83b5c8fc82
This commit is contained in:
Kyle Bayes
2023-01-27 00:19:14 -08:00
committed by Copybara-Service
parent 4b4ee27f64
commit c741dfce7d
5 changed files with 255 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/engine/engine_util_sparse_test.cc
+140
View File
@@ -16,12 +16,14 @@
#include "src/engine/engine_util_sparse.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "test/fixture.h"
namespace mujoco {
namespace {
using ::testing::ElementsAre;
using EngineUtilSparseTest = MujocoTest;
TEST_F(EngineUtilSparseTest, MjuDot) {
@@ -40,5 +42,143 @@ TEST_F(EngineUtilSparseTest, MjuDot) {
EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
}
TEST_F(EngineUtilSparseTest, MjuTranspose3by3) {
// 1 2 0 1 0 0
// 0 1 0 --> 2 1 3
// 0 3 0 0 0 0
mjtNum mat[] = {1, 2, 1, 3};
int colind[] = {0, 1, 1, 1};
int rownnz[] = {2, 1, 1};
int rowadr[] = {0, 2, 3};
mjtNum matT[] = {0, 0, 0, 0};
int colindT[] = {0, 0, 0, 0};
int rownnzT[] = {0, 0, 0};
int rowadrT[] = {0, 0, 0};
mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(matT, ElementsAre(1, 2, 1, 3));
EXPECT_THAT(colindT, ElementsAre(0, 0, 1, 2));
EXPECT_THAT(rownnzT, ElementsAre(1, 3, 0));
EXPECT_THAT(rowadrT, ElementsAre(0, 1, 4));
}
TEST_F(EngineUtilSparseTest, MjuTranspose1by3) {
// 1 0 3 1
// --> 0
// 3
mjtNum mat[] = {1, 3};
int colind[] = {0, 2};
int rownnz[] = {2};
int rowadr[] = {0};
mjtNum matT[] = {0, 0};
int colindT[] = {0, 0};
int rownnzT[] = {0, 0, 0};
int rowadrT[] = {0, 0, 0};
mju_transposeSparse(matT, mat, 1, 3, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(matT, ElementsAre(1, 3));
EXPECT_THAT(colindT, ElementsAre(0, 0));
EXPECT_THAT(rownnzT, ElementsAre(1, 0, 1));
EXPECT_THAT(rowadrT, ElementsAre(0, 1, 1));
}
TEST_F(EngineUtilSparseTest, MjuTranspose3by1) {
// 1 1 0 3
// 0 -->
// 3
mjtNum mat[] = {1, 3};
int colind[] = {0, 0};
int rownnz[] = {1, 0, 1};
int rowadr[] = {0, 1, 1};
mjtNum matT[] = {0, 0};
int colindT[] = {0, 0};
int rownnzT[] = {0};
int rowadrT[] = {0};
mju_transposeSparse(matT, mat, 3, 1, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(matT, ElementsAre(1, 3));
EXPECT_THAT(colindT, ElementsAre(0, 2));
EXPECT_THAT(rownnzT, ElementsAre(2));
EXPECT_THAT(rowadrT, ElementsAre(0));
}
TEST_F(EngineUtilSparseTest, MjuTransposeDense) {
// 1 2 3 1 4 7
// 4 5 6 --> 2 5 8
// 7 8 9 3 6 9
mjtNum mat[] = {1, 2, 3, 4, 5, 6, 7, 8, 9};
int colind[] = {0, 1, 2, 0, 1, 2, 0, 1, 2};
int rownnz[] = {3, 3, 3};
int rowadr[] = {0, 3, 6};
mjtNum matT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int colindT[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzT[] = {0, 0, 0};
int rowadrT[] = {0, 0, 0};
mju_transposeSparse(matT, mat, 3, 3, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(matT, ElementsAre(1, 4, 7, 2, 5, 8, 3, 6, 9));
EXPECT_THAT(colindT, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
EXPECT_THAT(rownnzT, ElementsAre(3, 3, 3));
EXPECT_THAT(rowadrT, ElementsAre(0, 3, 6));
}
TEST_F(EngineUtilSparseTest, MjuTranspose1by1) {
// 1 -> 1
mjtNum mat[] = {1};
int colind[] = {0};
int rownnz[] = {1};
int rowadr[] = {0};
mjtNum matT[] = {0};
int colindT[] = {0};
int rownnzT[] = {0};
int rowadrT[] = {0};
mju_transposeSparse(matT, mat, 1, 1, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(matT, ElementsAre(1));
EXPECT_THAT(colindT, ElementsAre(0));
EXPECT_THAT(rownnzT, ElementsAre(1));
EXPECT_THAT(rowadrT, ElementsAre(0));
}
TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) {
// 0 -> 0
mjtNum mat[] = {};
int colind[] = {};
int rownnz[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rowadr[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
mjtNum matT[] = {};
int colindT[] = {};
int rownnzT[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
int rowadrT[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
mju_transposeSparse(matT, mat, 10, 10, rownnzT, rowadrT, colindT, rownnz,
rowadr, colind);
EXPECT_THAT(rownnzT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0));
EXPECT_THAT(rowadrT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0));
}
} // namespace
} // namespace mujoco