Selectively enabling/disabling multithreading on CPU

[untom]

Hi!
By default, JAX uses multithreading when running on CPU. I have a usecase where I'd like to disable this for some of my code (I'd like to parallelize things myself in a more flexible way), while still maintaining it overall.

I'm aware that threads #743 and #1539 suggest setting the environmental var XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" to disable multithreading for the whole program. However, this is not flexible enough for me, as I'd only like to selectively disable/enable multithreading. An ideal API for me would be one where I could specify something along the lines of intra_op_parallelism_threads=1 when calling jax.jit.

[zhangqiaorjc]

Discussed with @hawkinsp , this may be possible with some changes to our runtime:

• in pjrt runtime, we can create a fake threadpool with only 1 thread, and pass it to XLA/CPU codegen'ed function whenever intra-op parallelism is disabled
• or teach xla about a thread limit that's separate to the pool limit
• or create a threadpool shim that has that effect

Then expose a locally scoped control via jax.config

[h3jia]

Hi, are there any updates on this? Ideally, it would be useful if we can limit the # of threads to some number (may be larger than 1). On shared cluster nodes usually one single user should not use all the available CPUs.