Merge pull request #2114 from CliMA/ck/adapt_cpu_gpu

Define cpu<->gpu conversions

Author: charleskawczynski

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"

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}

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

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

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

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()

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