Implement a more performant mju_transposeSparse that doesn't require dense memory allocation.
PiperOrigin-RevId: 505044194 Change-Id: Ibe0e39e3ad711b2b5bdad45180ca4f83b5c8fc82
This commit is contained in:
Kyle Bayes
committed by
Copybara-Service
Copybara-Service
parent
4b4ee27f64
commit
c741dfce7d
@@ -8,6 +8,10 @@ Upcoming version (not yet released)
|
||||
General
|
||||
^^^^^^^
|
||||
|
||||
- A more performant mju_transposeSparse has been implemented that doesn't require dense memory allocation.
|
||||
For a constraint Jacobian matrix from the
|
||||
`humanoid100.xml <https://github.com/deepmind/mujoco/blob/main/model/humanoid100/humanoid100.xml>`_ model,
|
||||
this function is 35% faster.
|
||||
- The function :ref:`mj_name2id` is now implemented using a hash function instead of a linear search for better
|
||||
performance.
|
||||
|
||||
|
||||
@@ -528,37 +528,48 @@ void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, i
|
||||
void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind) {
|
||||
// clear counters for transposed
|
||||
// clear number of non-zeros for each row of transposed
|
||||
memset(res_rownnz, 0, nc*sizeof(int));
|
||||
|
||||
// set uncompressed layout
|
||||
for (int rt=0; rt<nc; rt++) {
|
||||
res_rowadr[rt] = rt*nr;
|
||||
// total number of non-zeros of mat
|
||||
int nnz = rowadr[nr-1] + rownnz[nr-1];
|
||||
|
||||
// count the number of non-zeros for each row of the transposed matrix
|
||||
for (int i = 0; i<nnz; i++) {
|
||||
res_rownnz[colind[i]]++;
|
||||
}
|
||||
|
||||
// scan original, compute uncompressed
|
||||
for (int r=0; r<nr; r++) {
|
||||
for (int ci=0; ci<rownnz[r]; ci++) {
|
||||
// get rt=c
|
||||
int rt = colind[rowadr[r]+ci];
|
||||
|
||||
// record index ct=r, assuming uncompressed res_rowadr[rt]=rt*nr
|
||||
res_colind[rt*nr + res_rownnz[rt]] = r;
|
||||
|
||||
// copy data
|
||||
res[rt*nr + res_rownnz[rt]] = mat[rowadr[r]+ci];
|
||||
|
||||
// increase counter for rt
|
||||
res_rownnz[rt]++;
|
||||
}
|
||||
// compute the row addresses for the transposed matrix
|
||||
res_rowadr[0] = 0;
|
||||
for (int i = 1; i<nc; i++) {
|
||||
res_rowadr[i] = res_rowadr[i-1] + res_rownnz[i-1];
|
||||
}
|
||||
|
||||
// compress
|
||||
mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind);
|
||||
// r holds the current row in mat
|
||||
int r = 0;
|
||||
|
||||
// iterate through each non-zero entry of mat
|
||||
for (int i = 0; i<nnz; i++) {
|
||||
|
||||
// iterate to get to the current row (skipping rows with all zeros)
|
||||
while ((i-rowadr[r]) >= rownnz[r]) r++;
|
||||
|
||||
// row index becomes the column index
|
||||
res[res_rowadr[colind[i]]] = mat[i];
|
||||
|
||||
// column index becomes the row index (increment the rowadr)
|
||||
res_colind[res_rowadr[colind[i]]++] = r;
|
||||
}
|
||||
|
||||
// shift back row addresses
|
||||
for (int i = nc-1; i>0; i--) {
|
||||
res_rowadr[i] = res_rowadr[i-1];
|
||||
}
|
||||
|
||||
res_rowadr[0] = 0;
|
||||
}
|
||||
|
||||
|
||||
|
||||
// construct row supernodes
|
||||
void mju_superSparse(int nr, int* rowsuper,
|
||||
const int* rownnz, const int* rowadr, const int* colind) {
|
||||
|
||||
@@ -48,8 +48,8 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
|
||||
const int* colind, const int* rowsuper);
|
||||
|
||||
// compress layout of sparse matrix
|
||||
void mju_compressSparse(mjtNum* mat, int nr, int nc,
|
||||
int* rownnz, int* rowadr, int* colind);
|
||||
MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc,
|
||||
int* rownnz, int* rowadr, int* colind);
|
||||
|
||||
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
||||
int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
|
||||
@@ -61,9 +61,9 @@ void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNu
|
||||
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind);
|
||||
|
||||
// transpose sparse matrix
|
||||
void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind);
|
||||
MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
||||
int* res_rownnz, int* res_rowadr, int* res_colind,
|
||||
const int* rownnz, const int* rowadr, const int* colind);
|
||||
|
||||
// construct row supernodes
|
||||
void mju_superSparse(int nr, int* rowsuper,
|
||||
|
||||
@@ -26,9 +26,10 @@
|
||||
namespace mujoco {
|
||||
namespace {
|
||||
|
||||
using FuncPtr = decltype(&mju_combineSparse);
|
||||
using CombineFuncPtr = decltype(&mju_combineSparse);
|
||||
using TransposeFuncPtr = decltype(&mju_transposeSparse);
|
||||
|
||||
// number of steps to roll out before benhmarking
|
||||
// number of steps to roll out before benchmarking
|
||||
static const int kNumWarmupSteps = 500;
|
||||
|
||||
// copy array into vector
|
||||
@@ -38,6 +39,27 @@ std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
|
||||
|
||||
// ----------------------------- old functions --------------------------------
|
||||
|
||||
// transpose sparse matrix (uncompressed)
|
||||
void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline(
|
||||
mjtNum* res, const mjtNum* mat, int nr, int nc, int* res_rownnz,
|
||||
int* res_rowadr, int* res_colind, const int* rownnz, const int* rowadr,
|
||||
const int* colind) {
|
||||
memset(res_rownnz, 0, nc * sizeof(int));
|
||||
for (int rt = 0; rt < nc; rt++) {
|
||||
res_rowadr[rt] = rt * nr;
|
||||
}
|
||||
for (int r = 0; r < nr; r++) {
|
||||
for (int ci = 0; ci < rownnz[r]; ci++) {
|
||||
int rt = colind[rowadr[r] + ci];
|
||||
res_colind[rt * nr + res_rownnz[rt]] = r;
|
||||
res[rt * nr + res_rownnz[rt]] = mat[rowadr[r] + ci];
|
||||
res_rownnz[rt]++;
|
||||
}
|
||||
}
|
||||
|
||||
mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind);
|
||||
}
|
||||
|
||||
int compare_baseline(const int* vec1,
|
||||
const int* vec2,
|
||||
int n) {
|
||||
@@ -286,8 +308,9 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_MatVecSparse_1(
|
||||
}
|
||||
BENCHMARK(BM_MatVecSparse_1);
|
||||
|
||||
static void BM_combineSparse(benchmark::State& state, FuncPtr func) {
|
||||
static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
|
||||
static mjModel* m = LoadModelFromPath("humanoid/humanoid.xml");
|
||||
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
// warm-up rollout to get a typical state
|
||||
@@ -359,5 +382,52 @@ void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_combineSparse_old(
|
||||
}
|
||||
BENCHMARK(BM_combineSparse_old);
|
||||
|
||||
static void BM_transposeSparse(benchmark::State& state, TransposeFuncPtr func) {
|
||||
static mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml");
|
||||
|
||||
// force use of sparse matrices
|
||||
m->opt.jacobian = mjJAC_SPARSE;
|
||||
|
||||
mjData* d = mj_makeData(m);
|
||||
|
||||
// warm-up rollout to get a typical state
|
||||
while (d-> time < 2) {
|
||||
mj_step(m, d);
|
||||
}
|
||||
|
||||
mjMARKSTACK;
|
||||
|
||||
// need uncompressed layout
|
||||
mjtNum* res = mj_stackAlloc(d, m->nv * d->nefc);
|
||||
int* res_rownnz = (int*)mj_stackAlloc(d, m->nv);
|
||||
int* res_rowadr = (int*)mj_stackAlloc(d, m->nv);
|
||||
int* res_colind = (int*)mj_stackAlloc(d, m->nv * d->nefc);
|
||||
|
||||
// time benchmark
|
||||
for (auto s : state) {
|
||||
func(res, d->efc_J, d->nefc, m->nv, res_rownnz, res_rowadr, res_colind,
|
||||
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
|
||||
}
|
||||
|
||||
mjFREESTACK;
|
||||
mj_deleteData(d);
|
||||
state.SetItemsProcessed(state.iterations());
|
||||
}
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_new(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &mju_transposeSparse);
|
||||
}
|
||||
BENCHMARK(BM_transposeSparse_new);
|
||||
|
||||
void ABSL_ATTRIBUTE_NO_TAIL_CALL
|
||||
BM_transposeSparse_old(benchmark::State& state) {
|
||||
MujocoErrorTestGuard guard;
|
||||
BM_transposeSparse(state, &transposeSparse_baseline);
|
||||
}
|
||||
|
||||
BENCHMARK(BM_transposeSparse_old);
|
||||
|
||||
} // namespace
|
||||
} // namespace mujoco
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user