Remove DeviceIntervalTopology

Author: Charlie Kawczynski

ext/cuda/adapt.jl

@@ -42,7 +42,11 @@ Adapt.adapt_structure(to::CUDA.KernelAdaptor, space::Spaces.PointSpace) =
 Adapt.adapt_structure(
     to::CUDA.KernelAdaptor,
     topology::Topologies.IntervalTopology,
-) = Topologies.DeviceIntervalTopology(topology.boundaries)
+) = Topologies.IntervalTopology(
+    Adapt.adapt(to, ClimaComms.SingletonCommsContext(ClimaComms.device(topology.context))),
+    Adapt.adapt(to, topology.mesh),
+    Adapt.adapt(to, topology.boundaries),
+)
 
 Adapt.adapt_structure(
     to::CUDA.KernelAdaptor,

src/Domains/Domains.jl

@@ -32,12 +32,12 @@ struct IntervalDomain{CT, B} <: AbstractDomain where {
 } where {FT}
     coord_min::CT
     coord_max::CT
-    boundary_names::B
 end
 
-isperiodic(domain::IntervalDomain) = isnothing(domain.boundary_names)
-boundary_names(domain::IntervalDomain) =
-    isperiodic(domain) ? () : unique(domain.boundary_names)
+isperiodic(::IntervalDomain{CT, B}) where {CT, B} = B == nothing
+boundary_names(domain::IntervalDomain{CT, B}) where {CT, B} =
+    isperiodic(domain) ? () : unique(B)
+boundary_names_type(::IntervalDomain{CT, B}) where {CT, B} = B
 
 """
     IntervalDomain(coord⁻, coord⁺; periodic=true)
@@ -60,7 +60,13 @@ function IntervalDomain(
             ),
         )
     end
-    IntervalDomain(promote(coord_min, coord_max)..., boundary_names)
+    c = promote(coord_min, coord_max)
+    boundary_names = if isnothing(boundary_names)
+        boundary_names
+    else
+        Tuple(boundary_names)
+    end
+    IntervalDomain{eltype(c), boundary_names}(c...)
 end
 IntervalDomain(coords::IntervalSets.ClosedInterval; kwargs...) =
     IntervalDomain(coords.left, coords.right; kwargs...)
@@ -95,7 +101,7 @@ function print_interval(io::IO, domain::IntervalDomain{CT}) where {CT}
     if isperiodic(domain)
         print(io, "(periodic)")
     else
-        print(io, domain.boundary_names)
+        print(io, boundary_names_type(domain))
     end
 end
 function Base.show(io::IO, domain::IntervalDomain)

src/Grids/finitedifference.jl

@@ -73,7 +73,7 @@ function _FiniteDifferenceGrid(topology::Topologies.IntervalTopology)
     )
 
     return FiniteDifferenceGrid(
-        topology,
+        Adapt.adapt(ArrayType, topology),
         global_geometry,
         Adapt.adapt(ArrayType, center_local_geometry),
         Adapt.adapt(ArrayType, face_local_geometry),

src/InputOutput/writers.jl

@@ -166,10 +166,10 @@ function write_new!(
     )
     write_attribute(group, "coord_min", Geometry.component(domain.coord_min, 1))
     write_attribute(group, "coord_max", Geometry.component(domain.coord_max, 1))
-    !isnothing(domain.boundary_names) && write_attribute(
+    !isnothing(Domains.boundary_names(domain)) && write_attribute(
         group,
         "boundary_names",
-        [String(bname) for bname in domain.boundary_names],
+        [String(bname) for bname in Domains.boundary_names(domain)],
     )
     return name
 end

src/Meshes/interval.jl

@@ -26,6 +26,9 @@ struct IntervalMesh{S, I <: IntervalDomain, V <: AbstractVector, M} <:
     meta::M
 end
 
+using Adapt
+Adapt.@adapt_structure IntervalMesh
+
 # implies isequal
 Base.:(==)(mesh1::IntervalMesh, mesh2::IntervalMesh) =
     mesh1.domain == mesh2.domain && mesh1.faces == mesh2.faces
@@ -85,9 +88,9 @@ end
 function boundary_face_name(mesh::IntervalMesh, elem::Integer, face)
     if !Domains.isperiodic(mesh.domain)
         if elem == 1 && face == 1
-            return mesh.domain.boundary_names[1]
+            return Domains.boundary_names(mesh.domain)[1]
         elseif elem == nelements(mesh) && face == 2
-            return mesh.domain.boundary_names[2]
+            return Domains.boundary_names(mesh.domain)[2]
         end
     end
     return nothing
@@ -468,7 +471,7 @@ function truncate_mesh(
     new_domain = IntervalDomain(
         z_bottom,
         Geometry.ZPoint{FT}(z_top),
-        boundary_names = trunc_domain.boundary_names,
+        boundary_names = Domains.boundary_names_type(trunc_domain),
     )
     return IntervalMesh(new_domain, new_stretch; nelems = new_nelems)
 end

src/Topologies/interval.jl

@@ -18,15 +18,10 @@ end
 
 Adapt.@adapt_structure IntervalTopology
 
-## gpu
-struct DeviceIntervalTopology{B} <: AbstractIntervalTopology
-    boundaries::B
-end
-
-ClimaComms.context(topology::DeviceIntervalTopology) = DeviceSideContext()
-ClimaComms.device(topology::DeviceIntervalTopology) = DeviceSideDevice()
+ClimaComms.context(topology::IntervalTopology) = topology.context
+ClimaComms.device(topology::IntervalTopology) =
+    ClimaComms.device(topology.context)
 
-ClimaComms.device(topology::IntervalTopology) = topology.context.device
 ClimaComms.array_type(topology::IntervalTopology) =
     ClimaComms.array_type(topology.context.device)
 
@@ -46,14 +41,19 @@ function _IntervalTopology(
 )
     # currently only support SingletonCommsContext
     @assert context isa ClimaComms.SingletonCommsContext
-    if Domains.isperiodic(mesh.domain)
+    domain = mesh.domain
+    if Domains.isperiodic(domain)
         boundaries = NamedTuple()
-    elseif mesh.domain.boundary_names[1] == mesh.domain.boundary_names[2]
-        boundaries = NamedTuple{(mesh.domain.boundary_names[1],)}(1)
     else
-        boundaries = NamedTuple{mesh.domain.boundary_names}((1, 2))
+        bn = Domains.boundary_names_type(domain)
+        boundaries = if bn[1] == bn[2]
+            NamedTuple{(bn[1],)}(1)
+        else
+            NamedTuple{bn}((1, 2))
+        end
     end
-    IntervalTopology(context, mesh, boundaries)
+    ArrayType = ClimaComms.array_type(ClimaComms.device(context))
+    IntervalTopology(context, Adapt.adapt(ArrayType, mesh), boundaries)
 end
 
 IntervalTopology(device::ClimaComms.AbstractDevice, mesh::Meshes.IntervalMesh) =

test/Operators/finitedifference/unit_columnwise.jl

@@ -30,13 +30,11 @@ import ClimaCore.MatrixFields
 import ClimaCore.Spaces
 import ClimaCore.Fields
 
-Operators.fd_shmem_is_supported(bc::Base.Broadcast.Broadcasted) = false
-ClimaCore.Operators.use_fd_shmem() = false
+# Operators.fd_shmem_is_supported(bc::Base.Broadcast.Broadcasted) = false
+# ClimaCore.Operators.use_fd_shmem() = false
 # The existing implementation limits our ability to apply
 # the same expressions from within kernels
-ClimaComms.device(topology::Topologies.DeviceIntervalTopology) =
-    ClimaComms.CUDADevice()
-Fields.error_mismatched_spaces(::Type, ::Type) = nothing # causes unsupported dynamic function invocation
+# Fields.error_mismatched_spaces(::Type, ::Type) = nothing # causes unsupported dynamic function invocation
 
 const C1 = Geometry.Covariant1Vector
 const C2 = Geometry.Covariant2Vector
@@ -91,12 +89,14 @@ Base.Broadcast.broadcastable(x::RayleighSponge) = tuple(x)
 
 function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s)
     s isa Nothing && return NullBroadcasted()
-    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
-    zmax = z_max(axes(ᶠz))
+    ᶜz = Fields.coordinate_field(axes(ᶜuₕ)).z
+    ᶠz = Fields.coordinate_field(Spaces.face_space(axes(ᶜuₕ))).z
+    # (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
+    zmax = Spaces.z_max(axes(ᶠz))
     return @. lazy(-β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ)
 end
 
-function compute_kinetic(uₕ::Fields.Field, uᵥ::Fields.Field)
+function compute_kinetic(uₕ, uᵥ)
     @assert eltype(uₕ) <: Union{C1, C2, C12}
     @assert eltype(uᵥ) <: C3
     FT = Spaces.undertype(axes(uₕ))
@@ -167,9 +167,9 @@ function ᶜimplicit_tendency_bc(ᶜY, ᶠY, p, t)
     ᶠu³ = @. lazy(ᶠuₕ³ + CT3(ᶠu₃))
     tend_ρ_1 = @. lazy(ᶜdivᵥ(ᶠwinterp(ᶜJ, ᶜρ) * ᶠuₕ³))
     tend_ρe_tot_1 = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, dt, Val(:none))
-    ᶜuₕ₀ = (zero(eltype(ᶜuₕ)),)
+    ᶜuₕ₀ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge)
 
-    return @. lazy(ᶜtendencies(-tend_ρ_1, - ᶜuₕ₀, tend_ρe_tot_1))
+    return @. lazy(ᶜtendencies(-tend_ρ_1, ᶜuₕ₀, tend_ρe_tot_1))
 end
 
 function ᶠimplicit_tendency_bc(ᶜY, ᶠY, p, t)

test/Spaces/unit_spaces.jl

@@ -330,8 +330,8 @@ end
 
     @static if on_gpu
         adapted_space = Adapt.adapt(CUDA.KernelAdaptor(), c_space)
-        @test ClimaComms.context(adapted_space) == DeviceSideContext()
-        @test ClimaComms.device(adapted_space) == DeviceSideDevice()
+        @test ClimaComms.context(adapted_space) == ClimaComms.context(c_space)
+        @test ClimaComms.device(adapted_space) == ClimaComms.device(c_space)
 
         adapted_hspace = Adapt.adapt(CUDA.KernelAdaptor(), hspace)
         @test ClimaComms.context(adapted_hspace) == DeviceSideContext()