From 3084a39685ea92da402d96c41e9fbe6aeffeec4b Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Sat, 15 Jun 2024 16:06:08 -0400 Subject: [PATCH] Add and fix DataLayouts opt test --- .buildkite/pipeline.yml | 4 +++ src/DataLayouts/broadcast.jl | 33 +++++++++++++++------ src/Grids/finitedifference.jl | 18 +++++------- src/Hypsography/Hypsography.jl | 4 +-- test/DataLayouts/opt_similar.jl | 51 +++++++++++++++++++++++++++++++++ 5 files changed, 88 insertions(+), 22 deletions(-) create mode 100644 test/DataLayouts/opt_similar.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2301c9b3ae..f2bac5473a 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -81,6 +81,10 @@ steps: key: unit_data_fill command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill.jl" + - label: "Unit: data_opt_similar" + key: data_opt_similar + command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/opt_similar.jl" + - label: "Unit: data_ndims" key: unit_data_ndims command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_ndims.jl" diff --git a/src/DataLayouts/broadcast.jl b/src/DataLayouts/broadcast.jl index 64985c96ec..432bfc2bb9 100644 --- a/src/DataLayouts/broadcast.jl +++ b/src/DataLayouts/broadcast.jl @@ -329,34 +329,49 @@ function Base.similar( return IF{Eltype, Ni}(array) end +Base.similar( + bc::BroadcastedUnionVF{<:Any, Nv}, + ::Type{Eltype}, +) where {Nv, Eltype} = Base.similar(bc, Eltype, Val(Nv)) + function Base.similar( bc::BroadcastedUnionVF{<:Any, Nv, A}, ::Type{Eltype}, - dims = nothing, -) where {Nv, A, Eltype} - (_, _, _, newNv, Nh) = isnothing(dims) ? size(bc) : dims + ::Val{newNv}, +) where {Nv, A, Eltype, newNv} + (_, _, _, _, Nh) = size(bc) PA = parent_array_type(A) array = similar(PA, (newNv, typesize(eltype(A), Eltype))) return VF{Eltype, newNv}(array) end +Base.similar( + bc::BroadcastedUnionVIFH{<:Any, Nv}, + ::Type{Eltype}, +) where {Nv, Eltype} = Base.similar(bc, Eltype, Val(Nv)) + function Base.similar( bc::BroadcastedUnionVIFH{<:Any, Nv, Ni, A}, ::Type{Eltype}, - dims = nothing, -) where {Nv, Ni, A, Eltype} - (_, _, _, newNv, Nh) = isnothing(dims) ? size(bc) : dims + ::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)) return VIFH{Eltype, newNv, Ni}(array) end +Base.similar( + bc::BroadcastedUnionVIJFH{<:Any, Nv, Nij, A}, + ::Type{Eltype}, +) where {Nv, Nij, A, Eltype} = similar(bc, Eltype, Val(Nv)) + function Base.similar( bc::BroadcastedUnionVIJFH{<:Any, Nv, Nij, A}, ::Type{Eltype}, - dims = nothing, -) where {Nv, Nij, A, Eltype} - (_, _, _, newNv, Nh) = isnothing(dims) ? size(bc) : dims + ::Val{newNv}, +) where {Nv, Nij, A, Eltype, newNv} + (_, _, _, _, 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) diff --git a/src/Grids/finitedifference.jl b/src/Grids/finitedifference.jl index e1a991343a..fea82453e8 100644 --- a/src/Grids/finitedifference.jl +++ b/src/Grids/finitedifference.jl @@ -61,10 +61,8 @@ function _FiniteDifferenceGrid(topology::Topologies.IntervalTopology) for v in 1:Nv_face face_coordinates[v] = mesh.faces[v] end - center_local_geometry, face_local_geometry = fd_geometry_data( - face_coordinates; - periodic = Topologies.isperiodic(topology), - ) + center_local_geometry, face_local_geometry = + fd_geometry_data(face_coordinates, Val(Topologies.isperiodic(topology))) return FiniteDifferenceGrid( topology, @@ -76,19 +74,17 @@ end # called by the FiniteDifferenceGrid constructor, and the ExtrudedFiniteDifferenceGrid constructor with Hypsography function fd_geometry_data( - face_coordinates::DataLayouts.AbstractData{Geometry.ZPoint{FT}}; - periodic, -) where {FT} + face_coordinates::DataLayouts.AbstractData{Geometry.ZPoint{FT}}, + ::Val{periodic}, +) where {FT, periodic} CT = Geometry.ZPoint{FT} AIdx = (3,) LG = Geometry.LocalGeometry{AIdx, CT, FT, SMatrix{1, 1, FT, 1}} (Ni, Nj, Nk, Nv, Nh) = size(face_coordinates) Nv_face = Nv - periodic Nv_cent = Nv - 1 - center_local_geometry = - similar(face_coordinates, LG, (Ni, Nj, Nk, Nv_cent, Nh)) - face_local_geometry = - similar(face_coordinates, LG, (Ni, Nj, Nk, Nv_face, Nh)) + center_local_geometry = similar(face_coordinates, LG, Val(Nv_cent)) + face_local_geometry = similar(face_coordinates, LG, Val(Nv_face)) c1(args...) = Geometry.component(face_coordinates[CartesianIndex(args...)], 1) for h in 1:Nh, k in 1:Nk, j in 1:Nj, i in 1:Ni diff --git a/src/Hypsography/Hypsography.jl b/src/Hypsography/Hypsography.jl index 23d336782c..75ede63f9c 100644 --- a/src/Hypsography/Hypsography.jl +++ b/src/Hypsography/Hypsography.jl @@ -213,8 +213,8 @@ function _ExtrudedFiniteDifferenceGrid( ArrayType = ClimaComms.array_type(horizontal_grid.topology) # currently only works on Arrays (center_z_local_geometry, face_z_local_geometry) = Grids.fd_geometry_data( - Adapt.adapt(Array, face_z); - periodic = Topologies.isperiodic(vertical_grid.topology), + Adapt.adapt(Array, face_z), + Val(Topologies.isperiodic(vertical_grid.topology)), ) center_z_local_geometry = Adapt.adapt(ArrayType, center_z_local_geometry) diff --git a/test/DataLayouts/opt_similar.jl b/test/DataLayouts/opt_similar.jl new file mode 100644 index 0000000000..f69383a34c --- /dev/null +++ b/test/DataLayouts/opt_similar.jl @@ -0,0 +1,51 @@ +#= +julia --project +ENV["CLIMACOMMS_DEVICE"] = "CPU" +using Revise; include(joinpath("test", "DataLayouts", "opt_similar.jl")) +=# +using Test +using ClimaCore.DataLayouts +using ClimaCore: DataLayouts, Geometry +import ClimaComms +using StaticArrays: SMatrix +ClimaComms.@import_required_backends +using JET + +function test_similar!(data) + if data isa VF || data isa VIFH || data isa VIJFH + FT = eltype(parent(data)) + CT = Geometry.ZPoint{FT} + AIdx = (3,) + LG = Geometry.LocalGeometry{AIdx, CT, FT, SMatrix{1, 1, FT, 1}} + (_, _, _, Nv, _) = size(data) + similar(data, LG, Val(Nv)) + @test_opt similar(data, LG, Val(Nv)) + else + s = similar(data) # test callable + @test_opt similar(data) + end +end + +@testset "similar" begin + device = ClimaComms.device() + device_zeros(args...) = ClimaComms.array_type(device)(zeros(args...)) + FT = Float64 + S = FT + Nf = 1 + Nv = 4 + Nij = 3 + Nh = 5 + Nk = 6 +#! format: off + data = DataF{S}(device_zeros(FT,Nf)); test_similar!(data) + data = IJFH{S, Nij}(device_zeros(FT,Nij,Nij,Nf,Nh)); test_similar!(data) + data = IFH{S, Nij}(device_zeros(FT,Nij,Nf,Nh)); test_similar!(data) + data = IJF{S, Nij}(device_zeros(FT,Nij,Nij,Nf)); test_similar!(data) + data = IF{S, Nij}(device_zeros(FT,Nij,Nf)); test_similar!(data) + data = VF{S, Nv}(device_zeros(FT,Nv,Nf)); test_similar!(data) + data = VIJFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nij,Nf,Nh)); test_similar!(data) + data = VIFH{S, Nv, Nij}(device_zeros(FT,Nv,Nij,Nf,Nh)); test_similar!(data) +#! format: on + # data = DataLayouts.IJKFVH{S, Nij, Nk}(device_zeros(FT,Nij,Nij,Nk,Nf,Nv,Nh)); test_similar!(data) # TODO: test + # data = DataLayouts.IH1JH2{S, Nij}(device_zeros(FT,2*Nij,3*Nij)); test_similar!(data) # TODO: test +end