Move in-kernel specific adapt functions to cuda ext

Author: charleskawczynski

ext/ClimaCoreCUDAExt.jl

@@ -20,6 +20,7 @@ import ClimaCore.RecursiveApply:
 import ClimaCore.DataLayouts: get_N, get_Nv, get_Nij, get_Nij, get_Nh
 import ClimaCore.DataLayouts: UniversalSize
 
+include(joinpath("cuda", "adapt.jl"))
 include(joinpath("cuda", "cuda_utils.jl"))
 include(joinpath("cuda", "data_layouts.jl"))
 include(joinpath("cuda", "fields.jl"))

ext/cuda/adapt.jl

@@ -0,0 +1,44 @@
+import CUDA, Adapt
+import ClimaCore
+import ClimaCore: Grids, Spaces, Topologies, Devices
+
+Adapt.adapt_structure(
+    to::CUDA.KernelAdaptor,
+    grid::Grids.ExtrudedFiniteDifferenceGrid,
+) = Grids.DeviceExtrudedFiniteDifferenceGrid(
+    Adapt.adapt(to, Grids.vertical_topology(grid)),
+    Adapt.adapt(to, grid.horizontal_grid.quadrature_style),
+    Adapt.adapt(to, grid.global_geometry),
+    Adapt.adapt(to, grid.center_local_geometry),
+    Adapt.adapt(to, grid.face_local_geometry),
+)
+
+Adapt.adapt_structure(
+    to::CUDA.KernelAdaptor,
+    grid::Grids.FiniteDifferenceGrid,
+) = Grids.DeviceFiniteDifferenceGrid(
+    Adapt.adapt(to, grid.topology),
+    Adapt.adapt(to, grid.global_geometry),
+    Adapt.adapt(to, grid.center_local_geometry),
+    Adapt.adapt(to, grid.face_local_geometry),
+)
+
+Adapt.adapt_structure(
+    to::CUDA.KernelAdaptor,
+    grid::Grids.SpectralElementGrid2D,
+) = Grids.DeviceSpectralElementGrid2D(
+    Adapt.adapt(to, grid.quadrature_style),
+    Adapt.adapt(to, grid.global_geometry),
+    Adapt.adapt(to, grid.local_geometry),
+)
+
+Adapt.adapt_structure(to::CUDA.KernelAdaptor, space::Spaces.PointSpace) =
+    Spaces.PointSpace(
+        ClimaCore.DeviceSideContext(),
+        Adapt.adapt(to, Spaces.local_geometry_data(space)),
+    )
+
+Adapt.adapt_structure(
+    to::CUDA.KernelAdaptor,
+    topology::Topologies.IntervalTopology,
+) = Topologies.DeviceIntervalTopology(topology.boundaries)

src/Grids/extruded.jl

@@ -155,15 +155,6 @@ local_geometry_type(
     ::Type{DeviceExtrudedFiniteDifferenceGrid{VT, Q, GG, CLG, FLG}},
 ) where {VT, Q, GG, CLG, FLG} = eltype(CLG) # calls eltype from DataLayouts
 
-Adapt.adapt_structure(to, grid::ExtrudedFiniteDifferenceGrid) =
-    DeviceExtrudedFiniteDifferenceGrid(
-        Adapt.adapt(to, vertical_topology(grid)),
-        Adapt.adapt(to, grid.horizontal_grid.quadrature_style),
-        Adapt.adapt(to, grid.global_geometry),
-        Adapt.adapt(to, grid.center_local_geometry),
-        Adapt.adapt(to, grid.face_local_geometry),
-    )
-
 quadrature_style(grid::DeviceExtrudedFiniteDifferenceGrid) =
     grid.quadrature_style
 vertical_topology(grid::DeviceExtrudedFiniteDifferenceGrid) =

src/Grids/finitedifference.jl

@@ -190,14 +190,6 @@ local_geometry_type(
     ::Type{DeviceFiniteDifferenceGrid{T, GG, CLG, FLG}},
 ) where {T, GG, CLG, FLG} = eltype(CLG) # calls eltype from DataLayouts
 
-Adapt.adapt_structure(to, grid::FiniteDifferenceGrid) =
-    DeviceFiniteDifferenceGrid(
-        Adapt.adapt(to, grid.topology),
-        Adapt.adapt(to, grid.global_geometry),
-        Adapt.adapt(to, grid.center_local_geometry),
-        Adapt.adapt(to, grid.face_local_geometry),
-    )
-
 topology(grid::DeviceFiniteDifferenceGrid) = grid.topology
 vertical_topology(grid::DeviceFiniteDifferenceGrid) = grid.topology
 

src/Grids/spectralelement.jl

@@ -597,13 +597,6 @@ end
 ClimaComms.context(grid::DeviceSpectralElementGrid2D) = DeviceSideContext()
 ClimaComms.device(grid::DeviceSpectralElementGrid2D) = DeviceSideDevice()
 
-Adapt.adapt_structure(to, grid::SpectralElementGrid2D) =
-    DeviceSpectralElementGrid2D(
-        Adapt.adapt(to, grid.quadrature_style),
-        Adapt.adapt(to, grid.global_geometry),
-        Adapt.adapt(to, grid.local_geometry),
-    )
-
 ## aliases
 const RectilinearSpectralElementGrid2D =
     SpectralElementGrid2D{<:Topologies.RectilinearTopology2D}

src/Spaces/pointspace.jl

@@ -38,10 +38,6 @@ function PointSpace(
     return PointSpace(context, Adapt.adapt(ArrayType, local_geometry_data))
 end
 
-
-Adapt.adapt_structure(to, space::PointSpace) =
-    PointSpace(DeviceSideContext(), Adapt.adapt(to, local_geometry_data(space)))
-
 function PointSpace(
     context::ClimaComms.AbstractCommsContext,
     coord::Geometry.Abstract1DPoint{FT},

src/Topologies/interval.jl

@@ -20,8 +20,6 @@ end
 struct DeviceIntervalTopology{B} <: AbstractIntervalTopology
     boundaries::B
 end
-Adapt.adapt_structure(to, topology::IntervalTopology) =
-    DeviceIntervalTopology(topology.boundaries)
 
 ClimaComms.context(topology::DeviceIntervalTopology) = DeviceSideContext()
 ClimaComms.device(topology::DeviceIntervalTopology) = DeviceSideDevice()

test/Spaces/unit_spaces.jl

@@ -197,12 +197,12 @@ end
     @test length(Spaces.unique_nodes(hspace)) == 4
     @test length(Spaces.all_nodes(hspace)) == 4
 
-    if on_gpu
-        adapted_space = adapt(c_space)(c_space)
+    @static if on_gpu
+        adapted_space = adapt(CUDA.KernelAdaptor(), c_space)
         @test ClimaComms.context(adapted_space) == DeviceSideContext()
         @test ClimaComms.device(adapted_space) == DeviceSideDevice()
 
-        adapted_hspace = adapt(hspace)(hspace)
+        adapted_hspace = adapt(CUDA.KernelAdaptor(), hspace)
         @test ClimaComms.context(adapted_hspace) == DeviceSideContext()
         @test ClimaComms.device(adapted_hspace) == DeviceSideDevice()
     end
@@ -244,8 +244,8 @@ end
     local_geometry_slab = slab(Spaces.local_geometry_data(space), 1)
     dss_weights_slab = slab(Spaces.local_dss_weights(space), 1)
 
-    if on_gpu
-        adapted_space = adapt(space)(space)
+    @static if on_gpu
+        adapted_space = adapt(CUDA.KernelAdaptor(), space)
         @test ClimaComms.context(adapted_space) == DeviceSideContext()
         @test ClimaComms.device(adapted_space) == DeviceSideDevice()
     end