自定义CUDA实现PyTorch算子的四种简单方法
背景
在探索新的深度学习算法的时候，我们可能会遇到PyTorch提供的算子不能满足需求的情况，这时候就需要自定义PyTorch算子，将我们的算法集成到PyTorch的工作流中。同时，为了提高运算效率，算子往往都需要使用CUDA实现。所幸，PyTorch及很多其他Python库都提供了简化这一过程的方法，完全不需要PyTorch库源文件等其他代码（Tensorflow：🤧）。接下来本文会介绍自定义CUDA实现PyTorch算子的四种简单方法，包括使用Taichi库、Numba库、Cupy库和PyTorch的cpp_extension模块，以及其他的一些和GPU编程相关的内容。
CUDA编程的一些工具
工欲善其事，必先利其器。在研究CUDA实现算子之前，先看看几个实用的CUDA工具，这些工具可以帮助检查CUDA代码里的访存错误和研究CUDA代码的性能瓶颈，而且对Native和Python程序都适用。
检查各算子运行时间
检查各算子的运行时间，可以使用nvprof工具，这个工具是安装CUDA时自带的。nvprof的gpu trace功能能够检测出一个程序内部各CUDA算子的运行时间、Kernel数目、使用的SRAM大小、寄存器数量等信息，可以观察CUDA程序整体的运行时间主要花在哪，是计算相关的算子，还是数据搬运相关的算子，得到一个整体的结果。
gpu trace功能的命令行如下（以运行Python脚本为例）：
nvprof --print-gpu-trace python test.py
在新版本的CUDA中，nvprof的这一功能被搬到Nsight System工具包里了，所以命令行变为：
nsys nvprof --print-gpu-trace python test.py
注意Nsight System是需要另外安装的，以上两条命令都需要管理员或root权限。
检查各算子的运行参数
对于一个运行慢的算子，我们需要知道它哪里慢了，是计算次数太多，还是访存次数太多，还是Cache命中率太低，都需要一一检查。想要获得这些参数，需要使用nvprof的metrics功能，比如我想知道各算子的内存访问总大小，可以用下面的命令行：
nvprof --metrics dram_read_bytes python .\test.py
nvprof文档给出了metrics功能支持的参数类型。可以按照需要修改命令行。
在新版本的CUDA中，nvprof的这一功能被搬到Nsight Compute工具包里了，这个工具包同样需要另外安装，以上的命令行变为：
ncu --metrics dram__bytes_read.sum python3 test.py
注意Nsight Compute使用的参数名称和nvprof不一样，Nsight Compute文档给出了两者参数名称的对应。以上的命令都需要管理员或root权限，此外，在Docker容器里执行还需要在docker run命令里添加--privileged参数，否则这两条命令无法执行。
检查算子的访存错误
由于CUDA程序基本都是直接对native的数组进行操作，因此各种访存越界、溢出等错误都有可能发生，所幸，CUDA提供了Compute Sanitizer工具帮我们在运行时检测这类错误。假设要检查名为test的算子中的访存错误，则命令行为：
compute-sanitizer --kernel-regex kns=test --tool memcheck python test.py
就会自动检测发生访存错误的指令地址、非法访问的地址以及可能的合法内存区域的地址，可以根据这个检查是越了上界还是下界。注意用Compute Santizer检测时程序会运行得非常慢，所以最好把要检测的函数提取出来做单元测试。
另外，当CUDA内部运行错误时，Python的报错代码行有时候会和实际代码行大相径庭，这时就需要在运行时指定CUDA_LAUNCH_BLOCKING环境变量为1，即：
CUDA_LAUNCH_BLOCKING=1 python test.py
这样至少更容易看到是哪个算子出错了，为进一步缩小错误范围打下基础。
测试需求
设float32类型的张量src的大小为n*m*p，unsigned char类型张量idx的大小为m*c，且每个数的大小都在[0, p)这个区间里，现在需要按照idx对src进行索引，得到一个大小为n*m*c的float32类型张量dst。用Python代码表示是这样的：
for i in range(n):
for j in range(m):
for k in range(c):
dst[i, j, k] = src[i, j, idx[j, k]]
这个需求在Pytorch里可以用gather方法实现：
_idx = torch.tile(idx.to(dtype=torch.int64).unsqueeze(0), (n, 1, 1))
dst = torch.gather(src, 2, _idx)
设n=512，m=64，p=256，c=512，以上方法在Colab的免费显卡上运行时间为1.61毫秒。很明显，这种方式存在两个问题：
gather只支持int64类型的索引。因此还需要对idx进行类型转换，浪费了转换的时间以及传输大数组的时间；gather要求idx和src的维度数量完全一致，不支持广播。因此还需要使用tile方法对idx进行复制，浪费了复制的时间以及传输大数组的时间。
因此理论上我们对这个需求用CUDA自定义一个算子，能够提高程序的运行效率。
使用Taichi库自定义算子
Taichi库是一种支持使用Python语言编写高性能算子的库（准确来说是使用Taichi语言编写，但算子直接写在Python文件里，且Taichi语言的大部分语法都和Python一致，所以可以大致认为是使用Python语言编写）。它的特色在于编写时不用考虑Kernel级的CUDA编程细节，只需要像平常写程序那样写，Taichi运行时就会自动将其编译成CUDA程序了。
根据上面的需求，可以写出以下程序：
import taichi as ti
# 定义算子
@ti.kernel
def taichi_gather(src: ti.types.ndarray(), idx: ti.types.ndarray(), dst: ti.types.ndarray(), n: ti.int32, m: ti.int32, c: ti.int32):
for i, j, k in ti.ndrange(n, m, c):
dst[i, j, k] = src[i, j, idx[j, k]]
# 执行算子
dst = torch.empty((n, m, c), device='cuda')
taichi_gather(src, idx, dst, n, m, c)
可以看到写起来和写普通Python代码没什么区别，当然像Numpy、Pytorch库里的算子是不能在Taichi算子里调用的。运行时间是0.89毫秒（注意算子编译是在第一次执行算子的时候，所以测时间需要测第二次之后的执行时间，这是JIT（即时编译）的特性，后面基于Numba、Cupy的方法也一样），可以看到比用Pytorch的gather方法快了接近一倍，主要是省去了tile的开销和int64类型的传输带宽。
注意Taichi是直接在Pytorch张量所在的显存上操作，因此需要确保张量的显存空间是连续的，可以通过张量的is_contiguous方法确认是否连续，通过contiguous方法得到一个连续张量（如果原张量本就连续，该方法返回原张量的引用，否则返回一个新的连续张量）。
使用Taichi库自定义算子非常简单方便，基本不用考虑什么细节，不过毕竟是由编译器自动将串行程序转换成并行程序。像我们的需求不存在数据依赖，因此很容易就可以编译出高效率的程序；但实际需求中数据依赖还是很复杂的，编译器行为偏保守，能够保证程序的正确性，但效率就没办法保证了。如果用过FPGA的HLS（高层次综合），就知道要“说服”编译器优化某个地方是多么棘手，因此在需要极致优化的时候反而不如用更底层的工具。
使用Numba库自定义算子
Numba库也支持编写CUDA算子直接操作Pytorch张量的显存空间，和Taichi相比，它需要从Kernel级编写算子了，但和传统的CUDA编程相比，它又支持使用Python编写CUDA程序。
同样，根据上面的需求，可以写出下面的程序：
from numba import cuda
# 定义算子
@cuda.jit
def numba_gather(src, idx, dst, n, m, c):
for i in range(cuda.blockIdx.x, n, cuda.gridDim.x):
for j in range(cuda.blockIdx.y, m, cuda.gridDim.y):
for k in range(cuda.threadIdx.x, c, cuda.blockDim.x):
dst[i, j, k] = src[i, j, idx[j, k]]
# 执行算子
dst = torch.empty((n, m, c), device='cuda')
numba_gather[(64, 32), 128](cuda.as_cuda_array(src), cuda.as_cuda_array(idx), cuda.as_cuda_array(dst), n, m, c)
可以看到，作为Kernel级的算子，就需要考虑每个GPU thread需要计算的内容。这里我将枚举n映射到GPU block的第一维，枚举m映射到GPU block的第二维，枚举c映射到GPU thread。循环的起始点和步长体现了GPU编程的一个原则，就是尽量让所有GPU thread在同一时间内并行访问的数据连续，这样可以更好地利用缓存实现并行读取。另外，相比于Taichi，Numba需要显示地将张量用cuda.as_cuda_array方法包装后才能传入算子，同时还需要自己定义GPU block各维和GPU thread各维的个数，这就需要反复尝试才能找到最优组合。我在Colab的免费显卡上得到的最优组合是这个，其他GPU可能还不太一样。
运行时间是0.83毫秒，略快于Taichi，当然如果是更复杂的算子效率提升也会更大，代价就是需要耗费更多的精力去想Kernel级的程序怎么写，还要根据GPU架构尝试各种写法和参数搭配才能实现更高的性能。
使用Cupy库自定义算子
Cupy库原本主打的是做“带Cuda优化的Numpy”，API都尽可能做到和Numpy相同，不过Pytorch也基本能做到这一点，所以大家基本都用Pytorch。但是Cupy集成了编译CUDA算子的字符串的功能，这一点还是让它有用武之地，而且它也支持直接操作Pytorch张量显存。相比于Numba，CUDA算子是用C++实现的，因此支持所有CUDA的原生API，包括像#pragma unroll这样的优化指令，这些都是Numba没有的。
根据上面的需求写出的程序如下：
import cupy as cp
# 定义算子
cupy_gather = '''
extern "C" __global__ void cupy_gather(float* src, unsigned char * idx, float *dst, int n, int m, int p, int c) {
#pragma unroll
for (int i = blockIdx.x; i < n; i += gridDim.x) {
#pragma unroll
for (int j = blockIdx.y; j < m; j += gridDim.y) {
#pragma unroll
for (int k = threadIdx.x; k < c; k += blockDim.x) {
dst[i * m * c + j * c + k] = src[i * m * p + j * p + idx[j * c + k]];
}
}
}
}
'''
cupy_gather = cp.RawKernel(cupy_gather, 'cupy_gather')
# 执行算子
dst = torch.empty((n, m, c), device='cuda')
cupy_gather(grid=(64, 32,), block=(256,), args=[src.data_ptr(), idx.data_ptr(), dst.data_ptr(), n, m, p, c])
运行时间0.50毫秒，可以看到比Numba又快了不少。优化指令的使用是一方面的原因，但并不是主要的，因为我循环边界啥的都是变量，循环展开肯定做得很不全面；更主要的原因是C++程序就没有那么多运行时检查了，而且数组的访问地址也是显式计算的，因此可以全速执行。代价就是调试的时候很酸爽，遇到非法地址的运行错误就很难受。所以在算子实现上可以先用低效率的Numba实现一遍，毕竟Numba会把数组越界之类的错误清楚地显示出来，等bug调完后再上Cupy，当然Compute Sanitizer用得熟练的也可以直接上Cupy，看具体需求。
另外，Cupy版本的算子每个block下thread的数量为256是最优配置，不知道为什么和Numba不同，自己写Kernel就很容易遇到这种玄学问题。另外Cupy是使用data_ptr方法获得张量的显存地址。
使用Pytorch的cpp_extension模块编译算子
使用cpp_extension模块的方法应该不算是JIT的范畴了，本质思路应该是先将C++和CUDA代码编译成一个动态库，然后再由Python代码加载这个动态库执行算子。但相比于那种需要定义一大堆代码，写一大堆配置，还需要把整个库的源码都下载下来才能编译的方法已经非常简单了，而且是“热更新”的，改了C++和CUDA代码之后重新运行一下Python脚本，就会自动重新编译，所以也很方便。
上面的需求实现代码分三部分：首先是CUDA代码，写在extension_gather.cu中：
// extension_gather.cu
__global__ void cuda_extension_gather(float* src, unsigned char * idx, float *dst, int n, int m, int p, int c) {
#pragma unroll
for (int i = blockIdx.x; i < n; i += gridDim.x) {
#pragma unroll
for (int j = blockIdx.y; j < m; j += gridDim.y) {
#pragma unroll
for (int k = threadIdx.x; k < c; k += blockDim.x) {
dst[i * m * c + j * c + k] = src[i * m * p + j * p + idx[j * c + k]];
}
}
}
}
void launch_extension_gather(float* src, unsigned char * idx, float *dst, int n, int m, int p, int c) {
dim3 grid(64, 32), block(256);
cuda_extension_gather<<<grid, block>>>(src, idx, dst, n, m, p, c);
}
可以看到上面Kernel和Cupy定义的是一模一样的，下面就是传统的CUDA Kernel的启动函数。然后需要写extension_gather.cpp：
// extension_gather.cpp
#include <torch/extension.h>
void launch_extension_gather(float* src, unsigned char * idx, float *dst, int n, int m, int p, int c);
torch::Tensor extension_gather(torch::Tensor &src, torch::Tensor &idx, int n, int m, int p, int c) {
assert(src.is_contiguous());
assert(idx.is_contiguous());
auto dst = torch::empty({n, m, c}, torch::device(src.device()));
launch_extension_gather((float *)src.data_ptr(), (unsigned char *)idx.data_ptr(), (float *)dst.data_ptr(), n, m, p, c);
return dst;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("extension_gather", &extension_gather, "our gather");
}
这里我引用了torch/extension.h，这样就可以使用Pytorch的C++前端了。可以看到这里也是用data_ptr方法获取张量的显存地址，同时我把创建dst张量的过程放在函数里了，这样就可以用和调用Pytorch其他方法一样的方式调用自定义算子。最后用Pytorch自带的pybind11注册函数的Python接口。Python端执行算子的过程就很简单了：
from torch.utils.cpp_extension import load
# 编译、加载算子
extension_gather = load(name='extension_gather', sources=['extension_gather.cu', 'extension_gather.cpp'], verbose=True)
# 执行算子
dst = extension_gather.extension_gather(src, idx, n, m, p, c)
运行时间为0.49毫秒，比Cupy微微快一点，原因可能是Cupy调用算子的开销要更大一些，或者是把原本在Python里创建dst的代码移到C++程序里会减少一点Python解释代码的开销。这种方法的好处主要还是代码好管理，如果要写很多算子，或者是需要引用库文件，或者是要用到什么C++高级特性的时候，可能编写完整的CUDA和C++源文件要更方便一些。
总结
作为目前最流行的深度学习库之一，Pytorch一直在推出各种方便用户使用的功能，其他高性能计算库也纷纷增加了对Pytorch的兼容和支持，这让我们自定义CUDA实现PyTorch算子变得非常简单和便捷。不过，CUDA编程博大精深，想要实现高效率的算子，并不是一件简单的事情。因此我们一方面还是应该尽可能使用Pytorch提供的算子，毕竟经过这么长时间的更新和迭代，性能肯定是最优的；另一方面也需要认真学习CUDA编程，了解GPU架构，在迫不得已需要自己实现算子的情况下善用各种优化技术和工具，才能不断榨取“性能之油”。
