吹拉弹唱


  • Home
  • Archive
  • Categories
  • Tags
  • Books
  •  

© 2022 Kleon

Theme Typography by Makito

Proudly published with Hexo

机器学习系统 1-9 - 后端硬件编译器

Posted at 2021-01-26Updated at 2021-01-26 机器学习  机器学习 FPGA 

编译,广义上是将一种表达映射成另一种表达,维持功能不变。深度学习后端硬件的编译,是将DSL描述的计算图映射成后端硬件指令。

传统编译器直接解析文本,通过词法与语法分析,构建AST1。之后在AST上执行一系列的优化,称为优化pipeline。

优化后的AST可以转换为IR(比如LLVM2 IR),由LLVM进一步优化,并生成后端代码。如果是添加新后端并且复用LLVM,需要实现对应Codegen逻辑。AST也可以不借助LLVM,直接编写codegen逻辑生成不同层级的code或指令。

随着深度学习模型越来越复杂,在生产中广泛使用,对计算性能和成本的需求日趋强烈,深度学习编译器在这样的背景下逐步发展壮大。

在深度学习编译器提出前,广泛使用的优化方法是手工fuse和手工codegen。但这样的劣势显而易见,模型中可能存在的子图pattern千奇百怪,并且各种自定义长尾算子层出不求,如果一直依靠手工优化,开发时间和人力成本都不太能令人接受。

因此,借鉴传统编译器的思路,通过有限的基础表达符号表示任意计算图结构,由编译器的codegen层统一做后端代码生成。

# 分类

编译器根据编译发生时间的不同可以分为AOT3(Ahead Of Time)和JIT4(Just In Time)。

AOT编译通常用于standalone编译程序,由用户提前运行编译生成指令,并传输到运行设备上(非标设备/专用设备,比如ARM和定制NPU),通过runtime执行。

一次编译,到处运行,这允许AOT编译消耗更多的时间,探索更优结果,但也限制了灵活性。每当模型发生改动就需要触发一次编译行为,通常编译设备和运行设备是分立的,指令文件和runtime版本的同步会更复杂。

JIT在运行时条件触发编译,runtime和compiler捆绑,不存在版本不匹配的问题。JIT通常见于解释型程序,因为解释效率通常低于编译,通过运行前编译提高执行效率。

由于在运行时编译,编译时间相对运行时间必须要合理,这意味着不能引入过于复杂的优化逻辑,另外为了避免重复编译,可以加入cache(程序内/远程)提高效率。

# 前端

编译器前端需要读取各个框架的模型文件,比如Tensorflow的SavedModel或者PyTorch的TorchScript。

也有试图将模型表达统一的工作出现,比如ONNX,试图在Caffe、PyTorch、MXNet等框架间构建转换桥梁。愿景很好,方便用户在不同框架迁移,不被框架锁死,但最终没有被大规模使用。

首先,Tensorflow出自Google,并不买账,ONNX是其他大公司的联盟产物,然而框架上只有PyTorch不断发展,和Tensorflow二分天下。

统一模型表达,还是在统一社区,也就是统一算子集。

Tensorflow曾凭借其绝对的垄断地位十分强势,后端软件栈都主动接入,接受Tensorflow的核心算子集为golden标准。然而不同框架的算子集相差很多,基本上只有最主流的算子可以无缝转换,换个复杂的模型很有可能就挂了。尤其是Tensorflow的计算图,控制结构十分复杂,连官方的优化工具都可能出现bug。如果模型定制程度高,不如直接改脚本来的快。

当然对大多数算法工程师来说,没有被框架锁定这件事,只有看哪个框架的生态的轮子能符合业务需求。

# 优化

编译优化方法泛指一大类模型优化方法,比如Tensorflow自带的后端XLA,基于profiling搜索的TVM,基于整数线性规划的多面体方法等。

Tensorflow自带的前端计算图优化Grappler也可以算作编译器的优化功能。在图层面执行一般优化方法,比如CSE(公共子表达式消除),Constant Folding(常量折叠),去除无用节点(Dead Nodes)等。

对于硬件后端编译器,主要优化步骤如下:

  1. 硬件无关优化。
  2. 尽可能从子图中匹配更多符合硬件结构的pattern。
  3. 匹配可以转换为符合硬件结构pattern的子图,并实现转换逻辑(硬件相关优化),注意功能正确性。
  4. 尽量将匹配到的子图融合,使之最大化连通。
  5. 设定不连通子图最小阈值,防止主存和硬件内存的context switch开销过高,可以通过算子数量或者算力评估确定阈值,标记所有符合阈值的子图算子。
  6. 融合所有标记的子图,生成对应的模型IR或者硬件IR。
  7. 可以通过自定义Op的形式,将生成的IR序列化(作为Op参数),使用统一的Op(调用runtime)。
  8. 如果模型可以全量支持,则可以直接导出IR,使用runtime运行。

进一步的,编译器可以独立为TensorRT的形式,允许用户使用API直接构建运行,并且可以通过JIT codegen dispatch手工优化的code。

编译器相比于计算库更关注图层面优化,底层会协调计算库的不同实现使全局性能最优。


  • [1] Abstract syntax tree
  • [2] The LLVM Compiler Infrastructure
  • [3] Ahead-of-time compilation
  • [4] Just-in-time compilation

Share 

 Previous post: 机器学习系统 1-10 - 后端硬件框架接入 Next post: 机器学习系统 1-8 - 后端硬件运行时 

© 2022 Kleon

Theme Typography by Makito

Proudly published with Hexo