如何设计一个高内聚低耦合的模块——MegEngine 中自定义 Op 系统的实践经验

背景介绍

在算法研究的过程中，算法同学们可能经常会尝试定义各种新的神经网络层（neural network layer），比如 Layer Norm，Deformable Conv 等。

为了实现这些层以进行实验，算法同学可以使用神经网络框架或者 numpy 中提供的基础操作（如张量/标量的加减乘除等）去组合出所需的层的功能。

然而这通常会造成这些层的性能断崖式的下跌，大大影响了算法同学们尝试新算法的效率。所以很多情况下，算法同学们会选择为自己定义的层实现高性能的 kernel，并希望可以将之集成入框架作为框架中的 Op（Operator） 使用。

然而在一般情况下，算法同学必须要对框架本身十分了解，才可以灵活自由的将我们的 kernel 接入框架中去使用。但这并不是一件简单的事，神经网络框架作为一个规模庞大的系统，其结构之复杂远超常规的软件项目。

为了保证这样的机器学习系统的可维护性以及可拓展性，系统中往往会做各种各样的层次，模块设计，并对各种概念（比如 Op）进行抽象，各个层次各个模块间的交互又非常复杂。

以 MegEngine 中的 Op 系统为例，图 1 中展示 Op 这一最基本的概念在 MegEngine 中不同层的各种抽象。

图 1：MegEngine  中不同层次的  Operator  的抽象

• 在底层的 MegDNN 算子库中，Op 被抽象成了 MegDNNOpr 类，其中封装了各个 Op 在 x86，nvidia gpu 等硬件平台上的具体 kernel 实现以及相关硬件的 context 管理。

• 在静态图（graph runtime）中，Op 被抽象成了 OpNode 类，其主要目的并非是用于计算而是图优化，故而这个数据结构的设计上又有了相当多这方面的考量。

• 在动态图（imperative runtime）中，Op 又会被抽象为 OpDef 类，配合动态图解释器进行任务的调度。

• 在 python 中，Op 会被封装成 functional 和 module，这才是符合一般算法同学认知的 Op。

而从 python 中执行一个操作时，这些 Op 会逐层向下调用，分别在每一层完成一部分工作，直到最后才调用了 MegDNN 算子库中具体的 kernel。这个过程中，任何一个层次的 Op 概念都是缺一不可的。

其实不只是 Op，包括 Tensor 在内的很多其他概念，在 MegEngine 系统中都存在着类似的多种抽象。学习了解这样的一个框架设计需要花费大量的时间和精力，这个代价往往是算法同学难以接受的。

然而即使算法同学牺牲了大把时间和头发，学习了解了 MegEngine 的系统设计，完成了自己 kernel 到 MegEngine 的集成，事情还远远没有结束。kernel 的集成过程通常与框架本身是高度耦合的。

其构建 Op 的过程需要获取框架的所有源码，修改编译框架中的绝大多数模块。如果之后框架内部的相关抽象发生变化，则之前构建的 Op 则又变为不可用的状态。

为了允许把算法同学的 kernel 快速的集成入框架去进行使用，并且集成出来的 Op 既可以与框架内的原生 Op 有着一致的行为，同时其又与框架本身相解耦，MegEngine 提出了一套工具 Custom Op。

其可以很简单便捷的将算法同学自己编写的 c++/cuda kernel 封装成 Op 并自动化的编译成动态链接库并集成入 MegEngine 中。

然而，编写高性能的 c++/cuda kernel 对于一般没有体系结构/并行计算背景的算法研究人员依然是一件很困难的事情。

所以为了避免算法同学自行编写 kernel 的种种问题，MegEngine 基于 Custom Op 进一步提出了 Custom Op Generator，尝试利用神经网络编译器代码生成的方式去自动化端到端的生成 kernel 和 Custom Op 代码，并将之集成入 MegEngine，使算法同学无需编写任何 c++ 代码，即可在 MegEngine 中添加高性能的 kernel 并使用。

正常的 Op 集成与 Custom Op

为了便于理解 Custom Op 设计，我们首先对传统 Op 的集成过程和 Custom Op 的集成过程进行对比，建立起 Custom Op 的初步印象。

一般而言，我们的算法同学想将自己编写的 c++/cuda kernel 集成为 MegEngine 的 Op，那么他首先必须了解：

• MegEngine 整个的系统结构。

• MegEngine 中各种层次模块的功能。

• 诸如 Op，Tensor 等概念在不同层次模块的设计目的以及实质含义。

在对这样一个系统有了十足了解之后，其需要：

• 在 MegDNN 算子层中对自己的 kernel 进行封装，将之封装成 MegDNNOpr 类。

• 基于 MegEngine 中静态图中的相关组件将自己的 Op 封装成 OpNode 类。