Skip to content

Commit c02ea6a

Browse files
charleskawczynskicharleskawczynski
committed
Pass some unit tests, cleanup
1 parent 7c62c52 commit c02ea6a

File tree

4 files changed

+90
-127
lines changed

4 files changed

+90
-127
lines changed

‎src/DataLayouts/DataLayouts.jl

+38-69
Original file line numberDiff line numberDiff line change
@@ -259,62 +259,23 @@ end
259259
_getproperty(data, Val{name}())
260260
end
261261

262-
# Base.@propagate_inbounds function Base.getproperty(
263-
# data::AbstractData{S},
264-
# i::Integer,
265-
# ) where {S}
266-
# array = parent(data)
267-
# T = eltype(array)
268-
# SS = fieldtype(S, i)
269-
# offset = fieldtypeoffset(T, S, i)
270-
# nbytes = typesize(T, SS)
271-
# fdim = field_dim(data)
272-
# Ipre = ntuple(i -> Colon(), Val(fdim - 1))
273-
# Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim))
274-
# dataview =
275-
# @inbounds view(array, Ipre..., (offset + 1):(offset + nbytes), Ipost...)
276-
# union_all(data){SS, Base.tail(type_params(data))...}(dataview)
277-
# end
278-
279262
Base.@propagate_inbounds function Base.getproperty(
280263
data::AbstractData{S},
281264
i::Integer,
282265
) where {S}
283-
generic_property_view(data, i)
284-
end
285-
286-
# # In the past, we've sometimes needed a generated function
287-
# # for inference and constant propagation:
288-
# Base.@propagate_inbounds @generated function _property_view(
289-
# data::AD,
290-
# ::Val{Idx},
291-
# ) where {S, Idx, AD <: AbstractData{S}}
292-
# SS = fieldtype(S, Idx)
293-
# T = eltype(parent_array_type(AD))
294-
# offset = fieldtypeoffset(T, S, Val(Idx))
295-
# nbytes = typesize(T, SS)
296-
# fdim = field_dim(AD)
297-
# Ipre = ntuple(i -> Colon(), Val(fdim - 1))
298-
# Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim))
299-
# field_byterange = (offset + 1):(offset + nbytes)
300-
# return :($(union_all(AD)){$SS, $(Base.tail(type_params(AD)))...}(
301-
# @inbounds view(parent(data), $Ipre..., $field_byterange, $Ipost...)
302-
# ))
303-
# end
266+
P = Base.tail(type_params(data))
267+
SS = fieldtype(S, i)
268+
return union_all(data){SS, P...}(generic_property_view(data, i))
269+
end
304270

305271
@inline function generic_property_view(
306272
data::AbstractData{S},
307273
i::Integer,
308274
) where {S}
309275
fa = parent(data)
310276
T = eltype(fa)
311-
@show i
312-
@show S
313-
@show T
314277
SS = fieldtype(S, i)
315-
@show SS
316278
offset = fieldtypeoffset(T, S, i)
317-
@show offset
318279
nbytes = typesize(T, SS)
319280
field_byterange = (offset + 1):(offset + nbytes)
320281
return FieldArray{field_dim(data)}(ntuple(jf -> parent(data).arrays[offset + jf], nbytes))
@@ -634,6 +595,7 @@ function DataF{S}(::Type{ArrayType}) where {S, ArrayType}
634595
T = eltype(ArrayType)
635596
DataF{S}(ArrayType(undef, typesize(T, S)))
636597
end
598+
DataF{S}(fa::FieldArray) where {S} = DataF{S, typeof(fa)}(fa)
637599

638600
function DataF(x::T) where {T}
639601
if is_valid_basetype(Float64, T)
@@ -722,6 +684,8 @@ function IJF{S, Nij}(array::AbstractArray{T, 3}) where {S, Nij, T}
722684
IJF{S, Nij, typeof(fa)}(fa)
723685
end
724686

687+
IJF{S, Nij}(fa::FieldArray) where {S, Nij} = IJF{S, Nij, typeof(fa)}(fa)
688+
725689
function IJF{S, Nij}(::Type{MArray}, ::Type{T}) where {S, Nij, T}
726690
Nf = typesize(T, S)
727691
array = MArray{Tuple{Nij, Nij, Nf}, T, 3, Nij * Nij * Nf}(undef)
@@ -881,6 +845,8 @@ struct VF{S, Nv, FA <: FieldArray} <: DataColumn{S, Nv}
881845
array::FA
882846
end
883847

848+
VF{S, Nv}(fa::FieldArray) where {S, Nv} = VF{S, Nv, typeof(fa)}(fa)
849+
884850
parent_array_type(::Type{VF{S, Nv, A}}) where {S, Nv, A} = parent_array_type(A)
885851

886852
function VF{S, Nv}(array::AbstractArray{T, 2}) where {S, Nv, T}
@@ -907,8 +873,8 @@ Base.lastindex(data::VF) = length(data)
907873

908874
nlevels(::VF{S, Nv}) where {S, Nv} = Nv
909875

910-
Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) =
911-
_property_view(data, Val(i))
876+
# Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) =
877+
# generic_property_view(data, i)
912878

913879
@inline function Base.getindex(data::VF{S, Nv}, I::CartesianIndex) where {S, Nv}
914880
v = I.I[4]
@@ -1129,14 +1095,15 @@ end
11291095
# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7)
11301096
@inline function column(data::VIFH{S, Nv, Ni, Nh}, i, h) where {S, Nv, Ni, Nh}
11311097
array = parent(data)
1132-
@boundscheck (1 <= i <= Ni && 1 <= h <= Nh) ||
1133-
throw(BoundsError(data, (i, h)))
1098+
@boundscheck (1 <= i <= Ni && 1 <= h <= Nh) || throw(BoundsError(data, (i, h)))
11341099
Nf = ncomponents(data)
1135-
dataview = @inbounds SubArray(
1136-
array,
1137-
(Base.Slice(Base.OneTo(Nv)), i, Base.Slice(Base.OneTo(Nf)), h),
1138-
)
1139-
VF{S, Nv}(dataview)
1100+
dataview = @inbounds ntuple(Nf) do f
1101+
SubArray(
1102+
array.arrays[f],
1103+
(Base.Slice(Base.OneTo(Nv)), i, h),
1104+
)
1105+
end
1106+
VF{S, Nv}(FieldArray{field_dim(VF)}(dataview))
11401107
end
11411108

11421109
@inline function column(
@@ -1449,6 +1416,25 @@ type parameters.
14491416
@inline farray_size(data::VIJFH{S, Nv, Nij, Nh}) where {S, Nv, Nij, Nh} = (Nv, Nij, Nij, ncomponents(data), Nh)
14501417
@inline farray_size(data::VIFH{S, Nv, Ni, Nh}) where {S, Nv, Ni, Nh} = (Nv, Ni, ncomponents(data), Nh)
14511418

1419+
"""
1420+
float_type(data::AbstractData)
1421+
1422+
This is an internal function, please do not use outside of ClimaCore.
1423+
1424+
Returns the underlying float type of the backing array.
1425+
"""
1426+
@inline float_type(::Type{IJKFVH{S, Nij, Nk, Nv, Nh, FA}}) where {S, Nij, Nk, Nv, Nh, FA} = eltype(FA)
1427+
@inline float_type(::Type{IJFH{S, Nij, Nh, FA}}) where {S, Nij, Nh, FA} = eltype(FA)
1428+
@inline float_type(::Type{IFH{S, Ni, Nh, FA}}) where {S, Ni, Nh, FA} = eltype(FA)
1429+
@inline float_type(::Type{DataF{S, FA}}) where {S, FA} = eltype(FA)
1430+
@inline float_type(::Type{IJF{S, Nij, FA}}) where {S, Nij, FA} = eltype(FA)
1431+
@inline float_type(::Type{IF{S, Ni, FA}}) where {S, Ni, FA} = eltype(FA)
1432+
@inline float_type(::Type{VF{S, Nv, FA}}) where {S, Nv, FA} = eltype(FA)
1433+
@inline float_type(::Type{VIJFH{S, Nv, Nij, Nh, FA}}) where {S, Nv, Nij, Nh, FA} = eltype(FA)
1434+
@inline float_type(::Type{VIFH{S, Nv, Ni, Nh, FA}}) where {S, Nv, Ni, Nh, FA} = eltype(FA)
1435+
@inline float_type(::Type{IH1JH2{S, Nij, A}}) where {S, Nij, A} = eltype(A)
1436+
@inline float_type(::Type{IV1JH2{S, n1, Ni, A}}) where {S, n1, Ni, A} = eltype(A)
1437+
14521438
# Keep in sync with definition(s) in libs.
14531439
@inline slab_index(i, j) = CartesianIndex(i, j, 1, 1, 1)
14541440
@inline slab_index(i) = CartesianIndex(i, 1, 1, 1, 1)
@@ -1548,21 +1534,4 @@ include("mapreduce.jl")
15481534
include("non_extruded_broadcasted.jl")
15491535
include("has_uniform_datalayouts.jl")
15501536

1551-
1552-
float_type(
1553-
::Type{IJKFVH{S, Nij, Nk, Nv, Nh, FA}},
1554-
) where {S, Nij, Nk, Nv, Nh, FA} = eltype(FA)
1555-
float_type(::Type{IJFH{S, Nij, Nh, FA}}) where {S, Nij, Nh, FA} = eltype(FA)
1556-
float_type(::Type{IFH{S, Ni, Nh, FA}}) where {S, Ni, Nh, FA} = eltype(FA)
1557-
float_type(::Type{DataF{S, FA}}) where {S, FA} = eltype(FA)
1558-
float_type(::Type{IJF{S, Nij, FA}}) where {S, Nij, FA} = eltype(FA)
1559-
float_type(::Type{IF{S, Ni, FA}}) where {S, Ni, FA} = eltype(FA)
1560-
float_type(::Type{VF{S, Nv, FA}}) where {S, Nv, FA} = eltype(FA)
1561-
float_type(::Type{VIJFH{S, Nv, Nij, Nh, FA}}) where {S, Nv, Nij, Nh, FA} =
1562-
eltype(FA)
1563-
float_type(::Type{VIFH{S, Nv, Ni, Nh, FA}}) where {S, Nv, Ni, Nh, FA} =
1564-
eltype(FA)
1565-
float_type(::Type{IH1JH2{S, Nij, A}}) where {S, Nij, A} = eltype(A)
1566-
float_type(::Type{IV1JH2{S, n1, Ni, A}}) where {S, n1, Ni, A} = eltype(A)
1567-
15681537
end # module

‎src/DataLayouts/field_array.jl

+46-14
Original file line numberDiff line numberDiff line change
@@ -82,18 +82,45 @@ elsize(fa::FieldArray{FD, N, T}) where {FD, N, T} = size(fa.arrays[1])
8282
end
8383
end
8484

85+
@inline function Base.copyto!(x::FieldArray{FD, N, T}, y::AbstractArray{TA, NA}) where {FD, N, T, TA, NA}
86+
if ndims(T) == NA
87+
@inbounds for i in 1:N
88+
Base.copyto!(x.arrays[i], y)
89+
end
90+
elseif ndims(T) + 1 == NA
91+
@inbounds for I in CartesianIndices(y)
92+
x[I] = y[I]
93+
end
94+
end
95+
end
96+
8597
Base.similar(fa::FieldArray{FD, N, T}) where {FD, N, T} =
8698
FieldArray{FD, N, T}(ntuple(i -> similar(T, ndims(T)), N))
8799

100+
@inline insertafter(t::Tuple, i::Int, j::Int) =
101+
0 <= i <= length(t) ? _insertafter(t, i, j) : throw(BoundsError(t, i))
102+
@inline _insertafter(t::Tuple, i::Int, j::Int) =
103+
i == 0 ? (j, t...) : (t[1], _insertafter(Base.tail(t), i - 1, j)...)
104+
105+
original_dims(fa::FieldArray{FD, N, T}) where {FD, N, T} =
106+
insertafter(elsize(fa), FD-1, N)
107+
function Base.collect(fa::FieldArray{FD, N, T}) where {FD, N, T}
108+
odims = original_dims(fa)
109+
ND = ndims(T)+1
110+
a = similar(fa.arrays[1], eltype(T), odims)
111+
@inbounds for i in 1:N
112+
Ipre = ntuple(i -> Colon(), Val(FD - 1))
113+
Ipost = ntuple(i -> Colon(), Val(ND - FD))
114+
av = view(a, Ipre..., i, Ipost...)
115+
Base.copyto!(av, fa.arrays[i])
116+
end
117+
return a
118+
end
119+
88120
function Base.similar(::Type{<:FieldArray{FD, N, T}}, s) where {FD, N, T}
89121
FieldArray{FD, N}(ntuple(i -> similar(T, s), N))
90122
end
91123

92-
pre_post_colons(f, dim) = (;
93-
Ipre = ntuple(i -> Colon(), Val(length(size(f)[1:(dim - 1)]))),
94-
Ipost = ntuple(i -> Colon(), Val(length(size(f)[(dim + 1):end]))),
95-
)
96-
97124
field_array(array::AbstractArray, s::AbstractDataLayoutSingleton) =
98125
field_array(array, field_dim(s))
99126

@@ -114,7 +141,12 @@ function field_array(array::AbstractArray, fdim::Integer)
114141
arrays = ntuple(Nf) do i
115142
similar(farray)
116143
end
117-
(; Ipre, Ipost) = pre_post_colons(array, fdim)
144+
# Ipre = ntuple(i -> Colon(), Val(length(size(array)[1:(fdim - 1)])))
145+
# Ipost = ntuple(i -> Colon(), Val(length(size(array)[(fdim + 1):end])))
146+
ND = ndims(array) # TODO: confirm
147+
Ipre = ntuple(i -> Colon(), Val(fdim - 1))
148+
Ipost = ntuple(i -> Colon(), Val(ND - fdim))
149+
118150
S = map(x -> size(x), arrays)
119151
for i in 1:Nf
120152
arrays[i] .= array[Ipre..., i, Ipost...]
@@ -124,19 +156,19 @@ function field_array(array::AbstractArray, fdim::Integer)
124156
end
125157

126158
function Base.getindex(fa::FieldArray{FD}, I::CartesianIndex) where {FD}
127-
fdim = I.I[FD]
159+
FDI = I.I[FD]
128160
ND = length(I.I)
129-
Ipre = ntuple(i -> I.I[i], Val(fdim - 1))
130-
Ipost = ntuple(i -> I.I[i], Val(ND - fdim))
161+
Ipre = ntuple(i -> I.I[i], Val(FD - 1))
162+
Ipost = ntuple(i -> I.I[i], Val(ND - FD))
131163
IA = CartesianIndex((Ipre...,Ipost...))
132-
return fa.arrays[fdim][IA]
164+
return fa.arrays[FDI][IA]
133165
end
134166

135167
function Base.setindex!(fa::FieldArray{FD}, val, I::CartesianIndex) where {FD}
136-
fdim = I.I[FD]
168+
FDI = I.I[FD]
137169
ND = length(I.I)
138-
Ipre = ntuple(i -> I.I[i], Val(fdim - 1))
139-
Ipost = ntuple(i -> I.I[i], Val(ND - fdim))
170+
Ipre = ntuple(i -> I.I[i], Val(FD - 1))
171+
Ipost = ntuple(i -> I.I[i], Val(ND - FD))
140172
IA = CartesianIndex((Ipre...,Ipost...))
141-
return fa.arrays[fdim][IA]
173+
fa.arrays[FDI][IA] = val
142174
end

‎test/DataLayouts/unit_copyto.jl

+5-43
Original file line numberDiff line numberDiff line change
@@ -16,23 +16,21 @@ function test_copyto_float!(data)
1616
pdata = parent(data)
1717
Random.seed!(1234)
1818
# Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803
19-
rand_data = DataLayouts.rebuild(data, similar(pdata))
19+
rand_data = DataLayouts.rebuild(data, collect(pdata))
2020
ArrayType = ClimaComms.array_type(ClimaComms.device())
21-
parent(rand_data).arrays[1] .=
22-
ArrayType(rand(eltype(parent(data)), DataLayouts.farray_size(data)))
21+
Base.copyto!(parent(rand_data), ArrayType(rand(eltype(parent(data)), DataLayouts.elsize(parent(data)))))
2322
Base.copyto!(data, rand_data) # test copyto!(::AbstractData, ::AbstractData)
24-
@test all(pdata.arrays[1] .== parent(rand_data).arrays[1])
23+
@test all(pdata .== parent(rand_data))
2524
Base.copyto!(data, Base.Broadcast.broadcasted(+, rand_data, 1)) # test copyto!(::AbstractData, ::Broadcasted)
2625
@test all(pdata .== parent(rand_data) .+ 1)
2726
end
2827

2928
function test_copyto!(data)
3029
Random.seed!(1234)
3130
# Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803
32-
rand_data = DataLayouts.rebuild(data, similar(parent(data)))
31+
rand_data = DataLayouts.rebuild(data, collect(parent(data)))
3332
ArrayType = ClimaComms.array_type(ClimaComms.device())
34-
parent(rand_data) .=
35-
ArrayType(rand(eltype(parent(data)), DataLayouts.farray_size(data)))
33+
Base.copyto!(parent(rand_data), ArrayType(rand(eltype(parent(data)), DataLayouts.farray_size(data))))
3634
Base.copyto!(data, rand_data) # test copyto!(::AbstractData, ::AbstractData)
3735
@test all(parent(data.:1) .== parent(rand_data.:1))
3836
@test all(parent(data.:2) .== parent(rand_data.:2))
@@ -91,39 +89,3 @@ end
9189
# data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_copyto!(data) # TODO: test
9290
# data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_copyto!(data) # TODO: test
9391
end
94-
95-
@testset "copyto! views with Nf > 1" begin
96-
device = ClimaComms.device()
97-
device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...))
98-
data_view(data) = DataLayouts.rebuild(
99-
data,
100-
SubArray(
101-
parent(data),
102-
ntuple(
103-
i -> Base.Slice(Base.OneTo(DataLayouts.farray_size(data, i))),
104-
ndims(data),
105-
),
106-
),
107-
)
108-
FT = Float64
109-
S = Tuple{FT, FT}
110-
Nf = 2
111-
Nv = 4
112-
Nij = 3
113-
Nh = 5
114-
Nk = 6
115-
# Rather than using level/slab/column, let's just make views/SubArrays
116-
# directly so that we can easily test all cases:
117-
#! format: off
118-
data = IJFH{S, Nij, Nh}(device_zeros(FT,Nij,Nij,Nf,Nh)); test_copyto!(data_view(data))
119-
data = IFH{S, Nij, Nh}(device_zeros(FT,Nij,Nf,Nh)); test_copyto!(data_view(data))
120-
data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); test_copyto!(data_view(data))
121-
data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); test_copyto!(data_view(data))
122-
data = VF{S, Nv}(device_zeros(FT,Nv,Nf)); test_copyto!(data_view(data))
123-
data = VIJFH{S,Nv,Nij,Nh}(device_zeros(FT,Nv,Nij,Nij,Nf,Nh));test_copyto!(data_view(data))
124-
data = VIFH{S, Nv, Nij, Nh}(device_zeros(FT,Nv,Nij,Nf,Nh)); test_copyto!(data_view(data))
125-
#! format: on
126-
# TODO: test this
127-
# data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_copyto!(data) # TODO: test
128-
# data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_copyto!(data) # TODO: test
129-
end

‎test/DataLayouts/unit_fill.jl

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ function test_fill!(data, vals::Tuple{<:Any, <:Any})
1515
end
1616
function test_fill!(data, val::Real)
1717
fill!(data, val)
18-
@test all(parent(data).arrays[1] .== val)
18+
@test all(parent(data) .== val)
1919
end
2020

2121
@testset "fill! with Nf = 1" begin

0 commit comments

Comments
 (0)