作者：褚超群 | 旷视科技 MegEngine 架构师

背景介绍

在算法钻研的过程中，算法同学们可能常常会尝试定义各种新的神经网络层（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 类。
• 基于 MegEngine 中动态图中的相干组件将本人的 Op 封装成 OpDef 类。
• 编写 python 和 C++ 的交互代码，将本人的 Op 裸露到 python 环境中。

而为了简化这个过程，Custom Op 中提供了一套框架无关的非常简洁的 Op 的模型，算法用户在增加 Op 时无需对框架自身做任何理解。其惟一须要做的就是基于这套模型去设置一些 Op 的根本信息，比方 Op 有几个输入输出，调用哪个 kernel 等，从而建设起对于本人 Op 的形容，其画风个别如下代码所示。

CUSTOM_OP_REG(MatMulScale)              // 定义一个名为 MatMulScale 的 Op
     .add_inputs(2)                  // 两个输出 Tensor
     .add_outputs(1)                  // 一个输入 Tensor
     .add_param("scale", 1.0f)           // 一个名为 scale 的 Parameter，默认值为 1.0f
     .set_shape_infer(shape_infer)         // 设置这个 Op 的 shape 推导函数
     .set_compute("cuda", compute);        // 设置这个 Op 的 计算函数 

而这些设置过程个别应用几行至十几行的代码就能够表白，大大简化了用户集成 kernel 时的工作量。更多的对于 Custom Op 的应用能够参考 MegEngine Custom Op 应用阐明。

而后 Custom Op 会主动的将用户的 Op 封装成 MegEngine 中动态图与动态图中的 OpNode 与 OpDef，并为之生成和原生算子统一的 python 接口，从而使用户的 Custom Op 能够与零碎中的原生 Op 在接口和底层行为上放弃对立。

Custom Op 的设计

Custom Op 同时面向用户（即上述须要编写 kernel 的算法同学）与 MegEngine 零碎，使两者能够简略便捷的进行交互。而为了达到这个目标，Custom Op 须要具备以下个性：

• 面向用户，Custom Op 须要提供一套简洁对立的，与框架无关的，且编写 Op 所必须应用的抽象概念。用户能够应用这些形象接口去将本人的 c++/cuda kernel 封装成 Op。
• 面向零碎，Custom Op 须要给 MegEngine 提供一套齐备的 Op 的适配与管理工具，从而容许系统对 Custom Op 进行调用及保护治理。

基于此，咱们设计出了 Custom Op 的整体架构，如图 2 所示。

图 2：Custom Op 的整体构造

在个别算法用户的认知里，Op 就是个计算函数，承受一些输出，实现计算，而后失去一些输入。而这里的输入输出又分为张量型数据（即个别的输入输出 Tensor）和非张量型数据（即 Param，比方卷积中的 padding，stride 等等），如图 3 所示。

图 3：用户视角的 Op

为了符合算法用户对于 Op 的认知，Custom Op 次要向用户提供了三个形象 Tensor，Param 以及 Op。

• Tensor 是 Op kernel 次要计算和操作的对象，与 MegEngine python 中的 Tensor 有着基本一致的形象和行为。
• Param 是用于记录传输一些 Op 的非张量的输出（比方卷积中的 padding，stride 等等）。
• Op 是对用户的 c++/cuda kernel 计算函数的一个封装，同时记录着这个 Op 的输入输出 Tensor 以及 Param 的信息。

而面向 MegEngine 零碎，Custom Op 一方面给 MegEngine 提供了一套齐备的 Adaptors，能够依据 MegEngine 零碎中不同档次的须要，将用户的 Op 和 Tensor 适配成 MegEngine 的 runtime 能够解决的 Op 和 Tensor。另外一方面，Custom Op 同时还向 MegEngine 提供一套用户 Op 的 Managers，从而容许 MegEngine 对 Custom Op 进行保护治理。

上面咱们将别离介绍 Custom Op 的这些模块。

Tensor

在算法用户的视角中，Tensor 是一个多维数组，同时其还有着一些如形态，量化信息等属性，故而 Custom Op 中的 Tensor 也被设计为数据以及数据的相干属性的汇合。其中数据由一个指向数据存储空间的指针治理。而这些属性则通知咱们该如何去解析数据空间中的数据。如图 2 中 Tensor 局部所示，这些属性次要蕴含 Device，DType，Shape 这三者：

• Shape 代表的是 Tensor 维度信息。
• DType 对应 Tensor 中元素的数据类型，如 float32，uint8 等。
• Device 则示意这个 Tensor 在什么设施（cpu/gpu）上。

通过这些属性能够建设起对 Tensor 的齐备的形容。

事实上 Tensor 及其从属的这些属性在 MegEngine 零碎都会有另一套功能丰富，但对用户而言略显冗余的表白。为了升高应用难度，Custom Op 简化了这些概念，只留下编写 Op 时须要的性能。

具体来说，Shape 提供给用户的行为相似于 c++ 的原生数组，咱们能够应用如下的代码来构建和应用它：

Shape shape = {16, 3, 224, 224};    // 构建 shape
bool equal = (shape[3] == 224);     // 获取 shape 特定维度的值
shape = {128, 100};                 // 批改 shape 的值 

至于 Device 以及 DType，用户并不需要通晓其背地实现，只须要通过这些属性晓得数据在什么设施上，是什么类型就够了。故而 Custom Op 中的 Device 和 DType 的行为均相似于 string，用户能够以间接设置字符串值的模式去设置具体的 Device 和 DType 类型。

Device device = "x86";                 // 创立一个 x86 这种设施类型
device = "cuda";                     // 设施类型改为 cuda
bool equal = (device == "cuda");            // 判断某个 device 是否是 cuda

DType dtype = "float32";                       // 构建 dtype
bool equal = (dtype == "int8");           // 判断 dtype 是否相等 

而为了使这些类型的接口与 MegEngine 解耦，其实现时均应用了 pimpl 手法，暗藏了这些类型的内存布局，而用户通过一系列的接口去 set/get 其中的数据。

Param

Param 次要用于表白 Op 的非 Tensor 的输出（比方卷积中的 padding，stride 等等），然而不同的 Op 的这些非 Tensor 输出的差异往往十分大。可能 Op A 的 Param 是一系列的 string 而 Op B 的 Param 是一个 int 类型。所以咱们须要设计一套机制以将这些彼此差别很大的 Param 对立起来，供用户和 MegEngine 零碎进行应用。为了实现这个目标，Custom Op 中设计了 ParamVal。ParamVal 中会擦除各个 Op Param 的动态类型，使这些 Param 动态类型对立，以解决不同 Op 的 Param 类型不统一的问题，而后另外定义一套运行时的动静类型零碎去进行 Param 理论类型的治理。

说起来可能比较复杂，实际上其能够简化成上面的这个数据结构。其中应用 void* 进行类型擦除，并将擦除后的数据放在 data 中进行存储，而 type 中则记录着这个数据所对应的动静类型。

class ParamVal {
    void *data;             // 类型擦除后的数据
   DynamicDataType type;   // 数据的动静类型
};

这样的设计不仅能解决不同 Param 类型不对立的问题，同时这个动静类型的存在同时也大大缓解了 c++ 中没有反射所带来的一些不便。Custom Op 依据此动静类型零碎设计了一套对立的参数解析（Param Parse）和序列化机制，而无需用户为本人的 Param 去编写这部分代码。

同时为了进一步不便用户去应用，Custom Op 中还为 ParamVal 定义了十分多的运算符，以及其与动态类型之间互相转换的函数。最终展示在用户视角，ParamVal 的行为与 python 中的变量是很靠近的。

ParamVal a = 1.0, b = 2, c = false, d = {1, 2}; // 表白不同类型的数据
ParamVal e = a + b;                             // ParamVal 彼此间的计算
ParamVal f = e - 2;                             // ParamVal 与动态类型数据的计算
d = c; 

Op

在个别算法用户的视角里，Op 是对一个计算过程的形容而不记录保留任何数据信息。为了与算法用户的认知相对立，Custom Op 中的 Op，也被设计为无状态的，即 Op 中只保留相干的函数以及输入输出的一些布局信息，而不记录输入输出的具体值。

图 4：Op 及其组件

具体来说，在 Custom Op 中，Op 是输入输出 Tensor 信息（TensorInfo），Param 信息（ParamInfo），以及 Op 相干函数的汇合（Functions），如图 4 所示。TensorInfo 中记录了这个 Op 的输入输出 Tensor 的数量，名字，非法类型，维度以及内存调配策略等信息。ParamInfo 则记录着这个 Op 的各个 Param 的名字，默认类型以及默认值等。至于 Functions 其理论蕴含两局部，kernel 计算函数和 Tensor 属性推导函数。

• kernel 计算函数次要负责将 Tensor，Param 的值传递给用户的 c++/cuda kernel，并将计算结果返回。
• Tensor 属性推导函数则是依据输出 Tensor/Param 的一些属性在 kernel 执行前实现输入 Tensor 的数据布局的推导，从而将算子执行和算子内存调配解耦，以进行内存布局。

其中大部分函数 Custom Op 均提供了默认的实现，用户能够通过 override 这些函数的默认行为去定制化本人的 Op。

Manager 与 Adaptor

思考到 Custom Op 是一套与 MegEngine 零碎解耦的 Op 形象，依据 Custom Op 定义进去的 Op，MegEngine 并不能间接与之进行交互。为了解决这个问题，Custom Op 中额定设计了一组 Adaptors 和 Managers。其中 Adaptor 容许 MegEngine 应用 Custom Op，而 Manager 容许 MegEngine 感知和治理 Custom Op。

Adaptor 的设计目标次要蕴含两个方面：一方面其须要容许 MegEngine 去操作应用 Custom Op 去构建网络等，另一方面其须要容许 MegEngine 与 Custom Op 进行数据交互实现计算。

• 对于前者，Adaptors 能够将 Custom Op 包装成 MegEngine 中动态图与动态图中的 Op 形象，使之行为与 MegEngine 中内置 Op 保持一致。
• 对于后者，Adaptors 能够将 Custom Op 中面向用户的 Tensor 形象与 MegEngine 中的 Tensor 联合起来，使两者间能够相互转换，容许数据能够自在的在 MegEngine 与 Custom Op 之间流动。

而对于 Manager 模块，其提供了对 Custom Op 所编译成的动态链接库以及 Custom Op 自身的治理。具体来说，Custom Op 在应用时会以动态链接库的模式被加载入 MegEngine 零碎中，因而咱们基于 RAII 机制去治理这些动态链接库。动静库的加载与卸载和 Lib 类的结构与析构绑定起来，从而防止资源透露。而对于 Op 自身，Manager 提供了一些根本的增删改查的操作去容许 MegEngine 对 Custom Op 进行治理。

Custom Op 的编写

用户在应用 Custom Op 编写 Op 时，用户能够应用上述这些概念去将本人写好的 kernel 封装成 Custom Op，并应用 Custom Op 提供的构建工具将之编译成动态链接库，在运行时将之加载入 MegEngine 进行应用。

具体来说，如果当初咱们须要为 MegEngine 增加一个名为 MatmulScale 的算子，这个算子在计算时首先会对两个输出 Tensor，lhs 和 rhs 执行矩阵乘，而后再将这个矩阵乘的后果再乘以标量 Scale。

该算子数学上的执行过程的伪代码如下：

def MatMulScale(lhs, rhs, scale):
    result = lhs.dot(rhs)
    result = result * scale
    return result

对于这样的一个操作，假如咱们曾经为之写好了一份 cuda kernel 代码，并提供如下的接口函数用于调用：

void matmul_scale(const float *lhs, const float *rhs, float *result, size_t M, size_t K, size_t N, float scale);

这些的参数中，lhs，rhs，以及 result 是三个 float 类型的指针，别离代表这个 Op 的两个输出 Tensor 和一个输入 Tensor，其均须要指向一片曾经调配好的 cuda memory。而 M，K，N 是矩阵的维度信息，示意一个 M*K 的矩阵乘以一个 K*N 的矩阵。而 scale 则代表着矩阵乘的后果须要乘以的那个系数。

对于这种状况咱们能够编写如下的 c++ 代码，就能够将之封装成 MegEngine 的 Op。

void shape_infer(const std::vector<Shape> &inputs, const Param &params, std::vector<Shape> &outputs) {outputs[0] = {inputs[0][0], inputs[1][1]};
}

void compute(const std::vector<Tensor> &inputs, const Param &params, std::vector<Tensor> &outputs) {matmul_scale(inputs[0].data<float>(), inputs[1].data<float>(), outputs[0].data<float>(), ...); 
}

 CUSTOM_OP_REG(MatMulScale)              // 定义一个名为 MatMulScale 的 Op
     .add_inputs(2)                  // 两个输出 Tensor
     .add_outputs(1)                  // 一个输入 Tensor
     .add_param("scale", 1.0f)           // 一个名为 scale 的 Parameter，默认值为 1.0f
     .set_shape_infer(shape_infer)         // 设置这个 Op 的 shape 推导函数
     .set_compute("cuda", compute);        // 设置这个 Op 的 计算函数 

这段代码次要蕴含两个局部，第一个局部是一些函数的定义，包含输入 Tensor 属性推断函数和计算函数。其中前者会依据输出 Tensor 的属性（比方 shape）去推导输入 Tensor 的对应属性，而后者则是在其中调用 cuda kernel，实现计算。第二局部是 Op 的注册，次要用于定义 Op 有几个输入输出 Tensor，有几个 Param，并将下面定义的属性推断函数和计算函数的指针也注册给 Op。到此就实现了 Custom Op 的构建。

Custom Op Generator

通过 Custom Op 能够将用户编写好的 c++/cuda kernel 简略不便的集成入 MegEngine。然而，让用户本人去编写 kernel 总是最初的抉择，如何可能将用户编写 kernel 的这一步的工作也给省掉呢？MegEngine 正在尝试基于 AI 编译器提出一个工具 Custom Op Generator 去解决这个问题。

在 Custom Op Generator 中用户能够间接应用 AI 编译器提供的简略的 python 原语去建设其 Op 的表白，而不须要写任何 c++/cuda 的代码。而后 AI 编译器会主动生成 Op 所对应的 kernel 以及 Custom Op 的装璜代码将这个 kernel 封装成 MegEngine 的 Custom Op，并自动化的进行构建并将之集成入 MegEngine 中。全过程用户只需编写一些 python 代码即可，防止了用户本人编写 kernel 的问题。

一般而言，框架与编译器联合的 workflow 都是框架在前编译器在后，由框架去构建一个模型用于训练，而后将训练好的模型送给编译器进行优化部署。然而在 Custom Op Generator 中两者的地位恰恰相反，而在这样的一个 workflow 中，编译器在前而框架在后，编译器利用其代码生成的能力为框架提供拓展反对。从某种意义上来说，这是一种 AI 编译器与框架联合的一种新思路。

总结

Custom Op 作为沟通用户 kernel 和零碎 Op 的桥梁，其面向用户提供了一套简洁对立且与框架无关的 Op 形象，面向零碎提供一套齐备的 Op 适配与管理工具。为了这个目标，在设计实现 Custom Op 的过程中，咱们剖析用户编写 Op 时所必须的概念并进行形象，设计出了一套框架无关的 Op 模型并提供简洁的接口给用户，从而实现了用户侧与零碎侧的解耦，使用户编写的 Op 无需随着零碎的更新迭代而做对应变动。而同时为了容许 MegEngine 零碎去治理应用这些 Op，Custom Op 中又设计了相干治理与适配模块，其会自动化的将用户的 Op 封装适配成 MegEngine 动态图与动态图中的 Op，使用户 Op 在零碎中行为与原生 Op 保持一致，从而便于零碎的应用治理。通过这样的设计，用户在集成 Op 时无需理解 MegEngine 框架自身，只需二十行代码即可疾速的将 kernel 集成入 MegEngine 并应用，能够大大降低算法用户集成 kernel 时的难度与工作量。

原文地址：https://megengine.org.cn/blog…