详解PyTorch编译并调用自定义CUDA算子的三种方式

前言本文为一篇实操教程，作者用最为精简最容易理解的文字描述为大家讲解了用PyTorch编译并调用自定义CUDA算子的三种方式：JIT、Setuptools、CMake。

本文转载自算法码上来

作者 | godweiyang

仅用于学术分享，若侵权请联系删除

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

CV各大方向专栏与各个部署框架最全教程整理

计算机视觉入门1v3辅导班

本篇教程我们主要讲解如何「编译并调用」之前我们写好的CUDA算子，完整的代码还是放在了github仓库，欢迎大家star并fork：

https://github.com/godweiyang/torch-cuda-example

我保证，这是你网上简单「最为精简、最容易看懂」的一套代码了，因为我自己也是刚入门，复杂的我也看得累。

运行环境

NVIDIA Driver: 418.116.00
CUDA: 11.0
Python: 3.7.3
PyTorch: 1.7.0+cu110
CMake: 3.16.3
Ninja: 1.10.0
GCC: 8.3.0

这是我自己的运行环境，显卡是V100，其他环境不保证可以运行，但是大概率没问题，可能要做轻微修改。

代码结构

代码结构还是很清晰的。include文件夹用来放cuda算子的头文件（.h文件），里面是cuda算子的定义。kernel文件夹放cuda算子的具体实现（.cu文件）和cpp torch的接口封装（.cpp文件）。

最后是python端调用，我实现了两个功能。一是比较运行时间，上一篇教程详细讲过了；二是训练一个PyTorch模型，这个下一篇教程再来详细讲述。

编译cpp和cuda文件

JIT

JIT就是just-in-time，也就是即时编译，或者说动态编译，就是说在python代码运行的时候再去编译cpp和cuda文件。

JIT编译的方法上一篇教程已经演示过了，只需要在python端添加load代码即可：

需要注意的就是两个参数，extra_include_paths表示包含的头文件目录，sources表示需要编译的代码，一般就是.cpp和.cu文件。

cpp端用的是pybind11进行封装：

JIT编译看起来非常的简单，运行过程中也基本没有碰到坑，非常顺利。

运行成功的话可以看到Ninja调用了三条命令来编译：

由于输出太长，我省略了多数的参数信息，并精简了指令。可以看出先是调用nvcc编译了.cu，生成了add2_kernel.cuda.o；然后调用c++编译add2.cpp，生成了add2.o；最后调用c++生成动态链接库add2.so。

Setuptools

第二种编译的方式是通过Setuptools，也就是编写setup.py，具体代码如下：

编写方法也非常的常规，调用的是CUDAExtension。需要在include_dirs里加上头文件目录，不然会找不到头文件。

cpp端用的是pybind11进行封装：

接着执行：

这样就能生成动态链接库，同时将add2添加为python的模块了，可以直接import add2来调用。

如果执行正常的话，也是可以看到两条编译命令的：

然后会执行第三条：

最后同样生成了一个动态链接库，不过python端我们不需要加载这个动态链接库，因为setuptools已经帮我们把cuda算子调用的接口注册到python模块里了，直接import即可：

需要注意的是，这里我踩了一个坑，「.cpp和.cu文件名不要相同，也最好不要取容易与python自带库重复的名字」。此外要先import torch，然后再import add2，不然也会报错。

CMake

最后就是cmake编译的方式了，要编写一个CMakeLists.txt文件，代码如下：

这里踩了好几个大坑。首先是找不到nvcc的路径，于是第3行先设置了一下，当然如果你删了也能跑那就更好。然后是找不到python的几个头文件，于是加上了第11行，同样如果你删了也能跑那就更好。最后是一个巨坑，没有链接TORCH_PYTHON_LIBRARY，导致动态链接库生成成功了，但是调用执行一直报错，所以加上了第8行和第17行。

cpp端用的是TORCH_LIBRARY进行封装：