Additional threading functions and docs (#364)

Expose `MPI_Query_thread` and `MPI_Is_thread_main`, along with some doc cleanup.

Author: simonbyrne

docs/make.jl

@@ -33,6 +33,7 @@ for (example_title, example_md) in EXAMPLES
     end
 end
 
+DocMeta.setdocmeta!(MPI, :DocTestSetup, :(using MPI); recursive=true)
 
 makedocs(
     sitename = "MPI.jl",

docs/src/environment.md

@@ -6,12 +6,20 @@
 mpiexec
 ```
 
+## Enums
+
+```@docs
+MPI.ThreadLevel
+```
+
 ## Functions
 
 ```@docs
 MPI.Abort
 MPI.Init
 MPI.Init_thread
+MPI.Query_thread
+MPI.Is_thread_main
 MPI.Initialized
 MPI.Finalize
 MPI.Finalized

docs/src/installation.md

@@ -23,12 +23,12 @@ For other implementations, the build script will attempt to build a small C prog
 determine the appropriate type definitions and constants. This requires a compatible C
 compiler (`mpicc` by default).
 
-## Environment variables
+## [Environment variables](@id environment_variables)
 
 The following optional environment variables can be used to control certain aspects of the
 build script and other library behaviour:
 
-- `JULIA_MPI_ABI`: the application binary interface, a string matching an [`MPIABI`](@ref) value.
+- `JULIA_MPI_ABI`: the application binary interface, a string matching an [`MPI.MPIABI`](@ref) value.
 - `JULIA_MPI_PATH`: the top-level installation directory of MPI.
 - `JULIA_MPI_LIBRARY`: the path of the MPI shared library.
 - `JULIA_MPI_LIBRARY_PATH`: the directory containing the MPI library files.

docs/src/library.md

@@ -21,4 +21,5 @@ MPI.MPIABI
 
 ```@docs
 MPI.has_cuda
+MPI.identify_implementation
 ```

docs/src/topology.md

@@ -6,6 +6,7 @@ MPI.Cart_coords!
 MPI.Cart_create
 MPI.Cart_get
 MPI.Cart_rank
+MPI.Cart_shift
 MPI.Cart_sub
 MPI.Cartdim_get
 MPI.Dims_create!

src/environment.jl

@@ -55,30 +55,46 @@ function Init()
     end
 end
 
+"""
+    ThreadLevel
+
+An Enum denoting the level of threading support in the current process:
+
+ - `MPI.THREAD_SINGLE`: Only one thread will execute.
+
+ - `MPI.THREAD_FUNNELED`: The process may be multi-threaded, but the application must
+   ensure that only the main thread makes MPI calls. See [`Is_thread_main`](@ref).
+
+ - `MPI.THREAD_SERIALIZED`: The process may be multi-threaded, and multiple threads may
+   make MPI calls, but only one at a time (i.e. all MPI calls are serialized).
+
+ - `MPI.THREAD_MULTIPLE`: Multiple threads may call MPI, with no restrictions.
+
+# See also
+
+- [`Init_thread`](@ref)
+- [`Query_thread`](@ref)
+"""
 @enum ThreadLevel begin
     THREAD_SINGLE     = MPI_THREAD_SINGLE
     THREAD_FUNNELED   = MPI_THREAD_FUNNELED
     THREAD_SERIALIZED = MPI_THREAD_SERIALIZED
     THREAD_MULTIPLE   = MPI_THREAD_MULTIPLE
 end
-    
+
 
 """
     Init_thread(required::ThreadLevel)
 
 Initialize MPI and the MPI thread environment in the current process. The argument
-specifies the required thread level, which is one of the following:
+specifies the required level of threading support, see [`ThreadLevel`](@ref).
 
- - `MPI.THREAD_SINGLE`: Only one thread will execute.
- - `MPI.THREAD_FUNNELED`: The process may be multi-threaded, but the application must ensure that only the main thread makes MPI calls.
- - `MPI.THREAD_SERIALIZED`: The process may be multi-threaded, and multiple threads may make MPI calls, but only one at a time (i.e. all MPI calls are serialized).
- - `MPI.THREAD_MULTIPLE`: Multiple threads may call MPI, with no restrictions.
+The function will return the provided `ThreadLevel`, and values may be compared via
+inequalities, i.e.
 
-Tne function will return the provided `ThreadLevel`, and values may be compared via inequalities, i.e.
 ```julia
-if Init_thread(required) < required
-    error("Insufficient threading")
-end
+provided = Init_thread(required)
+@assert provided >= required
 ```
 
 All MPI programs must contain exactly one call to [`MPI.Init`](@ref) or
@@ -102,7 +118,7 @@ function Init_thread(required::ThreadLevel)
     if provided < required
         @warn "Thread level requested = $required, provided = $provided"
     end
-    
+
     REFCOUNT[] = 1
     atexit(refcount_dec)
 
@@ -112,6 +128,41 @@ function Init_thread(required::ThreadLevel)
     return provided
 end
 
+"""
+    Query_thread()
+
+Query the level of threading support in the current process.
+Returns a [`ThreadLevel`](@ref) value denoting
+
+# External links
+$(_doc_external("MPI_Query_thread"))
+"""
+function Query_thread()
+    r_provided = Ref{ThreadLevel}()
+
+    # int MPI_Query_thread(int *provided)
+    @mpichk ccall((:MPI_Query_thread, libmpi), Cint,
+                  (Ref{ThreadLevel},), r_provided)
+    return r_provided[]
+end
+
+"""
+    Is_thread_main()
+
+Queries whether the current thread is the main thread according to MPI. This can be called
+by any thread, and is useful for the  `THREAD_FUNNELED` [`ThreadLevel`](@ref).
+
+# External links
+$(_doc_external("MPI_Is_thread_main"))
+"""
+function Is_thread_main()
+    r_flag = Ref{Cint}()
+    # int MPI_Is_thread_main(int *flag)
+    @mpichk ccall((:MPI_Is_thread_main, libmpi), Cint,
+                  (Ref{Cint},), r_flag)
+    return r_flag[] != 0
+end
+
 
 """
     Finalize()
@@ -207,7 +258,7 @@ function Wtick()
         ccall((:MPI_Wtick, libmpi), stdcall, Cdouble, ())
     else
         ccall((:MPI_Wtick, libmpi), Cdouble, ())
-    end        
+    end
 end
 
 function Wtime()

src/implementations.jl

@@ -1,19 +1,19 @@
 
 function Get_library_version()
-    # There is no way to query at runtime what the length of the buffer should be.  
-    # https://github.com/mpi-forum/mpi-issues/issues/159    
+    # There is no way to query at runtime what the length of the buffer should be.
+    # https://github.com/mpi-forum/mpi-issues/issues/159
     # 8192 is the maximum value of MPI_MAX_LIBRARY_VERSION_STRING across known
     # implementations.
     buf = Array{UInt8}(undef, 8192)
     buflen = Ref{Cint}()
 
-    # Microsoft MPI uses stdcall calling convention    
-    libfilename, = split(basename(libmpi),'.')    
+    # Microsoft MPI uses stdcall calling convention
+    libfilename, = split(basename(libmpi),'.')
     if libfilename == "msmpi"
         ccall((:MPI_Get_library_version, libmpi), stdcall, Cint, (Ptr{UInt8}, Ref{Cint}), buf, buflen)
     else
         ccall((:MPI_Get_library_version, libmpi), Cint, (Ptr{UInt8}, Ref{Cint}), buf, buflen)
-    end        
+    end
     resize!(buf, buflen[])
     return String(buf)
 end
@@ -60,10 +60,15 @@ end
 
 An enum corresponding to known MPI Application Binary Interfaces (ABI)
 
-- `UnknownABI`: unable to determine MPI ABI. This 
-- `MPICHABI`: Compatible with [MPICH ABI Compatibility Initiative](https://www.mpich.org/abi/).
-- `OpenMPIABI`: Compatible with Open MPI
-- `MicrosftMPIABI`: Compatible with Microsoft MPI
+- `UnknownABI`: unable to determine MPI ABI. MPI.jl will attempt to build a small C
+  program to determine the necessary constants and type information.
+
+- `MPICHABI`: Compatible with [MPICH ABI Compatibility
+  Initiative](https://www.mpich.org/abi/).
+
+- `OpenMPIABI`: Compatible with [Open MPI](https://www.open-mpi.org/).
+
+- `MicrosftMPIABI`: Compatible with [Microsoft MPI](https://docs.microsoft.com/en-us/message-passing-interface/microsoft-mpi).
 """
 @enum MPIABI begin
     UnknownABI
@@ -80,7 +85,10 @@ Attempt to identify the MPI implementation based on
 
 - `impl`: a value of type [`MPIImpl`](@ref)
 - `version`: a `VersionNumber` of the library, or `nothing` if it cannot be determined.
-- `abi`: a value of [`MPIABI`](@ref). This can be overridden by the `JULIA_MPI_ABI` environment variable.
+- `abi`: a value of [`MPIABI`](@ref). This can be overridden by the `JULIA_MPI_ABI` [environment variable](@ref environment_variables).
+
+This function is only intended for internal use. Users should use [`MPI_LIBRARY`](@ref),
+[`MPI_LIBRARY_VERSION`](@ref) or [`MPI_LIBRARY_ABI`](@ref).
 """
 function identify_implementation()
     impl = UnknownMPI
@@ -95,7 +103,7 @@ function identify_implementation()
             if version >= v"3.1"
                 abi = MPICHABI
             end
-        end        
+        end
 
     elseif startswith(MPI_LIBRARY_VERSION_STRING, "Open MPI")
         # Open MPI / Spectrum MPI
@@ -107,7 +115,7 @@ function identify_implementation()
         abi = OpenMPIABI
 
     elseif startswith(MPI_LIBRARY_VERSION_STRING, "Microsoft MPI")
-        impl = MicrosoftMPI        
+        impl = MicrosoftMPI
         # "Microsoft MPI %u.%u.%u.%u%S"
         # ignore last 2 (build numbers)
         if (m = match(r"^Microsoft MPI v(\d+.\d+)", MPI_LIBRARY_VERSION_STRING)) !== nothing
@@ -118,7 +126,7 @@ function identify_implementation()
     elseif startswith(MPI_LIBRARY_VERSION_STRING, "Intel")
         impl = IntelMPI
 
-        # TODO: figure out how to parse        
+        # TODO: figure out how to parse
         # "Intel(R) MPI Library 2019 Update 4 for Linux* OS"
         if (m = match(r"^Intel\(R\) MPI Library (\d+)", MPI_LIBRARY_VERSION_STRING)) !== nothing
             version = VersionNumber(m.captures[1])
@@ -141,11 +149,11 @@ function identify_implementation()
     if (abienv = get(ENV, "JULIA_MPI_ABI", nothing)) !== nothing
         for inst in instances(MPIABI)
             if String(Symbol(inst)) == abienv
-                abi = inst                
+                abi = inst
             end
         end
     end
-        
+
     return impl, version, abi
 end
 
@@ -154,7 +162,7 @@ const MPI_LIBRARY, MPI_LIBRARY_VERSION, MPI_LIBRARY_ABI = identify_implementatio
 """
     MPI_LIBRARY :: MPIImpl
 
-The current MPI implementation: this is determined by 
+The current MPI implementation: this is determined by
 
 # See also
 - [`MPIImpl`](@ref)

test/test_threads.jl

@@ -10,6 +10,10 @@ end
 
 provided = MPI.Init_thread(MPI.THREAD_MULTIPLE)
 
+@test MPI.THREAD_SINGLE <= provided <= MPI.THREAD_MULTIPLE
+@test MPI.Query_thread() == provided
+@test MPI.Is_thread_main()
+
 comm = MPI.COMM_WORLD
 size = MPI.Comm_size(comm)
 rank = MPI.Comm_rank(comm)
@@ -24,7 +28,7 @@ if provided == MPI.THREAD_MULTIPLE
     recv_arr = zeros(N)
 
     reqs = Array{MPI.Request}(undef, 2N)
-    
+
     Threads.@threads for i = 1:N
         reqs[N+i] = MPI.Irecv!(@view(recv_arr[i:i]), src, i, comm)
         reqs[i] = MPI.Isend(@view(send_arr[i:i]), dst, i, comm)
@@ -33,6 +37,6 @@ if provided == MPI.THREAD_MULTIPLE
     MPI.Waitall!(reqs)
 
     @test recv_arr == send_arr
-end        
+end
 
 MPI.Finalize()