From d3aae917d45323e7ff6f993f04dadd533e59dc08 Mon Sep 17 00:00:00 2001 From: nefrathenrici Date: Thu, 13 Nov 2025 15:08:12 -0800 Subject: [PATCH] fix vertical_transport with vanleer_limiter on GPU --- .buildkite/pipeline.yml | 2 +- .../model_configs/diagnostic_edmfx_aquaplanet_gpu.yml | 2 ++ src/prognostic_equations/advection.jl | 10 +++++----- src/prognostic_equations/implicit/implicit_tendency.jl | 2 ++ 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 137edc9c07..c7915965a6 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1281,7 +1281,7 @@ steps: CLIMACOMMS_DEVICE: "CUDA" CLIMA_NAME_CUDA_KERNELS_FROM_STACK_TRACE: "true" agents: - slurm_mem: 24GB + slurm_mem: 32GB slurm_gpus: 1 - group: "Flame graphs" diff --git a/config/model_configs/diagnostic_edmfx_aquaplanet_gpu.yml b/config/model_configs/diagnostic_edmfx_aquaplanet_gpu.yml index 6f5fcda806..5613ae34b5 100644 --- a/config/model_configs/diagnostic_edmfx_aquaplanet_gpu.yml +++ b/config/model_configs/diagnostic_edmfx_aquaplanet_gpu.yml @@ -9,6 +9,8 @@ edmfx_detr_model: "Generalized" edmfx_nh_pressure: true edmfx_sgs_mass_flux: true edmfx_sgs_diffusive_flux: true +edmfx_sgsflux_upwinding: "vanleer_limiter" +edmfx_tracer_upwinding: "vanleer_limiter" moist: equil cloud_model: "quadrature_sgs" precip_model: 0M diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 1790bff5ef..90c9446b44 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -265,7 +265,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name if !(ρχ_name in (@name(ρe_tot), @name(ρq_tot))) ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, FT(dt), tracer_upwinding) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, tracer_upwinding) @. ᶜρχₜ += vtt end end @@ -280,15 +280,15 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) specific(Y.c.ρe_tot, Y.c.ρ), ), ) - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), energy_q_tot_upwinding) - vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, FT(dt), Val(:none)) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, energy_q_tot_upwinding) + vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, Val(:none)) @. Yₜ.c.ρe_tot += vtt - vtt_central end if !(p.atmos.moisture_model isa DryModel) && energy_q_tot_upwinding != Val(:none) ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) - vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), energy_q_tot_upwinding) - vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, FT(dt), Val(:none)) + vtt = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, dt, energy_q_tot_upwinding) + vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜq_tot, dt, Val(:none)) @. Yₜ.c.ρq_tot += vtt - vtt_central end diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 69365901e1..b80cd42083 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -86,6 +86,8 @@ end function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:vanleer_limiter}) ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J + FT = eltype(ᶠJ) + dt = FT(dt) return @. lazy( -(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠlin_vanleer(ᶠu³, ᶜχ, dt))), )