文章目录
- 深度学习IR梳理
- 1. IR属性
- 2. DL前端发展
- 3. DL编译器
- 4. DL编程语言
- Relay的主要内容
- 一、Expression in Relay
- 1. Dataflow and Control Fragments
- 2. 变量
- 3. 函数
- 3.1 闭包
- 3.2 多态和类型关系
- 3.3. Call
- 4. 算子
- 5. ADT Constructors
- 6. Moudle和Global Function
- 7. 常量和元组
- 8. Let Binding
- 9. Graph Bindings
- 10. If-Then-Else
- 11. ADT Matching
- 12. TempExprs
- 二、Type System in Relay
- 1. Algebrabic Data Type
- 2. Pattern Matching in Match Expressions
- 三、Relay Core Tensor Operators
- Relay Matching in Relay
- 四、优化
- 算子融合
- 量化
- 加速器相关优化
- 编译和执行
- 1)编译流程
- 2)部分执行
TVM前端之前用的NNVM,现在用的Relay,后面会往Relax和Unity方向转。先简单介绍一下Relay: A High-Level Compiler for Deep Learning。Relay的解释比较杂乱,按照论文和官方文档的解释它算是一个编译器框架或着IR(Intermediate Representation)。说是编译器框架有些大,说是IR他不单单可以做算子表示,还可以支持函数、类型等编程逻辑。简单来说,Relay作为TVM的前端表示是一种高阶的IR,不仅对算子和类型做了表示外还支持复杂的编程逻辑,类似于DSL(Domain-specific language),这是不同于其他简单的IR。Relay中定义了许多节点类型和函数类型,支持闭包,方便地对计算图进行描述。在TVM的运行过程中,用户会提供各种不同格式的模型如ONNX,TorchScript或者TFlite等,然后由解析器将这些类型转化为Relay格式,TVM提供的所有图优化操作会在Relay这种IR上进行操作,然后在将Relay转化为TIR来描述硬件相关的信息,Relay是后端无关的IR,不描述硬件信息。
深度学习IR梳理
1. IR属性
深度学习IR有三个挑战:1)表达能力,IR应该可以直接表示带有控制流、一阶函数、数据结构。2)兼容性,IR应该可以直接添加和整合新的优化操作。3)拓展性,他应该可以直接接入到新的设备中。Relay提供如下设计解决如上问题。首先,Relay IR是一个面向Tensor、静态类型的函数式IR,可以表达控制流、数据结构和一阶函数,提高表达能力。其二,将ML框架中的通用操作转化为编译Pass,这样就可以把传统编译器中的研究结果作为优化Pass利用起来,提高兼容性。其三,Relay提供了一种硬件无关的算子表示和领域相关的优化操作,确保了硬件之间的拓展性。
2. DL前端发展
DL早期是通过一些科学计算库如Numpy提供的低阶算子辅助编程的。模型会被表示为计算图,图中节点表示算子,边表示算子之间的数据流向。随着DL的发展,各大公司有了自己的开发框架如Tensorflow,Pyorch和编译器如XLA、Glow和TVM。这些框架可以分为支持静态图(static computation graphs)和支持动态图(dynamic computation graphs)两类。支持静态图的框架可以叫做先定义后运行(define-and-run),支持动态图的框架叫做边定义边运行(define-by-run)。支持静态图的框架对控制流和动态维度的模型支持不太友好,支持动态图的框架如Pytorch是借助python的特性边执行边构建计算图的,具有较高的表达能力,但是每次执行时都会重新构图,重新优化消耗巨大。
3. DL编译器
早期低阶的tensor编译器重点在于编写高性能算子如计算密集型的算子。对于代码的生成,比较新颖的设计就是计算分离架构,由TVM采用和多面体框架,由Tensor Comprehension等编译器采用。早期算子编译器的代码生成局限于标量循环嵌套,只能表示整个程序的一部分,忽视了内存管理、数据结构、闭包、控制流等细节。
现在的深度学习框架采用了编译器来处理性能和拓展性的问题,如XLA,GLow,nGraph和ONNC。这些图编译器通过计算图IRs,只做高阶的优化操作然后降阶到各种硬件或厂商指定的库上。降阶过程TF采用了MLIR,Pytorch引入了TorchScript。MLIR是一个共享的框架用于构建一组IR方言来实现编译器的的功能。Tensorflow通过为MLIR引入TF IR方言实现优化过程。TorchScript是一种类似于python语法的高阶IR,并作为Pytorch JIT编译器的的首层使用。PyTorch可以将程序改写为TorchScript格式,该格式可以由TorchScript VM执行或着通过JIT方式编译到目标平台。对于动态行为,TorchScript有一个分析JIT模式,可以在执行期间识别一个稳定的程序运行轨迹,这些稳定的静态轨迹可以进一步被一些低阶编译器优化。
4. DL编程语言
目前,针对机器学习的编程语言越来越多如JAX,Swift for Tensorflow和Lantern。Lantern是最接近Relay的编程语言,是一个深度学习DSL,可以作为代码生成器将代码降阶为C++或者CUDA代码。但是Lantern还不支持硬件加速器,也不专注于完整的程序优化。这些编程语言都是面向用户的DL编程环境的,并通过编译器IR生成代码。
Relay的主要内容
Relay是一个函数式的可微的编程语言,作为机器学习系统的IR使用。Relay支持代数数据类型、闭包、控制流和递归,相较于基于计算图的IR可以直接表示复杂的模型。Relay还包括一种使用类型关系的依赖类型,以便处理对参数形状有复杂要求的运算符的形状分析。