What's the plan for XLA_FLAGS=--xla_cpu_use_thunk_runtime=false ?

[patrick-kidger]

Similar to #24501.

Over in patrick-kidger/diffrax#549 we're observing a 3x performance slowdown when upgrading the JAX version, and it looks like the new XLA runtime is the culprit.

I'm just curious what the plan is for the new runtime? (In particular just that we won't end up with only the current slower performance going forward, of course!)

[bemporad]

I'm the developer of jax_sysid (https://github.com/bemporad/jax-sysid) and experienced a similar issue with jax>0.4.31. Here below is minimal code to observe substantial differences between 0.4.31 and the current latest version (0.4.38):

import jax.numpy as jnp
import jax
import numpy as np
import time

print(jax.print_environment_info())

# Data generation
seed = 3  # for reproducibility of results
np.random.seed(seed)

n = 3 
N = 3000000
U = jnp.array(np.random.randn(N, n, 1))
x0 = jnp.zeros((n,1))

@jax.jit
def sys(x, u):
    x = .95*x + u
    return x, x

t_scan=time.time()
_, Y = jax.lax.scan(sys, x0, U) 
print("CPU time: ", time.time()-t_scan)

See logs:

jax:    0.4.31
jaxlib: 0.4.31
numpy:  2.2.1
python: 3.12.8 (main, Dec  3 2024, 18:42:41) [Clang 16.0.0 (clang-1600.0.26.4)]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Darwin', node='', release='24.1.0', version='Darwin Kernel Version 24.1.0: Thu Oct 10 21:06:57 PDT 2024; root:xnu-11215.41.3~3/RELEASE_ARM64_T6041', machine='arm64')

None
CPU time:  0.02417731285095215

jax:    0.4.38
jaxlib: 0.4.38
numpy:  2.2.1
python: 3.12.8 (main, Dec  3 2024, 18:42:41) [Clang 16.0.0 (clang-1600.0.26.4)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='', release='24.1.0', version='Darwin Kernel Version 24.1.0: Thu Oct 10 21:06:57 PDT 2024; root:xnu-11215.41.3~3/RELEASE_ARM64_T6041', machine='arm64')

None
CPU time:  0.1252748966217041

Setting XLA_FLAGS=--xla_cpu_use_thunk_runtime=false in 0.4.38 indeed reduces CPU time:

CPU time:  0.031823158264160156

but it is still worse than with jax 0.4.31.

I hope this further helps to bring jax.lax.scan back to its original outstanding performance!