Native FFI registration and concurrent compilation of kernels

[pl-fuchs]

I use OpenEquivariance through MACE-JAX for molecular dynamics simulations. For this purpose, I developed [chemtrain-deploy](https://github.com/tummfm/chemtrain) to export JAX models and evaluate the resulting XLA call modules from LAMMPS through a native C++ connector. The simulation runtime does not import OEQ through Python.

For this to work, the connector must register every OEQ typed-FFI handler used by the exported model with the native runtime. OEQ currently exposes those handlers as capsules through its nanobind module. Reproducing the handler names and private wrapper symbols in chemtrain-deploy would make the connector depend on an internal list that can drift whenever OEQ adds a target or changes its FFI stages.

I needed an OEQ-owned native interface for this deployment, so I refactored the registration on the [mace-generated-kernels](https://github.com/tummfm/OpenEquivariance/tree/mace-generated-kernels) branch. Commit [9de9a39](tummfm@9de9a39) makes OEQ own a native handler table that acts as the source of truth for registration. Nanobind builds the existing Python registration dictionary from that table. A native consumer can read the same entries through a versioned C ABI instead of reproducing OEQ’s handler inventory.

The public PJRT handler-registration entry point currently accepts a single typed handler rather than a complete staged handler bundle. A native consumer therefore still needs an appropriate bundle-registration path for staged handlers. The OEQ interface proposed here is intended to provide the handler and state metadata without requiring a second OEQ-specific inventory in the connector.

I encountered a related issue with runtime-generated kernels. A deployed MACE model can contain many independent OEQ kernels. In the preceding implementation, cold cache misses compile while holding a single global cache mutex, so independent compilations are serialized.

Commit [18eb954](tummfm@18eb954) extends the native interface to include the instantiate stage and the XLA FFI state type required by the handlers. By the initialization stage, the typed state already contains the static kernel description. The FFI supplies the device ordinal, which is used to determine the target architecture and schedule compilation. This allows independent kernel compilations to proceed concurrently.

The prototype follows this lifecycle:

1. Instantiate parses the immutable kernel description and associates the FFI call with a shared kernel record.
2. Initialize obtains the device ordinal, determines the target architecture, and creates or reuses an architecture-specific compilation record. If an equivalent compilation is not already in flight or complete, it schedules the work on a host worker pool with a bounded, configurable number of workers.
3. Execute waits for pending compilation when necessary, performs or reuses the one-time launcher loading for that architecture, and launches the kernel on the stream supplied by the XLA FFI runtime.

Equivalent states share in-flight and completed compilation work for the same target architecture. The supplied hash is used only to select an interning bucket. Sharing additionally requires an exact match of the kernel family and complete payload, so hash collisions do not cause different kernels to alias.

The shared architecture-specific record retains the compiled image and loaded launcher for subsequent executions. Each executable state also caches its most recently resolved architecture record. In the current prototype, a process-global cache retains shared kernels and their compiled artifacts without eviction, and the number of compilation workers is configurable. These ownership, cache-lifetime, and scheduling choices are implementation details rather than requirements of the interface. They could be adjusted to match OEQ’s preferred design.

I see the native registration interface and staged compilation as related but separable changes. The first gives native consumers an OEQ-owned description of the handlers and state metadata without duplicating OEQ’s internal inventory. The second uses the staged FFI lifecycle to move compilation out of the execution path where possible and to share compilation work between equivalent kernels.

The current connector prototype is not synchronized with this OEQ branch and is not a ready-to-build paired revision. I am linking the public project only as context for the native consumer and connector design.

Would there be interest in adopting either or both of these directions in OEQ? I am happy to adjust the implementation and split the changes according to the maintainers’ preferred boundaries.

[vbharadwaj-bk]

Hey Paul - first, we're thrilled people are using OpenEquivariance - and second, these changes would be welcome in the main repo if it helps ease your own workflow.

The two independent changes seem like reasonable additions - feel free to open PRs for them separately and we can run the test suite and merge. Our main JAX user is the PFT (phonon-fine-tuning) folks, so we'll probably test that we don't break their flow.

The second change in particular, this seems like a reasonable solution to the problem - I'd just like to benchmark the runtime overhead to make sure that at kernel dispatch (barring any warmup iterations), the calls are as fast as possible.

Something to note: we may switch at some point from the registration of three separate endings (forward, backward, double-backward) to the registration of a single kernel with a variadic signature (see, https://openxla.org/xla/custom_call#variadic_arguments_and_results), for example that will take in a derivative number. How much do you think this would affect your flow? It would help clean up some of the crud in our codebase associated with distinct registrations, but I don't want to mess up your work.

Feel free to stage a couple of PRs and we'll take a look.

[pl-fuchs]

Thank you for the quick feedback, Vivek.

I have split the changes as suggested. PR #214 provides the native FFI registration interface, and the chained PR tummfm#1 adds the concurrent compilation path on top of it.

I reworked both a little compared with the original prototype. The first PR now has an OEQ-owned, ABI-versioned handler table and a Python-independent RegisterFFi entry point which registers complete staged handler bundles. The second PR extends this with the executable-state type and instantiate stage, and uses a bounded worker pool to share and parallelize CUDA kernel compilation. I also included a small benchmark of the compilation speedup (see PR tummfm#1).

Regarding the possible change from separate forward/backward/double-backward targets to a single variadic target: from the native registration side, I think the first PR should be largely insensitive to such a change. The registration procedure is not specific to chemtrain-deploy: the connector does not reproduce or depend on OEQ's target inventory, but simply calls RegisterFFi, while OEQ owns the list of targets and their staged handlers. Changing that target layout should primarily be an internal OEQ change rather than requiring corresponding registration logic in chemtrain-deploy.

There is one compatibility consequence for deployed models, though. The exported XLA/StableHLO modules contain the custom-call target names and their call signatures. If those OEQ custom calls change to a new variadic target/signature and the old registrations are removed, already exported models would have to be re-exported. However, long-term compatibility in chemtrain-deploy remains an open issue.

[asglover]

Hello Paul, It's amazing to hear people are using OEQ! Thank you for developing this feature concept!

I have been interested in creating an alternate API for PyTorch to separate the compilation and registration from the runtime. For the purpose of eliminating the JIT compilation and runtime dispatching. Ultimately to make packaging easier for applications that are not executing in PyTorch. I think your goals are very similar to mine, except for the jax bindings. I want to review your proposals to see if there is some kind of common shared utilities that might be factored out

I'll try to take a look this week as well

• Austin