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

Optimize and refactor combineSparse and combineSparseCount.

Browse Source 
PiperOrigin-RevId: 558126391
Change-Id: Iead14c31136e789560f3c8124b84bd4e383158f4
This commit is contained in:
Kyle Bayes
2023-08-18 06:13:20 -07:00
committed by Copybara-Service
parent b1ecc943bb
commit ba12a1354b
7 changed files with 49 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+18 -18
View File
@@ -546,9 +546,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
// compute Jacobian: sparse or dense
if (issparse) {
NV = mju_combineSparse(jac[0], jac[1], nv, 1, -deriv,
NV, NV2, chain, chain2,
sparse_buf, buf_ind);
NV = mju_combineSparse(jac[0], jac[1], 1, -deriv, NV, NV2, chain,
chain2, sparse_buf, buf_ind);
} else {
mju_addToScl(jac[0], jac[1], -deriv, nv);
}
@@ -1310,26 +1309,27 @@ void mj_makeImpedance(const mjModel* m, mjData* d) {
// count the number of non-zeros in the sum of two sparse vectors
int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
int a = 0;
int b = 0;
int nnz = 0;
int a = 0, b = 0, c_nnz = 0;
// while there are elements remaining in both a_ind and b_ind
// count c_nnz: nonzero indices common to both a and b
while (a < a_nnz && b < b_nnz) {
// add the smaller element of either a_ind[a] or b_ind[b] to the combined nnz
++nnz;
// common index, increment everything
if (a_ind[a] == b_ind[b]) {
c_nnz++;
a++;
b++;
}
// if a_ind[a] == b_ind[b], increment both a and b so that we don't double count
// otherwise, increment the index pointing to the smaller element
int aa = a;
int bb = b;
if (a_ind[aa] <= b_ind[bb]) ++a;
if (a_ind[aa] >= b_ind[bb]) ++b;
// update smallest index
else if (a_ind[a] < b_ind[b]) {
a++;
} else {
b++;
}
}
// count remaining elements from the vector with larger nnz
nnz += (a_nnz - a) + (b_nnz - b);
return nnz;
// union minus the intersection
return a_nnz + b_nnz - c_nnz;
}
src/engine/engine_core_smooth.c
+3 -3
View File
@@ -566,9 +566,9 @@ void mj_tendon(const mjModel* m, mjData* d) {
mju_mulMatTVec(tmp, jacdif, dif, 3, NV);
// add to existing
rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, nv, 1, 1/divisor,
rownnz[i], NV, colind+rowadr[i], chain,
sparse_buf, buf_ind);
rownnz[i] = mju_combineSparse(J+rowadr[i], tmp, 1, 1/divisor,
rownnz[i], NV, colind+rowadr[i],
chain, sparse_buf, buf_ind);
}
// dense
src/engine/engine_support.c
+1 -1
View File
@@ -1053,7 +1053,7 @@ void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
// add to destination
for (int i=0; i < nv; i++) {
rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], nv, 1, 1,
rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
rownnz[i], M_rownnz[i], colind + rowadr[i],
M_colind + M_rowadr[i], sparse_buf, buf_ind);
}
src/engine/engine_util_solve.c
+4 -5
View File
@@ -190,7 +190,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
int c = colind[adr+i];
// mat(c,0:c) = mat(c,0:c) - mat(r,c) * mat(r,0:c)
int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], c + 1, 1, -mat[adr+i],
int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], 1, -mat[adr+i],
rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r],
sparse_buf, buf_ind);
@@ -278,7 +278,7 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
mat[adr+nnz-1] = r;
// update row: mat(r,1:r-1) = (mat(r,1:r-1) + s*x(1:r-1)) / c
int new_nnz = mju_combineSparse(mat + adr, x, n, 1 / c, (flg_plus ? s / c : -s / c),
int new_nnz = mju_combineSparse(mat + adr, x, 1 / c, (flg_plus ? s / c : -s / c),
nnz-1, i, colind + adr, x_ind,
sparse_buf, buf_ind);
@@ -288,9 +288,8 @@ int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
}
// update x: x(1:r-1) = c*x(1:r-1) - s*mat(r,1:r-1)
int new_x_nnz = mju_combineSparse(x, mat+adr, n, c, -s,
i, nnz-1, x_ind, colind+adr,
sparse_buf, buf_ind);
int new_x_nnz = mju_combineSparse(x, mat+adr, c, -s, i, nnz-1, x_ind,
colind+adr, sparse_buf, buf_ind);
// update i, correct for changing x
i = i - 1 + (new_x_nnz - i);
src/engine/engine_util_sparse.c
+19 -14
View File
@@ -208,7 +208,7 @@ static int mju_compare(const int* vec1, const int* vec2, int n) {
// 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,
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
mjtNum* buf, int* buf_ind) {
// check for identical pattern
@@ -226,40 +226,45 @@ int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
memcpy(buf_ind, dst_ind, dst_nnz*sizeof(int));
}
// prepare to merge buf and scr into dst
// prepare to merge buf and src into dst
int bi = 0, si = 0, nnz = 0;
int buf_nnz = dst_nnz;
int badr = bi < buf_nnz ? buf_ind[bi] : n+1;
int sadr = si < src_nnz ? src_ind[si] : n+1;
// merge vectors
while (bi < buf_nnz || si < src_nnz) {
// both
while (bi < buf_nnz && si < src_nnz) {
int badr = buf_ind[bi];
int sadr = src_ind[si];
if (badr == sadr) {
dst[nnz] = a*buf[bi++] + b*src[si++];
dst_ind[nnz++] = badr;
badr = bi < buf_nnz ? buf_ind[bi] : n+1;
sadr = si < src_nnz ? src_ind[si] : n+1;
}
// dst only
// buf only
else if (badr < sadr) {
dst[nnz] = a*buf[bi++];
dst_ind[nnz++] = badr;
badr = bi < buf_nnz ? buf_ind[bi] : n+1;
}
// src only
else {
dst[nnz] = b*src[si++];
dst_ind[nnz++] = sadr;
sadr = si < src_nnz ? src_ind[si] : n+1;
}
}
// the rest of src only
while (si < src_nnz) {
dst[nnz] = b*src[si];
dst_ind[nnz++] = src_ind[si++];
}
// the rest of buf only
while (bi < buf_nnz) {
dst[nnz] = a*buf[bi];
dst_ind[nnz++] = buf_ind[bi++];
}
return nnz;
}
src/engine/engine_util_sparse.h
+1 -1
View File
@@ -52,7 +52,7 @@ 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,
int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b,
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
mjtNum* buf, int* buf_ind);
test/benchmark/engine_util_sparse_benchmark_test.cc
+3 -3
View File
@@ -206,7 +206,7 @@ int compare_memcmp(const int* vec1,
}
int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
const mjtNum* src, int n,
const mjtNum* src,
mjtNum a, mjtNum b,
int dst_nnz, int src_nnz,
int* dst_ind,
@@ -223,7 +223,7 @@ int ABSL_ATTRIBUTE_NOINLINE combineSparse_baseline(mjtNum* dst,
}
int ABSL_ATTRIBUTE_NOINLINE combineSparse_new(mjtNum* dst,
const mjtNum* src, int n,
const mjtNum* src,
mjtNum a, mjtNum b,
int dst_nnz, int src_nnz,
int* dst_ind,
@@ -470,7 +470,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
// true arguments should be i+1 and colind+rowadr[r]
// but instead we repeat rownnz[c] and colind+rowadr[c]
// in order to trigger all if's in combineSparse
func(H+rowadr[c], H+rowadr[r], c+1, 1, -H[adr+i],
func(H+rowadr[c], H+rowadr[r], 1, -H[adr+i],
rownnz[c], rownnz[c],
colind+rowadr[c], colind+rowadr[c], NULL, NULL);
}