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相比，自动调度器，不需要手动模板。自动调度器能够在较短的时间内，生成具有更好性能的调度器。通过在搜索空间构建和搜索算法方面，进行创新实现这一目标。

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