Merge #2274 #2282

2274: separate precomputed quantities in 3 funcs r=trontrytel a=trontrytel

Following what we discussed `@szy21` , `@Sbozzolo`  this PR splits prognostic edmfx precomputed quantities into 3 steps
1) updating environment
2) updating updraft thermo state and BC
3) updating the closures.

I'm not sure if it actually makes anything better though. Or if you want to split the individual closures. 

2282: Increase job memory for perf jobs and lower allocations r=Sbozzolo a=Sbozzolo



Co-authored-by: Anna Jaruga <[email protected]>
Co-authored-by: Gabriele Bozzola <[email protected]>

.buildkite/pipeline.yml

@@ -994,14 +994,14 @@ steps:
           julia --color=yes --project=perf perf/benchmark.jl
           $PERF_CONFIG_PATH/bm_perf_target.yml
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":computer: Benchmark: perf target (Threaded)"
         command: >
           julia --color=yes --threads 8 --project=perf perf/benchmark.jl
           $PERF_CONFIG_PATH/bm_perf_target_threaded.yml
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
           slurm_cpus_per_task: 8
 
       # Flame graphs
@@ -1011,31 +1011,31 @@ steps:
           $PERF_CONFIG_PATH/flame_perf_target.yml
         artifact_paths: "flame_perf_target/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (with tracers)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_target_tracers.yml
         artifact_paths: "flame_perf_target_tracers/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (diagnostic edmfx)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_target_diagnostic_edmfx.yml
         artifact_paths: "flame_perf_target_diagnostic_edmfx/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (edmf)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_target_edmf.yml
         artifact_paths: "flame_perf_target_edmf/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (Threaded)"
         command: >
@@ -1044,47 +1044,47 @@ steps:
         artifact_paths: "flame_perf_target_threaded/*"
         agents:
           slurm_cpus_per_task: 8
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (Callbacks)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_target_callbacks.yml
         artifact_paths: "flame_perf_target_callbacks/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: gravity wave"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_gw.yml
         artifact_paths: "flame_perf_gw/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       - label: ":fire: Flame graph: perf target (diagnostics)"
         command: >
           julia --color=yes --project=perf perf/flame.jl
           $PERF_CONFIG_PATH/flame_perf_diagnostics.yml
         artifact_paths: "flame_perf_diagnostics/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       # Inference
       - label: ":rocket: JET n-failures (inference)"
         command: >
           julia --color=yes --project=perf perf/jet_test_nfailures.jl
           $PERF_CONFIG_PATH/jet_n_failures.yml
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
       # Latency
       - label: ":mag::rocket: Invalidations"
         command: >
           julia --color=yes --project=perf perf/invalidations.jl
         artifact_paths: "invalidations/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 24GB
 
   - wait: ~
     continue_on_failure: true

perf/flame.jl

@@ -63,11 +63,11 @@ allocs_limit["flame_perf_target"] = 4656
 allocs_limit["flame_perf_target_tracers"] = 204288
 allocs_limit["flame_perf_target_edmfx"] = 253440
 allocs_limit["flame_perf_diagnostics"] = 3016136
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 920512
-allocs_limit["flame_perf_target_edmf"] = 6386476224
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 862464
+allocs_limit["flame_perf_target_edmf"] = 4020913344
 allocs_limit["flame_perf_target_threaded"] = 5857808
-allocs_limit["flame_perf_target_callbacks"] = 46407936
-allocs_limit["flame_perf_gw"] = 4844555632
+allocs_limit["flame_perf_target_callbacks"] = 38058048
+allocs_limit["flame_perf_gw"] = 4844444272
 
 if allocs < allocs_limit[job_id] * buffer
     @info "TODO: lower `allocs_limit[$job_id]` to: $(allocs)"

src/cache/precomputed_quantities.jl

@@ -154,7 +154,7 @@ end
 """
     set_velocity_at_top!(Y, turbconv_model)
 
-Modifies `Y.f.u₃` so that `u₃` is 0 at the model top. 
+Modifies `Y.f.u₃` so that `u₃` is 0 at the model top.
 """
 function set_velocity_at_top!(Y, turbconv_model)
     top_u₃ = Fields.level(
@@ -340,7 +340,9 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     SurfaceConditions.update_surface_conditions!(Y, p, t)
 
     if turbconv_model isa PrognosticEDMFX
-        set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
+        set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
+        set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
+        set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
     end
 
     if turbconv_model isa DiagnosticEDMFX

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -5,28 +5,17 @@ import Thermodynamics as TD
 import ClimaCore: Spaces, Fields
 
 """
-    set_prognostic_edmf_precomputed_quantities!(Y, p, t)
+    set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
 
-Updates the precomputed quantities stored in `p` for edmfx.
+Updates the edmf environment precomputed quantities stored in `p` for edmfx.
 """
-function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
-    (; energy_form, moisture_model, turbconv_model) = p.atmos
-    #EDMFX BCs only support total energy as state variable
-    @assert energy_form isa TotalEnergy
-    @assert !(moisture_model isa DryModel)
-
-    FT = Spaces.undertype(axes(Y.c))
-    (; params) = p
-    (; dt) = p.simulation
-    thermo_params = CAP.thermodynamics_params(params)
-    n = n_mass_flux_subdomains(turbconv_model)
-    thermo_args = (thermo_params, energy_form, moisture_model)
+function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
+    @assert !(p.atmos.moisture_model isa DryModel)
 
-    (; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot, ᶜρ_ref) = p
+    thermo_params = CAP.thermodynamics_params(p.params)
+    (; turbconv_model) = p.atmos
+    (; ᶜp, ᶜΦ, ᶜh_tot) = p
     (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
-    (; ᶜmixing_length, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜK_u, ᶜK_h) = p
-    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p
-    (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
 
     @. ᶜρa⁰ = ρa⁰(Y.c)
     @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
@@ -49,6 +38,30 @@ function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
     @. ᶜK⁰ += ᶜtke⁰
     @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜh_tot⁰ - ᶜK⁰ - ᶜΦ, ᶜq_tot⁰)
     @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
+    return nothing
+end
+
+"""
+    set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
+
+Updates the draft thermo state and boundary conditions
+precomputed quantities stored in `p` for edmfx.
+"""
+function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
+    (; energy_form, moisture_model, turbconv_model) = p.atmos
+    #EDMFX BCs only support total energy as state variable
+    @assert energy_form isa TotalEnergy
+    @assert !(moisture_model isa DryModel)
+
+    FT = Spaces.undertype(axes(Y.c))
+    n = n_mass_flux_subdomains(turbconv_model)
+
+    (; params) = p
+    thermo_params = CAP.thermodynamics_params(params)
+
+    (; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot) = p
+    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p
+    (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
 
     for j in 1:n
         ᶜuʲ = ᶜuʲs.:($j)
@@ -135,6 +148,30 @@ function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
             $(FT(turbconv_params.surface_area)) *
             TD.air_density(thermo_params, ᶜtsʲ_int_val)
     end
+    return nothing
+end
+
+"""
+    set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
+
+Updates the precomputed quantities stored in `p` for edmfx closures.
+"""
+function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
+
+    (; moisture_model, turbconv_model) = p.atmos
+    @assert !(moisture_model isa DryModel)
+
+    (; params) = p
+    (; dt) = p.simulation
+    thermo_params = CAP.thermodynamics_params(params)
+
+    n = n_mass_flux_subdomains(turbconv_model)
+
+    (; ᶜspecific, ᶜp, ᶜρ_ref) = p
+    (; ᶜtke⁰, ᶜρa⁰, ᶜu⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) = p
+    (; ᶜmixing_length, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜK_u, ᶜK_h) = p
+    (; ᶜuʲs, ᶜtsʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p
+    (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
 
     ᶜz = Fields.coordinate_field(Y.c).z
     z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)