Add timers to _VectorNonlinearOracle (#502)

Author: odow

ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl

@@ -129,8 +129,6 @@ struct _VectorNonlinearOracle <: MOI.AbstractVectorSet
     eval_jacobian::Function
     hessian_lagrangian_structure::Vector{Tuple{Int,Int}}
     eval_hessian_lagrangian::Union{Nothing,Function}
-    # Temporary storage
-    x::Vector{Float64}
 
     function _VectorNonlinearOracle(;
         dimension::Int,
@@ -154,8 +152,6 @@ struct _VectorNonlinearOracle <: MOI.AbstractVectorSet
             eval_jacobian,
             hessian_lagrangian_structure,
             eval_hessian_lagrangian,
-            # Temporary storage
-            zeros(dimension),
         )
     end
 end
@@ -174,6 +170,18 @@ function Base.show(io::IO, s::_VectorNonlinearOracle)
     return
 end
 
+mutable struct _VectorNonlinearOracleCache
+    set::_VectorNonlinearOracle
+    x::Vector{Float64}
+    eval_f_timer::Float64
+    eval_jacobian_timer::Float64
+    eval_hessian_lagrangian_timer::Float64
+
+    function _VectorNonlinearOracleCache(set::_VectorNonlinearOracle)
+        return new(set, zeros(set.input_dimension), 0.0, 0.0, 0.0)
+    end
+end
+
 """
     Optimizer()
 
@@ -202,7 +210,7 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
     barrier_iterations::Int
     ad_backend::MOI.Nonlinear.AbstractAutomaticDifferentiation
     vector_nonlinear_oracle_constraints::Vector{
-        Tuple{MOI.VectorOfVariables,_VectorNonlinearOracle},
+        Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache},
     }
 
     function Optimizer()
@@ -228,7 +236,7 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
             nothing,
             0,
             MOI.Nonlinear.SparseReverseMode(),
-            Tuple{MOI.VectorOfVariables,_VectorNonlinearOracle}[],
+            Tuple{MOI.VectorOfVariables,_VectorNonlinearOracleCache}[],
         )
     end
 end
@@ -847,7 +855,8 @@ function MOI.add_constraint(
     s::S,
 ) where {F<:MOI.VectorOfVariables,S<:_VectorNonlinearOracle}
     model.inner = nothing
-    push!(model.vector_nonlinear_oracle_constraints, (f, s))
+    cache = _VectorNonlinearOracleCache(s)
+    push!(model.vector_nonlinear_oracle_constraints, (f, cache))
     n = length(model.vector_nonlinear_oracle_constraints)
     return MOI.ConstraintIndex{F,S}(n)
 end
@@ -859,10 +868,10 @@ function row(
     offset = length(model.qp_data)
     for i in 1:(ci.value-1)
         _, s = model.vector_nonlinear_oracle_constraints[i]
-        offset += s.output_dimension
+        offset += s.set.output_dimension
     end
     _, s = model.vector_nonlinear_oracle_constraints[ci.value]
-    return offset .+ (1:s.output_dimension)
+    return offset .+ (1:s.set.output_dimension)
 end
 
 function MOI.get(
@@ -890,7 +899,7 @@ function MOI.get(
     _jacobian_structure(J, 0, f, s)
     J_val = zeros(length(J))
     _eval_constraint_jacobian(J_val, 0, model.inner.x, f, s)
-    dual = zeros(MOI.dimension(s))
+    dual = zeros(MOI.dimension(s.set))
     # dual = λ' * J(x)
     col_to_index = Dict(x.value => j for (j, x) in enumerate(f.variables))
     for ((row, col), J_rc) in zip(J, J_val)
@@ -1184,14 +1193,14 @@ function _eval_constraint(
     offset::Int,
     x::AbstractVector,
     f::MOI.VectorOfVariables,
-    s::_VectorNonlinearOracle,
+    s::_VectorNonlinearOracleCache,
 )
-    for i in 1:s.input_dimension
+    for i in 1:s.set.input_dimension
         s.x[i] = x[f.variables[i].value]
     end
-    ret = view(g, offset .+ (1:s.output_dimension))
-    s.eval_f(ret, s.x)
-    return offset + s.output_dimension
+    ret = view(g, offset .+ (1:s.set.output_dimension))
+    s.eval_f_timer += @elapsed s.set.eval_f(ret, s.x)
+    return offset + s.set.output_dimension
 end
 
 function MOI.eval_constraint(model::Optimizer, g, x)
@@ -1211,12 +1220,12 @@ function _jacobian_structure(
     ret::AbstractVector,
     row_offset::Int,
     f::MOI.VectorOfVariables,
-    s::_VectorNonlinearOracle,
+    s::_VectorNonlinearOracleCache,
 )
-    for (i, j) in s.jacobian_structure
+    for (i, j) in s.set.jacobian_structure
         push!(ret, (row_offset + i, f.variables[j].value))
     end
-    return row_offset + s.output_dimension
+    return row_offset + s.set.output_dimension
 end
 
 function MOI.jacobian_structure(model::Optimizer)
@@ -1241,13 +1250,14 @@ function _eval_constraint_jacobian(
     offset::Int,
     x::AbstractVector,
     f::MOI.VectorOfVariables,
-    s::_VectorNonlinearOracle,
+    s::_VectorNonlinearOracleCache,
 )
-    for i in 1:s.input_dimension
+    for i in 1:s.set.input_dimension
         s.x[i] = x[f.variables[i].value]
     end
-    nnz = length(s.jacobian_structure)
-    s.eval_jacobian(view(values, offset .+ (1:nnz)), s.x)
+    nnz = length(s.set.jacobian_structure)
+    s.eval_jacobian_timer +=
+        @elapsed s.set.eval_jacobian(view(values, offset .+ (1:nnz)), s.x)
     return offset + nnz
 end
 
@@ -1267,9 +1277,9 @@ end
 function _hessian_lagrangian_structure(
     ret::AbstractVector,
     f::MOI.VectorOfVariables,
-    s::_VectorNonlinearOracle,
+    s::_VectorNonlinearOracleCache,
 )
-    for (i, j) in s.hessian_lagrangian_structure
+    for (i, j) in s.set.hessian_lagrangian_structure
         push!(ret, (f.variables[i].value, f.variables[j].value))
     end
     return
@@ -1291,16 +1301,17 @@ function _eval_hessian_lagrangian(
     μ::AbstractVector,
     μ_offset::Int,
     f::MOI.VectorOfVariables,
-    s::_VectorNonlinearOracle,
+    s::_VectorNonlinearOracleCache,
 )
-    for i in 1:s.input_dimension
+    for i in 1:s.set.input_dimension
         s.x[i] = x[f.variables[i].value]
     end
-    H_nnz = length(s.hessian_lagrangian_structure)
+    H_nnz = length(s.set.hessian_lagrangian_structure)
     H_view = view(H, H_offset .+ (1:H_nnz))
-    μ_view = view(μ, μ_offset .+ (1:s.output_dimension))
-    s.eval_hessian_lagrangian(H_view, s.x, μ_view)
-    return H_offset + H_nnz, μ_offset + s.output_dimension
+    μ_view = view(μ, μ_offset .+ (1:s.set.output_dimension))
+    s.eval_hessian_lagrangian_timer +=
+        @elapsed s.set.eval_hessian_lagrangian(H_view, s.x, μ_view)
+    return H_offset + H_nnz, μ_offset + s.set.output_dimension
 end
 
 function MOI.eval_hessian_lagrangian(model::Optimizer, H, x, σ, μ)
@@ -1359,7 +1370,7 @@ function _setup_model(model::Optimizer)
         !isempty(model.vector_nonlinear_oracle_constraints)
     has_hessian = :Hess in MOI.features_available(model.nlp_data.evaluator)
     for (_, s) in model.vector_nonlinear_oracle_constraints
-        if s.eval_hessian_lagrangian === nothing
+        if s.set.eval_hessian_lagrangian === nothing
             has_hessian = false
             break
         end
@@ -1403,8 +1414,8 @@ function _setup_model(model::Optimizer)
     end
     g_L, g_U = copy(model.qp_data.g_L), copy(model.qp_data.g_U)
     for (_, s) in model.vector_nonlinear_oracle_constraints
-        append!(g_L, s.l)
-        append!(g_U, s.u)
+        append!(g_L, s.set.l)
+        append!(g_U, s.set.u)
     end
     for bound in model.nlp_data.constraint_bounds
         push!(g_L, bound.lower)
@@ -1532,6 +1543,12 @@ function MOI.optimize!(model::Optimizer)
         end
         return true
     end
+    # Clear timers
+    for (_, s) in model.vector_nonlinear_oracle_constraints
+        s.eval_f_timer = 0.0
+        s.eval_jacobian_timer = 0.0
+        s.eval_hessian_lagrangian_timer = 0.0
+    end
     Ipopt.SetIntermediateCallback(inner, _moi_callback)
     Ipopt.IpoptSolve(inner)
     model.solve_time = time() - start_time
@@ -1696,7 +1713,7 @@ function row(
 )
     offset = length(model.qp_data)
     for (_, s) in model.vector_nonlinear_oracle_constraints
-        offset += s.output_dimension
+        offset += s.set.output_dimension
     end
     return offset + ci.value
 end

test/MOI_wrapper.jl

@@ -858,6 +858,7 @@ function test_vector_nonlinear_oracle()
         l = [0.0, 0.0],
         u = [0.0, 0.0],
         eval_f = (ret, x) -> begin
+            sleep(0.5)  # to make timing test more robust
             @test length(ret) == 2
             @test length(x) == 5
             ret[1] = x[1]^2 - x[4]
@@ -866,6 +867,7 @@ function test_vector_nonlinear_oracle()
         end,
         jacobian_structure = [(1, 1), (2, 2), (2, 3), (1, 4), (2, 5)],
         eval_jacobian = (ret, x) -> begin
+            sleep(0.5)  # to make timing test more robust
             @test length(ret) == 5
             @test length(x) == 5
             ret[1] = 2 * x[1]
@@ -877,6 +879,7 @@ function test_vector_nonlinear_oracle()
         end,
         hessian_lagrangian_structure = [(1, 1), (2, 2), (3, 3)],
         eval_hessian_lagrangian = (ret, x, u) -> begin
+            sleep(1.0)  # to make timing test more robust
             @test length(ret) == 3
             @test length(x) == 5
             @test length(u) == 2
@@ -916,6 +919,22 @@ function test_vector_nonlinear_oracle()
     @test isapprox(y_v, [x_v[1]^2, x_v[2]^2 + x_v[3]^3], atol = 1e-5)
     @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ [x_v; y_v]
     @test MOI.get(model, MOI.ConstraintDual(), c) ≈ zeros(5)
+    # Test timers with plenty of buffer to avoid a flakey test
+    for (_, cache) in model.vector_nonlinear_oracle_constraints
+        @show cache.eval_f_timer
+        @test 0.9 < cache.eval_f_timer < 2
+        @test 0.9 < cache.eval_jacobian_timer < 4
+        @test 0.9 < cache.eval_hessian_lagrangian_timer < 2
+    end
+    # Test that optimize! resets the timers. Upper bounds are chosen such that
+    # they're violated if times from both solves were added together.
+    MOI.optimize!(model)
+    for (_, cache) in model.vector_nonlinear_oracle_constraints
+        @show cache.eval_f_timer
+        @test 0.9 < cache.eval_f_timer < 2
+        @test 0.9 < cache.eval_jacobian_timer < 4
+        @test 0.9 < cache.eval_hessian_lagrangian_timer < 2
+    end
     return
 end