How to limit computation to one CPU core?

[louis-shao]

Hi, I use jax as the frontend to dump a xla hlo module, following the example here, then compile and execute it with this openxla library. Is there a way I can limit the computation to one CPU core? I tried setting the number of cpu devices and max number of inflight computations both to 1, but it didn't work. What I observe is, if I add the taskset N prefix to bind the execution to one core, the computation time is about 2x longer than the case without taskset N, which seem indicating that it is not running on one core.

Update:

• I tried HloModuleConfig::set_intra_op_parallelism_threads and setting asynchronous to False when creating the client. Neither works.
• Adding XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=0" according to this makes the execution slower, but still quite faster than adding taskset N.