Nonlinear diff: skip evaluator rebuild when structure is unchanged

src/NonLinearProgram/NonLinearProgram.jl

@@ -352,6 +352,48 @@ function MOI.empty!(model::Model)
     return
 end
 
+# The `Cache` built by `_cache_evaluator!` is a pure function of the *structure* of the
+# `Form` (variables, constraints, parameters, objective): the evaluator tapes, the
+# Jacobian/Hessian sparsity and the index mappings do not depend on parameter values (the
+# evaluator reads those live from `form.model.parameters`) nor on the primal/dual point
+# (which lives in `model.x`, `model.y`, `model.s`). It therefore stays valid until the
+# structure changes. The methods below intercept every structural mutation of the model
+# and invalidate the cache, so that `_cache_evaluator!` can return an existing cache
+# instead of rebuilding it on every differentiation call.
+
+function _invalidate_evaluator_cache!(model::Model)
+    model.cache = nothing
+    return
+end
+
+function MOI.add_variable(model::Model)
+    _invalidate_evaluator_cache!(model)
+    return MOI.add_variable(model.model)
+end
+
+function MOI.add_variables(model::Model, n)
+    _invalidate_evaluator_cache!(model)
+    return MOI.add_variables(model.model, n)
+end
+
+function MOI.add_constraint(
+    model::Model,
+    func::MOI.AbstractFunction,
+    set::MOI.AbstractSet,
+)
+    _invalidate_evaluator_cache!(model)
+    return MOI.add_constraint(model.model, func, set)
+end
+
+function MOI.set(
+    model::Model,
+    attr::Union{MOI.ObjectiveSense,MOI.ObjectiveFunction},
+    value,
+)
+    _invalidate_evaluator_cache!(model)
+    return MOI.set(model.model, attr, value)
+end
+
 include("nlp_utilities.jl")
 include("vno_bridge.jl")
 
@@ -449,7 +491,23 @@ _get_num_primal_vars(model::Model) = _get_num_primal_vars(model.model)
 _get_num_params(form::Form) = length(_all_params(form))
 _get_num_params(model::Model) = _get_num_params(model.model)
 
+"""
+    _cache_evaluator!(model::Model)
+
+Build the `Cache` for `model` (the `MOI.Nonlinear` evaluator with its AD tapes and
+Jacobian/Hessian sparsity, plus the primal/dual/bound index mappings), or return the
+existing `model.cache` when there is one.
+
+The cache only depends on the structure of the `Form`, and every structural mutation
+resets `model.cache` to `nothing` (see `_invalidate_evaluator_cache!`), so an existing
+cache is always current: rebuilding it would produce an identical evaluator and identical
+mappings. Returning it directly avoids paying the tape construction and sparsity
+detection on every differentiation call.
+"""
 function _cache_evaluator!(model::Model)
+    if model.cache !== nothing
+        return model.cache
+    end
     form = model.model
     # Retrieve and sort primal variables by NLP index
     params = sort(_all_params(model); by = x -> x.value)

test/nlp_program.jl

@@ -1126,6 +1126,141 @@ function test_reverse_bounds_upper()
     @test isapprox(dp, 2.88888; atol = 1e-4)
 end
 
+# Reach the `NonLinearProgram.Model` inside the differentiation backend of a JuMP model
+# built with `DiffOpt.nonlinear_diff_model`. Used to check evaluator-cache reuse.
+function _inner_nlp_diff_model(model::JuMP.Model)
+    diffopt = JuMP.unsafe_backend(model)
+    inner = diffopt.diff
+    while !(inner isa DiffOpt.NonLinearProgram.Model)
+        inner = inner.model
+    end
+    return inner
+end
+
+function _build_cache_reuse_model(P)
+    model = DiffOpt.nonlinear_diff_model(Ipopt.Optimizer)
+    set_silent(model)
+    @variable(model, p[i=1:P] in MOI.Parameter(1.0 + 0.2 * i))
+    @variable(model, x[1:P] >= 0)
+    @constraint(model, [i = 1:P], x[i] * (1 + 0.1 * sin(p[i])) >= p[i])
+    @constraint(model, sum(x) >= 0.4 * P)
+    @objective(
+        model,
+        Min,
+        sum((x[i] - p[i])^2 for i in 1:P) + 0.01 * sum(x[i]^4 for i in 1:P)
+    )
+    return model, p, x
+end
+
+function test_evaluator_cache_reused_across_differentiations()
+    P = 4
+    model, p, x = _build_cache_reuse_model(P)
+    optimize!(model)
+    @assert is_solved_and_feasible(model)
+    # The first differentiation builds the evaluator cache.
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model, x[i], cos(1.3 * i))
+    end
+    DiffOpt.reverse_differentiate!(model)
+    nlp = _inner_nlp_diff_model(model)
+    cache = nlp.cache
+    @test cache !== nothing
+    # Later differentiation calls on the same structure return the same Cache object
+    # instead of rebuilding tapes and sparsity.
+    seed = [sin(0.7 * i) for i in 1:P]
+    DiffOpt.empty_input_sensitivities!(model)
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model, x[i], seed[i])
+    end
+    DiffOpt.reverse_differentiate!(model)
+    @test nlp.cache === cache
+    dp_reused = [
+        MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value
+        for i in 1:P
+    ]
+    DiffOpt.empty_input_sensitivities!(model)
+    DiffOpt.set_forward_parameter(model, p[1], 1.0)
+    DiffOpt.forward_differentiate!(model)
+    @test nlp.cache === cache
+    dx_reused =
+        [MOI.get(model, DiffOpt.ForwardVariablePrimal(), x[i]) for i in 1:P]
+    # A fresh model whose only differentiation uses the same seeds must produce
+    # identical results: the rebuilt cache is the same pure function of the structure,
+    # so reuse changes no floating-point operation.
+    model2, p2, x2 = _build_cache_reuse_model(P)
+    optimize!(model2)
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model2, x2[i], seed[i])
+    end
+    DiffOpt.reverse_differentiate!(model2)
+    dp_fresh = [
+        MOI.get(model2, DiffOpt.ReverseConstraintSet(), ParameterRef(p2[i])).value for i in 1:P
+    ]
+    @test dp_reused == dp_fresh
+    DiffOpt.empty_input_sensitivities!(model2)
+    DiffOpt.set_forward_parameter(model2, p2[1], 1.0)
+    DiffOpt.forward_differentiate!(model2)
+    dx_fresh =
+        [MOI.get(model2, DiffOpt.ForwardVariablePrimal(), x2[i]) for i in 1:P]
+    @test dx_reused == dx_fresh
+    return
+end
+
+function test_evaluator_cache_invalidated_on_structural_change()
+    P = 3
+    model, p, x = _build_cache_reuse_model(P)
+    optimize!(model)
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model, x[i], 1.0)
+    end
+    DiffOpt.reverse_differentiate!(model)
+    nlp = _inner_nlp_diff_model(model)
+    @test nlp.cache !== nothing
+    # Every structural mutation of the inner model must drop the cache.
+    MOI.set(nlp, MOI.ObjectiveSense(), MOI.MIN_SENSE)
+    @test nlp.cache === nothing
+    DiffOpt.reverse_differentiate!(model)
+    @test nlp.cache !== nothing
+    MOI.add_variable(nlp)
+    @test nlp.cache === nothing
+    return
+end
+
+function test_evaluator_cache_after_public_structural_edit()
+    # Through the public API, a structural edit resets the whole differentiation model
+    # (`DiffOpt.Optimizer` sets `diff = nothing`), so differentiation after the edit
+    # builds a new cache and must match a fresh model built directly in the final state.
+    P = 3
+    model, p, x = _build_cache_reuse_model(P)
+    optimize!(model)
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model, x[i], sin(1.0 + i))
+    end
+    DiffOpt.reverse_differentiate!(model)
+    nlp_before = _inner_nlp_diff_model(model)
+    @constraint(model, sum(x) <= 10.0 * P)
+    optimize!(model)
+    DiffOpt.reverse_differentiate!(model)
+    nlp_after = _inner_nlp_diff_model(model)
+    @test nlp_after !== nlp_before
+    dp = [
+        MOI.get(model, DiffOpt.ReverseConstraintSet(), ParameterRef(p[i])).value
+        for i in 1:P
+    ]
+    model2, p2, x2 = _build_cache_reuse_model(P)
+    @constraint(model2, sum(x2) <= 10.0 * P)
+    optimize!(model2)
+    for i in 1:P
+        DiffOpt.set_reverse_variable(model2, x2[i], sin(1.0 + i))
+    end
+    DiffOpt.reverse_differentiate!(model2)
+    dp2 = [
+        MOI.get(model2, DiffOpt.ReverseConstraintSet(), ParameterRef(p2[i])).value for i in 1:P
+    ]
+    @test dp ≈ dp2 rtol = 1e-10
+    return
+end
+
 end # module
 
 TestNLPProgram.runtests()