I'm looking into using jax to write callbacks for a C library, but the call overhead seems to be prohibitively large for smaller inputs.

import jax
import numpy as np
from jax.lib import xla_client

def foo(x):
    return 2 * x**2

foo_jit = jax.jit(foo)

x = jax.device_put(np.zeros(10))
%timeit foo_jit(x).block_until_ready()
%timeit np.arange(10)
# 82.5 µs ± 1.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# 448 ns ± 19.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

So about 140x slower than a small numpy function, which already has a very high overhead compared to a smaller a C function.

I managed to get rid of most of the python overhead by compiling the xla function on my own:

comp = jax.xla_computation(foo)(x)
exe = comp.Compile()
buffer = xla_client.LocalBuffer.from_pyval(x)
%timeit exe.Execute([buffer,]).block_host_until_ready()
# 19.6 µs ± 863 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

So quite a bit faster, but still not exactly fast. Part (or possibly most) of the problem seems to be that xla is running the computation in a different thread. Is it possible to pass it some options that would prevent that? Looking into the output of perf it also seems that it is spending a lot of time allocating things.

In an ideal world, an interface similar to that of numba.cfunc (see here) would be fantastic. Or a way similar to numba to access the different compiled functions. A way to create a LocalBuffer without copying data would also help (if the other problems are solved): eg xla_client.LocalBuffer.from_buffer(y, copy=False).

If there was any way to get our hands on the c function pointer that xla creates somewhere hidden in it's depths, that would help enormously.