内存物理上是一维的。深度学习框架将张量表示为具有 $(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(乘加运算)单元完美对齐。下图说明了布局选择对通用卷积操作内存吞吐量的性能影响。{ "layout": { "title": "不同布局策略下的内存吞吐效率", "xaxis": { "title": "布局格式", "showgrid": false }, "yaxis": { "title": "有效带宽 (GB/s)", "showgrid": true, "gridcolor": "#dee2e6" }, "plot_bgcolor": "white", "width": 600, "height": 400, "font": { "family": "Arial, sans-serif", "color": "#495057" }, "margin": { "l": 60, "r": 30, "t": 50, "b": 50 } }, "data": [ { "x": [ "NCHW (标准)", "NHWC (通道在后)", "NCHW16c (分块)" ], "y": [ 450, 780, 820 ], "type": "bar", "marker": { "color": [ "#a5d8ff", "#748ffc", "#4dabf7" ] }, "text": [ "步长访问开销", "合并访问", "向量对齐" ], "textposition": "auto" } ] }不同张量布局的有效内存带宽利用率比较。分块和通道在后布局通过最小化缓存行逐出,在现代密集架构上显著优于平面 $NCHW$ 格式。布局传播改变单个操作符的布局很简单;但处理对整个图的影响则很复杂。如果卷积从 $NCHW$ 转换为 $NHWC$,其输入必须进行转置。如果后续操作(例如,批归一化或ReLU)期望 $NCHW$,则输出必须转置回来。在每个节点前后插入显式的 Transpose 操作会违背优化的目的。内存带宽成本高昂,转置大型张量是一种内存密集型操作,会消耗大量时间和能量。为了解决这个问题,编译器实现了一个名为布局传播的阶段。在这个阶段,编译器根据硬件目标,为特定的“锚定”操作符(通常是卷积或矩阵乘法)分配首选布局。然后它遍历图以将此布局要求传播给相邻操作符。诸如ReLU、Add或Sigmoid等逐元素操作与布局无关;它们可以像处理 $NCHW$ 数据一样轻松地处理 $NHWC$ 数据,而无需改变其数学定义。编译器通过这些与布局无关的节点推动布局转换,直到遇到必须改变布局的边界(例如,一个Reshape操作或外部输出)。这使得编译器可以将整个子图转换为目标布局,从而将图边界处必要的转置次数降至最低。下图描绘了布局传播阶段中子图的转换过程。digraph G { rankdir=LR; bgcolor="transparent"; node [fontname="Arial", shape=box, style=filled, color="#dee2e6", fillcolor="white"]; edge [fontname="Arial", color="#868e96"]; subgraph cluster_0 { label = "原始图 (NCHW)"; style = dashed; color = "#adb5bd"; fontcolor = "#868e96"; node_input [label="输入\n(NCHW)", fillcolor="#e7f5ff", color="#74c0fc"]; node_conv [label="Conv2D\n(NCHW)", fillcolor="#ffe3e3", color="#ff8787"]; node_relu [label="ReLU", fillcolor="#e9ecef"]; node_input -> node_conv; node_conv -> node_relu; } subgraph cluster_1 { label = "优化图 (NHWC)"; style = solid; color = "#495057"; fontcolor = "#495057"; opt_input [label="输入\n(NCHW)", fillcolor="#e7f5ff", color="#74c0fc"]; opt_trans_in [label="布局转换\nNCHW -> NHWC", shape=hexagon, fillcolor="#fff9db", color="#fab005"]; opt_conv [label="Conv2D\n(NHWC)", fillcolor="#d0bfff", color="#9775fa"]; opt_relu [label="ReLU\n(原地操作)", fillcolor="#e9ecef"]; opt_trans_out [label="布局转换\nNHWC -> NCHW", shape=hexagon, fillcolor="#fff9db", color="#fab005"]; opt_input -> opt_trans_in; opt_trans_in -> opt_conv; opt_conv -> opt_relu [label="布局已传播"]; opt_relu -> opt_trans_out; } }布局传播的可视化。编译器识别出Conv2D需要NHWC布局。编译器没有只在Conv2D周围进行转置,而是将NHWC布局通过与布局无关的ReLU传播出去,将耗时的逆转换移至序列的末尾。实现限制实现布局转换需要细致处理操作符属性。当布局改变时,编译器不仅要更新输入张量,还要更新操作符的静态属性。权重转换: 卷积层的权重是常量张量。当从 $NCHW$ 切换到 $NHWC$ 时,权重必须在编译时(离线)永久重塑。如果编译器未能离线完成此操作,运行时将在每次推理过程中承担转置权重的巨大开销。属性映射: 操作通常具有沿特定轴定义的属性,例如 stride(步长)、padding(填充)或 dilation(膨胀)。如果布局旋转,这些属性必须重新索引。例如,[2, 2] 的步长通常适用于 $(H, W)$。如果布局改变,使得空间维度从索引2和3移动到索引1和2,则步长属性向量必须进行置换以匹配。目标是确保操作的语义含义保持一致,同时物理执行适应硬件的优势。通过自动化此过程,编译器将模型定义与硬件实现解耦,使相同的高级Keras或PyTorch代码能够在移动CPU(使用 $NCHW4c$)和数据中心GPU(使用 $NHWC$)上高效运行。