diff --git a/src/moi_wrapper.jl b/src/moi_wrapper.jl index 8d14a5f2..a19e8562 100644 --- a/src/moi_wrapper.jl +++ b/src/moi_wrapper.jl @@ -91,6 +91,12 @@ 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} @@ -98,8 +104,16 @@ mutable struct Optimizer{OT<:MOI.ModelLike} <: MOI.AbstractOptimizer 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 @@ -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 @@ -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( @@ -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}, @@ -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 @@ -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 diff --git a/test/moi_wrapper.jl b/test/moi_wrapper.jl index 2b92ee4b..24fa66b0 100644 --- a/test/moi_wrapper.jl +++ b/test/moi_wrapper.jl @@ -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