Exposing the XLA communicator through the FFI interface

[therooler]

I wrote package called JAXMg for doing distributed linear algebra using cuSOLVERMg: https://github.com/flatironinstitute/jaxmg
that works with jax>=0.6.1

This required me to go through the FFI interface and synchronise the different threads/processes on the C++ side in order to set up and orchestrate the solver. My current solution is to use shared memory to enable barriers between the threads/processes.

Ideally, I would access the XLA communicator that JAX uses directly on the C++ side. This would enable me to set up barriers between processes and would work across multiple device hosts. I eventually want to code up "JAXMp", that relies on cuSOLVERMp (more powerful, works for any number of GPUs), which would for sure require access to the XLA communicator.

Is it currently possible to access the XLA communicator somehow on the C++ side with FFI?