概述

最近在实现一个通过PyTorch Extension扩展PyTorch算子的C++算法,需要分析代码的运行瓶颈进行针对性优化。Intel VTune就是一个能从汇编级和源码级分析CPU运行瓶颈的工具。由于不明原因我没在服务器上跑通命令行版的VTune,所以把程序搬到Windows下分析了,因此记录一下Windows上使用VTune分析PyTorchExtension调用的Cpp程序的全配置流程。

跑通Python程序

第一步是跑通Python程序,主要是配置工具链实现对C++算子的JIT编译。Python这边的配置就不说了,按错误提示补全没装的库(如Ninja)就行,这里说一下Windows下编译器工具链的安装。由于Windows下PyTorch是用微软的编译器MSVC编译的,所以要编译我们的算子也得用MSVC,首先就得安装(如果之前装过Visual Studio的C++编译套件就可以跳过这一步):

  1. 下载Visual Studio Installer,和系统适应的版本就行。
  2. 在Visual Studio Installer选择要安装的工具,在“单个组件”里找到MSVC生成工具和Windows SDK,选最新版本的即可。

MSVC生成工具包含编译器、链接器等工具,Windows SDK则包含了操作系统相关的库。这也纠正了我之前的观点,我以为C/C++和其他现代语言一样,标准库和操作系统无关,编译工具链装上就自动有标准库了,实际上并不是,对于那些C相关的库,MSVC生成工具里就只有里面啥也没有的cstdiocstdlib这些头文件,只有装了Windows SDK,它们才能引用到真正包含库函数的stdio.hstdlib.h这些头文件,就和Linux里单装gcc不行,还得装上glibc才能编译C/C++程序。

  1. 安装完之后就可以运行了,在此之前如果程序是从Linux里搬过来的,需要改一下编译参数,因为MSVC的编译参数和gcc不太一样,如-std=c++17需要改成/std:c++17-O2改成/O2。以及关于AVX的使用,gcc需要具体指定使用的AVX指令的类别,而MSVC只有AVX、AVX2、AVX512三个级别依次升高的选项。比方说我要用_mm256_fmadd_ps这个函数,如果是MSVC直接用/arch:AVX2就可以了,而gcc则要用-mfma指定具体用的是AVX里的FMA系列指令,AVX的具体指令类别可参考Intel的文档,gcc的相关编译参数可参考gcc的文档
  2. 从开始菜单进入MSVC的命令行环境x64 Native Tools Command Prompt for VS 2022,cd到Python脚本所在位置,然后运行脚本,如果可以顺利运行说明配置完成。

VTune配置

VTune里创建项目然后点Configure Analysis就能配置运行,填上程序(python)、参数(脚本文件名)和工作目录(脚本运行时的当前路径)即可。

不过不能直接运行,因为MSVC的环境变量是通过脚本配置的,前面之所以要进入专门的MSVC命令行环境也是这个原因。所幸VTune提供Wrapper Script功能(在Advace Option里面),可以写一个批处理脚本让它在启动Python之前先运行MSVC的脚本配置环境变量:

call "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
%*

第一行就是运行用来配置环境变量的脚本,脚本的路径视Visual Studio的安装路径而定,第二行表示运行VTune分析的程序。在VTune的Wrapper Script里加载这个脚本,点左下角的运行按钮,就可以顺利运行了。

分析的过程中需要对照汇编代码和C++代码,这就需要给程序添加调试信息,只需要在Python脚本里添加编译参数-Z7添加链接参数/DEBUG即可,即:

extension = load(
        name='extension',
        sources=['extension.cpp'],
        extra_cflags=['/std:c++17', '/O2', '-Z7', '/arch:AVX2'],
        extra_ldflags=['/DEBUG'],
        verbose=True)

这样就能在VTune里浏览C++代码了,再次强调/DEBUG是加在链接参数(extra_ldflags)里面的,放错位置会导致C++代码不显示。

总结

微软的工具链特点就是如果全用它的工具开发(Visual Studio一把梭)就非常方便,傻瓜式操作,但如果想要将它的工具和其他工具协同使用,往往就需要颇费一般功夫,Powershell的二进制管道设计也体现了这一点。不过好在微软的工具链文档也是业界最全的,想搜的东西基本都能搜得到。