2、TVM基础概念:计算图、算子、张量、调度(Schedule)、编译流程(Relay/TIR/代码生成)概览
在深入TVM Micro端侧优化之前,必须先建立一套统一的“语言体系”。TVM的核心理念是将算法描述与硬件实现解耦:开发者用高层抽象描述“做什么”,TVM负责决定“怎么做”。本章将围绕五个核心概念——计算图、算子、张量、调度、编译流程——展开,帮助你在脑海中搭建起TVM的完整工作框架。理解这些概念是后续进行算子优化、内存规划和代码裁剪的基石。
上图展示了TVM的三层架构:前端表示层负责将模型转换为计算图、算子和张量;优化与变换层通过Relay IR和TIR进行图优化与算子变换,并应用调度策略;后端生成层则针对目标硬件(尤其是Micro级别的MCU或嵌入式CPU)生成高效代码。下面我们逐一拆解这些概念。
2.1 核心概念精讲
2.1.1 计算图:算法的结构化表达
核心原理:计算图(Computational Graph)是一种有向无环图(DAG),其中节点代表算子(如卷积、矩阵乘、激活函数),边代表张量(数据)的流动方向。它将复杂的神经网络或数值计算任务分解为一系列可独立优化的子任务。在TVM中,计算图是优化的起点,它决定了数据依赖关系和并行化潜力。
判断方法:识别一个计算图的关键在于:① 节点是否是无状态的(即输出仅由输入决定);② 边是否明确标注了数据的形状(shape)和数据类型(dtype)。在TVM中,计算图通常以Relay IR的形式存在,可以通过 relay.ir_pass 系列函数进行遍历和修改。
2.1.2 算子:计算的基本单元
核心原理:算子是计算图的最小执行单位,例如 nn.conv2d、add、multiply。在TVM中,算子分为两类:计算定义(Compute)和调度实现(Schedule)。计算定义描述“做什么”,调度实现描述“怎么做”。例如,一个矩阵乘法的计算定义是 C[i,j] = sum(A[i,k] * B[k,j]),而调度则决定循环的拆分、重排、向量化等。
操作流程:在TVM中注册一个自定义算子需要三步:① 使用 te.compute 定义计算逻辑;② 使用 te.schedule 创建默认调度;③ 通过 tvm.build 编译为可调用函数。对于Micro端,算子实现必须考虑寄存器数量、指令集特性(如是否支持SIMD)以及内存对齐。
| 算子类型 | 计算定义示例 | Micro端优化关注点 |
|---|---|---|
| 逐元素算子 | add, relu, clip | 向量化宽度、循环展开 |
| 归约算子 | sum, max, matmul | 累加器精度、分块策略 |
| 数据移动算子 | reshape, transpose | 避免实际数据拷贝,使用视图 |
| 复杂算子 | conv2d, depthwise | im2col变换、Winograd、间接卷积 |
2.1.3 张量:数据存储与访问的抽象
核心原理:张量(Tensor)是多维数组的抽象,包含形状(shape)、数据类型(dtype)和存储布局(layout,如NCHW或NHWC)。在TVM中,张量不仅是数据容器,还承载了作用域(如全局内存、共享内存、寄存器)和生命周期信息。对于Micro端,张量的内存分配策略直接影响性能——频繁的动态分配会导致堆碎片化,因此TVM Micro通常采用静态内存规划。
判断方法:在编写TVM算子时,需要明确张量的形状和步长(stride)。例如,一个形状为 (1, 3, 32, 32) 的NCHW张量,其通道维度的步长是 32*32=1024。错误的步长假设会导致越界访问。在Micro端,建议使用 tvm.tir.decl_buffer 显式声明缓冲区,以便编译器进行别名分析和内存复用。
2.2 调度(Schedule):性能优化的核心杠杆
核心原理:调度是TVM最强大的特性之一,它允许开发者在不改变计算逻辑的前提下,通过变换循环结构来适配硬件。常见的调度原语包括:split(拆分循环)、reorder(重排循环)、vectorize(向量化)、unroll(展开循环)、parallel(并行化)以及cache_read/cache_write(数据缓存)。
操作流程:一个典型的调度优化流程如下:
- 分析瓶颈:通过profiling确定是计算密集还是访存密集。
- 选择原语:计算密集则尝试向量化和展开;访存密集则尝试分块和缓存。
- 迭代调优:使用
tvm.autotvm或ansor自动搜索最佳调度参数。
# 示例:对矩阵乘进行分块和向量化调度
C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))
s = te.create_schedule(C.op)
# 拆分 i 维度为外层块和内层线程
i_outer, i_inner = s[C].split(C.op.axis[0], factor=4)
# 向量化内层循环
s[C].vectorize(i_inner)
案例拆解(小型场景):假设在Cortex-M7(支持双精度FPU)上执行一个 3x3 卷积。初始调度直接嵌套循环,每次计算一个输出像素。通过调度变换:① 将输入通道维度的循环展开(unroll),减少循环控制开销;② 将输出通道的循环向量化(利用FPU的SIMD能力)。结果:性能提升约2.3倍,代码体积仅增加8%。但代价是寄存器压力增大,需要仔细检查是否发生溢出(spill)。
cache_read 将数据加载到局部缓冲区,再进行计算。同时,避免过度展开循环导致代码膨胀,因为Micro端的Flash容量通常限制在1MB以内。
2.3 编译流程:从模型到机器码的旅程
核心原理:TVM的编译流程分为三个阶段:Relay(图优化)、TIR(张量表达式优化)和代码生成(后端实现)。Relay负责图级别的优化(如算子融合、常量折叠、布局转换);TIR负责循环级别的优化(如调度、内存规划);代码生成则根据目标架构生成C代码或机器码。
操作流程:以Micro端为例,完整流程如下:
| 阶段 | 输入 | 输出 | 关键Pass |
|---|---|---|---|
| Relay优化 | Relay IR | 优化后的Relay IR | 算子融合、布局转换、常量折叠 |
| Lowering | Relay IR | TIR | 算子Lowering为te.compute表达式 |
| TIR优化 | TIR | 优化后的TIR | 调度应用、存储重写、向量化 |
| 代码生成 | TIR | C源文件 / 目标文件 | CodegenC、CodegenLLVM、Micro专用后端 |
案例拆解(完整闭环):假设我们要在STM32G474(Cortex-M4,192KB SRAM)上部署一个MobileNetV1的深度可分离卷积层。
- 背景与约束:输入特征图
112x112x32,卷积核3x3,SRAM仅能容纳约192KB / 4B = 48K个float32数据,而输入本身就有112*112*32 = 401K个数据,必须分块处理。 - 初始现象:直接编译生成的代码运行耗时 45ms,且内存溢出(OOM)。
- 分析判断:通过TVM的
print_ir发现,默认调度将整个输入加载到SRAM,导致溢出。需要手动分块。 - 实际操作:使用
te.schedule将空间维度112x112拆分为16x16的块,每次仅加载一个块到SRAM,计算完成后写回Flash。同时,将通道维度的循环展开4次,利用Cortex-M4的SIMD指令。 - 结果对比:优化后耗时降至 12ms,内存占用降至 40KB,满足约束。
- 风险与代价:分块引入了边界处理逻辑,代码量增加约15%;且Flash写入次数增加,但MCU的Flash擦写寿命通常为10万次,本场景可接受。
- 复盘经验:Micro端编译优化的核心是“以空间换时间,以时间换空间”的权衡。必须先用
tvm.target明确指定MCU的寄存器数量和缓存大小,TVM才能做出合理的调度决策。
vectorize 宽度为8(针对AVX),而Cortex-M4仅支持128位SIMD(即4个float32)。必须通过 tvm.target.arm_cpu 指定 -mcpu=cortex-m4 和 -mattr=+neon 来生成正确的向量化代码。
2.4 总结与复盘
本章的核心逻辑是:计算图定义了“做什么”,算子和张量定义了“操作什么”,调度定义了“怎么做”,而编译流程则负责将这一切转化为可执行的机器码。在Micro端,这四个环节的约束条件(内存、Flash、指令集)比服务器端严格得多,因此必须深入理解每个概念背后的硬件含义。
- 计算图是优化入口,算子融合是Micro端最有效的图优化手段。
- 张量的布局(NCHW vs NHWC)直接影响内存访问效率,需根据MCU的DMA和缓存特性选择。
- 调度是性能调优的核心,但必须结合硬件参数(寄存器数、SIMD宽度)进行。
- 编译流程中,Relay负责“聪明”的图变换,TIR负责“精细”的循环变换,代码生成负责“适配”目标指令集。
下一阶段,我们将深入Relay IR的算子融合原理,并展示如何在Micro端通过 relay.ir_pass 自定义融合规则,进一步压缩内存占用和计算开销。