tvm
TVM 常被概括为“把模型编译成高性能 kernel 的系统”,但这句话如果不拆开,往往只会留下几个名词:Relay、compute、schedule、AutoTVM。真正重要的是理解这些层次为什么要分开,以及它们如何共同决定最终性能。
这篇文章尝试沿着一条完整主线来解释 TVM:模型如何进入编译器,为什么会被拆成图、计算表达和调度策略,schedule 为什么是性能核心,以及 TVM 适合在哪些场景发挥价值。
从模型到 kernel:TVM 的核心链路
先看一个非常简单的例子:
1 | y = relu(x @ W + b) |
从深度学习框架的角度看,这是一行普通代码;从编译器的角度看,它对应的是一条由多个算子组成的数据流:
1 | matmul -> add -> relu |
TVM 的核心工作,就是把这条高层语义链路逐层下沉,最终变成目标硬件可以高效执行的代码。这个过程可以概括为:
1 | 模型 |
这条链路里最关键的设计思想,是把“算什么”和“怎么高效地算”分开表达。前者对应 compute,后者对应 schedule。
第一层:Relay IR 负责描述模型结构
TVM 首先需要把模型转换成一种编译器内部可以分析和变换的表示。在传统 TVM 语境里,这一层通常是 Relay IR。Relay 的职责不是直接生成 kernel,而是准确描述模型中有哪些算子、它们之间如何连接、张量如何流动。
对于前面的例子,可以把 Relay 理解成这样一种图结构:
1 | x ----+ |
如果写成接近代码的形式,大致会像下面这样:
1 | x = relay.var("x", shape=(1, 768)) |
Relay 的重点不在于“快”,而在于“清楚”。它描述的是模型结构、算子依赖和张量关系,因此非常适合做图层面的优化,例如:
- 常量折叠
- 死代码消除
- 算子融合
- 子图划分
- 布局变换
也正因为如此,Relay 不直接回答“这个 matmul 应该用 shared memory 还是寄存器”,那是更低一层的工作。
第二层:compute 负责描述算子在数学上怎么算
从计算图继续向下,TVM 会把图中的算子转成张量级的计算表达。这个阶段关心的是数学定义,而不是执行细节。
以前面的 matmul -> add -> relu 为例,可以拆成:
1 | matmul: |
这就是 compute 的本质:用张量表达式定义结果如何由输入推导出来。它回答的是“算什么”,而不是“如何映射到硬件”。
这一层的抽象非常重要,因为一旦数学定义稳定下来,同一个 compute 理论上就可以对应多种不同的执行方式。CPU、GPU、NPU 的差异,不应该反过来污染算子本身的数学表达。
第三层:schedule 负责把数学表达映射到硬件
只有 compute 还远远不够。下面这个矩阵乘表达式,虽然已经定义了结果,但仍然无法直接高效运行:
1 | C[i, j] = sum_k A[i, k] * B[k, j] |
原因很简单:它只说明了计算关系,没有说明执行方式。真正决定性能的,是 schedule。
schedule 需要回答的问题包括:
- 循环如何切块(tiling)
- 哪些维度并行执行
- 数据是否放入 shared memory 或寄存器
- 访问顺序是否连续、是否利于 cache
- 是否做向量化
- 是否展开循环、是否双缓冲
这也是 TVM 与手写 CUDA 在思想上最接近的地方。很多 CUDA kernel 优化,本质上都可以理解为 schedule 决策。
如果把同一个 matmul 分别用“差的执行方式”和“好的执行方式”实现,差别会非常明显。
差的情况通常像这样:
1 | for i |
这种实现的问题是数据复用差、访存开销大、并行度利用不足。
而更好的实现往往会引入:
- block 级切块
- shared memory staging
- 连续访存与向量化加载
- 更合适的线程映射
因此可以说,compute 决定正确性,schedule 决定性能上限。
TVM 如何表达 schedule
TVM 的 schedule 并不是直接手写一整个 kernel,而是对循环和数据访问施加一系列变换。例如:
1 | i, j = C.op.axis |
这段代码表达的是把原本的大循环拆成更适合块化执行的小循环,再重新组织迭代顺序。
进一步还可以加入并行、cache 和向量化策略:
1 | s[C].parallel(io) |
这些变换背后的语义分别对应:
split:把大问题切成小块reorder:调整循环次序parallel:暴露并行执行机会cache_read:显式管理数据在内存层级中的位置compute_at:控制中间结果何时产生、何时复用vectorize:让计算更贴近 SIMD 或向量指令
从这个角度看,TVM 并不是绕开硬件,而是试图用更系统化的方式表达“如何让硬件跑得更好”。
为什么 TVM 要把 compute 和 schedule 分开
这个分层设计不是为了抽象而抽象,而是直接服务于可移植性与优化空间。
如果把数学定义和性能优化完全写死在一起,那么每适配一种硬件,就可能要重写一套实现。手写 CUDA 往往就是这种模式:一个 kernel 同时承载了计算逻辑和优化细节,因此它在目标硬件上可能极致高效,但复用性有限。
TVM 的思路则是:
- compute 尽量保持稳定
- schedule 随硬件变化
这样带来的直接好处是,同一个算子定义可以绑定到不同 schedule,再进一步生成适配 CPU、GPU 或其他后端的代码。换句话说,TVM 试图把“通用的数学结构”和“特定硬件的优化经验”拆成两个可组合的层次。
schedule 之后,才是最终代码生成
当 compute 和 schedule 都确定下来后,TVM 才会继续 lower,生成目标代码。最终结果可能是:
- CUDA kernel
- 向量化 CPU 代码
- 某类 NPU 的特定指令序列
- 调用外部高性能算子库的封装代码
在 GPU 场景下,最后落地的形态通常已经接近:
1 | __global__ void matmul_kernel(...) { |
到这里,模型才真正走完了从高层表达到底层执行的完整路径。
schedule 的上限,取决于搜索
理解 TVM 不能停留在“schedule 很重要”这句话上。更进一步地说,现代 AI 编译器的核心竞争力,往往在于能否在巨大 schedule 空间中找到足够好的解。
以 matmul 为例,单是下面这些维度就足以构成一个巨大的组合空间:
- tile size
- thread/block 映射方式
- cache 策略
- vector width
- unroll 深度
- layout 变换方式
- 是否使用特殊 intrinsic
这些选择相互耦合,很难靠人工一次性写出全局最优解。因此,TVM 不仅提供 schedule 表达能力,也发展出了自动搜索能力。
常见的路径包括:
- 手写 schedule:适合理解机制,也适合构造基线
- AutoTVM:通过模板和搜索空间进行自动调优
- Ansor / AutoScheduler:自动生成候选 schedule,再通过 cost model 与实测反馈搜索最优方案
其核心流程大致是:
1 | compute |
因此,TVM 中的 schedule 不只是“写优化规则”,更是“把硬件上的数据流优化问题转化为可搜索问题”。
TVM 的代价:用编译时间换运行效率
TVM 的优势从来不是“零成本”。尤其一旦引入 autotuning,编译时间会显著上升。
大致可以分成三种情况:
1. 普通编译
如果只是把模型从 Relay 编译到目标后端,不做大规模调优,开销通常是秒级到几十秒,和常见推理编译器相比并不夸张。
2. 自动调优
一旦进入 AutoTVM 或 Ansor 这类搜索流程,时间开销就可能上升到分钟甚至小时级。因为系统需要在真实硬件上反复试跑候选方案,才能知道哪种 schedule 更优。
3. 离线调优复用
工业落地时更常见的方式是离线调优:在固定模型、固定硬件上完成一次搜索,然后保存 tuning log,后续重复部署时直接复用。这也是 TVM “编译时间换运行时间”最现实的工程化方式。
这件事之所以合理,是因为很多推理系统属于“一次编译,多次执行”。如果模型会被长期、大规模调用,那么一次性较高的编译成本往往可以被持续收益摊薄。
TVM 适合什么场景
TVM 最典型的价值,并不一定出现在所有人最熟悉的通用 GPU 训练场景,而更常出现在异构、碎片化、需要深度定制的部署环境中。
边缘设备
边缘侧硬件种类非常多,CPU、GPU、DSP、NPU 的组合复杂,且不同设备的最佳执行方式差异明显。TVM 在这里的优势,是能把模型表示、图优化、算子调优和后端适配组织成一套统一框架。
AI 芯片和加速器
对于芯片厂商来说,TVM 往往不是一个现成推理框架,而更像一个编译器骨架。前端负责接收模型,后端则可以接入自定义 codegen、intrinsic、runtime 和外部算子库,从而适配自家硬件。
云端推理中的特定场景
在云端,TVM 并不是唯一主角。很多主流场景会优先选择 TensorRT、vLLM 之类更贴近特定 workload 的系统。但在 CPU 推理、长尾模型优化、异构适配或研究性系统中,TVM 依然有明显价值。
为什么 TVM 没有统治所有推理场景
TVM 很强,但它不是所有问题的统一答案。原因主要有三类。
第一,动态性问题。TVM 更擅长结构相对稳定、shape 相对可分析的场景;而很多大模型推理系统面对的是动态序列长度、动态 batch、复杂 runtime 状态管理,这部分挑战并不完全是静态编译问题。
第二,调优成本高。自动搜索能带来高性能,但也意味着更高的时间和工程成本,不适合所有快速迭代场景。
第三,工程门槛高。要真正把 TVM 用好,往往需要同时理解模型图、编译器 IR、底层代码生成、硬件特性和性能调优,这对团队能力提出了较高要求。
因此,TVM 更像是“适合在正确问题上发挥巨大作用”的系统,而不是一把对所有推理任务都同样高效的万能工具。
当目标是新硬件时,TVM 关注的是 target 与 backend
讨论 TVM 的一个高频场景,是“如何适配一类新硬件”。这个问题的重点不在于“支不支持 GPU”,而在于能否把硬件能力抽象成 TVM 能理解和优化的 target、后端和 runtime。
一条更完整的适配路线通常是:
1 | 硬件特性分析 |
1. 先理解硬件偏好
在做后端适配前,首先要回答几个问题:
- 计算单元是什么,偏向标量、向量还是矩阵引擎
- 支持哪些数据类型,如 fp16、bf16、int8、int4
- 内存层级如何组织,是否存在 local memory、scratchpad 或显式搬运
- 并行模型更像 CUDA、SIMD CPU,还是命令队列式 NPU
- 是否有现成 runtime API、intrinsic 或厂商算子库
这些信息会直接决定 schedule 空间该如何设计,也决定 codegen 的边界在哪里。
2. 决定接入层级
常见有两条路。
一条是 external codegen:如果厂商已经提供高性能算子库,可以让 TVM 在图层面识别并切分子图,再把相应部分下发给厂商库执行。
另一条是 native backend:如果要让 TVM 原生生成这类硬件的代码,就需要补齐 target、runtime、DeviceAPI、codegen、intrinsic lowering 等能力。
前者更适合已有成熟生态的专用加速器,后者则适合需要深度原生支持的新平台。
3. 把硬件能力转成可搜索的 schedule 空间
适配新硬件的关键,不是简单增加一个后端名字,而是把硬件特征翻译成真正可优化的 knobs。例如:
- tile 大小
- vector width
- cache_read / cache_write 策略
- local memory 使用方式
- unroll 深度
- double buffer
- layout transform
- tensor intrinsic
如果硬件存在专用矩阵指令,那么还需要通过 tensorize 等机制,把高层 pattern lower 到这些 intrinsic 上。否则,即便图和 compute 都是正确的,最终生成的也可能只是性能普通的标量或向量代码。
4. runtime 同样决定成败
很多时候,问题不只在 kernel 本身。设备内存分配、host-device copy、stream 或 queue、同步机制、多设备调度、图执行流程,都会影响最终可用性。一个没有 runtime 配套的后端,很难真正落地。
一个总结:TVM 的价值到底在哪里
如果只看表面,TVM 像是在“自动生成高性能 kernel”;但从编译器角度看,它真正提供的是一整套分层组织复杂性的方式:
- 用 Relay 表达模型结构
- 用 compute 表达数学计算
- 用 schedule 表达执行策略
- 用搜索在巨大优化空间中逼近最优解
- 用后端与 runtime 把这些能力落到具体硬件
这套分层的意义在于,模型语义、硬件特性和性能优化不再混成一团,而是可以分别建模、分别演化,再通过编译链路组合起来。
因此,理解 TVM 的关键,不是记住几个术语,而是建立这样一种认识:模型变成高性能 kernel,并不是一步完成的;它要先变成图,再变成张量级计算,再变成面向硬件的数据流组织方式,最后才落成真正可执行的代码。TVM 所做的,就是把这条路径系统化。
- Title: tvm
- Author: Ikko
- Created at : 2025-01-15 14:38:07
- Updated at : 2026-05-05 21:16:21
- Link: http://ikko-debug.github.io/2025/01/15/tvm/
- License: This work is licensed under CC BY-NC-SA 4.0.