2.4.2　添加C++编写的自定义算子

算子是构建神经网络的基础，也称为低级API；通过算子的封装可以实现各类神经网络层，当开发神经网络层遇到内置算子无法满足时，可以通过自定义算子实现。以MindSpore为例，实现一个GPU算子需要如下步骤：

（1）原语（Primitive）注册：算子原语是构建网络模型的基础单元，用户可以直接或者间接调用算子原语搭建一个神经网络模型。

（2）GPU算子开发：GPU Kernel用于调用GPU实现加速计算。

（3）GPU算子注册：算子注册用于将GPU Kernel及必要信息注册给框架，由框架完成对GPU Kernel的调用。

1．注册算子原语

算子原语通常包括算子名、算子输入、算子属性[初始化时需要填的参数，如卷积的步长（stride）、填充（padding）]、输入数据合法性校验、输出数据类型推导和维度推导。假设需要编写加法算子，主要内容如下：

（1）算子名：TensorAdd。

（2）算子属性：在构造函数__init__中初始化属性，因加法没有属性，因此__init__不需要额外输入。

（3）算子输入／输出及合法性校验：infer_shape方法中约束两个输入维度必须相同，输出的维度和输入维度相同。infer_dtype方法中约束两个输入数据必须是float32类型，输出的数据类型和输入数据类型相同。

（4）算子输出。

MindSpore中实现注册TensorAdd，如代码2.15所示。

代码2.15　MindSpore实现注册TensorAdd

在mindspore/ops/operations/math_ops.py文件内注册加法算子原语后，需要在mindspore/ops/operations/__init__中导出，方便Python导入模块时候调用，如代码2.16所示。

代码2.16　导出注册算子

2．GPU算子开发

继承GPU Kernel，实现加法使用类模板定义TensorAddGpuKernel，需要实现以下方法：

（1）Init():用于完成GPU Kernel的初始化，通常包括记录算子输入／输出维度，完成Launch前的准备工作；因此在此记录Tensor元素个数。

（2）GetInputSizeList():向框架反馈输入Tensor需要占用的显存字节数；返回了输入Tensor需要占用的字节数，TensorAdd有两个输入，每个输入占用字节数为element_num∗sizeof(T)。

（3）GetOutputSizeList():向框架反馈输出Tensor需要占用的显存字节数；返回了输出Tensor需要占用的字节数，TensorAdd有一个输出，占用element_num∗sizeof(T)字节。

（4）GetWorkspaceSizeList():向框架反馈工作空间（Workspace）字节数，工作空间是用于计算过程中存放临时数据的空间；由于TensorAdd不需要工作空间，因此GetWorkspace-SizeList()返回空的std::vector<size_t>。

（5）Launch():通常调用CUDA Kernel（CUDA Kernel是基于Nvidia GPU的并行计算架构开发的核函数），或者cuDNN接口等方式，完成算子在GPU上加速；Launch()接收输入、输出在显存的地址，接着调用TensorAdd完成加速。

GPU算子开发参见代码2.17。

代码2.17　GPU算子开发

TensorAdd中调用了CUDA kernelTensorAddKernel来实现element_num个元素的并行相加，如代码2.18所示。

代码2.18　实现并行相加

3．GPU算子注册

GPU算子信息包含：①Primive；②Input dtype，output dtype；③GPU Kernel class；④CUDA内置数据类型。框架会根据Primive和Input dtype、output dtype，调用以CUDA内置数据类型实例化GPU Kernel class模板类。代码2.19中分别注册了支持float（浮点型数）和int（整型数）的TensorAdd算子。

代码2.19　GPU算子注册

完成上述三步工作后，需要把MindSpore重新编译，在源码的根目录执行bash build.sh-e gpu，最后使用算子进行验证。