diff --git a/ext/cuda/topologies_dss.jl b/ext/cuda/topologies_dss.jl index 0d7952fb82..c84eb39c60 100644 --- a/ext/cuda/topologies_dss.jl +++ b/ext/cuda/topologies_dss.jl @@ -1,5 +1,5 @@ import ClimaCore: DataLayouts, Topologies, Spaces, Fields -import ClimaCore.DataLayouts: getindex_field, setindex_field! +import ClimaCore.DataLayouts: CartesianFieldIndex using CUDA import ClimaCore.Topologies import ClimaCore.Topologies: perimeter_vertex_node_index @@ -44,13 +44,13 @@ function dss_load_perimeter_data_kernel!( (nperimeter, _, _, nlevels, nelems) = size(perimeter_data) nfidx = DataLayouts.ncomponents(perimeter_data) sizep = (nlevels, nperimeter, nfidx, nelems) # assume VIFH order - CI = CartesianIndex + CI = CartesianFieldIndex if gidx ≤ prod(sizep) (level, p, fidx, elem) = cart_ind(sizep, gidx).I (ip, jp) = perimeter[p] - val = getindex_field(data, CI(ip, jp, fidx, level, elem)) - setindex_field!(perimeter_data, val, CI(p, 1, fidx, level, elem)) + perimeter_data[CI(p, 1, fidx, level, elem)] = + data[CI(ip, jp, fidx, level, elem)] end return nothing end @@ -84,13 +84,13 @@ function dss_unload_perimeter_data_kernel!( (nperimeter, _, _, nlevels, nelems) = size(perimeter_data) nfidx = DataLayouts.ncomponents(perimeter_data) sizep = (nlevels, nperimeter, nfidx, nelems) # assume VIFH order - CI = CartesianIndex + CI = CartesianFieldIndex if gidx ≤ prod(sizep) (level, p, fidx, elem) = cart_ind(sizep, gidx).I (ip, jp) = perimeter[p] - val = getindex_field(perimeter_data, CI(p, 1, fidx, level, elem)) - setindex_field!(data, val, CI(ip, jp, fidx, level, elem)) + data[CI(ip, jp, fidx, level, elem)] = + perimeter_data[CI(p, 1, fidx, level, elem)] end return nothing end @@ -139,7 +139,7 @@ function dss_local_kernel!( nlocalfaces = length(interior_faces) (nperimeter, _, _, nlevels, _) = size(perimeter_data) nfidx = DataLayouts.ncomponents(perimeter_data) - CI = CartesianIndex + CI = CartesianFieldIndex if gidx ≤ nlevels * nfidx * nlocalvertices # local vertices sizev = (nlevels, nfidx, nlocalvertices) (level, fidx, vertexid) = cart_ind(sizev, gidx).I @@ -149,17 +149,12 @@ function dss_local_kernel!( for idx in st:(en - 1) (lidx, vert) = local_vertices[idx] ip = perimeter_vertex_node_index(vert) - sum_data += - getindex_field(perimeter_data, CI(ip, 1, fidx, level, lidx)) + sum_data += perimeter_data[CI(ip, 1, fidx, level, lidx)] end for idx in st:(en - 1) (lidx, vert) = local_vertices[idx] ip = perimeter_vertex_node_index(vert) - setindex_field!( - perimeter_data, - sum_data, - CI(ip, 1, fidx, level, lidx), - ) + perimeter_data[CI(ip, 1, fidx, level, lidx)] = sum_data end elseif gidx ≤ nlevels * nfidx * (nlocalvertices + nlocalfaces) # interior faces nfacedof = div(nperimeter - 4, 4) @@ -176,11 +171,9 @@ function dss_local_kernel!( ip2 = inc2 == 1 ? first2 + i - 1 : first2 - i + 1 idx1 = CI(ip1, 1, fidx, level, lidx1) idx2 = CI(ip2, 1, fidx, level, lidx2) - val = - getindex_field(perimeter_data, idx1) + - getindex_field(perimeter_data, idx2) - setindex_field!(perimeter_data, val, idx1) - setindex_field!(perimeter_data, val, idx2) + val = perimeter_data[idx1] + perimeter_data[idx2] + perimeter_data[idx1] = val + perimeter_data[idx2] = val end end @@ -353,7 +346,7 @@ function dss_local_ghost_kernel!( FT = eltype(parent(perimeter_data)) (nperimeter, _, _, nlevels, _) = size(perimeter_data) nfidx = DataLayouts.ncomponents(perimeter_data) - CI = CartesianIndex + CI = CartesianFieldIndex nghostvertices = length(ghost_vertex_offset) - 1 if gidx ≤ nlevels * nfidx * nghostvertices sizev = (nlevels, nfidx, nghostvertices) @@ -365,19 +358,14 @@ function dss_local_ghost_kernel!( isghost, lidx, vert = ghost_vertices[idx] if !isghost ip = perimeter_vertex_node_index(vert) - sum_data += - getindex_field(perimeter_data, CI(ip, 1, fidx, level, lidx)) + sum_data += perimeter_data[CI(ip, 1, fidx, level, lidx)] end end for idx in st:(en - 1) isghost, lidx, vert = ghost_vertices[idx] if !isghost ip = perimeter_vertex_node_index(vert) - setindex_field!( - perimeter_data, - sum_data, - CI(ip, 1, fidx, level, lidx), - ) + perimeter_data[CI(ip, 1, fidx, level, lidx)] = sum_data end end end @@ -421,14 +409,13 @@ function fill_send_buffer_kernel!( (_, _, _, nlevels, nelems) = size(perimeter_data) nfid = DataLayouts.ncomponents(perimeter_data) sizet = (nlevels, nfid, nsend) - CI = CartesianIndex + CI = CartesianFieldIndex if gidx ≤ nlevels * nfid * nsend (level, fidx, isend) = cart_ind(sizet, gidx).I lidx = send_buf_idx[isend, 1] ip = send_buf_idx[isend, 2] idx = level + ((fidx - 1) + (isend - 1) * nfid) * nlevels - send_data[idx] = - getindex_field(perimeter_data, CI(ip, 1, fidx, level, lidx)) + send_data[idx] = perimeter_data[CI(ip, 1, fidx, level, lidx)] end return nothing end @@ -527,27 +514,20 @@ function dss_ghost_kernel!( (_, _, _, nlevels, _) = size(perimeter_data) nfidx = DataLayouts.ncomponents(perimeter_data) nghostvertices = length(ghost_vertex_offset) - 1 - CI = CartesianIndex + CI = CartesianFieldIndex if gidx ≤ nlevels * nfidx * nghostvertices (level, fidx, ghostvertexidx) = cart_ind((nlevels, nfidx, nghostvertices), gidx).I idxresult, lvertresult = repr_ghost_vertex[ghostvertexidx] ipresult = perimeter_vertex_node_index(lvertresult) - result = getindex_field( - perimeter_data, - CI(ipresult, 1, fidx, level, idxresult), - ) + result = perimeter_data[CI(ipresult, 1, fidx, level, idxresult)] st, en = ghost_vertex_offset[ghostvertexidx], ghost_vertex_offset[ghostvertexidx + 1] for vertexidx in st:(en - 1) isghost, eidx, lvert = ghost_vertices[vertexidx] if !isghost ip = perimeter_vertex_node_index(lvert) - setindex_field!( - perimeter_data, - result, - CI(ip, 1, fidx, level, eidx), - ) + perimeter_data[CI(ip, 1, fidx, level, eidx)] = result end end end diff --git a/lib/ClimaCoreTempestRemap/src/onlineremap.jl b/lib/ClimaCoreTempestRemap/src/onlineremap.jl index 4ea8abb2fd..a2e42ad5ed 100644 --- a/lib/ClimaCoreTempestRemap/src/onlineremap.jl +++ b/lib/ClimaCoreTempestRemap/src/onlineremap.jl @@ -1,4 +1,5 @@ using ClimaCore.DataLayouts +using ClimaCore.DataLayouts: CartesianFieldIndex using ClimaComms @@ -42,11 +43,9 @@ function remap!( R::LinearMap, source::IJFH{S, Nqs}, ) where {S, Nqt, Nqs} - source_array = parent(source) - target_array = parent(target) - - fill!(target_array, zero(eltype(target_array))) - Nf = size(target_array, 3) + fill!(target, zero(eltype(target))) + Nf = DataLayouts.ncomponents(target) + CI = CartesianFieldIndex # ideally we would use the tempestremap dgll (redundant node) representation # unfortunately, this doesn't appear to work quite as well (for out_type = dgll) as the cgll @@ -62,7 +61,7 @@ function remap!( view(R.target_local_idxs[3], n)[1], ) for f in 1:Nf - target_array[it, jt, f, et] += wt * source_array[is, js, f, es] + target[CI(it, jt, f, 1, et)] += wt * source[CI(is, js, f, 1, es)] end end @@ -91,11 +90,12 @@ function remap!(target::Fields.Field, R::LinearMap, source::Fields.Field) @assert Spaces.topology(axes(source)).context isa ClimaComms.SingletonCommsContext - target_array = parent(target) - source_array = parent(source) + CI = CartesianFieldIndex + target_values = Fields.field_values(target) + source_values = Fields.field_values(source) - fill!(target_array, zero(eltype(target_array))) - Nf = size(target_array, 3) + fill!(target, zero(eltype(target))) + Nf = DataLayouts.ncomponents(target) # ideally we would use the tempestremap dgll (redundant node) representation # unfortunately, this doesn't appear to work quite as well (for out_type = dgll) as the cgll @@ -118,8 +118,9 @@ function remap!(target::Fields.Field, R::LinearMap, source::Fields.Field) # only use local weights - i.e. et, es != 0 if (et != 0) for f in 1:Nf - target_array[it, jt, f, et] += - wt * source_array[is, js, f, es] + ci_src = CI(is, js, f, 1, es) + ci_tar = CI(it, jt, f, 1, et) + target_values[ci_tar] += wt * source_values[ci_src] end end end diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index 879133fabc..165a9df991 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -1491,6 +1491,28 @@ end ) end +""" + CartesianFieldIndex{N} <: Base.AbstractCartesianIndex{N} + +A CartesianIndex wrapper to call `getindex` / `setindex!` on +specific field variables in a datalayout. +""" +struct CartesianFieldIndex{N} <: Base.AbstractCartesianIndex{N} + CI::CartesianIndex{N} +end +CartesianFieldIndex(I...) = CartesianFieldIndex(CartesianIndex(I...)) + +Base.ndims(::CartesianFieldIndex{N}) where {N} = N +Base.@propagate_inbounds Base.getindex( + data::AbstractData, + CI::CartesianFieldIndex, +) = getindex_field(data, CI.CI) +Base.@propagate_inbounds Base.setindex!( + data::AbstractData, + val::Real, + CI::CartesianFieldIndex, +) = setindex_field!(data, val, CI.CI) + """ getindex_field(data, ci::CartesianIndex{5}) diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl index 281c6a2dbc..46cffb62df 100644 --- a/src/Topologies/dss.jl +++ b/src/Topologies/dss.jl @@ -1,5 +1,5 @@ using DocStringExtensions -using .DataLayouts: getindex_field, setindex_field! +using .DataLayouts: CartesianFieldIndex """ DSSBuffer{G, D, A, B} @@ -582,13 +582,12 @@ function fill_send_buffer!( Nf = DataLayouts.ncomponents(perimeter_data) nsend = size(send_buf_idx, 1) ctr = 1 - CI = CartesianIndex + CI = CartesianFieldIndex @inbounds for i in 1:nsend lidx = send_buf_idx[i, 1] ip = send_buf_idx[i, 2] for f in 1:Nf, v in 1:Nv - send_data[ctr] = - getindex_field(perimeter_data, CI(ip, 1, f, v, lidx)) + send_data[ctr] = perimeter_data[CI(ip, 1, f, v, lidx)] ctr += 1 end end @@ -612,14 +611,13 @@ function load_from_recv_buffer!( Nf = DataLayouts.ncomponents(perimeter_data) nrecv = size(recv_buf_idx, 1) ctr = 1 - CI = CartesianIndex + CI = CartesianFieldIndex @inbounds for i in 1:nrecv lidx = recv_buf_idx[i, 1] ip = recv_buf_idx[i, 2] for f in 1:Nf, v in 1:Nv ci = CI(ip, 1, f, v, lidx) - val = getindex_field(perimeter_data, ci) + recv_data[ctr] - setindex_field!(perimeter_data, val, ci) + perimeter_data[ci] += recv_data[ctr] ctr += 1 end end diff --git a/test/DataLayouts/unit_getindex_field.jl b/test/DataLayouts/unit_cartesian_field_index.jl similarity index 88% rename from test/DataLayouts/unit_getindex_field.jl rename to test/DataLayouts/unit_cartesian_field_index.jl index 17dd7f783b..7ad99d3c10 100644 --- a/test/DataLayouts/unit_getindex_field.jl +++ b/test/DataLayouts/unit_cartesian_field_index.jl @@ -1,10 +1,10 @@ #= julia --project -using Revise; include(joinpath("test", "DataLayouts", "unit_getindex_field.jl")) +using Revise; include(joinpath("test", "DataLayouts", "unit_cartesian_field_index.jl")) =# using Test using ClimaCore.DataLayouts -using ClimaCore.DataLayouts: getindex_field, setindex_field! +using ClimaCore.DataLayouts: CartesianFieldIndex using ClimaCore.DataLayouts: to_data_specific_field, singleton import ClimaCore.Geometry import ClimaComms @@ -31,15 +31,18 @@ function test_copyto_float!(data) ArrayType = ClimaComms.array_type(ClimaComms.device()) FT = eltype(parent(data)) parent(rand_data) .= ArrayType(rand(FT, DataLayouts.farray_size(data))) - # For a float, getindex and getindex_field return the same thing + # For a float, CartesianIndex and CartesianFieldIndex return the same thing for I in CartesianIndices(universal_axes(data)) - @test getindex_field(data, I) == getindex(data, I) + CI = CartesianFieldIndex(I.I) + @test data[CI] == data[I] end for I in CartesianIndices(universal_axes(data)) - setindex_field!(data, FT(prod(I.I)), I) + CI = CartesianFieldIndex(I.I) + data[CI] = FT(prod(I.I)) end for I in CartesianIndices(universal_axes(data)) - @test getindex_field(data, I) == prod(I.I) + CI = CartesianFieldIndex(I.I) + @test data[CI] == prod(I.I) end end @@ -55,7 +58,7 @@ function test_copyto!(data) for f in 1:DataLayouts.ncomponents(data) UFI = universal_field_index(I, f) DSI = CartesianIndex(to_data_specific_field(singleton(data), UFI.I)) - @test getindex_field(data, UFI) == parent(data)[DSI] + @test data[CartesianFieldIndex(UFI)] == parent(data)[DSI] end end @@ -64,13 +67,13 @@ function test_copyto!(data) UFI = universal_field_index(I, f) DSI = CartesianIndex(to_data_specific_field(singleton(data), UFI.I)) val = parent(data)[DSI] - setindex_field!(data, val + 1, UFI) + data[CartesianFieldIndex(UFI)] = val + 1 @test parent(data)[DSI] == val + 1 end end end -@testset "copyto! with Nf = 1" begin +@testset "CartesianFieldIndex with Nf = 1" begin device = ClimaComms.device() ArrayType = ClimaComms.array_type(device) FT = Float64 @@ -99,7 +102,7 @@ end # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto_float!(data) # TODO: test end -@testset "copyto! with Nf > 1" begin +@testset "CartesianFieldIndex with Nf > 1" begin device = ClimaComms.device() ArrayType = ClimaComms.array_type(device) FT = Float64 @@ -129,7 +132,7 @@ end # data = DataLayouts.IH1JH2{S}(ArrayType{FT}, zeros; Nij); test_copyto!(data) # TODO: test end -@testset "copyto! views with Nf > 1" begin +@testset "CartesianFieldIndex views with Nf > 1" begin device = ClimaComms.device() ArrayType = ClimaComms.array_type(device) data_view(data) = DataLayouts.rebuild( diff --git a/test/Fields/unit_field.jl b/test/Fields/unit_field.jl index 63a1f8f701..2fe4d5855b 100644 --- a/test/Fields/unit_field.jl +++ b/test/Fields/unit_field.jl @@ -57,9 +57,11 @@ end n1 = n2 = 1 Nh = n1 * n2 space = spectral_space_2D(n1 = n1, n2 = n2, Nij = Nij) + device = ClimaComms.device(space) + ArrayType = ClimaComms.array_type(device) - field = - Fields.Field(IJFH{ComplexF64, Nij}(ones(Nij, Nij, 2, n1 * n2)), space) + data = IJFH{ComplexF64}(ArrayType{Float64}, ones; Nij, Nh = n1 * n2) + field = Fields.Field(data, space) @test sum(field) ≈ Complex(1.0, 1.0) * 8.0 * 10.0 rtol = 10eps() @test sum(x -> 3.0, field) ≈ 3 * 8.0 * 10.0 rtol = 10eps()