Compile Times

[willtebbutt]

Compile times aren't completely horrible, but they do seem currently to be a little on the high side for e.g. Turing.jl (@yebai if you have a particular example in mind, please feel free to provide it below).

No effort has yet been made to make compile times good, so there is probably a lot of low-hanging fruit here for someone who knows what they're doing.

[yebai]

Here is one simple example illustrating the overhead from the compiling the AD compiler step, which is subsequently shared by all differentiations. The first call to sample uses roughly 23 seconds, while subsequent calls only use 2.6 seconds on my machine. These two sample calls use the identical model but are defined twice to invalidate possible caches within Tapir, so the extra time is most likely for compiling the AD compiler. Given that this is a generic step, we could use PrecompileTools.jl to coach Tapir cache the compilation result.

julia> using Turing, ADTypes, Mooncake

julia> @model demo1() = x ~ Normal()
demo1 (generic function with 2 methods)

julia> @time sample(demo1(), NUTS(;adtype=AutoMooncake(; config=nothing)), 2000);
 23.154330 seconds (57.13 M allocations: 3.820 GiB, 2.04% gc time, 97.11% compilation time)

julia> @model demo2() = x ~ Normal()
demo2 (generic function with 2 methods)

julia> @time sample(demo2(), NUTS(;adtype=AutoMooncake(; config=nothing)), 2000);
  2.635073 seconds (7.59 M allocations: 514.863 MiB, 2.32% gc time, 88.99% compilation time)

Related issue: #93

[willtebbutt]

Thanks for this Hong.

Re PrecompileTools.jl -- we'll have to do some work to figure out if this is the right approach. It's a really effective tool if there's just a handful of functions that you need to precompile, but at the minute I suspect that we're over-specialising in a lot of places, which would make using PrecompileTools.jl etc less effective. We'll need to first figure out where most of the time is spent, and the figure out what the right solution is.

[yebai]

I'm closing this since the compilation time is often reasonable. Please feel free to reopen if this is observed problematic again.