如何写一个深度学习编译器-六虎

作者：旷视 MegEngine 结构开发师 – 李明鑫

编译器本质上是一种进步开发功率的工具，将高档言语转化为初级言语（通常是二进制机器码），使得程序员不需求徒手写二进制。转化进程中，首要任务是保证正确性，一起需求进行优化以提高程序的运行功率。传统意义上的编译器的输入通常是某种高档言语，输出是可履行程序。在实践作业中触摸到了深度学习编译器开发，其设计思想与传统编译器十分类似，所以本文以深度学习编译器的开发、结合咱们实践开发的深度学习编译器 MegCC 为例，来阐明如何写一个编译器。本文首要分为以下两个部分：

介绍深度学习编译器，要点介绍编译器中前端和后端的优化办法。
以 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。

node-level optimizations。节点层面的优化首要是消除一些不必要的节点以及将某些节点替换为价值更小的节点。比如运用矩阵 A 与一个 0 维矩阵相加，则可消除该加法操作。
block-level optimizations。块层面的优化首要有代数简化和算子交融。 a. 代数简化，例如 A^T 和 B^T 进行矩阵乘，则可运用 B 与 A 矩阵乘之后进行转置进行替换，可节约一次转置运算。 b. 算子交融是常见的深度学习的优化手法。算子交融虽然不能削减核算量，但是可以削减访存量，进步核算访存比，然后提高功能。
dataflow-level optimizations。数据流层面的优化首要有静态内存规划等。 a. 静态内存规划通过在不发生内存堆叠的前提下尽或许复用内存，使得程序运行时所运用的内存尽或许小。

后端优化办法

后端通用的优化有循环展开、循环交融、掩盖访存等；别的依据硬件的不同，可运用基于硬件的指令映射、向量化等并行核算以及手工编写汇编 kernel 等手法进行针对性优化。图 1 展示了常用的后端优化办法 ¹

图 1 后端常用优化办法

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 首要流程是：

依靠 MegEngine (我司开源深度学习结构)进行模型的导入和静态图优化（block-level optimizations，算子交融等）。
将优化后的模型转化为基于 mlir 自界说的 MGB IR。
MGB IR 通过一系列 pass 通过 Abstract Kernel IR 终究转化到 Kernel IR。
将 Kernel IR 导出为 runtime model 和 runtime kernel，供 MegCC 的 runtime 部分运用。

图 2 MegCC compiler 流程

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 入门文档。

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"])
```

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″。

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 现已开源，仓库地址：github.com/MegEngine/M…，欢迎试用、star、issue。

附：

更多 MegEngine 信息获取，您可以：检查文档、和 GitHub 项目，或加入 MegEngine 用户沟通 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。

参考文献：

The Deep Learning Compiler: A Comprehensive Survey. MINGZHEN LI, YI LIU, etc. 2020. ↩

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

深度学习编译器简介

前端优化办法

后端优化办法

MegCC

MegCC简介

MegCC compiler

MegCC 中的 IR

MegCC 中首要的 Pass

Kernel 生成

附：

参考文献：

Footnotes

相关文章

用 Compose 写 App 可以多快？

肘子的 Swift 周报 #012 ｜ 用新博客告别 2023 年

近邻搜索算法浅析

【爆肝强烈推荐】 十多款2023年必备国内外王炸级AI工具 (免费 精品 好用) 让你秒变装逼佬,感受人工智能带来的生产力!

作者信息

肘子的 Swift 周报 #012 ｜用新博客告别 2023 年

【爆肝强烈推荐】十多款2023年必备国内外王炸级AI工具 (免费精品好用) 让你秒变装逼佬,感受人工智能带来的生产力!