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

IREE HLO项目介绍

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

图1：IREE项目MLIRIR示意图

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

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

图4：不同Dialect的operation进行conversion的过程

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

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

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

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

实现算子链的算法优化

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

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

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

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

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

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

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

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

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

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

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

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

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

图7：fusion的判断过程示意图

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

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

结论与思考

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

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

参考文献

1、Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infrastructure for the end of moore’s law, 2020

2、MLIR：https://mlir.llvm.org/

3、Chris Lattner, "Thoughts on Tensor Code Generation in MLIR", 2020-01-24, Talks and Related Publications - MLIR (llvm.org)

4、https://www.tensorflow.org/mlir/xla_gpu_codegen