Skip to content

Commit

Permalink
Refactor into AbstractBackends
Browse files Browse the repository at this point in the history
  • Loading branch information
@odow
odow committed Mar 28, 2022
1 parent d313db1 commit 393a260
Show file tree
Hide file tree
Showing 3 changed files with 172 additions and 98 deletions.
6 changes: 0 additions & 6 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
241 changes: 155 additions & 86 deletions src/nonlinear_oracle.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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
23 changes: 17 additions & 6 deletions test/runtests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ using JuMP

import Ipopt
import PowerModels
import SparseArrays
import SymbolicAD
import Test

Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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(
Expand Down

0 comments on commit 393a260

Please sign in to comment.