基于IREE HLO项目看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

上一个: 图神经网络加速器深度调研(上)

下一个: 片上系统设计案例分析—Xbox主机

近期文章

通用AI模型的未来:深度强化学习(deep reinforcement learning)

近年来,AI模型开始涌现出超越人类的潜力,在传统的围棋游戏以及拥有复杂规则和系统的电竞游戏(星际争霸2,Dota 2等)中都有体现。随着ChatGPT的出现,人们开始意识到语言模型成为通用人工智能的可能性,而这些模型的核心都是深度强化学习。

2023-05-08

“为了全人类更好地交流”:通用语音识别

人机交互的第一步往往由人发起,从你说出第一句话开始,计算机如何可以应答并能进而和你自然畅谈?在之前文章里我们分享了当前利用AI进行语音识别的关键步骤和做法,这次我们将随着技术的进化发展,领略当前通用语音模型的进化高度。

2023-04-24

查看更多