Add universe_size function (#280)

Exposes `MPI_UNIVERSE_SIZE` attribute of `MPI_COMM_WORLD`
JuliaParallel · Jul 1, 2019 · 7de5bf3 · 7de5bf3
1 parent 4e6ad89
commit 7de5bf3
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 1 deletion.
diff --git a/deps/consts_msmpi.jl b/deps/consts_msmpi.jl
@@ -82,6 +82,7 @@ const MPI_TAG_UB = reinterpret(Cint, 0x64400001)
 const MPI_COMM_TYPE_SHARED = Cint(1)
 const MPI_ORDER_C = Cint(56)
 const MPI_ORDER_FORTRAN = Cint(57)
+const MPI_UNIVERSE_SIZE = reinterpret(Cint, 0x64400009)
 
 const MPI_BOTTOM = reinterpret(SentinelPtr, 0)
 const MPI_IN_PLACE = reinterpret(SentinelPtr, -1)

diff --git a/deps/gen_consts.jl b/deps/gen_consts.jl
@@ -95,6 +95,7 @@ MPI_Cints = [
     :MPI_COMM_TYPE_SHARED,
     :MPI_ORDER_C,
     :MPI_ORDER_FORTRAN,
+    :MPI_UNIVERSE_SIZE
 ]
 
 MPI_pointers = [

diff --git a/src/comm.jl b/src/comm.jl
@@ -95,3 +95,21 @@ function Intercomm_merge(intercomm::Comm, flag::Bool)
     return newcomm
 end
 
+"""
+    universe_size()
+
+The total number of available slots, or `nothing` if it is not defined. This is determined by the `MPI_UNIVERSE_SIZE` attribute of `COMM_WORLD`.
+
+This is typically dependent on the MPI implementation: for MPICH-based implementations, this is specified by the `-usize` argument. OpenMPI defines a default value based on the number of processes available.
+"""
+function universe_size()
+    flag = Ref{Cint}()
+    result = Ref(Ptr{Cint}(C_NULL))
+    # int MPI_Comm_get_attr(MPI_Comm comm, int comm_keyval, void *attribute_val, int *flag)
+    @mpichk ccall((:MPI_Comm_get_attr, libmpi), Cint,
+        (MPI_Comm, Cint, Ptr{Cvoid}, Ptr{Cint}), MPI.COMM_WORLD, MPI_UNIVERSE_SIZE, result, flag)
+    if flag[] == 0
+        return nothing
+    end
+    Int(unsafe_load(result[]))
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -24,7 +24,6 @@ function runtests()
     testdir = dirname(@__FILE__)
     istest(f) = endswith(f, ".jl") && startswith(f, "test_")
     testfiles = sort(filter(istest, readdir(testdir)))
-
     extra_args = []
     @static if !Sys.iswindows()
         if occursin( "OpenRTE", read(`$mpiexec --version`, String))
@@ -34,6 +33,7 @@ function runtests()
 
     nfail = 0
     printstyled("Running MPI.jl tests\n"; color=:white)
+
     for f in testfiles
         coverage_opt = coverage_opts[Base.JLOptions().code_coverage]
         if f ∈ singlefiles

diff --git a/test/test_universe_size.jl b/test/test_universe_size.jl
@@ -0,0 +1,9 @@
+using Test
+using MPI
+
+MPI.Init()
+
+usize = MPI.universe_size()
+@test usize === nothing || usize >= 1
+
+MPI.Finalize()