使用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函数。

其他调用方法

前面提到，我们也可以在C程序中直接调用这个CUDA函数。例如在上面我们编译好libhello.so的CUDA动态链接库之后，用一个C文件去调用动态链接库：

#include <dlfcn.h>

int main()
{
    void* handle = dlopen("./libhello.so", RTLD_LAZY);
    int (*helloFromGPU)();
    helloFromGPU = dlsym(handle, "main");
    helloFromGPU();
    dlclose(handle);
    return 1;
}

然后使用如下指令编译成一个可执行文件：

$ gcc -o hello_c hello_c.c -ldl

再执行这个可执行文件，可以得到跟Python程序类似的打印结果。如果需要将这部分的C代码也编译成动态链接库供Python调用的话，可以使用如下指令分两次编译：

$ gcc -fPIC -c hello_c.c -o hello_c.o
$ gcc -shared hello_c.o -o libhello_c.so

路径问题

如果需要从其他路径（例如动态链接库的公共存储位置）来引入相关的动态链接库的话，可以在Python中获取环境变量。例如，我们先在环境中配置一个LD_LIBRARY_TEST_PATH：

$ export LD_LIBRARY_TEST_PATH=$(pwd)

这里是把当前路径的绝对路径设置为环境变量中的一个地址。然后在Python中可以查询到这个环境变量：

In [1]: import os

In [2]: os.getenv('LD_LIBRARY_TEST_PATH')
Out[2]: '/home/mindsponge/tests/toolkits'

这样在很多时候可以从外部解决库的引用地址的问题，不需要每次都去手动修改或者配置。

总结概要

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