[OptimizationMOI] improve performance of Jacobian and Hessian accesses

[odow]

cc @ccoffrin

Checklist

• Appropriate tests were added. (Don't know a good way to test this.)
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

The problem is that evaluator.J and evaluator.H in the hot loops involve a call to:

Optimization.jl/lib/OptimizationMOI/src/nlp.jl

Lines 23 to 26 in d6bea20

	function Base.getproperty(evaluator::MOIOptimizationNLPEvaluator, x::Symbol)
	if x in fieldnames(Optimization.ReInitCache)
	return getfield(evaluator.reinit_cache, x)
	end

this is expensive because it involves fieldnames.

Benchmark

julia> import ForwardDiff

julia> import Ipopt

julia> import Optimization

julia> import OptimizationMOI

julia> import Test

julia> function _run_optimization(N, automatic_differentiation)
           h = 1 / N
           alpha = 350
           x_offset = N + 1
           u_offset = 2(N + 1)
           function objective_fn(x, p)
               return sum(
                   0.5 * h * (x[x_offset+i+1]^2 + x[x_offset+i]^2) +
                   0.5 * alpha * h * (cos(x[i+1]) + cos(x[i])) for i in 1:N
               )
           end
           function constraint_fn(res, x, p)
               for i in 1:N
                   res[i] = x[x_offset+i+1] - x[x_offset+i] - 0.5 * h * (sin(x[i+1]) + sin(x[i]))
               end
               for i in 1:N
                   res[N+i] = x[i+1] - x[i] - 0.5 * h * x[u_offset+i+1] - 0.5 * h * x[u_offset+i]
               end
               return
           end
           prob = Optimization.OptimizationProblem(
               Optimization.OptimizationFunction(
                   objective_fn,
                   automatic_differentiation;
                   cons = constraint_fn,
               ),
               zeros(3 * (N + 1)),
               nothing;
               lb = vcat(fill(-1.0, N+1), fill(-0.05, N+1), fill(-Inf, N+1)),
               ub = vcat(fill(1.0, N+1), fill(0.05, N+1), fill(Inf, N+1)),
               lcons = zeros(2 * N),
               ucons = zeros(2 * N),
           )
           sol = Optimization.solve(prob, Ipopt.Optimizer(); print_level = 0)
           Test.@test ≈(sol.objective, 350.0; atol = 1e-6)
           Test.@test ≈(sol.u, zeros(3 * (N + 1)); atol = 1e-6)
           return
       end
_run_optimization (generic function with 1 method)

julia> run_forward_diff(N) = _run_optimization(N, Optimization.AutoForwardDiff())
run_forward_diff (generic function with 1 method)

julia> @time run_forward_diff(100);

******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
 Ipopt is released as open source code under the Eclipse Public License (EPL).
         For more information visit https://github.com/coin-or/Ipopt
******************************************************************************

 13.697718 seconds (9.53 M allocations: 751.142 MiB, 7.05% gc time, 92.11% compilation time: <1% of which was recompilation)

julia> @time run_forward_diff(100);
  0.564060 seconds (2.62 M allocations: 299.192 MiB, 17.04% gc time)

julia> # Revise to go back to old code

julia> @time run_forward_diff(100);
  0.986995 seconds (3.23 M allocations: 319.710 MiB, 4.80% gc time, 14.28% compilation time)

julia> @time run_forward_diff(100);
  0.815753 seconds (3.17 M allocations: 315.911 MiB, 1.76% gc time)

Profile

Before

After

[codecov]

Codecov Report

Attention: 12 lines in your changes are missing coverage. Please review.

Comparison is base (d6bea20) 3.25% compared to head (54085d2) 8.73%.

Files	Patch %	Lines
lib/OptimizationMOI/src/nlp.jl	0.00%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##           master    #680      +/-   ##
=========================================
+ Coverage    3.25%   8.73%   +5.47%     
=========================================
  Files          23      31       +8     
  Lines        1506    2519    +1013     
=========================================
+ Hits           49     220     +171     
- Misses       1457    2299     +842     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[odow]

Format-check seems unrelated. Also failing on master: https://github.com/SciML/Optimization.jl/actions/runs/7493571366/job/20399569099

[odow]

This makes little difference to the Optimization.AutoSparseReverseDiff() case because the main cost is computing the Hessian sparsity

[Vaibhavdixit02]

Yeah the sparsity detection is the main contributor across the board. Even with the MTK implementation that would be the major part. We need a more performant one 😅

[Vaibhavdixit02]

Is there a future where the MOI reverse mode becomes a package? 😄

[odow]

It's just part of MOI, sooo it's already a package?

The main blocker is that you'd need to convert the symbolic form into MOI's expression format. It also operates on a single model with all the constraints, not function-by-function.