在之前的文章中，我們探討了如何通過 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 激活函數。完整代碼見這裏。
圖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 參數的形式進行指定。