基于MLIR实现GEMM编译优化

发布时间:2021-07-31 19:00

摘要

 

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/256KL1/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编译框架,为编译器的开发者,在多层中间表达转换以及算法优化等方面提供强大的基础设施支持,降低开发编译器的门槛。本人也希望通过本文,可以让读者对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] 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.

上一个: 机器学习方法在数字芯片设计中的应用

下一个: 深度生成网络新思路:扩散概率模型

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多