丁丁的博客

原创 基于 NVRTC 和 CUPY 的 Python CUDA 扩展

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