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

Remove unnecessary diagnum argument in mj_solveLD and mj_factorI

Browse Source 
PiperOrigin-RevId: 758669364
Change-Id: Icfefc4a7a1e7d28da373725358ccf9d5f589f689
This commit is contained in:
Yuval Tassa
2025-05-14 07:12:19 -07:00
committed by Copybara-Service
parent b0c36c6b33
commit d8bebdc675
9 changed files with 33 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+16 -21
View File
@@ -1653,7 +1653,7 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
void mj_factorM(const mjModel* m, mjData* d) {
TM_START;
mju_copy(d->qLD, d->M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind);
TM_ADD(mjTIMER_POS_INERTIA);
}
@@ -1661,32 +1661,28 @@ void mj_factorM(const mjModel* m, mjData* d) {
// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
const int* rownnz, const int* rowadr, const int* colind) {
// backward loop over rows
for (int k=nv-1; k >= 0; k--) {
// get row k's address, diagonal index, inverse diagonal value
int rowadr_k = rowadr[k];
int diag_k = rowadr_k + rownnz[k] - 1;
mjtNum invD = 1 / mat[diag_k];
int start = rowadr[k];
int diag = rownnz[k] - 1;
int end = start + diag;
mjtNum invD = 1 / mat[end];
if (diaginv) diaginv[k] = invD;
// skip if simple
if (diagnum[k]) {
continue;
}
// update triangle above row k, inclusive
for (int adr=diag_k - 1; adr >= rowadr_k; adr--) {
// update triangle above row k
for (int adr=end - 1; adr >= start; adr--) {
// tmp = L(k, i) / L(k, k)
mjtNum tmp = mat[adr] * invD;
// update row i < k: L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k)
int i = colind[adr];
mju_addToScl(mat + rowadr[i], mat + rowadr_k, -tmp, rownnz[i]);
// update ith element of row k: L(k, i) /= L(k, k)
mat[adr] = tmp;
mju_addToScl(mat + rowadr[i], mat + start, -tmp, rownnz[i]);
}
// update row k: L(k, :) /= L(k, k)
mju_scl(mat + start, mat + start, invD, diag);
}
}
@@ -1807,11 +1803,11 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n,
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
const int* rownnz, const int* rowadr, const int* colind) {
// x <- L^-T x
for (int i=nv-1; i > 0; i--) {
// skip diagonal rows
if (diagnum[i]) {
if (rownnz[i] == 1) {
continue;
}
@@ -1862,8 +1858,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv,
// 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
if (rownnz[i] == 1) {
continue;
}
@@ -1895,7 +1890,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
mju_copy(x, y, n*m->nv);
}
mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
}
src/engine/engine_core_smooth.h
+2 -2
View File
@@ -63,7 +63,7 @@ MJAPI void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M,
// sparse L'*D*L factorizaton of inertia-like matrix
MJAPI void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
const int* rownnz, const int* rowadr, const int* colind);
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -75,7 +75,7 @@ MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n,
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// handle n vectors at once
MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
const int* rownnz, const int* rowadr, const int* colind);
// sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d
MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
src/engine/engine_forward.c
+4 -4
View File
@@ -873,14 +873,14 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
}
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind);
}
// solve
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
mju_copy(qacc, qfrc, m->nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
}
// advance state and time
@@ -1056,13 +1056,13 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_gather(d->qH, MhB, d->mapM2M, nC);
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, d->M_colind);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
} else {
mjERROR("integrator must be implicit or implicitfast");
src/engine/engine_solver.c
+1 -1
View File
@@ -1071,7 +1071,7 @@ static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) {
else {
mju_copy(ctx->Mgrad, ctx->grad, nv);
mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind);
}
}
test/benchmark/factorI_benchmark_test.cc
+1 -1
View File
@@ -59,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
} else {
mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
}
}
}
test/benchmark/inertia_benchmark_test.cc
+2 -2
View File
@@ -73,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
case SolveType::kCsr:
mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
}
}
}
test/benchmark/solveLD_benchmark_test.cc
+1 -1
View File
@@ -64,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv);
} else {
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
}
}
}
test/engine/engine_core_smooth_test.cc
+4 -4
View File
@@ -722,7 +722,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv; i++) {
@@ -757,7 +757,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
// expect vectors to match up to floating point precision
for (int i=0; i < nv*n; i++) {
@@ -795,7 +795,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
// expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
for (int i=0; i < n; i++) {
@@ -834,7 +834,7 @@ TEST_F(CoreSmoothTest, FactorIs) {
vector<mjtNum> qLDiagInv(nv, 0);
mj_factorI(qLD.data(), qLDiagInv.data(), nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
// expect outputs to match to floating point precision
EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));
test/engine/engine_derivative_test.cc
+2 -2
View File
@@ -436,7 +436,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
}
mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
// A = [dt*Ac; Ac]
mju_transpose(A, Ac, 2*nv, nv);
@@ -464,7 +464,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
d->moment_rowadr, d->moment_colind);
mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
d->M_rownnz, d->M_rowadr, d->M_colind);
mju_transpose(BcT, Bc, nu, nv);
mju_scl(B, BcT, dt*dt, nu*nv);
mju_scl(B+nu*nv, BcT, dt, nu*nv);