From 8a5f0920816ece01bba7b23265db8927a1c77845 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 8 Jan 2025 03:30:41 -0800 Subject: [PATCH] Allow CSR back-substitution to handle multiple vectors. PiperOrigin-RevId: 713229673 Change-Id: I7a5b43fe966cf9e482bd41e2eea6c30dd3ffa1d4 --- src/engine/engine_core_smooth.c | 97 ++++++++++++++++++------ src/engine/engine_core_smooth.h | 4 +- test/benchmark/inertia_benchmark_test.cc | 2 +- test/benchmark/solveLD_benchmark_test.cc | 2 +- test/engine/engine_core_smooth_test.cc | 40 +++++++++- 5 files changed, 116 insertions(+), 29 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 81e276f6..a3e255c7 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -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); + } + } } } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index f4733b08..a710416c 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -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 diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index 8ba580b7..7ddc3468 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -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); } } diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index f2dc39e1..18ebe69b 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -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); } } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 717fde98..e5f66c38 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -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 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 vec(nv*n); + vector 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];