4、内存复用基础：片上SRAM、寄存器、局部内存与全局内存的层次结构，内存复用对性能与功耗的影响

在TVM Micro算子（Micro Operator）的调度优化中，内存层次结构是决定性能与功耗的基石。与桌面GPU或服务器端AI加速器不同，微控制器（MCU）和嵌入式AI芯片的片上存储资源极为有限，通常仅有几十KB到几MB的SRAM，而外部DRAM（全局内存）的访问延迟可能是片上SRAM的10到100倍，能耗差异更是高达两个数量级。理解寄存器、局部内存（Local Memory）、片上SRAM（共享内存/缓存）与全局内存（Global Memory）之间的层级关系，是设计高效调度策略的前提。本章将深入剖析这一层次结构，并系统阐述内存复用（Memory Reuse）如何通过减少对全局内存的访问来同时提升性能与降低功耗。

上图展示了TVM Micro算子执行时的内存层次结构。从下往上看，全局内存（DRAM）容量最大但访问代价最高；片上SRAM作为中间层，是数据复用和调度的主战场；寄存器文件则位于最顶层，提供零延迟访问。右侧列出了内存复用的核心策略，这些策略的目标都是让数据尽可能在高层级（寄存器、SRAM）被多次使用，从而减少对低层级（DRAM）的访问。

4.1 核心原理：内存层次结构与访问代价模型

在嵌入式AI芯片（如ARM Cortex-M系列、RISC-V核、NPU微控制器）中，内存层次的设计直接决定了算子的执行效率。TVM Micro的代码生成器需要针对目标硬件的具体内存大小和带宽参数，生成最优的调度方案。理解这一层次结构，首先需要建立访问代价（Cost Model）的量化概念。

从上表可以看出，访问一次全局内存的能耗大约是访问SRAM的20倍，延迟则是10倍以上。因此，内存复用的核心目标就是：最大化数据在寄存器与SRAM中的重用次数，最小化对DRAM的访问频率。TVM的调度原语（如compute_at、storage_align、double_buffer）正是围绕这一目标设计的。

在TVM Micro中，local_memory通常映射到片上SRAM，而global_memory则对应外部DRAM。调度器通过cache_read和cache_write指令将数据从全局内存搬运到局部内存，再通过循环分块（Tiling）让计算密集部分在SRAM内完成。这种层次化设计使得每个数据元素从DRAM加载一次后，可以在SRAM中被多次读取，从而摊薄了访问成本。

4.2 判断方法与操作流程：如何分析算子的内存复用潜力

在实际开发中，面对一个具体的Micro算子（如卷积、矩阵乘、深度可分离卷积），我们需要系统性地分析其内存复用潜力。以下是标准化的分析流程：

1. 数据流分析：绘制算子的数据流图，标注每个张量的生产者（Producer）和消费者（Consumer）。识别哪些数据被多次读取（如卷积核权重、输入特征图的滑动窗口重叠区域）。
2. 循环结构分析：检查嵌套循环的迭代空间。如果内层循环的某个数据在外层循环中保持不变，则存在寄存器复用的机会；如果数据在相邻迭代中部分重叠，则存在SRAM复用的机会。
3. 存储容量估算：计算目标数据块的大小。例如，一个3x3卷积核在FP16精度下仅占18字节，完全可以放入寄存器；而一个64x64的特征图块（8KB）则需要放入SRAM。
4. 访问频率计算：统计每个数据元素从全局内存加载的次数。理想情况下，每个元素只加载一次，然后在SRAM/寄存器中被反复使用。
5. 调度策略选择：根据分析结果，选择合适的TVM调度原语。若数据块适合寄存器，使用vectorize和unroll；若适合SRAM，使用cache_read和compute_at。

以下是一个典型的TVM Micro调度代码片段，展示了如何通过cache_read实现SRAM复用：

# 假设A为输入矩阵，B为权重矩阵，C为输出
C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k))
# 将A和B缓存到局部内存（SRAM）
A_local = s.cache_read(A, "local", [C])
B_local = s.cache_read(B, "local", [C])
# 将计算节点绑定到SRAM层级
s[C].compute_at(s[C], s[C].op.axis[0])
# 对输出循环进行分块，使数据块适配SRAM大小
xo, xi = s[C].split(s[C].op.axis[0], factor=16)
yo, yi = s[C].split(s[C].op.axis[1], factor=16)
s[C].reorder(xo, yo, xi, yi)
# 对A_local和B_local应用compute_at，使其在分块内复用
s[A_local].compute_at(s[C], yo)
s[B_local].compute_at(s[C], yo)

这段代码的核心逻辑是：将大矩阵乘法分解为16x16的小块，每个小块的计算过程中，A和B的对应子块被加载到SRAM中，并在该块内被多次复用。通过这种方式，对DRAM的访问次数从M×N×K次降低到(M/16)×(N/16)×K次，减少了约256倍。

4.3 案例拆解：深度可分离卷积的内存复用优化

为了更具体地说明内存复用的实际效果，我们以一个在MCU上常见的深度可分离卷积（Depthwise Separable Convolution）为例。该算子由深度卷积（Depthwise Conv）和逐点卷积（Pointwise Conv）组成，是MobileNet等轻量级网络的核心算子。

背景与约束：目标硬件为ARM Cortex-M7内核，主频400MHz，片上SRAM为512KB，全局DRAM为8MB。输入特征图为112x112x32（FP16），深度卷积核为3x3x32，逐点卷积核为1x1x32x64。初始实现直接使用TVM默认调度，未做任何内存优化。

初始性能数据：

• 总执行时间：45.2ms
• DRAM访问次数：约1.2亿次
• SRAM利用率：15%
• 功耗：85mW（其中DRAM访问占60%）

分析判断过程：通过TVM的debug_keep_code和性能分析器，发现主要瓶颈在于深度卷积阶段。由于深度卷积的每个输出像素仅依赖3x3x1的输入区域，但默认调度将整个输入特征图从DRAM加载到SRAM，导致大量数据未被复用就被丢弃。同时，逐点卷积的权重（32x64x2字节=4KB）被反复从DRAM加载。

优化操作步骤：

1. 深度卷积分块：将输入特征图按通道维度分块，每块处理8个通道。这样，每个3x3卷积核在SRAM中加载一次，可被该通道块内的所有空间位置复用。
2. 权重驻留：将逐点卷积的权重（4KB）通过cache_read固定到SRAM中，并在整个逐点卷积阶段保持驻留。
3. 双缓冲：对深度卷积的输入使用double_buffer，使DMA搬运下一块数据的同时，当前块正在计算，隐藏DRAM访问延迟。
4. 寄存器缓存：在深度卷积的内层循环中，将3x3卷积核的9个权重值存入寄存器，避免每次乘加都从SRAM读取。

优化后结果对比：

风险与代价：优化后代码体积增加了约2KB（用于双缓冲和分块控制逻辑），且SRAM占用从76KB增加到348KB，但仍未超出512KB的限制。需要注意的是，双缓冲会占用额外的SRAM空间，在内存紧张时需要权衡缓冲块大小。

复盘经验：本案例的核心经验是——内存复用不是简单的缓存，而是需要结合算子的数据流特征，主动设计数据在SRAM中的生命周期。深度卷积的通道分块利用了通道间数据独立性，权重驻留利用了权重的全局复用性，双缓冲则隐藏了DRAM延迟。三者结合，才实现了3.5倍的性能提升和50%的功耗降低。

4.4 常见误区与总结复盘

在内存复用的实际应用中，开发者常陷入以下误区：

• 误区一：盲目使用cache_read。并非所有数据都适合缓存到SRAM。如果数据只被读取一次（如某些激活函数的输出），缓存反而增加了SRAM压力。应通过数据流分析确认数据的复用次数。
• 误区二：忽略寄存器复用。许多开发者只关注SRAM级别的复用，而忽略了寄存器。实际上，对于3x3卷积核、1x1卷积等小数据块，寄存器复用能减少50%以上的SRAM访问。
• 误区三：分块过大导致SRAM溢出。分块大小必须根据目标硬件的SRAM容量动态调整。TVM的te.autotvm或meta_schedule可以自动搜索最优分块参数，但在Micro场景下，手动指定往往更可靠。
• 误区四：忽视DMA与计算的并行。内存复用的最终目标是让数据搬运与计算重叠。如果只做分块而不使用double_buffer，CPU仍会等待DMA完成，无法充分利用带宽。

总结复盘：内存复用是TVM Micro算子调度的核心主题。本章从内存层次结构出发，建立了访问代价模型，给出了分析复用潜力的标准流程，并通过深度可分离卷积案例展示了从45.2ms到12.8ms的优化过程。核心要点可归纳为：

1. 数据分层：明确哪些数据应放在寄存器、哪些放在SRAM、哪些留在DRAM。
2. 生命周期管理：通过compute_at和cache_read控制数据在SRAM中的存活时间。
3. 访问重叠：使用双缓冲或流水线技术，让DMA搬运与计算并行执行。
4. 容量约束：始终以目标硬件的SRAM容量为硬约束，分块大小需留有余量。