diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl index 03ea1d8..fa7b158 100644 --- a/src/forward_over_reverse.jl +++ b/src/forward_over_reverse.jl @@ -38,7 +38,7 @@ function _eval_hessian( scale::Float64, nzcount::Int, )::Int - if ex.linearity == LINEAR + if ex.expr.linearity == LINEAR @assert length(ex.hess_I) == 0 return 0 end @@ -128,13 +128,9 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T} _reinterpret_unsafe(T, d.subexpression_forward_values_ϵ) for i in ex.dependent_subexpressions subexpr = d.subexpressions[i] - subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( - d, - subexpr, - _reinterpret_unsafe(T, subexpr.partials_storage_ϵ), - ) + subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(d, subexpr, T) end - _forward_eval_ϵ(d, ex, _reinterpret_unsafe(T, d.partials_storage_ϵ)) + _forward_eval_ϵ(d, ex.expr, T) # do a reverse pass subexpr_reverse_values_ϵ = _reinterpret_unsafe(T, d.subexpression_reverse_values_ϵ) @@ -144,9 +140,8 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T} end _reverse_eval_ϵ( output_ϵ, - ex, + ex.expr, _reinterpret_unsafe(T, d.storage_ϵ), - _reinterpret_unsafe(T, d.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, 1.0, @@ -159,7 +154,6 @@ function _hessian_slice_inner(d, ex, ::Type{T}) where {T} output_ϵ, subexpr, _reinterpret_unsafe(T, d.storage_ϵ), - _reinterpret_unsafe(T, subexpr.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, d.subexpression_reverse_values[j], @@ -173,8 +167,8 @@ end _forward_eval_ϵ( d::NLPEvaluator, ex::Union{_FunctionStorage,_SubexpressionStorage}, - partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, - ) where {N,T} + ::Type{P}, + ) where {N,T,P<:ForwardDiff.Partials{N,T}} Evaluate the directional derivatives of the expression tree in `ex`. @@ -186,10 +180,11 @@ This assumes that `_reverse_model(d, x)` has already been called. """ function _forward_eval_ϵ( d::NLPEvaluator, - ex::Union{_FunctionStorage,_SubexpressionStorage}, - partials_storage_ϵ::AbstractVector{P}, + ex::_SubexpressionStorage, + ::Type{P}, ) where {N,T,P<:ForwardDiff.Partials{N,T}} storage_ϵ = _reinterpret_unsafe(P, d.storage_ϵ) + partials_storage_ϵ = _reinterpret_unsafe(P, ex.partials_storage_ϵ) x_values_ϵ = _reinterpret_unsafe(P, d.input_ϵ) subexpression_values_ϵ = _reinterpret_unsafe(P, d.subexpression_forward_values_ϵ) @@ -370,14 +365,15 @@ end # to compute hessian-vector products. function _reverse_eval_ϵ( output_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, - ex::Union{_FunctionStorage,_SubexpressionStorage}, + ex::_SubexpressionStorage, reverse_storage_ϵ, - partials_storage_ϵ, subexpression_output, subexpression_output_ϵ, scale::T, scale_ϵ::ForwardDiff.Partials{N,T}, ) where {N,T} + partials_storage_ϵ = + _reinterpret_unsafe(ForwardDiff.Partials{N,T}, ex.partials_storage_ϵ) @assert length(reverse_storage_ϵ) >= length(ex.nodes) @assert length(partials_storage_ϵ) >= length(ex.nodes) if ex.nodes[1].type == Nonlinear.NODE_VARIABLE diff --git a/src/mathoptinterface_api.jl b/src/mathoptinterface_api.jl index 1766c25..21c6804 100644 --- a/src/mathoptinterface_api.jl +++ b/src/mathoptinterface_api.jl @@ -65,11 +65,12 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) for k in d.subexpression_order # Only load expressions which actually are used d.subexpression_forward_values[k] = NaN - subex = _SubexpressionStorage( - d.data.expressions[k], - d.subexpression_linearity, + expr = d.data.expressions[k] + subex, _ = _subexpression_and_linearity( + expr, moi_index_to_consecutive_index, - d.want_hess, + Float64[], + d, ) d.subexpressions[k] = subex d.subexpression_linearity[k] = subex.linearity @@ -101,43 +102,54 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) end end max_chunk = 1 + shared_partials_storage_ϵ = Float64[] if d.data.objective !== nothing + expr = something(d.data.objective) + subexpr, linearity = _subexpression_and_linearity( + expr, + moi_index_to_consecutive_index, + shared_partials_storage_ϵ, + d, + ) objective = _FunctionStorage( - main_expressions[1], - something(d.data.objective).values, + subexpr, N, coloring_storage, d.want_hess, d.subexpressions, individual_order[1], - d.subexpression_linearity, subexpression_edgelist, subexpression_variables, - moi_index_to_consecutive_index, + linearity, ) - max_expr_length = max(max_expr_length, length(objective.nodes)) + max_expr_length = max(max_expr_length, length(expr.nodes)) max_chunk = max(max_chunk, size(objective.seed_matrix, 2)) d.objective = objective end for (k, (_, constraint)) in enumerate(d.data.constraints) idx = d.data.objective !== nothing ? k + 1 : k + expr = constraint.expression + subexpr, linearity = _subexpression_and_linearity( + expr, + moi_index_to_consecutive_index, + shared_partials_storage_ϵ, + d, + ) push!( d.constraints, _FunctionStorage( - main_expressions[idx], - constraint.expression.values, + subexpr, N, coloring_storage, d.want_hess, d.subexpressions, individual_order[idx], - d.subexpression_linearity, subexpression_edgelist, subexpression_variables, - moi_index_to_consecutive_index, + linearity, ), ) - max_expr_length = max(max_expr_length, length(d.constraints[end].nodes)) + max_expr_length = max(max_expr_length, length(expr.nodes)) max_chunk = max(max_chunk, size(d.constraints[end].seed_matrix, 2)) end max_chunk = min(max_chunk, MAX_CHUNK) @@ -146,7 +158,8 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) d.input_ϵ = zeros(max_chunk * N) d.output_ϵ = zeros(max_chunk * N) # - d.partials_storage_ϵ = zeros(max_chunk * max_expr_length) + resize!(shared_partials_storage_ϵ, max_chunk * max_expr_length) + fill!(shared_partials_storage_ϵ, 0.0) d.storage_ϵ = zeros(max_chunk * max_expr_with_sub_length) # len = max_chunk * length(d.subexpressions) @@ -178,7 +191,7 @@ function MOI.eval_objective(d::NLPEvaluator, x) error("No nonlinear objective.") end _reverse_mode(d, x) - return something(d.objective).forward_storage[1] + return something(d.objective).expr.forward_storage[1] end function MOI.eval_objective_gradient(d::NLPEvaluator, g, x) @@ -194,7 +207,7 @@ end function MOI.eval_constraint(d::NLPEvaluator, g, x) _reverse_mode(d, x) for i in 1:length(d.constraints) - g[i] = d.constraints[i].forward_storage[1] + g[i] = d.constraints[i].expr.forward_storage[1] end return end @@ -345,11 +358,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) subexpr_forward_values_ϵ = reinterpret(T, d.subexpression_forward_values_ϵ) for i in d.subexpression_order subexpr = d.subexpressions[i] - subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( - d, - subexpr, - reinterpret(T, subexpr.partials_storage_ϵ), - ) + subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(d, subexpr, T) end # we only need to do one reverse pass through the subexpressions as well subexpr_reverse_values_ϵ = reinterpret(T, d.subexpression_reverse_values_ϵ) @@ -358,16 +367,11 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) fill!(d.storage_ϵ, 0.0) fill!(output_ϵ, zero(T)) if d.objective !== nothing - _forward_eval_ϵ( - d, - something(d.objective), - reinterpret(T, d.partials_storage_ϵ), - ) + _forward_eval_ϵ(d, something(d.objective).expr, T) _reverse_eval_ϵ( output_ϵ, - something(d.objective), - reinterpret(T, d.storage_ϵ), - reinterpret(T, d.partials_storage_ϵ), + something(d.objective).expr, + _reinterpret_unsafe(T, d.storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, σ, @@ -375,12 +379,11 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) ) end for (i, con) in enumerate(d.constraints) - _forward_eval_ϵ(d, con, reinterpret(T, d.partials_storage_ϵ)) + _forward_eval_ϵ(d, con.expr, T) _reverse_eval_ϵ( output_ϵ, - con, + con.expr, reinterpret(T, d.storage_ϵ), - reinterpret(T, d.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, μ[i], @@ -394,7 +397,6 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) output_ϵ, subexpr, reinterpret(T, d.storage_ϵ), - reinterpret(T, subexpr.partials_storage_ϵ), d.subexpression_reverse_values, subexpr_reverse_values_ϵ, d.subexpression_reverse_values[j], diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl index a9c9187..84cc358 100644 --- a/src/reverse_mode.jl +++ b/src/reverse_mode.jl @@ -39,20 +39,20 @@ function _reverse_mode(d::NLPEvaluator, x) _forward_eval(d.subexpressions[k], d, x) end if d.objective !== nothing - _forward_eval(d.objective::_FunctionStorage, d, x) + _forward_eval(something(d.objective).expr, d, x) end for con in d.constraints - _forward_eval(con, d, x) + _forward_eval(con.expr, d, x) end # Phase II for k in d.subexpression_order _reverse_eval(d.subexpressions[k]) end if d.objective !== nothing - _reverse_eval(d.objective::_FunctionStorage) + _reverse_eval(something(d.objective).expr) end for con in d.constraints - _reverse_eval(con) + _reverse_eval(con.expr) end # If a JuMP model uses the legacy nonlinear interface, then JuMP constructs # a NLPEvaluator at the start of a call to `JuMP.optimize!` and it passes in @@ -81,7 +81,7 @@ end """ _forward_eval( - f::Union{_FunctionStorage,_SubexpressionStorage}, + f::_SubexpressionStorage, d::NLPEvaluator, x::AbstractVector{T}, ) where {T} @@ -98,10 +98,7 @@ Forward-mode evaluation of an expression tree given in `f`. associate storage with each edge of the DAG. """ function _forward_eval( - # !!! warning - # This Union depends upon _FunctionStorage and _SubexpressionStorage - # having similarly named fields. - f::Union{_FunctionStorage,_SubexpressionStorage}, + f::_SubexpressionStorage, d::NLPEvaluator, x::AbstractVector{T}, )::T where {T} @@ -290,19 +287,14 @@ function _forward_eval( end """ - _reverse_eval(f::Union{_FunctionStorage,_SubexpressionStorage}) + _reverse_eval(f::_SubexpressionStorage) Reverse-mode evaluation of an expression tree given in `f`. * This function assumes `f.partials_storage` is already updated. * This function assumes that `f.reverse_storage` has been initialized with 0.0. """ -function _reverse_eval( - # !!! warning - # This Union depends upon _FunctionStorage and _SubexpressionStorage - # having similarly named fields. - f::Union{_FunctionStorage,_SubexpressionStorage}, -) +function _reverse_eval(f::_SubexpressionStorage) @assert length(f.reverse_storage) >= length(f.nodes) @assert length(f.partials_storage) >= length(f.nodes) # f.nodes is already in order such that parents always appear before @@ -361,9 +353,15 @@ end function _extract_reverse_pass_inner( output::AbstractVector{T}, - # !!! warning - # This Union depends upon _FunctionStorage and _SubexpressionStorage - # having similarly named fields. + f::_FunctionStorage, + subexpressions::AbstractVector{T}, + scale::T, +) where {T} + return _extract_reverse_pass_inner(output, f.expr, subexpressions, scale) +end + +function _extract_reverse_pass_inner( + output::AbstractVector{T}, f::Union{_FunctionStorage,_SubexpressionStorage}, subexpressions::AbstractVector{T}, scale::T, diff --git a/src/types.jl b/src/types.jl index fc599ac..5dad1b8 100644 --- a/src/types.jl +++ b/src/types.jl @@ -15,68 +15,75 @@ struct _SubexpressionStorage linearity::Linearity function _SubexpressionStorage( - expr::Nonlinear.Expression, - subexpression_linearity, - moi_index_to_consecutive_index, - want_hess::Bool, + nodes::Vector{Nonlinear.Node}, + adj::SparseArrays.SparseMatrixCSC{Bool,Int}, + const_values::Vector{Float64}, + partials_storage_ϵ::Vector{Float64}, + linearity::Linearity, ) - nodes = - _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index) - adj = Nonlinear.adjacency_matrix(nodes) N = length(nodes) - linearity = if want_hess - _classify_linearity(nodes, adj, subexpression_linearity)[1] - else - NONLINEAR - end return new( nodes, adj, - expr.values, + const_values, zeros(N), # forward_storage, zeros(N), # partials_storage, zeros(N), # reverse_storage, - Float64[], + partials_storage_ϵ, linearity, ) end end +# We don't need to store the full vector of `linearity` but we return +# it because it is needed in `compute_hessian_sparsity`. +function _subexpression_and_linearity( + expr::Nonlinear.Expression, + moi_index_to_consecutive_index, + partials_storage_ϵ::Vector{Float64}, + d, +) + nodes = _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index) + adj = Nonlinear.adjacency_matrix(nodes) + linearity = if d.want_hess + _classify_linearity(nodes, adj, d.subexpression_linearity) + else + [NONLINEAR] + end + return _SubexpressionStorage( + nodes, + adj, + expr.values, + partials_storage_ϵ, + linearity[1], + ), + linearity +end + struct _FunctionStorage - nodes::Vector{Nonlinear.Node} - adj::SparseArrays.SparseMatrixCSC{Bool,Int} - const_values::Vector{Float64} - forward_storage::Vector{Float64} - partials_storage::Vector{Float64} - reverse_storage::Vector{Float64} + expr::_SubexpressionStorage grad_sparsity::Vector{Int} # Nonzero pattern of Hessian matrix hess_I::Vector{Int} hess_J::Vector{Int} rinfo::Coloring.RecoveryInfo # coloring info for hessians seed_matrix::Matrix{Float64} - linearity::Linearity # subexpressions which this function depends on, ordered for forward pass. dependent_subexpressions::Vector{Int} function _FunctionStorage( - nodes::Vector{Nonlinear.Node}, - const_values, + expr::_SubexpressionStorage, num_variables, coloring_storage::Coloring.IndexedSet, want_hess::Bool, subexpressions::Vector{_SubexpressionStorage}, dependent_subexpressions, - subexpression_linearity, subexpression_edgelist, subexpression_variables, - moi_index_to_consecutive_index, + linearity::Vector{Linearity}, ) - nodes = _replace_moi_variables(nodes, moi_index_to_consecutive_index) - adj = Nonlinear.adjacency_matrix(nodes) - N = length(nodes) empty!(coloring_storage) - _compute_gradient_sparsity!(coloring_storage, nodes) + _compute_gradient_sparsity!(coloring_storage, expr.nodes) for k in dependent_subexpressions _compute_gradient_sparsity!( coloring_storage, @@ -86,10 +93,9 @@ struct _FunctionStorage grad_sparsity = sort!(collect(coloring_storage)) empty!(coloring_storage) if want_hess - linearity = _classify_linearity(nodes, adj, subexpression_linearity) edgelist = _compute_hessian_sparsity( - nodes, - adj, + expr.nodes, + expr.adj, linearity, coloring_storage, subexpression_edgelist, @@ -102,34 +108,22 @@ struct _FunctionStorage ) seed_matrix = Coloring.seed_matrix(rinfo) return new( - nodes, - adj, - const_values, - zeros(N), # forward_storage, - zeros(N), # partials_storage, - zeros(N), # reverse_storage, + expr, grad_sparsity, hess_I, hess_J, rinfo, seed_matrix, - linearity[1], dependent_subexpressions, ) else return new( - nodes, - adj, - const_values, - zeros(N), # forward_storage, - zeros(N), # partials_storage, - zeros(N), # reverse_storage, + expr, grad_sparsity, Int[], Int[], Coloring.RecoveryInfo(), Array{Float64}(undef, 0, 0), - NONLINEAR, dependent_subexpressions, ) end @@ -176,7 +170,6 @@ mutable struct NLPEvaluator <: MOI.AbstractNLPEvaluator # so the length should be multiplied by the maximum number of epsilon components disable_2ndorder::Bool # don't offer Hess or HessVec want_hess::Bool - partials_storage_ϵ::Vector{Float64} # (longest expression excluding subexpressions) storage_ϵ::Vector{Float64} # (longest expression including subexpressions) input_ϵ::Vector{Float64} # (number of variables) output_ϵ::Vector{Float64} # (number of variables)