Add support for nonuniform data structures

Author: dennisYatunin

src/DataLayouts/DataLayouts.jl

@@ -66,12 +66,11 @@ module DataLayouts
 import Base: Base, @propagate_inbounds
 import StaticArrays: SOneTo, MArray, SArray
 import ClimaComms
-import UnrolledUtilities: unrolled_map, unrolled_all
 import MultiBroadcastFusion as MBF
 import Adapt
-import UnrolledUtilities: unrolled_foreach, unrolled_all, unrolled_findfirst
+using UnrolledUtilities
 
-import ..Utilities: PlusHalf, unionall_type
+import ..Utilities: PlusHalf, unionall_type, replace_type_parameter
 import ..DebugOnly: call_post_op_callback, post_op_callback
 import ..slab, ..slab_args, ..column, ..column_args, ..level, ..level_args
 export slab,
@@ -131,6 +130,7 @@ struct IJHMask{B, NT, V} <: AbstractMask
 end
 
 include("struct.jl")
+include("parent_arrays.jl")
 
 abstract type AbstractData{S} end
 
@@ -395,7 +395,7 @@ end
 
 function replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T}
     array = parent(data)
-    S′ = replace_basetype(eltype(array), T, S)
+    S′ = replace_type_parameter(S, eltype(array), T)
     return union_all(singleton(data)){S′, Base.tail(type_params(data))...}(
         similar(array, T),
     )
@@ -878,17 +878,9 @@ function DataF{S}(::Type{ArrayType}, fun = similar;) where {S, ArrayType}
 end
 
 function DataF(x::T) where {T}
-    if is_valid_basetype(Float64, T)
-        d = DataF{T}(Array{Float64})
-        d[] = x
-        return d
-    elseif is_valid_basetype(Float32, T)
-        d = DataF{T}(Array{Float32})
-        d[] = x
-        return d
-    else
-        check_basetype(Float64, T)
-    end
+    d = DataF{T}(Array{default_basetype(T)})
+    d[] = x
+    return d
 end
 
 Base.@propagate_inbounds function Base.getindex(data::DataF{S}) where {S}
@@ -2270,8 +2262,6 @@ include("fused_copyto.jl")
 include("fill.jl")
 include("mapreduce.jl")
 
-include("struct_linear_indexing.jl")
-
 
 """
     set_mask_maps!(mask)

src/DataLayouts/parent_arrays.jl

@@ -0,0 +1,49 @@
+"""
+    parent_array_type(::Type{<:AbstractArray})
+
+Returns the parent array type underlying any wrapper types, with all
+dimensionality information removed.
+"""
+parent_array_type(::Type{<:Array{T}}) where {T} = Array{T}
+parent_array_type(::Type{<:MArray{S, T, N, L}}) where {S, T, N, L} =
+    MArray{S, T}
+parent_array_type(::Type{<:SubArray{T, N, A}}) where {T, N, A} =
+    parent_array_type(A)
+
+# ReshapedArray is needed for converting between arrays and fields for RRTMGP:
+parent_array_type(::Type{<:Base.ReshapedArray{T, N, P}}) where {T, N, P} =
+    parent_array_type(P)
+
+"""
+    promote_parent_array_type(::Type{<:AbstractArray}, ::Type{<:AbstractArray})
+
+Given two parent array types (without any dimensionality information), promote
+both the element types and the array types themselves.
+"""
+promote_parent_array_type(::Type{Array{T1}}, ::Type{Array{T2}}) where {T1, T2} =
+    Array{promote_type(T1, T2)}
+promote_parent_array_type(
+    ::Type{MArray{S, T1}},
+    ::Type{MArray{S, T2}},
+) where {S, T1, T2} = MArray{S, promote_type(T1, T2)}
+promote_parent_array_type(
+    ::Type{MArray{S, T1}},
+    ::Type{Array{T2}},
+) where {S, T1, T2} = MArray{S, promote_type(T1, T2)}
+promote_parent_array_type(
+    ::Type{Array{T1}},
+    ::Type{MArray{S, T2}},
+) where {S, T1, T2} = MArray{S, promote_type(T1, T2)}
+# Ditch sizes (they're never actually used!)
+promote_parent_array_type(
+    ::Type{MArray{S1, T1}},
+    ::Type{MArray{S2, T2}},
+) where {S1, T1, S2, T2} = MArray{S, promote_type(T1, T2)} where {S}
+promote_parent_array_type(
+    ::Type{MArray{S1, T1} where {S1}},
+    ::Type{MArray{S2, T2}},
+) where {T1, S2, T2} = MArray{S, promote_type(T1, T2)} where {S}
+promote_parent_array_type(
+    ::Type{MArray{S1, T1}},
+    ::Type{MArray{S2, T2} where {S2}},
+) where {S1, T1, T2} = MArray{S, promote_type(T1, T2)} where {S}