From d8955a5767bd841760f096f0834405eefcd826e6 Mon Sep 17 00:00:00 2001 From: fpacaud Date: Fri, 27 Aug 2021 14:50:36 +0200 Subject: [PATCH] Avoid unecessary allocations by forcing specialization --- src/MadNLP.jl | 4 +-- src/interiorpointsolver.jl | 56 ++++++++++++++++++++++++++------------ src/kktsystem.jl | 4 +-- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/src/MadNLP.jl b/src/MadNLP.jl index b1dbe3a33..384e695bc 100644 --- a/src/MadNLP.jl +++ b/src/MadNLP.jl @@ -32,8 +32,8 @@ include("options.jl") include("nonlinearprogram.jl") include("matrixtools.jl") include(joinpath("LinearSolvers","linearsolvers.jl")) -include(joinpath("kktsystem.jl")) -include(joinpath("interiorpointsolver.jl")) +include("kktsystem.jl") +include("interiorpointsolver.jl") include(joinpath("Interfaces","interfaces.jl")) # Initialize diff --git a/src/interiorpointsolver.jl b/src/interiorpointsolver.jl index 6a9a20efd..b27a3e9dc 100644 --- a/src/interiorpointsolver.jl +++ b/src/interiorpointsolver.jl @@ -255,7 +255,7 @@ function eval_f_wrapper(ipp::Solver, x::Vector{Float64}) nlp = ipp.nlp cnt = ipp.cnt @trace(ipp.logger,"Evaluating objective.") - cnt.eval_function_time += @elapsed obj_val = nlp.obj(view(x,1:nlp.n)) + cnt.eval_function_time += @elapsed obj_val = nlp.obj(view(x,1:nlp.n))::Float64 cnt.obj_cnt+=1 cnt.obj_cnt==1 && (is_valid(obj_val) || throw(InvalidNumberException())) return obj_val*ipp.obj_scale[] @@ -1093,17 +1093,26 @@ end # KKT system updates ------------------------------------------------------- # Set diagonal +# TODO: Temporary solution --> create auxiliary functions to force specialization +function _set_aug_diagonal_reduced!(pr_diag, du_diag, x, xl, xu, zl, zu) + pr_diag .= zl./(x.-xl) .+ zu./(xu.-x) + fill!(du_diag, 0.0) +end +function _set_aug_diagonal_unreduced!(pr_diag, du_diag, l_lower, u_lower, l_diag, u_diag, zl_r, zu_r, xl_r, xu_r, x_lr, x_ur) + pr_diag .= 0.0 + du_diag .= 0.0 + l_lower .= .-sqrt.(zl_r) + u_lower .= .-sqrt.(zu_r) + l_diag .= xl_r .- x_lr + u_diag .= x_ur .- xu_r +end function set_aug_diagonal!(kkt::AbstractKKTSystem, ips::Solver) - copyto!(kkt.pr_diag, ips.zl./(ips.x.-ips.xl) .+ ips.zu./(ips.xu.-ips.x)) - fill!(kkt.du_diag, 0.0) + _set_aug_diagonal_reduced!(kkt.pr_diag, kkt.du_diag, ips.x, ips.xl, ips.xu, ips.zl, ips.zu) end function set_aug_diagonal!(kkt::SparseUnreducedKKTSystem, ips::Solver) - kkt.pr_diag .= 0.0 - kkt.du_diag .= 0.0 - kkt.l_lower .= .-sqrt.(ips.zl_r) - kkt.u_lower .= .-sqrt.(ips.zu_r) - kkt.l_diag .= ips.xl_r .- ips.x_lr - kkt.u_diag .= ips.x_ur .- ips.xu_r + _set_aug_diagonal_unreduced!(kkt.pr_diag, kkt.du_diag, kkt.l_lower, kkt.u_lower, kkt.l_diag, kkt.u_diag, + ips.zl_r, ips.zu_r, ips.xl_r, ips.xu_r, ips.x_lr, ips.x_ur, + ) end # Robust restoration @@ -1121,15 +1130,24 @@ function set_aug_RR!(kkt::SparseUnreducedKKTSystem, ips::Solver, RR::RobustResto end # Set RHS +# TODO: Temporary solution --> create auxiliary functions to force specialization +function _set_aug_rhs_reduced!(px, pl, x, xl, xu, c, f, jacl, mu) + px.=.-f.+mu./(x.-xl).-mu./(xu.-x).-jacl + pl.=.-c +end +function _set_aug_rhs_unreduced!(px, pl, pzl, pzu, c, f, zl, zu, jacl, xl_r, xu_r, x_lr, x_ur, l_lower, u_lower, mu) + px.=.-f.+zl.-zu.-jacl + pl.=.-c + pzl.=(xl_r-x_lr).*l_lower .+ mu./l_lower + pzu.=(xu_r-x_ur).*u_lower .- mu./u_lower +end function set_aug_rhs!(ips::Solver, kkt::AbstractKKTSystem, c) - ips.px.=.-ips.f.+ips.mu./(ips.x.-ips.xl).-ips.mu./(ips.xu.-ips.x).-ips.jacl - ips.pl.=.-c + _set_aug_rhs_reduced!(ips.px, ips.pl, ips.x, ips.xl, ips.xu, ips.c, ips.f, ips.jacl, ips.mu) end function set_aug_rhs!(ips::Solver, kkt::SparseUnreducedKKTSystem, c) - ips.px.=.-ips.f.+ips.zl.-ips.zu.-ips.jacl - ips.pl.=.-c - ips.pzl.=(ips.xl_r-ips.x_lr).*kkt.l_lower.+ips.mu./kkt.l_lower - ips.pzu.=(ips.xu_r-ips.x_ur).*kkt.u_lower.-ips.mu./kkt.u_lower + _set_aug_rhs_unreduced!(ips.px, ips.pl, ips.pzl, ips.pzu, c, ips.f, ips.zl, ips.zu, + ips.jacl, ips.xl_r, ips.xu_r, ips.x_lr, ips.x_ur, + kkt.l_lower, kkt.u_lower, ips.mu) end # Set RHS RR @@ -1149,10 +1167,14 @@ function set_aug_rhs_RR!( end # Finish +function _finish_aug_solve_reduced!(dzl, dzu, zl_r, zu_r, dx_lr, dx_ur, x_lr, x_ur, xl_r, xu_r, mu) + dzl.= (mu.-zl_r.*dx_lr)./(x_lr.-xl_r).-zl_r + dzu.= (mu.+zu_r.*dx_ur)./(xu_r.-x_ur).-zu_r +end function finish_aug_solve!(ips::Solver, kkt::AbstractKKTSystem, mu) - ips.dzl.= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r - ips.dzu.= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r + _finish_aug_solve_reduced!(ips.dzl, ips.dzu, ips.zl_r, ips.zu_r, ips.dx_lr, ips.dx_ur, ips.x_lr, ips.x_ur, ips.xl_r, ips.xu_r, ips.mu) end +# Note: No need to force specialization here. function finish_aug_solve!(ips::Solver, kkt::SparseUnreducedKKTSystem, mu) ips.dzl.*=.-kkt.l_lower ips.dzu.*=kkt.u_lower diff --git a/src/kktsystem.jl b/src/kktsystem.jl index a8b3c2697..7559fa9b3 100644 --- a/src/kktsystem.jl +++ b/src/kktsystem.jl @@ -84,8 +84,8 @@ function build_kkt!(kkt::AbstractSparseKKTSystem{T, MT}) where {T, MT<:SparseMat treat_fixed_variable!(kkt) end -function set_jacobian_scaling!(kkt::AbstractSparseKKTSystem, constraint_scaling::AbstractVector) - nnzJ = length(kkt.jac) +function set_jacobian_scaling!(kkt::AbstractSparseKKTSystem{T, MT}, constraint_scaling::AbstractVector) where {T, MT} + nnzJ = length(kkt.jac)::Int for i in 1:nnzJ index = kkt.jac_raw.I[i] kkt.jacobian_scaling[i] = constraint_scaling[index]