Allocation-free trunk

[paraynaud]

@dpo, here are the results related to PartitionedVector PartitiallyseparableNLPModels PR.
It shows allocations during solve!(::TrunkSolver).
I compared different kind of NLPModels on arwhead a very partially-separable problem.

using ADNLPModels, NLPModels, NLPModelsModifiers
using OptimizationProblems, OptimizationProblems.PureJuMP, OptimizationProblems.ADNLPProblems
using JSOSolvers
using BenchmarkTools

using PartiallySeparableNLPModels

n=1000
adnlp = ADNLPProblems.arwhead(; n)
lbfgsnlp = LBFGSModel(adnlp)
psenlp = PSENLPModel(adnlp)
# Partitioned structure summary:
#            element functions: ████████████████████ 1999
#   distinct element functions: █⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 3
#    Element statistics:
#         constant: █⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 1               convex: ████████████████████ 999
#           linear: ████████████████████ 999            concave: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
#        quadratic: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0              general: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
#            cubic: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0
#          general: ████████████████████ 999
#   Element function dimensions:                         Variable overlaps:
#              min: ⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 0.0                min: █⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 2.0
#             mean: ███████████████⋅⋅⋅⋅⋅ 1.4992            mean: █⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅⋅ 2.997
#              max: ████████████████████ 2.0                max: ████████████████████ 999.0

trunksolver_ad = TrunkSolver(adnlp)
trunksolver_lbfgs = TrunkSolver(lbfgsnlp)
trunksolver_pse = TrunkSolver(psenlp)

ges_ad = solve!(trunksolverad, adnlp)
ges_lbfgs = solve!(trunksolverlbfgs, lbfgsnlp)
ges_pse = solve!(trunksolver, psenlp)

ges_ad.iter #5
ges_lbfgs.iter #12
ges_pse.iter #7

@allocated solve!(trunksolver_ad, adnlp)       #2372459088
@allocated solve!(trunksolver_lbfgs, lbfgsnlp) #677888016
@allocated solve!(trunksolver_pse, psenlp)     #1928080

@benchmark solve!(trunksolver_ad, adnlp)
# BenchmarkTools.Trial: 5 samples with 1 evaluation.
#  Range (min … max):  1.150 s …   1.314 s  ┊ GC (min … max): 21.65% … 20.28%
#  Time  (median):     1.181 s              ┊ GC (median):    21.47%
#  Time  (mean ± σ):   1.219 s ± 71.363 ms  ┊ GC (mean ± σ):  20.83% ±  0.85%
#   █        ██                                █            █  
#   █▁▁▁▁▁▁▁▁██▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
#   1.15 s         Histogram: frequency by time        1.31 s <
#  Memory estimate: 2.21 GiB, allocs estimate: 12901555.
@benchmark solve!(trunksolver_lbfgs, lbfgsnlp)
# BenchmarkTools.Trial: 16 samples with 1 evaluation.
#  Range (min … max):  301.923 ms … 360.639 ms  ┊ GC (min … max): 19.31% … 20.71%
#  Time  (median):     317.710 ms               ┊ GC (median):    18.33%
#  Time  (mean ± σ):   322.940 ms ±  18.404 ms  ┊ GC (mean ± σ):  18.51% ±  1.42%
#   ▁   ▁▁██        █  ▁    ▁         ▁       ▁  ▁     ▁        ▁
#   █▁▁▁████▁▁▁▁▁▁▁▁█▁▁█▁▁▁▁█▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁█▁▁█▁▁▁▁▁█▁▁▁▁▁▁▁▁█ ▁
#   302 ms           Histogram: frequency by time          361 ms <
#  Memory estimate: 646.48 MiB, allocs estimate: 6116979.
@benchmark solve!(trunksolver_pse, psenlp)
# BenchmarkTools.Trial: 218 samples with 1 evaluation.
#  Range (min … max):  18.153 ms … 48.549 ms  ┊ GC (min … max): 0.00% … 45.41%
#  Time  (median):     21.991 ms              ┊ GC (median):    0.00%
#  Time  (mean ± σ):   22.971 ms ±  4.321 ms  ┊ GC (mean ± σ):  0.78% ±  4.38%
#    ▁▆▆█▄ ▁  ▇█▂▁▆▂▁
#   ▄███████▃▆███████▆██▄▄▄▇▅▄▄▅▃▁▃▃▅▃▃▇▁▁▁▁▁▁▁▃▃▃▅▃▁▁▁▁▃▁▁▁▁▁▄ ▄
#   18.2 ms         Histogram: frequency by time        37.2 ms <
#  Memory estimate: 1.84 MiB, allocs estimate: 90024.

I used PProf to check where the last allocations were, and it comes from the objective evaluation, which is the same as the obj(adnlp).

Overall the results are great.
The last issue may be the performance of algorithms used to perform partial separability detection, and

solver = Dict{Symbol,Function}(
  :PBFGS => ((prob;kwargs...) -> JSOSolvers.trunk(PBFGSNLPModel(prob); monotone=true, max_eval, verbose=0, kwargs...)),
  :PS => ((prob;kwargs...) -> JSOSolvers.trunk(PSNLPModel(prob); monotone=true, max_eval, verbose=0, kwargs...)),
  :LBFGS => ((prob;kwargs...) -> JSOSolvers.trunk(LBFGSModel(prob); monotone=true, max_eval, verbose=0, kwargs...)),
  :TRUNK => ((prob;kwargs...) -> JSOSolvers.trunk(prob; monotone=true, max_eval, verbose=0, kwargs...)),
)

will count the partial separability detection in bmark_solver.
I am working on it.

After the [email protected] release, I will push the changes on the PR mentioned at the start.

PS: @tmigot, it might interest you

[paraynaud]

Iteration profile for n=500
iter_profile.pdf

Time profile for n=500
time_profile.pdf

[dpo]

Do you mean that this line allocates?
https://github.com/JuliaSmoothOptimizers/OptimizationProblems.jl/blob/main/src/ADNLPProblems/arwhead.jl#L8

[paraynaud]

Yes, I checked it with PProf.
See the result of @benchmark for the obj method:

n=100
adnlp = ADNLPProblems.arwhead(; n)
x = rand(n)
@benchmark NLPModels.obj(adnlp, x)
# BenchmarkTools.Trial: 10000 samples with 1 evaluation.
#  Range (min … max):  13.300 μs … 220.300 μs  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     17.700 μs               ┊ GC (median):    0.00%
#  Time  (mean ± σ):   19.992 μs ±   9.997 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%
#   ▄ ▁▆██▅▄▄▄▄▃▂▁ ▂▁                                            ▂
#   █████████████████▇▇▆▆▆▆▆▇▇▆▇▇▆▅▅▄▅▃▄▄▄▃▄▄▂▃▅▄▅▆▆▇▅▆▆▅▃▄▅▄▄▄▅ █
#   13.3 μs       Histogram: log(frequency) by time        64 μs <
#  Memory estimate: 14.02 KiB, allocs estimate: 893.

I talked about it to @tmigot, and he is aware of that behaviour.

[tmigot]

Hey @paraynaud ! I am aware of this problem, but not entirely sure why it allocates :). Could you test something?
I tried:

function f(x)
           n = length(x)
           return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
end
@allocated f(x)
@allocated f(x) # this one returns 16, which is expected

however

function f(x)
           return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
end
@allocated f(x)
@allocated f(x) # this one returns 14320, which is the same as `@allocated obj(adnlp, x)`

So, the allocations would come from the global variable !?

[dpo]

The nobj_eval counter?

[paraynaud]

I tried the same code than you @tmigot :

n = 100
function f(x)
  return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
end

function f2(x)
  n = length(x)
  return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
end

adnlp1 = ADNLPModel(f, rand(n))
adnlp2 = ADNLPModel(f2, rand(n))

@code_warntype NLPModels.obj(adnlp1, adnlp1.meta.x0)
@code_warntype NLPModels.obj(adnlp2, adnlp2.meta.x0)

@code_warntype adnlp1.f(adnlp1.meta.x0)
@code_warntype adnlp2.f(adnlp2.meta.x0)

x = rand(n)

@benchmark adnlp1.f(x) # Memory estimate: 14.55 KiB, allocs estimate: 895.
@benchmark adnlp2.f(x) # Memory estimate: 592 bytes, allocs estimate: 3.

@benchmark NLPModels.obj(adnlp1, x) #  Memory estimate: 13.98 KiB, allocs estimate: 893.
@benchmark NLPModels.obj(adnlp2, x) #  Memory estimate: 16 bytes, allocs estimate: 1.

from which I got the same result.

[paraynaud]

After 30min, I questionned the scope of n in the definition of arwhead, and google comforted me in that way: see this discussion, where the scope of a variable in nested functions is not separated.
So, I tested a new way to define arwhead, where this time the nested function f doesn't use n, see below:

n = 100
x = rand(n)

# equal to default arwhead
function my_arwhead1(; n::Int = default_nvar, type::Val{T} = Val(Float64), kwargs...) where {T}
  n < 2 && @warn("arwhead: number of variables must be ≥ 2")
  n = max(2, n)
  function f(x)
    n = length(x)
    return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
  end
  x0 = ones(T, n)
  return ADNLPModels.ADNLPModel(f, x0, name = "arwhead"; kwargs...)
end

# replacing n by _n in f
function my_arwhead2(; n::Int = default_nvar, type::Val{T} = Val(Float64), kwargs...) where {T} 
  n < 2 && @warn("arwhead: number of variables must be ≥ 2")
  n = max(2, n)
  function f(x)
    _n = length(x)
    return sum((x[i]^2 + x[_n]^2)^2 - 4 * x[i] + 3 for i = 1:(_n - 1))
  end
  x0 = ones(T, n)
  return ADNLPModels.ADNLPModel(f, x0, name = "arwhead"; kwargs...)
end

adnlp1 = my_arwhead1(;n)
adnlp2 = my_arwhead2(;n)

@benchmark NLPModels.obj(adnlp1, x) # Memory estimate: 14.02 KiB, allocs estimate: 893.
@benchmark NLPModels.obj(adnlp2, x) # Memory estimate: 16 bytes, allocs estimate: 1.

I think that solves our issue.
Tell me if you can reproduce it!

[tmigot]

I see... So, it would also be sufficient to do local n = length(x) in the function definition or do something like

function f(x; n = n)
  return sum((x[i]^2 + x[n]^2)^2 - 4 * x[i] + 3 for i = 1:(n - 1))
end

[tmigot]

Thanks @paraynaud ! Could you open a PR to fix this for arwhead in OptimizationProblems.jl? I will prepare a script to check the other problems.

[paraynaud]

You're welcome! (it helps me too :p)

What do you prefer

function f(x; n=n)
  ...
end

or

function f(x; n=length(x))
  ...
end

?

[tmigot]

The second one looks more generic.

[dpo]

That's a trap! Thanks for finding that!