MLIR编译器调度与优化点滴
MLIR编译框架下软硬协同设计的思考
自从AI芯片成为热门的研究课题，众多关于AI芯片架构探索的学术文章不断涌现，大家从不同的角度对AI芯片进行架构分析及性能优化。MLIR是谷歌团队推出的开源编译器框架，颇受瞩目，灵活的编译器架构提升了其在众多领域应用的潜力。通过自定义IR的衔接，可以在架构探索和MLIR之间架起一座桥梁，在编译的过程中，自动进行硬件架构的探索和软件的优化编译，甚至生成硬件的代码，实现软硬协同设计。
架构探索方法的介绍
近十年，AI领域专用芯片的演进极大地促进了架构探索(指架构定义及性能分析)的发展，先后出现了众多的分析方法，这些分析方法针对AI计算过程中关键算子以及网络模型进行建模分析，从PPA(Power-Performance-Area)三个角度评估硬件性能。与此同时，伴随着AI编译框架的发展，尤其受益于MLIR编译器框架的可复用及可扩展性(详见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。相同的数据流架构可以采用类似的方法进行分析。
图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 2x2kernel的卷积在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迁移到LLVMIR。
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, split及reorder调度原语(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的编译器框架开展软硬协同设计的巨大潜力。
基于MLIR实现GEMM编译优化
GEMM(General Matrix Multiplication)即通用矩阵乘法运算，由于其计算行为具有一定的复杂性以及规律性，是编译算法研究的绝佳场景。MLIR是近期非常热门的一个编译器软件框架，是工业界及科研界研究的一个热点，其提供了一套灵活的软件基础设施，对中间表达式（IR）及其相互之间的转换进行规范的管理，是一个非常友好的编译器开发平台[1][2]。即是分析在MLIR框架下，实现GEMM优化的内容，以及对MLIR在这一方面的实现优势的讨论。
GEMM优化策略介绍
矩阵乘法运算，由于其过程会包含大量的乘加操作，并且伴随大量的数据读写，因而如何充分利用好底层硬件的存储资源以及计算资源，是编译器对其性能优化的关键。目前，已有的一些优化策略主要包括：
1.矩阵分块（Tile）
当前的处理器性能主要受限于内存墙，即计算速度要大于数据存储的速度。为了打破内存墙的约束，各类硬件包括CPU及其他专用处理器，会设置不同层次的存储单元，而这些不同层级的存储单元往往大小以及读写速度不同，一般越靠近计算单元的存储其存储容量也越小但访问的速度也越快。如果可以将计算过程的数据局部化分块，而这些分块的数据可以独立完成计算，那么分块的数据就可以放在层次化的存储中，然后通过不同存储间建立Ping-Pong的数据传输方式，将数据存储与计算解耦，从而可以有效得隐藏存储墙的问题，提高计算效率。矩阵运算就有这种特点，因而可以通过矩阵分块来加速运算，如下图1所示，假设有两层存储，将输入矩阵A和B，以及输出矩阵C，根据存储大小划分成相应的小块，即m->mc，n->nc，k->kc，每次将Ac(mc, kc), Bc(kc,nc), Cc(mc, nc)送入到离计算单元更近的存储模块内，完成局部的计算后再进行下一次的计算。
图1 矩阵运算的Tile操作示意图
在不同的底层硬件中，由于存储的层次以及不同层次的存储的容量大小不一样，分块的大小也会不一样。比如，文章[3]中对CPU而言，（Ac, Bc, Cc）划块的大小与cache大小一致，而为了充分利用register的资源，还会对（Ac, Bc, Cc）再进一步细划块成（Ar, Br, Cr），其尺寸大小与寄存器的数量一致。
2.向量化（Vectorize）
向量化的操作，主要是利用硬件的向量化指令或者SIMD（单指令多数）指令的特性，实现一个指令周期对多个值操作的能力。如下图2所示，通过将4个数据组成向量，利用处理器可以处理4个元素的新向量的计算能力，可以将4个指令周期的处理时间，压缩成1个指令周期的处理时间，从而极大提高运算处理能力。
图2 vectorize操作示意图
3.循环展开（Unroll）
由于矩阵乘法有多层循环构成，如果底层硬件有一定的并行化能力，包括多线程多进程处理能力等，那么可以对循环进行适当展开，从而提高计算的并行度，增加并发执行的机会。如下图3所示，将一个次数为1024的循环，展开成256次循环，新的循环内又包含4条可以并行执行的展开计算，如果处理器能够并行处理循环内部的展开计算，那么通过对原来的循环展开，可以获得接近4倍的性能提升。
图3 循环展开操作示意图
矩阵乘法的运算也包括其他的优化策略，比如数据重排等，但总体而言，各类编译器都是利用这些策略，充分利用硬件的存储及计算资源，达到最佳的运算效率。一些传统的计算库，如：OpenBLAS, BLIS, MKL等，开发时间长，性能也有比较优秀的表现。
MLIR实现GEMM优化
MLIR基于多层中间表示的方言（Dialect）思想，提供了一整套完善的编译器基础框架，可以帮助开发者快速实现编译策略想法的编译器。主要参考论文[4]，分析GEMM运算在MLIR中的实现，对应的硬件Target是因特尔i7-8700K处理器，每个核包含有32/256KB L1/L2 Cache以及多核共享的12MB L3 Cache，处理器支持AVX-2指令（256bit),优化目标是一个2088x2048xf64与2048x2048xf64的矩阵乘。
首先，其在高层次的Dialect上定义了一个矩阵运算的算子，这个算子的参数包含了输入矩阵（A，B）以及输出矩阵（C），同时为这个算子添加了tile/unroll 的尺寸等属性。如下图4所示，其中（M_C, K_C, M_R, N_R）属于Tile尺寸，K_U属于Unroll的大小。这里面（M_C, K_C）的选择是使得M_CxK_C大小的A矩阵块能够在L2 cache中复用，（K_C, N_R）的选择是使得K_CxN_R大小的B矩阵块能够在L1 cache中复用，（M_R, N_R）的选择是使得M_RxN_R大小的输出矩阵块能够在CPU Register中复用，这些值是根据硬件计算或者tunning出来的，在这里面的测试取了一个经验值。这些属性可以协助转换到更低一层的算子的策略实现，而选择哪些属性，则是跟编译算法以及编译的底层硬件对象有关，这些属性也是协助转换成下一层跟硬件更贴近的中间表示的实现，因而可以根据实际需要，灵活使用。
图4 GEMM算子的高层次定义
其次，MLIR的特点就是通过统一的多层中间表示，来实现对算子的层层低层化（lower）到具体的硬件目标上。针对上述高层次上定义的矩阵乘法算子，通过利用其所携带的优化属性，以及底层硬件的特点，设计了多条转换的路径（Pass），从而进一步把该算子lower到MLIR框架提供的中间辅助层（此中选择了Affine, Linalg，和Standard Dialect）。在这一层的转换过程中，基本包含了所有的策略，如：Tile，定制化复制，unroll，vectorize等。然后再将中间的辅组层的Dialect，进一步lower到LLVM Dialect上，如下图5所示。
图5 GEMM算子Lowing的层次化Dialect示意图
最后，通过mlir提供的mlir-cpu-runner工具，可以运行最后生成的LLVM Dialect的结果。总体优化及运行测试的命令，如下图6所示。其中，“-hopt”，“-hopt-vect”等，是从高层的算子（hop.matmul）到中间辅组层的转换路径，每一条路径都包含有相应的编译策略，可以根据需要，灵活添加以及改变，“-convert-linalg-to-loops”, “-lower-affine”等时中间辅助层之间的转换，最后转换成LLVM Dialect。
图6 MLIR运行GEMM的命令示意图
总体上，一个GEMM运算通过在MLIR框架下实现，并重写优化策略的路径，可以得到如图7所示的结果，其中箭头1对应包含了所有重写优化策略的MLIR实现，可以看到其能达到的计算速率为61.94GFLOPS，离理论上的峰值计算速率（75.2GFLOPS）比较接近，跟传统的计算库相比（如：BLIS，OpenBLAS，MKL等），也有着可以媲美的结果，其中的差距可能是传统的计算库有tunning的机制以及在编译器后端生成汇编指令及机器码有更成熟且高效的优化，因而可以得到更好的优化结果。总体而言，用MLIR重写的GEMM优化算法有着非常良好的表现。
图7 MLIR编译运行结果与其他计算库的对比示意图
另一方面，MLIR框架提供了非常完善的C++以及Python接口，因而可以很方便接入已有的计算库，进行联合优化。在[4]文中尝试了用MLIR+BLIS的方法，将MLIR放在外侧（提供手动优化功能），BLIS则作为micro-kernel放在内侧（提供auto tunning功能），最终的结果如图7中箭头2所示。可以看出，对于DGEMM（双精度），通过MLIR与BLIS的联合优化，也可以达到接近峰值的性能，而其性能要比单独的MLIR或者BLIS优化要差一点。但其实在SGEMM（单精度）的测试中，MLIR+BLIS的优化又要比单独的MLIR或者BLIS优化要好一些，因而其中的性能在差异还需要进一步分析。总体而言，MLIR提供了非常完善的支持，可以融合已有的计算库的优化算法，去实现联合的编译优化。
MLIR实现GEMM优化的优势
通过上面对MLIR实现GEMM优化算法的编译的介绍，可以看出MLIR在其中有着非常突出的优势。
首先，MLIR框架提供了非常完善的编译器基础设施，可以让开发者不需要花费太多精力在编译器周边的实现，从而更加专注于编译算法的开发。同时，其基于多层中间表达的方式，可以让编译器更加模块化，可以让编译算法利用不同层次的中间表达的抽象信息，在不同的层次中逐步具体化，从而使得算法实现更加层次化，更加易于实现及管理。
其次，MLIR框架提供了一直到最底层硬件的表示支持，由于其可以层次化在不同的中间表示层实现编译算法，可以在高层次的中间表示中实现不依赖于底层硬件的通用算法，而在接近硬件底层中，开发针对性的路径实现相应的编译算法，因而可以很方便地针对不同硬件目标开发统一的编译环境。本人认为，这也是MLIR相对于一些现有的AI编译器，如：TVM等，最有优势的地方之一，由于其框架可以根据需要自行扩展Dialect，同时这些Dialect又在系统中遵循一套统一的范式进行管理，因而对不同的编译目标（硬件target）会有很强的扩展性，同时编译器的工程管理又可以做到非常好的统一性。
另外，MLIR框架提供了完善的C++/Python接口，可以很方便地接入已有的优化算法，快速实现算法迁移。
主要介绍了矩阵乘法运算在MLIR编译器框架实现的主要过程及内容，以及其在优化算法的多层次实现，以及接入已有优化算法的能力等方面的优势。MLIR编译框架，为编译器的开发者，在多层中间表达转换以及算法优化等方面提供强大的基础设施支持，降低开发编译器的门槛。
参考文献链接
https://mp.weixin.qq.com/s/s5_tA28L94arLdm5UijkZg
https://mp.weixin.qq.com/s/A1h4pJSJ8VF97DrZksNULg
[1] Y. Chen, J. Emer, and V. Sze, “Eyeriss: A spatial architecture for energy efficient 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)
[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] Tze Meng Low, etc. Analytical Modeling Is Enough for High-Performance BLIS. 2016.
[4] UdayBondhugula, High Performance Code Generation in MLIR: An Early Case Study With GEMM. 2020.