From 6ff39ad3032802ac2a942a11adb3fd7f1aa78125 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 9 Jan 2025 11:05:03 -0500 Subject: [PATCH] Define CPU<->GPU adaptations Try to fix downgrade ci --- Project.toml | 4 +- src/Fields/fieldvector.jl | 5 ++ src/Grids/extruded.jl | 2 + src/Grids/finitedifference.jl | 2 +- src/Grids/spectralelement.jl | 4 + test/Fields/unit_field.jl | 151 +++++++++++++++++++++++++++++++++- test/Spaces/unit_spaces.jl | 8 +- 7 files changed, 168 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 61aa70e07b..71c2cf394e 100644 --- a/Project.toml +++ b/Project.toml @@ -37,7 +37,7 @@ ClimaCoreCUDAExt = "CUDA" KrylovExt = "Krylov" [compat] -Adapt = "3, 4" +Adapt = "3.2.0, 4" Aqua = "0.8" ArgParse = "1" AssociatedLegendrePolynomials = "1" @@ -68,7 +68,7 @@ NVTX = "0.3" PkgVersion = "0.1, 0.2, 0.3" PrettyTables = "2" Random = "1" -RecursiveArrayTools = "3.1" +RecursiveArrayTools = "3.2" RootSolvers = "0.3, 0.4" SafeTestsets = "0.1" SparseArrays = "1" diff --git a/src/Fields/fieldvector.jl b/src/Fields/fieldvector.jl index d11e5ecebf..bdb799bcec 100644 --- a/src/Fields/fieldvector.jl +++ b/src/Fields/fieldvector.jl @@ -23,6 +23,11 @@ struct FieldVector{T, M} <: BlockArrays.AbstractBlockVector{T} end FieldVector{T}(values::M) where {T, M} = FieldVector{T, M}(values) +function Adapt.adapt_structure(to, fv::FieldVector) + pn = propertynames(fv) + vals = map(key -> Adapt.adapt(to, getproperty(fv, key)), pn) + return FieldVector(; NamedTuple{pn}(vals)...) +end """ Fields.ScalarWrapper(val) <: AbstractArray{T,0} diff --git a/src/Grids/extruded.jl b/src/Grids/extruded.jl index 4486b3e3f8..105cd887ca 100644 --- a/src/Grids/extruded.jl +++ b/src/Grids/extruded.jl @@ -38,6 +38,8 @@ mutable struct ExtrudedFiniteDifferenceGrid{ face_local_geometry::FLG end +Adapt.@adapt_structure ExtrudedFiniteDifferenceGrid + local_geometry_type( ::Type{ExtrudedFiniteDifferenceGrid{H, V, A, GG, CLG, FLG}}, ) where {H, V, A, GG, CLG, FLG} = eltype(CLG) # calls eltype from DataLayouts diff --git a/src/Grids/finitedifference.jl b/src/Grids/finitedifference.jl index fbb743cda1..402400568b 100644 --- a/src/Grids/finitedifference.jl +++ b/src/Grids/finitedifference.jl @@ -40,7 +40,7 @@ mutable struct FiniteDifferenceGrid{ center_local_geometry::CLG face_local_geometry::FLG end - +Adapt.@adapt_structure FiniteDifferenceGrid function FiniteDifferenceGrid(topology::Topologies.IntervalTopology) get!(Cache.OBJECT_CACHE, (FiniteDifferenceGrid, topology)) do diff --git a/src/Grids/spectralelement.jl b/src/Grids/spectralelement.jl index 75504b9c2f..7e20872fd5 100644 --- a/src/Grids/spectralelement.jl +++ b/src/Grids/spectralelement.jl @@ -21,6 +21,8 @@ mutable struct SpectralElementGrid1D{ dss_weights::D end +Adapt.@adapt_structure SpectralElementGrid1D + local_geometry_type( ::Type{SpectralElementGrid1D{T, Q, GG, LG}}, ) where {T, Q, GG, LG} = eltype(LG) # calls eltype from DataLayouts @@ -140,6 +142,8 @@ mutable struct SpectralElementGrid2D{ enable_bubble::Bool end +Adapt.@adapt_structure SpectralElementGrid2D + local_geometry_type( ::Type{SpectralElementGrid2D{T, Q, GG, LG, D, IS, BS}}, ) where {T, Q, GG, LG, D, IS, BS} = eltype(LG) # calls eltype from DataLayouts diff --git a/test/Fields/unit_field.jl b/test/Fields/unit_field.jl index d2f779289b..d0e7958395 100644 --- a/test/Fields/unit_field.jl +++ b/test/Fields/unit_field.jl @@ -1,6 +1,6 @@ #= julia --check-bounds=yes --project -julia --project +julia --project=.buildkite using Revise; include(joinpath("test", "Fields", "unit_field.jl")) =# using Test @@ -702,6 +702,155 @@ end nothing end +using ClimaCore.CommonSpaces +using ClimaCore.Grids +using Adapt + +function test_adapt(cpu_space_in) + test_adapt_space(cpu_space_in) + cpu_f_in = Fields.Field(Float64, cpu_space_in) + cpu_f_out = Adapt.adapt(Array, cpu_f_in) + @test parent(Spaces.local_geometry_data(axes(cpu_f_out))) isa Array + @test parent(Fields.field_values(cpu_f_out)) isa Array + + @static if ClimaComms.device() isa ClimaComms.CUDADevice + # cpu -> gpu + gpu_f_out = Adapt.adapt(CUDA.CuArray, cpu_f_in) + @test parent(Fields.field_values(gpu_f_out)) isa CUDA.CuArray + # gpu -> gpu + cpu_f_out = Adapt.adapt(Array, gpu_f_out) + @test parent(Fields.field_values(cpu_f_out)) isa Array + end +end + +function test_adapt_fieldvector(fv_in) + cpu_fv_out = Adapt.adapt(Array, fv_in) + @test parent(Spaces.local_geometry_data(axes(cpu_fv_out.c))) isa Array + @test parent(Spaces.local_geometry_data(axes(cpu_fv_out.f))) isa Array + @test parent(Fields.field_values(cpu_fv_out.c)) isa Array + @test parent(Fields.field_values(cpu_fv_out.f)) isa Array + + @static if ClimaComms.device() isa ClimaComms.CUDADevice + # cpu -> gpu + gpu_fv_out = Adapt.adapt(CUDA.CuArray, cpu_fv_out) + @test parent(Fields.field_values(gpu_fv_out.c)) isa CUDA.CuArray + @test parent(Fields.field_values(gpu_fv_out.f)) isa CUDA.CuArray + # gpu -> gpu + cpu_fv_out = Adapt.adapt(Array, gpu_fv_out) + @test parent(Fields.field_values(cpu_fv_out.c)) isa Array + @test parent(Fields.field_values(cpu_fv_out.f)) isa Array + end +end + +function test_adapt_space(cpu_space_in) + # cpu -> cpu + cpu_space_out = Adapt.adapt(Array, cpu_space_in) + @test parent(Spaces.local_geometry_data(cpu_space_out)) isa Array + + @static if ClimaComms.device() isa ClimaComms.CUDADevice + # cpu -> gpu + gpu_space_out = Adapt.adapt(CUDA.CuArray, cpu_space_in) + @test parent(Spaces.local_geometry_data(gpu_space_out)) isa CUDA.CuArray + # gpu -> gpu + cpu_space_out = Adapt.adapt(Array, gpu_space_out) + @test parent(Spaces.local_geometry_data(cpu_space_out)) isa Array + end +end + +@testset "Test Adapt" begin + space = ExtrudedCubedSphereSpace(; + device = ClimaComms.CPUSingleThreaded(), + z_elem = 10, + z_min = 0, + z_max = 1, + radius = 10, + h_elem = 10, + n_quad_points = 4, + staggering = Grids.CellCenter(), + ) + test_adapt(space) + + space = CubedSphereSpace(; + device = ClimaComms.CPUSingleThreaded(), + radius = 10, + n_quad_points = 4, + h_elem = 10, + ) + test_adapt(space) + + space = ColumnSpace(; + device = ClimaComms.CPUSingleThreaded(), + z_elem = 10, + z_min = 0, + z_max = 10, + staggering = CellCenter(), + ) + test_adapt(space) + + space = Box3DSpace(; + device = ClimaComms.CPUSingleThreaded(), + z_elem = 10, + x_min = 0, + x_max = 1, + y_min = 0, + y_max = 1, + z_min = 0, + z_max = 10, + periodic_x = false, + periodic_y = false, + n_quad_points = 4, + x_elem = 3, + y_elem = 4, + staggering = CellCenter(), + ) + test_adapt(space) + + space = SliceXZSpace(; + device = ClimaComms.CPUSingleThreaded(), + z_elem = 10, + x_min = 0, + x_max = 1, + z_min = 0, + z_max = 1, + periodic_x = false, + n_quad_points = 4, + x_elem = 4, + staggering = CellCenter(), + ) + test_adapt(space) + + space = RectangleXYSpace(; + device = ClimaComms.CPUSingleThreaded(), + x_min = 0, + x_max = 1, + y_min = 0, + y_max = 1, + periodic_x = false, + periodic_y = false, + n_quad_points = 4, + x_elem = 3, + y_elem = 4, + ) + test_adapt(space) + + # FieldVector + cspace = ExtrudedCubedSphereSpace(; + device = ClimaComms.CPUSingleThreaded(), + z_elem = 10, + z_min = 0, + z_max = 1, + radius = 10, + h_elem = 10, + n_quad_points = 4, + staggering = Grids.CellCenter(), + ) + fspace = Spaces.face_space(cspace) + c = Fields.zeros(cspace) + f = Fields.zeros(fspace) + fv = Fields.FieldVector(; c, f) + test_adapt_fieldvector(fv) +end + @testset "Memoization of spaces" begin space1 = spectral_space_2D() space2 = spectral_space_2D() diff --git a/test/Spaces/unit_spaces.jl b/test/Spaces/unit_spaces.jl index 3d4fc1f6aa..313d137506 100644 --- a/test/Spaces/unit_spaces.jl +++ b/test/Spaces/unit_spaces.jl @@ -5,7 +5,7 @@ using Revise; include(joinpath("test", "Spaces", "unit_spaces.jl")) using Test using ClimaComms using StaticArrays, IntervalSets, LinearAlgebra -using Adapt +import Adapt ClimaComms.@import_required_backends import ClimaCore: @@ -198,11 +198,11 @@ end @test length(Spaces.all_nodes(hspace)) == 4 @static if on_gpu - adapted_space = adapt(CUDA.KernelAdaptor(), c_space) + adapted_space = Adapt.adapt(CUDA.KernelAdaptor(), c_space) @test ClimaComms.context(adapted_space) == DeviceSideContext() @test ClimaComms.device(adapted_space) == DeviceSideDevice() - adapted_hspace = adapt(CUDA.KernelAdaptor(), hspace) + adapted_hspace = Adapt.adapt(CUDA.KernelAdaptor(), hspace) @test ClimaComms.context(adapted_hspace) == DeviceSideContext() @test ClimaComms.device(adapted_hspace) == DeviceSideDevice() end @@ -245,7 +245,7 @@ end dss_weights_slab = slab(Spaces.local_dss_weights(space), 1) @static if on_gpu - adapted_space = adapt(CUDA.KernelAdaptor(), space) + adapted_space = Adapt.adapt(CUDA.KernelAdaptor(), space) @test ClimaComms.context(adapted_space) == DeviceSideContext() @test ClimaComms.device(adapted_space) == DeviceSideDevice() end