Yuval Tassa
7eb8231fda
PiperOrigin-RevId: 712488529 Change-Id: Iad91c72654376539791d7856765a0d0ac9088251
915 lines
22 KiB
C
915 lines
22 KiB
C
// Copyright 2021 DeepMind Technologies Limited
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "engine/engine_util_sparse.h"
|
|
#include "engine/engine_util_sparse_avx.h" // IWYU pragma: keep
|
|
|
|
#include <string.h>
|
|
|
|
#include <mujoco/mjdata.h>
|
|
#include <mujoco/mjsan.h> // IWYU pragma: keep
|
|
#include <mujoco/mjtnum.h>
|
|
#include "engine/engine_io.h"
|
|
#include "engine/engine_util_blas.h"
|
|
#include "engine/engine_util_misc.h"
|
|
|
|
|
|
//------------------------------ sparse operations -------------------------------------------------
|
|
|
|
|
|
// dot-productX3, first vector is sparse; supernode of size 3
|
|
void mju_dotSparseX3(mjtNum* res0, mjtNum* res1, mjtNum* res2,
|
|
const mjtNum* vec10, const mjtNum* vec11, const mjtNum* vec12,
|
|
const mjtNum* vec2, int nnz1, const int* ind1) {
|
|
#ifdef mjUSEAVX
|
|
mju_dotSparseX3_avx(res0, res1, res2, vec10, vec11, vec12, vec2, nnz1, ind1);
|
|
#else
|
|
int i = 0;
|
|
|
|
// clear result
|
|
mjtNum RES0 = 0;
|
|
mjtNum RES1 = 0;
|
|
mjtNum RES2 = 0;
|
|
|
|
for (; i < nnz1; i++) {
|
|
mjtNum v2 = vec2[ind1[i]];
|
|
|
|
RES0 += vec10[i] * v2;
|
|
RES1 += vec11[i] * v2;
|
|
RES2 += vec12[i] * v2;
|
|
}
|
|
|
|
// copy result
|
|
*res0 = RES0;
|
|
*res1 = RES1;
|
|
*res2 = RES2;
|
|
#endif // mjUSEAVX
|
|
}
|
|
|
|
|
|
|
|
// dot-product, both vectors are sparse
|
|
// flg_unc2: is vec2 memory layout uncompressed
|
|
mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, int nnz1, const int* ind1, int nnz2,
|
|
const int* ind2, int flg_unc2) {
|
|
int i1 = 0, i2 = 0;
|
|
mjtNum res = 0;
|
|
|
|
// check for empty array
|
|
if (!nnz1 || !nnz2) {
|
|
return 0;
|
|
}
|
|
|
|
while (i1 < nnz1 && i2 < nnz2) {
|
|
// get current indices
|
|
int adr1 = ind1[i1], adr2 = ind2[i2];
|
|
|
|
// match: accumulate result, advance both
|
|
if (adr1 == adr2) {
|
|
if (flg_unc2) {
|
|
res += vec1[i1++] * vec2[adr2];
|
|
i2++;
|
|
} else {
|
|
res += vec1[i1++] * vec2[i2++];
|
|
}
|
|
}
|
|
|
|
// otherwise advance smaller
|
|
else if (adr1 < adr2) {
|
|
i1++;
|
|
} else {
|
|
i2++;
|
|
}
|
|
}
|
|
|
|
return res;
|
|
}
|
|
|
|
|
|
|
|
// convert matrix from dense to sparse
|
|
// nnz is size of res and colind, return 1 if too small, 0 otherwise
|
|
int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
|
int* rownnz, int* rowadr, int* colind, int nnz) {
|
|
if (nnz <= 0) {
|
|
return 1;
|
|
}
|
|
|
|
int adr = 0;
|
|
|
|
// find non-zeros and construct sparse
|
|
for (int r=0; r < nr; r++) {
|
|
// init row
|
|
rownnz[r] = 0;
|
|
rowadr[r] = adr;
|
|
|
|
// find non-zeros
|
|
for (int c=0; c < nc; c++) {
|
|
if (mat[r*nc+c]) {
|
|
// check for out of bounds
|
|
if (adr >= nnz) {
|
|
return 1;
|
|
}
|
|
|
|
// record index and count
|
|
colind[adr] = c;
|
|
rownnz[r]++;
|
|
|
|
// copy element
|
|
res[adr++] = mat[r*nc+c];
|
|
}
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
|
|
|
|
// convert matrix from sparse to dense
|
|
void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
|
|
const int* rownnz, const int* rowadr, const int* colind) {
|
|
// clear
|
|
mju_zero(res, nr*nc);
|
|
|
|
// copy non-zeros
|
|
for (int r=0; r < nr; r++) {
|
|
for (int i=0; i < rownnz[r]; i++) {
|
|
res[r*nc + colind[rowadr[r]+i]] = mat[rowadr[r]+i];
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// 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) {
|
|
#ifdef mjUSEAVX
|
|
mju_mulMatVecSparse_avx(res, mat, vec, nr, rownnz, rowadr, colind, rowsuper);
|
|
#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], /*flg_unc1=*/0);
|
|
}
|
|
#endif // mjUSEAVX
|
|
}
|
|
|
|
|
|
|
|
// multiply transposed sparse matrix and dense vector: res = mat' * vec.
|
|
void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc,
|
|
const int* rownnz, const int* rowadr, const int* colind) {
|
|
// clear res
|
|
mju_zero(res, nc);
|
|
|
|
for (int i=0; i < nr; i++) {
|
|
mjtNum scl = vec[i];
|
|
|
|
// skip if 0
|
|
if (!scl) continue;
|
|
|
|
// add row scaled by the corresponding vector element
|
|
int nnz = rownnz[i];
|
|
int adr = rowadr[i];
|
|
const int* ind = colind + adr;
|
|
const mjtNum* row = mat + adr;
|
|
for (int j=0; j < nnz; j++) {
|
|
res[ind[j]] += row[j] * scl;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// res = res*scl1 + vec*scl2
|
|
static void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) {
|
|
#ifdef mjUSEAVX
|
|
mju_addToSclScl_avx(res, vec, scl1, scl2, n);
|
|
#else
|
|
for (int i=0; i < n; i++) {
|
|
res[i] = res[i]*scl1 + vec[i]*scl2;
|
|
}
|
|
#endif // mjUSEAVX
|
|
}
|
|
|
|
|
|
|
|
// return 1 if vec1==vec2, 0 otherwise
|
|
static int mju_compare(const int* vec1, const int* vec2, int n) {
|
|
#ifdef mjUSEAVX
|
|
return mju_compare_avx(vec1, vec2, n);
|
|
#else
|
|
return !memcmp(vec1, vec2, n*sizeof(int));
|
|
#endif // mjUSEAVX
|
|
}
|
|
|
|
|
|
|
|
// 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, b = 0, c_nnz = 0;
|
|
|
|
// count c_nnz: nonzero indices common to both a and b
|
|
while (a < a_nnz && b < b_nnz) {
|
|
// common index, increment everything
|
|
if (a_ind[a] == b_ind[b]) {
|
|
c_nnz++;
|
|
a++;
|
|
b++;
|
|
}
|
|
|
|
// update smallest index
|
|
else if (a_ind[a] < b_ind[b]) {
|
|
a++;
|
|
} else {
|
|
b++;
|
|
}
|
|
}
|
|
|
|
// union minus the intersection
|
|
return a_nnz + b_nnz - c_nnz;
|
|
}
|
|
|
|
|
|
|
|
// combine two sparse vectors: dst = a*dst + b*src, return nnz of result
|
|
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
|
|
if (dst_nnz == src_nnz) {
|
|
if (mju_compare(dst_ind, src_ind, dst_nnz)) {
|
|
// combine mjtNum data directly
|
|
mju_addToSclScl(dst, src, a, b, dst_nnz);
|
|
return dst_nnz;
|
|
}
|
|
}
|
|
|
|
// copy dst into buf
|
|
if (dst_nnz) {
|
|
mju_copy(buf, dst, dst_nnz);
|
|
mju_copyInt(buf_ind, dst_ind, dst_nnz);
|
|
}
|
|
|
|
// prepare to merge buf and src into dst
|
|
int bi = 0, si = 0, nnz = 0;
|
|
int buf_nnz = dst_nnz;
|
|
|
|
// merge vectors
|
|
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;
|
|
}
|
|
|
|
// buf only
|
|
else if (badr < sadr) {
|
|
dst[nnz] = a*buf[bi++];
|
|
dst_ind[nnz++] = badr;
|
|
}
|
|
|
|
// src only
|
|
else {
|
|
dst[nnz] = b*src[si++];
|
|
dst_ind[nnz++] = sadr;
|
|
}
|
|
}
|
|
|
|
// 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;
|
|
}
|
|
|
|
|
|
|
|
// incomplete combine sparse: dst = a*dst + b*src at common indices
|
|
void mju_combineSparseInc(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
|
|
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind) {
|
|
// check for identical pattern
|
|
if (dst_nnz == src_nnz) {
|
|
if (mju_compare(dst_ind, src_ind, dst_nnz)) {
|
|
// combine mjtNum data directly
|
|
mju_addToSclScl(dst, src, a, b, dst_nnz);
|
|
return;
|
|
}
|
|
}
|
|
|
|
// scale dst by a
|
|
if (a != 1) {
|
|
mju_scl(dst, dst, a, dst_nnz);
|
|
}
|
|
|
|
// prepare to merge
|
|
int di = 0, si = 0;
|
|
int dadr = di < dst_nnz ? dst_ind[di] : n+1;
|
|
int sadr = si < src_nnz ? src_ind[si] : n+1;
|
|
|
|
// add src*b at common indices
|
|
while (di < dst_nnz) {
|
|
// both
|
|
if (dadr == sadr) {
|
|
dst[di++] += b*src[si++];
|
|
|
|
dadr = di < dst_nnz ? dst_ind[di] : n+1;
|
|
sadr = si < src_nnz ? src_ind[si] : n+1;
|
|
}
|
|
|
|
// dst only
|
|
else if (dadr < sadr) {
|
|
di++;
|
|
dadr = di < dst_nnz ? dst_ind[di] : n+1;
|
|
}
|
|
|
|
// src only
|
|
else {
|
|
si++;
|
|
sadr = si < src_nnz ? src_ind[si] : n+1;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// dst += scl*src, only at common non-zero indices
|
|
void mju_addToSclSparseInc(mjtNum* dst, const mjtNum* src,
|
|
int nnzdst, const int* inddst,
|
|
int nnzsrc, const int* indsrc, mjtNum scl) {
|
|
if (!nnzdst || !nnzsrc) {
|
|
return;
|
|
}
|
|
|
|
int adrs = 0, adrd = 0, inds = indsrc[0], indd = inddst[0];
|
|
while (1) {
|
|
// common non-zero index
|
|
if (inds == indd) {
|
|
// add
|
|
dst[adrd] += scl * src[adrs];
|
|
|
|
// advance src
|
|
if (++adrs < nnzsrc) {
|
|
inds = indsrc[adrs];
|
|
} else {
|
|
return;
|
|
}
|
|
|
|
// advance dst
|
|
if (++adrd < nnzdst) {
|
|
indd = inddst[adrd];
|
|
} else {
|
|
return;
|
|
}
|
|
}
|
|
|
|
// src non-zero index smaller: advance src
|
|
else if (inds < indd) {
|
|
if (++adrs < nnzsrc) {
|
|
inds = indsrc[adrs];
|
|
} else {
|
|
return;
|
|
}
|
|
}
|
|
|
|
// dst non-zero index smaller: advance dst
|
|
else {
|
|
if (++adrd < nnzdst) {
|
|
indd = inddst[adrd];
|
|
} else {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// add to sparse matrix: dst = dst + scl*src, return nnz of result
|
|
int mju_addToSparseMat(mjtNum* dst, const mjtNum* src, int n, int nrow, mjtNum scl,
|
|
int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
|
|
mjtNum* buf, int* buf_ind) {
|
|
// check for identical pattern
|
|
if (dst_nnz == src_nnz) {
|
|
if (dst_nnz == 0) {
|
|
return 0;
|
|
}
|
|
if (mju_compare(dst_ind, src_ind, dst_nnz)) {
|
|
// combine mjtNum data directly
|
|
mju_addToScl(dst, src, scl, nrow*dst_nnz);
|
|
return dst_nnz;
|
|
}
|
|
}
|
|
|
|
// prepare to merge scr and dst into buf^T
|
|
int si = 0, di = 0, nnz = 0;
|
|
int sadr = src_nnz ? src_ind[0] : n+1;
|
|
int dadr = dst_nnz ? dst_ind[0] : n+1;
|
|
|
|
// merge matrices
|
|
while (si < src_nnz || di < dst_nnz) {
|
|
// both
|
|
if (sadr == dadr) {
|
|
for (int k=0; k < nrow; k++) {
|
|
buf[nrow*nnz + k] = dst[di + k*dst_nnz] + scl*src[si + k*src_nnz];
|
|
}
|
|
|
|
buf_ind[nnz++] = sadr;
|
|
si++;
|
|
di++;
|
|
sadr = si < src_nnz ? src_ind[si] : n+1;
|
|
dadr = di < dst_nnz ? dst_ind[di] : n+1;
|
|
}
|
|
|
|
// dst only
|
|
else if (dadr < sadr) {
|
|
for (int k=0; k < nrow; k++) {
|
|
buf[nrow*nnz + k] = dst[di + k*dst_nnz];
|
|
}
|
|
|
|
buf_ind[nnz++] = dadr;
|
|
di++;
|
|
dadr = di < dst_nnz ? dst_ind[di] : n+1;
|
|
}
|
|
|
|
// src only
|
|
else {
|
|
for (int k=0; k < nrow; k++) {
|
|
buf[nrow*nnz + k] = scl*src[si + k*src_nnz];
|
|
}
|
|
|
|
buf_ind[nnz++] = sadr;
|
|
si++;
|
|
sadr = si < src_nnz ? src_ind[si] : n+1;
|
|
}
|
|
}
|
|
|
|
// copy transposed buf into dst
|
|
mju_transpose(dst, buf, nnz, nrow);
|
|
mju_copyInt(dst_ind, buf_ind, nnz);
|
|
|
|
return nnz;
|
|
}
|
|
|
|
|
|
|
|
// add(merge) two chains
|
|
int mju_addChains(int* res, int n, int NV1, int NV2,
|
|
const int* chain1, const int* chain2) {
|
|
// check for identical pattern
|
|
if (NV1 == NV2) {
|
|
if (NV1 == 0) {
|
|
return 0;
|
|
}
|
|
if (mju_compare(chain1, chain2, NV1)) {
|
|
mju_copyInt(res, chain1, NV1);
|
|
return NV1;
|
|
}
|
|
}
|
|
|
|
// prepare to merge
|
|
int i1 = 0, i2 = 0, NV = 0;
|
|
int adr1 = NV1 ? chain1[0] : n+1;
|
|
int adr2 = NV2 ? chain2[0] : n+1;
|
|
|
|
// merge chains
|
|
while (i1 < NV1 || i2 < NV2) {
|
|
// both
|
|
if (adr1 == adr2) {
|
|
res[NV++] = adr1;
|
|
i1++;
|
|
i2++;
|
|
adr1 = i1 < NV1 ? chain1[i1] : n+1;
|
|
adr2 = i2 < NV2 ? chain2[i2] : n+1;
|
|
}
|
|
|
|
// chain1 only
|
|
else if (adr1 < adr2) {
|
|
res[NV++] = adr1;
|
|
i1++;
|
|
adr1 = i1 < NV1 ? chain1[i1] : n+1;
|
|
}
|
|
|
|
// chain2 only
|
|
else {
|
|
res[NV++] = adr2;
|
|
i2++;
|
|
adr2 = i2 < NV2 ? chain2[i2] : n+1;
|
|
}
|
|
}
|
|
|
|
return NV;
|
|
}
|
|
|
|
|
|
|
|
// compress layout of sparse matrix
|
|
void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind) {
|
|
rowadr[0] = 0;
|
|
int adr = rownnz[0];
|
|
for (int r=1; r < nr; r++) {
|
|
// save old rowadr, record new
|
|
int rowadr1 = rowadr[r];
|
|
rowadr[r] = adr;
|
|
|
|
// shift mat and mat_colind
|
|
for (int adr1=rowadr1; adr1 < rowadr1+rownnz[r]; adr1++) {
|
|
mat[adr] = mat[adr1];
|
|
colind[adr] = colind[adr1];
|
|
adr++;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// 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) {
|
|
// clear number of non-zeros for each row of transposed
|
|
mju_zeroInt(res_rownnz, nc);
|
|
|
|
// count the number of non-zeros for each row of the transposed matrix
|
|
for (int r = 0; r < nr; r++) {
|
|
int start = rowadr[r];
|
|
int end = start + rownnz[r];
|
|
for (int j = start; j < end; j++) {
|
|
res_rownnz[colind[j]]++;
|
|
}
|
|
}
|
|
|
|
// 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];
|
|
}
|
|
|
|
// iterate through each non-zero entry of mat
|
|
for (int r = 0; r < nr; r++) {
|
|
int start = rowadr[r];
|
|
int end = start + rownnz[r];
|
|
for (int i = start; i < end; i++) {
|
|
// swap rows with columns and increment res_rowadr
|
|
int c = res_rowadr[colind[i]]++;
|
|
res_colind[c] = r;
|
|
if (res) {
|
|
res[c] = mat[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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) {
|
|
// no rows: nothing to do
|
|
if (!nr) {
|
|
return;
|
|
}
|
|
|
|
// find match to child
|
|
for (int r=0; r < nr-1; r++) {
|
|
// different number of nonzeros: cannot be a match
|
|
if (rownnz[r] != rownnz[r+1]) {
|
|
rowsuper[r] = 0;
|
|
}
|
|
|
|
// same number of nonzeros: compare colind vectors
|
|
else {
|
|
rowsuper[r] = mju_compare(colind+rowadr[r], colind+rowadr[r+1], rownnz[r]);
|
|
}
|
|
}
|
|
|
|
// clear last (by definition)
|
|
rowsuper[nr-1] = 0;
|
|
|
|
// accumulate in reverse
|
|
for (int r=nr-2; r >= 0; r--) {
|
|
if (rowsuper[r]) {
|
|
rowsuper[r] += rowsuper[r+1];
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
|
|
void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
|
|
const int* rownnz, const int* rowadr, const int* colind,
|
|
const int* rownnzT, const int* rowadrT, const int* colindT,
|
|
const int* rowsuperT, mjData* d, int flg_upper) {
|
|
mj_markStack(d);
|
|
int* chain = mjSTACKALLOC(d, 2*nr, int);
|
|
int nchain = 0;
|
|
int* res_colind = NULL;
|
|
|
|
for (int r=0; r < nr; r++) {
|
|
// supernode; copy everything to next row
|
|
if (rowsuperT && r > 0 && rowsuperT[r-1] > 0) {
|
|
res_rownnz[r] = res_rownnz[r - 1];
|
|
|
|
// fill in upper triangle
|
|
if (flg_upper) {
|
|
for (int j=0; j < nchain; j++) {
|
|
res_rownnz[res_colind[j]]++;
|
|
}
|
|
}
|
|
|
|
// update chain with diagonal
|
|
if (rownnzT[r]) {
|
|
res_colind[nchain++] = r;
|
|
res_rownnz[r]++;
|
|
}
|
|
} else {
|
|
int inew = 0, iold = nr;
|
|
nchain = 0;
|
|
|
|
for (int i=0; i < rownnzT[r]; i++) {
|
|
int c = colindT[rowadrT[r] + i];
|
|
|
|
int adr = inew;
|
|
inew = iold;
|
|
iold = adr;
|
|
|
|
int nnewchain = 0;
|
|
adr = 0;
|
|
int end = rowadr[c] + rownnz[c];
|
|
for (int adr1=rowadr[c]; adr1 < end; adr1++) {
|
|
int col_mat = colind[adr1];
|
|
while (adr < nchain && chain[iold + adr] < col_mat && chain[iold + adr] <= r) {
|
|
chain[inew + nnewchain++] = chain[iold + adr++];
|
|
}
|
|
|
|
// skip upper triangle
|
|
if (col_mat > r) {
|
|
break;
|
|
}
|
|
|
|
if (adr < nchain && chain[iold + adr] == col_mat) {
|
|
adr++;
|
|
}
|
|
chain[inew + nnewchain++] = col_mat;
|
|
}
|
|
|
|
while (adr < nchain && chain[iold + adr] <= r) {
|
|
chain[inew + nnewchain++] = chain[iold + adr++];
|
|
}
|
|
nchain = nnewchain;
|
|
}
|
|
|
|
// only computed for lower triangle
|
|
res_rownnz[r] = nchain;
|
|
res_colind = chain + inew;
|
|
|
|
// update upper triangle
|
|
if (flg_upper) {
|
|
int nchain_end = nchain;
|
|
|
|
// avoid double counting
|
|
if (nchain > 0 && res_colind[nchain-1] == r) {
|
|
nchain_end = nchain - 1;
|
|
}
|
|
|
|
for (int j=0; j < nchain_end; j++) {
|
|
res_rownnz[res_colind[j]]++;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
res_rowadr[0] = 0;
|
|
for (int r = 1; r < nr; r++) {
|
|
res_rowadr[r] = res_rowadr[r-1] + res_rownnz[r-1];
|
|
}
|
|
|
|
mj_freeStack(d);
|
|
}
|
|
|
|
|
|
// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
|
|
void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc) {
|
|
for (int r=0; r < nc; r++) {
|
|
res_rowadr[r] = r*nc;
|
|
}
|
|
}
|
|
|
|
|
|
|
|
// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout
|
|
// res_rowadr is required to be precomputed
|
|
void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
|
|
const mjtNum* diag, int nr, int nc,
|
|
int* res_rownnz, const int* res_rowadr, int* res_colind,
|
|
const int* rownnz, const int* rowadr,
|
|
const int* colind, const int* rowsuper,
|
|
const int* rownnzT, const int* rowadrT,
|
|
const int* colindT, const int* rowsuperT,
|
|
mjData* d, int flg_upper) {
|
|
// allocate space for accumulation buffer and matT
|
|
mj_markStack(d);
|
|
|
|
// a dense row buffer that stores the current row in the resulting matrix
|
|
mjtNum* buffer = mjSTACKALLOC(d, nc, mjtNum);
|
|
|
|
// these mark the currently set columns in the dense row buffer,
|
|
// used for when creating the resulting sparse row
|
|
int* markers = mjSTACKALLOC(d, nc, int);
|
|
|
|
for (int i=0; i < nc; i++) {
|
|
int rowadr_i = res_rowadr[i];
|
|
int* cols = res_colind + rowadr_i;
|
|
|
|
res_rownnz[i] = 0;
|
|
buffer[i] = 0;
|
|
markers[i] = 0;
|
|
|
|
// if rowsuper, use the previous row sparsity structure
|
|
if (rowsuperT && i > 0 && rowsuperT[i-1]) {
|
|
res_rownnz[i] = res_rownnz[i-1];
|
|
mju_copyInt(cols, res_colind+res_rowadr[i-1], res_rownnz[i]);
|
|
}
|
|
|
|
// iterate through each row of M'
|
|
int adrT = rowadrT[i];
|
|
int end_r = adrT + rownnzT[i];
|
|
for (int r = adrT; r < end_r; r++) {
|
|
int t = colindT[r];
|
|
int adr = rowadr[t];
|
|
int end_c = adr + rownnz[t];
|
|
for (int c=adr; c < end_c; c++) {
|
|
int cc = colind[c];
|
|
|
|
// ignore upper triangle
|
|
if (cc > i) {
|
|
break;
|
|
}
|
|
|
|
// add value to buffer
|
|
if (diag) {
|
|
buffer[cc] += matT[r] * diag[t] * mat[c];
|
|
} else {
|
|
buffer[cc] += matT[r] * mat[c];
|
|
}
|
|
|
|
// only need to insert nnz if not marked
|
|
if (!markers[cc]) {
|
|
markers[cc] = 1;
|
|
|
|
// since i is the rightmost column, it can be inserted at the end
|
|
if (cc == i) {
|
|
cols[res_rownnz[i]++] = cc;
|
|
continue;
|
|
}
|
|
|
|
// insert col in order via binary search
|
|
int l = 0, h = res_rownnz[i];
|
|
while (l < h) {
|
|
int m = (l + h) >> 1;
|
|
if (cols[m] < cc) {
|
|
l = m + 1;
|
|
} else {
|
|
h = m;
|
|
}
|
|
}
|
|
|
|
// cc is the rightmost column so far, it can be inserted at the end
|
|
if (l == res_rownnz[i]) {
|
|
cols[l] = cc;
|
|
res_rownnz[i]++;
|
|
continue;
|
|
}
|
|
|
|
// move the cols to the right
|
|
h = res_rownnz[i];
|
|
while (l < h) {
|
|
cols[h] = cols[h-1];
|
|
h--;
|
|
}
|
|
|
|
// insert
|
|
cols[l] = cc;
|
|
res_rownnz[i]++;
|
|
}
|
|
}
|
|
}
|
|
|
|
end_r = res_rownnz[i];
|
|
|
|
// rowsuperT: reuse sparsity, copy into res
|
|
if (rowsuperT && rowsuperT[i]) {
|
|
for (int r=0; r < end_r; r++) {
|
|
int c = cols[r];
|
|
res[rowadr_i + r] = buffer[c];
|
|
buffer[c] = 0;
|
|
}
|
|
}
|
|
|
|
// clear out buffers, sparsity cannot be reused
|
|
else {
|
|
for (int r=0; r < end_r; r++) {
|
|
int c = cols[r];
|
|
int adr = rowadr_i + r;
|
|
res[adr] = buffer[c];
|
|
res_colind[adr] = c;
|
|
buffer[c] = 0;
|
|
markers[c] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
// fill upper triangle
|
|
if (flg_upper) {
|
|
for (int i=0; i < nc; i++) {
|
|
int start = res_rowadr[i];
|
|
int end = start + res_rownnz[i] - 1;
|
|
for (int j=start; j < end; j++) {
|
|
int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
|
|
res[adr] = res[j];
|
|
res_colind[adr] = i;
|
|
}
|
|
}
|
|
}
|
|
|
|
mj_freeStack(d);
|
|
}
|
|
|
|
// compute row non-zeros of reverse-Cholesky factor L, return total non-zeros
|
|
// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
|
|
// note: reads pattern from upper triangle
|
|
int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
|
|
int n, mjData* d) {
|
|
mj_markStack(d);
|
|
int* parent = mjSTACKALLOC(d, n, int);
|
|
int* flag = mjSTACKALLOC(d, n, int);
|
|
|
|
// loop over rows in reverse order
|
|
for (int r = n - 1; r >= 0; r--) {
|
|
parent[r] = -1;
|
|
flag[r] = r;
|
|
L_rownnz[r] = 1; // start with 1 for diagonal
|
|
|
|
// loop over non-zero columns
|
|
int start = rowadr[r];
|
|
int end = start + rownnz[r];
|
|
for (int c = start; c < end; c++) {
|
|
int i = colind[c];
|
|
if (i > r) {
|
|
// traverse from i to ancestor, stop when row is flagged
|
|
while (flag[i] != r) {
|
|
// if not yet set, set parent to current row
|
|
if (parent[i] == -1) {
|
|
parent[i] = r;
|
|
}
|
|
|
|
// increment non-zeros, flag row i, advance to parent
|
|
L_rownnz[i]++;
|
|
flag[i] = r;
|
|
i = parent[i];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
mj_freeStack(d);
|
|
|
|
// sum up all row non-zeros
|
|
int nnz = 0;
|
|
for (int r = 0; r < n; r++) {
|
|
nnz += L_rownnz[r];
|
|
}
|
|
|
|
return nnz;
|
|
}
|