CUDA-aware MPI (#286)

Buffers are passed using an `MPIPtr` type, so CuArrays can be passed using implicit conversion.

Author: kose-y

Project.toml

@@ -7,6 +7,7 @@ version = "0.9.0"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 

src/MPI.jl

@@ -1,6 +1,7 @@
 module MPI
 
 using Libdl, Serialization
+using Requires
 
 macro mpichk(expr)
     @assert expr isa Expr && expr.head == :call && expr.args[1] == :ccall
@@ -24,7 +25,6 @@ end
 primitive type SentinelPtr
     Sys.WORD_SIZE
 end
-Base.cconvert(::Type{Ptr{T}}, sptr::SentinelPtr) where {T} = reinterpret(Ptr{T}, sptr)
 
 function _doc_external(fname)
 """
@@ -65,6 +65,8 @@ function __init__()
     if filesize(dlpath(libmpi)) != libmpi_size
         error("MPI library has changed, re-run Pkg.build(\"MPI\")")
     end
+
+    @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("cuda.jl")
 end
 
 end

src/collective.jl

@@ -0,0 +1,11 @@
+import .CuArrays: CuArray
+import .CuArrays.CUDAdrv: CuPtr, synchronize
+import .CuArrays.CUDAdrv.Mem: DeviceBuffer
+
+function Base.cconvert(::Type{MPIPtr}, buf::CuArray{T}) where T
+    Base.cconvert(CuPtr{T}, buf) # returns DeviceBuffer
+end
+
+function Base.unsafe_convert(::Type{MPIPtr}, buf::DeviceBuffer)
+    reinterpret(MPIPtr, buf.ptr)
+end

src/cuda.jl

@@ -15,6 +15,24 @@ 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
+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/datatypes.jl

@@ -1,3 +1,4 @@
+import Base: eltype
 
 # definition of the `Status` struct.
 # TODO: this bakes in a lot of assumptions about ordering and padding
@@ -94,12 +95,12 @@ end
 Complete a blocking send of `count` elements of type `datatype` from `buf` to MPI
 rank `dest` of communicator `comm` using the message tag `tag`
 """
-function Send(buf::MPIBuffertype{T}, count::Integer, datatype::Union{Datatype, MPI_Datatype},
-              dest::Integer, tag::Integer, comm::Comm) where T
+function Send(buf, count::Integer, datatype::Union{Datatype, MPI_Datatype},
+              dest::Integer, tag::Integer, comm::Comm)
     # int MPI_Send(const void* buf, int count, MPI_Datatype datatype, int dest,
     #              int tag, MPI_Comm comm)
     @mpichk ccall((:MPI_Send, libmpi), Cint,
-          (Ptr{T}, Cint, MPI_Datatype, Cint, Cint, MPI_Comm),
+          (MPIPtr, Cint, MPI_Datatype, Cint, Cint, MPI_Comm),
           buf, count, datatype, dest, tag, comm)
 end
 
@@ -110,18 +111,18 @@ end
 Complete a blocking send of `count` elements of `buf` to MPI rank `dest`
 of communicator `comm` using with the message tag `tag`
 """
-function Send(buf::MPIBuffertype{T}, count::Integer, dest::Integer,
-              tag::Integer, comm::Comm) where T
-    Send(buf, count, mpitype(T), dest, tag, comm)
+function Send(buf, count::Integer, dest::Integer,
+              tag::Integer, comm::Comm)
+    Send(buf, count, mpitype(eltype(buf)), dest, tag, comm)
 end
 
 """
-    Send(buf::Array{T}, dest::Integer, tag::Integer, comm::Comm) where T
+    Send(buf::AbstractArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
 
 Complete a blocking send of `buf` to MPI rank `dest` of communicator `comm`
 using with the message tag `tag`
 """
-function Send(buf::Array{T}, dest::Integer, tag::Integer, comm::Comm) where T
+function Send(buf::AbstractArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
     Send(buf, length(buf), dest, tag, comm)
 end
 
@@ -167,13 +168,13 @@ MPI rank `dest` of communicator `comm` using with the message tag `tag`
 
 Returns the commication `Request` for the nonblocking send.
 """
-function Isend(buf::MPIBuffertype{T}, count::Integer, datatype::Union{Datatype, MPI_Datatype},
-               dest::Integer, tag::Integer, comm::Comm) where T
+function Isend(buf, count::Integer, datatype::Union{Datatype, MPI_Datatype},
+               dest::Integer, tag::Integer, comm::Comm)
     req = Request()
     # int MPI_Isend(const void* buf, int count, MPI_Datatype datatype, int dest,
     #               int tag, MPI_Comm comm, MPI_Request *request)
     @mpichk ccall((:MPI_Isend, libmpi), Cint,
-          (Ptr{T}, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{MPI_Request}),
+          (MPIPtr, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{MPI_Request}),
                   buf, count, datatype, dest, tag, comm, req)
     req.buffer = buf
     return req
@@ -188,9 +189,9 @@ of communicator `comm` using with the message tag `tag`
 
 Returns the commication `Request` for the nonblocking send.
 """
-function Isend(buf::MPIBuffertype{T}, count::Integer,
-               dest::Integer, tag::Integer, comm::Comm) where T
-    Isend(buf, count, mpitype(T), dest, tag, comm)
+function Isend(buf, count::Integer,
+               dest::Integer, tag::Integer, comm::Comm)
+    Isend(buf, count, mpitype(eltype(buf)), dest, tag, comm)
 end
 
 """
@@ -201,8 +202,8 @@ using with the message tag `tag`
 
 Returns the commication `Request` for the nonblocking send.
 """
-function Isend(buf::Array{T}, dest::Integer, tag::Integer, comm::Comm) where T
-    Isend(buf, length(buf), dest, tag, comm)
+function Isend(buf::AbstractArray{T}, dest::Integer, tag::Integer, comm::Comm) where T
+    Isend(buf, length(buf), mpitype(T), dest, tag, comm)
 end
 
 """
@@ -216,7 +217,7 @@ 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), dest, tag, comm)
+    Isend(buf, length(buf), mpitype(T), dest, tag, comm)
 end
 
 """
@@ -254,13 +255,13 @@ from MPI rank `src` of communicator `comm` using with the message tag `tag`
 
 Returns the `Status` of the receive
 """
-function Recv!(buf::MPIBuffertype{T}, count::Integer, datatype::Union{Datatype,MPI_Datatype}, src::Integer,
-               tag::Integer, comm::Comm) where T
+function Recv!(buf, count::Integer, datatype::Union{Datatype,MPI_Datatype}, src::Integer,
+               tag::Integer, comm::Comm)
     stat_ref = Ref{Status}()
     # int MPI_Recv(void* buf, int count, MPI_Datatype datatype, int source,
     #              int tag, MPI_Comm comm, MPI_Status *status)
     @mpichk ccall((:MPI_Recv, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{Status}),
+                  (MPIPtr, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{Status}),
                   buf, count, datatype, src, tag, comm, stat_ref)
     return stat_ref[]
 end
@@ -274,9 +275,9 @@ Completes a blocking receive of up to `count` elements into `buf` from MPI rank
 
 Returns the `Status` of the receive
 """
-function Recv!(buf::MPIBuffertype{T}, count::Integer, src::Integer,
-               tag::Integer, comm::Comm) where T
-    Recv!(buf, count, mpitype(T), src, tag, comm)
+function Recv!(buf, count::Integer, src::Integer,
+               tag::Integer, comm::Comm)
+    Recv!(buf, count, mpitype(eltype(buf)), src, tag, comm)
 end
 
 
@@ -288,7 +289,7 @@ Completes a blocking receive into `buf` from MPI rank `src` of communicator
 
 Returns the `Status` of the receive
 """
-function Recv!(buf::Array{T}, src::Integer, tag::Integer, comm::Comm) where T
+function Recv!(buf::AbstractArray{T}, src::Integer, tag::Integer, comm::Comm) where T
     Recv!(buf, length(buf), src, tag, comm)
 end
 
@@ -329,13 +330,13 @@ from MPI rank `src` of communicator `comm` using with the message tag `tag`
 
 Returns the communication `Request` for the nonblocking receive.
 """
-function Irecv!(buf::MPIBuffertype{T}, count::Integer, datatype::Union{Datatype, MPI_Datatype},
+function Irecv!(buf, count::Integer, datatype::Union{Datatype, MPI_Datatype},
                     src::Integer, tag::Integer, comm::Comm) where T
     req = Request()
     # int MPI_Irecv(void* buf, int count, MPI_Datatype datatype, int source,
     #               int tag, MPI_Comm comm, MPI_Request *request)
     @mpichk ccall((:MPI_Irecv, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{MPI_Request}),
+                  (MPIPtr, Cint, MPI_Datatype, Cint, Cint, MPI_Comm, Ptr{MPI_Request}),
                   buf, count, datatype, src, tag, comm, req)
     req.buffer = buf
     return req
@@ -350,9 +351,9 @@ from MPI rank `src` of communicator `comm` using with the message tag `tag`
 
 Returns the communication `Request` for the nonblocking receive.
 """
-function Irecv!(buf::MPIBuffertype{T}, count::Integer,
-                    src::Integer, tag::Integer, comm::Comm) where T
-    Irecv!(buf, count, mpitype(T), src, tag, comm)
+function Irecv!(buf, count::Integer,
+                    src::Integer, tag::Integer, comm::Comm) 
+    Irecv!(buf, count, mpitype(eltype(buf)), src, tag, comm)
 end
 
 """
@@ -363,9 +364,9 @@ Starts a nonblocking receive into `buf` from MPI rank `src` of communicator
 
 Returns the communication `Request` for the nonblocking receive.
 """
-function Irecv!(buf::Array{T}, src::Integer, tag::Integer,
+function Irecv!(buf::AbstractArray{T}, src::Integer, tag::Integer,
                              comm::Comm) where T
-    Irecv!(buf, length(buf), src, tag, comm)
+    Irecv!(buf, length(buf), mpitype(T), src, tag, comm)
 end
 
 """
@@ -380,7 +381,7 @@ 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), src, tag, comm)
+    Irecv!(buf, length(buf), mpitype(T), src, tag, comm)
 end
 
 function irecv(src::Integer, tag::Integer, comm::Comm)

src/pointtopoint.jl

@@ -21,12 +21,12 @@ This is a collective call over `comm`.
 
 [`MPI.free`](@ref) should be called on the `Win` object once operations have been completed.
 """
-function Win_create(base::Array{T}, comm::Comm; infokws...) where T
+function Win_create(base::AbstractArray{T}, comm::Comm; infokws...) where T
     win = Win()
     # int MPI_Win_create(void *base, MPI_Aint size, int disp_unit, MPI_Info info,
     #                    MPI_Comm comm, MPI_Win *win)
     @mpichk ccall((:MPI_Win_create, libmpi), Cint,
-                  (Ptr{T}, Cptrdiff_t, Cint, MPI_Info, MPI_Comm, Ptr{MPI_Win}),
+                  (MPIPtr, Cptrdiff_t, Cint, MPI_Info, MPI_Comm, Ptr{MPI_Win}),
                   base, Cptrdiff_t(length(base)*sizeof(T)), sizeof(T), Info(infokws...), comm, win)
     refcount_inc()
     finalizer(free, win)
@@ -106,17 +106,17 @@ function Win_shared_query(win::Win, owner_rank::Int)
     out_len[], out_sizeT[], out_baseptr[]
 end
 
-function Win_attach(win::Win, base::Array{T}) where T
+function Win_attach(win::Win, base::AbstractArray{T}) where T
     # int MPI_Win_attach(MPI_Win win, void *base, MPI_Aint size)
     @mpichk ccall((:MPI_Win_attach, libmpi), Cint,
-                  (MPI_Win, Ptr{T}, Cptrdiff_t),
+                  (MPI_Win, MPIPtr, Cptrdiff_t),
                   win, base, Cptrdiff_t(sizeof(base)))
 end
 
-function Win_detach(win::Win, base::Array{T}) where T
+function Win_detach(win::Win, base::AbstractArray{T}) where T
     # int MPI_Win_detach(MPI_Win win, const void *base)
     @mpichk ccall((:MPI_Win_detach, libmpi), Cint,
-                  (MPI_Win, Ptr{T}),
+                  (MPI_Win, MPIPtr),
                   win, base)
 end
 
@@ -147,66 +147,73 @@ function Win_unlock(rank::Integer, win::Win)
     @mpichk ccall((:MPI_Win_unlock, libmpi), Cint, (Cint, MPI_Win), rank, win)
 end
 
-function Get(origin_buffer::MPIBuffertype{T}, count::Integer, target_rank::Integer, target_disp::Integer, win::Win) where T
+function Get(origin_buffer, count::Integer, target_rank::Integer, target_disp::Integer, win::Win)
+    T = eltype(origin_buffer)
     # int MPI_Get(void *origin_addr, int origin_count,
     #             MPI_Datatype origin_datatype, int target_rank,
     #             MPI_Aint target_disp, int target_count,
     #             MPI_Datatype target_datatype, MPI_Win win)
     @mpichk ccall((:MPI_Get, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Win),
+                  (MPIPtr, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Win),
                   origin_buffer, count, mpitype(T), target_rank, Cptrdiff_t(target_disp), count, mpitype(T), win)
 end
-function Get(origin_buffer::Array{T}, target_rank::Integer, win::Win) where T
+function Get(origin_buffer::AbstractArray{T}, target_rank::Integer, win::Win) where T
     count = length(origin_buffer)
     Get(origin_buffer, count, target_rank, 0, win)
 end
 function Get(origin_value::Ref{T}, target_rank::Integer, win::Win) where T
     Get(origin_value, 1, target_rank, 0, win)
 end
 
-function Put(origin_buffer::MPIBuffertype{T}, count::Integer, target_rank::Integer, target_disp::Integer, win::Win) where T
+function Put(origin_buffer, count::Integer, target_rank::Integer, target_disp::Integer, win::Win)
     # int MPI_Put(const void *origin_addr, int origin_count,
     #             MPI_Datatype origin_datatype, int target_rank,
     #             MPI_Aint target_disp, int target_count,
     #             MPI_Datatype target_datatype, MPI_Win win)
+    T = eltype(origin_buffer)
     @mpichk ccall((:MPI_Put, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Win),
+                  (MPIPtr, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Win),
                   origin_buffer, count, mpitype(T), target_rank, Cptrdiff_t(target_disp), count, mpitype(T), win)
 end
-function Put(origin_buffer::Array{T}, target_rank::Integer, win::Win) where T
+function Put(origin_buffer::AbstractArray{T}, target_rank::Integer, win::Win) where T
     count = length(origin_buffer)
     Put(origin_buffer, count, target_rank, 0, win)
 end
 function Put(origin_value::Ref{T}, target_rank::Integer, win::Win) where T
     Put(origin_value, 1, target_rank, 0, win)
 end
 
-function Fetch_and_op(sourceval::MPIBuffertype{T}, returnval::MPIBuffertype{T}, target_rank::Integer, target_disp::Integer, op::Op, win::Win) where T
+function Fetch_and_op(sourceval, returnval, target_rank::Integer, target_disp::Integer, op::Op, win::Win)
     # int MPI_Fetch_and_op(const void *origin_addr, void *result_addr,
     #                      MPI_Datatype datatype, int target_rank, MPI_Aint target_disp,
     #                      MPI_Op op, MPI_Win win)
+    @assert eltype(sourceval) == eltype(returnval)
+    T = eltype(sourceval)
     @mpichk ccall((:MPI_Fetch_and_op, libmpi), Cint,
-                  (Ptr{T}, Ptr{T}, MPI_Datatype, Cint, Cptrdiff_t, MPI_Op, MPI_Win),
+                  (MPIPtr, MPIPtr, MPI_Datatype, Cint, Cptrdiff_t, MPI_Op, MPI_Win),
                   sourceval, returnval, mpitype(T), target_rank, target_disp, op, win)
 end
 
-function Accumulate(origin_buffer::MPIBuffertype{T}, count::Integer, target_rank::Integer, target_disp::Integer, op::Op, win::Win) where T
+function Accumulate(origin_buffer, count::Integer, target_rank::Integer, target_disp::Integer, op::Op, win::Win)
     # int MPI_Accumulate(const void *origin_addr, int origin_count,
     #                    MPI_Datatype origin_datatype, int target_rank,
     #                    MPI_Aint target_disp, int target_count,
     #                    MPI_Datatype target_datatype, MPI_Op op, MPI_Win win)
+    T = eltype(origin_buffer)
     @mpichk ccall((:MPI_Accumulate, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Op, MPI_Win),
+                  (MPIPtr, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Op, MPI_Win),
                   origin_buffer, count, mpitype(T), target_rank, Cptrdiff_t(target_disp), count, mpitype(T), op, win)
 end
 
-function Get_accumulate(origin_buffer::MPIBuffertype{T}, result_buffer::MPIBuffertype{T}, count::Integer, target_rank::Integer, target_disp::Integer, op::Op, win::Win) where T
+function Get_accumulate(origin_buffer, result_buffer, count::Integer, target_rank::Integer, target_disp::Integer, op::Op, win::Win)
     # int MPI_Get_accumulate(const void *origin_addr, int origin_count,
     #                        MPI_Datatype origin_datatype, void *result_addr,
     #                        int result_count, MPI_Datatype result_datatype,
     #                        int target_rank, MPI_Aint target_disp, int target_count,
     #                        MPI_Datatype target_datatype, MPI_Op op, MPI_Win win)
+    @assert eltype(origin_buffer) == eltype(result_buffer)
+    T = eltype(origin_buffer)
     @mpichk ccall((:MPI_Get_accumulate, libmpi), Cint,
-                  (Ptr{T}, Cint, MPI_Datatype, Ptr{T}, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Op, MPI_Win),
+                  (MPIPtr, Cint, MPI_Datatype, MPIPtr, Cint, MPI_Datatype, Cint, Cptrdiff_t, Cint, MPI_Datatype, MPI_Op, MPI_Win),
                   origin_buffer, count, mpitype(T), result_buffer, count, mpitype(T), target_rank, Cptrdiff_t(target_disp), count, mpitype(T), op, win)
 end