优化 MKL 在 AMD CPU 上的性能

发表于

问题

实验室有一些 AMD EPYC 7713 的服务器，采购的原因是组里有一些人的程序有非常高的 CPU 负载（我也不知道是什么负载，为什么不能跑在 GPU 上，我也没有精力去逐个帮助解决），框框多的 AMD 处理器非常适合这种需求。

不过 AMD 的处理器虽然香，用在炼丹实验室会有额外的问题：Anaconda 安装的 numpy 和 PyTorch 默认都使用了 MKL 作为 BLAS 的实现，MKL 的 library function 也是大部分高 CPU 负载程序的热点，但 MKL 会判断自己是否在 Intel CPU 上运行，如果不是，则没有优化效果。

由于这是炼丹实验室，大家很少有足够的 HPC 基础去自己编译适合的 numpy 和 PyTorch 版本，也很难脱离 Anaconda，对于 MKL 的依赖因此很难去除。为此需要一个对一般用户无感知的解决方案。

解决方案

通过搜索引擎可以搜索到一个广为流传解决方案：设置环境变量 MKL_DEBUG_CPU_TYPE=5。这是个曾经有效的解决方案，但对于 MKL 2020 及之后的版本不再有效。

最终我在此处找到了更巧妙的解决方案。

MKL 会调用一个 mkl_serv_intel_cpu_true() 函数以检查自己是否运行在 Intel CPU 上，只要提供一个虚假的、始终返回 1 的 mkl_serv_intel_cpu_true()，即可欺骗 MKL 让它认为自己在 Intel CPU 上运行。

为此，可以利用 Linux 的 LD_PRELOAD 机制。LD_PRELOAD 指向的动态链接库有最高的加载优先级，只要编译一个想要的 mkl_serv_intel_cpu_true() 函数为 so 文件，并用 LD_PRELOAD 指向它，即可抢先完成此函数的加载。

笔者也经常有耳闻 LD_PRELOAD 机制被用于库函数劫持攻击，此处算是一种妙用。

具体实施

新建 mkl_trick.c:

1
2
3
int mkl_serv_intel_cpu_true() {
    return 1;
}

使用 gcc -shared -fPIC -o libmkl_trick.so mkl_trick.c 编译，并将生成的 libmkl_trick.so 复制到 /usr/local/lib。

在 Shell 的全局初始化文件中加入：

1
2
3
export MKL_DEBUG_CPU_TYPE=5  # 兼容旧版本 MKL
export MKL_ENABLE_INSTRUCTIONS=AVX2  # 可选，指明 MKL 可以使用 AVX2
export LD_PRELOAD=/usr/local/lib/libmkl_trick.so

实验室的同学有的用 Bash 也有的用 ZSH，所以两者都要修改：

  • Bash: 新建文件 /etc/profile.d/mkl.sh 并添加上述内容
  • ZSH: 添加到 /etc/zsh/zshenv

参考