Add debugging info

config/model_configs/box_analytic_cosine_mountain_test.yml

@@ -9,7 +9,7 @@ x_elem: 30
 y_elem: 30
 z_elem: 100
 z_stretch: false
-dt: "8secs"
+dt: "6secs"
 t_end: "24hours"
 use_reference_state: false
 rayleigh_sponge: true

config/model_configs/plane_analytic_cosine_mountain_test.yml

@@ -7,7 +7,7 @@ z_max: 21e3
 x_elem: 30
 z_elem: 100
 z_stretch: false
-dt: "8secs"
+dt: "6secs"
 t_end: "24hours"
 use_reference_state: false
 rayleigh_sponge: true

src/disable_topography_dss.jl → examples/hybrid/disable_topography_dss.jl

@@ -2,16 +2,14 @@ import ClimaCore:
     DataLayouts,
     Geometry,
     Topologies,
-    Quadratures,
     Grids,
     Hypsography,
     Spaces,
     Fields,
     Operators
 import ClimaTimeSteppers
 import ClimaComms
-using StaticArrays: SVector, SMatrix
-using LinearAlgebra: adjoint, Adjoint, ldiv!, DenseMatrix, lu, norm
+import LinearAlgebra: adjoint, ldiv!, DenseMatrix, lu, norm
 
 ClimaTimeSteppers.NVTX.@annotate function ClimaTimeSteppers.solve_newton!(
     alg::ClimaTimeSteppers.NewtonsMethod,

examples/hybrid/driver.jl

@@ -10,13 +10,8 @@ import ClimaAtmos as CA
 import Random
 Random.seed!(1234)
 
-# TODO: Fix the bug that's making this necessary for 2D topography runs.
-import ClimaCore.DataLayouts: promote_parent_array_type
-import CUDA
-promote_parent_array_type(
-    ::Type{CUDA.CuArray{T1, N, B} where {N}},
-    ::Type{Array{T2, N} where {N}},
-) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), N, B} where {N}
+include("fix_1d_spectral_space_on_gpu.jl")
+# include("disable_topography_dss.jl")
 
 if !(@isdefined config)
     (; config_file, job_id) = CA.commandline_kwargs()
@@ -130,8 +125,10 @@ if config.parsed_args["analytic_check"]
     (; predicted_steady_state) = integrator.p
     @assert !isnothing(predicted_steady_state)
 
-    ᶜuₕ_error = Y_end.c.uₕ .- CA.C12.(integrator.p.predicted_steady_state.ᶜu)
-    ᶠu₃_error = Y_end.f.u₃ .- CA.C3.(integrator.p.predicted_steady_state.ᶠu)
+    ᶜuₕ_error =
+        Geometry.UVWVector.(Y_end.c.uₕ .- CA.C12.(predicted_steady_state.ᶜu))
+    ᶠu₃_error =
+        Geometry.UVWVector.(Y_end.f.u₃ .- CA.C3.(predicted_steady_state.ᶠu))
 
     ᶜsponge_mask = ifelse.(Fields.coordinate_field(Y_end.c).z .< 13e3, 1.0, 0.0)
     ᶠsponge_mask = ifelse.(Fields.coordinate_field(Y_end.f).z .< 13e3, 1.0, 0.0)

examples/hybrid/fix_1d_spectral_space_on_gpu.jl

@@ -0,0 +1,231 @@
+import ClimaCore:
+    ClimaCore,
+    DataLayouts,
+    Geometry,
+    Topologies,
+    Quadratures,
+    Grids,
+    Spaces,
+    Fields,
+    Operators,
+    Hypsography,
+    slab
+import ClimaComms
+import CUDA
+import StaticArrays: SMatrix
+
+ClimaCoreCUDAExt = Base.get_extension(ClimaCore, :ClimaCoreCUDAExt)
+
+# TODO: Move this to ClimaCore.jl/ext/ClimaCoreCUDAExt/data_layouts.jl
+function Base.copyto!(
+    dest::DataLayouts.IFH{S, Ni},
+    bc::Union{
+        DataLayouts.IFH{S, Ni, A},
+        Base.Broadcast.Broadcasted{DataLayouts.IFHStyle{Ni, A}},
+    },
+) where {S, Ni, A <: ClimaCoreCUDAExt.CuArrayBackedTypes}
+    _, _, _, _, Nh = size(bc)
+    if Nh > 0
+        ClimaCoreCUDAExt.auto_launch!(
+            ClimaCoreCUDAExt.knl_copyto!,
+            (dest, bc),
+            dest;
+            threads_s = (Ni, 1),
+            blocks_s = (Nh, 1),
+        )
+    end
+    return dest
+end
+function Base.copyto!(
+    dest::DataLayouts.VIFH{S, Nv, Ni},
+    bc::Union{
+        DataLayouts.VIFH{S, Nv, Ni, A},
+        Base.Broadcast.Broadcasted{DataLayouts.VIFHStyle{Nv, Ni, A}},
+    },
+) where {S, Nv, Ni, A <: ClimaCoreCUDAExt.CuArrayBackedTypes}
+    _, _, _, _, Nh = size(bc)
+    if Nv > 0 && Nh > 0
+        Nv_per_block = min(Nv, fld(256, Ni))
+        Nv_blocks = cld(Nv, Nv_per_block)
+        ClimaCoreCUDAExt.auto_launch!(
+            ClimaCoreCUDAExt.knl_copyto!,
+            (dest, bc),
+            dest;
+            threads_s = (Ni, Nv_per_block),
+            blocks_s = (Nh, Nv_blocks),
+        )
+    end
+    return dest
+end
+
+# TODO: Move this to ClimaCore.jl/ext/ClimaCoreCUDAExt/operators_sem_shmem.jl.jl
+Base.@propagate_inbounds function ClimaCoreCUDAExt.operator_shmem(
+    space,
+    ::Val{Nvt},
+    op::Operators.Gradient{(1,)},
+    arg,
+) where {Nvt}
+    FT = Spaces.undertype(space)
+    QS = Spaces.quadrature_style(space)
+    Nq = Quadratures.degrees_of_freedom(QS)
+    # allocate temp output
+    IT = eltype(arg)
+    input = CUDA.CuStaticSharedArray(IT, (Nq, Nvt))
+    return (input,)
+end
+Base.@propagate_inbounds function ClimaCoreCUDAExt.operator_fill_shmem!(
+    op::Operators.Gradient{(1,)},
+    (input,),
+    space,
+    ij,
+    slabidx,
+    arg,
+)
+    vt = threadIdx().z
+    i, _ = ij.I
+    input[i, vt] = arg
+end
+
+# TODO: Move this to ClimaCore.jl/src/Grids/spectralelement.jl
+function Grids._SpectralElementGrid1D(
+    topology::Topologies.IntervalTopology,
+    quadrature_style::Quadratures.QuadratureStyle,
+)
+    DA = ClimaComms.array_type(topology)
+    global_geometry = Geometry.CartesianGlobalGeometry()
+    CoordType = Topologies.coordinate_type(topology)
+    AIdx = Geometry.coordinate_axis(CoordType)
+    FT = eltype(CoordType)
+    nelements = Topologies.nlocalelems(topology)
+    Nq = Quadratures.degrees_of_freedom(quadrature_style)
+    LG = Geometry.LocalGeometry{AIdx, CoordType, FT, SMatrix{1, 1, FT, 1}}
+    quad_points, quad_weights =
+        Quadratures.quadrature_points(FT, quadrature_style)
+
+    # Initialize the local_geometry with Array{FT} as the backing array type.
+    local_geometry = DataLayouts.IFH{LG, Nq}(Array{FT}, nelements)
+
+    for elem in 1:nelements
+        local_geometry_slab = slab(local_geometry, elem)
+        for i in 1:Nq
+            ξ = quad_points[i]
+            vcoords = Topologies.vertex_coordinates(topology, elem)
+            x = Geometry.linear_interpolate(vcoords, ξ)
+            ∂x∂ξ =
+                (
+                    Geometry.component(vcoords[2], 1) -
+                    Geometry.component(vcoords[1], 1)
+                ) / 2
+            J = abs(∂x∂ξ)
+            WJ = J * quad_weights[i]
+            local_geometry_slab[i] = Geometry.LocalGeometry(
+                x,
+                J,
+                WJ,
+                Geometry.AxisTensor(
+                    (
+                        Geometry.LocalAxis{AIdx}(),
+                        Geometry.CovariantAxis{AIdx}(),
+                    ),
+                    ∂x∂ξ,
+                ),
+            )
+        end
+    end
+
+    # TODO: Why was this code setting dss_weights = 1 / DSS(1)? That's just 1??
+    # dss_weights = copy(local_geometry.J)
+    # dss_weights .= one(FT)
+    # Topologies.dss_1d!(topology, dss_weights)
+    # dss_weights = one(FT) ./ dss_weights
+
+    dss_weights = copy(local_geometry.J)
+    Topologies.dss_1d!(topology, dss_weights)
+    dss_weights .= local_geometry.J ./ dss_weights
+
+    # If needed, move the local_geometry and dss_weights onto the GPU.
+    return Grids.SpectralElementGrid1D(
+        topology,
+        quadrature_style,
+        global_geometry,
+        DataLayouts.rebuild(local_geometry, DA),
+        DataLayouts.rebuild(dss_weights, DA),
+    )
+end
+
+# TODO: Move this to ClimaCore.jl/src/Hypsography/Hypsography.jl
+function Grids._ExtrudedFiniteDifferenceGrid(
+    horizontal_grid::Grids.AbstractGrid,
+    vertical_grid::Grids.FiniteDifferenceGrid,
+    adaption::Grids.HypsographyAdaption,
+    global_geometry::Geometry.AbstractGlobalGeometry,
+)
+    @assert Spaces.grid(axes(adaption.surface)) == horizontal_grid
+    z_surface = Fields.field_values(adaption.surface)
+
+    face_z_ref =
+        Grids.local_geometry_data(vertical_grid, Grids.CellFace()).coordinates
+    vertical_domain = Topologies.domain(vertical_grid)
+    z_top = vertical_domain.coord_max
+
+    face_z =
+        modified_ref_z_to_physical_z.(
+            Ref(typeof(adaption)),
+            face_z_ref,
+            z_surface,
+            Ref(z_top),
+            Ref(adaption_parameters(adaption)),
+        )
+
+    return Grids._ExtrudedFiniteDifferenceGrid(
+        horizontal_grid,
+        vertical_grid,
+        adaption,
+        global_geometry,
+        face_z,
+    )
+end
+
+adaption_parameters(::Grids.Flat) = (;)
+adaption_parameters(::Hypsography.LinearAdaption) = (;)
+adaption_parameters((; ηₕ, s)::Hypsography.SLEVEAdaption) = (; ηₕ, s)
+
+function modified_ref_z_to_physical_z(
+    ::Type{<:Grids.Flat},
+    z_ref::Geometry.ZPoint,
+    z_surface::Geometry.ZPoint,
+    z_top::Geometry.ZPoint,
+    _,
+)
+    return z_ref
+end
+function modified_ref_z_to_physical_z(
+    ::Type{<:Hypsography.LinearAdaption},
+    z_ref::Geometry.ZPoint,
+    z_surface::Geometry.ZPoint,
+    z_top::Geometry.ZPoint,
+    _,
+)
+    Geometry.ZPoint(z_ref.z + (1 - z_ref.z / z_top.z) * z_surface.z)
+end
+function modified_ref_z_to_physical_z(
+    ::Type{<:Hypsography.SLEVEAdaption},
+    z_ref::Geometry.ZPoint,
+    z_surface::Geometry.ZPoint,
+    z_top::Geometry.ZPoint,
+    (; ηₕ, s),
+)
+    if s * z_top.z <= z_surface.z
+        error("Decay scale (s*z_top) must be higher than max surface elevation")
+    end
+
+    η = z_ref.z / z_top.z
+    if η <= ηₕ
+        return Geometry.ZPoint(
+            η * z_top.z +
+            z_surface.z * (sinh((ηₕ - η) / s / ηₕ)) / (sinh(1 / s)),
+        )
+    else
+        return Geometry.ZPoint(η * z_top.z)
+    end
+end

src/ClimaAtmos.jl

@@ -1,7 +1,5 @@
 module ClimaAtmos
 
-# include("disable_topography_dss.jl")
-
 using NVTX, Colors
 import Thermodynamics as TD
 

src/analytic_solutions/analytic_solutions.jl

@@ -43,9 +43,6 @@ function initial_thermo_state(params, coord)
     T_min = FT(100)
     z_iso = log((T_min / T_sfc - a) / (1 - a)) / β
     p_iso = p_sfc * ((1 - a) / (1 - a * T_sfc / T_min))^(cp_d / R_d)
-    z > z_iso &&
-        @warn "Extending constant-N profile with isothermal profile above T = \
-               $T_min (z = $z_iso)" maxlog = 1
 
     # The perturbation in approximate_FΔuvw is computed around a background
     # state that is not hydrostatically balanced and is unphysical above 30 km,

src/utils/common_spaces.jl

@@ -120,6 +120,14 @@ function make_hybrid_spaces(
             hypsography =
                 Hypsography.LinearAdaption(Geometry.ZPoint.(z_surface))
         end
+        @show device
+        @show ClimaComms.device(Fields.coordinate_field(h_space))
+        @show ClimaComms.device(z_grid)
+        @show ClimaComms.device(z_surface)
+        @show ClimaComms.device(Geometry.ZPoint.(z_surface))
+        @show typeof(h_space)
+        @show typeof(z_grid)
+        @show typeof(z_surface)
     end
     grid = Grids.ExtrudedFiniteDifferenceGrid(h_grid, z_grid, hypsography; deep)
     # TODO: return the grid