Make cuda an extension

.buildkite/pipeline.yml

@@ -44,6 +44,8 @@ steps:
 
       - label: "GPU unit tests"
         command: "julia --color=yes --project=test test/runtests_gpu.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_ntasks: 1
           slurm_gres: "gpu:1"
@@ -95,12 +97,16 @@ steps:
 
       - label: "GPU benchmarks"
         command: "julia --color=yes --project=test perf/benchmark.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_ntasks: 1
           slurm_gres: "gpu:1"
 
       - label: "GPU All-sky DYAMOND benchmark"
         command: "julia --color=yes --project=test test/all_sky_dyamond_gpu_benchmark.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_ntasks: 1
           slurm_gres: "gpu:1"

Project.toml

@@ -6,7 +6,6 @@ version = "0.13.2"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 GaussQuadrature = "d54b0c1a-921d-58e0-8e36-89d8069c0969"
@@ -15,16 +14,18 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 
 [weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 
 [extensions]
 CreateParametersExt = "ClimaParams"
+RRTMGPCUDAExt = "CUDA"
 
 [compat]
 Adapt = "3.3, 4"
 Artifacts = "1"
 CUDA = "4, 5"
-ClimaComms = "0.5.1"
+ClimaComms = "0.5.1, 0.6"
 DocStringExtensions = "0.8, 0.9"
 GaussQuadrature = "0.5"
 Random = "1"

ext/RRTMGPCUDAExt.jl

@@ -0,0 +1,47 @@
+module RRTMGPCUDAExt
+
+import ClimaComms
+import RRTMGP.Parameters as RP
+import RRTMGP.Fluxes: FluxLW, FluxSW
+import RRTMGP.Fluxes: add_to_flux!
+import RRTMGP.Fluxes: set_flux_to_zero!
+import RRTMGP.Sources: SourceLWNoScat
+import RRTMGP.Sources: SourceLW2Str
+import RRTMGP.Sources: SourceSW2Str
+import RRTMGP.BCs: LwBCs, SwBCs
+import RRTMGP.Optics: OneScalar, TwoStream
+import RRTMGP.Optics
+import RRTMGP.Optics: compute_optical_props!
+import RRTMGP.Optics: compute_col_gas!
+import RRTMGP.Optics: compute_col_gas_kernel!
+import RRTMGP.AtmosphericStates: setup_gray_as_pr_grid!
+import RRTMGP.AtmosphericStates: setup_gray_as_pr_grid_kernel!
+import RRTMGP.AtmosphericStates: GrayAtmosphericState
+import RRTMGP.AtmosphericStates: AtmosphericState
+import RRTMGP.AtmosphericStates: CloudState
+import RRTMGP.AtmosphericStates
+import RRTMGP.GrayUtils: update_profile_lw!
+import RRTMGP.GrayUtils: compute_gray_heating_rate!
+import RRTMGP.GrayUtils: compute_gray_heating_rate_kernel!
+import RRTMGP.GrayUtils: update_profile_lw_kernel!
+import RRTMGP.LookUpTables: LookUpLW, LookUpCld, PadeCld, LookUpSW
+import RRTMGP.RTESolver: rte_lw_noscat_solve!
+import RRTMGP.RTESolver: rte_lw_noscat!
+import RRTMGP.RTESolver: rte_lw_2stream_solve!
+import RRTMGP.RTESolver: rte_lw_2stream!
+import RRTMGP.RTESolver: rte_sw_2stream_solve!
+import RRTMGP.RTESolver: rte_sw_2stream!
+import RRTMGP.RTESolver: rte_sw_noscat!
+import RRTMGP.RTESolver: rte_sw_noscat_solve!
+import CUDA: threadIdx, blockIdx, blockDim, @cuda
+
+include(joinpath("cuda", "gray_atmospheric_states.jl"))
+include(joinpath("cuda", "rte_longwave_2stream.jl"))
+include(joinpath("cuda", "optics.jl"))
+include(joinpath("cuda", "rte_shortwave_1scalar.jl"))
+include(joinpath("cuda", "optics_gray_utils.jl"))
+include(joinpath("cuda", "rte_shortwave_2stream.jl"))
+include(joinpath("cuda", "rte_longwave_1scalar.jl"))
+
+
+end

ext/cuda/gray_atmospheric_states.jl

@@ -0,0 +1,15 @@
+
+function setup_gray_as_pr_grid!(::ClimaComms.CUDADevice, ncol, args...)
+    max_threads = 256
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    @cuda always_inline = true threads = (tx) blocks = (bx) _setup_gray_as_pr_grid_kernel!(ncol, args...)
+end
+
+function _setup_gray_as_pr_grid_kernel!(ncol, args...)
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    if gcol ≤ ncol
+        setup_gray_as_pr_grid_kernel!(args..., gcol)
+    end
+    return nothing
+end

ext/cuda/optics.jl

@@ -0,0 +1,32 @@
+
+function compute_col_gas!(
+    device::ClimaComms.CUDADevice,
+    p_lev::AbstractArray{FT, 2},
+    col_dry::AbstractArray{FT, 2},
+    param_set::RP.ARP,
+    vmr_h2o::Union{AbstractArray{FT, 2}, Nothing} = nothing,
+    lat::Union{AbstractArray{FT, 1}, Nothing} = nothing,
+    max_threads::Int = Int(256),
+) where {FT}
+    nlay, ncol = size(col_dry)
+    mol_m_dry = RP.molmass_dryair(param_set)
+    mol_m_h2o = RP.molmass_water(param_set)
+    avogadro = RP.avogad(param_set)
+    helmert1 = RP.grav(param_set)
+    args = (p_lev, mol_m_dry, mol_m_h2o, avogadro, helmert1, vmr_h2o, lat)
+    tx = min(nlay * ncol, max_threads)
+    bx = cld(nlay * ncol, tx)
+    @cuda always_inline = true threads = (tx) blocks = (bx) compute_col_gas_CUDA!(col_dry, args...)
+    return nothing
+end
+
+function compute_col_gas_CUDA!(col_dry, args...)
+    glx = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlay, ncol = size(col_dry)
+    if glx ≤ nlay * ncol
+        glay = (glx % nlay == 0) ? nlay : (glx % nlay)
+        gcol = cld(glx, nlay)
+        compute_col_gas_kernel!(col_dry, args..., glay, gcol)
+    end
+    return nothing
+end

ext/cuda/optics_gray_utils.jl

@@ -0,0 +1,59 @@
+
+function update_profile_lw!(
+    device::ClimaComms.CUDADevice,
+    sbc,
+    t_lay,
+    t_lev,
+    flux_dn,
+    flux_net,
+    hr_lay,
+    flux_grad,
+    T_ex_lev,
+    Δt,
+    nlay,
+    nlev,
+    ncol,
+)
+    max_threads = 256
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    @cuda always_inline = true threads = (tx) blocks = (bx) _update_profile_lw_kernel!(
+        sbc,
+        ncol,
+        t_lay,
+        t_lev,
+        hr_lay,
+        flux_dn,
+        flux_net,
+        flux_grad,
+        T_ex_lev,
+        Δt,
+        nlay,
+        nlev,
+    )
+end
+
+function _update_profile_lw_kernel!(sbc, ncol, args...)
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    if gcol ≤ ncol
+        update_profile_lw_kernel!(sbc, args..., gcol)
+    end
+    return nothing
+end
+
+function compute_gray_heating_rate!(device::ClimaComms.CUDADevice, hr_lay, p_lev, ncol, nlay, flux_net, cp_d_, grav_)
+    args = (hr_lay, flux_net, p_lev, grav_, cp_d_, nlay)
+    max_threads = 256
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    @cuda always_inline = true threads = (tx) blocks = (bx) compute_gray_heating_rate_CUDA!(ncol, args...)
+    return nothing
+end
+
+function compute_gray_heating_rate_CUDA!(ncol, args...)
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    if gcol ≤ ncol
+        compute_gray_heating_rate_kernel!(args..., gcol)
+    end
+    return nothing
+end

ext/cuda/rte_longwave_1scalar.jl

@@ -0,0 +1,101 @@
+function rte_lw_noscat_solve!(
+    device::ClimaComms.CUDADevice,
+    flux_lw::FluxLW,
+    src_lw::SourceLWNoScat,
+    bcs_lw::LwBCs,
+    op::OneScalar,
+    max_threads,
+    as::GrayAtmosphericState,
+)
+    nlay, ncol = AtmosphericStates.get_dims(as)
+    nlev = nlay + 1
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux_lw, src_lw, bcs_lw, op, nlay, ncol, as)
+    @cuda always_inline = true threads = (tx) blocks = (bx) rte_lw_noscat_solve_CUDA!(args...)
+    return nothing
+end
+
+function rte_lw_noscat_solve_CUDA!(
+    flux_lw::FluxLW,
+    src_lw::SourceLWNoScat,
+    bcs_lw::LwBCs,
+    op::OneScalar,
+    nlay,
+    ncol,
+    as::GrayAtmosphericState,
+)
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlev = nlay + 1
+    igpt, ibnd = 1, 1
+    τ = op.τ
+    Ds = op.angle_disc.gauss_Ds
+    (; flux_up, flux_dn, flux_net) = flux_lw
+    if gcol ≤ ncol
+        compute_optical_props!(op, as, src_lw, gcol)
+        rte_lw_noscat!(src_lw, bcs_lw, op, gcol, flux_lw, igpt, ibnd, nlay, nlev)
+        @inbounds for ilev in 1:nlev
+            flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
+        end
+    end
+    return nothing
+end
+
+function rte_lw_noscat_solve!(
+    device::ClimaComms.CUDADevice,
+    flux::FluxLW,
+    flux_lw::FluxLW,
+    src_lw::SourceLWNoScat,
+    bcs_lw::LwBCs,
+    op::OneScalar,
+    max_threads,
+    as::AtmosphericState,
+    lookup_lw::LookUpLW,
+    lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+)
+    nlay, ncol = AtmosphericStates.get_dims(as)
+    nlev = nlay + 1
+    tx = min(ncol, max_threads)
+    bx = cld(ncol, tx)
+    args = (flux, flux_lw, src_lw, bcs_lw, op, nlay, ncol, as, lookup_lw, lookup_lw_cld)
+    @cuda always_inline = true threads = (tx) blocks = (bx) rte_lw_noscat_solve_CUDA!(args...)
+    return nothing
+end
+
+function rte_lw_noscat_solve_CUDA!(
+    flux::FluxLW,
+    flux_lw::FluxLW,
+    src_lw::SourceLWNoScat,
+    bcs_lw::LwBCs,
+    op::OneScalar,
+    nlay,
+    ncol,
+    as::AtmosphericState,
+    lookup_lw::LookUpLW,
+    lookup_lw_cld::Union{LookUpCld, PadeCld, Nothing} = nothing,
+)
+    gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
+    nlev = nlay + 1
+    (; major_gpt2bnd) = lookup_lw.band_data
+    n_gpt = length(major_gpt2bnd)
+    τ = op.τ
+    Ds = op.angle_disc.gauss_Ds
+    if gcol ≤ ncol
+        flux_up_lw = flux_lw.flux_up
+        flux_dn_lw = flux_lw.flux_dn
+        flux_net_lw = flux_lw.flux_net
+        @inbounds for igpt in 1:n_gpt
+            ibnd = major_gpt2bnd[igpt]
+            igpt == 1 && set_flux_to_zero!(flux_lw, gcol)
+            compute_optical_props!(op, as, src_lw, gcol, igpt, lookup_lw, lookup_lw_cld)
+            rte_lw_noscat!(src_lw, bcs_lw, op, gcol, flux, igpt, ibnd, nlay, nlev)
+            add_to_flux!(flux_lw, flux, gcol)
+        end
+        @inbounds begin
+            for ilev in 1:nlev
+                flux_net_lw[ilev, gcol] = flux_up_lw[ilev, gcol] - flux_dn_lw[ilev, gcol]
+            end
+        end
+    end
+    return nothing
+end