2、TVM基础回顾：计算图、算子、调度、编译流程，以及Micro扩展的动机

在深入TVM Micro的代码生成与算子注册之前，我们必须先建立统一的语言和认知框架。TVM并非一个单一的编译器，而是一个端到端的机器学习编译栈，其核心抽象包括计算图（Computational Graph）、算子（Operator）、调度（Schedule）与编译流程（Compilation Pipeline）。理解这些基础概念，是理解Micro扩展如何在资源受限的嵌入式设备上落地的前提。本章将系统梳理这些核心概念，并剖析为何在微控制器（MCU）场景下，传统的TVM编译与运行时策略需要被重新审视与扩展。

上图展示了TVM高层框架与Micro扩展之间的映射关系。左侧是TVM的标准抽象，右侧是Micro场景下的具体化。可以看到，Micro扩展并非推翻TVM，而是将TVM的通用机制（如算子注册、调度原语）约束在嵌入式硬件能力范围内，并引入新的代码生成策略（如AOT编译）与轻量级运行时（MicroTVM Runtime）。

2.1 核心原理：从计算图到可执行代码的抽象层次

TVM将深度学习模型的编译过程分解为多个层次的中间表示（IR），每一层都服务于特定的优化目标。理解这些层次是掌握代码生成的关键。

• 计算图（Relay IR）：高层图表示，描述数据流与算子之间的依赖关系。它负责图级别的优化，如算子融合（Operator Fusion）、常量折叠（Constant Folding）和布局转换（Layout Transform）。
• 张量表达式（TE / TIR）：低层表示，将每个算子展开为具体的循环嵌套和内存访问模式。TIR（Tensor IR）是TVM进行调度优化的核心载体，它允许开发者显式控制循环顺序、并行度、向量化等底层细节。
• 目标代码（Target Code）：最终生成的C/C++源码或机器码。对于Micro场景，通常生成C源码，再由交叉编译器编译为MCU可执行的二进制。

这种分层设计使得TVM能够在前端（如PyTorch、TensorFlow）与后端（如x86、ARM、RISC-V）之间建立解耦，而Micro扩展正是针对“后端”这一层进行了深度定制。

2.2 算子注册机制：TVM的“插件化”核心

算子（Operator）是计算图的基本执行单元。TVM通过一个全局的算子注册表（Op Registry）来管理所有可用的算子。每个算子都关联了其计算定义（Compute）、调度规则（Schedule）以及代码生成策略（Codegen）。

在Micro场景下，算子注册的约束更为严格。例如，一个标准的conv2d算子可能依赖浮点运算，但许多MCU（如Cortex-M0+）没有FPU（浮点单元）。因此，Micro扩展需要注册一个“定点数版本”的conv2d算子，或者通过查表法（Lookup Table）近似激活函数。

2.3 调度（Schedule）在Micro场景下的特殊约束

调度是TVM中性能优化的核心手段。在服务器GPU上，调度主要关注线程块划分、共享内存使用和向量化。但在MCU上，调度逻辑完全不同：

1. 内存层次：MCU通常只有几十KB的SRAM，没有DRAM，更没有缓存。因此，调度必须最小化中间张量的内存占用，尽量采用“计算-存储”原地更新的策略。
2. 循环展开：为了减少循环控制开销，可以适度进行循环展开（Loop Unrolling），但展开过大会导致代码体积膨胀，超出Flash容量。
3. 算子融合：在Micro场景下，算子融合不仅是减少内核启动开销，更是为了减少中间结果的内存写回。例如，将卷积+ReLU融合，可以避免将卷积结果写回SRAM再读出的过程。

# 标准TVM调度（面向GPU）
s = te.create_schedule(conv2d.op)
s[conv2d].split(axis, factor=16)  # 分块
s[conv2d].bind(block_idx, thread_idx)  # 绑定线程

# Micro调度（面向Cortex-M4）
s = te.create_schedule(conv2d.op)
s[conv2d].reorder(y, x, k)  # 重排循环，确保输出按行连续写入
s[conv2d].unroll(k, factor=4)  # 小因子展开，减少循环开销
# 注意：不进行分块，因为MCU没有缓存，分块反而增加索引计算开销

上述对比清晰地展示了调度策略的差异。在Micro中，调度的首要目标是“减少内存足迹”和“保证确定性执行”，而非最大化并行度。

2.4 编译流程：从模型到二进制镜像的完整链路

TVM的编译流程可以概括为“图优化 → 算子Lowering → 代码生成 → 运行时打包”。在Micro扩展中，这一流程被进一步强化了静态分析和AOT（Ahead-Of-Time）编译。

• 图优化阶段：使用Relay Pass Manager执行融合、剪枝等操作。在Micro中，需要额外检查算子是否支持“零拷贝”融合。
• Lowering阶段：将Relay IR转换为TIR。此阶段会调用算子注册表中的Compute函数，并应用默认调度。Micro扩展会在此阶段插入“内存规划”Pass，为每个中间张量分配静态内存地址。
• 代码生成阶段：调用Target的Codegen模块，生成C源码。Micro的Codegen需要避免生成依赖malloc的代码，并确保所有数组大小在编译期已知。
• 运行时打包：生成一个自包含的C文件，包含模型权重、图结构描述（通常为线性化的指令序列）以及MicroTVM Runtime的轻量实现。

2.5 案例拆解：将MobileNetV2的深度可分离卷积部署到Cortex-M7

为了将上述概念串联起来，我们以一个实际案例来展示Micro扩展的动机与实现路径。假设我们需要在STM32F746（Cortex-M7，主频216MHz，SRAM 320KB，Flash 1MB）上部署一个图像分类模型。

初始现象与目标：直接使用标准TVM编译MobileNetV2，生成的代码在PC上模拟运行正常，但交叉编译到MCU后，程序启动即崩溃。分析发现，标准TVM生成的代码中包含了动态内存分配（malloc）和浮点运算，而Cortex-M7虽然支持单精度浮点，但malloc在裸机环境下不可用。

关键输入参数：模型输入为224x224x3，第一层卷积输出通道为32。标准TVM为中间张量分配了约224*224*32*4 = 6.4MB的临时缓冲区，远超SRAM容量。

分析判断过程：

1. 内存瓶颈：标准TVM的调度假设有足够的内存，因此会为每个算子分配独立的输出缓冲区。在MCU上，必须复用缓冲区。
2. 算子支持：MobileNetV2中的深度可分离卷积（Depthwise Conv2D）在TVM的算子注册表中存在，但其默认调度针对GPU，未考虑MCU的寄存器压力。
3. 浮点处理：虽然Cortex-M7有FPU，但为了降低功耗和提升速度，我们决定将权重从FP32量化到INT8，并使用定点运算。

实际操作步骤：

# 1. 使用TVM的量化工具将模型转换为INT8
# 2. 注册自定义的INT8 DepthwiseConv2D算子
# 3. 编写自定义调度：将输出缓冲区与输入缓冲区重叠（in-place计算）
# 4. 使用AOT编译生成C代码，并禁用所有动态内存分配
# 5. 在STM32CubeIDE中编译并链接MicroTVM Runtime

结果对比：优化后，SRAM占用从6.4MB降至约180KB（满足320KB限制），Flash占用从1.2MB降至约800KB（满足1MB限制）。推理时间从无法运行变为每帧约350ms，虽然不算快，但已满足低功耗场景下的非实时要求。

风险与代价：定点化导致精度损失约0.5%的Top-1准确率。此外，手动编写调度需要深入理解MCU架构，开发周期增加了约3天。

复盘经验：Micro部署的核心不是“优化计算”，而是“管理内存”。算子注册和调度必须围绕“静态内存规划”这一核心约束展开。

2.6 常见误区与总结

总结本章，TVM Micro扩展的动机源于三个不可调和的矛盾：动态图与静态内存的矛盾、浮点计算与硬件资源的矛盾、通用运行时与裸机环境的矛盾。解决这些矛盾的关键在于：

• 深入理解TVM的算子注册机制，并针对MCU扩展新的算子实现。
• 将调度目标从“吞吐量”转变为“内存峰值最小化”。
• 采用AOT编译，彻底摒弃解释执行和动态分配。