Polyhedral Model—AI芯片软硬件优化利器（二）
作者简介：
作者要术甲杰（笔名），博士，研究领域为Polyhedral model，目前主要从事AI芯片上Poly优化等工作。
文章中的内容仅代表笔者本人观点。文中难免有错误的见解和一些笔误等情况，还望各位读者能够见谅。
笔者非常感谢唐杉博士对这篇文章提出的意见，使这篇文章获得更多的受用读者。作者本人也非常欢迎对Poly有兴趣的研究机构、大学和企业能够共同合作从事这方面的工作，如对Poly感兴趣可在知乎私信联系作者（知乎账号https://www.zhihu.com/people/yao-zhu-jia-jie/activities）。
由于芯片算力的提升，基于深度学习的人工智能掀起了计算技术领域的新一轮变革。在深度学习领域，由通用芯片与人工智能领域专用芯片构成的异构架构成为提升深度学习算法性能、实现人工智能目标的基础。然而，目前市场上各处理器厂商自主研发的加速芯片的架构非常复杂，给深度学习算法在芯片上的部署和调度提出了巨大的挑战，如何将深度学习算法有效地部署在当前的异构架构加速芯片上是人工智能和高性能计算领域的一个重要课题。
在这样的背景下，深度学习编译软件栈成为解决这个问题的一种有效手段，包括Google MLIR[1] 、TVM[2]、Facebook TensorComprehensions[3]、Intel Stripe[4]、Nvidia Diesel[5]以及MIT Halide[6, 7]和Tiramisu[8]等。这些工作的一个共同点是都已经或正在借助Polyhedral model[9-11]（以下简称Poly）来实现算子层的循环优化，这得益于Poly强大的循环优化能力。因此，许多处理器厂商和研究人员开始对Poly产生了兴趣。可惜的是，Poly在国内的发展基本处于空白。所以，在这篇文章里，笔者就自己对Poly的一点了解和前期的一些经验，简单介绍一下Poly的基本原理，并讨论一下Poly在深度学习编译软件栈和AI芯片上发挥的作用，以及一些开放性的问题。不足之处，还望各位一起讨论，多加批评。
全文分为三个部分连载，包括：
二. Poly在深度学习领域中发挥的作用（本文）
三. AI芯片上利用Poly进行软硬件优化的一些问题
Poly在深度学习领域中发挥的作用
从对上层应用的约束角度来看，作为一种通用程序设计语言的编译优化模型，Poly本身对应用是敏感的，只能处理满足一定约束条件的、规则的应用。Poly要求被分析的应用中，循环边界、数组下标都是仿射表达式，而且控制流必须是静态可判定的，我们暂且把这种对应用的要求称为静态仿射约束。实际上，对于通用语言而言，静态仿射约束的限制对程序的要求不算低，但是深度学习领域的大部分核心计算却恰好满足这种静态仿射约束，所以许多深度学习编译软件栈利用Poly来实现循环优化。
而从充分发挥底层AI芯片架构的能力角度来讲，Poly也非常适合，这得益于Poly能够自动判定和实现上层应用中循环的tiling/blocking（分块）变换并自动将软件循环映射到并行硬件上。本系列文章第一篇（Polyhedral Model—AI芯片软硬件优化利器（一））中图12和图13就是Poly在GPU上自动实现分块并将分块后对应的循环维度映射到GPU的线程块和线程两级并行硬件抽象上的实例。
为什么Poly需要自动实现分块？这是由底层AI芯片的架构导致的。以GPU为例，图14[15]所示是GPU的架构示意图。每个GPU上拥有自己的全局缓存（Global/Device Memory），然后每个线程块也有自己的局部缓存（Shared/Local Memory）。缓存越靠近计算单元，访存的速度越快，但是缓存空间越小。因此，当计算数据量大于缓存空间的时候，就需要通过将原来的数据进行分块的方式存储到缓存上，以此来适应目标架构的硬件特征。
图14 GPU架构示意图
而专用AI芯片的架构可能更复杂，如图15[16]所示是TPU v2和TPU v3的架构示意图，每个TPU有多种不同类型的计算单元，包括标量、向量以及矩阵计算单元，这些不同的计算单元对应地可能会有各自不同的缓存空间，这就给分块提出了更高的要求。只有通过对应用的正确分块才能充分利用好芯片上的架构特征。
图15 TPU架构示意图
当前一部分深度学习编译软件栈采用了手工调度和映射的方式来将上层应用部署到底层芯片上。以TVM为例，图16[17]中给出了一个TVM的调度示例。其中，调度过程首先将计算s进行分块（对应图16中的split操作），然后将分块后的维度映射到GPU的线程块和线程上（对应图16中的bind操作）。
图16 TVM调度示例
这部分工作需要由熟悉底层芯片架构的人员来编写，并且要人工分析分块的合法性，映射也需要手工完成。而Poly的作用就是将上述手工调度的过程自动实现。为了实现自动调度，许多深度学习编译软件栈开始采用Poly来实现上述功能。那么，Poly在深度学习软件栈上发挥的作用如何呢？
首先，Poly能够计算精确的数据流信息。Poly通过将传统的编译器中语句之间的依赖关系细化到语句实例的粒度，分析的结果比传统的方法更精确。计算精确的数据流信息有以下三点好处。
1. 计算精确的缓存搬移数据量。Poly不仅能自动计算出从管理核心（如CPU）到加速芯片（如GPU）之间传输的数据总量，还负责计算加速芯片上多级缓存之间的数据搬移总量，例如从GPU的global memory到shared memory上的数据搬移。“存储墙”问题给我们揭示了一个道理：数据搬移是程序性能提升的关键，尤其是现在市场上越来越复杂的多级缓存架构上，数据总量计算是否精确对程序性能的影响更加明显。
2. 降低内存空间使用。通过计算精确的数据流信息，Poly可以计算出临时tensor变量，这些临时变量的声明可在对应的缓存级别上实现，从而降低加速芯片上数据的访存开销。例如，图9所示的代码段中，tensor b就可以看作是一个临时tensor变量。
3. 自动实现缓存上的数据部署。以华为刚公布的昇腾AI处理器芯片为例，如图17[18]是该芯片的AI Core架构示意图。其中，UnifiedBuffer（输出缓冲区）和L1 Buffer（输入缓冲区）是低级缓存，离计算单元较远；BufferA L0/B L0/C L0是高级缓存，靠近计算单元。在低级缓存上，Poly可以借助标记节点，将不同计算单元所需的数据分别流向UnifiedBuffer和L1 Buffer；同时，当数据到达高级缓存时，Poly仍然可以借助标记节点将数据自动部署到BufferA L0/B L0/C L0。（注：这里描述的是如何通过Poly来实现这样的数据分流，只是为了说明Poly能够实现这样的自动数据部署功能，与具体实现无关。至于昇腾AI处理器芯片的编译团队是否使用了Poly，或者是否使用了这种方法来实现数据的自动部署还请以官方公布为准。）
图17 昇腾AI处理器的DaVinci Core架构示意图
其次，Poly能够实现几乎全部的循环变换。Poly通过仿射函数来实现几乎所有循环变换及其组合，这种仿射函数的计算过程不仅要考虑应用程序的并行性和局部性，还要考虑底层加速芯片的硬件特征。从循环变换角度来讲，Poly对编译软件栈的贡献包括以下几个方面。
1. Poly中的调度算法[19-22]能够根据依赖关系分析的结果自动计算出变换后循环的并行性、循环维度是否可以实施分块等特征，这些特征为后面硬件上的计算任务分配、缓存上的循环变换提供了理论依据。（这些信息保存在band节点（下面会介绍）的属性中）而部分循环变换如skewing/shifting（倾斜/偏移）、interchange（交换）等都可以在调度阶段自动完成。我们仍然以图9中所示的例子来说明。对于图10（Polyhedral Model—AI芯片软硬件优化利器（一））生成的代码，Poly计算出来的调度用其中间表示（schedule tree）[23]后得到的结果如图18所示，而图11（Polyhedral Model—AI芯片软硬件优化利器（一））生成的代码对应的调度如图19所示。（注：为方便说明，这里的schedule tree可能和实际在Poly中使用的有所不同，我们只是为了更直观地表示schedule tree的表示方式。）其中，domain节点包含所有的语句实例集合，sequence节点表示其子节点按序执行，而“[]”包含的节点称为band节点，可以想象成循环。
图18 图10对应的schedule tree表示
图19 图11对应的schedule tree表示
2. 自动实现深度学习应用中最关键的tiling/blocking（分块）和fusion（合并）变换。分块的目的是为了充分利用加速芯片上的缓存，而合并的目的是为了生成更多的临时缓存变量，降低访存开销。而且，Poly通过数学变换，能够自动实现更复杂的、手工难以实现的分块形状[6, 24-26]。其中，合并可根据调度选项在调度变换过程实现，分块则是在调度变换之后根据循环维度是否可分块等特征来实现。如图18和19就是根据不同的编译选项实现的合并策略对应的schedule tree，其中合并已经通过sequence节点实现，而分块的实现在Poly上很简单，只需要将band节点中的仿射函数进行修改就可以得到分块对应的schedule tree。如图20是图18经过分块之后的调度树，图19的分块也可以同样的方式得到，我们就不再赘述了。
图20 图18分块之后的schedule tree
3. 通过代码生成方式自动实现不改变语句顺序、但只改变循环结构的变换。这类循环变换包括peeling（剥离）、unrolling（展开）等。因为这些循环变换不改变语句的执行顺序，而只是对循环的结构进行修改来实现。这些循环变换对特殊加速芯片上的代码生成有十分重要的作用，例如一些架构可能并不喜欢循环上下界中有min/max这样的操作，此时就需要实现这类循环变换。这类循环变换可以通过在schedule tree中的band节点上添加特殊的options属性来实现。（注：我们的图中没有标出options，但实际使用的schedule tree中有options，而options中的内容是一个集合或者映射表达式，计算起来也很方便。）
第三，Poly能够自动实现存储系统的管理。在越来越复杂的加速芯片架构上，复杂的存储系统是实现芯片上计算部署的难点，即便是硬件开发人员来手工实现程序在存储结构上的管理，也是一个十分耗时且易出错的任务。而Poly借助中间表示自动实现了在多级缓存结构上的存储管理[27]，使得底层优化和硬件开发人员从这些琐碎的工作中脱离出来。这种自动管理存储系统的实现包括以下两个方面。
1. 自动计算缓存之间传递数据需要插入的位置。由于数据传输指令在原程序中是不存在的，所以Poly要能够实现这种从无到有的指令生成过程，并且正确计算出相应的位置。Poly借助schedule tree上的特殊节点和仿射函数，实现了数据传输指令位置的准确计算和自动插入。
2. 自动生成数据传输指令的循环信息。确定数据传输指令的位置后，Poly可以根据数学关系计算出当前指令所在循环的层次和维度信息，并自动为数据传输指令计算对应的调度关系，然后交给后端代码生成器生成代码。
如图21所示，是图20的schedule tree经过插入特殊的extension节点之后，得到的带有数据传输指令的中间表示。其中，kernel0和kernel1分别对应图20中最上面sequence节点下的两棵子树，而to_device_B和to_device_a表示从CPU的内存上拷贝tensor B和a到GPU的global memory，这两个语句在计算之前。from_device_c表示将GPU上的tensor c从global memory传输回CPU内存上，这个语句在计算之后。Poly并没有传输tensor b，而是在GPU的global memory上创建和使用了tensor b。（注：to_device_B和to_device_a也可以颠倒顺序执行，为了便于说明我们在这里假设按序执行。）
图21 图20的schedule tree插入数据传输指令之后的中间表示
最后，Poly还能够自动计算出变换之后循环到硬件上的映射。在提供多级并行硬件抽象和按计算的类型提供不同计算单元的加速芯片上，软件循环要实现到硬件上的映射，而这种映射关系也可以借助Poly的仿射函数和schedule tree上的标记来自动实现。这可以通过在kernel0和kernel1的子树内的band节点上添加特殊标记来实现。（注：图中未标出。）
--待续--
Reference:
[1] Lattner, Chris, andJacques Pienaar. "MLIR Primer: A Compiler Infrastructure for the End of Moore’s Law." (2019).
[2] Chen, Tianqi, ThierryMoreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan etal. "TVM: An automated end-to-end optimizing compiler for deeplearning." In 13th USENIX Symposium on Operating Systems Design andImplementation (OSDI 18), pp. 578-594. 2018.
[3] Vasilache, Nicolas,Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, WilliamS. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. "Tensorcomprehensions: Framework-agnostic high-performance machine learning abstractions."arXiv preprint arXiv:1802.04730 (2018).
[4] Zerrell, Tim, andJeremy Bruestle. "Stripe: Tensor Compilation via the Nested PolyhedralModel." arXiv preprint arXiv:1903.06498 (2019).
[5] Elango, Venmugil, NormRubin, Mahesh Ravishankar, Hariharan Sandanagobalane, and Vinod Grover."Diesel: DSL for linear algebra and neural net computations on GPUs."In Proceedings of the 2nd ACM SIGPLAN International Workshop on MachineLearning and Programming Languages (MAPL), pp. 42-51. ACM, 2018.
[6] Mullapudi, Ravi Teja,Vinay Vasista, and Uday Bondhugula. "Polymage: Automatic optimization forimage processing pipelines." In Proceedings of the Twentieth InternationalConference on Architectural Support for Programming Languages and OperatingSystems (ASPLOS), pp. 429-443. ACM, 2015.
[7] Mullapudi, Ravi Teja,Andrew Adams, Dillon Sharlet, Jonathan Ragan-Kelley, and Kayvon Fatahalian."Automatically scheduling halide image processing pipelines." ACMTransactions on Graphics (TOG) 35, no. 4 (2016): 83.
[8] Baghdadi, Riyadh,Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, YunmingZhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. "Tiramisu: Apolyhedral compiler for expressing fast and portable code." In Proceedingsof the 2019 IEEE/ACM International Symposium on Code Generation andOptimization (CGO), pp. 193-205. IEEE Press, 2019.
[9] Feautrier, Paul, andChristian Lengauer. "Polyhedron model." Encyclopedia of parallelcomputing (2011): 1581-1592.
[10]https://polyhedral.info. 2019.
[11] 赵捷, 李颖颖, 赵荣彩. 基于多面体模型的编译“黑魔法”. 软件学报, 2018, 29(8): 2371−2396.
[12] https://www.csa.iisc.ac.in/~udayb/slides/uday-polyhedral-opt.pdf
[13] Kennedy, Ken, and John R. Allen.Optimizing compilers for modern architectures: a dependence-based approach.Morgan Kaufmann Publishers Inc., 2001.
[14] Grosser, Tobias. "A decoupledapproach to high-level loop optimization: tile shapes, polyhedral buildingblocks and low-level compilers." PhD diss., 2014.
[15]http://www.pgroup.com/lit/articles/insider/v2n1ah.htm
[16]https://cloud.google.com/tpu/docs/system-architecture
[17] Cohen, Albert. " Polyhedral Compilation as a Design Pattern for Compiler Construction."(2019). https://pliss2019.github.io/albert_cohen_slides.pdf
[18] Liao Heng, "DaVinci: A Scalable Architecture for Neural Network Computing", Hot Chips 2019
[19] Bondhugula, Uday, Albert Hartono,Jagannathan Ramanujam, and Ponnuswamy Sadayappan. "A practical automaticpolyhedral parallelizer and locality optimizer." In 29th ACM SIGPLANConference on Programming Language Design and Implementation (PLDI), pp.101-113. ACM, 2008.
[20] Bondhugula, Uday, Aravind Acharya, andAlbert Cohen. "The pluto+ algorithm: A practical approach forparallelization and locality optimization of affine loop nests." ACMTransactions on Programming Languages and Systems (TOPLAS) 38, no. 3 (2016):12.
[21] Acharya, Aravind, Uday Bondhugula, andAlbert Cohen. "Polyhedral auto-transformation with no integer linear programming."In 39th ACM SIGPLAN Conference on Programming Language Design andImplementation (PLDI), pp. 529-542. ACM, 2018.
[22] Zinenko, Oleksandr, Sven Verdoolaege,Chandan Reddy, Jun Shirako, Tobias Grosser, Vivek Sarkar, and Albert Cohen."Modeling the conflicting demands of parallelism and temporal/spatiallocality in affine scheduling." In Proceedings of the 27th InternationalConference on Compiler Construction (CC), pp. 3-13. ACM, 2018.
[23] Grosser, Tobias, Sven Verdoolaege, andAlbert Cohen. "Polyhedral AST generation is more than scanningpolyhedra." ACM Transactions on Programming Languages and Systems (TOPLAS)37, no. 4 (2015): 12.
[24] Krishnamoorthy,Sriram, Muthu Baskaran, Uday Bondhugula, Jagannathan Ramanujam, Atanas Rountev,and Ponnuswamy Sadayappan. "Effective automatic parallelization of stencilcomputations." In Proceedings of the 28th ACM SIGPLAN Conference onProgramming Language Design and Implementation (PLDI), pp. 235-244. ACM, 2007.
[25] Bondhugula, Uday, Vinayaka Bandishti,and Irshad Pananilath. "Diamond tiling: Tiling techniques to maximizeparallelism for stencil computations." IEEE Transactions on Parallel andDistributed Systems (TPDS) 28, no. 5 (2016): 1285-1298.
[26] Grosser, Tobias, Albert Cohen, JustinHolewinski, Ponuswamy Sadayappan, and Sven Verdoolaege. "Hybridhexagonal/classical tiling for GPUs." In Proceedings of Annual IEEE/ACMInternational Symposium on Code Generation and Optimization (CGO), p. 66. ACM,2014.
[27] Verdoolaege, Sven, Juan Carlos Juega,Albert Cohen, Jose Ignacio Gomez, Christian Tenllado, and Francky Catthoor."Polyhedral parallel code generation for CUDA." ACM Transactions onArchitecture and Code Optimization (TACO) 9, no. 4 (2013): 54.
[28] https://github.com/tensorflow/mlir/blob/master/g3doc/RationaleSimplifiedPolyhedralForm.md
题图来自网络，版权归原作者所有
本文为个人兴趣之作，仅代表本人观点，与就职单位无关