布局转换

内存物理上是一维的。深度学习框架将张量表示为具有 $(N, C, H, W)$ 等形状的多维对象，但硬件看到的是字节的平坦序列。逻辑张量索引与物理内存地址之间的映射决定了内存布局。编译器选择内存布局是图级优化中非常重要的决定之一，因为它直接影响数据局部性、缓存使用率以及专用硬件指令的使用能力。

张量布局与硬件固有要求不匹配会导致分散的内存访问模式。这会导致高缓存未命中率，并阻止使用像Tensor Cores或AVX-512这样的高吞吐量指令。布局转换是编译器的一个阶段，它负责重写图以确保数据按计算引擎期望的顺序存储在内存中。

步长和局部性的影响

操作效率与内存访问步长相关。考虑一个标准的2D卷积。该操作涉及输入通道和滤波器权重之间的点积。

如果我们以 $NCHW$ 格式（批次、通道、高度、宽度）存储数据，则最内层维度是宽度。图像中空间上相邻（宽度方向）的值在内存中也是相邻的。然而，在不同通道但相同空间位置的值之间，相隔 $H \times W$ 个元素。

与此相对，考虑 $NHWC$ 格式（批次、高度、宽度、通道）。在这种格式下，最内层维度是通道。对应于所有通道中相同像素的值被连续存储。由于现代卷积实现通常在通道维度上进行归约（沿深度方向累积输入），$NHWC$ 允许硬件在单次事务中加载密集向量的通道数据。这显著提高了GPU上的内存合并效率，并使CPU上的向量化成为可能。

我们可以将 $NCHW$ 布局中4D张量索引 $(n, c, h, w)$ 的内存地址 $A$ 定义为：

$A_{NCHW}(n, c, h, w) = n \cdot (C H W) + c \cdot (H W) + h \cdot (W) + w$

对于 $NHWC$ 布局，映射变为：

$A_{NHWC}(n, c, h, w) = n \cdot (H W C) + h \cdot (W C) + w \cdot (C) + c$

当编译器将图转换为特定目标时，它会向后端查询首选布局。对于使用Tensor Cores的NVIDIA GPU（通过cuDNN或Cutlass），首选布局几乎总是 $NHWC$ （通常称为通道在后）。对于依赖SIMD指令的CPU，最佳布局通常是分块格式，它创建的数据块大小刚好适合向量寄存器。

分块布局和向量化

对于CPU目标，像 $NCHW$ 或 $NHWC$ 这样的标准布局通常不足以最大化算术强度。为了充分使用SIMD单元（例如AVX-512或ARM Neon），编译器经常使用布局分块（也称为平铺或打包）。

分块是指将一个维度拆分为一个外维度和一个固定大小 $k$ 的内维度。例如，我们可以将通道维度 $C$ 转换为 $C_{out} \times k$。这将4D张量 $NCHW$ 转换为5D张量 $NCHWc$，其中小写 $c$ 表示一个连续存储的通道块（例如8或16）。

如果我们选择块大小为16，布局变为 $NCHW16c$。内存地址计算公式变为：

$A_{NCHW16c}(n, c_{out}, h, w, c_{in}) = n \cdot (C_{out} H W \cdot 16) + c_{out} \cdot (H W \cdot 16) + h \cdot (W \cdot 16) + w \cdot (16) + c_{in}$

这种结构确保当CPU处理一个空间像素时，它在一条指令中将精确的16个通道加载到512位向量寄存器中。这消除了对收集/分散指令的需求，并确保数据与FMA（乘加运算）单元完美对齐。

下图说明了布局选择对通用卷积操作内存吞吐量的性能影响。

不同张量布局的有效内存带宽利用率比较。分块和通道在后布局通过最小化缓存行逐出，在现代密集架构上显著优于平面 $NCHW$ 格式。

布局传播

改变单个操作符的布局很简单；但处理对整个图的影响则很复杂。如果卷积从 $NCHW$ 转换为 $NHWC$，其输入必须进行转置。如果后续操作（例如，批归一化或ReLU）期望 $NCHW$，则输出必须转置回来。

在每个节点前后插入显式的 Transpose 操作会违背优化的目的。内存带宽成本高昂，转置大型张量是一种内存密集型操作，会消耗大量时间和能量。为了解决这个问题，编译器实现了一个名为布局传播的阶段。

在这个阶段，编译器根据硬件目标，为特定的“锚定”操作符（通常是卷积或矩阵乘法）分配首选布局。然后它遍历图以将此布局要求传播给相邻操作符。诸如ReLU、Add或Sigmoid等逐元素操作与布局无关；它们可以像处理 $NCHW$ 数据一样轻松地处理 $NHWC$ 数据，而无需改变其数学定义。

编译器通过这些与布局无关的节点推动布局转换，直到遇到必须改变布局的边界（例如，一个Reshape操作或外部输出）。这使得编译器可以将整个子图转换为目标布局，从而将图边界处必要的转置次数降至最低。

下图描绘了布局传播阶段中子图的转换过程。

布局传播的可视化。编译器识别出Conv2D需要NHWC布局。编译器没有只在Conv2D周围进行转置，而是将NHWC布局通过与布局无关的ReLU传播出去，将耗时的逆转换移至序列的末尾。

实现限制

实现布局转换需要细致处理操作符属性。当布局改变时，编译器不仅要更新输入张量，还要更新操作符的静态属性。

1. 权重转换： 卷积层的权重是常量张量。当从 $NCHW$ 切换到 $NHWC$ 时，权重必须在编译时（离线）永久重塑。如果编译器未能离线完成此操作，运行时将在每次推理过程中承担转置权重的巨大开销。
2. 属性映射： 操作通常具有沿特定轴定义的属性，例如 stride（步长）、padding（填充）或 dilation（膨胀）。如果布局旋转，这些属性必须重新索引。例如，[2, 2] 的步长通常适用于 $(H, W)$。如果布局改变，使得空间维度从索引2和3移动到索引1和2，则步长属性向量必须进行置换以匹配。

目标是确保操作的语义含义保持一致，同时物理执行适应硬件的优势。通过自动化此过程，编译器将模型定义与硬件实现解耦，使相同的高级Keras或PyTorch代码能够在移动CPU（使用 $NCHW4c$）和数据中心GPU（使用 $NHWC$）上高效运行。