Abstract away use of ClimaComms mpicomm

CliMA · Jan 6, 2025 · 951f4f4 · 951f4f4
1 parent 81a8fa1
commit 951f4f4
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 4 deletions.
diff --git a/src/ClimaCore.jl b/src/ClimaCore.jl
@@ -27,5 +27,6 @@ include("CommonGrids/CommonGrids.jl")
 include("CommonSpaces/CommonSpaces.jl")
 
 include("deprecated.jl")
+include("compat.jl")
 
 end # module
diff --git a/src/InputOutput/InputOutput.jl b/src/InputOutput/InputOutput.jl
@@ -13,6 +13,7 @@ import ..Geometry,
     ..Hypsography
 import ..VERSION
 import ..Utilities: PlusHalf, half
+import ..climacomms_mpicomm
 
 include("writers.jl")
 include("readers.jl")

diff --git a/src/InputOutput/readers.jl b/src/InputOutput/readers.jl
@@ -95,7 +95,7 @@ function HDF5Reader(
     if context isa ClimaComms.SingletonCommsContext
         file = h5open(filename, "r")
     else
-        file = h5open(filename, "r", context.mpicomm)
+        file = h5open(filename, "r", climacomms_mpicomm(context))
     end
     if !haskey(attrs(file), "ClimaCore version")
         error("Not a ClimaCore HDF5 file")

diff --git a/src/InputOutput/writers.jl b/src/InputOutput/writers.jl
@@ -52,7 +52,7 @@ function HDF5Writer(
     if context isa ClimaComms.SingletonCommsContext
         file = h5open(filename, mode)
     else
-        file = h5open(filename, mode, context.mpicomm)
+        file = h5open(filename, mode, climacomms_mpicomm(context))
     end
     # Add an attribute to the file if it doesn't already exist
     if haskey(attributes(file), "ClimaCore version")

diff --git a/src/compat.jl b/src/compat.jl
@@ -0,0 +1,3 @@
+import ClimaComms
+
+climacomms_mpicomm(ctx::ClimaComms.MPICommsContext) = ctx.mpicomm
diff --git a/test/Operators/spectralelement/benchmark_utils.jl b/test/Operators/spectralelement/benchmark_utils.jl
@@ -15,6 +15,7 @@ import ClimaCore.Geometry as Geometry
 import ClimaCore.Quadratures as Quadratures
 import ClimaCore.Fields as Fields
 import ClimaCore as CC
+import ClimaCore: climacomms_mpicomm
 import ClimaComms
 ClimaComms.@import_required_backends
 
@@ -160,9 +161,9 @@ function setup_kernel_args(ARGS::Vector{String} = ARGS)
        device isa ClimaComms.CUDADevice
         # assign GPUs based on local rank
         local_comm = ClimaComms.MPI.Comm_split_type(
-            context.mpicomm,
+            climacomms_mpicomm(context),
             ClimaComms.MPI.COMM_TYPE_SHARED,
-            ClimaComms.MPI.Comm_rank(context.mpicomm),
+            ClimaComms.MPI.Comm_rank(climacomms_mpicomm(context)),
         )
         CUDA.device!(
             ClimaComms.MPI.Comm_rank(local_comm) % length(CUDA.devices()),
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		import ClimaComms

		climacomms_mpicomm(ctx::ClimaComms.MPICommsContext) = ctx.mpicomm