内存布局转换的代价：NCHW 到 NHWC 的重排折损与消除

在视觉模型与多模态大模型（Vision-Language Models）的编译优化中，张量的**内存排布格式（Memory Layout / Data Format）**是影响硬件计算效率的最关键因素之一。

很多刚接触深度学习底层优化的开发者，经常在模型代码中随意插入 permute、transpose 或 reshape 操作，以为这只是简单的几行维度变换。然而在底层硬件层面，一次看似不起眼的从 NCHW（通道在前）到 NHWC（通道在后）的内存重排（Layout Transformation），往往会引发巨额的全局内存搬运与缓存颠簸（Cache Thrashing）。

理解内存排布背后的硬件原理，并在编译器 Pass 中实施全局布局优化（Layout Propagation & Elimination），是消除这类隐形性能杀手的核心手段。

+--------------------------------------------------------------------------+
|                     NCHW vs NHWC 物理内存排布对比                          |
+--------------------------------------------------------------------------+
| NCHW (PyTorch Eager 默认):                                               |
| [R0, R1, R2... (全图的红通道)] -> [G0, G1... (全图绿通道)] -> [B0, B1...] |
| -> 特点: 空间局部性集中在同一通道内部，对向量化点积计算极度不友好               |
+--------------------------------------------------------------------------+
| NHWC (Tensor Core / CPU SIMD 偏好):                                      |
| [R0, G0, B0] -> [R1, G1, B1] -> [R2, G2, B2] ... (像素级多通道紧凑打包)  |
| -> 特点: 同一像素的各通道连续存放，天然契合 128-bit 向量化加载与 Tensor Core |
+--------------------------------------------------------------------------+

1. 为什么硬件极度偏爱 NHWC 布局

在 CPU（AVX2/AVX-512/Neon）和现代 GPU（NVIDIA Tensor Core）上执行卷积或线性变换时：

• 计算的核心算子是沿通道维度（Channel 维）执行点积累加：$Y = \sum X_c \times W_c$；
• 在 NHWC 格式下，同一个空间坐标 $(h, w)$ 上的所有通道数据在物理内存中是完全连续存放的。硬件可以通过一条向量加载指令（如 LDG.E.128 或 _mm256_load_ps）一次性把 8 个或 16 个通道数据直接填满向量寄存器，立刻进行 FMA 乘加；
• 在 NCHW 格式下，相邻通道的数据在物理内存中间隔了整整 $H \times W$ 个浮点数（跨度可能达几万字节）。要获取同一像素的多通道数据，CPU/GPU 必须执行离散的跨步读取（Strided Gather），这会导致 L1/L2 Cache Line 命中率暴跌，内存带宽被大量无用的填充字节吞噬。

因此，cuDNN 和 TensorRT 在底层算子库中，全面要求输入张量采用 NHWC（或类似 NC/32HW32 的打包格式）以释放 Tensor Core 的峰值算力。

2. 局部转置的灾难：跨步搬运的带宽折损

如果在计算图中，前端框架输出了 NCHW 张量，而底层算子需要 NHWC，系统就必须在两个算子之间插入一个 Layout Transform（转置）算子。

这个转置算子是一个纯粹的 Memory-Bound 算子：

• 算术强度（FLOPs/Byte）为 0；
• 必须将整个几十兆的张量从全局显存读入片上，进行坐标解算，再非连续地写回全局显存；
• 在一个包含数十层卷积的网络中，如果存在频繁的 NCHW $\leftrightarrow$ NHWC 来回转换，这些转置算子消耗的时间甚至能占到整网总延迟的 30% 以上！

3. AI 编译器的全局布局传播与消除算法

为了消除这些冗余的转置开销，现代 AI 编译器（如 TVM、MLIR）引入了全局布局重写 Pass（Layout Rewriting & Propagation Pass）。

该 Pass 的执行逻辑分为三步：

1. 算子偏好标记（Sink/Source Tagging）：
   编译器遍历计算图，为每个算子标记其在当前目标硬件上的最优布局（例如标记 Conv2D 为 NHWC，标记 Elementwise 为 Any）；
2. 布局传播与下推（Layout Pushdown）：
   编译器沿着数据流 DAG，将转置算子尽可能向计算图的两端（输入端与输出端）推导。如果转置算子遇到了无偏好的逐元素算子（如 ReLU、Add），编译器直接将这些算子的执行布局重写为 NHWC，让转置算子穿透过去；
3. 互逆转置对消（Layout Cancellation）：
   当两个互逆的转置算子（如 NCHW -> NHWC 紧接着 NHWC -> NCHW）在下推过程中相遇时，编译器直接将这两个算子在计算图中同时剪枝消除。

// 优化前：频繁转置
Conv(NCHW) -> [Trans: NCHW->NHWC] -> ReLU -> [Trans: NHWC->NCHW] -> Conv(NCHW)

// 优化后：全图布局统一，转置算子完全消除
Conv(NHWC) -> ReLU(NHWC) -> Conv(NHWC)

通过编译器在图级别的全局布局对齐，我们不仅让底层硬件跑在了最舒服的 NHWC 连续向量访存模式下，而且彻底消除了中间所有无意义的显存搬运损耗。