TVM | 基本概念

TVM 概念整理

1 Relay IR

• 功能：Relay IR 是 TVM 的高层中间表示（IR），专为机器学习模型设计。它支持函数式编程范式、静态类型系统和自动微分，能够表达复杂的模型结构（如控制流、递归、闭包）。
• 核心特性：
  • 代数数据类型（ADT）：用于表示树状结构（如递归神经网络）。
  • 闭包（Closures）：支持高阶函数，允许函数作为参数传递或返回值。
  • 全局与局部变量：全局变量以 @前缀标识（如 @func），局部变量以 %前缀标识（如 %x）。
  • 模块化（Module）：管理多个全局函数，支持相互递归调用。

2 Relay Pass

• 功能：对 Relay IR 进行优化的转换流程，包括常量折叠、死代码消除、算子融合等。
• 类型：
  • 数据流优化：针对纯计算图部分（如算子融合）。
  • 控制流优化：处理条件分支（if-else）、ADT 匹配等复杂逻辑。

3 TE （Tensor Express）

TE（Tensor Expression，张量表达式） 是一种用于定义和优化张量计算的领域特定语言（DSL）。它位于 TVM 栈的中间层，介于高层计算图（如 Relay IR）和底层硬件指令（如 TIR）之间，核心作用是为算子（如向量加法、矩阵乘法）提供声明式描述，并通过调度优化实现高性能硬件适配。

声明计算逻辑

用户通过 TE 的 API（如 te.placeholder和 te.compute）描述张量操作，例如向量加法：

n = te.var("n")
A = te.placeholder((n,), name="A")  # 输入张量
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")  # 计算规则[1,3](@ref)

编译流程中的角色

TE 需通过 te.create_prim_func转换为 TIR（Tensor IR），再编译为硬件代码（如 LLVM IR 或 CUDA），无法直接生成可执行文件。

4 调度原语（Schedule Primitives）

• 功能：在 TE 层定义计算如何映射到硬件执行，通过调整循环嵌套、内存访问和并行策略提升性能。
• 常见原语：
  • 循环变换：split（拆分循环）、reorder（重排序）、unroll（循环展开）。
  • 并行化：parallel（多线程）、vectorize（向量化）。
  • 内存优化：cache_read/cache_write（缓存数据）、bind（绑定到 GPU 线程块）。

4.1 循环调度原语（Loop Scheduling Primitives）

控制循环结构的执行方式，优化计算效率和内存访问模式：

split（循环拆分）

• 将大循环拆分为嵌套小循环，便于并行化或向量化。
• 示例：将循环轴拆为outer和inner两层，用于GPU线程块划分。

bx, tx = s[C].split(C.op.axis[0], factor=64)   # 拆分为64大小的块

fuse（循环融合）

• 合并多个连续循环，减少循环开销。
• 适用场景：减少小规模循环的控制流开销。

reorder（循环重排序）

• 调整嵌套循环的顺序，优化缓存局部性。
• 示例：将内存连续访问的循环移至内层。

unroll（循环展开）

• 展开循环体，减少分支预测开销。
• 硬件适配：CPU中利用指令级并行（ILP），GPU中减少线程同步

4.2 内存调度原语（Memory Scheduling Primitives）

优化数据存储布局和传输效率：

cache_read/ cache_write（缓存读写）

• 将数据缓存到快速内存（如GPU共享内存），减少全局内存访问延迟。
• 示例：卷积计算中缓存输入/输出张量。

storage_align（存储对齐）

• 调整内存地址对齐，适配硬件SIMD指令要求（如AVX-512需64字节对齐）。

pad（边界填充）

• 为张量添加填充，避免内存访问越界并适配硬件向量化。

4.3 AutoTVM

• 功能：自动化搜索最优调度参数（如循环分块大小、并行策略），避免手动调优的耗时。
• 工作流程：
  1. 定义搜索空间：通过 cfg.define_split()等 API 指定参数范围。
  2. 选择调优器：如随机搜索（RandomTuner）、遗传算法（GATuner）或基于模型的调优（XGBTuner）。
  3. 评估配置：在目标硬件上测量不同配置的性能，选择最佳方案。

4.4 计算图变换原语（Computation Graph Transformation）

针对算子间依赖的全局优化：

compute_inline（内联计算）

• 将小算子内联到调用处，消除临时内存分配。
• 适用场景：融合Conv2D + ReLU等连续算子。

compute_at（指定计算位置）

• 将一个算子的计算嵌入到另一个算子的循环中，提升数据局部性。

5 TIR（Tensor IR）

• 功能：TVM 的低层中间表示，将 TE 或 Relay 生成的计算转换为硬件可执行代码。支持显式内存分配、循环嵌套和硬件特性（如 GPU 线程同步）。
• 应用场景：作为 Relay 或 TE 的编译目标，最终生成 LLVM IR、CUDA 等后端代码。

6 Runtime 系统

• 功能：部署优化后的模型到多样硬件（CPU/GPU/嵌入式设备），支持动态执行和异构计算。
• 组件：
  • NDArray：统一的多维数据容器，跨硬件后端共享内存。
  • Virtual Machine：执行 Relay 或 TIR 编译的字节码，支持动态形状和控制流。

核心概念对比概览

概念	定位	关键能力
Relay IR	高层模型表示	支持控制流、闭包、ADT，适合复杂模型（如RNN、动态图）。
TE	算子级优化	声明式计算 + 调度原语，聚焦单算子性能优化。
AutoTVM	自动化调度优化	搜索最佳硬件参数，提升跨平台泛化能力。
TIR	低层代码生成	显式内存管理、硬件指令映射，衔接后端编译器（如LLVM）。
Runtime	部署与执行	异构硬件支持、动态加载模型，提供统一接口（如NDArray）。

总结

TVM 的架构覆盖了从模型表示（Relay IR）、算子优化（TE + 调度）、自动化调优（AutoTVM）到代码生成（TIR）和运行时部署的全流程。其核心价值在于：

1. 灵活性：通过 Relay 支持复杂模型逻辑，突破传统计算图的限制。
2. 性能可移植性：借助调度原语和 AutoTVM 适应多样硬件。
3. 编译优化一体化：从高层 IR 到底层代码生成无缝衔接，实现端到端优化。