diff --git a/config/model_configs/box_analytic_cosine_mountain_test.yml b/config/model_configs/box_analytic_cosine_mountain_test.yml index cbccbe47d88..2c634847959 100644 --- a/config/model_configs/box_analytic_cosine_mountain_test.yml +++ b/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 diff --git a/config/model_configs/plane_analytic_cosine_mountain_test.yml b/config/model_configs/plane_analytic_cosine_mountain_test.yml index cdfa35f157d..308cd9fb054 100644 --- a/config/model_configs/plane_analytic_cosine_mountain_test.yml +++ b/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 diff --git a/src/disable_topography_dss.jl b/examples/hybrid/disable_topography_dss.jl similarity index 97% rename from src/disable_topography_dss.jl rename to examples/hybrid/disable_topography_dss.jl index a34b37cdbd7..79bbbbd8e19 100644 --- a/src/disable_topography_dss.jl +++ b/examples/hybrid/disable_topography_dss.jl @@ -2,7 +2,6 @@ import ClimaCore: DataLayouts, Geometry, Topologies, - Quadratures, Grids, Hypsography, Spaces, @@ -10,8 +9,7 @@ import ClimaCore: 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, diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index eee538127de..db25901af4f 100644 --- a/examples/hybrid/driver.jl +++ b/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) diff --git a/examples/hybrid/fix_1d_spectral_space_on_gpu.jl b/examples/hybrid/fix_1d_spectral_space_on_gpu.jl new file mode 100644 index 00000000000..80452c8abbf --- /dev/null +++ b/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 diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index a589dfc6644..d67f87a1e0e 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -1,7 +1,5 @@ module ClimaAtmos -# include("disable_topography_dss.jl") - using NVTX, Colors import Thermodynamics as TD diff --git a/src/analytic_solutions/analytic_solutions.jl b/src/analytic_solutions/analytic_solutions.jl index b88d309ea59..15eb0093bf1 100644 --- a/src/analytic_solutions/analytic_solutions.jl +++ b/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, diff --git a/src/utils/common_spaces.jl b/src/utils/common_spaces.jl index 7feb5f38059..f12e89dc3b4 100644 --- a/src/utils/common_spaces.jl +++ b/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