wip

Author: charleskawczynski

ext/cuda/data_layouts.jl

@@ -13,10 +13,16 @@ import CUDA
 parent_array_type(::Type{<:CUDA.CuArray{T, N, B} where {N}}) where {T, B} =
     CUDA.CuArray{T, N, B} where {N}
 
+# Can we remove this?
+# parent_array_type(
+#     ::Type{<:CUDA.CuArray{T, N, B} where {N}},
+#     ::Val{ND},
+# ) where {T, B, ND} = CUDA.CuArray{T, ND, B}
+
 parent_array_type(
     ::Type{<:CUDA.CuArray{T, N, B} where {N}},
-    ::Val{ND},
-) where {T, B, ND} = CUDA.CuArray{T, ND, B}
+    as::ArraySize,
+) where {T, B} = CUDA.CuArray{T, ndims(as), B}
 
 # Ensure that both parent array types have the same memory buffer type.
 promote_parent_array_type(

src/DataLayouts/DataLayouts.jl

@@ -51,6 +51,12 @@ abstract type AbstractData{S} end
 @inline Base.size(data::AbstractData, i::Integer) = size(data)[i]
 @inline Base.size(data::AbstractData) = universal_size(data)
 
+struct ArraySize{FD, Nf, S} end
+@inline ArraySize(data::AbstractData, i::Integer) = ArraySize(data)[i]
+@inline ArraySize(data::AbstractData) = ArraySize{field_dim(data), ncomponents(data), farray_size(data)}()
+@inline Base.ndims(::ArraySize{FD, Nf, S}) where {FD, Nf, S} = length(S)
+@inline Base.ndims(::Type{ArraySize{FD, Nf, S}}) where {FD, Nf, S} = length(S)
+
 include("field_array.jl")
 include("struct.jl")
 
@@ -296,39 +302,6 @@ end
     ))
 end
 
-for DL in (
-    :IJKFVH,
-    :IJFH,
-    :IFH,
-    # :DataF, # we want to MethodError for DataF
-    :IJF,
-    :IF,
-    :VF,
-    :VIJFH,
-    :VIFH,
-    :IH1JH2,
-    :IV1JH2,
-)
-    @eval @propagate_inbounds function Base.getindex(
-        data::$DL{S},
-        I::Integer,
-    ) where {S}
-        @inbounds get_struct(parent(data), S, Val(field_dim(data)), I)
-    end
-    @eval @propagate_inbounds function Base.setindex!(
-        data::$DL{S},
-        val,
-        I::Integer,
-    ) where {S}
-        @inbounds set_struct!(
-            parent(data),
-            convert(S, val),
-            Val(field_dim(data)),
-            I,
-        )
-    end
-end
-
 function replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T}
     array = parent(data)
     S′ = replace_basetype(eltype(array), T, S)
@@ -540,6 +513,7 @@ Base.@propagate_inbounds slab(data::IFH, v::Integer, h::Integer) = slab(data, h)
 @inline function column(data::IFH{S, Ni}, i, h) where {S, Ni}
     @boundscheck (1 <= h <= get_Nh(data) && 1 <= i <= Ni) ||
                  throw(BoundsError(data, (i, h)))
+    fa = parent(data)
     dataview = @inbounds FieldArray{field_dim(DataF)}(
         ntuple(jf -> view(parent(fa.arrays[jf]), i, h), ncomponents(fa)),
     )
@@ -731,6 +705,7 @@ end
 @inline function column(data::IJF{S, Nij}, i, j) where {S, Nij}
     @boundscheck (1 <= j <= Nij && 1 <= i <= Nij) ||
                  throw(BoundsError(data, (i, j)))
+    fa = parent(data)
     dataview = @inbounds FieldArray{field_dim(DataF)}(
         ntuple(jf -> view(parent(fa.arrays[jf]), i, j), ncomponents(fa)),
     )
@@ -905,8 +880,13 @@ end
 
 @inline function level(data::VF{S}, v) where {S}
     @boundscheck (1 <= v <= nlevels(data)) || throw(BoundsError(data, (v)))
-    array = parent(data)
-    dataview = @inbounds view(array, v, :)
+    fa = parent(data)
+    dataview = @inbounds FieldArray{field_dim(DataF)}(
+        ntuple(ncomponents(fa)) do jf
+            view(parent(fa.arrays[jf]), v, :)
+        end,
+    )
+
     DataF{S}(dataview)
 end
 
@@ -1247,6 +1227,9 @@ Adapt.adapt_structure(to, data::AbstractData{S}) where {S} =
 rebuild(data::AbstractData, array::AbstractArray) =
     union_all(data){type_params(data)...}(array)
 
+rebuild(data::AbstractData, fa::FieldArray) =
+    union_all(data){type_params(data)...}(fa)
+
 empty_kernel_stats(::ClimaComms.AbstractDevice) = nothing
 empty_kernel_stats() = empty_kernel_stats(ClimaComms.device())
 
@@ -1459,6 +1442,40 @@ type parameters.
 
 #! format: on
 
+# Skip DataF here, since we want that to MethodError.
+for DL in (:IJKFVH, :IJFH, :IFH, :IJF, :IF, :VF, :VIJFH, :VIFH)
+    @eval @propagate_inbounds Base.getindex(data::$(DL), I::Integer) =
+        linear_getindex(data, I)
+    @eval @propagate_inbounds Base.setindex!(data::$(DL), val, I::Integer) =
+        linear_setindex!(data, val, I)
+end
+
+# Datalayouts
+@propagate_inbounds function linear_getindex(
+    data::AbstractData{S},
+    I::Integer,
+) where {S}
+    s_array = farray_size(data)
+    ss = StaticSize(s_array, field_dim(data))
+    @inbounds get_struct_linear(parent(data), S, Val(field_dim(data)), ss, I)
+end
+@propagate_inbounds function linear_setindex!(
+    data::AbstractData{S},
+    val,
+    I::Integer,
+) where {S}
+    s_array = farray_size(data)
+    ss = StaticSize(s_array, field_dim(data))
+    @inbounds set_struct_linear!(
+        parent(data),
+        convert(S, val),
+        Val(field_dim(data)),
+        ss,
+        I,
+    )
+end
+
+
 Base.ndims(data::AbstractData) = Base.ndims(typeof(data))
 Base.ndims(::Type{T}) where {T <: AbstractData} =
     Base.ndims(field_array_type(T))

src/DataLayouts/broadcast.jl

@@ -19,7 +19,7 @@ Data0DStyle(::Type{DataFStyle{A}}) where {A} = DataFStyle{A}
 abstract type DataColumnStyle <: DataStyle end
 struct VFStyle{Nv, A} <: DataColumnStyle end
 DataStyle(::Type{VF{S, Nv, A}}) where {S, Nv, A} =
-    VFStyle{Nv, to_parent_array_type(A)}()
+    VFStyle{Nv, parent_array_type(A)}()
 DataColumnStyle(::Type{VFStyle{Nv, A}}) where {Nv, A} = VFStyle{Nv, A}
 
 abstract type Data1DStyle{Ni} <: DataStyle end
@@ -42,7 +42,7 @@ DataStyle(::Type{IJF{S, Nij, A}}) where {S, Nij, A} =
 abstract type Data2DStyle{Nij} <: DataStyle end
 struct IJFHStyle{Nij, Nh, A} <: Data2DStyle{Nij} end
 DataStyle(::Type{IJFH{S, Nij, Nh, A}}) where {S, Nij, Nh, A} =
-    IJFHStyle{Nij, Nh, to_parent_array_type(A)}()
+    IJFHStyle{Nij, Nh, parent_array_type(A)}()
 DataSlab2DStyle(::Type{IJFHStyle{Nij, Nh, A}}) where {Nij, Nh, A} =
     IJFStyle{Nij, A}
 
@@ -378,8 +378,14 @@ function Base.similar(
 ) where {Nv, Nij, Nh, A, Eltype, newNv}
     T = eltype(A)
     Nf = typesize(eltype(A), Eltype)
-    fat = field_array_type(A, Val(field_dim(VIJFH)), Val(Nf), Val(4))
-    array = similar(fat, Base.Dims((newNv, Nij, Nij, Nh)))
+    # fat = rebuild_type(A, Val(field_dim(VIJFH)), Val(Nf), Val(4))
+    _size = (newNv, Nij, Nij, Nh)
+    as = ArraySize{field_dim(VIJFH), Nf, _size}()
+    # fat = if A isa AbstractArray
+    #     field_array_type(A, as)
+    # else
+    # end
+    array = similar(rebuild_field_array_type(A, as), _size)
     vd = VIJFH{Eltype, newNv, Nij, Nh}(array)
     return vd
 end

src/DataLayouts/field_array.jl

@@ -46,19 +46,34 @@ float_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = eltype(T)
 parent_array_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} =
     parent_array_type(T)
 # field_array_type(::Type{FieldArray{N,T}}, ::Val{Nf}) where {N,T, Nf} = FieldArray{Nf, parent_array_type(T, Val(ndims(T)))}
-field_array_type(
+# field_array_type(
+#     ::Type{FieldArray{FD, N, T}},
+#     ::Val{Nf},
+#     ::Val{ND},
+# ) where {FD, N, T, Nf, ND} = FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
+
+rebuild_type(
+    ::Type{FieldArray{FD, N, T}},
+    as::ArraySize{FD,Nf},
+) where {FD, N, T, Nf} = FieldArray{FD, Nf, parent_array_type(T, Val(ndims(as)))}
+
+rebuild_field_array_type(
     ::Type{FieldArray{FD, N, T}},
-    ::Val{Nf},
-    ::Val{ND},
-) where {FD, N, T, Nf, ND} = FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
-# field_array_type(::Type{T}, ::Val{Nf}) where {T <: AbstractArray, Nf} = FieldArray{Nf, parent_array_type(T, Val(ndims(T)))}
-field_array_type(
+    as::ArraySize{FD,Nf},
+) where {FD, N, T, Nf} = FieldArray{FD, Nf, parent_array_type(T, Val(ndims(as)))}
+
+rebuild_field_array_type(
     ::Type{T},
-    ::Val{FD},
-    ::Val{Nf},
-    ::Val{ND},
-) where {T <: AbstractArray, FD, Nf, ND} =
-    FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
+    as::ArraySize{FD,Nf,S},
+) where {FD, T<:AbstractArray, Nf, S} = FieldArray{FD, Nf, parent_array_type(T, as)}
+
+# field_array_type(
+#     ::Type{T},
+#     ::Val{FD},
+#     ::Val{Nf},
+#     ::Val{ND},
+# ) where {T <: AbstractArray, FD, Nf, ND} =
+#     FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
 Base.ndims(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = Base.ndims(T) + 1
 Base.eltype(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = eltype(T)
 array_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = T
@@ -101,8 +116,16 @@ end
     end
 end
 
-Base.similar(fa::FieldArray{FD, N, T}) where {FD, N, T} =
-    FieldArray{FD, N, T}(ntuple(i -> similar(T, ndims(T)), N))
+Base.similar(fa::FieldArray{FD, N, T}, dims) where {FD, N, T} =
+    FieldArray{FD, N, T}(ntuple(i -> similar(T, dims), N))
+function Base.similar(::Type{FieldArray{FD, N, T}}, dims) where {FD, N, T}
+    FieldArray{FD, N, T}(ntuple(i -> similar(T, dims), N))
+end
+
+function Base.similar(::Type{FieldArray{FD, N, T}}) where {FD, N, T}
+    isconcretetype(T) || error("Array type $T is not concrete, pass `dims` to similar or use concrete array type.")
+    FieldArray{FD, N, T}(ntuple(i -> similar(T), N))
+end
 
 @inline insertafter(t::Tuple, i::Int, j::Int) =
     0 <= i <= length(t) ? _insertafter(t, i, j) : throw(BoundsError(t, i))
@@ -124,13 +147,12 @@ function Base.collect(fa::FieldArray{FD, N, T}) where {FD, N, T}
     return a
 end
 
-function Base.similar(::Type{<:FieldArray{FD, N, T}}, s) where {FD, N, T}
-    FieldArray{FD, N}(ntuple(i -> similar(T, s), N))
-end
-
 field_array(array::AbstractArray, s::AbstractDataLayoutSingleton) =
     field_array(array, field_dim(s))
 
+field_arrays(fa::FieldArray) = getfield(fa, :arrays)
+field_arrays(data::AbstractData) = field_arrays(parent(data))
+
 FieldArray{FD, N}(
     fa::FA,
 ) where {FD, N, T <: AbstractArray, FA <: NTuple{N, T}} =
@@ -164,7 +186,13 @@ function field_array(array::AbstractArray, fdim::Integer)
     return FieldArray{fdim, Nf, eltype(arrays)}(arrays)
 end
 
-function Base.getindex(fa::FieldArray{FD}, I::CartesianIndex) where {FD}
+function Base.:(==)(fa::FieldArray, array::AbstractArray)
+    return collect(fa) == array
+end
+
+Base.getindex(fa::FieldArray, I::Integer...) = getindex(fa, CartesianIndex(I))
+
+function Base.getindex(fa::FieldArray{FD}, I::CartesianIndex{5}) where {FD}
     FDI = I.I[FD]
     ND = length(I.I)
     Ipre = ntuple(i -> I.I[i], Val(FD - 1))
@@ -173,7 +201,7 @@ function Base.getindex(fa::FieldArray{FD}, I::CartesianIndex) where {FD}
     return fa.arrays[FDI][IA]
 end
 
-function Base.setindex!(fa::FieldArray{FD}, val, I::CartesianIndex) where {FD}
+function Base.setindex!(fa::FieldArray{FD}, val, I::CartesianIndex{5}) where {FD}
     FDI = I.I[FD]
     ND = length(I.I)
     Ipre = ntuple(i -> I.I[i], Val(FD - 1))

src/DataLayouts/non_extruded_broadcasted.jl

@@ -127,34 +127,3 @@ Base.IndexStyle(::Type{<:NonExtrudedBroadcasted{<:Any, <:Tuple{Any}}}) = IndexLi
 # ============================================================
 
 #! format: on
-# Datalayouts
-@propagate_inbounds function linear_getindex(
-    data::AbstractData{S},
-    I::Integer,
-) where {S}
-    s_array = farray_size(data)
-    ss = StaticSize(s_array, field_dim(data))
-    @inbounds get_struct_linear(parent(data), S, Val(field_dim(data)), ss, I)
-end
-@propagate_inbounds function linear_setindex!(
-    data::AbstractData{S},
-    val,
-    I::Integer,
-) where {S}
-    s_array = farray_size(data)
-    ss = StaticSize(s_array, field_dim(data))
-    @inbounds set_struct_linear!(
-        parent(data),
-        convert(S, val),
-        Val(field_dim(data)),
-        ss,
-        I,
-    )
-end
-
-for DL in (:IJKFVH, :IJFH, :IFH, :IJF, :IF, :VF, :VIJFH, :VIFH) # Skip DataF, since we want that to MethodError.
-    @eval @propagate_inbounds Base.getindex(data::$(DL), I::Integer) =
-        linear_getindex(data, I)
-    @eval @propagate_inbounds Base.setindex!(data::$(DL), val, I::Integer) =
-        linear_setindex!(data, val, I)
-end

src/DataLayouts/struct.jl

@@ -87,6 +87,7 @@ 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{<:Array{T}}, as::ArraySize) where {T} = Array{T, ndims(as)}
 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} =

test/DataLayouts/data0d.jl

@@ -129,9 +129,9 @@ end
     FT = Float64
     array = rand(FT, 2, 1)
     data = IF{FT, 2}(array)
-    @test DataLayouts.data2array(data) == reshape(parent(data), :)
-    @test parent(DataLayouts.array2data(DataLayouts.data2array(data), data)) ==
-          parent(data)
+    @test DataLayouts.data2array(data) == reshape(collect(parent(data)), :)
+    @test size(collect(parent(DataLayouts.array2data(DataLayouts.data2array(data), data)))) ==
+          size(collect(parent(data)))
 end
 
 @testset "broadcasting DataF + IFH data object => IFH" begin
@@ -148,9 +148,9 @@ end
     FT = Float64
     array = rand(FT, 2, 1, Nh)
     data = IFH{FT, 2, Nh}(array)
-    @test DataLayouts.data2array(data) == reshape(parent(data), :)
-    @test parent(DataLayouts.array2data(DataLayouts.data2array(data), data)) ==
-          parent(data)
+    @test DataLayouts.data2array(data) == reshape(collect(parent(data)), :)
+    @test size(collect(parent(DataLayouts.array2data(DataLayouts.data2array(data), data)))) ==
+          size(collect(parent(data)))
 end
 
 @testset "broadcasting DataF + IJF data object => IJF" begin
@@ -165,9 +165,9 @@ end
     FT = Float64
     array = rand(FT, 2, 2, 1)
     data = IJF{FT, 2}(array)
-    @test DataLayouts.data2array(data) == reshape(parent(data), :)
-    @test parent(DataLayouts.array2data(DataLayouts.data2array(data), data)) ==
-          parent(data)
+    @test DataLayouts.data2array(data) == reshape(collect(parent(data)), :)
+    @test size(collect(parent(DataLayouts.array2data(DataLayouts.data2array(data), data)))) ==
+          size(collect(parent(data)))
 end
 
 @testset "broadcasting DataF + IJFH data object => IJFH" begin
@@ -184,9 +184,9 @@ end
     Nh = 3
     array = rand(FT, 2, 2, 1, Nh)
     data = IJFH{FT, 2, Nh}(array)
-    @test DataLayouts.data2array(data) == reshape(parent(data), :)
-    @test parent(DataLayouts.array2data(DataLayouts.data2array(data), data)) ==
-          parent(data)
+    @test DataLayouts.data2array(data) == reshape(collect(parent(data)), :)
+    @test size(collect(parent(DataLayouts.array2data(DataLayouts.data2array(data), data)))) ==
+          size(collect(parent(data)))
 end
 
 @testset "broadcasting DataF + VIFH data object => VIFH" begin