Skip to content

Reuse the differentiation backend #377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joaquimg wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Reuse the differentiation backend #377

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 59 additions & 8 deletions src/moi_wrapper.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,15 +91,29 @@ mutable struct Optimizer{OT<:MOI.ModelLike} <: MOI.AbstractOptimizer
# instantiated differentiation backend from the options above
diff::Any

# the backend last used for `diff`, and the `(constructor, parametric)` key
# it was built for. `diff` is invalidated by essentially every mutation of
# the model, but the bridge graph of the backend is not: see `_get_diff`
diff_backend::Any
diff_backend_key::Any

# mapping between the `optimizer` and the `diff` models
index_map::Union{Nothing,MOI.Utilities.IndexMap}

# sensitivity input cache using MOI-like sparse format
input_cache::InputCache

function Optimizer(optimizer::OT) where {OT<:MOI.ModelLike}
output =
new{OT}(optimizer, Any[], nothing, nothing, nothing, InputCache())
output = new{OT}(
optimizer,
Any[],
nothing,
nothing,
nothing,
nothing,
nothing,
InputCache(),
)
add_all_model_constructors(output)
add_default_factorization(output)
return output
Expand Down Expand Up @@ -248,6 +262,9 @@ function MOI.empty!(model::Optimizer)
model.diff = nothing
model.index_map = nothing
empty!(model.input_cache)
# `model.diff_backend` is deliberately kept: `_get_diff` empties it before
# handing it back, so it holds no problem data, and keeping it preserves the
# resolved bridge graph. See `_get_diff`.
return
end

Expand Down Expand Up @@ -777,17 +794,23 @@ function _add_bridges(instantiated_model)
return model
end

function _instantiate_diff(model::Optimizer, constructor)
function _is_parametric(model::Optimizer)
# parametric_diff = MOI.supports_constraint(
# model,
# MOI.VariableIndex,
# MOI.Parameter{Float64},
# )
list = MOI.get(
# `NumberOfConstraints` rather than `ListOfConstraintIndices`: this runs on
# every `_diff` and only the emptiness matters, so there is no reason to
# allocate the whole index vector
n = MOI.get(
model,
MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.Parameter{Float64}}(),
MOI.NumberOfConstraints{MOI.VariableIndex,MOI.Parameter{Float64}}(),
)
parametric_diff = !isempty(list)
return !iszero(n)
end

function _instantiate_diff(constructor, parametric_diff::Bool)
model_instance = MOI.instantiate(constructor)
needs_poi =
!MOI.supports_add_constrained_variable(
Expand All @@ -801,6 +824,26 @@ function _instantiate_diff(model::Optimizer, constructor)
return model_bridged
end

# Return an empty differentiation model, reusing `model.diff_backend` when it
# was built for the same `(constructor, parametric_diff)`.
#
# `_diff` is re-entered after every `optimize!`, and instantiating a
# differentiation model resolves the bridge graph of a fresh
# `LazyBridgeOptimizer` from scratch. That graph depends only on the structure
# of the problem, so it is identical every time. `MOI.empty!` on a bridge
# optimizer clears the model and the bridge maps but keeps `graph` and
# `cached_bridge_type` (only `add_bridge`/`remove_bridge` reset those), so
# emptying and reusing the backend drops everything solution-dependent while
# keeping the graph.
function _get_diff(model::Optimizer, constructor, parametric_diff::Bool)
if model.diff_backend !== nothing &&
model.diff_backend_key === (constructor, parametric_diff)
MOI.empty!(model.diff_backend)
return model.diff_backend
end
return _instantiate_diff(constructor, parametric_diff)
end

function _diff(
model::Optimizer,
attr::Union{ForwardDifferentiate,ReverseDifferentiate},
Expand All @@ -811,20 +854,26 @@ function _diff(
elseif isnothing(model.diff)
_check_termination_status(model)
model_constructor = MOI.get(model, ModelConstructor())
parametric_diff = _is_parametric(model)
if isnothing(model_constructor)
for constructor in model.model_constructors
model.diff = _instantiate_diff(model, constructor)
model.diff = _get_diff(model, constructor, parametric_diff)
try
model.index_map = MOI.copy_to(model.diff, model.optimizer)
catch err
if err isa MOI.UnsupportedConstraint ||
err isa MOI.UnsupportedAttribute
# the copy stopped part-way; empty it so it is not left
# in a half-copied state
MOI.empty!(model.diff)
model.diff = nothing
else
rethrow(err)
end
end
if !isnothing(model.diff)
model.diff_backend = model.diff
model.diff_backend_key = (constructor, parametric_diff)
break
end
end
Expand All @@ -838,8 +887,10 @@ function _diff(
)
end
else
model.diff = _instantiate_diff(model, model_constructor)
model.diff = _get_diff(model, model_constructor, parametric_diff)
model.index_map = MOI.copy_to(model.diff, model.optimizer)
model.diff_backend = model.diff
model.diff_backend_key = (model_constructor, parametric_diff)
end
_copy_dual(model.diff, model.optimizer, model.index_map)
end
Expand Down
88 changes: 88 additions & 0 deletions test/moi_wrapper.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,94 @@ function test_forward_or_reverse_without_optimizing_throws()
return
end

# `model.diff` is invalidated by `MOI.optimize!`, so a solve/differentiate loop
# re-enters `_diff` on every iteration. The differentiation backend must be
# reused across those rebuilds, otherwise its bridge graph - which depends only
# on the problem structure - is resolved from scratch every time.
function test_diff_model_is_reused_across_solves()
# min x s.t. x >= p, so that x == p and dx/dp == 1
# model building
model = DiffOpt.diff_optimizer(HiGHS.Optimizer)
MOI.set(model, MOI.Silent(), true)
p, cp = MOI.add_constrained_variable(model, MOI.Parameter(3.0))
x = MOI.add_variable(model)
MOI.add_constraint(
model,
MOI.ScalarAffineFunction(
MOI.ScalarAffineTerm.([1.0, -1.0], [x, p]),
0.0,
),
MOI.GreaterThan(0.0),
)
MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
MOI.set(
model,
MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0),
)
@test model.diff_backend === nothing
sensitivities, backends = Float64[], Any[]
# modify optimize loop
for (p_value, direction) in ((3.0, 1.0), (5.0, 2.0), (7.0, -1.0))
MOI.set(model, MOI.ConstraintSet(), cp, MOI.Parameter(p_value))
MOI.optimize!(model)
DiffOpt.empty_input_sensitivities!(model)
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, direction)
DiffOpt.reverse_differentiate!(model)
push!(
sensitivities,
MOI.get(model, DiffOpt.ReverseConstraintSet(), cp).value,
)
# core of the test: cache the backends for later verification
push!(backends, model.diff)
# the solve really did move, so reusing the backend is not vacuous
@test MOI.get(model, MOI.VariablePrimal(), x) ≈ p_value atol = ATOL
end
# dx/dp == 1, so the reverse sensitivity is the seeded direction
@test sensitivities ≈ [1.0, 2.0, -1.0] atol = ATOL rtol = RTOL
@test all(b === backends[1] for b in backends)
@test model.diff_backend === backends[1]
return
end

# `_instantiate_diff` wraps the differentiation model in POI only when the
# problem has parameters, and a problem can gain them after it has already been
# differentiated - so that flag has to be part of the key the backend is kept
# under, otherwise the wrong wrapping is reused.
function test_diff_backend_distinguishes_parametric_models()
model = DiffOpt.diff_optimizer(HiGHS.Optimizer)
MOI.set(model, MOI.Silent(), true)
x = MOI.add_variable(model)
MOI.add_constraint(model, x, MOI.GreaterThan(1.0))
MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), x)
MOI.optimize!(model)
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, 1.0)
DiffOpt.reverse_differentiate!(model)
non_parametric = model.diff
@test model.diff_backend === non_parametric
# now make the same model parametric
p, cp = MOI.add_constrained_variable(model, MOI.Parameter(2.0))
MOI.add_constraint(
model,
MOI.ScalarAffineFunction(
MOI.ScalarAffineTerm.([1.0, -1.0], [x, p]),
0.0,
),
MOI.GreaterThan(0.0),
)
MOI.optimize!(model)
DiffOpt.empty_input_sensitivities!(model)
MOI.set(model, DiffOpt.ReverseVariablePrimal(), x, 1.0)
DiffOpt.reverse_differentiate!(model)
@test MOI.get(model, MOI.VariablePrimal(), x) ≈ 2.0 atol = ATOL
@test MOI.get(model, DiffOpt.ReverseConstraintSet(), cp).value ≈ 1.0 atol =
ATOL
# rebuilt under the parametric key, not a reuse of the non-parametric one
@test model.diff !== non_parametric
return
end

struct TestSolver end

# always use IterativeSolvers
Expand Down
Loading