MindSpore导入CUDA算子

技术背景

当今众多的基于Python的AI框架（如MindSpore、PyTorch等）给了开发者非常便利的编程的条件，我们可以用Python的简单的语法写代码，然后由框架在后端自动编译成可以在GPU上高效计算的程序。而对于一些定制化比较高的算法，MindSpore也支持了相关的接口，允许开发者自己开发相应的CUDA算子（需要统一接口），然后编译成.so动态链接库，再用MindSpore内置的函数加载为本地算子。本文针对这种方案写一个简单的示例。

程序结构

本地自己手写一个CUDA算子，一般至少需要两个文件和一个nvcc的环境，最好是在安装完成MindSpore的GPU版本之后，再尝试CUDA算子的引入。具体MindSpore的安装方法，可以参考MindSpore官网，这里不做赘述。我这里使用的环境是10.1版本的nvcc：

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243
$ python3 -m pip show mindspore
Name: mindspore
Version: 2.1.0
Summary: MindSpore is a new open source deep learning training/inference framework that could be used for mobile, edge and cloud scenarios.
Home-page: https://www.mindspore.cn
Author: The MindSpore Authors
Author-email: contact@mindspore.cn
License: Apache 2.0
Location: /home/dechin/anaconda3/envs/mindspore-latest/lib/python3.7/site-packages
Requires: numpy, protobuf, pillow, astunparse, scipy, psutil, asttokens, packaging
Required-by: 

需要准备的两个文件，一个是CUDA算子本身的.cu文件，另一个是用来调用CUDA算子的.py文件。操作流程是：先按照自己的需求写好CUDA算子，然后用nvcc进行编译，编译输出为.so的动态链接库，然后在python脚本中使用mindspore.ops.Custom生成相应的算子。在MindSpore2.1之后的版本中，对于本地CUDA算子的调用定义了统一的接口，其格式为：

extern "C" int CustomFunc(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);

具体方法可以参考官网的这一篇文档说明。这样的话，在.cu文件中至少有两个函数，一个是原本用于计算的Kernel函数，另一个是用于统一标准接口的CustomFunc函数。需要说明的是，旧版本的MindSpore是没有这样的规范的，所以旧版本的算子没有CustomFunc函数也能够用nvcc编译，但是无法在新版本的MindSpore中调用。

一维张量求和

我们用一个一维的张量求和的示例来演示一下如何在本地写一个可以用MindSpore来调用的CUDA算子，一维张量求和的算法是比较简单的：

\[C_i=A_i+B_i \]

那么对应的CUDA算子的代码如下所示：

// custom_add.cu
// nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu
// 常量，一般可以放在.cuh头文件中
constexpr int THREADS = 1024;
// 用于CUDA计算的Kernel函数
__global__ void CustomAddKernel(float *input1, float *input2, float *output, size_t size) {
    auto idx = blockIdx.x * THREADS + threadIdx.x;
    if (idx < size) {
        // 逐元素操作，CUDA算子的基本写法
        output[idx] = input1[idx] + input2[idx];
    }
}
// 标准算子接口
extern "C" int CustomAdd(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream,
                         void *extra) {
    cudaStream_t custream = static_cast(stream);
    // 输出变量的位置
    constexpr int OUTPUT_INDEX = 2;
    // 传入的参数都是指针形式
    float *input1 = static_cast<float *>(params[0]);
    float *input2 = static_cast<float *>(params[1]);
    float *output = static_cast<float *>(params[2]);
    // 获取输入张量的大小
    int size = shapes[OUTPUT_INDEX][0];
    // GPU运算中的block和thread，一般要求block*thread大于或者等于size即可
    int blocks = ceil(size / THREADS) + 1;
    // 调用Kernel函数
    CustomAddKernel<<0, custream>>>(input1, input2, output, size);
    return 0;
}

值得注意的是，上述CustomAdd函数中的params有3个输入，但是实际上其中一个是返回值，这也是MindSpore对于标准接口的设定，不需要我们额外传入一个变量。保存好上述的CUDA算子代码之后，可以用如下指令直接编译成python可以调用的动态链接库：

$ nvcc --shared -Xcompiler -fPIC -o custom_add.so custom_add.cu

编译完成后，会在当前目录下生成一个新的.so文件，然后就可以在python代码中进行调用：

# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
 
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
                       out_shape=[t1.shape[0]],
                       out_dtype=ms.float32,
                       func_type="aot"
                       )
res = CustomAdd(t1, t2)
ops.print_(res)

上述的CustomAdd就是我们导入的基于CUDA算子来写的本地MindSpore算子，并且这种算子还可以使用MindSpore进行注册，这样就不需要每次使用都去加载这个动态链接库，感兴趣的童鞋可以自己研究一下算子注册的方法。上述Python代码的运行结果如下：

$ python3 test_custom_ops.py 
Tensor(shape=[3], dtype=Float32, value= [ 4.00000000e+00,  4.00000000e+00,  4.00000000e+00])

可见跟我们预期的结果是一致的，那么就完成了一个本地CUDA算子的实现和调用。

自定义算子反向传播

在前面的章节里面我们已经实现了一个本地的一维张量求和算子，但是这还不是一个完整的算子实现，因为我们在AI框架中一般要求算子可自动微分，光实现一个算子是不完整的，因此这里我们再通过bprop函数，来实现一下自定义算子的反向传播：

# test_custom_ops.py
# python3 test_custom_ops.py
import mindspore as ms
from mindspore import ops, Tensor, context
from mindspore import nn, grad
ms.set_context(device_target="GPU", mode=context.GRAPH_MODE)
 
t1 = Tensor([1., 2., 3.], ms.float32)
t2 = Tensor([3., 2., 1.], ms.float32)
CustomAdd = ops.Custom("./custom_add.so:CustomAdd",
                       out_shape=[t1.shape[0]],
                       out_dtype=ms.float32,
                       func_type="aot"
                       )
为了自动微分，我们需要定义一个Cell类来封装我们的自定义算子：
class Add(nn.Cell):
    # 反向传播函数
    def bprop(self, x, y, out, dout):
        return (y, )
    # 计算函数
    def construct(self, x, y):
        return CustomAdd(x, y)
# 把Cell类加载为custom_add函数
custom_add = Add()
# 计算求和结果
res = custom_add(t1, t2)
# 计算自动微分结果
res_g = grad(custom_add, grad_position=(0, ))(t1, t2)
print(res)
print(res_g)

在这个代码中，主要就是增加了一个Cell类和两个新的函数bprop和construct，分别用于计算函数的反向传播和正向值，代码运行的结果如下：

$ python3 test_custom_ops.py 
[4. 4. 4.]
[3. 2. 1.]

当然，这里我们没有再额外写一个用于返回反向传播值的CUDA算子，但是原则上对于较为复杂的函数，是需要自己手动写一个用于求微分数值的CUDA算子的。

总结概要

本文介绍了在MindSpore标准格式下进行CUDA算子开发的方法和流程，可以让开发者在现有的AI框架下仍然可以调用基于CUDA实现的高性能的算子。并且，除了常规的数值计算之外，在MindSpore框架下，我们还可以通过实现一个bprop函数，使得我们手写的这个CUDA算子也可以使用MindSpore框架本身自带的自动微分-端到端微分技术。

版权声明

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

作者ID：DechinPhy

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

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