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

Allow CSR back-substitution to handle multiple vectors.

Browse Source 
PiperOrigin-RevId: 713229673
Change-Id: I7a5b43fe966cf9e482bd41e2eea6c30dd3ffa1d4
This commit is contained in:
Yuval Tassa
2025-01-08 03:30:41 -08:00
committed by Copybara-Service
parent 4d82ab5762
commit 8a5f092081
5 changed files with 116 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+73 -24
View File
@@ -1607,39 +1607,88 @@ void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// like mj_solveLD, but using the CSR representation of L
void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
void mj_solveLDs(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
// x <- L^-T x
for (int i=nv-1; i > 0; i--) {
// skip diagonal (simple) rows, exploit sparsity of input vector
if (diagnum[i] || x[i] == 0) {
continue;
// single vector
if (n == 1) {
// x <- L^-T x
for (int i=nv-1; i > 0; i--) {
// skip diagonal rows, zero elements in input vector
mjtNum x_i = x[i];
if (x_i == 0 || diagnum[i]) {
continue;
}
int start = rowadr[i];
int end = start + rownnz[i] - 1;
for (int adr=start; adr < end; adr++) {
x[colind[adr]] -= qLDs[adr] * x_i;
}
}
int d = rownnz[i] - 1;
int adr_i = rowadr[i];
mjtNum x_i = x[i];
for (int j=0; j < d; j++) {
int adr = adr_i + j;
x[colind[adr]] -= qLDs[adr] * x_i;
// x <- D^-1 x
for (int i=0; i < nv; i++) {
x[i] *= qLDiagInv[i];
}
// x <- L^-1 x
for (int i=1; i < nv; i++) {
// skip diagonal rows
if (diagnum[i]) {
i += diagnum[i] - 1; // iterating forward: skip ahead, adjust i
continue;
}
int adr = rowadr[i];
x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0);
}
}
// x(i) /= D(i,i)
for (int i=0; i < nv; i++) {
x[i] *= qLDiagInv[i];
}
// multiple vectors
else {
// x <- L^-T x
for (int i=nv-1; i > 0; i--) {
// skip diagonal rows
if (diagnum[i]) {
continue;
}
// x <- L^-1 x
for (int i=1; i < nv; i++) {
// skip diagonal (simple) rows
if (diagnum[i]) {
i += diagnum[i] - 1; // when iterating forward we can skip ahead
continue;
int start = rowadr[i];
int end = start + rownnz[i] - 1;
for (int adr=start; adr < end; adr++) {
int j = colind[adr];
mjtNum val = qLDs[adr];
for (int offset=0; offset < n*nv; offset+=nv) {
mjtNum x_i;
if ((x_i = x[i+offset])) {
x[j+offset] -= val * x_i;
}
}
}
}
int adr = rowadr[i];
x[i] -= mju_dotSparse(qLDs+adr, x, rownnz[i] - 1, colind+adr, /*flg_unc1=*/0);
// x <- D^-1 x
for (int i=0; i < nv; i++) {
mjtNum invD_i = qLDiagInv[i];
for (int offset=0; offset < n*nv; offset+=nv) {
x[i+offset] *= invD_i;
}
}
// x <- L^-1 x
for (int i=1; i < nv; i++) {
// skip diagonal rows
if (diagnum[i]) {
i += diagnum[i] - 1; // iterating forward: skip ahead, adjust i
continue;
}
int adr = rowadr[i];
int d = rownnz[i] - 1;
for (int offset=0; offset < n*nv; offset+=nv) {
x[i+offset] -= mju_dotSparse(qLDs+adr, x+offset, d, colind+adr, /*flg_unc1=*/0);
}
}
}
}
src/engine/engine_core_smooth.h
+2 -2
View File
@@ -64,8 +64,8 @@ MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
const mjtNum* qLD, const mjtNum* qLDiagInv);
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// like mj_solveLD, but using the CSR representation of L
MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv,
// handle n vectors at once
MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
// sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d
test/benchmark/inertia_benchmark_test.cc
+1 -1
View File
@@ -72,7 +72,7 @@ static void BM_solve(benchmark::State& state, SolveType type) {
mj_factorIs(LDs, d->qLDiagInv, m->nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mju_copy(res, vec, m->nv);
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
}
}
test/benchmark/solveLD_benchmark_test.cc
+1 -1
View File
@@ -63,7 +63,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
mj_solveM(m, d, res, vec, 1);
} else {
mju_copy(res, vec, m->nv);
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv,
mj_solveLDs(res, LDs, d->qLDiagInv, m->nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
}
}
test/engine/engine_core_smooth_test.cc
+39 -1
View File
@@ -495,7 +495,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0;
mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv,
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision
@@ -507,6 +507,44 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mj_deleteModel(m);
}
TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
char error[1024];
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
mjData* d = mj_makeData(m);
mj_forward(m, d);
int nv = m->nv;
int nC = m->nC;
// copy LD into LDs: CSR format
vector<mjtNum> LDs(nC);
for (int i=0; i < nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
}
// compare n LD and LDs vector solve
int n = 3;
vector<mjtNum> vec(nv*n);
vector<mjtNum> vec2(nv*n);
for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i;
for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0;
mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv*n; i++) {
EXPECT_FLOAT_EQ(vec[i], vec2[i]);
}
mj_deleteData(d);
mj_deleteModel(m);
}
TEST_F(CoreSmoothTest, FactorIs) {
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
char error[1024];