diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7b1d59ca..1567fb19 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,6 +54,13 @@ jobs: ${{ runner.os }}-test-${{ env.cache-name }}- ${{ runner.os }}-test- ${{ runner.os }}- + - name: MOI + shell: julia --project=@. {0} + run: | + using Pkg + Pkg.add([ + PackageSpec(name="MathOptInterface", rev="bl/linearity"), + ]) - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 diff --git a/Project.toml b/Project.toml index c8defac2..841c2456 100644 --- a/Project.toml +++ b/Project.toml @@ -11,7 +11,15 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +[weakdeps] +ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50" + +[extensions] +NLPModelsJuMPArrayDiffExt = "ArrayDiff" + [compat] +ArrayDiff = "0.1" +ExaModels = "0.11" JuMP = "1.25" LinearAlgebra = "1.10" MathOptInterface = "1.46" @@ -25,10 +33,11 @@ Test = "1.10" julia = "1.10" [extras] +ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NLPModelsTest = "7998695d-6960-4d3a-85c4-e1bceb8cd856" Percival = "01435c0c-c90d-11e9-3788-63660f8fbccc" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["LinearAlgebra", "NLPModelsTest", "Percival", "Test"] +test = ["ExaModels", "LinearAlgebra", "NLPModelsTest", "Percival", "Test"] diff --git a/ext/NLPModelsJuMPArrayDiffExt.jl b/ext/NLPModelsJuMPArrayDiffExt.jl new file mode 100644 index 00000000..2f3b9989 --- /dev/null +++ b/ext/NLPModelsJuMPArrayDiffExt.jl @@ -0,0 +1,110 @@ +module NLPModelsJuMPArrayDiffExt + +import NLPModelsJuMP +import ArrayDiff +import MathOptInterface as MOI +import NLPModels +import LinearAlgebra + +NLPModelsJuMP._nonlinear_model(ad::ArrayDiff.Mode) = ArrayDiff.model(ad) + +NLPModelsJuMP._supports_vector_objective(::ArrayDiff.Mode) = true + +# Detect `(...).^2` (broadcast `:^` with exponent 2) and return the residual `...`. +function NLPModelsJuMP._detect_squared_residual(inner::ArrayDiff.ArrayNonlinearFunction) + if inner.head !== :^ || !inner.broadcasted + return nothing + end + if length(inner.args) != 2 + return nothing + end + exponent = inner.args[2] + if !(exponent isa Number) || exponent != 2 + return nothing + end + return inner.args[1] +end + +mutable struct ArrayDiffNLSModel{T, V <: AbstractVector{T}, R} <: NLPModels.AbstractNLSModel{T, V} + meta::NLPModels.NLPModelMeta{T, V} + nls_meta::NLPModels.NLSMeta{T, V} + counters::NLPModels.NLSCounters + evaluator::ArrayDiff.Evaluator{T, R} +end + +function NLPModelsJuMP._build_nls_from_residual( + moimodel::MOI.ModelLike, + residual::ArrayDiff.ArrayNonlinearFunction, + ad::ArrayDiff.Mode{S}, +) where {S <: AbstractVector{<:Real}} + T = eltype(S) + V = S + _, nvar, lvar, uvar, x0 = NLPModelsJuMP.parser_variables(moimodel) + lvar = convert(V, lvar) + uvar = convert(V, uvar) + x0 = convert(V, x0) + model = ArrayDiff.model(ad) + ArrayDiff.set_residual!(model, residual) + vars = MOI.get(moimodel, MOI.ListOfVariableIndices()) + evaluator = MOI.Nonlinear.Evaluator(model, ad, vars) + MOI.initialize(evaluator, [:Grad, :Jac, :JacVec]) + nresid = ArrayDiff.residual_dimension(evaluator) + meta = NLPModels.NLPModelMeta{T, V}( + nvar; + x0 = x0, + lvar = lvar, + uvar = uvar, + minimize = MOI.get(moimodel, MOI.ObjectiveSense()) == MOI.MIN_SENSE, + islp = false, + name = "ArrayDiffNLS", + hprod_available = false, + hess_available = false, + ) + nls_meta = NLPModels.NLSMeta{T, V}( + nresid, + nvar; + x0 = x0, + nnzj = nresid * nvar, + nnzh = 0, + jac_residual_available = false, + hess_residual_available = false, + jprod_residual_available = true, + jtprod_residual_available = true, + hprod_residual_available = false, + ) + return ArrayDiffNLSModel(meta, nls_meta, NLPModels.NLSCounters(), evaluator) +end + +function NLPModels.residual!( + nls::ArrayDiffNLSModel, + x::AbstractVector, + Fx::AbstractVector, +) + NLPModels.increment!(nls, :neval_residual) + ArrayDiff.eval_residual!(nls.evaluator, Fx, x) + return Fx +end + +function NLPModels.jprod_residual!( + nls::ArrayDiffNLSModel, + x::AbstractVector, + v::AbstractVector, + Jv::AbstractVector, +) + NLPModels.increment!(nls, :neval_jprod_residual) + ArrayDiff.eval_residual_jprod!(nls.evaluator, Jv, x, v) + return Jv +end + +function NLPModels.jtprod_residual!( + nls::ArrayDiffNLSModel, + x::AbstractVector, + v::AbstractVector, + Jtv::AbstractVector, +) + NLPModels.increment!(nls, :neval_jtprod_residual) + ArrayDiff.eval_residual_jtprod!(nls.evaluator, Jtv, x, v) + return Jtv +end + +end diff --git a/src/MOI_wrapper.jl b/src/MOI_wrapper.jl index 99dc3abb..717204fa 100644 --- a/src/MOI_wrapper.jl +++ b/src/MOI_wrapper.jl @@ -3,11 +3,19 @@ import SolverCore mutable struct Optimizer <: MOI.AbstractOptimizer options::Dict{String, Any} silent::Bool + ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation solver - nlp::Union{Nothing, MathOptNLPModel} + nlp::Union{Nothing, AbstractNLPModel} stats::Union{Nothing, SolverCore.GenericExecutionStats} function Optimizer() - return new(Dict{String, Any}(), false, nothing, nothing, nothing) + return new( + Dict{String, Any}(), + false, + MOI.Nonlinear.SparseReverseMode(), + nothing, + nothing, + nothing, + ) end end @@ -48,6 +56,32 @@ end MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent +function MOI.supports( + optimizer::Optimizer, + ::MOI.ObjectiveFunction{F}, +) where {F <: MOI.AbstractVectorFunction} + return _supports_vector_objective(optimizer.ad_backend) +end + +### +### MOI.AutomaticDifferentiationBackend +### + +MOI.supports(::Optimizer, ::MOI.AutomaticDifferentiationBackend) = true + +function MOI.get(optimizer::Optimizer, ::MOI.AutomaticDifferentiationBackend) + return optimizer.ad_backend +end + +function MOI.set( + optimizer::Optimizer, + ::MOI.AutomaticDifferentiationBackend, + backend::MOI.Nonlinear.AbstractAutomaticDifferentiation, +) + optimizer.ad_backend = backend + return +end + ### ### MOI.AbstractModelAttribute ### @@ -89,12 +123,32 @@ function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike) "No solver specified, use for instance `using Percival; JuMP.set_attribute(model, \"solver\", PercivalSolver)`", ) end - dest.nlp, index_map = nlp_model(src) + nls = _try_nls_model(src, dest.ad_backend) + if nls !== nothing + dest.nlp = nls + dest.solver = dest.options["solver"](dest.nlp) + return parser_variables(src)[1] + end + if _route_all_to_evaluator(dest.ad_backend) + dest.nlp, index_map = evaluator_nlp_model(src; ad_backend = dest.ad_backend) + else + dest.nlp, index_map = nlp_model(src; ad_backend = dest.ad_backend) + end dest.solver = dest.options["solver"](dest.nlp) return index_map end function MOI.optimize!(model::Optimizer) + if model.nlp === nothing + # Direct mode: build NLPModel from the optimizer itself + if !haskey(model.options, "solver") + error( + "No solver specified, use for instance `using Percival; JuMP.set_attribute(model, \"solver\", PercivalSolver)`", + ) + end + model.nlp, _ = nlp_model(model; ad_backend = model.ad_backend) + model.solver = model.options["solver"](model.nlp) + end options = Dict{Symbol, Any}( Symbol(key) => model.options[key] for key in keys(model.options) if key != "solver" ) diff --git a/src/NLPModelsJuMP.jl b/src/NLPModelsJuMP.jl index 5e726539..e4d72565 100644 --- a/src/NLPModelsJuMP.jl +++ b/src/NLPModelsJuMP.jl @@ -3,6 +3,7 @@ module NLPModelsJuMP include("utils.jl") include("moi_nlp_model.jl") include("moi_nls_model.jl") +include("moi_evaluator_model.jl") include("MOI_wrapper.jl") end diff --git a/src/moi_evaluator_model.jl b/src/moi_evaluator_model.jl new file mode 100644 index 00000000..85cf3d6d --- /dev/null +++ b/src/moi_evaluator_model.jl @@ -0,0 +1,419 @@ +export EvaluatorNLPModel + +""" + _route_all_to_evaluator(backend::MOI.Nonlinear.AbstractAutomaticDifferentiation) + +Hook for AD-backend extensions: return `true` if all supported constraints +(affine, quadratic, nonlinear, and vector-nonlinear-oracle) should be routed +to the backend's nonlinear model built by `MOI.Nonlinear.model(backend)`, +instead of parsing the affine and quadratic constraints into +[`LinearConstraints`](@ref) and [`QuadraticConstraints`](@ref). + +Backends that exploit the structure of affine and quadratic constraints +themselves (for example `ExaModels.SIMDMode`) opt in by overriding +`MOI.Nonlinear.exploits_structure`; the model is then built with +[`evaluator_nlp_model`](@ref). +""" +function _route_all_to_evaluator( + backend::MOI.Nonlinear.AbstractAutomaticDifferentiation, +) + return MOI.Nonlinear.exploits_structure(backend) +end + +""" + EvaluatorNLPModel <: AbstractNLPModel{Float64, Vector{Float64}} + +An `AbstractNLPModel` whose rows are all served by a single +`MOI.AbstractNLPEvaluator` built from an AD backend's model stack, instead of +the [`LinearConstraints`](@ref)/[`QuadraticConstraints`](@ref)/evaluator split +of [`MathOptNLPModel`](@ref). + +The `lin`/`nln` split of the meta is derived from +`MOI.Nonlinear.constraint_linearity`, and the constant Jacobian block of the +linear rows is materialized once at construction, so `cons_lin!` and +`jac_lin_coord!` never trigger the evaluator. The unsplit methods (`cons!`, +`jac_coord!`, `hess_coord!`) are single full passes over the evaluator, so +fused evaluators pay no slicing cost on them. +""" +mutable struct EvaluatorNLPModel{E <: MOI.AbstractNLPEvaluator} <: + AbstractNLPModel{Float64, Vector{Float64}} + meta::NLPModelMeta{Float64, Vector{Float64}} + evaluator::E + has_jacvec::Bool + has_hessvec::Bool + # Full Jacobian and Hessian structures, in evaluator order. + jac_rows::Vector{Int} + jac_cols::Vector{Int} + hess_rows::Vector{Int} + hess_cols::Vector{Int} + # Indices of the Jacobian entries that belong to linear and to nonlinear + # rows, and the corresponding structures with rows renumbered to 1:nlin and + # 1:nnln. + lin_jac_index::Vector{Int} + nln_jac_index::Vector{Int} + lin_jac_rows::Vector{Int} + lin_jac_cols::Vector{Int} + nln_jac_rows::Vector{Int} + nln_jac_cols::Vector{Int} + # The materialized linear block: rows `meta.lin` are `A * x + b`, and the + # values of their Jacobian entries are the constant `lin_vals`. + lin_A::SparseMatrixCSC{Float64, Int} + lin_b::Vector{Float64} + lin_vals::Vector{Float64} + # Scratch buffers for the full-pass-and-gather methods. + g_buffer::Vector{Float64} + j_buffer::Vector{Float64} + counters::Counters +end + +""" + evaluator_nlp_model(moimodel::MOI.ModelLike; ad_backend, hessian, name) + +Build an [`EvaluatorNLPModel`](@ref) by feeding the objective and every +constraint of `moimodel` into `MOI.Nonlinear.model(ad_backend)`. + +The rows of the model are the rows of the resulting evaluator: the layers' +rows (for example vector-nonlinear-oracle constraints) come first, followed +by the scalar constraints in the order they appear in `moimodel`. +""" +function evaluator_nlp_model( + moimodel::MOI.ModelLike; + ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation, + hessian::Bool = true, + name::String = "Generic", +) + index_map, nvar, lvar, uvar, x0 = parser_variables(moimodel) + nlp_model = MOI.Nonlinear.model(ad_backend) + vars = MOI.get(moimodel, MOI.ListOfVariableIndices()) + # Constraints. The oracle rows come before the scalar rows, whatever the + # order in `moimodel`, because the layers' rows come first. + oracle_cis = Tuple{MOI.ConstraintIndex, Int}[] + scalar_cis = MOI.ConstraintIndex[] + oracle_rows = 0 + for (F, S) in MOI.get(moimodel, MOI.ListOfConstraintTypesPresent()) + if F == MOI.VariableIndex + continue # Variable bounds are handled by `parser_variables`. + elseif F == MOI.VectorOfVariables && S <: MOI.VectorNonlinearOracle + for ci in MOI.get(moimodel, MOI.ListOfConstraintIndices{F, S}()) + func = MOI.get(moimodel, MOI.ConstraintFunction(), ci) + set = MOI.get(moimodel, MOI.ConstraintSet(), ci) + MOI.Nonlinear.add_constraint(nlp_model, func, set) + push!(oracle_cis, (ci, oracle_rows)) + oracle_rows += set.output_dimension + end + else + for ci in MOI.get(moimodel, MOI.ListOfConstraintIndices{F, S}()) + func = MOI.get(moimodel, MOI.ConstraintFunction(), ci) + set = MOI.get(moimodel, MOI.ConstraintSet(), ci) + MOI.Nonlinear.add_constraint(nlp_model, func, set) + push!(scalar_cis, ci) + end + end + end + for (ci, offset) in oracle_cis + index_map[ci] = typeof(ci)(offset + 1) + end + for (k, ci) in enumerate(scalar_cis) + index_map[ci] = typeof(ci)(oracle_rows + k) + end + # Objective + F = MOI.get(moimodel, MOI.ObjectiveFunctionType()) + sense = MOI.get(moimodel, MOI.ObjectiveSense()) + if sense != MOI.FEASIBILITY_SENSE + MOI.Nonlinear.set_objective( + nlp_model, + MOI.get(moimodel, MOI.ObjectiveFunction{F}()), + ) + end + # Evaluator + evaluator = MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars) + features = MOI.features_available(evaluator) + requested = [:Grad, :Jac, :JacVec, :HessVec] + if hessian + push!(requested, :Hess) + end + requested = intersect(requested, features) + MOI.initialize(evaluator, requested) + has_hess = :Hess in requested + bounds = MOI.Nonlinear.constraint_bounds(evaluator) + ncon = length(bounds) + lcon = [b.lower for b in bounds] + ucon = [b.upper for b in bounds] + linearity = MOI.Nonlinear.constraint_linearity(evaluator) + lin = if linearity === nothing + Int[] + else + findall(==(MOI.Nonlinear.LINEAR), linearity) + end + # Jacobian structure, partitioned by row linearity. + jac_structure = MOI.jacobian_structure(evaluator) + nnzj = length(jac_structure) + jac_rows = [r for (r, _) in jac_structure] + jac_cols = [c for (_, c) in jac_structure] + lin_pos = zeros(Int, ncon) + for (k, r) in enumerate(lin) + lin_pos[r] = k + end + nln = findall(iszero, lin_pos) + nln_pos = zeros(Int, ncon) + for (k, r) in enumerate(nln) + nln_pos[r] = k + end + lin_jac_index = findall(k -> lin_pos[jac_rows[k]] > 0, 1:nnzj) + nln_jac_index = findall(k -> lin_pos[jac_rows[k]] == 0, 1:nnzj) + lin_jac_rows = [lin_pos[jac_rows[k]] for k in lin_jac_index] + lin_jac_cols = jac_cols[lin_jac_index] + nln_jac_rows = [nln_pos[jac_rows[k]] for k in nln_jac_index] + nln_jac_cols = jac_cols[nln_jac_index] + # Materialize the linear block: its Jacobian entries are x-independent, so + # one evaluation gives the constant values, and one evaluation of the + # constraints at zero gives the constants. + j_buffer = zeros(nnzj) + g_buffer = zeros(ncon) + lin_b = zeros(length(lin)) + lin_vals = zeros(length(lin_jac_index)) + lin_A = SparseMatrixCSC{Float64, Int}(spzeros(length(lin), nvar)) + if !isempty(lin) + MOI.eval_constraint_jacobian(evaluator, j_buffer, x0) + lin_vals .= j_buffer[lin_jac_index] + lin_A = sparse(lin_jac_rows, lin_jac_cols, lin_vals, length(lin), nvar) + MOI.eval_constraint(evaluator, g_buffer, zeros(nvar)) + lin_b .= g_buffer[lin] + end + # Hessian structure, forced to the lower triangle. + hess_structure = if has_hess + MOI.hessian_lagrangian_structure(evaluator) + else + Tuple{Int, Int}[] + end + hess_rows = [max(r, c) for (r, c) in hess_structure] + hess_cols = [min(r, c) for (r, c) in hess_structure] + meta = NLPModelMeta( + nvar, + x0 = x0, + lvar = lvar, + uvar = uvar, + ncon = ncon, + y0 = zeros(ncon), + lcon = lcon, + ucon = ucon, + nnzj = nnzj, + nnzh = length(hess_structure), + lin = lin, + lin_nnzj = length(lin_jac_index), + nln_nnzj = length(nln_jac_index), + minimize = sense != MOI.MAX_SENSE, + islp = false, + name = name, + ) + nlp = EvaluatorNLPModel( + meta, + evaluator, + :JacVec in requested, + :HessVec in requested, + jac_rows, + jac_cols, + hess_rows, + hess_cols, + lin_jac_index, + nln_jac_index, + lin_jac_rows, + lin_jac_cols, + nln_jac_rows, + nln_jac_cols, + lin_A, + lin_b, + lin_vals, + g_buffer, + j_buffer, + Counters(), + ) + return nlp, index_map +end + +function NLPModels.obj(nlp::EvaluatorNLPModel, x::AbstractVector) + increment!(nlp, :neval_obj) + return MOI.eval_objective(nlp.evaluator, x) +end + +function NLPModels.grad!(nlp::EvaluatorNLPModel, x::AbstractVector, g::AbstractVector) + increment!(nlp, :neval_grad) + MOI.eval_objective_gradient(nlp.evaluator, g, x) + return g +end + +function NLPModels.cons!(nlp::EvaluatorNLPModel, x::AbstractVector, c::AbstractVector) + increment!(nlp, :neval_cons) + MOI.eval_constraint(nlp.evaluator, c, x) + return c +end + +function NLPModels.cons_lin!(nlp::EvaluatorNLPModel, x::AbstractVector, c::AbstractVector) + increment!(nlp, :neval_cons_lin) + NLPModels.coo_prod!(nlp.lin_jac_rows, nlp.lin_jac_cols, nlp.lin_vals, x, c) + c .+= nlp.lin_b + return c +end + +function NLPModels.cons_nln!(nlp::EvaluatorNLPModel, x::AbstractVector, c::AbstractVector) + increment!(nlp, :neval_cons_nln) + MOI.eval_constraint(nlp.evaluator, nlp.g_buffer, x) + c .= view(nlp.g_buffer, nlp.meta.nln) + return c +end + +function NLPModels.jac_structure!( + nlp::EvaluatorNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= nlp.jac_rows + cols .= nlp.jac_cols + return rows, cols +end + +function NLPModels.jac_lin_structure!( + nlp::EvaluatorNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= nlp.lin_jac_rows + cols .= nlp.lin_jac_cols + return rows, cols +end + +function NLPModels.jac_nln_structure!( + nlp::EvaluatorNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= nlp.nln_jac_rows + cols .= nlp.nln_jac_cols + return rows, cols +end + +function NLPModels.jac_coord!(nlp::EvaluatorNLPModel, x::AbstractVector, vals::AbstractVector) + increment!(nlp, :neval_jac) + MOI.eval_constraint_jacobian(nlp.evaluator, vals, x) + return vals +end + +function NLPModels.jac_lin_coord!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + vals::AbstractVector, +) + increment!(nlp, :neval_jac_lin) + vals .= nlp.lin_vals + return vals +end + +function NLPModels.jac_nln_coord!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + vals::AbstractVector, +) + increment!(nlp, :neval_jac_nln) + MOI.eval_constraint_jacobian(nlp.evaluator, nlp.j_buffer, x) + vals .= view(nlp.j_buffer, nlp.nln_jac_index) + return vals +end + +function NLPModels.jprod!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + v::AbstractVector, + Jv::AbstractVector, +) + increment!(nlp, :neval_jprod) + if nlp.has_jacvec + MOI.eval_constraint_jacobian_product(nlp.evaluator, Jv, x, v) + else + MOI.eval_constraint_jacobian(nlp.evaluator, nlp.j_buffer, x) + NLPModels.coo_prod!(nlp.jac_rows, nlp.jac_cols, nlp.j_buffer, v, Jv) + end + return Jv +end + +function NLPModels.jtprod!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + v::AbstractVector, + Jtv::AbstractVector, +) + increment!(nlp, :neval_jtprod) + if nlp.has_jacvec + MOI.eval_constraint_jacobian_transpose_product(nlp.evaluator, Jtv, x, v) + else + MOI.eval_constraint_jacobian(nlp.evaluator, nlp.j_buffer, x) + NLPModels.coo_prod!(nlp.jac_cols, nlp.jac_rows, nlp.j_buffer, v, Jtv) + end + return Jtv +end + +function NLPModels.hess_structure!( + nlp::EvaluatorNLPModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + rows .= nlp.hess_rows + cols .= nlp.hess_cols + return rows, cols +end + +function NLPModels.hess_coord!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + y::AbstractVector, + vals::AbstractVector; + obj_weight::Real = 1.0, +) + increment!(nlp, :neval_hess) + MOI.eval_hessian_lagrangian(nlp.evaluator, vals, x, obj_weight, y) + return vals +end + +function NLPModels.hess_coord!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + vals::AbstractVector; + obj_weight::Real = 1.0, +) + increment!(nlp, :neval_hess) + MOI.eval_hessian_lagrangian( + nlp.evaluator, + vals, + x, + obj_weight, + zeros(nlp.meta.ncon), + ) + return vals +end + +function NLPModels.hprod!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + y::AbstractVector, + v::AbstractVector, + hv::AbstractVector; + obj_weight::Real = 1.0, +) + increment!(nlp, :neval_hprod) + if !nlp.has_hessvec + error( + "The AD backend's evaluator does not support Hessian-vector products.", + ) + end + MOI.eval_hessian_lagrangian_product(nlp.evaluator, hv, x, v, obj_weight, y) + return hv +end + +function NLPModels.hprod!( + nlp::EvaluatorNLPModel, + x::AbstractVector, + v::AbstractVector, + hv::AbstractVector; + obj_weight::Real = 1.0, +) + return NLPModels.hprod!(nlp, x, zeros(nlp.meta.ncon), v, hv, obj_weight = obj_weight) +end diff --git a/src/moi_nlp_model.jl b/src/moi_nlp_model.jl index a4975aa3..b952a013 100644 --- a/src/moi_nlp_model.jl +++ b/src/moi_nlp_model.jl @@ -2,7 +2,7 @@ export MathOptNLPModel mutable struct MathOptNLPModel <: AbstractNLPModel{Float64, Vector{Float64}} meta::NLPModelMeta{Float64, Vector{Float64}} - eval::MOI.Nonlinear.Evaluator + eval::MOI.AbstractNLPEvaluator lincon::LinearConstraints quadcon::QuadraticConstraints nlcon::NonLinearStructure @@ -29,12 +29,17 @@ function MathOptNLPModel(moimodel::MOI.ModelLike; kws...) return nlp_model(moimodel; kws...)[1] end -function nlp_model(moimodel::MOI.ModelLike; hessian::Bool = true, name::String = "Generic") +function nlp_model( + moimodel::MOI.ModelLike; + hessian::Bool = true, + name::String = "Generic", + ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation = _get_ad_backend(moimodel), +) index_map, nvar, lvar, uvar, x0 = parser_variables(moimodel) nlin, lincon, lin_lcon, lin_ucon, quadcon, quad_lcon, quad_ucon = parser_MOI(moimodel, index_map, nvar) - nlp_data = _nlp_block(moimodel) + nlp_data = _nlp_block(moimodel, ad_backend) nlcon = parser_NL(nlp_data, hessian = hessian) oracles = parser_oracles(moimodel) counters = Counters() diff --git a/src/utils.jl b/src/utils.jl index 9013716c..ec434cbc 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -517,10 +517,12 @@ function parser_MOI(moimodel, index_map, nvar) return nlin, lincon, lin_lcon, lin_ucon, quadcon, quad_lcon, quad_ucon end -# Affine or quadratic, nothing to do -_nlp_model(::MOI.Nonlinear.Model, ::MOI.ModelLike, ::Type, ::Type) = false +# Affine or quadratic, nothing to do. +# The first argument is either a `MOI.Nonlinear.Model`, possibly wrapped in +# `MOI.Nonlinear` layers, or a custom AD backend's model, so it is untyped. +_nlp_model(dest, ::MOI.ModelLike, ::Type, ::Type) = false -function _nlp_model(dest::MOI.Nonlinear.Model, src::MOI.ModelLike, F::Type{SNF}, S::Type) +function _nlp_model(dest, src::MOI.ModelLike, F::Type{SNF}, S::Type) has_nonlinear = false for ci in MOI.get(src, MOI.ListOfConstraintIndices{F, S}()) MOI.Nonlinear.add_constraint( @@ -533,8 +535,52 @@ function _nlp_model(dest::MOI.Nonlinear.Model, src::MOI.ModelLike, F::Type{SNF}, return has_nonlinear end -function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model} - nlp_model = MOI.Nonlinear.Model() +# Delegates to `MOI.Nonlinear.model` so that custom AD backends receive the +# model type they declared, wrapped in the layers they did not opt out of. +_nonlinear_model(backend::MOI.Nonlinear.AbstractAutomaticDifferentiation) = MOI.Nonlinear.model(backend) + +""" + _detect_squared_residual(inner) + +Hook for AD extensions: given the `inner` function under a `:sum` root +(i.e., the `?` in `sum(?)`), return the residual whose square sum is being +minimized — typically the first argument of a broadcast `:^` with exponent 2. + +Default returns `nothing` (no NLS routing). Extensions for AD backends that +carry vector-function types (e.g. `ArrayDiff.Mode`) override this. +""" +_detect_squared_residual(::Any) = nothing + +""" + _build_nls_from_residual(moimodel, residual, ad_backend) + +Hook for AD extensions: build an `AbstractNLSModel` that evaluates the given +`residual` (a vector function) using `ad_backend`. Default returns `nothing`, +which makes the optimizer fall back to `MathOptNLPModel`. +""" +_build_nls_from_residual(::Any, ::Any, ::MOI.Nonlinear.AbstractAutomaticDifferentiation) = nothing + +function _try_nls_model( + moimodel::MOI.ModelLike, + ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation, +) + F = MOI.get(moimodel, MOI.ObjectiveFunctionType()) + if !(F <: SNF) + return nothing + end + obj = MOI.get(moimodel, MOI.ObjectiveFunction{F}()) + if obj.head !== :sum || length(obj.args) != 1 + return nothing + end + residual = _detect_squared_residual(obj.args[1]) + if residual === nothing + return nothing + end + return _build_nls_from_residual(moimodel, residual, ad_backend) +end + +function _nlp_model(model::MOI.ModelLike, ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation) + nlp_model = _nonlinear_model(ad_backend) has_nonlinear = false for attr in MOI.get(model, MOI.ListOfModelAttributesSet()) if attr isa MOI.UserDefinedFunction @@ -550,6 +596,11 @@ function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model} if F <: SNF MOI.Nonlinear.set_objective(nlp_model, MOI.get(model, MOI.ObjectiveFunction{F}())) has_nonlinear = true + elseif F <: MOI.AbstractVectorFunction && _supports_vector_objective(ad_backend) + # ArrayNonlinearFunction or similar: return the function directly. + # The ad_backend from the model will build the evaluator. + func = MOI.get(model, MOI.ObjectiveFunction{F}()) + return func end if !has_nonlinear return nothing @@ -557,11 +608,42 @@ function _nlp_model(model::MOI.ModelLike)::Union{Nothing, MOI.Nonlinear.Model} return nlp_model end -function _nlp_block(model::MOI.ModelLike) +""" + _supports_vector_objective(backend::MOI.Nonlinear.AbstractAutomaticDifferentiation) + +Hook for AD-backend extensions: return `true` if the backend can evaluate a +vector-valued objective function (for example, `ArrayDiff.Mode` with an +`ArrayNonlinearFunction` objective). +""" +_supports_vector_objective(::MOI.Nonlinear.AbstractAutomaticDifferentiation) = false + +function _get_ad_backend(model::MOI.ModelLike) + if MOI.supports(model, MOI.AutomaticDifferentiationBackend()) + try + return MOI.get(model, MOI.AutomaticDifferentiationBackend()) + catch err + # For example, a `CachingOptimizer` in state `NO_OPTIMIZER` claims + # support but cannot answer the query. + if !(err isa MOI.GetAttributeNotAllowed) + rethrow() + end + end + end + return MOI.Nonlinear.SparseReverseMode() +end + +function _nlp_block( + model::MOI.ModelLike, + ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation = _get_ad_backend(model), +) # Old interface with `@NL...` - nlp_data = MOI.get(model, MOI.NLPBlock()) + nlp_data = if MOI.NLPBlock() in MOI.get(model, MOI.ListOfModelAttributesSet()) + MOI.get(model, MOI.NLPBlock()) + else + nothing + end # New interface with `@constraint` and `@objective` - nlp_model = _nlp_model(model) + nlp_model = _nlp_model(model, ad_backend) vars = MOI.get(model, MOI.ListOfVariableIndices()) if isnothing(nlp_data) if isnothing(nlp_model) @@ -569,8 +651,7 @@ function _nlp_block(model::MOI.ModelLike) MOI.Nonlinear.Evaluator(MOI.Nonlinear.Model(), MOI.Nonlinear.SparseReverseMode(), vars) nlp_data = MOI.NLPBlockData(evaluator) else - backend = MOI.Nonlinear.SparseReverseMode() - evaluator = MOI.Nonlinear.Evaluator(nlp_model, backend, vars) + evaluator = MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars) nlp_data = MOI.NLPBlockData(evaluator) end else diff --git a/test/runtests.jl b/test/runtests.jl index 7588866f..fc128103 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,6 +23,7 @@ include("test_moi_nlp_model.jl") include("test_moi_nls_model.jl") include("test_moi_nlp_oracle.jl") +include("test_evaluator_model.jl") include("nlp_consistency.jl") include("nls_consistency.jl") diff --git a/test/test_evaluator_model.jl b/test/test_evaluator_model.jl new file mode 100644 index 00000000..25857422 --- /dev/null +++ b/test/test_evaluator_model.jl @@ -0,0 +1,115 @@ +using Test +import ExaModels +import NLPModels +import NLPModelsJuMP +import MathOptInterface as MOI +using JuMP + +@testset "EvaluatorNLPModel with ExaModels.SIMDMode" begin + # hs071-flavoured model with a linear and a quadratic constraint mixed in. + jm = Model() + @variable(jm, 1 <= x[1:4] <= 5) + set_start_value.(x, [1.0, 5.0, 5.0, 1.0]) + @objective(jm, Min, x[1] * x[4] * (x[1] + x[2] + x[3]) + x[3]) + @constraint(jm, c_nl, x[1] * x[2] * x[3] * x[4] >= 25.0) + @constraint(jm, c_q, sum(x[i]^2 for i in 1:4) == 40.0) + @constraint(jm, c_l, x[1] + 2x[2] + 3x[3] <= 30.0) + moi = backend(jm) + mode = ExaModels.SIMDMode() + @test NLPModelsJuMP._route_all_to_evaluator(mode) + nlp, _ = NLPModelsJuMP.evaluator_nlp_model(moi; ad_backend = mode) + @test nlp isa NLPModelsJuMP.EvaluatorNLPModel + meta = nlp.meta + @test meta.nvar == 4 + @test meta.ncon == 3 + @test meta.minimize + # Rows follow the order the constraints appear in the MOI model, with the + # linear row identified by constraint_linearity. + @test length(meta.lin) == 1 + @test length(meta.nln) == 2 + xv = [1.0, 4.9, 5.0, 1.1] + @test NLPModels.obj(nlp, xv) ≈ xv[1] * xv[4] * (xv[1] + xv[2] + xv[3]) + xv[3] + g = NLPModels.grad(nlp, xv) + @test g ≈ [ + xv[4] * (2xv[1] + xv[2] + xv[3]), + xv[1] * xv[4], + xv[1] * xv[4] + 1.0, + xv[1] * (xv[1] + xv[2] + xv[3]), + ] + c = NLPModels.cons(nlp, xv) + expected = Dict( + :nl => prod(xv) - 25.0, + :q => sum(xv .^ 2), + :l => xv[1] + 2xv[2] + 3xv[3], + ) + @test sort(c) ≈ sort([expected[:nl], expected[:q], expected[:l]]) + # The linear row served by the materialized block matches the evaluator. + c_lin = zeros(1) + NLPModels.cons_lin!(nlp, xv, c_lin) + @test c_lin[1] ≈ expected[:l] + c_nln = zeros(2) + NLPModels.cons_nln!(nlp, xv, c_nln) + @test sort(c_nln) ≈ sort([expected[:nl], expected[:q]]) + # Jacobian: dense comparison. + J = NLPModels.jac(nlp, xv) + Jd = Matrix(J) + rows = Dict{Symbol,Int}() + for r in 1:3 + if Jd[r, :] ≈ [1.0, 2.0, 3.0, 0.0] + rows[:l] = r + elseif Jd[r, :] ≈ 2 .* xv + rows[:q] = r + else + rows[:nl] = r + end + end + @test length(rows) == 3 + @test Jd[rows[:nl], :] ≈ [prod(xv) / xv[i] for i in 1:4] # constant does not affect J + @test rows[:l] in meta.lin + # jprod/jtprod agree with the dense Jacobian. + v = [1.0, -2.0, 0.5, 3.0] + @test NLPModels.jprod(nlp, xv, v) ≈ Jd * v + w = [1.0, -1.0, 2.0] + @test NLPModels.jtprod(nlp, xv, w) ≈ Jd' * w + # Hessian of the Lagrangian: dense symmetric comparison against hprod. + y = [2.0, -3.0, 4.0] + σ = 1.5 + H = zeros(4, 4) + hrows, hcols = NLPModels.hess_structure(nlp) + hvals = NLPModels.hess_coord(nlp, xv, y; obj_weight = σ) + for (r, cc, val) in zip(hrows, hcols, hvals) + H[r, cc] += val + if r != cc + H[cc, r] += val + end + end + hv = NLPModels.hprod(nlp, xv, y, v; obj_weight = σ) + @test hv ≈ H * v + # Solve through the MOI wrapper with Percival, comparing against the + # default SparseReverseMode path on a Percival-friendly model. + import Percival + results = Dict{Symbol,Any}() + for (key, backend) in + (:classic => MOI.Nonlinear.SparseReverseMode(), :exa => mode) + jm2 = Model(NLPModelsJuMP.Optimizer) + set_silent(jm2) + set_attribute(jm2, "solver", Percival.PercivalSolver) + set_attribute(jm2, MOI.AutomaticDifferentiationBackend(), backend) + @variable(jm2, 0 <= z[1:3] <= 5, start = 1.0) + @objective(jm2, Min, (z[1] - 1)^2 + (z[2] - 2)^2 + (z[3] - 3)^2) + @constraint(jm2, z[1] + 2z[2] <= 3.0) + @constraint(jm2, z[1] * z[2] >= 0.25) + @constraint(jm2, exp(z[3]) <= 10.0) + optimize!(jm2) + results[key] = + (termination_status(jm2), objective_value(jm2), value.(z)) + end + # Percival's behavior on this problem is environment-dependent (the + # pristine SparseReverseMode path gives the same result), so the test is + # that the ExaModels path agrees exactly with the classic path. + @test results[:classic][1] == results[:exa][1] + @test isequal(results[:classic][2], results[:exa][2]) || + isapprox(results[:classic][2], results[:exa][2]; atol = 1e-4) + @test isequal(results[:classic][3], results[:exa][3]) || + isapprox(results[:classic][3], results[:exa][3]; atol = 1e-3) +end