关于深度学习:如何写一个深度学习编译器

编译器实质上是一种进步开发效率的工具，将高级语言转换为低级语言（通常是二进制机器码），使得程序员不须要徒手写二进制。转换过程中，首要任务是保障正确性，同时须要进行优化以晋升程序的运行效率。传统意义上的编译器的输出通常是某种高级语言，输入是可执行程序。在理论工作中接触到了深度学习编译器开发，其设计思维与传统编译器十分相似，所以本文以深度学习编译器的开发、联合咱们理论开发的深度学习编译器 MegCC 为例，来阐明如何写一个编译器。本文次要分为以下两个局部：

1. 介绍深度学习编译器，重点介绍编译器中前端和后端的优化办法。
2. 以 MegCC 为例介绍如何开发一个深度学习编译器。

深度学习编译器简介

与传统编译器不同，深度学习编译器的输出是神经网络模型、输入是可运行在不同平台的表白了输出的神经网络模型的计算过程的可执行程序。但深度学习编译器又与传统编译器相似，都分为前端和后端，前端负责执行硬件无关的优化，后端负责执行硬件相干的优化。对编译器来说，最重要的两个概念是 IR(intermediate representation, 两头示意)和 Pass。对于人类来说，形象是了解复杂事物的一种重要形式，IR 就是对编译过程两头产物的形象，IR 通常有多级，越高级的 IR 越形象，越低级的 IR 越具体。Pass 定义了如何将高级 IR 逐渐 lowering 到低级 IR，并负责进行优化。上面依据前端和后端进行分类，介绍优化的办法。

前端优化办法

前端首先须要依据输出的模型构建计算图，生成 high-level IR，而后进行一系列的优化。因为优化是基于计算图的，并不波及具体计算，所以该优化是后端无关的。常见的优化伎俩有可分为三类：node-level optimizations；block-level optimizations；dataflow-level optimizations。

1. node-level optimizations。节点层面的优化次要是打消一些不必要的节点以及将某些节点替换为代价更小的节点。比方应用矩阵 A 与一个 0 维矩阵相加，则可打消该加法操作。
2. block-level optimizations。块层面的优化次要有代数简化和算子交融。a. 代数简化，例如 A^T 和 B^T 进行矩阵乘，则可应用 B 与 A 矩阵乘之后进行转置进行替换，可节约一次转置运算。b. 算子交融是常见的深度学习的优化伎俩。算子交融尽管不能缩小计算量，然而能够缩小访存量，进步计算访存比，从而晋升性能。
3. dataflow-level optimizations。数据流层面的优化次要有动态内存布局等。a. 动态内存布局通过在不产生内存重叠的前提下尽可能复用内存，使得程序运行时所应用的内存尽可能小。

后端优化办法

后端通用的优化有循环展开、循环交融、覆盖访存等；另外依据硬件的不同，可应用基于硬件的指令映射、向量化等并行计算以及手工编写汇编 kernel 等伎俩进行针对性优化。图 1 展现了罕用的后端优化办法 ¹

<p align=center> 图 1 后端罕用优化办法 </p>

MegCC

接下来就以 MegCC 为例概括介绍一下基于 MLIR 实现一个深度学习编译器，其要害就是如何依据需要定义一系列 IR，以及定义 Pass 将高级 IR lowering 到低级 IR，同时进行上述优化。

MegCC 简介

MegCC 实现的原理是：深度学习模型在推理时候，每一个 Operator 都会对应一个计算 kernel 并实现计算，所以整个深度学习模型在推理时就是一次执行所有 Operator 的计算 kernel，执行实现之后就能够取得最终推理的后果。传统深度学习推理框架在运行时会做以下几件事件：

• 计算图优化 —– 次要和模型相干。
• Kernel 抉择 —– 为模型的每个 Operator 依据参数抉择适合的 Kernel 进行计算。
• 内存调配 —– 由模型以及模型中每个 Operator 执行的 Kernel 决定内存调配的大小。
• 执行每个 Operator 的 Kernel —– 和推理的数据强相干。

在上述传统深度学习推理须要实现的事件中，图优化，Kernel 抉择，内存调配都是只和训练好的模型相干和推理时候的输出数据不相干，因而这些工作都能够放在模型编译时实现，运行时仅仅执行每一个 Operator 的 Kernel 就能够实现推理。MegCC 就是将下面图优化，Kernel 抉择，内存调配都放在 MegCC 的编译阶段实现，将 Operator 的 Kernel 计算才放到 Runtime 中进行计算，这样有以下劣势：

• Runtime 十分轻量，比起传统的推理框架小一个数量级，因为 Runtime 只蕴含了模型中所必须的 Kernel，不相干的不会被编译进去。
• 晋升性能，因为 Runtime 只做 kernel 计算，所以防止了不必要的开销。
• Kernel 性能优化，因为每一个 Kernel 都是针对每一个 Operator 定制的，因而能够依据 Operator 的参数进行更加深刻的优化。
• 解决 Operator fuse 之后的算子长尾问题，比方对 conv 之后交融的 activation 的品种和数量没有限度，能够反对更多的 fuse，也不造成 Runtime 的大小有显著的扭转。
• 另外 MegCC 的 runtime 应用纯 C 实现，能够轻松移植到其余的嵌入式芯片中。

MegCC 次要蕴含两局部，一部分是 compiler 局部，另外一部分是 runtime 局部，上面重点介绍与编译相干的 compiler 局部。

MegCC compiler

Compiler 次要流程是：

1. 依赖 MegEngine (我司开源深度学习框架)进行模型的导入和动态图优化（block-level optimizations，算子交融等）。
2. 将优化后的模型转换为基于 mlir 自定义的 MGB IR。
3. MGB IR 通过一系列 pass 通过 Abstract Kernel IR 最终转换到 Kernel IR。
4. 将 Kernel IR 导出为 runtime model 和 runtime kernel，供 MegCC 的 runtime 局部应用。

​

<p align=center> 图 2 MegCC compiler 流程 </p>

MegCC 中的 IR

MegCC 基于 MLIR 定义了一系列的 IR。MLIR 的 IR 定义须要用户定义 Dialect（详见官网文档），而后由 TableGen 在程序编译阶段转换成 C++ 示意。

• MGB IR：定义为和 MegEngine 中 Operator 一一对应，是 MegCC 导入进 mlir 零碎的入口 IR，它蕴含了每个 Opr 的类型以及这个 Opr 对应的参数，其每一个输入输出变量都是 Tensor，并且是单赋值 (SSA) 的。详见 GitHub MegCC MGB IR。
• Abstract Kernel IR：形象 Kernel 层 IR，次要下面 MGB IR 通过转换之后失去，该 IR 中的输入输出曾经 lowering 到 Buffer 了，因而不再是 SSA，另外 Opr 的属性也由 MegEngine 中定义的枚举值，转变成为了字符串。详见 GitHub MegCC Abstract Kernel IR。
• Kernel IR：示意曾经生成 Kernel 之后的 IR 模式，其曾经没有 Opr 的概念，整个计算图通过一个个对应的 Kernel 链接在一起，Opr 的参数等都固化在了定义好的 Kernel 中。详见 GitHub MegCC Kernel IR。

MegCC 中次要的 Pass

• MGBToKernelPass：这个 Pass 次要将 MGB IR 转换为 Abstract Kernel IR，转换过程中次要实现几件事件：
• 将 MGB IR 中的所有输入输出 Tensor 类型转换为 Buffer 类型。
• 将 MGB IR 中的所有枚举参数转换为对应的字符，这样 Abstract Kernel IR 就能够齐全和 MegEngine 解耦。
• 将一些内存搬运相干的 Opr 全副转换为 Relayout，如：Concat，SetSubtensor 等 Opr（node-level optimizations）。
• 将判断 Opr 是动态 shape 还是动静 shape，动静 shape 就是输出 tensor 的 shape 须要依赖输出的值能力计算出来的，如：输入一个 tensor 中所有大于 1 的数。如果是动态 shape 间接转换到 Abstract Kernel IR，如果是动静 shape 间接转换到 Kernel IR 的 Instruction 中。
• MGBFuseKernelPass：利用在 MGB IR 上，基于 mlir 的模板匹配的办法尽可能的实现 kernel 的交融，比方间断两个 typecvt 合并成为一个 typecvt 等（block-level optimizations，算子交融）。
• MemoryForwardingPass：将遍历 Abstract Kernel IR 所有可能不必计算，间接 share 输出内存的 Opr，如果这些 Opr 的确不必计算，则间接 forward memory，如果这些 Opr 须要进行内存搬运，则会用 Relayout Opr 替换原来的 Opr（node-level optimizations）。
• KernelMaterializationPass：将所有 Abstract Kernel IR 都装载上真正 Kernel code 并转化为 KernelCall，而后增加对应的 KernelDef。KernelCall 和 KernelDef 之间通过 symbol 进行匹配。
• StaticMemoryPlanningPass：将所有动态 shape 的 memref 进行内存布局，内存布局算法应用改良的 MegEngine 的内存布局算法 –PushDown 算法，可能极大水平的压缩运行时内存使用量。同时将 mlir 的 memref.Alloc 替换为 Kernel IR 的 MemPlan，MemPlan 中次要记录了内存布局的一整块 memref 以及该 Tensor 在布局的内存中的偏移量（dataflow-level optimizations，动态内存布局）。

下面的 Pass 就实现模型的图优化、内存布局以及 Kernel 生成，上文提到的后端优化即在 Kernel 生成阶段体现，目前 MegCC 次要应用人工优化的 Kernel 模版。最终能够依据 Runtime 中定义的模型格局 dump 编译之后的模型，以及生成计算模型所需的 Kernel 文件。上面以一个简略的模型为例，应用 MegCC 的辅助工具(下载 Release 包) mgb-importer 和 megcc-opt，察看通过各个 Pass 的解决 IR 的变动。也可应用 mgb-to-tinynn 工具间接实现模型的编译过程，详见 MegCC 入门文档。

1. dump 模型（应用 megengine）

import megengine.functional as F
import megengine.module as M
import megengine.optimizer as optim
from megengine import jit
import megengine
import numpy as np

class MulAddNet(M.Module):
def __init__(self):
    super().__init__()

def forward(self, input):
    x = input * 2.
    x = x + 1.5

    return x

model = MulAddNet()
model.eval()

@jit.trace(symbolic=True, capture_as_const=True)
def infer_func(data, *, model):
    pred = model(data)
    return pred

data = megengine.Tensor([[1., 2.], [3., 4.]])
output = infer_func(data, model=model)
print(output)

infer_func.dump("MulAdd.mge", arg_names=["data"])
```

2. importer 模型到 MGB IR

./bin/mgb-importer MulAdd.mge mulAdd.mlir
cat mulAdd.mlir
output:

module {"MGB.ParamStorage"() {sym_name = "const<2>[2]", sym_visibility = "private", type = tensor<1xf32>, user_count = 1 : i32, value = dense<2.000000e+00> : tensor<1xf32>} : () -> ()
  "MGB.ParamStorage"() {sym_name = "const<1.5>[4]", sym_visibility = "private", type = tensor<1xf32>, user_count = 1 : i32, value = dense<1.500000e+00> : tensor<1xf32>} : () -> ()
  func @mulAdd(%arg0: tensor<2x2xf32> {mgb.func_arg_name = "data"}) -> (tensor<2x2xf32> {mgb.func_result_name = "FUSE_MUL_ADD3(const<2>[2],data,const<1.5>[4])[14]"}) {%0 = "MGB.Reshape"(%arg0) {axis = 7 : i32} : (tensor<2x2xf32>) -> tensor<2x2xf32>
    %1 = "MGB.ParamProvider"() {name = @"const<1.5>[4]"} : () -> tensor<1xf32>
    %2 = "MGB.ParamProvider"() {name = @"const<2>[2]"} : () -> tensor<1xf32>
    %3 = "MGB.Elemwise"(%2, %0, %1) {mode = 35 : i32} : (tensor<1xf32>, tensor<2x2xf32>, tensor<1xf32>) -> tensor<2x2xf32>
    return %3 : tensor<2x2xf32>
  }
}

能够看到，在 importer 的过程中，乘法运算和加法运算被交融成了 ”FUSE_MUL_ADD3″。

3. MGBToKernelPass、MemoryForwardingPass 和 StaticMemoryPlanningPass

./bin/megcc-opt --MGB-to-Kernel --memory-forwarding --static-memory-planning mulAdd.mlir > mulAdd_final.mlir
cat mulAdd_final.mlir
output:

#map = affine_map<(d0, d1) -> (d0 * 2 + d1)>
module {"Kernel.WeightStorage"() {sym_name = "const<2>[2]", type = tensor<1xf32>, user_count = 1 : i32, value = dense<2.000000e+00> : tensor<1xf32>} : () -> ()
  "Kernel.WeightStorage"() {sym_name = "const<1.5>[4]", type = tensor<1xf32>, user_count = 1 : i32, value = dense<1.500000e+00> : tensor<1xf32>} : () -> ()
  func @mulAdd(%arg0: memref<2x2xf32> {mgb.func_arg_name = "data"}, %arg1: memref<16xi8> {mgb.func_arg_name = "kGlobalBuffer"}) -> (memref<2x2xf32, #map> {mgb.func_result_name = "FUSE_MUL_ADD3(const<2>[2],data,const<1.5>[4])[14]"}) {%0 = "Kernel.Reshape"(%arg0) {axis = 7 : i32, determined = true} : (memref<2x2xf32>) -> memref<2x2xf32, #map>
    %1 = "Kernel.GetWeight"() {name = @"const<1.5>[4]"} : () -> memref<1xf32>
    %2 = "Kernel.GetWeight"() {name = @"const<2>[2]"} : () -> memref<1xf32>
    %3 = "Kernel.MemPlan"(%arg1) : (memref<16xi8>) -> memref<2x2xf32, #map>
    "Kernel.FUSE_MUL_ADD3"(%2, %0, %1, %3) : (memref<1xf32>, memref<2x2xf32, #map>, memref<1xf32>, memref<2x2xf32, #map>) -> ()
    return %3 : memref<2x2xf32, #map>
  }
}

通过下面几个 Pass，MGB IR 被转换为了 Kernel IR 并进行了内存布局。感兴趣的话能够更细粒度地看每个 Pass 做的事件，应用 megcc-opt 的参数管制应用哪些 Pass。

Kernel 生成

MegCC Compiler 会为模型中的每个 Operator 生成一个对应的 Kernel 来实现计算。目前 MegCC 中大多数 Kernel 为人工优化并提前写好的 Kernel 模板，这些模板会依据具体的 Operator 参数生成对应的 Kernel。大多数为人工优化的 Kernel 的起因是：目前在 CPU 上不搜参的状况下，mlir 生成的 Kernel 性能和手写的 Kernel 还有肯定的间隔，然而主动生成 Kernel 的办法长期来看是比拟可取的。

MegCC 现已开源，仓库地址：https://github.com/MegEngine/…，欢送试用、star、issue。

附：

更多 MegEngine 信息获取，您能够：查看文档、和 GitHub 我的项目，或退出 MegEngine 用户交换 QQ 群：1029741705。欢送参加 MegEngine 社区奉献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。

参考文献：