Merge #1222

1222: Deprecate comm_context, use ClimaComms.context r=charleskawczynski a=charleskawczynski

This PR deprecates
 - `comm_context` in favor of using `ClimaComms.context`
and methods that assume the device is internally determined using `CUDA.functional()`, which is problematic if CUDA is available but we want to run on the CPU:
 - `HDF5Reader`
 - `HDF5Writer`
 - `Topology2D`
 - (fixes `create_dss_buffer` to use the topology context)

A step towards #1170.

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

Author: bors[bot]

src/ClimaCore.jl

@@ -3,8 +3,6 @@ module ClimaCore
 using PkgVersion
 const VERSION = PkgVersion.@Version
 
-function comm_context end
-
 include("interface.jl")
 include("Utilities/Utilities.jl")
 include("RecursiveApply/RecursiveApply.jl")

src/Fields/Fields.jl

@@ -1,6 +1,5 @@
 module Fields
 
-import ..comm_context
 import ClimaComms
 import ..enable_threading
 import ..slab, ..slab_args, ..column, ..column_args, ..level
@@ -38,18 +37,18 @@ Field(values::V, space::S) where {V <: AbstractData, S <: AbstractSpace} =
 Field(::Type{T}, space::S) where {T, S <: AbstractSpace} =
     Field(similar(Spaces.coordinates_data(space), T), space)
 
-comm_context(field::Field) = comm_context(axes(field))
+ClimaComms.context(field::Field) = ClimaComms.context(axes(field))
 
-comm_context(space::Spaces.ExtrudedFiniteDifferenceSpace) =
-    comm_context(space.horizontal_space)
-comm_context(space::Spaces.SpectralElementSpace2D) =
-    comm_context(space.topology)
-comm_context(space::S) where {S <: Spaces.AbstractSpace} =
-    ClimaComms.SingletonCommsContext()
+ClimaComms.context(space::Spaces.ExtrudedFiniteDifferenceSpace) =
+    ClimaComms.context(space.horizontal_space)
+ClimaComms.context(space::Spaces.SpectralElementSpace2D) =
+    ClimaComms.context(space.topology)
+ClimaComms.context(space::S) where {S <: Spaces.AbstractSpace} =
+    ClimaComms.context(space.topology)
 
-comm_context(topology::Topologies.Topology2D) = topology.context
-comm_context(topology::T) where {T <: Topologies.AbstractTopology} =
-    ClimaComms.SingletonCommsContext()
+ClimaComms.context(topology::Topologies.Topology2D) = topology.context
+ClimaComms.context(topology::T) where {T <: Topologies.AbstractTopology} =
+    topology.context
 
 Adapt.adapt_structure(to, field::Field) = Field(
     Adapt.adapt(to, Fields.field_values(field)),

src/Fields/mapreduce.jl

@@ -43,7 +43,7 @@ function Base.sum(
     field::Union{Field, Base.Broadcast.Broadcasted{<:FieldStyle}},
     ::ClimaComms.CPUDevice,
 )
-    context = comm_context(axes(field))
+    context = ClimaComms.context(axes(field))
     data_sum = DataLayouts.DataF(local_sum(field))
     ClimaComms.allreduce!(context, parent(data_sum), +)
     return data_sum[]
@@ -62,7 +62,7 @@ Approximate maximum of `v` or `f.(v)` over the domain.
 If `v` is a distributed field, this uses a `ClimaComms.allreduce` operation.
 """
 function Base.maximum(fn, field::Field, ::ClimaComms.CPUDevice)
-    context = comm_context(axes(field))
+    context = ClimaComms.context(axes(field))
     data_max = DataLayouts.DataF(mapreduce(fn, max, todata(field)))
     ClimaComms.allreduce!(context, parent(data_max), max)
     return data_max[]
@@ -74,7 +74,7 @@ Base.maximum(fn, field::Field) =
 Base.maximum(field::Field) = Base.maximum(field, ClimaComms.device(field))
 
 function Base.minimum(fn, field::Field, ::ClimaComms.CPUDevice)
-    context = comm_context(axes(field))
+    context = ClimaComms.context(axes(field))
     data_min = DataLayouts.DataF(mapreduce(fn, min, todata(field)))
     ClimaComms.allreduce!(context, parent(data_min), min)
     return data_min[]
@@ -109,7 +109,7 @@ function Statistics.mean(
     ::ClimaComms.CPUDevice,
 )
     space = axes(field)
-    context = comm_context(space)
+    context = ClimaComms.context(space)
     data_combined =
         DataLayouts.DataF((local_sum(field), Spaces.local_area(space)))
     ClimaComms.allreduce!(context, parent(data_combined), +)

src/InputOutput/readers.jl

@@ -92,9 +92,14 @@ struct HDF5Reader{C <: ClimaComms.AbstractCommsContext}
     space_cache::Dict{Any, Any}
 end
 
+@deprecate HDF5Reader(filename::AbstractString) HDF5Reader(
+    filename,
+    ClimaComms.SingletonCommsContext(),
+)
+
 function HDF5Reader(
     filename::AbstractString,
-    context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(),
+    context::ClimaComms.AbstractCommsContext,
 )
     if context isa ClimaComms.SingletonCommsContext
         file = h5open(filename, "r")

src/InputOutput/writers.jl

@@ -37,9 +37,14 @@ struct HDF5Writer{C <: ClimaComms.AbstractCommsContext} <: AbstractWriter
     cache::Dict{String, String}
 end
 
+@deprecate HDF5Writer(filename::AbstractString) HDF5Writer(
+    filename,
+    ClimaComms.SingletonCommsContext(),
+)
+
 function HDF5Writer(
     filename::AbstractString,
-    context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(),
+    context::ClimaComms.AbstractCommsContext,
 )
     if context isa ClimaComms.SingletonCommsContext
         file = h5open(filename, "w")

src/Spaces/Spaces.jl

@@ -67,7 +67,7 @@ weights ``W_i`` multiplied by the Jacobian determinants ``J_i``:
 If `space` is distributed, this uses a `ClimaComms.allreduce` operation.
 """
 area(space::Spaces.AbstractSpace) =
-    ClimaComms.allreduce(comm_context(space), local_area(space), +)
+    ClimaComms.allreduce(ClimaComms.context(space), local_area(space), +)
 
 ClimaComms.array_type(space::AbstractSpace) =
     ClimaComms.array_type(ClimaComms.device(space))

src/Spaces/dss.jl

@@ -58,9 +58,7 @@ function create_dss_buffer(
     local_geometry = nothing,
     local_weights = nothing,
 ) where {S, Nij}
-    context =
-        topology isa Topologies.Topology2D ? topology.context :
-        ClimaComms.SingletonCommsContext()
+    context = topology.context
     DA = ClimaComms.array_type(topology)
     convert_to_array = DA isa Array ? false : true
     (_, _, _, Nv, nelems) = Base.size(data)

src/Spaces/finitedifference.jl

@@ -212,12 +212,14 @@ Base.@propagate_inbounds function level(
     v::PlusHalf,
 )
     @inbounds local_geometry = level(local_geometry_data(space), v.i + 1)
-    PointSpace(local_geometry)
+    context = ClimaComms.context(space)
+    PointSpace(context, local_geometry)
 end
 Base.@propagate_inbounds function level(
     space::CenterFiniteDifferenceSpace,
     v::Int,
 )
     local_geometry = level(local_geometry_data(space), v)
-    PointSpace(local_geometry)
+    context = ClimaComms.context(space)
+    PointSpace(context, local_geometry)
 end

src/Spaces/pointspace.jl

@@ -7,20 +7,35 @@ local_geometry_data(space::AbstractPointSpace) = space.local_geometry
 
 A zero-dimensional space.
 """
-struct PointSpace{LG} <: AbstractPointSpace
+struct PointSpace{C <: ClimaComms.AbstractCommsContext, LG} <:
+       AbstractPointSpace
+    context::C
     local_geometry::LG
 end
 
-ClimaComms.device(space::PointSpace) = ClimaComms.CPUDevice()
+ClimaComms.device(space::PointSpace) = ClimaComms.device(space.context)
+ClimaComms.context(space::PointSpace) = space.context
 
-function PointSpace(local_geometry::LG) where {LG <: Geometry.LocalGeometry}
+@deprecate PointSpace(x::Geometry.LocalGeometry) PointSpace(
+    ClimaComms.SingletonCommsContext(ClimaComms.CPUDevice()),
+    x,
+) false
+
+function PointSpace(
+    context::ClimaComms.AbstractCommsContext,
+    local_geometry::LG,
+) where {LG <: Geometry.LocalGeometry}
     FT = Geometry.undertype(LG)
+    # TODO: inherit array type
     local_geometry_data = DataLayouts.DataF{LG}(Array{FT})
     local_geometry_data[] = local_geometry
-    return PointSpace(local_geometry_data)
+    return PointSpace(context, local_geometry_data)
 end
 
-function PointSpace(coord::Geometry.Abstract1DPoint{FT}) where {FT}
+function PointSpace(
+    context::ClimaComms.AbstractCommsContext,
+    coord::Geometry.Abstract1DPoint{FT},
+) where {FT}
     CoordType = typeof(coord)
     AIdx = Geometry.coordinate_axis(CoordType)
     local_geometry = Geometry.LocalGeometry(
@@ -32,5 +47,5 @@ function PointSpace(coord::Geometry.Abstract1DPoint{FT}) where {FT}
             FT(1.0),
         ),
     )
-    return PointSpace(local_geometry)
+    return PointSpace(context, local_geometry)
 end

src/Spaces/spectralelement.jl

@@ -644,14 +644,16 @@ Base.@propagate_inbounds slab(space::AbstractSpectralElementSpace, h) =
 
 Base.@propagate_inbounds function column(space::SpectralElementSpace1D, i, h)
     local_geometry = column(local_geometry_data(space), i, h)
-    PointSpace(local_geometry)
+    context = ClimaComms.context(space)
+    PointSpace(context, local_geometry)
 end
 Base.@propagate_inbounds column(space::SpectralElementSpace1D, i, j, h) =
     column(space, i, h)
 
 Base.@propagate_inbounds function column(space::SpectralElementSpace2D, i, j, h)
     local_geometry = column(local_geometry_data(space), i, j, h)
-    PointSpace(local_geometry)
+    context = ClimaComms.context(space)
+    PointSpace(context, local_geometry)
 end
 
 # XXX: this cannot take `space` as it must be constructed beforehand so