GPU-friendly param structs for many parameters

[odunbar]

Issue

Details found in #278 . Here we seek to address the issue of a parameter struct ps_gpu:

• whose items can contain float-values and flat (possibly large) vectors of float-values.
• can live in GPU device memory
• can store O(>100s) items
• is readable through the current functional interface of getting values via alpha(ps_gpu)
• can handle the derived parameters e.g. gamma(ps_gpu) = alpha(ps_gpu) / beta(ps_gpu)

The current setup

• Various parameter box structs (created in each repository) under the filenames parameters.jl or similar, (e.g., ClimaLand, Thermodynamics, ClimaAtmos etc.) this file contains
  • The build and specification of the parameter structs from file

struct ps::ParamSet{FT}
    alpha::FT
    beta::Vector{FT}
end

• The methods obtaining parameters from the struct via a call

alpha(ps::ParamSet) = ps.alpha
beta(ps::ParamSet) = ps.beta

• The methods to obtain derived parameters

gamma(ps::ParamSet) = alpha(ps) ./ beta(ps)

A Centralized Solution (?)

The one benefit of conversion, is that it could be doable from one centralized location
Converting the ps -> ps_gpu would require, for each ps.
Getting the values, and their sizes
Flattening

• Create a new value-shift struct

struct ParamsGPU <: ParamSet
    values  :: CuArray{Float32,1}
    offsets :: CuArray{Int32,1}
end

For these new parameter structs to be passed around, they would need to be subtypes of the original type:
Redefine the methods for the new types

alpha(ps_gpu::ParamsGPU) = view(ps_gpu.values, ps_gpu.offsets[1]:ps_gpu.offsets[2]-1)
beta(ps_gpu::ParamsGPU)  = view(ps_gpu.values, ps_gpu.offsets[2]:ps_gpu.offsets[3]-1)
gamma(ps_gpu::ParamsGPU) = alpha(ps_gpu)/beta(ps_gpu)

Question: Can such a process be automated?

A Distributed solution (?)

One other method, could be, more of a switch at the point of building the ParamSet. In the case of running Clima on GPU, we simply define the counterpart ParamsGPU struct and alpha(), beta(), gamma() methods. This is clean as we never build the originally defined ParamSet, nor need to convert anything. However, this is distributed, as all ParamSets are defined in every package.

In the respective parameters.jl file, we would define the GPU parameter structs replacing/alongside the current structs/methods. As we have one entry point for each package, this should be a straightforward switch, but requires a lot of communication between the teams in Clima for the update.

[odunbar]

Two "Broadcastable"-based implementations

[simple, distributed - dependent on compiler optimization]

In every repo a function

broadcastable(ps::ParamSet) = tuple(ps)

is defined, which is what causes the default behaviour to dump into parameter memory. We redefine this for 1-d Cu-Array

# already created
struct ps::ParamSet{FT}
    alpha::FT
    beta::Tuple{FT}
end
param = ParamSet(1.0,2.0)

# new addition
broadcastable(ps::ParamSet) = CuArray([ps])

f.(x, param)

function f(x, params)
alpha = alpha(params)
end

Sticking point. Does julia read in the entirity of params into register memory during the compute, or does it just only read the alpha as it needs.

[more complicated, but quite centralized - manually bypassing default struct reading]

replace_basetype(cu_array, T) = # define somewhere in ClimaCore.DataLayouts

# in the constructor for ClimaAtmosParameters:
    return ClimaCore.DataLayouts.DataF{typeof(ps)}(replace_basetype(CuArray([ps]), eltype(ps)))

[odunbar]

Current Status:

1. Merge of ClimaCore #2464 should resolve this issue at the core. [DONE]
2. Propagate the change to ClimaAtmos [AWAITING ClimaAtmos Update]

[odunbar]

CliMA/ClimaAtmos.jl#4336