DataLayouts

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 inefficient or fail to compile on GPUs, but it works on CPUs).

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 one-dimensional array. It can be constructed either from the parent array type or from the parent array itself.

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.

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. It can be constructed either from the parent array type or from the parent array itself, though using a type requires passing Nh_dynamic if Nh is set to nothing.

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.

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. It can be constructed either from the parent array type or from the parent array itself, though using a type requires passing Nh_dynamic if Nh is set to nothing.

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. It can be constructed either from the parent array type or from the parent array itself, though using a type requires passing Nh_dynamic if Nh is set to nothing.

layout_type(D)
layout_type(data)

Return the type of a DataLayout, stripped of all its type parameters.

parent_type(D)
parent_type(data)

Return the type of the parent array used by a DataLayout, or a similar abstract type if the concrete type is unavailable.

f_dim(D)
f_dim(data)

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

shape_params(D)
shape_params(data)

Return a NamedTuple with all shape-related parameters of a DataLayout. This excludes its element type, its parent array type, and its DataScope.

inferred_size(D)
inferred_size(data)

Return the size of a DataLayout, with dimensions that cannot be inferred from its type set to nothing.

has_inferred_size(D)
has_inferred_size(data)

Return whether every dimension of a DataLayout can be inferred from its type.

vijh_params(D)
vijh_params(data)

Return a NamedTuple with Nv, Ni, Nj, and Nh, the lengths of the V, I, J, and H axes in a DataLayout. Axes that the layout does not have are reported with length 1. Like inferred_size, this returns nothing for dimensions that cannot be inferred from the type.

nlevels(D)
nlevels(data)

Return the length of the V axis in a DataLayout.

nquadpoints(D)
nquadpoints(data)

Return the product of the lengths of the I and J axes in a DataLayout.

nelems(D)
nelems(data)

Return the length of the H axis in a DataLayout. When the length cannot be inferred from the type, a concrete instance must be provided instead; passing a type throws an ArgumentError.

ncomponents(D)
ncomponents(data)

Return the length of the hidden F axis in a DataLayout, or 1 if there is no separate F axis.

layout_constructor(D, [T]; [params...])
layout_constructor(data, [T]; [params...])

Return a 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.

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.

reassign(D, scope)
reassign(data, scope)

Assign a new DataScope to a DataLayout, or determine the result type of performing such an assignment for a layout of type D.

replace_basetype(data::DataLayout, ::Type{B})

Return an uninitialized DataLayout like data whose parent array has element type B in place of the current base type, with the element type of data rewritten accordingly (e.g. a layout of Tuple{Float64, Float64} backed by a Float64 array becomes a layout of Tuple{Float32, Float32} backed by a Float32 array when B = Float32).

DataScope(A)
DataScope(args...)

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

Extended help

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

• partition.
• num_threads or num_partitions (only one needs a method).
• thread_rank or partition_rank (only one needs a method).
• parallelize_over and synchronize.
• scoped_array and scoped_static_array.
• strided_access.

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:

• ThisKernel, if the DataLayout was generated by level, slab, column, or a single-point view.
• Either ThisBlock, ThisSubBlock, or ThisThread, if it was generated by foreach_level, foreach_slab, or foreach_column.
• ThisThread, if it was generated by foreach_point.

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

ThisThread()

DataScope that represents the currently running thread.

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.

partition(scope)

Return a DataScope whose threads are a subset of scope. This is 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.

is_subscope(subscope, scope)

Return whether subscope is equal to the DataScope scope, or is a partition of it, or a partition of a partition, and so on.

num_threads(scope)

Return the number of threads that are part of a DataScope.

num_partitions(scope)

Return the number of partitions (results of partition) in a DataScope.

thread_rank(scope)

Return an integer between 1 and num_threads that identifies the calling thread within a DataScope.

partition_rank(scope)

Return an integer between 1 and num_partitions that identifies the calling thread's partition within a DataScope.

parallelize_over(f, scope)

Call f() from each thread in a DataScope. Code that appears outside of this instruction is not necessarily parallelized over all available threads.

resolve_pool_threads()

Return the 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 in task-local storage, 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 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 would keep pointwise loops from staying allocation free.

release_pool_threads()

Give back the threads that resolve_pool_threads claimed for the current loop, and return nothing. Every loop that resolves a ThisThreadPool must call this once it has finished, whether or not it was given any threads.

synchronize(scope)

Synchronize 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 it.

scoped_array(scope, T, dims; [buffer])

Return an array with element type T and size dims, 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; the default is buffer = false.

scoped_static_array(scope, T, dims)

Return a statically sized array with element type T and size dims, whose values can be modified by every thread in a DataScope. If a DataScope does not provide its own method, this falls back to calling scoped_array.

Two allocations with equal element types and sizes may be assigned to the same memory (this is what happens for block-level allocations on GPUs, which are stored in shared memory), so their lifetimes must not overlap unless they are guaranteed to hold identical values.

strided_access(scope)

Return whether subscope_indices should return a strided range of indices for the given DataScope, rather than a contiguous range. The default is true.

subscope_indices(subscope, scope, indices)

Return the part of indices (either linear or Cartesian) assigned to the calling subscope when the indices are divided among the subscopes 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.

each_slice_index(op, arg)

Return an iterator over the indices of the slices that the slice operator op produces from arg, generalizing eachindex to level, slab, column, and view (for single-point slices). The result is always an iterator of Cartesian indices, whose scalar offsets permit SIMD optimization (a view at a linear index wraps its parent in a 1-dimensional ReshapedArray, which blocks SIMD in pointwise loops).

slice_subscope(scope, op, args...)

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

scope itself is only used when its thread count is a compile-time constant (see static_num_threads), since scoped_slice_loop gives every slice a scope with a statically known number of threads.

foreach_slice(op, f, args...; mask = NoMask(), enumerate = Val(false))
foreach_slice(scope, op, f, args...; mask = NoMask(), enumerate = Val(false))

Apply f to slices of every DataLayout or similarly indexable argument, as a generalization of eachslice/mapslices. The slice operator op can be any of the following:

• level, but only when nelems is statically inferrable.
• slab or column.
• view (for single-point slices).

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

Statements in f execute in order at the slice's scope: same-shaped pointwise broadcasts assign every point to the same thread in each statement, and spectral operators publish values that cross threads through synchronized buffers, so each statement may read the results of the statements before it.

By default, f is called as f(slices...). Setting enumerate to Val(true) makes this f(index, slices...), like in a loop over Base.enumerate(arg).

foreach_point(f, args...; mask = NoMask(), enumerate = Val(false))

Run foreach_slice with view as the slice operator.

foreach_level(f, args...; mask = NoMask(), enumerate = Val(false))

Run foreach_slice with level as the slice operator.

foreach_slab(f, args...; mask = NoMask(), enumerate = Val(false))

Run foreach_slice with slab as the slice operator.

foreach_column(f, args...; mask = NoMask(), enumerate = Val(false))

Run foreach_slice with column as the slice operator.

reduce_points(op, arg; mask = NoMask(), [init])
reduce_points(scope, op, arg; mask = NoMask(), [init])

Combine all values of a DataLayout or similarly indexable argument with op, as a generalization of reduce, using the threads of scope, by default the largest available DataScope that can access the argument. A DataMask may be used to skip a particular subset of points. The init value must be specified if the mask disables every point, if there are no points in arg to begin with, or when reducing with any mask on a GPU, where a mask can leave whole blocks without active points.

column_reduce!(op, dest, arg; mask = NoMask(), flip = Val(false), [init])

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

DataMask

Abstract type for masks that mark points in a discretized domain as active or inactive. Subtypes: NoMask and IJHMask.

NoMask()

A DataMask that marks every point in a discretized domain as active.

IJHMask(data)

A DataMask that marks the columns of a VIJFH or VIJHF layout as active or inactive. The constructor marks every column as active; modify is_active and call set_mask_maps! to change the mask.

Fields

• is_active: A layout similar to level(data, 1) that holds the boolean mask.
• N: A one-element array that holds the total number of active columns.
• i_map: An array that holds the i-index of each active column.
• j_map: An array that holds the j-index of each active column.
• h_map: An array that holds the h-index of each active column.

is_active(mask::IJHMask)

The boolean layout that marks each column of the masked data as active (true) or inactive (false); it has the shape of level(data, 1). Modify it in place and call set_mask_maps! to apply the change.

set_mask_maps!(mask)

Update mask.N, mask.i_map, mask.j_map, and mask.h_map in an IJHMask based on the values in mask.is_active, and return mask. This allocates memory when using GPUs, so it should only be called infrequently.

should_compute(mask, index)

Return whether a DataMask marks the point at index as active.

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

Convert 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). This is 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 the num_indices entries of an array that get_struct reads at index. 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 (an optimization that matters in GPU kernels, where reads from global memory are 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 shorter 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 are 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(T)

Return 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.

check_basetype(B, T)

Check whether set_struct! and get_struct can use values of type B to store a value of type T. Return nothing if so; otherwise throw an ArgumentError that names a specific field which cannot use B as a basetype.

checked_valid_basetype(B, T)

Return either B or the default_basetype of T, depending on whether B satisfies check_basetype for T.

num_basetypes(B, T)

Return the number of values of type B that set_struct! and get_struct use to store a single value of type T.

struct_field_view(array, T, Val(i), [Val(F)])

Return 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.

set_struct!(array, value, [index, Val(F)])
set_struct!(array, value, [index, stride])

Populate array with data that represents any isbits value, using bitcast_struct to convert value into entries of the array, and return 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 value'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

get_struct(array, T, [index, Val(F)])
get_struct(array, T, [index, stride])

Load 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 value'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)

view_struct(array, T, [index, Val(F)])

Return a view of the struct data that get_struct would load, instead of the value itself. The value can be accessed with get_struct(struct_view, T), and it can be updated with set_struct!(struct_view, new_value).

DataStyle(D)

BroadcastStyle for a DataLayout of type D, which stores the layout_type and its corresponding value of ndims as type parameters.

LazyDataLayout{D}

A DataStyle broadcast expression whose layout_type is D.

layout_args(bc)

Return a tuple of every DataLayout and LazyDataLayout among the arguments of a broadcast expression.

modify_args(f, bc, f_args...)

Return a copy of a broadcast expression in which each of the layout_args is replaced with f(layout_arg, f_args...).