这可能是创建自定义C++ Operator最简单的方式 - MobulaOP使用说明

大家好，我想在这里给大家介绍我的一个项目：MobulaOP. MobulaOP是一个简单且灵活的跨框架算子创建工具。使用MobulaOP, 不需要重新编译深度学习框架的源码，就可以创建自定义的C++算子。而且只需要一份C++代码和简单的定义，自定义算子就可以在CPU和GPU上运行。

之所以建立这个项目，是因为我发现MXNet创建自定义算子的方法不太方便，其他深度学习框架也存在这个问题。当前，创建自定义算子的方法主要为：

重新编译深度学习框架的源码重新编译源码耗时过长。需要了解对应框架的算子实现形式，编写出的代码不适用于其他框架。
使用运行时编译(Run-Time Compilation)API 需要编写对应的CUDA代码，编写过程较复杂，无法在CPU环境下进行调试。
加载动态文件需要了解对应框架的动态加载实现形式，编写较复杂，一份代码不适用于多个框架。

因此，我设计了MobulaOP项目，希望能解决这些问题。 MobulaOP项目当前的特性有：

项目实现精简，不需要重新编译深度学习框架，就可以实现自定义的C++ operator;
只需要编写一份代码，就可以让自定义算子运行在不同设备(CPU/GPU)，以及不同的深度学习框架(如MXNet, PyTorch)或数值计算库NumPy上；
在编写自定义层的过程中，用户有更多的注意力关注在运算的实现上；
对MXNet有更多的支持，使用MobulaOP可以更方便地创建自定义算子(Custom Operator).

下面，我想简单地介绍一下MobulaOP的使用方法。

配置MobulaOP

在终端下输入以下命令

# 将MobulaOP项目拷贝下来
git clone https://github.com/wkcn/MobulaOP
# 进入项目文件夹
cd MobulaOP
# 安装依赖库numpy, pyyaml和easydict
pip install -r requirements.txt
# 进行编译，如果需要在GPU下使用，在选项中输入y
sh build.sh
# 将MobulaOP文件夹加入PYTHONPATH环境变量中
export PYTHONPATH=$PYTHONPATH:$(pwd)

当执行完以上命令后，在项目目录外打开Python交互界面，输入import mobula，如果没有提示，则配置成功。

核函数

配置好MobulaOP后，就可以使用C++编写算子(operator)的运算函数了。这里把并行计算的运算函数称为核函数。以创建一个逐位乘法算子为例，它的实现为：

template <typename T>
MOBULA_KERNEL mul_elemwise_kernel(const int n, const T* a, const T* b, T* out) {
    parfor(n, [&](int i) {
        out[i] = a[i] * b[i];
    });
}

没错，定义一个逐位乘法函数只需要6行代码，并且它支持在CPU和GPU下运行。其中，MOBULA_KERNEL宏声明了这个函数是一个核函数，核函数不需要定义返回值，同时核函数的函数名末尾为_kernel. 对于参数列表，MobulaOP要求第一个参数为并行计算的线程数。MobulaOP会自动将参数列表中const T*类型的参数识别为输入数组的指针，将T*类型的参数识别为输出数组的指针。在函数中，有一个parfor函数，代表并行地执行这个函数。这个函数的第一个参数为循环体的总迭代数，第二个参数为一个接收迭代下标的函数，这里是一个匿名函数。MobulaOP会根据运行设备对parfor进行不同的展开。当这段代码在CPU下运行时，MobulaOP会将这段函数展开为：

for (int i = 0; i < n; ++i) {
	out[i] = a[i] * b[i];
}

MobulaOP会自动地使用多线程、OpenMP、CUDA等方法并行地执行这个循环。需要注意的是：

MOBULA_KERNEL核函数的第一个参数为调用这个函数进行并行计算的线程数；
核函数内部语句均为并行执行，编写核函数时要注意线程安全问题，当前MobulaOP提供了CPU/GPU下的atomic_add原子加函数。
允许在一个核函数内设置多次调用parfor, 这些parfor的总迭代数可以不同，但实际使用的线程数是相同的;
如果要在核函数中调用其他函数，被调用的函数的声明前需要添加宏MOBULA_DEVICE, 这里需要声明返回类型，例子：

template <typename T>
MOBULA_DEVICE T maximum(const T a, const T b) {
        return a >= b ? a : b;
}

接下来，使用MobulaOP执行这段代码。 MobulaOP能够自动分析、生成代码，并调用编译器将代码编译为动态链接库。

把C++函数保存为MulElemWise.cpp文件，放在如下的文件目录结构中:

tutorial
└── MulElemWise
    └─── MulElemWise.cpp

在tutorial文件夹下创建test_mul_elemwise.py文件，在这个文件中编写Python代码：

import mobula
mobula.op.load('MulElemWise')

import mxnet as mx
a = mx.nd.array([1,2,3])
b = mx.nd.array([4,5,6])
out = mx.nd.empty(a.shape)
mobula.func.mul_elemwise(a.size, a, b, out)
print (out)  # [4, 10, 18]

这段代码中，与MobulaOP相关的一共有三行(第1、2、8行) 第1行代码将导入MobulaOP包。第2行代码将加载MulElemWise模块。MobulaOP将会搜索MulElemWise文件夹中是否存在同名的cpp或py文件，以及__init__.py文件。若找到这些文件，将会对文件进行编译或加载。mobula.op.load也支持指定搜索目录，如mobula.op.load('MulElemWise', os.path.dirname(__file__)). 第8行调用核函数mul_elemwise，与函数声明MOBULA_KERNEL mul_elemwise_kernel(const int n, const T* a, const T* b, T* out);对比，在Python中调用的函数名比C++中的函数名少了后缀_kernel. MobulaOP把加载后的核函数添加到mobula.func中，调用mobula.func.<核函数名>即可调用C++函数，并且MobulaOP能够自动转换参数的类型, 包括处理非内存连续数组、根据参数的输入输出类型自动调用wait_to_read、wait_to_write函数等功能。

如何将核函数封装成一个算子(operator)呢，MobulaOP提供了一个简单的声明方法。在tutorial/MulElemWise文件夹下创建文件MulElemWise.py, 输入以下代码：

import mobula

@mobula.op.register
class MulElemWise:
    def forward(self, a, b):
        mobula.func.mul_elemwise(a.size, a, b, self.y)
    def backward(self, dy):
        self.dX[0][:] = self.F.multiply(dy, self.X[1])
        mobula.func.mul_elemwise(dy.size, dy, self.X[0], self.dX[1])
    def infer_shape(self, in_shape):
        assert in_shape[0] == in_shape[1]
        return in_shape, [in_shape[0]]

第3行的@mobula.op.register为一个Python装饰器，它将其下面的类注册为算子。一个算子类需要定义forward, backward以及infer_shape函数。在forward函数的参数列表中，a和b是算子前向传播的输入；在backward函数的参数列表中，dy为算子后向传播时输入的导数。 MobulaOP会根据forward函数得到算子的输入个数和名称，根据backward得到输出个数。 infer_shape函数传入的是元组(tuple)的列表，分别表示各输入的尺寸(shape). infer_shape的返回值有两个值，第一个值是各个输入的尺寸，第二个值是各个输出的尺寸。infer_shape和MXNet自定义层里的infer_shape是等价的。

在算子的forward和backward函数中，定义了一些变量：

变量名	描述
self.F	当前环境，假如使用MXNet, self.F = mx.nd
self.X[k]	第k个输入
self.Y[k]	第k个输出
self.dX[k]	第k个输入的导数
self.dY[k]	第k个输出的导数
self.x	第1个输入
self.y	第1个输出
self.dx	第1个输入的导数
self.dy	第1个输出的导数
self.req[k]	第k个输入/输出的处理模式(null/write/add/replace)

需要注意的是，赋值的时候需要加上[:]，如self.X[0][:] = data 也可以使用内置的assign函数进行赋值，如self.assign(self.X[0], self.req[0], data), 这里和MXNet是一致的。

wkcn/MobulaOP-tutorial-cn.md

这可能是创建自定义C++ Operator最简单的方式 - MobulaOP使用说明

配置MobulaOP

核函数