Precompilation in the SnoopPrecompile/Julia 1.9+ world

[timholy]

Because of its popularity, I'd like to use JuMP as one of several "showcases" for the impact of pkgimages in Julia 1.9 (CC @vchuravy, @vtjnash, @KristofferC). It turns out that to really showcase this work, JuMP might need a few tweaks. Normally I just submit PRs, but due to the "orthogonality" of solvers, JuMP presents some interesting challenges, and so I decided to instead open this issue.

JuMP and its ecosystem have had a lot of nice work done on invalidations and precompilation already, and these lay the foundation and make everything I'm about to show much easier. (Thanks!) The main remaining gap is due, I think, to the fact that the precompilation work occurred before the arrival of SnoopPrecompile and pkgimages in Julia 1.9.

Let me begin by showing that there's opportunity for substantial further improvement. All tests were conducted on a reasonably up-to-date Julia master:

julia> @time using JuMP, GLPK
  7.286941 seconds (9.35 M allocations: 604.153 MiB, 4.49% gc time, 0.79% compilation time)

julia> @time @eval begin
               let
                   model = Model(GLPK.Optimizer)
                   @variable(model, x >= 0)
                   @variable(model, 0 <= y <= 3)
                   @objective(model, Min, 12x + 20y)
                   @constraint(model, c1, 6x + 8y >= 100)
                   @constraint(model, c2, 7x + 12y >= 120)
                   optimize!(model)
               end
           end;
  8.465601 seconds (10.10 M allocations: 1.385 GiB, 4.05% gc time, 99.78% compilation time)

Now, let me create a new package, StartupJuMP, purely for the purpose of extra precompilation. (That's a viable strategy on Julia 1.8 and higher, with the big impact arriving in Julia 1.9.) The source code is a single file, src/StartupJuMP.jl, with contents:

module StartupJuMP

using GLPK
using JuMP
using SnoopPrecompile

@precompile_all_calls begin
    # Because lots of the work is done by macros, and macros are expanded
    # at lowering time, not much of this would get precompiled without `@eval`
    @eval begin
        let
            model = Model(GLPK.Optimizer)
            @variable(model, x >= 0)
            @variable(model, 0 <= y <= 3)
            @objective(model, Min, 12x + 20y)
            @constraint(model, c1, 6x + 8y >= 100)
            @constraint(model, c2, 7x + 12y >= 120)
            optimize!(model)
        end
    end
end

end # module StartupJuMP

Now:

julia> @time using JuMP, GLPK, StartupJuMP
  6.297161 seconds (9.80 M allocations: 630.934 MiB, 4.43% gc time, 0.35% compilation time)

julia> @time @eval begin
               let
                   model = Model(GLPK.Optimizer)
                   @variable(model, x >= 0)
                   @variable(model, 0 <= y <= 3)
                   @objective(model, Min, 12x + 20y)
                   @constraint(model, c1, 6x + 8y >= 100)
                   @constraint(model, c2, 7x + 12y >= 120)
                   optimize!(model)
               end
           end;
  0.331432 seconds (154.94 k allocations: 10.237 MiB, 97.93% compilation time)

You can see a decrease in load time, which TBH I find puzzling (I expected a modest increase). But more importantly, you see a massive decrease in time-to-first-execution (TTFX), to the point where TTFX just doesn't feel like a problem anymore. And keep in mind that this is on top of all the nice work the JuMP ecosystem has already done to minimize TTFX: this small SnoopPrecompile workload improves the quality of precompilation substantially.

Now, ordinarily I'd just suggest putting that @precompile_all_calls block in JuMP. However, a big issue is that this precompilation workload relies on GLPK, and I'm guessing you don't want to precompile GLPK "into" JuMP. I'm not sufficiently familiar with JuMP to pick a good alternative, but in very rough terms I imagine there are at least three potential paths, which might be used alone or in combination:

• run this precompile workload with some kind of default/dummy solver, which doesn't actually perform optimization but at least allows completion (optimize! might return nothing, for example). This might (?) precompile a lot of the machinery.
• identify the missing precompiles, and add them to the current precompile code. One issue, though, is that this seems a bit more fragile to internal Julia changes than using SnoopPrecompile. For instance, this precompile directive will cause precompilation failure on Julia 1.9, because Julia 1.9 will eliminate kwfuncs. The advantage of SnoopPrecompile is that you write everything in terms of the public interface and dispatch will generate the right precompiles on each different version of Julia.
• use the upcoming weakdeps infrastructure to create the analog of my StartupJuMP for each solver. The advantage of weakdeps compared to similar solutions like Requires.jl is that the "extension packages" get precompiled and cached.

I'm happy to help, but given the issues I think it would be ideal if a JuMP developer helped choose the approach and shepherd the changes through.

[blegat]

Thanks ! I would suggest a mix of the first and third option (but without the weakdeps).
The methods compiled are of two types:

JuMP methods

Methods with JuMP.Model and other JuMP types in the signature. These do not depend on the type of the solver since the type JuMP.Model is not parametrized by the solver type so these method should get precompiled by a script with a dummy solver Model(() -> MOI.Utilities.MockOptimizer(MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())).

MOI methods

Methods with MOI.CachingOptimizer{...,MOI.LazyBridgeOptimizer{SolverType}. These are methods called on the MOI backend of the JuMP model for which the signature depends on the SolverType. For these, could build exactly the same MOI model that JuMP builds with:

optimizer = MOI.instantiate(GLPK.Optimizer; with_bridge_type = Float64)
cache = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())
moi_backend = MOI.Utilities.CachingOptimizer(cache, optimizer)

and precompiles the methods JuMP uses, e.g., by running some of the tests in MOI.Test. This would miss the JuMP function called on the MOI backend (usually these are the _moi_...) functions. We could get these with the weakdeps infrastructure, I would be interested to know how much of the TTFX is gained by this part, it might only be a small fractions as these are quite simple functions usually.

[timholy]

Thanks! I have a starter PR in jump-dev/JuMP.jl#3193. Feel free to add to it or suggest improvements.

[timholy]

My current question is "what else needs to be added after jump-dev/JuMP.jl#3193?" Below I pose some questions that I can't really answer myself, I'm hoping that others can. But more generally, since I won't have time or expertise to improve everything myself, it's probably best to show you how I go about this kind of analysis. So I do the following:

using JuMP, GLPK
using SnoopCompileCore
tinf = @snoopi_deep @eval begin
    whatever_example_I_want_to_be_fast
end;
using SnoopCompile
using ProfileView
ProfileView.view(flamegraph(tinf))

For this case (and again, already exploiting jump-dev/JuMP.jl#3193) I get a graph like this:

This page describes the elements & interaction you can use, but briefly:

• time runs horizontally, call-depth vertically
• each "stack" represents a separate runtime-dispatched entrance into type inference
• to see the method that was being inferred upon entrance to inference, click on the bottom of bars (left-click prints to REPL, right-click opens in your ENV["EDITOR"])
• fat bars mean "expensive type inference"
• empty spaces typically mean LLVM codegen (the wider the gap, the longer LLVM is taking). @snoopi_deep doesn't capture data on what LLVM is doing (though @snoopl can). Usually you can make some guesses about what's being codegened by which inference bars came right before, so I typically don't use @snoopl.
• red bars indicate a combination of types not available in the package owning the method being optimized. E.g., anything GLPK-dependent will be red. If above the red you get non-red, that indicates an opportunity to precompile everything above it.

Some interesting findings:

• Early in the trace, MutableArithmetics._rewrite(::Bool, ::Bool, ::Expr, ::Nothing, ::Vector{Any}, ::Vector{Any}, ::Symbol) appears more than once (noteworthy because typically we only need to infer once) and there are yellow bars above it. This indicates that it's being inferred multiple times due to constant-propagation, i.e., Julia is specializing and compiling separate code for, e.g., minus=false versus minus=true. I am guessing you don't want that: while the inference time is small (the bars are skinny), the LLVM time after each of these is nontrivial. You might consider using Base.@constprop :none, or Compat.@constprop :none if you need to support older versions of Julia, on that function. Is that something you want?
• consider whether you might want to put @nospecialize around some ::Type and ::Function arguments. If the code needs to be specialized for performance reasons, that would be a bad idea, but I'm guessing that much of JuMP is about problem setup (unlikely to be performance sensitive) and not, e.g., about evaluating the objective function as quickly as possible (certainly performance sensitive). But I could easily be wrong (this is a place where expertise with the package is really essential). If they are safe to @nospecialize, as a bonus, then you can precompile them once (from JuMP) with a generic optimizer and not have to worry about specializing on each solver. Some examples of functions that might be candidates showing up in this analysis are Model(optimizer_factory; add_bridges::Bool = true), function instantiate(optimizer_constructor; with_bridge_type::Union{Nothing,Type} = nothing), etc.
• the "fattest" (most time-consuming) stack of inference bars is MOI.optimize!(m::CachingOptimizer). Presumably that would have to be handled by weakdeps mechanism? Note, though, that eventually a lot of bars reach non-red; e.g., the widest of the non-red bars is for MathOptInterface.get(::MathOptInterface.Utilities.UniversalFallback{MathOptInterface.Utilities.Model{Float64}}, ::MathOptInterface.ListOfConstraintIndices{MathOptInterface.VectorOfVariables}). Does the MockOptimizer not reach these? If so, what can we can add as a workload to MOI to incorporate them?
• the largest gap (presumably most LLVM time) is right after JuMP._moi_add_variable(::MathOptInterface.Utilities.CachingOptimizer{MathOptInterface.Bridges.LazyBridgeOptimizer{GLPK.Optimizer}, MathOptInterface.Utilities.UniversalFallback{MathOptInterface.Utilities.Model{Float64}}}, ::Model, ::ScalarVariable{Int64, Float64, Float64, Float64}, ::String). Is that another weakdeps case?

[odow]

but I'm guessing that much of JuMP is about problem setup (unlikely to be performance sensitive) and not, e.g., about evaluating the objective function as quickly as possible (certainly performance sensitive). But I could easily be wrong (this is a place where expertise with the package is really essential)

There's a trade-off here that I played around with a while ago. For small problems, it doesn't matter. But we also want to support building problems with 10^6+ variables and constraints, and then it really pays to add the explicit ::F) where {F<:Function} in parts of MOI.

Some examples of functions that might be candidates showing up in this analysis are

Yeah, these ones could be changed.

the "fattest" (most time-consuming) stack of inference bars is MOI.optimize!(m::CachingOptimizer). Presumably that would have to be handled by weakdeps mechanism?

The right place to add precompilation for these sorts of thing is either in MOI, or in the solvers.

Is that another weakdeps case?

My preference is not to add any weakdeps to the JuMP ecosystem just yet, until a few Julia releases have happened and we can assess how it works. Maintaining them in JuMP for a bunch of solvers seems like a bit of work, especially if we can get the TTFX down via other approaches.

[odow]

Okay, question. How should I setup precompile statements for solvers which set a global constant in __init__?

For example:

https://github.com/jump-dev/HiGHS.jl/blob/9ec05c31c3e73c02661f4de0d0def5b3b664a991/src/HiGHS.jl#L12-L15

Is it okay to do something like this?

import SnoopPrecompile

SnoopPrecompile.@precompile_setup begin
    SnoopPrecompile.@precompile_all_calls begin
        __init__()
        model = MOI.instantiate(HiGHS.Optimizer; with_bridge_type = Float64)
    end
end

[odow]

I have it working for HiGHS: jump-dev/HiGHS.jl#147.

There's just one problem left:

julia> using JuMP, HiGHS

julia> using SnoopCompileCore

julia> tinf = @snoopi_deep @eval begin
         model = Model(HiGHS.Optimizer)
         @variable(model, x >= 0)
         @variable(model, 0 <= y <= 3)
         @objective(model, Min, 12x + 20y)
         @constraint(model, c1, 6x + 8y >= 100)
         @constraint(model, c2, 7x + 12y >= 120)
         optimize!(model)
       end;
Running HiGHS 1.4.0 [date: 1970-01-01, git hash: bcf6c0b22]
Copyright (c) 2022 ERGO-Code under MIT licence terms
Presolving model
2 rows, 2 cols, 4 nonzeros
2 rows, 2 cols, 4 nonzeros
Presolve : Reductions: rows 2(-0); columns 2(-0); elements 4(-0) - Not reduced
Problem not reduced by presolve: solving the LP
Using EKK dual simplex solver - serial
  Iteration        Objective     Infeasibilities num(sum)
          0     0.0000000000e+00 Pr: 2(220) 0s
          2     2.0500000000e+02 Pr: 0(0) 0s
Model   status      : Optimal
Simplex   iterations: 2
Objective value     :  2.0500000000e+02
HiGHS run time      :          0.00

julia> using SnoopCompile

julia> using ProfileView

julia> ProfileView.view(flamegraph(tinf))

The red bars are because of

MathOptInterface.jl/src/Utilities/copy.jl

Lines 487 to 490 in 40b81c5

	Any[
	_try_constrain_variables_on_creation(dest, src, index_map, S)
	for S in sorted_variable_sets_by_cost(dest, src)
	]

which calls

MathOptInterface.jl/src/Utilities/copy.jl

Lines 165 to 173 in 40b81c5

	function _try_constrain_variables_on_creation(
	dest::MOI.ModelLike,
	src::MOI.ModelLike,
	index_map::IndexMap,
	::Type{S},
	) where {S<:MOI.AbstractVectorSet}
	not_added = MOI.ConstraintIndex{MOI.VectorOfVariables,S}[]
	for ci_src in
	MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}())

but because sorted_variable_sets_by_cost(dest, src) is a Vector{<:Type}, it fails to infer and tries to compile

/Users/oscar/.julia/dev/MathOptInterface/src/Utilities/universalfallback.jl:445, MethodInstance for MathOptInterface.get(::MathOptInterface.Utilities.UniversalFallback{MathOptInterface.Utilities.Model{Float64}}, ::MathOptInterface.ListOfConstraintIndices{MathOptInterface.VectorOfVariables})

despite the fact that this method doesn't exist and won't get called at runtime. We also can't annotate the type, and adding mixtures of precompile directives to MOI and HiGHS didn't seem to fix the problem. Any ideas on how to resolve?

[odow]

I'm about to add a PR to MOI that gets it down to:

Nice progress!

It brings HiGHS down to

julia> @time @eval begin
           let
               model = Model(HiGHS.Optimizer)
               @variable(model, x >= 0)
               @variable(model, 0 <= y <= 3)
               @objective(model, Min, 12x + 20y)
               @constraint(model, c1, 6x + 8y >= 100)
               @constraint(model, c2, 7x + 12y >= 120)
               optimize!(model)
           end
       end;
Running HiGHS 1.4.0 [date: 1970-01-01, git hash: bcf6c0b22]
Copyright (c) 2022 ERGO-Code under MIT licence terms
Presolving model
2 rows, 2 cols, 4 nonzeros
2 rows, 2 cols, 4 nonzeros
Presolve : Reductions: rows 2(-0); columns 2(-0); elements 4(-0) - Not reduced
Problem not reduced by presolve: solving the LP
Using EKK dual simplex solver - serial
  Iteration        Objective     Infeasibilities num(sum)
          0     0.0000000000e+00 Pr: 2(220) 0s
          2     2.0500000000e+02 Pr: 0(0) 0s
Model   status      : Optimal
Simplex   iterations: 2
Objective value     :  2.0500000000e+02
HiGHS run time      :          0.00
  0.689926 seconds (465.37 k allocations: 31.314 MiB, 98.62% compilation time)

(Considering we started at ~6seconds before, and even more before you made the change to JuMP.)

[odow]

With jump-dev/JuMP.jl#3195, the graph is now

[odow]

The _rewrite bits would be greatly improved by jump-dev/JuMP.jl#3125:

and it gets us to less than 0.5s:

julia> @time @eval begin
           let
               model = Model(HiGHS.Optimizer)
               @variable(model, x >= 0)
               @variable(model, 0 <= y <= 3)
               @objective(model, Min, 12x + 20y)
               @constraint(model, c1, 6x + 8y >= 100)
               @constraint(model, c2, 7x + 12y >= 120)
               optimize!(model)
           end
       end;
Running HiGHS 1.4.0 [date: 1970-01-01, git hash: bcf6c0b22]
Copyright (c) 2022 ERGO-Code under MIT licence terms
Presolving model
2 rows, 2 cols, 4 nonzeros
2 rows, 2 cols, 4 nonzeros
Presolve : Reductions: rows 2(-0); columns 2(-0); elements 4(-0) - Not reduced
Problem not reduced by presolve: solving the LP
Using EKK dual simplex solver - serial
  Iteration        Objective     Infeasibilities num(sum)
          0     0.0000000000e+00 Pr: 2(220) 0s
          2     2.0500000000e+02 Pr: 0(0) 0s
Model   status      : Optimal
Simplex   iterations: 2
Objective value     :  2.0500000000e+02
HiGHS run time      :          0.00
  0.483362 seconds (309.34 k allocations: 20.749 MiB, 98.07% compilation time)

So I'm close to calling this a win.

The left red tower is the _moi_add_variable which makes sense. That's the JuMP->HiGHS transition point that we can't get without a weakdep. The right tower is the _try_add_constrained_variables method which is pretty horrible. There are probably some more small improvements that could be made.