Feature request: Building jaxlib (cpu) with threading disabled?

[josephrocca]

Discussed in #11130

^{Originally posted by josephrocca June 17, 2022}
Hello, I'm currently trying to get JAX (cpu only) working in the browser. This requires compiling with Emscripten, and apparently I need to compile without threading.

IIUC it's possible to restrict JAX to a single thread at run time, but I wasn't able to find any info on compiling without threading.

I'm wondering if there's a flag for this that I can set somewhere? If not, is there any viable (even if very hacky) way to do it?

Thanks!

[zhangqiaorjc]

@josephrocca do you mean compiling without any references to threading? I don't think that's easily doable in JAX.

If you just want fewer threads, there might need to be a few source-level changes possible

1. tfrt::CreateMultiThreadedWorkQueue(1, 1)

https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.cc;l=133-135?q=tfrt_cpu_pjrt_client&ss=tensorflow

2. Set env var NPROC=1

https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/pjrt/utils.cc;l=272?q=DefaultThreadPoolSize&ss=tensorflow

Note that you'd need to re-compile jaxlib from source.
https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source

But even these won't get you to exactly 1 thread.