From 99d861a911c759dbf91260f812958ad5b30540d8 Mon Sep 17 00:00:00 2001 From: Erik Schnetter Date: Tue, 8 Sep 2026 16:11:01 -0400 Subject: [PATCH] Allow user-defined operators on all architectures via a trampoline pool Many architectures do not support cfunctions wrapping a closure. This makes it impossible to create user-defined reduction operators at run time. This PR introduces 128 pre-compiled functions as a trampoline pool. This allows using closures on all architectures. (I'm using the same approach in the [mpif](https://github.com/eschnett/mpif) Fortran MPI bindings.) --- docs/examples/03-reduce.jl | 6 +- docs/src/knownissues.md | 16 ++-- src/operators.jl | 153 +++++++++++++++++++++++++++++++++---- test/common.jl | 7 ++ test/test_allreduce.jl | 9 +-- test/test_op_pool.jl | 138 +++++++++++++++++++++++++++++++++ test/test_reduce.jl | 9 --- 7 files changed, 300 insertions(+), 38 deletions(-) create mode 100644 test/test_op_pool.jl diff --git a/docs/examples/03-reduce.jl b/docs/examples/03-reduce.jl index cfd3f9c99..48533433b 100644 --- a/docs/examples/03-reduce.jl +++ b/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]' diff --git a/docs/src/knownissues.md b/docs/src/knownissues.md index 660651dfc..5faf88ab7 100644 --- a/docs/src/knownissues.md +++ b/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. diff --git a/src/operators.jl b/src/operators.jl index 9f3fed798..11740b4c9 100644 --- a/src/operators.jl +++ b/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) @@ -99,22 +107,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) @@ -127,12 +250,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) diff --git a/test/common.jl b/test/common.jl index f9385a1d1..c7ef4abaa 100644 --- a/test/common.jl +++ b/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) diff --git a/test/test_allreduce.jl b/test/test_allreduce.jl index 6f23a2e45..41a20c0aa 100644 --- a/test/test_allreduce.jl +++ b/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] diff --git a/test/test_op_pool.jl b/test/test_op_pool.jl new file mode 100644 index 000000000..92de88a3c --- /dev/null +++ b/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() diff --git a/test/test_reduce.jl b/test/test_reduce.jl index b99a7aa20..5582eea1e 100644 --- a/test/test_reduce.jl +++ b/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