Skip to content

Commit

Permalink
Add DataLayouts convenience constructors
Browse files Browse the repository at this point in the history
  • Loading branch information
@charleskawczynski
charleskawczynski committed Oct 10, 2024
1 parent 0b6aa34 commit cf7776a
Show file tree
Hide file tree
Showing 16 changed files with 427 additions and 339 deletions.
133 changes: 130 additions & 3 deletions src/DataLayouts/DataLayouts.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -302,12 +302,20 @@ function replace_basetype(data::AbstractData{S}, ::Type{T}) where {S, T}
)
end

maybe_populate!(array, ::typeof(similar)) = nothing
maybe_populate!(array, ::typeof(ones)) = fill!(array, 1)
maybe_populate!(array, ::typeof(zeros)) = fill!(array, 0)
function maybe_populate!(array, ::typeof(rand))
parent(array) .= typeof(array)(rand(eltype(array), size(array)))
end

# ==================
# Data3D DataLayout
# ==================

"""
IJKFVH{S, Nij, Nk}(array::AbstractArray{T, 6}) <: Data3D{S, Nij, Nk}
IJKFVH{S}(ArrayType[, ones | zeros | rand]; Nij, Nk, Nv)
A 3D DataLayout. TODO: Add more docs
"""
Expand All @@ -327,6 +335,20 @@ function IJKFVH{S, Nij, Nk, Nv}(
IJKFVH{S, Nij, Nk, Nv, typeof(array)}(array)
end

function IJKFVH{S}(
::Type{ArrayType},
fun = similar;
Nv::Integer,
Nij::Integer,
Nk::Integer,
Nh::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nij, Nij, Nk, Nf, Nv, Nh)
maybe_populate!(array, fun)
IJKFVH{S, Nij, Nk, Nv}(array)
end

@inline universal_size(data::IJKFVH{S, Nij, Nk, Nv}) where {S, Nij, Nk, Nv} =
(Nij, Nij, Nk, Nv, get_Nh_dynamic(data))

Expand All @@ -337,6 +359,7 @@ end
"""
IJFH{S, Nij, A} <: Data2D{S, Nij}
IJFH{S,Nij}(ArrayType, nelements)
IJFH{S}(ArrayType[, ones | zeros | rand]; Nij, Nh)
Backing `DataLayout` for 2D spectral element slabs.
Expand All @@ -360,6 +383,18 @@ function IJFH{S, Nij}(array::AbstractArray{T, 4}) where {S, Nij, T}
IJFH{S, Nij, typeof(array)}(array)
end

function IJFH{S}(
::Type{ArrayType},
fun = similar;
Nij::Integer,
Nh::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nij, Nij, Nf, Nh)
maybe_populate!(array, fun)
IJFH{S, Nij}(array)
end

@inline universal_size(data::IJFH{S, Nij}) where {S, Nij} =
(Nij, Nij, 1, 1, get_Nh_dynamic(data))

Expand Down Expand Up @@ -408,6 +443,7 @@ Base.length(data::Data1D) = get_Nh_dynamic(data)
"""
IFH{S,Ni,Nh,A} <: Data1D{S, Ni}
IFH{S,Ni,Nh}(ArrayType)
IFH{S}(ArrayType[, ones | zeros | rand]; Ni, Nh)
Backing `DataLayout` for 1D spectral element slabs.
Expand All @@ -431,6 +467,18 @@ function IFH{S, Ni}(array::AbstractArray{T, 3}) where {S, Ni, T}
IFH{S, Ni, typeof(array)}(array)
end

function IFH{S}(
::Type{ArrayType},
fun = similar;
Ni::Integer,
Nh::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Ni, Nf, Nh)
maybe_populate!(array, fun)
IFH{S, Ni}(array)
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))
Expand Down Expand Up @@ -465,6 +513,7 @@ Base.length(data::Data0D) = 1

"""
DataF{S, A} <: Data0D{S}
DataF{S}(ArrayType[, ones | zeros | rand])
Backing `DataLayout` for 0D point data.
"""
Expand All @@ -478,9 +527,11 @@ function DataF{S}(array::AbstractVector{T}) where {S, T}
DataF{S, typeof(array)}(array)
end

function DataF{S}(::Type{ArrayType}) where {S, ArrayType}
T = eltype(ArrayType)
DataF{S}(ArrayType(undef, typesize(T, S)))
function DataF{S}(::Type{ArrayType}, fun = similar;) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nf)
maybe_populate!(array, fun)
DataF{S}(array)
end

function DataF(x::T) where {T}
Expand Down Expand Up @@ -544,6 +595,7 @@ end

"""
IJF{S, Nij, A} <: DataSlab2D{S, Nij}
IJF{S}(ArrayType[, ones | zeros | rand]; Nij)
Backing `DataLayout` for 2D spectral element slab data.
Expand All @@ -563,6 +615,17 @@ function IJF{S, Nij}(array::AbstractArray{T, 3}) where {S, Nij, T}
IJF{S, Nij, typeof(array)}(array)
end

function IJF{S}(
::Type{ArrayType},
fun = similar;
Nij::Integer,
) where {S, ArrayType}
Nf = typesize(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)
array = MArray{Tuple{Nij, Nij, Nf}, T, 3, Nij * Nij * Nf}(undef)
Expand Down Expand Up @@ -598,6 +661,7 @@ end

"""
IF{S, Ni, A} <: DataSlab1D{S, Ni}
IF{S}(ArrayType[, ones | zeros | rand]; Ni)
Backing `DataLayout` for 1D spectral element slab data.
Expand All @@ -615,6 +679,18 @@ function IF{S, Ni}(array::AbstractArray{T, 2}) where {S, Ni, T}
@assert size(array, 2) == typesize(T, S)
IF{S, Ni, typeof(array)}(array)
end

function IF{S}(
::Type{ArrayType},
fun = similar;
Ni::Integer,
) where {S, ArrayType}
Nf = typesize(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)
array = MArray{Tuple{Ni, Nf}, T, 2, Ni * Nf}(undef)
Expand All @@ -639,6 +715,7 @@ Base.length(data::DataColumn) = get_Nv(data)

"""
VF{S, A} <: DataColumn{S, Nv}
VF{S}(ArrayType[, ones | zeros | rand]; Nv)
Backing `DataLayout` for 1D FV column data.
Expand All @@ -657,6 +734,17 @@ function VF{S, Nv}(array::AbstractArray{T, 2}) where {S, Nv, T}
VF{S, Nv, typeof(array)}(array)
end

function VF{S}(
::Type{ArrayType},
fun = similar;
Nv::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nv, Nf)
maybe_populate!(array, fun)
VF{S, Nv}(array)
end

function VF{S, Nv}(array::AbstractVector{T}) where {S, Nv, T}
check_basetype(T, S)
@assert typesize(T, S) == 1
Expand Down Expand Up @@ -697,6 +785,7 @@ end

"""
VIJFH{S, Nij, A} <: Data2DX{S, Nij}
VIJFH{S}(ArrayType[, ones | zeros | rand]; Nv, Nij, Nh)
Backing `DataLayout` for 2D spectral element slab + extruded 1D FV column data.
Expand All @@ -715,6 +804,19 @@ function VIJFH{S, Nv, Nij}(array::AbstractArray{T, 5}) where {S, Nv, Nij, T}
VIJFH{S, Nv, Nij, typeof(array)}(array)
end

function VIJFH{S}(
::Type{ArrayType},
fun = similar;
Nv::Integer,
Nij::Integer,
Nh::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nv, Nij, Nij, Nf, Nh)
maybe_populate!(array, fun)
VIJFH{S, Nv, Nij, typeof(array)}(array)
end

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

@inline universal_size(data::VIJFH{<:Any, Nv, Nij}) where {Nv, Nij} =
Expand Down Expand Up @@ -778,6 +880,7 @@ end

"""
VIFH{S, Nv, Ni, A} <: Data1DX{S, Nv, Ni}
VIFH{S}(ArrayType[, ones | zeros | rand]; Nv, Ni, Nh)
Backing `DataLayout` for 1D spectral element slab + extruded 1D FV column data.
Expand All @@ -796,6 +899,19 @@ function VIFH{S, Nv, Ni}(array::AbstractArray{T, 4}) where {S, Nv, Ni, T}
VIFH{S, Nv, Ni, typeof(array)}(array)
end

function VIFH{S}(
::Type{ArrayType},
fun = similar;
Nv::Integer,
Ni::Integer,
Nh::Integer,
) where {S, ArrayType}
Nf = typesize(eltype(ArrayType), S)
array = similar(ArrayType, Nv, Ni, Nf, Nh)
maybe_populate!(array, fun)
VIFH{S, Nv, Ni, typeof(array)}(array)
end

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

@inline universal_size(data::VIFH{<:Any, Nv, Ni}) where {Nv, Ni} =
Expand Down Expand Up @@ -844,6 +960,7 @@ end

"""
IH1JH2{S, Nij}(data::AbstractMatrix{S})
IH1JH2{S}(ArrayType[, ones | zeros | rand]; Nij)
Stores a 2D field in a matrix using a column-major format.
The primary use is for interpolation to a regular grid for ex. plotting / field output.
Expand All @@ -858,6 +975,16 @@ function IH1JH2{S, Nij}(array::AbstractMatrix{S}) where {S, Nij}
IH1JH2{S, Nij, typeof(array)}(array)
end

function IH1JH2{S}(
::Type{ArrayType},
fun = similar;
Nij::Integer,
) where {S, ArrayType}
array = similar(ArrayType, 2 * Nij, 3 * Nij)
maybe_populate!(array, fun)
IH1JH2{S, Nij}(array)
end

@inline universal_size(data::IH1JH2{S, Nij}) where {S, Nij} =
(Nij, Nij, 1, 1, div(array_length(data), Nij * Nij))

Expand Down
25 changes: 12 additions & 13 deletions test/DataLayouts/benchmark_copyto.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,38 +40,37 @@ end

@testset "copyto! with Nf = 1" begin
device = ClimaComms.device()
device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...))
ArrayType = ClimaComms.array_type(device)
FT = Float64
S = FT
Nf = 1
Nv = 63
Nij = 4
Ni = Nij = 4
Nh = 30 * 30 * 6
Nk = 6
bm = Benchmark(; float_type = FT, device_name)
data = DataF{S}(device_zeros(FT, Nf))
data = DataF{S}(ArrayType{FT})
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)
data = IJFH{S, Nij}(device_zeros(FT, Nij, Nij, Nf, Nh))
data = IJFH{S}(ArrayType{FT}; Nij, Nh)
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)
data = IFH{S, Nij}(device_zeros(FT, Nij, Nf, Nh))
data = IFH{S}(ArrayType{FT}; Nij, Nh)
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)
# The parent array of IJF and IF datalayouts are MArrays, and can therefore not bm, be passed into CUDA kernels on the RHS.
# data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3)
# data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3)
data = VF{S, Nv}(device_zeros(FT, Nv, Nf))
# data = IJF{S}(ArrayType{FT}; Nij); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3)
# data = IF{S}(ArrayType{FT}; Nij); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3)
data = VF{S}(ArrayType{FT}; Nv)
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)
data = VIJFH{S, Nv, Nij}(device_zeros(FT, Nv, Nij, Nij, Nf, Nh))
data = VIJFH{S}(ArrayType{FT}; Nv, Nij, Nh)
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)
data = VIFH{S, Nv, Nij}(device_zeros(FT, Nv, Nij, Nf, Nh))
data = VIFH{S}(ArrayType{FT}; Nv, Nij, Nh)
benchmarkcopyto!(bm, device, data, 3)
@test all(parent(data) .== 3)

# data = IJKFVH{S}(device_zeros(FT,Nij,Nij,Nk,Nf,Nh)); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test
# data = IH1JH2{S}(device_zeros(FT,Nij,Nij,Nk,Nf,Nh)); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test
# data = IJKFVH{S}(ArrayType{FT}; Nij,Nk,Nh); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test
# data = IH1JH2{S}(ArrayType{FT}; Nij,Nk,Nh); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test
tabulate_benchmark(bm)
end
Loading

0 comments on commit cf7776a

Please sign in to comment.