Add a public XLA FFI registration provider

[pl-fuchs]

This PR addresses the native FFI registration discussed in #213. I separated the XLA FFI handlers and shared handler table from the nanobind-specific registration, while keeping the existing JAX interface unchanged. I then added a Python-independent provider that registers the complete handler bundles through a stable C entry point. Any XLA-based application can now register the same OEQ-owned handlers using its local XLA_FFI_Api, without depending on nanobind or maintaining a duplicate handler list.

[vbharadwaj-bk]

Interesting set of changes, thanks!

The bulk of my review concentrated on the file now called ffi_handlers.cpp, and those diffs look fine. I'm not so familiar with XLA's FFI beyond the nanobind tutorial, but if it passes the build CI and the tests, I presume it's fine, and the code style also looks unobjectionable.

Comments / Questions

1. Can you point us to some code or documentation on the FFI handler / table registration functions, etc. that you used? If we need to modify / maintain in the future, we will need some references to go off.

2. We now have four files ffi_handler-table, ffi_handlers, ffi_provider, and ffi_registration. Can you provide some clarity about why the code was sectioned in this way?

3. I presume that XLA's preferred build system is Bazel, hence the introduction of Bazel build file. I think this is fine if it is too difficult to make this an additional CMakeLists.txt target, but any comment on that? I suppose it doesn't hurt, and we do want this used in LAMMPS as painlessly as possible.

Other than these questions, seems good.

[pl-fuchs]

1. The main reference I used is XLA's FFI C API:

https://github.com/openxla/xla/blob/main/xla/ffi/api/c_api.h

In particular, the implementation follows the staged handler model defined there (instantiate, prepare, initialize, and execute) and the native registration path through XLA_FFI_Handler_Bundle / XLA_FFI_Handler_Register.

The OEQ-side handler table is intentionally general enough to represent these stages rather than only the execute handler currently exposed through JAX. This change is also needed for the follow-up PR tummfm#1, which attaches compilation state to the XLA executable.

2. The four files separate the handler implementation from the two different registration mechanisms.

ffi_handlers.cpp contains the actual OEQ FFI handler implementations and owns the table of handlers.

ffi_handler_table.h defines a small, OEQ-owned, Python-independent C ABI for exposing that table. It contains the handler name, the function pointers for the different XLA stages, and the handler traits. Both registration paths consume this same table, so there is only one authoritative list of handlers.

jax_ffi_registration.cpp implements the existing JAX-specific registration path. This part depends on nanobind and turns the handlers into the capsules expected by jax.ffi.register_ffi_target.

ffi_provider.cpp implements the native, Python-independent registration path. It consumes the same handler table and registers each handler bundle directly through XLA_FFI_Handler_Register using an XLA_FFI_Api supplied by the caller.

The intention of the separation is therefore to keep nanobind/JAX entirely out of the native XLA/PJRT path. The two registration mechanisms are disjoint, while the handlers themselves are shared.

3. A CMake target for the standalone provider should also be possible. The exported provider interface is a C ABI, so the resulting shared library can be loaded independently of Python/nanobind.

The main reason I added the Bazel target first is my intended integration with chemtrain-deploy and its Bazel build path. It also allows the provider to build directly against the XLA and CUDA repositories selected by the consuming Bazel workspace.

[vbharadwaj-bk]

Sounds reasonable, thanks for the quick responses, and CI passes. Am running the JAX test suite on my end overnight, should get back to you in a day.

[vbharadwaj-bk]

https://github.com/PASSIONLab/OpenEquivariance/tree/jax_ffi_modifications

@pl-fuchs There are some errors for earlier JAX versions that were corrected by changing some of the ffi:Ffi function calls. With the modifications on the branch above, the test suite passes. Could you take a look and modify your branch accordingly? After that, this should be good to go.

[pl-fuchs]

Thanks for catching this issue. I merged the fix and updated the branch.

[vbharadwaj-bk]

✅