Make Nv a type parameter

[charleskawczynski]

This PR adds Nv, the number of vertical levels, to the type space in all datalayouts with vertical indices.

This is a step towards #11, and is required for #1746.

[charleskawczynski]

Note to self: import nlevels outside of DataLayouts

[charleskawczynski]

Perhaps this closes #11. I think we don’t actually need Nf because we can compute it at compile time with S and T using typesize.

[charleskawczynski]

I know that there's variation in CI timings, but this nevertheless looks interesting:

        main        this PR
[ Info: 237.470 ns  222.863 ns    <=> t_ave[(:FiniteDifferenceSpace, op_GradientF2C!, :none)]
[ Info: 249.621 ns  248.128 ns    <=> t_ave[(:FiniteDifferenceSpace, op_GradientF2C!, :SetValue, :SetValue)]
[ Info: 265.366 ns  265.147 ns    <=> t_ave[(:FiniteDifferenceSpace, op_GradientC2F!, :SetGradient, :SetGradient)]
[ Info: 260.755 ns  249.049 ns    <=> t_ave[(:FiniteDifferenceSpace, op_GradientC2F!, :SetValue, :SetValue)]
[ Info: 868.150 ns  898.061 ns    <=> t_ave[(:FiniteDifferenceSpace, op_DivergenceF2C!, :none)]
[ Info: 946.190 ns  1.009 μs      <=> t_ave[(:FiniteDifferenceSpace, op_DivergenceF2C!, :Extrapolate, :Extrapolate)]
[ Info: 1.838 μs    2.026 μs      <=> t_ave[(:FiniteDifferenceSpace, op_DivergenceC2F!, :SetDivergence, :SetDivergence)]
[ Info: 233.450 ns  231.981 ns    <=> t_ave[(:FiniteDifferenceSpace, op_InterpolateF2C!, :none)]
[ Info: 254.957 ns  255.297 ns    <=> t_ave[(:FiniteDifferenceSpace, op_InterpolateC2F!, :SetValue, :SetValue)]
[ Info: 266.017 ns  256.917 ns    <=> t_ave[(:FiniteDifferenceSpace, op_InterpolateC2F!, :Extrapolate, :Extrapolate)]
[ Info: 540.719 ns  552.730 ns    <=> t_ave[(:FiniteDifferenceSpace, op_broadcast_example2!, :none)]
[ Info: 170.575 ns  212.516 ns    <=> t_ave[(:FiniteDifferenceSpace, op_LeftBiasedC2F!, :SetValue)]
[ Info: 131.155 ns  125.212 ns    <=> t_ave[(:FiniteDifferenceSpace, op_LeftBiasedF2C!, :none)]
[ Info: 169.429 ns  172.763 ns    <=> t_ave[(:FiniteDifferenceSpace, op_LeftBiasedF2C!, :SetValue)]
[ Info: 183.726 ns  193.870 ns    <=> t_ave[(:FiniteDifferenceSpace, op_RightBiasedC2F!, :SetValue)]
[ Info: 128.836 ns  118.281 ns    <=> t_ave[(:FiniteDifferenceSpace, op_RightBiasedF2C!, :none)]
[ Info: 131.492 ns  129.221 ns    <=> t_ave[(:FiniteDifferenceSpace, op_RightBiasedF2C!, :SetValue)]
[ Info: 2.000 μs    1.688 μs      <=> t_ave[(:FiniteDifferenceSpace, op_CurlC2F!, :SetCurl, :SetCurl)]
[ Info: 1.617 μs    1.514 μs      <=> t_ave[(:FiniteDifferenceSpace, op_CurlC2F!, :SetValue, :SetValue)]
[ Info: 1.743 μs    1.771 μs      <=> t_ave[(:FiniteDifferenceSpace, op_UpwindBiasedProductC2F!, :SetValue, :SetValue)]
[ Info: 1.752 μs    1.876 μs      <=> t_ave[(:FiniteDifferenceSpace, op_UpwindBiasedProductC2F!, :Extrapolate, :Extrapolate)]
[ Info: 2.451 μs    2.452 μs      <=> t_ave[(:FiniteDifferenceSpace, op_divUpwind3rdOrderBiasedProductC2F!, :ThirdOrderOneSided, :ThirdOrderOneSided, :SetValue, :SetValue)]
[ Info: 1.120 μs    1.058 μs      <=> t_ave[(:FiniteDifferenceSpace, op_divgrad_CC!, :SetValue, :SetValue, :none)]
[ Info: 820.049 ns  806.683 ns    <=> t_ave[(:FiniteDifferenceSpace, op_divgrad_FF!, :none, :SetDivergence, :SetDivergence)]
[ Info: 655.541 ns  658.332 ns    <=> t_ave[(:FiniteDifferenceSpace, op_div_interp_CC!, :SetValue, :SetValue, :none)]
[ Info: 620.839 ns  632.233 ns    <=> t_ave[(:FiniteDifferenceSpace, op_div_interp_FF!, :none, :SetValue, :SetValue)]
[ Info: 1.655 μs    1.670 μs      <=> t_ave[(:FiniteDifferenceSpace, op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)]
[ Info: 1.601 μs    1.689 μs      <=> t_ave[(:FiniteDifferenceSpace, op_divgrad_uₕ!, :none, :SetValue, :SetValue)]
[ Info: 4.928 ms    1.669 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_GradientF2C!, :none)]
[ Info: 4.705 ms    1.852 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_GradientF2C!, :SetValue, :SetValue)]
[ Info: 4.718 ms    1.834 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_GradientC2F!, :SetGradient, :SetGradient)]
[ Info: 4.741 ms    2.105 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_GradientC2F!, :SetValue, :SetValue)]
[ Info: 15.411 ms   6.415 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_DivergenceF2C!, :none)]
[ Info: 15.779 ms   6.458 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_DivergenceF2C!, :Extrapolate, :Extrapolate)]
[ Info: 15.116 ms   6.688 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_DivergenceC2F!, :SetDivergence, :SetDivergence)]
[ Info: 5.124 ms    1.485 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_InterpolateF2C!, :none)]
[ Info: 5.116 ms    1.753 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_InterpolateC2F!, :SetValue, :SetValue)]
[ Info: 5.701 ms    1.499 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_InterpolateC2F!, :Extrapolate, :Extrapolate)]
[ Info: 3.957 ms    1.182 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_broadcast_example0!, :none)]
[ Info: 15.379 ms   5.837 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_broadcast_example1!, :none)]
[ Info: 15.261 ms   5.435 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_broadcast_example2!, :none)]
[ Info: 5.595 ms    1.464 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_LeftBiasedC2F!, :SetValue)]
[ Info: 5.584 ms    1.924 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_LeftBiasedF2C!, :none)]
[ Info: 4.788 ms    1.427 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_LeftBiasedF2C!, :SetValue)]
[ Info: 4.713 ms    1.887 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_RightBiasedC2F!, :SetValue)]
[ Info: 4.586 ms    1.643 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_RightBiasedF2C!, :none)]
[ Info: 4.628 ms    1.413 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_RightBiasedF2C!, :SetValue)]
[ Info: 11.524 ms   4.870 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_CurlC2F!, :SetCurl, :SetCurl)]
[ Info: 11.448 ms   4.747 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_CurlC2F!, :SetValue, :SetValue)]
[ Info: 8.319 ms    3.621 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_UpwindBiasedProductC2F!, :SetValue, :SetValue)]
[ Info: 9.108 ms    3.922 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_UpwindBiasedProductC2F!, :Extrapolate, :Extrapolate)]
[ Info: 12.548 ms   5.747 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_divUpwind3rdOrderBiasedProductC2F!, :ThirdOrderOneSided, :ThirdOrderOneSided, :SetValue, :SetValue)]
[ Info: 11.766 ms   4.361 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_divgrad_CC!, :SetValue, :SetValue, :none)]
[ Info: 11.666 ms   4.705 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_divgrad_FF!, :none, :SetDivergence, :SetDivergence)]
[ Info: 9.799 ms    3.452 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_div_interp_CC!, :SetValue, :SetValue, :none)]
[ Info: 9.956 ms    3.724 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_div_interp_FF!, :none, :SetValue, :SetValue)]
[ Info: 18.960 ms   7.432 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_divgrad_uₕ!, :none, :SetValue, :Extrapolate)]
[ Info: 17.221 ms   7.677 ms      <=> t_ave[(:ExtrudedFiniteDifferenceSpace, op_divgrad_uₕ!, :none, :SetValue, :SetValue)]

(these are CPU benchmarks)

[charleskawczynski]

GPU target pipeline was run here: https://buildkite.com/clima/climaatmos-target-gpu-simulations/builds/308.

So, only 1-2% faster on the GPU, perhaps within the noise. After all, this PR is was not intended to introduce a performance improvement, it was only intended to allow us to rebuild DataLayouts inside broadcast expressions with MArrays (#1746).

[charleskawczynski]

I don't think/consider the type parameters of our DataLayouts to be public-facing, so I'd consider these changes to be non-breaking. However, ClimaAtmos has some functions that broke due to this PR. The functions are:

'''
    array2field(array, space)

Wraps `array` in a `ClimaCore` `Field` that is defined over `space`. Can be used
to simplify the process of getting and setting values in an `RRTMGPModel`; e.g.
'''
    array2field(model.center_temperature, center_space) .=
        center_temperature_field
    face_flux_field .= array2field(model.face_flux, face_space)
'''

The dimensions of `array` are assumed to be `([number of vertical nodes], number
of horizontal nodes)`. Also, `array` must represent a `Field` of scalars, so
that the struct type of the resulting `Field` is the same as the element type of
`array`. If this restriction were removed, one would also need to pass the
desired `Field` struct type as an argument to `array2field`, which would then
need to permute the dimensions of `array` to match the target `DataLayout`.
'''
array2field(array, space) =
    Fields.Field(array2data(array, Spaces.local_geometry_data(space)), space)
array2data(
    array::AbstractArray{T, 1},
    ::DataLayouts.IF{<:Any, Ni},
) where {T, Ni} = DataLayouts.IF{T, Ni}(reshape(array, Ni, 1))
array2data(
    array::AbstractArray{T, 1},
    ::DataLayouts.IFH{<:Any, Ni},
) where {T, Ni} = DataLayouts.IFH{T, Ni}(reshape(array, Ni, 1, :))
array2data(
    array::AbstractArray{T, 1},
    ::DataLayouts.IJF{<:Any, Nij},
) where {T, Nij} = DataLayouts.IJF{T, Nij}(reshape(array, Nij, Nij, 1))
array2data(
    array::AbstractArray{T, 1},
    ::DataLayouts.IJFH{<:Any, Nij},
) where {T, Nij} = DataLayouts.IJFH{T, Nij}(reshape(array, Nij, Nij, 1, :))
array2data(
    array::AbstractArray{T, 2},
    ::DataLayouts.VF{<:Any, Nv},
) where {T, Nv} = DataLayouts.VF{T, Nv}(reshape(array, size(array, 1), 1))
array2data(
    array::AbstractArray{T, 2},
    ::DataLayouts.VIFH{<:Any, Nv, Ni},
) where {T, Nv, Ni} =
    DataLayouts.VIFH{T, Nv, Ni}(reshape(array, size(array, 1), Ni, 1, :))
array2data(
    array::AbstractArray{T, 2},
    ::DataLayouts.VIJFH{<:Any, Nv, Nij},
) where {T, Nv, Nij} = DataLayouts.VIJFH{T, Nv, Nij}(
    reshape(array, size(array, 1), Nij, Nij, 1, :),
)

Arguably, these should belong in ClimaCore.