基于 NVRTC 和 CUPY 的 Python CUDA 扩展

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

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

注： pytorch 0.4+ 提供了新的辅助工具，简化了 C++/CUDA 扩展，参见 tutorial。

1. 基本流程

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

图1 扩展流程

2. 示例

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