search

CliMA / ClimaTimeSteppers.jl

A CPU- and GPU-friendly package for solving ordinary differential equations
Apache License 2.0
50 stars 5 forks source link

Compile fewer methods #320

Open charleskawczynski opened 3 weeks ago

charleskawczynski commented 3 weeks ago

Inspired by , this PR attempts to reduce the number of compile methods (and hopefully overall compilation time) by using Base.Cartesian utilities:

using Revise; include("perf/jet.jl")
using MethodAnalysis
import ClimaTimeSteppers as CTS
length(methodinstances(CTS.update_stage!))
length(methodinstances(CTS.unrolled_foreach))

Main

julia> using Revise; include("perf/jet.jl")
Precompiling ClimaTimeSteppers
  1 dependency successfully precompiled in 3 seconds. 180 already precompiled.
step_allocs = 12768
Test Summary:     | Pass  Total  Time
JET / allocations |    2      2  3.1s
Test.DefaultTestSet("JET / allocations", ^R
^[[Any[], 2, false, false, true, 1.728926011863406e9, 1.728926014999377e9, false, "/Users/charliekawczynski/Dropbox/Caltech/work/dev/CliMA/ClimaTimeSteppers.jl/perf/jet.jl")

julia> using MethodAnalysis

julia> import ClimaTimeSteppers as CTS

julia> length(methodinstances(CTS.update_stage!))
5

julia> length(methodinstances(CTS.unrolled_foreach))
10

This branch

julia> using Revise; include("perf/jet.jl")
Precompiling ClimaTimeSteppers
  1 dependency successfully precompiled in 5 seconds. 180 already precompiled.
step_allocs = 12768
Test Summary:     | Pass  Total  Time
JET / allocations |    2      2  3.7s
Test.DefaultTestSet("JET / allocations", Any[], 2, false, false, true, 1.728925925530902e9, 1.728925929182172e9, false, "/Users/charliekawczynski/Dropbox/Caltech/work/dev/CliMA/ClimaTimeSteppers.jl/perf/jet.jl")
julia> using MethodAnalysis

julia> import ClimaTimeSteppers as CTS

julia> length(methodinstances(CTS.update_stage!))
2

julia> length(methodinstances(CTS.unrolled_foreach))
1

I think I'd better use SnoopCompile to incorporate the importance of timings in real-world cases.