第3章：TVM框架入门

说实话，我第一次接触TVM时，心里是有点抵触的。那时候手头有个模型要在嵌入式设备上跑，TensorFlow Lite跑不动，手写优化又太累。后来硬着头皮啃了TVM，才发现——嗯，这东西确实有两把刷子。

这一章，咱们就聊聊TVM的核心架构。我会从三个层面展开：整体架构长什么样、中间表示Relay IR怎么用、以及一条完整的编译流水线是怎么走的。你跟着我走一遍，基本就能上手了。

3.1 TVM架构概览

TVM的架构，说白了就是一套“把模型描述变成高效可执行代码”的工具链。它不关心你模型是用PyTorch还是TensorFlow写的，它只关心你能不能转成它的中间表示。

我习惯把TVM分成三层来看：

  • 前端层：负责把各种框架的模型“吃进来”，转成统一的IR。比如从ONNX、TensorFlow、PyTorch导入。
  • 优化层：这是TVM的核心。它会对IR做各种变换——算子融合、常量折叠、内存规划等等。说白了，就是让计算图跑得更快。
  • 后端层：把优化后的IR“吐出去”，生成目标硬件能跑的代码。可以是LLVM、CUDA、OpenCL，甚至是Vulkan。

你想想看，这三层分工明确，各管各的。前端只管“翻译”，优化只管“改造”，后端只管“生成”。这种设计，让我在项目里调试起来特别方便——哪一层出了问题，直接定位就行。

核心要点：TVM的架构设计遵循“前端-优化-后端”三层分离。每一层都可以独立替换或扩展。我在实际项目中，就曾把前端从ONNX换成PyTorch，后端从LLVM换成CUDA，中间优化层几乎没动。

下面这张图，是我自己画的TVM整体架构流程。你看一眼，心里就有数了。

前端层 PyTorch / TF / ONNX → Relay IR 优化层 算子融合 常量折叠 内存规划 自动调度 后端层 LLVM / CUDA / OpenCL → 可执行代码 TVM三层架构：前端吃模型，优化做变换，后端吐代码 输入模型 Relay IR 优化后IR 目标代码

3.2 Relay IR详解

Relay IR是TVM的中间表示，也是整个框架的“灵魂”。我刚开始学的时候，觉得它就是个计算图。后来用多了才发现，它比普通计算图要灵活得多。

Relay IR有几个关键特性：

  • 函数式表达：每个节点都是一个函数调用，支持高阶函数和闭包。说白了，你可以把一段计算当成一个函数传来传去。
  • 类型系统：Relay有完整的类型推断机制。你写代码时不用显式声明类型，编译器会自动推导。这个特性在调试时特别有用——类型不对，编译期就报错，不会等到运行时才崩。
  • 模块化：一个Relay程序由多个模块组成，每个模块包含多个函数定义。这种设计让代码复用变得很简单。

我的经验：刚开始接触Relay IR时，我建议你多打印一下IR的文本表示。TVM提供了relay.ir_builder和relay.transform工具，你可以直观地看到每一步变换后的IR长什么样。我曾经靠这个手段，排查出一个算子融合没生效的bug。

来看一个简单的Relay IR示例。假设我们有一个加法操作：

import tvm
from tvm import relay

# 创建两个变量节点
x = relay.var("x", shape=(3, 4), dtype="float32")
y = relay.var("y", shape=(3, 4), dtype="float32")

# 加法操作
z = relay.add(x, y)

# 构建函数
func = relay.Function([x, y], z)

# 打印IR
print(func)

输出会是这样：

fn (%x: Tensor[(3, 4), float32], %y: Tensor[(3, 4), float32]) -> Tensor[(3, 4), float32] {
  %0 = add(%x, %y)
  %0
}

你看，Relay IR的文本形式非常直观。每个变量都有类型标注，每个操作都有明确的输入输出。这种设计，让我在做算子融合时能清晰地看到哪些操作可以合并。

3.3 TVM的编译流程

TVM的编译流程，我习惯分成三个阶段来讲：前端导入、中间优化、后端生成。每个阶段都有对应的Pass（编译优化步骤）来处理。

3.3.1 前端：从模型到Relay IR

前端的工作，就是把各种框架的模型“翻译”成Relay IR。TVM提供了多个前端接口：

框架 导入接口 注意事项
TensorFlow relay.frontend.from_tensorflow 需要冻结图（freeze graph）
PyTorch relay.frontend.from_pytorch 支持torch.jit.trace导出的模型
ONNX relay.frontend.from_onnx 算子支持最全，推荐使用
Keras relay.frontend.from_keras 需要h5格式的模型文件

避坑指南：我曾经在导入PyTorch模型时，遇到一个算子TVM不支持。折腾了半天，最后发现可以用relay.frontend.from_pytorch的custom_convert_map参数，手动注册一个自定义转换函数。所以，遇到不支持的算子别慌，先查文档，看看能不能自己写转换逻辑。

3.3.2 优化：Pass管线

优化阶段，是TVM最精彩的部分。它通过一系列Pass，把Relay IR变成更高效的版本。我常用的Pass有这些：

  • 算子融合（FuseOps）：把多个连续的小算子合并成一个。比如Conv2D + ReLU + BatchNorm，可以融合成一个算子。这样能减少内存访问次数，提升缓存命中率。
  • 常量折叠（FoldConstant）：把编译期就能算出来的常量表达式提前计算好。比如add(2, 3)直接变成5。
  • 消除死代码（EliminateDeadCode）：删掉那些输出不会被用到的节点。这个Pass看起来简单，但在大模型里能省不少计算量。
  • 布局变换（AlterOpLayout）：调整张量的内存布局，让它更适配目标硬件。比如把NCHW变成NHWC。

这些Pass的执行顺序，我建议你按“先化简、再融合、最后布局调整”来排。为什么？因为先做常量折叠和死代码消除，能让图变得更干净，后面的融合和布局调整才能做得更彻底。

# 一个典型的Pass管线配置
from tvm import relay

# 创建Pass管理器
with tvm.transform.PassContext(opt_level=3):
    # 应用一系列Pass
    mod = relay.transform.FuseOps()(mod)
    mod = relay.transform.FoldConstant()(mod)
    mod = relay.transform.EliminateDeadCode()(mod)
    mod = relay.transform.AlterOpLayout()(mod)

3.3.3 后端：生成目标代码

后端阶段，TVM会把优化后的Relay IR编译成目标硬件的代码。这里涉及两个关键组件：

  • TE（Tensor Expression）：张量表达式，用来描述算子的计算逻辑。你可以用TE手写算子，然后让TVM自动调度。
  • AutoTVM / AutoScheduler：自动调优工具。它会搜索最优的调度参数，比如循环展开因子、线程块大小等。

我个人的习惯是，先用默认调度跑一遍，看看性能瓶颈在哪。然后针对热点算子，用AutoTVM做自动调优。这样既省时间，又能拿到不错的性能。

# 后端编译示例
from tvm import relay
import tvm

# 假设mod已经是优化后的Relay IR
# 指定目标硬件
target = tvm.target.Target("llvm -mcpu=core-avx2")

# 编译成可执行模块
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target)

# 保存编译后的库
lib.export_library("deploy.so")

关键点：TVM的编译流程是“前端导入 → 优化Pass → 后端生成”。每一步都可以独立调试和替换。我在项目中，经常把优化Pass的顺序调来调去，看看哪种组合效果最好。这种灵活性，是TVM相比其他编译器的一大优势。

好了，这一章的内容就到这。TVM的架构、Relay IR、编译流程，这三个东西你串起来理解，就能明白TVM是怎么工作的了。下一章，咱们会深入算子融合的具体实现，到时候我会拿实际案例来演示。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →