Skip to content

Commit

Permalink
Reduce more code duplication
Browse files Browse the repository at this point in the history
  • Loading branch information
charleskawczynski committed Aug 14, 2024
1 parent 4800990 commit dc72795
Showing 1 changed file with 21 additions and 140 deletions.
161 changes: 21 additions & 140 deletions src/DataLayouts/DataLayouts.jl
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ abstract type AbstractData{S} end
@inline Base.size(data::AbstractData) = universal_size(data)

"""
struct UniversalSize{Nv, Nij, Nh} end
struct UniversalSize{Ni, Nj, Nv, Nh} end
UniversalSize(data::AbstractData)
A struct containing static dimensions, universal to all datalayouts:
Expand All @@ -66,7 +66,7 @@ struct UniversalSize{Ni, Nj, Nv, Nh} end

@inline function UniversalSize(data::AbstractData)
us = universal_size(data)
UniversalSize{us[1], us[2], us[3], us[5]}()
UniversalSize{us[1], us[2], us[4], us[5]}()
end

@inline array_length(data::AbstractData) = prod(size(parent(data)))
Expand Down Expand Up @@ -267,6 +267,25 @@ Base.@propagate_inbounds function Base.getproperty(
union_all(data){SS, Base.tail(type_params(data))...}(dataview)
end

# In the past, we've sometimes needed a generated function
# for inference and constant propagation:
Base.@propagate_inbounds @generated function _property_view(
data::AD,
::Val{Idx},
) where {S, Idx, AD <: AbstractData{S}}
SS = fieldtype(S, Idx)
T = eltype(parent_array_type(AD))
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
fdim = field_dim(AD)
Ipre = ntuple(i -> Colon(), Val(fdim - 1))
Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim))
field_byterange = (offset + 1):(offset + nbytes)
return :($(union_all(AD)){$SS, $(Base.tail(type_params(AD)))...}(
@inbounds view(parent(data), $Ipre..., $field_byterange, $Ipost...)
))
end

function replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T}
array = parent(data)
S′ = replace_basetype(eltype(array), T, S)
Expand All @@ -275,24 +294,6 @@ function replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T}
)
end

# @generated function _property_view(
# data::AD,
# ::Val{Idx},
# ) where {S, Idx, AD <: AbstractData{S}}
# SS = fieldtype(S, Idx)
# T = eltype(parent_array_type(AD))
# offset = fieldtypeoffset(T, S, Val(Idx))
# nbytes = typesize(T, SS)
# fdim = field_dim(AD)
# Ipre = ntuple(i -> Colon(), Val(fdim - 1))
# Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim))
# field_byterange = (offset + 1):(offset + nbytes)
# return :($(union_all(AD)){$SS, $(type_params(AD))...}(
# @inbounds view(parent(data), $Ipre..., $field_byterange, $Ipost...)
# ))
# end


# ==================
# Data3D DataLayout
# ==================
Expand Down Expand Up @@ -323,20 +324,6 @@ function IJKFVH{S, Nij, Nk, Nv, Nh}(
IJKFVH{S, Nij, Nk, Nv, Nh, typeof(array)}(array)
end

@generated function _property_view(
data::IJKFVH{S, Nij, Nk, Nv, Nh, A},
::Val{Idx},
) where {S, Nij, Nk, Nv, Nh, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(IJKFVH{$SS, $Nij, $Nk, $Nv, $Nh}(
@inbounds view(parent(data), :, :, :, $field_byterange, :, :)
))
end

@inline universal_size(
::IJKFVH{S, Nij, Nk, Nv, Nh},
) where {S, Nij, Nk, Nv, Nh} = (Nij, Nij, Nk, Nv, Nh)
Expand Down Expand Up @@ -417,20 +404,6 @@ end

Base.length(data::IJFH) = get_Nh(data)

@generated function _property_view(
data::IJFH{S, Nij, Nh, A},
::Val{Idx},
) where {S, Nij, Nh, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(IJFH{$SS, $Nij, $Nh}(
@inbounds view(parent(data), :, :, $field_byterange, :)
))
end

@inline function slab(data::IJFH{S, Nij}, h::Integer) where {S, Nij}
@boundscheck (1 <= h <= get_Nh(data)) || throw(BoundsError(data, (h,)))
dataview = @inbounds view(parent(data), :, :, :, h)
Expand Down Expand Up @@ -528,20 +501,6 @@ end
Base.@propagate_inbounds column(data::IFH{S, Ni}, i, j, h) where {S, Ni} =
column(data, i, h)

@generated function _property_view(
data::IFH{S, Ni, Nh, A},
::Val{Idx},
) where {S, Ni, Nh, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(IFH{$SS, $Ni, $Nh}(
@inbounds view(parent(data), :, $field_byterange, :)
))
end

@inline function Base.getindex(data::IFH{S}, I::CartesianIndex{5}) where {S}
@inbounds get_struct(
parent(data),
Expand Down Expand Up @@ -608,18 +567,6 @@ function DataF(x::T) where {T}
end
end

@generated function _property_view(
data::DataF{S, A},
::Val{Idx},
) where {S, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(DataF{$SS}(@inbounds view(parent(data), $field_byterange)))
end

Base.@propagate_inbounds function Base.getindex(data::DataF{S}) where {S}
@inbounds get_struct(
parent(data),
Expand Down Expand Up @@ -706,20 +653,6 @@ end

@inline universal_size(::IJF{S, Nij}) where {S, Nij} = (Nij, Nij, 1, 1, 1)

@generated function _property_view(
data::IJF{S, Nij, A},
::Val{Idx},
) where {S, Nij, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(IJF{$SS, $Nij}(
@inbounds view(parent(data), :, :, $field_byterange)
))
end

@inline function Base.getindex(
data::IJF{S, Nij},
I::CartesianIndex,
Expand Down Expand Up @@ -811,18 +744,6 @@ function SArray(data::IF{S, Ni, <:MArray}) where {S, Ni}
IF{S, Ni}(SArray(parent(data)))
end

@generated function _property_view(
data::IF{S, Ni, A},
::Val{Idx},
) where {S, Ni, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(IF{$SS, $Ni}(@inbounds view(parent(data), :, $field_byterange)))
end

@inline function Base.getindex(data::IF{S, Ni}, I::CartesianIndex) where {S, Ni}
i = I.I[1]
@boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
Expand Down Expand Up @@ -904,18 +825,6 @@ Base.lastindex(data::VF) = length(data)

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

@generated function _property_view(
data::VF{S, Nv, A},
::Val{Idx},
) where {S, Nv, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(VF{$SS, $Nv}(@inbounds view(parent(data), :, $field_byterange)))
end

Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) =
_property_view(data, Val(i))

Expand Down Expand Up @@ -1004,20 +913,6 @@ end

Base.length(data::VIJFH) = get_Nv(data) * get_Nh(data)

@generated function _property_view(
data::VIJFH{S, Nv, Nij, Nh, A},
::Val{Idx},
) where {S, Nv, Nij, Nh, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(VIJFH{$SS, $Nv, $Nij, $Nh}(
@inbounds view(parent(data), :, :, :, $field_byterange, :)
))
end

# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7)
@inline function slab(data::VIJFH{S, Nv, Nij, Nh}, v, h) where {S, Nv, Nij, Nh}
array = parent(data)
Expand Down Expand Up @@ -1141,20 +1036,6 @@ Base.copy(data::VIFH{S, Nv, Ni, Nh}) where {S, Nv, Ni, Nh} =

Base.length(data::VIFH) = nlevels(data) * get_Nh(data)

@generated function _property_view(
data::VIFH{S, Nv, Ni, Nh, A},
::Val{Idx},
) where {S, Nv, Ni, Nh, A, Idx}
SS = fieldtype(S, Idx)
T = eltype(A)
offset = fieldtypeoffset(T, S, Val(Idx))
nbytes = typesize(T, SS)
field_byterange = (offset + 1):(offset + nbytes)
return :(VIFH{$SS, $Nv, $Ni, $Nh}(
@inbounds view(parent(data), :, :, $field_byterange, :)
))
end

# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7)
@inline function slab(data::VIFH{S, Nv, Ni, Nh}, v, h) where {S, Nv, Ni, Nh}
array = parent(data)
Expand Down

0 comments on commit dc72795

Please sign in to comment.