How to inference scan function faster in PJRT on CPU?

[puct9]

I am using the PJRT C API to inference a scan function.

def scan_fn(carry, x):
    ...
    return new_carry, y

My recipe for doing so involves:

1. Export the function from Jax

# Write this value to a file
jax.export.export(jax.jit(scan_fn))(dummy_carry, dummy_x).mlir_module_serialized

2. From C++, use the PJRT C API after building the pjrt_c_api_cpu_plugin.so from XLA repo.

void* hande = dlopen("pjrt_c_api_cpu_plugin.so", RTLD_LAZY);
...

3. Infer the function. Inference code borrowed from here: https://github.com/LeelaChessZero/lc0/blob/96c0bebdaffcd65d554d35fd1cd110a25e6561cc/src/neural/backends/xla/pjrt.cc#L150

The benchmark results show that a single invocation of the scan function to take about 15 microseconds. For the function I had (GRU), this is relatively high, as a hand-crafed function using intrinsics in native C++ achieves 3 microseconds, and an ONNX function with the same scan function design achieves 8 microseconds.

I suspect that one of the inefficiencies may be that PJRT_LoadedExecutable_Execute necessarily allocates new buffers for every invocation, along with overhead from the threading model, though I cannot prove this.

I am also happy for the scan function's new carry return variable to mutate the input carry variable in-place. However, I do not know if this is possible. After rewriting scan_fn to take Refs, it seems that jax.export.export does not support this.

• How to inference this scan function faster?
• Is there intuition for how the runtimes may scale as the computation gets more expensive? (these figures are with a GRU with very few parameters). Maybe this method just has a high baseline cost but scales better to larger computations and data?
• Are there alternatives that may be better suited for this task?