GitHub

DataLayouts

Data layouts

ClimaCore.DataLayouts.DataLayout — Type
DataLayout{T, N, F, S, A}

An N-dimensional AbstractArray containing values of type T, stored in a parent array of type A whose memory layout is determined by the layout's type. Every value can be identified by four indices: a vertical level v, horizontal quadrature points i and j, and a horizontal element h. The components of each value are optionally stored along a hidden field axis F of the parent array, leading to a hybrid of the traditional "array-of-structs" (F = 1) and "struct-of-arrays" (F = ndims(A)) approaches to storing non-scalar data. The DataScope S determines how loops and reductions over the values are parallelized on CPUs and GPUs, and it dictates which array types are allocated.

Several layouts are available, named after the order of their parent axes:

  • DataF is a 0-dimensional array that stores a single value, with an Nf-element parent array (used in place of a Ref)
  • VIJFH is an Nv × Ni × Nj × Nh array that stores spatially varying data, with each value spread along the fourth parent axis
  • VIJHF is like VIJFH with the F and H axes swapped, which permits linear indexing and improves performance for operators that only access one field at a time
  • VIJHWithF generalizes VIJFH and VIJHF to any F axis position, with F = nothing removing the axis altogether
  • VIH1 and IH1JH2 store vertical and horizontal planes of interpolated data for plotting, whose ih1 and jh2 indices combine i and j with h1 and h2 (orthogonal components of h in rectangular domains)
julia> data = VIJFH{Tuple{Int64, Float64, Int128}, 10, 5, 5, nothing}(Array{Int64}, 20);

julia> size(data), size(parent(data)) # Nh = 20 elements, Nf = 4 Int64 storage values
((10, 5, 5, 20), (10, 5, 5, 4, 20))

julia> data[1, 2, 3, 4] = (0, 1.0, 2); data.:1[1, 2, 3, 4], data[1, 2, 3, 4].:2
(0, 1.0)

Extended Help

DataLayouts also provide the following functionality for ClimaCore:

  • Assigning a DataScope to every batch of data, and automatically partitioning data across nestable multithreaded operations
  • Storing specific array dimensions as type parameters, and allocating static arrays in place of regular arrays when every dimension can be inferred
  • Using linear indices in place of Cartesian indices where doing so may improve performance, including in getindex and view operations
  • Automatic nested broadcasting over Tuple and NamedTuple values (or other supported iterator types), along with broadcasting over array indices
  • Checking for type stability before evaluating operations like broadcasts and reductions, avoiding inefficient CPU behavior and GPU compilation errors
  • Falling back to built-in AbstractArray methods when specialized ClimaCore code is not available (this may be highly inefficient or fail to compile on GPUs, but it should generally work on CPUs)
source
ClimaCore.DataLayouts.DataF — Type
DataF{T, [S]}(A)
DataF{T, [S]}(array)

DataLayout representing a single value of type T, which can be stored across multiple array indices. This is used in place of a Ref to wrap data that is stored in any array. May be constructed either from the parent array type or the parent array itself.

source
ClimaCore.DataLayouts.VIJHWithF — Type
VIJHWithF{T, Nv, Ni, Nj, Nh, F, [S]}(A, [Nh_dynamic])

Generalization of a VIJFH and a VIJHF, which supports any value of the parameter F between 1 and 5, representing FVIJH, VFIJH, and so on. The parameter can also be nothing, which drops the F axis altogether.

source
ClimaCore.DataLayouts.VIJFH — Type
VIJFH{T, Nv, Ni, Nj, Nh, [S]}(A, [Nh_dynamic])
VIJFH{T, Nv, Ni, Nj, Nh, [S]}(array)

DataLayout representing values of type T stored across Nv vertical levels, Nh horizontal elements, and Ni × Nj quadrature points per element. The parameters Nv, Ni, and Nj must be integers, but Nh may be set to nothing and obtained at runtime from the array size. Each value of type T can be stored across multiple indices along the fourth array axis. May be constructed either from the parent array type or the parent array itself, though using a type requires passing an additional integer if Nh is set to nothing.

source
ClimaCore.DataLayouts.VIJHF — Type
VIJHF{T, Nv, Ni, Nj, Nh, [S]}(A, [Nh_dynamic])
VIJHF{T, Nv, Ni, Nj, Nh, [S]}(array)

DataLayout similar to VIJFH, but with the last two axes of the parent array swapped. Offers better performance than VIJFH for operations that only access one field from each value of type T.

source
ClimaCore.DataLayouts.VIH1 — Type
VIH1{T, Nv, Ni, Nh, [S]}(A, [Nh_dynamic])
VIH1{T, Nv, Ni, Nh, [S]}(array)

DataLayout representing values of type T stored across Nv vertical levels and Ni × Nh1 horizontal quadrature points. This ignores the second horizontal direction, which spans Nj × Nh2 quadrature points (Nh is given by Nh1 × Nh2). The parameters Nv and Ni must be integers, but Nh may be set to nothing and obtained at runtime from the array size; when it is not nothing, Nh can only be set to 1. May be constructed either from the parent array type or the parent array itself, though using a type requires passing an additional integer if Nh is set to nothing.

source
ClimaCore.DataLayouts.IH1JH2 — Type
IH1JH2{T, Ni, Nj, Nh, [S]}(A, [Nh_dynamic])
IH1JH2{T, Ni, Nj, Nh, [S]}(array)

DataLayout representing values of type T stored across Ni × Nh1 quadrature points along one horizontal direction and Nj × Nh2 quadrature points along the other horizontal direction (Nh is given by Nh1 × Nh2). This ignores the vertical direction, which spans Nv levels. The parameters Ni and Nj must be integers, but Nh may be set to nothing and obtained at runtime from the array size; when it is not nothing, Nh can only be set to 1. May be constructed either from the parent array type or the parent array itself, though using a type requires passing an additional integer if Nh is set to nothing.

source

Layout properties

ClimaCore.DataLayouts.f_dim — Function
f_dim(D)
f_dim(data)

Index of the F axis in a parent array for a DataLayout, or nothing if there is no separate F axis. The value of nothing is chosen instead of missing because GPUCompiler.jl compares type parameters with ==, which returns the non-boolean missing whenever one of its arguments is missing.

source
ClimaCore.DataLayouts.layout_constructor — Function
layout_constructor(D, [T]; [params...])
layout_constructor(data, [T]; [params...])

Constructor for a similar DataLayout that can be applied as constructor(array), with the element type optionally replaced with T, and with any subset of the shape_params optionally replaced with params.

source
ClimaCore.DataLayouts.rebuild — Function
rebuild(data, A, [T]; [params...])
rebuild(data, array, [T]; [params...])

Reconstruct a DataLayout with a modified parent array, either converting its parent array to some type A, or replacing it with another array. As in layout_constructor, a new element type and new shape_params may also be specified.

The new array can be stored on a different device (e.g., Array vs CuArray), so the DataScope is modified if it is inconsistent with the new array.

source

Data scopes

ClimaCore.DataLayouts.DataScope — Type
DataScope(A)
DataScope(args...)

Singleton type that represents a computational unit responsible for updating the values in an AbstractArray of type A. May also be constructed using an instance of an array or any similarly indexable argument, or by combining the DataScopes from multiple arguments (always selecting the smallest scope).

Extended Help

DataScopes can be compared using is_subscope, and they define methods for the following functions:

Every DataLayout is assigned a specific DataScope, but the scope of a generic AbstractArray must be inferred from its type. While some types of arrays can only be assigned to one scope (e.g., a StaticArray can only be accessed from ThisThread), this is not always the case. For example, a SubArray view of a CUDA.CuDeviceArray that serves as the parent of a DataLayout inside a GPU kernel can be assigned to many different scopes:

In general, the DataScope of a non-DataLayout array will include all threads that are able to access it, so the SubArray in the example above would be assigned to ThisKernel. If a smaller scope is required, the array needs to be wrapped in a DataLayout.

source
ClimaCore.DataLayouts.ThisThreadPool — Type
ThisThreadPool()

DataScope that represents threads from the default thread pool on a CPU.

Loops that run at the same time divide the pool between them, with each loop's share of the pool determined by resolve_pool_threads when it starts. A loop that cannot claim more than one thread — because the pool is busy, because Julia was started with a single thread, or because the loop is nested in a multithreaded loop located outside of ClimaCore — runs on the thread that launched it.

source
ClimaCore.DataLayouts.partition — Function
partition(scope)

DataScope whose threads are a subset of the specified scope. Acts as a statically inferrable, device-agnostic generalization of the tiled_partition function from CUDA's cooperative_groups extension. By default, the entire scope is placed in a single partition.

source
ClimaCore.DataLayouts.resolve_pool_threads — Function
resolve_pool_threads()

Number of threads from the default thread pool that the current loop may use, which is 1 when the loop has to run on the thread that started it. A loop that is given more than one thread records the count, so that num_threads and thread_rank report the same division of the loop's indices for as long as it runs.

Every loop that calls this must give its threads back with release_pool_threads once it has finished.

The count is deliberately not returned as a DataScope, and jl_in_threaded_region is only read here, once per loop. A scope whose type is only known at run time becomes a union in every call below it, and reading the process-global threaded-region flag more than once per loop lets an unrelated task flip it in between, either of which would keep pointwise loops from staying allocation free.

source
ClimaCore.DataLayouts.synchronize — Function
synchronize(scope)

Synchronizes all threads in a DataScope, so that no thread can begin executing code that comes after this instruction until all other threads have finished executing the code that came before this instruction.

source
ClimaCore.DataLayouts.scoped_array — Function
scoped_array(scope, T, dims; [buffer])

Array with the specified element type and size, whose values can be modified by every thread in a DataScope. When buffer = true, a task-local buffer is reused instead of allocating a new array.

source
ClimaCore.DataLayouts.subscope_indices — Function
subscope_indices(subscope, scope, indices)

Divides a collection of indices (either linear or Cartesian) among subsets of a DataScope. The result is a strided range if strided_access is true for scope, or a contiguous range if it is false.

Contiguous ranges are generated by partitioning the indices into chunks whose lengths differ from each other by at most 1, which guarantees that every subset of the scope gets a nonempty chunk whenever there are at least as many indices as subsets. In contrast, always assigning cld(length(indices), n_subsets) indices to each subset can lead to one or more empty subsets.

source

Loops and reductions

ClimaCore.DataLayouts.each_slice_index — Function
each_slice_index(op, args...)

Generalization of eachindex for the slice operators level, slab, column, and view (for creating single-point slices). The result is always an iterator of Cartesian indices, whose scalar offsets are simple enough for SIMD optimization (a view at a linear index wraps its parent in a 1-dimensional ReshapedArray, which blocks SIMD in pointwise loops).

The arguments' axes are combined with stable_combine_axes, which expands singleton and 0-dimensional axes like broadcasting does and never throws, since this function is called from GPU kernels, where an error path either fails to compile or traps with an unrelated CUDA error. Arguments whose axes are genuinely incompatible are rejected on the host by foreach_slice before any kernel is launched.

source
ClimaCore.DataLayouts.slice_subscope — Function
slice_subscope(scope, op, args...)

DataScope that foreach_slice assigns to slices of the given arguments when parallelizing over scope. By default, this is the smallest subset of scope that does not require any thread to process more than one point from the largest slice returned by op. When no such subset is available, the largest subset is used in order to minimize the number of points per thread.

source
ClimaCore.DataLayouts.foreach_slice — Function
foreach_slice(op, f, args...; [mask])
foreach_slice(scope, op, f, args...; mask)

Generalization of eachslice/mapslices that applies f to slices of every DataLayout or similarly indexable argument, where the slice operator op can be any of the following:

Each slice is assigned to a slice_subscope of scope, which by default is the largest available DataScope that can access every argument. A DataMask may also be used to skip over a particular subset of slices.

The mask is only given a default of NoMask when no scope is specified, since that is the only method a loop starts from. Every method that takes a scope requires a mask, so that a loop which passes its keyword arguments on cannot quietly drop a mask and compute over the points it excludes.

source
ClimaCore.DataLayouts.reduce_points — Function
reduce_points(op, arg; [mask], [init])
reduce_points(scope, op, arg; mask, [init])

Generalization of reduce that uses op to combine values stored in a DataLayout or similarly indexable argument.

This combines all values in the given argument that are assigned to scope, which by default is the largest available DataScope that can access the argument. A DataMask may also be used to skip over a particular subset of points. If the mask disables every point, or if there are no points in arg to begin with, the init value must be specified.

As in foreach_slice, the mask is only given a default of NoMask when no scope is specified, so that a reduction which passes its keyword arguments on cannot quietly drop a mask.

source
ClimaCore.DataLayouts.column_reduce! — Function
column_reduce!(op, dest, arg; [mask], [flip], [init])

Use foreach_column to combine the levels of each column of arg with op, storing the results in corresponding columns of dest. Setting flip to true changes the order of reduction from left-associative (default) to right-associative, and init seeds the fold when it is given.

source

Masks

ClimaCore.DataLayouts.IJHMask — Type
IJHMask(data)

A DataMask that marks the columns of a VIJFH or VIJHF layout as active or inactive, using the following cached values:

  • is_active, a layout similar to level(data, 1) representing a boolean mask
  • N, an array that contains the total number of active columns
  • i_map, an array that contains the i-index of each active column
  • j_map, an array that contains the j-index of each active column
  • h_map, an array that contains the h-index of each active column
source

Struct storage

ClimaCore.DataLayouts.bitcast_struct — Function
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
source
ClimaCore.DataLayouts.default_basetype — Function
default_basetype(T)

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

source
ClimaCore.DataLayouts.struct_field_view — Function
struct_field_view(array, T, Val(i), [Val(F)])

Creates a view of the data in array that corresponds to a particular field of T, assuming that array has been populated by set_struct!. The field is specified through a Val that contains its index i, 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, with the dimension identified by a Val that contains its index F. When there is no such dimension, F may be replaced with nothing.

source
ClimaCore.DataLayouts.set_struct! — Function
set_struct!(array, value, [index, Val(F)])
set_struct!(array, value, [index, stride])

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, with the dimension specified as Val(F). The target values's index should be a CartesianIndex that contains its coordinate along every dimension except F. When there is no such dimension, F may be replaced with nothing.

Arrays that support linear indexing can also be accessed using two integers, where one corresponds to the start of a value, and another corresponds to the stride along the F axis between consecutive components of the value.

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, 4, 2), (Int32(2), Int32(0), Int128(1)), 5, Val(1))
4×2 Matrix{Int64}:
 0  2
 0  0
 0  1
 0  0

julia> set_struct!(zeros(Int64, 3, 4), (Int32(2), Int32(0), Int128(1)), 2, 3)
3×4 Matrix{Int64}:
 0  0  0  0
 2  0  1  0
 0  0  0  0
source
ClimaCore.DataLayouts.get_struct — Function
get_struct(array, T, [index, Val(F)])
get_struct(array, T, [index, stride])

Loads a value of type T 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, with the dimension specified as Val(F). The target values's index should be a CartesianIndex that contains its coordinate along every dimension except F. When there is no such dimension, F may be replaced with nothing.

Arrays that support linear indexing can also be accessed using two integers, where one corresponds to the start of a value, and another corresponds to the stride along the F axis between consecutive components of the value.

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 2; 0 0; 0 1; 0 0], Tuple{Int32, Int32, Int128}, 5, Val(1))
(2, 0, 1)

julia> get_struct([0 0 0 0; 2 0 1 0; 0 0 0 0], Tuple{Int32, Int32, Int128}, 2, 3)
(2, 0, 1)
source
ClimaCore.DataLayouts.view_struct — Function
view_struct(array, T, [index, Val(F)])

Analogous to get_struct, but for a view of the struct data instead of the value itself. The value may be accessed with get_struct(struct_view, T), and it can be updated with set_struct!(struct_view, new_value).

source

Broadcasting