FFI API tutorial

[olivier-peltre]

I'm in the process of refactoring a few (recent) kernel bindings we had working, since the XLA custom_call API docs vanished a week ago in favour of the newer ffi_call API.

Thanks for the tutorial and relifted API, which will probably help reduce boilerplate. Let me however a few points of improvement which would really help:

• would it be possible to avoid auto keyword in the tutorial code samples? e.g. here

• would that be within reach to include some kind of polymorphic / templated example? (there was a rather clever macro in the former custom call tutorial to dispatch on buffer dtypes)

Thanks a lot in advance, I think this might help future readers too! 🙏

(cc @dfm, hoping you have some bandwidth for this, apologies if not)

[olivier-peltre]

PS : on a side note,

• how stable do you expect copy/paste of xla/ffi API headers to be in the future? should we make sure to update them regularly to avoid disruption?

[dfm]

Great suggestions all around! I'm not so keen to add dtype dispatching to the tutorial, but I think it would be very sensible to add it to the example code. I'll do that!

I'd be happy to review a PR removing auto from the tutorial (I think the one you linked is probably the only one we can remove!).

how stable do you expect copy/paste of xla/ffi API headers to be in the future? should we make sure to update them regularly to avoid disruption?

The headers don't really change in breaking ways frequently. You should be safe to copy them. We'll announce on the changelog when there is a version bump, and provide details about compatibility.

[dfm]

The rms_norm example code now includes an example of how I do dtype dispatching with the FFI. Here's the relevant PR: #26607

[olivier-peltre]

Hi again @dfm, thanks a lot for the dispatch addition in the tutorial! 🙏

I'm only finding the time to update my JAX bindings now and this is very useful.

However, you might still want to fix the header

This tutorial requires JAX v0.4.31 or newer

With some kind of disclaimer, as with JAX 0.4.34 I still have the following error:

>>> from jax import ffi
ModuleNotFoundError: No module named jax.ffi

Since ideally the kernels should be useable across a range of JAX versions in downstream projects, I'm also wondering how far jax.extend.ffi can be exchanged with jax.ffi as a fallback... So I'll try this for now and see how it goes 🤞.

[olivier-peltre]

Also, it's not very clear to me yet where all the former jax.core.Primitive and mlir.register_lowering boilerplate went... I'm assuming it's useless now, which sounds great! 🙏

I appreciate JAX and perhaps especially this FFI facet are not 100% stable yet, though keeping some form of transition guides on the topic would really help in the future! TBH maintaining cuda bindings for JAX feels a bit scary for now ^^