TVM前端研究--Relay

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还包括一种使用类型关系的依赖类型，以便处理对参数形状有复杂要求的运算符的形状分析。

一、Expression in Relay</