MLIR应用技术分析

MLIR应用技术分析

7.1.MLIR多层编译框架实现全同态加密的讨论

摘要

MLIR是谷歌推出的开源编译框架。基于论文：SyFER-MLIR: Integrating Fully Homomorphic Encryption Into the MLIR Compiler Framework，介绍将全同态加密技术集成到MLIR编译框架中，达到一次加密，多设备编译的目的。这个实例对于如何使用MLIR，具有非常好的启发性。分为两部分，第一部分是MLIR相关技术介绍，另一部分是介绍同态加密技术及其在MLIR中的集成。

MLIR简介

MLIR是谷歌团队针对异构系统提出的开源编译框架，旨在打造具有可复用性和可扩展性的编译器基础设施。正如其名，MLIR主要特点在于多层级表达，在不同层级间实现转换，优化以及代码生成。MLIR可以针对领域专用架构架构（DSA：Domain Specific Architecture），比如现在出现的很多ML加速器，快速搭建编译器；同时也能够连接现有的编译技术，比如LLVM。MLIR源起于Google Tensor Flow框架的编译生态中存在的诸多挑战，如图1所示。 

 

图1 Tensor Flow 编译生态

 

MLIR的核心是为构建现代编译器提供丰富的基础组件，这些基础组件用于不同层级的中间表示的转化及优化中。不同于传统编译器通过一层的中间表达（IR: Intermediate Representation）直接翻译为可执行指令，MLIR通过多层级的中间表达实现转换，在每层级的IR表示中，MLIR提供丰富的优化组件进行优化。用户也可以自己开发优化组件，最终翻译为可执行指令，映射到不同的设备。MLIR的设计理念使得编译器整体框架从抽象到具体结构清晰，层次分明，任务明确。MLIR项目文档中的几个Rationale文章是理解其设计思想的很好参考。

MLIR将中间表达也称之为方言(dialect)，方言都符合相同的语法格式，如图2。在MLIR语法中，采用的基于SSA(Static Single Assignment)的数据结构，Operation（下文翻译为操作）是一等公民(The first class citizen)，每个Op都有自己的标识符(Identifier)，Op接受输入参数(Value)同时也可以返回参数，Op具有属性(Attribute)，属性在编译阶段需要保持为常量，用于在编译阶段告知编译器Op的基本信息。具体可以参考后文。 

 

图2 MLIR 语法格式

 

图3 MLIR 代码生成的流水线(Codegen Pipeline)

图3表示的是MLIR方言不断下译的流水线，图中左侧部分表示在每层方言中进行的优化操作，中间部分表示方言之间的转换。前端的输入来自TensorFlow，对应于编译器的结构，可以看成为MLIR的前端。在方言的不断转换的过程中，体现的是MLIR的重要特点渐进下译。以下以图3为例介绍MLIR的lowering生成代码的过程。

1.TF的描述转换为HLO(High-Level Optimizer)方言，HLO实际上是XLA（Accelerated Linear Algebra）的IR表示，里面的Op基本和XLA一一对应。

2.HLO转换为LHLO(Late HLO)，HLO和LHLO的区别在于HLO注重的是tensor的表达，不考虑到内存的分配，比如tensor<8x32x16xfp32>，仅仅表示为tensor，没有具体的内存信息，LHLO会为tensor开辟内存空间，也就是图中的buffer assign，buffer assign相当于传统编译器中的内存分配。

3.LHLO被转换为Linalg(Linear Algebra) 方言和Affine方言，Linalg为线性代数方言，Affine是针对多面体编译的方言。在此层方言表达中完成tiling展开，for-loop循环优化操作等。

4.针对不同的硬件设备产生代码，在此层转换成LLVM IR，然后再通过LLVM的codegen生成代码。

MLIR提供了一个比较灵活和规范的框架，相应的，也看到很多基于MLIR的工作，逐渐形成一定的生态，甚至超出了ML编译器的范畴。但是，对于MLIR的争论和质疑也一直存在。一个近期的讨论可以参考知乎问题，如何评价MLIR项目中Linalg Dialect的设计思想？（https://www.zhihu.com/question/442964082），这里不再赘述。

抛开这些争议，从实际的应用上，MLIR的出现确实起到了辅助编译器开发的作用。大家可以更快地搭建解决特定问题的编译工具，或者将一些新的特色引入到编译器中，比如下面将要讨论的基于MLIR编译器框架实现同态加密，非常具有启发性。

同态加密技术介绍

同态加密（HE: Homomorphic Encryption)是一项直接在加密数据上进行运算而无需先解密数据再运算的加密技术，最终通过解密计算结果就可以得到在非加密状态下运行的结果。同态加密技术广泛应用在涉及到敏感数据处理的场景，比如医疗数据分析和基因分析等。

 

图4 客户端-服务端的同态加密场景（C代表客户端，S代表服务端）

下面以图4的客户端-服务端同态加密使用场景介绍同态加密解密过程：

1. 客户端将数据进行同态加密处理得到加密数据Enc(m)，m表示信息，Enc表示加密方法。

2. 客户端将加密的数据发送到服务端。

3. 当客户端想对自己的数据执行一个函数操作f()，会将操作的名称发送给服务端，比如图中的query。

4. 服务端接收到命令后，在加密的数据上进行f()操作。

5. 服务端返回f()操作的结果。

6. 客户端通过解密的方法获得原始数据的计算结果。
在上文关于MLIR技术介绍中提到了在MLIR的代码生成的流程中采用不同层级的中间表达和优化，那么也理应可以在中间表达和优化中引入关于同态加密的技术。

同态加密技术集成到MLIR框架中

同态加密分为4种，所分析的论文针对的是全同态加密(FHE:Full  Homomorphic Encryption)，以下统称为同态加密。常见的FHE方法包括Brakerski-Gentry-Vaikuntanathan(BGV)，Fan-Vercauten(B/FV)和Cheon-Kim-Kim-Song(CKKS)等。针对这些FHE方法，现有的实现是通过调用库的方式实现，比如MicrosoftSEAL。但是目前的HFE库存在几点局限：1.只能进行有限的跨操作优化(cross-operation optimization)，跨操作优化可以理解为移除掉不必要的指令或者调整操作顺序，参考后文图7。库虽然提供了一系列的原语(primitives)，这些原语操作都是单独优化，不能进行跨操作优化; 2.不能进行重写操作(rewrite)以及高层次的优化，比如NAND(x,x,parameters)不能重写为NOT(x,parameters) ; 3. 不能进行低层次的优化，比如嵌套循环优化和循环调度; 4.缺少模块化，这样导致更改加密方法时，需要重写代码。

针对以上基于FHE库实现方式存在的缺点，论文中提出了利用MLIR的编译框架产生加密的程序，其核心思想在于将FHE库改写为方言，类似前文提及到的将XLA改为HLO，目前实现的是Gentry-Sahai-Waters(GSW)和B/FV，具体流程如下图5所示。

 

 

图5 GSW和B/FV逐级下译管道的简易框图 

正常的MLIR编译流程如图5中的左图所示，基于GSW和BFV的加密编译流程如图5中中图和右图所示。普通代码作为输入，采用MLIR的逐级下译的功能将代码转换为较低级别的IR，这些IR在经过自开发的lowering规则转换为加密的IR，同态加密的方言转换过程为Linalg->Affine/SCF(StructureControl Flow)->Standard -> GSW/BFVIR，实现了同态加密的IR然后按照正常的MLIR编译流程编译，最后在LLVM层级可以根据不同的硬件设备产生不同的代码，这样带来了一次加密，多设备编译的便捷性。

 

图6 GSW和B/FVLowering的详细过程 

从GSW和BFV向下lowering的详细过程如图6所示，GSW首先lowering到S-GSW(Simple GSW)、Standard以及LinAlg方言，S-GSW然后lowering到Affine，SCF和standard方言中，最后Affine，Standard，LinAlg和SCF都lowering到LLVM。图中的箭头表示的Pass，实现的作用是从一种方言向另一种方言的转换，颜色越深表示实现难度越大，蓝色是难度最低的，MLIR的框架中提供相关的Pass实现，用户也可以自己开发Pass。可惜文章没有给出GSW和S-GSW以及BFV和S-BFV之间的区别。

在同态加密的方法中，每个加密的函数操作能表示为一个DAG图（Directed Acyclic Graph），每个节点表示一个原语，每条边表示一个原语的输出作为另外一个原语的输入。

针对选择的GSW和BVF，它们的原语都很简单，GSW的原语是二进制的门电路(Binary Gate)，BVF的原语是加法和乘法，简单的原语极大地降低了优化的难度。

在Lowering的过程中，论文使用了三类优化技术。第一，使用重写机制完成优化操作，重写操作可以简单地理解为使用一种表示方式替代一种或者多种的表示，目的在于使用高性能的Op替代低性能的Op，比如前文提到的使用NOT(x,parameters)替代NAND(x,x,parameters)，此外还提及到了复杂的重写机制，比如修剪代码，删除常量和优化两输入和三输入的子图，但文中没给出具体的实例说明。第二，实现跨操作（cross-operation）优化，跨操作优化的目的是删除不必要的指令，比如文中举例合并数论转换(NTT:Number Theoretic Transform)加法的操作，如图7所示。第三，针对不同的硬件进行循环优化，这也是MLIR本身具有的特性。

 

图7 多项式乘加的优化 

的核心在于将FHE两种的库GSW和BFV改写为MLIR中的方言，并实现了各个层级的转换和优化，最终将同态加密技术集成到了MLIR编译框架中。整体论文的实现难度相对较低，也比较遗憾文中没给出实验结果，但作者的思路是值得肯定的。

在AI应用领域中，对于端侧敏感数据的处理往往采用本地私有化部署的方式，在端侧完成推理，但端侧设备的性能有时会成为瓶颈。或许存在下面两种可能，在端设备上使用集成了同态加密的MLIR编译器，将编译后的指令传到云端执行，然后云端将执行结果返回给端设备；云端部署同态加密MLIR编译器，端设备将同态加密的数据发送到云端处理，云端部署编译器在于能够适配不同的加密方法，具有更高的编译性能。

7.2.MLIR编译框架下软硬协同设计的思考

摘要 

自从AI芯片成为热门的研究课题，众多关于AI芯片架构探索的学术文章不断涌现，大家从不同的角度对AI芯片进行架构分析及性能优化。MLIR是谷歌团队推出的开源编译器框架，颇受瞩目，灵活的编译器架构提升了其在众多领域应用的潜力。通过自定义IR的衔接，可以在架构探索和MLIR之间架起一座桥梁，在编译的过程中，自动进行硬件架构的探索和软件的优化编译，甚至生成硬件的代码，实现软硬协同设计。

架构探索方法的介绍

近十年，AI领域专用芯片的演进极大地促进了架构探索（指的是架构定义及性能分析）的发展，先后出现了众多的分析方法，这些分析方法针对AI计算过程中关键算子以及网络模型进行建模分析，从PPA(Power-Performance-Area)三个角度评估硬件性能。与此同时，伴随着AI编译框架的发展，尤其受益于MLIR编译器框架的可复用及可扩展性，将这些分析方法融入到MLIR框架中也变得十分可能，从而使用编译器对硬件架构进行探索。 

架构分析中关注三个方面的表达，分别是计算架构（Computation Element），存储结构(Memory Hierarchy )和互联结构(Interconnect)。对硬件架构进行性能分析时，数据流是搭建分析方法的基础，根据数据流的表达，将workload的计算过程映射到硬件架构的三类实现中。在学术研究中，Eyeriss 是较早将数据流引入到AI芯片的性能分析中，根据定义，AI的数据流可以分为三类，输出静止（Output Stationary），权重静止（Weight Stationary）和行静止（Row Stationary）。随后的研究中，MAGNet将其扩种为更多的描述方式，如图1 所示，但还是围绕OS，WS和RS展开。根据数据流的划分，AI架构既可以分为这三类，比如NVDLA属于WS，Shi-dinanao属于OS，Eyeriss属于RS。相同的数据流架构可以采用类似的方法进行分析。

 

       图1 不同数据流对应的for-loop表示

围绕数据流表示和硬件映射的表达上，可以归为三类，分别是以计算为中心(computation-centric)的Timeloop, 以数据流为中心(data-centric)的MAESTRO和以关系为中心（relation-centric）的TENET。以计算为中心的表示方法关注的是for-loop表达在时间维度上映射到硬件架构；以数据流为中心的表达关注的是数据映射（data mapping）和复用（reuse）；以关系为中心的表达关注循环表达和计算单元及调度之间的关系。将对第二种data-centric的表达方式展开。

在MAESTRO的工作中，将data mapping和reuse作为一等公民，关注的是数据在时间和空间两个维度的复用。对于WS的计算架构，weight在时间维度上复用（相当于保持不变），中间计算结果是在空间维度上复用，其复用如图2所示。

 

图2 2x2 kernel的卷积在WS类型加速器数据复用的表示

关于时间和空间数据复用的表达，文中提出了一种IR的表示方式，称之为时域映射(Temporal Map)和空域映射(Spatial Map)。时域映射表示特定的维度与单个PE之间的映射关系，空域映射表示的是特定的维度与多个PE之间的映射关系，具体的表示如下：

1.T(size, offset)α： α表示的特定的维度，比如权重的weight，width及channel等，size表示表示单个时间步长（time step）下α所在维度的index映射到单个PE的尺寸，offset表示的是相邻的时间步长的index偏移。对应的for循环表达如图3所示。

2.S(size, offset)：α表示特定的维度，size表示维度α映射到每个PE的index的尺寸，offset表示映射到相邻PE的index偏移。

 

图3 时域和空域映射与循环表达之间的对应关系

假设一个计算架构有3个PE，卷积的权重大小为6，输入元素个数为17，步进为1，计算过程可以通过图4表示。在图中，标签1表示for循环的表达，标签2表示在时域和空域的IR表达，标签3表示数据在PE的分布及时间上的计算过程，图中可以看出cycle1到cycle4复用S中的index(1)，也就是weight保持静止。标签4表示空域映射、时域映射以及计算顺序，其中t表示按照所示的箭头方向依次计算。基于这样的IR表达及时间上的计算过程，就可以表示出一个WS架构的计算过程。 

 

图4 1D卷积操作在时域和空域的表示演示图

基于IR的性能分析方法 

Aladdin是较早开展基于编译的方式进行硬件的性能分析，将性能分析提前到RTL代码之前，避免了RTL代码及C-model大量的开发工作，基本的思路是将计算任务lowering到动态数据依赖图（DDDG：Dynamic Data Dependence Graph）级别，DDDG是针对特定架构的中间表达（Intermediate Representation）的表示，如图5所示。针对特定的硬件架构，分析DDDG的动态执行过程，即可评估出性能和功耗的数据，他们基于ILDJIT compiler IR。

  

图5 DDDG的计算表示

 基于GEM5的工作，他们将其扩展为GEM5-Aladdin，用于对加速器系统级的性能分析，涵盖了SoC的接口通信开销，从而实现加速器架构和通信的协同设计。GEM5负责CPU和内存系统的性能分析，Aladdin负责加速器的性能分析。DDDG的表示从ILDJIT IR迁移到LLVM IR。

Interstellar是将Halide语言用于AI架构的性能分析，数据流表达的方式属于computation-centric，核心工作是将和计算及数据流相关的for-loop转换到Halide的scheduling language，同时显性表达存储和计算。其中，关于架构和数据流是在Halide编译过程中的IR表达中引入，同时和Halide语言中的hardware primitive对应起来，将整个计算过程拆解到IR级别，然后映射到硬件结构，最后根据数据流的计算过程评估硬件的性能，整体过程如图6所示。最终采用调用硬件语言代码库的方式生成硬件设计。

 

图6 标签1为Halide语言描述conv操作；标签2表示Halide Lowering过程中对in, compute_at, split及reorder调度原语(scheduling primitives)的IR表示；标签3表示调度原语和硬件架构的对应关系

架构级别的IR

Micro-IR文章的核心思想是将加速器的架构表示为一个并发的结构图（Concurrent Structural Graph），每个组件就是一个架构级别的硬件单元，比如计算单元、网络或者存储器。结构图中显性地表达了加速器的构成组件，以及不同组件之间的数据流动，最终回归到数据流的表达和实现上。定义架构级别IR的好处在于1)将算法的表达和硬件架构解耦，2)将硬件的优化和RTL的代码实现解耦。这样一来，硬件架构IR层的优化工作可以单独展开。

整个编译的架构基于LLVM的实现，前端接入为AI framework，然后编译到LLVM IR，LLVM IR再对接到Micro-IR，在Micro-IR优化的PASS中聚焦就是前文提及到的关于数据流的映射、调度，tiling以及映射到硬件的intrinsic。最后对接到chisel的IR FIRTL，生成可综合的硬件语言。

 

图7 Micro-IR的编译流程

对于架构的表达，也是围绕数据流、存储和互联的展开，如图8所示，将一个简单的奇偶乘法翻译到IR图层，再翻译到IR的具体表达。

 

图8 Micro-IR的编译表示

MLIR中引入架构探索的可能性和挑战

简单总结一下上述工作中的关键点：

1.经过上述章节的分析发现现有的性能分析方法的研究工作都有IR表示的思想，而且基于数据流的表示思想具有较好的理论基础，从时域和空域两个维度展开，也有很好的IR具体实现。

2.基于IR性能分析的方法也处于不断演进的过程中，从ILDJIT到LLVM再到Halide，都证实了基于IR进行架构探索的可行性。同时不同的表示方式具有不同的有点，比如Halide中突出调度的思想，可以将该思想引入到MLIR中，形成schedule IR。

3.关于硬件架构IR表示的文章也较多，比如Spatial，文中举例的micro-IR 是比较典型的标准，与MLIR都基于LLVM的编译流程，将其引入到MLIR中作为硬件架构存在可能性。

4.Union是将MAESTRO性能分析的工作引入到MLIR的框架中，但是MAESTRO是作为架构探索的工具使用，没有接入到MLIR的编译流程中。

挑战与可能性同在，大致分为如下几点：

１.目前的架构探索都是基于相对规则的架构展开，没有涉及到复杂的工业界的芯片，存有一定的局限性，将其方法应用到工业界还有很大的隔阂。

２.定义一个通用型的架构IR比较困难。架构是比较分散的，不同的任务需求有不同的架构设计，虽然架构设计从大的层面分为计算、存储和互联，但通过IR精准地刻画架构充满挑战，比如对于架构IR控制流的表示，Micro-IR中关于控制流的表达没有进行详细的阐述。

３.在编译过程中，如何将软件任务能够自动翻译到架构IR上，同时能够对硬件架构进行自动调整和优化，这也是很大的挑战。目前是针对特定的已知架构，将计算任务映射到硬件。

 结论与思考

总结了现有的针对AI架构的数据流分析方法，以及基于数据流分析方法构建的架构探索工具，同时介绍了现有的硬件架构的IR。这些丰富的分析方法和IR表示为架构探索引入到MLIR提供了可能性，也让看到了基于MLIR的编译器框架开展软硬协同设计的巨大潜力。

7.3.外部借力：Pluto助力MLIR编译器的多面体优化

摘要

多面体编译是一项成熟的编译优化技术，演进了几十年，在传统的编译器中常作为一种优化工具使用，比如LLVM中使用的Polly，在GCC中使用的GRAPHITE。近些年来，多面体技术也引入到AI编译器中，进行循环优化及算子融合优化等。将关注在MLIR中以类插件的形式引入多面体优化技术，补充其多面体优化能力。

多面体模型的介绍

多面体模型(Polyhedral)主要关注的是程序设计中的循环优化问题，两层循环的循环变量的取值范围可以构成一个平面，三层循环的循环变量可以组成一个长方体，如图1所示，因此得名多面体模型。

 

图1 不同循环层次的多面体表示

多面体编译优化关注的是在确保程序执行正确的前提下重组多重循环的结构，实现性能的最优化。比如图2的循环中，左图表示的原始的二维迭代空间，蓝色箭头表示数据（黑点）之间存在依赖关系，对角线的绿色表示数据没有依赖关系，经过变基操作之后变为右图的表达式及迭代空间，从形状看像是把多面体进行了变形，形象地体现出多面体优化的过程。当然，变形的目的是为了实现并行计算，达到更好的性能。

 

图2 多面体的变基转换及Affine_map表示

MLIR中的多面体表示

MLIR关于多面体的设计重表达轻优化，也就是说MLIR充分利用其IR的特性定义了表示多面体的方言Affine，而没有进行多面体的优化实现。这样做的目的符合MLIR重在编译器基础设施搭建的特性，而具体的实现可以自行定义，好比MLIR向用户交付了布局合理的毛坯房，内部装修各取所好。另外，关于多面体优化的工具也很成熟，各种开源工具齐全，比如isl，Polly，Pluto，以及CLooG，也为将编译优化工具引入到MLIR提供了便利。

MLIR中Affine方言的定义使用具有多面体特征的循环和条件判断来表示，显示地表示静态控制部分（SCoP：Static Control Part），比如affine.for, affine.if, affine.parallel等，具体可以参见官方文档。在Affine的表达中，使用Dimension和Symbol两类标识符，二者在MLIR语法中均为index类型，同时MLIR也对这两类表示进行约束有助于提升分析和转换能力。从表示形式上看，Dimension以圆括号来声明，Symbol用方括号声明，Dimension即字面意思表示仿射对象的维度信息，比如映射，集合或者具体的loop循环以及一个tensor，Symbol表示的是多面体中的参数，在编译阶段是未知的。在准线性的分析表达上，Symbol当作常量对待，因此Symbol和Dimension之间可以进行乘加等线性操作，但Dimension之间的操作是非法的。另外，Dimension和symbol都遵从SSA赋值。

MLIR的Affine重表达体现在通过具体的映射可以表示出多面体的变换，比如图1中的变基操作，通过affine_map语句就能够体现出来。矩阵计算中常用的关于内存的tiling操作，也可以通过affine_map表示。通过#tiled_2d_128x256 = affine_map<(d0, d1) -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)>可以表示如下图所示的tiling切分。X维度按照128个元素，Y维度按照256维度进行tiling切分，形成了x.outer和y.outer，内部的取模运算计算在单个tiling内部的偏移量，形成x.inner和y.inner。图4中表示采用Affine方言表示多项式乘法C[i+j] += A[i] * B[j]，从中可以明确看到和多面体相关的操作表达。

 

图3. 平铺操作

 

图4 多项式乘法的仿射方言表示

MLIR中引入多面体优化

Polygeist文章沿用LLVM引入Polly，GCC引入GRAPHITE的思路，将开源的Pluto引入到MLIR中实现多面体优化。整体的编译采用LLVM的编译流程，前端Clang分析输入的C语言代码，转换到MLIR中的Affine方言，然后从Affine转换到Pluto，二者之间的交互采用多面体技术中常用的数据格式OpenScop，优化后的代码再次转换到MLIR，然后走LLVM的编译流程。

在整体的转换中，使用的方法就是遍历语法树（AST）。Clang前端遍历Clang的AST语法树，将node映射到MLIR的操作中；从Pluto到MLIR的转换过程中使用CLooG生成初始循环的AST，然后遍历AST中的node，创建与循环和条件判断相对应的MLIR的操作。这种做法可以产生和现有编译器的应用程序二进制接口（Application Binary Interface）兼容的代码，免去重新构建基础设施。

Polygeist出彩的地方的是将C/C++引入到了MLIR，实现了clang前端对MLIR的对接，利用MLIR的IR变换能力对接到开源的Pluto工具进行多面体的优化。从整体的测试效果看，优化的性能没有得到特别显著的提升，原因在于针对benchmark中的实例，现有的编译工具已经优化的很好，比如LLVM对于Clang前端的IR能够减少特定的load操作，而LLVM对MLIR的IR还不支持此优化功能，MLIR还是相对比较新鲜的事物。正如文中所言，使用MLIR生成比多面体工具更加优化的代码不是该工作的兴趣点，而是证实在MLIR中引入多面体优化工具的灵活性。

当然，整体的转换工作不是一蹴而就的，需要涉及到C语言和MLIR之间类型的转换，目前支持的数据类型如图5所示，不支持struct的数据类型，同时也不支持C语言中的break操作。在MLIR中没有指针的概念，关于内存的表示只有memref，MLIR本身也不支持在memref中嵌套memref，为了对接C语言中的指针的指针，文章增加对memref的嵌套使用，也就是修改了原有MLIR的特性。

 

图5 C语言，LLVMIR和MLIR数据类型的对应关系 

整体实现的流程如图6所示，前端搭建了C代码到MLIR的连接，通过遍历Clang的AST语法树，将每个访问的节点映射到MLIR中的SCF或者standard方言中的操作。MLIR起到表达控制流的作用，在方言的表达中直接查找到循环，减少在Pluto的CFG (Control Flow Graph) 中查找loops循环的必要。但是实际上，Pluto还是参与了C/C++中非线性for循环的查找。AST中C语言的数据类型先是转换到LLVM的数据类型，然后转换到MLIR中Standard方言中的数据类型。 

多面体相关的Affine code转换通过识别标识符将#pragma scop和#pragma endscop表示的code直接转化为affine.for 循环。循环约束以affine_map的形式表示，比如(affine_map<()[s0]->(s0)>[%bound])，()中表示的是表示Dimension，[ ]表示的是Symbol，如前文所言。 

前端输入的C语言经过转换到MLIR-SCF方言层级，通过raise-affine的PASS转换为Affine方言，具体实现的功能是将standard 方言中的load, store，SCF方言中的for和if转换到affine方言中对应的操作。利用Pluto的能力进行优化处理，然后再通过low-affine PASS转回到MLIR-SCF，此处借助CLooG进行语法树分析，然后走MLIR的LLVM编译流程。

 

图6 Polygeist的编译流程

结论与思考 

现有的多面体优化的库，比如isl，Polly，主要用于C语言的source-source转换，聚焦于底层级的优化，无法直接用在MLIR的设计中，因为底层级的表示无法还原完整的高层级的表达。Polygeist的工作将MLIR的多面体表示和现有的高层级的优化工具结合起来，采用MLIR Affine方言和OpenScop数据格式的双向转换方便开发者搭建基于MLIR的编译流程，然后使用现有的多面体优化工具优化，最后返回到MLIR进行进一步的转换并最终生成代码。对于启发在于可以在MLIR中引入其他优化工具助力编译优化，根据需求补足MLIR中缺失的能力。 

多面体优化是一项成熟的技术，但也受限于对仿射变换的依赖，对无法进行仿射的循环的优化能力较弱，存在一定的局限性，因此无法在工业界得到广泛应用。同时，多面体优化技术理论相对复杂难懂，从事相关研究的人员较少，难以进行落地。尽管如此，多面体技术在解决特定的问题方面尤其独特的作用，比如在深度学习领域，对多算子融合和多层循环优化方面有着极大的帮助，可以将现有的多面体技术引入到AI编译器中，进行特定功能的优化。 

由于水平有限，文中存在不足的地方请各位读者批评指正，也欢迎大家一起参与讨论。

多模态模型—AI芯片软硬件优化利器

由于芯片算力的提升，基于深度学习的人工智能掀起了计算技术领域的新一轮变革。在深度学习领域，由通用芯片与人工智能领域专用芯片构成的异构架构成为提升深度学习算法性能、实现人工智能目标的基础。然而，目前市场上各处理器厂商自主研发的加速芯片的架构非常复杂，给深度学习算法在芯片上的部署和调度提出了巨大的挑战，如何将深度学习算法有效地部署在当前的异构架构加速芯片上是人工智能和高性能计算领域的一个重要课题。

在这样的背景下，深度学习编译软件栈成为解决这个问题的一种有效手段，包括Google MLIR、TVM、Facebook TensorComprehensions、Intel Stripe、Nvidia Diesel，以及MIT Halide和Tiramisu等。这些工作的一个共同点是都已经或正在借助Polyhedral model（以下简称Poly）来实现算子层的循环优化，这得益于Poly强大的循环优化能力。因此，许多处理器厂商和研究人员开始对Poly产生了兴趣。可惜的是，Poly在国内的发展基本处于空白。所以，在这篇文章里，笔者就自己对Poly的一点了解和前期的一些经验，简单介绍一下Poly的基本原理，并讨论一下Poly在深度学习编译软件栈和AI芯片上发挥的作用，以及一些开放性的问题。不足之处，还望各位一起讨论，多加批评。

将分为三个部分连载，包括：

一. Poly基本原理及卷积分析示例

二. Poly在深度学习领域中发挥的作用

三. AI芯片上利用Poly进行软硬件优化的一些问题

Poly 是并行编译领域的一种数学抽象模型，利用空间几何的仿射变换来实现循环优化。作为一种通用程序设计语言的编译优化模型，Poly在过去30年左右的时间里，从发掘程序的并行性和局部性角度出发，逐步形成了一套稳定、完善的流程，也开始从实验室逐步走向开源软件和应用到商用编译器中。尽管如此，从事Poly研究的团队和社区在国内仍然是凤毛麟角，一方面原因是因为国内科研和企业的价值导向；而更重要的一点，也确实是因为Poly是个小众的研究领域，即便在国际上从事这方面研究的团队和企业也为数不多。

Poly基本原理介绍 

考虑到许多读者可能对Poly并不了解，而且许多Poly文献读起来也比较抽象，先简单介绍一下Poly的工作原理。力图用最简单的代数与几何描述来解释Poly的基本原理。这部分内容参考了文献的图片，通过解读这些图片来解释其中原理。

首先，假设有如图1所示的一段简单的循环嵌套，其中N为常数。循环嵌套内语句通过对A[i-1][j]和A[i][j-1]存储数据的引用来更新A[i][j]位置上的数据。如果把语句在循环内的每一次迭代实例抽象成空间上的一个点，那么可以构造一个以(i,j)为基的二维空间，如图2所示。图中每个黑色的点表示写入A[i][j]的一次语句的迭代实例，从而可以构造出一个所有黑色的点构成的一个矩形，这个矩形就可以看作是二维空间上的一个Polyhedron（多面体），这个空间称为该计算的迭代空间。

 

图1 一段简单的代码示例

 

图2 图1示例代码的迭代空间

可以用代数中的集合来对这个二维空间上的Polyhedron进行表示，即{[i,j]: 1 <= i <= N - 1 与 1 <= j <= N–1}，其中[i,j]是一个二元组，“:”后面的不等式表示这个集合的区间。可以给这个二元组做一个命名，叫做S，表示一个语句，那么这个语句的Polyhedron就可以表示成{S[i,j]: 1 <= i <= N - 1 and 1 <= j <= N – 1}。

由于语句S是先迭代i循环再迭代j循环，因此可以给语句S定义一个调度（顺序），这个调度用映射表示，即{ S[i,j] -> [i,j] }，表示语句S[i,j] 先按i的顺序迭代再按照j的顺序迭代。

接下来，来分析语句和它访存的数组之间的关系，在代数中用映射来表示关系。图1中语句S对数组A进行读和写，那么可以用Poly来计算出S和A之间的读访存关系，可以表示成{ S[i,j] -> A[i - 1, j]: 1 <= i <= N -1与1 <= j <= N- 1; S[i,j] -> A[i,j - 1]: 1 <= i <= N - 1与1 <= j <= N -1 } 。同样地，写访存关系可以表示成{ S[i,j] -> A[i,j]: 1 <= i <= N - 1与1 <= j <= N -1 }。

基于这个读写访存关系，Poly就可以计算出这个循环嵌套内的依赖关系，这个依赖关系可以表示成另外一种映射关系，即{ S[i,j] -> S[i, 1 + j]: 1 <= i <= N - 1与1 <= j <=N - 2; S[i,j] -> S[i + 1, j]: 1 <= i <= N - 2与1 <= j <= N- 1 }。

可以注意到，Poly对程序的表示都是用集合和映射来完成的。当把语句实例之间的依赖关系用蓝色箭头表示在迭代空间内时，就可以得到如图3所示的形式。根据依赖的基本定理，没有依赖关系的语句实例之间是可以并行执行的，而图中绿色带内（对角线上）的所有点之间没有依赖关系，所以这些点之间可以并行执行。但是发现这个二维空间的基是(i,j)，即对应i和j两层循环，无法标记可以并行的循环，因为这个绿色带与任何一根轴都不平行。所以Poly利用仿射变换把基(i,j)进行变换，使绿色带能与空间基的某根轴能够平行，这样轴对应的循环就能并行，所以可以将图3所示的空间转化成如图4所示的形式。

此时，语句S的调度就可以表示成{ S[i,j] -> [i + j, j]}的形式。所以Poly的变换过程也称为调度变换过程，而调度变换的过程就是变基过程、实现循环变换的过程。

 

图3 带依赖关系的迭代空间

 

图4 变基之后的迭代空间

图4中绿色带和j轴平行，这样在代码中表示起来就方便了。说Poly做循环变换的过程就是将基(i,j)变成(i + j, j)的一个过程，也就是说，Poly的底层原理就是求解一个系数矩阵，这个系数矩阵能够将向量(i,j)转换成向量(i + j, j)。

根据这样的调度，Poly就可以利用它的代码生成器，生成如图5所示的代码。此时，内层循环就可以并行了。（注：这里示意的是源到源翻译的Poly编译器，也就是Poly生成的代码还需要交给基础编译器如GCC、ICC、LLVM等编译成机器码才能运行。也有内嵌在基础编译中的Poly工具。）

 

图5 Poly变换后生成的代码

当然，这里举的例子是一个很简单的例子，在实际应用中还有很多复杂的情况要考虑。Poly几乎考虑了所有的循环变换，包括Interchange（交换）、Skewing/Shifting（倾斜/偏移）、Reversal（反转）、Tiling/Blocking（分块）、Stripe-mining、Fusion（合并）、Fission/Distribution（分布）、Peeling（剥离）、Unrolling（展开）、Unswitching、Index-set splitting、Coalescing/Linearization等，图6～8[14]中给出了几种Poly中实现的循环变换示意图，右上角的代码表示原输入循环嵌套，右下角的代码表示经过Poly变换后生成的代码。图中左边的集合和映射关系的含义分别为：J代表原程序语句的迭代空间，S表示输入程序时的调度，T表示目标调度，ST就是Poly要计算的调度变换。

 

图6 Poly中偏移变换示意图

 

图7 Poly中融合变换示意图

 

图8 Poly中平铺变换示意图

深度学习应用的Poly优化

让以图9中所示的二维卷积运算（矩阵乘法）为例来简单介绍Poly是如何优化深度学习应用的。

 

图9 一个2D卷积示例

Poly会将循环嵌套内的计算抽象成一个语句。例如图9中S1语句表示卷积初始化，S2代表卷积归约；而S0和S3则分别可以看作卷积操作前后的一些操作，比如S0可以想象成是量化语句，而S3可以看作是卷积后的relu操作等。

为了便于理解，以CPU上的OpenMP程序为目标对图9中的示例进行变换。Poly在对这样的二维卷积运算进行变换的时候，会充分考虑程序的并行性和局部性。如果对变换后的程序并行性的要求大于局部性的要求，那么Poly会自动生成如图10所示的OpenMP代码；如果对局部性的要求高于并行性，那么Poly会自动生成如图11所示的OpenMP代码。（注：不同的Poly编译器生成的代码可能会因采用的调度算法、编译选项、代码生成方式等因素而不同。）

 

图10 Poly生成的OpenMP代码——并行性大于局部性

 

图11 Poly生成的OpenMP代码——局部性大于并行性

通过对比图10和图11，两种生成的代码采用的循环fusion（合并）策略不同：图10中所示的代码采用了({S0}, {S1, S2, S3})的合并策略，图11中生成的代码则使用了({S0,S1, S2, S3})的合并策略，但是必须通过对S2向右偏移99次、S3向右偏移148次，以及循环层次的interchange（交换）来实现这样的合并。显然，图11所示的代码局部性更好。而并行性上，仔细研究后不难发现，图11生成的代码中，只有最外层c0循环是可以并行的，而图10代码中，S0语句的c0、c1循环都可以并行，并且包含S1、S2、 S3三条语句的循环嵌套的c0、c1循环也都可以并行，相对于图11代码，图10生成的代码可并行循环的维度更多。

当然，在面向CPU生成OpenMP代码时，多维并行的优势没有那么明显，但是当目标架构包含多层并行硬件抽象时，图9中的代码能够更好地利用底层加速芯片。例如，当面向GPU生成CUDA代码时，而图10对应的CUDA代码（如图12所示）由于合并成了两个部分，会生成2个kernel，但是每个kernel内c0维度的循环被映射到GPU的线程块上，而c1维度的循环被映射到GPU的线程上；图11对应的CUDA代码（如图13所示）只有1个kernel，但是只有c0维度的循环被映射到GPU的线程块和线程两级并行抽象上。为了便于阅读，并未开启GPU上共享内存和私有内存自动生成功能。从图中也不难发现，Poly也可以自动生成线程之间的同步语句。（注：图中循环分块大小为32，图12中线程块上线程布局为32*16，图13中为32*4*4。）

 

图12 Poly生成的CUDA代码——并行性大于局部性

 

图13 Poly生成的CUDA代码——局部性大于并行性

值得注意的是，为了充分挖掘程序的并行性和局部性，Poly会自动计算出一些循环变换来实现有利于并行性和局部性的变换。例如，为了能够达到图11和图13中所有语句的合并，Poly会自动对S2和S3进行shifting（偏移）和interchange（交换）。

7.4.基于MLIR实现GEMM编译优化

摘要

GEMM(General Matrix Multiplication)即通用矩阵乘法运算，由于其计算行为具有一定的复杂性以及规律性，是编译算法研究的绝佳场景。MLIR是近期非常热门的一个编译器软件框架，是工业界及科研界研究的一个热点，其提供了一套灵活的软件基础设施，对中间表达式（IR）及其相互之间的转换进行规范的管理，是一个非常友好的编译器开发平台。即是分析在MLIR框架下，实现GEMM优化的内容，以及对MLIR在这一方面的实现优势的讨论。

GEMM优化策略介绍

矩阵乘法运算，由于其过程会包含大量的乘加操作，并且伴随大量的数据读写，因而如何充分利用好底层硬件的存储资源以及计算资源，是编译器对其性能优化的关键。目前，已有的一些优化策略主要包括：

1.矩阵分块（Tile）

当前的处理器性能主要受限于内存墙，即计算速度要大于数据存储的速度。为了打破内存墙的约束，各类硬件包括CPU及其他专用处理器，会设置不同层次的存储单元，而这些不同层级的存储单元往往大小以及读写速度不同，一般越靠近计算单元的存储其存储容量也越小但访问的速度也越快。如果可以将计算过程的数据局部化分块，而这些分块的数据可以独立完成计算，那么分块的数据就可以放在层次化的存储中，然后通过不同存储间建立Ping-Pong的数据传输方式，将数据存储与计算解耦，从而可以有效得隐藏存储墙的问题，提高计算效率。矩阵运算就有这种特点，因而可以通过矩阵分块来加速运算，如下图1所示，假设有两层存储，将输入矩阵A和B，以及输出矩阵C，根据存储大小划分成相应的小块，即m->mc，n->nc，k->kc，每次将Ac(mc, kc), Bc(kc,nc), Cc(mc, nc)送入到离计算单元更近的存储模块内，完成局部的计算后再进行下一次的计算。

 

图1 矩阵运算的瓦片操作示意图

在不同的底层硬件中，由于存储的层次以及不同层次的存储的容量大小不一样，分块的大小也会不一样。比如，对于CPU而言，（Ac, Bc, Cc）划块的大小与cache大小一致，而为了充分利用register的资源，还会对（Ac, Bc, Cc）再进一步细划块成（Ar, Br, Cr），其尺寸大小与寄存器的数量一致。

2.向量化（Vectorize）

向量化的操作，主要是利用硬件的向量化指令或者SIMD（单指令多数）指令的特性，实现一个指令周期对多个值操作的能力。如下图2所示，通过将4个数据组成向量，利用处理器可以处理4个元素的新向量的计算能力，可以将4个指令周期的处理时间，压缩成1个指令周期的处理时间，从而极大提高运算处理能力。

 

图2.向量化操作示意图

3.循环展开（Unroll）

由于矩阵乘法有多层循环构成，如果底层硬件有一定的并行化能力，包括多线程多进程处理能力等，那么可以对循环进行适当展开，从而提高计算的并行度，增加并发执行的机会。如下图3所示，将一个次数为1024的循环，展开成256次循环，新的循环内又包含4条可以并行执行的展开计算，如果处理器能够并行处理循环内部的展开计算，那么通过对原来的循环展开，可以获得接近4倍的性能提升。

 

 图3.循环展开操作示意图

矩阵乘法的运算也包括其他的优化策略，比如数据重排等，但总体而言，各类编译器都是利用这些策略，充分利用硬件的存储及计算资源，达到最佳的运算效率。一些传统的计算库，如：OpenBLAS, BLIS, MKL等，开发时间长，性能也有比较优秀的表现。

MLIR实现GEMM优化

MLIR基于多层中间表示的方言（Dialect）思想，提供了一整套完善的编译器基础框架，可以帮助开发者快速实现编译策略想法的编译器。分析GEMM运算在MLIR中的实现，对应的硬件Target是因特尔i7-8700K处理器，每个核包含有32/256KB L1/L2 Cache以及多核共享的12MB L3 Cache，处理器支持AVX-2指令（256bit),优化目标是一个2088x2048xf64与2048x2048xf64的矩阵乘。

首先，其在高层次的Dialect上定义了一个矩阵运算的算子，这个算子的参数包含了输入矩阵（A，B）以及输出矩阵（C），同时为这个算子添加了tile/unroll 的尺寸等属性。如下图4所示，其中（M_C, K_C, M_R, N_R）属于Tile尺寸，K_U属于Unroll的大小。这里面（M_C, K_C）的选择是使得M_CxK_C大小的A矩阵块能够在L2 cache中复用，（K_C, N_R）的选择是使得K_CxN_R大小的B矩阵块能够在L1 cache中复用，（M_R, N_R）的选择是使得M_RxN_R大小的输出矩阵块能够在CPU Register中复用，这些值是根据硬件计算或者tunning出来的，在这里面的测试取了一个经验值。这些属性可以协助转换到更低一层的算子的策略实现，而选择哪些属性，则是跟编译算法以及编译的底层硬件对象有关，这些属性也是协助转换成下一层跟硬件更贴近的中间表示的实现，因而可以根据实际需要，灵活使用。

 

图4 GEMM算子的高层次定义 

其次，MLIR的特点就是通过统一的多层中间表示，来实现对算子的层层低层化（lower）到具体的硬件目标上。针对上述高层次上定义的矩阵乘法算子，通过利用其所携带的优化属性，以及底层硬件的特点，设计了多条转换的路径（Pass），从而进一步把该算子lower到MLIR框架提供的中间辅助层（此中选择了Affine, Linalg，和Standard Dialect）。在这一层的转换过程中，基本包含了所有的策略，如：Tile，定制化复制，unroll，vectorize等。然后再将中间的辅组层的Dialect，进一步lower到LLVM Dialect上，如下图5所示。

图5 GEMM算子下译的层次化Dialect示意图

最后，通过mlir提供的mlir-cpu-runner工具，可以运行最后生成的LLVM Dialect的结果。总体优化及运行测试的命令，如下图6所示。其中，-hopt，-hopt-vect等，是从高层的算子（hop.matmul）到中间辅组层的转换路径，每一条路径都包含有相应的编译策略，可以根据需要，灵活添加以及改变，-convert-linalg-to-loops, -lower-affine等时中间辅助层之间的转换，最后转换成LLVM Dialect。 

 

图6 MLIR运行GEMM的命令示意图

总体上，一个GEMM运算通过在MLIR框架下实现，并重写优化策略的路径，可以得到如图7所示的结果，其中箭头1对应包含了所有重写优化策略的MLIR实现，可以看到其能达到的计算速率为61.94GFLOPS，离理论上的峰值计算速率（75.2GFLOPS）比较接近，跟传统的计算库相比（如：BLIS，OpenBLAS，MKL等），也有着可以媲美的结果，其中的差距可能是传统的计算库有tunning的机制以及在编译器后端生成汇编指令及机器码有更成熟且高效的优化，因而可以得到更好的优化结果。总体而言，用MLIR重写的GEMM优化算法有着非常良好的表现。

 

图7.MLIR编译运行结果与其他计算库的对比示意图

另一方面，MLIR框架提供了非常完善的C++以及Python接口，因而可以很方便接入已有的计算库，进行联合优化。尝试了用MLIR+BLIS的方法，将MLIR放在外侧（提供手动优化功能），BLIS则作为micro-kernel放在内侧（提供auto tunning功能），最终的结果如图7中箭头2所示。可以看出，对于DGEMM（双精度），通过MLIR与BLIS的联合优化，也可以达到接近峰值的性能，而其性能要比单独的MLIR或者BLIS优化要差一点。但其实在SGEMM（单精度）的测试中，MLIR+BLIS的优化又要比单独的MLIR或者BLIS优化要好一些，因而其中的性能在差异还需要进一步分析。总体而言，MLIR提供了非常完善的支持，可以融合已有的计算库的优化算法，去实现联合的编译优化。

MLIR实现GEMM优化的优势 

通过上面对MLIR实现GEMM优化算法的编译的介绍，可以看出MLIR在其中有着非常突出的优势。

首先，MLIR框架提供了非常完善的编译器基础设施，可以让开发者不需要花费太多精力在编译器周边的实现，从而更加专注于编译算法的开发。同时，其基于多层中间表达的方式，可以让编译器更加模块化，可以让编译算法利用不同层次的中间表达的抽象信息，在不同的层次中逐步具体化，从而使得算法实现更加层次化，更加易于实现及管理。

其次，MLIR框架提供了一直到最底层硬件的表示支持，由于其可以层次化在不同的中间表示层实现编译算法，可以在高层次的中间表示中实现不依赖于底层硬件的通用算法，而在接近硬件底层中，开发针对性的路径实现相应的编译算法，因而可以很方便地针对不同硬件目标开发统一的编译环境。本人认为，这也是MLIR相对于一些现有的AI编译器，如：TVM等，最有优势的地方之一，由于其框架可以根据需要自行扩展Dialect，同时这些Dialect又在系统中遵循一套统一的范式进行管理，因而对不同的编译目标（硬件target）会有很强的扩展性，同时编译器的工程管理又可以做到非常好的统一性。

另外，MLIR框架提供了完善的C++/Python接口，可以很方便地接入已有的优化算法，快速实现算法迁移。

结论与思考

主要介绍了矩阵乘法运算在MLIR编译器框架实现的主要过程及内容，以及其在优化算法的多层次实现，以及接入已有优化算法的能力等方面的优势。MLIR编译框架，为编译器的开发者，在多层中间表达转换以及算法优化等方面提供强大的基础设施支持，降低开发编译器的门槛。

7.6.MLIR编译技术应对硬件设计挑战？—初探CIRCT项目

  后摩尔时代，芯片架构向专用化的方向发展，打造一套统一的编译器变得势在必行，顶层抽象，底层优化，然后映射到各个不同的专用芯片；同理，在芯片设计领域，顶层支持多种硬件表达，底层抽象、优化、映射到不同的芯片架构变得十分可能。CIRCT项目正是基于MLIR编译技术打造的全新的统一的EDA (Electronic Design Automation)框架，类似于软件中的TensorFlow或者PyTorch，旨在为EDA设计工具提供完整的灵活可配的基础设施，通过层层抽象的方式在一套框架中支持多种硬件设计模式。同时，CIRCT项目与MLIR相结合，做到硬件设计和软件编译的同源，实现软件生态和硬件设计生态的一致化。

CIRCT背景介绍

CIRCT项目的全称是Circuit IR Compiler and Tools，基于电路中间表示的编译器和设计工具，将DSL/IR/Compile等软件开发的思想应用到开源硬件设计领域（可参考知乎文章Google的『泛芯片』科技蓝图），加速硬件设计的流程，同时也寻求解决EDA工具的零碎化及封闭化的缺陷，培育统一的硬件设计社区。目前，CIRCT项目处于探索阶段，主要集中于底层抽象IR的定义和开发。

在目前的EDA开源工具方面，有不少出色的工作，比如综合工具Yosys，Verilog仿真工具VERILATOR, 新的硬件设计工具CHISEL, 基于Python的硬件设计工具Magma等等。但是这些开源硬件设计工具独自成生态，彼此独立不兼容，导致了设计工具的碎片化。如果将这些设计力量统一起来打造一个统一的开源的硬件设计工具，其前景令人充满想象。为此，Christ等人推出了基于MLIR的开源的EDA框架CIRCT，其核心思想是将软件的编译模式搬移到硬件设计中，顶层重用MLIR的高级抽象IR，底层重新引入新的描述电路IR（Verilog/VHDL也是描述电路的IR），打造具有模块化特点的基础设施，能够更好地适配到各种硬件设计中，比如CPU，GPU，CPU+xPU等。

在ASPLOS 2021大会上，Chris Lattner博士做了The Golden Age of Compiler Design in an Eraof HW/SW Co-design的演讲，阐述了异构计算黄金时代的今天，也迎来了编译器的黄金时代。文中提到了针对新硬件架构下的硬件设计工具，需打破现有的碎片化状态，打造统一的硬件设计框架。

在机器学习硬件领域，向着专用化程度越来越高的方向发展，如图1所示，这样将硬件设计推向了两个维度，一个是固有的处理器，如CPU及GPU等；一个是专用化的处理器，如FPGA及ASIC等。在这两个维度中，分散着针对不同环节的开发工具，如图2所示，而CIRCT项目目的在于统一这些零碎的工具，将其融合到一个统一的EDA框架中。从MLIR整个大的生态来看，基于MLIR的软件框架可以涵盖软件开发工具，如图2中红线的左半部分；CIRCT可以涵盖硬件设计工具，如图2中红线的右半部分。MLIR将软件设计和硬件设计彼此关联起来。 

 

图1 机器学习硬件分布图 

 

图2 软件开发工具及硬件设计工具

CIRCT项目每周举行一次例会，邀请各个领域的技术达人进行分享，努力从更多的开源项目中汲取力量，这也就是集百家工具之所长。

CIRCT软件框架的介绍

CIRCT软件栈如图6所示，这是目前构想的软件栈，尚未开发完成。软件栈分为两部分，蓝色部分表示的是MLIR的基础设施，灰色部分表示的是CIRCT项目独有的基础设施，从图中可以看出整个MLIR的基础设施具有非常好的复用性，同时整个软件栈也体现了MLIR层层下译的特点。主要关注于灰色部分，灰色部分也包含目前CIRCT中的所有方言，但这些方言目前处于相对零碎的状态，没有统一起来。

 

 

图3 CIRCT项目的软件栈

1、Handshake。Handshake方言主要描述独立的非同步的数据传输，数据传输的通道基于FIFO实现。Handshake方言中定义了握手模型和控制逻辑，比如fork，join，mux及demux等，如图4所示，通过这些抽象的操作可以表示出程序中常用到的for循环逻辑，如图5所示。该项工作主要由Xilinx的研究人员开发。 

 

 

图4 Handshake方言中的主要操作

 

图5 Handshake表示for循环结构

2、HIR: An MLIR-based Intermediate Representation for Hardware Accelerator Description。HIR是最新加入到CIRCT项目中，由印度理工的研究人员开发，主要在高层次描述时序电路的行为，比如带有延迟的计算、流水线、状态机，同时也有具体调度行为和并行计算的表示等。延迟信息的表示可以参考图6中的Op: hir.delay. 

 

图6 HIR中矩阵乘的表示 

3、ESI (Elastic Silicon Interconnect): 主要是表示片上和片外的互联关系，用于实现点对点双向的数据流通路，数据流通路中含有数据类型，数据类型支持所有MLIR中定义的类型。这些数据类型可以通过接口提供给软件的API，high-level的开发者可以直接获取到数据通路上的数据类型，打破了硬件描述语言中的wire-level protocol和high-level编程的统一。ESI项目源自于微软，目前是CIRCT项目中的一个方言。

4、FIR-RTL选用来自Chisel编译器的IR，是比Verilog更高一级的表达。

5、RTL/Seq/Comb RTL用于描述module和端口，Seq描述时序电路，comb描述组合电路，目前处于开发阶段。

6、LLHD LLHD是受LLVM启发开发的针对硬件设计RTL代码的多层级IR表示。整个LLHD工程分为行为级的表示和结构级的表示。构想的设计中前端支持目前多种硬件语言，比如Chisel，MyHDL，SystemVerilog等，后端可以加入不同硬件厂商的综合工具，最终输出网表级别的IR，设计的流程如图9所示。另外，LLHD对于CIRCT的价值还在于能够基于LLHD的IR实现仿真。在目前的实现中，LLHD的前端只支持SystemVerilog。

 

图7 LLHD设计的flow

CIRCT方言转换 

图8表示了各个方言之间的转换关系。从图中可以看出，目前支持两条到Verilog的转换路径：

 

图8 各个方言之间的转换关系

1、输入为FIR或者Handshake：具体的转换路径为FIR -> high FIRRTL -> mid FIRRTL -> low FIRRTL -> Verilog

2、输入为SystemVerilog：具体的转换路径为SystemVerilog -> behavior LLHD -> structure LLHD ->LLHD-sim

 

CIRCT项目工程中test文件夹提供了查看各个dialect之间转换的窗口，下面以conversion和EmitVerilog两个文件为例介绍Handshake->FIRRTL->Verilog的转换路径。

1、Conversion: 测试各个方言之间的转换。以Handshake转换为FIRRTL代码为例，将Handshake中的branch操作转换为FIRRTL的实现。图9中上半部分代码为Handshake的表达，%arg0表示为输入的数据，%arg1表示为输入的控制逻辑，如果控制逻辑为false，返回%arg0以及对应的控制逻辑；下半部分代码为FIRRTL的表达。执行命令为：circt-opt-lower-handshake-to-firrtl-split-input-file test_branch.mlir。 

 

 

 

图9 Handshake到FIRRTL的转换

2、EmitVerilog： 将各个方言转换为Verilog代码输出。图10表示由FIRRTL到Verilog代码的转换。执行命令为：circt-translate firrtl-dialect.mlir -emit-verilog。 

 

图10 FIRRTL到RTL代码的转换

总结与思考 

CIRCT项目是MLIR生态中一个非常具有启发性的项目，将编译技术应用到硬件设计中，颠覆了对传统硬件设计工具的认识，寻求利用编译技术统一硬件设计工具，打造EDA框架。同时，CIRCT项目也为MLIR的生态拼上了硬件设计的板块，让人对MLIR整个生态充满期待，但不得不说CIRCT项目宏大，充满挑战，目前还处于起步阶段，聚焦在底层IR的抽象表达，全线打通还有很长的路要走。当下，对于电路IR的抽象存在各种挑战，新的IR抽象是否能够完备地描述电路行为，是否能够解决现有Verilog在硬件描述方便的缺陷，比如学习门槛高，代码跟踪困难等局限；未来，能否将芯片验证及芯片设计方法学引入到CIRCT中，能否加快芯片设计和验证的周期。CIRCT项目征途漫漫，令人期待。

7.7.基于IREE HLO项目看MLIR编译器实现的过程及优势

MLIR(Multi-Level Intermediate Representation）是谷歌团队开发的一个非常热门的开源编译器框架，其提供了一套灵活的软件基础设施，来规范中间表达式（IR）及其相互之间的转换，建立了一个非常友好的编译器开发平台，一些比较好的对MLIR框架解读可以参考。IREE项目也是谷歌推出的，基于MLIR进行统一的IR规范管理的一个从tensorflow模型到不同硬件加速器平台的端对端项目，其前端的核心内容就是对XLA对应的HLO算子进行MLIR实现。MLIR在编译器实现方面，有一些非常优秀的特点，即是基于HLO项目对此做一个抛砖引玉的简单讨论。 

IREE HLO项目介绍

MLIR HLO项目是IREE的重要的组成部分，IREE项目的IR结构如图1所示。MLIR HLO的主体其实就是XLA（Accelerated Linear Algebra）编译器高性能算子组成的IR表示，其Operation基本上是和XLA中相对应的HLO是一一对应关系。MLIR HLO部分的输入是由上端translator工具将TFgraph parse成TF IR并转换成相应的HLO IR，并在其内进行一定的处理，包括Dialect之间的转换，以及主要继承于XLA相应的算法处理，最后输出或者下译到后端的硬件target相对应的IR。

 

图1：IREE项目MLIRIR示意图

MLIR HLO项目的实现内容包含了许多工作，其主要包括：1.不同Dialect的定义，及其算子（Operation）内容的定义；2.不同Dialect之间算子转换；3.Dialect内或之间的算子链的算法处理（比如：多面体处理，融合操作）。接下来内容，即从这三方面进行分析讨论。

定义Dialect算子内容

一般编译器IR定义的算子，其内容主要包括：1.输入变量；2.输出的数据类型及形状；3.其他附属信息（用来表征一些优化信息或者高层次信息）。而如果不同层次的IR去进行独立的开发，可能会使得IR的API混乱，并且部分的软件基础设施也会重复开发。而MLIR的Dialect可以很好解决这些问题，每一层的IR可以包含一种或者多种Dialect表示，而每一个Dialect包含各种自定义的算子（Operation）。MLIR的Dialect有比较统一的写法格式，并且由其提供的table_gen工具，最终可以为Dialect的各个算子生成统一的C++接口。

 

 

图2：HLO Dialect算子的定义

一个典型Dialect算子的table_gen的定义方式，如图2所示，主要包含两部分：1.特殊属性部分（traits，图2黄色部分）；2.以及主体部分（图2蓝色部分）。特殊属性部分定义了算子的特殊操作类，主要以interface方式呈现，在特定算子组合的优化方面有着重要的作用。主体部分即是MLIR定义的统一属性部分，其包括：arguments，results等内容。其中arguments的内容，又可以根据需要，通过在build函数里面的addOperands及addAttribute函数（如图3的伪代码），被分配到Operands（相当于op的input）或者Attributes属性里面。Operands，Attributes，Results这三个属性，在MLIR软件框架里，又为其提供了不同的访问函数（参见后面算子链的算法优化实现关于基于operand的get_using函数介绍），可以极大提高对算子转换及优化实现的效率。 

 

图3：build函数的伪代码 

总体而言，MLIR在这一方面，为算子的访问，表示，及验证等方面提供了非常完整的基础设施，使得开发人员可以专注于算子的内容定义，而不用花很多的精力去关心算子的周边接口表示及其他内容的实现。 

实现Dialect间算子转换

Dialect之间的下译过程，即是对其相应的算子进行转换过程。MLIR框架中的算子，包含了丰富的接口函数，可以利用这些接口函数，可以快速得到需要转换的算子的信息，包括：1. 函数getOperand得到operand内容；2.函数getAttribute得到所有的Attribute内容；3.函数getResult得到result内容；以及其他函数得到相对应其他的内容。再根据这些内容，计算出新的一个或者多个operations所需要的operands/attributes及其他的内容信息，并组装成新的operations。具体流程图如图4所示，包括：

 

图4：不同Dialect的算子进行转换的过程 

1、对于旧的operation，通过相应的get函数得到attributes, results, operands等信息；

2、步骤1中得到的信息通过一定的转换，得到新的attributes, operands等信息；

3、创建新的operations，并将转换后的信息填充到这些operations内，同时移除旧的operation。

总体而言，MLIR在这一方面，也是提供了比较完备的API访问接口，以及算子创建及移除机制，可以让开发人员更加专注于转换的算法实现，而省去转换过程中周边的一些适配工作。

 

实现算子链的算法优化

 

在对算子从高层次到低层次转换过程中，对一些算子根据计算特点进行算法优化以达到最好的性能效果，也是编译器重要的工作之一。MLIR HLO项目中，其算法优化主要继承于XLA HLO项目中已开发的一些优化内容。其中非常重要的一个内容是在HLO方言层面，对算子进行融合处理，即kLoop，KInput算法处理。图5即是一个kLoop的一个例子，其主要的思想是将所有由相同输出形状的连接算子组成的算子链（目前只支持element-wise性质的算子），融合到一起进行运算。

 

 

 

图5：MLIR HLO融合例子，上图是输入的mlir，下图是融合的结果

对于该例子，其整个算法的流程示意图，如图6所示，主要的内容包括：

 

图6：MLIR HLO融合过程的示意图（对应图5中的例子）

1、基于MLIR形式的op_list信息，去建立一个有向的连接图；

2、通过深度优先算法（DFS，Depth First Search），对有向连接图进行查找，通过一定的判断条件，得到并融合可以融合的pattern；

3、不断迭代上一步的步骤，直到遍历所有的节点，最后得到最终的融合结果（图5中的例子，最终是将整个的operation融合成一个大的融合操作）；

4、将融合结果的有向连接图和未参与融合的其他部分内容连接在一起，并输出新的MLIR形式的op_list列表。

首先，在建立有向连接图过程中，其主要内容是获取每个节点的输入及输出节点（即：in和out），MLIR对此提供了两个函数：

• operand.getDefiningOp()：可以获取operand所对应的输入operation节点；
• result.getUses()：可以获取result所对应的输出operation节点。

基于上述两个功能函数，可以非常方便得获取有向连接图节点的输入及输出（in和out）,建立算子的有向连接图。

其次，MLIR的operation定义有特殊属性部分内容（即traits，主要以interface形式呈现），这些内容可以在某些特定的实现方面提供支持。HLO项目定义了InferFusibilityOpInterface这样的一个特殊属性接口，该属性接口包含了多个功能函数，都是用来判断两个连接的operations是否可以fuse，而需要参与融合算法的operations则需要添加该特殊属性。HLO 融合算法判断的过程如图7所示。判断的条件包括：operation本身是否是可以融合的，operation与其输入节点（withOperand）或者输出节点（withConsumer）是否可以融合，两个operations的输出形状是否相等。这样，通过这些特殊属性，可以很方便建立operation的判断过程。

 

图7：融合的判断过程示意图

最后，将融合完成后的有向图，恢复到MLIR的op_list形式，MLIR也提供了相应的一些功能函数，极大地方便了这一过程的实现。

总体而言，MLIR在对新算法的开发过程，其提供的基础设施框架可以减少周边适配的工作，让开发者更多专注于算法本身的实现，而对已有算法的移植也提供了非常友好的兼容性。 

结论与思考

通过结合MLIR HLO项目，分析了基于MLIR框架编译器的一些关键点的代码实现以及优势。总体而言，基于MLIR的框架，为编译器多IR的实现提供了丰富且统一的软件API接口，以及其他非常高效实用的基础设施软件框架支持，可以让常规的编译器开发更为便捷和高效。当然，也有一些观点认为MLIR并没有解决一些编译器的根本问题，比如auto-tiling及auto-schedule等，而且过多的Dialect表示层级划分反而可能加剧IR层级的碎片化问题，同时高层次的IR表示存在算子表述的不完备性等问题。但本人认为，这些也都是传统编译器的老大难问题，MLIR则只是提供了一个非常强大且完备的框架，可以为算子的表示，转换以及算法优化等方面提供强大的基础设施支持，降低开发编译器的门槛。本人也希望通过，可以让读者对MLIR的工程实现过程更加清晰一些，从中得到一点启发。

 

参考文献链接

https://www.birentech.com/Research_nstitute_details/18087818.html

https://www.birentech.com/Research_nstitute_details/18087805.html

https://www.birentech.com/Research_nstitute_details/18087790.html

https://www.birentech.com/Research_nstitute_details/18087956.html

https://www.birentech.com/Research_nstitute_details/18087965.html

https://www.birentech.com/Research_nstitute_details/18087976.html

https://www.birentech.com/Research_nstitute_details/18087816.html