[论文解读] Equality Saturation for Tensor Graph Superoptimization

Yichen Yang, Phitchaya Mangpo Phothilimtha+4 more|arXiv (Cornell University)|Jan 5, 2021

Parallel Computing and Optimization Techniques参考文献 29被引用 31

一句话总结

本论文介绍 Tensat，一种使用等价饱和的张量图超优化器，通过 e-graph 同时应用所有重写，实现最高可达 16% 的运行时加速，且优化速度比先前的串行方法快最多 48x。

ABSTRACT

One of the major optimizations employed in deep learning frameworks is graph rewriting. Production frameworks rely on heuristics to decide if rewrite rules should be applied and in which order. Prior research has shown that one can discover more optimal tensor computation graphs if we search for a better sequence of substitutions instead of relying on heuristics. However, we observe that existing approaches for tensor graph superoptimization both in production and research frameworks apply substitutions in a sequential manner. Such sequential search methods are sensitive to the order in which the substitutions are applied and often only explore a small fragment of the exponential space of equivalent graphs. This paper presents a novel technique for tensor graph superoptimization that employs equality saturation to apply all possible substitutions at once. We show that our approach can find optimized graphs with up to 16% speedup over state-of-the-art, while spending on average 48x less time optimizing.

研究动机与目标

• 动机：说明在张量图优化中串行重写策略的次优性。
• 引入带有 e-图的等价饱和，以探索更大尺度的等价图空间。
• 扩展等价饱和以处理张量图中的非局部和多模式重写。
• 在标准的 ML 模型上展示实际的加速，并分析优化时间。
• 提供可扩展的提取技术（ILP）和循环过滤策略，以维持可控的搜索。

提出的方法

• 用一个基于运算符的有向无环图来表示张量图，并扩展以支持等价饱和。
• 通过 e-图应用一组重写规则，在不进行破坏性替换的前提下生成所有等价图。
• 执行饱和 e-图的探索，使用多模式重写处理和规范的 S 表达式匹配方法。
• 从根的 e-类中提取最低成本的图，使用贪婪方法或带拓扑约束的基于 ILP 的形式化。
• 使用基于硬件（GPU）上的算子运行时间的成本模型来指导提取。
• 在提取阶段，可选地求解 ILP 以选择最优节点，并通过拓扑排序实现循环避免机制。

实验结果

研究问题

• RQ1 等价饱和是否能够克服张量图重写中的相位排序，找到全局最优图？
• RQ2 如何有效扩展 e-图以适应非局部和多模式的张量重写？
• RQ3 针对张量图，哪些实际的提取方法（贪婪 vs ILP）可以最小化运行时成本？
• RQ4 应用多模式重写时的性能权衡（加速与优化时间）？
• RQ5 在多种模型上，Tensat 与最先进的串行搜索（TASO）相比如何？

主要发现

模型	搜索时间（s）TASO	运行时加速（%）TASO	搜索时间（s）Tensat	运行时加速（%）Tensat
NasRNN	177.3	0.5	45.4	68.9
BERT	13.6	1.4	8.5	9.2
ResNeXt-50	25.3	0.7	5.5	8.8
NasNet-A	1226	10.6	1.9	7.3
SqueezeNet	16.4	0.3	6.7	24.5
VGG-19	8.9	0.4	8.9	8.9
Inception-v3	68.6	5.1	6.3	10.0

• 在所评估的模型上，Tensat 的运行时间比最先进的 TASO 快最多 16%。
• 由于对大量等价图的紧凑 e-图表示，优化时间显著降低（最高快 300 倍）。
• 等价饱和能够覆盖比串行回溯更大的搜索空间，在大多数基准中带来更好的加速。
• 在某些情况下，增加多模式迭代可以进一步提高加速，但也带来 e-图爆炸和 ILP 超时的风险；循环过滤可以缓解。
• Tensat 在加速与优化时间之间提供有竞争力的折中，并可实现在编译工作流中的实际整合。
• 在 NasRNN、SqueezeNet、Inception-v3 等模型上的评估展示了广泛的适用性。