tweak MPIPtr, modify and enable CUDA tests, add some docs

- A few small tweaks to the MPIPtr implementation (which we can also use for SubArray contiguous checks)
- Modify the tests to support CuArrays where possible
- Enable CuArray tests on the buildbot
- Add some docs

Author: simonbyrne

.gitlab-ci.yml

@@ -20,12 +20,16 @@ include:
     - popd
   script:
     - export JULIA_MPI_PATH="${HOME}/mpi"
+    - export JULIA_PROJECT="test/cudaenv"
     - ${JULIA_MPI_PATH}/bin/ompi_info
     - julia -e 'using InteractiveUtils;
                 versioninfo()'
-    - julia --project -e 'using Pkg;
-                          Pkg.build();
-                          Pkg.test(coverage=true)'
+    - julia --color=yes -e 'using Pkg;
+                            Pkg.develop(PackageSpec(path=pwd()));
+                            Pkg.instantiate();
+                            Pkg.build()'
+    - julia --color=yes test/runtests.jl
+
 .gputest:
   extends: .projecttest
   variables:

Project.toml

@@ -16,9 +16,8 @@ julia = "1"
 
 [extras]
 DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DoubleFloats", "LinearAlgebra", "Pkg", "Test"]
+test = ["DoubleFloats", "Pkg", "Test"]

docs/src/usage.md

@@ -22,6 +22,13 @@ Hello world, I am rank 2 of 3
 Hello world, I am rank 1 of 3
 ```
 
+## CUDA-aware MPI support
+
+If your MPI implementation has been compiled with CUDA support, then `CuArray`s (from the
+[CuArrays.jl](https://github.com/JuliaGPU/CuArrays.jl) package) can be passed directly as
+send and receive buffers for point-to-point and collective operations (they may also work
+with one-sided operations, but these are not often supported).
+
 ## Finalizers
 
 In order to ensure MPI routines are called in the correct order at finalization time,

src/MPI.jl

@@ -22,10 +22,15 @@ function deserialize(x)
     Serialization.deserialize(s)
 end
 
-primitive type SentinelPtr
-    Sys.WORD_SIZE
+primitive type SentinelPtr Sys.WORD_SIZE
 end
 
+primitive type MPIPtr Sys.WORD_SIZE
+end
+Base.cconvert(::Type{MPIPtr}, x::SentinelPtr) = x
+Base.unsafe_convert(::Type{MPIPtr}, x::SentinelPtr) = reinterpret(MPIPtr, x)
+
+
 function _doc_external(fname)
 """
 - `$fname` man page: [OpenMPI](https://www.open-mpi.org/doc/current/man3/$fname.3.php), [MPICH](https://www.mpich.org/static/docs/latest/www3/$fname.html)

src/collective.jl

@@ -35,11 +35,6 @@ function Bcast!(buffer::AbstractArray{T}, root::Integer, comm::Comm) where T
     Bcast!(buffer, length(buffer), root, comm)
 end
 
-function Bcast!(buffer::SubArray{T}, root::Integer, comm::Comm) where T
-    @assert Base.iscontiguous(buffer)
-    Bcast!(buffer, length(buffer), root, comm)
-end
-
 #=
 function Bcast{T}(obj::T, root::Integer, comm::Comm)
     buf = [T]
@@ -112,25 +107,20 @@ To specify the output buffer, see [`Reduce!`](@ref).
 
 To perform the reduction in place, see [`Reduce_in_place!`](@ref).
 """
-function Reduce(sendbuf::MPIBuffertype{T}, count::Integer,
+function Reduce(sendbuf, count::Integer,
                 op, root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
-    recvbuf = Array{T}(undef, isroot ? count : 0)
+    recvbuf = similar(sendbuf, isroot ? count : 0)
     Reduce!(sendbuf, recvbuf, count, op, root, comm)
 end
 
 function Reduce(sendbuf::AbstractArray{T,N}, op,
-    root::Integer, comm::Comm) where {T,N}
+                root::Integer, comm::Comm) where {T,N}
     isroot = Comm_rank(comm) == root
     recvbuf = similar(sendbuf, isroot ? size(sendbuf) : Tuple(zeros(Int, ndims(sendbuf))))
     Reduce!(sendbuf, recvbuf, length(sendbuf), op, root, comm)
 end
 
-function Reduce(sendbuf::SubArray{T}, op, root::Integer, comm::Comm) where T
-    @assert Base.iscontiguous(sendbuf)
-    Reduce(sendbuf, length(sendbuf), op, root, comm)
-end
-
 function Reduce(object::T, op
                 , root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
@@ -238,12 +228,6 @@ output buffer in all processes of the group.
 
 To specify the output buffer or perform the operation in pace, see [`Allreduce!`](@ref).
 """
-function Allreduce(sendbuf::MPIBuffertype{T}, op, comm::Comm) where T
-
-  recvbuf = similar(sendbuf)
-  Allreduce!(sendbuf, recvbuf, length(recvbuf), op, comm)
-end
-
 function Allreduce(sendbuf::AbstractArray{T, N}, op, comm::Comm) where {T, N}
     recvbuf = similar(sendbuf, size(sendbuf))
     Allreduce!(sendbuf, recvbuf, length(sendbuf), op, comm)
@@ -257,14 +241,6 @@ function Allreduce(obj::T, op, comm::Comm) where T
     outref[]
 end
 
-# Deprecation warning for lowercase allreduce that was used until v. 0.7.2
-# Should be removed at some point in the future
-function allreduce(sendbuf::MPIBuffertype{T}, op,
-                   comm::Comm) where T
-    @warn "`allreduce` is deprecated, use `Allreduce` instead."
-    Allreduce(sendbuf, op, comm)
-end
-
 
 """
     Scatter!(sendbuf, recvbuf, count, root, comm)
@@ -456,10 +432,10 @@ Each process sends the first `count` elements of the buffer `sendbuf` to the
 `root` process. The `root` allocates the output buffer and stores elements in
 rank order.
 """
-function Gather(sendbuf::MPIBuffertype{T}, count::Integer,
+function Gather(sendbuf, count::Integer,
                 root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
-    recvbuf = Array{T}(undef, isroot ? Comm_size(comm) * count : 0)
+    recvbuf = similar(sendbuf, isroot ? Comm_size(comm) * count : 0)
     Gather!(sendbuf, recvbuf, count, root, comm)
 end
 
@@ -469,11 +445,6 @@ function Gather(sendbuf::AbstractArray{T}, root::Integer, comm::Comm) where T
     Gather!(sendbuf, recvbuf, length(sendbuf), root, comm)
 end
 
-function Gather(sendbuf::SubArray{T}, root::Integer, comm::Comm) where T
-    @assert Base.iscontiguous(sendbuf)
-    Gather(sendbuf, length(sendbuf), root, comm)
-end
-
 function Gather(object::T, root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
     sendbuf = T[object]
@@ -549,25 +520,17 @@ function Allgather!(buf, count::Integer,
 end
 
 """
-    Allgather(sendbuf, count, comm)
+    Allgather(sendbuf[, count=length(sendbuf)], comm)
 
 Each process sends the first `count` elements of `sendbuf` to the
 other processes, who store the results in rank order allocating
 the output buffer.
 """
-function Allgather(sendbuf::MPIBuffertype{T}, count::Integer,
-                   comm::Comm) where T
-    recvbuf = Array{T}(undef, Comm_size(comm) * count)
+function Allgather(sendbuf, count::Integer, comm::Comm)
+    recvbuf = similar(sendbuf, Comm_size(comm) * count)
     Allgather!(sendbuf, recvbuf, count, comm)
 end
-
-function Allgather(sendbuf::AbstractArray{T}, comm::Comm) where T
-    recvbuf = similar(sendbuf, Comm_size(comm) * length(sendbuf))
-    Allgather!(sendbuf, recvbuf, length(sendbuf), comm)
-end
-
-function Allgather(sendbuf::SubArray{T}, comm::Comm) where T
-    @assert Base.iscontiguous(sendbuf)
+function Allgather(sendbuf::AbstractArray, comm::Comm)
     Allgather(sendbuf, length(sendbuf), comm)
 end
 
@@ -613,7 +576,7 @@ in rank order.
 function Gatherv(sendbuf, counts::Vector{Cint},
                  root::Integer, comm::Comm)
     isroot = Comm_rank(comm) == root
-    recvbuf = Array{T}(undef, isroot ? sum(counts) : 0)
+    recvbuf = similar(sendbuf, isroot ? sum(counts) : 0)
     Gatherv!(sendbuf, recvbuf, counts, root, comm)
 end
 
@@ -797,24 +760,61 @@ function Scan(sendbuf, count::Integer,
     recvbuf
 end
 
-function Scan(object::T, op::Union{Op,MPI_Op}, comm::Comm) where T
+
+function Scan!(sendbuf, recvbuf, count::Integer,
+               op::Union{Op,MPI_Op}, comm::Comm)
+    T = eltype(sendbuf)
+    # int MPI_Scan(const void* sendbuf, void* recvbuf, int count,
+    #              MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
+    @mpichk ccall((:MPI_Scan, libmpi), Cint,
+                  (MPIPtr, MPIPtr, Cint, MPI_Datatype, MPI_Op, MPI_Comm),
+                  sendbuf, recvbuf, count, mpitype(T), op, comm)
+    recvbuf
+end
+function Scan!(sendbuf, recvbuf, count::Integer, opfunc, comm::Comm)
+    Scan!(sendbuf, recvbuf, count, Op(opfunc, eltype(sendbuf)), comm)
+end
+function Scan!(sendbuf::AbstractArray, recvbuf, op, comm::Comm)
+    Scan!(sendbuf, recvbuf, length(sendbuf), op, comm)
+end
+
+function Scan(sendbuf, count::Integer, op, comm::Comm)
+    Scan!(sendbuf, similar(sendbuf, count), count, op, comm)
+end
+
+function Scan(sendbuf::AbstractArray, op, comm::Comm)
+    Scan(sendbuf, length(sendbuf), op, comm)
+end
+function Scan(object::T, op, comm::Comm) where T
     sendbuf = T[object]
     Scan(sendbuf,1,op,comm)
 end
 
-function Exscan(sendbuf, count::Integer,
+function Exscan!(sendbuf, recvbuf, count::Integer,
                 op::Union{Op,MPI_Op}, comm::Comm)
     T = eltype(sendbuf)
-    recvbuf = similar(sendbuf, count)
     # int MPI_Exscan(const void* sendbuf, void* recvbuf, int count,
     #                MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
     @mpichk ccall((:MPI_Exscan, libmpi), Cint,
           (MPIPtr, MPIPtr, Cint, MPI_Datatype, MPI_Op, MPI_Comm),
           sendbuf, recvbuf, count, mpitype(T), op, comm)
     recvbuf
 end
+function Exscan!(sendbuf, recvbuf, count::Integer, opfunc, comm::Comm)
+    Exscan!(sendbuf, recvbuf, count, Op(opfunc, eltype(sendbuf)), comm)
+end
+function Exscan!(sendbuf::AbstractArray, recvbuf, op, comm::Comm)
+    Exscan!(sendbuf, recvbuf, length(sendbuf), op, comm)
+end
+
+function Exscan(sendbuf, count::Integer, op, comm::Comm)
+    Exscan!(sendbuf, similar(sendbuf, count), count, op, comm)
+end
 
-function Exscan(object::T, op::Union{Op,MPI_Op}, comm::Comm) where T
+function Exscan(sendbuf::AbstractArray, op, comm::Comm)
+    Exscan(sendbuf, length(sendbuf), op, comm)
+end
+function Exscan(object::T, op, comm::Comm) where T
     sendbuf = T[object]
     Exscan(sendbuf,1,op,comm)
 end

src/datatypes.jl

@@ -15,24 +15,16 @@ MPIBuffertype{T} = Union{Ptr{T}, Array{T}, SubArray{T}, Ref{T}}
 
 MPIBuffertypeOrConst{T} = Union{MPIBuffertype{T}, SentinelPtr}
 
-if sizeof(Ptr{Cvoid}) == 8
-    primitive type MPIPtr 64 end
-else
-    primitive type MPIPtr 32 end
+Base.cconvert(::Type{MPIPtr}, x::Union{Ptr{T}, Array{T}, Ref{T}}) where T = Base.cconvert(Ptr{T}, x)
+function Base.cconvert(::Type{MPIPtr}, x::SubArray{T}) where T
+    @assert Base.iscontiguous(x)
+    Base.cconvert(Ptr{T}, x)
 end
-
-MPIPtr(x::Cint) where T = reinterpret(MPIPtr, x)
-
-Base.cconvert(::Type{MPIPtr}, x::MPIBuffertype{T}) where T = Base.cconvert(Ptr{T}, x)
-
 function Base.unsafe_convert(::Type{MPIPtr}, x::MPIBuffertype{T}) where T
     ptr = Base.unsafe_convert(Ptr{T}, x)
     reinterpret(MPIPtr, ptr)
 end
 
-Base.cconvert(::Type{MPIPtr}, x::SentinelPtr) = x
-Base.unsafe_convert(::Type{MPIPtr}, x::SentinelPtr) = reinterpret(MPIPtr, x)
-
 fieldoffsets(::Type{T}) where {T} = Int[fieldoffset(T, i) for i in 1:length(fieldnames(T))]
 
 # Define a function mpitype(T) that returns the MPI datatype code for

src/pointtopoint.jl

@@ -126,17 +126,6 @@ function Send(buf::AbstractArray{T}, dest::Integer, tag::Integer, comm::Comm) wh
     Send(buf, length(buf), dest, tag, comm)
 end
 
-"""
-    Send(buf::SubArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
-
-Complete a blocking send of `SubArray` `buf` to MPI rank `dest` of communicator
-`comm` using with the message tag `tag`. Note that the `buf` must be contiguous.
-"""
-function Send(buf::SubArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
-    @assert Base.iscontiguous(buf)
-    Send(buf, length(buf), dest, tag, comm)
-end
-
 """
     Send(obj::T, dest::Integer, tag::Integer, comm::Comm) where T
 
@@ -206,20 +195,6 @@ function Isend(buf::AbstractArray{T}, dest::Integer, tag::Integer, comm::Comm) w
     Isend(buf, length(buf), mpitype(T), dest, tag, comm)
 end
 
-"""
-    Isend(buf::SubArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
-
-Starts a nonblocking send of `SubArray` `buf` to MPI rank `dest` of communicator
-`comm` using with the message tag `tag`. Note that the `buf` must be contiguous.
-
-Returns the commication `Request` for the nonblocking send.
-"""
-function Isend(buf::SubArray{T}, dest::Integer, tag::Integer,
-               comm::Comm) where T
-    @assert Base.iscontiguous(buf)
-    Isend(buf, length(buf), mpitype(T), dest, tag, comm)
-end
-
 """
     Isend(obj::T, dest::Integer, tag::Integer, comm::Comm) where T
 
@@ -293,20 +268,6 @@ function Recv!(buf::AbstractArray{T}, src::Integer, tag::Integer, comm::Comm) wh
     Recv!(buf, length(buf), src, tag, comm)
 end
 
-"""
-    Recv!(buf::SubArray{T}, src::Integer, tag::Integer, comm::Comm) where T
-
-Completes a blocking receive into `SubArray` `buf` from MPI rank `src` of
-communicator `comm` using with the message tag `tag`. Note that `buf` must be
-contiguous.
-
-Returns the `Status` of the receive
-"""
-function Recv!(buf::SubArray{T}, src::Integer, tag::Integer, comm::Comm) where T
-    @assert Base.iscontiguous(buf)
-    Recv!(buf, length(buf), src, tag, comm)
-end
-
 function Recv(::Type{T}, src::Integer, tag::Integer, comm::Comm) where T
     buf = Ref{T}()
     stat = Recv!(buf, 1, src, tag, comm)
@@ -369,21 +330,6 @@ function Irecv!(buf::AbstractArray{T}, src::Integer, tag::Integer,
     Irecv!(buf, length(buf), mpitype(T), src, tag, comm)
 end
 
-"""
-    Irecv!(buf::SubArray{T}, src::Integer, tag::Integer, comm::Comm) where T
-
-Starts a nonblocking receive into `SubArray` `buf` from MPI rank `src` of
-communicator `comm` using with the message tag `tag`. Note that `buf` must be
-contiguous.
-
-Returns the communication `Request` for the nonblocking receive.
-"""
-function Irecv!(buf::SubArray{T}, src::Integer, tag::Integer,
-                             comm::Comm) where T
-    @assert Base.iscontiguous(buf)
-    Irecv!(buf, length(buf), mpitype(T), src, tag, comm)
-end
-
 function irecv(src::Integer, tag::Integer, comm::Comm)
     (flag, stat) = Iprobe(src, tag, comm)
     if !flag

test/cudaenv/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
+DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"