[UnoSolver.jl] Re-use existing AD data structure on re-solve for efficient `set_normalized_rhs` calls

[franckgaga]

Apply logics for affine constraints similar to https://github.com/jump-dev/Ipopt.jl/pull/516, to support efficient JuMP.set_normalized_rhs calls.

Long story short, modifying the constant $b$ vector of affine constraints $Ax \le b$ with JuMP.set_normalized_rhs will invalidate the inner model, thus requiring re-initialization of the automatic differentiation engine, which is not necessary nor efficient. Here's a not-so-minimal reproducible example:

using ModelPredictiveControl, JuMP
using UnoSolver

N = 35 # number of JuMP.optimize! calls

function f!(ẋ, x, u, _ , p)
    g, L, K, m = p          # [m/s²], [m], [kg/s], [kg]
    θ, ω = x[1], x[2]       # [rad], [rad/s]
    τ  = u[1]               # [Nm]
    ẋ[1] = ω
    ẋ[2] = -g/L*sin(θ) - K/m*ω + τ/m/L^2
    return nothing
end
h!(y, x, _ , _ ) = (y[1] = 180/π*x[1]; nothing) # [°]
p = [9.8, 0.4, 1.2, 0.3]
nu, nx, ny, Ts = 1, 2, 1, 0.1
model = NonLinModel(f!, h!, Ts, nu, nx, ny; p)
p_plant = copy(p); p_plant[3] = p[3]*1.25
plant = NonLinModel(f!, h!, Ts, nu, nx, ny; p=p_plant)
Hp, Hc, Mwt, Nwt = 20, 2, [0.5], [2.5]
α=0.01; σQ=[0.1, 1.0]; σR=[5.0]; nint_u=[1]; σQint_u=[0.1]
σQint_ym = zeros(0)
umin, umax = [-1.5], [+1.5]
transcription = MultipleShooting()
optim = Model(() -> UnoSolver.Optimizer(preset="filtersqp")); 
oracle = true
hessian = true
nmpc = NonLinMPC(model; 
    Hp, Hc, Mwt, Nwt, Cwt=Inf, transcription, oracle, hessian, optim,
    α, σQ, σR, nint_u, σQint_u, σQint_ym
)
nmpc = setconstraint!(nmpc; umin, umax)
unset_time_limit_sec(nmpc.optim)
sim!(nmpc, N, [180.0]; plant=plant, x_0=[0, 0], x̂_0=[0, 0, 0])
@time sim!(nmpc, N, [180.0]; plant=plant, x_0=[0, 0], x̂_0=[0, 0, 0])
@profview sim!(nmpc, N, [180.0]; plant=plant, x_0=[0, 0], x̂_0=[0, 0, 0])

We can see a significant initialize section in the profile (the "useless" re-initialization of the AD engine):

This portion will completely diseapear of the flame graph if something similar to https://github.com/jump-dev/Ipopt.jl/pull/516 is applied here.

[cvanaret]

Thanks for the tip @franckgaga! Uno doesn't support re-solve yet, it's been on my to-do list for a while (in particular in the context of MINLP), but I have other fish to fry at the moment :)

[franckgaga]

Great! But not sure what you meant by "doesn't support re-solve yet". Re-solving a JuMP.Model instantiated with UnoSolver.jl works well, at least for NonLinMPC objects from ModelPredictiveControl.jl.

In this specific context, the only things that needs to be updated before re-solving is the constant $b$ vector in the linear inequality constraints (the RHS of $Ax \le b$). And, similarly to Ipopt.jl, the updating does work but it is just inefficient right now.

[cvanaret]

Mh I thought so because Uno stores the model + some reformulations of it internally. But IIRC everything is stored as references/pointers, so it may well be that modifying the RHS on the Julia side automatically updates the C++ model 👍 Of course, that doesn't hold anymore if e.g. you add constraints because the dimensions of the problem change.

[franckgaga]

Yes, it seems that it just automatically work in Julia since the values at the C++ pointers are automatically updated.

Of course, that doesn't hold anymore if e.g. you add constraints because the dimensions of the problem change.

Of course. Here the number of linear constraints are constant between the re-solves.

[cvanaret]

On second thoughts: Uno is actually solving the subsequent models from scratch. This is because each call to uno_optimize allocates fresh memory and computes the reformulations of the model. So not quite "re-solving" :)

I think I need to decouple passing the model to the solver, and solving it (like HiGHS does: https://ergo-code.github.io/HiGHS/dev/interfaces/cpp/library/). This way, we'd have a finer control over things.

[franckgaga]

I don't know if the MathOptInterface codebase in UnoSolver.jl is more similar to Ipopt.jl or KNITRO.jl codebase, but, for the records, a similar solution was implemented in KNITRO.jl yesterday. The changes are implemented in this PR: https://github.com/jump-dev/KNITRO.jl/pull/378

[odow]

It's more similar to Ipopt. The wrapper could cache the nonlinear Evaluator to avoid MOI.initialize, but Charlie is saying that there will still be some overhead within Uno (just like there is with Ipopt).

[franckgaga]

I think the best approach to performance improvements is incremental. "Real re-solving" (without re-allocating) may be a long-term goal, but the little tweak mentioned by @odow would be a first step in reducing the overhead when re-calling optimize! on a model with UnoSolver.jl and set_normilized_rhs calls.

edit: BTW the re-allocation overhead seems to be quite small here. The time spent in Uno C++ code (including the re-allocations) is the dark grey region in the flame graph above. It's quite low in terms of runtime (less than half the runtime of Ipopt.jl) and, clearly, the initialize portion dominates the flame graph above.

[franckgaga]

Hello @amontoison, does it seems like a big task ? Is there anyway that I can help (even If I have almost no knowledge in MOI internals XP) ?

I would like to compare performances to Ipopt.jl on other problems like moving horizon estimator, but it's no longer significantly competitive to Ipopt, because of the useless re-initialization above triggered by set_normalized_rhs.

[amontoison]

It is no a lot of work but I just came back to Chicago yesterday and I need to do other things before.

I will try to do it one evening of the week.