TVM OSDI 2018 — 把 Halide 思想搬到深度学习

待复核

是什么

TVM 是 2018 年 OSDI 的一篇论文，提出了深度学习的端到端优化编译器。给它一个 PyTorch / TensorFlow 训出来的模型，它能编出在 CPU、GPU、嵌入式芯片、自研加速器上都跑得不慢的代码。

日常类比：传统 vendor 出算子库（cuDNN / MKL / TensorRT），相当于”每种厨房（硬件）都得请一位专属大厨手写菜谱”。TVM 是一位会调度的总厨——一份菜谱（模型）扔过去，它先看整盘菜的搭配（图级优化），再为每口锅单独写执行计划（算子级 schedule），最后让一个机器学习模型试 1000 种火候组合挑最快的（autotune）。

它的核心论点：编译器，不是手工算子库，才是 ML 系统正确的抽象层。

为什么重要

不理解这篇论文，下面这些事都没法解释：

为什么 PyTorch 2.0 的 torch.compile / Google IREE / Meta Glow 都走”IR + autotune”路线——TVM 是这条路的开山参考
为什么深度学习落地手机/嵌入式不再依赖 vendor SDK——通用编译器代替了手工算子库
为什么 cuDNN / MKL 不是不可超越——人手只能调几十个 case，机器能搜上千组合
为什么 MLIR 这种”多层 IR”思想会火——TVM 已经验证过分层 IR + 后端可插拔的工程价值

核心要点

TVM 的设计是三层栈：

图级 IR（论文用 NNVM，后来换成 Relay）：把模型表达成”算子组成的计算图”。在这层做算子融合（conv + bias + relu 合一次）、布局变换、常量折叠。类比：先看整盘菜的搭配——哪些工序能合并、哪些原料能预处理。
算子级 schedule（继承 Halide）：单个算子（如矩阵乘）先描述”算什么”（compute），再单独写”怎么算”（schedule：tile 切块多大、循环展开、绑到哪几个线程、用 SIMD——一次算一排数的指令——还是张量原语）。Halide 2013 把这个分离做到图像处理里的 stencil（滑动窗口滤波）；TVM 把它推广到张量。
ML 自动调优（AutoTVM）：人没法手调 1000 种 schedule 组合。给一个 schedule 模板和参数空间（split 因子、tile 尺寸等），用 XGBoost 学”每种组合在这块硬件上大概多快”，挑前 N 个实测，迭代搜索。这是论文最具突破性的一步——在深度学习算子调度搜索里，较早把 ML cost model 嵌进编译器。

三层叠起来叫 end-to-end 编译：上接框架，下接硬件，中间靠 IR 解耦。

论文还顺带提出了 VTA（Versatile Tensor Accelerator）——一个可编程的开源张量加速器，用来证明这套栈对全新硬件也能定制：RTL、运行时、到 TVM 后端 codegen 都打通。新硬件不必等 vendor SDK，自己接 TVM 即可。

实践案例

案例 1：同一个模型编到不同硬件

import tvm
from tvm import relay

mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)

# 编到 ARM 树莓派
target_arm = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu -mcpu=cortex-a72")
lib_arm = relay.build(mod, target=target_arm, params=params)

# 编到 Nvidia GPU
target_cuda = tvm.target.Target("cuda -arch=sm_75")
lib_cuda = relay.build(mod, target=target_cuda, params=params)

逐部分解释：

from_onnx：把已训好的模型读成图级 IR（现代栈用 Relay）
Target(...)：告诉编译器”锅是哪一口”——CPU 还是 GPU、什么指令集
relay.build：同一份 mod，换 target 就换优化 pass、切块和向量化；上层模型一行不动

案例 2：写一个调度，看出 schedule 的力量

from tvm import te

A = te.placeholder((1024, 1024), name="A")
B = te.placeholder((1024, 1024), name="B")
k = te.reduce_axis((0, 1024), name="k")
C = te.compute((1024, 1024), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))

s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], 32, 32)
s[C].vectorize(yi)
s[C].parallel(xo)

逐部分解释：

te.compute：只写”算什么”（矩阵乘），不写怎么并行
tile(..., 32, 32)：把大循环切成 32×32 小块，方便塞进缓存
vectorize(yi)：内层用 SIMD 一次算一排；parallel(xo)：外层多核并行
结果数值不变，只改执行顺序——换硬件换 schedule，性能可差几十倍

案例 3：AutoTVM 让机器自己搜

@autotvm.template("conv2d_nchw")
def conv2d_template(N, H, W, CO, CI, KH, KW):
    cfg = autotvm.get_config()
    # 教学示意：axis_x 来自算子某条空间轴（如输出高），由 te.compute 的 axis 提供
    axis_x = ...  # e.g. s[out].op.axis[2]
    cfg.define_split("tile_x", axis_x, num_outputs=3)  # 拆成 3 段可搜切分
    cfg.define_knob("unroll", [0, 1])

tuner = autotvm.tuner.XGBTuner(task)
tuner.tune(n_trial=1000)

逐部分解释：

@autotvm.template：声明”可搜的菜谱骨架”；axis_x 是要搜的那条循环轴
define_split / XGBTuner：画出 (tile, unroll) 空间，用 XGBoost 估速再实测
论文报告：在 Mali / Nvidia / ARM 上，调出的算子可与 cuDNN / MKL / TensorRT 可比甚至更优

踩过的坑

autotune 慢——单算子搜 1000 候选要几十分钟到几小时；2020 年后 Ansor / MetaSchedule 才把这块磨平到可接受。
cost model 换硬件就要重训——XGBoost 学的是”这块 GPU 上时延”，换张卡分布就漂；冷启动阶段常被人工库回打。
schedule DSL 仍偏底层——“算法 vs 调度解耦”听着优雅，写出来的 schedule 像汇编技巧合集，调一组 tile 因子要试很多次。
早期图级 IR 偏静态 shape——输入长度会变的模型（聊天里越聊越长的 KV cache）要等到后来的 Relax / TIR 才好用。
新算子覆盖度——遇到没见过的算子（如新 attention 变体），TVM 没现成 schedule 模板，fallback 路径慢，得手写补。

适用 vs 不适用场景

适用：

把训好的模型部署到非主流硬件（嵌入式 CPU、Mali GPU、自研 NPU / FPGA）
想榨干特定算子在特定硬件上的最后 20% 性能
多硬件统一部署栈——避免每个平台一套 SDK
学术 / 公司内部 ML 编译器原型——TVM 是公认参考实现

不适用：

训练（TVM 主打推理；训练有 PyTorch / JAX 自家编译器）
极致动态 shape（早期偏静态；近年 Relax 才补齐）
想要”开箱即用 5 分钟跑通”——autotune 那一步躲不过
vendor 已经把硬件吃透的场景（Nvidia 服务端 fp16 推理，TensorRT 通常更香）

历史小故事（可跳过）

2013：MIT 的 Ragan-Kelley 等发表 Halide，第一次把”算法 vs schedule”分离做到工业级，但只面向图像处理 stencil。
2015-2017：Tianqi Chen（XGBoost / MXNet 作者）在做 NNVM 时意识到——框架内置算子库不可持续，不如做成编译器。
2018：OSDI 论文发表；同年主流框架仍在堆 vendor 算子库，TVM 走通用编译器路线，并带上 VTA 证明硬件可定制。
2019-2022：Ansor / MetaSchedule 把搜索效率提升 10×；MLIR / IREE / TorchInductor 吸收其设计。
今天：Apache TVM 仍是开源参考；分层 IR + 自动调优已渗透现代 DL 编译栈。

学到什么

分离”算什么”与”怎么算”——Halide 留给 TVM、再留给整个 ML 编译器领域的核心 idea
调度空间太大就让机器搜——人手调几十个 case，cost model 能搜上千个，结果常超人手
分层 IR 是后端可插拔的关键——图级做高层优化、算子级做低层调度，新硬件只需补 codegen
ML 系统的瓶颈常是工程——TVM 把”模型 → 任意硬件”的部署成本从月级压到天级

关联

halide —— 算法 vs schedule 分离的开山之作；TVM 把它从图像处理推到深度学习
llvm —— TVM 算子最终通过 LLVM 后端落到机器码
xla-compiler —— 同时代另一条 ML 编译器路线（XLA 偏静态、TVM 偏可调度）
tensorflow-osdi-2016 —— 同 OSDI 框架背景；TVM 与之互补
mlir —— 后续吸收 TVM 分层 IR 思想
feautrier-polyhedral —— 多面体调度；TVM 的循环变换与之有思想相通
attention —— 现代模型核心算子，常被当 TVM 调优 benchmark

反向链接

passnet-graph-compiler —— PassNet — 让大模型给图编译器写优化 pass
taso-2019 —— TASO — 让机器自己发现深度学习图重写规则

TVM OSDI 2018 — 把 Halide 思想搬到深度学习

是什么

为什么重要

核心要点

实践案例

案例 1：同一个模型编到不同硬件

案例 2：写一个调度，看出 schedule 的力量

案例 3：AutoTVM 让机器自己搜

踩过的坑

适用 vs 不适用场景

历史小故事（可跳过）

学到什么

延伸阅读

关联

反向链接