TVM自动调度器

TVM自动调度器

随着模型大小，算子多样性和硬件异构性的不断增长，优化深度神经网络的执行速度非常困难。从计算的角度来看，深度神经网络只是张量计算的一层又一层。这些张量计算（例如matmul和conv2d）可以通过数学表达式轻松描述。在现代硬件上为其提供高性能的实现可能会非常具有挑战性。必须应用各种低级优化，利用特殊的硬件内在函数来实现高性能。建立线性代数和神经网络加速库（如CuBLAS，CuDNN，oneMKL和oneDNN）需要大量的工程工作。

如果可以编写数学表达式，将其神奇地转化为有效的代码实现，那么生活将会更加轻松。深度学习编译器TVM及其搜索模块AutoTVM被构建为实现此目标的第一步。AutoTVM使用基于模板的搜索算法来查找给定张量计算的有效实现。这是基于模板的方法，仍然需要域专家为每个平台上的每个算子实施非平凡的手动模板。TVM代码存储库中的这些模板有超过15,000行代码。除了很难开发之外，这些模板通常效率低下，搜索空间有限，无法获得最佳性能。

为了解决AutoTVM的局限性，启动了Ansor项目，该项目旨在实现用于生成张量计算代码的全自动自动调度器。Ansor自动调度器仅将张量表达式作为输入，无需手动模板即可生成高性能代码。在搜索空间构建和搜索算法方面进行了创新。自动调度器可以以更自动化的方式以更少的搜索时间获得更好的性能。

Ansor自动调度器现已作为tvm.auto_scheduler软件包集成到Apache TVM中。这是来自加州大学伯克利分校，阿里巴巴，AWS和OctoML的合作者的共同努力。可在TVM网站上获得有关Intel CPU，ARM CPU，NVIDIA GPU和Mali GPU的详细教程。本文将进行简要介绍显示一些基准测试结果。

系统总览

AutoTVM与自动调度器

表1.工作流比较

表1比较了在AutoTVM和自动调度器中为算子生成代码的工作流。在AutoTVM中，开发人员必须经历三个步骤。在步骤1中，开发人员必须使用TVM的张量表达语言编写计算定义。这部分相对容易，因为TVM的张量表达式语言看起来就像数学表达式一样。在步骤2中，开发人员必须编写一个调度模板，该模板通常由20-100行棘手的DSL代码组成。这部分需要目标硬件体系结构和算子语义的领域专业知识，很困难。最后的步骤（步骤3）通过搜索算法自动执行。

在自动调度程序中，通过自动搜索空间构建消除了最困难的步骤2，使用更好的搜索算法来加速步骤3。通过自动搜索空间构建，不仅消除了巨大的人工工作，而且还可以探索更多的优化组合。这种自动化非免费提供，因为仍然需要设计规则来生成搜索空间。这些规则非常笼统。基于张量表达式的静态分析。设计一些通用规则，应用于深度学习中几乎所有的张量计算。

搜索过程

图1.搜索过程概述

图1.显示了优化整个神经网络时自动调度器的搜索过程。该系统将深度学习模型作为输入。然后，使用Relay的运算符融合遍历将大模型划分为小子图。任务调度器用于分配时间资源，优化许多子图。在每次迭代中，都会选择一个最有可能提高端到端性能的子图。对于此子图，分析其张量表达式为其生成几个草图。然后，使用学习成本模型运行进化搜索，以获取一批优化程序。优化程序将发送到实际硬件进行测量。测量完成后，分析结果将用作反馈来更新系统的所有组件。反复重复，直到优化收敛或用完时间预算为止。

值得注意的是，由于自动调度器从头开始生成调度器，会重用TOPI中的现有计算定义，但不会重用调度器模板。

基准结果

对AutoTVM和Auto-scheduler的性能进行基准测试。CPU基准测试是在AWS c5.9xlarge上完成的，该c5.9xlarge配备了Intel 18核skylake 8124-m CPU。GPU基准测试是在配备了NVIDIA T4 GPU的AWS g4dn.4xlarge上完成的。所有的基准代码，原始数据，调整日志都可以在此仓库中找到。

生成代码的性能

在三个网络上对fp32单批推理延迟进行了基准测试。图2显示了相对于AutoTVM的自动调度器的相对加速。在所有情况下，自动调度器的性能都比AutoTVM高出1.02倍至8.95倍。这是因为自动调度器会探索更大的搜索空间，其中涵盖了TOPI手动模板中遗漏的更有效的优化组合。BERT-base @ GPU是极端情况，手动模板的设计非常糟糕。换句话说，对于BERT模型中的形状，用于密集层的手动模板效果不佳。

图2.代码性能比较（越高越好）

搜索时间

众所周知，基于搜索的方法可能非常耗时，也关心搜索时间。让搜索收敛到单个神经网络通常需要几个小时。图3比较了AutoTVM和自动计划程序的搜索时间。尽管自动搜索程序的搜索空间较大，但在大多数情况下其收敛所需的时间要少得多。这主要是因为自动调度器具有更好的成本模型和任务调度器。

图3.搜索时间比较（越低越好）

更多结果

上面的存储库用作TVM的内部基准测试工具，仅比较了最新的AutoTVM和AutoScheduler。还尝试了在Apple M1芯片上进行自动调度，获得了一些不错的结果。

结论

构建了TVM自动调度器，该系统会自动为张量表达式生成高性能代码。与以前的AutoTVM相比，自动调度器不需要手动模板。此外，自动调度器能够在较短的时间内生成具有更好性能的调度器。通过在搜索空间构建和搜索算法方面进行创新来实现这一目标。

对自动调度器的当前性能感到兴奋。有兴趣扩展自动调度器的功能，以更好地支持稀疏运算符，低精度运算符和动态形状。