diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index dc2bcbebe6..12a96de49c 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1601,6 +1601,20 @@ steps: agents: slurm_gpus: 1 + - label: ":computer: Float 32 3D sphere baroclinic wave (ρe) HF datalayout GPU" + key: "gpu_baroclinic_wave_rho_e_float32_hf" + command: + - "julia --color=yes --project=.buildkite examples/hybrid/driver.jl" + artifact_paths: + - "examples/hybrid/sphere/output/baroclinic_wave_rhoe_hf/Float32/*" + env: + TEST_NAME: "sphere/baroclinic_wave_rhoe_hf" + FLOAT_TYPE: "Float32" + HorizontalLayout: "IJHF" + CLIMACOMMS_DEVICE: "CUDA" + agents: + slurm_gpus: 1 + - label: ":computer: 3D Box limiters advection slotted spheres" key: "cpu_box_advection_limiter_slotted_spheres" command: @@ -1870,6 +1884,17 @@ steps: TEST_NAME: "sphere/baroclinic_wave_rhoe" FLOAT_TYPE: "Float64" + - label: ":computer: Float 64 3D sphere baroclinic wave (ρe) HF datalayout" + key: "cpu_baroclinic_wave_rho_e_float64_hf" + command: + - "julia --color=yes --project=.buildkite examples/hybrid/driver.jl" + artifact_paths: + - "examples/hybrid/sphere/output/baroclinic_wave_rhoe_hf/Float64/*" + env: + TEST_NAME: "sphere/baroclinic_wave_rhoe_hf" + FLOAT_TYPE: "Float64" + HorizontalLayout: "IJHF" + - label: ":computer: 3D sphere baroclinic wave (ρe)" key: "cpu_baroclinic_wave_rho_e" command: diff --git a/NEWS.md b/NEWS.md index ce8a0b9412..963737d8a1 100644 --- a/NEWS.md +++ b/NEWS.md @@ -4,11 +4,15 @@ ClimaCore.jl Release Notes main ------- - - Fixed world-age issue on Julia 1.11 issue [Julia#54780](https://github.com/JuliaLang/julia/issues/54780), PR [#2034](https://github.com/CliMA/ClimaCore.jl/pull/2034). + - We've added new datalayouts: `VIJHF`,`IJHF`,`IHF`,`VIHF`, to explore their performance compared to our existing datalayouts: `VIJFH`,`IJFH`,`IFH`,`VIFH`. PR [#2055](https://github.com/CliMA/ClimaCore.jl/pull/2053), PR [#2052](https://github.com/CliMA/ClimaCore.jl/pull/2055). + - We've refactored some modules to use less internals. PR [#2053](https://github.com/CliMA/ClimaCore.jl/pull/2053), PR [#2052](https://github.com/CliMA/ClimaCore.jl/pull/2052), [#2051](https://github.com/CliMA/ClimaCore.jl/pull/2051), [#2049](https://github.com/CliMA/ClimaCore.jl/pull/2049). + - Some work was done in attempt to reduce specializations and compile time. PR [#2042](https://github.com/CliMA/ClimaCore.jl/pull/2042), [#2041](https://github.com/CliMA/ClimaCore.jl/pull/2041) v0.14.19 ------- + - Fixed world-age issue on Julia 1.11 issue [Julia#54780](https://github.com/JuliaLang/julia/issues/54780), PR [#2034](https://github.com/CliMA/ClimaCore.jl/pull/2034). + ### ![][badge-🐛bugfix] Fix undefined behavior in `DataLayout`s PR [#2034](https://github.com/CliMA/ClimaCore.jl/pull/2034) fixes some undefined diff --git a/docs/src/api.md b/docs/src/api.md index d4f9df8c3e..7d7aad6434 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -32,6 +32,10 @@ DataLayouts.IFH DataLayouts.IJFH DataLayouts.VIFH DataLayouts.VIJFH +DataLayouts.IHF +DataLayouts.IJHF +DataLayouts.VIHF +DataLayouts.VIJHF ``` ## Geometry diff --git a/examples/common_spaces.jl b/examples/common_spaces.jl index 971e79690f..9efa7f355a 100644 --- a/examples/common_spaces.jl +++ b/examples/common_spaces.jl @@ -35,6 +35,7 @@ function make_horizontal_space( mesh, npoly, context::ClimaComms.SingletonCommsContext, + HorizontalLayout = DataLayouts.IJFH, ) quad = Quadratures.GLL{npoly + 1}() if mesh isa Meshes.AbstractMesh1D @@ -42,7 +43,7 @@ function make_horizontal_space( space = Spaces.SpectralElementSpace1D(topology, quad) elseif mesh isa Meshes.AbstractMesh2D topology = Topologies.Topology2D(context, mesh) - space = Spaces.SpectralElementSpace2D(topology, quad) + space = Spaces.SpectralElementSpace2D(topology, quad; HorizontalLayout) end return space end @@ -51,13 +52,14 @@ function make_horizontal_space( mesh, npoly, comms_ctx::ClimaComms.MPICommsContext, + HorizontalLayout = DataLayouts.IJFH, ) quad = Quadratures.GLL{npoly + 1}() if mesh isa Meshes.AbstractMesh1D error("Distributed mode does not work with 1D horizontal spaces.") elseif mesh isa Meshes.AbstractMesh2D topology = Topologies.Topology2D(comms_ctx, mesh) - space = Spaces.SpectralElementSpace2D(topology, quad) + space = Spaces.SpectralElementSpace2D(topology, quad; HorizontalLayout) end return space end diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl index a713318089..519724e691 100644 --- a/examples/hybrid/driver.jl +++ b/examples/hybrid/driver.jl @@ -30,6 +30,7 @@ ClimaComms.@import_required_backends import SciMLBase const comms_ctx = ClimaComms.context() is_distributed = comms_ctx isa ClimaComms.MPICommsContext +using ClimaCore: DataLayouts using Logging @@ -91,7 +92,16 @@ if haskey(ENV, "RESTART_FILE") ᶠlocal_geometry = Fields.local_geometry_field(Y.f) else t_start = FT(0) - h_space = make_horizontal_space(horizontal_mesh, npoly, comms_ctx) + HorizontalLayouts = Dict() + HorizontalLayouts["IJFH"] = DataLayouts.IJFH + HorizontalLayouts["IJHF"] = DataLayouts.IJHF + HorizontalLayout = HorizontalLayouts[get(ENV, "HorizontalLayout", "IJFH")] + h_space = make_horizontal_space( + horizontal_mesh, + npoly, + comms_ctx, + HorizontalLayout, + ) center_space, face_space = make_hybrid_spaces(h_space, z_max, z_elem; z_stretch) ᶜlocal_geometry = Fields.local_geometry_field(center_space) diff --git a/examples/hybrid/sphere/baroclinic_wave_rhoe_hf.jl b/examples/hybrid/sphere/baroclinic_wave_rhoe_hf.jl new file mode 100644 index 0000000000..9f720bc63a --- /dev/null +++ b/examples/hybrid/sphere/baroclinic_wave_rhoe_hf.jl @@ -0,0 +1,40 @@ +using ClimaCorePlots, Plots +using ClimaCore.DataLayouts + +include("baroclinic_wave_utilities.jl") + +const sponge = false + +# Variables required for driver.jl (modify as needed) +horizontal_mesh = cubed_sphere_mesh(; radius = R, h_elem = 4) +npoly = 4 +z_max = FT(30e3) +z_elem = 10 +t_end = FT(60 * 60 * 24 * 10) +dt = FT(400) +dt_save_to_sol = FT(60 * 60 * 24) +dt_save_to_disk = FT(0) # 0 means don't save to disk +ode_algorithm = CTS.SSP333 +jacobian_flags = (; ∂ᶜ𝔼ₜ∂ᶠ𝕄_mode = :no_∂ᶜp∂ᶜK, ∂ᶠ𝕄ₜ∂ᶜρ_mode = :exact) + +additional_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) = merge( + hyperdiffusion_cache(ᶜlocal_geometry, ᶠlocal_geometry; κ₄ = FT(2e17)), + sponge ? rayleigh_sponge_cache(ᶜlocal_geometry, ᶠlocal_geometry, dt) : (;), +) +function additional_tendency!(Yₜ, Y, p, t) + hyperdiffusion_tendency!(Yₜ, Y, p, t) + sponge && rayleigh_sponge_tendency!(Yₜ, Y, p, t) +end + +center_initial_condition(local_geometry) = + center_initial_condition(local_geometry, Val(:ρe)) +function postprocessing(sol, output_dir) + @info "L₂ norm of ρe at t = $(sol.t[1]): $(norm(sol.u[1].c.ρe))" + @info "L₂ norm of ρe at t = $(sol.t[end]): $(norm(sol.u[end].c.ρe))" + + anim = Plots.@animate for Y in sol.u + ᶜv = Geometry.UVVector.(Y.c.uₕ).components.data.:2 + Plots.plot(ᶜv, level = 3, clim = (-6, 6)) + end + Plots.mp4(anim, joinpath(output_dir, "v.mp4"), fps = 5) +end diff --git a/ext/cuda/data_layouts.jl b/ext/cuda/data_layouts.jl index 6af88f897d..471006a25b 100644 --- a/ext/cuda/data_layouts.jl +++ b/ext/cuda/data_layouts.jl @@ -1,8 +1,10 @@ import ClimaCore.DataLayouts: AbstractData import ClimaCore.DataLayouts: FusedMultiBroadcast -import ClimaCore.DataLayouts: IJKFVH, IJFH, VIJFH, VIFH, IFH, IJF, IF, VF, DataF +import ClimaCore.DataLayouts: + IJKFVH, IJFH, IJHF, VIJFH, VIJHF, VIFH, VIHF, IFH, IHF, IJF, IF, VF, DataF import ClimaCore.DataLayouts: IJFHStyle, VIJFHStyle, VFStyle, DataFStyle +import ClimaCore.DataLayouts: IJHFStyle, VIJHFStyle import ClimaCore.DataLayouts: promote_parent_array_type import ClimaCore.DataLayouts: parent_array_type import ClimaCore.DataLayouts: isascalar diff --git a/ext/cuda/data_layouts_mapreduce.jl b/ext/cuda/data_layouts_mapreduce.jl index 569d702b9d..272216f7c7 100644 --- a/ext/cuda/data_layouts_mapreduce.jl +++ b/ext/cuda/data_layouts_mapreduce.jl @@ -21,7 +21,13 @@ end function mapreduce_cuda( f, op, - data::Union{DataLayouts.VF, DataLayouts.IJFH, DataLayouts.VIJFH}; + data::Union{ + DataLayouts.VF, + DataLayouts.IJFH, + DataLayouts.IJHF, + DataLayouts.VIJFH, + DataLayouts.VIJHF, + }; weighted_jacobian = OnesArray(parent(data)), opargs..., ) diff --git a/ext/cuda/data_layouts_threadblock.jl b/ext/cuda/data_layouts_threadblock.jl index a221a7f02b..6ff4967855 100644 --- a/ext/cuda/data_layouts_threadblock.jl +++ b/ext/cuda/data_layouts_threadblock.jl @@ -47,24 +47,33 @@ bounds to ensure that the result of function is_valid_index end ##### VIJFH -@inline function partition(data::DataLayouts.VIJFH, n_max_threads::Integer) +@inline function partition( + data::Union{DataLayouts.VIJFH, DataLayouts.VIJHF}, + n_max_threads::Integer, +) (Nij, _, _, Nv, Nh) = DataLayouts.universal_size(data) Nv_thread = min(Int(fld(n_max_threads, Nij * Nij)), Nv) Nv_blocks = cld(Nv, Nv_thread) @assert prod((Nv_thread, Nij, Nij)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nv_thread, Nij, Nij))),$n_max_threads)" return (; threads = (Nv_thread, Nij, Nij), blocks = (Nv_blocks, Nh)) end -@inline function universal_index(::DataLayouts.VIJFH) +@inline function universal_index(::Union{DataLayouts.VIJFH, DataLayouts.VIJHF}) (tv, i, j) = CUDA.threadIdx() (bv, h) = CUDA.blockIdx() v = tv + (bv - 1) * CUDA.blockDim().x return CartesianIndex((i, j, 1, v, h)) end -@inline is_valid_index(::DataLayouts.VIJFH, I::CI5, us::UniversalSize) = - 1 ≤ I[4] ≤ DataLayouts.get_Nv(us) +@inline is_valid_index( + ::Union{DataLayouts.VIJFH, DataLayouts.VIJHF}, + I::CI5, + us::UniversalSize, +) = 1 ≤ I[4] ≤ DataLayouts.get_Nv(us) ##### IJFH -@inline function partition(data::DataLayouts.IJFH, n_max_threads::Integer) +@inline function partition( + data::Union{DataLayouts.IJFH, DataLayouts.IJHF}, + n_max_threads::Integer, +) (Nij, _, _, _, Nh) = DataLayouts.universal_size(data) Nh_thread = min( Int(fld(n_max_threads, Nij * Nij)), @@ -75,31 +84,40 @@ end @assert prod((Nij, Nij)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nij, Nij))),$n_max_threads)" return (; threads = (Nij, Nij, Nh_thread), blocks = (Nh_blocks,)) end -@inline function universal_index(::DataLayouts.IJFH) +@inline function universal_index(::Union{DataLayouts.IJFH, DataLayouts.IJHF}) (i, j, th) = CUDA.threadIdx() (bh,) = CUDA.blockIdx() h = th + (bh - 1) * CUDA.blockDim().z return CartesianIndex((i, j, 1, 1, h)) end -@inline is_valid_index(::DataLayouts.IJFH, I::CI5, us::UniversalSize) = - 1 ≤ I[5] ≤ DataLayouts.get_Nh(us) +@inline is_valid_index( + ::Union{DataLayouts.IJFH, DataLayouts.IJHF}, + I::CI5, + us::UniversalSize, +) = 1 ≤ I[5] ≤ DataLayouts.get_Nh(us) ##### IFH -@inline function partition(data::DataLayouts.IFH, n_max_threads::Integer) +@inline function partition( + data::Union{DataLayouts.IFH, DataLayouts.IHF}, + n_max_threads::Integer, +) (Ni, _, _, _, Nh) = DataLayouts.universal_size(data) Nh_thread = min(Int(fld(n_max_threads, Ni)), Nh) Nh_blocks = cld(Nh, Nh_thread) @assert prod((Ni, Nh_thread)) ≤ n_max_threads "threads,n_max_threads=($(prod((Ni, Nh_thread))),$n_max_threads)" return (; threads = (Ni, Nh_thread), blocks = (Nh_blocks,)) end -@inline function universal_index(::DataLayouts.IFH) +@inline function universal_index(::Union{DataLayouts.IFH, DataLayouts.IHF}) (i, th) = CUDA.threadIdx() (bh,) = CUDA.blockIdx() h = th + (bh - 1) * CUDA.blockDim().y return CartesianIndex((i, 1, 1, 1, h)) end -@inline is_valid_index(::DataLayouts.IFH, I::CI5, us::UniversalSize) = - 1 ≤ I[5] ≤ DataLayouts.get_Nh(us) +@inline is_valid_index( + ::Union{DataLayouts.IFH, DataLayouts.IHF}, + I::CI5, + us::UniversalSize, +) = 1 ≤ I[5] ≤ DataLayouts.get_Nh(us) ##### IJF @inline function partition(data::DataLayouts.IJF, n_max_threads::Integer) @@ -126,21 +144,27 @@ end @inline is_valid_index(::DataLayouts.IF, I::CI5, us::UniversalSize) = true ##### VIFH -@inline function partition(data::DataLayouts.VIFH, n_max_threads::Integer) +@inline function partition( + data::Union{DataLayouts.VIFH, DataLayouts.VIHF}, + n_max_threads::Integer, +) (Ni, _, _, Nv, Nh) = DataLayouts.universal_size(data) Nv_thread = min(Int(fld(n_max_threads, Ni)), Nv) Nv_blocks = cld(Nv, Nv_thread) @assert prod((Nv_thread, Ni)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nv_thread, Ni))),$n_max_threads)" return (; threads = (Nv_thread, Ni), blocks = (Nv_blocks, Nh)) end -@inline function universal_index(::DataLayouts.VIFH) +@inline function universal_index(::Union{DataLayouts.VIFH, DataLayouts.VIHF}) (tv, i) = CUDA.threadIdx() (bv, h) = CUDA.blockIdx() v = tv + (bv - 1) * CUDA.blockDim().x return CartesianIndex((i, 1, 1, v, h)) end -@inline is_valid_index(::DataLayouts.VIFH, I::CI5, us::UniversalSize) = - 1 ≤ I[4] ≤ DataLayouts.get_Nv(us) +@inline is_valid_index( + ::Union{DataLayouts.VIFH, DataLayouts.VIHF}, + I::CI5, + us::UniversalSize, +) = 1 ≤ I[4] ≤ DataLayouts.get_Nv(us) ##### VF @inline function partition(data::DataLayouts.VF, n_max_threads::Integer) diff --git a/ext/cuda/topologies_dss.jl b/ext/cuda/topologies_dss.jl index c84eb39c60..f70cd48956 100644 --- a/ext/cuda/topologies_dss.jl +++ b/ext/cuda/topologies_dss.jl @@ -2,6 +2,7 @@ import ClimaCore: DataLayouts, Topologies, Spaces, Fields import ClimaCore.DataLayouts: CartesianFieldIndex using CUDA import ClimaCore.Topologies +import ClimaCore.Topologies: DSSDataTypes, DSSPerimeterTypes, DSSWeightTypes import ClimaCore.Topologies: perimeter_vertex_node_index _max_threads_cuda() = 256 @@ -18,7 +19,7 @@ _configure_threadblock(nitems) = function Topologies.dss_load_perimeter_data!( ::ClimaComms.CUDADevice, dss_buffer::Topologies.DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, perimeter::Topologies.Perimeter2D, ) (; perimeter_data) = dss_buffer @@ -36,8 +37,8 @@ function Topologies.dss_load_perimeter_data!( end function dss_load_perimeter_data_kernel!( - perimeter_data::DataLayouts.AbstractData, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, perimeter::Topologies.Perimeter2D{Nq}, ) where {Nq} gidx = threadIdx().x + (blockIdx().x - Int32(1)) * blockDim().x @@ -57,7 +58,7 @@ end function Topologies.dss_unload_perimeter_data!( ::ClimaComms.CUDADevice, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, dss_buffer::Topologies.DSSBuffer, perimeter, ) @@ -76,8 +77,8 @@ function Topologies.dss_unload_perimeter_data!( end function dss_unload_perimeter_data_kernel!( - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - perimeter_data::AbstractData, + data::DSSDataTypes, + perimeter_data::DSSPerimeterTypes, perimeter::Topologies.Perimeter2D{Nq}, ) where {Nq} gidx = threadIdx().x + (blockIdx().x - Int32(1)) * blockDim().x @@ -97,7 +98,7 @@ end function Topologies.dss_local!( ::ClimaComms.CUDADevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Topologies.Perimeter2D, topology::Topologies.Topology2D, ) @@ -127,7 +128,7 @@ function Topologies.dss_local!( end function dss_local_kernel!( - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, local_vertices::AbstractVector{Tuple{Int, Int}}, local_vertex_offset::AbstractVector{Int}, interior_faces::AbstractVector{Tuple{Int, Int, Int, Int, Bool}}, @@ -182,11 +183,11 @@ end function Topologies.dss_transform!( device::ClimaComms.CUDADevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, perimeter::Topologies.Perimeter2D, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, + local_geometry::DSSDataTypes, + weight::DSSWeightTypes, localelems::AbstractVector{Int}, ) nlocalelems = length(localelems) @@ -217,11 +218,11 @@ function Topologies.dss_transform!( end function dss_transform_kernel!( - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, perimeter::Topologies.Perimeter2D, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, + local_geometry::DSSDataTypes, + weight::DSSWeightTypes, localelems::AbstractVector{Int}, ::Val{nlocalelems}, ) where {nlocalelems} @@ -248,9 +249,9 @@ end function Topologies.dss_untransform!( device::ClimaComms.CUDADevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, + local_geometry::DSSDataTypes, perimeter::Topologies.Perimeter2D, localelems::AbstractVector{Int}, ) @@ -280,9 +281,9 @@ function Topologies.dss_untransform!( end function dss_untransform_kernel!( - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, + local_geometry::DSSDataTypes, perimeter::Topologies.Perimeter2D, localelems::AbstractVector{Int}, ::Val{nlocalelems}, @@ -309,7 +310,7 @@ end # TODO: Function stubs, code to be implemented, needed only for distributed GPU runs function Topologies.dss_local_ghost!( ::ClimaComms.CUDADevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Topologies.Perimeter2D, topology::Topologies.AbstractTopology, ) @@ -337,7 +338,7 @@ function Topologies.dss_local_ghost!( end function dss_local_ghost_kernel!( - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, ghost_vertices, ghost_vertex_offset, perimeter::Topologies.Perimeter2D, @@ -402,7 +403,7 @@ end function fill_send_buffer_kernel!( send_data::AbstractArray{FT, 1}, send_buf_idx::AbstractArray{I, 2}, - perimeter_data::AbstractData, + perimeter_data::DSSPerimeterTypes, ::Val{nsend}, ) where {FT <: AbstractFloat, I <: Int, nsend} gidx = threadIdx().x + (blockIdx().x - Int32(1)) * blockDim().x @@ -444,7 +445,7 @@ function Topologies.load_from_recv_buffer!( end function load_from_recv_buffer_kernel!( - perimeter_data::AbstractData, + perimeter_data::DSSPerimeterTypes, recv_data::AbstractArray{FT, 1}, recv_buf_idx::AbstractArray{I, 2}, ::Val{nrecv}, @@ -474,7 +475,7 @@ end function Topologies.dss_ghost!( ::ClimaComms.CUDADevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Topologies.Perimeter2D, topology::Topologies.Topology2D, ) @@ -503,7 +504,7 @@ function Topologies.dss_ghost!( end function dss_ghost_kernel!( - perimeter_data::AbstractData, + perimeter_data::DSSPerimeterTypes, ghost_vertices, ghost_vertex_offset, repr_ghost_vertex, diff --git a/lib/ClimaCorePlots/src/ClimaCorePlots.jl b/lib/ClimaCorePlots/src/ClimaCorePlots.jl index 40a7887212..d6e60987e7 100644 --- a/lib/ClimaCorePlots/src/ClimaCorePlots.jl +++ b/lib/ClimaCorePlots/src/ClimaCorePlots.jl @@ -431,8 +431,15 @@ function _unfolded_pannel_matrix(field, interpolate) # TODO: inefficient memory wise, but good enough for now panels = [fill(NaN, (panel_size * dof, panel_size * dof)) for _ in 1:6] - interpolated_data = DataLayouts.IJFH{FT, interpolate}(Array{FT}, nelem) field_data = Fields.field_values(field) + fdim = DataLayouts.field_dim(DataLayouts.singleton(field_data)) + interpolated_data_type = if fdim == ndims(field_data) + DataLayouts.IJHF + else + DataLayouts.IJFH + end + interpolated_data = + interpolated_data_type{FT, interpolate}(Array{FT}, nelem) Operators.tensor_product!(interpolated_data, field_data, Imat) diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index 65d6f0d1c0..859b51a3ce 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -7,13 +7,17 @@ TODO: Add links to these datalayouts - `IJKFVH` - `IJFH` + - `IJHF` - `IFH` + - `IHF` - `DataF` - `IJF` - `IF` - `VF` - `VIJFH` + - `VIJHF` - `VIFH` + - `VIHF` - `IH1JH2` - `IV1JH2` @@ -38,7 +42,21 @@ import MultiBroadcastFusion as MBF import Adapt import ..slab, ..slab_args, ..column, ..column_args, ..level -export slab, column, level, IJFH, IJF, IFH, IF, VF, VIJFH, VIFH, DataF +export slab, + column, + level, + IJFH, + IJHF, + IJF, + IFH, + IHF, + IF, + VF, + VIJFH, + VIJHF, + VIFH, + VIHF, + DataF # Internal types for managing CPU/GPU dispatching abstract type AbstractDispatchToDevice end @@ -327,13 +345,17 @@ end abstract type AbstractDataSingleton end struct IJKFVHSingleton <: AbstractDataSingleton end struct IJFHSingleton <: AbstractDataSingleton end +struct IJHFSingleton <: AbstractDataSingleton end struct IFHSingleton <: AbstractDataSingleton end +struct IHFSingleton <: AbstractDataSingleton end struct DataFSingleton <: AbstractDataSingleton end struct IJFSingleton <: AbstractDataSingleton end struct IFSingleton <: AbstractDataSingleton end struct VFSingleton <: AbstractDataSingleton end struct VIJFHSingleton <: AbstractDataSingleton end +struct VIJHFSingleton <: AbstractDataSingleton end struct VIFHSingleton <: AbstractDataSingleton end +struct VIHFSingleton <: AbstractDataSingleton end struct IH1JH2Singleton <: AbstractDataSingleton end struct IV1JH2Singleton <: AbstractDataSingleton end @@ -493,6 +515,99 @@ function gather( end end +""" + IJHF{S, Nij, A} <: Data2D{S, Nij} + IJHF{S,Nij}(ArrayType, nelements) + + +Backing `DataLayout` for 2D spectral element slabs. + +Element nodal point (I,J) data is contiguous for each datatype `S` struct field (F), +for each 2D mesh element slab (H). + +The `ArrayType`-constructor constructs a IJHF 2D Spectral +DataLayout given the backing `ArrayType`, quadrature degrees +of freedom `Nij × Nij`, and the number of mesh elements `nelements`. + + IJHF{S}(ArrayType[, Base.ones | zeros | rand]; Nij, Nh) + +The keyword constructor returns a `IJHF` given +the `ArrayType` and (optionally) an initialization +method (one of `Base.ones`, `Base.zeros`, `Random.rand`) +and the keywords: + - `Nij` quadrature degrees of freedom per horizontal direction + - `Nh` number of mesh elements + +!!! note + Objects made with the keyword constructor accept integer + keyword inputs, so they are dynamically created. You may + want to use a different constructor if you're making the + object in a performance-critical section, and if you know + the type parameters at compile time. +""" +struct IJHF{S, Nij, A} <: Data2D{S, Nij} + array::A +end + +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) + IJHF{S, Nij, typeof(array)}(array) +end + +function IJHF{S}( + ::Type{ArrayType}, + fun = similar; + Nij::Integer, + Nh::Integer, +) where {S, ArrayType} + Nf = typesize(eltype(ArrayType), S) + array = similar(ArrayType, Nij, Nij, Nh, Nf) + maybe_populate!(array, fun) + IJHF{S, Nij}(array) +end + +@inline universal_size(data::IJHF{S, Nij}) where {S, Nij} = + (Nij, Nij, 1, 1, get_Nh_dynamic(data)) + +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))) +end + +Base.length(data::IJHF) = get_Nh_dynamic(data) + +Base.@propagate_inbounds slab(data::IJHF, h::Integer) = slab(data, 1, h) + +@inline function slab(data::IJHF{S, Nij}, v::Integer, h::Integer) where {S, Nij} + @boundscheck (v >= 1 && 1 <= h <= get_Nh_dynamic(data)) || + throw(BoundsError(data, (v, h))) + dataview = @inbounds view(parent(data), :, :, h, :) + IJF{S, Nij}(dataview) +end + +@inline function column(data::IJHF{S, Nij}, i, j, h) where {S, Nij} + @boundscheck ( + 1 <= j <= Nij && 1 <= i <= Nij && 1 <= h <= get_Nh_dynamic(data) + ) || throw(BoundsError(data, (i, j, h))) + dataview = @inbounds view(parent(data), i, j, h, :) + DataF{S}(dataview) +end + +function gather( + ctx::ClimaComms.AbstractCommsContext, + data::IJHF{S, Nij}, +) where {S, Nij} + gatherdata = ClimaComms.gather(ctx, parent(data)) + if ClimaComms.iamroot(ctx) + IJHF{S, Nij}(gatherdata) + else + nothing + end +end + # ================== # Data1D DataLayout # ================== @@ -578,6 +693,85 @@ end Base.@propagate_inbounds column(data::IFH{S, Ni}, i, j, h) where {S, Ni} = column(data, i, h) +""" + IHF{S,Ni,Nh,A} <: Data1D{S, Ni} + IHF{S,Ni,Nh}(ArrayType) + +Backing `DataLayout` for 1D spectral element slabs. + +Element nodal point (I) data is contiguous for each +datatype `S` struct field (F), for each 1D mesh element (H). + + +The `ArrayType`-constructor makes a IHF 1D Spectral +DataLayout given the backing `ArrayType`, quadrature +degrees of freedom `Ni`, and the number of mesh elements +`Nh`. + + IHF{S}(ArrayType[, ones | zeros | rand]; Ni, Nh) + +The keyword constructor returns a `IHF` given +the `ArrayType` and (optionally) an initialization +method (one of `Base.ones`, `Base.zeros`, `Random.rand`) +and the keywords: + - `Ni` quadrature degrees of freedom in the horizontal direction + - `Nh` number of mesh elements + +!!! note + Objects made with the keyword constructor accept integer + keyword inputs, so they are dynamically created. You may + want to use a different constructor if you're making the + object in a performance-critical section, and if you know + the type parameters at compile time. +""" +struct IHF{S, Ni, A} <: Data1D{S, Ni} + array::A +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) + IHF{S, Ni, typeof(array)}(array) +end + +function IHF{S}( + ::Type{ArrayType}, + fun = similar; + Ni::Integer, + Nh::Integer, +) where {S, ArrayType} + Nf = typesize(eltype(ArrayType), S) + array = similar(ArrayType, Ni, Nh, Nf) + maybe_populate!(array, fun) + IHF{S, Ni}(array) +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))) +end + +@inline universal_size(data::IHF{S, Ni}) where {S, Ni} = + (Ni, 1, 1, 1, get_Nh_dynamic(data)) + +@inline function slab(data::IHF{S, Ni}, h::Integer) where {S, Ni} + @boundscheck (1 <= h <= get_Nh_dynamic(data)) || + throw(BoundsError(data, (h,))) + dataview = @inbounds view(parent(data), :, h, :) + IF{S, Ni}(dataview) +end +Base.@propagate_inbounds slab(data::IHF, v::Integer, h::Integer) = slab(data, h) + +@inline function column(data::IHF{S, Ni}, i, h) where {S, Ni} + @boundscheck (1 <= h <= get_Nh_dynamic(data) && 1 <= i <= Ni) || + throw(BoundsError(data, (i, h))) + dataview = @inbounds view(parent(data), i, h, :) + DataF{S}(dataview) +end +Base.@propagate_inbounds column(data::IHF{S, Ni}, i, j, h) where {S, Ni} = + column(data, i, h) + # ====================== # Data0D DataLayout # ====================== @@ -1011,6 +1205,113 @@ function gather( end end +""" + VIJHF{S, Nij, A} <: Data2DX{S, Nij} + +Backing `DataLayout` for 2D spectral element slab + extruded 1D FV column data. + +Column levels (V) are contiguous for every element nodal point (I, J) +for each `S` datatype struct field (F), for each 2D mesh element slab (H). + + VIJHF{S}(ArrayType[, ones | zeros | rand]; Nv, Nij, Nh) + +The keyword constructor returns a `VIJHF` given +the `ArrayType` and (optionally) an initialization +method (one of `Base.ones`, `Base.zeros`, `Random.rand`) +and the keywords: + - `Nv` number of vertical degrees of freedom + - `Nij` quadrature degrees of freedom per horizontal direction + - `Nh` number of horizontal elements + +!!! note + Objects made with the keyword constructor accept integer + keyword inputs, so they are dynamically created. You may + want to use a different constructor if you're making the + object in a performance-critical section, and if you know + the type parameters at compile time. +""" +struct VIJHF{S, Nv, Nij, A} <: Data2DX{S, Nv, Nij} + array::A +end + +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) + VIJHF{S, Nv, Nij, typeof(array)}(array) +end + +function VIJHF{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, Nh, Nf) + maybe_populate!(array, fun) + VIJHF{S, Nv, Nij, typeof(array)}(array) +end + +nlevels(::VIJHF{S, Nv}) where {S, Nv} = Nv + +@inline universal_size(data::VIJHF{<:Any, Nv, Nij}) where {Nv, Nij} = + (Nij, Nij, 1, Nv, get_Nh_dynamic(data)) + +Base.length(data::VIJHF) = get_Nv(data) * get_Nh_dynamic(data) + +# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7) +@inline function slab(data::VIJHF{S, Nv, Nij}, v, h) where {S, Nv, Nij} + array = parent(data) + @boundscheck (1 <= v <= Nv && 1 <= h <= get_Nh_dynamic(data)) || + throw(BoundsError(data, (v, h))) + Nf = ncomponents(data) + dataview = @inbounds view( + array, + v, + Base.Slice(Base.OneTo(Nij)), + Base.Slice(Base.OneTo(Nij)), + h, + Base.Slice(Base.OneTo(Nf)), + ) + IJF{S, Nij}(dataview) +end + +# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7) +@inline function column(data::VIJHF{S, Nv, Nij}, i, j, h) where {S, Nv, Nij} + array = parent(data) + @boundscheck ( + 1 <= i <= Nij && 1 <= j <= Nij && 1 <= h <= get_Nh_dynamic(data) + ) || throw(BoundsError(data, (i, j, h))) + Nf = ncomponents(data) + dataview = @inbounds SubArray( + array, + (Base.Slice(Base.OneTo(Nv)), i, j, h, Base.Slice(Base.OneTo(Nf))), + ) + VF{S, Nv}(dataview) +end + +@inline function level(data::VIJHF{S, Nv, Nij}, v) where {S, Nv, Nij} + array = parent(data) + @boundscheck (1 <= v <= Nv) || throw(BoundsError(data, (v,))) + dataview = @inbounds view(array, v, :, :, :, :) + IJHF{S, Nij}(dataview) +end + +function gather( + ctx::ClimaComms.AbstractCommsContext, + data::VIJHF{S, Nv, Nij}, +) where {S, Nv, Nij} + gatherdata = ClimaComms.gather(ctx, parent(data)) + if ClimaComms.iamroot(ctx) + VIJHF{S, Nv, Nij}(gatherdata) + else + nothing + end +end + # ====================== # Data1DX DataLayout # ====================== @@ -1107,6 +1408,98 @@ end IFH{S, Nij}(dataview) end +""" + VIHF{S, Nv, Ni, A} <: Data1DX{S, Nv, Ni} + +Backing `DataLayout` for 1D spectral element slab + extruded 1D FV column data. + +Column levels (V) are contiguous for every element nodal point (I) +for each datatype `S` struct field (F), for each 1D mesh element slab (H). + + VIHF{S}(ArrayType[, ones | zeros | rand]; Nv, Ni, Nh) + +The keyword constructor returns a `VIHF` given +the `ArrayType` and (optionally) an initialization +method (one of `Base.ones`, `Base.zeros`, `Random.rand`) +and the keywords: + - `Nv` number of vertical degrees of freedom + - `Ni` quadrature degrees of freedom in the horizontal direction + - `Nh` number of horizontal elements + +!!! note + Objects made with the keyword constructor accept integer + keyword inputs, so they are dynamically created. You may + want to use a different constructor if you're making the + object in a performance-critical section, and if you know + the type parameters at compile time. +""" +struct VIHF{S, Nv, Ni, A} <: Data1DX{S, Nv, Ni} + array::A +end + +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) + VIHF{S, Nv, Ni, typeof(array)}(array) +end + +function VIHF{S}( + ::Type{ArrayType}, + fun = similar; + Nv::Integer, + Ni::Integer, + Nh::Integer, +) where {S, ArrayType} + Nf = typesize(eltype(ArrayType), S) + array = similar(ArrayType, Nv, Ni, Nh, Nf) + maybe_populate!(array, fun) + VIHF{S, Nv, Ni, typeof(array)}(array) +end + +nlevels(::VIHF{S, Nv}) where {S, Nv} = Nv + +@inline universal_size(data::VIHF{<:Any, Nv, Ni}) where {Nv, Ni} = + (Ni, 1, 1, Nv, get_Nh_dynamic(data)) + +Base.length(data::VIHF) = nlevels(data) * get_Nh_dynamic(data) + +# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7) +@inline function slab(data::VIHF{S, Nv, Ni}, v, h) where {S, Nv, Ni} + array = parent(data) + @boundscheck (1 <= v <= Nv && 1 <= h <= get_Nh_dynamic(data)) || + throw(BoundsError(data, (v, h))) + Nf = ncomponents(data) + dataview = @inbounds SubArray( + array, + (v, Base.Slice(Base.OneTo(Ni)), h, Base.Slice(Base.OneTo(Nf))), + ) + IF{S, Ni}(dataview) +end + +Base.@propagate_inbounds column(data::VIHF, i, h) = column(data, i, 1, h) + +# Note: construct the subarray view directly as optimizer fails in Base.to_indices (v1.7) +@inline function column(data::VIHF{S, Nv, Ni}, i, j, h) where {S, Nv, Ni} + array = parent(data) + @boundscheck (1 <= i <= Ni && j == 1 && 1 <= h <= get_Nh_dynamic(data)) || + throw(BoundsError(data, (i, j, h))) + Nf = ncomponents(data) + dataview = @inbounds SubArray( + array, + (Base.Slice(Base.OneTo(Nv)), i, h, Base.Slice(Base.OneTo(Nf))), + ) + VF{S, Nv}(dataview) +end + +@inline function level(data::VIHF{S, Nv, Nij}, v) where {S, Nv, Nij} + array = parent(data) + @boundscheck (1 <= v <= Nv) || throw(BoundsError(data, (v,))) + dataview = @inbounds view(array, v, :, :, :) + IHF{S, Nij}(dataview) +end + # ========================================= # Special DataLayouts for regular gridding # ========================================= @@ -1247,9 +1640,13 @@ empty_kernel_stats() = empty_kernel_stats(ClimaComms.device()) #! format: off @inline get_Nij(::IJKFVH{S, Nij}) where {S, Nij} = Nij @inline get_Nij(::IJFH{S, Nij}) where {S, Nij} = Nij +@inline get_Nij(::IJHF{S, Nij}) where {S, Nij} = Nij @inline get_Nij(::VIJFH{S, Nv, Nij}) where {S, Nv, Nij} = Nij +@inline get_Nij(::VIJHF{S, Nv, Nij}) where {S, Nv, Nij} = Nij @inline get_Nij(::VIFH{S, Nv, Nij}) where {S, Nv, Nij} = Nij +@inline get_Nij(::VIHF{S, Nv, Nij}) where {S, Nv, Nij} = Nij @inline get_Nij(::IFH{S, Nij}) where {S, Nij} = Nij +@inline get_Nij(::IHF{S, Nij}) where {S, Nij} = Nij @inline get_Nij(::IJF{S, Nij}) where {S, Nij} = Nij @inline get_Nij(::IF{S, Nij}) where {S, Nij} = Nij @@ -1266,23 +1663,31 @@ type parameters. """ @inline field_dim(::IJKFVHSingleton) = 4 @inline field_dim(::IJFHSingleton) = 3 +@inline field_dim(::IJHFSingleton) = 4 @inline field_dim(::IFHSingleton) = 2 +@inline field_dim(::IHFSingleton) = 3 @inline field_dim(::DataFSingleton) = 1 @inline field_dim(::IJFSingleton) = 3 @inline field_dim(::IFSingleton) = 2 @inline field_dim(::VFSingleton) = 2 @inline field_dim(::VIJFHSingleton) = 4 +@inline field_dim(::VIJHFSingleton) = 5 @inline field_dim(::VIFHSingleton) = 3 +@inline field_dim(::VIHFSingleton) = 4 @inline field_dim(::Type{IJKFVH}) = 4 @inline field_dim(::Type{IJFH}) = 3 +@inline field_dim(::Type{IJHF}) = 4 @inline field_dim(::Type{IFH}) = 2 +@inline field_dim(::Type{IHF}) = 3 @inline field_dim(::Type{DataF}) = 1 @inline field_dim(::Type{IJF}) = 3 @inline field_dim(::Type{IF}) = 2 @inline field_dim(::Type{VF}) = 2 @inline field_dim(::Type{VIJFH}) = 4 +@inline field_dim(::Type{VIJHF}) = 5 @inline field_dim(::Type{VIFH}) = 3 +@inline field_dim(::Type{VIHF}) = 4 """ h_dim(::AbstractDataSingleton) @@ -1297,26 +1702,38 @@ type parameters. """ @inline h_dim(::IJKFVHSingleton) = 5 @inline h_dim(::IJFHSingleton) = 4 +@inline h_dim(::IJHFSingleton) = 3 @inline h_dim(::IFHSingleton) = 3 +@inline h_dim(::IHFSingleton) = 2 @inline h_dim(::VIJFHSingleton) = 5 +@inline h_dim(::VIJHFSingleton) = 4 @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 to_data_specific_field(::DataFSingleton, I::Tuple) = (I[3],) @inline to_data_specific_field(::VFSingleton, I::Tuple) = (I[4], I[3]) @inline to_data_specific_field(::IFSingleton, I::Tuple) = (I[1], I[3]) @inline to_data_specific_field(::IJFSingleton, I::Tuple) = (I[1], I[2], I[3]) @inline to_data_specific_field(::IJFHSingleton, I::Tuple) = (I[1], I[2], I[3], I[5]) +@inline to_data_specific_field(::IJHFSingleton, I::Tuple) = (I[1], I[2], I[5], I[3]) @inline to_data_specific_field(::IFHSingleton, I::Tuple) = (I[1], I[3], I[5]) +@inline to_data_specific_field(::IHFSingleton, I::Tuple) = (I[1], I[5], I[3]) @inline to_data_specific_field(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], I[3], I[5]) +@inline to_data_specific_field(::VIJHFSingleton, I::Tuple) = (I[4], I[1], I[2], I[5], I[3]) @inline to_data_specific_field(::VIFHSingleton, I::Tuple) = (I[4], I[1], I[3], I[5]) +@inline to_data_specific_field(::VIHFSingleton, I::Tuple) = (I[4], I[1], I[5], I[3]) """ bounds_condition(data::AbstractData, I::Tuple) @@ -1345,13 +1762,17 @@ type parameters. @inline type_params(data::AbstractData) = type_params(typeof(data)) @inline type_params(::Type{IJKFVH{S, Nij, Nk, Nv, A}}) where {S, Nij, Nk, Nv, A} = (S, Nij, Nk, Nv) @inline type_params(::Type{IJFH{S, Nij, A}}) where {S, Nij, A} = (S, Nij) +@inline type_params(::Type{IJHF{S, Nij, A}}) where {S, Nij, A} = (S, Nij) @inline type_params(::Type{IFH{S, Ni, A}}) where {S, Ni, A} = (S, Ni) +@inline type_params(::Type{IHF{S, Ni, A}}) where {S, Ni, A} = (S, Ni) @inline type_params(::Type{DataF{S, A}}) where {S, A} = (S,) @inline type_params(::Type{IJF{S, Nij, A}}) where {S, Nij, A} = (S, Nij) @inline type_params(::Type{IF{S, Ni, A}}) where {S, Ni, A} = (S, Ni) @inline type_params(::Type{VF{S, Nv, A}}) where {S, Nv, A} = (S, Nv) @inline type_params(::Type{VIJFH{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = (S, Nv, Nij) +@inline type_params(::Type{VIJHF{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = (S, Nv, Nij) @inline type_params(::Type{VIFH{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = (S, Nv, Ni) +@inline type_params(::Type{VIHF{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = (S, Nv, Ni) @inline type_params(::Type{IH1JH2{S, Nij, A}}) where {S, Nij, A} = (S, Nij) @inline type_params(::Type{IV1JH2{S, n1, Ni, A}}) where {S, n1, Ni, A} = (S, n1, Ni) @@ -1370,13 +1791,17 @@ type parameters. """ @inline union_all(::IJKFVHSingleton) = IJKFVH @inline union_all(::IJFHSingleton) = IJFH +@inline union_all(::IJHFSingleton) = IJHF @inline union_all(::IFHSingleton) = IFH +@inline union_all(::IHFSingleton) = IHF @inline union_all(::DataFSingleton) = DataF @inline union_all(::IJFSingleton) = IJF @inline union_all(::IFSingleton) = IF @inline union_all(::VFSingleton) = VF @inline union_all(::VIJFHSingleton) = VIJFH +@inline union_all(::VIJHFSingleton) = VIJHF @inline union_all(::VIFHSingleton) = VIFH +@inline union_all(::VIHFSingleton) = VIHF @inline union_all(::IH1JH2Singleton) = IH1JH2 @inline union_all(::IV1JH2Singleton) = IV1JH2 @@ -1395,13 +1820,17 @@ type parameters. @inline array_size(data::AbstractData, i::Integer) = array_size(data)[i] @inline array_size(data::IJKFVH{S, Nij, Nk, Nv}) where {S, Nij, Nk, Nv} = (Nij, Nij, Nk, 1, Nv, get_Nh_dynamic(data)) @inline array_size(data::IJFH{S, Nij}) where {S, Nij} = (Nij, Nij, 1, get_Nh_dynamic(data)) +@inline array_size(data::IJHF{S, Nij}) where {S, Nij} = (Nij, Nij, get_Nh_dynamic(data), 1) @inline array_size(data::IFH{S, Ni}) where {S, Ni} = (Ni, 1, get_Nh_dynamic(data)) +@inline array_size(data::IHF{S, Ni}) where {S, Ni} = (Ni, get_Nh_dynamic(data), 1) @inline array_size(data::DataF{S}) where {S} = (1,) @inline array_size(data::IJF{S, Nij}) where {S, Nij} = (Nij, Nij, 1) @inline array_size(data::IF{S, Ni}) where {S, Ni} = (Ni, 1) @inline array_size(data::VF{S, Nv}) where {S, Nv} = (Nv, 1) @inline array_size(data::VIJFH{S, Nv, Nij}) where {S, Nv, Nij} = (Nv, Nij, Nij, 1, get_Nh_dynamic(data)) +@inline array_size(data::VIJHF{S, Nv, Nij}) where {S, Nv, Nij} = (Nv, Nij, Nij, get_Nh_dynamic(data), 1) @inline array_size(data::VIFH{S, Nv, Ni}) where {S, Nv, Ni} = (Nv, Ni, 1, get_Nh_dynamic(data)) +@inline array_size(data::VIHF{S, Nv, Ni}) where {S, Nv, Ni} = (Nv, Ni, get_Nh_dynamic(data), 1) """ farray_size(data::AbstractData) @@ -1417,13 +1846,17 @@ type parameters. @inline farray_size(data::AbstractData, i::Integer) = farray_size(data)[i] @inline farray_size(data::IJKFVH{S, Nij, Nk, Nv}) where {S, Nij, Nk, Nv} = (Nij, Nij, Nk, ncomponents(data), Nv, get_Nh_dynamic(data)) @inline farray_size(data::IJFH{S, Nij}) where {S, Nij} = (Nij, Nij, ncomponents(data), get_Nh_dynamic(data)) +@inline farray_size(data::IJHF{S, Nij}) where {S, Nij} = (Nij, Nij, get_Nh_dynamic(data), ncomponents(data)) @inline farray_size(data::IFH{S, Ni}) where {S, Ni} = (Ni, ncomponents(data), get_Nh_dynamic(data)) +@inline farray_size(data::IHF{S, Ni}) where {S, Ni} = (Ni, get_Nh_dynamic(data), ncomponents(data)) @inline farray_size(data::DataF{S}) where {S} = (ncomponents(data),) @inline farray_size(data::IJF{S, Nij}) where {S, Nij} = (Nij, Nij, ncomponents(data)) @inline farray_size(data::IF{S, Ni}) where {S, Ni} = (Ni, ncomponents(data)) @inline farray_size(data::VF{S, Nv}) where {S, Nv} = (Nv, ncomponents(data)) @inline farray_size(data::VIJFH{S, Nv, Nij}) where {S, Nv, Nij} = (Nv, Nij, Nij, ncomponents(data), get_Nh_dynamic(data)) +@inline farray_size(data::VIJHF{S, Nv, Nij}) where {S, Nv, Nij} = (Nv, Nij, Nij, get_Nh_dynamic(data), ncomponents(data)) @inline farray_size(data::VIFH{S, Nv, Ni}) where {S, Nv, Ni} = (Nv, Ni, ncomponents(data), get_Nh_dynamic(data)) +@inline farray_size(data::VIHF{S, Nv, Ni}) where {S, Nv, Ni} = (Nv, Ni, get_Nh_dynamic(data), ncomponents(data)) # Keep in sync with definition(s) in libs. @inline slab_index(i::T, j::T) where {T} = CartesianIndex(i, j, T(1), T(1), T(1)) @@ -1445,13 +1878,17 @@ type parameters. # Equivalent to: # @generated parent_array_type(::Type{A}) where {A <: AbstractData} = Tuple(A.parameters)[end] @inline parent_array_type(::Type{IFH{S, Ni, A}}) where {S, Ni, A} = A +@inline parent_array_type(::Type{IHF{S, Ni, A}}) where {S, Ni, A} = A @inline parent_array_type(::Type{DataF{S, A}}) where {S, A} = A @inline parent_array_type(::Type{IJF{S, Nij, A}}) where {S, Nij, A} = A @inline parent_array_type(::Type{IF{S, Ni, A}}) where {S, Ni, A} = A @inline parent_array_type(::Type{VF{S, Nv, A}}) where {S, Nv, A} = A @inline parent_array_type(::Type{VIJFH{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = A +@inline parent_array_type(::Type{VIJHF{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = A @inline parent_array_type(::Type{VIFH{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = A +@inline parent_array_type(::Type{VIHF{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = A @inline parent_array_type(::Type{IJFH{S, Nij, A}}) where {S, Nij, A} = A +@inline parent_array_type(::Type{IJHF{S, Nij, A}}) where {S, Nij, A} = A @inline parent_array_type(::Type{IH1JH2{S, Nij, A}}) where {S, Nij, A} = A @inline parent_array_type(::Type{IV1JH2{S, n1, Ni, A}}) where {S, n1, Ni, A} = A @inline parent_array_type(::Type{IJKFVH{S, Nij, Nk, Nv, A}}) where {S, Nij, Nk, Nv, A} = A @@ -1463,7 +1900,7 @@ Base.ndims(::Type{T}) where {T <: AbstractData} = Base.ndims(parent_array_type(T)) @inline function Base.getindex( - data::Union{IJF, IJFH, IFH, VIJFH, VIFH, VF, IF}, + data::Union{IJF, IJFH, IJHF, IFH, IHF, VIJFH, VIJHF, VIFH, VIHF, VF, IF}, I::CartesianIndex, ) @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) @@ -1477,7 +1914,7 @@ Base.ndims(::Type{T}) where {T <: AbstractData} = end @inline function Base.setindex!( - data::Union{IJF, IJFH, IFH, VIJFH, VIFH, VF, IF}, + data::Union{IJF, IJFH, IJHF, IFH, IHF, VIJFH, VIJHF, VIFH, VIHF, VF, IF}, val, I::CartesianIndex, ) @@ -1524,7 +1961,20 @@ The universal index order is `CartesianIndex(i, j, f, v, h)`, see see the notation in [`DataLayouts`](@ref) for more information. """ @inline function getindex_field( - data::Union{DataF, IJF, IJFH, IFH, VIJFH, VIFH, VF, IF}, + data::Union{ + DataF, + IJF, + IJFH, + IJHF, + IFH, + IHF, + VIJFH, + VIJHF, + VIFH, + VIHF, + VF, + IF, + }, I::CartesianIndex, # universal index ) @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) @@ -1544,7 +1994,20 @@ The universal index order is `CartesianIndex(i, j, f, v, h)`, see see the notation in [`DataLayouts`](@ref) for more information. """ @inline function setindex_field!( - data::Union{DataF, IJF, IJFH, IFH, VIJFH, VIFH, VF, IF}, + data::Union{ + DataF, + IJF, + IJFH, + IJHF, + IFH, + IHF, + VIJFH, + VIJHF, + VIFH, + VIHF, + VF, + IF, + }, val::Real, I::CartesianIndex, # universal index ) @@ -1556,12 +2019,27 @@ see the notation in [`DataLayouts`](@ref) for more information. ) 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, IFH, VIJFH, VIFH, VF, IF}, + data::Union{ + IJF, + IJFH, + IJHF, + IFH, + IHF, + VIJFH, + VIJHF, + VIFH, + VIHF, + VF, + IF, + }, I::Vararg{Int, N}, ) where {N} = Base.getindex( data, @@ -1569,7 +2047,19 @@ if VERSION ≥ v"1.11.0-beta" ) @inline Base.setindex!( - data::Union{IJF, IJFH, IFH, VIJFH, VIFH, VF, IF}, + data::Union{ + IJF, + IJFH, + IJHF, + IFH, + IHF, + VIJFH, + VIJHF, + VIFH, + VIHF, + VF, + IF, + }, val, I::Vararg{Int, N}, ) where {N} = Base.setindex!( @@ -1594,6 +2084,28 @@ if VERSION ≥ v"1.11.0-beta" @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) + s = singleton(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) + s = singleton(data) + @inbounds set_struct_linear!(parent(data), convert(S, val), I, s_array) + end + end """ @@ -1609,10 +2121,17 @@ Also, this assumes that `eltype(data) <: Real`. """ function data2array end -data2array(data::Union{IF, IFH}) = reshape(parent(data), :) -data2array(data::Union{IJF, IJFH}) = reshape(parent(data), :) -data2array(data::Union{VF{S, Nv}, VIFH{S, Nv}, VIJFH{S, Nv}}) where {S, Nv} = - reshape(parent(data), Nv, :) +data2array(data::Union{IF, IFH, IHF}) = reshape(parent(data), :) +data2array(data::Union{IJF, IJFH, IJHF}) = reshape(parent(data), :) +data2array( + data::Union{ + VF{S, Nv}, + VIFH{S, Nv}, + VIHF{S, Nv}, + VIJFH{S, Nv}, + VIJHF{S, Nv}, + }, +) where {S, Nv} = reshape(parent(data), Nv, :) """ array2data(array, ::AbstractData) @@ -1643,32 +2162,45 @@ device_dispatch(x::MArray) = ToCPU() @inline singleton(@nospecialize(::IJKFVH)) = IJKFVHSingleton() @inline singleton(@nospecialize(::IJFH)) = IJFHSingleton() +@inline singleton(@nospecialize(::IJHF)) = IJHFSingleton() @inline singleton(@nospecialize(::IFH)) = IFHSingleton() +@inline singleton(@nospecialize(::IHF)) = IHFSingleton() @inline singleton(@nospecialize(::DataF)) = DataFSingleton() @inline singleton(@nospecialize(::IJF)) = IJFSingleton() @inline singleton(@nospecialize(::IF)) = IFSingleton() @inline singleton(@nospecialize(::VF)) = VFSingleton() @inline singleton(@nospecialize(::VIJFH)) = VIJFHSingleton() +@inline singleton(@nospecialize(::VIJHF)) = VIJHFSingleton() @inline singleton(@nospecialize(::VIFH)) = VIFHSingleton() +@inline singleton(@nospecialize(::VIHF)) = VIHFSingleton() @inline singleton(@nospecialize(::IH1JH2)) = IH1JH2Singleton() @inline singleton(@nospecialize(::IV1JH2)) = IV1JH2Singleton() @inline singleton(::Type{IJKFVH}) = IJKFVHSingleton() @inline singleton(::Type{IJFH}) = IJFHSingleton() +@inline singleton(::Type{IJHF}) = IJHFSingleton() @inline singleton(::Type{IFH}) = IFHSingleton() +@inline singleton(::Type{IHF}) = IHFSingleton() @inline singleton(::Type{DataF}) = DataFSingleton() @inline singleton(::Type{IJF}) = IJFSingleton() @inline singleton(::Type{IF}) = IFSingleton() @inline singleton(::Type{VF}) = VFSingleton() @inline singleton(::Type{VIJFH}) = VIJFHSingleton() +@inline singleton(::Type{VIJHF}) = VIJHFSingleton() @inline singleton(::Type{VIFH}) = VIFHSingleton() +@inline singleton(::Type{VIHF}) = VIHFSingleton() @inline singleton(::Type{IH1JH2}) = IH1JH2Singleton() @inline singleton(::Type{IV1JH2}) = IV1JH2Singleton() +include("non_extruded_broadcasted.jl") +include("has_uniform_datalayouts.jl") + include("copyto.jl") include("fused_copyto.jl") include("fill.jl") include("mapreduce.jl") +include("struct_linear_indexing.jl") + end # module diff --git a/src/DataLayouts/broadcast.jl b/src/DataLayouts/broadcast.jl index ca7d921e83..6fc2e1b13e 100644 --- a/src/DataLayouts/broadcast.jl +++ b/src/DataLayouts/broadcast.jl @@ -26,6 +26,9 @@ abstract type Data1DStyle{Ni} <: DataStyle end struct IFHStyle{Ni, A} <: Data1DStyle{Ni} end DataStyle(::Type{IFH{S, Ni, A}}) where {S, Ni, A} = IFHStyle{Ni, parent_array_type(A)}() +struct IHFStyle{Ni, A} <: Data1DStyle{Ni} end +DataStyle(::Type{IHF{S, Ni, A}}) where {S, Ni, A} = + IHFStyle{Ni, parent_array_type(A)}() abstract type DataSlab1DStyle{Ni} <: DataStyle end DataSlab1DStyle(::Type{IFHStyle{Ni, A}}) where {Ni, A} = IFStyle{Ni, A} @@ -45,6 +48,11 @@ DataStyle(::Type{IJFH{S, Nij, A}}) where {S, Nij, A} = IJFHStyle{Nij, parent_array_type(A)}() DataSlab2DStyle(::Type{IJFHStyle{Nij, A}}) where {Nij, A} = IJFStyle{Nij, A} +struct IJHFStyle{Nij, A} <: Data2DStyle{Nij} end +DataStyle(::Type{IJHF{S, Nij, A}}) where {S, Nij, A} = + IJHFStyle{Nij, parent_array_type(A)}() +DataSlab2DStyle(::Type{IJHFStyle{Nij, A}}) where {Nij, A} = IJFStyle{Nij, A} + abstract type Data1DXStyle{Nv, Ni} <: DataStyle end struct VIFHStyle{Nv, Ni, A} <: Data1DXStyle{Nv, Ni} end DataStyle(::Type{VIFH{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = @@ -54,6 +62,14 @@ Data1DXStyle(::Type{VIFHStyle{Nv, Ni, A}}) where {Ni, Nv, A} = DataColumnStyle(::Type{VIFHStyle{Nv, Ni, A}}) where {Ni, Nv, A} = VFStyle{Nv, A} DataSlab1DStyle(::Type{VIFHStyle{Nv, Ni, A}}) where {Ni, Nv, A} = IFStyle{Ni, A} +struct VIHFStyle{Nv, Ni, A} <: Data1DXStyle{Nv, Ni} end +DataStyle(::Type{VIHF{S, Nv, Ni, A}}) where {S, Nv, Ni, A} = + VIHFStyle{Nv, Ni, parent_array_type(A)}() +Data1DXStyle(::Type{VIHFStyle{Nv, Ni, A}}) where {Ni, Nv, A} = + VIHFStyle{Nv, Ni, A} +DataColumnStyle(::Type{VIHFStyle{Nv, Ni, A}}) where {Ni, Nv, A} = VFStyle{Nv, A} +DataSlab1DStyle(::Type{VIHFStyle{Nv, Ni, A}}) where {Ni, Nv, A} = IFStyle{Ni, A} + abstract type Data2DXStyle{Nv, Nij} <: DataStyle end struct VIJFHStyle{Nv, Nij, A} <: Data2DXStyle{Nv, Nij} end DataStyle(::Type{VIJFH{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = @@ -65,19 +81,33 @@ DataColumnStyle(::Type{VIJFHStyle{Nv, Nij, A}}) where {Nv, Nij, A} = DataSlab2DStyle(::Type{VIJFHStyle{Nv, Nij, A}}) where {Nv, Nij, A} = IJFStyle{Nij, A} +struct VIJHFStyle{Nv, Nij, A} <: Data2DXStyle{Nv, Nij} end +DataStyle(::Type{VIJHF{S, Nv, Nij, A}}) where {S, Nv, Nij, A} = + VIJHFStyle{Nv, Nij, parent_array_type(A)}() +Data2DXStyle(::Type{VIJHFStyle{Nv, Nij, A}}) where {Nv, Nij, A} = + VIJHFStyle{Nv, Nij, A} +DataColumnStyle(::Type{VIJHFStyle{Nv, Nij, A}}) where {Nv, Nij, A} = + VFStyle{Nv, A} +DataSlab2DStyle(::Type{VIJHFStyle{Nv, Nij, A}}) where {Nv, Nij, A} = + IJFStyle{Nij, A} + ##### ##### Union styles ##### #! format: off const BroadcastedUnionIJFH{S, Nij, A} = Union{Base.Broadcast.Broadcasted{IJFHStyle{Nij, A}}, IJFH{S, Nij, A}} +const BroadcastedUnionIJHF{S, Nij, A} = Union{Base.Broadcast.Broadcasted{IJHFStyle{Nij, A}}, IJHF{S, Nij, A}} const BroadcastedUnionIFH{S, Ni, A} = Union{Base.Broadcast.Broadcasted{IFHStyle{Ni, A}}, IFH{S, Ni, A}} -const BroadcastedUnionIJF{S, Nij, A} = Union{Base.Broadcast.Broadcasted{IJFStyle{Nij, A}}, IJF{S, Nij, A}} -const BroadcastedUnionIF{S, Ni, A} = Union{Base.Broadcast.Broadcasted{IFStyle{Ni, A}}, IF{S, Ni, A}} +const BroadcastedUnionIHF{S, Ni, A} = Union{Base.Broadcast.Broadcasted{IHFStyle{Ni, A}}, IHF{S, Ni, A}} +const BroadcastedUnionIJF{S, Nij, A} = Union{Base.Broadcast.Broadcasted{IJFStyle{Nij, A}}, IJF{S, Nij, A}} +const BroadcastedUnionIF{S, Ni, A} = Union{Base.Broadcast.Broadcasted{IFStyle{Ni, A}}, IF{S, Ni, A}} const BroadcastedUnionVIFH{S, Nv, Ni, A} = Union{Base.Broadcast.Broadcasted{VIFHStyle{Nv, Ni, A}}, VIFH{S, Nv, Ni, A}} +const BroadcastedUnionVIHF{S, Nv, Ni, A} = Union{Base.Broadcast.Broadcasted{VIHFStyle{Nv, Ni, A}}, VIHF{S, Nv, Ni, A}} const BroadcastedUnionVIJFH{S, Nv, Nij, A} = Union{Base.Broadcast.Broadcasted{VIJFHStyle{Nv, Nij, A}}, VIJFH{S, Nv, Nij, A}} -const BroadcastedUnionVF{S, Nv, A} = Union{Base.Broadcast.Broadcasted{VFStyle{Nv, A}}, VF{S, Nv, A}} -const BroadcastedUnionDataF{S, A} = Union{Base.Broadcast.Broadcasted{DataFStyle{A}}, DataF{S, A}} +const BroadcastedUnionVIJHF{S, Nv, Nij, A} = Union{Base.Broadcast.Broadcasted{VIJHFStyle{Nv, Nij, A}}, VIJHF{S, Nv, Nij, A}} +const BroadcastedUnionVF{S, Nv, A} = Union{Base.Broadcast.Broadcasted{VFStyle{Nv, A}}, VF{S, Nv, A}} +const BroadcastedUnionDataF{S, A} = Union{Base.Broadcast.Broadcasted{DataFStyle{A}}, DataF{S, A}} #! format: on abstract type Data3DStyle <: DataStyle end @@ -111,11 +141,20 @@ Base.Broadcast.BroadcastStyle( ::IFHStyle{Ni, A1}, ::IFHStyle{Ni, A2}, ) where {Ni, A1, A2} = IFHStyle{Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::IHFStyle{Ni, A1}, + ::IHFStyle{Ni, A2}, +) where {Ni, A1, A2} = IHFStyle{Ni, promote_parent_array_type(A1, A2)}() Base.Broadcast.BroadcastStyle( ::VIFHStyle{Nv, Ni, A1}, ::VIFHStyle{Nv, Ni, A2}, ) where {Nv, Ni, A1, A2} = VIFHStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VIHFStyle{Nv, Ni, A1}, + ::VIHFStyle{Nv, Ni, A2}, +) where {Nv, Ni, A1, A2} = + VIHFStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() Base.Broadcast.BroadcastStyle( ::IJFStyle{Nij, A1}, ::IJFStyle{Nij, A2}, @@ -124,11 +163,20 @@ Base.Broadcast.BroadcastStyle( ::IJFHStyle{Nij, A1}, ::IJFHStyle{Nij, A2}, ) where {Nij, A1, A2} = IJFHStyle{Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::IJHFStyle{Nij, A1}, + ::IJHFStyle{Nij, A2}, +) where {Nij, A1, A2} = IJHFStyle{Nij, promote_parent_array_type(A1, A2)}() Base.Broadcast.BroadcastStyle( ::VIJFHStyle{Nv, Nij, A1}, ::VIJFHStyle{Nv, Nij, A2}, ) where {Nv, Nij, A1, A2} = VIJFHStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VIJHFStyle{Nv, Nij, A1}, + ::VIJHFStyle{Nv, Nij, A2}, +) where {Nv, Nij, A1, A2} = + VIJHFStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() Base.Broadcast.BroadcastStyle( ::DataFStyle{A1}, @@ -150,59 +198,117 @@ Base.Broadcast.BroadcastStyle( ::IFHStyle{Ni, A2}, ) where {Ni, A1, A2} = IFHStyle{Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::DataFStyle{A1}, + ::IHFStyle{Ni, A2}, +) where {Ni, A1, A2} = IHFStyle{Ni, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::DataFStyle{A1}, ::IJFHStyle{Nij, A2}, ) where {Nij, A1, A2} = IJFHStyle{Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::DataFStyle{A1}, + ::IJHFStyle{Nij, A2}, +) where {Nij, A1, A2} = IJHFStyle{Nij, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::DataFStyle{A1}, ::VIFHStyle{Nv, Ni, A2}, ) where {Nv, Ni, A1, A2} = VIFHStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::DataFStyle{A1}, + ::VIHFStyle{Nv, Ni, A2}, +) where {Nv, Ni, A1, A2} = + VIHFStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::DataFStyle{A1}, ::VIJFHStyle{Nv, Nij, A2}, ) where {Nv, Nij, A1, A2} = VIJFHStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::DataFStyle{A1}, + ::VIJHFStyle{Nv, Nij, A2}, +) where {Nv, Nij, A1, A2} = + VIJHFStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::VFStyle{Nv, A1}, ::IFHStyle{Ni, A2}, ) where {Nv, Ni, A1, A2} = VIFHStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VFStyle{Nv, A1}, + ::IHFStyle{Ni, A2}, +) where {Nv, Ni, A1, A2} = + VIHFStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::VFStyle{Nv, A1}, ::IJFHStyle{Nij, A2}, ) where {Nv, Nij, A1, A2} = VIJFHStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VFStyle{Nv, A1}, + ::IJHFStyle{Nij, A2}, +) where {Nv, Nij, A1, A2} = + VIJHFStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::VFStyle{Nv, A1}, ::VIFHStyle{Nv, Ni, A2}, ) where {Nv, Ni, A1, A2} = VIFHStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VFStyle{Nv, A1}, + ::VIHFStyle{Nv, Ni, A2}, +) where {Nv, Ni, A1, A2} = + VIHFStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::VFStyle{Nv, A1}, ::VIJFHStyle{Nv, Nij, A2}, ) where {Nv, Nij, A1, A2} = VIJFHStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::VFStyle{Nv, A1}, + ::VIJHFStyle{Nv, Nij, A2}, +) where {Nv, Nij, A1, A2} = + VIJHFStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::IFHStyle{Ni, A1}, ::VIFHStyle{Nv, Ni, A2}, ) where {Nv, Ni, A1, A2} = VIFHStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::IFHStyle{Ni, A1}, + ::VIHFStyle{Nv, Ni, A2}, +) where {Nv, Ni, A1, A2} = + VIHFStyle{Nv, Ni, promote_parent_array_type(A1, A2)}() + Base.Broadcast.BroadcastStyle( ::IJFHStyle{Nij, A1}, ::VIJFHStyle{Nv, Nij, A2}, ) where {Nv, Nij, A1, A2} = VIJFHStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() +Base.Broadcast.BroadcastStyle( + ::IJHFStyle{Nij, A1}, + ::VIJHFStyle{Nv, Nij, A2}, +) where {Nv, Nij, A1, A2} = + VIJHFStyle{Nv, Nij, promote_parent_array_type(A1, A2)}() + Base.Broadcast.broadcastable(data::AbstractData) = data Base.@propagate_inbounds function slab( @@ -302,6 +408,16 @@ function Base.similar( return IJFH{Eltype, Nij}(array) end +function Base.similar( + bc::BroadcastedUnionIJHF{<:Any, Nij, A}, + ::Type{Eltype}, + (_, _, _, _, Nh) = size(bc), +) where {Nij, A, Eltype} + PA = parent_array_type(A) + array = similar(PA, (Nij, Nij, Nh, typesize(eltype(A), Eltype))) + return IJHF{Eltype, Nij}(array) +end + function Base.similar( bc::BroadcastedUnionIFH{<:Any, Ni, A}, ::Type{Eltype}, @@ -312,6 +428,16 @@ function Base.similar( return IFH{Eltype, Ni}(array) end +function Base.similar( + bc::BroadcastedUnionIHF{<:Any, Ni, A}, + ::Type{Eltype}, + (_, _, _, _, Nh) = size(bc), +) where {Ni, A, Eltype} + PA = parent_array_type(A) + array = similar(PA, (Ni, Nh, typesize(eltype(A), Eltype))) + return IHF{Eltype, Ni}(array) +end + function Base.similar( ::BroadcastedUnionIJF{<:Any, Nij, A}, ::Type{Eltype}, @@ -346,7 +472,7 @@ function Base.similar( end Base.similar( - bc::BroadcastedUnionVIFH{<:Any, Nv}, + bc::Union{BroadcastedUnionVIFH{<:Any, Nv}, BroadcastedUnionVIHF{<:Any, Nv}}, ::Type{Eltype}, ) where {Nv, Eltype} = Base.similar(bc, Eltype, Val(Nv)) @@ -361,11 +487,27 @@ function Base.similar( return VIFH{Eltype, newNv, Ni}(array) end +function Base.similar( + bc::BroadcastedUnionVIHF{<:Any, Nv, Ni, A}, + ::Type{Eltype}, + ::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))) + return VIHF{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)) +Base.similar( + bc::BroadcastedUnionVIJHF{<: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}, @@ -377,6 +519,17 @@ function Base.similar( return VIJFH{Eltype, newNv, Nij}(array) end +function Base.similar( + bc::BroadcastedUnionVIJHF{<:Any, Nv, Nij, A}, + ::Type{Eltype}, + ::Val{newNv}, +) where {Nv, Nij, A, Eltype, newNv} + (_, _, _, _, 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) +end + # ============= FusedMultiBroadcast isascalar( diff --git a/src/DataLayouts/copyto.jl b/src/DataLayouts/copyto.jl index ec82a918f1..01c5ee2dc2 100644 --- a/src/DataLayouts/copyto.jl +++ b/src/DataLayouts/copyto.jl @@ -1,11 +1,35 @@ ##### ##### Dispatching and edge cases ##### - -Base.copyto!( - dest::AbstractData, - @nospecialize(bc::Union{AbstractData, Base.Broadcast.Broadcasted}), -) = Base.copyto!(dest, bc, device_dispatch(parent(dest))) +if VERSION ≥ v"1.11.0-beta" + # + function Base.copyto!( + dest::AbstractData{S}, + bc::Union{AbstractData, Base.Broadcast.Broadcasted}, + ) where {S} + return Base.copyto!(dest, bc, device_dispatch(parent(dest))) + end +else + function Base.copyto!( + dest::AbstractData{S}, + bc::Union{AbstractData, Base.Broadcast.Broadcasted}, + ) where {S} + dev = device_dispatch(parent(dest)) + if dev isa ToCPU && + has_uniform_datalayouts(bc) && + dest isa EndsWithField && + !(dest isa DataF) + # Specialize on linear indexing when possible: + bc′ = Base.Broadcast.instantiate(to_non_extruded_broadcasted(bc)) + @inbounds @simd for I in 1:get_N(UniversalSize(dest)) + dest[I] = convert(S, bc′[I]) + end + else + Base.copyto!(dest, bc, device_dispatch(parent(dest))) + end + return dest + end +end # Specialize on non-Broadcasted objects function Base.copyto!(dest::D, src::D) where {D <: AbstractData} @@ -44,8 +68,8 @@ function Base.copyto!( end function Base.copyto!( - dest::IJFH{S, Nij}, - bc::BroadcastedUnionIJFH{S, Nij}, + dest::Union{IJFH{S, Nij}, IJHF{S, Nij}}, + bc::Union{BroadcastedUnionIJFH{S, Nij}, BroadcastedUnionIJHF{S, Nij}}, ::ToCPU, ) where {S, Nij} (_, _, _, _, Nh) = size(dest) @@ -57,8 +81,8 @@ function Base.copyto!( end function Base.copyto!( - dest::IFH{S, Ni}, - bc::BroadcastedUnionIFH{S, Ni}, + dest::Union{IFH{S, Ni}, IHF{S, Ni}}, + bc::Union{BroadcastedUnionIFH{S, Ni}, BroadcastedUnionIHF{S, Ni}}, ::ToCPU, ) where {S, Ni} (_, _, _, _, Nh) = size(dest) @@ -121,8 +145,8 @@ function Base.copyto!( end function Base.copyto!( - dest::VIFH{S, Nv, Ni}, - bc::BroadcastedUnionVIFH{S, Nv, Ni}, + dest::Union{VIFH{S, Nv, Ni}, VIHF{S, Nv, Ni}}, + bc::Union{BroadcastedUnionVIFH{S, Nv, Ni}, BroadcastedUnionVIHF{S, Nv, Ni}}, ::ToCPU, ) where {S, Nv, Ni} # copy contiguous columns @@ -135,8 +159,11 @@ function Base.copyto!( end function Base.copyto!( - dest::VIJFH{S, Nv, Nij}, - bc::BroadcastedUnionVIJFH{S, Nv, Nij}, + dest::Union{VIJFH{S, Nv, Nij}, VIJHF{S, Nv, Nij}}, + bc::Union{ + BroadcastedUnionVIJFH{S, Nv, Nij}, + BroadcastedUnionVIJHF{S, Nv, Nij}, + }, ::ToCPU, ) where {S, Nv, Nij} # copy contiguous columns diff --git a/src/DataLayouts/fill.jl b/src/DataLayouts/fill.jl index c165e61b2a..4dfd1da735 100644 --- a/src/DataLayouts/fill.jl +++ b/src/DataLayouts/fill.jl @@ -1,3 +1,10 @@ +function Base.fill!(dest::EndsWithField, val, ::ToCPU) + @inbounds @simd for I in 1:get_N(UniversalSize(dest)) + dest[I] = val + end + return dest +end + function Base.fill!(data::IJFH, val, ::ToCPU) (_, _, _, _, Nh) = size(data) @inbounds for h in 1:Nh @@ -19,27 +26,6 @@ function Base.fill!(data::DataF, val, ::ToCPU) return data end -function Base.fill!(data::IJF{S, Nij}, val, ::ToCPU) where {S, Nij} - @inbounds for j in 1:Nij, i in 1:Nij - data[CartesianIndex(i, j, 1, 1, 1)] = val - end - return data -end - -function Base.fill!(data::IF{S, Ni}, val, ::ToCPU) where {S, Ni} - @inbounds for i in 1:Ni - data[CartesianIndex(i, 1, 1, 1, 1)] = val - end - return data -end - -function Base.fill!(data::VF, val, ::ToCPU) - Nv = nlevels(data) - @inbounds for v in 1:Nv - data[CartesianIndex(1, 1, 1, v, 1)] = val - end - return data -end function Base.fill!(data::VIJFH, val, ::ToCPU) (Ni, Nj, _, Nv, Nh) = size(data) diff --git a/src/DataLayouts/fused_copyto.jl b/src/DataLayouts/fused_copyto.jl index 93010cc888..f397342897 100644 --- a/src/DataLayouts/fused_copyto.jl +++ b/src/DataLayouts/fused_copyto.jl @@ -23,7 +23,7 @@ end function fused_copyto!( fmbc::FusedMultiBroadcast, - dest1::VIJFH{S1, Nv1, Nij}, + dest1::Union{VIJFH{S1, Nv1, Nij}, VIJHF{S1, Nv1, Nij}}, ::ToCPU, ) where {S1, Nv1, Nij} for (dest, bc) in fmbc.pairs @@ -40,7 +40,7 @@ end function fused_copyto!( fmbc::FusedMultiBroadcast, - dest1::IJFH{S, Nij}, + dest1::Union{IJFH{S, Nij}, IJHF{S, Nij}}, ::ToCPU, ) where {S, Nij} # copy contiguous columns @@ -58,7 +58,7 @@ end function fused_copyto!( fmbc::FusedMultiBroadcast, - dest1::VIFH{S, Nv1, Ni}, + dest1::Union{VIFH{S, Nv1, Ni}, VIHF{S, Nv1, Ni}}, ::ToCPU, ) where {S, Nv1, Ni} # copy contiguous columns diff --git a/src/DataLayouts/has_uniform_datalayouts.jl b/src/DataLayouts/has_uniform_datalayouts.jl new file mode 100644 index 0000000000..b3f56e2aaa --- /dev/null +++ b/src/DataLayouts/has_uniform_datalayouts.jl @@ -0,0 +1,74 @@ +@inline function first_datalayout_in_bc(args::Tuple, rargs...) + x1 = first_datalayout_in_bc(args[1], rargs...) + x1 isa AbstractData && return x1 + return first_datalayout_in_bc(Base.tail(args), rargs...) +end + +@inline first_datalayout_in_bc(args::Tuple{Any}, rargs...) = + first_datalayout_in_bc(args[1], rargs...) +@inline first_datalayout_in_bc(args::Tuple{}, rargs...) = nothing +@inline first_datalayout_in_bc(x) = nothing +@inline first_datalayout_in_bc(x::AbstractData) = x + +@inline first_datalayout_in_bc(bc::Base.Broadcast.Broadcasted) = + first_datalayout_in_bc(bc.args) + +@inline _has_uniform_datalayouts_args(truesofar, start, args::Tuple, rargs...) = + truesofar && + _has_uniform_datalayouts(truesofar, start, args[1], rargs...) && + _has_uniform_datalayouts_args(truesofar, start, Base.tail(args), rargs...) + +@inline _has_uniform_datalayouts_args( + truesofar, + start, + args::Tuple{Any}, + rargs..., +) = truesofar && _has_uniform_datalayouts(truesofar, start, args[1], rargs...) +@inline _has_uniform_datalayouts_args(truesofar, _, args::Tuple{}, rargs...) = + truesofar + +@inline function _has_uniform_datalayouts( + truesofar, + start, + bc::Base.Broadcast.Broadcasted, +) + return truesofar && _has_uniform_datalayouts_args(truesofar, start, bc.args) +end +for DL in ( + :IJKFVH, + :IJFH, + :IJHF, + :IFH, + :IHF, + :DataF, + :IJF, + :IF, + :VF, + :VIJFH, + :VIJHF, + :VIFH, + :VIHF, +) + @eval begin + @inline _has_uniform_datalayouts(truesofar, ::$(DL), ::$(DL)) = true + end +end +@inline _has_uniform_datalayouts(truesofar, _, x::AbstractData) = false +@inline _has_uniform_datalayouts(truesofar, _, x) = truesofar + +""" + has_uniform_datalayouts +Find the first datalayout in the broadcast expression (BCE), +and compares against every other datalayout in the BCE. Returns + - `true` if the broadcasted object has only a single kind of datalayout (e.g. VF,VF, VIJFH,VIJFH) + - `false` if the broadcasted object has multiple kinds of datalayouts (e.g. VIJFH, VIFH) +Note: a broadcasted object can have different _types_, + e.g., `VIFJH{Float64}` and `VIFJH{Tuple{Float64,Float64}}` + but not different kinds, e.g., `VIFJH{Float64}` and `VF{Float64}`. +""" +function has_uniform_datalayouts end + +@inline has_uniform_datalayouts(bc::Base.Broadcast.Broadcasted) = + _has_uniform_datalayouts_args(true, first_datalayout_in_bc(bc), bc.args) + +@inline has_uniform_datalayouts(bc::AbstractData) = true diff --git a/src/DataLayouts/mapreduce.jl b/src/DataLayouts/mapreduce.jl index a67f86941f..648a23aa76 100644 --- a/src/DataLayouts/mapreduce.jl +++ b/src/DataLayouts/mapreduce.jl @@ -15,7 +15,10 @@ end function Base.mapreduce( fn::F, op::Op, - bc::BroadcastedUnionIJFH{<:Any, Nij, A}, + bc::Union{ + BroadcastedUnionIJFH{<:Any, Nij, A}, + BroadcastedUnionIJHF{<:Any, Nij, A}, + }, ) where {F, Op, Nij, A} # mapreduce across DataSlab2D (_, _, _, _, Nh) = size(bc) @@ -29,7 +32,10 @@ end function Base.mapreduce( fn::F, op::Op, - bc::BroadcastedUnionIFH{<:Any, Ni, A}, + bc::Union{ + BroadcastedUnionIFH{<:Any, Ni, A}, + BroadcastedUnionIHF{<:Any, Ni, A}, + }, ) where {F, Op, Ni, A} # mapreduce across DataSlab1D (_, _, _, _, Nh) = size(bc) @@ -77,7 +83,10 @@ end function Base.mapreduce( fn::F, op::Op, - bc::BroadcastedUnionVIFH{<:Any, Nv, Ni, A}, + bc::Union{ + BroadcastedUnionVIFH{<:Any, Nv, Ni, A}, + BroadcastedUnionVIHF{<:Any, Nv, Ni, A}, + }, ) where {F, Op, Nv, Ni, A} # mapreduce across columns (_, _, _, _, Nh) = size(bc) @@ -91,7 +100,10 @@ end function Base.mapreduce( fn::F, op::Op, - bc::BroadcastedUnionVIJFH{<:Any, Nv, Nij, A}, + bc::Union{ + BroadcastedUnionVIJFH{<:Any, Nv, Nij, A}, + BroadcastedUnionVIJHF{<:Any, Nv, Nij, A}, + }, ) where {F, Op, Nv, Nij, A} # mapreduce across columns (_, _, _, _, Nh) = size(bc) diff --git a/src/DataLayouts/non_extruded_broadcasted.jl b/src/DataLayouts/non_extruded_broadcasted.jl new file mode 100644 index 0000000000..ca40f698ed --- /dev/null +++ b/src/DataLayouts/non_extruded_broadcasted.jl @@ -0,0 +1,129 @@ +#! format: off +# ============================================================ Adapted from Base.Broadcast (julia version 1.10.4) +import Base.Broadcast: BroadcastStyle +struct NonExtrudedBroadcasted{ + Style <: Union{Nothing, BroadcastStyle}, + Axes, + F, + Args <: Tuple, +} <: Base.AbstractBroadcasted + style::Style + f::F + args::Args + axes::Axes # the axes of the resulting object (may be bigger than implied by `args` if this is nested inside a larger `NonExtrudedBroadcasted`) + + NonExtrudedBroadcasted(style::Union{Nothing, BroadcastStyle}, f::Tuple, args::Tuple) = + error() # disambiguation: tuple is not callable + function NonExtrudedBroadcasted( + style::Union{Nothing, BroadcastStyle}, + f::F, + args::Tuple, + axes = nothing, + ) where {F} + # using Core.Typeof rather than F preserves inferrability when f is a type + return new{typeof(style), typeof(axes), Core.Typeof(f), typeof(args)}( + style, + f, + args, + axes, + ) + end + function NonExtrudedBroadcasted(f::F, args::Tuple, axes = nothing) where {F} + NonExtrudedBroadcasted(combine_styles(args...)::BroadcastStyle, f, args, axes) + end + function NonExtrudedBroadcasted{Style}(f::F, args, axes = nothing) where {Style, F} + return new{Style, typeof(axes), Core.Typeof(f), typeof(args)}( + Style()::Style, + f, + args, + axes, + ) + end + function NonExtrudedBroadcasted{Style, Axes, F, Args}( + f, + args, + axes, + ) where {Style, Axes, F, Args} + return new{Style, Axes, F, Args}(Style()::Style, f, args, axes) + end +end + +@inline to_non_extruded_broadcasted(bc::Base.Broadcast.Broadcasted) = + NonExtrudedBroadcasted(bc.style, bc.f, to_non_extruded_broadcasted(bc.args), bc.axes) +@inline to_non_extruded_broadcasted(x) = x +NonExtrudedBroadcasted(bc::Base.Broadcast.Broadcasted) = to_non_extruded_broadcasted(bc) + +@inline to_non_extruded_broadcasted(args::Tuple) = ( + to_non_extruded_broadcasted(args[1]), + to_non_extruded_broadcasted(Base.tail(args))..., +) +@inline to_non_extruded_broadcasted(args::Tuple{Any}) = + (to_non_extruded_broadcasted(args[1]),) +@inline to_non_extruded_broadcasted(args::Tuple{}) = () + +@inline _checkbounds(bc, _, I) = nothing # TODO: fix this case +@inline _checkbounds(bc, ::Tuple, I) = Base.checkbounds(bc, I) +@inline function Base.getindex( + bc::NonExtrudedBroadcasted, + I::Union{Integer, CartesianIndex}, +) + @boundscheck _checkbounds(bc, axes(bc), I) # is this really the only issue? + @inbounds _broadcast_getindex(bc, I) +end + +# --- here, we define our own bounds checks +@inline function Base.checkbounds(bc::NonExtrudedBroadcasted, I::Integer) + # Base.checkbounds_indices(Bool, axes(bc), (I,)) || Base.throw_boundserror(bc, (I,)) # from Base + Base.checkbounds_indices(Bool, (Base.OneTo(n_dofs(bc)),), (I,)) || Base.throw_boundserror(bc, (I,)) +end + +import StaticArrays +to_tuple(t::Tuple) = t +to_tuple(t::NTuple{N, <: Base.OneTo}) where {N} = map(x->x.stop, t) +to_tuple(t::NTuple{N, <: StaticArrays.SOneTo}) where {N} = map(x->x.stop, t) +n_dofs(bc::NonExtrudedBroadcasted) = prod(to_tuple(axes(bc))) +# --- + +Base.@propagate_inbounds _broadcast_getindex( + A::Union{Ref, AbstractArray{<:Any, 0}, Number}, + I::Integer, +) = A[] # Scalar-likes can just ignore all indices +Base.@propagate_inbounds _broadcast_getindex( + ::Ref{Type{T}}, + I::Integer, +) where {T} = T +# Tuples are statically known to be singleton or vector-like +Base.@propagate_inbounds _broadcast_getindex(A::Tuple{Any}, I::Integer) = A[1] +Base.@propagate_inbounds _broadcast_getindex(A::Tuple, I::Integer) = A[I[1]] +# Everything else falls back to dynamically dropping broadcasted indices based upon its axes +# Base.@propagate_inbounds _broadcast_getindex(A, I) = A[newindex(A, I)] +Base.@propagate_inbounds _broadcast_getindex(A, I::Integer) = A[I] +Base.@propagate_inbounds function _broadcast_getindex( + bc::NonExtrudedBroadcasted{<:Any, <:Any, <:Any, <:Any}, + I::Integer, +) + args = _getindex(bc.args, I) + return _broadcast_getindex_evalf(bc.f, args...) +end +@inline _broadcast_getindex_evalf(f::Tf, args::Vararg{Any, N}) where {Tf, N} = + f(args...) # not propagate_inbounds +Base.@propagate_inbounds _getindex(args::Tuple, I) = + (_broadcast_getindex(args[1], I), _getindex(Base.tail(args), I)...) +Base.@propagate_inbounds _getindex(args::Tuple{Any}, I) = + (_broadcast_getindex(args[1], I),) +Base.@propagate_inbounds _getindex(args::Tuple{}, I) = () + +@inline Base.axes(bc::NonExtrudedBroadcasted) = _axes(bc, bc.axes) +_axes(::NonExtrudedBroadcasted, axes::Tuple) = axes +@inline _axes(bc::NonExtrudedBroadcasted, ::Nothing) = Base.Broadcast.combine_axes(bc.args...) +_axes(bc::NonExtrudedBroadcasted{<:Base.Broadcast.AbstractArrayStyle{0}}, ::Nothing) = () +@inline Base.axes(bc::NonExtrudedBroadcasted{<:Any, <:NTuple{N}}, d::Integer) where {N} = + d <= N ? axes(bc)[d] : OneTo(1) +Base.IndexStyle(::Type{<:NonExtrudedBroadcasted{<:Any, <:Tuple{Any}}}) = IndexLinear() +@inline _axes(::NonExtrudedBroadcasted, axes) = axes +@inline Base.eltype(bc::NonExtrudedBroadcasted) = Base.Broadcast.combine_axes(bc.args...) + + +# ============================================================ + +#! format: on diff --git a/src/DataLayouts/struct_linear_indexing.jl b/src/DataLayouts/struct_linear_indexing.jl new file mode 100644 index 0000000000..a24e736d33 --- /dev/null +++ b/src/DataLayouts/struct_linear_indexing.jl @@ -0,0 +1,130 @@ +##### +##### Linear indexing +##### + +abstract type _Size end +struct DynamicSize <: _Size end +struct StaticSize{S_array, FD} <: _Size + function StaticSize{S, FD}() where {S, FD} + new{S::Tuple{Vararg{Int}}, FD}() + end +end + +Base.@pure StaticSize(s::Tuple{Vararg{Int}}, FD) = StaticSize{s, FD}() + +# Some @pure convenience functions for `StaticSize` +s_field_dim_1(::Type{StaticSize{S, FD}}) where {S, FD} = + ntuple(j -> j == FD ? 1 : S[j], length(S)) +s_field_dim_1(::StaticSize{S, FD}) where {S, FD} = + ntuple(j -> j == FD ? 1 : S[j], length(S)) + +Base.@pure get(::Type{StaticSize{S}}) where {S} = S +Base.@pure get(::StaticSize{S}) where {S} = S +Base.@pure Base.getindex(::StaticSize{S}, i::Int) where {S} = + i <= length(S) ? S[i] : 1 +Base.@pure Base.ndims(::StaticSize{S}) where {S} = length(S) +Base.@pure Base.ndims(::Type{StaticSize{S}}) where {S} = length(S) +Base.@pure Base.length(::StaticSize{S}) where {S} = prod(S) + +Base.@propagate_inbounds cart_inds(n::NTuple) = + @inbounds CartesianIndices(map(x -> Base.OneTo(x), n)) +Base.@propagate_inbounds linear_inds(n::NTuple) = + @inbounds LinearIndices(map(x -> Base.OneTo(x), n)) + +@inline function offset_index_linear( + base_index::Integer, + field_offset, + array_size, +) + @inbounds begin + li = base_index + prod(array_size[1:(end - 1)]) * field_offset + end + return li +end + +Base.@propagate_inbounds @generated function get_struct_linear( + array::AbstractArray{T}, + ::Type{S}, + start_index::Integer, + array_size, + base_index = start_index, +) where {T, S} + tup = :(()) + for i in 1:fieldcount(S) + push!( + tup.args, + :(get_struct_linear( + array, + fieldtype(S, $i), + offset_index_linear( + base_index, + $(fieldtypeoffset(T, S, Val(i))), + array_size, + ), + array_size, + base_index, + )), + ) + 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, + base_index = start_index, +) 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, + base_index = start_index, +) 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( + base_index, + $(fieldtypeoffset(T, S, Val(i))), + array_size, + ), + array_size, + base_index, + )), + ) + 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, + base_index = start_index, +) where {S} + @inbounds array[start_index] = val + val +end diff --git a/src/Grids/spectralelement.jl b/src/Grids/spectralelement.jl index c304b2e86f..abfb1d732a 100644 --- a/src/Grids/spectralelement.jl +++ b/src/Grids/spectralelement.jl @@ -29,30 +29,34 @@ local_geometry_type( # this means that if data is saved in two different files, reloading will give fields which live on the same grid function SpectralElementGrid1D( topology::Topologies.IntervalTopology, - quadrature_style::Quadratures.QuadratureStyle, + quadrature_style::Quadratures.QuadratureStyle; + HorizontalLayout = DataLayouts.IFH, ) get!( Cache.OBJECT_CACHE, (SpectralElementGrid1D, topology, quadrature_style), ) do - _SpectralElementGrid1D(topology, quadrature_style) + _SpectralElementGrid1D(topology, quadrature_style, HorizontalLayout) end end _SpectralElementGrid1D( topology::Topologies.IntervalTopology, quadrature_style::Quadratures.QuadratureStyle, + HorizontalLayout = DataLayouts.IFH, ) = _SpectralElementGrid1D( topology, quadrature_style, Val(Topologies.nlocalelems(topology)), + HorizontalLayout, ) function _SpectralElementGrid1D( topology::Topologies.IntervalTopology, quadrature_style::Quadratures.QuadratureStyle, ::Val{Nh}, -) where {Nh} + ::Type{HorizontalLayout}, +) where {Nh, HorizontalLayout} global_geometry = Geometry.CartesianGlobalGeometry() CoordType = Topologies.coordinate_type(topology) AIdx = Geometry.coordinate_axis(CoordType) @@ -60,7 +64,7 @@ function _SpectralElementGrid1D( Nq = Quadratures.degrees_of_freedom(quadrature_style) LG = Geometry.LocalGeometry{AIdx, CoordType, FT, SMatrix{1, 1, FT, 1}} - local_geometry = DataLayouts.IFH{LG, Nq}(Array{FT}, Nh) + local_geometry = HorizontalLayout{LG, Nq}(Array{FT}, Nh) quad_points, quad_weights = Quadratures.quadrature_points(FT, quadrature_style) @@ -137,7 +141,7 @@ local_geometry_type( ) where {T, Q, GG, LG, D, IS, BS} = eltype(LG) # calls eltype from DataLayouts """ - SpectralElementSpace2D(topology, quadrature_style; enable_bubble) + SpectralElementSpace2D(topology, quadrature_style; enable_bubble, HorizontalLayout = DataLayouts.IJFH) Construct a `SpectralElementSpace2D` instance given a `topology` and `quadrature`. The flag `enable_bubble` enables the `bubble correction` for more accurate element areas. @@ -146,6 +150,7 @@ flag `enable_bubble` enables the `bubble correction` for more accurate element a - topology: Topology2D - quadrature_style: QuadratureStyle - enable_bubble: Bool +- HorizontalLayout: Type{<:AbstractData} The idea behind the so-called `bubble_correction` is that the numerical area of the domain (e.g., the sphere) is given by the sum of nodal integration weights @@ -171,13 +176,25 @@ Note: This is accurate only for cubed-spheres of the [`Meshes.EquiangularCubedSp function SpectralElementGrid2D( topology::Topologies.Topology2D, quadrature_style::Quadratures.QuadratureStyle; + HorizontalLayout = DataLayouts.IJFH, enable_bubble::Bool = false, ) get!( Cache.OBJECT_CACHE, - (SpectralElementGrid2D, topology, quadrature_style, enable_bubble), + ( + SpectralElementGrid2D, + topology, + quadrature_style, + enable_bubble, + HorizontalLayout, + ), ) do - _SpectralElementGrid2D(topology, quadrature_style; enable_bubble) + _SpectralElementGrid2D( + topology, + quadrature_style, + HorizontalLayout; + enable_bubble, + ) end end @@ -193,22 +210,32 @@ end _SpectralElementGrid2D( topology::Topologies.Topology2D, - quadrature_style::Quadratures.QuadratureStyle; + quadrature_style::Quadratures.QuadratureStyle, + HorizontalLayout = DataLayouts.IJFH; enable_bubble::Bool, ) = _SpectralElementGrid2D( topology, quadrature_style, - Val(Topologies.nlocalelems(topology)); + Val(Topologies.nlocalelems(topology)), + HorizontalLayout; enable_bubble, ) function _SpectralElementGrid2D( topology::Topologies.Topology2D, quadrature_style::Quadratures.QuadratureStyle, - ::Val{Nh}; + ::Val{Nh}, + ::Type{HorizontalLayout}; enable_bubble::Bool, -) where {Nh} - +) where {Nh, HorizontalLayout} + @assert HorizontalLayout <: Union{DataLayouts.IJHF, DataLayouts.IJFH} + SurfaceLayout = if HorizontalLayout <: DataLayouts.IJFH + DataLayouts.IFH + elseif HorizontalLayout <: DataLayouts.IJHF + DataLayouts.IHF + else + error("Uncaught case") + end # 1. compute localgeom for local elememts # 2. ghost exchange of localgeom # 3. do a round of dss on WJs @@ -241,7 +268,7 @@ function _SpectralElementGrid2D( LG = Geometry.LocalGeometry{AIdx, CoordType2D, FT, SMatrix{2, 2, FT, 4}} - local_geometry = DataLayouts.IJFH{LG, Nq}(Array{FT}, Nh) + local_geometry = HorizontalLayout{LG, Nq}(Array{FT}, Nh) quad_points, quad_weights = Quadratures.quadrature_points(FT, quadrature_style) @@ -399,7 +426,7 @@ function _SpectralElementGrid2D( if quadrature_style isa Quadratures.GLL internal_surface_geometry = - DataLayouts.IFH{SG, Nq}(Array{FT}, length(interior_faces)) + SurfaceLayout{SG, Nq}(Array{FT}, length(interior_faces)) for (iface, (lidx⁻, face⁻, lidx⁺, face⁺, reversed)) in enumerate(interior_faces) internal_surface_geometry_slab = @@ -438,7 +465,7 @@ function _SpectralElementGrid2D( boundary_faces = Topologies.boundary_faces(topology, boundarytag) boundary_surface_geometry = - DataLayouts.IFH{SG, Nq}(Array{FT}, length(boundary_faces)) + SurfaceLayout{SG, Nq}(Array{FT}, length(boundary_faces)) for (iface, (elem, face)) in enumerate(boundary_faces) boundary_surface_geometry_slab = slab(boundary_surface_geometry, iface) diff --git a/src/Spaces/dss.jl b/src/Spaces/dss.jl index 607a1487c7..21420f8c1d 100644 --- a/src/Spaces/dss.jl +++ b/src/Spaces/dss.jl @@ -10,7 +10,11 @@ import ..Topologies: dss_local_ghost!, dss_ghost!, fill_send_buffer!, - load_from_recv_buffer! + load_from_recv_buffer!, + DSSTypesAll, + DSSDataTypes, + DSSPerimeterTypes, + DSSWeightTypes perimeter(space::AbstractSpectralElementSpace) = Topologies.Perimeter2D( @@ -27,7 +31,12 @@ perimeter(space::AbstractSpectralElementSpace) = Topologies.Perimeter2D( Creates a [`DSSBuffer`](@ref) for the field data corresponding to `data` """ function create_dss_buffer( - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::Union{ + DataLayouts.IJFH, + DataLayouts.IJHF, + DataLayouts.VIJFH, + DataLayouts.VIJHF, + }, hspace::SpectralElementSpace2D, ) create_dss_buffer( @@ -39,7 +48,12 @@ function create_dss_buffer( end function create_dss_buffer( - data::Union{DataLayouts.IFH, DataLayouts.VIFH}, + data::Union{ + DataLayouts.IFH, + DataLayouts.IHF, + DataLayouts.VIFH, + DataLayouts.VIHF, + }, hspace::SpectralElementSpace1D, ) nothing @@ -49,9 +63,13 @@ end function weighted_dss!( data::Union{ DataLayouts.IFH, + DataLayouts.IHF, DataLayouts.VIFH, + DataLayouts.VIHF, DataLayouts.IJFH, + DataLayouts.IJHF, DataLayouts.VIJFH, + DataLayouts.VIJHF, }, space::Union{ AbstractSpectralElementSpace, @@ -71,12 +89,7 @@ It comprises of the following steps: 3). [`Spaces.weighted_dss_ghost!`](@ref) """ function weighted_dss!( - data::Union{ - DataLayouts.IFH, - DataLayouts.VIFH, - DataLayouts.IJFH, - DataLayouts.VIJFH, - }, + data::DSSTypesAll, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, dss_buffer::Union{DSSBuffer, Nothing}, ) @@ -92,7 +105,7 @@ function weighted_dss_prepare!(data, space, dss_buffer::Nothing) end function weighted_dss_prepare!( - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, space::Union{ Spaces.SpectralElementSpace2D, Spaces.ExtrudedFiniteDifferenceSpace, @@ -128,9 +141,13 @@ cuda_synchronize(device::ClimaComms.AbstractDevice; kwargs...) = nothing weighted_dss_start!( data::Union{ DataLayouts.IFH, + DataLayouts.IHF, DataLayouts.VIFH, + DataLayouts.VIHF, DataLayouts.IJFH, + DataLayouts.IJHF, DataLayouts.VIJFH, + DataLayouts.VIJHF, }, space::Union{ AbstractSpectralElementSpace, @@ -155,7 +172,7 @@ representative ghost vertices which store result of "ghost local" DSS are loaded 4). Start DSS communication with neighboring processes """ function weighted_dss_start!( - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, space::Union{ Spaces.SpectralElementSpace2D, Spaces.ExtrudedFiniteDifferenceSpace, @@ -176,9 +193,13 @@ weighted_dss_start!(data, space, dss_buffer::Nothing) = nothing weighted_dss_internal!( data::Union{ DataLayouts.IFH, + DataLayouts.IHF, DataLayouts.VIFH, + DataLayouts.VIHF, DataLayouts.IJFH, + DataLayouts.IJHF, DataLayouts.VIJFH, + DataLayouts.VIJHF, }, space::Union{ AbstractSpectralElementSpace, @@ -195,24 +216,14 @@ and perimeter elements to facilitate overlapping of communication with computati 3). [`Spaces.dss_local!`](@ref) computes the weighted DSS on local vertices and faces. """ weighted_dss_internal!( - data::Union{ - DataLayouts.IFH, - DataLayouts.VIFH, - DataLayouts.IJFH, - DataLayouts.VIJFH, - }, + data::DSSTypesAll, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, dss_buffer::Union{DSSBuffer, Nothing}, ) = weighted_dss_internal!(data, space, horizontal_space(space), dss_buffer) function weighted_dss_internal!( - data::Union{ - DataLayouts.IFH, - DataLayouts.VIFH, - DataLayouts.IJFH, - DataLayouts.VIJFH, - }, + data::DSSTypesAll, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, hspace::AbstractSpectralElementSpace, dss_buffer::Union{DSSBuffer, Nothing}, @@ -260,9 +271,13 @@ end weighted_dss_ghost!( data::Union{ DataLayouts.IFH, + DataLayouts.IHF, DataLayouts.VIFH, + DataLayouts.VIHF, DataLayouts.IJFH, + DataLayouts.IJHF, DataLayouts.VIJFH, + DataLayouts.VIJHF, }, space::Union{ AbstractSpectralElementSpace, @@ -283,12 +298,7 @@ This transforms the DSS'd local vectors back to Covariant12 vectors, and copies `perimeter_data` to `data`. """ weighted_dss_ghost!( - data::Union{ - DataLayouts.IFH, - DataLayouts.VIFH, - DataLayouts.IJFH, - DataLayouts.VIJFH, - }, + data::DSSTypesAll, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, dss_buffer::Union{DSSBuffer, Nothing}, ) = weighted_dss_ghost!(data, space, horizontal_space(space), dss_buffer) @@ -296,7 +306,7 @@ weighted_dss_ghost!( function weighted_dss_ghost!( - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, hspace::SpectralElementSpace2D, dss_buffer::DSSBuffer, diff --git a/src/Spaces/spectralelement.jl b/src/Spaces/spectralelement.jl index cbc0990bbf..a1b02e424d 100644 --- a/src/Spaces/spectralelement.jl +++ b/src/Spaces/spectralelement.jl @@ -47,6 +47,12 @@ end # 1D """ SpectralElementSpace1D(grid::SpectralElementGrid1D) + SpectralElementSpace1D( + topology::Topologies.IntervalTopology, + quadrature_style::Quadratures.QuadratureStyle; + kwargs... + ) + """ struct SpectralElementSpace1D{G} <: AbstractSpectralElementSpace grid::G @@ -63,15 +69,21 @@ local_geometry_type(::Type{SpectralElementSpace1D{G}}) where {G} = function SpectralElementSpace1D( topology::Topologies.IntervalTopology, - quadrature_style::Quadratures.QuadratureStyle, + quadrature_style::Quadratures.QuadratureStyle; + kwargs..., ) - grid = Grids.SpectralElementGrid1D(topology, quadrature_style) + grid = Grids.SpectralElementGrid1D(topology, quadrature_style; kwargs...) SpectralElementSpace1D(grid) end # 2D """ SpectralElementSpace2D(grid::SpectralElementGrid1D) + SpectralElementSpace2D( + topology::Topologies.Topology2D, + quadrature_style::Quadratures.QuadratureStyle; + kwargs..., + ) """ struct SpectralElementSpace2D{G} <: AbstractSpectralElementSpace grid::G diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl index 46cffb62df..a78b7090be 100644 --- a/src/Topologies/dss.jl +++ b/src/Topologies/dss.jl @@ -1,6 +1,25 @@ using DocStringExtensions using .DataLayouts: CartesianFieldIndex +const DSSTypesAll = Union{ + DataLayouts.IFH, + DataLayouts.IHF, + DataLayouts.VIFH, + DataLayouts.VIHF, + DataLayouts.IJFH, + DataLayouts.IJHF, + DataLayouts.VIJFH, + DataLayouts.VIJHF, +} +const DSSDataTypes = Union{ + DataLayouts.IJFH, + DataLayouts.IJHF, + DataLayouts.VIJFH, + DataLayouts.VIJHF, +} +const DSSPerimeterTypes = Union{DataLayouts.VIFH, DataLayouts.VIHF} +const DSSWeightTypes = Union{DataLayouts.IJFH, DataLayouts.IJHF} + """ DSSBuffer{G, D, A, B} @@ -11,7 +30,7 @@ struct DSSBuffer{S, G, D, A, B, VI} "ClimaComms graph context for communication" graph_context::G """ - Perimeter `DataLayout` object: typically a `VIFH{TT,Nv,Np,Nh}`, where `TT` is the + Perimeter `DataLayout` object: typically a `VIFH{TT,Nv,Np,Nh}` or `VIHF{TT,Nv,Np,Nh}`, where `TT` is the transformed type, `Nv` is the number of vertical levels, and `Np` is the length of the perimeter """ perimeter_data::D @@ -31,20 +50,35 @@ end """ create_dss_buffer( - data::Union{DataLayouts.IJFH{S}, DataLayouts.VIJFH{S}}, + data::Union{DataLayouts.IJFH{S}, DataLayouts.IJHF{S}, DataLayouts.VIJFH{S}, DataLayouts.VIJHF{S}}, topology::Topology2D, - local_geometry = nothing, - local_weights = nothing, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF, Nothing} = nothing, + local_weights::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF, Nothing} = nothing, ) where {S} Creates a [`DSSBuffer`](@ref) for the field data corresponding to `data` """ +create_dss_buffer( + data::DSSDataTypes, + topology::Topology2D, + local_geometry::Union{DSSDataTypes, Nothing} = nothing, + local_weights::Union{DSSDataTypes, Nothing} = nothing, +) = create_dss_buffer( + data, + topology, + DataLayouts.VIFH, + local_geometry, + local_weights, +) + function create_dss_buffer( - data::Union{DataLayouts.IJFH{S}, DataLayouts.VIJFH{S}}, + data::DSSDataTypes, topology::Topology2D, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH, Nothing} = nothing, - local_weights::Union{DataLayouts.IJFH, DataLayouts.VIJFH, Nothing} = nothing, -) where {S} + ::Type{PerimeterLayout}, + local_geometry::Union{DSSDataTypes, Nothing} = nothing, + local_weights::Union{DSSDataTypes, Nothing} = nothing, +) where {PerimeterLayout} + S = eltype(data) Nij = DataLayouts.get_Nij(data) Nij_lg = isnothing(local_geometry) ? Nij : DataLayouts.get_Nij(local_geometry) @@ -70,7 +104,18 @@ function create_dss_buffer( else _transformed_type(data, local_geometry, local_weights, DA) # extract transformed type end - perimeter_data = DataLayouts.VIFH{TS, Nv, Np}(DA{T}(undef, Nv, Np, Nf, Nh)) + + perimeter_data = if !isnothing(local_geometry) + fdim = DataLayouts.field_dim(DataLayouts.singleton(local_geometry)) + if fdim == ndims(local_geometry) + DataLayouts.VIHF{TS, Nv, Np}(DA{T}(undef, Nv, Np, Nh, Nf)) + else + DataLayouts.VIFH{TS, Nv, Np}(DA{T}(undef, Nv, Np, Nf, Nh)) + end + else + PerimeterLayout{TS, Nv, Np}(DA{T}(undef, Nv, Np, Nf, Nh)) + end + if context isa ClimaComms.SingletonCommsContext graph_context = ClimaComms.SingletonGraphContext(context) send_data, recv_data = T[], T[] @@ -128,14 +173,14 @@ assert_same_eltype(::DataLayouts.AbstractData{S}, ::DSSBuffer{S}) where {S} = assert_same_eltype(::DataLayouts.AbstractData, ::Nothing) = nothing """ - function dss_transform!( + dss_transform!( device::ClimaComms.AbstractDevice, dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, - perimeter::AbstractPerimeter, - localelems::Vector{Int}, + data::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, + weight::Union{DataLayouts.IJFH, DataLayouts.IJHF}, + perimeter::Perimeter2D, + localelems::AbstractVector{Int}, ) Transforms vectors from Covariant axes to physical (local axis), weights the data at perimeter nodes, @@ -156,9 +201,9 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). function dss_transform!( device::ClimaComms.AbstractDevice, dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, + data::DSSDataTypes, + local_geometry::DSSDataTypes, + weight::DSSWeightTypes, perimeter::Perimeter2D, localelems::AbstractVector{Int}, ) @@ -200,8 +245,8 @@ to `UVVector` if `T <: UVVector`. """ function dss_transform!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::Union{DataLayouts.VIFH, DataLayouts.VIHF}, + data::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, perimeter::AbstractPerimeter, bc, localelems::Vector{Int}, @@ -222,11 +267,11 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). """ function dss_transform!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, perimeter::Perimeter2D{Nq}, local_geometry, - weight, + weight::DSSWeightTypes, localelems::Vector{Int}, ) where {Nq} (_, _, _, nlevels, _) = DataLayouts.universal_size(perimeter_data) @@ -254,8 +299,8 @@ end dss_untransform!( device::ClimaComms.AbstractDevice, dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, perimeter::AbstractPerimeter, ) @@ -276,8 +321,8 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). function dss_untransform!( device::ClimaComms.AbstractDevice, dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, + local_geometry::DSSDataTypes, perimeter::Perimeter2D, localelems::AbstractVector{Int}, ) @@ -297,8 +342,8 @@ end """ function dss_untransform!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::Union{DataLayouts.VIFH, DataLayouts.VIHF}, + data::Union{DataLayouts.IJFH, DataLayouts.IJHF, DataLayouts.VIJFH, DataLayouts.VIJHF}, local_geometry, localelems::Vector{Int}, ) @@ -317,9 +362,9 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). function dss_untransform!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter_data::DSSPerimeterTypes, + data::DSSDataTypes, + local_geometry::DSSDataTypes, perimeter::Perimeter2D, localelems::Vector{Int}, ) @@ -342,7 +387,7 @@ end function dss_load_perimeter_data!( ::ClimaComms.AbstractCPUDevice, dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, perimeter::Perimeter2D, ) (; perimeter_data) = dss_buffer @@ -359,7 +404,7 @@ end function dss_unload_perimeter_data!( ::ClimaComms.AbstractCPUDevice, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + data::DSSDataTypes, dss_buffer::DSSBuffer, perimeter::Perimeter2D, ) @@ -389,7 +434,7 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). """ function dss_local!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Perimeter2D, topology::Topology2D, ) @@ -408,7 +453,7 @@ end Apply dss to local vertices. """ function dss_local_vertices!( - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Perimeter2D, topology::Topology2D, ) @@ -438,7 +483,7 @@ function dss_local_vertices!( end function dss_local_faces!( - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::Perimeter2D, topology::Topology2D, ) @@ -479,7 +524,7 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). """ function dss_local_ghost!( ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::AbstractPerimeter, topology::AbstractTopology, ) @@ -536,7 +581,7 @@ Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). """ function dss_ghost!( device::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, + perimeter_data::DSSPerimeterTypes, perimeter::AbstractPerimeter, topology::AbstractTopology, ) @@ -629,10 +674,8 @@ end Computed unweighted/pure DSS of `data`. """ -function dss!( - data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, <:Any, Nij}}, - topology::Topology2D, -) where {S, Nij} +function dss!(data::DSSDataTypes, topology::Topology2D) + Nij = DataLayouts.get_Nij(data) length(parent(data)) == 0 && return nothing device = ClimaComms.device(topology) perimeter = Perimeter2D(Nij) diff --git a/test/DataLayouts/benchmark_copyto.jl b/test/DataLayouts/benchmark_copyto.jl index ead88942f8..415d7adbd0 100644 --- a/test/DataLayouts/benchmark_copyto.jl +++ b/test/DataLayouts/benchmark_copyto.jl @@ -54,9 +54,15 @@ end data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) benchmarkcopyto!(bm, device, data, 3) @test all(parent(data) .== 3) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + benchmarkcopyto!(bm, device, data, 3) + @test all(parent(data) .== 3) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) benchmarkcopyto!(bm, device, data, 3) @test all(parent(data) .== 3) + data = IHF{S}(ArrayType{FT}, zeros; Ni, 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}(ArrayType{FT}, zeros; Nij); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # data = IF{S}(ArrayType{FT}, zeros; Ni); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) @@ -66,9 +72,15 @@ end data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) benchmarkcopyto!(bm, device, data, 3) @test all(parent(data) .== 3) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + benchmarkcopyto!(bm, device, data, 3) + @test all(parent(data) .== 3) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) benchmarkcopyto!(bm, device, data, 3) @test all(parent(data) .== 3) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + benchmarkcopyto!(bm, device, data, 3) + @test all(parent(data) .== 3) # data = IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nh); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test # data = IH1JH2{S}(ArrayType{FT}, zeros; Nij,Nk,Nh); benchmarkcopyto!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test diff --git a/test/DataLayouts/benchmark_fill.jl b/test/DataLayouts/benchmark_fill.jl index 0ec876eb93..36bb3cc9b7 100644 --- a/test/DataLayouts/benchmark_fill.jl +++ b/test/DataLayouts/benchmark_fill.jl @@ -30,7 +30,7 @@ function benchmarkfill!(bm, device, data, val) kernel_time_s = t_min, nreps = nreps, caller, - problem_size = size(data), + problem_size = DataLayouts.array_size(data), n_reads_writes, ) end @@ -51,9 +51,15 @@ end data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) benchmarkfill!(bm, device, data, 3) @test all(parent(data) .== 3) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + benchmarkfill!(bm, device, data, 3) + @test all(parent(data) .== 3) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) benchmarkfill!(bm, device, data, 3) @test all(parent(data) .== 3) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + benchmarkfill!(bm, device, data, 3) + @test all(parent(data) .== 3) data = IJF{S}(ArrayType{FT}, zeros; Nij) benchmarkfill!(bm, device, data, 3) @test all(parent(data) .== 3) @@ -66,9 +72,15 @@ end data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) benchmarkfill!(bm, device, data, 3) @test all(parent(data) .== 3) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + benchmarkfill!(bm, device, data, 3) + @test all(parent(data) .== 3) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) benchmarkfill!(bm, device, data, 3) @test all(parent(data) .== 3) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + benchmarkfill!(bm, device, data, 3) + @test all(parent(data) .== 3) # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); benchmarkfill!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); benchmarkfill!(bm, device, data, 3); @test all(parent(data) .== 3) # TODO: test diff --git a/test/DataLayouts/opt_similar.jl b/test/DataLayouts/opt_similar.jl index 2ff0b053a3..0f1b7eb15a 100644 --- a/test/DataLayouts/opt_similar.jl +++ b/test/DataLayouts/opt_similar.jl @@ -39,8 +39,12 @@ end test_similar!(data) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_similar!(data) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_similar!(data) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_similar!(data) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_similar!(data) data = IJF{S}(ArrayType{FT}, zeros; Nij) test_similar!(data) data = IF{S}(ArrayType{FT}, zeros; Ni) @@ -49,8 +53,12 @@ end test_similar!(data) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_similar!(data) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_similar!(data) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_similar!(data) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_similar!(data) # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_similar!(data) # TODO: test # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_similar!(data) # TODO: test end diff --git a/test/DataLayouts/opt_universal_size.jl b/test/DataLayouts/opt_universal_size.jl index 27a5e69b0b..c4462a3970 100644 --- a/test/DataLayouts/opt_universal_size.jl +++ b/test/DataLayouts/opt_universal_size.jl @@ -61,8 +61,12 @@ end test_universal_size(data) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_universal_size(data) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_universal_size(data) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_universal_size(data) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_universal_size(data) data = IJF{S}(ArrayType{FT}, zeros; Nij) test_universal_size(data) data = IF{S}(ArrayType{FT}, zeros; Ni) @@ -71,8 +75,12 @@ end test_universal_size(data) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_universal_size(data) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_universal_size(data) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_universal_size(data) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_universal_size(data) # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_universal_size(data) # TODO: test # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_universal_size(data) # TODO: test end diff --git a/test/DataLayouts/unit_cartesian_field_index.jl b/test/DataLayouts/unit_cartesian_field_index.jl index 7ad99d3c10..a5c6156576 100644 --- a/test/DataLayouts/unit_cartesian_field_index.jl +++ b/test/DataLayouts/unit_cartesian_field_index.jl @@ -86,8 +86,12 @@ end 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) @@ -96,8 +100,12 @@ end 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 diff --git a/test/DataLayouts/unit_copyto.jl b/test/DataLayouts/unit_copyto.jl index 6bf260bb08..5759bf04b3 100644 --- a/test/DataLayouts/unit_copyto.jl +++ b/test/DataLayouts/unit_copyto.jl @@ -54,8 +54,12 @@ end 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) @@ -64,8 +68,12 @@ end 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 @@ -83,8 +91,12 @@ end 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) @@ -93,8 +105,12 @@ end 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 @@ -123,8 +139,12 @@ end # 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) @@ -133,8 +153,12 @@ end 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 diff --git a/test/DataLayouts/unit_data2array.jl b/test/DataLayouts/unit_data2array.jl index 8ce5084064..806226d8d4 100644 --- a/test/DataLayouts/unit_data2array.jl +++ b/test/DataLayouts/unit_data2array.jl @@ -29,6 +29,10 @@ end @test DataLayouts.data2array(data) == reshape(parent(data), :) @test is_data2array2data_identity(data) + data = IHF{FT}(ArrayType{FT}, rand; Ni, Nh) + @test DataLayouts.data2array(data) == reshape(parent(data), :) + @test is_data2array2data_identity(data) + data = IJF{FT}(ArrayType{FT}, rand; Nij) @test DataLayouts.data2array(data) == reshape(parent(data), :) @test is_data2array2data_identity(data) @@ -37,11 +41,23 @@ end @test DataLayouts.data2array(data) == reshape(parent(data), :) @test is_data2array2data_identity(data) + data = IJHF{FT}(ArrayType{FT}, rand; Nij, Nh) + @test DataLayouts.data2array(data) == reshape(parent(data), :) + @test is_data2array2data_identity(data) + data = VIFH{FT}(ArrayType{FT}, rand; Nv, Ni, Nh) @test DataLayouts.data2array(data) == reshape(parent(data), Nv, :) @test is_data2array2data_identity(data) + data = VIHF{FT}(ArrayType{FT}, rand; Nv, Ni, Nh) + @test DataLayouts.data2array(data) == reshape(parent(data), Nv, :) + @test is_data2array2data_identity(data) + data = VIJFH{FT}(ArrayType{FT}, rand; Nv, Nij, Nh) @test DataLayouts.data2array(data) == reshape(parent(data), Nv, :) @test is_data2array2data_identity(data) + + data = VIJHF{FT}(ArrayType{FT}, rand; Nv, Nij, Nh) + @test DataLayouts.data2array(data) == reshape(parent(data), Nv, :) + @test is_data2array2data_identity(data) end diff --git a/test/DataLayouts/unit_fill.jl b/test/DataLayouts/unit_fill.jl index 472a30a296..f439de5a22 100644 --- a/test/DataLayouts/unit_fill.jl +++ b/test/DataLayouts/unit_fill.jl @@ -31,8 +31,12 @@ end test_fill!(data, 3) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_fill!(data, 3) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_fill!(data, 3) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_fill!(data, 3) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_fill!(data, 3) data = IJF{S}(ArrayType{FT}, zeros; Nij) test_fill!(data, 3) data = IF{S}(ArrayType{FT}, zeros; Ni) @@ -41,8 +45,12 @@ end test_fill!(data, 3) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_fill!(data, 3) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_fill!(data, 3) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_fill!(data, 3) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_fill!(data, 3) # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_fill!(data, 3) # TODO: test # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_fill!(data, 3) # TODO: test @@ -62,8 +70,12 @@ end test_fill!(data, (2, 3)) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_fill!(data, (2, 3)) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_fill!(data, (2, 3)) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_fill!(data, (2, 3)) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_fill!(data, (2, 3)) data = IJF{S}(ArrayType{FT}, zeros; Nij) test_fill!(data, (2, 3)) data = IF{S}(ArrayType{FT}, zeros; Ni) @@ -72,8 +84,12 @@ end test_fill!(data, (2, 3)) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_fill!(data, (2, 3)) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_fill!(data, (2, 3)) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_fill!(data, (2, 3)) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_fill!(data, (2, 3)) # TODO: test this # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_fill!(data, (2,3)) # TODO: test @@ -104,8 +120,12 @@ end data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_fill!(data_view(data), (2, 3)) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_fill!(data_view(data), (2, 3)) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_fill!(data_view(data), (2, 3)) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_fill!(data_view(data), (2, 3)) data = IJF{S}(ArrayType{FT}, zeros; Nij) test_fill!(data_view(data), (2, 3)) data = IF{S}(ArrayType{FT}, zeros; Ni) @@ -114,8 +134,12 @@ end test_fill!(data_view(data), (2, 3)) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_fill!(data_view(data), (2, 3)) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_fill!(data_view(data), (2, 3)) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_fill!(data_view(data), (2, 3)) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_fill!(data_view(data), (2, 3)) # TODO: test this # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_fill!(data, (2,3)) # TODO: test @@ -136,7 +160,8 @@ end # a parent-similar array. data = data.:2 array₀ = DataLayouts.data2array(data) - @test typeof(array₀) <: Base.ReshapedArray + endswith_field = data isa Union{VIJHF, VIHF, IJHF, IHF} + endswith_field || @test typeof(array₀) <: Base.ReshapedArray rdata = DataLayouts.array2data(array₀, data) newdata = DataLayouts.rebuild( data, @@ -149,9 +174,9 @@ end ), ) rarray = parent(parent(newdata)) - @test typeof(rarray) <: Base.ReshapedArray + endswith_field || @test typeof(rarray) <: Base.ReshapedArray subarray = parent(rarray) - @test typeof(subarray) <: Base.SubArray + endswith_field || @test typeof(subarray) <: Base.SubArray array = parent(subarray) newdata end @@ -165,15 +190,23 @@ end data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_fill!(reshaped_array(data), 2) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_fill!(reshaped_array(data), 2) data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) test_fill!(reshaped_array(data), 2) + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + test_fill!(reshaped_array(data), 2) # data = IJF{S}(ArrayType{FT}, zeros; Nij); test_fill!(reshaped_array(data), 2) # data = IF{S}(ArrayType{FT}, zeros; Ni); test_fill!(reshaped_array(data), 2) # data = VF{S}(ArrayType{FT}, zeros; Nv); test_fill!(reshaped_array(data), 2) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_fill!(reshaped_array(data), 2) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_fill!(reshaped_array(data), 2) data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) test_fill!(reshaped_array(data), 2) + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + test_fill!(reshaped_array(data), 2) # TODO: test this # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_fill!(reshaped_array(data), 2) # TODO: test diff --git a/test/DataLayouts/unit_has_uniform_datalayouts.jl b/test/DataLayouts/unit_has_uniform_datalayouts.jl new file mode 100644 index 0000000000..ef21d3eab4 --- /dev/null +++ b/test/DataLayouts/unit_has_uniform_datalayouts.jl @@ -0,0 +1,46 @@ +#= +julia --project +using Revise; include(joinpath("test", "DataLayouts", "unit_has_uniform_datalayouts.jl")) +=# +using Test +using ClimaCore.DataLayouts +import ClimaCore.Geometry +import ClimaComms +import LazyBroadcast: @lazy +using StaticArrays +import Random +Random.seed!(1234) + +@testset "has_uniform_datalayouts" begin + device = ClimaComms.device() + ArrayType = ClimaComms.array_type(device) + FT = Float64 + S = FT + Nf = 1 + Nv = 4 + Ni = Nij = 3 + Nh = 5 + Nk = 6 + data_DataF = DataF{S}(ArrayType{FT}, zeros) + data_IJFH = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) + data_IFH = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) + data_IJF = IJF{S}(ArrayType{FT}, zeros; Nij) + data_IF = IF{S}(ArrayType{FT}, zeros; Ni) + data_VF = VF{S}(ArrayType{FT}, zeros; Nv) + data_VIJFH = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + data_VIFH = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + + bc = @lazy @. data_VIFH + data_VIFH + @test DataLayouts.has_uniform_datalayouts(bc) + bc = @lazy @. data_IJFH + data_VF + @test !DataLayouts.has_uniform_datalayouts(bc) + + data_VIJFHᶜ = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + data_VIJFHᶠ = VIJFH{S}(ArrayType{FT}, zeros; Nv = Nv + 1, Nij, Nh) + + # This is not a valid broadcast expression, + # but these two datalayouts can exist in a + # valid broadcast expression (e.g., interpolation). + bc = @lazy @. data_VIJFHᶜ + data_VIJFHᶠ + @test DataLayouts.has_uniform_datalayouts(bc) +end diff --git a/test/DataLayouts/unit_linear_indexing.jl b/test/DataLayouts/unit_linear_indexing.jl new file mode 100644 index 0000000000..f95272d255 --- /dev/null +++ b/test/DataLayouts/unit_linear_indexing.jl @@ -0,0 +1,171 @@ +#= +julia --check-bounds=yes --project +using Revise; include(joinpath("test", "DataLayouts", "unit_linear_indexing.jl")) +=# +using Test +using ClimaCore.DataLayouts +using ClimaCore.DataLayouts: get_struct_linear +import ClimaCore.Geometry +using StaticArrays +import Random +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)) + +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 = DataLayouts.StaticSize(s_array, 3) + @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, + ) +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 = DataLayouts.StaticSize(s_array, 3) + @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, + ) +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 = DataLayouts.StaticSize(s_array, 3) + @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, + ) +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 = DataLayouts.StaticSize(s_array, 5) + + @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, + ) +end + +# # TODO: add set_struct! diff --git a/test/DataLayouts/unit_mapreduce.jl b/test/DataLayouts/unit_mapreduce.jl index b1d69471ea..a0bc4d33f6 100644 --- a/test/DataLayouts/unit_mapreduce.jl +++ b/test/DataLayouts/unit_mapreduce.jl @@ -77,6 +77,8 @@ end test_mapreduce_1!(context, data) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_mapreduce_1!(context, data) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_mapreduce_1!(context, data) # data = IFH{S}(ArrayType{FT}, zeros; Ni,Nh); test_mapreduce_1!(context, data) # data = IJF{S}(ArrayType{FT}, zeros; Nij); test_mapreduce_1!(context, data) # data = IF{S}(ArrayType{FT}, zeros; Ni); test_mapreduce_1!(context, data) @@ -84,6 +86,8 @@ end test_mapreduce_1!(context, data) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_mapreduce_1!(context, data) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_mapreduce_1!(context, data) # data = VIFH{S}(ArrayType{FT}, zeros; Nv,Nij,Nh); test_mapreduce_1!(context, data) # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_mapreduce_1!(context, data) # TODO: test # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_mapreduce_1!(context, data) # TODO: test @@ -101,6 +105,8 @@ end test_mapreduce_2!(context, data) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_mapreduce_2!(context, data) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_mapreduce_2!(context, data) # data = IFH{S}(ArrayType{FT}, zeros; Ni,Nh); test_mapreduce_2!(context, data) # data = IJF{S}(ArrayType{FT}, zeros; Nij); test_mapreduce_2!(context, data) # data = IF{S}(ArrayType{FT}, zeros; Ni); test_mapreduce_2!(context, data) @@ -108,6 +114,8 @@ end test_mapreduce_2!(context, data) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_mapreduce_2!(context, data) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_mapreduce_2!(context, data) # data = VIFH{S}(ArrayType{FT}, zeros; Nv,Nij,Nh); test_mapreduce_2!(context, data) # TODO: test this # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_mapreduce_2!(context, data) # TODO: test @@ -138,6 +146,8 @@ end test_mapreduce_2!(context, data_view(data)) data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) test_mapreduce_2!(context, data_view(data)) + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + test_mapreduce_2!(context, data_view(data)) # data = IFH{S}(ArrayType{FT}, zeros; Ni,Nh); test_mapreduce_2!(context, data_view(data)) # data = IJF{S}(ArrayType{FT}, zeros; Nij); test_mapreduce_2!(context, data_view(data)) # data = IF{S}(ArrayType{FT}, zeros; Ni); test_mapreduce_2!(context, data_view(data)) @@ -145,6 +155,8 @@ end test_mapreduce_2!(context, data_view(data)) data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) test_mapreduce_2!(context, data_view(data)) + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_mapreduce_2!(context, data_view(data)) # data = VIFH{S}(ArrayType{FT}, zeros; Nv,Nij,Nh); test_mapreduce_2!(context, data_view(data)) # TODO: test this # data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij,Nk,Nv,Nh); test_mapreduce_2!(context, data_view(data)) # TODO: test diff --git a/test/DataLayouts/unit_ndims.jl b/test/DataLayouts/unit_ndims.jl index 50114d514d..1255839bd5 100644 --- a/test/DataLayouts/unit_ndims.jl +++ b/test/DataLayouts/unit_ndims.jl @@ -29,18 +29,30 @@ ClimaComms.@import_required_backends data = IFH{S}(ArrayType{FT}, zeros; Ni, Nh) @test ndims(data) == 3 @test ndims(typeof(data)) == 3 + data = IHF{S}(ArrayType{FT}, zeros; Ni, Nh) + @test ndims(data) == 3 + @test ndims(typeof(data)) == 3 data = IJF{S}(ArrayType{FT}, zeros; Nij) @test ndims(data) == 3 @test ndims(typeof(data)) == 3 data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) @test ndims(data) == 4 @test ndims(typeof(data)) == 4 + data = IJHF{S}(ArrayType{FT}, zeros; Nij, Nh) + @test ndims(data) == 4 + @test ndims(typeof(data)) == 4 data = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) @test ndims(data) == 4 @test ndims(typeof(data)) == 4 + data = VIHF{S}(ArrayType{FT}, zeros; Nv, Ni, Nh) + @test ndims(data) == 4 + @test ndims(typeof(data)) == 4 data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) @test ndims(data) == 5 @test ndims(typeof(data)) == 5 + data = VIJHF{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + @test ndims(data) == 5 + @test ndims(typeof(data)) == 5 data = DataLayouts.IJKFVH{S}(ArrayType{FT}, zeros; Nij, Nk, Nv, Nh) @test ndims(data) == 6 @test ndims(typeof(data)) == 6 diff --git a/test/Fields/unit_field.jl b/test/Fields/unit_field.jl index 1aa74e24c0..9ffa40f1ed 100644 --- a/test/Fields/unit_field.jl +++ b/test/Fields/unit_field.jl @@ -12,6 +12,7 @@ using OrderedCollections using StaticArrays, IntervalSets import ClimaCore import ClimaCore.Utilities: PlusHalf +import ClimaCore.DataLayouts import ClimaCore.DataLayouts: IJFH import ClimaCore: Fields, @@ -852,6 +853,39 @@ end nothing end +@testset "HF fields" begin + FT = Float32 + device = ClimaComms.device() + context = ClimaComms.context(device) + radius = FT(128) + zelem = 63 + zlim = (0, 1) + vertdomain = Domains.IntervalDomain( + Geometry.ZPoint{FT}(zlim[1]), + Geometry.ZPoint{FT}(zlim[2]); + boundary_names = (:bottom, :top), + ) + vertmesh = Meshes.IntervalMesh(vertdomain, nelems = zelem) + vtopology = Topologies.IntervalTopology(context, vertmesh) + vspace = Spaces.CenterFiniteDifferenceSpace(vtopology) + + hdomain = Domains.SphereDomain(radius) + helem = 4 + hmesh = Meshes.EquiangularCubedSphere(hdomain, helem) + htopology = Topologies.Topology2D(context, hmesh) + Nq = 4 + quad = Quadratures.GLL{Nq}() + hspace = Spaces.SpectralElementSpace2D( + htopology, + quad; + HorizontalLayout = DataLayouts.IJHF, + ) + cspace = Spaces.ExtrudedFiniteDifferenceSpace(hspace, vspace) + + f = fill(FT(0), cspace) + @. f += 1 +end + include("unit_field_multi_broadcast_fusion.jl") nothing diff --git a/test/runtests.jl b/test/runtests.jl index 9b5f61fd58..b1b405f830 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,6 +12,8 @@ UnitTest("DataLayouts ndims" ,"DataLayouts/unit_ndims.jl") UnitTest("DataLayouts array<->data" ,"DataLayouts/unit_data2array.jl"), UnitTest("DataLayouts get_struct" ,"DataLayouts/unit_struct.jl"), UnitTest("DataLayouts get/set_index_field" ,"DataLayouts/unit_cartesian_field_index.jl"), +UnitTest("DataLayouts has_uniform_datalayouts" ,"DataLayouts/unit_has_uniform_datalayouts.jl"), +UnitTest("DataLayouts linear indexing" ,"DataLayouts/unit_linear_indexing.jl"), UnitTest("Recursive" ,"RecursiveApply/unit_recursive_apply.jl"), UnitTest("PlusHalf" ,"Utilities/unit_plushalf.jl"), UnitTest("DataLayouts 0D" ,"DataLayouts/data0d.jl"),