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

Speed improvement with no AVX. 40% on Linux Intel Xeon and 60% on ARM Mac.

Browse Source 
PiperOrigin-RevId: 488360660
Change-Id: I423269cd362fdcc2434393b691638ac6ed778ef2
This commit is contained in:
Alessio Quaglino
2022-11-14 07:30:51 -08:00
committed by Copybara-Service
parent ae29a22d55
commit 7b0fbc63f8
7 changed files with 310 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_sparse.c
+16
View File
@@ -73,6 +73,22 @@ mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
high64 = _mm_unpackhi_pd(vlow, vlow);
res = _mm_cvtsd_f64(_mm_add_sd(vlow, high64));
}
#else
int n_4 = nnz1 - 4;
mjtNum res0 = 0;
mjtNum res1 = 0;
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]];
}
res = (res0 + res2) + (res1 + res3);
#endif
// scalar part
src/engine/engine_util_sparse.h
+6 -5
View File
@@ -16,6 +16,7 @@
#define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPARSE_H_
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjtnum.h>
#ifdef __cplusplus
@@ -25,8 +26,8 @@ extern "C" {
//------------------------------ sparse operations -------------------------------------------------
// dot-product, first vector is sparse
mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1);
MJAPI mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
const int nnz1, const int* ind1);
// dot-product, both vectors are sparse
mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
@@ -42,9 +43,9 @@ void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);
// multiply sparse matrix and dense vector: res = mat * vec
void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int nr, const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper);
MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int nr, const int* rownnz, const int* rowadr,
const int* colind, const int* rowsuper);
// compress layout of sparse matrix
void mju_compressSparse(mjtNum* mat, int nr, int nc,