use cfunction closures for user-defined operators (#284)

* use cfunction closures for user-defined operators

* add tests for sum of DoubleFloats

* precompile test environment to avoid race conditions

Author: simonbyrne

Project.toml

@@ -14,8 +14,10 @@ Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 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 = ["LinearAlgebra", "Test"]
+test = ["DoubleFloats", "LinearAlgebra", "Pkg", "Test"]

src/MPI.jl

@@ -31,8 +31,8 @@ include("handle.jl")
 include("info.jl")
 include("comm.jl")
 include("environment.jl")
-include("operators.jl")
 include("datatypes.jl")
+include("operators.jl")
 include("pointtopoint.jl")
 include("collective.jl")
 include("topology.jl")

src/collective.jl

@@ -92,12 +92,11 @@ end
 
 # Convert user-provided functions to MPI.Op
 Reduce!(sendbuf::MPIBuffertype{T}, recvbuf::MPIBuffertype{T},
-        count::Integer, opfunc::Function, root::Integer,
-        comm::Comm) where {T} =
-    Reduce!(sendbuf, recvbuf, count, user_op(opfunc), root, comm)
+        count::Integer, opfunc, root::Integer, comm::Comm) where {T} =
+    Reduce!(sendbuf, recvbuf, count, Op(opfunc, T), root, comm)
 
 function Reduce!(sendbuf::MPIBuffertype{T}, recvbuf::MPIBuffertype{T},
-                 op::Union{Op, Function}, root::Integer, comm::Comm) where T
+                 op, root::Integer, comm::Comm) where T
     Reduce!(sendbuf, recvbuf, length(sendbuf), op, root, comm)
 end
 
@@ -113,25 +112,26 @@ 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,
-                op::Union{Op, Function}, root::Integer, comm::Comm) where T
+                op, root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
     recvbuf = Array{T}(undef, isroot ? count : 0)
     Reduce!(sendbuf, recvbuf, count, op, root, comm)
 end
 
-function Reduce(sendbuf::Array{T,N}, op::Union{Op,Function},
+function Reduce(sendbuf::Array{T,N}, op,
     root::Integer, comm::Comm) where {T,N}
     isroot = Comm_rank(comm) == root
     recvbuf = Array{T,N}(undef, isroot ? size(sendbuf) : Tuple(zeros(Int, ndims(sendbuf))))
     Reduce!(sendbuf, recvbuf, length(sendbuf), op, root, comm)
 end
 
-function Reduce(sendbuf::SubArray{T}, op::Union{Op,Function}, root::Integer, comm::Comm) where T
+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::Union{Op,Function}, root::Integer, comm::Comm) where T
+function Reduce(object::T, op
+                , root::Integer, comm::Comm) where T
     isroot = Comm_rank(comm) == root
     sendbuf = T[object]
     recvbuf = Reduce(sendbuf, op, root, comm)
@@ -174,9 +174,9 @@ function Reduce_in_place!(buf::MPIBuffertype{T}, count::Integer,
 end
 
 # Convert to MPI.Op
-Reduce_in_place!(buf::MPIBuffertype{T}, count::Integer, op::Function,
+Reduce_in_place!(buf::MPIBuffertype{T}, count::Integer, op,
                  root::Integer, comm::Comm) where T =
-                    Reduce_in_place!(buf, count, user_op(op), root, comm)
+                    Reduce_in_place!(buf, count, Op(op,T), root, comm)
 
 """
     Allreduce!(sendbuf, recvbuf[, count=length(sendbuf)], op, comm)
@@ -206,11 +206,11 @@ end
 
 # Convert user-provided functions to MPI.Op
 Allreduce!(sendbuf::MPIBuffertypeOrConst{T}, recvbuf::MPIBuffertype{T},
-           count::Integer, opfunc::Function, comm::Comm) where {T} =
-    Allreduce!(sendbuf, recvbuf, count, user_op(opfunc), comm)
+           count::Integer, opfunc, comm::Comm) where {T} =
+    Allreduce!(sendbuf, recvbuf, count, Op(opfunc,T), comm)
 
 function Allreduce!(sendbuf::MPIBuffertypeOrConst{T}, recvbuf::MPIBuffertype{T},
-                   op::Union{Op,Function}, comm::Comm) where T
+                   op, comm::Comm) where T
     Allreduce!(sendbuf, recvbuf, length(recvbuf), op, comm)
 end
 
@@ -222,7 +222,7 @@ the results on all the processes in the group.
 
 Equivalent to calling `Allreduce!(MPI.IN_PLACE, buf, op, comm)`
 """
-function Allreduce!(buf::MPIBuffertype{T}, op::Union{Op, Function}, comm::Comm) where T
+function Allreduce!(buf::MPIBuffertype{T}, op, comm::Comm) where T
     Allreduce!(MPI.IN_PLACE, buf, length(buf), op, comm)
 end
 
@@ -234,18 +234,18 @@ 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::Union{Op,Function}, comm::Comm) where T
+function Allreduce(sendbuf::MPIBuffertype{T}, op, comm::Comm) where T
 
   recvbuf = similar(sendbuf)
   Allreduce!(sendbuf, recvbuf, length(recvbuf), op, comm)
 end
 
-function Allreduce(sendbuf::Array{T, N}, op::Union{Op, Function}, comm::Comm) where {T, N}
+function Allreduce(sendbuf::Array{T, N}, op, comm::Comm) where {T, N}
     recvbuf = Array{T,N}(undef, size(sendbuf))
     Allreduce!(sendbuf, recvbuf, length(sendbuf), op, comm)
 end
 
-function Allreduce(obj::T, op::Union{Op,Function}, comm::Comm) where T
+function Allreduce(obj::T, op, comm::Comm) where T
     objref = Ref(obj)
     outref = Ref{T}()
     Allreduce!(objref, outref, 1, op, comm)
@@ -255,7 +255,7 @@ 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::Union{Op,Function},
+function allreduce(sendbuf::MPIBuffertype{T}, op,
                    comm::Comm) where T
     @warn "`allreduce` is deprecated, use `Allreduce` instead."
     Allreduce(sendbuf, op, comm)

src/datatypes.jl

@@ -3,6 +3,10 @@
 const DATATYPE_NULL = _Datatype(MPI_DATATYPE_NULL)
 Datatype() = Datatype(DATATYPE_NULL.val)
 
+const MPIInteger = Union{Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64}
+const MPIFloatingPoint = Union{Float32, Float64}
+const MPIComplex = Union{ComplexF32, ComplexF64}
+
 const MPIDatatype = Union{Char,
                             Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64,
                             UInt64,

src/handle.jl

@@ -8,6 +8,8 @@ macro mpi_handle(def, extrafields...)
         Base.@__doc__ mutable struct $(esc(def))
             val::$mpiname
             $(extrafields...)
+            
+            $(esc(name))(val::$mpiname, $(extrafields...)) = new(val, $(extrafields...))
         end
 
         # const initializer
@@ -29,6 +31,5 @@ macro mpi_handle(def, extrafields...)
         function Base.unsafe_convert(::Type{Ptr{$mpiname}}, obj::$name)
             convert(Ptr{$mpiname}, pointer_from_objref(obj))
         end
-
     end
 end

src/operators.jl

@@ -1,81 +1,87 @@
-@mpi_handle Op
+"""
+    Op
 
-const OP_NULL = _Op(MPI_OP_NULL)
-const BAND    = _Op(MPI_BAND)
-const BOR     = _Op(MPI_BOR)
-const BXOR    = _Op(MPI_BXOR)
-const LAND    = _Op(MPI_LAND)
-const LOR     = _Op(MPI_LOR)
-const LXOR    = _Op(MPI_LXOR)
-const MAX     = _Op(MPI_MAX)
-const MIN     = _Op(MPI_MIN)
-const NO_OP   = _Op(MPI_NO_OP)
-const PROD    = _Op(MPI_PROD)
-const REPLACE = _Op(MPI_REPLACE)
-const SUM     = _Op(MPI_SUM)
+An MPI reduction operator, for use with [`Reduce!`](@ref)/[`Reduce`](@ref),
+[`Reduce_in_place!`](@ref), [`Allreduce!`](@ref)/[`Allreduce`](@ref), [`Scan`](@ref) or [`Exscan`](@ref).
 
+## Usage
 
+    Op(op, T=Any; iscommutative=false)
 
-import Base.Threads: nthreads, threadid
+Wrap the Julia reduction function `op` for arguments of type `T`. `op` is assumed to be
+associative, and if `iscommutative` is true, assumed to be commutative as well.
+"""
+@mpi_handle Op fptr
 
-# Implement user-defined MPI reduction operations, by passing Julia
-# functions as callbacks to MPI.
+const OP_NULL = _Op(MPI_OP_NULL, nothing)
+const BAND    = _Op(MPI_BAND, nothing)
+const BOR     = _Op(MPI_BOR, nothing)
+const BXOR    = _Op(MPI_BXOR, nothing)
+const LAND    = _Op(MPI_LAND, nothing)
+const LOR     = _Op(MPI_LOR, nothing)
+const LXOR    = _Op(MPI_LXOR, nothing)
+const MAX     = _Op(MPI_MAX, nothing)
+const MIN     = _Op(MPI_MIN, nothing)
+const PROD    = _Op(MPI_PROD, nothing)
+const REPLACE = _Op(MPI_REPLACE, nothing)
+const SUM     = _Op(MPI_SUM, nothing)
 
-# Unfortunately, MPI_Op_create takes a function that does not accept
-# a void* "thunk" parameter, making it impossible to fully simulate
-# a closure.  So, we have to use a global variable instead.   (Since
-# the reduction functions are barriers, being re-entrant is probably
-# not important in practice, fortunately.)   For MPI_THREAD_MULTIPLE
-# using Julia native threading, however, we do make this global thread-local
-const _user_functions = Array{Function}(undef, 1) # resized to nthreads() at runtime
-const _user_op = _Op(MPI_OP_NULL)
+if @isdefined(MPI_NO_OP)
+    const NO_OP   = _Op(MPI_NO_OP, nothing)
+end
+
+Op(::typeof(min), ::Type{T}; iscommutative=true) where {T<:Union{MPIInteger,MPIFloatingPoint}} = MIN
+Op(::typeof(max), ::Type{T}; iscommutative=true) where {T<:Union{MPIInteger,MPIFloatingPoint}} = MAX
+Op(::typeof(+), ::Type{T}; iscommutative=true) where {T<:Union{MPIInteger,MPIFloatingPoint,MPIComplex}} = SUM
+Op(::typeof(*), ::Type{T}; iscommutative=true) where {T<:Union{MPIInteger,MPIFloatingPoint,MPIComplex}} = PROD
+Op(::typeof(&), ::Type{T}; iscommutative=true) where {T<:MPIInteger} = BAND
+Op(::typeof(|), ::Type{T}; iscommutative=true) where {T<:MPIInteger} = BOR
+Op(::typeof(⊻), ::Type{T}; iscommutative=true) where {T<:MPIInteger} = BXOR
+
+
+function free(op::Op)
+    if op.val != OP_NULL.val
+        @mpichk ccall((:MPI_Op_free, libmpi), Cint, (Ptr{MPI_Op},), op)
+        refcount_dec()
+    end
+    return nothing
+end
 
-# C callback function corresponding to MPI_User_function
-function _mpi_user_function(_a::Ptr{Nothing}, _b::Ptr{Nothing}, _len::Ptr{Cint}, t::Ptr{MPI_Datatype})
+struct OpWrapper{F,T}
+    f::F
+end
+
+function (w::OpWrapper{F,T})(_a::Ptr{Cvoid}, _b::Ptr{Cvoid}, _len::Ptr{Cint}, t::Ptr{MPI_Datatype}) where {F,T}
     len = unsafe_load(_len)
-    T = mpitype_dict_inverse[unsafe_load(t)]
-    a = Ptr{T}(_a)
-    b = Ptr{T}(_b)
-    f = _user_functions[threadid()]
+    if isconcretetype(T)
+        S = T
+    else
+        S = mpitype_dict_inverse[unsafe_load(t)]
+    end
+    a = Ptr{S}(_a)
+    b = Ptr{S}(_b)
     for i = 1:len
-        unsafe_store!(b, f(unsafe_load(a,i), unsafe_load(b,i)), i)
+        unsafe_store!(b, w.f(unsafe_load(a,i), unsafe_load(b,i)), i)
     end
     return nothing
 end
 
-function user_op(opfunc::Function)
+
+function Op(f, T=Any; iscommutative=false)
     if Sys.iswindows() && Sys.WORD_SIZE == 32
-        error("Custom reduction operators are not supported on 32-bit Windows.\nSee https://github.com/JuliaParallel/MPI.jl/issues/246 for more details.")
+        error("User-defined reduction operators are not supported on 32-bit Windows.\nSee https://github.com/JuliaParallel/MPI.jl/issues/246 for more details.")
     end
+    w = OpWrapper{typeof(f),T}(f)
+    fptr = @cfunction($w, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}, Ptr{MPI_Datatype}))
 
-    # we must initialize these at runtime, but it can't be done in __init__
-    # since MPI.Init is not called yet.  So we do it lazily here:
-    if _user_op.val == OP_NULL.val
-        # FIXME: to be thread-safe, there should really be a mutex lock
-        # of some sort so that this initialization only occurs once.
-        # To do when native threading in Julia stabilizes (and is documented).
-        resize!(_user_functions, nthreads())
-        user_function = @cfunction(_mpi_user_function, Nothing, (Ptr{Nothing}, Ptr{Nothing}, Ptr{Cint}, Ptr{MPI_Datatype}))
-        @mpichk ccall((:MPI_Op_create, libmpi), Cint,
-             (Ptr{Cvoid}, Ref{Cint}, Ptr{MPI_Op}),
-             user_function, false, _user_op)
-    end
+    op = Op(OP_NULL.val, fptr)
+    # int MPI_Op_create(MPI_User_function* user_fn, int commute, MPI_Op* op)
+    @mpichk ccall((:MPI_Op_create, libmpi), Cint,
+                  (Ptr{Cvoid}, Cint, Ptr{MPI_Op}),
+                  fptr, iscommutative, op)
 
-    _user_functions[threadid()] = opfunc
-    return _user_op
+    refcount_inc()
+    finalizer(free, op)
+    return op
 end
 
-
-
-# Match predefined ops to Julia functions
-for (f,op) in [
-    (+, :SUM),
-    (*, :PROD),
-    (min, :MIN),
-    (max, :MAX),
-    (&, :BAND),
-    (|, :BOR),
-    (⊻, :BXOR),
-]
-    @eval user_op(::$(typeof(f))) = $op.val
-end

test/runtests.jl

@@ -1,3 +1,6 @@
+using Pkg
+pkg"precompile"
+
 using MPI
 using Test
 

test/test_allreduce.jl

@@ -59,5 +59,7 @@ for typ=[Int]
 end
 
 MPI.Barrier( MPI.COMM_WORLD )
+
+GC.gc()
 MPI.Finalize()
 @test MPI.Finalized()

test/test_reduce.jl

@@ -1,7 +1,7 @@
 using Test
 
 using MPI
-using LinearAlgebra
+using LinearAlgebra, DoubleFloats
 
 MPI.Init()
 
@@ -86,7 +86,21 @@ for typ=[Int]
     end
 end
 
-
 MPI.Barrier( MPI.COMM_WORLD )
+
+if !(Sys.iswindows() && Sys.WORD_SIZE == 32)
+
+    send_arr = [Double64(i)/10 for i = 1:10]
+
+    if rank == root
+        @test MPI.Reduce(send_arr, +, root, MPI.COMM_WORLD) ≈ [Double64(sz*i)/10 for i = 1:10] rtol=sz*eps(Double64)
+    else
+        @test MPI.Reduce(send_arr, +, root, MPI.COMM_WORLD) === nothing
+    end
+
+    MPI.Barrier( MPI.COMM_WORLD )
+end
+
+GC.gc()
 MPI.Finalize()
 @test MPI.Finalized()