From a0ccabd0353351186e89c01a2c832ca406e44c93 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Fri, 15 Dec 2023 12:05:31 -0500
Subject: [PATCH 1/3] Assert atmos to not have UnionAll types

---
 src/ClimaAtmos.jl          |  1 +
 src/solver/type_getters.jl |  1 +
 src/utils/gpu_compat.jl    | 30 ++++++++++++++++++++++++++++++
 3 files changed, 32 insertions(+)
 create mode 100644 src/utils/gpu_compat.jl

diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
index 51b7e197e8..2afa960619 100644
--- a/src/ClimaAtmos.jl
+++ b/src/ClimaAtmos.jl
@@ -6,6 +6,7 @@ include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP
 
 include(joinpath("utils", "abbreviations.jl"))
+include(joinpath("utils", "gpu_compat.jl"))
 include(joinpath("utils", "common_spaces.jl"))
 include(joinpath("solver", "types.jl"))
 include(joinpath("solver", "cli_options.jl"))
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 7cf2aff14e..f4e35e213a 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -78,6 +78,7 @@ function get_atmos(config::AtmosConfig, params)
         surface_model = get_surface_model(parsed_args),
         numerics = get_numerics(parsed_args),
     )
+    @assert !@any_reltype(atmos, (UnionAll,))
 
     @info "AtmosModel: \n$(summary(atmos))"
     return atmos
diff --git a/src/utils/gpu_compat.jl b/src/utils/gpu_compat.jl
new file mode 100644
index 0000000000..b5f065fa05
--- /dev/null
+++ b/src/utils/gpu_compat.jl
@@ -0,0 +1,30 @@
+"""
+    @any_reltype(::Any, t::Tuple, warn=true)
+
+Returns a Bool (and prints warnings) if the given
+data structure has an instance of any types in `t`.
+"""
+function any_reltype(found, obj, name, ets, pc = (); warn = true)
+    for pn in propertynames(obj)
+        prop = getproperty(obj, pn)
+        pc_full = (pc..., ".", pn)
+        pc_string = name * string(join(pc_full))
+        for et in ets
+            if prop isa et
+                warn && @warn "$pc_string::$(typeof(prop)) is a DataType"
+                found = true
+            end
+        end
+        found = found || any_reltype(found, prop, name, ets, pc_full; warn)
+    end
+    return found
+end
+macro any_reltype(obj, ets, warn = true)
+    return :(any_reltype(
+        false,
+        $(esc(obj)),
+        $(string(obj)),
+        $(esc(ets));
+        warn = $(esc(warn)),
+    ))
+end

From a0a092af1b856cb6597fd333038f87a97528cac3 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Fri, 15 Dec 2023 11:59:28 -0500
Subject: [PATCH 2/3] Fix 2438

---
 src/cache/cache.jl         | 8 +++++---
 src/solver/type_getters.jl | 5 ++---
 src/solver/types.jl        | 2 ++
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
index 4e7cd73b0d..2ab1ae2842 100644
--- a/src/cache/cache.jl
+++ b/src/cache/cache.jl
@@ -129,9 +129,11 @@ function build_cache(Y, atmos, params, surface_setup, dt, start_date)
     net_energy_flux_toa = [Geometry.WVector(FT(0))]
     net_energy_flux_sfc = [Geometry.WVector(FT(0))]
 
-    limiter =
-        isnothing(atmos.numerics.limiter) ? nothing :
-        atmos.numerics.limiter(similar(Y.c, FT))
+    limiter = if isnothing(atmos.numerics.limiter)
+        nothing
+    elseif atmos.numerics.limiter isa QuasiMonotoneLimiter
+        Limiters.QuasiMonotoneLimiter(similar(Y.c, FT))
+    end
 
     numerics = (; limiter)
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index f4e35e213a..91be9a500b 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -78,7 +78,7 @@ function get_atmos(config::AtmosConfig, params)
         surface_model = get_surface_model(parsed_args),
         numerics = get_numerics(parsed_args),
     )
-    @assert !@any_reltype(atmos, (UnionAll,))
+    @assert !@any_reltype(atmos, (UnionAll, DataType))
 
     @info "AtmosModel: \n$(summary(atmos))"
     return atmos
@@ -97,8 +97,7 @@ function get_numerics(parsed_args)
     edmfx_sgsflux_upwinding =
         Val(Symbol(parsed_args["edmfx_sgsflux_upwinding"]))
 
-    limiter =
-        parsed_args["apply_limiter"] ? Limiters.QuasiMonotoneLimiter : nothing
+    limiter = parsed_args["apply_limiter"] ? CA.QuasiMonotoneLimiter() : nothing
 
     # wrap each upwinding mode in a Val for dispatch
     numerics = AtmosNumerics(;
diff --git a/src/solver/types.jl b/src/solver/types.jl
index e1b8cc8afc..314d88c0fa 100644
--- a/src/solver/types.jl
+++ b/src/solver/types.jl
@@ -270,6 +270,8 @@ abstract type AbstractTimesteppingMode end
 struct Explicit <: AbstractTimesteppingMode end
 struct Implicit <: AbstractTimesteppingMode end
 
+struct QuasiMonotoneLimiter end # For dispatching to use the ClimaCore QuasiMonotoneLimiter.
+
 Base.@kwdef struct AtmosNumerics{
     EN_UP,
     TR_UP,

From 50aae9576a0dfcd5553b6e8655cfe1e7a0607bdf Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 18 Dec 2023 10:32:27 -0500
Subject: [PATCH 3/3] Increase memory for gpu implicit bw

---
 .buildkite/pipeline.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 7ae8ae60c0..c55f2ed248 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -751,6 +751,7 @@ steps:
         artifact_paths: "target_gpu_implicit_baroclinic_wave/*"
         agents:
           slurm_gpus: 1
+          slurm_mem: 32G
 
       - label: "GPU: GPU dry baroclinic wave - 4 gpus"
         key: "target_gpu_implicit_baroclinic_wave_4process"