3. 内存子系统优化：让数据在芯片里跑得更顺畅
内存优化这事儿，说白了就是解决一个核心矛盾——计算单元饿不死，但数据喂不饱。我刚开始做模型移植时，总觉得算子优化到位就万事大吉，结果一跑起来，性能惨不忍睹。后来才明白，很多时候瓶颈不在算力，而在内存。
这一章，我们重点聊三个方向：内存层级结构怎么用、数据布局怎么摆、内存怎么复用。搞懂了这些，你的模型在芯片上才能真正跑起来。
3.1 内存层级结构分析
先看一张图，帮你快速建立整体认知。
这张图你应该不陌生。但我想强调一个容易被忽略的点：每一层之间的带宽差异，往往在10倍以上。我见过不少项目，算子本身优化得挺好，但数据在DDR和SRAM之间来回倒腾，性能直接腰斩。
3.2 数据布局优化：NCHW vs NHWC
数据布局，说白了就是张量在内存里怎么排排坐。两种主流格式：NCHW和NHWC。别看只是维度顺序不同，对性能的影响天差地别。
| 对比维度 | NCHW | NHWC |
|---|---|---|
| 内存访问模式 | 通道连续，空间不连续 | 空间连续，通道不连续 |
| 卷积计算 | 适合GPU cuDNN，向量化友好 | 适合CPU，缓存命中率高 |
| BN/激活层 | 通道维度连续，计算高效 | 需要转置或特殊处理 |
| 硬件偏好 | NVIDIA GPU、寒武纪 | 高通、联发科、ARM CPU |
我个人习惯，在项目开始前先搞清楚目标芯片的偏好。比如我之前移植一个分类模型到高通平台，默认用了NCHW，结果跑起来比预期慢了一倍。改成NHWC后，性能直接翻回来。为什么？因为高通DSP的向量单元，对空间连续的数据访问效率更高。
核心原则：数据布局没有的好坏，关键看你的计算单元怎么访存。如果卷积核是1x1，NHWC反而更优，因为通道维度可以并行计算。
来看一个实际转换的例子：
// PyTorch默认NCHW: [1, 3, 224, 224]
// 转成NHWC: [1, 224, 224, 3]
import torch
tensor_nchw = torch.randn(1, 3, 224, 224)
tensor_nhwc = tensor_nchw.permute(0, 2, 3, 1).contiguous()
// 检查内存布局
print(tensor_nhwc.stride()) // (150528, 1, 672, 224)
注意那个contiguous()调用。嗯，这里有个坑——permute只是改变了视图，内存还是原来的顺序。必须显式调用contiguous，才能真正重排数据。我曾经因为这个疏忽，调试了整整一个下午。
3.3 内存池与复用机制
模型推理时，每一层都要分配和释放内存。频繁的malloc/free，不仅慢，还会产生碎片。你想想看，嵌入式设备内存本来就紧张，碎片一多，后面的大张量可能就分配不出来了。
解决方案就是内存池。提前申请一大块连续内存，然后在池子里做分配和回收。
我的经验：内存池的大小，建议设为模型最大中间张量总和的1.2倍。太小了容易溢出，太大了浪费。我一般先跑一遍profile，统计所有层的输出大小，再定池大小。
内存复用的核心思路：生命周期不重叠的张量，可以共用同一块内存。比如卷积层的输出，在进入激活函数后就不再需要了，那激活函数的输出就可以复用这块内存。
来看一个简单的内存池实现思路：
// 伪代码：内存池管理器
class MemoryPool {
void* base_ptr; // 池基地址
size_t pool_size; // 池大小
size_t offset; // 当前分配偏移
void* alloc(size_t size) {
// 对齐到64字节（缓存行对齐）
size = align_up(size, 64);
if (offset + size > pool_size) return nullptr;
void* ptr = base_ptr + offset;
offset += size;
return ptr;
}
void reset() {
offset = 0; // 一轮推理结束，重置偏移
}
};
实际项目中，我更喜欢用双缓冲机制。什么意思？就是准备两块内存，一块给当前层读，一块给下一层写。这样读写可以流水线并行，不用等上一块内存释放。
注意：双缓冲虽然快，但内存占用翻倍。如果芯片内存吃紧，建议用单缓冲+原地计算。比如某些激活函数（ReLU、LeakyReLU）可以直接修改输入数据，不需要额外内存。
最后，分享一个我踩过的坑。有一次做端侧部署，模型里有个超大特征图（1x128x128x256），我按常规思路做了内存池，结果跑起来总是崩。排查了半天，发现是内存对齐的问题——芯片的DMA要求地址64字节对齐，我分配时没注意。加上对齐后，问题解决。
所以，内存对齐不是小事。不同芯片要求不同，有的32字节，有的128字节。建议在内存池的alloc函数里，统一做对齐处理。
好了，这一章的内容就这些。内存优化是个细活，需要你对芯片的架构、数据的流向、内存的分配都有清晰的认识。多跑profile，多看内存分布图，慢慢你就能找到感觉。