内存布局转换

在硬件上处理高维张量，需要将逻辑维度映射到线性的物理内存地址空间。尽管框架出于用户方便或历史原因常采用特定布局，但选定的内存布局对卷积和矩阵乘法核的性能有显著作用。编译器通过改写图，使其使用与硬件内存层次结构和向量指令集相符的内存布局，从而优化执行。

逻辑布局与物理布局的差异

一个表示一批图像的4D张量，其维度通常表示为 $N$ (批次), $C$ (通道), $H$ (高), 和 $W$ (宽)。从数学角度看，访问 $(n, c, h, w)$ 处的元素是抽象的。实际上，这个元素存在于RAM一维数组中的一个特定偏移量处。

步长配置决定了这个偏移量。深度学习中存在两种主要的布局：

• NCHW (通道优先): 数据以平面方式存储。第一个通道的所有像素连续存放，随后是第二个通道的所有像素。这是PyTorch的默认设置。
• NHWC (通道最后): 数据以交错方式存储。特定像素所有通道的值一起存放。这是许多后端推理引擎和TensorFlow的默认设置。

它们的地址计算方式不同。对于形状为 $(N, C, H, W)$ 的张量，元素的线性地址计算如下：

NCHW 地址计算: $\text{偏移量} = n \cdot (C \cdot H \cdot W) + c \cdot (H \cdot W) + h \cdot W + w$

NHWC 地址计算: $\text{偏移量} = n \cdot (H \cdot W \cdot C) + h \cdot (W \cdot C) + w \cdot C + c$

布局的选择决定了计算所需的相邻元素之间的内存距离。如果卷积核需要对特定像素位置的所有通道进行求和，NHWC会将其值紧邻地放在内存中。在NCHW中，这些值会根据图像的空间大小 ($H \times W$) 进行跨步，可能引起缓存抖动。

内存步长的可视化

为了弄清编译器为何转换布局，需要考虑数据在线性内存中的存放方式。在下面的图表中，我们展示一个简化的 $1 \times 3 \times 2 \times 2$ 张量（1幅图像，3个通道，2x2空间分辨率）的内存排列。

NCHW 与 NHWC 格式的线性内存放置对比。

在NCHW示例中，访问像素 $(0,0)$ 的红色、绿色和蓝色需要跳过所有其他空间位置。在NHWC中，像素 $(0,0)$ 的R、G、B是相邻的。

硬件匹配与向量化

布局转换的首要目的是提高硬件效率。现代CPU和GPU主要依靠SIMD（单指令多数据）指令和张量核来执行算术运算。这些单元在加载连续数据块时效率最高。

1. 空间局部性: 当卷积在图像上滑动时，它执行权重与输入通道之间的点积。如果布局是NHWC，内循环会遍历 $C$。由于这些值是连续的，CPU可以在一个周期内加载8或16个浮点数的向量。如果布局是NCHW，向量加载会从不连续的内存位置收集数据，这会慢很多。
2. 张量核: 专用矩阵加速单元（如NVIDIA张量核）通常要求数据采用特定的分块格式（例如，分成小块的矩阵）以执行矩阵乘法。编译器必须确保输入张量在馈入这些加速器之前具有正确的形状。

布局转换处理阶段

当ML编译器分析计算图时，它会查找在执行时间中占据主导地位的卷积和池化算子。如果目标硬件（例如ARM CPU或NVIDIA GPU）倾向于特定布局，编译器就会启动一个转换处理阶段。

这个过程不只涉及转置张量。它需要重写算子本身。如果原始图中包含为NCHW配置的Conv2D算子，编译器会用针对NHWC优化的Conv2D变体替换它。

然而，仅仅替换算子是不够的，因为来自用户或前一层的输入数据可能仍是原始格式。编译器必须插入布局转换节点，通常称为 Transpose 或 Permute，以使数据匹配。

设想一个子图，其中卷积后接着ReLU激活函数：

1. 原始: 输入 (NCHW) $\rightarrow$ Conv2D (NCHW) $\rightarrow$ ReLU $\rightarrow$ 输出
2. 转换后: 输入 (NCHW) $\rightarrow$ 转置(至 NHWC) $\rightarrow$ Conv2D (NHWC) $\rightarrow$ ReLU $\rightarrow$ 转置(至 NCHW) $\rightarrow$ 输出

编译器还会执行布局传播。如果多个卷积按顺序发生，在每层之间来回转置数据会降低效率。编译器通过逐元素运算（如ReLU，它与布局无关）传播首选布局，以减少转置操作的数量。

分块布局与打包

除了标准的NCHW和NHWC格式，编译器经常为特定后端目标使用分块布局（也称为平铺或打包布局）。这些是混合格式，旨在使数据精确地适配特定CPU架构（如AVX-512）的向量寄存器或矩阵单元。

常见的块状布局是 NCHWc，通常表示为 $N C/k H W k c$。在这里，通道维度 $C$ 被分成一个外部维度和一个大小为 $k$ 的内部块。例如，如果 $k=16$，该布局将16个通道分组在一起。这确保了内部维度总是向量长度的倍数，从而消除了在核的最内层循环中进行复杂边界检查的必要。

使用分块布局需要编译器完全重写消费核的索引逻辑。尽管这增加了代码生成的复杂程度，但在CPU上的性能提升是可观的，通常比标准NCHW执行快2到3倍。

评估成本模型

布局转换并非总是有利。编译器插入的Transpose操作会产生内存带宽开销。编译器使用成本模型来权衡其益处。

如果一个图包含单个轻量级卷积，后跟内存密集型操作，则数据重排的开销可能超过优化卷积核带来的速度提升。现代ML编译器使用启发式方法或自动调优结果来判断布局转换是否对特定模型和硬件组合来说是全局优化的。