DNMY: experimental testing for fast resolves of NLP

This is NOT safe to merge because doesn't update the expression
graphs and so will break AmplNLWriter etc.

src/JuMP.jl

@@ -152,6 +152,8 @@ mutable struct Model <: AbstractModel
     # A flag to track whether we have modified the model after calling
     # optimize!.
     is_model_dirty::Bool
+    # A flag to track whether we need to rebuild the NLP model before optimize!
+    is_nlp_model_dirty::Bool
     # Enable extensions to attach arbitrary information to a JuMP model by
     # using an extension-specific symbol as a key.
     ext::Dict{Symbol,Any}
@@ -256,6 +258,7 @@ function direct_model(backend::MOI.ModelLike)
         Dict{Symbol,Any}(),
         0,
         false,
+        true,
         Dict{Symbol,Any}(),
         true,
     )

src/nlp.jl

@@ -25,6 +25,7 @@ end
 function _init_NLP(model::Model)
     if model.nlp_model === nothing
         model.nlp_model = MOI.Nonlinear.Model()
+        model.is_nlp_model_dirty = true
     end
     return
 end
@@ -161,6 +162,7 @@ function set_nonlinear_objective(model::Model, sense::MOI.OptimizationSense, x)
     _init_NLP(model)
     set_objective_sense(model, sense)
     MOI.Nonlinear.set_objective(model.nlp_model, x)
+    model.is_nlp_model_dirty = true
     return
 end
 
@@ -206,6 +208,7 @@ end
 function add_nonlinear_parameter(model::Model, value::Real)
     _init_NLP(model)
     p = MOI.Nonlinear.add_parameter(model.nlp_model, Float64(value))
+    model.is_nlp_model_dirty = true
     return NonlinearParameter(model, p.value)
 end
 
@@ -312,6 +315,7 @@ subexpression[1]: x + x ^ 2.0
 function add_nonlinear_expression(model::Model, ex)
     _init_NLP(model)
     index = MOI.Nonlinear.add_expression(model.nlp_model, ex)
+    model.is_nlp_model_dirty = true
     return NonlinearExpression(model, index.value)
 end
 
@@ -439,6 +443,7 @@ function add_nonlinear_constraint(model::Model, ex::Expr)
     _init_NLP(model)
     f, set = _expr_to_constraint(ex)
     c = MOI.Nonlinear.add_constraint(model.nlp_model, f, set)
+    model.is_nlp_model_dirty = true
     return ConstraintRef(model, c, ScalarShape())
 end
 
@@ -465,6 +470,7 @@ function delete(model::Model, c::NonlinearConstraintRef)
     _init_NLP(model)
     index = MOI.Nonlinear.ConstraintIndex(c.index.value)
     MOI.Nonlinear.delete(model.nlp_model, index)
+    model.is_nlp_model_dirty = true
     return
 end
 
@@ -647,6 +653,7 @@ function register(
     end
     _init_NLP(model)
     MOI.Nonlinear.register_operator(model.nlp_model, op, dimension, f)
+    model.is_nlp_model_dirty = true
     return
 end
 
@@ -725,6 +732,7 @@ function register(
         end
         MOI.Nonlinear.register_operator(model.nlp_model, op, dimension, f, ∇f)
     end
+    model.is_nlp_model_dirty = true
     return
 end
 
@@ -780,6 +788,7 @@ function register(
     end
     _init_NLP(model)
     MOI.Nonlinear.register_operator(model.nlp_model, op, dimension, f, ∇f, ∇²f)
+    model.is_nlp_model_dirty = true
     return
 end
 

src/optimizer_interface.jl

@@ -160,13 +160,15 @@ function optimize!(
 )
     # The nlp_model is not kept in sync, so re-set it here.
     # TODO: Consider how to handle incremental solves.
-    if nonlinear_model(model) !== nothing
+    if nonlinear_model(model) !== nothing && model.is_nlp_model_dirty
+        @warn("Setting NLPBLock")
         evaluator = MOI.Nonlinear.Evaluator(
             nonlinear_model(model),
             _differentiation_backend,
             index.(all_variables(model)),
         )
         MOI.set(model, MOI.NLPBlock(), MOI.NLPBlockData(evaluator))
+        model.is_nlp_model_dirty = false
     end
     # If the user or an extension has provided an optimize hook, call
     # that instead of solving the model ourselves

src/variables.jl

@@ -1062,6 +1062,7 @@ function add_variable end
 
 function add_variable(model::Model, v::ScalarVariable, name::String = "")
     model.is_model_dirty = true
+    model.is_nlp_model_dirty = true
     return _moi_add_variable(backend(model), model, v, name)
 end