领域编译器发展的前世今生 • 综述
张朔铭,博士研究生,正在中国科学院计算技术研究所崔慧敏研究员指导下攻读计算机系统结构博士学位,目前主要的研究方向是AI编译。
zhangshuoming17@mails.ucas.ac.cn
本文分为两个部分,第一部分为综述,第二部分重点讨论AI编译技术
近年来,随着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等。之后交由通用编译器完成汇编生成,链接最终生成可执行程序。这一设计不需要关注编译的后端部分,因而可以把重点放在利用领域专有知识对程序的优化上,通用优化则可以交给通用编译器来完成。
在过去几十年中,很多领域都提出了代表性的领域编程语言及编译器,这里我们对几个代表性的工作进行简单的介绍。
Halide是MIT计算机科学和人工智能实验室的研究人员开发的一门新的编程语言,旨在改变图像处理领域编程的现状。Halide语言将算法描述与算法中数据存储与计算执行等调度方面的内容解耦合。从而算法的描述和性能的优化可以分开进行,可以在不改动算法的情况下对算法进行调优。它将代码分为Algorithm和Schedule两个部分,Algorithm部分仅仅描述算法的功能,不进行实际的实现;Schedule部分指定算法在何时何处进行计算。使得比起用传统语言编写的代码,用Halide编写的程序不仅更容易阅读、编写和修改,而且 Halide能够自动优化代码,而人工对代码进行优化加速,这个过程需要花费数小时来完成。针对几种常用的图像处理算法,将Halide的版本与人工优化的版本进行比较,在不同平台上的运行结果如图所示。可见,在前3个算法中Halide版本的算法不仅简化了代码量,同时性能有着极大的提升。
Halide DSL: small code size & high performance[4]
TACO的张量向量乘法示例与生成的C代码[12]
[1] ORACLE CORPORATION. 2017. API Index. https://docs.oracle.com/database/121/SQLRF/intro002.htm#SQLRF50933
[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模型生成