领域编译器发展的前世今生 • 面向AI的编译技术

作者简介：

张朔铭，博士研究生，正在中国科学院计算技术研究所崔慧敏研究员指导下攻读计算机系统结构博士学位，目前主要的研究方向是AI编译。

本文分为两个部分，第一部分为综述（领域编译器发展的前世今生 • 综述）；这部分重点讨论面向AI领域的编译技术。

0. 前言

随着人工智能时代的来临，AI领域应用的大量出现也促进着领域编译的发展，最突出的表现就是多种AI编译器的普及和应用。AI领域有几个重要的特征使得AI编译器面临很多新的机遇和挑战：一是AI领域中编程框架对计算图与算子分离的设计机制为编译优化提供了更多的机会和更广阔的空间；二是AI领域中对张量的抽象为编译优化提供了具有鲜明领域特征的语义信息；三是以Python为主的动态解释器语言前端为其与AI编译器的衔接带了挑战；四是面向AI的领域专用架构为应用的可移植性带来了挑战。在这些因素的驱动下，近年来学术界和工业界在AI编译方面提出了一系列创新性的方法，也为编译这一基础学科的发展注入了新的活力。

1. 图算分离与图算融合的优化探索

为了让开发者使用方便，框架前端(图层)会尽量对Tensor计算进行抽象封装，开发者只要关注逻辑意义上的模型和算子；而在后端算子层的性能优化时，又可以打破算子的边界，从更细粒度的循环调度等维度，结合不同的硬件特点完成优化。这种图算分离的解耦设计大大简化了AI复杂系统的设计，因此，多层IR设计无疑是较好的选择，目前的主流IR设计也是分为图（TVM Relay，XLA HLO，MindSpore MindIR等）和算子（TVM tir，XLA LLO，MindSpore AKG等）两层。以主流AI编译器TVM[13]和TensorFlow XLA[14]为例。TVM 和 XLA 上层都采用了数据流图的中间表示，用图结点来表示计算，边表示数据流的依赖。在下层TVM 和 XLA 都针对编译器自动生成不同平台的高效代码进行了设计。其中TVM底层针对深度学习核心的张量处理设计的中间表示tir，它借鉴了Halide的中间表示来描述结点内的计算，可以针对不同目标平台定制调度策略，从而实现平台相关的深度优化。TensorFlow XLA 则提出了一种基于代数表示的中间表示（XLA HLO），高层的数据流图被转换为XLA HLO 的中间表示，在此中间表示上可以实施支持jit的算子融合、内存操作消除等优化，优化后的XLA HLO 可以被翻译为LLVM 中间表示或直接映射到TPU 平台。TVM和TensorFlow XLA 的图与张量（或代数）中间表示相结合的方法，一方面能够适配人工智能领域用数据流图来描述应用的需求，另一方面又能够兼顾应用在不同硬件平台之间的移植和优化。

但图层和算子层独立优化无法充分发挥芯片性能。近来面向图算融合的优化也日益成为学术界重要的研究方向。EasyView[15]提出了针对在网络实现中高频出现的tensor view类算子的端到端在线编译自动融合方法，包含view lowering，内存活动追踪，读写关系一致的算子拓扑序列获取，以及计算内存优化策略等内容。Apollo[16]设计了一个开放式多层规约式融合架构以实现不同算子融合方式的协同组合。将不同的融合方式实现为不同的Layer，在各级layer分别做基于polyhedral优化的循环融合，通过计算图算子级别依赖和元素级别依赖的分析对访存密集型算子尽可能融合，识别无计算依赖的算子并行化等优化，然后通过对不同Layer进行逐层规约合并，从而得到最终的融合算子子图，并获得最佳的融合性能收益。

Apollo架构：子图切分，融合，优化[16]

2. 面向张量的极致编译优化

AI编译器的核心抽象是张量（矩阵的高维推广）。在AI领域，各种数据，如图片，文字，视频等，都被抽象成张量，而原本对这些数据的处理也被抽象成对张量的计算操作，如卷积，转置，池化，归一化等，这些对张量的操作按照顺序组合就组成了张量计算的数据流图。做这层抽象的意义在笔者看来是因为传统编译器的IR，如LLVM IR太底层了，其对数据的处理粒度在标量至多向量级别，而编译器针对如此底层的IR分析能力有限。毕竟通过嵌套在多层循环里的load，store和alu操作已经很难还原出其原本的计算信息（如矩阵转置等）。从而像是矩阵转置再转置这样的pattern就基本无法在LLVM IR的粒度分析出来，而通过张量和张量计算的抽象层，这种优化就可以很容易实现。

AI编译器的优化目标主要是为了提升AI模型的端到端性能，这个性能会受到包括计算访存比，并行性，资源占用率等多方面因素影响，因而很难通过一个通用的策略涵盖大量不同的后端而都能达到非常优秀的性能。AI编译器通常通过搜索调度空间的方式，来寻找适配后端的极致张量优化策略。这里以TVM为例。TVM将计算和调度分离，计算通过张量表达式表示，张量表达式在设计上借鉴了 Halide、Darkroom 和 TACO。调度则是针对计算的一系列变换，为了在许多后端实现高性能，必须要支持足够多的调度原语来涵盖不同硬件后端的各种优化而包括tile/fuse/reorder/bind/compute_at等等，通过调度可以挖掘张量计算在特定硬件后端下的极致性能。TVM陆续发展了从基于模版的AutoTVM到基于搜索的Ansor，再到通过DSL tensorIR结合二者优点的Meta Schedule逐渐发展起来的自动调度搜索策略，可以针对不同的目标硬件平台来自动搜索的更优卸载方式，从而实现平台相关的深度优化。

TVM 张量表达式优化流程，多后端调度支持[17]

3. Python为主的动态解释器语言前端衔接

与传统编译器不同，AI编译器通常不需要Lexer/Parser，而是基于前端语言（主要是Python）的AST将模型解析并构造为计算图IR，侧重于保留shape、layout等Tensor计算特征信息，当然部分编译器还能保留控制流的信息。这里的难点在于，Python是一种灵活度极高的解释执行的语言，AI编译器需要把它转到静态的IR上。针对这一难点，目前一个比较突出的工作是pytorch2.0[18]提出的TorchDynamo[19]这一 JIT 编译接口，传统pytorch编程无论是trace还是eager模式都没有办法简单地通过python代码获取模型的图结构，导致模型导出、算子融合优化、模型量化等工作都出现困难。TorchDynamo支持在运行时修改python动态执行逻辑，修改的时机是在CPython 解释器的 ByteCode 执行前，从而可以使通过在运行时设置一个自定义的 Frame，该 Frame 中的 ByteCode 支持 CallBack 到 Python 层去修改。供用户自定义计算图。利用这一机制TorchDynamo支持了动态图的特性，即只需要通过python的执行机制即可自动调用后方对应的静态子模型。以下图为例，TorchDynamo方式写的bar模块具有一个动态信息的条件分支（对b.sum()<0的判断），这是传统trace和eager模式都无法支持的描述，但通过TorchDynamo方式，python执行到条件分支时可以根据动态信息自动根据条件调用有if语句的子图或没有if语句的子图，从而完成了对python if语句描述的动态图信息的支持。

通过TorchDynamo支持动态的控制依赖计算图[19]

4. DSA芯片架构的支持

随着AI应用的算力需求日益增长，以CPU为主的通用计算算力越来越难以满足大规模AI模型的性能和时延需求，因此AI应用往往需要通过高性能的针对AI的DSA架构的加速器后端来加速算子性能。这不仅包括对基于FPGA、ASIC、类脑计算等专用加速器的，也有在GPGPU上扩展的tensor core等专用加速核心。这些加速器在计算通用性，可配置性，功耗，性能，性价比，易用性等方面各有千秋，同时在实现架构上差异巨大，根据这些架构约定的软硬件接口层次的不同，如是否将内部cache操作暴露给编译器，是否硬件自动内存管理等，编译器能做的优化层次也不同。这也导致AI应用部署在不同的后端架构时需要生成出各种粒度的最终计算负载，这种粒度对应硬件/runtime规定的接口。如在GPU后端下，一个或多个图上的算子，在经过算子融合等优化步骤后，最终匹配上了cudnn/miopen等GPU算子库的一个库函数实现，最终编译器通过codegen生成以cuMalloc，cuMemcpy，cuLaunchKernel等cuda driver级别的api调用，借助CUDA runtime软件栈分别完成GPU内存空间管理，CPU到GPU的数据传输，算子的底层库函数调用，GPU数据回传等操作，最终完成了计算步骤。由于DSA架构的差异性和多样性，每一种架构都可能对应一种codegen方式和runtime支持。因而如何设计扩展DSA架构以实现对新的DSA架构的快速支持也是AI编译器的难点。

不同类型AI芯片后端体系结构的比较[20]

5. 面向神经网络的全局优化

AI应用的优化主要是为了提升端到端的模型性能，其受到包括计算访存比，并行性，资源占用率等多方面因素影响。在通用的编译器优化（如LLVM的优化pass）以外，AI编译器引入的面向神经网络应用的特定优化主要包括三类：张量优化，自动微分，自动并行。在AI领域，计算被抽象成张量的计算，基于张量的计算原语和计算优化是AI编译器的重要关注点。自动微分是支持AI网络训练的重要基础，因为模型训练的基础就是梯度下降和误差反向传播，都需要自动微分技术的支持。目前主要有基于计算图的自动微分、基于Tape和运算符重载的自动微分、基于source2source的自动微分等主流方案。自动并行技术则主要包括：数据并行、算子级模型并行、Pipeline模型并行、优化器模型并行和重计算等，以提高整体并行度的方式优化网络性能。

6. 领域定制架构下的编译基础设施

随着传统通用处理器的性能提升越来越困难，近年来领域定制硬件（DSA）成为体系结构设计中新的增长点，也是计算机体系结构黄金时代重要的发展。多种类型的加速器不断涌现。例如GPU、NPU、FGPA、ASIC 等。而为了适配大量新生的DSA，领域编译开发需要通过共用来降低开发成本，并重点关注于核心商业逻辑的实现，这使得领域编译也像电力，网络，公共云等技术一样，走向了技术演进的自然终点：基础设施化。其中最有代表性的基础设施是MLIR[21，24]。

MLIR是由Google 提出的一个能够快速构建领域编译器的基础设施，提出了一种构建可重用、可扩展编译器基础结构的新方法。其核心思想是利用多层次中间表示来解决软件的碎片化问题，减少构建Domain Specific Compiler的开销。MLIR虽然目前主要用于机器学习领域，但在设计上是通用的编译器框架，比如也有FLANG（llvm中的FORTRAN编译器）[22]，CIRCT（用于硬件设计）[23]等与ML无关的项目。MLIR 提供一系列可复用的易扩展的基础组件，从而使得不同领域的编译开发人员能够快速的搭建领域专用编译器，而且不同领域的编译分析及优化可以被复用。

与 LLVM IR 唯一的中间表示不同，MLIR能通过多层方言（dialect）的设计，表示更高层次的结构和操作，比如神经网络的图结构，张量的计算等。MLIR将一些领域的特性抽象为方言并允许用户自定义新的方言，与此同时，MLIR 提供了不同方言间的转换机制来实现不同方言上编译分析和优化的复用。MLIR执行过程和LLVM一样，IR会通过由Pass组成的优化Pipeline，不断地方言内，方言间变换直到生成最终的IR，然后被lower到底层的通用IR上进行代码生成。MLIR不仅仅是一个中间表示，而是一个新的编译器基础设施。近年来，学术界和工业界也在MLIR 上开展了很多领域编译优化的工作。如下图所示的将TensorFlow XLA 接入到了MLIR 的例子。上层的模型输入为TF Graph，在MLIR架构下逐层变换到HLO，LHLO，Affine，Vector等更低层次的方言上，在每级方言上都有对应层次和粒度的优化和调度，如在最高层的HLO适合做融合等。最终生成LLVM IR或SPIR-V等通用中间表示，再调用后端通用编译器完成最终代码生成。这部分更详细的讨论可以参考“编译器与IR的思考: LLVM IR，SPIR-V到MLIR”

MLIR在Tensorflow XLA上的实现[25]

7. 国内学界和业界的工作

AI生态的火热也促使着学术界和工业界在AI编译系统和AI异构加速器等方面全面发展。国内在AI生态的建设方面开始较早，投入众多，催生着大量的学术成果和工业界基础设施的发展。

一方面，国内学术界在AI编译系统和AI加速器体系结构等方面有很多研究突破，包括但不限于语言，编译，软硬件系统设计等方面。例如：针对目前的张量优化只考虑了完全等价变换，为张量程序优化引入了部分等价变换的优化和对应的纠正机制的PET[26]；通过语言层面引入对tensor的细粒度控制，包括tensor的不规则索引等，避免了大量冗余计算的FreeTensor[27]；通过特定于GNN网络的图运算符抽象做优化的uGrapher[28]；将硬件抽象设计为IR以支持更多intrinsic 原语的AMOS[29]；通过基于残差的精度细化技术，控制量化误差，在性能和精度之间进行权衡的QUANTENSOR[30]；根据优化目标配置实现云-移动部署的性能功耗综合优化的DNNTune[31]；面向任意精度计算（Arbitrary Precision Computing：APC）的Cambricon-P[32]体系结构，等等。

另一方面，工业界结合各厂家自身业务的需求，在AI基础设施和系统技术上不断优化和探索，在不同的维度持续发力，也贡献了很多开源项目。在AI编译基础设施方面，华为的MindSpore社区[33]提供了一个主要面向华为的昇腾处理器后端的云边端全场景开放AI推理和训练框架。其在包括图算联合优化，大规模分布式自动并行，自动算子生成等多项技术上做出了探索和贡献。阿里巴巴的PAI团队也专注于编译优化，探索了XLA，TVM，MLIR等多条技术路线，目前在大颗粒算子融合技术，以及GPU上访存密集型算子的融合优化上也取得不错的效果，并在MLIR这条技术路线上扩充了框架应对动态输入shape上的能力[34]。此外，国内的大型互联网公司，AI技术和芯片公司等都在面向AI的编译技术上有越来越多的投入，极大推进了相关技术的发展。

8. 小结

领域编译器作为通用编译器的重要补充，在发挥极致性能和提升开发效率方面一直发挥着重要的作用，近年来在AI领域更是备受关注。随着AI技术的快速发展，DSA硬件在AI计算中的大量使用，AI软件栈也相应的日趋复杂，对编译技术提出更高的要求，这也大大促进了编译学科的快速发展。我们相信，在强烈的需求驱动下，通过学界和业界的共同努力，领域编译技术在各种类型的计算系统中将扮演越来越重要的角色。

- 全文完 -

参考文献

[1] ORACLE CORPORATION. 2017. API Index. https://docs.oracle.com/database/121/SQLRF/intro002.htm#SQLRF50933

[2] Jonathan Ragan-Kelley, Andrew Adams, Dillon Sharlet, Connelly Barnes, Sylvain Paris, Marc Levoy, Saman Amarasinghe, and Frédo Durand. 2017. Halide: decoupling algorithms from schedules for high-performance image processing. Commun. ACM 61, 1 (January 2018), 106–115. https://doi.org/10.1145/3150211

[3] LLVM ORGANIZATION. 2023. API Index. https://lldb.llvm.org/

[4] Jonathan Ragan-Kelley, Andrew Adams, Dillon Sharlet, Connelly Barnes, Sylvain Paris, Marc Levoy, Saman Amarasinghe, and Frédo Durand. 2017. Halide: decoupling algorithms from schedules for high-performance image processing. Commun. ACM 61, 1 (January 2018), 106–115. https://doi.org/10.1145/3150211

[5] Justin Holewinski, Louis-Noël Pouchet, and P. Sadayappan. 2012. High-performance code generation for stencil computations on GPU architectures. In Proceedings of the 26th ACM international conference on Supercomputing (ICS '12). Association for Computing Machinery, New York, NY, USA, 311–320. https://doi.org/10.1145/2304576.2304619

[6] A Xiao. 2020. What is stencil computation?https://www.zhihu.com/question/302053357

[7] Jia Guo, Ganesh Bikshandi, Basilio B. Fraguela, Maria J. Garzaran, and David Padua. 2008. Programming with tiles. In Proceedings of the 13th ACM SIGPLAN Symposium on Principles and practice of parallel programming (PPoPP '08). Association for Computing Machinery, New York, NY, USA, 111–122. https://doi.org/10.1145/1345206.1345225

[8] Jia Guo, Ganesh Bikshandi, Basilio B. Fraguela, Maria J. Garzaran, and David Padua. 2008. Programming with tiles. In Proceedings of the 13th ACM SIGPLAN Symposium on Principles and practice of parallel programming (PPoPP '08). Association for Computing Machinery, New York, NY, USA, 111–122. https://doi.org/10.1145/1345206.1345225

[9] Yunming Zhang, Mengjiao Yang, Riyadh Baghdadi, Shoaib Kamil, Julian Shun, and Saman Amarasinghe. 2018. GraphIt: a high-performance graph DSL. Proc. ACM Program. Lang. 2, OOPSLA, Article 121 (November 2018), 30 pages. https://doi.org/10.1145/3276491

[10] Yunming Zhang, Mengjiao Yang, Riyadh Baghdadi, Shoaib Kamil, Julian Shun, and Saman Amarasinghe. 2018. GraphIt: a high-performance graph DSL. Proc. ACM Program. Lang. 2, OOPSLA, Article 121 (November 2018), 30 pages. https://doi.org/10.1145/3276491

[11] Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017. The tensor algebra compiler. Proc. ACM Program. Lang. 1, OOPSLA, Article 77 (October 2017), 29 pages. https://doi.org/10.1145/3133901

[12] Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017. The tensor algebra compiler. Proc. ACM Program. Lang. 1, OOPSLA, Article 77 (October 2017), 29 pages. https://doi.org/10.1145/3133901

[13] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: an automated end-to-end optimizing compiler for deep learning. In Proceedings of the 13th USENIX conference on Operating Systems Design and Implementation (OSDI'18). USENIX Association, USA, 579–594.

[14] GOOGLE LLC. 2022. API Index. https://www.tensorflow.org/xla

[15] Lijuan Jiang , Ping Xu , Qianchao Zhu , Xiuhong Li , Shengen Yan , Xingcheng Zhang , Dahua Lin , Wenjing Ma , Zhouyang Li , Jun Liu , Jinming Ma , Minxi Jin , Chao Yang ,"EasyView: Enabling and Scheduling Tensor Views in Deep Learning Compilers". ICPP '22: Proceedings of the 51st International Conference on Parallel ProcessingAugust 2022 Article No.: 54Pages 1–11 https://doi.org/10.1145/3545008.3545037

[16] Jie Zhao, Xiong Gao, Ruijie Xia, Zhaochuang Zhang, Deshi Chen, Lei Chen, Renwei Zhang, Zhen Geng, Bin Cheng, Xuefeng Jin, "Apollo: Automatic Partition-based Operator Fusion through Layer by Layer Optimization", Proceedings of Machine Learning and Systems 4 (2022): 1-19.

[17] Chen, Tianqi, et al. "TVM: end-to-end optimization stack for deep learning." arXiv preprint arXiv:1802.04799 11.20 (2018).

[18] PyTorch 2.0 | PyTorch

[19] PYTORCH ORGANIZATION. 2022. API Index. https://pytorch.org/tutorials/intermediate/dynamo_tutorial.html

[20]Beibei Wei, 2022. Levelling up: Opportunities and Challenges of China's AI Chips Industry. https://equalocean.com/analysis/2022091518910

[21] Lattner, Chris, et al. "MLIR: A compiler infrastructure for the end of Moore's law." arXiv preprint arXiv:2002.11054 (2020).

[22] LLVM ORGANIZATION. 2023. API Index. https://flang.llvm.org/docs/

[23] LLVM ORGANIZATION. 2023. API Index. https://circt.llvm.org/

[24] Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. 2021. MLIR: scaling compiler infrastructure for domain specific computation. In Proceedings of the 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO '21). IEEE Press, 2–14. https://doi.org/10.1109/CGO51591.2021.9370308

[25] LLVM ORGANIZATION. 2023. API Index. https://mlir.llvm.org/docs/Dialects/Vector/

[26] Wang, H., Zhai, J., Gao, M., Ma, Z., Tang, S., Zheng, L., ... & Jia, Z. (2021, July). PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections. In OSDI (pp. 37-54).

[27] Shizhi Tang, Jidong Zhai, Haojie Wang, Lin Jiang, Liyan Zheng, Zhenhao Yuan, and Chen Zhang. 2022. FreeTensor: a free-form DSL with holistic optimizations for irregular tensor programs. In Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation (PLDI 2022). Association for Computing Machinery, New York, NY, USA, 872–887. https://doi.org/10.1145/3519939.3523448

[28] Yangjie Zhou, Jingwen Leng, Yaoxu Song, Shuwen Lu, Mian Wang, Chao Li, Minyi Guo, Wenting Shen, Yong Li, Wei Lin, Xiangwen Liu, and Hanqing Wu. 2023. UGrapher: High-Performance Graph Operator Computation via Unified Abstraction for Graph Neural Networks. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS 2023). Association for Computing Machinery, New York, NY, USA, 878–891. https://doi.org/10.1145/3575693.3575723

[29] Size Zheng, Renze Chen, Anjiang Wei, Yicheng Jin, Qin Han, Liqiang Lu, Bingyang Wu, Xiuhong Li, Shengen Yan, and Yun Liang. 2022. AMOS: enabling automatic mapping for tensor computations on spatial accelerators with hardware abstraction. In Proceedings of the 49th Annual International Symposium on Computer Architecture (ISCA '22). Association for Computing Machinery, New York, NY, USA, 874–887. https://doi.org/10.1145/3470496.3527440

[30] G. Li et al., "Unleashing the Low-Precision Computation Potential of Tensor Cores on GPUs," 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), Seoul, Korea (South), 2021, pp. 90-102, doi: 10.1109/CGO51591.2021.9370335.

[31] Chunwei Xia, Jiacheng Zhao, Huimin Cui, Xiaobing Feng, and Jingling Xue. 2019. DNNTune: Automatic Benchmarking DNN Models for Mobile-cloud Computing. ACM Trans. Archit. Code Optim. 16, 4, Article 49 (December 2019), 26 pages. https://doi.org/10.1145/3368305

[32] Y. Hao et al., "Cambricon-P: A Bitflow Architecture for Arbitrary Precision Computing," 2022 55th IEEE/ACM International Symposium on Microarchitecture (MICRO), Chicago, IL, USA, 2022, pp. 57-72, doi: 10.1109/MICRO56248.2022.00016.

[33]HUAWEI Corporation Limited. 2023. API Index. https://www.mindspore.cn/

[34] Zhu, K., Zhao, W. Y., Zheng, Z., Guo, T. Y., Zhao, P. Z., Bai, J. J., ... & Lin, W. (2021, April). DISC: A dynamic shape compiler for machine learning workloads. In Proceedings of the 1st Workshop on Machine Learning and Systems (pp. 89-95).

题图由stable diffusion 2模型生成

继续阅读

阅读原文