内存感知布局转换

通过算子合并和简化重构算子序列，优化张量数据在内存中的物理布局是一种重要的图级别优化技术。张量维度在线性内存中的排序方式直接影响数据局部性、缓存利用率、向量 (vector)化效率，并最终影响计算核心在特定硬件上的性能表现。内存感知布局转换分析计算图和目标硬件，以选择合适的数据布局，仅在必要且有利时才插入显式转置操作。

数据布局的重要性

计算机视觉中常用的4D张量，包括批次 (N)、通道 (C)、高度 (H) 和宽度 (W)。这些张量主要有两种内存布局：

1. NCHW（通道优先）： 数据沿宽度维度连续存储，其次是高度、通道，最后是批次。tensor[n][c][h][w]元素的内存地址可计算为 offset = n*C*H*W + c*H*W + h*W + w。这种布局在PyTorch等框架中常见，并常被cuDNN等GPU库在特定操作中偏爱。
2. NHWC（通道最后）： 数据沿通道维度连续存储，其次是宽度、高度，最后是批次。tensor[n][c][h][w]的内存地址计算为 offset = n*H*W*C + h*W*C + w*C + c。这种布局是TensorFlow中的默认设置，对于CPU执行和某些硬件加速器可能有利。

NCHW和NHWC之间（或更高维度的其他可能布局）的选择并非随意。它对性能有深远影响：

• 缓存局部性： 卷积等操作涉及在空间维度（H，W）上滑动滤波器。如果空间上一起访问的数据在内存中也相邻，缓存命中率就会提高。NHWC通常为典型的CPU卷积循环结构提供更好的空间局部性，因为元素 [h][w][c=0]、[h][w][c=1]、... [h][w][c=C-1]是连续的。NCHW提供更好的通道局部性。
• 内存访问合并（GPU）： GPU通过让warp中的线程同时访问连续内存位置（合并访问）来获得高内存带宽。针对NCHW设计的核函数在多线程并发处理给定空间位置的多个通道时，可能实现更好的合并访问。
• 向量 (vector)化（SIMD）： CPU SIMD指令（如AVX、NEON）对数据向量进行操作。如果要在同时处理的数据元素位于连续内存中，编译器通常可以更有效地向量化 (quantization)操作。对于一次处理多个通道的操作，NHWC的连续通道数据有利于向量化。
• 硬件专用化： 某些硬件，如NVIDIA的Tensor Cores或Google的TPU，可能具有内部结构或指令集，能以特定布局（通常类似NHWC或平铺变体）的数据实现最高效的运行。

执行布局转换

机器学习 (machine learning)编译器将布局转换实现为图遍历。该过程通常包含：

1. 硬件性能分析/建模： 编译器需要有关目标硬件对不同操作（例如，卷积、池化、矩阵乘法）首选布局的信息。这可能来源于内置启发式规则、性能模型，甚至经验性性能分析数据。
2. 算子标注： 识别图中对布局敏感的算子（例如，卷积）与对布局不敏感的算子（例如，ReLU等逐元素操作）。对布局不敏感的算子通常会传播其输入或输出的布局偏好。
3. 布局偏好传播： 从对布局敏感的算子和硬件偏好开始，编译器尝试在图中传播所需的布局。例如，如果GPU卷积偏好NCHW输入和输出，则此偏好会传递给相邻的算子。
4. 冲突解决和转置插入： 当相邻算子偏好不同布局，或当一个算子接收到具有冲突布局要求的输入时，就会出现冲突。编译器必须决定在哪里插入显式转置节点（例如，NCHW_to_NHWC 或 NHWC_to_NCHW）。目标是尽量减少转置操作的数量，因为它们会带来计算开销和内存流量。成本模型用于权衡在首选布局下执行算子的收益与所需转置操作的成本。
5. 与合并的互动： 布局转换通常在算子合并之前进行或与算子合并交错进行。合并共享相同首选布局的算子可以消除中间张量以及它们之间潜在的转置操作。反之，选择特定布局可能会促成或阻碍某些合并机会。

示例：Conv -> ReLU -> Conv 的布局选择

思考一个简单的序列：Conv1 -> ReLU -> Conv2。

• 假设目标是GPU，cuDNN对卷积偏好NCHW。
• ReLU 对布局不敏感。

编译器可能这样分析：

• Conv1 偏好NCHW输入和输出。
• ReLU 可以在NCHW数据上高效运行，并传播该布局。
• Conv2 偏好NCHW输入。

在这种情况下，如果图的输入已经是NCHW（或者可以在开始时以较低成本进行转置），编译器很可能会在此序列中全程保持NCHW布局，避免任何内部转置。

示例：保持 NCHW 布局的图转换。

现在，假设 Conv2（可能由于其特定维度或不同的内核实现）在该硬件上使用NHWC输入时性能更优。

编译器的成本模型会评估：

1. Conv1(NCHW) -> ReLU(NCHW) -> 转置(NCHW->NHWC) -> Conv2(NHWC) 的成本。
2. 转置(输入->NHWC) -> Conv1(NHWC) -> ReLU(NHWC) -> Conv2(NHWC) 的成本（假设 Conv1 可以在NHWC中运行，但可能较慢）。
3. 其他组合。

如果 Conv2 在NHWC中的性能收益超过了转置成本，编译器就会插入转置操作。

示例：插入转置操作以优化 Conv2 执行的图转换。

挑战

选择最合适的布局很复杂：

• 转置成本： 转置操作并非没有开销。最大限度地减少其插入很重要。
• 硬件多样性： CPU、GPU 和专用加速器之间，最佳布局差异很大。
• 算子差异： 即使在相同的硬件上，不同的算子（或具有不同参数 (parameter)的相同算子）也可能偏好不同的布局。
• 全局最优与局部最优： 为整个复杂图找到全局最优的布局分配在计算上是困难的；编译器通常依赖启发式方法。

内存感知布局转换是一种强大的图优化技术。通过根据硬件特性和算子序列智能选择数据布局，编译器可以大幅提高缓存性能、内存带宽利用率和整体执行速度，经常与算子合并配合，以最大限度地减少开销并提高核函数效率。MLIR 等框架提供基础设施（例如，方言属性、接口）来表示布局信息并系统地实现这些转换。