LLM2D
基于等式饱和和蒙特卡洛树搜索的张量计算图优化
Optimizing Tensor Computation Graphs with Equality Saturation and Monte Carlo Tree Search
作者: Jakob Hartmann, Guoliang He, Eiko Yoneki
发布日期: 10/10/2024
arXiv ID: oai:arXiv.org:2410.05534v1

摘要

深度神经网络在现实世界中的有效性往往取决于其延迟,因此需要优化技术来减少模型的推理时间,同时保持其性能。一种流行的方法是通过替换单个子图,将输入计算图依次重写为等效但更快的图。这种方法导致了所谓的相位排序问题,其中应用一个重写规则可能会消除随后应用更佳规则的可能性。最近的研究表明,来自编译器优化的相等饱和技术可以通过首先构建一个中间表示(IR)来缓解这个问题,该表示在第二步中提取最佳解决方案之前有效地存储输入程序的多个优化版本。然而,在实践中,内存约束会阻止 IR 捕获所有优化版本,从而在构建阶段重新引入相位排序问题。在本文中,我们提出了一种张量图重写方法,该方法使用蒙特卡罗树搜索来构建更优的 IR,方法是识别最有希望的重写规则。我们还介绍了一种新颖的提取算法,该算法可以提供 IR 中表示的张量程序的快速准确的运行时估计。与现有方法相比,我们的方法将神经网络的推理加速提高了高达 11%。