开启XLA优化的TensorFlow无法充分利用CPU算力问题咨询

• XLA循环融合优化压缩了可并行算子数量
  你的核心计算逻辑是tf.while_loop内部的向量乘法，XLA开启后会默认执行循环融合优化，把循环内部的多个小算子合并为单个融合算子，虽然大幅降低了算子调度开销，但也将原本可以被TensorFlow运行时拆分到多线程执行的多个独立计算单元，打包为少数几个大计算节点，天然降低了可并行度，导致CPU利用率无法打满。
• XLA的CPU调度逻辑不复用TensorFlow原生线程池
  原生TF图模式下的inter_op_parallelism、intra_op_parallelism参数仅控制TF运行时的线程池策略，XLA编译后生成的二进制会自行管理计算调度逻辑，上述TF侧参数对XLA内部的并行策略影响非常有限，所以即使调大参数也无法提升CPU利用率。
• 向量化优化拉高了单核心执行效率
  如果你的向量乘法计算粒度较小，XLA编译时会自动调用AVX/AVX2等SIMD指令集做向量化优化，单个CPU核心的执行效率远高于原生TF模式，少量核心就能覆盖全部计算需求，也会表现为总CPU利用率降低。

• 关闭XLA循环融合优化
  在程序启动前设置环境变量即可生效：

export XLA_FLAGS="--xla_cpu_enable_loop_fusion=false --xla_cpu_enable_fusion=false"

该操作会保留XLA其他优化能力，同时保留更多可并行的独立算子，提升多核心利用率，你需要实测整体吞吐变化，部分小计算量场景可能会因为调度开销上升出现性能下降。

• 拆分tf.while_loop的并行维度
  蒙特卡洛模拟的多路径计算本身是完全独立的，不需要把所有模拟路径都塞到同一个tf.while_loop内，你可以把路径拆分为4~8个互不依赖的tf.while_loop分支，XLA编译后会生成多个可并行执行的计算单元，即可充分利用全部CPU核心。
• 调整XLA侧的并行控制参数
  XLA有独立的CPU并行配置，你可以设置如下环境变量强制开启多线程并行：

export XLA_FLAGS="--xla_cpu_multi_thread_eigen=true --xla_cpu_parallelism_threshold=0 --xla_max_parallelism=8"

调低并行触发阈值后，小计算量的算子也会触发多线程执行，同时限制最大并行度为你的CPU核心数8，避免线程过多带来的调度开销。

• 升级TensorFlow版本
  TensorFlow 2.5属于比较早期的2.x版本，后续版本对XLA CPU的多线程调度做了大量优化，升级到2.10及以上版本后，多数场景下XLA的CPU利用率都有明显提升。

内容的提问来源于stack exchange，提问作者photon1981