LoopTune：基于强化学习的张量计算优化
机器学习
2023-11-09 v3 编程语言
摘要
先进的编译器技术对于使机器学习应用能够在新型硬件上运行至关重要，但传统编译器无法提供性能，流行的自动调优器搜索时间长，而专家优化库则带来不可持续的成本。为此，我们开发了 LoopTune，一种深度强化学习编译器，用于优化深度学习模型在 CPU 上的张量计算。LoopTune 在利用超快轻量级代码生成器 LoopNest 执行硬件特定优化的同时，优化张量遍历顺序。借助一种新颖的基于图的表示与动作空间，LoopTune 将 LoopNest 加速 3.2 倍，生成的代码比 TVM 快一个数量级，比 MetaSchedule 快 2.8 倍，比 AutoTVM 快 1.08 倍，并持续达到手工调优库 Numpy 的水平。此外，LoopTune 在数秒量级内完成代码调优。
引用
@article{arxiv.2309.01825,
title = {LoopTune: Optimizing Tensor Computations with Reinforcement Learning},
author = {Dejan Grubisic and Bram Wasti and Chris Cummins and John Mellor-Crummey and Aleksandar Zlateski},
journal= {arXiv preprint arXiv:2309.01825},
year = {2023}
}