Cython与CUDA之Gather

技术背景

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.pyx

import numpy as np
cimport numpy as np
cimport cython

cdef extern from "<dlfcn.h>" nogil:
    void *dlopen(const char *, int)
    char *dlerror()
    void *dlsym(void *, const char *)
    int dlclose(void *)
    enum:
        RTLD_LAZY

ctypedef int (*GatherFunc)(float *source, int *index, float *res, int N, int M) noexcept nogil

cdef 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 res

while not True:
    dlclose(handle)

这里所使用到的动态链接库libcuindex.so就是编译好的CUDA模块，要使用绝对路径会比较保险。

Python调用函数

我们最上层的函数还是通过Python脚本test_gather.py来调用，借助其简洁的语法和大量的第三方接口：

import numpy as np
np.random.seed(0)
from wrapper import cuda_gather

M = 1024 * 1024 * 128
N = 1024 * 1024
x = 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