From 2dd518734fec79179efe33409054258b9f611c34 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 11 Oct 2024 05:08:27 -0700 Subject: [PATCH] Stop over-allocating memory in Newton solver. PiperOrigin-RevId: 684797464 Change-Id: I3389416cb69a8564b5119f6e3944b579893ea511 --- doc/changelog.rst | 6 +- model/humanoid/100_humanoids.xml | 50 ++++++++ src/engine/engine_print.c | 38 +++--- src/engine/engine_solver.c | 189 ++++++++++++++++++++--------- src/engine/engine_util_solve.c | 23 +--- src/engine/engine_util_solve.h | 7 +- test/sample/testspeed_test.sh | 3 +- test/xml/xml_native_writer_test.cc | 3 +- 8 files changed, 220 insertions(+), 99 deletions(-) create mode 100644 model/humanoid/100_humanoids.xml diff --git a/doc/changelog.rst b/doc/changelog.rst index 04cb26a5..6c575e38 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -8,9 +8,13 @@ Upcoming version (not yet released) General ^^^^^^^ +- The Newton solver no longer requires ``nv*nv`` memory allocation, allowing for much larger models. See e.g., + `100_humanoids.xml `__. + Two quadratic-memory allocations still remain to be fully sparsified: ``mjData.actuator_moment`` and the matrices used + by the PGS solver. - Removed the :at:`solid` and :at:`membrane` plugins and moved the associated computations into the engine. See `3D example model `__ and `2D example model - `__ for examples of flex objects + `__ for examples of flex objects that previously required these plugins. - Replaced the function ``mjs_setActivePlugins`` with :ref:`mjs_activatePlugin`. diff --git a/model/humanoid/100_humanoids.xml b/model/humanoid/100_humanoids.xml new file mode 100644 index 00000000..356a23d5 --- /dev/null +++ b/model/humanoid/100_humanoids.xml @@ -0,0 +1,50 @@ + + + + diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index b9edf0e4..f0384f93 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -756,7 +756,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, mjERROR("attempting to print mjData when stack is in use"); } - mjtNum *M; + mjtNum *M = NULL; mj_markStack(d); // check format string @@ -780,8 +780,10 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, return; } - // allocate full inertia - M = mj_stackAllocNum(d, m->nv*m->nv); + // allocate full inertia if it's small + if (m->nv <= 200) { + M = mj_stackAllocNum(d, m->nv*m->nv); + } #ifdef MEMORY_SANITIZER // If memory sanitizer is active, d->buffer will be marked as poisoned, even @@ -969,13 +971,15 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, printArray("ACTUATOR_MOMENT", m->nu, m->nv, d->actuator_moment, fp, float_format); printArray("CRB", m->nbody, 10, d->crb, fp, float_format); - // construct and print full M matrix - mj_fullM(m, M, d->qM); - printArray("QM", m->nv, m->nv, M, fp, float_format); + if (M) { + // construct and print full M matrix + mj_fullM(m, M, d->qM); + printArray("QM", m->nv, m->nv, M, fp, float_format); - // construct and print full LD matrix - mj_fullM(m, M, d->qLD); - printArray("QLD", m->nv, m->nv, M, fp, float_format); + // construct and print full LD matrix + mj_fullM(m, M, d->qLD); + printArray("QLD", m->nv, m->nv, M, fp, float_format); + } printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format); @@ -1064,14 +1068,16 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, } fprintf(fp, "\n\n"); - // print qDeriv - mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind); - printArray("QDERIV", m->nv, m->nv, M, fp, float_format); + if (M) { + // print qDeriv + mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind); + printArray("QDERIV", m->nv, m->nv, M, fp, float_format); - // print qLU - mju_sparse2dense(M, d->qLU, m->nv, m->nv, d->D_rownnz, d->D_rowadr, - d->D_colind); - printArray("QLU", m->nv, m->nv, M, fp, float_format); + // print qLU + mju_sparse2dense(M, d->qLU, m->nv, m->nv, d->D_rownnz, d->D_rowadr, + d->D_colind); + printArray("QLU", m->nv, m->nv, M, fp, float_format); + } // contact fprintf(fp, "CONTACT\n"); diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 99c02368..7f73642d 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -830,11 +830,29 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, // Hessian (Newton only) ctx->flg_Newton = flg_Newton; - if (flg_Newton && mj_isSparse(m)) { - d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int)); - if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr"); - d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int)); - if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz"); + if (flg_Newton) { + // sparse: allocate L_rowadr, L_rownnz + if (mj_isSparse(m)) { + d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int)); + if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr"); + d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int)); + if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz"); + + // zero nnzL, clear pointers (compute and allocate later in HessianDirect) + d->nnzL = 0; + d->L_colind = NULL; + d->L = NULL; + d->Lcone = NULL; + } + + // dense: allocate L + else if (d->nnzL != nv*nv) { + d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nv*nv, _Alignof(mjtNum)); + if (!d->L) mjERROR("failed to allocate L"); + + // set dense nnzL + d->nnzL = nv*nv; + } } } @@ -1362,6 +1380,16 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) { static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { int nv = m->nv, nefc = d->nefc; mjtNum local[36]; + + // allocate Lcone if required + if (!d->Lcone) { + d->Lcone = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nnzL, _Alignof(mjtNum)); + } + if (!d->Lcone) mjERROR("failed to allocate Lcone"); + + // start with Hcone = H + mju_copy(d->Lcone, d->L, d->nnzL); + mj_markStack(d); // storage for L'*J @@ -1369,9 +1397,6 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { mjtNum* LTJ_row = mj_stackAllocNum(d, nv); int* LTJ_ind = mj_stackAllocInt(d, nv); - // start with Hcone = H - mju_copy(d->Lcone, d->L, d->nnzL); - // add contributions for (int i=0; i < nefc; i++) { if (d->efc_state[i] == mjCNSTRSTATE_CONE) { @@ -1440,21 +1465,6 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) { // TODO: b/295296178 - add island support to Newton solver static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) { int nv = m->nv, nefc = d->nefc; - - // allocate Hessian on arena if not already allocated - if (!d->nnzL) { - int nnz = nv*nv; - if (mj_isSparse(m)) { - d->L_colind = mj_arenaAllocByte(d, sizeof(int) * nnz, _Alignof(int)); - if (!d->L_colind) mjERROR("failed to allocate L_colind"); - } - d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz, _Alignof(mjtNum)); - if (!d->L) mjERROR("failed to allocate L"); - d->Lcone = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz, _Alignof(mjtNum)); - if (!d->Lcone) mjERROR("failed to allocate Lcone"); - d->nnzL = nnz; - } - mj_markStack(d); // compute D corresponding to quad states @@ -1469,44 +1479,109 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) { // sparse if (mj_isSparse(m)) { - // fill-in reduced sparse inertia matrix C (no off-diagonals for simple dofs) - int nC = m->nC; - mjtNum* C = mj_stackAllocNum(d, nC); - for (int i=0; i < nC; i++) { + // copy values of reduced sparse inertia matrix C, get nnz + int nnz_C = m->nC; + mjtNum* C = mj_stackAllocNum(d, nnz_C); + for (int i=0; i < nnz_C; i++) { C[i] = d->qM[d->mapM2C[i]]; } - // compute H = J'*D*J + // allocate and initialize Hessian rowadr, rownnz; get nnz for J'*J + int* H_rowadr = mj_stackAllocInt(d, nv); + int* H_rownnz = mj_stackAllocInt(d, nv); + mju_sqrMatTDSparseInit(H_rownnz, H_rowadr, nv, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, + d); + int nnz_JTJ = H_rowadr[nv-1] + H_rownnz[nv-1]; - // TODO(b/266802572): remove uncompressed layout - mju_sqrMatTDUncompressedInit(d->L_rowadr, nv); - mju_sqrMatTDSparse(d->L, d->efc_J, d->efc_JT, D, nefc, nv, - d->L_rownnz, d->L_rowadr, d->L_colind, - d->efc_J_rownnz, d->efc_J_rowadr, - d->efc_J_colind, NULL, - d->efc_JT_rownnz, d->efc_JT_rowadr, - d->efc_JT_colind, d->efc_JT_rowsuper, d); - - // compute H = M + J'*D*J - mj_addMSparse(m, d, d->L, d->L_rownnz, d->L_rowadr, d->L_colind, - C, d->C_rownnz, d->C_rowadr, d->C_colind); - - // factorize H, uncompressed layout - int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL, - d->L_rownnz, d->L_rowadr, d->L_colind, d); - - // rank-defficient, SHOULD NOT OCCUR - if (rank != nv) { - mjERROR("rank-defficient Hessian"); + // shift H rowadr to make room for C + int shift = 0; + for (int r = 0; r < nv - 1; r++) { + shift += d->C_rownnz[r]; + H_rowadr[r + 1] += shift; } - // compress layout of H - mju_compressSparse(d->L, nv, nv, d->L_rownnz, d->L_rowadr, d->L_colind); + // allocate Hessian H, colind + int nnz_H = nnz_C + nnz_JTJ; + mjtNum* H = mj_stackAllocNum(d, nnz_H); + int* H_colind = mj_stackAllocInt(d, nnz_H); - // count nnz - d->nnzL = d->L_rowadr[nv-1] + d->L_rownnz[nv-1]; - if (d->nnzL > nv*nv) { // SHOULD NOT OCCUR - mjERROR("more nonzero values than elements in sparse direct-solver Hessian"); + // compute H = J'*D*J + mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, nefc, nv, + H_rownnz, H_rowadr, H_colind, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, + d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, + d); + + // add mass matrix; H = J'*D*J + C + mj_addMSparse(m, d, H, H_rownnz, H_rowadr, H_colind, + C, d->C_rownnz, d->C_rowadr, d->C_colind); + + // count row and total non-zeros of reverse-Cholesky factor L + int* parent = mj_stackAllocInt(d, nv); + int* flag = mj_stackAllocInt(d, nv); + int nnz_L = mju_cholFactorNNZ(d->L_rownnz, parent, flag, H_rownnz, H_rowadr, H_colind, nv); + + // allocate L_colind, L on arena if required + if (!d->nnzL) { + // nnzL is 0 but pointers are allocated; SHOULD NOT OCCUR + if (d->L_colind || d->L) { + mjERROR("nnzL is 0 but L_colind or L or Lcone are allocated"); + } + + // allocate on arena + d->L_colind = mj_arenaAllocByte(d, sizeof(int) * nnz_L, _Alignof(int)); + if (!d->L_colind) mjERROR("failed to allocate L_colind"); + d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz_L, _Alignof(mjtNum)); + if (!d->L) mjERROR("failed to allocate L"); + + // set nnzL + d->nnzL = nnz_L; + } else if (d->nnzL != nnz_L) { + // nnzL is nonzero but not equal to computed value; SHOULD NOT OCCUR + mjERROR("nnzL is nonzero but not equal to computed value"); + } + + // compute L row adresses: L_rowadr = cumsum(L_rownnz) + d->L_rowadr[0] = 0; + for (int r=1; r < nv; r++) { + d->L_rowadr[r] = d->L_rowadr[r-1] + d->L_rownnz[r-1]; + } + + // copy H lower-triangle into L + for (int r = 0; r < nv; r++) { + // count H non-zeros up to diagonal (inclusive) for row r + const int* colind = H_colind + H_rowadr[r]; + int rownnz = 1; + while (rownnz < nv && colind[rownnz - 1] < r) { + rownnz++; + } + + // last row element is not the diagonal; SHOULD NOT OCCUR + if (colind[rownnz - 1] != r) { + mjERROR("Newton solver Hessian has zero diagonal on row %d", r); + } + + // copy values and column indices + mju_copy(d->L + d->L_rowadr[r], H + H_rowadr[r], rownnz); + mju_copyInt(d->L_colind + d->L_rowadr[r], H_colind + H_rowadr[r], rownnz); + + // set L_rownnz + d->L_rownnz[r] = rownnz; + } + + // in-place sparse factorization L = chol(H) + int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL, d->L_rownnz, d->L_rowadr, d->L_colind, d); + + // rank-deficient; SHOULD NOT OCCUR + if (rank != nv) { + mjERROR("rank-deficient Hessian"); + } + + // pre-counted nnzL does not match post-factorization nnzL; SHOULD NOT OCCUR + if (d->nnzL != d->L_rowadr[nv-1] + d->L_rownnz[nv-1]) { + mjERROR("mismatch between pre-counted and post-factorization L nonzeros"); } } @@ -1518,9 +1593,6 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) { // factorize H mju_cholFactor(d->L, nv, mjMINVAL); - - // set nnz - d->nnzL = nv*nv; } mj_freeStack(d); @@ -1538,8 +1610,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) { // incremental update to Hessian // TODO: b/295296178 - add island support to Newton solver -static void HessianIncremental(const mjModel* m, mjData* d, - mjCGContext* ctx, const int* oldstate) { +static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, const int* oldstate) { int rank, nv = m->nv, nefc = d->nefc; mj_markStack(d); diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index f51928b7..dd6c7843 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -141,28 +141,15 @@ int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus) { //---------------------------- sparse Cholesky ----------------------------------------------------- // sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank' -// mat must have uncompressed layout; rownnz is modified to end at diagonal +// mat must be lower-triangular, have preallocated space for fill-in int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, - int* rownnz, int* rowadr, int* colind, + int* rownnz, const int* rowadr, int* colind, mjData* d) { int rank = n; mj_markStack(d); + mjtNum* buf = mj_stackAllocNum(d, n); int* buf_ind = mj_stackAllocInt(d, n); - mjtNum* sparse_buf = mj_stackAllocNum(d, n); - - // shrink rows so that rownnz ends at diagonal - for (int r=0; r < n; r++) { - // shrink - while (rownnz[r] > 0 && colind[rowadr[r]+rownnz[r]-1] > r) { - rownnz[r]--; - } - - // check - if (rownnz[r] == 0 || colind[rowadr[r]+rownnz[r]-1] != r) { - mjERROR("matrix must have non-zero diagonal"); - } - } // backpass over rows for (int r=n-1; r >= 0; r--) { @@ -191,7 +178,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, // 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], 1, -mat[adr+i], rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r], - sparse_buf, buf_ind); + buf, buf_ind); // assign new nnz to row c rownnz[c] = nnz_c; @@ -251,7 +238,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int // sparse reverse-order Cholesky rank-one update: L'*L +/- x*x'; return rank // x is sparse, change in sparsity pattern of mat is not allowed int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus, - int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind, + const int* rownnz, const int* rowadr, int* colind, int x_nnz, int* x_ind, mjData* d) { mj_markStack(d); int* buf_ind = mj_stackAllocInt(d, n); diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index a5cf8697..308e2d77 100644 --- a/src/engine/engine_util_solve.h +++ b/src/engine/engine_util_solve.h @@ -33,8 +33,9 @@ MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus); // sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank' -// mat must have uncompressed layout; rownnz is modified to end at diagonal -int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, int* rownnz, int* rowadr, int* colind, +// mat must be lower-triangular, have preallocated space for fill-in +int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, + int* rownnz, const int* rowadr, int* colind, mjData* d); // sparse reverse-order Cholesky solve @@ -44,7 +45,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int // sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank // x is sparse, change in sparsity pattern of mat is not allowed int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus, - int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind, + const int* rownnz, const int* rowadr, int* colind, int x_nnz, int* x_ind, mjData* d); // band-dense Cholesky decomposition diff --git a/test/sample/testspeed_test.sh b/test/sample/testspeed_test.sh index fdfd9d95..c5a76445 100755 --- a/test/sample/testspeed_test.sh +++ b/test/sample/testspeed_test.sh @@ -28,7 +28,8 @@ test_model() { local iterations=10 # for particularly slow models, only run 2 steps under ASAN, or skip. if [[ ${TESTSPEED_ASAN:-0} != 0 ]]; then - if [[ "$model" == */composite/particle.xml || + if [[ "$model" == */humanoid/100_humanoids.xml || + "$model" == */composite/particle.xml || "$model" == */replicate/bunnies.xml || "$model" == */replicate/leaves.xml || "$model" == */replicate/particle.xml || diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index 9f5a8622..ce1d93e2 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -1312,7 +1312,8 @@ TEST_F(XMLWriterTest, WriteReadCompare) { std::string xml = p.path().string(); // if file is meant to fail, skip it - if (absl::StrContains(p.path().string(), "malformed_") || + if (absl::StrContains(p.path().string(), "100_humanoids") || + absl::StrContains(p.path().string(), "malformed_") || absl::StrContains(p.path().string(), "touch_grid") || absl::StrContains(p.path().string(), "gmsh_") || absl::StrContains(p.path().string(), "shark_") ||