pytorch的C++扩展编译器
“关于编译器的一个小说明:由于 ABI 版本控制问题,用于构建C++扩展的编译器必须与构建 PyTorch 时使用的编译器 ABI 兼容。在实践中,这意味着您必须在 Linux 上使用 GCC 版本 4.9 及更高版本。对于Ubuntu 16.04和其他更新的Linux发行版,这应该已经是默认编译器了。在 MacOS 上,您必须使用 clang(没有任何 ABI 版本控制问题)。在最坏的情况下,您可以使用编译器从源代码构建PyTorch,然后使用相同的编译器构建扩展。” ——pytorch官网
为什么需要自定义C++和CUDA扩展?
“在某些情况下,有进一步改进性能的空间。最明显的原因是PyTorch不了解您正在实现的算法。它只知道您用于编写算法的单个操作。因此,PyTorch必须一个接一个地单独执行您的操作。由于对操作的实现(或内核)的每个单独调用(可能涉及启动 CUDA 内核)都具有一定的开销,因此在许多函数调用中,此开销可能会变得很大。此外,运行我们代码的Python解释器本身可能会减慢我们的程序速度。因此,加快速度的一个明确方法是重写C++(或CUDA)中的部分,并融合特定的操作组。融合意味着将许多函数的实现组合到一个函数中,这可以从更少的内核启动以及我们可以执行的其他优化中受益,从而提高全局数据流的可见性。” ——pytorch官网
具体流程
c++ built “ahead of time”
- setup.py里调用setuptools
- C++头文件#include <torch/extension.h>
- 使用pybind11将C++函数或类绑定到Python中
- 终端运行python setup.py install
C++ built “just in time”
具体略,提前built比较常见。
CUDA
- 首先编写一个C++文件,定义将从 Python 调用的函数。
- 并使用 pybind11 将这些函数绑定到 Python
- 声明在 CUDA () 文件中定义的函数
- 在 CUDA 文件中,编写实际的 CUDA 内核
- setup.py集成