Intel NumPy einsum无法充分利用CPU全部核心问题咨询

环境配置 #

• 代码大量使用np.einsum调用
• 虚拟环境已安装Intel优化版NumPy，安装命令：

  $ pip install mkl
  $ pip install intel-numpy 
  

• np.show_config()输出显示已启用MKL相关配置，且支持AVX2等SIMD扩展
• 硬件：Intel单插槽8核16线程处理器

已尝试的多线程设置方式 #

1. WSL环境中先配置环境变量，再运行脚本：

   export MKL_NUM_THREADS=8
   python script.py
   

2. 在脚本导入NumPy前设置环境变量：

   import os
   os.environ["MKL_NUM_THREADS"] = '8' # 也曾尝试设为16
   

问题现象 #

通过htop监控系统资源时，始终只有1个核心处于100%利用率状态；仅在np.einsum调用的瞬间，偶尔有2-3个其他核心出现约3%的短暂负载，之后迅速回落至0%。

最小可复现代码 #

N_thetas = 20
N_rs = 600
N_phis = 18

Legendre_matr = np.random.rand(N_thetas,  2*N_thetas-1,  N_thetas)
after_first_int = np.random.rand(N_rs - 1, 2*N_thetas - 1, N_thetas)
eigenvectors_memory = np.random.rand(N_rs - 1, N_rs - 1, N_thetas)
normaliz = np.random.rand(N_rs - 1, )
Cilmoft = np.random.rand(2*N_thetas-1, N_rs-1, N_thetas)
Psi_m = np.random.rand( 2*N_thetas - 1,    N_rs - 1,    N_thetas )
arange_for_m_values = np.arange(-N_thetas+1, N_thetas)
phis = np.linspace(0.0, 2*np.pi, N_phis)    

after_second_int = np.einsum('ijk,rjk->rij', Legendre_matr, after_first_int, optimize=True)

Psi_lm = np.einsum('jkl,j,mkl->jml', eigenvectors_memory, normaliz , Cilmoft, optimize=True) 

Psi_on_grid_from_Cilmoft = np.einsum('xjk,xn->jkn',   Psi_m,  np.exp(1j * np.outer(arange_for_m_values, phis)), optimize=True )

疑问 #

这种核心利用率极低的情况是否正常？还是我的操作存在疏漏？

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