From 09e48c6a6d3cef30175e2f08a6e2258877cc1227 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 18 Oct 2024 09:44:51 -0400 Subject: [PATCH] Define getindex_field and setindex_field --- .buildkite/pipeline.yml | 4 + src/DataLayouts/DataLayouts.jl | 51 +++++++ test/DataLayouts/unit_getindex_field.jl | 170 ++++++++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 226 insertions(+) create mode 100644 test/DataLayouts/unit_getindex_field.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b7bcba57f0..dc49fa1fe0 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -86,6 +86,10 @@ steps: key: unit_data_copyto command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl" + - label: "Unit: getindex_field" + key: unit_data_getindex_field + command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_getindex_field.jl" + - label: "Unit: mapreduce" key: unit_data_mapreduce command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl" diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index bece0535a9..879133fabc 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -1309,6 +1309,15 @@ type parameters. @inline to_data_specific(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], 1, I[5]) @inline to_data_specific(::VIFHSingleton, I::Tuple) = (I[4], I[1], 1, I[5]) +@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(::IFHSingleton, I::Tuple) = (I[1], I[3], I[5]) +@inline to_data_specific_field(::VIJFHSingleton, I::Tuple) = (I[4], I[1], I[2], I[3], I[5]) +@inline to_data_specific_field(::VIFHSingleton, I::Tuple) = (I[4], I[1], I[3], I[5]) + """ bounds_condition(data::AbstractData, I::Tuple) @@ -1482,6 +1491,48 @@ end ) end +""" + getindex_field(data, ci::CartesianIndex{5}) + +Returns the value of the data at universal index `ci`, +for the specific field `f` in the `CartesianIndex`. + +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}, + I::CartesianIndex, # universal index +) + @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) + @inbounds Base.getindex( + parent(data), + CartesianIndex(to_data_specific_field(singleton(data), I.I)), + ) +end + +""" + setindex_field!(data, val::Real, ci::CartesianIndex{5}) + +Stores the value `val` of the data at universal index `ci`, +for the specific field `f` in the `CartesianIndex`. + +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}, + val::Real, + I::CartesianIndex, # universal index +) + @boundscheck bounds_condition(data, I) || throw(BoundsError(data, I)) + @inbounds Base.setindex!( + parent(data), + val, + CartesianIndex(to_data_specific_field(singleton(data), I.I)), + ) +end + 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, diff --git a/test/DataLayouts/unit_getindex_field.jl b/test/DataLayouts/unit_getindex_field.jl new file mode 100644 index 0000000000..17dd7f783b --- /dev/null +++ b/test/DataLayouts/unit_getindex_field.jl @@ -0,0 +1,170 @@ +#= +julia --project +using Revise; include(joinpath("test", "DataLayouts", "unit_getindex_field.jl")) +=# +using Test +using ClimaCore.DataLayouts +using ClimaCore.DataLayouts: getindex_field, setindex_field! +using ClimaCore.DataLayouts: to_data_specific_field, singleton +import ClimaCore.Geometry +import ClimaComms +using StaticArrays +ClimaComms.@import_required_backends +import Random +Random.seed!(1234) + +universal_axes(data) = + map(size(data)) do i + s = + i == DataLayouts.field_dim(singleton(data)) ? + DataLayouts.ncomponents(data) : i + Base.OneTo(s) + end + +universal_field_index(I::CartesianIndex, f) = + CartesianIndex(map(i -> i == 3 ? f : i, I.I)) + +function test_copyto_float!(data) + Random.seed!(1234) + # Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803 + rand_data = DataLayouts.rebuild(data, similar(parent(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 I in CartesianIndices(universal_axes(data)) + @test getindex_field(data, I) == getindex(data, I) + end + for I in CartesianIndices(universal_axes(data)) + setindex_field!(data, FT(prod(I.I)), I) + end + for I in CartesianIndices(universal_axes(data)) + @test getindex_field(data, I) == prod(I.I) + end +end + +function test_copyto!(data) + Random.seed!(1234) + # Normally we'd use `similar` here, but https://github.com/CliMA/ClimaCore.jl/issues/1803 + rand_data = DataLayouts.rebuild(data, similar(parent(data))) + ArrayType = ClimaComms.array_type(ClimaComms.device()) + FT = eltype(parent(data)) + parent(rand_data) .= ArrayType(rand(FT, DataLayouts.farray_size(data))) + + for I in CartesianIndices(universal_axes(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] + end + end + + for I in CartesianIndices(universal_axes(data)) + for f in 1:DataLayouts.ncomponents(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) + @test parent(data)[DSI] == val + 1 + end + end +end + +@testset "copyto! with Nf = 1" begin + device = ClimaComms.device() + ArrayType = ClimaComms.array_type(device) + FT = Float64 + S = FT + Nv = 4 + Ni = Nij = 3 + Nh = 5 + Nk = 6 + data = DataF{S}(ArrayType{FT}, zeros) + test_copyto_float!(data) + data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) + test_copyto_float!(data) + data = IFH{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) + test_copyto_float!(data) + data = VF{S}(ArrayType{FT}, zeros; Nv) + test_copyto_float!(data) + data = VIJFH{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 = 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 + +@testset "copyto! with Nf > 1" begin + device = ClimaComms.device() + ArrayType = ClimaComms.array_type(device) + FT = Float64 + S = Tuple{FT, FT} + Nv = 4 + Ni = Nij = 3 + Nh = 5 + Nk = 6 + data = DataF{S}(ArrayType{FT}, zeros) + test_copyto!(data) + data = IJFH{S}(ArrayType{FT}, zeros; Nij, Nh) + test_copyto!(data) + data = IFH{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) + test_copyto!(data) + data = VF{S}(ArrayType{FT}, zeros; Nv) + test_copyto!(data) + data = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh) + test_copyto!(data) + data = VIFH{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 +end + +@testset "copyto! views with Nf > 1" begin + device = ClimaComms.device() + ArrayType = ClimaComms.array_type(device) + data_view(data) = DataLayouts.rebuild( + data, + SubArray( + parent(data), + ntuple( + i -> Base.Slice(Base.OneTo(DataLayouts.farray_size(data, i))), + ndims(data), + ), + ), + ) + FT = Float64 + S = Tuple{FT, FT} + Nv = 4 + Ni = Nij = 3 + Nh = 5 + Nk = 6 + # Rather than using level/slab/column, let's just make views/SubArrays + # directly so that we can easily test all cases: + data = IJFH{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 = IJF{S}(ArrayType{FT}, zeros; Nij) + test_copyto!(data_view(data)) + data = IF{S}(ArrayType{FT}, zeros; Ni) + test_copyto!(data_view(data)) + data = VF{S}(ArrayType{FT}, zeros; Nv) + test_copyto!(data_view(data)) + data = VIJFH{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)) + # 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 +end diff --git a/test/runtests.jl b/test/runtests.jl index 5b79c5b70b..a7af3ba5df 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ UnitTest("DataLayouts fill" ,"DataLayouts/unit_fill.jl"), 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_getindex_field.jl"), UnitTest("Recursive" ,"RecursiveApply/unit_recursive_apply.jl"), UnitTest("PlusHalf" ,"Utilities/unit_plushalf.jl"), UnitTest("DataLayouts 0D" ,"DataLayouts/data0d.jl"),