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

Allow dot products with sparse vectors to specify that one vector uses uncompressed memory.

Browse Source 
PiperOrigin-RevId: 561951337
Change-Id: I1310d7c09d9a85ad9b43877155844a9d0ce6edab
This commit is contained in:
Yuval Tassa
2023-09-01 07:39:27 -07:00
committed by Copybara-Service
parent 70c5fa50f2
commit 29aa5e4a41
7 changed files with 143 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+4 -4
View File
@@ -420,7 +420,7 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d,
int Jrowadr = d->efc_J_rowadr[row];
int* Jind = d->efc_J_colind + Jrowadr;
mjtNum* J = d->efc_J + Jrowadr;
res[i] = mju_dotSparse2(vec, J, vecnnz, vecind, Jnnz, Jind);
res[i] = mju_dotSparse2(vec, J, vecnnz, vecind, Jnnz, Jind, /*flg_unc2=*/0);
}
}
@@ -428,7 +428,7 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d,
else {
int nv = m->nv;
for (int i=0; i < resnnz; i++) {
res[i] = mju_dotSparse(vec, d->efc_J + nv*resind[i], vecnnz, vecind);
res[i] = mju_dotSparse(vec, d->efc_J + nv*resind[i], vecnnz, vecind, /*flg_unc1=*/0);
}
}
}
@@ -481,7 +481,7 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
int JTrowadr = d->efc_JT_rowadr[row];
int* JTind = d->efc_JT_colind + JTrowadr;
mjtNum* JT = d->efc_JT + JTrowadr;
res[i] = mju_dotSparse2(vec, JT, vecnnz, vecind, JTnnz, JTind);
res[i] = mju_dotSparse2(vec, JT, vecnnz, vecind, JTnnz, JTind, /*flg_unc2=*/0);
}
}
@@ -489,7 +489,7 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
else {
int nefc = d->nefc;
for (int i=0; i < resnnz; i++) {
res[i] = mju_dotSparse(vec, d->efc_JT + nefc*resind[i], vecnnz, vecind);
res[i] = mju_dotSparse(vec, d->efc_JT + nefc*resind[i], vecnnz, vecind, /*flg_unc1=*/0);
}
}
}
src/engine/engine_solver.c
+2 -1
View File
@@ -168,7 +168,8 @@ static void residual(const mjModel* m, mjData* d, mjtNum* res, int i, int dim, i
for (int j=0; j < dim; j++) {
res[j] = d->efc_b[i+j] + mju_dotSparse(d->efc_AR + d->efc_AR_rowadr[i+j],
d->efc_force, d->efc_AR_rownnz[i+j],
d->efc_AR_colind + d->efc_AR_rowadr[i+j]);
d->efc_AR_colind + d->efc_AR_rowadr[i+j],
/*flg_unc1=*/0);
}
}
src/engine/engine_util_solve.c
+1 -1
View File
@@ -235,7 +235,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// x(i) -= sum_j L(i,j)*x(j), j=0:i-1
if (nnz > 1) {
res[i] -= mju_dotSparse(mat+adr, res, nnz-1, colind+adr);
res[i] -= mju_dotSparse(mat+adr, res, nnz-1, colind+adr, /*flg_unc1=*/0);
// modulo AVX, the above line does
// for (int j=0; j<nnz-1; j++)
// res[i] -= mat[adr+j]*res[colind[adr+j]];
src/engine/engine_util_sparse.c
+35 -12
View File
@@ -27,10 +27,11 @@
//------------------------------ sparse operations -------------------------------------------------
// dot-product, first vector is sparse
// flg_unc1: is vec1 memory layout uncompressed
mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1) {
const int nnz1, const int* ind1, int flg_unc1) {
#ifdef mjUSEAVX
return mju_dotSparse_avx(vec1, vec2, nnz1, ind1);
return mju_dotSparse_avx(vec1, vec2, nnz1, ind1, flg_unc1);
#else
int i = 0;
mjtNum res = 0;
@@ -40,17 +41,33 @@ mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
mjtNum res2 = 0;
mjtNum res3 = 0;
for (; i <= n_4; i+=4) {
res0 += vec1[i+0] * vec2[ind1[i+0]];
res1 += vec1[i+1] * vec2[ind1[i+1]];
res2 += vec1[i+2] * vec2[ind1[i+2]];
res3 += vec1[i+3] * vec2[ind1[i+3]];
if (flg_unc1) {
for (; i <= n_4; i+=4) {
res0 += vec1[ind1[i+0]] * vec2[ind1[i+0]];
res1 += vec1[ind1[i+1]] * vec2[ind1[i+1]];
res2 += vec1[ind1[i+2]] * vec2[ind1[i+2]];
res3 += vec1[ind1[i+3]] * vec2[ind1[i+3]];
}
} else {
for (; i <= n_4; i+=4) {
res0 += vec1[i+0] * vec2[ind1[i+0]];
res1 += vec1[i+1] * vec2[ind1[i+1]];
res2 += vec1[i+2] * vec2[ind1[i+2]];
res3 += vec1[i+3] * vec2[ind1[i+3]];
}
}
res = (res0 + res2) + (res1 + res3);
// scalar part
for (; i < nnz1; i++) {
res += vec1[i] * vec2[ind1[i]];
if (flg_unc1) {
for (; i < nnz1; i++) {
res += vec1[ind1[i]] * vec2[ind1[i]];
}
} else {
for (; i < nnz1; i++) {
res += vec1[i] * vec2[ind1[i]];
}
}
return res;
@@ -91,9 +108,10 @@ void mju_dotSparseX3(mjtNum* res0, mjtNum* res1, mjtNum* res2,
// dot-product, both vectors are sparse
// flg_unc2: is vec2 memory layout uncompressed
mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1,
const int nnz2, const int* ind2) {
const int nnz2, const int* ind2, int flg_unc2) {
int i1 = 0, i2 = 0;
mjtNum res = 0;
@@ -108,7 +126,12 @@ mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
// match: accumulate result, advance both
if (adr1 == adr2) {
res += vec1[i1++] * vec2[i2++];
if (flg_unc2) {
res += vec1[i1++] * vec2[adr2];
i2++;
} else {
res += vec1[i1++] * vec2[i2++];
}
}
// otherwise advance smaller
@@ -176,7 +199,7 @@ void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
#else
// regular sparse dot-product
for (int r=0; r < nr; r++) {
res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
res[r] = mju_dotSparse(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc1=*/0);
}
#endif // mjUSEAVX
}
src/engine/engine_util_sparse.h
+6 -6
View File
@@ -25,14 +25,14 @@ extern "C" {
//------------------------------ sparse operations -------------------------------------------------
// dot-product, first vector is sparse
// dot-product, vec1 is sparse, can be uncompressed
MJAPI mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1);
const int nnz1, const int* ind1, int flg_unc1);
// dot-product, both vectors are sparse
mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1,
const int nnz2, const int* ind2);
// dot-product, both vectors are sparse, vec2 can be uncompressed
MJAPI mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1,
const 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,
src/engine/engine_util_sparse_avx.h
+46 -16
View File
@@ -30,9 +30,10 @@
//------------------------------ sparse operations using avx ---------------------------------------
// dot-product, first vector is sparse
// flg_unc1: is vec1 memory layout uncompressed
static inline
mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1) {
const int nnz1, const int* ind1, int flg_unc1) {
int i = 0;
mjtNum res = 0;
int nnz1_4 = nnz1 - 4;
@@ -47,20 +48,43 @@ mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2,
vec2[ind1[2]],
vec2[ind1[1]],
vec2[ind1[0]]);
val1 = _mm256_loadu_pd(vec1);
if (flg_unc1) {
val1 = _mm256_set_pd(vec1[ind1[3]],
vec1[ind1[2]],
vec1[ind1[1]],
vec1[ind1[0]]);
} else {
val1 = _mm256_loadu_pd(vec1);
}
sum = _mm256_mul_pd(val1, val2);
i = 4;
// parallel computation
while (i<=nnz1_4) {
val1 = _mm256_loadu_pd(vec1+i);
val2 = _mm256_set_pd(vec2[ind1[i+3]],
vec2[ind1[i+2]],
vec2[ind1[i+1]],
vec2[ind1[i+0]]);
prod = _mm256_mul_pd(val1, val2);
sum = _mm256_add_pd(sum, prod);
i += 4;
if (flg_unc1) {
while (i<=nnz1_4) {
val1 = _mm256_set_pd(vec1[ind1[i+3]],
vec1[ind1[i+2]],
vec1[ind1[i+1]],
vec1[ind1[i+0]]);
val2 = _mm256_set_pd(vec2[ind1[i+3]],
vec2[ind1[i+2]],
vec2[ind1[i+1]],
vec2[ind1[i+0]]);
prod = _mm256_mul_pd(val1, val2);
sum = _mm256_add_pd(sum, prod);
i += 4;
}
} else {
while (i<=nnz1_4) {
val1 = _mm256_loadu_pd(vec1+i);
val2 = _mm256_set_pd(vec2[ind1[i+3]],
vec2[ind1[i+2]],
vec2[ind1[i+1]],
vec2[ind1[i+0]]);
prod = _mm256_mul_pd(val1, val2);
sum = _mm256_add_pd(sum, prod);
i += 4;
}
}
// reduce
@@ -72,8 +96,14 @@ mjtNum mju_dotSparse_avx(const mjtNum* vec1, const mjtNum* vec2,
}
// scalar part
for (; i<nnz1; i++) {
res += vec1[i] * vec2[ind1[i]];
if (flg_unc1) {
for (; i < nnz1; i++) {
res += vec1[ind1[i]] * vec2[ind1[i]];
}
} else {
for (; i < nnz1; i++) {
res += vec1[i] * vec2[ind1[i]];
}
}
return res;
@@ -179,7 +209,7 @@ void mju_mulMatVecSparse_avx(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
if (!rowsuper) {
// regular sparse dot-product
for (int r=0; r<nr; r++) {
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc2=*/0);
}
return;
@@ -202,7 +232,7 @@ void mju_mulMatVecSparse_avx(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
// handle remaining rows
while (rs>0) {
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc2=*/0);
r++;
rs--;
@@ -213,7 +243,7 @@ void mju_mulMatVecSparse_avx(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
}
else {
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r]);
res[r] = mju_dotSparse_avx(mat+rowadr[r], vec, rownnz[r], colind+rowadr[r], /*flg_unc2=*/0);
}
}
}
test/engine/engine_util_sparse_test.cc
+49 -10
View File
@@ -30,19 +30,58 @@ using ::testing::ElementsAre;
using EngineUtilSparseTest = MujocoTest;
TEST_F(EngineUtilSparseTest, MjuDot) {
mjtNum a[] = {1, 2, 3, 4, 5, 6, 7};
mjtNum b[] = {7, 0, 6, 0, 0, 5, 0, 0, 0, 4, 0, 0, 0, 3, 0, 0, 2, 0, 1};
mjtNum a[] = {2, 3, 4, 5, 6, 7, 8};
mjtNum u[] = {2, 1, 3, 1, 1, 4, 1, 1, 1, 5, 1, 1, 1, 6, 1, 1, 7, 1, 8};
mjtNum b[] = {8, 1, 7, 1, 1, 6, 1, 1, 1, 5, 1, 1, 1, 4, 1, 1, 3, 1, 2};
int i[] = {0, 2, 5, 9, 13, 16, 18};
// test various vector lengths as mju_dotSparse adds numbers in groups of four
EXPECT_EQ(mju_dotSparse(a, b, 0, i), 0);
EXPECT_EQ(mju_dotSparse(a, b, 1, i), 7);
EXPECT_EQ(mju_dotSparse(a, b, 2, i), 7 + 2*6);
EXPECT_EQ(mju_dotSparse(a, b, 3, i), 7 + 2*6 + 3*5);
EXPECT_EQ(mju_dotSparse(a, b, 4, i), 7 + 2*6 + 3*5 + 4*4);
EXPECT_EQ(mju_dotSparse(a, b, 5, i), 7 + 2*6 + 3*5 + 4*4 + 5*3);
EXPECT_EQ(mju_dotSparse(a, b, 6, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2);
EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
// a is compressed
int flg_unc1 = 0;
EXPECT_EQ(mju_dotSparse(a, b, 0, i, flg_unc1), 0);
EXPECT_EQ(mju_dotSparse(a, b, 1, i, flg_unc1), 2*8);
EXPECT_EQ(mju_dotSparse(a, b, 2, i, flg_unc1), 2*8 + 3*7);
EXPECT_EQ(mju_dotSparse(a, b, 3, i, flg_unc1), 2*8 + 3*7 + 4*6);
EXPECT_EQ(mju_dotSparse(a, b, 4, i, flg_unc1), 2*8 + 3*7 + 4*6 + 5*5);
EXPECT_EQ(mju_dotSparse(a, b, 5, i, flg_unc1), 2*8 + 3*7 + 4*6 + 5*5 + 6*4);
EXPECT_EQ(mju_dotSparse(a, b, 6, i, flg_unc1),
2*8 + 3*7 + 4*6 + 5*5 + 6*4 + 7*3);
EXPECT_EQ(mju_dotSparse(a, b, 7, i, flg_unc1),
2*8 + 3*7 + 4*6 + 5*5 + 6*4 + 7*3 + 8*2);
// u is compressed
flg_unc1 = 1;
EXPECT_EQ(mju_dotSparse(u, b, 0, i, flg_unc1), 0);
EXPECT_EQ(mju_dotSparse(u, b, 1, i, flg_unc1), 2*8);
EXPECT_EQ(mju_dotSparse(u, b, 2, i, flg_unc1), 2*8 + 3*7);
EXPECT_EQ(mju_dotSparse(u, b, 3, i, flg_unc1), 2*8 + 3*7 + 4*6);
EXPECT_EQ(mju_dotSparse(u, b, 4, i, flg_unc1), 2*8 + 3*7 + 4*6 + 5*5);
EXPECT_EQ(mju_dotSparse(u, b, 5, i, flg_unc1), 2*8 + 3*7 + 4*6 + 5*5 + 6*4);
EXPECT_EQ(mju_dotSparse(u, b, 6, i, flg_unc1),
2*8 + 3*7 + 4*6 + 5*5 + 6*4 + 7*3);
EXPECT_EQ(mju_dotSparse(u, b, 7, i, flg_unc1),
2*8 + 3*7 + 4*6 + 5*5 + 6*4 + 7*3 + 8*2);
}
TEST_F(EngineUtilSparseTest, MjuDot2) {
constexpr int annz = 6;
constexpr int bnnz = 5;
int ia[annz] = {0, 2, 5, 6, 7};
mjtNum a[annz] = {2, 3, 4, 5, 6};
int ib[bnnz] = { 1, 2, 3, 5, 7};
mjtNum b[bnnz] = { 8, 7, 6, 5, 4};
mjtNum u[] = {1, 8, 7, 6, 1, 5, 1, 4};
// test various vector lengths as mju_dotSparse adds numbers in groups of four
// a is compressed
int flg_unc2 = 0;
EXPECT_EQ(mju_dotSparse2(a, b, annz, ia, bnnz, ib, flg_unc2), 3*7+4*5+6*4);
// u is uncompressed
flg_unc2 = 1;
EXPECT_EQ(mju_dotSparse2(a, u, annz, ia, bnnz, ib, flg_unc2), 3*7+4*5+6*4);
}
TEST_F(EngineUtilSparseTest, CombineSparseCount) {