fix vertical_transport with vanleer_limiter on GPU

Author: nefrathenrici

.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"

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

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
 

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))),
         )