Use MPIX_Op_create_x for user-defined operators when available

[vchuravy]

Summary

MPICH ≥ 4.3 ships the experimental MPIX_Op_create_x, which passes a void *extra_state context pointer to the user-defined reduction function (the implementation of the MPI Forum proposal mpi-forum/mpi-issues#839):

typedef void (MPIX_User_function_x)(void *a, void *b, MPI_Count len, MPI_Datatype datatype, void *extra_state);
typedef void (MPIX_Destructor_function)(void *extra_state);
int MPIX_Op_create_x(MPIX_User_function_x *user_fn_x, MPIX_Destructor_function *destructor_fn,
                     int commute, void *extra_state, MPI_Op *op);

With it, MPI.Op(f, T) no longer needs a closure @cfunction: the OpWrapper is passed as extra_state and recovered with unsafe_pointer_to_objref inside a plain @cfunction. This removes the LLVM-trampoline dependency and should make user-defined reduction operators work on aarch64 / ppc64le / arm (#404) when MPICH is the backend. Other backends keep using the existing closure path unchanged.

Changes

• MPI.API.MPIX_Op_create_x binding and MPI.API.HAS_MPIX_Op_create_x, determined with dlsym at precompile time (same mechanism @mpichk uses for its version checks).
• Op(f, T) takes the new path when available, otherwise the existing closure-cfunction path with its platform errors.
• OpWrapper becomes mutable struct so its address can be taken; both callback signatures share the reduction kernel _reduce!. The T = Any case now receives the MPI_Datatype by value.
• A no-op destructor is passed: MPICH's generated argument checks reject destructor_fn = NULL with MPI_ERR_ARG, even though MPIR_Op_free_impl tolerates it.
• Tests: custom-op tests in test_reduce.jl are enabled whenever HAS_MPIX_Op_create_x, plus a new test with a state-capturing closure.
• Docs: known-issues section on custom reduction operators.

Testing

Locally on x86_64 Linux, 3 ranks:

• MPICH_jll 5.0.1 (HAS_MPIX_Op_create_x == true): test_reduce.jl, test_allreduce.jl, test_scan.jl, test_exscan.jl, test_onesided.jl pass. Also checked Op(f, Any) with Int32 and Vector{Int} inputs, and freeing the op followed by GC.
• OpenMPI_jll 5.0.10 (HAS_MPIX_Op_create_x == false): test_reduce.jl passes via the unchanged fallback path.

Not tested: an actual aarch64 machine. CI on macOS aarch64 with MPICH_jll would be the real check for this, which is the main reason this is a draft.

Open questions

• Should HAS_MPIX_Op_create_x be surfaced in MPI.versioninfo()?
• Should the fallback path's architecture error be kept as is, or should Op suggest switching to MPICH?
• The Op.fptr field now holds either a closure CFunction or an OpWrapper; a rename would be a (minor) breaking change for anyone poking at it.

🤖 Generated with Claude Code

[vchuravy]

We must GC preserve w globally

[eschnett]

Yes, w must be preserved until we are sure the operator is not used any more. I don't think we can tell when this is the case, since MPI might keep the operator around even after we call MPI_Op_free. I think we can never free the closure.