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芯片上发挥的作用，以及一些开放性的问题。不足之处，还望各位一起讨论，多加批评。
全文分为三个部分连载，包括：
三. AI芯片上利用Poly进行软硬件优化的一些问题（本文）
综合前文的论述，Poly能够在深度学习编译软件栈中实现许多复杂的、有意义的工作，这也是为什么现如今许多处理器厂商和深度学习框架的研究人员对Poly有所需求的原因。笔者希望通过介绍Poly已经在深度学习编译软件栈中能够实现的功能，引起国内厂商和研究机构对Poly的关注。但这并不代表着Poly能够解决现有的全部问题，或者现有的Poly技术可以不经过任何改进的情况下很好地处理深度学习和AI加速芯片上需要解决的问题。通过前期的研究和经验，笔者认为未来在深度学习编译软件栈中利用Poly来实现优化，可能需要考虑或者值得研究和实现的问题有以下几个方面。
1. 关于IR
关于深度学习领域的IR之争介绍，唐杉博士在之前的公众号文章“Deep Learning的IR“之争”” 已经进行了介绍。我们在这篇文章中主要关注的是Poly的IR。在之前的内容中，我们主要基于Poly传统的schedule tree表示介绍了如何实现AI芯片上的软硬件优化。Google MLIR[29]针对schedule tree的不足，提出了一种简化的Poly IR。以图22中所示的代码为例，用schedule tree对该部分代码进行表示，可以得到如图23所示的IR形式，而MLIR则是将其表示成图24的形式。（注：图23中的schedule tree与我们之前的第二部分中图18-21表示的内容一样，只不过这里用文字的形式表示出图的内容，图23中schedule后的标量维度可以对应为图18-21中的sequence节点。）
图22 另一个简单用例
图23 图22代码的scheduletree表示
图24MLIR对图22中的简化Poly IR
通过对比可以看出，MLIR的Poly IR对循环进行了更加显式的表达，而省略了schedule tree中的domain、schedule等信息。这种Poly IR简化了Poly在实现循环变换之后的代码生成过程。例如，在实现倾斜变换时，MLIR可以直接基于图24生成如图25所示的Poly IR。
图25 经过倾斜后的MLIR的Poly IR
但是相比于传统的schedule tree表示，循环变换的过程更复杂了。在schedule tree上，循环变换，如这里的倾斜，可以直接修改schedule的仿射函数来实现，可参考图6；但在MLIR中却要对应地修改显式表达的循环变量及对应的下标信息。
笔者对这种方案存在两点疑问。一是在Poly的整个流程中，虽然代码生成也比较复杂，但是循环变换的时间开销可能比代码生成的开销更高，虽然简化了代码生成，但是循环变换更加复杂，不知道这样的代价是否值当？当然，Google MLIR团队集结了编译领域最顶级的专家和最熟悉Poly的研究团队，笔者相信他们提出这种简化的Poly IR肯定是经过深思熟虑的，可能后期需要向Google MLIR的Poly专家再请教之后才可能解答这个疑惑。二是这种简化的Poly IR是为了简化从ML Function到CFGFunction的代码生成过程，那如果Poly变换之后的输出不是基于LLVM IR的框架是否还有必要采用这种简化的Poly IR？毕竟，目前深度学习框架的“IR之争”还没有结束。
2. 关于循环变换
Poly的调度算法[19-22]基于线性整数规划来求解新的调度仿射函数，而这个过程中会考虑到几乎所有的循环变换及多个循环变换之间的组合。以图26[30]为例是线性整数规划问题求解的示意图，其中蓝色的点表示整个空间上的整数，而图中的斜边可以看作是循环边界等信息给出的约束，这些约束构成了一个可行解区间（图中绿色部分）。那么调度问题可以抽象成在这个绿色的解空间内寻找一个目标问题（红色箭头）的最优解（在Poly 里，就是寻找按字典序最小的整数解）。
图26 线性整数规划问题示意图
但是，一个重要的前提条件是Poly是面向通用程序设计语言的编译数学模型，如果我们将Poly应用到如深度学习这样的特定领域，是否需要考虑和通用语言一样的循环变换集合？一个很简单的例子是对于一个卷积算子，卷积核的循环嵌套会嵌套在输入图像的循环嵌套内部，而卷积核的循环维度范围可能会比输入图像的循环维度范围小很多。当Poly计算新的调度时，输入图像的循环维度和卷积核的循环维度可能发生倾斜变换，但这种倾斜似乎对卷积计算后面的变形、代码生成等问题都不太友好。所以，Poly调度算法考虑的循环变换，在深度学习领域是否都需要，还是只需要比较核心的、对性能提升比较关键的几种变换？如果我们减少了需要考虑的循环变换个数及其组合，也就是说在图26中我们缩小了可行解的区间，那么求解起来是不是会更高效一些？
如果上面问题的答案是肯定的，那么笔者认为目前而言Poly能实现的循环变换中，对深度学习应用最关键的循环变换应该是分块和合并（可能有待商榷）。假设只考虑分块和合并这两种循环变换，这种情况下问题似乎简单一些。但是编译优化中还有一个比较关键的问题就是如何决定实现的循环变换的顺序。是先做合并后做分块，还是先做分块再做合并？事实上，对于循环变换的顺序判定问题，传统的Poly中间表示没有给出明确的答案，而不幸的是，MLIR也没有解决这个问题。当然，这只是极简情况下的假设。只有分块和合并显然是不够的，因为循环变换后的代码生成还要借助distribution（分布）来保证向量化等问题的发掘。
3. 关于分块
我们前面针对GPU、TPU以及昇腾910的架构都进行了分析，（多级）缓存是目前市场上AI芯片采用的架构趋势，而专用AI芯片如TPU、昇腾910等专用计算单元的设计似乎也引领了AI芯片的另一种方向。可以说，在当前的AI芯片上，分块是软件栈必须实现的一种优化手段了。
而针对分块这一种变换而言，仍然还有很多值得研究的问题。比如，以图27[31]中的二维卷积为例，卷积核（kernel）通过在输入图像（input）上进行“滑动”来计算输出图像（output）的结果。当卷积核针对输入图像的某一个像素点（图中深蓝色的方块）进行计算时，需要通过对其周边特定区域的像素点（图中浅蓝色的方块）进行加权（即卷积操作）后得到输出图像上的一个像素点（图中红色方块）。由于卷积核需要在输入图像上进行滑动，这种滑动的过程在大多数情况下会导致输入图像的数据被多次访问。以图27为例，当滑动的步长（stride）为1时，除输入图像上第一列和最后一列的像素点之外所有的像素点都会被重复计算。如果我们按照卷积核的大小对输入图像进行分块，那么分块之后输入图像的每个分块之间都会存在overlap（数据重叠）问题。如何利用Poly在深度学习应用中自动实现这种满足数据重叠的分块？
图27 卷积操作示例
一种方式是采用PolyMage[6]类似的方法利用Poly的调度来求解这样的overlap的区间，但是这种方式有可能会导致过多的冗余计算，而且用调度来求解分块的形状在某种程度上会使Poly的过程变得更加复杂，代码生成亦如是；另一种方式是在schedule tree上利用特殊的节点来实现，但是目前这种方式的代码实现都还没有公开。
另外一个问题是在上一期的内容中有读者提问到关于分块和冗余计算的问题。冗余计算的确会给性能的提升带来一定的影响，但是这种冗余计算的引入是为了实现分块之间的并行。我们在前文提到过，并行性和局部性有的时候是冲突的，为了达到两者之间的平衡，往往是需要作出一些其它的牺牲来达到目的[32]。而更重要的是，这种带有冗余计算的分块形状是目前几种分块形状中，实现降低内存开销最有效的一种形状。以图28[6]为例，图中列出来三种不同的分块形状，其中最左侧的梯形分块引入了冗余计算，但是这种分块在一次分块计算完成（水平方向）后，分块内需要传递给下一次计算的活跃变量（红色圆圈）总数最少，而其它形状如中间的分裂分块和最右侧的平行四边形分块剩余的活跃变量总数都很多，无法实现有效降低内存开销的目的。（注：图中未分析钻石分块[25]和六角形分块[26]，但是这两种分块可以看作是分裂分块的一种特殊形式。）
图28 不同分块形状分块内活跃变量的分析对比
除此之外，分块面临的问题还有很多。比如，Poly中实现的分块都是计算的分块，而数据分块只是通过计算分块和计算与数据之间的仿射函数来计算得到，这种结果能够保证数据的分块是最优的吗？在分布式结构上呢？而针对TPU和昇腾910等专用AI加速芯片，多级分块应该如何实现才能更好的发挥这些加速芯片的特征呢？
4. 关于合并
循环的合并是一个挖掘局部性的过程。我们仍然想强调的问题是，局部性和并行性是加速芯片上两个非常重要的变换目标，但是这两个目标有的时候是互斥的，就如我们在图10和图11中所示的例子一样。合并的循环越多，破坏计算并行性的可能性越大；而如果要保持计算的并行性，可能就要放弃一些循环的合并。
然而，在不同的架构上哪些合并是最优的，似乎静态判定是不太可能的。就如我们在第一部分分享的那样，在CPU上生成OpenMP代码可能一层并行就足矣，这时局部性的效果可能就比并行性的效果更好；而在GPU上，由于有两层并行硬件的抽象，可能并行性的收益比局部性的效果更佳。所以，现在许多深度学习软件栈也采用了Auto-tuning的方式来通过实际的多次运行来判定哪种策略是最优的。然而，即便是Auto-tuning的方式，能够保证遍历到所有的合并形式吗？如何选择一个合适的合并策略，是必须要通过调优的方式来确定吗？利用静态分析的方式来遍历所有合并策略的工作也研究过[33]，但是这种动态规划的方式是不是又会带来时间复杂度的问题？
5. 关于Poly时间复杂度和对Poly的扩展
关于Poly的时间复杂度问题，我们在上文中已经提到Poly的调度实质是线性整数规划问题的求解过程，而实际上Poly的代码生成过程也会涉及到线性整数规划问题的求解。我们在讨论深度学习领域是否需要所有的循环变换及其组合的时候，设想从减少循环变换的个数来减小解空间，以此来加速调度的过程；另外，MLIR的初衷也是为了降低代码生成的复杂度。而文献[34, 35]也试图对特定情况或者将线性整数规划问题简化为线性规划问题来降低时间复杂度。不过，诸此种种都没有从实质上解决掉问题的关键，因为问题的实质仍然是NP级别的难题。想要从质上改变这个现状，可能还需要一段比较长的时间，其它的计算机科学领域的方法比如constraint programming说不定也能是一个解决的方法。
另外，Poly的静态仿射约束对稀疏tensor等领域的扩展也提出了挑战。关于稀疏tensor的工作目前也有了一定的研究[36, 37]，但Poly无法直接应用于含有非规则下标的tensor的情况，怎么样在这个领域对Poly进行扩展也可能是深度学习利用Poly优化的另一个需要解决的问题。因为Poly在解决稀疏矩阵问题的研究时，有了一定的进展[38-41]，这说明Poly的non-affine扩展还是可行的，而深度学习框架的可定制性给这个问题也创造了更多的机会。
后记
笔者前后用近半个月左右的时间，对这个系列的内容进行了梳理、归纳和分析，这期间也遇到了不少的问题，在撰写的过程当中笔者自己也强化了不少这方面的知识。感谢唐杉博士能够提供这样的平台让Polyhedral model能够走进这么多人的视野。文章在撰写的过程当中也得到了许多其他朋友的帮助，让这个系列的文章能够更好地展现在大家面前。也感谢许多读者提出的疑问和建议，帮助笔者更好地梳理行文。
文章在成文发表之前，笔者和唐杉博士对文章进行了多次的修改和仔细阅读，力图将最好的文字和图片呈现在大家面前，但是仍然没有能避免一些小的问题。作者和公众号平台的后台人员发现的一些小的错误包括内容一中图6-8种T和ST表示的含义说反了、内容二中图18和20下schedule tree的S0语句band节点的下标误写成了S1、内容二的参考文献[28]没有引用等，这给读者带来了不少麻烦，在此也向各位读者表示歉意。
文章内容的知识点形成离不开许多各个领域专家、前辈的多年积累，衷心感谢为包括Polyhedral model、深度学习、AI芯片等技术发展作出贡献的研究人员，尤其是像昇腾910等国产芯片的设计和相关架构知识的介绍，从中文进行了详细的描述，给作者对架构的理解带来了不少的启迪。
经过作者和唐杉博士的共同努力，我们终于在新中国70周年国庆节前夕完成了该系列文章的连载工作。我们希望通过自己的努力，无论是从国产芯片还是国产系统软件等角度来为祖国的建设事业贡献自己的一些微薄之力。最后，祝祖国繁荣昌盛，祝所有的读者国庆节快乐！
Reference:
[1] Lattner, Chris, andJacques Pienaar. "MLIR Primer: A Compiler Infrastructure for the End ofMoore’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 deep learning."In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI18), 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 learningabstractions." 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 JohnR. Allen. Optimizing compilers for modern architectures: a dependence-basedapproach. Morgan Kaufmann Publishers Inc., 2001.
[14] Grosser, Tobias."A decoupled approach to high-level loop optimization: tile shapes, polyhedralbuilding blocks 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 aDesign Pattern for Compiler Construction." (2019). https://pliss2019.github.io/albert_cohen_slides.pdf
[18] https://www.ithome.com/0/440/590.htm
[19] Bondhugula, Uday,Albert Hartono, Jagannathan Ramanujam, and Ponnuswamy Sadayappan. "Apractical automatic polyhedral parallelizer and locality optimizer." In29th ACM SIGPLAN Conference on Programming Language Design and Implementation(PLDI), pp. 101-113. ACM, 2008.
[20] Bondhugula, Uday,Aravind Acharya, and Albert Cohen. "The pluto+ algorithm: A practicalapproach for parallelization and locality optimization of affine loopnests." ACM Transactions on Programming Languages and Systems (TOPLAS) 38,no. 3 (2016): 12.
[21] Acharya, Aravind, UdayBondhugula, and Albert Cohen. "Polyhedral auto-transformation with no integerlinear programming." In 39th ACM SIGPLAN Conference on ProgrammingLanguage Design and Implementation (PLDI), pp. 529-542. ACM, 2018.
[22] Zinenko, Oleksandr,Sven Verdoolaege, Chandan Reddy, Jun Shirako, Tobias Grosser, Vivek Sarkar, andAlbert Cohen. "Modeling the conflicting demands of parallelism andtemporal/spatial locality in affine scheduling." In Proceedings of the27th International Conference on Compiler Construction (CC), pp. 3-13. ACM,2018.
[23] Grosser, Tobias, Sven Verdoolaege,and Albert 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://mp.weixin.qq.com/s/0iDVjaucRUpn2UrVBuQ-oQ
[29] https://github.com/tensorflow/mlir/blob/master/g3doc/RationaleSimplifiedPolyhedralForm.md
[30] Stege,Ulrike. "The impact of parameterized complexity to interdisciplinaryproblem solving." The multivariate algorithmic revolution and beyond.Springer, Berlin, Heidelberg, 2012. 56-68.
[31] http://intellabs.github.io/RiverTrail/tutorial/
[32] Ragan-Kelley,Jonathan, et al. "Halide: a language and compiler for optimizingparallelism, locality, and recomputation in image processing pipelines." In 33th ACM SIGPLAN Conference on ProgrammingLanguage Design and Implementation (PLDI), pp. 519-530. ACM, 2013.
[33] Jangda, Abhinav, andUday Bondhugula. "An effective fusion and tile size model for optimizingimage processing pipelines." In ACM SIGPLAN symposium on Principles andPractice of Parallel Programming (PPoPP), pp. 261-275, ACM, 2018.
[34] Upadrasta, Ramakrishna,and Albert Cohen. "Sub-polyhedral scheduling using (unit-)two-variable-per-inequality polyhedra." In Proceedings of the 40th AnnualACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages (POPL). pp.484-496, ACM, 2013.
[35] Acharya, Aravind, UdayBondhugula, and Albert Cohen. "Polyhedral auto-transformation with nointeger linear programming." In Proceedings of the 39th
ACMSIGPLAN Conference onProgramming Language Design and Implementation (PLDI). pp. 529-542. ACM, 2018.
[36] Kjolstad,Fredrik, et al. "The tensor algebra compiler." Proceedings of the ACMon Programming Languages 1.OOPSLA (2017): 77.
[37] Chou,Stephen, Fredrik Kjolstad, and Saman Amarasinghe. "Format abstraction forsparse tensor algebra compilers." Proceedings of the ACM on ProgrammingLanguages 2.OOPSLA (2018): 123.
[38] Benabderrahmane,Mohamed-Walid, et al. "Thepolyhedral model is more widely applicable than you think." InternationalConference on Compiler Construction (CC). pp. 283-303. Springer, Berlin, Heidelberg,2010.
[39]Venkat, Anand, et al. "Non-affine extensions to polyhedral codegeneration." Proceedings ofAnnual IEEE/ACM International Symposium on Code Generation and Optimization(CGO). pp. 185-194. ACM, 2014.
[40] Venkat,Anand, Mary Hall, and Michelle Strout. "Loop and data transformations forsparse matrix code." In Proceedings of the 36th ACM SIGPLAN Conference onProgramming Language Design and Implementation (PLDI). pp. 521-532. ACM, 2015.
[41] Zhao, Jie,Michael Kruse, and Albert Cohen. "A polyhedral compilation framework forloops with dynamic data-dependent bounds." Proceedings of the 27thInternational Conference on Compiler Construction (CC). pp. 14-24, ACM, 2018.
题图来自网络，版权归原作者所有
本文为个人兴趣之作，仅代表本人观点，与就职单位无关