2、TVM与Relay基础：TVM架构、Relay IR核心概念、计算图与算子、Relay与TIR的关系

TVM（Tensor Virtual Machine）是一个端到端的深度学习编译栈，其核心设计目标是将深度学习模型从各种前端框架（PyTorch、TensorFlow、ONNX等）统一转换到中间表示（IR），并通过多层优化与代码生成，最终部署到CPU、GPU、NPU乃至FPGA等异构硬件上。在TVM的体系结构中，Relay IR扮演着“高级图优化层”的角色，而TIR（Tensor Intermediate Representation）则承担“底层算子优化与代码生成”的职责。理解Relay与TIR的分工与衔接，是掌握TVM编译流程、进行图优化和端侧部署的关键前提。本章将系统梳理TVM的整体架构、Relay IR的核心概念、计算图与算子的表达方式，以及Relay与TIR之间的转换关系。

上图展示了TVM编译栈的核心分层：前端框架经过Relay转换器统一为计算图，Relay层负责图级别的优化（如算子融合、布局转换），随后通过Lowering过程将Relay表达式转换为TIR，TIR层执行底层循环变换、存储优化和线程映射，最终由后端代码生成器输出目标机器码。整个流程中，Relay与TIR的边界清晰：Relay关注“图”和“数据流”，TIR关注“循环”和“存储”。下面我们逐一展开。

2.1 TVM架构总览：从模型到部署的编译流水线

TVM的架构设计遵循“前端无关、后端无关”的原则。其核心组件包括：前端解析器（Frontend）、Relay IR、TIR、自动调优模块（AutoTVM / Ansor）以及运行时（Runtime）。前端解析器负责将不同框架的模型转换为Relay IR，这一过程称为“Relay转换”。例如，PyTorch模型通过torchvision.models加载后，使用from_pytorch接口转换为Relay计算图。Relay IR是一种函数式、可扩展的图表示，它支持显式的数据流依赖、类型推导和函数式变换，这使得图优化Pass可以安全地重写计算图。

TVM的编译流程可以概括为以下阶段：

1. 模型导入：将ONNX/PyTorch/TensorFlow模型解析为Relay IR。
2. 图优化：在Relay层执行算子融合（如Conv+BN+ReLU融合）、常量折叠、死代码消除、布局转换等Pass。
3. Lowering：将优化后的Relay表达式转换为TIR，这一过程会为每个Relay算子生成对应的TIR计算函数。
4. 底层优化：在TIR层进行循环分块、向量化、缓存重排、线程绑定等操作。
5. 代码生成：通过LLVM、CUDA等后端生成目标代码，并封装为可调用的函数。
6. 运行时部署：将编译产物与TVM Runtime集成，支持在端侧设备上执行推理。

这种分层设计使得TVM能够灵活地适配新硬件：只需为新的后端实现TIR到机器码的代码生成器，而无需修改前端和Relay优化逻辑。

2.2 Relay IR核心概念：函数式图表示与类型系统

Relay IR是一种基于表达式树的强类型语言，其核心数据结构包括：Expr（表达式）、Function（函数）、Call（调用节点）、Var（变量）、Constant（常量）以及Tuple（元组）。每个Relay表达式都带有静态类型信息，这为编译期优化提供了基础。例如，一个卷积层的Relay表示如下：

# Relay 示例：Conv2D + BiasAdd + ReLU
fn (%x: Tensor[(1, 3, 224, 224), float32],
    %w: Tensor[(64, 3, 7, 7), float32],
    %b: Tensor[(64), float32]) -> Tensor[(1, 64, 112, 112), float32] {
  %0 = nn.conv2d(%x, %w, strides=[2, 2], padding=[3, 3, 3, 3]);
  %1 = add(%0, %b);
  %2 = nn.relu(%1);
  %2
}

Relay的类型系统支持形状推导和数据类型检查，这有助于在编译早期发现维度不匹配等错误。此外，Relay还支持AD（自动微分）扩展，但推理部署中更常用的是其图变换能力。Relay的Pass基础设施允许开发者编写自定义优化，例如SimplifyInference用于将BatchNorm折叠到卷积中，FoldScaleAxis用于合并缩放因子。

理解Relay的关键在于：它是一等公民的图IR，允许将子图作为函数传递，这为算子融合提供了便利。例如，将Conv2D和ReLU融合为一个函数，可以避免中间张量的内存往返。

2.3 计算图与算子：数据流表达与算子语义

计算图是深度学习中模型的结构化表示，Relay中的计算图由Call节点和Var节点构成的有向无环图（DAG）。每个Call节点对应一个算子（Op），算子具有名称、属性和输入输出类型。TVM内置了丰富的算子库，包括nn.conv2d、nn.dense、nn.batch_norm、nn.max_pool2d等。算子属性（如卷积的stride、padding、dilation）在Relay中以Attrs对象存储，并在Lowering时传递给TIR。

计算图的优化通常围绕“算子融合”和“数据布局”展开。以MobileNet为例，其深度可分离卷积包含DepthwiseConv2D和PointwiseConv2D，在Relay中可以通过FuseOps Pass将两者融合为一个算子组，减少内核启动开销。此外，布局转换（如NCHW到NHWC）在端侧NPU上往往能显著提升性能，Relay的AlterOpLayout Pass可以自动插入布局转换节点。

算子语义的准确性至关重要。例如，nn.conv2d的data_layout和kernel_layout属性必须与输入张量匹配，否则Lowering阶段会报错。在编写自定义算子时，需要同时提供Relay注册信息和TIR计算函数，确保图优化和代码生成都能正确处理。

2.4 Relay与TIR的关系：从图到循环的Lowering机制

Relay与TIR的关系可以概括为“图级抽象”与“算子级实现”的关系。Relay负责描述“做什么”（计算图结构），TIR负责描述“怎么做”（循环、存储、线程）。Lowering过程将每个Relay Call节点转换为对应的TIR函数，这一过程由relay.backend.lower接口驱动。例如，Relay中的nn.conv2d会被Lower为TIR中的tir.conv2d实现，该实现包含多层嵌套循环、滑动窗口计算和累加器初始化。

Lowering并非简单的“一对一”映射，而是包含一系列TIR变换Pass：

• 循环分块（Loop Partition）：将大循环拆分为适合缓存的小块。
• 向量化（Vectorize）：将标量运算转换为SIMD指令。
• 存储重排（Storage Rewrite）：调整数据在共享内存/寄存器中的布局。
• 线程绑定（Thread Binding）：将循环轴映射到GPU的blockIdx/threadIdx。

一个典型的Lowering流程如下：Relay图 → LowerToTE（生成TE表达式） → Schedule（应用调度原语） → Build（生成TIR） → 后端代码生成。其中，TE（Tensor Expression）是TIR的前身，它允许开发者用类似NumPy的语法描述计算，并通过调度原语（如split、reorder、vectorize）控制循环结构。

2.5 常见误区与总结复盘

在学习和使用TVM/Relay过程中，开发者常陷入以下误区：

• 混淆Relay与TIR的职责：试图在Relay层进行循环优化，或在TIR层进行图结构修改，导致代码复杂且难以维护。正确做法是：图优化在Relay层，循环优化在TIR层。
• 忽略类型推导：Relay的强类型系统是优化和安全性的基础。如果自定义算子未正确注册类型，可能导致Pass无法匹配，甚至编译崩溃。
• 过度依赖自动调优：AutoTVM虽然能自动搜索最佳调度，但在端侧资源受限时，手动编写关键算子的TIR调度往往更可控。例如，在DSP上部署时，手动指定循环展开因子比自动搜索更高效。
• 忽视布局转换开销：在Relay中插入布局转换节点虽然能提升算子效率，但转换本身也有内存拷贝成本。需要评估净收益，必要时使用layout_free标记算子。

总结复盘：本章的核心是建立TVM编译栈的整体认知。Relay作为图IR，提供了高层次的优化空间；TIR作为底层IR，提供了精细的硬件适配能力。理解两者的关系，有助于在端侧部署时快速定位性能瓶颈——是图结构冗余，还是循环实现低效。下一阶段的学习应聚焦于如何编写自定义Relay Pass和TIR调度，以应对实际部署中的特殊需求。