marius311 / CMBLensing.jl

The automatically differentiable and GPU-compatible toolkit for CMB analysis.
https://cosmicmar.com/CMBLensing.jl
Other
52 stars 10 forks source link

Adds precompile statements for way faster startup times #82

Closed marius311 closed 1 year ago

marius311 commented 1 year ago

On 1.8 w/o this branch:

julia> @time using CMBLensing
 13.048558 seconds (30.55 M allocations: 1.851 GiB, 4.54% gc time, 24.28% compilation time: 36% of which was recompilation)

julia> (;f, ϕ, ds) = @time load_sim(;pol=:P, Nside=16, θpix=1);
 24.746780 seconds (74.11 M allocations: 3.693 GiB, 2.86% gc time, 91.50% compilation time: 14% of which was recompilation)

julia> @time logpdf(ds; f, ϕ)
  2.559928 seconds (7.86 M allocations: 389.767 MiB, 2.87% gc time, 99.65% compilation time: 25% of which was recompilation)
7929.795f0

julia> @time gradient((f, ϕ) -> logpdf(ds; f, ϕ), f, ϕ);
 28.532642 seconds (72.92 M allocations: 3.707 GiB, 3.82% gc time, 99.88% compilation time)

On 1.9 w/ this branch:

julia> @time using CMBLensing
 14.241679 seconds (28.90 M allocations: 1.816 GiB, 6.89% gc time, 0.90% compilation time: 100% of which was recompilation)

julia> (;f, ϕ, ds) = @time load_sim(;pol=:P, Nside=16, θpix=1);
  2.841666 seconds (3.06 M allocations: 231.382 MiB, 4.16% gc time)

julia> @time logpdf(ds; f, ϕ)
  0.018066 seconds (3.31 k allocations: 305.770 KiB)
7932.341f0

# gradient gets some boost but not as much
julia> @time gradient((f, ϕ) -> logpdf(ds; f, ϕ), f, ϕ)
 17.620470 seconds (27.34 M allocations: 1.727 GiB, 7.52% gc time)

It does massively increase the precompile time, which in theory is a one-time cost when installing the package but in practice since this undergoes so much development is a bit of a negative. Need to explore ways to mitigate this.

julia> @time using CMBLensing
[ Info: Precompiling CMBLensing [b60c06c0-7e54-11e8-3788-4bd722d65317]
668.719248 seconds (28.94 M allocations: 1.819 GiB, 0.16% gc time, 0.02% compilation time: 100% of which was recompilation)