Add and fix DataLayouts opt test

.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"

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)

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

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)

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