From f75329612a4a77bff2ccb5f97d84743fa285d897 Mon Sep 17 00:00:00 2001 From: Anna Jaruga Date: Wed, 11 Dec 2024 13:26:45 -0800 Subject: [PATCH] Add ClimaCore GPU test --- .buildkite/pipeline.yml | 9 ++ test/Project.toml | 1 + test/gpu_clima_core_test.jl | 175 ++++++++++++++++++++++++++++++++++++ 3 files changed, 185 insertions(+) create mode 100644 test/gpu_clima_core_test.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 883489975..bb423c9f5 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -40,3 +40,12 @@ steps: CLIMACOMMS_DEVICE: "CUDA" agents: slurm_gpus: 1 + + - label: ":cyclone: ClimaCore + GPU unit tests" + key: "clima_core_gpu_unittests" + command: + - "julia --project=test --color=yes test/gpu_clima_core_test.jl" + env: + CLIMACOMMS_DEVICE: "CUDA" + agents: + slurm_gpus: 1 diff --git a/test/Project.toml b/test/Project.toml index c5c244f47..11b954321 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" +ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884" ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c" CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" diff --git a/test/gpu_clima_core_test.jl b/test/gpu_clima_core_test.jl new file mode 100644 index 000000000..fd6475c60 --- /dev/null +++ b/test/gpu_clima_core_test.jl @@ -0,0 +1,175 @@ +#ENV["CLIMACOMMS_DEVICE"] = "CUDA" + +import ClimaComms +ClimaComms.@import_required_backends +import ClimaCore as CC + +import ClimaParams as CP +import CloudMicrophysics.Parameters as CMP +import CloudMicrophysics.MicrophysicsNonEq as CMN + +""" + A helper function to create a ClimaCore 1d column space +""" +function make_column(::Type{FT}) where {FT} + + context = ClimaComms.SingletonCommsContext(ClimaComms.CUDADevice()) + + vert_domain = CC.Domains.IntervalDomain( + CC.Geometry.ZPoint{FT}(FT(0)), + CC.Geometry.ZPoint{FT}(FT(1000)); + boundary_names = (:bottom, :top), + ) + vert_mesh = CC.Meshes.IntervalMesh(vert_domain; nelems = 1000) + vert_topology = CC.Topologies.IntervalTopology(context, vert_mesh) + vert_center_space = CC.Spaces.CenterFiniteDifferenceSpace(vert_topology) + return vert_center_space +end + +""" + A helper function to create a ClimaCore extruded sphere space +""" +function make_extruded_sphere(::Type{FT}) where {FT} + + context = ClimaComms.SingletonCommsContext(ClimaComms.CUDADevice()) + + # Define vertical + # domain + vert_domain = CC.Domains.IntervalDomain( + CC.Geometry.ZPoint{FT}(FT(0)), + CC.Geometry.ZPoint{FT}(FT(1000)); + boundary_names = (:bottom, :top), + ) + # mesh + vert_mesh = CC.Meshes.IntervalMesh(vert_domain; nelems = 1000) + # topology + vert_topology = CC.Topologies.IntervalTopology(context, vert_mesh) + # grid + vert_grid = CC.Grids.FiniteDifferenceGrid(vert_topology) + #vert_center_space = CC.Spaces.CenterFiniteDifferenceSpace(vert_topology) + + # Define horizontal: + # domain + horz_domain = CC.Domains.SphereDomain(FT(30)) + # mesh + horz_mesh = CC.Meshes.EquiangularCubedSphere(horz_domain, 4) + # topology + horz_topology = CC.Topologies.Topology2D( + context, + horz_mesh, + CC.Topologies.spacefillingcurve(horz_mesh), + ) + # space + horz_space = CC.Spaces.SpectralElementSpace2D( + horz_topology, + CC.Quadratures.GLL{3 + 1}(); + enable_bubble = true, + ) + # grid + horz_grid = CC.Spaces.grid(horz_space) + + # Define surface + z_surface = zeros(horz_space) + hypsography = CC.Hypsography.Flat() + + # Define grid + deep = false + grid = CC.Grids.ExtrudedFiniteDifferenceGrid( + horz_grid, + vert_grid, + hypsography; + deep, + ) + + # Define 3D space + center_extruded_space = CC.Spaces.CenterExtrudedFiniteDifferenceSpace(grid) + return center_extruded_space +end + +""" + Try to reproduce the setup of how terminal velocity is used in Atmos +""" +function set_sedimentation_precomputed_quantities(Y, p, t) + + (; w) = p + (; params) = p + + @. w = CMN.terminal_velocity( + params.liquid, + params.Ch2022.rain, + Y.ρ, + max(0, Y.ρq / Y.ρ), + ) + return nothing +end + +function main_1d(::Type{FT}) where {FT} + + Ch2022 = CMP.Chen2022VelType(FT) + liquid = CMP.CloudLiquid(FT) + ice = CMP.CloudIce(FT) + rain = CMP.Rain(FT) + snow = CMP.Snow(FT) + + params = (; liquid, ice, Ch2022) + + space_1d_ρq = make_column(FT) + space_1d_ρ = make_column(FT) + space_1d_w = make_column(FT) + + ρq = CC.Fields.ones(space_1d_ρq) .* FT(1e-3) + ρ = CC.Fields.ones(space_1d_ρ) + w = CC.Fields.zeros(space_1d_w) + + Y = (; ρq, ρ) + p = (; w, params) + + t = 1 + + set_sedimentation_precomputed_quantities(Y, p, t) + return nothing +end + +function main_3d(::Type{FT}) where {FT} + + Ch2022 = CMP.Chen2022VelType(FT) + liquid = CMP.CloudLiquid(FT) + ice = CMP.CloudIce(FT) + rain = CMP.Rain(FT) + snow = CMP.Snow(FT) + + params = (; liquid, ice, Ch2022) + + space_3d_ρq = make_extruded_sphere(FT) + space_3d_ρ = make_extruded_sphere(FT) + space_3d_w = make_extruded_sphere(FT) + + ρq = CC.Fields.ones(space_3d_ρq) .* FT(1e-3) + ρ = CC.Fields.ones(space_3d_ρ) + w = CC.Fields.zeros(space_3d_w) + + Y = (; ρq, ρ) + p = (; w, params) + + t = 1 + + set_sedimentation_precomputed_quantities(Y, p, t) + return nothing +end + +using Test +@testset "GPU inference failure 1D Float64" begin + main_1d(Float64) +end + +@testset "GPU inference failure 3D Float64" begin + main_3d(Float64) +end + +@testset "GPU inference failure 1D Float32" begin + main_1d(Float32) +end + +@testset "GPU inference failure 3D Float32" begin + main_3d(Float32) +end