Tiramisu：一种基于Polyheral的深度学习模型编译器

1. 推理引擎编译模型的一般过程

业界主流的深度学习推理引擎，如TensorRT、Tensorflow Lite、TVM等，均由两个主要组件构成：模型编译器和推理运行时，前者负责将模型编译为可运行于目标设备的IR/代码，后者则负责在目标设备上导入编译后的IR/代码并执行推理过程。很明显，模型编译器是否能生成合理运用设备资源的代码，是保证模型推理性能的关键。模型编译器编译模型的一般过程，就是将模型由高级表示，转换为一系列中间表示（IR）的过程。如下图，以TVM为例，这个过程是：原始模型->Relay IR（图级别IR）-> Tensor IR（算子级别IR）-> Runtime IR（运行时IR）。

TVM框架模型编译过程

其中，由Relay IR转化为Tensor IR的步骤称为“调度生成”，对于TVM以外的推理引擎，该步骤的名称可能有差异，但都是必须的。调度的关键在于从时间、空间两个维度上彻底挖掘计算资源的极限潜能：

• 时间维度上，任意时刻应尽可能地并行化数据无依赖关系的计算（CPU多物理核、GPU流处理器、CPU/GPU高级指令集），以避免不必要的串行计算

• 空间维度上，尽可能将计算所需的数据放置于存取速度最快的区域（CPU L1 Cache、GPU L1 Cache、GPU C-Cache、GPU共享内存等），减少数据读写时间

综上可知，推理引擎的核心在于模型编译器，模型编译器的核心在于调度过程的生成技术。

2. 调度生成技术的分类

调度过程是否合理直接决定了CodeGen生成的运行时代码的执行效率，目前业界的调度编译技术可以归纳到下表的四个象限中：

调度生成技术分类

	调度自动生成	调度非自动生成
依赖分析	PolyMage, Tensor, Comprehensions	AlphaZ,CHiLL,URUK,Tiramisu
区间分析	AutoScheduler(TVM)	Halide,AutoTVM

如上表，编译技术根据循环嵌套分析算法的不同可以分为基于依赖分析和基于区间分析的两大类：依赖分析也即传统编译器的应用仿射变换优化循环嵌套代码的多面体分析技术，由于深度学习模型的算子在推理阶段的循环控制流是静态可判定的，因此非常适合应用该技术优化计算过程；相比依赖分析，区间分析针对图像处理领域的常用计算（针对图像矩阵的卷积、池化操作）简化了循环计算过程为循环轴对齐，即简化依赖分析的多面体抽象为长方体抽象，以牺牲一定的资源利用为代价简化常用算子的编译过程。

两者相比，基于依赖分析的Polyhedral模型的调度描述更加细化、表达力更强，理论上可以将优化做到极致，但缺点是算法原理相对复杂且优化分析的复杂度更高；而基于区间分析的调度模型的优势在于，其在图像处理领域的优化效果和前者相差无几，但优化分析的复杂度低很多，缺点则是对于图像处理领域外的代码调度表达力不足，难以优化运行代码到极致性能。

另一种分类方式是调度生成的自动化程度，非自动化生成调度的编译器通常会向用户提供一种领域特定语言（DSL），如TVM的Tensor Expression、Tiramisu的Tiramisu Language，用户使用DSL语言描述由算子的计算到具体调度的转化过程；而自动化生成调度的编译器则会内置一套或多套编译准则，这套准则根据用户定义的计算过程描述以及设备性能描述自动生成最优的调度过程。

两类方法相比，非自动化的方法需要用户对目标设备的体系结构有足够理解并提供调度生成模板（AutoTVM）或具体调度过程（Tiramisu），用户在自定义的模板/过程上可以调整调度参数以优化调度过程；而自动化的方法则是对编译准则的设计者在计算机体系结构、代码编译原理方面提出了很高的要求，以确保设计的编译准则可以根据给定算子以及运行设备信息生成高效的调度过程。

根据以上对比可知，“四象限表”列出的所有调度生成技术并不存在绝对的优劣之分，每一种技术都是根据自身需求在通用性/特定领域性能两个维度上做取舍。本文即将介绍的Tiramisu可以归类于“基于依赖分析的调度非自动化生成”中。

3. Tiramisu DSL

Tiramisu定义了一套领域专用语言（DSL），该语言以C++为基础，提供了一套API供用户调用。用户可基于Tiramisu DSL APIs定义循环优化、内存布局等转化规则以指导算子调度的生成过程，Tiramisu Compiler进而根据用户定义的规则将原始深度学习模型的所有算子转化为低级别IR，并最终生成运行于设备后端的优化代码。理解Tiramisu DSL的一个高效方法是了解其定义的数据结构和算法，下图展示了Tiramisu转化代码为设备代码的全过程。

3.1 Tiramisu DSL——算法篇

Tiramisu定义的算法更准确的称呼是调度命令，用于描述数据排布如何设置、多层循环如何做仿射变换、设备计算资源如何利用等信息。Tiramisu共定义了4种类型的调度命令：

• 循环嵌套变换命令：这一类型的调度命令包括常见的仿射变换，如循环展开、分割、移位等。

• 循环-硬件关联命令：该类型的调度命令包括循环并行、向量化以及绑定循环到指定计算资源的操作等。

• 数据操作命令：数据操作命令可以分为4种类型：(1) 分配Tensor空间命令 (2) 设置Tensor属性命令，如设置数据存储位置(host/device/shared) (3)数据拷贝命令 (4) 设置数据存取属性命令。如表所示，数据操作命令也有高级和低级之分，通常用户使用高级命令即可完成一般的调度规划，更细致的规划则需要低级命令参与

• 数据同步操作命令：Tiramisu相比其他Compiler比较有特色的命令，类似于MapReduce的思路。设计者考虑到一次计算的数据量非常大的情况下可能需要多节点共同计算，因此设计了send/recv的调度操作，籍此可以在多节点之间共享数据或数据片段。

以blur算法为例，原始的blur算法计算定义如下：

与之等价的计算过程为：

接下来可以应用Tiramisu定义的调度命令优化blur算法的计算过程，如下图(a)所示，tile()命令将by计算作了循环展开操作；同时compute_at()命令在j0循环开始的地方计算了bx，供后续by的计算过程调用；parallelize()命令在i0循环处并行计算i0对应的循环体（即i0循环对应的代码块）。

同样的blur算法，在GPU上优化的方式则大不相同，如上图(b)：首先，tile_gpu()命令展开by循环的计算并映射展开后的循环到GPU block上；compute_at()命令和(a)的功能相当；cache_shared_at()表示将bx的计算结果保存于共享内存中；store_in()指定了bx和by的入口函数，本例中表示bx和by的计算结果以SOA格式存储；最后的device_to_host_copy明显是计算后的结果拷贝（GPU到内存）。上图(c)考虑blur算法运行在分布式系统上，假设数组in[][][]在分布式系统中的所有节点上都已初始化，每个节点n根据数组in的一个chunk执行各自的计算任务，再用send()和recv()命令在节点间同步计算结果。

3.2 Tiramisu DSL——数据结构篇

数据结构方面Tiramisu定义了4层中间表示，分层式设计的目的在于解耦循环嵌套、内存排布以及分布式计算通信三类优化操作，以简化调度命令的设计过程。第一层IR用producer-consumer关系描述原始算法的计算过程（不考虑内存分配）；第二层IR指定算法涉及的所有子计算的执行顺序；第三层IR指定数据在被调用之前应当以何种布局被放置于哪里；第四层IR（可选）指定在分布式系统中各节点协同计算的方式。

理解基于多面体模型(Polyheral)优化过程以及后文的IR定义需要掌握两个基本概念：整数集合与映射。在Polyhedral模型中，整数集合代表的是迭代域，映射用于表示内存访问并转换迭代域和内存访问（应用循环嵌套和内存访问转换）。例如，以下的整数元组集合描述了一个两重循环：