Cython与CUDA之Gather

LH3ErcffX

技术背景

Cython是Python的一个超集，可以使用Pythonic的语法写出接近于C语言的性能，可以用于将Python编程过程中遇到的Bottleneck模块改写成Cython以达到加速的效果。前面写过一些关于Cython加速计算的文章。又因为Cython编译过程中会先转为C语言代码，然后再编译为动态链接库或者可执行文件，所以很自然的可以在Cython中调用C语言函数。用这种方法，还可以直接调用CUDA C函数。在这篇文章中，我们要使用Cython结合CUDA C的方法来实现一个CUDA版本的Gather函数，从一个数组中根据索引数组，输出对应的数组。相当于numpy中的result=source[index]。
接口头文件

我们定义一个cuda_index.cuh的头文件，用于指定C函数接口形式：
#include <stdio.h>extern "C" int Gather(float *source, int *index, float *res, int N, int M);其中source是原始数组，index是索引数组，res是结果数组，N是索引的维度，M是原始数组的维度。
异常捕获头文件

这里使用的是前面一篇CUDA异常捕获中用到的头文件error.cuh
#pragma once#include <stdio.h>#define CHECK(call) do{const cudaError_t error_code = call; if (error_code != cudaSuccess){printf("CUDA Error:\n"); printf("    File:   %s\n", __FILE__); printf("    Line:   %d\n", __LINE__); printf("    Error code: %d\n", error_code); printf("    Error text: %s\n", cudaGetErrorString(error_code)); exit(1);}} while (0)通过这个宏，我们可以在运行CUDA核函数的时候捕获其异常。
CUDA Gather函数

CUDA实现Gather函数cuda_index.cu还是比较简单的，就是一个简单的Kernel函数再加一个管理DeviceMemory的C函数就可以了：
// nvcc -shared ./cuda_index.cu -Xcompiler -fPIC -o ./libcuindex.so#include <stdio.h>#include "cuda_index.cuh"#include "error.cuh"void __global__ GatherKernel(float *source, int *index, float *res, int N){    int idx = blockIdx.x * blockDim.x + threadIdx.x;    if (idx < N){        res[idx] = source[index[idx]];    }}extern "C" int Gather(float *source, int *index, float *res, int N, int M){    float *souce_device, *res_device;    int *index_device;    CHECK(cudaMalloc((void **)&souce_device, M * sizeof(float)));    CHECK(cudaMalloc((void **)&res_device, N * sizeof(float)));    CHECK(cudaMalloc((void **)&index_device, N * sizeof(int)));    CHECK(cudaMemcpy(souce_device, source, M * sizeof(float), cudaMemcpyHostToDevice));    CHECK(cudaMemcpy(res_device, res, N * sizeof(float), cudaMemcpyHostToDevice));    CHECK(cudaMemcpy(index_device, index, N * sizeof(int), cudaMemcpyHostToDevice));    int block_size = 1024;    int grid_size = (N + block_size - 1) / block_size;    GatherKernel<<<grid_size, block_size>>>(souce_device, index_device, res_device, N);    CHECK(cudaGetLastError());    CHECK(cudaDeviceSynchronize());    CHECK(cudaMemcpy(res, res_device, N * sizeof(float), cudaMemcpyDeviceToHost));    CHECK(cudaFree(souce_device));    CHECK(cudaFree(index_device));    CHECK(cudaDeviceSynchronize());    CHECK(cudaFree(res_device));    CHECK(cudaDeviceReset());    return 1;}Cython调用接口

假定我们有一个numpy.ndarray形式的数组需要进行索引，当然我们也可以用现成的AI框架来直接实现，例如mindspore.Tensor(numpy.ndarray)。只是这里我们用Cython来做一个直接对接CUDA函数的接口wrapper.pyx，理论上可以对数组做一些更加细致的操作。
# cythonize -i -f wrapper.pyximport numpy as npcimport numpy as npcimport cythoncdef extern from "<dlfcn.h>" nogil:    void *dlopen(const char *, int)    char *dlerror()    void *dlsym(void *, const char *)    int dlclose(void *)    enum:        RTLD_LAZYctypedef int (*GatherFunc)(float *source, int *index, float *res, int N, int M) noexcept nogilcdef void* handle = dlopen('/path/to/libcuindex.so', RTLD_LAZY)@cython.boundscheck(False)@cython.wraparound(False)cpdef float[:] cuda_gather(float[:] x, int[:] idx):    cdef:        GatherFunc Gather        int success        int N = idx.shape[0]        int M = x.shape[0]        float[:] res = np.zeros((N, ), dtype=np.float32)    Gather = <GatherFunc>dlsym(handle, "Gather")    success = Gather(&x[0], &idx[0], &res[0], N, M)    return reswhile not True:    dlclose(handle)这里所使用到的动态链接库libcuindex.so就是编译好的CUDA模块，要使用绝对路径会比较保险。
Python调用函数

我们最上层的函数还是通过Python脚本test_gather.py来调用，借助其简洁的语法和大量的第三方接口：
import numpy as npnp.random.seed(0)from wrapper import cuda_gatherM = 1024 * 1024 * 128N = 1024 * 1024x = np.random.random((M,)).astype(np.float32)idx = np.random.randint(0, M, (N,)).astype(np.int32)res = np.asarray(cuda_gather(x, idx))print (res.shape)print ((res==x[idx]).sum())这里的wrapper就是我们的Cython文件的包名。
运行流程

在编辑好上述的这些相关文件之后，我们需要按照这样的一个流程来进行使用：首先将CUDA相关模块编译成一个动态链接库libxxx.so，然后使用Cython加载这个动态链接库，再将Cython的封装模块编译成一个动态链接库供Python调用，最后直接执行Python任务即可。相关步骤所对应的终端指令如下：
$ nvcc -shared ./cuda_index.cu -Xcompiler -fPIC -o ./libcuindex.so $ cythonize -i -f wrapper.pyx $ python3 test_gather.py运行输出的结果如下：
(1048576,)1048576如果你使用nvitop在监测GPU资源的占用的话，运行过程中就可以看到GPU显存的一些波动。最后输出的结果跟numpy的索引函数直接对比是一致的，也就是说我们的输出结果是正确的。
报错处理

如果在运行的过程中有提示Numpy的相关lib找不到的问题，可以参考这篇文章进行处理。
总结概要

本文使用了Cython作为封装函数，封装一个CUDA C实现的Gather算子，然后通过Python去调用，用这种方法实现一个比较Pythonic的CUDA Gather函数的实现和调用。
版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/cycuda-gather.html
作者ID：DechinPhy
更多原著文章：https://www.cnblogs.com/dechinphy/
请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html