TVM前端研究--Relay

182 阅读19分钟

2024-10-29 21-59-22 的屏幕截图.png 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

1. Dataflow and Control Fragments

Relay是一个面向表达式的语言,其提供的数据流和控制段相较于传统基于计算图的IR是非常有用。数据流段是Relay程序中一段不包含控制流的部分,而控制流会根据某个值而沿着计算图的拓扑改变执行顺序。一般控制流部分包含If-Then-Else表达式、ADT Matching表达式。

从计算图视角看,函数就是一个子图,函数参数是子图中的自由变量的名字。因此,如果函数体中只有数据流,那么该函数的调用就在数据流段中,即函数调用也是Relay数据流中的一部分,反之,如果函数体中包含控制流则不是dataflow fragment的一部分。

2. 变量

Relay中变量分为全局变量用@表示和局部变量用%表示。全局变量在全局可视环境中(module)被唯一定义。局部变量一般引用一个函数或者let表达式。

3. 函数

与其他编程语言中的函数类似,其可以表示为子图的泛化。Relay中函数是高阶的,支持闭包,可以作为参数传递给函数或者由函数返回。函数中的参数和返回值支持类型指定和类型推演。如果是递归函数的话,需要通过let绑定将函数绑定到局部变量上然后在递归处用该变量名。

3.1 闭包

闭包就是函数中定义了函数,在函数调用时可以访问函数内部定义的局部函数。

3.2 多态和类型关系

一个Relay 函数可以支持多种类型,具体类型在调用时指定,同时返回指定类型。这个称作类型多态。

3.3. Call

在Relay中,带函数类型的表达式是可调用的,结果存储在局部变量中。

4. 算子

算子是一个原始的操作如add、conv2d,但不是在Relay语言中定义的。算子是在C++的全局算子注册表中声明的,许多常规算子由TVM 的TOPI(Tensor Operator Inventory)库承担。

注册一个算子用户需要提供算子实现、算子类型和其他元数据。因为算子注册表是基于列存储的,其中算子就是key,任何元数据有可以被注册为新的列。从Relay类型系统角度来看,算子就是一个函数,算子可以像函数一样调用并且拥有函数类型,算子的类型是通过类型关系(type relation)注册的。算子在Relay程序中通过指针唯一标识,没有前缀符号。

5. ADT Constructors

Algebraic data types是一个给定的函数类型,需要在call节点(函数或算子)中使用。。。

6. Moudle和Global Function

Relay中有一个全局的数据结构module用于记录全局函数的定义。Module的用处就是允许全局函数引用其他全局函数或者递归引用自己。Relay中的module类似于计算图中保存子图记录的数据结构。全局函数与普通函数表达式行为上相同,不同点在于全局函数函数名前有@作为前缀,可以递归调用。

7. 常量和元组

常量保持算子参数,元组提供多个输出的功能。没啥可介绍的,详见参考文献吧。

8. Let Binding

let binding是一个不可变的局部变量绑定,允许用户给表达式绑定一个名字。let binding包含局部变量,可选的类型标注、值。let可以将函数表达式绑定到变量上,变量可以在函数内递归调用。一组序列的let binding可以视作数据流图,这些bindings可以视作子图。

9. Graph Bindings

graph binding允许在Relay程序中显式创建数据流图然后绑定到一个临时变量上。每一个对该变量的引用对应指向该数据流图的一条边。不同于let binding,graph binding不是表示为Relay中的AST节点而是引用它AST节点值的meta-variable。出于开发和优化目的,Relay引入了pass实现通过graph-bindings定义的数据流图和通过let binding的A-normal程序之间的转换。

10. If-Then-Else

如同c-like编程语言中的语法

11. ADT Matching

看不进去, 还挺重要的,以后在补充吧。。。

12. TempExprs

Relay中的程序转化可能要求给程序AST插入临时状态以引导进一步转换。TempExpr就是出于这个目的被设计的。所有继承自Tempxpr的节点不能直接出现在用户代码里,一般在pass中插入,然后在pass结束后删除。

二、Type System in Relay

Relay是一个静态类型、类型推断的语言,需要程序完全类型确定或者少部分做类型标注。静态类型在编译优化中非常有用,可以在不运行程序的前提下获取数据的属性信息如运行时shape、数据布局和存储等信息。Relay的Algebraic Data Types允许轻松灵活地组合类型,以构建可以归纳推理并用于编写递归函数的数据结构。

Relay的类型系统由一组为shape提供的彼此依赖的类型组成,依赖叫做类型关系用函数实现表示输入输出之间的符号关系,也就是说,类型系统会跟踪tensor shape的轨迹。将tensor shape视作类型,shape推理转换为类型推理,允许Relay在编译期间进行推理tensor的shape,比如根据算子输入推理输出shape。静态的shape推理允许Relay进行AOT编译,在编译阶段提供更多信息用于优化。

拥有的类型:

Tensor,Tuple,Type Parameter,Type Constraint,Function、Incomplete、Algebrabic Data Type.

1. Algebrabic Data Type

ADT是一组命名的构造器集合,每个构造器拥有指定类型的参数。一个ADT实例就是一个包含参数具体值和构造器名字的容器。参数值可以在match的方式解析到。因此,ADT有时也被称作“tagged unions”,实例中可能包含不同类型,构造器起到了tag的作用指示如何解释内容。也可以理解为是enum或者struct的泛化。类似于enum是ADT内部包含有限个值;类似于struct是ADT实例包含一组拥有指定类型的字段,不同的是类型系统允许相同类型按照语义编辑进不同的组别中。从类型系统的角度分析,ADT构造器参数可以是带类型的参数,因此不同参数类型的ADT实例可以视作不同的类型。其二,可递归。ADT的构造器可以是该ADT的实例,因此ADT会像棵树或者列表,可以引导式构建。

为了便于递归式定义ADT,ADT的定义会绑定一个全局类型变量来唯一指定它。ADT会有不同的名字,因此尽管在结构完全相同情况下,还是会被视作不同的类型。

通过不同构造器但是相同参数类型创建的ADT实例是类型相同的,参数类型不同的两个ADT实例类型不同。比如一个整型的列表与一个浮点tensor列表不是相同的类型。

List是ADT的名字,是全局变量,可以被递归调用;a是类型参数,Nil和Cons是两个构造函数,Nil定义是fn<a>()->List[a],Cons定义是fn<a>(a,List[a])->List[a],递归引用了List。

2. Pattern Matching in Match Expressions

在Relay中,match表达式可以匹配多种不同的模式。具体来说,match case中的pattern分为三类:

constructor patterns会匹配具体的ADT构造器,如果一个值与构造函数匹配,则构造函数的每个参数都将与嵌套模式匹配。

wildcard patterns通配符会匹配任何值但不会绑定到变量上

variable patterns会匹配任意值然后绑定到一个变量上。

三、Relay Core Tensor Operators

算子是Relay不同于一般编程语言的点,Relay使用不同的算子允许后端根据硬件选择不同的降阶策略。并且Realy的算子集是可拓展的,允许用户添加新的算子。Relay提供了一组核心的tensor操作算子,用户可以基于提供的算子表示模型。算子主要分为如下几类:

  1. 基础算子,可以表示多层全连接网络。

  2. 卷积算子,卷积相关的算子,位于tvm.relay.nn包下

  3. 数学和转化算子,

  4. 广播和规约算子

  5. 视觉或图像处理算子,位于tvm.relay.image/tvm.relay.vision包下

  6. 算法算子如argsort, topk

  7. 临时算子,支持广播算子的反向传播

  8. 方言算子,主要是支持量化的算子,位于tvm.relay.qnn.op包下

Relay Matching in Relay

在Realy程序中可以识别一个纯的数据流子图然后以某种方式转化如融合、量化、外部代码生成和设备相关优化。许多pass需要实现许多模板代码,需要用户从遍历器和AST匹配的角度思考。许多转化pass可以简单描述为图重写,因此需要一个语言在构建重写过程时需要首先描述一个可以匹配的pattern。这种语言不仅对构建重写过程有用,同时对已有的pass提供扩展点,比如fusion pass可以通过一组fusion pattern参数化,量化pass可以持有一组patterns用于确定那些算子要在给定平台量化。

dataflow_pattern其实就是根据tvm中的正则接口做了python端的封装。具体参考Language Referencetvm.relay.dataflow_pattern - tvm 0.19.dev0 documentation

四、优化

算子融合

融合算子可以共享计算,移除中间结果的分配,进一步优化循环嵌套。传统融合技术太过封闭,仅限于特定的算子或指定的硬件。比如,传统算子融合技术会首先选择一个用于融合的算子序列,然后用对应手写的融合后的算子来替代掉这组序列。一般这种融合后的实现是有厂商的算子库提供的,如果没有融合后的算子实现则不进行算子融合。当然,也有一些框架如XLA等为选择的待融合序列生成对应平台的代码。

Relay算子是以TVM计算表达式为支持的,该表达式以高阶的DSL描述操作过程,忽略了低阶的调度细节。TVM的计算与调度分离的策略为Relay算子融合提供了强大的支持,可以融合任意长度的算子链,生成shape确定的融合算子。TVM还可以在融合算法和优化算法结束后通过auto-tuning重新调度。

Relay的融合主要分两步:1)提取,首先识别满足融合条件的子表达式,然后将表达式分解为局部函数并标记为Primitive表示该函数为基本单元。Primitive函数最后会降阶为平台相关的代码。识别子表达式是通过构建后支配树获取的,这部分以后单独出一期再讲。2)降阶。Relay编译器会将primitive函数生成为平台和维度确定的代码。对于每一个算子,Relay首先用TVM表达式表示,然后将这些表达式融合为一个聚合的表达式以表示融合的算子。用TVM生成代码还需要一个调度策略,可以为单个算子采用默认调度策略但默认调度不支持融合。为了为融合后的表达式生成代码,需要基于待融合算子生成一个master调度,该调度会采取适当的调度行为生成融合后的代码。

量化

受限于硬件资源约束(内存),深度学习模型部署在边缘侧具有一定难度。直接降低数据位宽会显著降低模型准确率,量化神经网络通过使用更小的精度或非标准的数据类型来提升内存的利用和吞吐。不同量化技术的优劣取决于平台和模型类型。目前,大多DL框架选择了一组定点量化策略和数据类型,因为需要手动实现算子。

与以往不同的是,Relay引入一个基于编译器的量化流,支持多种量化策略,可以为每个算子自动生成代码。Relay提供了一个一般意义上的程序重写框架,可以为每个算子用指定类型和精度标注输入输出来量化。用户可以用新的量化策略如选用有符号或无符号整数、不同截断策略如floor、ceiling和stochastic rounding来替换Relay默认量化规则。

上图展示了整个重写规则,主要包括三个步骤:标注annotation、校准calibration和实现realization。其中,标注就是为每个算子按照标注规则插入模拟量化算子。每个要量化的输入输出会经过一个simQ算子,该算子模拟了量化过程如将float32转化为int8。simQ算子有一组参数需要被校准以正确量化计算图,主要包括位宽bits、缩放尺度scale和数据范围range。simQ模拟量化的过程就是将未量化的类型缩放到目标类型上,如上图中的公式所示。

对于校准这些量化参数,Relay提供了大量策略来设置这些参数。第一种策略就是超参暴力搜索,在所有scale中遍历直到找到不会益处的结果;或者可以为每个通道选择一个scale,也可以采用基于KL散度的loss优化scale。

当量化参数确定后,实现过程就是将simQ算子转化为如下量化算子

加速器相关优化

对于一些特定的加速器,可能无法完全执行Relay程序,需要通过一些特殊的优化操作将Relay编译到硬件加速器上。Axis scale folding是一种移除类卷积算子前后的scaling算子的优化操作。这个优化操作是为缺乏标量乘子的加速器适配的,旨在消除所有标量算子。Parallel convolution combination是一种融合多个2D卷积,共享相同输入的特定优化操作。该优化操作的目的在于为GPU生成一个大的kernel。

编译和执行

1)编译流程

编译Relay需要三个步骤,首先,前端转化输入格式为Relay IR,然后,Relay编译器做类型检查和优化程序。最后。Relay后端会将Relay程序转化为后端硬件可执行的格式,过程就是降阶算子为TVM表达式,然后为最终的表达式计算一个调度,然后降阶为可执行代码。

前端部分有多种方式生成Relay程序。用户可以通过C++或Python API构建,或者解析一个Relay的文本格式或者从从磁盘加载序列化的格式,亦或者从其他框架中导入并内部转化为Relay程序。

编译器部分主要执行一些优化操作(这个是Relay论文里这么说的,对这里的“编译器”不做强行解释,这里就是优化操作,没有编译格式上的转化)。Relay解析为抽象语法树AST,一组Pelay-to-Relay的Pass依次执行,或收集信息,或改变结构。

Relay后端通过解构代码生成程序为多个不同的阶段来生成机器指定的代码。Relay会把所有算子转化为TVM表达式以生成稠密的现行代数内核,结果是一个对象文件,包含所有算子硬件相关的实现。通过将算子表示为TVM表达式可以程序化的转化他们或者自动为转化的算子生成新的实现。

Relay程序的执行有多种方式:Relay解释器,一种JIT的编译方式;Relay 虚拟机;TVM graph runtime以及Relay AOT编译器,将Relay程序转化为C++以生成目标指定的二进制格式。

2)部分执行

现有深度学习IRs依赖于状态和常量的混合评估以优化用户程序。Partial evaluation是一种泛化的常量评估,可以减少部分常量程序。Partial evaluator允许使用高阶抽象,不用限制代码,可以编译到特殊的目标平台上。Relay是第一个将Partial evaluation技术应用到深度学习场景的编译器,当其与其他优化操作组合时会得到许多有用的操作。比如Partial evaluator可以执行循环拆解并进一步不融合而不用额外的编译操作。Partial evaluator是通过定义一个解释器执行的,其中值是静态的。evaluator需要把程序转化为A-normal的形式以确保正确有序的计算。

参考

Language Reference

Relay IR 简介 | Apache TVM 中文站