[论文解读] Tensor Program Optimization with Probabilistic Programs

Junru Shao, Xiyou Zhou+8 more|arXiv (Cornell University)|May 26, 2022

Parallel Computing and Optimization Techniques被引用 10

一句话总结

MetaSchedule 引入一种领域特定的概率编程语言，以模块化方式构建丰富的张量程序优化搜索空间，使领域专家能够轻松扩展优化技术。通过将硬件特定的转换模块与学习驱动的搜索相结合，其在端到端深度学习工作负载上的推理速度相比 TVM 提高了 48%。

ABSTRACT

Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efficient optimization relies on a rich search space and effective search. Most existing efforts adopt a search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, a domain-specific probabilistic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts to analyze the program, and easily propose stochastic choices in a modular way to compose program transformation accordingly. We also build an end-to-end learning-driven framework to find an optimized program for a given search space. Experimental results show that MetaSchedule can cover the search space used in the state-of-the-art tensor program optimization frameworks in a modular way. Additionally, it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.

研究动机与目标

• 将自动张量程序优化中的搜索空间构建与搜索算法解耦。
• 使领域专家能够在不修改核心框架的前提下，逐步扩展和自定义优化搜索空间。
• 开发一种学习驱动的框架，利用构建的搜索空间找到优化的张量程序。
• 在多种硬件后端上实现与最先进系统（如 TVM 和 PyTorch）相当或更优的性能。
• 证明模块化、可组合的搜索空间构建能够为真实世界深度学习工作负载带来可测量的性能提升。

提出的方法

• 将搜索空间表示为概率程序，以建模程序转换中的随机选择。
• 允许领域专家定义可组合的模块化转换规则作为概率程序。
• 使用可微搜索算法探索概率搜索空间，以找到低延迟的张量程序。
• 将多个转换模块（例如，循环分块、向量化、张量化）组合成统一的搜索空间。
• 集成硬件特定模块（例如，使用 Tensor Core）以利用 GPU Tensor Core 等专用硬件原语。
• 构建一个端到端系统，支持在 CPU 和 GPU 平台上自动调优，且对框架层级的改动极小。

实验结果

研究问题

• RQ1 概率编程抽象能否有效实现张量程序优化中搜索空间构建与搜索的解耦？
• RQ2 所提出的抽象在覆盖 TVM 和 PyTorch 等最先进框架中使用的优化技术方面，表达能力如何？
• RQ3 转换规则的模块化组合在多大程度上能提升优化性能？
• RQ4 该框架能否在端到端深度学习模型上实现与现有系统相当或更优的性能？
• RQ5 领域专家在多大程度上能高效地通过硬件特定优化扩展搜索空间？

主要发现

• MetaSchedule 在 CPU 和 GPU 上对多种算子的性能与 TVM（Ansor）相当或更优。
• 在端到端深度学习模型上，MetaSchedule 的性能与 TVM 和 PyTorch 相当或更优，在使用硬件特定模块时，BERT-Large 的推理速度比 TVM（AutoTVM）快 48%。
• 仅用两天时间编写 82 行硬件特定模块，即实现了相比 TVM（AutoTVM） 48% 的性能提升，证明了定制化的高度实用性。
• 逐步组合转换模块可带来一致的性能提升，表明模块化设计支持可扩展的搜索空间增长。
• MetaSchedule 的搜索空间表达能力足够强，可覆盖现有框架的优化空间，且无需对底层系统进行手术式修改。
• 该框架实现了领域专家与基于学习的搜索之间的高效协作，显著降低了引入新优化技术的工作量。