1. 关于IR
这篇文章重点讨论Poly的IR。此前的内容主要基于Poly传统的schedule tree表示，介绍了如何实现AI芯片上的软硬件优化。Google MLIR[29]针对schedule tree的不足，提出了一种简化的Poly IR。以图22中的代码为例，用schedule tree表示可以得到图23的形式，而MLIR则将其表达为图24。
图22 另一个简单用例
图23 图22代码的scheduletree表示
图24 MLIR对图22中的简化Poly IR
对比之下，MLIR的Poly IR对循环做了更显式的表达，省去了schedule tree中的domain、schedule等信息。这种简化让循环变换后的代码生成过程变得更直接。比如做倾斜变换时，MLIR可以直接基于图24生成图25所示的Poly IR。
图25 经过倾斜后的MLIR的Poly IR
不过，与传统schedule tree相比，循环变换本身反而复杂了。在schedule tree上做倾斜，直接修改schedule的仿射函数就行（参考图6）；但在MLIR中，需要对应地修改显式表达的循环变量及下标信息。
这里有两个值得细想的问题。第一，Poly的整个流程中，代码生成固然复杂，但循环变换的时间开销可能更高。简化代码生成的同时让变换变复杂，这笔账划不划算？当然，Google MLIR团队集结了编译领域最顶级的专家和Poly研究团队，这种简化肯定经过深思熟虑，不过具体权衡可能需要向MLIR的Poly专家请教才能解惑。第二，这种简化Poly IR是为了简化从ML Function到CFGFunction的代码生成，如果Poly变换后的输出不是基于LLVM IR的框架，是否有必要沿用这种简化？毕竟，深度学习框架的“IR之争”尚未定论。
2. 关于循环变换
Poly的调度算法[19-22]基于线性整数规划求解新的调度仿射函数，这个过程会考虑几乎所有循环变换及其组合。图26[30]展示了线性整数规划问题的求解示意：蓝色点表示整个空间上的整数，斜边是循环边界等约束，构成可行解区间（绿色部分）。调度问题就是在绿色解空间中寻找目标（红色箭头）的最优解——在Poly里，就是按字典序最小的整数解。
图26 线性整数规划问题示意图
一个重要前提：Poly本是面向通用程序设计语言的编译数学模型。如果将其应用到深度学习这样的特定领域，是否还需要和通用语言一样的循环变换集合？举一个简单的例子：对于卷积算子，卷积核的循环嵌套在输入图像的循环内部，而卷积核的循环维度范围可能比输入图像小很多。当Poly计算新调度时，输入图像和卷积核的循环维度可能发生倾斜变换，但这种倾斜对后续的变形、代码生成往往不太友好。那么，Poly调度算法考虑的循环变换，在深度学习领域是否全都要，还是只需保留对性能提升最核心的几种？如果减少需要考虑的循环变换个数及其组合——也就是在图26中缩小可行解区间——求解是不是会更高效？
如果答案是肯定的，那么目前Poly能实现的循环变换中，对深度学习应用最关键的可能是分块和合并（当然有待商榷）。假设只考虑分块和合并两种，问题似乎简单一些。但编译优化中还有一个关键问题：怎么决定循环变换的顺序？先合并后分块，还是先分块后合并？传统的Poly中间表示没有给出明确答案，不幸的是，MLIR也没有解决这个问题。当然这只是极简假设，只有分块和合并显然不够，因为循环变换后的代码生成还得借助distribution来保证向量化等问题的发掘。
3. 关于分块
前面针对GPU、TPU和昇腾910的架构分析显示，（多级）缓存是目前AI芯片的架构趋势，而TPU、昇腾910这类专用计算单元的设计也引领了AI芯片的另一个方向。可以说，在当前的AI芯片上，分块已经是软件栈必须实现的优化手段。
即便只针对分块这一种变换，也还有很多值得研究的问题。以图27[31]中的二维卷积为例：卷积核在输入图像上“滑动”计算输出图像的结果。当卷积核针对输入图像的某个像素点（深蓝色方块）计算时，需要对其周边特定区域的像素点（浅蓝色方块）加权后得到输出上的一个像素点（红色方块）。由于卷积核需要滑动，大多数情况下输入图像的数据会被多次访问。图27中，当步长为1时，除第一列和最后一列外，所有像素点都会被重复计算。如果按卷积核大小对输入图像分块，那么相邻分块之间就会存在数据重叠问题。如何利用Poly在深度学习应用中自动实现这种满足数据重叠的分块？
图27 卷积操作示例
一种方式是采用PolyMage[6]类似的方法，利用Poly的调度求解重叠区间，但这可能导致过多的冗余计算，而且用调度来求解分块形状会使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的调度本质上是线性整数规划问题求解，而代码生成过程也涉及线性整数规划。在讨论深度学习领域是否需要所有循环变换及其组合时，设想通过减少循环变换个数来缩小解空间以加速调度过程；另外，MLIR的初衷也是为了降低代码生成的复杂度。文献[34,35]试图在特定情况下将线性整数规划简化为线性规划来降低时间复杂度。但这些都没有从实质上解决问题——问题本质仍是NP难题。要质变&现状可能还需要较长时间，其他计算机科学方法如constraint programming说不定也能提供一个解决思路。
另外，Poly的静态仿射约束对稀疏tensor等领域的扩展提出了挑战。关于稀疏tensor的工作已有一些研究[36,37]，但Poly无法直接应用于含有非规则下标的tensor。如何在这个领域对Poly进行扩展，可能是深度学习利用Poly优化的另一个需要解决的问题。Poly在解决稀疏矩阵问题方面已有一定进展[38-41]，说明non-affine扩展是可行的，而深度学习框架的可定制性也给这个问题创造了更多机会。