AI编译器TVM与MLIR框架分析
Posted 吴建明
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了AI编译器TVM与MLIR框架分析相关的知识，希望对你有一定的参考价值。
AI编译器TVM与MLIR框架分析
面向ASIC设备的编译器框架：TVM or MLIR？
2019~2021年，“摩尔定律失效”这一关键词频频出现于各大技术网站，在此背景下，市面上多如牛毛的AI芯片公司不约而同地给出了通用CPU+专用ASIC芯片的方案，以应对日益增长的AI边、端侧推理计算需求。在AI DSA芯片的开发实践中，棘手的问题除了底层硬件的设计，更多的还是AI模型在DSA芯片上优化、部署执行这一过程所需软件栈的实现，也即“AI编译器”技术栈，在这一领域最常常被大家提起并衡短论长的，莫过于TVM和MLIR。
严格来说，MLIR和TVM并不适合在一起对比：TVM是面向深度学习的模型编译器，用户可借此可直接获得编译/优化模型为推理blob的能力（可以看做机器学习时代的GCC、Clang），有兴趣了解TVM的同学可以参考这里；MLIR则是编译器基础设施类软件（可以看做机器学习时代的LLVM），面向的是需要构建自定义编译器的用户，它的基本设想是通过MLIR的Dialect共享生态减少用户开发编译器的工作量。
以上图为例，TensorFlow本身的优化推理过程如蓝色和绿色部分所示：TensorFlow Graph首先转换为XLA的HLO IR，应用XLA的优化Pass后再lower到目标设备的IR，例如对于x86/arm可lower为LLVM IR等，这样做的问题在于：
- XLA的图优化Pass是封闭的，用户如果想在其他框架例如TVM下实现同样的优化，需要阅读XLA的源码和TVM的源码，并添加Pass代码到TVM中，开发成本很高
- XLA HLO到LLVM IR的跨度太大，实现开销大，此处的开销包括各类针对微架构的带宽、缓存、指令集的优化
MLIR解决问题的方式如上图红色部分所示：先把外部IR转化为MLIR格式的Dialect IR（Translation），再把该IR lower到其他的MLIR方言IR（Conversion），最后从Target Dialect IR转出为Target IR（Anti-Translation），这么做的好处是，如果lower过程中所需的IR Dialect以及Target Dialect已经由MLIR生态中的其他用户实现了，那么你就可以直接从这些Dialect及其对应的优化Pass中受益。假设方言Memory擅长做内存分配优化，那就convert过去执行优化；假设Tile方言做Conv Tiling优化很好，那就从Memory IR convert过去继续优化...
MLIR的设想虽然非常好，但具体实现方面MLIR框架并没有直接解决上面列出的问题，而是把问题摊派给了加入MLIR生态的用户们：
- TensorFlow/Onnx/Torch等IR转为MLIR的Translation工作（入口工程），需要各框架团队以及开源社区支持
- Dialect之间并不天生支持完美切换，用户有需要时需要自行定制开发，然后根据个人意愿开源
- Translation工程、Conversion工程都有可能因为上游IR版本变动引入新bug，也就是上游项目变动需要几何级数的下游项目联动
- 总之，MLIR的设想虽然很棒，但在其生态完全爆发前，其工程开发总量也会很惊人。
支持上层模型，谁更完备？
TVM：TVM通过其relay.frontend模块下的功能实现高层模型转为TVM Relay IR，目前TVM支持MXNet、Keras、Onnx、TFLite、CoreML、Caffe2、Tensorflow、Darknet、Pytorch、Caffe和Paddle模型的输入，转换后的Relay IR可以被配套的图优化Pass进一步优化。
MLIR：目前已知部分主流框架实现了Dialect入口项目以加入MLIR生态：Onnx-MLIR、Torch-MLIR、Tensorflow-MLIR，这些项目通过自定义Dialect将Framework IR转换为MLIR格式的IR（Translation），转换后的IR可以通过Convert到其他Dialect的方法逐步lower到目标IR（Conversion）的MLIR方言，再通过反向的格式翻译为Target IR，如LLVM IR，进而生成可执行文件。
- 对于上层主流框架TVM的支持更全面，且配套了常用的图优化Pass集合；MLIR相比起来支持框架没有那么全面，转换后的图优化Pass数量取决于对应Dialect内置的Pass数量。
适配底层硬件，谁更方便？
适配底层硬件通常有两方面考量：对底层硬件设计过程的支持；模型编译优化过程的支持。
- 支持硬件设计：传统硬件架构设计和上层编译器设计通常是“两头凑”模式，即硬件设计专注于在当前制程下优化架构，编译器则专注于基于架构设计IR的优化、编译过程；但在AI芯片领域，计算需求的变更周期更短，我们更希望两者协同性更好以适应快速的需求迭代，理想情况下给出DSA计算核心描述，即可立即获得硬件的RTL级描述代码。
- 支持编译优化：即从模型图IR到芯片可执行代码的编译器软件栈。
TVM目前开源的面向底层硬件的开发模块是VTA，全称为Versatile Tensor Accelerator，直译为灵活的张量加速器，下图是VTA在TVM框架整体技术栈中的位置：
从下向上看，基于FPGA的arch设计需要遵照VTA的微架构（取指、读、存、算），该架构下要实现VTA的ISA（读、存、GEMM、ALU），满足微架构和ISA要求的硬件即可运行VTA Runtime，支持算子执行；从上向下看，运行VTA运行时的ASIC，默认支持TVM技术栈的所有优化手段，包括图优化、算子优化以及AutoTune等编译优化和代码生成技术。VTA的微架构如下所示：
VTA的微架构包括四个必须实现的模块：Fetch模块负责从DRAM取指到其余三个模块配套的FIFO队列；LOAD模块根据当前指令从DRAM中读取数据；STORE模块和LOAD模块功能相反；Compute模块除了包括模块内部的寄存器文件和RISC OP缓存，还包含核心的两种计算单元——ALU和GEMM核心。
VTA的ISA是一种类CISC的规范，包含4个CISC指令：LOAD/STORE/GEMM/ALU，LOAD负责从DRAM中读取2D Tensor到输入buffer、权重buffer或寄存器文件中，或者读取micro-kernel到micro-op 缓存，此外还需支持输入/权重buffer的动态padding；Store指令负责将输出的2D Tensor写回DRAM；ALU会被分解为一个micro-op序列，在矩阵Tensor上执行Elementwise类计算；GEMM会被分解为一个micro-op序列，负责矩阵乘计算。
一些限制：VTA主推的硬件开发代码并非目前EDA领域主流的Verilog/SystemVerilog实现的RTL描述代码，而是Xilinx的HLS C++代码；ISA虽然支持一定级别的扩展，但ISA指令目前只可处理二维Tensor是很硬的限制。
- 小结：VTA配合TVM提供了一整套针对开源深度学习的专用硬件开发的硬件设计、编译器和上层优化软件的全栈工具链，虽然设计硬件时会受限于其ISA和MicroArch，但它的确降低了机器学习从业者开发DSA芯片的门槛。
MLIR对底层硬件设计的支持方面，最近活跃度比较高的项目是Clang创始人Chris Lattner发起的开源项目CIRCT，其全称是Circuit IR Compiler and Tools，即基于电路中间表示的编译器和设计工具，该项目将DSL/IR/Compile等软件开发的思想应用到开源硬件设计领域，以加速硬件设计的流程，同时也寻求解决EDA工具的零碎化及封闭化的缺陷。
下面的图是Chris在ASPLOS 2021大会用的一张slide，蓝色的框为通用处理器（CPU/GPU/TPU），灰色的为专用处理器，可以看出CIRCT的希望是配合现有的MLIR生态覆盖通用+专用处理器场景，统一这些场景下开发所需的零碎工具到一个统一的EDA框架下：基于MLIR的软件框架可以涵盖软件开发工具（红线左），而CIRCT覆盖硬件设计工具（红线右）。
CIRCT的软件栈如下图所示，大体上软件站分为两部分：蓝色的是MLIR Core Project里的基础设施，包括Linalg/Affine/Vector/Standard这些核心方言；灰色的是CIRCT整合的硬件开发相关的方言，这些方言目前尚无统一的组织结构，更像是NAS领域的NNI（方言沙拉），例如Handshake由Xilinx研究人员开发，用于描述独立的非同步数据传输，定义了握手模型和fork/join/mux等控制逻辑；HIR是印度理工的研究人员开发的高层次时序电路描述方言，用来表示带延迟的计算、流水线、状态机等；下一层的FIRRTL是Chisel编译器的IR，Seq描述时序电路，Comb描述组合电路；最底层的LLHD是受LLVM启发的针对硬件RTL代码的多层IR表示，计划支持Chisel, MyHDL, SystemVerilog。
下图是以上方言之间的部分转换关系，可以看出，若我们要将MLIR的Core Dialect转换为Verilog Dialect，可以选择的路线为：
FIR -> high FIRRTL -> mid FIRRTL -> low FIRRTL -> Verilog
编译器软件栈方面，MLIR生态中针对DSA较常见的做法是先把高层模型的IR转化为MLIR，再经过逐层lower下沉到目标硬件级MLIR，如LLVM Dialect，最终把LLVM Dialect转为LLVM IR，除了以上MLIR内部的“打洞”工作，基于LLVM/GCC支持生成芯片上可执行代码的工作包括：
- 后端指令描述（ISA）
- 后端特性支持，包括栈设计、寄存器描述
- 后端优化：指令调度支持、特殊支持生成、硬件loop、条件执行指令
- 特殊后端优化：循环展开、delay slot调度...
- 总结：TVM的VTA在硬件设计、编译器软件栈方面的开发支持很完备，但需要用户设计硬件时遵循其自主设计的ISA，对硬件设计的限制很大；MLIR没有这方面限制，但在DSA的HW/SW协同设计方面以及模型编译支持方面，开源生态中没有成熟的“全栈”方案。
一点看法
回到本篇主题，首先毫无疑问的，TVM-VTA方案在易用性和开发效率的优势应该是没有争议的，其开发流程最容易上手，开发者无论是设计硬件还是在硬件上优化算子实现的开发代价都比基于MLIR开发小很多。但对于实际开发DSA产品的团队，能做到完全不在意其微架构和指令集的限制几乎是不可能的;不考虑VTA的情况下，目前TVM和MLIR都没有完全支持定制DSA的设计、上层编译的成熟路线，也就是最终选型的结果对实际工程开发的影响不会很大，都需要开发者/团队深度定制大量代码。
谈谈AI编译软件栈、软硬件协同、TVM和MLIR的看法
1. 软硬件协同
2个核心要点，都是钱砸出来的感悟。
1）对于SIMD为主的DSA，指令的图灵完备决定了芯片的可编程性。例如某些算子不能使用NPU编程，还需要在另外一个cpu上编程，这种异构就会把现有的事情变得复杂。尤其是AI编译器就很难做的健壮，客观需要一个过程去逐步识别哪些是软件本身的问题，还是硬件本身就无法做到。
2）片上的内存和片内带宽。从DDR到计算部件的内存层级设计，非常关键，也决定了芯片的可编程性。如果数据从DDR到计算部件的开销比较大，自然是希望中间结果可以驻留在片上（片上的缓存最好也足够大），不溢出到DDR上，这样会让ai编译器变得异常复杂，那解决方案自然需要更长的时间成熟和稳定。
图1 memory hierachy for mobile vs. laptop vs. server
另外，片上的内存层级越少越好，能做成cache就做成cache，当然最好是scratch pad和cache让软件可以自己选择模式。scratch pad模式下，软件可以做一些比较极致的调度，把片上缓存成分利用，计算部件的算力才能充分发挥。
图2 Memory hierarchy of the GeForce GTX780 (Kepler)
2. AI编译栈
1）底层llvm编译器做好还是蛮难的。很多AI芯片选择VLIW架构，或者自己设计一套标量+向量张量的ISA，需要从头搭建一个llvm/gcc的编译器；对gpu公司来说，就是一个类cuda的编译器。国内能把这件事情做好的公司真的不多。
我们往往可以发现，不少公司的算子库，第一个版本都是基于汇编。不是大家不想用c/c++的异构语言，是因为支持这个语言的编译器本身就不好做，写出来的算子性能能够比上汇编算子，往往需要好几年时间的打磨。例如寒武纪，最开始是汇编算子，后来逐步到bang c，然后才有机会去做更高抽象的算子编程工具，例如基于tvm。
其他不少ai芯片公司，或者gpu公司，感觉也面临这样的问题。而对于选择riscv的公司，就省事很多，基于开源的riscv加上一堆intrinsic的支持，做少量的优化，就差不多了。
图3 RISC-V "V" Vector Extension Intrinsics
2）理解了1），算子库的状态就很好理解了。llvm编译器还不太理想的，只能堆更多人去写汇编算子。打磨几年，基于c/c++异构的算子团队才会起来，同时才有机会去做更高抽象的算子编程工具。但无论是写汇编还是写c/c++异构算子，对于最终用户自定义算子，门槛还是挺高的。
3）如何兼容cuda，是很多公司想做的事情。很多cuda代码已经在那里了，如果新的架构可以很容易的跑起来，性能也不会损失很多，对于芯片的推广，就是一个巨大的生态优势。AMD已经在这方面做很好的一个范例。但实际去分析AMD做的硬件和软件的事情，你会发现，这本身的设计和投入不是一个小的事情。
图4 移植CUDA到AMD的HIP，AMD提供复杂项目的移植服务，cuda-to-hip@amd.com
4）基于tensor的ai编译器，工作在这个领域已经5-6年了，感觉没什么特别的东西了，可能新的idea还在出来（例如基于micro kernel），但要解决的就是那些问题。
图5 常见ai compiler设计架构
图6 DL compiler比较：TVM, nGraph, TC, Glow, and XLA（2020年）
关键是基于芯片体系结构去解决问题，最终衡量标准是模型和程序的性能，如果基于达到手写或者硬件性能的90%以上了，应该也无所谓是基于tvm，还是mlir，或者是完全自研了。
图7 Google lasted work: MLIR Primer
当然，ai编译器的稳定性和泛化性还是很重要的，直接决定了芯片的推广范围有多大。
5）自定义算子。用户都想要，最好是能兼容cuda，或者更好用的编程方式。但都不太好做。提供c + intrinsic，用户基本很难用起来。值得关注类似triton或者tensorir，或者太极的基于tensor抽象的编程方式，但距离成熟可用，还有距离。
图8 OpenCL for Low-level Parallel Programing
6）训练和推理。推理感觉没啥好说的。训练对于硬件来说，必须搞定内存和片上带宽问题，另外还有片间和跨机互联问题。没有片上带宽问题，就不太需要特别极致的算子融合了。图层和算子就可以做的比较简单，就能利用xla，或者mhlo这种小集合的算子去表达tf或者pytorch的2k多算子。
如果硬件的指标不太好看，就会给软件带来更复杂的工作，例如需要更强的算子融合能力，可能一些现成的开源不一定能用了。这些资源投入都是十分巨大的，动辄就是几十人团队，需要2年以上时间去构建和打磨。
3. TVM和MLIR
我还是那个观点，两个不同的数据结构而已，怎么解决你的问题才最重要。
图9 AI模型在硬件上部署的一般流程
我们也看到了基于mlir的项目。总体来说，mlir是一个很符合传统编译器工程师的设计，dialect机制方便定制和扩展，从前端但最后生成机器指令，这个过程比较平缓。
图10 Chris@ASPLOS 2021
relay到tir算子expand的过程太陡峭，算子的负担可以稍微迁移一部分给图层，图层会有机会做更全局的优化。这些，对于交付来说，其实不痛不痒，mlir并没有带来巨大的收益。同时目前mlir还缺自定义算子的能力，类似tvm的dsl还没有构建出来，以及算子tiling或者tuning的工具也还没有。这些估计都需要时间去构建，估计2-3年？
图11 AutoTVM自动优化架构
总体来说，tvm当前更能满足产品交付的需求。mlir想做到相同的能力（算子能力的差距，图层能力其实差不多），还需要更长时间。当然二者都在不停演进，相互吸收优点。
参考文献链接
https://mp.weixin.qq.com/s/Ph2eqWIpbNutJylwTTVc7g
https://mp.weixin.qq.com/s/0V-FKZKeeysrIvdNQb3vlA
MLIR编译器调度与优化点滴
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.
以上是关于AI编译器TVM与MLIR框架分析的主要内容，如果未能解决你的问题，请参考以下文章