大家好，我想在这里给大家介绍我的一个项目：MobulaOP.
MobulaOP是一个简单且灵活的跨框架算子创建工具包。不需要重新编译深度学习框架的源码，就可以创建自定义的C++算子。而且只需要一份C++代码实现和简单的定义，自定义算子就可以在CPU和GPU上运行。
之所以建立这个项目，是因为我发现MXNet创建自定义算子的方法不太方便，其他深度学习框架也同样存在这个问题。
当前，创建自定义算子的方法主要有：
- 重新编译深度学习框架的源码
重新编译源码耗时过长。需要了解对应框架的算子实现形式，编写出的代码不适用于其他框架。
- 重新编译深度学习框架的源码
- 使用运行时编译(Run-Time Compilation)API
需要编写对应的CUDA代码，编写过程较复杂，无法在CPU环境下进行调试。
- 使用运行时编译(Run-Time Compilation)API
- 加载动态文件
需要了解对应框架的动态加载实现形式，编写较复杂，一份代码不适用于多个框架。
- 加载动态文件
因此，我设计了MobulaOP项目，希望能解决上述问题。
MobulaOP项目当前的特性有：
- 项目实现精简，不需要重新编译深度学习框架，就可以实现自定义的C++ operator;
- 只需要编写一份代码，就可以让自定义算子运行在不同设备(CPU/GPU)，以及不同的深度学习框架(如MXNet, PyTorch)或数值计算库NumPy上；
- 在编写自定义层的过程中，用户有更多的注意力关注在运算的实现上；
- 对MXNet有更多的支持，使用MobulaOP可以更方便地创建自定义算子(Custom Operator).
MobulaOP暂时只支持Linux系统，之后会加入对Windows等系统的支持。
下面，我想简单地介绍一下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循环函数。这个函数的第一个参数为循环体的总迭代数，第二个参数为一个接收迭代下标的函数，这里使用了匿名函数。下标i从0开始计数，满足0 <= i < n。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下单精度浮点数(float32)的
atomic_add原子加函数； - 在一个核函数内，允许多次调用
parfor函数, 这些parfor的总迭代数可以不同，但实际使用的线程数是相同的； -
parfor函数只允许在核函数内部进行调用； - 如果要在核函数中调用其他函数，被调用的函数的声明前需要添加宏
MOBULA_DEVICE, 并声明返回值类型。
例子：返回两个数中的最大值
template <typename T>
MOBULA_DEVICE T maximum(const T a, const T b) {
return a >= b ? a : b;
}
执行核函数
接下来，使用MobulaOP执行上述核函数。
MobulaOP能够自动分析、生成代码，并调用编译器将代码编译为动态链接库。
把上述核函数保存为MulElemWise.cpp文件，放在如下的文件目录结构中:
tutorial
└── MulElemWise
└─── MulElemWise.cpp
在tutorial文件夹下创建test_mul_func.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]
在终端中输入python test_mul_func.py即可执行。
这段代码中，与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)
如何将核函数封装成一个算子(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), 这里的assign函数和MXNet是一致的。
测试自定义算子
编写好MulElemWise算子的定义后，来测试一下吧。
在tutorial文件夹下创建文件test_mul_op.py, 输入代码：
import mobula
mobula.op.load('MulElemWise')
import mxnet as mx
a = mx.nd.array([1,2,3])
b = mx.nd.array([4,5,6])
a.attach_grad()
b.attach_grad()
with mx.autograd.record():
c = mobula.op.MulElemWise(a, b)
c.backward()
print (c) # [4, 10, 18]
print ('a.grad = {}'.format(a.grad.asnumpy())) # [4, 5, 6]
print ('b.grad = {}'.format(b.grad.asnumpy())) # [1, 2, 3]
同样，在终端输入python test_mul_op.py指令执行。
这里与MobulaOP有关的新代码是第11行: c = mobula.op.MulElemWise(a, b)
MobulaOP加载MulElemWise模块后，分析了MulElemWise文件夹下的MulElemWise.cpp文件，把核函数注册到mobula.func中；同时加载同一个文件夹下的MulElemWise.py文件，将算子注册到mobula.op中。这个过程没有发生编译。
当mobula.op.MulElemWise(a, b)执行时，MobulaOP会根据变量类型，自动编译所需要的动态链接库，并返回结果。
mobula.op.MulElemWise也可以接受MXNet的符号(Symbol)、Numpy数组或PyTorch Tensor.
例子：
MXNet的符号(Symbol):
a_sym = mx.sym.Variable('a')
b_sym = mx.sym.Variable('b')
c_sym = mobula.op.MulElemWise(a_sym, b_sym)
Numpy数组：
a_np = np.array([1,2,3])
b_np = np.array([4,5,6])
# 由于Numpy不支持记录梯度，因此需要一个实例记录梯度
op = mobula.op.MulElemWise[np.ndarray]()
c_np = op(a_np, b_np)
如何在Gluon内使用MobulaOP定义的算子呢？
我们可以这样写：
class MulElemWiseBlock(mx.gluon.nn.HybridBlock):
def hybrid_forward(self, F, a, b):
return mobula.op.MulElemWise(a, b)
这就是MobulaOP的简单使用介绍，上述代码可以在项目的文档部分(docs)查看。
希望MobulaOP能够对大家有帮助。
同时，欢迎大家对MobulaOP项目提Issue和PR. 谢谢！
