Skip to content

Commit eefb427

Browse files
authored
add SplitType (#584)
1 parent 5ee8aa1 commit eefb427

File tree

2 files changed

+59
-4
lines changed

2 files changed

+59
-4
lines changed

src/comm.jl

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,26 +164,57 @@ function Comm_dup(comm::Comm)
164164
end
165165

166166
"""
167-
Comm_split(comm::Comm, color::Integer, key::Integer)
167+
Comm_split(comm::Comm, color::Union{Integer,Nothing}, key::Integer)
168+
169+
Partition the communicator `comm`, one for each value of `color`, returning a
170+
new communicator. Within each group, the processes are ranked in the order of
171+
`key`, with ties broken by the order of `comm`.
172+
173+
`color` should be a non-negative integer, or `nothing`, in which case a null
174+
communicator is returned for that rank.
168175
169176
# External links
170177
$(_doc_external("MPI_Comm_split"))
171178
"""
172-
function Comm_split(comm::Comm, color::Integer, key::Integer)
179+
function Comm_split(comm::Comm, color::Union{Integer, Nothing}, key::Integer)
180+
if isnothing(color)
181+
color = Consts.MPI_UNDEFINED[]
182+
end
173183
newcomm = Comm()
174184
@mpichk ccall((:MPI_Comm_split, libmpi), Cint,
175185
(MPI_Comm, Cint, Cint, Ptr{MPI_Comm}), comm, color, key, newcomm)
176186
finalizer(free, newcomm)
177187
newcomm
178188
end
179189

190+
mutable struct SplitType
191+
val::Cint
192+
end
193+
const COMM_TYPE_SHARED = SplitType(-1)
194+
add_load_time_hook!(() -> COMM_TYPE_SHARED.val = Consts.MPI_COMM_TYPE_SHARED[])
195+
196+
180197
"""
181-
Comm_split_type(comm::Comm, split_type::Integer, key::Integer; kwargs...)
198+
Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)
199+
200+
Partitions the communicator `comm` based on `split_type`, returning a new
201+
communicator. Within each group, the processes are ranked in the order of
202+
`key`, with ties broken by the order of `comm`.
203+
204+
Currently only one `split_type` is provided:
205+
206+
- `MPI.COMM_TYPE_SHARED`: splits the communicator into subcommunicators, each of
207+
which can create a shared memory region.
182208
183209
# External links
184210
$(_doc_external("MPI_Comm_split_type"))
185211
"""
186-
function Comm_split_type(comm::Comm,split_type::Integer,key::Integer; kwargs...)
212+
function Comm_split_type(comm::Comm, split_type, key::Integer; kwargs...)
213+
if isnothing(split_type)
214+
split_type = Consts.MPI_UNDEFINED[]
215+
elseif split_type isa SplitType
216+
split_type = split_type.val
217+
end
187218
newcomm = Comm()
188219
@mpichk ccall((:MPI_Comm_split_type, libmpi), Cint,
189220
(MPI_Comm, Cint, Cint, MPI_Info, Ptr{MPI_Comm}),

test/test_comm.jl

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@ using MPI
44
MPI.Init()
55

66
comm = MPI.COMM_WORLD
7+
comm_rank = MPI.Comm_rank(comm)
8+
comm_size = MPI.Comm_size(comm)
9+
710
@test MPI.Comm_compare(comm, comm) == MPI.IDENT
11+
@test comm == comm
812
@test MPI.Comm_compare(comm, MPI.COMM_SELF) == (MPI.Comm_size(comm) == 1 ? MPI.CONGRUENT : MPI.UNEQUAL)
913
MPI.Barrier(comm)
1014
comm2 = MPI.Comm_dup(comm)
@@ -18,5 +22,25 @@ MPI.Barrier(comm3)
1822
# Don't free comm2
1923
MPI.free(comm3)
2024

25+
splitcomm1 = MPI.Comm_split(comm, comm_rank % 2, 0)
26+
@test div(comm_size, 2) <= MPI.Comm_size(splitcomm1) <= div(comm_size+1, 2)
27+
@test MPI.Comm_rank(splitcomm1) == div(comm_rank, 2)
28+
29+
splitcomm2 = MPI.Comm_split(comm, comm_rank < 2 ? 0 : nothing, 0)
30+
if comm_rank < 2
31+
@test MPI.Comm_size(splitcomm2) <= 2
32+
@test MPI.Comm_rank(splitcomm2) == comm_rank
33+
else
34+
@test splitcomm2 == MPI.COMM_NULL
35+
end
36+
37+
splitcomm3 = MPI.Comm_split_type(comm, comm_rank < 2 ? MPI.COMM_TYPE_SHARED : nothing, 0)
38+
if comm_rank < 2
39+
@test MPI.Comm_size(splitcomm3) <= 2
40+
@test MPI.Comm_rank(splitcomm3) < 2
41+
else
42+
@test splitcomm3 == MPI.COMM_NULL
43+
end
44+
2145
MPI.Finalize()
2246
@test MPI.Finalized()

0 commit comments

Comments
 (0)