From 736d49fe4fe238355a5d5319a29b0b3d85199184 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 14 Jul 2023 11:26:44 -0700 Subject: [PATCH 1/2] Fix some allocations and gpu moist example --- .../diagnostic_edmf_precomputed_quantities.jl | 28 +++++++++++-------- .../non_orographic_gravity_wave.jl | 2 +- .../microphysics/precipitation.jl | 7 ++--- .../radiation/radiation.jl | 2 +- .../nogw_test_mima.jl | 2 +- 5 files changed, 22 insertions(+), 19 deletions(-) diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index ab7062ae38..1bc850be52 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -343,18 +343,20 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t) scale_factor = FT(1e-6) @. u³ʲ_halflevel = ifelse( ( - u³ʲ_datau³ʲ_data < - (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) || - ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level) + ( + u³ʲ_datau³ʲ_data < + (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) + ) | (ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level)) ), CT3(FT(0)), CT3(sqrt(max(FT(0), u³ʲ_datau³ʲ_data))), ) @. ρaʲ_level = ifelse( ( - u³ʲ_datau³ʲ_data < - (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) || - ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level) + ( + u³ʲ_datau³ʲ_data < + (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) + ) | (ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level)) ), FT(0), ρaʲu³ʲ_data / sqrt(max(FT(0), u³ʲ_datau³ʲ_data)), @@ -378,9 +380,10 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t) ) @. h_totʲ_level = ifelse( ( - u³ʲ_datau³ʲ_data < - (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) || - ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level) + ( + u³ʲ_datau³ʲ_data < + (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) + ) | (ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level)) ), h_tot_level, ρaʲu³ʲ_datah_tot / ρaʲu³ʲ_data, @@ -404,9 +407,10 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t) ) @. q_totʲ_level = ifelse( ( - u³ʲ_datau³ʲ_data < - (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) || - ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level) + ( + u³ʲ_datau³ʲ_data < + (scale_factor / (∂x³∂ξ³_level * ∂x³∂ξ³_level)) + ) | (ρaʲu³ʲ_data < (scale_factor / ∂x³∂ξ³_level)) ), q_tot_level, ρaʲu³ʲ_dataq_tot / ρaʲu³ʲ_data, diff --git a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl index e56914c17a..9de73a34ae 100644 --- a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl +++ b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl @@ -73,7 +73,7 @@ function non_orographic_gravity_wave_cache( # This latitude depend source follows MiMA specs source_ampl = @. ifelse( - lat > ϕ0_n || lat < ϕ0_s, + (lat > ϕ0_n) | (lat < ϕ0_s), Bt_0 + Bt_n * FT(0.5) * (FT(1) + tanh((lat - ϕ0_n) / dϕ_n)) + Bt_s * FT(0.5) * (FT(1) + tanh((lat - ϕ0_s) / dϕ_s)), diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index ac8fd950f3..2a5ca81b94 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -94,10 +94,9 @@ function precipitation_tendency!( # update precip in cache for coupler's use # 3d rain and snow @. ᶜT[colidx] = TD.air_temperature(thermo_params, ᶜts[colidx]) - @. ᶜ3d_rain[colidx] = - ifelse(ᶜT[colidx] >= FT(273.15), ᶜS_ρq_tot[colidx], FT(0)) - @. ᶜ3d_snow[colidx] = - ifelse(ᶜT[colidx] < FT(273.15), ᶜS_ρq_tot[colidx], FT(0)) + T_fr = TD.Parameters.T_freeze(thermo_params) + @. ᶜ3d_rain[colidx] = ifelse(ᶜT[colidx] >= T_fr, ᶜS_ρq_tot[colidx], 0) + @. ᶜ3d_snow[colidx] = ifelse(ᶜT[colidx] < T_fr, ᶜS_ρq_tot[colidx], 0) Operators.column_integral_definite!( col_integrated_rain[colidx], ᶜ3d_rain[colidx], diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 31056aed82..08f1d98615 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -148,7 +148,7 @@ function radiation_model_cache( ), center_cloud_fraction = RRTMGPI.field2array( @. ifelse( - ᶜis_bottom_cloud || ᶜis_top_cloud, + ᶜis_bottom_cloud | ᶜis_top_cloud, FT(1), FT(0) * ᶜΔz, ) diff --git a/test/non_orographic_gravity_wave/nogw_test_mima.jl b/test/non_orographic_gravity_wave/nogw_test_mima.jl index ae755f5a23..793fa9fae2 100644 --- a/test/non_orographic_gravity_wave/nogw_test_mima.jl +++ b/test/non_orographic_gravity_wave/nogw_test_mima.jl @@ -46,7 +46,7 @@ function non_orographic_gravity_wave_cache( # source amplitude following MiMA: radical change between subtropics and the tropic source_ampl = @. ifelse( - lat > ϕ0_n || lat < ϕ0_s, + (lat > ϕ0_n) | (lat < ϕ0_s), Bt_0 + Bt_n * FT(0.5) * (FT(1) + tanh((lat - ϕ0_n) / dϕ_n)) + Bt_s * FT(0.5) * (FT(1) + tanh((lat - ϕ0_s) / dϕ_s)), From 54e3e060ec8c444d7abded8477b47a720df2ab86 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 14 Jul 2023 13:21:41 -0700 Subject: [PATCH 2/2] Remove gpu moist soft fail --- .buildkite/pipeline.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 17f7d00884..c4031b5c96 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -812,7 +812,6 @@ steps: - "julia --project -e 'using CUDA; CUDA.versioninfo()'" - "nsys profile --trace=nvtx,cuda --output=gpu_implicit_barowave_moist/report julia --color=yes --project=examples perf/benchmark_step.jl --job_id gpu_implicit_barowave_moist --z_elem 25 --h_elem 12 --initial_condition MoistBaroclinicWave --t_end 1mins --dt 1secs --dt_save_to_sol Inf --dt_save_to_disk Inf --moist equil --precip 0M --apply_limiter false" artifact_paths: "gpu_implicit_barowave_moist/*" - soft_fail: true agents: slurm_gpus: 1