From ba12a1354b485b7c940f95d3a159ddee802e4aaa Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Fri, 18 Aug 2023 06:13:20 -0700 Subject: [PATCH] Optimize and refactor combineSparse and combineSparseCount. PiperOrigin-RevId: 558126391 Change-Id: Iead14c31136e789560f3c8124b84bd4e383158f4 --- src/engine/engine_core_constraint.c | 36 +++++++++---------- src/engine/engine_core_smooth.c | 6 ++-- src/engine/engine_support.c | 2 +- src/engine/engine_util_solve.c | 9 +++-- src/engine/engine_util_sparse.c | 33 +++++++++-------- src/engine/engine_util_sparse.h | 2 +- .../engine_util_sparse_benchmark_test.cc | 6 ++-- 7 files changed, 49 insertions(+), 45 deletions(-) diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index a29c6993..354b4056 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -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; } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 53168dad..1609ab07 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -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 diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 4af5504c..34793d5e 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -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); } diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 2d6c5b90..fa2d869a 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -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); diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 6f9cc8f0..e9925646 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -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; } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 698bbbe9..a4852e63 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -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); diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index 45a13df5..de493ae3 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -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); }