From 393a2609182cf16171016b33f3617e999c74da5f Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 28 Mar 2022 14:36:49 +1300 Subject: [PATCH] Refactor into AbstractBackends --- README.md | 6 - src/nonlinear_oracle.jl | 241 ++++++++++++++++++++++++++-------------- test/runtests.jl | 23 +++- 3 files changed, 172 insertions(+), 98 deletions(-) diff --git a/README.md b/README.md index 2a9be80..51abe26 100644 --- a/README.md +++ b/README.md @@ -26,12 +26,6 @@ constraints. So something like is great because we compute the derivative as `cos(x[i])` and then we can fill in the 10_000 derivative functions without having to do other calculus. -There are a couple of options in the evaluation functions, each with their own -performance tradeoffs. If we want to consider parallel evaluation, then we -need to use the offsets to allow out-of-order execution. I've threaded on the -hash of the symbolic constraints because constraints based on the same symbolic -template share a temporary storage cache in the `_SymbolicsFunction`. - ## Installation Install SymbolicAD as follows: diff --git a/src/nonlinear_oracle.jl b/src/nonlinear_oracle.jl index 56ebaf1..a693ffe 100644 --- a/src/nonlinear_oracle.jl +++ b/src/nonlinear_oracle.jl @@ -133,22 +133,39 @@ function _SymbolicsFunction(f::_Function, features::Vector{Symbol}) return _SymbolicsFunction(f, g!, h!, g_cache, h_cache, ∇²f_structure) end -struct _ConstraintOffset - index::Int - ∇f_offset::Int - ∇²f_offset::Int -end -mutable struct _NonlinearOracle <: MOI.AbstractNLPEvaluator - use_threads::Bool +""" + AbstractNonlinearOracleBackend + +An abstract type to help dispatch on different implementations of the function +evaluations. +""" +abstract type AbstractNonlinearOracleBackend end + +""" + DefaultBackend() <: AbstractNonlinearOracleBackend + +A simple implementation that loops over the different constraints. +""" +struct DefaultBackend <: AbstractNonlinearOracleBackend end + +mutable struct _NonlinearOracle{B} <: MOI.AbstractNLPEvaluator + backend::B objective::Union{Nothing,_Function} constraints::Vector{_Function} symbolic_functions::Dict{UInt,_SymbolicsFunction} - offsets::Vector{Vector{_ConstraintOffset}} eval_constraint_timer::Float64 eval_constraint_jacobian_timer::Float64 eval_hessian_lagrangian_timer::Float64 end +function MOI.initialize( + ::_NonlinearOracle, + ::AbstractNonlinearOracleBackend, + ::Vector{Symbol}, +) + return +end + function MOI.initialize(oracle::_NonlinearOracle, features::Vector{Symbol}) oracle.eval_constraint_timer = 0.0 oracle.eval_constraint_jacobian_timer = 0.0 @@ -158,34 +175,14 @@ function MOI.initialize(oracle::_NonlinearOracle, features::Vector{Symbol}) f = oracle.constraints[findfirst(isequal(h), hashes)] oracle.symbolic_functions[h] = _SymbolicsFunction(f, features) end - obj_hess_offset = 0 if oracle.objective !== nothing f = oracle.objective h = f.expr.hash if !haskey(oracle.symbolic_functions, h) oracle.symbolic_functions[h] = _SymbolicsFunction(f, features) - obj_hess_offset = length(oracle.symbolic_functions[h].H) - end - end - if oracle.use_threads - offsets = Dict{UInt,Vector{_ConstraintOffset}}( - h => _ConstraintOffset[] for h in keys(oracle.symbolic_functions) - ) - index, ∇f_offset, ∇²f_offset = 1, 1, obj_hess_offset + 1 - for f in oracle.constraints - push!( - offsets[f.expr.hash], - _ConstraintOffset(index, ∇f_offset, ∇²f_offset), - ) - index += 1 - symbolic_function = oracle.symbolic_functions[f.expr.hash] - ∇f_offset += length(symbolic_function.g) - ∇²f_offset += length(symbolic_function.H) - end - for v in values(offsets) - push!(oracle.offsets, v) end end + MOI.initialize(oracle, oracle.backend, features) return end @@ -257,17 +254,8 @@ function MOI.eval_constraint( x::AbstractVector{Float64}, ) start = time() - if oracle.use_threads - Threads.@threads for offsets in oracle.offsets - for c in offsets - func = oracle.constraints[c.index] - @inbounds g[c.index] = _eval_function(oracle, func, x) - end - end - else - for row in 1:length(g) - g[row] = _eval_function(oracle, oracle.constraints[row], x) - end + for row in 1:length(g) + g[row] = _eval_function(oracle, oracle.constraints[row], x) end oracle.eval_constraint_timer += time() - start return @@ -290,24 +278,12 @@ function MOI.eval_constraint_jacobian( x::AbstractVector{Float64}, ) start = time() - if oracle.use_threads - Threads.@threads for offset in oracle.offsets - for c in offset - func = oracle.constraints[c.index] - g = _eval_gradient(oracle, func, x) - for i in 1:length(g) - @inbounds J[c.∇f_offset+i-1] = g[i] - end - end - end - else - k = 1 - for i in 1:length(oracle.constraints) - g = _eval_gradient(oracle, oracle.constraints[i], x) - for gi in g - J[k] = gi - k += 1 - end + k = 1 + for i in 1:length(oracle.constraints) + g = _eval_gradient(oracle, oracle.constraints[i], x) + for gi in g + J[k] = gi + k += 1 end end oracle.eval_constraint_jacobian_timer += time() - start @@ -351,42 +327,30 @@ function MOI.eval_hessian_lagrangian( end hessian_offset += length(h) end - if oracle.use_threads - Threads.@threads for offset in oracle.offsets - for c in offset - func = oracle.constraints[c.index] - h = _eval_hessian(oracle, func, x) - for i in 1:length(h) - @inbounds H[c.∇²f_offset+i-1] = μ[c.index] * h[i] - end - end - end - else - k = hessian_offset + 1 - for i in 1:length(oracle.constraints) - h = _eval_hessian(oracle, oracle.constraints[i], x) - for nzval in h - H[k] = μ[i] * nzval - k += 1 - end + k = hessian_offset + 1 + for i in 1:length(oracle.constraints) + h = _eval_hessian(oracle, oracle.constraints[i], x) + for nzval in h + H[k] = μ[i] * nzval + k += 1 end end oracle.eval_hessian_lagrangian_timer += time() - start return end -function _to_sparse(IJ, V) - I, J = [i for (i, _) in IJ], [j for (_, j) in IJ] - n = max(maximum(I), maximum(J)) - return SparseArrays.sparse(I, J, V, n, n) -end - """ - nlp_block_data(d::AbstractNLPEvaluator; use_threads::Bool = false) + nlp_block_data( + d::MOI.AbstractNLPEvaluator; + backend::AbstractNonlinearOracleBackend = DefaultBackend(), + ) Convert an AbstractNLPEvaluator into a SymbolicAD instance of MOI.NLPBlockData. """ -function nlp_block_data(d::MOI.AbstractNLPEvaluator; use_threads::Bool = false) +function nlp_block_data( + d::MOI.AbstractNLPEvaluator; + backend::AbstractNonlinearOracleBackend = DefaultBackend(), +) MOI.initialize(d, [:ExprGraph]) objective = d.has_nlobj ? _Function(MOI.objective_expr(d)) : nothing n = length(d.constraints) @@ -398,14 +362,119 @@ function nlp_block_data(d::MOI.AbstractNLPEvaluator; use_threads::Bool = false) bounds[i] = bound end oracle = _NonlinearOracle( - use_threads, + backend, objective, functions, Dict{UInt,_SymbolicsFunction}(), - Vector{_ConstraintOffset}[], 0.0, 0.0, 0.0, ) return MOI.NLPBlockData(bounds, oracle, d.has_nlobj) end + +### +### ThreadedBackend +### + +struct _ConstraintOffset + index::Int + ∇f_offset::Int + ∇²f_offset::Int +end + +struct ThreadedBackend <: AbstractNonlinearOracleBackend + offsets::Vector{Vector{_ConstraintOffset}} + ThreadedBackend() = new(Vector{_ConstraintOffset}[]) +end + +function MOI.initialize( + oracle::_NonlinearOracle, + backend::ThreadedBackend, + ::Vector{Symbol}, +) + obj_hess_offset = 0 + if oracle.objective !== nothing + h = oracle.objective.expr.hash + obj_hess_offset = length(oracle.symbolic_functions[h].H) + end + offsets = Dict{UInt,Vector{_ConstraintOffset}}( + h => _ConstraintOffset[] for h in keys(oracle.symbolic_functions) + ) + index, ∇f_offset, ∇²f_offset = 1, 1, obj_hess_offset + 1 + for f in oracle.constraints + push!(offsets[f.expr.hash], _ConstraintOffset(index, ∇f_offset, ∇²f_offset)) + index += 1 + symbolic_function = oracle.symbolic_functions[f.expr.hash] + ∇f_offset += length(symbolic_function.g) + ∇²f_offset += length(symbolic_function.H) + end + for v in values(offsets) + push!(backend.offsets, v) + end + return +end + +function MOI.eval_constraint( + oracle::_NonlinearOracle{ThreadedBackend}, + g::AbstractVector{Float64}, + x::AbstractVector{Float64}, +) + start = time() + Threads.@threads for offsets in oracle.backend.offsets + for c in offsets + func = oracle.constraints[c.index] + @inbounds g[c.index] = _eval_function(oracle, func, x) + end + end + oracle.eval_constraint_timer += time() - start + return +end + +function MOI.eval_constraint_jacobian( + oracle::_NonlinearOracle{ThreadedBackend}, + J::AbstractVector{Float64}, + x::AbstractVector{Float64}, +) + start = time() + Threads.@threads for offset in oracle.backend.offsets + for c in offset + func = oracle.constraints[c.index] + g = _eval_gradient(oracle, func, x) + for i in 1:length(g) + @inbounds J[c.∇f_offset+i-1] = g[i] + end + end + end + oracle.eval_constraint_jacobian_timer += time() - start + return +end + +function MOI.eval_hessian_lagrangian( + oracle::_NonlinearOracle{ThreadedBackend}, + H::AbstractVector{Float64}, + x::AbstractVector{Float64}, + σ::Float64, + μ::AbstractVector{Float64}, +) + start = time() + hessian_offset = 0 + if oracle.objective !== nothing + h = _eval_hessian(oracle, oracle.objective, x) + for i in 1:length(h) + H[i] = σ * h[i] + end + hessian_offset += length(h) + end + Threads.@threads for offset in oracle.backend.offsets + for c in offset + func = oracle.constraints[c.index] + h = _eval_hessian(oracle, func, x) + for i in 1:length(h) + @inbounds H[c.∇²f_offset+i-1] = μ[c.index] * h[i] + end + end + end + oracle.eval_hessian_lagrangian_timer += time() - start + return +end diff --git a/test/runtests.jl b/test/runtests.jl index 59db125..6cc30ee 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -4,6 +4,7 @@ using JuMP import Ipopt import PowerModels +import SparseArrays import SymbolicAD import Test @@ -127,9 +128,15 @@ function run_unit_benchmark(model::JuMP.Model) # For the Hessian, we won't compute identical sparsity patterns due to # duplicates, etc. What we need to do is check the full sparse matrix. + function _to_sparse(IJ, V) + I, J = [i for (i, _) in IJ], [j for (_, j) in IJ] + n = max(maximum(I), maximum(J)) + return SparseArrays.sparse(I, J, V, n, n) + end + Test.@test isapprox( - SymbolicAD._to_sparse(H, H_nz), - SymbolicAD._to_sparse(H1, H_nz_1); + _to_sparse(H, H_nz), + _to_sparse(H1, H_nz_1); atol = 1e-6, ) return @@ -152,8 +159,10 @@ function run_solution_benchmark(model::JuMP.Model) serial_model = Ipopt.Optimizer() MOI.copy_to(serial_model, model) - serial_nlp_block = - SymbolicAD.nlp_block_data(JuMP.NLPEvaluator(model); use_threads = false) + serial_nlp_block = SymbolicAD.nlp_block_data( + JuMP.NLPEvaluator(model); + backend = SymbolicAD.DefaultBackend(), + ) MOI.set(serial_model, MOI.NLPBlock(), serial_nlp_block) MOI.optimize!(serial_model) serial_solution = MOI.get( @@ -166,8 +175,10 @@ function run_solution_benchmark(model::JuMP.Model) threaded_model = Ipopt.Optimizer() MOI.copy_to(threaded_model, model) - threaded_nlp_block = - SymbolicAD.nlp_block_data(JuMP.NLPEvaluator(model); use_threads = true) + threaded_nlp_block = SymbolicAD.nlp_block_data( + JuMP.NLPEvaluator(model); + backend = SymbolicAD.ThreadedBackend(), + ) MOI.set(threaded_model, MOI.NLPBlock(), threaded_nlp_block) MOI.optimize!(threaded_model) threaded_solution = MOI.get(