diff --git a/.buildkite/Manifest.toml b/.buildkite/Manifest.toml index 649093b54c..2398f17082 100644 --- a/.buildkite/Manifest.toml +++ b/.buildkite/Manifest.toml @@ -373,7 +373,7 @@ weakdeps = ["CUDA", "MPI"] ClimaCommsMPIExt = "MPI" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "ClimaInterpolations", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"] +deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "ClimaInterpolations", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LLVM", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"] path = ".." uuid = "d414da3d-4745-48bb-8d80-42e94e092884" version = "0.14.50" diff --git a/Project.toml b/Project.toml index e465361a65..42254602c7 100644 --- a/Project.toml +++ b/Project.toml @@ -18,6 +18,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" +LLVM = "929cbde3-209d-540e-8aea-75f648917ca0" LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MultiBroadcastFusion = "c3c07f87-98de-43f2-a76f-835b330b2cbb" @@ -65,6 +66,7 @@ IntervalSets = "0.5, 0.6, 0.7" JET = "0.9" Krylov = "0.9, 0.10" KrylovKit = "0.6, 0.7, 0.8" +LLVM = "9" LazyBroadcast = "1" LinearAlgebra = "1" Logging = "1" diff --git a/docs/src/APIs/datalayouts_api.md b/docs/src/APIs/datalayouts_api.md index c81b3e9257..3c7e19b337 100644 --- a/docs/src/APIs/datalayouts_api.md +++ b/docs/src/APIs/datalayouts_api.md @@ -18,4 +18,14 @@ DataLayouts.IHF DataLayouts.IJHF DataLayouts.VIHF DataLayouts.VIJHF +DataLayouts.bitcast_struct +DataLayouts.default_basetype +DataLayouts.check_basetype +DataLayouts.checked_valid_basetype +DataLayouts.num_basetypes +DataLayouts.struct_field_view +DataLayouts.set_struct! +DataLayouts.get_struct +DataLayouts.parent_array_type +DataLayouts.promote_parent_array_type ``` diff --git a/docs/src/APIs/utilities_api.md b/docs/src/APIs/utilities_api.md index 211cea8c7b..ac1f0b91c3 100644 --- a/docs/src/APIs/utilities_api.md +++ b/docs/src/APIs/utilities_api.md @@ -7,6 +7,9 @@ CurrentModule = ClimaCore ```@docs Utilities.PlusHalf Utilities.half +Utilities.replace_type_parameter +Utilities.fieldtype_vals +Utilities.new ``` ## Utilities.Cache diff --git a/examples/column/tvd_advection.jl b/examples/column/tvd_advection.jl index 60a0688e8b..9fec523e1b 100644 --- a/examples/column/tvd_advection.jl +++ b/examples/column/tvd_advection.jl @@ -56,8 +56,6 @@ function tendency!(yₜ, y, parameters, t) method = limiter_method, ) - If = Operators.InterpolateC2F() - if limiter_method == "Zalesak" @. yₜ.q = -divf2c( @@ -68,8 +66,6 @@ function tendency!(yₜ, y, parameters, t) ), ) else - Δfluxₕ = @. w * If(y.q) - Δfluxₗ = @. upwind1(w, y.q) @. yₜ.q = -divf2c( upwind1(w, y.q) + diff --git a/ext/cuda/data_layouts.jl b/ext/cuda/data_layouts.jl index e55564227a..de54b23463 100644 --- a/ext/cuda/data_layouts.jl +++ b/ext/cuda/data_layouts.jl @@ -12,51 +12,27 @@ import ClimaCore.DataLayouts: fused_copyto! import Adapt import CUDA -parent_array_type(::Type{<:CUDA.CuArray{T, N, B} where {N}}) where {T, B} = - CUDA.CuArray{T, N, B} where {N} - -# allow on-device use of lazy broadcast objects +# Ensure that all CuArrays have the same memory buffer type. parent_array_type( - ::Type{<:CUDA.CuDeviceArray{T, N, A} where {N}}, -) where {T, A} = CUDA.CuDeviceArray{T, N, A} where {N} - -# Ensure that both parent array types have the same memory buffer type. -promote_parent_array_type( - ::Type{CUDA.CuArray{T1, N, B} where {N}}, - ::Type{CUDA.CuArray{T2, N, B} where {N}}, -) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), N, B} where {N} - -# allow on-device use of lazy broadcast objects + ::Type{<:CUDA.CuArray{<:Any, <:Any, B}}, + ::Type{T}, +) where {T, B} = CUDA.CuArray{T, <:Any, B} promote_parent_array_type( - ::Type{CUDA.CuDeviceArray{T1, N, B} where {N}}, - ::Type{CUDA.CuDeviceArray{T2, N, B} where {N}}, -) where {T1, T2, B} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N} + ::Type{CUDA.CuArray{T1, <:Any, B}}, + ::Type{CUDA.CuArray{T2, <:Any, B}}, +) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), <:Any, B} -# allow on-device use of lazy broadcast objects with different type params -promote_parent_array_type( - ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}}, - ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}}, -) where {T1, T2, B1, B2} = - CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B} - -# allow on-device use of lazy broadcast objects with different type params +# Allow on-device use of lazy broadcast objects. +parent_array_type(::Type{<:CUDA.CuDeviceArray}, ::Type{T}) where {T} = + CUDA.CuDeviceArray{T} promote_parent_array_type( ::Type{CUDA.CuDeviceArray{T1}}, - ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}}, -) where {T1, T2, B2} = - CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B} - -promote_parent_array_type( - ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}}, - ::Type{CUDA.CuDeviceArray{T2} where {N}}, -) where {T1, T2, B1} = - CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B} + ::Type{CUDA.CuDeviceArray{T2}}, +) where {T1, T2} = CUDA.CuDeviceArray{promote_type(T1, T2)} # Make `similar` accept our special `UnionAll` parent array type for CuArray. -Base.similar( - ::Type{CUDA.CuArray{T, N′, B} where {N′}}, - dims::Dims{N}, -) where {T, N, B} = similar(CUDA.CuArray{T, N, B}, dims) +Base.similar(::Type{CUDA.CuArray{T, <:Any, B}}, dims::Dims{N}) where {T, N, B} = + similar(CUDA.CuArray{T, N, B}, dims) unval(::Val{CI}) where {CI} = CI unval(CI) = CI diff --git a/ext/cuda/operators_sem_shmem.jl b/ext/cuda/operators_sem_shmem.jl index 6ad8b29a73..5a5a17ef03 100644 --- a/ext/cuda/operators_sem_shmem.jl +++ b/ext/cuda/operators_sem_shmem.jl @@ -92,7 +92,6 @@ Base.@propagate_inbounds function operator_shmem( Nq = Quadratures.degrees_of_freedom(QS) # allocate temp output RT = operator_return_eltype(op, eltype(arg)) - Nf = DataLayouts.typesize(FT, RT) WJv¹ = CUDA.CuStaticSharedArray(RT, (Nq, Nvt)) return (WJv¹,) end @@ -107,7 +106,6 @@ Base.@propagate_inbounds function operator_shmem( Nq = Quadratures.degrees_of_freedom(QS) # allocate temp output RT = operator_return_eltype(op, eltype(arg)) - Nf = DataLayouts.typesize(FT, RT) WJv¹ = CUDA.CuStaticSharedArray(RT, (Nq, Nq, Nvt)) WJv² = CUDA.CuStaticSharedArray(RT, (Nq, Nq, Nvt)) return (WJv¹, WJv²) diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index 9ac3296df6..42fdcc49d6 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -63,15 +63,16 @@ leverage efficient linear indexing: """ module DataLayouts -import Base: Base, @propagate_inbounds +import Base: Base, @propagate_inbounds, issingletontype import StaticArrays: SOneTo, MArray, SArray +import LLVM: unsafe_load 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, fieldtype_vals import ..DebugOnly: call_post_op_callback, post_op_callback import ..slab, ..slab_args, ..column, ..column_args, ..level, ..level_args export slab, @@ -130,7 +131,8 @@ struct IJHMask{B, NT, V} <: AbstractMask h_map::V end -include("struct.jl") +include("bitcast_struct.jl") +include("struct_storage.jl") abstract type AbstractData{S} end @@ -316,90 +318,45 @@ end Base.parent(data::AbstractData) = getfield(data, :array) -Base.similar(data::AbstractData{S}) where {S} = similar(data, S) - @inline Base.:(==)(data1::D, data2::D) where {D <: AbstractData} = parent(data1) == parent(data2) -@inline function ncomponents(data::AbstractData{S}) where {S} - typesize(eltype(parent(data)), S) +@inline ncomponents(data::AbstractData{S}) where {S} = + num_basetypes(eltype(parent(data)), S) + +@inline function field_index_view(data::AbstractData{S}, ::Val{F}) where {S, F} + 1 <= F <= fieldcount(S) || throw(ArgumentError("Type $S has no field $F")) + D = field_dim(singleton(data)) + params = Base.tail(type_params(data)) + array_view = @inbounds struct_field_view(parent(data), S, Val(F), Val(D)) + return union_all(singleton(data)){fieldtype(S, F), params...}(array_view) end -@inline _getproperty(data::AbstractData, ::Val{Name}) where {Name} = - _getproperty(data, Val(Name), Name) +@inline field_name_view(data::AbstractData{S}, ::Val{name}) where {S, name} = + name in fieldnames(S) ? + field_index_view(data, Val(unrolled_findfirst(==(name), fieldnames(S)))) : + throw(ArgumentError("Type $S has no field $name")) -@generated function _getproperty( - data::AbstractData{S}, - ::Val{Name}, - name, -) where {S, Name} - i = findfirst(isequal(Name), fieldnames(S)) - static_idx = Val{i}() - return :( - Base.@_inline_meta; DataLayouts._property_view(data, $static_idx, name) - ) -end +@inline Base.dotgetproperty(data::AbstractData, prop) = getproperty(data, prop) +@inline Base.getproperty(data::AbstractData, i::Integer) = + field_index_view(data, Val(i)) +@inline Base.getproperty(data::AbstractData, name::Symbol) = + field_name_view(data, Val(name)) -@inline function Base.getproperty(data::AbstractData{S}, name::Symbol) where {S} - _getproperty(data, Val{name}(), name) -end -@inline function Base.dotgetproperty( - data::AbstractData{S}, - name::Symbol, -) where {S} - _getproperty(data, Val{name}(), name) +function replace_storage(data::AbstractData, ::Type{S}, ::Type{T}) where {S, T} + D = field_dim(singleton(data)) + params = Base.tail(type_params(data)) + new_array_size = Base.setindex(size(parent(data)), num_basetypes(T, S), D) + new_array = similar(parent(data), T, new_array_size...) + return union_all(singleton(data)){S, params...}(new_array) end -Base.@propagate_inbounds function Base.getproperty( - data::AbstractData{S}, - i::Integer, -) where {S} - array = parent(data) - T = eltype(array) - SS = fieldtype(S, i) - offset = fieldtypeoffset(T, S, i) - nbytes = typesize(T, SS) - fdim = field_dim(singleton(data)) - Ipre = ntuple(i -> Colon(), Val(fdim - 1)) - Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim)) - dataview = - @inbounds view(array, Ipre..., (offset + 1):(offset + nbytes), Ipost...) - union_all(singleton(data)){SS, Base.tail(type_params(data))...}(dataview) -end - -@noinline _property_view( - data::AbstractData{S}, - ::Val{nothing}, - name, -) where {S} = error("Invalid field name `$name` for type `$(S)`.") - -# 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}, - name, -) 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.name.wrapper) - Ipre = ntuple(i -> Colon(), Val(fdim - 1)) - Ipost = ntuple(i -> Colon(), Val(ndims(data) - fdim)) - field_byterange = (offset + 1):(offset + nbytes) - return :($(AD.name.wrapper){$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) - return union_all(singleton(data)){S′, Base.tail(type_params(data))...}( - similar(array, T), - ) -end +replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T} = + replace_storage(data, replace_type_parameter(S, eltype(parent(data)), T), T) + +Base.similar(data::AbstractData{S}) where {S} = similar(data, S) +Base.similar(data::AbstractData, ::Type{S}) where {S} = + replace_storage(data, S, checked_valid_basetype(eltype(parent(data)), S)) maybe_populate!(array, ::typeof(similar)) = nothing maybe_populate!(array, ::typeof(ones)) = fill!(array, 1) @@ -466,7 +423,7 @@ function IJKFVH{S, Nij, Nk, Nv}( @assert size(array, 1) == Nij @assert size(array, 2) == Nij @assert size(array, 3) == Nk - @assert size(array, 4) == typesize(T, S) + @assert size(array, 4) == num_basetypes(T, S) @assert size(array, 5) == Nv IJKFVH{S, Nij, Nk, Nv, typeof(array)}(array) end @@ -479,7 +436,7 @@ function IJKFVH{S}( Nk::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nij, Nij, Nk, Nf, Nv, Nh) maybe_populate!(array, fun) IJKFVH{S, Nij, Nk, Nv}(array) @@ -530,7 +487,7 @@ function IJFH{S, Nij}(array::AbstractArray{T, 4}) where {S, Nij, T} check_basetype(T, S) @assert size(array, 1) == Nij @assert size(array, 2) == Nij - @assert size(array, 3) == typesize(T, S) + @assert size(array, 3) == num_basetypes(T, S) IJFH{S, Nij, typeof(array)}(array) end @@ -540,7 +497,7 @@ function IJFH{S}( Nij::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nij, Nij, Nf, Nh) maybe_populate!(array, fun) IJFH{S, Nij}(array) @@ -551,7 +508,7 @@ end function IJFH{S, Nij}(::Type{ArrayType}, Nh::Integer) where {S, Nij, ArrayType} T = eltype(ArrayType) - IJFH{S, Nij}(ArrayType(undef, Nij, Nij, typesize(T, S), Nh)) + IJFH{S, Nij}(ArrayType(undef, Nij, Nij, num_basetypes(T, S), Nh)) end Base.length(data::IJFH) = get_Nh_dynamic(data) @@ -623,7 +580,7 @@ function IJHF{S, Nij}(array::AbstractArray{T, 4}) where {S, Nij, T} check_basetype(T, S) @assert size(array, 1) == Nij @assert size(array, 2) == Nij - @assert size(array, 4) == typesize(T, S) + @assert size(array, 4) == num_basetypes(T, S) IJHF{S, Nij, typeof(array)}(array) end @@ -633,7 +590,7 @@ function IJHF{S}( Nij::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nij, Nij, Nh, Nf) maybe_populate!(array, fun) IJHF{S, Nij}(array) @@ -644,7 +601,7 @@ end function IJHF{S, Nij}(::Type{ArrayType}, Nh::Integer) where {S, Nij, ArrayType} T = eltype(ArrayType) - IJHF{S, Nij}(ArrayType(undef, Nij, Nij, Nh, typesize(T, S))) + IJHF{S, Nij}(ArrayType(undef, Nij, Nij, Nh, num_basetypes(T, S))) end Base.length(data::IJHF) = get_Nh_dynamic(data) @@ -722,7 +679,7 @@ end function IFH{S, Ni}(array::AbstractArray{T, 3}) where {S, Ni, T} check_basetype(T, S) @assert size(array, 1) == Ni - @assert size(array, 2) == typesize(T, S) + @assert size(array, 2) == num_basetypes(T, S) IFH{S, Ni, typeof(array)}(array) end @@ -732,7 +689,7 @@ function IFH{S}( Ni::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Ni, Nf, Nh) maybe_populate!(array, fun) IFH{S, Ni}(array) @@ -740,7 +697,7 @@ end function IFH{S, Ni}(::Type{ArrayType}, Nh::Integer) where {S, Ni, ArrayType} T = eltype(ArrayType) - IFH{S, Ni}(ArrayType(undef, Ni, typesize(T, S), Nh)) + IFH{S, Ni}(ArrayType(undef, Ni, num_basetypes(T, S), Nh)) end @inline universal_size(data::IFH{S, Ni}) where {S, Ni} = @@ -801,7 +758,7 @@ end function IHF{S, Ni}(array::AbstractArray{T, 3}) where {S, Ni, T} check_basetype(T, S) @assert size(array, 1) == Ni - @assert size(array, 3) == typesize(T, S) + @assert size(array, 3) == num_basetypes(T, S) IHF{S, Ni, typeof(array)}(array) end @@ -811,7 +768,7 @@ function IHF{S}( Ni::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Ni, Nh, Nf) maybe_populate!(array, fun) IHF{S, Ni}(array) @@ -819,7 +776,7 @@ end function IHF{S, Ni}(::Type{ArrayType}, Nh::Integer) where {S, Ni, ArrayType} T = eltype(ArrayType) - IHF{S, Ni}(ArrayType(undef, Ni, Nh, typesize(T, S))) + IHF{S, Ni}(ArrayType(undef, Ni, Nh, num_basetypes(T, S))) end @inline universal_size(data::IHF{S, Ni}) where {S, Ni} = @@ -866,60 +823,31 @@ end function DataF{S}(array::AbstractVector{T}) where {S, T} check_basetype(T, S) - @assert size(array, 1) == typesize(T, S) + @assert size(array, 1) == num_basetypes(T, S) DataF{S, typeof(array)}(array) end function DataF{S}(::Type{ArrayType}, fun = similar;) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nf) maybe_populate!(array, fun) DataF{S}(array) 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} - @inbounds get_struct( - parent(data), - S, - Val(field_dim(singleton(data))), - CartesianIndex(1), - ) -end - -@propagate_inbounds function Base.getindex(col::Data0D, I::CartesianIndex{5}) - @inbounds col[] -end +@propagate_inbounds Base.getindex(data::DataF) = + get_struct(parent(data), eltype(data)) +@propagate_inbounds Base.getindex(data::DataF, I::CartesianIndex{5}) = data[] -Base.@propagate_inbounds function Base.setindex!(data::DataF{S}, val) where {S} - @inbounds set_struct!( - parent(data), - convert(S, val), - Val(field_dim(singleton(data))), - CartesianIndex(1), - ) -end - -@propagate_inbounds function Base.setindex!( - col::Data0D, - val, - I::CartesianIndex{5}, -) - @inbounds col[] = val -end +@propagate_inbounds Base.setindex!(data::DataF{S}, val) where {S} = + set_struct!(parent(data), convert(eltype(data), val)) +@propagate_inbounds Base.setindex!(data::DataF, val, I::CartesianIndex{5}) = + data[] = val # ====================== # DataSlab2D DataLayout @@ -968,7 +896,7 @@ function IJF{S, Nij}(array::AbstractArray{T, 3}) where {S, Nij, T} @assert size(array, 1) == Nij @assert size(array, 2) == Nij check_basetype(T, S) - @assert size(array, 3) == typesize(T, S) + @assert size(array, 3) == num_basetypes(T, S) IJF{S, Nij, typeof(array)}(array) end @@ -977,14 +905,14 @@ function IJF{S}( fun = similar; Nij::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nij, Nij, Nf) maybe_populate!(array, fun) IJF{S, Nij}(array) end function IJF{S, Nij}(::Type{MArray}, ::Type{T}) where {S, Nij, T} - Nf = typesize(T, S) + Nf = num_basetypes(T, S) array = MArray{Tuple{Nij, Nij, Nf}, T, 3, Nij * Nij * Nf}(undef) IJF{S, Nij}(array) end @@ -1047,7 +975,7 @@ end function IF{S, Ni}(array::AbstractArray{T, 2}) where {S, Ni, T} @assert size(array, 1) == Ni check_basetype(T, S) - @assert size(array, 2) == typesize(T, S) + @assert size(array, 2) == num_basetypes(T, S) IF{S, Ni, typeof(array)}(array) end @@ -1056,14 +984,14 @@ function IF{S}( fun = similar; Ni::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Ni, Nf) maybe_populate!(array, fun) IF{S, Ni}(array) end function IF{S, Ni}(::Type{MArray}, ::Type{T}) where {S, Ni, T} - Nf = typesize(T, S) + Nf = num_basetypes(T, S) array = MArray{Tuple{Ni, Nf}, T, 2, Ni * Nf}(undef) IF{S, Ni}(array) end @@ -1115,7 +1043,7 @@ end function VF{S, Nv}(array::AbstractArray{T, 2}) where {S, Nv, T} check_basetype(T, S) @assert size(array, 1) == Nv - @assert size(array, 2) == typesize(T, S) + @assert size(array, 2) == num_basetypes(T, S) VF{S, Nv, typeof(array)}(array) end @@ -1124,7 +1052,7 @@ function VF{S}( fun = similar; Nv::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nv, Nf) maybe_populate!(array, fun) VF{S, Nv}(array) @@ -1132,23 +1060,20 @@ end function VF{S, Nv}(array::AbstractVector{T}) where {S, Nv, T} check_basetype(T, S) - @assert typesize(T, S) == 1 + @assert num_basetypes(T, S) == 1 VF{S, Nv}(reshape(array, (:, 1))) end function VF{S, Nv}(::Type{ArrayType}, nelements) where {S, Nv, ArrayType} T = eltype(ArrayType) check_basetype(T, S) - VF{S, Nv}(ArrayType(undef, nelements, typesize(T, S))) + VF{S, Nv}(ArrayType(undef, nelements, num_basetypes(T, S))) end Base.lastindex(data::VF) = length(data) nlevels(::VF{S, Nv}) where {S, Nv} = Nv -Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) = - _property_view(data, Val(i), i) - Base.@propagate_inbounds column(data::VF, i, h) = column(data, i, 1, h) @inline function column(data::VF, i, j, h) @@ -1201,7 +1126,7 @@ function VIJFH{S, Nv, Nij}(array::AbstractArray{T, 5}) where {S, Nv, Nij, T} check_basetype(T, S) @assert size(array, 1) == Nv @assert size(array, 2) == size(array, 3) == Nij - @assert size(array, 4) == typesize(T, S) + @assert size(array, 4) == num_basetypes(T, S) VIJFH{S, Nv, Nij, typeof(array)}(array) end @@ -1212,7 +1137,7 @@ function VIJFH{S}( Nij::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nv, Nij, Nij, Nf, Nh) maybe_populate!(array, fun) VIJFH{S, Nv, Nij, typeof(array)}(array) @@ -1308,7 +1233,7 @@ function VIJHF{S, Nv, Nij}(array::AbstractArray{T, 5}) where {S, Nv, Nij, T} check_basetype(T, S) @assert size(array, 1) == Nv @assert size(array, 2) == size(array, 3) == Nij - @assert size(array, 5) == typesize(T, S) + @assert size(array, 5) == num_basetypes(T, S) VIJHF{S, Nv, Nij, typeof(array)}(array) end @@ -1319,7 +1244,7 @@ function VIJHF{S}( Nij::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nv, Nij, Nij, Nh, Nf) maybe_populate!(array, fun) VIJHF{S, Nv, Nij, typeof(array)}(array) @@ -1419,7 +1344,7 @@ function VIFH{S, Nv, Ni}(array::AbstractArray{T, 4}) where {S, Nv, Ni, T} check_basetype(T, S) @assert size(array, 1) == Nv @assert size(array, 2) == Ni - @assert size(array, 3) == typesize(T, S) + @assert size(array, 3) == num_basetypes(T, S) VIFH{S, Nv, Ni, typeof(array)}(array) end @@ -1430,7 +1355,7 @@ function VIFH{S}( Ni::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nv, Ni, Nf, Nh) maybe_populate!(array, fun) VIFH{S, Nv, Ni, typeof(array)}(array) @@ -1511,7 +1436,7 @@ function VIHF{S, Nv, Ni}(array::AbstractArray{T, 4}) where {S, Nv, Ni, T} check_basetype(T, S) @assert size(array, 1) == Nv @assert size(array, 2) == Ni - @assert size(array, 4) == typesize(T, S) + @assert size(array, 4) == num_basetypes(T, S) VIHF{S, Nv, Ni, typeof(array)}(array) end @@ -1522,7 +1447,7 @@ function VIHF{S}( Ni::Integer, Nh::Integer, ) where {S, ArrayType} - Nf = typesize(eltype(ArrayType), S) + Nf = num_basetypes(eltype(ArrayType), S) array = similar(ArrayType, Nv, Ni, Nh, Nf) maybe_populate!(array, fun) VIHF{S, Nv, Ni, typeof(array)}(array) @@ -1781,17 +1706,18 @@ type parameters. @inline h_dim(::VIFHSingleton) = 4 @inline h_dim(::VIHFSingleton) = 3 -@inline to_data_specific(::VFSingleton, I::Tuple) = (I[4], 1) -@inline to_data_specific(::IFSingleton, I::Tuple) = (I[1], 1) -@inline to_data_specific(::IJFSingleton, I::Tuple) = (I[1], I[2], 1) -@inline to_data_specific(::IJFHSingleton, I::Tuple) = (I[1], I[2], 1, I[5]) -@inline to_data_specific(::IJHFSingleton, I::Tuple) = (I[1], I[2], I[5], 1) -@inline to_data_specific(::IFHSingleton, I::Tuple) = (I[1], 1, I[5]) -@inline to_data_specific(::IHFSingleton, I::Tuple) = (I[1], I[5], 1) -@inline to_data_specific(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], 1, I[5]) -@inline to_data_specific(::VIJHFSingleton, I::Tuple) = (I[4], I[1], I[2], I[5], 1) -@inline to_data_specific(::VIFHSingleton, I::Tuple) = (I[4], I[1], 1, I[5]) -@inline to_data_specific(::VIHFSingleton, I::Tuple) = (I[4], I[1], I[5], 1) +@inline data_specific_index(::DataF, I) = CartesianIndex() +@inline data_specific_index(::VF, I) = CartesianIndex(I[4]) +@inline data_specific_index(::IF, I) = CartesianIndex(I[1]) +@inline data_specific_index(::IJF, I) = CartesianIndex(I[1], I[2]) +@inline data_specific_index(::IJFH, I) = CartesianIndex(I[1], I[2], I[5]) +@inline data_specific_index(::IJHF, I) = CartesianIndex(I[1], I[2], I[5]) +@inline data_specific_index(::IFH, I) = CartesianIndex(I[1], I[5]) +@inline data_specific_index(::IHF, I) = CartesianIndex(I[1], I[5]) +@inline data_specific_index(::VIJFH, I) = CartesianIndex(I[4], I[1], I[2], I[5]) +@inline data_specific_index(::VIJHF, I) = CartesianIndex(I[4], I[1], I[2], I[5]) +@inline data_specific_index(::VIFH, I) = CartesianIndex(I[4], I[1], I[5]) +@inline data_specific_index(::VIHF, I) = CartesianIndex(I[4], I[1], I[5]) @inline to_data_specific_field(::DataFSingleton, I::Tuple) = (I[3],) @inline to_data_specific_field(::VFSingleton, I::Tuple) = (I[4], I[3]) @@ -1970,34 +1896,34 @@ Base.ndims(data::AbstractData) = Base.ndims(typeof(data)) Base.ndims(::Type{T}) where {T <: AbstractData} = Base.ndims(parent_array_type(T)) -@inline function Base.getindex( - data::Union{IJF, IJFH, IJHF, IFH, IHF, VIJFH, VIJHF, VIFH, VIHF, VF, IF}, - I::CartesianIndex, -) - @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) - s = singleton(data) - return get_struct( +@propagate_inbounds Base.getindex(data::AbstractData, I::CartesianIndex) = + get_struct( parent(data), eltype(data), - Val(field_dim(s)), - CartesianIndex(to_data_specific(s, I.I)), + data_specific_index(data, I), + Val(field_dim(singleton(data))), ) -end +@propagate_inbounds Base.getindex(data::AbstractData, i::Integer) = + get_struct(parent(data), eltype(data), i, Val(field_dim(singleton(data)))) +@propagate_inbounds Base.getindex(data::AbstractData, I::Integer...) = + getindex(data, CartesianIndex(I)) -@inline function Base.setindex!( - data::Union{IJF, IJFH, IJHF, IFH, IHF, VIJFH, VIJHF, VIFH, VIHF, VF, IF}, - val, - I::CartesianIndex, -) - @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) - s = singleton(data) - @inbounds set_struct!( +@propagate_inbounds Base.setindex!(data::AbstractData, val, I::CartesianIndex) = + set_struct!( parent(data), convert(eltype(data), val), - Val(field_dim(s)), - CartesianIndex(to_data_specific(s, I.I)), + data_specific_index(data, I), + Val(field_dim(singleton(data))), ) -end +@propagate_inbounds Base.setindex!(data::AbstractData, val, i::Integer) = + set_struct!( + parent(data), + convert(eltype(data), val), + i, + Val(field_dim(singleton(data))), + ) +@propagate_inbounds Base.setindex!(data::AbstractData, val, I::Integer...) = + setindex!(data, val, CartesianIndex(I)) """ CartesianFieldIndex{N} <: Base.AbstractCartesianIndex{N} @@ -2093,90 +2019,6 @@ end const EndsWithField{S} = Union{IJHF{S}, IHF{S}, IJF{S}, IF{S}, VF{S}, VIJHF{S}, VIHF{S}} -if VERSION ≥ v"1.11.0-beta" - ### --------------- Support for multi-dimensional indexing - # TODO: can we remove this? It's not needed for Julia 1.10, - # but seems needed in Julia 1.11. - @inline Base.getindex( - data::Union{ - IJF, - IJFH, - IJHF, - IFH, - IHF, - VIJFH, - VIJHF, - VIFH, - VIHF, - VF, - IF, - }, - I::Vararg{Int, N}, - ) where {N} = Base.getindex( - data, - CartesianIndex(to_universal_index(singleton(data), I)), - ) - - @inline Base.setindex!( - data::Union{ - IJF, - IJFH, - IJHF, - IFH, - IHF, - VIJFH, - VIJHF, - VIFH, - VIHF, - VF, - IF, - }, - val, - I::Vararg{Int, N}, - ) where {N} = Base.setindex!( - data, - val, - CartesianIndex(to_universal_index(singleton(data), I)), - ) - - # Certain datalayouts support special indexing. - # Like VF datalayouts with `getindex(::VF, v::Integer)` - #! format: off - @inline to_universal_index(::VFSingleton, I::NTuple{1, T}) where {T} = (T(1), T(1), T(1), I[1], T(1)) - @inline to_universal_index(::IFSingleton, I::NTuple{1, T}) where {T} = (I[1], T(1), T(1), T(1), T(1)) - @inline to_universal_index(::IFSingleton, I::NTuple{2, T}) where {T} = (I[1], T(1), T(1), T(1), T(1)) - @inline to_universal_index(::IFSingleton, I::NTuple{3, T}) where {T} = (I[1], T(1), T(1), T(1), T(1)) - @inline to_universal_index(::IFSingleton, I::NTuple{4, T}) where {T} = (I[1], T(1), T(1), T(1), T(1)) - @inline to_universal_index(::IFSingleton, I::NTuple{5, T}) where {T} = (I[1], T(1), T(1), T(1), T(1)) - @inline to_universal_index(::IJFSingleton, I::NTuple{2, T}) where {T} = (I[1], I[2], T(1), T(1), T(1)) - @inline to_universal_index(::IJFSingleton, I::NTuple{3, T}) where {T} = (I[1], I[2], T(1), T(1), T(1)) - @inline to_universal_index(::IJFSingleton, I::NTuple{4, T}) where {T} = (I[1], I[2], T(1), T(1), T(1)) - @inline to_universal_index(::IJFSingleton, I::NTuple{5, T}) where {T} = (I[1], I[2], T(1), T(1), T(1)) - @inline to_universal_index(::AbstractDataSingleton, I::NTuple{5}) = I - #! format: on - ### --------------- -else - # Only support datalayouts that end with fields, since those - # are the only layouts where we can efficiently compute the - # strides. - @propagate_inbounds function Base.getindex( - data::EndsWithField{S}, - I::Integer, - ) where {S} - s_array = farray_size(data) - @inbounds get_struct_linear(parent(data), S, I, s_array) - end - @propagate_inbounds function Base.setindex!( - data::EndsWithField{S}, - val, - I::Integer, - ) where {S} - s_array = farray_size(data) - @inbounds set_struct_linear!(parent(data), convert(S, val), I, s_array) - end - -end - """ data2array(::AbstractData) @@ -2270,8 +2112,6 @@ include("fused_copyto.jl") include("fill.jl") include("mapreduce.jl") -include("struct_linear_indexing.jl") - """ set_mask_maps!(mask) diff --git a/src/DataLayouts/bitcast_struct.jl b/src/DataLayouts/bitcast_struct.jl new file mode 100644 index 0000000000..61cf5190b5 --- /dev/null +++ b/src/DataLayouts/bitcast_struct.jl @@ -0,0 +1,157 @@ +# Treat array elements as if they are fields of a single tuple value +field_expr(f, value_expr) = :(Core.getfield($value_expr, $f)) +field_expr(f, array_expr, index_expr) = + :(@inbounds $array_expr[struct_index($f, $array_expr, $index_expr...)]) + +# Keep array element read instructions separate unless the full value is needed +full_value_expr(@nospecialize(S), value_expr) = value_expr +full_value_expr(@nospecialize(S), inputs...) = + Expr(:tuple, (field_expr(f, inputs...) for f in 1:fieldcount(S))...) + +# Manually inline all read instructions into a generated function body, instead +# of relying on Julia's compiler to inline them during the code lowering stage +bitcast_struct_expr(@nospecialize(T), @nospecialize(S), inputs...) = + if T === S + full_value_expr(S, inputs...) + elseif sizeof(T) != sizeof(S) + error("Cannot bitcast $S ($(sizeof(S)) bytes) to $T ($(sizeof(T)) bytes)") + elseif issingletontype(T) + T.instance + elseif iszero(fieldcount(T)) && iszero(fieldcount(S)) && T !== Bool + # Use Core.bitcast to convert primitive types, but not to turn non-Bools + # into Bools, since that zeros out their first seven bits; for example, + # Core.bitcast(Int8, Core.bitcast(Bool, Int8(3))) returns 1 instead of 3 + :(Core.bitcast($T, $(full_value_expr(S, inputs...)))) + elseif isone(count(!issingletontype, fieldtypes(S))) + # Ignore the singleton fields of S + f = findfirst(!issingletontype, fieldtypes(S)) + bitcast_struct_expr(T, fieldtype(S, f), field_expr(f, inputs...)) + elseif isone(count(!issingletontype, fieldtypes(T))) + # Use instances to get the singleton fields of T + T_expr = Expr(:new, T) + for field_type in fieldtypes(T) + field_value_expr = + issingletontype(field_type) ? field_type.instance : + bitcast_struct_expr(field_type, S, inputs...) + push!(T_expr.args, field_value_expr) + end + T_expr + else + # Use unsafe_load to convert composite types, and to losslessly turn + # non-Bools into Bools; implemented like getindex(v::MArray, i) from + # https://github.com/JuliaArrays/StaticArrays.jl/blob/v1.0.0/src/MArray.jl#L85, + # but with an LLVMPtr instead of a Ptr so the LLVM compiler can apply + # SROA/DCE (scalar replacement of aggregates and dead code elimination) + isbitstype(T) || error("Cannot allocate $T in stack memory for bitcast") + isbitstype(S) || error("Cannot allocate $S in stack memory for bitcast") + quote + stack_memory = Ref($(full_value_expr(S, inputs...))) + pointer = Core.LLVMPtr{$T, 0}(pointer_from_objref(stack_memory)) + GC.@preserve stack_memory unsafe_load(pointer) + end + end + +""" + bitcast_struct(T, value) + bitcast_struct(T, array, Val(num_indices), index...) + +Converts `value` into an `isbits` type `T` that spans the same number of bytes +(counting all bytes that are used as padding; see extended help for details). +Serves as a GPU-compatible generalization of the native `Core.bitcast` function, +losslessly converting between arbitrary data types, including composite types. + +Instead of converting a single value, it is also possible to convert a subset of +an array corresponding to the result of [`get_struct`](@ref). This is equivalent +to converting the array elements after first loading them into a tuple, but with +guaranteed inlining for arbitrary data types. Inlining is necessary for the +compiler's [`getfield_elim_pass!`](https://hackmd.io/bZz8k6SHQQuNUW-Vs7rqfw) to +eliminate reads of array elements for unused fields of `T` (a key optimization +in GPU kernels, where reads from global memory can be relatively expensive). + +# Examples + +```julia-repl +julia> bitcast_struct(NTuple{4, Int8}, Int32(1)) +(1, 0, 0, 0) + +julia> bitcast_struct(NTuple{6, Int32}, (2 * eps(0.0), eps(0.0), 0.0)) +(2, 0, 1, 0, 0, 0) + +julia> bitcast_struct(Tuple{Int32, Int32, Int128}, (2, 0, 1, 0)) +(2, 0, 1) +``` + +# Extended help + +The output of `bitcast_struct(T, value)` is similar to the output of +`reinterpret(T, value)`, with both functions interpreting sequential bytes in +[little-endian order](https://en.wikipedia.org/wiki/Endianness): + +```julia-repl +julia> reinterpret(NTuple{4, Int8}, Int32(1)) +(1, 0, 0, 0) + +julia> reinterpret(NTuple{6, Int32}, (2 * eps(0.0), eps(0.0), 0.0)) +(2, 0, 1, 0, 0, 0) + +julia> reinterpret(Tuple{Int32, Int32, Int128}, (2, 1, 0)) +(2, 0, 1) +``` + +As the last example shows, `bitcast_struct` and `reinterpret` can behave +differently when converting between data structures with nonuniform field sizes. +Specifically, they differ for data structures that are stored with +[padding](https://www.gnu.org/software/c-intro-and-ref/manual/html_node/Structure-Layout.html), +which the C code underlying Julia uses to ensure that fields are efficiently +aligned in stack memory. + +Unlike `reinterpret(T, value)`, which avoids mixing padding with non-padding +(it recursively traverses fields of `value` and `T`, introducing offsets when +their padding bytes are in different positions), `bitcast_struct(T, value)` +makes no distinction between padding and non-padding. Although `reinterpret` is +therefore less likely to produce unexpected outputs, it also performs runtime +allocations in heap memory, making it unsuitable for GPU kernels that do not +support such allocations. In contrast, `bitcast_struct` has a much simpler +implementation, with all of its allocations confined to stack memory. Moreover, +as long as `bitcast_struct` is only called within [`set_struct!`](@ref) and +[`get_struct`](@ref), potentially unexpected outputs will be hidden from users. + +In addition to the low-level method of `reinterpret` for `isbits` inputs, there +is another method for `AbstractArray` inputs that behaves exactly like +`bitcast_struct` when it comes to padding: + +```julia-repl +julia> reinterpret(reshape, NTuple{4, Int8}, Int32[1])[1] +(1, 0, 0, 0) + +julia> reinterpret(reshape, NTuple{6, Int32}, [2 * eps(0.0), eps(0.0), 0.0])[1] +(2, 0, 1, 0, 0, 0) + +julia> reinterpret(reshape, Tuple{Int32, Int32, Int128}, [2, 0, 1, 0])[1] +(2, 0, 1) +``` + +This method of `reinterpret` reads bytes from heap memory without distinguishing +padding and non-padding, in the same way as `bitcast_struct` reads bytes from +stack memory. So, while the method of `reinterpret` for `isbits` inputs can +construct the nonuniform type `Tuple{Int32, Int32, Int128}` from three `Int64`s, +`bitcast_struct` and the method for arrays both require a fourth `Int64`, +spanning the eight padding bytes inserted between the `Int32`s and the `Int128`. + +For more information about `reinterpret` and padding, see the following: +- https://discourse.julialang.org/t/reinterpret-returns-wrong-values +- https://discourse.julialang.org/t/reinterpret-vector-into-single-struct +- https://discourse.julialang.org/t/reinterpret-vector-of-mixed-type-tuples +""" +@generated bitcast_struct(::Type{T}, value::S) where {T, S} = + Expr(:block, :@inline, bitcast_struct_expr(T, S, :value)) + +@generated function bitcast_struct( + ::Type{T}, + array, + ::Val{num_indices}, + index..., +) where {T, num_indices} + S = NTuple{num_indices, eltype(array)} + return Expr(:block, :@inline, bitcast_struct_expr(T, S, :array, :index)) +end diff --git a/src/DataLayouts/broadcast.jl b/src/DataLayouts/broadcast.jl index 600bffa7d5..a43f205148 100644 --- a/src/DataLayouts/broadcast.jl +++ b/src/DataLayouts/broadcast.jl @@ -439,8 +439,8 @@ function Base.similar( bc::BroadcastedUnionDataF{<:Any, A}, ::Type{Eltype}, ) where {A, Eltype} - PA = parent_array_type(A) - array = similar(PA, (typesize(eltype(A), Eltype))) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (num_basetypes(eltype(PA), Eltype))) return DataF{Eltype}(array) end @@ -449,8 +449,8 @@ function Base.similar( ::Type{Eltype}, (_, _, _, _, Nh) = size(bc), ) where {Nij, A, Eltype} - PA = parent_array_type(A) - array = similar(PA, (Nij, Nij, typesize(eltype(A), Eltype), Nh)) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Nij, Nij, num_basetypes(eltype(PA), Eltype), Nh)) return IJFH{Eltype, Nij}(array) end @@ -459,8 +459,8 @@ function Base.similar( ::Type{Eltype}, (_, _, _, _, Nh) = size(bc), ) where {Nij, A, Eltype} - PA = parent_array_type(A) - array = similar(PA, (Nij, Nij, Nh, typesize(eltype(A), Eltype))) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Nij, Nij, Nh, num_basetypes(eltype(PA), Eltype))) return IJHF{Eltype, Nij}(array) end @@ -469,8 +469,8 @@ function Base.similar( ::Type{Eltype}, (_, _, _, _, Nh) = size(bc), ) where {Ni, A, Eltype} - PA = parent_array_type(A) - array = similar(PA, (Ni, typesize(eltype(A), Eltype), Nh)) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Ni, num_basetypes(eltype(PA), Eltype), Nh)) return IFH{Eltype, Ni}(array) end @@ -479,8 +479,8 @@ function Base.similar( ::Type{Eltype}, (_, _, _, _, Nh) = size(bc), ) where {Ni, A, Eltype} - PA = parent_array_type(A) - array = similar(PA, (Ni, Nh, typesize(eltype(A), Eltype))) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Ni, Nh, num_basetypes(eltype(PA), Eltype))) return IHF{Eltype, Ni}(array) end @@ -488,8 +488,9 @@ function Base.similar( ::BroadcastedUnionIJF{<:Any, Nij, A}, ::Type{Eltype}, ) where {Nij, A, Eltype} - Nf = typesize(eltype(A), Eltype) - array = MArray{Tuple{Nij, Nij, Nf}, eltype(A), 3, Nij * Nij * Nf}(undef) + T = checked_valid_basetype(eltype(A), Eltype) + Nf = num_basetypes(T, Eltype) + array = MArray{Tuple{Nij, Nij, Nf}, T, 3, Nij * Nij * Nf}(undef) return IJF{Eltype, Nij}(array) end @@ -497,8 +498,9 @@ function Base.similar( ::BroadcastedUnionIF{<:Any, Ni, A}, ::Type{Eltype}, ) where {Ni, A, Eltype} - Nf = typesize(eltype(A), Eltype) - array = MArray{Tuple{Ni, Nf}, eltype(A), 2, Ni * Nf}(undef) + T = checked_valid_basetype(eltype(A), Eltype) + Nf = num_basetypes(T, Eltype) + array = MArray{Tuple{Ni, Nf}, T, 2, Ni * Nf}(undef) return IF{Eltype, Ni}(array) end @@ -512,8 +514,8 @@ function Base.similar( ::Type{Eltype}, ::Val{newNv}, ) where {Nv, A, Eltype, newNv} - PA = parent_array_type(A) - array = similar(PA, (newNv, typesize(eltype(A), Eltype))) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (newNv, num_basetypes(eltype(PA), Eltype))) return VF{Eltype, newNv}(array) end @@ -528,8 +530,8 @@ function Base.similar( ::Val{newNv}, ) where {Nv, Ni, A, Eltype, newNv} (_, _, _, _, Nh) = size(bc) - PA = parent_array_type(A) - array = similar(PA, (newNv, Ni, typesize(eltype(A), Eltype), Nh)) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (newNv, Ni, num_basetypes(eltype(PA), Eltype), Nh)) return VIFH{Eltype, newNv, Ni}(array) end @@ -539,8 +541,8 @@ function Base.similar( ::Val{newNv}, ) where {Nv, Ni, A, Eltype, newNv} (_, _, _, _, Nh) = size(bc) - PA = parent_array_type(A) - array = similar(PA, (newNv, Ni, Nh, typesize(eltype(A), Eltype))) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (newNv, Ni, Nh, num_basetypes(eltype(PA), Eltype))) return VIHF{Eltype, newNv, Ni}(array) end @@ -555,25 +557,25 @@ Base.similar( ) where {Nv, Nij, A, Eltype} = similar(bc, Eltype, Val(Nv)) function Base.similar( - bc::BroadcastedUnionVIJFH{<:Any, Nv, Nij, A}, + bc::BroadcastedUnionVIJFH{<:Any, <:Any, Nij, A}, ::Type{Eltype}, - ::Val{newNv}, -) where {Nv, Nij, A, Eltype, newNv} + ::Val{Nv}, +) where {Nij, A, Eltype, Nv} (_, _, _, _, Nh) = size(bc) - PA = parent_array_type(A) - array = similar(PA, (newNv, Nij, Nij, typesize(eltype(A), Eltype), Nh)) - return VIJFH{Eltype, newNv, Nij}(array) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Nv, Nij, Nij, num_basetypes(eltype(PA), Eltype), Nh)) + return VIJFH{Eltype, Nv, Nij}(array) end function Base.similar( - bc::BroadcastedUnionVIJHF{<:Any, Nv, Nij, A}, + bc::BroadcastedUnionVIJHF{<:Any, <:Any, Nij, A}, ::Type{Eltype}, - ::Val{newNv}, -) where {Nv, Nij, A, Eltype, newNv} + ::Val{Nv}, +) where {Nij, A, Eltype, Nv} (_, _, _, _, Nh) = size(bc) - PA = parent_array_type(A) - array = similar(PA, (newNv, Nij, Nij, Nh, typesize(eltype(A), Eltype))) - return VIJHF{Eltype, newNv, Nij}(array) + PA = parent_array_type(A, checked_valid_basetype(eltype(A), Eltype)) + array = similar(PA, (Nv, Nij, Nij, Nh, num_basetypes(eltype(PA), Eltype))) + return VIJHF{Eltype, Nv, Nij}(array) end # ============= FusedMultiBroadcast diff --git a/src/DataLayouts/struct.jl b/src/DataLayouts/struct.jl deleted file mode 100644 index b58ae9819e..0000000000 --- a/src/DataLayouts/struct.jl +++ /dev/null @@ -1,276 +0,0 @@ -# 1) given a Vector{T}, T<:Number, construct an object of type S using the same underlying data -# e.g. -# array = [1.0,2.0,3.0], S = Tuple{Complex{Float64},Float64} => (1.0 + 2.0im, 3.0) - -# 2) for an Array{T,N}, and some 1 <= D <= N, -# reconstruct S from a slice of array# - -# The main difference between StructArrays and what we want is that we want 1 underlying (rather than # fields) -# there is no heterogeneity in the struct (can store recursive representations) of field types - -# get_offset(array, S, offset) => S(...), next_offset - -""" - is_valid_basetype(::Type{T}, ::Type{S}) - -Determines whether an object of type `S` can be stored in an array with elements -of type `T` by recursively checking whether the non-empty fields of `S` can be -stored in such an array. If `S` is empty, this is always true. -""" -function is_valid_basetype(::Type{T}, ::Type{S}) where {T, S} - sizeof(S) == 0 || - fieldcount(S) > 0 && - unrolled_all(s -> is_valid_basetype(T, s), fieldtypes(S)) -end -is_valid_basetype(::Type{T}, ::Type{<:T}) where {T} = true - -""" - check_basetype(::Type{T}, ::Type{S}) - -Check whether the types `T` and `S` have well-defined sizes, and whether an -object of type `S` can be stored in an array with elements of type `T`. -""" -function check_basetype(::Type{T}, ::Type{S}) where {T, S} - if @generated - if !isbitstype(T) - estr = "Base type $T has indeterminate size" - :(error($estr)) - elseif !isbitstype(S) - estr = "Struct type $S has indeterminate size" - :(error($estr)) - elseif !is_valid_basetype(T, S) - estr = "Struct type $S cannot be represented using base type $T" - :(error($estr)) - else - :(return nothing) - end - else - isbitstype(T) || error("Base type $T has indeterminate size") - isbitstype(S) || error("Struct type $S has indeterminate size") - is_valid_basetype(T, S) || - error("Struct type $S cannot be represented using base type $T") - return nothing - end -end - -""" - replace_basetype(::Type{T}, ::Type{T′}, ::Type{S}) - -Changes the type parameters of `S` to produce a new type `S′` such that, if -`is_valid_basetype(T, S)` is true, then so is `is_valid_basetype(T′, S′)`. -""" -replace_basetype(::Type{T}, ::Type{T′}, ::Type{S}) where {T, T′, S} = - length(S.parameters) == 0 ? S : - S.name.wrapper{replace_basetypes(T, T′, Tuple(S.parameters))...} -replace_basetype(::Type{T}, ::Type{T′}, ::Type{<:T}) where {T, T′} = T′ -replace_basetype(::Type{T}, ::Type{T′}, value) where {T, T′} = value -replace_basetypes(::Type{T}, ::Type{T′}, values) where {T, T′} = - unrolled_map(values) do value - replace_basetype(T, T′, value) - end -# TODO: This could potentially lead to some annoying bugs, since it replaces -# type parameters instead of field types. So, if `S` has `Float64` as a -# parameter, `replace_basetype(Float64, Float32, S)` will replace that parameter -# with `Float32`, regardless of whether the parameter corresponds to any values -# stored in an object of type `S`. In other words, if a user utilizes types as -# type parameters and specializes their code on those parameters, this may -# change the results of their code. -# Note that there is no way to write `replace_basetype` using field types -# instead of type parameters, since replacing the field types of an object will -# change that object's type parameters, but there is no general way to map field -# types to type parameters. - -""" - 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} - -""" - StructArrays.bypass_constructor(T, args) - -Create an instance of type `T` from a tuple of field values `args`, bypassing -possible internal constructors. `T` should be a concrete type. -""" -Base.@propagate_inbounds @generated function bypass_constructor( - ::Type{T}, - args, -) where {T} - vars = ntuple(_ -> gensym(), fieldcount(T)) - assign = [ - :(@inbounds $var::$(fieldtype(T, i)) = getfield(args, $i)) for - (i, var) in enumerate(vars) - ] - construct = Expr(:new, :T, vars...) - Expr(:block, assign..., construct) -end - -""" - fieldtypeoffset(T,S,i) - -Similar to `fieldoffset(S,i)`, but gives result in multiples of `sizeof(T)` instead of bytes. -""" -fieldtypeoffset(::Type{T}, ::Type{S}, i) where {T, S} = - Int(div(fieldoffset(S, i), sizeof(T))) - -@generated function fieldtypeoffset( - ::Type{T}, - ::Type{S}, - ::Val{i}, -) where {T, S, i} - return :(Int(div(fieldoffset(S, i), sizeof(T)))) -end - -""" - typesize(T,S) - -Similar to `sizeof(S)`, but gives the result in multiples of `sizeof(T)`. -""" -typesize(::Type{T}, ::Type{S}) where {T, S} = div(sizeof(S), sizeof(T)) - -@inline offset_index( - start_index::CartesianIndex{N}, - ::Val{D}, - offset, -) where {N, D} = CartesianIndex( - ntuple(n -> n == D ? start_index[n] + offset : start_index[n], N), -) - -""" - get_struct(array, S, Val(D), start_index) - -Construct an object of type `S` packed along the `D` dimension, from the values of `array`, -starting at `start_index`. -""" -Base.@propagate_inbounds @generated function get_struct( - array::AbstractArray{T}, - ::Type{S}, - ::Val{D}, - start_index::CartesianIndex, -) where {T, S, D} - # recursion base case: hit array type is the same as the struct leaf type - if T === S # Use Union-splitting for better latency - return quote - Base.@_propagate_inbounds_meta - @inbounds return array[start_index] - end - end - tup = :(()) - for i in 1:fieldcount(S) - push!( - tup.args, - :(get_struct( - array, - fieldtype(S, $i), - Val($D), - offset_index( - start_index, - Val($D), - $(fieldtypeoffset(T, S, Val(i))), - ), - )), - ) - end - return quote - Base.@_propagate_inbounds_meta - @inbounds bypass_constructor(S, $tup) - end -end - -""" - set_struct!(array, val::S, Val(D), start_index) - -Store an object `val` of type `S` packed along the `D` dimension, into `array`, -starting at `start_index`. -""" -Base.@propagate_inbounds @generated function set_struct!( - array::AbstractArray{T}, - val::S, - ::Val{D}, - start_index::CartesianIndex, -) where {T, S, D} - ex = quote - Base.@_propagate_inbounds_meta - end - for i in 1:fieldcount(S) - push!( - ex.args, - :(set_struct!( - array, - getfield(val, $i), - Val($D), - offset_index(start_index, Val($D), $(fieldtypeoffset(T, S, i))), - )), - ) - end - push!(ex.args, :(return val)) - return ex -end - -Base.@propagate_inbounds function set_struct!( - array::AbstractArray{S}, - val::S, - ::Val{D}, - index::CartesianIndex, -) where {S, D} - @inbounds array[index] = val - val -end - -# For complex nested types (ex. wrapped SMatrix) we hit a recursion limit and de-optimize -# We know the recursion will terminate due to the fact that bitstype fields -# cannot be self referential so there are no cycles in get/set_struct (bounded tree) -# TODO: enforce inference termination some other way -if hasfield(Method, :recursion_relation) - dont_limit = (args...) -> true - for m in methods(get_struct) - m.recursion_relation = dont_limit - end - for m in methods(set_struct!) - m.recursion_relation = dont_limit - end -end diff --git a/src/DataLayouts/struct_linear_indexing.jl b/src/DataLayouts/struct_linear_indexing.jl deleted file mode 100644 index 68e2bd10b0..0000000000 --- a/src/DataLayouts/struct_linear_indexing.jl +++ /dev/null @@ -1,121 +0,0 @@ -##### -##### Linear indexing -##### - -""" - offset_index_linear( - start_index::Integer, - field_offset, - array_size, - ) - -Compute the linear offset from a starting index, -the field offset, and the array size. - -This can be done more efficiently if the array -size is statically known, but we currently -do no store `Nh` in the type space. -""" -@inline function offset_index_linear( - start_index::Integer, - field_offset, - array_size, -) - # Assumes that the field index is _last_: - return @inbounds start_index + prod(array_size[1:(end - 1)]) * field_offset -end - -Base.@propagate_inbounds @generated function get_struct_linear( - array::AbstractArray{T}, - ::Type{S}, - start_index::Integer, - array_size, -) where {T, S} - tup = :(()) - for i in 1:fieldcount(S) - push!( - tup.args, - :(get_struct_linear( - array, - fieldtype(S, $i), - offset_index_linear( - start_index, - $(fieldtypeoffset(T, S, Val(i))), - array_size, - ), - array_size, - )), - ) - end - return quote - Base.@_propagate_inbounds_meta - @inbounds bypass_constructor(S, $tup) - end -end - -# recursion base case: hit array type is the same as the struct leaf type -Base.@propagate_inbounds function get_struct_linear( - array::AbstractArray{S}, - ::Type{S}, - start_index::Integer, - array_size, -) where {S} - @inbounds return array[start_index] -end - -""" - set_struct!(array, val::S, Val(D), start_index) -Store an object `val` of type `S` packed along the `D` dimension, into `array`, -starting at `start_index`. -""" -Base.@propagate_inbounds @generated function set_struct_linear!( - array::AbstractArray{T}, - val::S, - start_index::Integer, - array_size, -) where {T, S} - ex = quote - Base.@_propagate_inbounds_meta - end - for i in 1:fieldcount(S) - push!( - ex.args, - :(set_struct_linear!( - array, - getfield(val, $i), - offset_index_linear( - start_index, - $(fieldtypeoffset(T, S, Val(i))), - array_size, - ), - array_size, - )), - ) - end - push!(ex.args, :(return val)) - return ex -end - -Base.@propagate_inbounds function set_struct_linear!( - array::AbstractArray{S}, - val::S, - start_index::Integer, - array_size, -) where {S} - @inbounds array[start_index] = val - val -end - -# For complex nested types (ex. wrapped SMatrix) we hit a recursion limit and de-optimize -# We know the recursion will terminate due to the fact that bitstype fields -# cannot be self referential so there are no cycles in get/set_struct (bounded tree) -# TODO: enforce inference termination some other way -if hasfield(Method, :recursion_relation) - dont_limit = (args...) -> true - for m in methods(get_struct_linear) - m.recursion_relation = dont_limit - end - for m in methods(set_struct_linear!) - m.recursion_relation = dont_limit - end -end diff --git a/src/DataLayouts/struct_storage.jl b/src/DataLayouts/struct_storage.jl new file mode 100644 index 0000000000..26ac1e66d2 --- /dev/null +++ b/src/DataLayouts/struct_storage.jl @@ -0,0 +1,265 @@ +@inline default_basetype_size(::Type{S}) where {S} = + default_basetype_size(Val(S)) +@inline field_type_by_size(::Type{S}, ::Val{num_bytes}) where {S, num_bytes} = + field_type_by_size(Val(S), Val(num_bytes)) + +# Wrap each type in a Val to guarantee recursive inlining +@inline default_basetype_size(::Val{S}) where {S} = + issingletontype(S) || iszero(fieldcount(S)) ? sizeof(S) : + unrolled_mapreduce(default_basetype_size, gcd, fieldtype_vals(S)) +@inline field_type_by_size(::Val{S}, ::Val{num_bytes}) where {S, num_bytes} = + sizeof(S) == num_bytes ? S : + issingletontype(S) || iszero(fieldcount(S)) ? nothing : + unrolled_mapreduce( + Base.Fix2(field_type_by_size, Val(num_bytes)), + (option1, option2) -> isnothing(option2) ? option1 : option2, + fieldtype_vals(S), + ) + +""" + default_basetype(S) + +Finds a type that [`set_struct!`](@ref) and [`get_struct`](@ref) can use to +store either a value of type `S`, or any of the fields within such a value. If +possible, this type is found by recursively searching the `fieldtypes` of `S`; +otherwise, an unsigned integer type is selected based on the `fieldtype` sizes. +""" +@inline function default_basetype(::Type{S}) where {S} + issingletontype(S) && return UInt8 + T = field_type_by_size(S, Val(default_basetype_size(S))) + !isnothing(T) && return T + default_basetype_size(S) == 1 && return UInt8 + default_basetype_size(S) == 2 && return UInt16 + default_basetype_size(S) == 4 && return UInt32 + default_basetype_size(S) == 8 && return UInt64 + default_basetype_size(S) == 16 && return UInt128 +end + +@inline is_valid_basetype(::Type{T}, ::Type{S}) where {T, S} = + issingletontype(S) || iszero(sizeof(S) % sizeof(T)) + +@generated invalid_basetype_string(::Type{T}, ::Type{S}) where {T, S} = + "Cannot store value of type $S ($(sizeof(S)) bytes) using values of type \ + $T ($(sizeof(T)) bytes)" + +@inline function invalid_basetype_error(::Type{T}, ::Type{S}) where {T, S} + F = unrolled_findfirst(Base.Fix1(!is_valid_basetype, T), fieldtypes(S)) + isnothing(F) && return throw(ArgumentError(invalid_basetype_string(T, S))) + return invalid_basetype_error(T, fieldtype(S, F)) +end + +""" + check_basetype(T, S) + +Checks whether [`set_struct!`](@ref) and [`get_struct`](@ref) can use values of +type `T` to store a value of type `S`. Throws an error if this is not the case, +printing out an example of a specific field that cannot use `T` as a basetype. +""" +@inline check_basetype(::Type{T}, ::Type{S}) where {T, S} = + is_valid_basetype(T, S) ? nothing : invalid_basetype_error(T, S) + +""" + checked_valid_basetype(T, S) + +Returns either `T` or the [`default_basetype`](@ref) of `S`, depending on +whether `T` satisfies [`check_basetype`](@ref) for `S`. +""" +@inline checked_valid_basetype(::Type{T}, ::Type{S}) where {T, S} = + is_valid_basetype(T, S) ? T : default_basetype(S) + +""" + num_basetypes(T, S) + +Determines how many values of type `T` are required by [`set_struct!`](@ref) and +[`get_struct`](@ref) to store a single value of type `S`. +""" +@inline num_basetypes(::Type{T}, ::Type{S}) where {T, S} = sizeof(S) ÷ sizeof(T) + +""" + struct_field_view(array, S, Val(F), [Val(D)]) + +Creates a view of the data in `array` that corresponds to a particular field of +`S`, assuming that `array` has been populated by [`set_struct!`](@ref). The +field is specified through a `Val` that contains its index `F`, and it can be +loaded from the resulting view using [`get_struct`](@ref). + +For multidimensional arrays with values stored along a particular dimension, the +resulting view contains the specified field from each value. By default, values +are assumed to be stored along the last array dimension, but any other dimension +can be specified through a `Val` that contains its index `D`. +""" +@inline function struct_field_view( + array, + ::Type{S}, + ::Val{F}, + ::Val{D} = Val(ndims(array)), +) where {S, F, D} + num_D_indices = num_basetypes(eltype(array), fieldtype(S, F)) + last_D_index = num_basetypes(eltype(array), Tuple{fieldtypes(S)[1:F]...}) + D_indices = (last_D_index - num_D_indices + 1):last_D_index + all_indices = Base.setindex(axes(array), D_indices, D) + @boundscheck checkbounds(array, all_indices...) + return Base.unsafe_view(array, all_indices...) +end + +@inline check_struct_indices(array, ::Val{num_indices}) where {num_indices} = + checkbounds(array, 1:num_indices) +@inline function check_struct_indices( + array, + ::Val{num_indices}, + start::Integer, + ::Val{D} = Val(ndims(array)), +) where {num_indices, D} + step = prod(size(array)[1:(D - 1)]) + checkbounds(array, range(start; step, length = num_indices)) +end +@inline function check_struct_indices( + array, + ::Val{num_indices}, + index::CartesianIndex, + ::Val{D} = Val(ndims(array)), +) where {num_indices, D} + start = CartesianIndex(Tuple(index)[1:(D - 1)]..., 1, Tuple(index)[D:end]...) + checkbounds(array, start:Base.setindex(start, num_indices, D)) +end + +@inline struct_index(i, array) = i +@inline struct_index( + i, + array, + start::Integer, + ::Val{D} = Val(ndims(array)), +) where {D} = start + (i - 1) * prod(size(array)[1:(D - 1)]) +@inline struct_index( + i, + array, + index::CartesianIndex, + ::Val{D} = Val(ndims(array)), +) where {D} = + CartesianIndex(Tuple(index)[1:(D - 1)]..., i, Tuple(index)[D:end]...) + +""" + set_struct!(array, value, [index], [Val(D)]) + +Populates `array` with data that represents any `isbits` `value`, using +[`bitcast_struct`](@ref) to convert `value` into entries of the array. + +For multidimensional arrays with values stored along a particular dimension, an +index is used to identify the location of one value. By default, values will be +stored along the last array dimension, but any other dimension can be specified +as `Val(D)`. The target location's index should be either an integer that +corresponds to its start, or a `CartesianIndex` that contains its coordinate +along every dimension except `D`. + +# Examples + +```julia-repl +julia> set_struct!(zeros(Int8, 4), Int32(1)) +4-element Vector{Int8}: + 1 + 0 + 0 + 0 + +julia> set_struct!(zeros(Int64, 4), (Int32(2), Int32(0), Int128(1))) +4-element Vector{Int64}: + 2 + 0 + 1 + 0 + +julia> set_struct!(zeros(Int64, 2, 4), (Int32(2), Int32(0), Int128(1)), 2) +2×4 Matrix{Int64}: + 0 0 0 0 + 2 0 1 0 + +julia> set_struct!(zeros(Int64, 4, 2), (Int32(2), Int32(0), Int128(1)), 5, Val(1)) +4×2 Matrix{Int64}: + 0 2 + 0 0 + 0 1 + 0 0 +``` +""" +Base.@propagate_inbounds function set_struct!(array, value::S, index...) where {S} + num_indices = num_basetypes(eltype(array), S) + @boundscheck check_struct_indices(array, Val(num_indices), index...) + entries = bitcast_struct(NTuple{num_indices, eltype(array)}, value) + unrolled_foreach(enumerate(entries)) do (i, entry) + @inbounds array[struct_index(i, array, index...)] = entry + end + return array +end + +""" + get_struct(array, S, [index], [Val(D)]) + +Loads a value of type `S` that [`set_struct!`](@ref) has stored in `array`, +using [`bitcast_struct`](@ref) to convert entries of the array into this value. + +For multidimensional arrays with values stored along a particular dimension, an +index is used to identify the location of one value. By default, values are +assumed to be stored along the last array dimension, but any other dimension can +be specified as `Val(D)`. The target location's index should be either an +integer that corresponds to its start, or a `CartesianIndex` that contains its +coordinate along every dimension except `D`. + +# Examples + +```julia-repl +julia> get_struct(Int8[1, 0, 0, 0], Int32) +1 + +julia> get_struct([2, 0, 1, 0], Tuple{Int32, Int32, Int128}) +(2, 0, 1) + +julia> get_struct([0 0 0 0; 2 0 1 0], Tuple{Int32, Int32, Int128}, 2) +(2, 0, 1) + +julia> get_struct([0 2; 0 0; 0 1; 0 0], Tuple{Int32, Int32, Int128}, 5, Val(1)) +(2, 0, 1) +``` +""" +Base.@propagate_inbounds function get_struct(array, ::Type{S}, index...) where {S} + num_indices = num_basetypes(eltype(array), S) + @boundscheck check_struct_indices(array, Val(num_indices), index...) + return bitcast_struct(S, array, Val(num_indices), index...) +end + +""" + parent_array_type(A, [T]) + +Determines the array type underlying the wrapper type `A`, dropping all +parameters related to array dimensions. A new basetype `T` can be specified to +replace the original `eltype(A)`. +""" +parent_array_type(::Type{A}) where {A} = parent_array_type(A, eltype(A)) +parent_array_type(::Type{<:Array}, ::Type{T}) where {T} = Array{T} +parent_array_type(::Type{<:MArray}, ::Type{T}) where {T} = MArray{<:Any, T} +parent_array_type(::Type{<:SubArray{<:Any, <:Any, A}}, ::Type{T}) where {A, T} = + parent_array_type(A, T) +parent_array_type( + ::Type{<:Base.ReshapedArray{<:Any, <:Any, A}}, + ::Type{T}, +) where {A, T} = parent_array_type(A, T) + +""" + promote_parent_array_type(A1, A2) + +Promotes two array types `A1` and `A2` generated by [`parent_array_type`](@ref), +which includes promoting their basetypes `eltype(A1)` and `eltype(A2)`. +""" +promote_parent_array_type(::Type{Array{T1}}, ::Type{Array{T2}}) where {T1, T2} = + Array{promote_type(T1, T2)} +promote_parent_array_type( + ::Type{MArray{<:Any, T1}}, + ::Type{MArray{<:Any, T2}}, +) where {T1, T2} = MArray{<:Any, promote_type(T1, T2)} +promote_parent_array_type( + ::Type{MArray{<:Any, T1}}, + ::Type{Array{T2}}, +) where {T1, T2} = MArray{<:Any, promote_type(T1, T2)} +promote_parent_array_type( + ::Type{Array{T1}}, + ::Type{MArray{<:Any, T2}}, +) where {T1, T2} = MArray{<:Any, promote_type(T1, T2)} diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index 63211bedc9..6f2b6759a8 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -54,13 +54,6 @@ Field(values::V, space::S) where {V <: AbstractData, S <: AbstractSpace} = Field(::Type{T}, space::S) where {T, S <: AbstractSpace} = Field(similar(Spaces.coordinates_data(space), T), space) -function Field(::Type{Bool}, space::S) where {S <: AbstractSpace} - FT = Spaces.undertype(space) - data = similar(Spaces.coordinates_data(space), FT) - bool_data = DataLayouts.replace_basetype(data, Bool) - return Field(bool_data, space) -end - local_geometry_type(::Field{V, S}) where {V, S} = local_geometry_type(S) ClimaComms.context(field::Field) = ClimaComms.context(axes(field)) @@ -170,17 +163,16 @@ ClimaComms.device(field::Field) = ClimaComms.device(axes(field)) ClimaComms.array_type(field::Field) = ClimaComms.array_type(ClimaComms.device(field)) -# need to define twice to avoid ambiguities @inline Base.dotgetproperty(field::Field, prop) = Base.getproperty(field, prop) - +@inline Base.getproperty(field::Field, i::Integer) = Field( + DataLayouts.field_index_view(field_values(field), Val(i)), + axes(field), +) @inline Base.getproperty(field::Field, name::Symbol) = Field( - DataLayouts._getproperty(field_values(field), Val{name}()), + DataLayouts.field_name_view(field_values(field), Val(name)), axes(field), ) -@inline Base.getproperty(field::Field, name::Integer) = - Field(getproperty(field_values(field), name), axes(field)) - Base.eltype(::Type{<:Field{V}}) where {V} = eltype(V) Base.parent(field::Field) = parent(field_values(field)) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 78993a9933..cfb24a38f3 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -248,7 +248,7 @@ function get_internal_entry( field_offset_and_type(name_pair, T, S, name_pair) if isa(index_method, Val{:view}) @assert target_type <: T - band_element_size = DataLayouts.typesize(T, S) + band_element_size = DataLayouts.num_basetypes(T, S) singleton_datalayout = DataLayouts.singleton(Fields.field_values(entry)) scalar_band_type = band_matrix_row_type(outer_diagonals(eltype(entry))..., target_type) @@ -398,40 +398,30 @@ function field_offset_and_type( extract_first(name_pair[1]) in fieldnames(S) # index with first part of name_pair[1] remaining_field_chain = (drop_first(name_pair[1]), name_pair[2]) child_type = fieldtype(S, extract_first(name_pair[1])) - field_index = unrolled_filter( - i -> fieldname(S, i) == extract_first(name_pair[1]), - 1:fieldcount(S), - )[1] (remaining_offset, end_type, index_method) = field_offset_and_type( remaining_field_chain, T, child_type, full_key, ) - return ( - DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset, - end_type, - index_method, - ) + F = unrolled_findfirst(==(extract_first(name_pair[1])), fieldnames(S)) + offset_of_field_F = + DataLayouts.num_basetypes(T, Tuple{fieldtypes(S)[1:(F - 1)]...}) + return (remaining_offset + offset_of_field_F, end_type, index_method) elseif name_pair[2] != @name() && extract_first(name_pair[2]) in fieldnames(S) # index with first part of name_pair[2] remaining_field_chain = name_pair[1], drop_first(name_pair[2]) child_type = fieldtype(S, extract_first(name_pair[2])) - field_index = unrolled_filter( - i -> fieldname(S, i) == extract_first(name_pair[2]), - 1:fieldcount(S), - )[1] (remaining_offset, end_type, index_method) = field_offset_and_type( remaining_field_chain, T, child_type, full_key, ) - return ( - DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset, - end_type, - index_method, - ) + F = unrolled_findfirst(==(extract_first(name_pair[2])), fieldnames(S)) + offset_of_field_F = + DataLayouts.num_basetypes(T, Tuple{fieldtypes(S)[1:(F - 1)]...}) + return (remaining_offset + offset_of_field_F, end_type, index_method) elseif !any(isequal(@name()), name_pair) # implicit tensor structure optimization (remaining_offset, end_type, index_method) = field_offset_and_type( (drop_first(name_pair[1]), drop_first(name_pair[2])), diff --git a/src/Operators/columnwise.jl b/src/Operators/columnwise.jl index 42e4609f94..fc81027482 100644 --- a/src/Operators/columnwise.jl +++ b/src/Operators/columnwise.jl @@ -130,12 +130,12 @@ function columnwise_kernel!( ᶜus = DataLayouts.UniversalSize(ᶜY_fv) ᶠus = DataLayouts.UniversalSize(ᶠY_fv) (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(ᶠus) - ᶜTS = DataLayouts.typesize(FT, eltype(ᶜY_fv)) - ᶠTS = DataLayouts.typesize(FT, eltype(ᶠY_fv)) + ᶜTS = DataLayouts.num_basetypes(FT, eltype(ᶜY_fv)) + ᶠTS = DataLayouts.num_basetypes(FT, eltype(ᶠY_fv)) ᶜlg = Spaces.local_geometry_data(axes(_ᶜY)) ᶠlg = Spaces.local_geometry_data(axes(_ᶠY)) SLG = partial_lg_type(eltype(ᶜlg)) - ᶜTS_lg = DataLayouts.typesize(FT, SLG) + ᶜTS_lg = DataLayouts.num_basetypes(FT, SLG) ᶜui = universal_index_columnwise(device, UI, ᶜus) ᶠui = universal_index_columnwise(device, UI, ᶠus) diff --git a/src/Operators/spectralelement.jl b/src/Operators/spectralelement.jl index 55ea4bb736..a3b5f31f7c 100644 --- a/src/Operators/spectralelement.jl +++ b/src/Operators/spectralelement.jl @@ -1503,7 +1503,7 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) else error("invalid return type") end - for j in 1:Nq, i in 1:Nq + @inbounds for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) out[slab_index(i, j)] = diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl index bf2ac51099..3f0b1c03ba 100644 --- a/src/Topologies/dss.jl +++ b/src/Topologies/dss.jl @@ -105,7 +105,7 @@ function create_dss_buffer( else _transformed_type(data, local_geometry, dss_weights, DA) # extract transformed type end - Nf = DataLayouts.typesize(T, TS) + Nf = DataLayouts.num_basetypes(T, TS) perimeter_data = if !isnothing(local_geometry) fdim = DataLayouts.field_dim(DataLayouts.singleton(local_geometry)) diff --git a/src/Utilities/Utilities.jl b/src/Utilities/Utilities.jl index 0ff41630d4..e542224a0b 100644 --- a/src/Utilities/Utilities.jl +++ b/src/Utilities/Utilities.jl @@ -1,5 +1,7 @@ module Utilities +import UnrolledUtilities: unrolled_map + include("plushalf.jl") include("cache.jl") @@ -60,5 +62,47 @@ function unionall_type(::Type{T}) where {T} return T.name.wrapper end +""" + replace_type_parameter(T, P, P′) + +Recursively modifies the parameters of `T`, replacing every subtype of `P` with +`P′`. This is like constructing a value of type `T` and converting subfields of +type `P` to type `P′`, though no constructors are actually called or compiled. +""" +replace_type_parameter(T, P, P′) = replace_type_parameter(T, Val(Tuple{P, P′})) + +# Wrap the two constant types in a Val to guarantee recursive inlining +replace_type_parameter(not_a_type, _) = not_a_type +replace_type_parameter(::Type{<:P}, val::Val{Tuple{P, P′}}) where {P, P′} = P′ +replace_type_parameter(::Type{T}, val::Val{Tuple{P, P′}}) where {T, P, P′} = + isempty(T.parameters) ? T : + unionall_type(T){ + unrolled_map(Base.Fix2(replace_type_parameter, val), Tuple(T.parameters))..., + } + +""" + fieldtype_vals(T) + +Statically inferrable analogue of `Val.(fieldtypes(T))`. Functions of `Type`s +are specialized upon successful constant propagation, but functions of `Val`s +are always specialized, so `fieldtype_vals` can be used in place of `fieldtypes` +to ensure that recursive functions over nested types have inferrable outputs. +""" +@inline fieldtype_vals(::Type{T}) where {T} = + ntuple(Val ∘ Base.Fix1(fieldtype, T), Val(fieldcount(T))) + +""" + new(T, [fields]) + +Exposes the `new` pseudo-function that allocates a value of type `T` with the +specified fields. Can also be called without a second argument to leave the +allocated value with uninitialized fields. + +In contrast to the pseudo-function, this only asserts that all fields match the +`fieldtypes` of `T`, rather than automatically converting them to those types. +""" +@generated new(::Type{T}) where {T} = Expr(:new, :T) +@generated new(::Type{T}, fields) where {T} = + Expr(:splatnew, :T, :(fields::$(Tuple{fieldtypes(T)...}))) end # module diff --git a/test/DataLayouts/data0d.jl b/test/DataLayouts/data0d.jl index e1dd00bb77..dfb1f0ad21 100644 --- a/test/DataLayouts/data0d.jl +++ b/test/DataLayouts/data0d.jl @@ -7,8 +7,6 @@ using JET using ClimaComms using ClimaCore.DataLayouts -using StaticArrays -using ClimaCore.DataLayouts: get_struct, set_struct! TestFloatTypes = (Float32, Float64) device = ClimaComms.device() @@ -46,8 +44,8 @@ end @testset "DataF boundscheck" begin S = Tuple{Complex{Float64}, Float64} data = DataF{S}(ArrayType{Float64}, zeros) + @test data[1] == data[] @test data[][2] == zero(Float64) - @test_throws MethodError data[1] end @testset "DataF type safety" begin @@ -68,9 +66,7 @@ end @testset "DataF error messages" begin SA = (; a = 1.0) data = DataF{typeof(SA)}(ArrayType{Float64}) - @test_throws ErrorException( - "Invalid field name `oops` for type `$(typeof(SA))`.", - ) data.oops + @test_throws ArgumentError data.oops end @testset "DataF broadcasting between 0D data objects and scalars" begin diff --git a/test/DataLayouts/data1d.jl b/test/DataLayouts/data1d.jl index ff67ce3367..93938ef9e6 100644 --- a/test/DataLayouts/data1d.jl +++ b/test/DataLayouts/data1d.jl @@ -7,8 +7,7 @@ using JET using ClimaComms using ClimaCore.DataLayouts -using StaticArrays -using ClimaCore.DataLayouts: get_struct, set_struct!, vindex +using ClimaCore.DataLayouts: vindex TestFloatTypes = (Float32, Float64) device = ClimaComms.device() diff --git a/test/DataLayouts/data2d.jl b/test/DataLayouts/data2d.jl index d88de4b267..c89acee3fa 100644 --- a/test/DataLayouts/data2d.jl +++ b/test/DataLayouts/data2d.jl @@ -5,43 +5,34 @@ using Revise; include(joinpath("test", "DataLayouts", "data2d.jl")) using Test using ClimaComms using ClimaCore.DataLayouts -using StaticArrays -using ClimaCore.DataLayouts: check_basetype, get_struct, set_struct!, slab_index +using ClimaCore.DataLayouts: check_basetype, slab_index device = ClimaComms.device() ArrayType = ClimaComms.array_type(device) @testset "check_basetype" begin + @test_throws Exception check_basetype(Real, Real) @test_throws Exception check_basetype(Real, Float64) @test_throws Exception check_basetype(Float64, Real) @test isnothing(check_basetype(Float64, Float64)) - @test isnothing(check_basetype(Float64, Complex{Float64})) + @test isnothing(check_basetype(Float32, Float64)) + @test_throws Exception check_basetype(Float64, Float32) - @test_throws Exception check_basetype(Float32, Float64) - @test_throws Exception check_basetype(Float64, Complex{Float32}) - - @test isnothing(check_basetype(Float64, Tuple{})) @test isnothing(check_basetype(Tuple{}, Tuple{})) + @test isnothing(check_basetype(Float64, Tuple{})) @test_throws Exception check_basetype(Tuple{}, Float64) - @test isnothing(check_basetype(Int, Tuple{Int, Complex{Int}})) - @test isnothing(check_basetype(Float64, typeof(SA[1.0 2.0; 3.0 4.0]))) - S = typeof((a = ((1.0, 2.0f0), (3.0, 4.0f0)), b = (5.0, 6.0f0))) + @test isnothing(check_basetype(Float32, S)) + @test isnothing(check_basetype(Float64, S)) @test isnothing(check_basetype(Tuple{Float64, Float32}, S)) + @test_throws Exception check_basetype(NTuple{4, Float64}, S) S = typeof(((), (1.0 + 2.0im, NamedTuple()), 3.0 + 4.0im, ())) + @test isnothing(check_basetype(Float32, S)) @test isnothing(check_basetype(Float64, S)) @test isnothing(check_basetype(Complex{Float64}, S)) -end - -@testset "get_struct / set_struct!" begin - array = [1.0, 2.0, 3.0] - S = Tuple{Complex{Float64}, Float64} - @test get_struct(array, S, Val(1), CartesianIndex(1)) == (1.0 + 2.0im, 3.0) - set_struct!(array, (4.0 + 2.0im, 6.0), Val(1), CartesianIndex(1)) - @test array == [4.0, 2.0, 6.0] - @test get_struct(array, S, Val(1), CartesianIndex(1)) == (4.0 + 2.0im, 6.0) + @test_throws Exception check_basetype(NTuple{5, Float64}, S) end @testset "IJFH" begin diff --git a/test/DataLayouts/unit_copyto.jl b/test/DataLayouts/unit_copyto.jl index 5759bf04b3..327b0a2d7f 100644 --- a/test/DataLayouts/unit_copyto.jl +++ b/test/DataLayouts/unit_copyto.jl @@ -4,6 +4,7 @@ using Revise; include(joinpath("test", "DataLayouts", "unit_copyto.jl")) =# using Test using ClimaCore.DataLayouts +import ClimaCore.RecursiveApply: ⊞ import ClimaCore.Geometry import ClimaComms using StaticArrays @@ -11,155 +12,87 @@ ClimaComms.@import_required_backends import Random Random.seed!(1234) -function test_copyto_float!(data) - Random.seed!(1234) - # Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803 +all_layouts(ArrayType, S; Ni = 3, Nij = 3, Nv = 4, Nh = 5, Nk = 6) = ( + DataF{S}(ArrayType, zeros), + VF{S}(ArrayType, zeros; Nv), + IF{S}(ArrayType, zeros; Ni), + IJF{S}(ArrayType, zeros; Nij), + IFH{S}(ArrayType, zeros; Ni, Nh), + IHF{S}(ArrayType, zeros; Ni, Nh), + IJFH{S}(ArrayType, zeros; Nij, Nh), + IJHF{S}(ArrayType, zeros; Nij, Nh), + VIFH{S}(ArrayType, zeros; Nv, Ni, Nh), + VIHF{S}(ArrayType, zeros; Nv, Ni, Nh), + VIJFH{S}(ArrayType, zeros; Nv, Nij, Nh), + VIJHF{S}(ArrayType, zeros; Nv, Nij, Nh), + # DataLayouts.IJKFVH{S}(ArrayType, zeros; Nij, Nk, Nv, Nh), + # DataLayouts.IH1JH2{S}(ArrayType, zeros; Nij), +) + +function test_copyto_single_F!(data) + # Avoid using similar here due to https://github.com/CliMA/ClimaCore.jl/issues/1803 rand_data = DataLayouts.rebuild(data, similar(parent(data))) - ArrayType = ClimaComms.array_type(ClimaComms.device()) - parent(rand_data) .= - ArrayType(rand(eltype(parent(data)), DataLayouts.farray_size(data))) - Base.copyto!(data, rand_data) # test copyto!(::AbstractData, ::AbstractData) - @test all(parent(data) .== parent(rand_data)) - Base.copyto!(data, Base.Broadcast.broadcasted(+, rand_data, 1)) # test copyto!(::AbstractData, ::Broadcasted) - @test all(parent(data) .== parent(rand_data) .+ 1) + Random.rand!(parent(rand_data)) + to_data(array) = DataLayouts.bitcast_struct.(eltype(data), array) + + Base.copyto!(data, rand_data) + @test all(to_data(parent(data)) .== to_data(parent(rand_data))) + + Base.copyto!(data, Base.Broadcast.broadcasted(⊞, rand_data, 0x1)) + @test all(to_data(parent(data)) .== to_data(parent(rand_data)) .⊞ 0x1) end -function test_copyto!(data) - Random.seed!(1234) - # Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803 +function test_copyto_multiple_F!(data) + # Avoid using similar here due to https://github.com/CliMA/ClimaCore.jl/issues/1803 rand_data = DataLayouts.rebuild(data, similar(parent(data))) - ArrayType = ClimaComms.array_type(ClimaComms.device()) - parent(rand_data) .= - ArrayType(rand(eltype(parent(data)), DataLayouts.farray_size(data))) - Base.copyto!(data, rand_data) # test copyto!(::AbstractData, ::AbstractData) - @test all(parent(data.:1) .== parent(rand_data.:1)) - @test all(parent(data.:2) .== parent(rand_data.:2)) - @test all(parent(data) .== parent(rand_data)) - Base.copyto!(data.:1, Base.Broadcast.broadcasted(+, rand_data.:1, 1)) # test copyto!(::AbstractData, ::Broadcasted) - Base.copyto!(data.:2, Base.Broadcast.broadcasted(+, rand_data.:2, 1)) # test copyto!(::AbstractData, ::Broadcasted) - @test all(parent(data.:1) .== parent(rand_data.:1) .+ 1) - @test all(parent(data.:2) .== parent(rand_data.:2) .+ 1) -end + Random.rand!(parent(rand_data)) + to_data(array) = DataLayouts.bitcast_struct.(eltype(data.:1), array) -@testset "copyto! with Nf = 1" begin - device = ClimaComms.device() - ArrayType = ClimaComms.array_type(device) - FT = Float64 - S = FT - Nv = 4 - Ni = Nij = 3 - Nh = 5 - Nk = 6 - data = DataF{S}(ArrayType{FT}, zeros) - test_copyto_float!(data) - data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto_float!(data) - data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto_float!(data) - data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto_float!(data) - data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto_float!(data) - data = IJF{S}(ArrayType{FT}, zeros; Nij) - test_copyto_float!(data) - data = IF{S}(ArrayType{FT}, zeros; Ni) - test_copyto_float!(data) - data = VF{S}(ArrayType{FT}, zeros; Nv) - test_copyto_float!(data) - data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto_float!(data) - data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto_float!(data) - data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto_float!(data) - data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto_float!(data) - # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto_float!(data) # TODO: test - # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto_float!(data) # TODO: test -end + Base.copyto!(data, rand_data) + @test all(to_data(parent(data.:1)) .== to_data(parent(rand_data.:1))) + @test all(parent(data.:2) .== parent(rand_data.:2)) + # No need to convert the second component, since it has no internal padding -@testset "copyto! with Nf > 1" begin - device = ClimaComms.device() - ArrayType = ClimaComms.array_type(device) - FT = Float64 - S = Tuple{FT, FT} - Nv = 4 - Ni = Nij = 3 - Nh = 5 - Nk = 6 - data = DataF{S}(ArrayType{FT}, zeros) - test_copyto!(data) - data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto!(data) - data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto!(data) - data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto!(data) - data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto!(data) - data = IJF{S}(ArrayType{FT}, zeros; Nij) - test_copyto!(data) - data = IF{S}(ArrayType{FT}, zeros; Ni) - test_copyto!(data) - data = VF{S}(ArrayType{FT}, zeros; Nv) - test_copyto!(data) - data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto!(data) - data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto!(data) - data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto!(data) - data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto!(data) - # TODO: test this - # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto!(data) # TODO: test - # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto!(data) # TODO: test + Base.copyto!(data, Base.Broadcast.broadcasted(⊞, rand_data, 0x1)) + @test all(to_data(parent(data.:1)) .== to_data(parent(rand_data.:1)) .⊞ 0x1) + # Do not test the second component, since it spans multiple array indices end -@testset "copyto! views with Nf > 1" begin - device = ClimaComms.device() - ArrayType = ClimaComms.array_type(device) - data_view(data) = DataLayouts.rebuild( - data, - SubArray( - parent(data), - ntuple( - i -> Base.Slice(Base.OneTo(DataLayouts.farray_size(data, i))), - ndims(data), - ), +data_view(data) = DataLayouts.rebuild( + data, + SubArray( + parent(data), + ntuple( + i -> Base.Slice(Base.OneTo(DataLayouts.farray_size(data, i))), + ndims(data), ), - ) - FT = Float64 - S = Tuple{FT, FT} - Nv = 4 - Ni = Nij = 3 - Nh = 5 - Nk = 6 - # Rather than using level/slab/column, let's just make views/SubArrays - # directly so that we can easily test all cases: - data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto!(data_view(data)) - data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) - test_copyto!(data_view(data)) - data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto!(data_view(data)) - data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) - test_copyto!(data_view(data)) - data = IJF{S}(ArrayType{FT}, zeros; Nij) - test_copyto!(data_view(data)) - data = IF{S}(ArrayType{FT}, zeros; Ni) - test_copyto!(data_view(data)) - data = VF{S}(ArrayType{FT}, zeros; Nv) - test_copyto!(data_view(data)) - data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto!(data_view(data)) - data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) - test_copyto!(data_view(data)) - data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto!(data_view(data)) - data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) - test_copyto!(data_view(data)) - # TODO: test this - # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_copyto!(data) # TODO: test - # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto!(data) # TODO: test + ), +) + +@testset "copyto!" begin + ArrayType = ClimaComms.array_type(ClimaComms.device()){Float64} + @testset "Nf = 1 (uniform)" begin + for data in all_layouts(ArrayType, Float64) + test_copyto_single_F!(data) + test_copyto_single_F!(data_view(data)) + end + end + @testset "Nf = 1 (nonuniform)" begin + for data in all_layouts(ArrayType, Tuple{Int32, UInt8}) + test_copyto_single_F!(data) + test_copyto_single_F!(data_view(data)) + end + end + @testset "Nf = 3 (uniform)" begin + for data in all_layouts(ArrayType, Tuple{Float64, NTuple{2, Float64}}) + test_copyto_multiple_F!(data) + test_copyto_multiple_F!(data_view(data)) + end + end + @testset "Nf = 3 (nonuniform)" begin + for data in all_layouts(ArrayType, Tuple{Tuple{Int32, UInt8}, UInt128}) + test_copyto_multiple_F!(data) + test_copyto_multiple_F!(data_view(data)) + end + end end diff --git a/test/DataLayouts/unit_linear_indexing.jl b/test/DataLayouts/unit_linear_indexing.jl index 07122f8fa4..a7ba78a9bb 100644 --- a/test/DataLayouts/unit_linear_indexing.jl +++ b/test/DataLayouts/unit_linear_indexing.jl @@ -1,180 +1,54 @@ -#= -julia --check-bounds=yes --project -using Revise; include(joinpath("test", "DataLayouts", "unit_linear_indexing.jl")) -=# using Test -using ClimaCore.Geometry: Covariant3Vector -using ClimaCore.DataLayouts -using ClimaCore.DataLayouts: get_struct_linear -import ClimaCore.Geometry -using StaticArrays -import Random - -using IntervalSets, LinearAlgebra +using IntervalSets using ClimaComms ClimaComms.@import_required_backends -import ClimaCore: Domains, Meshes, Topologies, Spaces, Fields, Operators -import ClimaCore.DataLayouts: get_struct -import LazyBroadcast: @lazy - -Random.seed!(1234) - -import ClimaCore.DataLayouts as DL -field_dim_to_one(s, dim) = Tuple(map(j -> j == dim ? 1 : s[j], 1:length(s))) - -Base.@propagate_inbounds cart_ind(n::NTuple, i::Integer) = - @inbounds CartesianIndices(map(x -> Base.OneTo(x), n))[i] -Base.@propagate_inbounds linear_ind(n::NTuple, ci::CartesianIndex) = - @inbounds LinearIndices(map(x -> Base.OneTo(x), n))[ci] -Base.@propagate_inbounds linear_ind(n::NTuple, loc::NTuple) = - linear_ind(n, CartesianIndex(loc)) -function debug_get_struct_linear(args...; expect_test_throws = false) - if expect_test_throws - get_struct_linear(args...) - else - try - get_struct_linear(args...) - catch - get_struct_linear(args...) - end - end -end - -function one_to_n(a::Array) - for i in 1:length(a) - a[i] = i - end - return a -end -one_to_n(s::Tuple, ::Type{FT}) where {FT} = one_to_n(zeros(FT, s...)) -ncomponents(::Type{FT}, ::Type{S}) where {FT, S} = div(sizeof(S), sizeof(FT)) +import ClimaCore: Geometry, Domains, Meshes, Spaces +import ClimaCore.DataLayouts: get_struct struct Foo{T} x::T y::T end -Base.zero(::Type{Foo{T}}) where {T} = Foo{T}(0, 0) - @testset "get_struct - IHF indexing (float)" begin FT = Float64 S = FT - s_array = (3, 4, 1) - @test ncomponents(FT, S) == 1 - a = one_to_n(s_array, FT) - ss = s_array - @test debug_get_struct_linear(a, S, 1, ss) == 1.0 - @test debug_get_struct_linear(a, S, 2, ss) == 2.0 - @test debug_get_struct_linear(a, S, 3, ss) == 3.0 - @test debug_get_struct_linear(a, S, 4, ss) == 4.0 - @test debug_get_struct_linear(a, S, 5, ss) == 5.0 - @test debug_get_struct_linear(a, S, 6, ss) == 6.0 - @test debug_get_struct_linear(a, S, 7, ss) == 7.0 - @test debug_get_struct_linear(a, S, 8, ss) == 8.0 - @test debug_get_struct_linear(a, S, 9, ss) == 9.0 - @test debug_get_struct_linear(a, S, 10, ss) == 10.0 - @test debug_get_struct_linear(a, S, 11, ss) == 11.0 - @test debug_get_struct_linear(a, S, 12, ss) == 12.0 - @test_throws BoundsError debug_get_struct_linear( - a, - S, - 13, - ss; - expect_test_throws = true, - ) + a = reshape(FT.(1:12), 3, 4, 1) + for i in 1:12 + @test get_struct(a, S, i) == FT(i) + end + @test_throws BoundsError get_struct(a, S, 13) end @testset "get_struct - IHF indexing" begin FT = Float64 S = Foo{FT} - s_array = (3, 4, 2) - @test ncomponents(FT, S) == 2 - a = one_to_n(s_array, FT) - ss = s_array - @test debug_get_struct_linear(a, S, 1, ss) == Foo{FT}(1.0, 13.0) - @test debug_get_struct_linear(a, S, 2, ss) == Foo{FT}(2.0, 14.0) - @test debug_get_struct_linear(a, S, 3, ss) == Foo{FT}(3.0, 15.0) - @test debug_get_struct_linear(a, S, 4, ss) == Foo{FT}(4.0, 16.0) - @test debug_get_struct_linear(a, S, 5, ss) == Foo{FT}(5.0, 17.0) - @test debug_get_struct_linear(a, S, 6, ss) == Foo{FT}(6.0, 18.0) - @test debug_get_struct_linear(a, S, 7, ss) == Foo{FT}(7.0, 19.0) - @test debug_get_struct_linear(a, S, 8, ss) == Foo{FT}(8.0, 20.0) - @test debug_get_struct_linear(a, S, 9, ss) == Foo{FT}(9.0, 21.0) - @test debug_get_struct_linear(a, S, 10, ss) == Foo{FT}(10.0, 22.0) - @test debug_get_struct_linear(a, S, 11, ss) == Foo{FT}(11.0, 23.0) - @test debug_get_struct_linear(a, S, 12, ss) == Foo{FT}(12.0, 24.0) - @test_throws BoundsError debug_get_struct_linear( - a, - S, - 13, - ss; - expect_test_throws = true, - ) + a = reshape(FT.(1:24), 3, 4, 2) + for i in 1:12 + @test get_struct(a, S, i) == Foo{FT}(i, 12 + i) + end + @test_throws BoundsError get_struct(a, S, 13) end @testset "get_struct - IJF indexing" begin FT = Float64 S = Foo{FT} - s_array = (3, 4, 2) - @test ncomponents(FT, S) == 2 - s = field_dim_to_one(s_array, 3) - a = one_to_n(s_array, FT) - ss = s_array - @test debug_get_struct_linear(a, S, 1, ss) == Foo{FT}(1.0, 13.0) - @test debug_get_struct_linear(a, S, 2, ss) == Foo{FT}(2.0, 14.0) - @test debug_get_struct_linear(a, S, 3, ss) == Foo{FT}(3.0, 15.0) - @test debug_get_struct_linear(a, S, 4, ss) == Foo{FT}(4.0, 16.0) - @test debug_get_struct_linear(a, S, 5, ss) == Foo{FT}(5.0, 17.0) - @test debug_get_struct_linear(a, S, 6, ss) == Foo{FT}(6.0, 18.0) - @test debug_get_struct_linear(a, S, 7, ss) == Foo{FT}(7.0, 19.0) - @test debug_get_struct_linear(a, S, 8, ss) == Foo{FT}(8.0, 20.0) - @test debug_get_struct_linear(a, S, 9, ss) == Foo{FT}(9.0, 21.0) - @test debug_get_struct_linear(a, S, 10, ss) == Foo{FT}(10.0, 22.0) - @test debug_get_struct_linear(a, S, 11, ss) == Foo{FT}(11.0, 23.0) - @test debug_get_struct_linear(a, S, 12, ss) == Foo{FT}(12.0, 24.0) - @test_throws BoundsError debug_get_struct_linear( - a, - S, - 13, - ss; - expect_test_throws = true, - ) + a = reshape(FT.(1:24), 3, 4, 2) + for i in 1:12 + @test get_struct(a, S, i) == Foo{FT}(i, 12 + i) + end + @test_throws BoundsError get_struct(a, S, 13) end @testset "get_struct - VIJHF indexing" begin FT = Float64 S = Foo{FT} - s_array = (2, 2, 2, 2, 2) - @test ncomponents(FT, S) == 2 - s = field_dim_to_one(s_array, 5) - a = one_to_n(s_array, FT) - ss = s_array - - @test debug_get_struct_linear(a, S, 1, ss) == Foo{FT}(1.0, 17.0) - @test debug_get_struct_linear(a, S, 2, ss) == Foo{FT}(2.0, 18.0) - @test debug_get_struct_linear(a, S, 3, ss) == Foo{FT}(3.0, 19.0) - @test debug_get_struct_linear(a, S, 4, ss) == Foo{FT}(4.0, 20.0) - @test debug_get_struct_linear(a, S, 5, ss) == Foo{FT}(5.0, 21.0) - @test debug_get_struct_linear(a, S, 6, ss) == Foo{FT}(6.0, 22.0) - @test debug_get_struct_linear(a, S, 7, ss) == Foo{FT}(7.0, 23.0) - @test debug_get_struct_linear(a, S, 8, ss) == Foo{FT}(8.0, 24.0) - @test debug_get_struct_linear(a, S, 9, ss) == Foo{FT}(9.0, 25.0) - @test debug_get_struct_linear(a, S, 10, ss) == Foo{FT}(10.0, 26.0) - @test debug_get_struct_linear(a, S, 11, ss) == Foo{FT}(11.0, 27.0) - @test debug_get_struct_linear(a, S, 12, ss) == Foo{FT}(12.0, 28.0) - @test debug_get_struct_linear(a, S, 13, ss) == Foo{FT}(13.0, 29.0) - @test debug_get_struct_linear(a, S, 14, ss) == Foo{FT}(14.0, 30.0) - @test debug_get_struct_linear(a, S, 15, ss) == Foo{FT}(15.0, 31.0) - @test debug_get_struct_linear(a, S, 16, ss) == Foo{FT}(16.0, 32.0) - - @test_throws BoundsError debug_get_struct_linear( - a, - S, - 17, - ss; - expect_test_throws = true, - ) + a = reshape(FT.(1:32), 2, 2, 2, 2, 2) + for i in 1:16 + @test get_struct(a, S, i) == Foo{FT}(i, 16 + i) + end + @test_throws BoundsError get_struct(a, S, 17) end @testset "get_struct - example" begin @@ -183,33 +57,14 @@ end interval = Geometry.ZPoint(FT(0.0)) .. Geometry.ZPoint(FT(1.0)) domain = Domains.IntervalDomain(interval; boundary_names = (:left, :right)) mesh = Meshes.IntervalMesh(domain, stretch_fn, nelems = 5) - cs = Spaces.CenterFiniteDifferenceSpace(ClimaComms.device(), mesh) - fs = Spaces.FaceFiniteDifferenceSpace(cs) - faces = Fields.coordinate_field(fs) - lg_data = Spaces.local_geometry_data(axes(faces)) - lg1 = lg_data[CartesianIndex(1, 1, 1, 1, 1)] + space = Spaces.FaceFiniteDifferenceSpace(ClimaComms.device(), mesh) + lg_data = Spaces.local_geometry_data(space) a = parent(lg_data) S = eltype(lg_data) - ss = size(a) - @test get_struct_linear(a, S, 1, ss) == - get_struct(a, S, Val(2), CartesianIndex(1, 1)) - @test get_struct_linear(a, S, 2, ss) == - get_struct(a, S, Val(2), CartesianIndex(2, 1)) - @test get_struct_linear(a, S, 3, ss) == - get_struct(a, S, Val(2), CartesianIndex(3, 1)) - @test get_struct_linear(a, S, 4, ss) == - get_struct(a, S, Val(2), CartesianIndex(4, 1)) - @test get_struct_linear(a, S, 5, ss) == - get_struct(a, S, Val(2), CartesianIndex(5, 1)) - @test get_struct_linear(a, S, 6, ss) == - get_struct(a, S, Val(2), CartesianIndex(6, 1)) - @test_throws BoundsError debug_get_struct_linear( - a, - S, - 7, - ss; - expect_test_throws = true, - ) + for i in 1:6 + @test get_struct(a, S, i) == get_struct(a, S, CartesianIndex(i), Val(2)) + end + @test_throws BoundsError get_struct(a, S, 7) end # TODO: add set_struct! diff --git a/test/DataLayouts/unit_non_extruded_broadcast.jl b/test/DataLayouts/unit_non_extruded_broadcast.jl index 30531d00ab..66f306b282 100644 --- a/test/DataLayouts/unit_non_extruded_broadcast.jl +++ b/test/DataLayouts/unit_non_extruded_broadcast.jl @@ -37,7 +37,7 @@ end @test !DataLayouts.isascalar(bc) bc = DataLayouts.to_non_extruded_broadcasted(bc) @test !DataLayouts.isascalar(bc) - @test_throws MethodError bc[1] + @test bc[1] == bc[] @test bc[] == 10.0 end diff --git a/test/DataLayouts/unit_struct.jl b/test/DataLayouts/unit_struct.jl index 05b79ba1f9..c64c857be4 100644 --- a/test/DataLayouts/unit_struct.jl +++ b/test/DataLayouts/unit_struct.jl @@ -1,99 +1,39 @@ -#= -julia --check-bounds=yes --project -using Revise; include(joinpath("test", "DataLayouts", "unit_struct.jl")) -=# using Test -using ClimaCore.DataLayouts using ClimaCore.DataLayouts: get_struct -using StaticArrays - -function one_to_n(a::Array) - for i in 1:length(a) - a[i] = i - end - return a -end -one_to_n(s::Tuple, ::Type{FT}) where {FT} = one_to_n(zeros(FT, s...)) -ncomponents(::Type{FT}, ::Type{S}) where {FT, S} = div(sizeof(S), sizeof(FT)) -field_dim_to_one(s, dim) = Tuple(map(j -> j == dim ? 1 : s[j], 1:length(s))) -CI(s) = CartesianIndices(map(ξ -> Base.OneTo(ξ), s)) struct Foo{T} x::T y::T end -Base.zero(::Type{Foo{T}}) where {T} = Foo{T}(0, 0) - @testset "get_struct - IFH indexing" begin FT = Float64 S = Foo{FT} - s_array = (3, 2, 4) - @test ncomponents(FT, S) == 2 - s = field_dim_to_one(s_array, 2) - a = one_to_n(s_array, FT) - @test get_struct(a, S, Val(2), CI(s)[1]) == Foo{FT}(1.0, 4.0) - @test get_struct(a, S, Val(2), CI(s)[2]) == Foo{FT}(2.0, 5.0) - @test get_struct(a, S, Val(2), CI(s)[3]) == Foo{FT}(3.0, 6.0) - @test get_struct(a, S, Val(2), CI(s)[4]) == Foo{FT}(7.0, 10.0) - @test get_struct(a, S, Val(2), CI(s)[5]) == Foo{FT}(8.0, 11.0) - @test get_struct(a, S, Val(2), CI(s)[6]) == Foo{FT}(9.0, 12.0) - @test get_struct(a, S, Val(2), CI(s)[7]) == Foo{FT}(13.0, 16.0) - @test get_struct(a, S, Val(2), CI(s)[8]) == Foo{FT}(14.0, 17.0) - @test get_struct(a, S, Val(2), CI(s)[9]) == Foo{FT}(15.0, 18.0) - @test get_struct(a, S, Val(2), CI(s)[10]) == Foo{FT}(19.0, 22.0) - @test get_struct(a, S, Val(2), CI(s)[11]) == Foo{FT}(20.0, 23.0) - @test get_struct(a, S, Val(2), CI(s)[12]) == Foo{FT}(21.0, 24.0) - @test_throws BoundsError get_struct(a, S, Val(2), CI(s)[13]) + a = reshape(FT.(1:24), 3, 2, 4) + for I in CartesianIndices((3, 4)) + i = I[1] + 6 * (I[2] - 1) + @test get_struct(a, S, I, Val(2)) == Foo{FT}(i, i + 3) + end end @testset "get_struct - IJF indexing" begin FT = Float64 S = Foo{FT} - s_array = (3, 4, 2) - @test ncomponents(FT, S) == 2 - s = field_dim_to_one(s_array, 3) - a = one_to_n(s_array, FT) - @test get_struct(a, S, Val(3), CI(s)[1]) == Foo{FT}(1.0, 13.0) - @test get_struct(a, S, Val(3), CI(s)[2]) == Foo{FT}(2.0, 14.0) - @test get_struct(a, S, Val(3), CI(s)[3]) == Foo{FT}(3.0, 15.0) - @test get_struct(a, S, Val(3), CI(s)[4]) == Foo{FT}(4.0, 16.0) - @test get_struct(a, S, Val(3), CI(s)[5]) == Foo{FT}(5.0, 17.0) - @test get_struct(a, S, Val(3), CI(s)[6]) == Foo{FT}(6.0, 18.0) - @test get_struct(a, S, Val(3), CI(s)[7]) == Foo{FT}(7.0, 19.0) - @test get_struct(a, S, Val(3), CI(s)[8]) == Foo{FT}(8.0, 20.0) - @test get_struct(a, S, Val(3), CI(s)[9]) == Foo{FT}(9.0, 21.0) - @test get_struct(a, S, Val(3), CI(s)[10]) == Foo{FT}(10.0, 22.0) - @test get_struct(a, S, Val(3), CI(s)[11]) == Foo{FT}(11.0, 23.0) - @test get_struct(a, S, Val(3), CI(s)[12]) == Foo{FT}(12.0, 24.0) - @test_throws BoundsError get_struct(a, S, Val(3), CI(s)[13]) + a = reshape(FT.(1:24), 3, 4, 2) + for I in CartesianIndices((3, 4)) + i = I[1] + 3 * (I[2] - 1) + @test get_struct(a, S, I, Val(3)) == Foo{FT}(i, i + 12) + end end @testset "get_struct - VIJFH indexing" begin FT = Float64 S = Foo{FT} - s_array = (2, 2, 2, 2, 2) - s = field_dim_to_one(s_array, 4) - a = one_to_n(s_array, FT) - @test ncomponents(FT, S) == 2 - - @test get_struct(a, S, Val(4), CI(s)[1]) == Foo{FT}(1.0, 9.0) - @test get_struct(a, S, Val(4), CI(s)[2]) == Foo{FT}(2.0, 10.0) - @test get_struct(a, S, Val(4), CI(s)[3]) == Foo{FT}(3.0, 11.0) - @test get_struct(a, S, Val(4), CI(s)[4]) == Foo{FT}(4.0, 12.0) - @test get_struct(a, S, Val(4), CI(s)[5]) == Foo{FT}(5.0, 13.0) - @test get_struct(a, S, Val(4), CI(s)[6]) == Foo{FT}(6.0, 14.0) - @test get_struct(a, S, Val(4), CI(s)[7]) == Foo{FT}(7.0, 15.0) - @test get_struct(a, S, Val(4), CI(s)[8]) == Foo{FT}(8.0, 16.0) - @test get_struct(a, S, Val(4), CI(s)[9]) == Foo{FT}(17.0, 25.0) - @test get_struct(a, S, Val(4), CI(s)[10]) == Foo{FT}(18.0, 26.0) - @test get_struct(a, S, Val(4), CI(s)[11]) == Foo{FT}(19.0, 27.0) - @test get_struct(a, S, Val(4), CI(s)[12]) == Foo{FT}(20.0, 28.0) - @test get_struct(a, S, Val(4), CI(s)[13]) == Foo{FT}(21.0, 29.0) - @test get_struct(a, S, Val(4), CI(s)[14]) == Foo{FT}(22.0, 30.0) - @test get_struct(a, S, Val(4), CI(s)[15]) == Foo{FT}(23.0, 31.0) - @test get_struct(a, S, Val(4), CI(s)[16]) == Foo{FT}(24.0, 32.0) - @test_throws BoundsError get_struct(a, S, Val(4), CI(s)[17]) + a = reshape(FT.(1:32), 2, 2, 2, 2, 2) + for I in CartesianIndices((2, 2, 2, 2)) + i = I[1] + 2 * (I[2] - 1) + 4 * (I[3] - 1) + 16 * (I[4] - 1) + @test get_struct(a, S, I, Val(4)) == Foo{FT}(i, i + 8) + end end # TODO: add set_struct! diff --git a/test/Geometry/axistensors.jl b/test/Geometry/axistensors.jl index ca6d39bc2a..85898b9e33 100644 --- a/test/Geometry/axistensors.jl +++ b/test/Geometry/axistensors.jl @@ -73,7 +73,7 @@ import ClimaCore @test_throws DimensionMismatch M * x @test_throws DimensionMismatch M \ x - @test DataLayouts.typesize(Float64, typeof(x)) == 2 + @test DataLayouts.num_basetypes(Float64, typeof(x)) == 2 end @testset "Printing" begin diff --git a/test/MatrixFields/field_names.jl b/test/MatrixFields/field_names.jl index 64d5146817..3d0849ec5c 100644 --- a/test/MatrixFields/field_names.jl +++ b/test/MatrixFields/field_names.jl @@ -1,6 +1,6 @@ import LinearAlgebra: I import ClimaCore.RecursiveApply: rzero -import ClimaCore.DataLayouts: replace_basetype +import ClimaCore.Utilities: replace_type_parameter import ClimaCore.MatrixFields: @name, is_subset_that_covers_set include("matrix_field_test_utils.jl") @@ -717,7 +717,7 @@ end @testset "FieldNameDict Unit Tests" begin x = get_x() FT = Float64 - x_FT = convert(replace_basetype(Int, FT, typeof(x)), x) + x_FT = convert(replace_type_parameter(typeof(x), Int, FT), x) C3 = Geometry.Covariant3Vector{FT} C12 = Geometry.Covariant12Vector{FT} @@ -727,10 +727,10 @@ end CT3XC3 = typeof(zero(CT3) * zero(C3)') C12XCT12 = typeof(zero(C12) * zero(CT12)') CT3XCT12 = typeof(zero(CT3) * zero(CT12)') - x_C12 = rzero(replace_basetype(Int, C12, typeof(x))) - x_CT3 = rzero(replace_basetype(Int, CT3, typeof(x))) - x_C12XC3 = rzero(replace_basetype(Int, C12XC3, typeof(x))) - x_CT3XCT12 = rzero(replace_basetype(Int, CT3XCT12, typeof(x))) + x_C12 = rzero(replace_type_parameter(typeof(x), Int, C12)) + x_CT3 = rzero(replace_type_parameter(typeof(x), Int, CT3)) + x_C12XC3 = rzero(replace_type_parameter(typeof(x), Int, C12XC3)) + x_CT3XCT12 = rzero(replace_type_parameter(typeof(x), Int, CT3XCT12)) I_CT3XC3 = DiagonalMatrixRow(Geometry.AxisTensor(axes(CT3XC3), I)) I_C12XCT12 = DiagonalMatrixRow(Geometry.AxisTensor(axes(C12XCT12), I)) diff --git a/test/Operators/finitedifference/opt_examples.jl b/test/Operators/finitedifference/opt_examples.jl index 47056564c2..43b9641098 100644 --- a/test/Operators/finitedifference/opt_examples.jl +++ b/test/Operators/finitedifference/opt_examples.jl @@ -386,7 +386,7 @@ function alloc_test_nested_expressions_12(cfield, ffield, ntcfield, ntffield) p = @allocated begin @. cznt = cxnt * cynt * Ic(fynt) * Ic(fynt) * cϕnt * cψnt end - @test_broken p == 0 + @test p == 0 broken = using_cuda end end @@ -449,7 +449,7 @@ function alloc_test_nested_expressions_13( fψ end #! format: on - @test_broken p_i == 0 + @test p_i == 0 broken = using_cuda end end