跳转到内容
Jason's Study

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 的设计是三层栈：

  1. 图级 IR（Relay）：把模型表达成”算子组成的计算图”。在这层做算子融合（conv + bias + relu 合一次）、布局变换、常量折叠。类比：先看整盘菜的搭配——哪些工序能合并、哪些原料能预处理。

  2. 算子级 schedule（继承 Halide）：单个算子（如矩阵乘）先描述”算什么”（compute），再单独写”怎么算”（schedule：tile 多大、循环展开、绑到哪几个线程、用 SIMD 还是张量原语）。Halide 2013 把这个分离做到图像处理 stencil；TVM 把它推广到张量。

  3. ML 自动调优（AutoTVM）：人没法手调 1000 种 schedule 组合。给一个 schedule 模板和参数空间（split 因子、tile 尺寸等），用 XGBoost 学”每种组合在这块硬件上大概多快”，挑前 N 个实测，迭代搜索。这是论文最具突破性的一步——首次把 ML cost model 用进编译器搜索。

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

论文还顺带提出了 VTA（Versatile Tensor Accelerator）——一个可编程的开源张量加速器。VTA 不是论文核心，但用来证明这套栈对全新硬件也能定制：从 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)

target 一变，编译器走的优化 pass、生成的指令序列、循环切块、向量化方式全变；上层模型代码一行不动。

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

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)

tile / vectorize / parallel 是 schedule 原语——计算结果完全没变，只是改了执行顺序和并行度。同一份 compute 在不同硬件上换不同 schedule，性能差几十倍。

案例 3：AutoTVM 让机器自己搜

@autotvm.template("conv2d_nchw")
def conv2d_template(N, H, W, CO, CI, KH, KW):
    cfg = autotvm.get_config()
    cfg.define_split("tile_x", x, num_outputs=3)
    cfg.define_knob("unroll", [0, 1])


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

XGBoost 学”每种 (tile, unroll) 组合在这块 GPU 上大概多快”，论文报告：在 ARM Mali GPU、Nvidia GPU、ARM CPU 三类后端上，AutoTVM 调出来的算子打到与 cuDNN / MKL / TensorRT 可比甚至更优。

踩过的坑

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

适用 vs 不适用场景

适用：

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

不适用：

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

历史小故事（可跳过）

  • 2013：MIT 的 Ragan-Kelley 等发表 Halide，第一次把”算法 vs schedule”分离做到工业级，但只面向图像处理 stencil。
  • 2015-2017：华盛顿大学 Tianqi Chen（XGBoost / MXNet 作者）在做 NNVM 时意识到——框架内置算子库不可持续，每个新硬件都要重写一套，不如把它做成编译器。
  • 2018：OSDI 论文发表，开源到 Apache。同年 PyTorch / TensorFlow 都还在堆 vendor 算子库，TVM 是少数走”通用编译器”路线的项目。论文同时提出 VTA（Versatile Tensor Accelerator），证明栈对硬件可定制。
  • 2019-2022：Ansor / MetaSchedule 把搜索效率提升 10× 以上；MLIR / IREE / TorchInductor 一波 ML 编译器都吸收 TVM 的设计。
  • 今天：Apache TVM 仍是 ML 编译器领域的开源参考；其分层 IR + 自动调优思想已渗透到几乎所有现代 DL 编译栈。

学到什么

  1. 分离”算什么”与”怎么算”——这是 Halide 留给 TVM、再留给整个 ML 编译器领域的核心 idea
  2. 调度空间太大就让机器搜——人手能调几十个 case，cost model + 自动调优能搜上千个，结果常超人手
  3. 分层 IR 是后端可插拔的关键——图级 IR 做高层优化、算子级 IR 做低层调度，新硬件只需补底层 codegen
  4. ML 系统的瓶颈不是算法是工程——TVM 真正的价值是把”模型 → 任意硬件”的部署成本从月降到天
  5. 学术论文也能产生工业基础设施——这篇 OSDI 论文今天仍是 Apache TVM 项目主线，所有现代 DL 编译器都在引用

延伸阅读

关联

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