Skip to content

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

Open
eschnett wants to merge 1 commit into
masterfrom
claude/user-defined-reduction-arm64v8-a6d454
Open

eschnett wants to merge 1 commit into
masterfrom
claude/user-defined-reduction-arm64v8-a6d454

Conversation

@eschnett

@eschnett eschnett commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

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 Fortran MPI bindings.)

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.)
@vchuravy

vchuravy commented Sep 8, 2026

Copy link
Copy Markdown
Member

I am not a fan of the added complexity and I was hoping that RegisterOp covered most of the use-cases in particular since it allows for RegisterOp(F, Any) and then use the datatype to dispatch.

So the user-defined operator case should only be needed for true closures.

I was wondering if we could smuggle the additional information through a DataType, but I guess it would be a bit awkard.

Smuggle context through the datatype. The callback does receive MPI_Datatype*. Create a dedicated type per op with MPI_Type_dup or MPI_Type_contiguous(1, T), then attach the context with MPI_Type_set_attr using a keyval from MPI_Type_create_keyval, or just key a global dict on the handle. In the callback, MPI_Type_get_attr recovers it. Cost: the op must be paired with its own datatype, so you cannot reuce with a bare MPI_DOUBLE.

@eschnett

eschnett commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

From a user's perspective, having to modify datatypes is even more inconvenient than using a macro to define closures. I'd say: We either keep the current setup (requiring macros until LLVM figures things out), or we use the trampoline pool. Since I'm using a trampoline pool in other MPI-related libraries I didn't think this was too much added complexity.

If you want, we can remove the code that automatically extends the trampoline pool at run time. This would simplify things.

Of course there's also waiting for MPI 5.1 or MPI 6. They know that reduction operators have this problem.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants