wip

ext/ClimaCommsMPIExt.jl

@@ -271,4 +271,32 @@ function ClimaComms.finish(
     MPI.Waitall(ghost.send_reqs)
 end
 
+function Base.summary(io::IO, ctx::ClimaComms.MPICommsContext)
+    if !MPI.Initialized()
+        ClimaComms.iamroot(ctx) && @warn "MPI is not initialized."
+        return nothing
+    end
+
+    if ClimaComms.iamroot(ctx)
+        println(io, "Context: $(typeof(ctx).name.name)")
+        println(io, "Device: $(typeof(ctx.device))")
+        println(io, "Total Processes: $(ClimaComms.nprocs(ctx))")
+    end
+    ClimaComms.barrier(ctx)
+    rank = MPI.Comm_rank(ctx.mpicomm)
+    node_name = MPI.Get_processor_name()
+
+    if ctx.device isa ClimaComms.CUDADevice
+        # Match GPU rank assignment - should this live in a function to avoid code duplication?
+        local_comm = MPI.Comm_split_type(ctx.mpicomm, MPI.COMM_TYPE_SHARED, rank)
+        local_rank = MPI.Comm_rank(local_comm)
+        local_size = MPI.Comm_size(local_comm)
+        print(io, "Rank: $rank, Local Rank: $local_rank, Node: $node_name")
+
+        MPI.free(local_comm)
+    else
+        print(io, "Rank: $rank, Node: $node_name")
+    end
+end
+
 end

src/singleton.jl

@@ -49,3 +49,9 @@ graph_context(ctx::SingletonCommsContext, args...) = SingletonGraphContext(ctx)
 start(gctx::SingletonGraphContext) = nothing
 progress(gctx::SingletonGraphContext) = nothing
 finish(gctx::SingletonGraphContext) = nothing
+
+function Base.summary(io::IO, ctx::SingletonCommsContext) 
+    println(io)
+    println(io, "Context: $(typeof(ctx).name.name)")
+    println(io, "Device: $(typeof(ctx.device))")
+end