Add an ffi_call function with a similar signature to pure_callback

[dfm]

Currently, JAX users who want to use XLA custom calls must interact with private APIs (e.g. core.Primitive) and MLIR. This doesn’t provide a great developer experience, and it would be useful to provide some sort of public API. This has been previously discussed in several different contexts (including #12632), and this PR builds on ideas from this previous work to present a simple API that covers some core use cases.

There are more advanced use cases which would require finer-grained customization, and these would continue to rely on the private API. But, there do appear to be examples of use cases that would be satisfied by this simpler interface.

Example

The general idea is to provide a function called (something like) jax.extend.ffi.ffi_call with a signature that is similar to jax.pure_callback, that lowers to a custom call. For example, the existing implementation of lu_pivots_to_permutation on GPU (the only FFI custom call currently in jaxlib), could (to first approximation) be written as:

def ffi_call_lu_pivots_to_permutation(pivots, permutation_size):
  return jex.ffi.ffi_call(
      "cu_lu_pivots_to_permutation",

      # Output types are specified without reference to MLIR
      jax.ShapeDtypeStruct(
          shape=dims[:-1] + (permutation_size,),
          dtype=pivots.dtype,
      ),

      # Input arguments
      pivots,

      # Keyword arguments are passed to the FFI custom call as attributes
      permutation_size=np.int32(permutation_size),  # Note: np not jnp
  )

from jax.lib import xla_client
xla_client.register_custom_call_target(
    "cu_lu_pivots_to_permutation", ..., platform="CUDA", api_version=1)

Platform-dependent behavior should be handled in user code with the help of lax.platform_dependent. (Currently this doesn't work, but @gnecula is looking into it.) Like jax.pure_callback, this could be combined with custom_jvp or custom_vjp to support autodiff. vmap is discussed below.

Batching

This proof-of-concept implementation includes a vectorized parameter which has the same behavior as the vectorized parameter to jax.pure_callback (in fact it uses exactly the same batching rule). The tl;dr is that when vectorized is False, the base custom call is executed in a while loop, but when vectorized is True, the vmapped primitive calls the same custom call with an extra batch dimension on the left. This behavior could potentially work with the FFI interface since the input buffers include dimension metadata, but it’s a restrictive interface. Is there a better approach (don’t say custom_vmap! Or do...)?

Alternatives

If we’re not totally wedded to aligning with the jax.pure_callback interface, it’s possible that a "builder" interface would be more future proof. For example, the syntax for the demo from above would be something like:

do_call = jex.ffi.make_ffi_call("cu_lu_pivots_to_permutation")
do_call(
    jax.ShapeDtypeStruct(
        shape=dims[:-1] + (permutation_size,),
        dtype=pivots.dtype,
    ),
    pivots,
    batch_size=np.int64(batch_size),
    pivot_size=np.int32(pivot_size),
    permutation_size=np.int32(permutation_size),
)

This has the potential benefit that do_call could include metadata like a reference to the underlying core.Primitive so that users could use that for further customization.

[andportnoy]

Love the potential for reducing boilerplate!
How does this work when ffi_call is called multiple times over the life of a process? There seems to be only one ffi_call_p declared globally, does that mean ffi_call will overwrite the abstract eval rule, impl, etc.? Or are you planning to have a global ffi_call_p_registry or similar that adds a new primitive every time ffi_call is called? Or is that not necessary somehow?

[dfm]

Good question! There's actually only ever need for one primitive because all of the relevant parameters are passed at bind time. You'll see that all the rules take the target name as a parameter, for example. This means that the same primitive can be reused for any number of calls to the same (or many different) FFI targets.

This is actually the same pattern which is used for the callback primitives (here's the primitive for pure_callback, for example), and all other higher order primitives (custom_vjp, etc.).

[andportnoy]

I see thanks, so if I'm understanding this correctly, ffi_call_p assumes a new "persona" every time ffi_call is invoked and that works fine.

[dfm]

Exactly! If you inspect the jaxpr with an ffi_call in it, you'll see something like:

ffi_call[
  target_name=my_ffi_call_target_name
  result_avals=(ShapedArray(float32[5,3]), ...)
  vectorized=True
  a_parameter_that_i_set=1.5
]

[superbobry]

OOC why not jax.ShapeDtypeStruct as in pure_callback?

[dfm]

pure_callback is typed to take Any for result_shape_dtypes because it accepts a pytree, but the leaves are implicitly expected to satisfy DuckTypedArray. In both cases, the docstrings say that ShapeDtypeStruct is useful for constructing this argument, but the only requirement in each case is that the leaves have .shape and .dtype properties. With this in mind, this seemed right to use this annotation here, rather than the stricter ShapeDtypeStruct!

[gnecula]

Have you considered always returning multiple results: result_shape_dtypes should always be a sequence, and the restult is always a sequence. The typing declarations would be simpler and the usage also.

[dfm]

Good question! I did consider that, but decided that it was probably worth supporting single outputs since a large fraction of the uses of custom_call that I could find in Google code, and across GitHub only produce a single output. Furthermore, pure_callback also supports both single and multiple outputs (Note: that's not really a fair comparison since it also supports arbitrary pytree output), so I aimed to support as much of that interface as we can for an FFI. So, since the implementation is fairly straightforward, I came down on the side of supporting single or multiple outputs transparently, but I don't feel strongly if there is pressure otherwise!

[gnecula]

SGTM