[论文解读] A Learned Performance Model for Tensor Processing Units

Samuel J. Kaufman, Phitchaya Mangpo Phothilimthana+5 more|arXiv (Cornell University)|Aug 3, 2020

Parallel Computing and Optimization Techniques参考文献 26被引用 33

一句话总结

这篇论文提出了一个用于TPU的学习性能模型，通过图神经网络预测内核运行时间，在tile-size选择和操作融合方面优于分析模型，并在硬件访问受限时帮助自动调优。

ABSTRACT

Accurate hardware performance models are critical to efficient code generation. They can be used by compilers to make heuristic decisions, by superoptimizers as a minimization objective, or by autotuners to find an optimal configuration for a specific program. However, they are difficult to develop because contemporary processors are complex, and the recent proliferation of deep learning accelerators has increased the development burden. We demonstrate a method of learning performance models from a corpus of tensor computation graph programs for Tensor Processing Unit (TPU) instances. We show that our learned model outperforms a heavily-optimized analytical performance model on two tasks -- tile-size selection and operator fusion -- and that it helps an autotuner discover faster programs in a setting where access to TPUs is limited or expensive.

研究动机与目标

• 为TPU上的张量程序提供准确的运行时预测，以降低对真实硬件的依赖。
• 在多样的张量程序和优化任务上泛化，而无需大量特征工程。
• 实现可重新定向的性能建模，以适应不同编译器优化任务。
• 在TPU访问受限或成本高的情况下，展示改进的自动调优效率。

提出的方法

• 将张量程序表示为数据流图，包含操作节点和张量流边。
• 使用图神经网络（GraphSAGE）从局部邻域计算节点嵌入。
• 在节点嵌入之上可选地耦合一个序列模型（LSTM或Transformer），以捕捉长程依赖。
• 将操作码嵌入和每个节点/内核特征并入，形成模型输入向量。
• 用任务特定目标进行训练：tile-size选择的成对排序损失和对于融合的平方误差（对数变换目标）损失，以预测绝对运行时。
• 与高度手工调优的分析模型比较，并将学习模型集成到XLA自动调优器中，以引导配置搜索。

实验结果

研究问题

• RQ1 一个学习模型是否能够准确预测TPU内核运行时，以有效指导tile-size选择和操作融合？
• RQ2 基于图的表示、最少的特征工程，是否能够在未见过的张量程序和不同硬件代（TPU v2/v3）上泛化？
• RQ3 将学习模型集成到自动调优器中，是否可以减少找到快速配置所需的真实硬件评估次数？

主要发现

• 学习模型在tile-size和融合任务方面，与真实测量相比分别达到96.3%和95.5%的准确度。
• 学习模型在tile-size和融合任务上分别比最佳手工调优的分析模型提高了2.4%和26.6%的准确度。
• 在TPU v3上，学习模型在tile-size评估上达到3.8%的平均误差，平均相关性略低（0.65）。
• 消融研究显示边的方向、静态性能特征和内核嵌入选择显著影响融合任务的准确性，在tile-size任务准确性方面影响较小。
• 将学习模型整合到XLA自动调优器中，在硬件访问受限或昂贵时，能更快地发现高效配置。