Allow user-defined operators on all architectures via a trampoline pool

docs/examples/03-reduce.jl

@@ -31,9 +31,9 @@ function pool(S1::SummaryStat, S2::SummaryStat)
     SummaryStat(m,v,n)
 end
 
-# Register the custom reduction operator.  This is necessary only on platforms
-# where Julia doesn't support closures as cfunctions (e.g. ARM), but can be used
-# on all platforms for consistency.
+# Register the custom reduction operator.  This is optional -- any function can be
+# passed directly to the reduction -- but it builds the C callback ahead of time
+# instead of at run time, which is worth doing for an operator used repeatedly.
 MPI.@RegisterOp(pool, SummaryStat)
 
 X = randn(10,3) .* [1,3,7]'

docs/src/knownissues.md

@@ -204,11 +204,15 @@ After that, [this script](https://gist.github.com/luraess/c228ec08629737888a18c6
 
 ## Custom reduction operators
 
-It is not possible to use custom reduction operators [with 32-bit Microsoft MPI](https://github.com/JuliaParallel/MPI.jl/issues/246) on Windows and on [ARM CPUs](https://github.com/JuliaParallel/MPI.jl/issues/404) with any operating system.
-These issues are due to due how custom operators are currently implemented in MPI.jl, that is by using [closure cfunctions](https://docs.julialang.org/en/v1/manual/calling-c-and-fortran-code/index.html#Closure-cfunctions).
-However they have two limitations:
+It is not possible to use custom reduction operators [with 32-bit Microsoft MPI](https://github.com/JuliaParallel/MPI.jl/issues/246) on Windows.
+Custom operators are passed to MPI as [C-compatible function pointers](https://docs.julialang.org/en/v1/manual/calling-c-and-fortran-code/index.html#Creating-C-Compatible-Julia-Function-Pointers-1), which Julia cannot produce for the `stdcall` calling convention that 32-bit Microsoft MPI expects.
+[`MPI.@RegisterOp`](@ref) does not help here either.
 
-* [Julia's C-compatible function pointers](https://docs.julialang.org/en/v1/manual/calling-c-and-fortran-code/index.html#Creating-C-Compatible-Julia-Function-Pointers-1) cannot be used where the `stdcall` calling convention is expected, which is the case for 32-bit Microsoft MPI,
-* closure cfunctions in Julia are based on LLVM trampolines, which are not supported on ARM architecture.
+Custom operators do work on ARM and other non-x86 CPUs, where they used to be unavailable ([#404](https://github.com/JuliaParallel/MPI.jl/issues/404)).
+There, Julia cannot build [closure cfunctions](https://docs.julialang.org/en/v1/manual/calling-c-and-fortran-code/index.html#Closure-cfunctions) — they are based on LLVM trampolines, which only exist on x86 and x86-64 — so MPI.jl instead draws from a pool of statically compiled callbacks.
+This has two consequences on those platforms:
 
-As an alternative [`MPI.@RegisterOp`](@ref) may be used to statically register reduction operations.
+* each invocation of the callback costs one additional dynamic dispatch. MPI passes many elements per invocation, so this is amortized and usually not measurable;
+* each *distinct* operator permanently occupies a pool slot. `MPI_Op_free` only marks an operation for deallocation, and MPI may keep calling the user function until every operation referencing it has completed, so the function cannot be released. Identical operators share a slot, so reducing repeatedly with the same function is fine.
+
+[`MPI.@RegisterOp`](@ref) may be used to register a reduction operation statically. This avoids both the dispatch and the pool slot, and is worth doing for operators used in hot loops.

src/operators.jl

@@ -10,6 +10,14 @@ An MPI reduction operator, for use with [Reduce/Scan collective operations](@ref
 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.
 
+!!! note
+    On architectures where Julia cannot build closure cfunctions (anything other than x86
+    and x86-64), each *distinct* operator permanently occupies a slot in an internal pool
+    of C callbacks: `MPI_Op_free` only marks an operation for deallocation, so the wrapped
+    function has to stay callable for the rest of the process. Identical operators share a
+    slot, so reducing repeatedly with the same function is fine; [`@RegisterOp`](@ref)
+    avoids the pool entirely.
+
 ## See also
 
 - [`Reduce!`](@ref)/[`Reduce`](@ref)
@@ -96,22 +104,137 @@ function (w::OpWrapper{F,T})(_a::Ptr{Cvoid}, _b::Ptr{Cvoid}, _len::Ptr{Cint}, t:
     return nothing
 end
 
+# Closure cfunctions (`@cfunction($f, ...)`) are implemented with LLVM trampolines, which
+# only exist on x86 and x86-64. Everywhere else creating one does not throw, it aborts the
+# process, so this has to be an allowlist: an unknown architecture must take the
+# trampoline pool path below.
+const HAVE_CLOSURE_CFUNCTION = Sys.ARCH ∈ (:x86_64, :i686)
+
+@static if !HAVE_CLOSURE_CFUNCTION
+
+# Where closure cfunctions are unavailable, `Op` draws from a pool of statically defined
+# trampolines instead. A trampoline is an ordinary top-level function -- which `@cfunction`
+# accepts on every architecture -- forwarding to the `OpWrapper` held in its own `OpSlot`.
+# That forwarding call is a dynamic dispatch, but it happens once per invocation of the
+# callback rather than once per element, so it is amortized over the `len` elements MPI
+# passes each time.
+# See https://github.com/JuliaParallel/MPI.jl/issues/404
+mutable struct OpSlot
+    wrapper::Any
+    fptr::Ptr{Cvoid}
+end
+OpSlot() = OpSlot(nothing, C_NULL)
+
+const OP_POOL_SIZE = 128
+const OP_POOL = OpSlot[OpSlot() for _ in 1:OP_POOL_SIZE]
+# Number of slots handed out so far. Slots are never returned to the pool: `MPI_Op_free`
+# only marks an operation for deallocation, and MPI may go on calling the user function
+# until every operation referencing it has completed, so a wrapper handed to MPI has to
+# stay alive and callable for the rest of the process.
+const OP_POOL_USED = Ref(0)
+# Maps an `OpWrapper` to the trampoline already installed for it. `Reduce!` and friends
+# construct an `Op` on every call, so without this a loop reducing with the same operator
+# would consume a slot per iteration. Keyed by object identity: `OpWrapper` and Julia
+# closures are immutable, so `===` compares the captured values with `===`, which is
+# exactly the condition under which two wrappers are interchangeable.
+const OP_SLOT_CACHE = IdDict{Any,Ptr{Cvoid}}()
+const OP_POOL_LOCK = ReentrantLock()
+const OP_POOL_WARN_AT = 4 * OP_POOL_SIZE
+const OP_POOL_WARNED = Ref(false)
+
+# NOTE: each trampoline reaches its slot through the object interpolated into its body,
+# not through a global binding. `grow_op_pool!` builds trampolines the same way in an
+# already-running session, where defining a new global would be read in a world older than
+# the one that defines it.
+for i in 1:OP_POOL_SIZE
+    @eval function $(Symbol(:_op_trampoline_, i))(a::Ptr{Cvoid}, b::Ptr{Cvoid},
+                                                  len::Ptr{Cint}, t::Ptr{MPI_Datatype})
+        $(OP_POOL[i]).wrapper(a, b, len, t)
+        return nothing
+    end
+end
+
+# `@cfunction` pointers must not be taken from a precompiled image, so they are refreshed
+# at load time. This is idempotent, as it must be: load time hooks also run inside the
+# precompile workload.
+@eval function init_op_pool()
+    $([:($(OP_POOL[i]).fptr =
+             @cfunction($(Symbol(:_op_trampoline_, i)), Cvoid,
+                        (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}, Ptr{MPI_Datatype})))
+       for i in 1:OP_POOL_SIZE]...)
+    return nothing
+end
+add_load_time_hook!(init_op_pool)
+
+"""
+    grow_op_pool!()
+
+Add one trampoline to `OP_POOL` and return its slot. Compiling it costs on the
+order of 10ms, so this is a slow path, taken only once the static pool is exhausted.
+The caller must hold `OP_POOL_LOCK`.
+"""
+function grow_op_pool!()
+    slot = OpSlot()
+    tramp = Symbol(:_op_trampoline_, length(OP_POOL) + 1)
+    @eval function $tramp(a::Ptr{Cvoid}, b::Ptr{Cvoid}, len::Ptr{Cint}, t::Ptr{MPI_Datatype})
+        $(slot).wrapper(a, b, len, t)
+        return nothing
+    end
+    # A separate `eval`: a single `@eval begin ... end` is compiled as one thunk, so
+    # `$tramp` would not yet be defined when the `@cfunction` in it is resolved.
+    slot.fptr = @eval @cfunction($tramp, Cvoid,
+                                 (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}, Ptr{MPI_Datatype}))
+    push!(OP_POOL, slot)
+    return slot
+end
+
+"""
+    op_fptr(w::OpWrapper)
+
+Return a C function pointer to a trampoline that forwards to `w`, installing `w` in a
+fresh pool slot unless an identical wrapper already occupies one.
+"""
+function op_fptr(w)
+    @lock OP_POOL_LOCK begin
+        get!(OP_SLOT_CACHE, w) do
+            n = OP_POOL_USED[] + 1
+            slot = n <= length(OP_POOL) ? OP_POOL[n] : grow_op_pool!()
+            slot.wrapper = w
+            OP_POOL_USED[] = n
+            if n >= OP_POOL_WARN_AT && !OP_POOL_WARNED[]
+                OP_POOL_WARNED[] = true
+                @warn """
+                    $n distinct user-defined reduction operators have been created. On $(Sys.ARCH)
+                    each one permanently occupies a slot in MPI.jl's callback pool, and every slot
+                    past the first $OP_POOL_SIZE must be compiled at run time.
+
+                    Identical operators share a slot, so this usually means a loop is building a
+                    new operator each iteration. Hoist `op = MPI.Op(f, T)` out of the loop, or
+                    register the function once with `MPI.@RegisterOp(f, T)`.
+                    """
+            end
+            slot.fptr
+        end
+    end
+end
+
+end # @static if !HAVE_CLOSURE_CFUNCTION
+
 function Op(f, T=Any; iscommutative=false)
     @static if MPI_LIBRARY == "MicrosoftMPI" && Sys.WORD_SIZE == 32
+        # Julia's C-compatible function pointers cannot use the `stdcall` calling
+        # convention that 32-bit Microsoft MPI expects.
         error("""
             User-defined reduction operators are not supported on 32-bit Windows.
             See https://github.com/JuliaParallel/MPI.jl/issues/246 for more details.
         """)
-    elseif Sys.ARCH ∈ (:aarch64, :ppc64le, :powerpc64le) || startswith(lowercase(String(Sys.ARCH)), "arm")
-        error("""
-            User-defined reduction operators are currently not supported on non-Intel architectures.
-            See https://github.com/JuliaParallel/MPI.jl/issues/404 for more details.
-
-            You may want to use `@RegisterOp` to statically register `f`.
-            """)
     end
     w = OpWrapper{typeof(f),T}(f)
-    fptr = @cfunction($w, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}, Ptr{MPI_Datatype}))
+    fptr = @static if HAVE_CLOSURE_CFUNCTION
+        @cfunction($w, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint}, Ptr{MPI_Datatype}))
+    else
+        op_fptr(w)
+    end
 
     op = Op(OP_NULL.val, fptr)
     # int MPI_Op_create(MPI_User_function* user_fn, int commute, MPI_Op* op)
@@ -124,12 +247,14 @@ end
 """
     @RegisterOp(f, T)
 
-Register a custom operator [`Op`](@ref) using the function `f` statically.
-On platfroms like AArch64, Julia does not support runtime closures,
-being passed to C. The generic version of [`Op`](@ref) uses runtime closures
-to support arbitrary functions being passed as MPI reduction operators.
-`@RegisterOp` statically adds a function to the set of functions allowed as
-as an MPI operator.
+Statically register the function `f` as a reduction operator [`Op`](@ref) for arguments of
+type `T`.
+
+This is an optimization, not a requirement: [`Op`](@ref) accepts any function on any
+architecture. `@RegisterOp` builds the C callback for `f` when the enclosing module is
+compiled rather than at run time, which avoids a dynamic dispatch on each invocation of
+the callback and, on architectures without closure cfunctions, avoids permanently
+occupying a slot in MPI.jl's internal callback pool.
 
 ```julia
 function my_reduce(x, y)

test/common.jl

@@ -34,3 +34,10 @@ const MPITestTypes = setdiff([
     UInt8, UInt16, UInt32, UInt64,
     Float32, Float64, ComplexF32, ComplexF64
 ], skip_eltypes)
+
+# Custom reduction operators are wrapped in a C callback. That is not possible for 32-bit
+# Microsoft MPI, which expects the `stdcall` calling convention, and the callback runs on
+# the host, so it is only meaningful for host arrays.
+const can_do_closures =
+    ArrayType === Array &&
+    !(MPI.MPI_LIBRARY == "MicrosoftMPI" && Sys.WORD_SIZE == 32)

test/test_allreduce.jl

@@ -4,13 +4,10 @@ MPI.Init()
 
 comm_size = MPI.Comm_size(MPI.COMM_WORLD)
 
-if ArrayType != Array ||
-   MPI.MPI_LIBRARY == "MicrosoftMPI" && Sys.WORD_SIZE == 32 ||
-   Sys.ARCH === :powerpc64le || Sys.ARCH === :ppc64le ||
-   Sys.ARCH === :aarch64 || startswith(string(Sys.ARCH), "arm")
-    operators = [MPI.SUM, +]
-else
+if can_do_closures
     operators = [MPI.SUM, +, (x,y) -> 2x+y-x]
+else
+    operators = [MPI.SUM, +]
 end
 
 for T = [Int]

test/test_op_pool.jl

@@ -0,0 +1,138 @@
+include("common.jl")
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+sz = MPI.Comm_size(comm)
+rank = MPI.Comm_rank(comm)
+
+# Build a closure capturing `arr`. Written as a function so that the closures it returns
+# all share one type and differ only in the array they capture.
+mkop(arr) = (x, y) -> x + y + arr[1]
+
+# Reduction operators that are genuinely associative and commutative, so that the result
+# does not depend on the order in which MPI applies them, but that are still distinct
+# Julia objects.
+mkadd(k) = (x, y) -> x + y + k - k
+
+named_add(x, y) = x + y
+
+struct PoolStat
+    n::Float64
+    s::Float64
+end
+
+if can_do_closures
+
+    # A runtime closure capturing a local. Before MPI.jl grew the trampoline pool this
+    # threw on aarch64/ppc64le, and the underlying `@cfunction($w, ...)` aborted the
+    # process.
+    let zero = [0]
+        add = mkop(zero)
+
+        send = fill(rank + 1, 3)
+        recv = similar(send)
+        MPI.Allreduce!(send, recv, add, comm)
+        @test recv == fill(sum(1:sz), 3)
+
+        @test MPI.Allreduce(rank + 1, add, comm) == sum(1:sz)
+
+        root = sz - 1
+        red = MPI.Reduce(rank + 1, add, comm; root=root)
+        if rank == root
+            @test red == sum(1:sz)
+        else
+            @test red === nothing
+        end
+    end
+
+    # A closure over a user-defined isbits struct
+    let pool = (a::PoolStat, b::PoolStat) -> PoolStat(a.n + b.n, a.s + b.s)
+        tot = MPI.Allreduce(PoolStat(1.0, rank + 1.0), pool, comm)
+        @test tot.n == sz
+        @test tot.s == sum(1:sz)
+
+        # `T = Any` takes the runtime-dispatch branch of `OpWrapper`, which recovers the
+        # element type from the datatype MPI passes to the callback.
+        op = MPI.Op(pool, Any)
+        tot = MPI.Allreduce(PoolStat(1.0, rank + 1.0), op, comm)
+        @test tot.n == sz
+        @test tot.s == sum(1:sz)
+        MPI.free(op)
+    end
+
+end
+
+# The remaining tests are about the trampoline pool itself, which only exists on
+# architectures without closure cfunctions.
+if can_do_closures && !MPI.HAVE_CLOSURE_CFUNCTION
+
+    @test length(MPI.OP_POOL) >= MPI.OP_POOL_SIZE
+    @test length(unique(slot.fptr for slot in MPI.OP_POOL)) == length(MPI.OP_POOL)
+    @test all(slot -> slot.fptr != C_NULL, MPI.OP_POOL)
+
+    # An `Op` built by the pool carries the trampoline's pointer
+    let op = MPI.Op(mkadd(1), Int)
+        @test op.fptr isa Ptr{Cvoid}
+        @test op.fptr != C_NULL
+        MPI.free(op)
+    end
+
+    # Identical wrappers share a slot: `Reduce!` and friends build an `Op` on every call,
+    # so a loop reducing with the same operator must not consume a slot per iteration.
+    let arr = [0], send = fill(rank + 1, 3), recv = fill(0, 3)
+        used = MPI.OP_POOL_USED[]
+        for _ in 1:100
+            MPI.Allreduce!(send, recv, mkop(arr), comm)   # closure rebuilt each iteration
+        end
+        @test MPI.OP_POOL_USED[] == used + 1
+        @test recv == fill(sum(1:sz), 3)
+
+        # ... and likewise for a named function
+        used = MPI.OP_POOL_USED[]
+        for _ in 1:10
+            MPI.Allreduce!(send, recv, named_add, comm)
+        end
+        @test MPI.OP_POOL_USED[] == used + 1
+        @test recv == fill(sum(1:sz), 3)
+    end
+
+    # Wrappers over different arrays must NOT share a slot, even when the arrays are
+    # `==`-equal: the pool is keyed by identity precisely so the wrong array can never be
+    # installed in a trampoline.
+    let a1 = [0], a2 = [0]
+        @test a1 == a2 && a1 !== a2
+        w1 = MPI.OpWrapper{typeof(mkop(a1)),Int}(mkop(a1))
+        w2 = MPI.OpWrapper{typeof(mkop(a2)),Int}(mkop(a2))
+        p1 = MPI.op_fptr(w1)
+        p2 = MPI.op_fptr(w2)
+        @test p1 != p2
+        # rebuilding an identical wrapper reuses the slot
+        @test MPI.op_fptr(MPI.OpWrapper{typeof(mkop(a1)),Int}(mkop(a1))) == p1
+        # each slot holds its own wrapper
+        slot1 = MPI.OP_POOL[findfirst(s -> s.fptr == p1, MPI.OP_POOL)]
+        slot2 = MPI.OP_POOL[findfirst(s -> s.fptr == p2, MPI.OP_POOL)]
+        @test slot1.wrapper === w1
+        @test slot2.wrapper === w2
+        @test slot1.wrapper.f.arr === a1
+        @test slot2.wrapper.f.arr === a2
+    end
+
+    # Exhausting the static pool grows it at run time
+    let n = MPI.OP_POOL_SIZE + 2
+        ops = [MPI.Op(mkadd(k), Int) for k in 1:n]
+        @test length(unique(op.fptr for op in ops)) == n
+        @test MPI.OP_POOL_USED[] > MPI.OP_POOL_SIZE
+        @test length(MPI.OP_POOL) > MPI.OP_POOL_SIZE
+        # a grown trampoline works like any other
+        @test MPI.Allreduce(rank + 1, ops[end], comm) == sum(1:sz)
+        foreach(MPI.free, ops)
+    end
+
+end
+
+MPI.Barrier(MPI.COMM_WORLD)
+
+GC.gc()
+MPI.Finalize()
+@test MPI.Finalized()

test/test_reduce.jl

@@ -1,14 +1,5 @@
 include("common.jl")
 
-# Closures might not be supported by cfunction
-const can_do_closures =
-    ArrayType === Array &&
-    !(MPI.MPI_LIBRARY == "MicrosoftMPI" && Sys.WORD_SIZE == 32) &&
-    Sys.ARCH !== :powerpc64le &&
-    Sys.ARCH !== :ppc64le &&
-    Sys.ARCH !== :aarch64 &&
-    !startswith(string(Sys.ARCH), "arm")
-
 # a non-builtin isbits type to test generic MPI.reduce
 struct TestSum
     hi::Float64