使用Cython调用CUDA Kernel函数

技术背景

前面写过一篇关于Cython和C语言混合编程的文章，在Cython中可以使用非常Pythonic的方法去调用C语言中的函数。另外我们也曾在文章中介绍过Python中使用CUDA计算的一种方案。其实从Python中去调用CUDA有很多种解决方案，例如直接使用MindSpore、PyTorch、Jax等成熟的框架进行GPU加速，也可以使用Numba的CUDA即时编译JIT来直接对Python函数进行GPU加速，还可以使用PyCuda或者Cupy来实现Python中的GPU加速。但是不论是哪一种方案，其实本质上都绕不开到C语言和Kernel函数的转换。这里提供了另外一个思路，允许我们给出Python的API接口，又可以同时使用熟悉的Python函数和C函数、CUDA函数，以达到不同的目的，这就是Python+Cython+C/CUDA的方案。

构建思路

随着Python语言在机器学习领域的大力推广，Python编程对于众多的科技工作者而言总是一个较为靠前的选项。所以不论使用何种方式构建自己的项目，一般我们都会选择面向用户开放一个Python的API接口。如果做一个项目，项目的本身对于程序性能的要求并不高的情况下，我们可以直接使用Python进行相关的计算。但如果对计算性能要求比较高，那么C/CUDA会是一个更好的选择。那么剩下一个待解决的问题就是，如何构建Python API与C/CUDA之间的交互。其实这里本身有两个方案：1. 直接把C/CUDA代码编译成*.so动态链接库，然后在Python中加载动态链接库的函数作为接口。2. 只用C/CUDA执行计算，把C/CUDA代码编译成*.so动态链接库，从Cython中对两者函数进行调度和管理。

从用户面的角度来说，第一个方案虽然可能性能还不错，但凡是需要新增一些功能模块，都需要去修改C/CUDA代码。如果只是功能模块还好说，如果涉及到任务调度和接口传输的问题，那门槛就更高了。第二个方案在性能上有可能略有衰减，但是因为接口传输和调度都是在Cython中完成的，即使是只会Python的开发者，也可以以Cython为入口来开发自己需要的模块。

案例演示

这里我们演示一个简单的Hello World程序，首先我们可以写一个hello.cu的CUDA文件：

#include<stdio.h>

__global__ void HelloKernel(void){
    printf("Hello World From GPU!\n");
}

int main(){
    HelloKernel<<<1,3>>>();
    cudaDeviceReset();
    return 1;
}

然后使用nvcc对这个CUDA项目进行编译：

$ nvcc -c ./hello.cu -Xcompiler -fPIC -o ./libhello.so

得到了一个libhello.so的动态链接文件。然后在hello.pyx中引入这个动态链接库：

# cythonize -i hello.pyx
import ctypes
libcuda = ctypes.CDLL('./libhello.so')

cpdef int run_cuda():
    cdef int res
    res = libcuda.main()
    if res:
        print ('Load cuda function successfully!')
    else:
        print ('Failed to load cuda function.')
    return 1

这里我们可以使用cythonize指令编译pyx文件，也可以使用python的setup.py配置编译指令。为了方便我们就用cythonize演示一下这个案例：

$ cythonize -i hello.pyx

那么会得到一个hello.c文件和一个hello.cpython-37m-x86_64-linux-gnu.so动态链接文件。到这里相当于我们对CUDA文件中的main函数封装了一层cython的调用，然后就可以在python文件中调用这个cython的run_cuda()函数：

# $ python3 hello.py
from hello import run_cuda
run_cuda()

然后运行这个python文件，输出为：

$ python3 hello.py
Hello World From GPU!
Hello World From GPU!
Hello World From GPU!
Load cuda function successfully!

成功的从Python侧调用了CUDA Kernel函数。

总结概要

从Python接口调用GPU进行加速的方案有很多，包括Cupy和PyCuda以及之前介绍过的Numba，还可以使用MindSpore、PyTorch和Jax等成熟的深度学习框架，这里介绍了一种直接写CUDA Kernel函数的方案。为了能够做到CUDA-C和Python编程的分离，这里引入了Cython作为中间接口，这样一来Python开发者和C开发者可以去共同开发相应的高性能方法。

版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/cython-cuda.html

作者ID：DechinPhy

更多原著文章：https://www.cnblogs.com/dechinphy/

请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html