diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 4fb1224d..b03ad913 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -813,6 +813,8 @@ typedef struct { // Newton arrays, known-size (PrimalAllocate) mjtNum* D; // constraint inertia (nefc x 1) + mjtNum* cholupd; // scratch for rank-1 Cholesky updates (nv x 1) + mjtNum* LTJ; // L'*J for cone Cholesky updates (6 x nv) int* H_rowadr; // Hessian row addresses (nv x 1) int* H_rownnz; // Hessian row nonzeros (nv x 1) int* HT_rownnz; // Hessian transpose row nonzeros (nv x 1) @@ -988,6 +990,10 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) { // Newton only, known-size arrays if (flg_Newton) { ctx->D = mjSTACKALLOC(d, nefc, mjtNum); + ctx->cholupd = mjSTACKALLOC(d, nv, mjtNum); + if (ctx->is_elliptic) { + ctx->LTJ = mjSTACKALLOC(d, 6*nv, mjtNum); + } // sparse Newton only if (ctx->is_sparse) { @@ -1678,10 +1684,7 @@ static void HessianCone(mjData* d, mjPrimalContext* ctx) { // start with Hcone = H mju_copy(ctx->Lcone, ctx->L, ctx->nL); - mj_markStack(d); - - // storage for L'*J - mjtNum* LTJ = mjSTACKALLOC(d, 6*nv, mjtNum); + mjtNum* LTJ = ctx->LTJ; // add contributions for (int i=0; i < nefc; i++) { @@ -1737,18 +1740,13 @@ static void HessianCone(mjData* d, mjPrimalContext* ctx) { i += (dim-1); } } - - mj_freeStack(d); } // incremental update to Hessian factor due to changes in efc_state static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldstate) { int rank, nv = ctx->nv, nefc = ctx->nefc; - mj_markStack(d); - - // local space - mjtNum* vec = mjSTACKALLOC(d, nv, mjtNum); + mjtNum* cholupd = ctx->cholupd; // clear update counter ctx->nupdate = 0; @@ -1769,27 +1767,26 @@ static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldst // perform update if flagged if (flag_update != -1) { - // update with vec = J(i,:)*sqrt(D[i])) + // update with cholupd = J(i,:)*sqrt(D[i])) if (ctx->is_sparse) { // get nnz and adr of row i const int nnz = ctx->J_rownnz[i], adr = ctx->J_rowadr[i]; - // scale vec - mju_scl(vec, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz); + // scale cholupd + mju_scl(cholupd, ctx->J+adr, mju_sqrt(ctx->efc_D[i]), nnz); // sparse update or downdate - rank = mju_cholUpdateSparse(ctx->L, vec, nv, flag_update, + rank = mju_cholUpdateSparse(ctx->L, cholupd, nv, flag_update, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind, nnz, ctx->J_colind+adr, d); } else { - mju_scl(vec, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv); - rank = mju_cholUpdate(ctx->L, vec, nv, flag_update); + mju_scl(cholupd, ctx->J+i*nv, mju_sqrt(ctx->efc_D[i]), nv); + rank = mju_cholUpdate(ctx->L, cholupd, nv, flag_update); } ctx->nupdate++; // recompute H directly if accuracy lost if (rank < nv) { - mj_freeStack(d); FactorizeHessian(d, ctx, /*flg_recompute=*/1); // nothing else to do @@ -1802,8 +1799,6 @@ static void HessianIncremental(mjData* d, mjPrimalContext* ctx, const int* oldst if (ctx->ncone) { HessianCone(d, ctx); } - - mj_freeStack(d); }