Skip to content

Commit 7c62c52

Browse files
charleskawczynskicharleskawczynski
committed
wip
1 parent 90aaab1 commit 7c62c52

File tree

5 files changed

+125
-91
lines changed

5 files changed

+125
-91
lines changed

‎src/DataLayouts/DataLayouts.jl

Lines changed: 52 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -259,60 +259,72 @@ 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+
262279
Base.@propagate_inbounds function Base.getproperty(
263280
data::AbstractData{S},
264281
i::Integer,
265282
) 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-
279-
# In the past, we've sometimes needed a generated function
280-
# for inference and constant propagation:
281-
Base.@propagate_inbounds @generated function _property_view(
282-
data::AD,
283-
::Val{Idx},
284-
) where {S, Idx, AD <: AbstractData{S}}
285-
SS = fieldtype(S, Idx)
286-
T = eltype(parent_array_type(AD))
287-
offset = fieldtypeoffset(T, S, Val(Idx))
288-
nbytes = typesize(T, SS)
289-
fdim = field_dim(AD)
290-
Ipre = ntuple(i -> Colon(), Val(fdim - 1))
291-
Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim))
292-
field_byterange = (offset + 1):(offset + nbytes)
293-
return :($(union_all(AD)){$SS, $(Base.tail(type_params(AD)))...}(
294-
@inbounds view(parent(data), $Ipre..., $field_byterange, $Ipost...)
295-
))
296-
end
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
297304

298305
@inline function generic_property_view(
299306
data::AbstractData{S},
300307
i::Integer,
301308
) where {S}
302309
fa = parent(data)
303310
T = eltype(fa)
311+
@show i
312+
@show S
313+
@show T
304314
SS = fieldtype(S, i)
315+
@show SS
305316
offset = fieldtypeoffset(T, S, i)
317+
@show offset
306318
nbytes = typesize(T, SS)
307319
field_byterange = (offset + 1):(offset + nbytes)
308-
return FieldArray(ntuple(jf -> parent(data).arrays[offset + jf], nbytes))
320+
return FieldArray{field_dim(data)}(ntuple(jf -> parent(data).arrays[offset + jf], nbytes))
309321
end
310322

311323
@inline @generated function generic_property_view(
312324
data::AbstractData{S},
313325
::Val{Idx},
314326
) where {S, Idx}
315-
:(FieldArray(
327+
:(FieldArray{field_dim(data)}(
316328
ntuple(
317329
jf -> parent(data).arrays[fieldtypeoffset(
318330
eltype(parent(data)),
@@ -469,7 +481,7 @@ Base.length(data::IJFH) = get_Nh(data)
469481
sub_arrays = ntuple(ncomponents(fa)) do jf
470482
view(fa.arrays[jf], :, :, h)
471483
end
472-
dataview = FieldArray(sub_arrays)
484+
dataview = FieldArray{field_dim(IJF)}(sub_arrays)
473485
IJF{S, Nij, typeof(dataview)}(dataview)
474486
end
475487

@@ -483,7 +495,7 @@ end
483495
@boundscheck (1 <= j <= Nij && 1 <= i <= Nij && 1 <= h <= get_Nh(data)) ||
484496
throw(BoundsError(data, (i, j, h)))
485497
fa = parent(data)
486-
dataview = @inbounds FieldArray(
498+
dataview = @inbounds FieldArray{field_dim(DataF)}(
487499
ntuple(jf -> view(parent(fa.arrays[jf]), i, j, h), ncomponents(fa)),
488500
)
489501
DataF{S, typeof(dataview)}(dataview)
@@ -556,15 +568,15 @@ end
556568
i -> view(parent(data).arrays[i], :, h),
557569
ncomponents(data),
558570
)
559-
dataview = FieldArray(toa_view)
571+
dataview = FieldArray{field_dim(IF)}(toa_view)
560572
IF{S, Ni}(dataview)
561573
end
562574
Base.@propagate_inbounds slab(data::IFH, v::Integer, h::Integer) = slab(data, h)
563575

564576
@inline function column(data::IFH{S, Ni}, i, h) where {S, Ni}
565577
@boundscheck (1 <= h <= get_Nh(data) && 1 <= i <= Ni) ||
566578
throw(BoundsError(data, (i, h)))
567-
dataview = @inbounds FieldArray(
579+
dataview = @inbounds FieldArray{field_dim(DataF)}(
568580
ntuple(jf -> view(parent(fa.arrays[jf]), i, h), ncomponents(fa)),
569581
)
570582
DataF{S, typeof(dataview)}(dataview)
@@ -757,7 +769,7 @@ end
757769
@inline function column(data::IJF{S, Nij}, i, j) where {S, Nij}
758770
@boundscheck (1 <= j <= Nij && 1 <= i <= Nij) ||
759771
throw(BoundsError(data, (i, j)))
760-
dataview = @inbounds FieldArray(
772+
dataview = @inbounds FieldArray{field_dim(DataF)}(
761773
ntuple(jf -> view(parent(fa.arrays[jf]), i, j), ncomponents(fa)),
762774
)
763775
DataF{S, typeof(dataview)}(dataview)
@@ -843,7 +855,7 @@ end
843855
@inline function column(data::IF{S, Ni}, i) where {S, Ni}
844856
@boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
845857
fa = parent(data)
846-
dataview = @inbounds FieldArray(
858+
dataview = @inbounds FieldArray{field_dim(DataF)}(
847859
ntuple(jf -> view(parent(fa.arrays[jf]), i), ncomponents(fa)),
848860
)
849861
DataF{S, typeof(dataview)}(dataview)
@@ -1007,7 +1019,7 @@ end
10071019
@boundscheck (1 <= i <= Nij && 1 <= j <= Nij && 1 <= h <= Nh) ||
10081020
throw(BoundsError(data, (i, j, h)))
10091021
Nf = ncomponents(data)
1010-
dataview = @inbounds FieldArray(
1022+
dataview = @inbounds FieldArray{field_dim(VF)}(
10111023
ntuple(ncomponents(fa)) do jf
10121024
SubArray(parent(fa.arrays[jf]), (Base.Slice(Base.OneTo(Nv)), i, j, h))
10131025
end,

‎src/DataLayouts/broadcast.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ function Base.similar(
378378
) where {Nv, Nij, Nh, A, Eltype, newNv}
379379
T = eltype(A)
380380
Nf = typesize(eltype(A), Eltype)
381-
fat = field_array_type(A, Val(Nf), Val(4))
381+
fat = field_array_type(A, Val(field_dim(VIJFH)), Val(Nf), Val(4))
382382
array = similar(fat, Base.Dims((newNv, Nij, Nij, Nh)))
383383
vd = VIJFH{Eltype, newNv, Nij, Nh}(array)
384384
return vd

‎src/DataLayouts/field_array.jl

Lines changed: 54 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -26,66 +26,67 @@ end
2626
@inline field_dim(::VIJFHSingleton) = 4
2727
@inline field_dim(::VIFHSingleton) = 3
2828

29-
struct FieldArray{N, T <: AbstractArray}
29+
struct FieldArray{FD, N, T <: AbstractArray}
3030
arrays::NTuple{N, T}
3131
end
32-
function Adapt.adapt_structure(to, fa::FieldArray{N, T}) where {N, T}
32+
function Adapt.adapt_structure(to, fa::FieldArray{FD, N, T}) where {FD, N, T}
3333
arrays = ntuple(i -> Adapt.adapt(to, fa.arrays[i]), N)
34-
FieldArray{N, eltype(arrays)}(arrays)
34+
FieldArray{FD, N, eltype(arrays)}(arrays)
3535
end
36-
function rebuild(fa::FieldArray{N, T}, ::Type{DA}) where {N, T, DA}
36+
function rebuild(fa::FieldArray{FD, N, T}, ::Type{DA}) where {FD, N, T, DA}
3737
arrays = ntuple(i -> DA(fa.arrays[i]), N)
38-
FieldArray{N, eltype(arrays)}(arrays)
38+
FieldArray{FD, N, eltype(arrays)}(arrays)
3939
end
4040

41-
Base.length(fa::FieldArray{N, T}) where {N, T} = N
42-
Base.copy(fa::FieldArray{N, T}) where {N, T} =
43-
FieldArray{N, T}(ntuple(i -> copy(fa.arrays[i]), N))
41+
Base.length(fa::FieldArray{FD, N, T}) where {FD, N, T} = N
42+
Base.copy(fa::FieldArray{FD, N, T}) where {FD, N, T} =
43+
FieldArray{FD, N, T}(ntuple(i -> copy(fa.arrays[i]), N))
4444

45-
float_type(::Type{FieldArray{N, T}}) where {N, T} = eltype(T)
46-
parent_array_type(::Type{FieldArray{N, T}}) where {N, T} = parent_array_type(T)
45+
float_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = eltype(T)
46+
parent_array_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = parent_array_type(T)
4747
# field_array_type(::Type{FieldArray{N,T}}, ::Val{Nf}) where {N,T, Nf} = FieldArray{Nf, parent_array_type(T, Val(ndims(T)))}
4848
field_array_type(
49-
::Type{FieldArray{N, T}},
49+
::Type{FieldArray{FD, N, T}},
5050
::Val{Nf},
5151
::Val{ND},
52-
) where {N, T, Nf, ND} = FieldArray{Nf, parent_array_type(T, Val(ND))}
52+
) where {FD, N, T, Nf, ND} = FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
5353
# field_array_type(::Type{T}, ::Val{Nf}) where {T <: AbstractArray, Nf} = FieldArray{Nf, parent_array_type(T, Val(ndims(T)))}
5454
field_array_type(
5555
::Type{T},
56+
::Val{FD},
5657
::Val{Nf},
5758
::Val{ND},
58-
) where {T <: AbstractArray, Nf, ND} =
59-
FieldArray{Nf, parent_array_type(T, Val(ND))}
59+
) where {T <: AbstractArray, FD, Nf, ND} =
60+
FieldArray{FD, Nf, parent_array_type(T, Val(ND))}
6061
# Base.ndims(::Type{FieldArray{N,T}}) where {N,T} = Base.ndims(T)+1
61-
Base.eltype(::Type{FieldArray{N, T}}) where {N, T} = eltype(T)
62-
array_type(::Type{FieldArray{N, T}}) where {N, T} = T
63-
ncomponents(::Type{FieldArray{N, T}}) where {N, T} = N
64-
ncomponents(::FieldArray{N, T}) where {N, T} = N
62+
Base.eltype(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = eltype(T)
63+
array_type(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = T
64+
ncomponents(::Type{FieldArray{FD, N, T}}) where {FD, N, T} = N
65+
ncomponents(::FieldArray{FD, N, T}) where {FD, N, T} = N
6566
# Base.size(fa::FieldArray{N,T}) where {N,T} = size(fa.arrays[1])
66-
to_parent_array_type(::Type{FA}) where {N, T, FA <: FieldArray{N, T}} =
67-
FieldArray{N, parent_array_type(FA)}
67+
to_parent_array_type(::Type{FA}) where {FD, N, T, FA <: FieldArray{FD, N, T}} =
68+
FieldArray{FD, N, parent_array_type(FA)}
6869

6970
promote_parent_array_type(
70-
::Type{FieldArray{NA, A}},
71-
::Type{FieldArray{NB, B}},
72-
) where {NA, NB, A, B} =
71+
::Type{FieldArray{FDA, NA, A}},
72+
::Type{FieldArray{FDB, NB, B}},
73+
) where {FDA, FDB, NA, NB, A, B} =
7374
# FieldArray{N, promote_parent_array_type(A, B)} where {N}
7475
promote_parent_array_type(A, B)
7576

76-
elsize(fa::FieldArray{N, T}) where {N, T} = size(fa.arrays[1])
77+
elsize(fa::FieldArray{FD, N, T}) where {FD, N, T} = size(fa.arrays[1])
7778

78-
@inline function Base.copyto!(x::FA, y::FA) where {N, T, FA <: FieldArray{N, T}}
79+
@inline function Base.copyto!(x::FA, y::FA) where {FD, N, T, FA <: FieldArray{FD, N, T}}
7980
@inbounds for i in 1:N
8081
Base.copyto!(x.arrays[i], y.arrays[i])
8182
end
8283
end
8384

84-
Base.similar(fa::FieldArray{N, T}) where {N, T} =
85-
FieldArray{N, T}(ntuple(i -> similar(T, ndims(T)), N))
85+
Base.similar(fa::FieldArray{FD, N, T}) where {FD, N, T} =
86+
FieldArray{FD, N, T}(ntuple(i -> similar(T, ndims(T)), N))
8687

87-
function Base.similar(::Type{<:FieldArray{N, T}}, s) where {N, T}
88-
FieldArray{N}(ntuple(i -> similar(T, s), N))
88+
function Base.similar(::Type{<:FieldArray{FD, N, T}}, s) where {FD, N, T}
89+
FieldArray{FD, N}(ntuple(i -> similar(T, s), N))
8990
end
9091

9192
pre_post_colons(f, dim) = (;
@@ -96,8 +97,11 @@ pre_post_colons(f, dim) = (;
9697
field_array(array::AbstractArray, s::AbstractDataLayoutSingleton) =
9798
field_array(array, field_dim(s))
9899

99-
FieldArray{N}(fa::FA) where {N, T <: AbstractArray, FA <: NTuple{N, T}} =
100-
FieldArray{N, T}(fa)
100+
FieldArray{FD, N}(fa::FA) where {FD, N, T <: AbstractArray, FA <: NTuple{N, T}} =
101+
FieldArray{FD, N, T}(fa)
102+
103+
FieldArray{FD}(fa::FA) where {FD, N, T <: AbstractArray, FA <: NTuple{N, T}} =
104+
FieldArray{FD, N, T}(fa)
101105

102106
function field_array(array::AbstractArray, fdim::Integer)
103107
Nf = size(array, fdim)
@@ -116,5 +120,23 @@ function field_array(array::AbstractArray, fdim::Integer)
116120
arrays[i] .= array[Ipre..., i, Ipost...]
117121
end
118122
end
119-
return FieldArray{Nf, eltype(arrays)}(arrays)
123+
return FieldArray{fdim, Nf, eltype(arrays)}(arrays)
124+
end
125+
126+
function Base.getindex(fa::FieldArray{FD}, I::CartesianIndex) where {FD}
127+
fdim = I.I[FD]
128+
ND = length(I.I)
129+
Ipre = ntuple(i -> I.I[i], Val(fdim - 1))
130+
Ipost = ntuple(i -> I.I[i], Val(ND - fdim))
131+
IA = CartesianIndex((Ipre...,Ipost...))
132+
return fa.arrays[fdim][IA]
133+
end
134+
135+
function Base.setindex!(fa::FieldArray{FD}, val, I::CartesianIndex) where {FD}
136+
fdim = I.I[FD]
137+
ND = length(I.I)
138+
Ipre = ntuple(i -> I.I[i], Val(fdim - 1))
139+
Ipost = ntuple(i -> I.I[i], Val(ND - fdim))
140+
IA = CartesianIndex((Ipre...,Ipost...))
141+
return fa.arrays[fdim][IA]
120142
end

0 commit comments

Comments
 (0)