reduce allocation NonLinMPC

Author: franckgaga

src/controller/execute.jl

@@ -104,11 +104,12 @@ julia> round.(getinfo(mpc)[:Ŷ], digits=3)
 """
 function getinfo(mpc::PredictiveController{NT}) where NT<:Real
     info = Dict{Symbol, Union{JuMP._SolutionSummary, Vector{NT}, NT}}()
-    Ŷ, x̂, u     = similar(mpc.Ŷop), similar(mpc.estim.x̂), similar(mpc.estim.lastu0)
-    Ŷ, x̂end     = predict!(Ŷ, x̂, u, mpc, mpc.estim.model, mpc.ΔŨ)
+    Ŷ, u        = similar(mpc.Ŷop), similar(mpc.estim.lastu0)
+    x̂, x̂next    = similar(mpc.estim.x̂), similar(mpc.estim.x̂)
+    Ŷ, x̂end     = predict!(Ŷ, x̂, x̂next, u, mpc, mpc.estim.model, mpc.ΔŨ)
     U           = mpc.S̃*mpc.ΔŨ + mpc.T_lastu
     Ȳ, Ū        = similar(Ŷ), similar(U)
-    J           = obj_nonlinprog!(Ȳ, Ū, mpc, mpc.estim.model, Ŷ, mpc.ΔŨ)
+    J           = obj_nonlinprog!(U, Ȳ, Ū, mpc, mpc.estim.model, Ŷ, mpc.ΔŨ)
     info[:ΔU]   = mpc.ΔŨ[1:mpc.Hc*mpc.estim.model.nu]
     info[:ϵ]    = isinf(mpc.C) ? NaN : mpc.ΔŨ[end]
     info[:J]    = J
@@ -297,15 +298,15 @@ function linconstraint!(mpc::PredictiveController, ::SimModel)
 end
 
 @doc raw"""
-    predict!(Ŷ, x̂, _ , mpc::PredictiveController, model::LinModel, ΔŨ) -> Ŷ, x̂end
+    predict!(Ŷ, x̂, _ , _ , mpc::PredictiveController, model::LinModel, ΔŨ) -> Ŷ, x̂end
 
 Compute the predictions `Ŷ` and terminal states `x̂end` if model is a [`LinModel`](@ref).
 
 The method mutates `Ŷ` and `x̂` vector arguments. The `x̂end` vector is used for
 the terminal constraints applied on ``\mathbf{x̂}_{k-1}(k+H_p)``.
 """
 function predict!(
-    Ŷ, x̂, _ , mpc::PredictiveController, ::LinModel, ΔŨ::Vector{NT}
+    Ŷ, x̂, _ , _ , mpc::PredictiveController, ::LinModel, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     # in-place operations to reduce allocations :
     Ŷ .= mul!(Ŷ, mpc.Ẽ, ΔŨ) .+ mpc.F
@@ -315,19 +316,18 @@ function predict!(
 end
 
 @doc raw"""
-    predict!(Ŷ, x̂, u, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
+    predict!(Ŷ, x̂, x̂next, u, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
 
 Compute both vectors if `model` is not a [`LinModel`](@ref). 
     
-The method mutates `Ŷ`, `x̂` and `u` arguments.
+The method mutates `Ŷ`, `x̂`, `x̂next` and `u` arguments.
 """
 function predict!(
-    Ŷ, x̂, u, mpc::PredictiveController, model::SimModel, ΔŨ::Vector{NT}
+    Ŷ, x̂, x̂next, u, mpc::PredictiveController, model::SimModel, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     nu, ny, nd, Hp, Hc = model.nu, model.ny, model.nd, mpc.Hp, mpc.Hc
     u0 = u
     x̂  .= mpc.estim.x̂
-    x̂next = similar(x̂) # TODO: avoid this allocation if possible
     u0 .= mpc.estim.lastu0
     d0  = @views mpc.d0[1:end]
     for j=1:Hp
@@ -346,20 +346,20 @@ function predict!(
 end
 
 """
-    obj_nonlinprog!( _ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ)
+    obj_nonlinprog!(U , _ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ)
 
 Nonlinear programming objective function when `model` is a [`LinModel`](@ref).
 
 The function is called by the nonlinear optimizer of [`NonLinMPC`](@ref) controllers. It can
 also be called on any [`PredictiveController`](@ref)s to evaluate the objective function `J`
-at specific input increments `ΔŨ` and predictions `Ŷ` values.
+at specific input increments `ΔŨ` and predictions `Ŷ` values. It mutates the `U` argument.
 """
 function obj_nonlinprog!(
-    _ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ::Vector{NT}
+    U , _ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     J = obj_quadprog(ΔŨ, mpc.H̃, mpc.q̃) + mpc.p[]
     if !iszero(mpc.E)
-        U = mpc.S̃*ΔŨ .+ mpc.T_lastu
+        U .= mul!(U, mpc.S̃, ΔŨ) .+ mpc.T_lastu
         UE = [U; U[(end - model.nu + 1):end]]
         ŶE = [mpc.ŷ; Ŷ]
         J += mpc.E*mpc.JE(UE, ŶE, mpc.D̂E)
@@ -368,14 +368,15 @@ function obj_nonlinprog!(
 end
 
 """
-    obj_nonlinprog!(Ȳ, Ū, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ)
+    obj_nonlinprog!(U, Ȳ, Ū, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ)
 
 Nonlinear programming objective function when `model` is not a [`LinModel`](@ref). The
 function `dot(x, A, x)` is a performant way of calculating `x'*A*x`. This method mutates
-`Ȳ` and `Ū` vector arguments (output and input setpoint tracking error, respectively).
+`U`, `Ȳ` and `Ū` arguments (input over `Hp`, and output and input setpoint tracking error, 
+respectively).
 """
 function obj_nonlinprog!(
-    Ȳ, Ū, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ::Vector{NT}
+    U, Ȳ, Ū, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     # --- output setpoint tracking term ---
     Ȳ  .= mpc.R̂y .- Ŷ
@@ -384,7 +385,7 @@ function obj_nonlinprog!(
     JΔŨ = dot(ΔŨ, mpc.Ñ_Hc, ΔŨ)
     # --- input over prediction horizon ---
     if !mpc.noR̂u || !iszero(mpc.E)
-        U = mpc.S̃*ΔŨ + mpc.T_lastu
+        U .= mul!(U, mpc.S̃, ΔŨ) .+ mpc.T_lastu
     end
     # --- input setpoint tracking term ---
     if !mpc.noR̂u

src/controller/explicitmpc.jl

@@ -197,7 +197,9 @@ The solution is ``\mathbf{ΔŨ = - H̃^{-1} q̃}``, see [`init_quadprog`](@ref)
 optim_objective!(mpc::ExplicitMPC) = lmul!(-1, ldiv!(mpc.ΔŨ, mpc.H̃_chol, mpc.q̃))
 
 "Compute the predictions but not the terminal states if `mpc` is an [`ExplicitMPC`](@ref)."
-function predict!(Ŷ, x̂, _ , mpc::ExplicitMPC, ::LinModel, ΔŨ::Vector{NT}) where {NT<:Real}
+function predict!(
+    Ŷ, x̂, _ , _ , mpc::ExplicitMPC, ::LinModel, ΔŨ::Vector{NT}
+) where {NT<:Real}
     # in-place operations to reduce allocations :
     Ŷ .= mul!(Ŷ, mpc.Ẽ, ΔŨ) .+ mpc.F
     x̂ .= NaN

src/controller/nonlinmpc.jl

@@ -302,49 +302,55 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
     nu, ny, nx̂, Hp, ng = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp, length(con.i_g)
     # inspired from https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/tips_and_tricks/#User-defined-operators-with-vector-outputs
     Jfunc, gfunc = let mpc=mpc, model=model, ng=ng, nΔŨ=nΔŨ, nŶ=Hp*ny, nx̂=nx̂, nu=nu, nU=Hp*nu
+        Nc = nΔŨ + 3
         last_ΔŨtup_float, last_ΔŨtup_dual = nothing, nothing
-        Ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
-        g_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), nΔŨ + 3)
-        x̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), nΔŨ + 3)
-        u_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nu), nΔŨ + 3)
-        Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
-        Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nU), nΔŨ + 3)
+        Ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
+        U_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
+        g_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, ng), Nc)
+        x̂_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nx̂), Nc)
+        x̂next_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), Nc)
+        u_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu), Nc)
+        Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
+        Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
         function Jfunc(ΔŨtup::JNT...)
             ΔŨ1 = ΔŨtup[begin]
             Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_float
-                x̂, u = get_tmp(x̂_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
+                x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+                u = get_tmp(u_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, mpc, model, ΔŨ)
                 g = get_tmp(g_cache, ΔŨ1)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
-            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
-            return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
+            U, Ȳ, Ū = get_tmp(U_cache, ΔŨ1), get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
+            return obj_nonlinprog!(U, Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
         end
         function Jfunc(ΔŨtup::ForwardDiff.Dual...)
             ΔŨ1 = ΔŨtup[begin]
             Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_dual
-                x̂, u = get_tmp(x̂_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
+                x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+                u = get_tmp(u_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, mpc, model, ΔŨ)
                 g = get_tmp(g_cache, ΔŨ1)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end
-            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
-            return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
+            U, Ȳ, Ū = get_tmp(U_cache, ΔŨ1), get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
+            return obj_nonlinprog!(U, Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
         end
         function gfunc_i(i, ΔŨtup::NTuple{N, JNT}) where N
             ΔŨ1 = ΔŨtup[begin]
             g = get_tmp(g_cache, ΔŨ1)
             if ΔŨtup !== last_ΔŨtup_float
-                x̂, u = get_tmp(x̂_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
                 Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
+                x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+                u = get_tmp(u_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
@@ -354,10 +360,11 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
             ΔŨ1 = ΔŨtup[begin]
             g = get_tmp(g_cache, ΔŨ1)
             if ΔŨtup !== last_ΔŨtup_dual
-                x̂, u = get_tmp(x̂_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
                 Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
+                x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+                u = get_tmp(u_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end

src/estimator/execute.jl

@@ -15,7 +15,7 @@ function remove_op!(estim::StateEstimator, u, ym, d)
 end
 
 @doc raw"""
-    f̂!(x̂next, estim::StateEstimator, model::SimModel, x̂, u, d) -> x̂next
+    f̂!(x̂next, estim::StateEstimator, model::SimModel, x̂, u, d) -> nothing
 
 Mutating state function ``\mathbf{f̂}`` of the augmented model.
 
@@ -33,24 +33,25 @@ function f̂!(x̂next, estim::StateEstimator, model::SimModel, x̂, u, d)
     # `@views` macro avoid copies with matrix slice operator e.g. [a:b]
     @views x̂d, x̂s = x̂[1:model.nx], x̂[model.nx+1:end]
     @views x̂d_next, x̂s_next = x̂next[1:model.nx], x̂next[model.nx+1:end]
-    T = promote_type(eltype(x̂), eltype(u), eltype(d))
+    T = promote_type(eltype(x̂), eltype(u))
     u_us = Vector{T}(undef, model.nu) # TODO: avoid this allocation if possible
     u_us .= u .+ mul!(u_us, estim.Cs_u, x̂s)
     f!(x̂d_next, model, x̂d, u_us, d)
     mul!(x̂s_next, estim.As, x̂s)
-    return x̂next
+    return nothing
 end
+
 "Use the augmented model matrices if `model` is a [`LinModel`](@ref)."
 function f̂!(x̂next, estim::StateEstimator, ::LinModel, x̂, u, d)
     x̂next .= 0
     mul!(x̂next, estim.Â,  x̂, 1, 1)
     mul!(x̂next, estim.B̂u, u, 1, 1)
     mul!(x̂next, estim.B̂d, d, 1, 1)
-    return x̂next
+    return nothing
 end
 
 @doc raw"""
-    ĥ!(ŷ, estim::StateEstimator, model::SimModel, x̂, d) -> ŷ
+    ĥ!(ŷ, estim::StateEstimator, model::SimModel, x̂, d) -> nothing
 
 Mutating output function ``\mathbf{ĥ}`` of the augmented model, see [`f̂!`](@ref).
 """
@@ -59,14 +60,14 @@ function ĥ!(ŷ, estim::StateEstimator, model::SimModel, x̂, d)
     @views x̂d, x̂s = x̂[1:model.nx], x̂[model.nx+1:end]
     h!(ŷ, model, x̂d, d)
     mul!(ŷ, estim.Cs_y, x̂s, 1, 1)
-    return ŷ
+    return nothing
 end
 "Use the augmented model matrices if `model` is a [`LinModel`](@ref)."
 function ĥ!(ŷ, estim::StateEstimator, ::LinModel, x̂, d)
     ŷ .= 0
     mul!(ŷ, estim.Ĉ,  x̂, 1, 1)
     mul!(ŷ, estim.D̂d, d, 1, 1)
-    return ŷ
+    return nothing
 end
 
 

src/estimator/mhe/construct.jl

@@ -992,11 +992,12 @@ function init_optimization!(
     nV̂, nX̂, ng = He*estim.nym, He*estim.nx̂, length(con.i_g)
     # see init_optimization!(mpc::NonLinMPC, optim) for details on the inspiration
     Jfunc, gfunc = let estim=estim, model=model, nZ̃=nZ̃ , nV̂=nV̂, nX̂=nX̂, ng=ng, nx̂=estim.nx̂
+        Nc = nZ̃ + 3
         last_Z̃tup_float, last_Z̃tup_dual = nothing, nothing
-        V̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nV̂), nZ̃ + 3)
-        g_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), nZ̃ + 3)
-        X̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nX̂), nZ̃ + 3)
-        x̄_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), nZ̃ + 3)
+        V̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nV̂), Nc)
+        g_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), Nc)
+        X̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nX̂), Nc)
+        x̄_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), Nc)
         function Jfunc(Z̃tup::JNT...)
             Z̃1 = Z̃tup[begin]
             V̂ = get_tmp(V̂_cache, Z̃1)

src/model/linmodel.jl

@@ -222,7 +222,7 @@ function steadystate!(model::LinModel, u, d)
 end
 
 """
-    f!(xnext, model::LinModel, x, u, d) -> xnext
+    f!(xnext, model::LinModel, x, u, d) -> nothing
 
 Evaluate `xnext = A*x + Bu*u + Bd*d` in-place when `model` is a [`LinModel`](@ref).
 """
@@ -231,20 +231,20 @@ function f!(xnext, model::LinModel, x, u, d)
     mul!(xnext, model.A,  x, 1, 1)
     mul!(xnext, model.Bu, u, 1, 1)
     mul!(xnext, model.Bd, d, 1, 1)
-    return xnext
+    return nothing
 end
 
 
 """
-    h!(y, model::LinModel, x, d) -> y
+    h!(y, model::LinModel, x, d) -> nothing
 
 Evaluate `y = C*x + Dd*d` in-place when `model` is a [`LinModel`](@ref).
 """
 function h!(y, model::LinModel, x, d)
     y .= 0
     mul!(y, model.C,  x, 1, 1)
     mul!(y, model.Dd, d, 1, 1)
-    return y
+    return nothing
 end
 
 typestr(model::LinModel) = "linear"