【pytorch】关于C++和CUDA扩展

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集成