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

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

  • 发布时间:2022-01-03 07:00
  • 访问量:

【概要描述】软硬协同设计是AI芯片设计中关于性能分析及优化的重要手段,本文将探讨将现有的架构探索的方法放在MLIR框架下实现的可行性,从而实现基于MLIR的软硬协同设计。

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

【概要描述】软硬协同设计是AI芯片设计中关于性能分析及优化的重要手段,本文将探讨将现有的架构探索的方法放在MLIR框架下实现的可行性,从而实现基于MLIR的软硬协同设计。

  • 分类:研究院
  • 作者:
  • 来源:
  • 发布时间:2022-01-03 07:00
  • 访问量:
详情

摘要

 

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

 

架构探索方法的介绍

 

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

 

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

 

https://mmbiz.qpic.cn/mmbiz_png/O9VicLCaMLWXic7Yg0UxXmQUYFSGkzj6UQCP7V7sskibUPKrWmIqRtAz4695G1JKibGS8CUkc4qGfsUytZBehreribQ/640?wx_fmt=png

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

 

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

 

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

 

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

 

 

关于时间和空间数据复用的表达,文中提出了一种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 时域和空域映射与循环表达之间的对应关系[4]

 

假设一个计算架构有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[6]是较早开展基于编译的方式进行硬件的性能分析,将性能分析提前到RTL代码之前,避免了RTL代码及C-model大量的开发工作,基本的思路是将计算任务lowering到动态数据依赖图(DDDG:Dynamic Data Dependence Graph)级别,DDDG是针对特定架构的中间表达(Intermediate Representation)的表示,如图5所示。针对特定的硬件架构,分析DDDG的动态执行过程,即可评估出性能和功耗的数据,他们基于ILDJIT compiler IR[7]。

 

图5 DDDG的计算表示[6]

 

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

 

Interstellar[8]是将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, splitreorder调度原语(scheduling primitives)的IR表示;标签3表示调度原语和硬件架构的对应关系[8]

 

架构级别的IR

 

Micro-IR[9]文章的核心思想是将加速器的架构表示为一个并发的结构图(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的编译流程[9]

 

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

 

图8 Micro-IR的编译表示[9]

 

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

 

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

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

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

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

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

 

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

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

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

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

 

结论与思考

 

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

 

参考文献

[1] Y. Chen, J. Emer, and V. Sze, “Eyeriss: A spatial architecture for energyefficient

dataflow for convolutional neural networks,” in Proc. ISCA,2016.

[2] R. Venkatesan, Y. S. Shao, M. Wang, J. Clemons, S. Dai, M. Fojtik, B. Keller, A. Klinefelter, N. R. Pinckney, P. Raina et al., “MAGNet: A Modular Accelerator Generator for Neural Networks,” in ICCAD, 2019

[3] A. Parashar, P. Raina, Y. S. Shao, Y. Chen, V. A. Ying, A. Mukkara, R. Venkatesan, B. Khailany, S. W. Keckler, and J. Emer, “Timeloop: A Systematic Approach to DNN Accelerator Evaluation,” in 2019 IEEE International Symposium on Performance Analysis of Systems and

Software, 2019

[4] H. Kwon, P. Chatarasi, V. Sarkar, T. Krishna, M. Pellauer, and A. Parashar, “Maestro: A data-centric approach to understand reuse, performance, and hardware cost of dnn mappings,” IEEE Micro, 2020.

[5] L. Lu, N. Guan, Y. Wang, L. Jia, Z. Luo, J. Yin, J. Cong, and Y. Liang, “TENET: A Framework for Modeling Tensor Dataflow Based on Relation-centric Notation,” in 2021 ACM/IEEE 48rd Annual International Symposium on Computer Architecture, 2021.

[6] S. Shao, B. Reagen, G.-Y. Wei, and D. Brooks, “Aladdin: A Pre-RTL, Power-Performance Accelerator Simulator Enabling Large Design Space Exploration of Customized Architectures,” in ISCA, 2014.

[7] S. Campanoni, G. Agosta, S. Crespi-Reghizzi, and A. D. Biagio, “A highly flexible, parallel virtual machine: Design and experience of ildjit,” Software Practice Expererience, 2010.

[8] X. Yang, M. Gao, Q. Liu, J. Setter, J. Pu, A. Nayak, S. Bell, K. Cao, H. Ha, P. Raina, C. Kozyrakis, and M. Horowitz, “Interstellar: Using halide’s scheduling language to analyze dnn accelerators,” in Proceedings of the Twenty-Fifth International Conference on Architectural

Support for Programming Languages and Operating Systems (ASPLOS), 2020.

[9] Sharifian, Amirali & Hojabr, Reza & Rahimi, Navid & Liu, Sihao & Guha, Apala & Nowatzki, Tony & Shriraman, Arrvindh. (2019). μIR -An intermediate representation for transforming and optimizing the microarchitecture of application accelerators. 940-953. 10.1145/3352460.3358292.

[10] David Koeplinger, MatthewFeldman, Raghu Prabhakar, Yaqi Zhang, Stefan Hadjis, Ruben Fiszel, Tian Zhao, Luigi Nardi, Ardavan Pedram, Christos Kozyrakis, and Kunle Olukotun. 2018. Spatial: A Language and Compiler for Application Accelerators. In Proceedings of the 39th ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI 2018).

[11] Geonhwa Jeong, Gokcen Kestor, Prasanth Chatarasi, Angshuman Parashar, Po-An Tsai, Sivasankaran Rajamanickam, Roberto Gioiosa, Tushar Krishna: Union: A Unified HW-SW Co-Design Ecosystem in MLIR for Evaluating Tensor Operations on Spatial Accelerators. CoRR abs/2109.07419 (2021)

 

近期文章

应对大规模训练中的通信挑战:同步与拓扑

应对大规模训练中的通信挑战:同步与拓扑

随着各领域数据的不断积累和规范化,越来越多的研究者提出或使用大模型,如何更好地应对分布式训练系统的通信问题已成为研究热点。本文首先介绍分布式训练中的通信问题,然后从减少用于同步的通信和采用更利于通信的拓扑结构两个方面介绍应对通信问题的方法。
2022-01-24
众包训练:另类的分布式异构深度模型训练方法

众包训练:另类的分布式异构深度模型训练方法

众包是互联网时代发展出来的一种任务工作模式。众包计算在科学研究中已经得到了广泛应用。本文介绍了将众包用于深度模型训练,特别是大模型训练的一些研究和实践,总结了众包训练面临的限制和不同的实现策略,并讨论了众包学习的应用。
2022-01-17

联系我们

这是描述信息

招贤纳士

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

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