TVM这篇论文搞深度学习部署的人应该都不陌生。它是编译器和机器学习系统领域绕不开的一篇工作讲的是怎么把深度学习模型从各种框架里端到端地优化并部署到不同硬件上。这篇是系列的第三篇前两篇已经把论文的动机、整体架构和核心抽象讲得差不多了这篇继续往核心走重点拆解论文里计算图优化、算子级调度优化和自动调优这三块内容。先说清楚一件事标题写的是“全文翻译”但这篇不是逐字逐句的译文更像是“精读笔记”。原文是写给系统研究者看的直接照搬翻译过来很多地方读者反而看不懂。所以我按主题拆开讲哪里是原文的意思哪里是基于实操经验的补充我会在行文里说清楚。适合正在做模型部署优化、想弄懂AI编译器内部逻辑、或者准备把论文思路转化成实际工作的读者。1. 先明确TVM到底解决了什么问题1.1 碎片化部署带来的重复劳动论文开头描述了一个非常现实的问题深度学习推理要从训练框架迁移到五花八门的硬件上执行。每套框架有自己的算子库和优化方法每类硬件CPU、GPU、各种加速器又有不同的指令集、存储结构和并行模型。你在这个设备上调好的算子实现换一个设备基本作废。我自己做部署时也深有体会最耗时间的往往不是模型结构本身而是“换一块板子就要重写一遍优化代码”。论文里把这种状态称为碎片化意思很直白同一个卷积算子在NVIDIA GPU上可以用cuDNN路径到了嵌入式设备的NPU上就得手写而且手写的过程几乎不可迁移。TVM要做的是让这套优化过程变成一条端到端的自动化管线。1.2 两阶段优化的核心分离思路TVM架构上最关键的决策是把优化拆成两个阶段。第一个阶段基于计算图做高层优化这时只关心模型里的算子怎么组合、数据怎么流动不关心底层硬件长什么样第二个阶段进入算子内部针对具体硬件的存储层次、向量宽度、并行能力来做调度优化。这个分离的逻辑其实很朴素模型的大结构相对稳定但算子内部实现高度依赖硬件。你不能指望一套手写代码通吃所有设备那就得把“模型要做什么”和“每个算子怎么做”分开。前者是高层图优化后者是算子级调度。论文里专门设计了一套中间表示来衔接两个阶段把计算图里的算子展开成可以进一步变换的循环结构这样后面的调度才有操作空间。这个两段式思路本质上是编译器里“前端-中端-后端”分层的延续只不过换到了深度学习场景里。在我后来的项目里这个经验同样适用优化任何复杂系统第一件事不是动手改代码而是先把“做什么”和“怎么做”之间的接口定义清楚。2. 计算图优化层先让模型“想清楚再干活”2.1 图级三板斧算子融合、常量折叠、代数化简论文里基于计算图的优化非常多但核心可以归纳成三个手段算子融合、常量折叠、代数化简。算子融合是指把多个相邻操作合并成一个操作最典型的是把卷积、批归一化、激活函数融合成一个算子。这样做直接减少了两次数据读写。深度学习推理里带宽往往才是瓶颈数据在内存和外存之间搬一次的成本可能比计算本身还高。融合的本质就是减少中间结果的物化让数据尽量留在缓存里。常量折叠是把编译期能算出来的东西提前算掉。比如一个子图里如果全部输入都是常量就没必要在每次推理时重复计算。批归一化在推理阶段可以把均值和方差折进卷积核里这既是代数化简也属于常量折叠的一种体现。代数化简则更宽泛包括把注意力计算里的缩放因子提前乘进矩阵、把连续两个转置抵消掉等等。这三种手段在TVM里并不是简单规则拼接而是靠一套基于图重写的机制统一驱动。每条规则都定义“什么结构可以匹配、匹配后替换成什么结构”框架不断重复应用规则直到结果稳定。优化类型核心目的典型例子算子融合减少中间数据读写conv bn relu 合成单算子常量折叠减少重复计算把所有常量输入预先算完代数化简降低单次计算量转置抵消、缩放因子前移2.2 布局转换最容易忽略的一环图优化里还有一个特别容易被初学者跳过去的内容数据布局转换。深度学习里数据维度顺序可以有不同的排法比如NCHW和NHWC两种常见布局有些硬件对特定布局有特殊优化路径。论文在高层优化阶段会主动识别哪些操作需要统一布局并在合适位置插入转换节点。为什么要专门做这件事因为布局不只是“排列顺序”它直接影响访存的局部性和向量化效率。图像数据在内存里是连续排列的NHWC布局让同一像素的所有通道紧挨着适合某些SIMD处理器一次性读出多个通道而NCHW布局让每个通道的所有像素连成一片在卷积计算里对通道维做累加时更友好。没有一个布局能同时满足所有算子所以布局转换实际上是一个搜索和权衡的过程。2.3 一个融合案例的直观计算拿最常见的convbnrelu融合来说我算过一笔账。假设输入是224x224大小的特征图通道256批量8一次推理这个子图的数据量大概是8乘以224乘以224乘以256乘以4字节大约411MB。如果不融合卷积写完一批中间结果、批归一化读一遍写一遍、激活再读一遍光中间数据的流量就翻了好几倍。融合之后这几百MB的中间流量直接被省掉了。这也是为什么论文把算子融合放在图优化第一优先级。它不从“少算几次乘加”出发而从“少搬几次数据”出发这个视角在真实硬件上往往更关键。我做过一个简单的对比实验同样的模型开融合和不开融合在CPU上推理延迟能差出30%到50%主要就差在访存上。3. 算子级调度优化精细打磨“具体怎么做”3.1 为什么不直接写高性能算子库图优化做完之后计算图里的每个算子还是抽象描述真正要落到硬件上还需要决定怎么切分循环、怎么处理向量化、怎么安排计算顺序。这时候如果直接依赖现成的算子库一是算子覆盖有限二是很难针对特定模型结构做定制三是新硬件出现时算子库根本来不及适配。论文给出的回答是提供一套调度原语让你以代码生成的方式自己“拼”出高性能算子实现。相当于把硬件优化知识从具体算子实现里抽离出来变成一组可以组合的变换操作。这套原语并不是论文原创发明的它借鉴了Halide的思想但TVM把它扩展到了深度学习场景并且跟后面的自动调优系统打通。理解调度原语的关键是心里始终有一条“循环嵌套”的线。任何一个算子都可以表示成多层循环套叠的计算调度就是变换这些循环的切分、合并、顺序和并行方式。3.2 经典调度原语逐个拆调度原语里最常用的几个我逐个说一下。split是把一个循环维度拆成两个维度。比如原来循环范围是32你可以拆成4和8这样做的目的通常是让外层循环可以被并行执行内层循环可以做向量化。fuse是split的逆操作把两层循环合并为一层有时候是为了减少循环开销有时候是为了配合某些硬件特性。vectorize是把一个内层循环转换成一次向量指令。现代CPU几乎都支持SIMD一条指令能同时处理多个数据。如果你不告诉编译器循环可以向量化或者你的循环结构让它无法向量化性能会差很多。unroll是把循环体展开减少循环分支判断增加指令级并行的机会但展开太狠会让指令缓存压力变大需要平衡。tile则是对多维循环做分块先按块遍历再在块内遍历。做tile的核心目的是改善访存局部性。比如一个矩阵乘如果完全按行遍历内层每次访问的内存跨度可能很大缓存命中率就很低。分块之后一个块内要用的数据能尽量留在缓存里。parallel是把某个循环分发到多个线程执行这是多核CPU上最直接的加速手段。选择哪个循环做并行很重要原则是让并行粒度足够大、线程之间不需要频繁同步。原语作用使用场景split循环拆分成两层为并行/向量化创造结构fuse合并两层循环减少循环开销vectorize转成向量指令支持SIMD的CPUunroll展开循环体减少分支、提升IPCtile多维分块改善缓存命中率parallel多线程并行多核CPU3.3 怎么组合这些原语才算“好”调度单独知道每个原语干什么远远不够。真正有价值的是知道面对一个具体算子时怎么组合这些原语。论文里强调调度的目标通常有三个方面一是最大化数据复用让数据在缓存里多待几轮二是最大化向量化和并行化程度把硬件的计算能力榨出来三是合理控制寄存器压力不要产生大量溢出。这三个目标之间经常有矛盾。tile的分块尺寸设得太大数据复用是好了但寄存器可能不够用反而在寄存器溢出上浪费时间设得太小缓存局部性又差了。所以调度本质上是一个多种因素之间的权衡过程。这里有一个我自己反复踩过的坑看到别人的调度模板很复杂就以为步骤越多越好。实际上一开始应该从最简单的调度开始先并行化外层的batch维度再看内层能否向量化然后逐步加上tile。每加一层变换就测一次性能而不是一次性把六七个原语全堆上去。堆多了不仅性能说不清楚出了问题也根本不知道是哪一步引起的。4. 自动调优把“玄学”变成可搜索的工程问题4.1 为什么手工调度无法推广上一节说的调度原语经验丰富的工程师确实可以手工拼出一个不错的算子实现。但问题在于硬件的种类太多、算子的形态更多一个模型里几十个算子每个算子在不同硬件上的最优调度都不一样。手工调根本调不过来而且极度依赖个人经验。论文针对这个问题提出了自动调优方案AutoTVM核心思路是把调度过程建模成一个搜索问题。定义一个调度模板模板里有若干可配置参数然后让算法去自动搜索最合适的参数组合。这里的调度模板和最终写死的调度代码不一样它是一段带有“占位符”的代码类似参数模板里面哪些地方可以换split的因子、哪些地方可以调整循环顺序、哪些地方可以开关unroll都明确标注出来。4.2 搜索策略必须搭配代价模型如果只把参数列出来然后一个个试问题就大了。因为搜索空间巨大而真实测量一个配置需要在目标硬件上跑一遍推理时间成本极高。论文的做法是用一个机器学习模型来“预测”某个配置的性能这个模型叫代价模型。先用一小部分真实测量的数据训练它然后让搜索算法在代价模型的指导下探索新的配置每隔一段时间把最有希望的配置再拿去做真实测量用新数据更新模型如此循环迭代。这个思路跟超参数优化里的贝叶斯优化有相似之处但TVM的代价模型输入是经过编码的调度配置和硬件特征而不是一个简单的向量。值得一提的是论文里利用了跨硬件迁移。在硬件A上学到的调度经验对硬件B上类似算子的搜索仍然有参考价值这能明显减少从零开始搜索的时间。我没有验证过所有细节但在自己的实验里把一台机器上调好的配置迁移到同体系的新型号设备上做二次搜索收敛速度确实快不少。4.3 论文里经典卷积调度模板的启示论文里展示的卷积调度模板非常有代表性我简化一下它的核心结构。模板里至少包含几个可调维度一个决定怎么对输出通道做tile一个决定怎么对输入通道做tile一个决定空间维的分块方式一个控制是否对某个维度unroll一个控制并行使用多少线程。这样算下来组合数量非常惊人如果没有代价模型想都不用想。但奇怪的是自动搜索出来的配置里有些设计决策是“反常识”的。比如某个平台上最好的配置并没有把循环unroll到极致而是选择了适度的展开。原因可能是过深的unroll导致指令缓存溢出完全抵消了分支减少带来的收益甚至更差。这就说明在没有实测反馈的情况下光靠理论推导很难触及真正的最优配置。我记得第一次跟着论文思路跑自动调优时心里还有点抵触觉得这不就是“暴力搜索嘛”。后来看到一个搜索找到的配置比手工调了几天还快20%才真正理解这件事的价值所在。手工调优依赖直觉和经验自动调优依赖的是高质量反馈和数据积累后者在硬件迭代频繁的时代明显更可持续。5. 论文翻译中几个棘手术语的处理5.1 最容易被误译的几个词既然是翻译系列我专门把论文里几个容易翻译出歧义的术语拿出来说一下。schedule这个词很多人会直接翻成“调度”这个没错但要清楚它在这里不是指“任务调度”的那种调度。它不是决定“谁先执行谁后执行”而是决定“计算循环怎么变换、怎么映射到硬件资源”。我翻的时候更倾向于把它理解成“计算方案的组织方式”。lowering这个词直译是“降低”在编译器语境里它指的是把高层的抽象表示转换成更接近硬件的底层表示。我认为翻成“下降”或“降级”都容易产生误解更准确的表达应该是“底层化转换”但这也拗口。通讯群里大家约定俗成还叫“lowering”我这篇里有时也会直接用英文词因为译错比不译更糟。tune这个词在论文里指的是自动调优过程中的参数搜索阶段跟一般的“调整”不太一样。它强调的是反复迭代试错、根据反馈修正的过程。所以我在正文里大部分地方直接用“调优”个别强调搜索过程时会说“自动调优”或“参数搜索”。英文术语含义建议译法schedule计算循环的组织方式调度注明非任务调度lowering高层表示转底层表示底层化转换 / 保留英文tuning基于反馈反复搜索参数调优op fusion多个算子合并算子融合layout transform数据维度顺序变换布局转换5.2 几个容易混淆的概念对比graph-level optimization和op-level optimization算一个经常混的地方。前者是针对计算图整体或子图做变换不进入某个算子内部后者是针对单个算子的循环结构做调整。很多人会把布局转换当成算子级优化实际上图优化阶段已经会插入布局转换节点算子级优化只是在给定布局下进一步打磨实现。还有一个容易混的是compile time和runtime. TVM的很多优化是在编译阶段完成的也就是模型部署前一次性完成运行时不承担重活。但有些动态决策比如根据输入形状变化来选择不同调度是在运行时才做的。区分这两个阶段很重要因为如果你把编译期能干的活留到运行期干推理延迟会很难看。6. 读这篇论文以及做部署优化的实用心得6.1 这篇论文应该怎么读说实话直接从头到尾读一遍论文原文体验并不好。论文信息密度很高而且有很多系统背景的隐含知识。我的建议是先读Introduction和架构总览部分理解它“把优化分成两个阶段”的核心思想就够了。然后跳到实验部分看一眼搞清楚它的对比对象是什么评估指标又是什么。最后再回头读图优化和调度相关的章节。原文里最需要细读的是调度原语的定义部分这里没有捷径。如果不把split、fuse、tile这些原语的语义彻底弄清楚后面自动调优那部分看得会云里雾里。我建议读到这里时配合一个简单的算子手动写一遍调度尝试比如写一个矩阵乘的简单调度分别尝试不同的tile尺寸跑一下看看性能差异。这个操作比读十遍论文都管用。6.2 论文实验设计里值得学习的点论文的实验部分也有不少可以学的细节。它不仅在多个硬件上进行评测还专门对比了“手工调优的专家算子库”“自动调优搜索的结果”“理想选择的上界”这三档结果。这种对比方式可以清楚地区分系统本身的搜索能力有多强代价模型的拟合有多准以及自动调优跟理论上限之间还剩多少距离。我后来做自己的部署项目时借鉴了这个做法。不管优化手段多花哨先把“当前默认配置性能”“手工优化后性能”“自动调优后性能”“一个理论上限估计值”这四列数字列出来。有了这四列数字哪一块还有空间就一目了然了。如果没有这个表格很容易出现一个问题费了很大力气调完发现性能只提升5%但实际上默认配置距离最优值还有40%的差距意味着你优化的方向或者手段本身就选错了。6.3 部署实践中几条避坑经验最后顺手整理几条近年在工程里反复验证过的经验。第一不要为了用自动调优而自动调优。算子数量少、部署周期短的场景下手工调度可能一小时就完成了自动调优启动搜索反而要跑很久。自动调优适合的是“一次调优、大量部署”的场景比如一个模型要部署到大量同类设备上值得前期投入搜索时间。第二调优之前先确认硬件状态一致。同一型号设备散热条件不同、降频策略不同甚至同一块板子上跑的系统版本不同都可能让搜索结果偏离。我遇到过一次在实验室设备上调优效果很好、换到现场设备性能大幅回落的情况排查半天发现是设备降频策略导致的。第三调度模板的设计决定了调优上限。自动调优只能搜索模板允许的配置范围如果你在模板里根本没有给出tile维度那搜索再怎么跑也找不到tile带来的收益。模板设计的合理性代表了你把硬件经验编码进系统的水平。把手工已知的好策略先放进去再让搜索去拓展这在初期是相当务实的策略。第四度量性能时不要只看单次推理时间至少跑足一定轮次取稳定值。深度学习推理中有缓存预热、频率爬升等效应前几次推理的耗时会明显不正常。取稳定状态的数据才有比较意义这个道理虽然简单但确实经常被人忽略。
