MPI Distributed Graph API (#597)

* Add MPI distributed graph API
* Add neighbor collectives

Co-authored-by: Mosè Giordano <giordano@users.noreply.github.com>

Author: termi-official

docs/examples/09-graph_communication.jl

@@ -0,0 +1,125 @@
+using Test
+using MPI
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+size = MPI.Comm_size(comm)
+rank = MPI.Comm_rank(comm)
+
+#
+# Setup the following communication graph
+#
+#   +-----+
+#   |     |
+#   v     v
+#   0<-+  3
+#   ^  |  ^
+#   |  |  |
+#   v  |  v
+#   1  +--2
+#   ^     |
+#   |     |
+#   +-----+
+#
+#
+
+if rank == 0
+    dest   = Cint[1,3]
+    degree = Cint[length(dest)]
+elseif rank == 1
+    dest   = Cint[0]
+    degree = Cint[length(dest)]
+elseif rank == 2
+    dest   = Cint[3,0,1]
+    degree = Cint[length(dest)]
+elseif rank == 3
+    dest   = Cint[0,2,1]
+    degree = Cint[length(dest)]
+end
+
+source = Cint[rank]
+graph_comm = MPI.Dist_graph_create(comm, source, degree, dest)
+
+#
+# Now send the rank across the edges.
+#
+# Version 1: use allgather primitive
+#
+
+send = [rank]
+if rank == 0
+    recv = [-1, -1, -1]
+elseif rank == 1
+    recv = [-1, -1, -1]
+elseif rank == 2
+    recv = [-1]
+elseif rank == 3
+    recv = [-1, -1]
+end
+
+MPI.Neighbor_allgather!(send, recv, graph_comm);
+
+println("rank = $(rank): $(recv)")
+
+#
+# Version 2: use alltoall primitive
+#
+
+if rank == 0
+    send   = [rank, rank]
+    recv   = [-1, -1, -1]
+elseif rank == 1
+    send = [rank]
+    recv   = [-1, -1, -1]
+elseif rank == 2
+    send = [rank, rank, rank]
+    recv   = [-1]
+elseif rank == 3
+    send = [rank, rank, rank]
+    recv   = [-1, -1]
+end
+
+MPI.Neighbor_alltoall!(UBuffer(send,1), UBuffer(recv,1), graph_comm);
+
+println("rank = $(rank): $(recv)")
+
+#
+# Now send the rank exactly rank times across the edges.
+#
+# Rank i receives i+1 values from each adjacent process
+if rank == 0
+    send       = [rank, rank,
+                  rank, rank, rank, rank]
+    send_count = [2, 4]
+
+    recv   = [-1, -1, -1]
+    recv_count = [1, 1, 1]
+elseif rank == 1
+    send       = [rank]
+    send_count = [1]
+
+    recv       = [-1, -1, -1, -1, -1, -1]
+    recv_count = [2, 2, 2]
+elseif rank == 2
+    send       = [rank, rank, rank, rank,
+                  rank,
+                  rank,rank]
+    send_count = [4, 1, 2]
+
+    recv       = [-1, -1, -1]
+    recv_count = [3]
+elseif rank == 3
+    send       = [rank, 
+                  rank, rank,rank,
+                  rank, rank]
+    send_count = [1, 3, 2]
+
+    recv       = [-1, -1, -1, -1, -1, -1, -1, -1]
+    recv_count = [4, 4]
+end
+
+MPI.Neighbor_alltoallv!(VBuffer(send,send_count), VBuffer(recv,recv_count), graph_comm);
+println("rank = $(rank): $(recv)")
+
+MPI.Finalize()

docs/make.jl

@@ -12,6 +12,7 @@ EXAMPLES = [
     "Scatterv and Gatherv" => "examples/06-scatterv.md",
     "Active RMA" => "examples/07-rma_active.md",
     "Passive RMA" => "examples/08-rma_passive.md",
+    "Graph Communication" => "examples/09-graph_communication.md",
 ]
 
 examples_md_dir = joinpath(@__DIR__,"src/examples")
@@ -36,9 +37,9 @@ for (example_title, example_md) in EXAMPLES
         println(mdfile)
 
         println(mdfile, "```")
-        println(mdfile, "> mpiexecjl -n 3 julia $example_jl")
+        println(mdfile, "> mpiexecjl -n 4 julia $example_jl")
         cd(@__DIR__) do
-            write(mdfile, mpiexec(cmd -> read(`$cmd -n 3 $(Base.julia_cmd()) --project $example_jl`)))
+            write(mdfile, mpiexec(cmd -> read(`$cmd -n 4 $(Base.julia_cmd()) --project $example_jl`)))
         end
         println(mdfile, "```")
     end

docs/src/reference/collective.md

@@ -25,6 +25,8 @@ MPI.Gatherv!
 MPI.Allgather!
 MPI.Allgather
 MPI.Allgatherv!
+MPI.Neighbor_allgather!
+MPI.Neighbor_allgatherv!
 ```
 
 ### Scatter
@@ -41,6 +43,8 @@ MPI.Scatterv!
 MPI.Alltoall!
 MPI.Alltoall
 MPI.Alltoallv!
+MPI.Neighbor_alltoall!
+MPI.Neighbor_alltoallv!
 ```
 
 ## Reduce/Scan

docs/src/reference/topology.md

@@ -9,4 +9,8 @@ MPI.Cart_rank
 MPI.Cart_shift
 MPI.Cart_sub
 MPI.Cartdim_get
+MPI.Dist_graph_create
+MPI.Dist_graph_create_adjacent
+MPI.Dist_graph_neighbors_count
+MPI.Dist_graph_neighbors!
 ```

src/collective.jl

@@ -839,3 +839,119 @@ Exscan(sendbuf::AbstractArray, op, comm::Comm) =
     Exscan!(sendbuf, similar(sendbuf), op, comm)
 Exscan(object::T, op, comm::Comm) where {T} =
     Exscan!(Ref(object), Ref{T}(), op, comm)[]
+
+"""
+    Neighbor_alltoall!(sendbuf::UBuffer, recvbuf::UBuffer, comm::Comm)
+
+Perform an all-to-all communication along the directed edges of the graph with fixed size messages.
+
+See also [`MPI.Alltoall!`](@ref).
+
+# External links
+$(_doc_external("MPI_Neighbor_alltoall"))
+"""
+function Neighbor_alltoall!(sendbuf::UBuffer, recvbuf::UBuffer, graph_comm::Comm)
+    # int MPI_Neighbor_alltoall(const void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf,
+    #       int recvcount, MPI_Datatype recvtype, MPI_Comm graph_comm)
+    @mpichk ccall((:MPI_Neighbor_alltoall, libmpi), Cint,
+                (MPIPtr, Cint, MPI_Datatype, MPIPtr, Cint, MPI_Datatype, MPI_Comm),
+                sendbuf.data, sendbuf.count, sendbuf.datatype,
+                recvbuf.data, recvbuf.count, recvbuf.datatype,
+                graph_comm) v"3.0"
+    return recvbuf.data
+end
+
+Neighbor_alltoall!(sendbuf::InPlace, recvbuf::UBuffer, graph_comm::Comm) =
+    Neighbor_alltoall!(UBuffer(IN_PLACE), recvbuf, graph_comm)
+Neighbor_alltoall!(sendrecvbuf::UBuffer, graph_comm::Comm) =
+    Neighbor_alltoall!(IN_PLACE, sendrecvbuf, comm)
+Neighbor_alltoall(sendbuf::UBuffer, graph_comm::Comm) =
+    Neighbor_alltoall!(sendbuf, similar(sendbuf), graph_comm)
+
+"""
+    Neighbor_alltoallv!(sendbuf::VBuffer, recvbuf::VBuffer, graph_comm::Comm)
+
+Perform an all-to-all communication along the directed edges of the graph with variable size messages.
+
+See also [`MPI.Alltoallv!`](@ref).
+
+# External links
+$(_doc_external("MPI_Neighbor_alltoallv"))
+"""
+function Neighbor_alltoallv!(sendbuf::VBuffer, recvbuf::VBuffer, graph_comm::Comm)
+    # int MPI_Neighbor_alltoallv!(const void* sendbuf, const int sendcounts[],
+    #       const int sdispls[], MPI_Datatype sendtype, void* recvbuf,
+    #       const int recvcounts[], const int rdispls[],
+    #       MPI_Datatype recvtype, MPI_Comm comm)
+    @mpichk ccall((:MPI_Neighbor_alltoallv, libmpi), Cint,
+                    (MPIPtr, Ptr{Cint}, Ptr{Cint}, MPI_Datatype,
+                    MPIPtr, Ptr{Cint}, Ptr{Cint}, MPI_Datatype,
+                    MPI_Comm),
+                    sendbuf.data, sendbuf.counts, sendbuf.displs, sendbuf.datatype,
+                    recvbuf.data, recvbuf.counts, recvbuf.displs, recvbuf.datatype,
+                    graph_comm) v"3.0"
+    return recvbuf.data
+end
+
+"""
+    Neighbor_allgather!(sendbuf::Buffer, recvbuf::UBuffer, comm::Comm)
+
+Perform an all-gather communication along the directed edges of the graph.
+
+See also [`MPI.Allgather!`](@ref).
+
+# External links
+$(_doc_external("MPI_Neighbor_allgather"))
+"""
+function Neighbor_allgather!(sendbuf::Buffer, recvbuf::UBuffer, graph_comm::Comm)
+    # int MPI_Neighbor_allgather(const void* sendbuf, int sendcount,
+    #       MPI_Datatype sendtype, void* recvbuf, int recvcount,
+    #       MPI_Datatype recvtype, MPI_Comm comm)
+    @mpichk ccall((:MPI_Neighbor_allgather, libmpi), Cint,
+                (MPIPtr, Cint, MPI_Datatype, MPIPtr, Cint, MPI_Datatype, MPI_Comm),
+                sendbuf.data, sendbuf.count, sendbuf.datatype,
+                recvbuf.data, recvbuf.count, recvbuf.datatype, graph_comm) v"3.0"
+    
+    return recvbuf.data
+end
+Neighbor_allgather!(sendbuf, recvbuf::UBuffer, graph_comm::Comm) =
+    Neighbor_allgather!(Buffer_send(sendbuf), recvbuf, graph_comm)
+
+Neighbor_allgather!(sendbuf::Union{Ref,AbstractArray}, recvbuf::AbstractArray, graph_comm::Comm) =
+    Neighbor_allgather!(sendbuf, UBuffer(recvbuf, length(sendbuf)), graph_comm)
+
+
+function Neighbor_allgather!(sendrecvbuf::UBuffer, graph_comm::Comm)
+    Neighbor_allgather!(IN_PLACE, sendrecvbuf, graph_comm)
+end
+
+"""
+    Neighbor_allgatherv!(sendbuf::Buffer, recvbuf::VBuffer, comm::Comm)
+
+Perform an all-gather communication along the directed edges of the graph with variable sized data.
+
+See also [`MPI.Allgatherv!`](@ref).
+
+# External links
+$(_doc_external("MPI_Neighbor_allgatherv"))
+"""
+function Neighbor_allgatherv!(sendbuf::Buffer, recvbuf::VBuffer, graph_comm::Comm)
+    # int MPI_Neighbor_allgatherv(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
+    #                             void *recvbuf, const int recvcounts[], const int displs[],
+    #                             MPI_Datatype recvtype, MPI_Comm comm)
+    @mpichk ccall((:MPI_Neighbor_allgatherv, libmpi), Cint,
+                (MPIPtr, Cint, MPI_Datatype, MPIPtr, Ptr{Cint}, Ptr{Cint}, MPI_Datatype, MPI_Comm),
+                sendbuf.data, sendbuf.count, sendbuf.datatype,
+                recvbuf.data, recvbuf.counts, recvbuf.displs, recvbuf.datatype, graph_comm) v"3.0"
+    return recvbuf.data
+end
+Neighbor_allgatherv!(sendbuf, recvbuf::VBuffer, graph_comm::Comm) =
+    Neighbor_allgatherv!(Buffer_send(sendbuf), recvbuf, graph_comm)
+
+Neighbor_allgatherv!(sendbuf::Union{Ref,AbstractArray}, recvbuf::AbstractArray, graph_comm::Comm) =
+    Neighbor_allgatherv!(sendbuf, VBuffer(recvbuf, length(sendbuf)), graph_comm)
+
+
+function Neighbor_allgatherv!(sendrecvbuf::VBuffer, graph_comm::Comm)
+    Neighbor_allgatherv!(IN_PLACE, sendrecvbuf, graph_comm)
+end

src/consts/mpich.jl

@@ -224,8 +224,8 @@ const MPI_Win_errhandler_fn = MPI_Win_errhandler_function
 @const_ref MPI_ARGV_NULL      Ptr{Cvoid} C_NULL
 @const_ref MPI_ARGVS_NULL     Ptr{Cvoid} C_NULL
 
-@const_ref MPI_UNWEIGHTED     Ptr{Cint} cglobal((:MPI_UNWEIGHTED, libmpi), Ptr{Cint})
-@const_ref MPI_WEIGHTS_EMPTY  Ptr{Cint} cglobal((:MPI_WEIGHTS_EMPTY, libmpi), Ptr{Cint})
+@const_ref MPI_UNWEIGHTED     Ptr{Cint} unsafe_load(cglobal((:MPI_UNWEIGHTED, libmpi), Ptr{Cint}))
+@const_ref MPI_WEIGHTS_EMPTY  Ptr{Cint} unsafe_load(cglobal((:MPI_WEIGHTS_EMPTY, libmpi), Ptr{Cint}))
 @const_ref MPI_BOTTOM         Ptr{Cvoid} C_NULL
 @const_ref MPI_IN_PLACE       Ptr{Cvoid} -1
 

src/error.jl

@@ -28,7 +28,7 @@ macro mpichk(expr, min_version=nothing)
         fn = expr.args[2].args[1].value
         if isnothing(dlsym(libmpi_handle, fn; throw_error=false))
             return quote
-                throw(FeatureLevelError($fn, $min_version))
+                throw(FeatureLevelError($(QuoteNode(fn)), $min_version))
             end
         end
     end