3. TVM基础

TVM 作为端到端深度学习编译器，其核心价值在于将不同前端框架（PyTorch、TensorFlow、ONNX 等）描述的模型，统一转换为可优化、可部署到多种硬件后端（CPU、GPU、NPU、FPGA）的中间表示，并最终生成高效代码。理解 TVM 的架构分层与 IR 体系，是掌握 NPU 后端开发的前提。本章将系统拆解 TVM 的模块化设计、Relay IR 与 TE/TIR 的职责边界、Pass 基础设施的工作机制，以及从图到机器码的完整生成流程。

上图展示了 TVM 编译栈的核心分层与 IR 流转路径。从高层前端到低层硬件，每一层都承担明确的职责，并通过 Pass 基础设施进行优化。下面我们将逐一深入每个模块。

3.1 TVM 整体架构与设计哲学

核心原理：三层分离与统一中间表示

TVM 的架构设计遵循“前端无关、后端可插拔”的原则。其核心思想是将深度学习编译过程划分为三个独立层次：图级别（Graph Level）、张量级别（Tensor Level）和指令级别（Instruction Level）。每一层都有对应的中间表示（IR），层与层之间通过明确的 lowering 操作衔接。

• 图级别（Relay IR）：负责处理模型整体结构，包括数据流、控制流、类型推导和算子融合。它不关心具体算子的内部实现，只关注算子之间的连接关系。
• 张量级别（TE/TIR）：TE 用于描述算子的计算逻辑和调度策略（如循环顺序、并行度），TIR 则是经过 lowering 后的底层表示，包含显式的循环、内存访问和硬件原语。
• 指令级别（Codegen）：将 TIR 转换为目标硬件的源码或指令，例如 CUDA C、LLVM IR、NPU 汇编或专用指令集。

这种分层设计带来的直接好处是：新增一个硬件后端时，只需关注 TIR 到目标代码的转换以及 TIR Pass 的定制，而无需修改前端解析和 Relay 优化逻辑。

判断方法：如何定位一个编译问题属于哪一层

在实际开发中，快速定位问题所在层级是调试的关键。以下是一些判断依据：

现象	所属层级	典型排查手段
模型转换失败、类型不匹配、算子无法识别	Relay IR 层	打印 Relay IR（print(mod)），检查算子是否被正确映射
性能不达标、循环顺序不合理、内存占用过高	TE/TIR 层	查看 TIR 代码，分析循环嵌套和内存访问模式
生成代码编译错误、指令非法、寄存器溢出	Codegen 层	检查生成的源码或汇编，确认目标后端约束

案例拆解：一个卷积算子在不同层的表现

以 3x3 卷积为例，在 Relay 中它只是一个 nn.conv2d 节点，包含输入、权重、步长、填充等属性。经过 TE 层，它被展开为嵌套循环：for co, ci, kh, kw, h, w。在 TIR 层，这些循环会被进一步细化，例如将 kh 和 kw 循环展开，或者将 h 和 w 循环向量化。对于 NPU 后端，TIR 可能会被转换为调用专用硬件指令的形式，例如“矩阵乘累加指令”。

3.2 Relay IR：函数式图表示与优化

核心原理：可扩展的图 IR

Relay IR 是一种支持函数式编程语义的图表示。它不仅仅是简单的 DAG（有向无环图），还支持 let 绑定、if 条件、match 模式匹配等高级特性，这使得它能够表达动态模型（如循环神经网络中的变长序列）。Relay 中的每个节点都是一个 Expr，包括 Constant、Call、Function、Tuple 等。

Relay 的核心优势在于其类型系统。每个表达式都有明确的类型（如 Tensor[(1, 3, 224, 224), float32]），编译器可以利用类型信息进行优化，例如自动推断内存布局、检查算子输入是否合法。

操作流程：从前端模型到 Relay IR

1. 使用 from_torch、from_onnx 等接口导入模型。
2. TVM 将前端模型转换为 Relay 的 Module，其中包含一个主函数（main）和可能的子函数。
3. 执行一系列 Relay Pass，如 FoldConstant（常量折叠）、EliminateCommonSubexpr（公共子表达式消除）、FuseOps（算子融合）。

import tvm
from tvm import relay
# 假设 torch_model 是 PyTorch 模型
# shape_dict = {"input": (1, 3, 224, 224)}
# mod, params = relay.frontend.from_pytorch(torch_model, shape_dict)
# 打印 Relay IR
# print(mod)

常见误区：混淆 Relay 与 TE 的优化边界

3.3 TE/TIR：张量表达式与底层中间表示

核心原理：计算与调度分离

TE（Tensor Expression）是 TVM 中描述算子计算的 DSL（领域特定语言）。其核心设计是计算（compute）与调度（schedule）分离。计算部分只描述“要算什么”，例如 C = A + B；调度部分描述“怎么算”，例如循环分块（split）、向量化（vectorize）、并行化（parallel）、绑定线程（bind）。

TIR（Tensor Intermediate Representation）是 TE 经过 lowering 后的产物。它更接近硬件，包含显式的 For 循环、Buffer 访问、LetStmt 变量绑定等。TIR 是 Pass 优化和代码生成的直接操作对象。

操作流程：TE 到 TIR 的 lowering

1. 使用 te.compute 定义算子计算逻辑。
2. 使用 te.create_schedule 创建默认调度。
3. 通过 schedule 原语（如 split, reorder, vectorize）修改循环结构。
4. 调用 tvm.lower 将 TE 转换为 TIR。

import tvm
from tvm import te
# 定义矩阵加法
A = te.placeholder((1024, 1024), name="A")
B = te.placeholder((1024, 1024), name="B")
C = te.compute((1024, 1024), lambda i, j: A[i, j] + B[i, j], name="C")
s = te.create_schedule(C.op)
# 调度：将 i 循环分为两个循环，外层并行
i_outer, i_inner = s[C].split(C.op.axis[0], factor=4)
s[C].parallel(i_outer)
# 生成 TIR
ir = tvm.lower(s, [A, B, C], simple_mode=True)
print(ir)

案例拆解：为 NPU 定制 TIR 生成

假设目标 NPU 支持 16x16 的矩阵乘累加指令。在 TE 层，我们可以将卷积的矩阵乘法部分显式地表达为 tir.call_intrin，或者通过 TIR Pass 将标准的循环结构识别并替换为硬件指令。例如，在 TIR 中，一个 16x16 的循环块可以被替换为一条 mma.sync 指令（类似 GPU 的 Tensor Core）。对于 NPU，这可能是 npu.matrix_mul 指令。

3.4 Pass 基础设施：优化管线的管理引擎

核心原理：模块化与可组合的优化

TVM 的 Pass 基础设施提供了一套统一的机制来注册、调度和执行优化。一个 Pass 可以是一个函数，接收一个 IR 模块并返回一个新的 IR 模块。Pass 之间可以声明依赖关系，Pass 管理器负责按照依赖顺序执行它们。这种设计使得优化管线高度可配置，用户可以根据目标后端选择不同的 Pass 序列。

Pass 分为三类：Relay Pass（操作 Relay IR）、TIR Pass（操作 TIR）、目标后端 Pass（如 NPU 特有的内存分配 Pass）。

操作流程：编写并注册一个自定义 Pass

1. 定义 Pass 函数，输入输出均为 IRModule。
2. 使用 @tvm.ir.transform.module_pass 装饰器注册。
3. 将 Pass 添加到 Pass 序列中，通过 tvm.transform.Sequential 执行。

from tvm import ir
@ir.transform.module_pass(opt_level=1, name="MyNPUOpt")
class MyNPUOpt:
    def transform_module(self, mod, ctx):
        # 对 mod 进行变换
        return mod
# 使用
seq = tvm.transform.Sequential([MyNPUOpt()])
with tvm.transform.PassContext(opt_level=3):
    mod = seq(mod)

常见误区：Pass 执行顺序的陷阱

3.5 代码生成流程：从 TIR 到可执行模块

核心原理：目标描述与代码生成器

代码生成是编译流程的最后一步。TVM 通过 Target 对象描述目标硬件（如 target = tvm.target.Target("npu")）。代码生成器（Codegen）根据 Target 的配置，将 TIR 转换为目标源码。对于 NPU，这通常意味着生成 C++ 代码，其中嵌入对 NPU 驱动或运行时 API 的调用。

生成流程包括：TIR 预处理（如存储展平、寄存器分配）、源码生成（将 TIR 语句转换为目标语言）、编译打包（调用外部编译器如 gcc 或 LLVM，生成最终可执行文件或动态库）。

操作流程：完整的编译与部署链路

1. 构建 Relay 模型并执行图优化。
2. 将 Relay 转换为 TE 并执行调度。
3. Lower 到 TIR 并执行 TIR Pass。
4. 调用 tvm.build 生成目标模块。
5. 使用运行时（如 tvm.contrib.graph_executor）加载模块并执行推理。

# 假设 mod 是 Relay 模块，params 是参数字典
# target = tvm.target.Target("npu")
# with tvm.transform.PassContext(opt_level=3):
#     lib = relay.build(mod, target=target, params=params)
# 导出库
# lib.export_library("deploy.so")

案例拆解：NPU 后端的代码生成调试

假设在 NPU 上部署一个 MobileNet 模型，编译后推理结果错误。调试步骤：

1. 检查 TIR：打印 TIR，确认卷积循环是否被正确转换为 NPU 指令调用。
2. 检查生成源码：查看生成的 C++ 代码，确认 NPU 指令的参数（如地址、形状）是否正确。
3. 运行时日志：启用 NPU 驱动日志，查看指令执行是否报错。

通过对比 CPU 后端和 NPU 后端的 TIR 差异，可以快速定位是调度问题还是 Codegen 问题。

3.6 总结与复盘

本章梳理了 TVM 的核心架构。理解 Relay、TE/TIR、Pass 和 Codegen 的职责边界，是进行 NPU 后端开发的基础。关键要点：

• 分层思维：始终明确当前操作属于哪一层，避免跨层优化。
• Pass 是核心工具：NPU 后端的定制化优化主要通过编写 TIR Pass 实现。
• 调试方法：善用 print 和 PassContext 的 instrument 功能，观察 IR 变化。