ENH: Delegating from JAX to CuPy under the JIT

[steppi]

lazy_apply works very well for delegating to the NumPy backend on CPU within JAX JIT-ed functions. I've had it on my wish-list for a while to be able to do the same thing with CuPy on GPU. lazy_apply uses pure_callback which inherently brings data to host, so that's not a viable approach for what I'd like to do. I've started tinkering with the JAX FFI which allows calling out to external code within the JAX JIT and have come up with a working prototype which allows for such delegation. It can be found here, https://github.com/steppi/jax_cupy_bridge/tree/main. It offers jax_cupy_bridge.cupy_lazy_apply which works much like lazy_apply with a few limitations:

• func must have an out argument for specifying the output array or arrays.
• func cannot do any dynamic cuda allocations, if it needs workspace buffers, it needs to take them as inputs. (This is actually the reason behind the first limitation).
• It can only have array arguments. I don't think this is an inherent limitation though.

There's still a lot that can be done within these limitations though, for instance, all CuPy ufuncs backed by a kernel from xsf can be supported in JAX under the JIT this way. I've created a discussion on the jax-ml discussion board here, jax-ml/jax#34732, to try to get a sense of whether what I'm doing here is kosher, or if there may be better ways to accomplish this objective.

For now, I'm just making people aware of this. My hope, if what I've done is not a cursed abomination, is that this sort of thing could be integrated into lazy_apply so it can be made to seamlessly handle JAX to CuPy delegation if requested. Or perhaps, it would be better to be more explicit about things. I'm not too settled on any particular API, I just want to be able to delegate from JAX to CuPy under the JIT.

cc @rgommers @mdhaber @ogrisel who I briefly mentioned this idea to in a project meeting earlier this week.

[rgommers]

This sounds reasonable at first sight, and potentially very useful. However, beyond the "will JAX devs say this is okay and future-proof", the elephant in the room is distribution. You have a C++ extension there which depends on XLA headers. That's probably a nonstarter for SciPy itself as well as for array-api-extra. Is there any way this could be done with the JIT compilation support in JAX and CuPy?

[steppi]

This sounds reasonable at first sight, and potentially very useful. However, beyond the "will JAX devs say this is okay and future-proof", the elephant in the room is distribution. You have a C++ extension there which depends on XLA headers. That's probably a nonstarter for SciPy itself as well as for array-api-extra. Is there any way this could be done with the JIT compilation support in JAX and CuPy?

So my thought is that this JAX to CuPy bridge would be a separate package with a JAX and CuPy dependency. It would be an optional dependency of array-api-extra just as JAX and CuPy are, with import guards as appropriate. It only gets used if it (and JAX and CuPy) are installed. We document the installation process and advertise this as an experimental feature. I've found that getting JAX and CuPy installed alongside each other can be a little finicky (on my Linux workstation with CUDA 12 it works if I pip install cupy-cuda12x first and then jax[cuda12] but not the other way around), but if users can get them installed together they can benefit.

[rgommers]

That sounds doable in principle. It'll be hard for users to get this installed reliably, so I wouldn't expect a ton of usage, but it could work.

Next hairy question: is there any kind of ABI stability? Those FFI docs don't mention anything.

[steppi]

That sounds doable in principle. It'll be hard for users to get this installed reliably, so I wouldn't expect a ton of usage, but it could work.

Yeah, this seems like the kind of thing that will only be used by people who have a real need, and thus would be willing to go through the hoops. I think it's something I would have a personal interest in using in the future though, because it makes a lot of statistical stuff involving special functions viable with the JAX JIT on GPU.

Next hairy question: is there any kind of ABI stability? Those FFI docs don't mention anything.

From here, https://openxla.org/xla/custom_call, the answer seems to be not yet.

Caution: The custom-call API/ABI uses PJRT-style versioning (major, minor), however at this point it is still experimental and can be broken at any time. Once API/ABI is finalized we intend to provide stability guarantees similar to PJRT.

My understanding is that the FFI recently got API stability, but for now the Jax CuPy bridge would need to be recompiled every time jaxlib is updated. I'm not sure when they plan to impose ABI stability. I hope to get some insight on that in the JAX discussion.

[rgommers]

Might want to keep this source-only then, at least for now, with jaxlib releases every month, and then there might also be the CUDA 12/13 issue.

[lucascolley]

since there is nothing actionable for this repo yet, going to convert this to a discussion.