DataLayouts

ClimaCore.DataLayouts

Defines the following DataLayouts (see individual docs for more info):

TODO: Add links to these datalayouts

• IJKFVH
• IJFH
• IJHF
• IFH
• IHF
• DataF
• IJF
• IF
• VF
• VIJFH
• VIJHF
• VIFH
• VIHF
• IH1JH2
• IV1JH2

Notation:

• i,j are horizontal node indices within an element
• k is the vertical node index within an element
• f is the field index (1 if field is scalar, >1 if it is a vector field)
• v is the vertical element index in a stack
• h is the element stack index

Data layout is specified by the order in which they appear, e.g. IJKFVH indexes the underlying array as [i,j,k,f,v,h]

Datalayouts that end with the field index

One of the fundamental features of datalayouts is to be able to store multiple variables in the same array, and then access those variables by name. As such, we occasionally must index into multiple variables when performing operations with a datalayout.

We can efficiently support linear indexing with datalayouts whose field index (f) is first or last. This is for the same reason as https://docs.julialang.org/en/v1/devdocs/subarrays/#Linear-indexing:

Linear indexing can be implemented efficiently when the entire array
has a single stride that separates successive elements, starting from
some offset.

Therefore, we provide special handling for these datalayouts where possible to leverage efficient linear indexing.

Here are some references containing relevant discussions and efforts to leverage efficient linear indexing:

• https://github.com/CliMA/ClimaCore.jl/issues/1889
• https://github.com/JuliaLang/julia/issues/28126
• https://github.com/JuliaLang/julia/issues/32051
• https://github.com/maleadt/StaticCartesian.jl
• https://github.com/JuliaGPU/GPUArrays.jl/pull/454#issuecomment-1431575721
• https://github.com/JuliaGPU/GPUArrays.jl/pull/520
• https://github.com/JuliaGPU/GPUArrays.jl/pull/464

DataF{S, A} <: Data0D{S}

Backing DataLayout for 0D point data.

DataF{S}(ArrayType[, ones | zeros | rand])

The ArrayType constructor returns a DataF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand).

IF{S, Ni, A} <: DataSlab1D{S, Ni}

Backing DataLayout for 1D spectral element slab data.

Nodal element data (I) are contiguous for each S datatype struct field (F) for a single element slab.

A DataSlab1D view can be returned from other Data1D objects by calling slab(data, idx...).

IF{S}(ArrayType[, ones | zeros | rand]; Ni)

The keyword constructor returns a IF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Ni quadrature degrees of freedom in the horizontal direction

IJF{S, Nij, A} <: DataSlab2D{S, Nij}

Backing DataLayout for 2D spectral element slab data.

Nodal element data (I,J) are contiguous for each S datatype struct field (F) for a single element slab.

A DataSlab2D view can be returned from other Data2D objects by calling slab(data, idx...).

IJF{S}(ArrayType[, ones | zeros | rand]; Nij)

The keyword constructor returns a IJF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nij quadrature degrees of freedom per horizontal direction

VF{S, A} <: DataColumn{S, Nv}

Backing DataLayout for 1D FV column data.

Column level data (V) are contiguous for each S datatype struct field (F).

A DataColumn view can be returned from other Data1DX, Data2DX objects by calling column(data, idx...).

VF{S}(ArrayType[, ones | zeros | rand]; Nv)

The keyword constructor returns a VF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nv number of vertical degrees of freedom

IFH{S,Ni,Nh,A} <: Data1D{S, Ni}
IFH{S,Ni,Nh}(ArrayType)

Backing DataLayout for 1D spectral element slabs.

Element nodal point (I) data is contiguous for each datatype S struct field (F), for each 1D mesh element (H).

The ArrayType-constructor makes a IFH 1D Spectral DataLayout given the backing ArrayType, quadrature degrees of freedom Ni, and the number of mesh elements Nh.

IFH{S}(ArrayType[, ones | zeros | rand]; Ni, Nh)

The keyword constructor returns a IFH given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Ni quadrature degrees of freedom in the horizontal direction
• Nh number of mesh elements

IJFH{S, Nij, A} <: Data2D{S, Nij}
IJFH{S,Nij}(ArrayType, nelements)

Backing DataLayout for 2D spectral element slabs.

Element nodal point (I,J) data is contiguous for each datatype S struct field (F), for each 2D mesh element slab (H).

The ArrayType-constructor constructs a IJFH 2D Spectral DataLayout given the backing ArrayType, quadrature degrees of freedom Nij × Nij, and the number of mesh elements nelements.

IJFH{S}(ArrayType[, Base.ones | zeros | rand]; Nij, Nh)

The keyword constructor returns a IJFH given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nij quadrature degrees of freedom per horizontal direction
• Nh number of mesh elements

VIFH{S, Nv, Ni, A} <: Data1DX{S, Nv, Ni}

Backing DataLayout for 1D spectral element slab + extruded 1D FV column data.

Column levels (V) are contiguous for every element nodal point (I) for each datatype S struct field (F), for each 1D mesh element slab (H).

VIFH{S}(ArrayType[, ones | zeros | rand]; Nv, Ni, Nh)

The keyword constructor returns a VIFH given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nv number of vertical degrees of freedom
• Ni quadrature degrees of freedom in the horizontal direction
• Nh number of horizontal elements

VIJFH{S, Nij, A} <: Data2DX{S, Nij}

Backing DataLayout for 2D spectral element slab + extruded 1D FV column data.

Column levels (V) are contiguous for every element nodal point (I, J) for each S datatype struct field (F), for each 2D mesh element slab (H).

VIJFH{S}(ArrayType[, ones | zeros | rand]; Nv, Nij, Nh)

The keyword constructor returns a VIJFH given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nv number of vertical degrees of freedom
• Nij quadrature degrees of freedom per horizontal direction
• Nh number of horizontal elements

IHF{S,Ni,Nh,A} <: Data1D{S, Ni}
IHF{S,Ni,Nh}(ArrayType)

Backing DataLayout for 1D spectral element slabs.

Element nodal point (I) data is contiguous for each datatype S struct field (F), for each 1D mesh element (H).

The ArrayType-constructor makes a IHF 1D Spectral DataLayout given the backing ArrayType, quadrature degrees of freedom Ni, and the number of mesh elements Nh.

IHF{S}(ArrayType[, ones | zeros | rand]; Ni, Nh)

The keyword constructor returns a IHF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Ni quadrature degrees of freedom in the horizontal direction
• Nh number of mesh elements

IJHF{S, Nij, A} <: Data2D{S, Nij}
IJHF{S,Nij}(ArrayType, nelements)

Backing DataLayout for 2D spectral element slabs.

Element nodal point (I,J) data is contiguous for each datatype S struct field (F), for each 2D mesh element slab (H).

The ArrayType-constructor constructs a IJHF 2D Spectral DataLayout given the backing ArrayType, quadrature degrees of freedom Nij × Nij, and the number of mesh elements nelements.

IJHF{S}(ArrayType[, Base.ones | zeros | rand]; Nij, Nh)

The keyword constructor returns a IJHF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nij quadrature degrees of freedom per horizontal direction
• Nh number of mesh elements

VIHF{S, Nv, Ni, A} <: Data1DX{S, Nv, Ni}

Backing DataLayout for 1D spectral element slab + extruded 1D FV column data.

Column levels (V) are contiguous for every element nodal point (I) for each datatype S struct field (F), for each 1D mesh element slab (H).

VIHF{S}(ArrayType[, ones | zeros | rand]; Nv, Ni, Nh)

The keyword constructor returns a VIHF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nv number of vertical degrees of freedom
• Ni quadrature degrees of freedom in the horizontal direction
• Nh number of horizontal elements

VIJHF{S, Nij, A} <: Data2DX{S, Nij}

Backing DataLayout for 2D spectral element slab + extruded 1D FV column data.

Column levels (V) are contiguous for every element nodal point (I, J) for each S datatype struct field (F), for each 2D mesh element slab (H).

VIJHF{S}(ArrayType[, ones | zeros | rand]; Nv, Nij, Nh)

The keyword constructor returns a VIJHF given the ArrayType and (optionally) an initialization method (one of Base.ones, Base.zeros, Random.rand) and the keywords:

• Nv number of vertical degrees of freedom
• Nij quadrature degrees of freedom per horizontal direction
• Nh number of horizontal elements

bitcast_struct(T, value)
bitcast_struct(T, array, Val(num_indices), index...)

Converts value into an isbits type T that spans the same number of bytes (counting all bytes that are used as padding; see extended help for details). Serves as a GPU-compatible generalization of the native Core.bitcast function, losslessly converting between arbitrary data types, including composite types.

Instead of converting a single value, it is also possible to convert a subset of an array corresponding to the result of get_struct. This is equivalent to converting the array elements after first loading them into a tuple, but with guaranteed inlining for arbitrary data types. Inlining is necessary for the compiler's getfield_elim_pass! to eliminate reads of array elements for unused fields of T (a key optimization in GPU kernels, where reads from global memory can be relatively expensive).

Examples

julia> bitcast_struct(NTuple{4, Int8}, Int32(1))
(1, 0, 0, 0)

julia> bitcast_struct(NTuple{6, Int32}, (2 * eps(0.0), eps(0.0), 0.0))
(2, 0, 1, 0, 0, 0)

julia> bitcast_struct(Tuple{Int32, Int32, Int128}, (2, 0, 1, 0))
(2, 0, 1)

Extended help

The output of bitcast_struct(T, value) is similar to the output of reinterpret(T, value), with both functions interpreting sequential bytes in little-endian order:

julia> reinterpret(NTuple{4, Int8}, Int32(1))
(1, 0, 0, 0)

julia> reinterpret(NTuple{6, Int32}, (2 * eps(0.0), eps(0.0), 0.0))
(2, 0, 1, 0, 0, 0)

julia> reinterpret(Tuple{Int32, Int32, Int128}, (2, 1, 0))
(2, 0, 1)

As the last example shows, bitcast_struct and reinterpret can behave differently when converting between data structures with nonuniform field sizes. Specifically, they differ for data structures that are stored with padding, which the C code underlying Julia uses to ensure that fields are efficiently aligned in stack memory.

Unlike reinterpret(T, value), which avoids mixing padding with non-padding (it recursively traverses fields of value and T, introducing offsets when their padding bytes are in different positions), bitcast_struct(T, value) makes no distinction between padding and non-padding. Although reinterpret is therefore less likely to produce unexpected outputs, it also performs runtime allocations in heap memory, making it unsuitable for GPU kernels that do not support such allocations. In contrast, bitcast_struct has a much simpler implementation, with all of its allocations confined to stack memory. Moreover, as long as bitcast_struct is only called within set_struct! and get_struct, potentially unexpected outputs will be hidden from users.

In addition to the low-level method of reinterpret for isbits inputs, there is another method for AbstractArray inputs that behaves exactly like bitcast_struct when it comes to padding:

julia> reinterpret(reshape, NTuple{4, Int8}, Int32[1])[1]
(1, 0, 0, 0)

julia> reinterpret(reshape, NTuple{6, Int32}, [2 * eps(0.0), eps(0.0), 0.0])[1]
(2, 0, 1, 0, 0, 0)

julia> reinterpret(reshape, Tuple{Int32, Int32, Int128}, [2, 0, 1, 0])[1]
(2, 0, 1)

This method of reinterpret reads bytes from heap memory without distinguishing padding and non-padding, in the same way as bitcast_struct reads bytes from stack memory. So, while the method of reinterpret for isbits inputs can construct the nonuniform type Tuple{Int32, Int32, Int128} from three Int64s, bitcast_struct and the method for arrays both require a fourth Int64, spanning the eight padding bytes inserted between the Int32s and the Int128.

For more information about reinterpret and padding, see the following:

• https://discourse.julialang.org/t/reinterpret-returns-wrong-values
• https://discourse.julialang.org/t/reinterpret-vector-into-single-struct
• https://discourse.julialang.org/t/reinterpret-vector-of-mixed-type-tuples

default_basetype(S)

Finds a type that set_struct! and get_struct can use to store either a value of type S, or any of the fields within such a value. If possible, this type is found by recursively searching the fieldtypes of S; otherwise, an unsigned integer type is selected based on the fieldtype sizes.

check_basetype(T, S)

Checks whether set_struct! and get_struct can use values of type T to store a value of type S. Throws an error if this is not the case, printing out an example of a specific field that cannot use T as a basetype.

checked_valid_basetype(T, S)

Returns either T or the default_basetype of S, depending on whether T satisfies check_basetype for S.

num_basetypes(T, S)

Determines how many values of type T are required by set_struct! and get_struct to store a single value of type S.

struct_field_view(array, S, Val(F), [Val(D)])

Creates a view of the data in array that corresponds to a particular field of S, assuming that array has been populated by set_struct!. The field is specified through a Val that contains its index F, and it can be loaded from the resulting view using get_struct.

For multidimensional arrays with values stored along a particular dimension, the resulting view contains the specified field from each value. By default, values are assumed to be stored along the last array dimension, but any other dimension can be specified through a Val that contains its index D.

set_struct!(array, value, [index], [Val(D)])

Populates array with data that represents any isbits value, using bitcast_struct to convert value into entries of the array.

For multidimensional arrays with values stored along a particular dimension, an index is used to identify the location of one value. By default, values will be stored along the last array dimension, but any other dimension can be specified as Val(D). The target location's index should be either an integer that corresponds to its start, or a CartesianIndex that contains its coordinate along every dimension except D.

Examples

julia> set_struct!(zeros(Int8, 4), Int32(1))
4-element Vector{Int8}:
 1
 0
 0
 0

julia> set_struct!(zeros(Int64, 4), (Int32(2), Int32(0), Int128(1)))
4-element Vector{Int64}:
 2
 0
 1
 0

julia> set_struct!(zeros(Int64, 2, 4), (Int32(2), Int32(0), Int128(1)), 2)
2×4 Matrix{Int64}:
 0  0  0  0
 2  0  1  0

julia> set_struct!(zeros(Int64, 4, 2), (Int32(2), Int32(0), Int128(1)), 5, Val(1))
4×2 Matrix{Int64}:
 0  2
 0  0
 0  1
 0  0

get_struct(array, S, [index], [Val(D)])

Loads a value of type S that set_struct! has stored in array, using bitcast_struct to convert entries of the array into this value.

For multidimensional arrays with values stored along a particular dimension, an index is used to identify the location of one value. By default, values are assumed to be stored along the last array dimension, but any other dimension can be specified as Val(D). The target location's index should be either an integer that corresponds to its start, or a CartesianIndex that contains its coordinate along every dimension except D.

Examples

julia> get_struct(Int8[1, 0, 0, 0], Int32)
1

julia> get_struct([2, 0, 1, 0], Tuple{Int32, Int32, Int128})
(2, 0, 1)

julia> get_struct([0 0 0 0; 2 0 1 0], Tuple{Int32, Int32, Int128}, 2)
(2, 0, 1)

julia> get_struct([0 2; 0 0; 0 1; 0 0], Tuple{Int32, Int32, Int128}, 5, Val(1))
(2, 0, 1)

parent_array_type(A, [T])

Determines the array type underlying the wrapper type A, dropping all parameters related to array dimensions. A new basetype T can be specified to replace the original eltype(A).

parent_array_type(data::AbstractData)

This is an internal function, please do not use outside of ClimaCore.

Returns the the backing array type.

This function is helpful for writing generic code, when reconstructing new datalayouts with new type parameters.

promote_parent_array_type(A1, A2)

Promotes two array types A1 and A2 generated by parent_array_type, which includes promoting their basetypes eltype(A1) and eltype(A2).