作者简介：
张朔铭，博士研究生，正在中国科学院计算技术研究所崔慧敏研究员指导下攻读计算机系统结构博士学位，目前主要的研究方向是AI编译。
[email protected]
本文分为两个部分，第一部分为综述，第二部分重点讨论AI编译技术
0. 前言
近年来，随着GPU和DSA架构在不同领域的广泛应用，特别是AI系统相关技术的飞速发展，对于编译器的需求越来越强烈。编译器已经从一个相对小众的研究领域，变为学界和业界都高度关注并大量投入的方向。与此同时，编译器的开发人员也从芯片研发团队开始延伸到更上层的软件层面。在很多领域的软件系统中，都开始引入编译技术来实现提升开发效率或运行效率等目标。本文从领域编译器的角色着眼，来讨论领域编译器发展的前世今生。
通用编译器 vs. 领域编译器 
编译器作为芯片配套的重要基础软件，它向上封装高级语言从而提高应用软件的开发效率和兼容性，向下适应体系结构而生成高效的可执行代码，是沟通软硬件之间的桥梁。一方面，通用编译器负责将某种编程语言映射到某种特定的芯片，例如Intel平台的C语言编译器负责将C语言的程序映射为Intel的X86指令集。通用编译器在设计编译优化时，会将通用性作为最重要的出发点之一。为了更便捷高效的支持更多语言（前端）以及更多目标芯片（后端），GCC、LLVM等框架逐渐成为重要的编译基础设施，它们都集成了非常多的通用编译优化组件，如循环优化、指令级优化、存储层次优化等。另一方面，领域编译器面向某些特定的领域应用程序而设计，旨在大幅提升该领域研发人员的编程效率或程序的运行性能。因此领域编译器在设计编译优化时，会将该领域的应用特征进行一定的总结抽象，形成领域特定的一些计算模式，并根据这些模式来设计针对性的优化，例如面向Stencil计算特征的优化等。
结合领域编程语言的领域编译器 
早期的领域编译技术伴随领域编程语言（DSL）开始发展，作为解决通用编译技术难以在特定领域生成高性能代码的解决方案出现。DSL会针对领域的特征设计易于描述领域问题的编程语言（创建新语言或从现有语言扩展），通常使用专用的IR中间表示（一层或多层）来对领域知识进行表示，这些领域知识的引入可以极大的方便编译器的分析和优化。领域编程语言的提出有两个目标：一是为了大大的简化领域开发用户的编程复杂性，这方面的代表性语言是数据库查询语言SQL[1]，它通过对数据库常用操作进行抽象为用户提供了非常便捷的编程接口，然后利用编译器将其映射到底层数据库引擎所提供的操作接口；二是为了大幅提升领域应用的运行性能，这方面的代表性语言是Halide[2]，它为图像处理领域提供了一种抽象层面的编程语言，基于这一语言编译器以调度的形式来抽象底层优化，从而避免了传统编译器中为了通用程序的正确性而引入的复杂的分析机制，使其可以非常便捷的在调度空间中通过搜索从而得到一个性能很高的目标代码。
领域编译器相比较于通用编译器，它仍然是一个解析高级语言，生成中间表示并优化，最后代码生成的系统。它在设计上会有更多的选择理念，除去领域无关的通用优化外，它还需要增加领域知识的专用表示，面向领域特定特征的专门优化等等。在系统实现上，有的是将领域编译实现在通用编译器中，有的是在通用编译器的前端增加一个高层次的领域编译器。前者以LLDB[3]为例，LLDB是clang/llvm架构的调试器，使用类似GDB的调试语句来调试程序运行，调试的需求主要在动态调控程序运行，获得调试信息，而这些都是通用clang/llvm架构中包含的，因而LLDB的实现上套用了clang/llvm的基础设施。像LLDB这样能够紧密结合在通用编译器基础设施的领域编译器还是比较少的，更多的DSL编译器的做法则是采用了将高层次的领域编译器作为前端插在通用编译器，如LLVM之前。DSL程序通过领域编译器的前端将DSL做词法语法解析生成高层次IR（一层或多层），该IR包含领域知识的数据结构，在高层IR上可以做专用的领域知识优化，也可以做通用的编译优化，然后生成通用编程语言的输出，如C/CUDA/OpenCL/ LLVM IR等。之后交由通用编译器完成汇编生成，链接最终生成可执行程序。这一设计不需要关注编译的后端部分，因而可以把重点放在利用领域专有知识对程序的优化上，通用优化则可以交给通用编译器来完成。
在过去几十年中，很多领域都提出了代表性的领域编程语言及编译器，这里我们对几个代表性的工作进行简单的介绍。
1. 面向图像处理领域的Halide
Halide是MIT计算机科学和人工智能实验室的研究人员开发的一门新的编程语言，旨在改变图像处理领域编程的现状。Halide语言将算法描述与算法中数据存储与计算执行等调度方面的内容解耦合。从而算法的描述和性能的优化可以分开进行，可以在不改动算法的情况下对算法进行调优。它将代码分为Algorithm和Schedule两个部分，Algorithm部分仅仅描述算法的功能，不进行实际的实现；Schedule部分指定算法在何时何处进行计算。使得比起用传统语言编写的代码，用Halide编写的程序不仅更容易阅读、编写和修改，而且 Halide能够自动优化代码，而人工对代码进行优化加速，这个过程需要花费数小时来完成。针对几种常用的图像处理算法，将Halide的版本与人工优化的版本进行比较，在不同平台上的运行结果如图所示。可见，在前3个算法中Halide版本的算法不仅简化了代码量，同时性能有着极大的提升。
Halide DSL: small code size & high performance[4]
2. 面向Stencil计算的领域编程语言
Stencil计算是数值模拟程序中常见的循环运算模式，其特点是遍历计算区域，每个位置均执行相同的计算操作指令模板（stencil），所以将这种计算称作stencil计算。stencil计算大量出现在科学计算领域，并且通常负载延时敏感应用。因而使用者通常希望将其卸载到高性能设备如GPU上，以提升端到端性能，但在GPU上，片外内存访问和GPU计算单元的调度是不同于CPU的，这是编译stencil计算程序到GPU执行不可忽略的问题。俄亥俄州大的团队设计了一款面向stencil计算的DSL[5]，用户只需要简单的高层stencil操作描述，通过DSL编译器分析stencil操作类型，访存模式，根据GPU编程模型，自动生成tile size，shared memory优化，block间同步等优化。这些优化减少了访存带宽瓶颈，增加了处理单元utility，最终生成了在GPU上运行性能更高的stencil程序。
Stencil计算常见模板[6]
3. 面向矩阵计算的领域编程语言HTA
Tiling算法是挖掘并行性，优化数据局部性重要的方法。HTA(Hierarchically Tiled Arrays)[7]是一个在现代OOP语言上扩展的面向数组tiling描述的数据结构。通过HTA，许多科学计算算法，如LU分解，分块矩阵计算，stencil计算等，都容易通过tile来表示，而HTA的数据布局，tile管理都可以自动进行，而无需涉及到最底层每个元素的计算，从而大大减少了下标计算的描述，简化了函数接口，极大地降低了描述并行程序的复杂性，同时仍然达到了库级别的高性能。HTA还引入了dynamic partitioning和overlapped tiling的语言特性，进一步增强了tiling描述的灵活性和普适性。
HTA的tiling描述[8]
4. 面向图计算的领域编程语言GraphIt
GraphIt [9]是一种用于图计算的新 DSL，它可以为在具有不同大小和结构的图上运行的具有不同性能特征的算法生成快速实现。GraphIt 将计算的内容（算法）与计算的方式（调度）分开。程序员使用算法语言指定算法，而性能优化使用单独的调度语言指定。调度语言使程序员能够通过将大量边遍历和顶点数据布局优化组合在一起，轻松地搜索这个复杂的权衡空间。GraphIt 的性能比当时最先进的共享内存框架（包括Ligra、Green-Marl、GraphMat、Galois、Gemini 和 Grazelle）快最多4.8 倍，同时代码行数最多减少了一个数量级。
Graphlt与多种框架的图计算性能比较[10]
5. 面向张量的领域编译器TACO
张量代数编译器 (Tensor Algebra COmpiler：TACO)[11] 是一个用于计算稀疏和密集张量上的张量代数表达式的代码生成工具，特别针对稀疏张量做了大量优化。TACO通过简单的format描述做稀疏张量存储的自动生成，通过迭代图中间表示描述复合张量表达式中稀疏操作数的多级索引数据结构，通过用户指定循环生成的代码可以直接做并行化，并且支持自动算子融合。TACO在广泛使用的稀疏张量代数和稀疏线性代数库中获得与手动优化内核相当的性能。
TACO的张量向量乘法示例与生成的C代码[12]
6. 面向AI领域的编译技术
随着人工智能时代的来临，AI领域应用的大量出现也促进着领域编译的发展，最突出的表现就是多种AI编译器的普及和应用。AI领域有几个重要的特征使得AI编译器面临很多新的机遇和挑战：一是AI领域中编程框架对计算图与算子分离的设计机制为编译优化提供了更多的机会和更广阔的空间；二是AI领域中对张量的抽象为编译优化提供了具有鲜明领域特征的语义信息；三是以Python为主的动态解释器语言前端为其与AI编译器的衔接带了挑战；四是面向AI的领域专用架构为应用的可移植性带来了挑战。在这些因素的驱动下，近年来学术界和工业界在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] 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.
[16] Jinho Jung, Hong Hu, Joy Arulraj, Taesoo Kim, and Woonhak Kang. 2019. APOLLO: automatic detection and diagnosis of performance regressions in database systems. Proc. VLDB Endow. 13, 1 (September 2019), 57–70. https://doi.org/10.14778/3357377.3357382
[17] Jinho Jung, Hong Hu, Joy Arulraj, Taesoo Kim, and Woonhak Kang. 2019. APOLLO: automatic detection and diagnosis of performance regressions in database systems. Proc. VLDB Endow. 13, 1 (September 2019), 57–70. https://doi.org/10.14778/3357377.3357382
[18] Chen, Tianqi, et al. "TVM: end-to-end optimization stack for deep learning." arXiv preprint arXiv:1802.04799 11.20 (2018).
[19] PYTORCH ORGANIZATION. 2022. API Index. https://pytorch.org/tutorials/intermediate/dynamo_tutorial.html
[20] PYTORCH ORGANIZATION. 2022. API Index. https://pytorch.org/tutorials/intermediate/dynamo_tutorial.html
[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] LLVM ORGANIZATION. 2023. API Index. https://flang.llvm.org/docs/
[25] 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
[26] LLVM ORGANIZATION. 2023. API Index. https://mlir.llvm.org/docs/Dialects/Vector/
题图由stable diffusion 2模型生成
继续阅读
阅读原文
关键词
编译器
代码
语言
性能
框架