IREE HLO与MLIR编译器

IREE HLO与MLIR编译器

MLIR(Multi-Level Intermediate Representation）是谷歌团队开发的开源编译器框架，提供了一套灵活的软件基础设施，以便规范中间表达式（IR）及其相互之间的转换，建立了一个友好的编译器开发平台，一些比较好的对MLIR框架解读可以参考。IREE项目也是谷歌推出的，基于MLIR进行统一的IR的从Tensorflow模型到不同硬件加速器平台的端对端项目，其前端的核心内容就是对XLA对应的HLO算子进行MLIR实现。 

IREE HLO项目介绍

MLIR HLO项目是IREE的重要的组成部分，IREE项目的IR结构如图1所示。MLIR HLO的主体是XLA（Accelerated Linear Algebra）编译器高性能算子组成的IR表示，其操作是与XLA中相对应的HLO。MLIR HLO的输入是由上端翻译工具将TFgraph解析成TF IR，并转换成相应的HLO IR，并进行一定的处理，包括方言间的转换，以及继承XLA相应的算法处理，最后输出或者下译到后端的硬件对应的IR。

 

图1.IREE项目MLIR示意图

MLIR HLO项目的实现内容包含了许多工作，其主要包括：

1）不同方言、算子的定义；

2）不同方言之间算子转换；

3）方言内或之间的算子链的算法处理（比如：多面体处理，融合操作）。

接下来，从这三方面进行分析讨论。

定义Dialect算子内容

一般编译器IR定义的算子，其内容主要包括：

1）输入变量；

2）输出数据类型及形状；

3）其他附属信息（用来表征一些优化信息或者高层次信息）。而如果不同层次的IR去进行独立的开发，可能会使得IR的API混乱，并且部分的软件基础设施也会重复开发。而MLIR的Dialect可以很好解决这些问题，每一层的IR可以包含一种或者多种Dialect表示，而每一个Dialect包含各种自定义的算子。MLIR的Dialect有统一的格式，并提供table_gen工具，最终可以为方言算子生成统一的C++接口。 

 

图2. HLO Dialect算子的定义

一个典型Dialect算子的table_gen的定义方式，如图2所示，主要包含两部分：

1）特殊属性部分（特征，图2黄色部分）；

2）主体部分（图2蓝色部分）。

特殊属性部分定义了算子的特殊操作类，主要以接口方式呈现，在算子优化方面有着重要的作用。主体部分即是MLIR定义的统一属性部分，其包括：参数、结果等。其中参数，通过在构建函数里面的addOperands及addAttribute函数（如图3的伪代码），被分配到操作数（相当于op的输入）或者属性里。操作数、属性、结果这三个属性，MLIR提供了不同的访问函数，可以极大提高对算子转换及优化实现的效率。 

 

图3.构建函数的伪代码

总之，MLIR为算子的访问、表示、验证等方面提供了完整的基础设施，使得开发人不用关心算子的接口表示及其他实现。 

实现Dialect间算子转换

Dialect之间的下译过程，即是对算子进行转换过程。MLIR框架中的算子，包含了丰富的接口函数，可以快速得到需要转换的算子的信息，包括：

1）函数getOperand得到操作数；

2）函数getAttribute得到所有的属性；

3）函数getResult得到结果；以及其他函数得到相对应其他的内容。再根据这些内容，计算操作数/属性及其他的内容信息，并组装成新的操作。具体流程图如图4所示，包括： 

 

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

1）对于旧的操作，通过get函数得到属性、结果、操作数等信息；

2）通过一定的转换，得到新的属性、操作数等信息；

3）创建新属性，并将转换后的信息填充到这些操作内，同时移除旧的操作。

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

实现算子链的算法优化

在对算子从高层次到低层次转换过程中，对一些算子优化以达到最好的性能。MLIR HLO项目中，算法优化主要继承XLA HLO项目中已开发的一些优化内容。其中一个内容是在HLO方言层面，对算子进行融合处理，即kLoop，K输入算法处理。图5即是一个kLoop的一个例子，主要是将所有由相同输出形状的连接算子组成的算子链（目前只支持元素级性质的算子），融合到一起进行运算。 

 

 

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

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

 

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

1）基于MLIR的op_list信息，建立一个有向连接图；

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

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

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

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

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

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

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

 

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

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

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

小结

通过结合MLIR HLO项目，分析了基于MLIR框架编译器的一些关键点的代码实现以及优势。基于MLIR的框架，为编译器多IR的实现提供了丰富且统一的软件API接口，以及其他非常高效实用的基础设施软件框架支持，可以让常规的编译器开发更为便捷与高效。但是，MLIR并没有解决一些编译器的根本问题，比如自动平铺及自动调度等，而且过多的Dialect表示层级划分，反而可能加剧IR层级的碎片化问题，同时高级IR表示存在算子表述的不完备性等问题。这些也都是传统编译器的老大难问题，MLIR提供了一个非常强大且完备的框架，可以为算子表示、转换以及算法优化等方面提供强大的基础设施支持，降低开发编译器的门槛。