diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index d5314d1da1..3f70b535c1 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -86,7 +86,7 @@ export slab, VIHF, DataF -# Internal types for managing CPU/GPU dispatching +# Internal types for managing CPU/GPU dispatching / conversions abstract type AbstractDispatchToDevice end struct ToCPU <: AbstractDispatchToDevice end struct ToCUDA <: AbstractDispatchToDevice end diff --git a/test/Spaces/unit_spaces.jl b/test/Spaces/unit_spaces.jl index 0c4ec37f32..9b3732a30e 100644 --- a/test/Spaces/unit_spaces.jl +++ b/test/Spaces/unit_spaces.jl @@ -324,6 +324,45 @@ end end end +using ClimaCore.CommonSpaces +using ClimaCore.Grids +using ClimaCore.DataLayouts: ToCPU, ToCUDA +using Adapt +@testset "Adapt between CPU<->CPU" begin + cpu_space_in = 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() + ) + cpu_space_out = Adapt.adapt(ToCPU(), cpu_space_in) + @test cpu_space_in === cpu_space_out +end + +if ClimaComms.device() isa ClimaComms.CUDADevice + @testset "Adapt between CPU->GPU->CPU" begin + cpu_space_in = 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() + ) + cpu_space_out = Adapt.adapt(ToCPU(), cpu_space_in) + @test cpu_space_in === cpu_space_out + gpu_space_out = Adapt.adapt(ToCUDA(), cpu_space_in) + gpu_array_type = ClimaComms.array_type(ClimaComms.CUDADevice()) + @test parent(Spaces.coordinates_data(space)) isa gpu_array_type + + end +end #= @testset "dss on 2×2 rectangular mesh (unstructured)" begin