Skip to content

Commit

Permalink
Add kernel benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
charleskawczynski committed Jan 19, 2024
1 parent c5a6676 commit 76a8bb4
Show file tree
Hide file tree
Showing 5 changed files with 201 additions and 2 deletions.
17 changes: 17 additions & 0 deletions .buildkite/pipeline.yml
Original file line number Diff line number Diff line change
Expand Up @@ -87,3 +87,20 @@ steps:
slurm_ntasks: 1
slurm_gpus: 1
timeout_in_minutes: 60

- label: "KernelAbstraction GPU kernel benchmarks"
command:
- "julia --project=gpuenv perf/kernel_bm.jl CuArray"
agents:
slurm_ntasks: 1
slurm_gpus: 1
timeout_in_minutes: 60

- label: "KernelAbstraction CPU kernel benchmarks"
command:
- "julia --project=gpuenv perf/kernel_bm.jl"
agents:
slurm_ntasks: 1
slurm_gpus: 1
timeout_in_minutes: 60

7 changes: 7 additions & 0 deletions perf/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
Expand All @@ -12,3 +13,9 @@ RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"

[compat]
CLIMAParameters = "0.8"
KernelAbstractions = "0.9"
RootSolvers = "0.4"
julia = "1.7"
171 changes: 171 additions & 0 deletions perf/kernel_bm.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,171 @@
#=
using Revise; "CuArray" in ARGS || push!(ARGS, "CuArray"); include("perf/kernel_bm.jl")
using Revise; include("perf/kernel_bm.jl")
=#
using Test

using KernelAbstractions
using BenchmarkTools
import KernelAbstractions as KA
using Random
using LinearAlgebra
import RootSolvers as RS

import Thermodynamics as TD
import Thermodynamics.Parameters as TP
import CLIMAParameters as CP

if get(ARGS, 1, "Array") == "CuArray"
import CUDA
ArrayType = CUDA.CuArray
CUDA.allowscalar(false)
else
ArrayType = Array
end

const param_set_Float64 = TP.ThermodynamicsParameters(Float64)
const param_set_Float32 = TP.ThermodynamicsParameters(Float32)
parameter_set(::Type{Float64}) = param_set_Float64
parameter_set(::Type{Float32}) = param_set_Float32

@show ArrayType

@kernel function test_thermo_kernel!(
param_set,
dst::AbstractArray{FT},
e_int,
ρ,
p,
q_tot,
) where {FT}
i = @index(Global)
@inbounds begin

param_set = parameter_set(FT)
ts = TD.PhaseEquil_ρeq(param_set, FT(ρ[i]), FT(e_int[i]), FT(q_tot[i]))

dst[1, i] = TD.air_pressure(param_set, ts)
dst[2, i] = TD.air_temperature(param_set, ts)
dst[3, i] = TD.air_density(param_set, ts)
dst[4, i] = TD.specific_volume(param_set, ts)
dst[5, i] = TD.soundspeed_air(param_set, ts)
dst[6, i] = TD.total_specific_humidity(param_set, ts)
dst[7, i] = TD.liquid_specific_humidity(param_set, ts)
dst[8, i] = TD.ice_specific_humidity(param_set, ts)
dst[9, i] = TD.vapor_specific_humidity(param_set, ts)
dst[10, i] = TD.total_energy(param_set, ts, FT(0), FT(0))
dst[11, i] = TD.internal_energy(param_set, ts)
dst[12, i] = TD.internal_energy_sat(param_set, ts)
dst[13, i] = TD.internal_energy_dry(param_set, ts)
dst[14, i] = TD.internal_energy_vapor(param_set, ts)
dst[15, i] = TD.internal_energy_liquid(param_set, ts)
dst[16, i] = TD.internal_energy_ice(param_set, ts)
dst[17, i] = TD.cp_m(param_set, ts)
dst[18, i] = TD.cv_m(param_set, ts)
dst[19, i] = TD.gas_constant_air(param_set, ts)
# dst[20, i] = TD.gas_constants(param_set, ts)
dst[21, i] = TD.latent_heat_vapor(param_set, ts)
dst[22, i] = TD.latent_heat_sublim(param_set, ts)
dst[23, i] = TD.latent_heat_fusion(param_set, ts)
dst[24, i] = TD.latent_heat_liq_ice(param_set, ts)
dst[25, i] = TD.saturation_vapor_pressure(param_set, ts, TD.Liquid())
# dst[26, i] = TD.q_vap_saturation_generic(param_set, ts)
dst[27, i] = TD.q_vap_saturation(param_set, ts)
dst[28, i] = TD.q_vap_saturation_liquid(param_set, ts)
dst[29, i] = TD.q_vap_saturation_ice(param_set, ts)
dst[30, i] = TD.saturation_excess(param_set, ts)
dst[31, i] = TD.supersaturation(param_set, ts, TD.Liquid())
dst[32, i] = TD.liquid_fraction(param_set, ts)
dst[33, i] = TD.PhasePartition_equil(param_set, ts).tot
dst[34, i] = TD.dry_pottemp(param_set, ts)
dst[35, i] = TD.virtual_pottemp(param_set, ts)
dst[36, i] = TD.virtual_dry_static_energy(param_set, ts, FT(0))
dst[37, i] = TD.exner(param_set, ts)
# dst[38, i] = TD.shum_to_mixing_ratio(param_set, ts)
dst[39, i] = TD.mixing_ratios(param_set, ts).tot
dst[40, i] = TD.vol_vapor_mixing_ratio(param_set, ts)
dst[41, i] = TD.liquid_ice_pottemp(param_set, ts)
dst[42, i] = TD.liquid_ice_pottemp_sat(param_set, ts)
dst[43, i] = TD.relative_humidity(param_set, ts)
dst[44, i] = TD.virtual_temperature(param_set, ts)
dst[45, i] = TD.condensate(param_set, ts)
dst[46, i] = TD.has_condensate(param_set, ts)
dst[47, i] = TD.specific_enthalpy(param_set, ts)
dst[48, i] = TD.total_specific_enthalpy(param_set, ts, FT(0))
dst[49, i] = TD.moist_static_energy(param_set, ts, FT(0))
dst[50, i] = TD.specific_entropy(param_set, ts)
dst[51, i] = TD.saturated(param_set, ts)

end
end

# Since we use `rand` to generate the ProfileSet,
# just initialize on the CPU, and provide convert
# function to move arrays to the GPU.
convert_profile_set(ps::TD.TestedProfiles.ProfileSet, ArrayType, slice) =
TD.TestedProfiles.ProfileSet(
ArrayType(ps.z[slice]),
ArrayType(ps.T[slice]),
ArrayType(ps.p[slice]),
ArrayType(ps.RS[slice]),
ArrayType(ps.e_int[slice]),
ArrayType(ps.h[slice]),
ArrayType(ps.ρ[slice]),
ArrayType(ps.θ_liq_ice[slice]),
ArrayType(ps.q_tot[slice]),
ArrayType(ps.q_liq[slice]),
ArrayType(ps.q_ice[slice]),
TD.PhasePartition.(ps.q_tot[slice], ps.q_liq[slice], ps.q_ice[slice]),
ArrayType(ps.RH[slice]),
ArrayType(ps.e_pot[slice]),
ArrayType(ps.u[slice]),
ArrayType(ps.v[slice]),
ArrayType(ps.w[slice]),
ArrayType(ps.e_kin[slice]),
ps.phase_type,
)

function test_thermo!(param_set, d_dst, profiles)
(; e_int, ρ, p, q_tot) = profiles
n_profiles = length(profiles.z)
ndrange = (n_profiles,)
backend = KA.get_backend(d_dst)
kernel! = test_thermo_kernel!(backend)
kernel!(param_set, d_dst, e_int, ρ, p, q_tot; ndrange = ndrange)
KA.synchronize(backend)
return nothing
end

@testset "Thermodynamics - kernels" begin
FT = Float32
param_set = parameter_set(FT)
profiles = TD.TestedProfiles.PhaseEquilProfiles(param_set, Array)
slice = Colon()
profiles = convert_profile_set(profiles, ArrayType, slice)

n_profiles = length(profiles.z)
n_vars = length(propertynames(profiles))
d_dst = ArrayType(Array{FT}(undef, 51, n_profiles))
fill!(d_dst, 0)

test_thermo!(param_set, d_dst, profiles) # compile first
trial =
BenchmarkTools.@benchmark test_thermo!($param_set, $d_dst, $profiles)
show(stdout, MIME("text/plain"), trial)

(; e_int, ρ, p, q_tot) = profiles
# Test
ts_cpu =
TD.PhaseEquil_ρeq.(
param_set,
Array{FT}(ρ),
Array{FT}(e_int),
Array{FT}(q_tot),
)
@test all(Array(d_dst)[1, :] .≈ TD.air_pressure.(param_set, ts_cpu))
@test all(Array(d_dst)[2, :] .≈ TD.air_temperature.(param_set, ts_cpu))

end

rm(joinpath(@__DIR__, "logfilepath_Float32.toml"); force = true)
rm(joinpath(@__DIR__, "logfilepath_Float64.toml"); force = true)
6 changes: 5 additions & 1 deletion src/relations.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1385,7 +1385,7 @@ function PhasePartition_equil(
return PhasePartition_equil(param_set, T, ρ, q_tot, p_vap_sat, λ)
end

PhasePartition_equil(param_set::APS, ts::AbstractPhaseNonEquil) =
PhasePartition_equil(param_set::APS, ts::ThermodynamicState) =

Check warning on line 1388 in src/relations.jl

View check run for this annotation

Codecov / codecov/patch

src/relations.jl#L1388

Added line #L1388 was not covered by tests
PhasePartition_equil(
param_set,
air_temperature(param_set, ts),
Expand Down Expand Up @@ -2179,6 +2179,8 @@ function latent_heat_liq_ice(
LH_s0::FT = TP.LH_s0(param_set)
return LH_v0 * q.liq + LH_s0 * q.ice
end
latent_heat_liq_ice(param_set::APS, ts::ThermodynamicState) =

Check warning on line 2182 in src/relations.jl

View check run for this annotation

Codecov / codecov/patch

src/relations.jl#L2182

Added line #L2182 was not covered by tests
latent_heat_liq_ice(param_set, PhasePartition(param_set, ts))


"""
Expand Down Expand Up @@ -2735,6 +2737,8 @@ function vol_vapor_mixing_ratio(
q_vap = vapor_specific_humidity(q)
return molmass_ratio * shum_to_mixing_ratio(q_vap, q.tot)
end
vol_vapor_mixing_ratio(param_set, ts::ThermodynamicState) =

Check warning on line 2740 in src/relations.jl

View check run for this annotation

Codecov / codecov/patch

src/relations.jl#L2740

Added line #L2740 was not covered by tests
vol_vapor_mixing_ratio(param_set, PhasePartition(param_set, ts))

"""
relative_humidity(param_set, T, p, phase_type, q::PhasePartition)
Expand Down
2 changes: 1 addition & 1 deletion test/runtests_gpu.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ convert_profile_set(ps::TD.TestedProfiles.ProfileSet, ArrayType, slice) =
ndrange = (n_profiles,)
backend = KA.get_backend(d_dst)
kernel! = test_thermo_kernel!(backend)
event = kernel!(param_set, d_dst, e_int, ρ, p, q_tot, ndrange = ndrange)
kernel!(param_set, d_dst, e_int, ρ, p, q_tot; ndrange = ndrange)
KA.synchronize(backend)

ts_cpu =
Expand Down

0 comments on commit 76a8bb4

Please sign in to comment.