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

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

  • 发布时间:2021-04-10 11:03
  • 访问量:

【概要描述】MLIR是近期非常热门的一个编译器软件框架,是工业界及科研界研究的一个热点。本文结合IREEHLO项目,对MLIR项目实现的过程及优势做一个总结分析,希望对MLIR的使用有一定的参考作用。

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

【概要描述】MLIR是近期非常热门的一个编译器软件框架,是工业界及科研界研究的一个热点。本文结合IREEHLO项目,对MLIR项目实现的过程及优势做一个总结分析,希望对MLIR的使用有一定的参考作用。

  • 分类:研究院
  • 作者:壁仞科技研究院
  • 来源:
  • 发布时间:2021-04-10 11:03
  • 访问量:
详情

摘要

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

近期文章

流体力学与物理导引神经网络

流体力学与物理导引神经网络

现代科学中众多复杂的关键应用依赖着对流体运动的精确预测。然而,由于流体力学方程数值求解的复杂性,基于传统数值方法的高分辨率长周期数值模拟运算量大且难以保证数值稳定性。本文将深入介绍PINN方法在流体计算中的应用并分析其与传统数值方法的差别。
2021-10-25
物理导引神经网络方法分析

物理导引神经网络方法分析

随着GPU能力的提升,支撑深度学习的软硬件生态得到了快速发展。通过深度学习来解决科学计算问题成了一种趋势,其中用深度学习来求解偏微分方程的方法也逐渐兴起。尤其引人注意的是一种称为物理导引神经网络的方法,其为科学计算领域注入了新的活力。
2021-10-18

联系我们

这是描述信息

招贤纳士

公众号二维码
这是描述信息

 版权所有   ©   上海壁仞科技有限公司   |    沪ICP备19047354号