基于 NVRTC 和 CUPY 的 Python CUDA 扩展

在之前的文章中，我们探讨了如何通过 cffi，扩展 pytthon （pytorch）。利用 cffi 方法，我们需要单独的 C 和 CUDA 源文件，还需要事先进行编译，不但过程比较繁琐，代码结构也稍显凌乱。对于一些简单的 CUDA 扩展（代码量不大，没有复杂的库依赖），显得不够友好。

这里，我们介绍如何通过在线编译的方式，直接为 pytorch 提供 CUDA 扩展（当然，也可以是纯 C 的扩展）。

1. 基本流程

这里， 我们尝试为 python （具体的为 pytorch） 添加 CUDA 扩展。

基本的，我们利用 pynvrtc （NVRTC 的官方python封装） 在线编译 CUDA 代码，利用 cupy（Chainer 的低层计算引擎，我们只利用其封闭 CUDA 函数的功能），来对编译后的 CUDA 代码提供 python 调用接口。完整流程如图1所示。



图1 扩展流程

2. 示例

做为例子，我们试着用 CUDA 实现 ReLU 激活函数。完整代码见这里。

ReLU(x)=max(0,x)(2)(2)ReLU(x)=max(0,x)



图2. ReLU 激活函数【src】

2.1 编写 CUDA 代码

基于 numpy 的实现如下：

def relu_forward(x):
y = x.copy()
y[x < 0] = 0
return y

def relu_backward(y_grad, x):
x_grad = y_grad.copy()
x_grad[x < 0] = 0
return x_grad

前后向都只涉及 element-wise 的操作，适合于 GPU 并行实现。对应的 CUDA 代码如下：

kernel = '''
extern "C"
__global__ void relu_forward(float *output, const float *input, int num)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.x;
for (; tid < num; tid += stride) {
output[tid] = input[tid] >= 0 ? input[tid] : 0;
}
}

extern "C"
__global__ void relu_backward(float *input_grad, const float *output_grad, const float *input, int num)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.x;
for (; tid < num; tid += stride) {
input_grad[tid] = input[tid] >= 0 ? output_grad[tid] : 0;
}
}
'''

这里， 我们将 CUDA 做为 python 的字符串，直接定在文件中，无需单独的 .cu 文件。

2.2 在线编译 CUDA 源码

我们使用 pynvrtc 提供的高层接口来编译上面定义的 CUDA 代码（更多低层接口，详细官方文档）。。编译过各如下：

from pynvrtc.compiler import Program

program = Program(kernel, 'relu.cu')
ptx = program.compile()

这里， 我们将 CUDA 源码编译为 PTX （GPU上的汇编语言）。这实际运行中，GPU 驱动会负责将 PTX 翻译为机器码进行执行。

2.3 封装 CUDA 函数

为了方便在 python 程序中直接调用，我们需要将 PTX 函数进行封装。这个可以借助 cupy 方便的实现。方法如下：

from cupy.cuda import function

m = function.Module()
m.load(bytes(ptx))

self.relu_forward = m.get_function('relu_forward')
self.relu_backward = m.get_function('relu_backward')

2.4 调用 CUDA 函数

已经有了 python 接口，通过传入 GPU 指针，可以进行函数调用。具体方法如下：

y = x.new(*x.size())

###
batch_size, hidden_size = x.size()
num = batch_size * hidden_size
grid_hidden_size = min(num, 512)
grid = (int(math.ceil(num / grid_hidden_size)), 1)

# CUDA syntax: relu_forward<<<grid, block, 0, stream>>>(...)
self.relu_forward(grid=grid, block=(grid_hidden_size, 1),
stream=stream,
args=[y.data_ptr(), x.data_ptr(), num])

对照 CUDA 调用的语法，可以看到，cupy 的封装将 CUDA 所需参数都以 python 参数的形式进行指定。

结语

可以使用 pycuda 实现 nvrtc + cupy 类似的功能，但 pycuda 社区似乎并不是特别活跃，项目更新也比较慢。

这里没有讨论 CUDA 的 stream 参数, 得到 stream 具体方法可以参见完整的代码， 及pytorch 的相关文档。