add SplitType (#584)

Author: simonbyrne

src/comm.jl

@@ -164,26 +164,57 @@ function Comm_dup(comm::Comm)
 end
 
 """
-    Comm_split(comm::Comm, color::Integer, key::Integer)
+    Comm_split(comm::Comm, color::Union{Integer,Nothing}, key::Integer)
+
+Partition the communicator `comm`, one for each value of `color`, returning a
+new communicator. Within each group, the processes are ranked in the order of
+`key`, with ties broken by the order of `comm`.
+
+`color` should be a non-negative integer, or `nothing`, in which case a null
+communicator is returned for that rank.
 
 # External links
 $(_doc_external("MPI_Comm_split"))
 """
-function Comm_split(comm::Comm, color::Integer, key::Integer)
+function Comm_split(comm::Comm, color::Union{Integer, Nothing}, key::Integer)
+    if isnothing(color)
+        color = Consts.MPI_UNDEFINED[]
+    end
     newcomm = Comm()
     @mpichk ccall((:MPI_Comm_split, libmpi), Cint,
         (MPI_Comm, Cint, Cint, Ptr{MPI_Comm}), comm, color, key, newcomm)
     finalizer(free, newcomm)
     newcomm
 end
 
+mutable struct SplitType
+    val::Cint
+end
+const COMM_TYPE_SHARED = SplitType(-1)
+add_load_time_hook!(() -> COMM_TYPE_SHARED.val     = Consts.MPI_COMM_TYPE_SHARED[])
+
+
 """
-    Comm_split_type(comm::Comm, split_type::Integer, key::Integer; kwargs...)
+    Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)
+
+Partitions the communicator `comm` based on `split_type`, returning a new
+communicator. Within each group, the processes are ranked in the order of
+`key`, with ties broken by the order of `comm`.
+
+Currently only one `split_type` is provided:
+
+- `MPI.COMM_TYPE_SHARED`: splits the communicator into subcommunicators, each of
+  which can create a shared memory region.
 
 # External links
 $(_doc_external("MPI_Comm_split_type"))
 """
-function Comm_split_type(comm::Comm,split_type::Integer,key::Integer; kwargs...)
+function Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)
+    if isnothing(split_type)
+        split_type = Consts.MPI_UNDEFINED[]
+    elseif split_type isa SplitType
+        split_type = split_type.val
+    end
     newcomm = Comm()
     @mpichk ccall((:MPI_Comm_split_type, libmpi), Cint,
           (MPI_Comm, Cint, Cint, MPI_Info, Ptr{MPI_Comm}),

test/test_comm.jl

@@ -4,7 +4,11 @@ using MPI
 MPI.Init()
 
 comm = MPI.COMM_WORLD
+comm_rank = MPI.Comm_rank(comm)
+comm_size = MPI.Comm_size(comm)
+
 @test MPI.Comm_compare(comm, comm) == MPI.IDENT
+@test comm == comm
 @test MPI.Comm_compare(comm, MPI.COMM_SELF) == (MPI.Comm_size(comm) == 1 ? MPI.CONGRUENT : MPI.UNEQUAL)
 MPI.Barrier(comm)
 comm2 = MPI.Comm_dup(comm)
@@ -18,5 +22,25 @@ MPI.Barrier(comm3)
 # Don't free comm2
 MPI.free(comm3)
 
+splitcomm1 = MPI.Comm_split(comm, comm_rank % 2, 0)
+@test div(comm_size, 2) <= MPI.Comm_size(splitcomm1) <= div(comm_size+1, 2)
+@test MPI.Comm_rank(splitcomm1) == div(comm_rank, 2)
+
+splitcomm2 = MPI.Comm_split(comm, comm_rank < 2 ? 0 : nothing, 0)
+if comm_rank < 2
+    @test MPI.Comm_size(splitcomm2) <= 2
+    @test MPI.Comm_rank(splitcomm2) == comm_rank
+else
+    @test splitcomm2 == MPI.COMM_NULL
+end
+
+splitcomm3 = MPI.Comm_split_type(comm, comm_rank < 2 ? MPI.COMM_TYPE_SHARED : nothing, 0)
+if comm_rank < 2
+    @test MPI.Comm_size(splitcomm3) <= 2
+    @test MPI.Comm_rank(splitcomm3) < 2
+else
+    @test splitcomm3 == MPI.COMM_NULL
+end
+
 MPI.Finalize()
 @test MPI.Finalized()