Merge #71

71: Use concrete phase types r=charleskawczynski a=charleskawczynski This should help with allocations, since we make closures over these types, which are Unionall right now, this PR makes them concrete. Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>
CliMA · Dec 12, 2021 · fcd8f07 · fcd8f07 · charleskawczynski · Dec 12, 2021
2 parents c18f29b + 4709361
commit fcd8f07
Show file tree

Hide file tree

Showing 10 changed files with 92 additions and 87 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -62,5 +62,5 @@ steps:
       slurm_gres: "gpu:1"
     env:
       JULIA_VERSION: "1.7.0"
-      CUDA_VERSION: "11.2"
+      CUDA_VERSION: "10.2"
     timeout_in_minutes: 60
diff --git a/Project.toml b/Project.toml
@@ -1,12 +1,11 @@
 name = "Thermodynamics"
 uuid = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 authors = ["Climate Modeling Alliance"]
-version = "0.5.8"
+version = "0.5.9"
 
 [deps]
 CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
@@ -16,5 +15,4 @@ CLIMAParameters = "0.1, 0.2, 0.3"
 DocStringExtensions = "0.8.1"
 KernelAbstractions = "0.5, 0.6, 0.7"
 RootSolvers = "0.2"
-ExprTools = "0.1 - 0.1.3"
 julia = "1.4"
diff --git a/perf/alloc_per_constructor.jl b/perf/alloc_per_constructor.jl
@@ -8,15 +8,12 @@ include("common.jl")
 constructor = ENV["ALLOCATION_CONSTRUCTOR"]
 @info "Recording allocations for $constructor"
 
-
-if constructor == "ρeq"
-    thermo_state = thermo_state_ρeq_newton
-else
-end
-if constructor == "pθq"
-    thermo_state = thermo_state_pθq
-else
-end
+thermo_state_map = Dict(
+    "ρeq" => thermo_state_ρeq_newton,
+    "pθq" => thermo_state_pθq,
+    "pTq" => thermo_state_pTq,
+)
+thermo_state = thermo_state_map[constructor]
 
 thermo_state() # compile first
 Profile.clear_malloc_data()

diff --git a/perf/allocs.jl b/perf/allocs.jl
@@ -20,7 +20,7 @@ all_dirs_to_monitor = [
 # https://github.com/jheinen/GR.jl/issues/278#issuecomment-587090846
 ENV["GKSwstype"] = "nul"
 
-constructors = ["ρeq", "pθq"]
+constructors = ["ρeq", "pθq", "pTq"]
 
 allocs = Dict()
 for constructor in constructors
@@ -101,6 +101,10 @@ function plot_allocs(constructor, allocs_per_case, n_unique_bytes)
     end
 
     all_bytes = all_bytes ./ 10^3
+    if isempty(all_bytes)
+        @info "$constructor: 0 allocations! 🎉"
+        return nothing
+    end
     max_bytes = maximum(all_bytes)
     @info "$constructor: $all_bytes"
     xtick_name(filename, linenumber) = "$filename, line number: $linenumber"

diff --git a/perf/benchmark.jl b/perf/benchmark.jl
@@ -10,5 +10,7 @@ include("common.jl")
     show(stdout, MIME("text/plain"), trial)
     trial = BenchmarkTools.@benchmark $thermo_state_pθq()
     show(stdout, MIME("text/plain"), trial)
+    trial = BenchmarkTools.@benchmark $thermo_state_pTq()
+    show(stdout, MIME("text/plain"), trial)
 
 end
diff --git a/perf/common.jl b/perf/common.jl
@@ -19,7 +19,7 @@ const param_set = EarthParameterSet()
 
 ArrayType = Array{Float64}
 profiles = TD.TestedProfiles.PhaseEquilProfiles(param_set, ArrayType)
-UnPack.@unpack e_int, ρ, p, θ_liq_ice, q_tot = profiles
+UnPack.@unpack e_int, T, ρ, p, θ_liq_ice, q_tot = profiles
 
 function thermo_state_ρeq_newton()
     ts =
@@ -57,3 +57,7 @@ end
 function thermo_state_pθq()
     ts = TD.PhaseEquil_pθq.(param_set, p, θ_liq_ice, q_tot)
 end
+
+function thermo_state_pTq()
+    ts = TD.PhaseEquil_pTq.(param_set, p, T, q_tot)
+end
diff --git a/src/TestedProfiles.jl b/src/TestedProfiles.jl
@@ -161,7 +161,7 @@ end
 Returns a `ProfileSet` used to test dry thermodynamic states.
 """
 function PhaseDryProfiles(param_set::APS, ::Type{ArrayType}) where {ArrayType}
-    phase_type = TD.PhaseDry
+    phase_type = TD.PhaseDry{eltype(ArrayType), typeof(param_set)}
 
     z_range, relative_sat, T_surface, T_min = input_config(ArrayType)
     z, T_virt, p, RS =
@@ -221,7 +221,7 @@ end
 Returns a `ProfileSet` used to test moist states in thermodynamic equilibrium.
 """
 function PhaseEquilProfiles(param_set::APS, ::Type{ArrayType}) where {ArrayType}
-    phase_type = TD.PhaseEquil
+    phase_type = TD.PhaseEquil{eltype(ArrayType), typeof(param_set)}
 
     # Prescribe z_range, relative_sat, T_surface, T_min
     z_range, relative_sat, T_surface, T_min = input_config(ArrayType)

diff --git a/src/config_numerical_method.jl b/src/config_numerical_method.jl
@@ -30,8 +30,8 @@ function sa_numerical_method(
     ρ::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.NewtonsMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.NewtonsMethod, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     T_init =
         max(T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
@@ -47,8 +47,8 @@ function sa_numerical_method(
     ρ::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.NewtonsMethodAD}
+    ::Type{phase_type},
+) where {FT, NM <: RS.NewtonsMethodAD, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     T_init =
         max(T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
@@ -61,8 +61,8 @@ function sa_numerical_method(
     ρ::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.SecantMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.SecantMethod, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
     T_2 = air_temperature(param_set, e_int, q_pt)
@@ -77,8 +77,8 @@ function sa_numerical_method(
     ρ::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.RegulaFalsiMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.RegulaFalsiMethod, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
     T_2 = air_temperature(param_set, e_int, q_pt)
@@ -97,8 +97,8 @@ function sa_numerical_method_ρpq(
     ρ::FT,
     p::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.NewtonsMethodAD}
+    ::Type{phase_type},
+) where {FT, NM <: RS.NewtonsMethodAD, phase_type <: PhaseEquil}
     q_pt = PhasePartition(q_tot)
     T_init = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt)
     return RS.NewtonsMethodAD(T_init)
@@ -110,8 +110,8 @@ function sa_numerical_method_ρpq(
     ρ::FT,
     p::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.RegulaFalsiMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.RegulaFalsiMethod, phase_type <: PhaseEquil}
     q_pt = PhasePartition(q_tot)
     T_1 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) - 5
     T_2 = air_temperature_from_ideal_gas_law(param_set, p, ρ, q_pt) + 5
@@ -128,8 +128,8 @@ function sa_numerical_method_peq(
     p::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.NewtonsMethodAD}
+    ::Type{phase_type},
+) where {FT, NM <: RS.NewtonsMethodAD, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     T_init =
         max(T_min, air_temperature(param_set, e_int, PhasePartition(q_tot))) # Assume all vapor
@@ -142,8 +142,8 @@ function sa_numerical_method_peq(
     p::FT,
     e_int::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.SecantMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.SecantMethod, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     q_pt = PhasePartition(q_tot, FT(0), q_tot) # Assume all ice
     T_2 = air_temperature(param_set, e_int, q_pt)
@@ -162,8 +162,8 @@ function sa_numerical_method_pθq(
     p::FT,
     θ_liq_ice::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.RegulaFalsiMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.RegulaFalsiMethod, phase_type <: PhaseEquil}
     _T_min::FT = CPP.T_min(param_set)
     _T_max::FT = CPP.T_max(param_set)
     air_temp(q) = air_temperature_given_pθq(param_set, p, θ_liq_ice, q)
@@ -179,8 +179,8 @@ function sa_numerical_method_pθq(
     p::FT,
     θ_liq_ice::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.SecantMethod}
+    ::Type{phase_type},
+) where {FT, NM <: RS.SecantMethod, phase_type <: PhaseEquil}
     _T_min::FT = CPP.T_min(param_set)
     air_temp(q) = air_temperature_given_pθq(param_set, p, θ_liq_ice, q)
     T_1 = max(_T_min, air_temp(PhasePartition(q_tot))) # Assume all vapor
@@ -195,8 +195,8 @@ function sa_numerical_method_pθq(
     p::FT,
     θ_liq_ice::FT,
     q_tot::FT,
-    phase_type::Type{<:PhaseEquil},
-) where {FT, NM <: RS.NewtonsMethodAD}
+    ::Type{phase_type},
+) where {FT, NM <: RS.NewtonsMethodAD, phase_type <: PhaseEquil}
     T_min::FT = CPP.T_min(param_set)
     air_temp(q) = air_temperature_given_pθq(param_set, p, θ_liq_ice, q)
     T_init = max(T_min, air_temp(PhasePartition(q_tot))) # Assume all vapor