Add MPI.Dist_graph_neighbors (#707)

This patch adds `MPI.Dist_graph_neighbors` as a convenience wrapper
around `MPI.Dist_graph_neighbors_count` and `MPI.Dist_graph_neighbors!`
that allocates the required result vectors automatically.

Author: fredrikekre

docs/src/reference/topology.md

@@ -13,4 +13,5 @@ MPI.Dist_graph_create
 MPI.Dist_graph_create_adjacent
 MPI.Dist_graph_neighbors_count
 MPI.Dist_graph_neighbors!
+MPI.Dist_graph_neighbors
 ```

src/topology.jl

@@ -377,3 +377,29 @@ function Dist_graph_neighbors!(graph_comm::Comm, sources::Vector{Cint}, destinat
     destination_weights = Array{Cint}(undef,0)
     Dist_graph_neighbors!(graph_comm, sources::Vector{Cint}, source_weights, destinations::Vector{Cint}, destination_weights)
 end
+
+"""
+    Dist_graph_neighbors(graph_comm::Comm)
+
+Return `(sources, source_weights, destinations, destination_weights)` of the graph
+communicator `graph_comm`. For unweighted graphs `source_weights` and `destination_weights`
+are `nothing`.
+
+This function is a wrapper around [`MPI.Dist_graph_neighbors_count`](@ref) and
+[`MPI.Dist_graph_neighbors!`](@ref) that automatically handles the allocation of the result
+vectors.
+"""
+function Dist_graph_neighbors(graph_comm::Comm)
+    indegree, outdegree, weighted = Dist_graph_neighbors_count(graph_comm)
+    sources = Vector{Cint}(undef, indegree)
+    destinations = Vector{Cint}(undef, outdegree)
+    if weighted
+        source_weights = Vector{Cint}(undef, indegree)
+        destination_weights = Vector{Cint}(undef, outdegree)
+        Dist_graph_neighbors!(graph_comm, sources, source_weights, destinations, destination_weights)
+        return sources, source_weights, destinations, destination_weights
+    else
+        Dist_graph_neighbors!(graph_comm, sources, destinations)
+        return sources, nothing, destinations, nothing
+    end
+end

test/test_neighbor_comm.jl

@@ -0,0 +1,57 @@
+using Test
+using MPI
+
+MPI.Init()
+
+const comm = MPI.COMM_WORLD
+const rank = MPI.Comm_rank(comm)
+const comm_size = MPI.Comm_size(comm)
+
+# Generate a ring graph with optional weights: 0 -> 1 -> 2 -> ... -> comm_size -> 0
+const prev_rank = (rank + comm_size - 1) % comm_size
+const next_rank = (rank + 1) % comm_size
+function ring_graph(; weighted)
+    sources = Cint[rank]
+    degrees = Cint.(length.(sources))
+    destinations = Cint[next_rank]
+    weights = weighted ? Cint[rank + comm_size] : MPI.UNWEIGHTED
+    return MPI.Dist_graph_create(comm, sources, degrees, destinations; weights=weights)
+end
+
+# Unweighted graph
+let
+    graph_comm = ring_graph(; weighted=false)
+    indeg, outdeg, weighted = MPI.Dist_graph_neighbors_count(graph_comm)
+    @test indeg == outdeg == 1
+    @test !weighted
+    src = Vector{Cint}(undef, indeg)
+    dst = Vector{Cint}(undef, outdeg)
+    MPI.Dist_graph_neighbors!(graph_comm, src, dst)
+    src2, srcw2, dst2, dstw2 = MPI.Dist_graph_neighbors(graph_comm)
+    @test src == src2 == Cint[prev_rank]
+    @test dst == dst2 == Cint[next_rank]
+    @test srcw2 === dstw2 === nothing
+end
+
+# Weighted graph
+let
+    graph_comm = ring_graph(; weighted=true)
+    indeg, outdeg, weighted = MPI.Dist_graph_neighbors_count(graph_comm)
+    @test indeg == outdeg == 1
+    @test weighted
+    src = Vector{Cint}(undef, indeg)
+    dst = Vector{Cint}(undef, outdeg)
+    MPI.Dist_graph_neighbors!(graph_comm, src, dst)
+    src2 = Vector{Cint}(undef, indeg)
+    srcw2 = Vector{Cint}(undef, indeg)
+    dst2 = Vector{Cint}(undef, outdeg)
+    dstw2 = Vector{Cint}(undef, outdeg)
+    MPI.Dist_graph_neighbors!(graph_comm, src2, srcw2, dst2, dstw2)
+    src3, srcw3, dst3, dstw3 = MPI.Dist_graph_neighbors(graph_comm)
+    @test src == src2 == src3 == Cint[prev_rank]
+    @test dst == dst2 == dst3 == Cint[next_rank]
+    @test srcw2 == srcw3 == Cint[prev_rank + comm_size]
+end
+
+MPI.Finalize()
+@test MPI.Finalized()