diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index c097cbc77d2..92e7a5893c9 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -23,11 +23,6 @@ steps: command: - "echo $$JULIA_DEPOT_PATH" - - echo "--- Instantiate project" - - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project -e 'using Pkg; Pkg.precompile()'" - - "julia --project -e 'using Pkg; Pkg.status()'" - - echo "--- Instantiate examples" - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=examples -e 'using Pkg; Pkg.precompile()'" diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index bb223353e29..98a6be08bdb 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -2,6 +2,7 @@ module ClimaAtmos using NVTX, Colors import Thermodynamics as TD +import ClimaCore.DataLayouts: @fused include(joinpath("parameters", "Parameters.jl")) import .Parameters as CAP diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index 2b327783df1..2826664b7fa 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -17,17 +17,19 @@ function set_precipitation_precomputed_quantities!(Y, p, t) cmp = CAP.microphysics_params(p.params) # compute the precipitation terminal velocity [m/s] - @. ᶜwᵣ = CM1.terminal_velocity( - cmp.pr, - cmp.tv.rain, - Y.c.ρ, - abs(Y.c.ρq_rai / Y.c.ρ), - ) - @. ᶜwₛ = CM1.terminal_velocity( - cmp.ps, - cmp.tv.snow, - Y.c.ρ, - abs(Y.c.ρq_sno / Y.c.ρ), - ) + @fused begin + @. ᶜwᵣ = CM1.terminal_velocity( + cmp.pr, + cmp.tv.rain, + Y.c.ρ, + abs(Y.c.ρq_rai / Y.c.ρ), + ) + @. ᶜwₛ = CM1.terminal_velocity( + cmp.ps, + cmp.tv.snow, + Y.c.ρ, + abs(Y.c.ρq_sno / Y.c.ρ), + ) + end return nothing end diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 34b6f88b405..ab434e4d929 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -454,8 +454,13 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ # TODO: We should think more about these increments before we use them. end - @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) - @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + (; ᶜh_tot) = p.precomputed + @fused begin + @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) + @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + @. ᶜh_tot = + TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) + end if turbconv_model isa AbstractEDMF @. p.precomputed.ᶜgradᵥ_θ_virt = @@ -466,8 +471,6 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) end - (; ᶜh_tot) = p.precomputed - @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) if !isnothing(p.sfc_setup) SurfaceConditions.update_surface_conditions!(Y, p, t) diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index df85dbcb16a..4335dd68fe6 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -19,27 +19,31 @@ function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³ (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed - @. ᶜρa⁰ = ρa⁰(Y.c) - @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) - @. ᶜmse⁰ = divide_by_ρa( - Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_tot⁰ = divide_by_ρa( - Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_tot, - Y.c.ρ, - turbconv_model, - ) + @fused begin + @. ᶜρa⁰ = ρa⁰(Y.c) + @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) + @. ᶜmse⁰ = divide_by_ρa( + Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρ * (ᶜh_tot - ᶜK), + Y.c.ρ, + turbconv_model, + ) + @. ᶜq_tot⁰ = divide_by_ρa( + Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρq_tot, + Y.c.ρ, + turbconv_model, + ) + end set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ - @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + @fused begin + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) + @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + end return nothing end @@ -77,9 +81,6 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³) - @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 - @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) - @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) # EDMFX boundary condition: @@ -109,49 +110,60 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1)) ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1)) ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1)) - @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc( - ᶜz_int_val - z_sfc_val, - ᶜρ_int_val, - ᶜh_tot_int_val - ᶜK_int_val, - buoyancy_flux_val, - ρ_flux_h_tot_val, - ustar_val, - obukhov_length_val, - sfc_local_geometry_val, - ) - # ... and the first interior point for EDMFX ᶜq_totʲ. ᶜq_tot_int_val = Fields.field_values(Fields.level(ᶜspecific.q_tot, 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) - @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( - ᶜz_int_val - z_sfc_val, - ᶜρ_int_val, - ᶜq_tot_int_val, - buoyancy_flux_val, - ρ_flux_q_tot_val, - ustar_val, - obukhov_length_val, - sfc_local_geometry_val, - ) - # Then overwrite the prognostic variables at first inetrior point. ᶜΦ_int_val = Fields.field_values(Fields.level(ᶜΦ, 1)) ᶜtsʲ_int_val = Fields.field_values(Fields.level(ᶜtsʲ, 1)) - @. ᶜtsʲ_int_val = TD.PhaseEquil_phq( - thermo_params, - ᶜp_int_val, - ᶜmseʲ_int_val - ᶜΦ_int_val, - ᶜq_totʲ_int_val, - ) + sgsʲs_ρ_int_val = Fields.field_values(Fields.level(ᶜρʲs.:($j), 1)) sgsʲs_ρa_int_val = Fields.field_values(Fields.level(Y.c.sgsʲs.:($j).ρa, 1)) turbconv_params = CAP.turbconv_params(params) - @. sgsʲs_ρ_int_val = TD.air_density(thermo_params, ᶜtsʲ_int_val) - @. sgsʲs_ρa_int_val = - $(FT(turbconv_params.surface_area)) * - TD.air_density(thermo_params, ᶜtsʲ_int_val) + + @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 + @fused begin + @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) + @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) + end + + @fused begin + @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc( + ᶜz_int_val - z_sfc_val, + ᶜρ_int_val, + ᶜh_tot_int_val - ᶜK_int_val, + buoyancy_flux_val, + ρ_flux_h_tot_val, + ustar_val, + obukhov_length_val, + sfc_local_geometry_val, + ) + + @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( + ᶜz_int_val - z_sfc_val, + ᶜρ_int_val, + ᶜq_tot_int_val, + buoyancy_flux_val, + ρ_flux_q_tot_val, + ustar_val, + obukhov_length_val, + sfc_local_geometry_val, + ) + + # Then overwrite the prognostic variables at first inetrior point. + @. ᶜtsʲ_int_val = TD.PhaseEquil_phq( + thermo_params, + ᶜp_int_val, + ᶜmseʲ_int_val - ᶜΦ_int_val, + ᶜq_totʲ_int_val, + ) + @. sgsʲs_ρ_int_val = TD.air_density(thermo_params, ᶜtsʲ_int_val) + @. sgsʲs_ρa_int_val = + $(FT(turbconv_params.surface_area)) * + TD.air_density(thermo_params, ᶜtsʲ_int_val) + end end return nothing end diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index 2cf7ccb3c85..10cc342f52d 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -293,17 +293,19 @@ function edmfx_entr_detr_tendency!( for j in 1:n - @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += - Y.c.sgsʲs.:($$j).ρa[colidx] * - (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx]) + @fused begin + @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += + Y.c.sgsʲs.:($$j).ρa[colidx] * + (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx]) - @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += - ᶜentrʲs.:($$j)[colidx] * - (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx]) + @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += + ᶜentrʲs.:($$j)[colidx] * + (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx]) - @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += - ᶜentrʲs.:($$j)[colidx] * - (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx]) + @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += + ᶜentrʲs.:($$j)[colidx] * + (ᶜq_tot⁰[colidx] - Y.c.sgsʲs.:($$j).q_tot[colidx]) + end @. Yₜ.f.sgsʲs.:($$j).u₃[colidx] += ᶠinterp(ᶜentrʲs.:($$j)[colidx]) * diff --git a/src/prognostic_equations/edmfx_precipitation.jl b/src/prognostic_equations/edmfx_precipitation.jl index 7b0dcc2cd69..e9fa2546ad4 100644 --- a/src/prognostic_equations/edmfx_precipitation.jl +++ b/src/prognostic_equations/edmfx_precipitation.jl @@ -28,20 +28,22 @@ function edmfx_precipitation_tendency!( for j in 1:n - @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += - Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜS_q_totʲs.:($$j)[colidx] + @fused begin + @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] += + Y.c.sgsʲs.:($$j).ρa[colidx] * ᶜS_q_totʲs.:($$j)[colidx] - @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += - ᶜS_q_totʲs.:($$j)[colidx] * ( - e_tot_0M_precipitation_sources_helper( - thermo_params, - ᶜtsʲs.:($$j)[colidx], - ᶜΦ[colidx], - ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx]) - ) + @. Yₜ.c.sgsʲs.:($$j).mse[colidx] += + ᶜS_q_totʲs.:($$j)[colidx] * ( + e_tot_0M_precipitation_sources_helper( + thermo_params, + ᶜtsʲs.:($$j)[colidx], + ᶜΦ[colidx], + ) - TD.internal_energy(thermo_params, ᶜtsʲs.:($$j)[colidx]) + ) - @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += - ᶜS_q_totʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] += + ᶜS_q_totʲs.:($$j)[colidx] * (1 - Y.c.sgsʲs.:($$j).q_tot[colidx]) + end end return nothing end diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 10dba82f650..17e71a892bf 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -26,12 +26,15 @@ function edmfx_sgs_mass_flux_tendency!( ᶠu³_diff_colidx = p.scratch.ᶠtemp_CT3[colidx] ᶜa_scalar_colidx = p.scratch.ᶜtemp_scalar[colidx] for j in 1:n - @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] - @. ᶜa_scalar_colidx = - ( - Y.c.sgsʲs.:($$j).mse[colidx] + ᶜKʲs.:($$j)[colidx] - - ᶜh_tot[colidx] - ) * draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx]) + @fused begin + @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] + @. ᶜa_scalar_colidx = + ( + Y.c.sgsʲs.:($$j).mse[colidx] + ᶜKʲs.:($$j)[colidx] - + ᶜh_tot[colidx] + ) * + draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx]) + end vertical_transport!( Yₜ.c.ρe_tot[colidx], ᶜJ[colidx], @@ -42,10 +45,12 @@ function edmfx_sgs_mass_flux_tendency!( edmfx_sgsflux_upwinding, ) end - @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx] - @. ᶜa_scalar_colidx = - (ᶜmse⁰[colidx] + ᶜK⁰[colidx] - ᶜh_tot[colidx]) * - draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx]) + @fused begin + @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx] + @. ᶜa_scalar_colidx = + (ᶜmse⁰[colidx] + ᶜK⁰[colidx] - ᶜh_tot[colidx]) * + draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx]) + end vertical_transport!( Yₜ.c.ρe_tot[colidx], ᶜJ[colidx], @@ -59,10 +64,17 @@ function edmfx_sgs_mass_flux_tendency!( if !(p.atmos.moisture_model isa DryModel) # specific humidity for j in 1:n - @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] - @. ᶜa_scalar_colidx = - (Y.c.sgsʲs.:($$j).q_tot[colidx] - ᶜspecific.q_tot[colidx]) * - draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx]) + @fused begin + @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] + @. ᶜa_scalar_colidx = + ( + Y.c.sgsʲs.:($$j).q_tot[colidx] - + ᶜspecific.q_tot[colidx] + ) * draft_area( + Y.c.sgsʲs.:($$j).ρa[colidx], + ᶜρʲs.:($$j)[colidx], + ) + end vertical_transport!( Yₜ.c.ρq_tot[colidx], ᶜJ[colidx], @@ -73,10 +85,12 @@ function edmfx_sgs_mass_flux_tendency!( edmfx_sgsflux_upwinding, ) end - @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx] - @. ᶜa_scalar_colidx = - (ᶜq_tot⁰[colidx] - ᶜspecific.q_tot[colidx]) * - draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx]) + @fused begin + @. ᶠu³_diff_colidx = ᶠu³⁰[colidx] - ᶠu³[colidx] + @. ᶜa_scalar_colidx = + (ᶜq_tot⁰[colidx] - ᶜspecific.q_tot[colidx]) * + draft_area(ᶜρa⁰[colidx], ᶜρ⁰[colidx]) + end vertical_transport!( Yₜ.c.ρq_tot[colidx], ᶜJ[colidx], @@ -253,8 +267,10 @@ function edmfx_sgs_diffusive_flux_tendency!( ) @. ᶜρχₜ_diffusion[colidx] = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜq_tot⁰[colidx]))) - @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] - @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + @fused begin + @. Yₜ.c.ρq_tot[colidx] -= ᶜρχₜ_diffusion[colidx] + @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + end end # momentum diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 477a2a36eee..8d1e1f6e3c2 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -36,11 +36,13 @@ function edmfx_tke_tendency!( if use_prognostic_tke(turbconv_model) # shear production - @. Yₜ.c.sgs⁰.ρatke[colidx] += - 2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx] - # buoyancy production - @. Yₜ.c.sgs⁰.ρatke[colidx] -= - ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx] + @fused begin + @. Yₜ.c.sgs⁰.ρatke[colidx] += + 2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx] + # buoyancy production + @. Yₜ.c.sgs⁰.ρatke[colidx] -= + ᶜρa⁰[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx] + end # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation for j in 1:n diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl index 07a752e7fb2..6b14544cb51 100644 --- a/src/prognostic_equations/gm_sgs_closures.jl +++ b/src/prognostic_equations/gm_sgs_closures.jl @@ -54,18 +54,20 @@ function compute_gm_mixing_length!(ᶜmixing_length, Y, p) compute_strain_rate_center!(ᶜstrain_rate, ᶠu) ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_2 - @. ᶜprandtl_nvec = turbulent_prandtl_number( - params, - obukhov_length, - ᶜlinear_buoygrad, - norm_sqr(ᶜstrain_rate), - ) + @fused begin + @. ᶜprandtl_nvec = turbulent_prandtl_number( + params, + obukhov_length, + ᶜlinear_buoygrad, + norm_sqr(ᶜstrain_rate), + ) - @. ᶜmixing_length = smagorinsky_lilly_length( - CAP.c_smag(params), - sqrt(max(ᶜlinear_buoygrad, 0)), #N_eff - ᶜdz, - ᶜprandtl_nvec, - norm_sqr(ᶜstrain_rate), - ) + @. ᶜmixing_length = smagorinsky_lilly_length( + CAP.c_smag(params), + sqrt(max(ᶜlinear_buoygrad, 0)), #N_eff + ᶜdz, + ᶜprandtl_nvec, + norm_sqr(ᶜstrain_rate), + ) + end end diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index cf046c05d9c..877962787db 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -92,8 +92,10 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) - C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j))))) @. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse)) - @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j)) - @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j)) + @fused begin + @. ᶜ∇²uₕʲs.:($$j) = C12(ᶜ∇²uʲs.:($$j)) + @. ᶜ∇²uᵥʲs.:($$j) = C3(ᶜ∇²uʲs.:($$j)) + end end end end diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index 39f03bb0316..33ab3651b14 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -95,9 +95,17 @@ function vertical_diffusion_boundary_layer_tendency!( ᶠgradᵥ(ᶜχ[colidx]) ), ) - @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] + # @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] + # if !(χ_name in (:q_rai, :q_sno)) + # @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + # end if !(χ_name in (:q_rai, :q_sno)) - @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + @fused begin + @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] + @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx] + end + else + @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx] end end end diff --git a/src/prognostic_equations/zero_gridscale_tendency.jl b/src/prognostic_equations/zero_gridscale_tendency.jl index 04e113eaee4..5126933ff8c 100644 --- a/src/prognostic_equations/zero_gridscale_tendency.jl +++ b/src/prognostic_equations/zero_gridscale_tendency.jl @@ -5,10 +5,12 @@ function zero_gridscale_tendency!(Yₜ, Y, p, t, colidx) # turn off all grid-scale tendencies in the simple edmfx test if !p.atmos.gs_tendency - @. Yₜ.c.ρ[colidx] = 0 - @. Yₜ.c.uₕ[colidx] = C12(0, 0) - @. Yₜ.f.u₃[colidx] = C3(0) - @. Yₜ.c.ρe_tot[colidx] = 0 + @fused begin + @. Yₜ.c.ρ[colidx] = 0 + @. Yₜ.c.uₕ[colidx] = C12(0, 0) + @. Yₜ.f.u₃[colidx] = C3(0) + @. Yₜ.c.ρe_tot[colidx] = 0 + end for ρχ_name in filter(is_tracer_var, propertynames(Y.c)) @. Yₜ.c.:($$ρχ_name)[colidx] = 0 end diff --git a/src/prognostic_equations/zero_velocity.jl b/src/prognostic_equations/zero_velocity.jl index a73162733b3..27bf0392241 100644 --- a/src/prognostic_equations/zero_velocity.jl +++ b/src/prognostic_equations/zero_velocity.jl @@ -8,8 +8,10 @@ function zero_velocity_tendency!(Yₜ, Y, p, t, colidx) FT = eltype(Y) n = n_mass_flux_subdomains(p.atmos.turbconv_model) - @. Yₜ.c.uₕ[colidx] = C12(FT(0), FT(0)) - @. Yₜ.f.u₃[colidx] = Geometry.Covariant3Vector(FT(0)) + @fused begin + @. Yₜ.c.uₕ[colidx] = C12(FT(0), FT(0)) + @. Yₜ.f.u₃[colidx] = Geometry.Covariant3Vector(FT(0)) + end if p.atmos.turbconv_model isa PrognosticEDMFX for j in 1:n @. Yₜ.f.sgsʲs.:($$j).u₃[colidx] =