第 6 章：数据布局优化

本章深入探讨 AI 编译器中的数据布局优化技术。数据布局直接影响内存访问模式、缓存利用率和硬件加速器效率，是实现高性能推理和训练的关键。我们将学习如何根据硬件特性选择最优布局、最小化布局转换开销、优化内存对齐，以及在混合精度计算中管理不同精度的数据布局。

学习目标：

• 理解不同数据布局对性能的影响机制
• 掌握布局选择的定量分析方法
• 学会设计布局转换最小化算法
• 理解硬件对齐要求与优化策略
• 掌握混合精度场景下的布局管理

6.1 NCHW vs NHWC 选择策略

6.1.1 布局定义与内存访问模式

在深度学习中，四维张量的两种主要布局方式为：

• NCHW（批次-通道-高度-宽度）：也称为通道优先布局
• NHWC（批次-高度-宽度-通道）：也称为通道最后布局

对于张量 $T \in \mathbb{R}^{N \times C \times H \times W}$，内存地址计算公式为：

NCHW布局： $$\text{addr}(n,c,h,w) = \text{base} + \text{sizeof}(T) \cdot (n \cdot CHW + c \cdot HW + h \cdot W + w)$$ NHWC布局： $$\text{addr}(n,h,w,c) = \text{base} + \text{sizeof}(T) \cdot (n \cdot HWC + h \cdot WC + w \cdot C + c)$$

6.1.2 硬件特性分析

不同硬件架构对布局有不同的偏好：

GPU（NVIDIA）特性：

• Tensor Core 要求特定的内存布局和对齐
• Warp 级别的合并访存偏好连续的通道数据
• cuDNN 大部分算子针对 NCHW 优化

CPU（x86/ARM）特性：

• SIMD 指令（AVX-512, NEON）更适合处理连续的小块数据
• 缓存行大小（通常 64 字节）影响最优布局选择
• NHWC 在小通道数时缓存友好

移动端 NPU 特性：

• 通常偏好 NHWC 布局（如 TensorFlow Lite 默认）
• 硬件设计针对逐像素处理优化

6.1.3 性能建模与决策

定义布局选择的代价函数： $$C_{\text{layout}} = \alpha \cdot C_{\text{compute}} + \beta \cdot C_{\text{memory}} + \gamma \cdot C_{\text{transform}}$$ 其中：

• $C_{\text{compute}}$：计算效率损失
• $C_{\text{memory}}$：内存访问代价
• $C_{\text{transform}}$：布局转换开销
• $\alpha, \beta, \gamma$：权重系数（硬件相关）

卷积操作的内存访问分析：

对于卷积核 $K \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times K_h \times K_w}$：

NCHW 布局下的缓存未命中率： $$M_{\text{NCHW}} = 1 - \frac{\min(L_1, C_{\text{in}} \cdot K_h \cdot K_w \cdot \text{sizeof}(T))}{C_{\text{in}} \cdot K_h \cdot K_w \cdot \text{sizeof}(T)}$$ NHWC 布局下的缓存未命中率： $$M_{\text{NHWC}} = 1 - \frac{\min(L_1, W \cdot C_{\text{in}} \cdot \text{sizeof}(T))}{W \cdot C_{\text{in}} \cdot \text{sizeof}(T)}$$ 其中 $L_1$ 为 L1 缓存大小。

6.1.4 启发式选择算法

算法 6.1：数据布局选择
输入：计算图 G，硬件规格 H
输出：每个张量的布局 L

1. 初始化所有张量为硬件偏好布局
2. for 每个算子 op in G:
3.     计算 op 在不同布局下的性能 P_NCHW, P_NHWC
4.     if P_NCHW / P_NHWC > threshold:
5.         标记 op 的输入输出为 NCHW
6.     else if P_NHWC / P_NCHW > threshold:
7.         标记 op 的输入输出为 NHWC
8. 运行布局传播算法解决冲突
9. return 最终布局分配 L

6.2 布局转换最小化

6.2.1 转换代价模型

布局转换（transpose）的代价包括：

1. 内存带宽消耗： $$B_{\text{transpose}} = 2 \times \text{TensorSize} \times \text{sizeof}(T)$$

2. 缓存污染： $$C_{\text{pollution}} = \frac{\text{TensorSize}}{\text{CacheSize}} \times P_{\text{reuse}}$$

3. 计算开销： $$T_{\text{compute}} = \frac{\text{TensorSize}}{\text{Bandwidth}} + \text{LatencyOverhead}$$

6.2.2 图级优化算法

将布局优化建模为图着色问题：

定义增广图 $G' = (V', E')$：

• 节点 $V'$：原图中的张量
• 边 $E'$：张量间的数据依赖
• 边权重：转换代价

目标函数： $$\min \sum_{e \in E'} w_e \cdot \mathbb{1}[\text{layout}(v_i) \neq \text{layout}(v_j)]$$ 这是一个 NP-hard 问题，使用近似算法求解。

6.2.3 动态规划解法

对于链式结构的计算图，可使用动态规划： $$dp[i][l] = \min_{l' \in \{NCHW, NHWC\}} \{dp[i-1][l'] + C_{\text{op}}(i, l) + C_{\text{trans}}(l', l)\}$$ 其中：

• $dp[i][l]$：前 $i$ 个算子，第 $i$ 个算子使用布局 $l$ 的最小代价
• $C_{\text{op}}(i, l)$：算子 $i$ 在布局 $l$ 下的计算代价
• $C_{\text{trans}}(l', l)$：从布局 $l'$ 转换到 $l$ 的代价

6.2.4 贪心插入策略

算法 6.2：贪心转换插入
输入：计算图 G，初始布局 L
输出：带转换的计算图 G'

1. 计算每条边的布局不匹配代价
2. 按代价降序排序边集合
3. for 每条高代价边 e = (u, v):
4.     计算插入转换的收益 R
5.     if R > 0:
6.         在边 e 上插入转换节点
7.         更新相关边的代价
8. return 修改后的图 G'

6.3 Padding 与对齐优化

6.3.1 硬件对齐要求

不同硬件的对齐要求：

| 硬件类型 | 对齐要求 | 原因 |

6.3.2 Padding 策略

最小 Padding 算法：

给定原始维度 $d$ 和对齐要求 $a$： $$d_{\text{padded}} = \lceil \frac{d}{a} \rceil \times a$$ Padding 开销： $$\text{Overhead} = \frac{d_{\text{padded}} - d}{d} \times 100\%$$ 多维 Padding 优化：

对于张量 $T \in \mathbb{R}^{N \times C \times H \times W}$，需要考虑：

1. 最内层维度优先： - NCHW：优先 pad W 维度 - NHWC：优先 pad C 维度

2. 访存模式匹配：

if (硬件支持向量化):
    align_size = vector_width
else:
    align_size = cache_line_size

6.3.3 动态 Padding 决策

考虑计算与内存的权衡： $$\text{PaddingBenefit} = T_{\text{saved}} - M_{\text{wasted}}$$ 其中：

• $T_{\text{saved}}$：对齐带来的计算时间节省
• $M_{\text{wasted}}$：额外内存开销的代价

自适应 Padding 算法：

算法 6.3：自适应 Padding
输入：张量维度 D，硬件规格 H
输出：padded 维度 D'

1. for 每个维度 d in D:
2.     候选集 = {d, align_ceil(d, 16), align_ceil(d, 32), ...}
3.     for 每个候选 c in 候选集:
4.         计算性能收益 perf_gain(c)
5.         计算内存开销 mem_cost(c)
6.     选择最优 padding: d' = argmax(perf_gain - λ * mem_cost)
7. return D'

6.3.4 卷积的隐式 Padding

对于卷积操作，padding 还涉及边界处理： $$\text{output_size} = \lfloor \frac{\text{input_size} + 2 \times \text{pad} - \text{kernel_size}}{\text{stride}} \rfloor + 1$$ 优化目标：

1. 保持输出尺寸为硬件友好值
2. 最小化边界填充的计算开销
3. 考虑感受野的对称性

6.4 混合精度布局

6.4.1 精度类型与存储格式

常见精度类型的存储特性：

| 类型 | 位宽 | 指数位 | 尾数位 | 范围 | 硬件支持 |

内存布局影响：

对于混合精度张量，内存地址计算需要考虑不同精度的对齐： $$\text{addr}_{\text{mixed}}(i) = \text{base} + \sum_{j=0}^{i-1} \text{align}(\text{sizeof}(T_j), a)$$ 其中 $T_j$ 为第 $j$ 个元素的类型，$a$ 为硬件对齐要求。

6.4.2 精度转换策略

量化感知布局：

对于量化操作 $Q: \mathbb{R} \rightarrow \mathbb{Z}$： $$Q(x) = \text{round}\left(\frac{x - z}{s}\right)$$ 其中 $s$ 为缩放因子，$z$ 为零点。

布局选择影响量化效率：

• 通道量化：NCHW 布局下按 C 维度量化更高效
• 逐层量化：NHWC 布局下全局量化访存连续

6.4.3 混合精度图优化

定义精度分配问题：

给定计算图 $G = (V, E)$，为每个节点 $v \in V$ 分配精度 $p_v \in \{FP32, FP16, INT8\}$，优化目标： $$\min \sum_{v \in V} T_v(p_v) + \sum_{(u,v) \in E} C_{uv}(p_u, p_v)$$ 其中：

• $T_v(p_v)$：节点 $v$ 在精度 $p_v$ 下的计算时间
• $C_{uv}(p_u, p_v)$：精度转换代价

动态规划求解： $$dp[v][p] = T_v(p) + \min_{p' \in P} \{dp[\text{pred}(v)][p'] + C(p', p)\}$$

6.4.4 存储格式优化

打包存储（Packed Storage）：

对于 INT4 等亚字节类型，需要打包存储：

原始布局：[a0][a1][a2][a3][a4][a5][a6][a7]（每个4位）
打包布局：[a0|a1][a2|a3][a4|a5][a6|a7]（每个字节）

访问公式： $$\text{packed_addr}(i) = \text{base} + \lfloor i / k \rfloor$$ $$\text{offset}(i) = (i \bmod k) \times \text{bits_per_elem}$$

其中 $k$ 为每个存储单元的元素数。

混合精度分块：

算法 6.4：混合精度分块布局
输入：张量 T，精度映射 P
输出：优化后的内存布局 L

1. 将张量按精度类型分组
2. for 每个精度组 g:
3.     计算最优块大小 b = hardware_preferred_size(g.precision)
4.     将 g 划分为大小为 b 的块
5.     for 每个块 block:
6.         应用硬件特定的布局（如 tensor core 布局）
7. 按访问模式排列各精度块
8. return 分块布局 L

6.4.5 自动驾驶场景的精度布局

在自动驾驶场景中，不同任务对精度要求不同：

| 任务类型 | 推荐精度 | 布局偏好 | 原因 |

多模态融合的布局策略：

摄像头 + LiDAR + Radar 融合时：

1. 统一坐标系后的点云：FP32 + 稀疏格式
2. 图像特征：INT8/FP16 + NHWC
3. 融合特征：FP16 + 自定义打包格式

本章小结

数据布局优化是 AI 编译器实现高性能的关键技术。本章介绍了：

1. 布局选择：基于硬件特性和访存模式的 NCHW/NHWC 选择策略
2. 转换最小化：通过图优化算法减少布局转换开销
3. 对齐优化：平衡硬件对齐要求与内存开销
4. 混合精度：不同精度类型的布局管理和优化

关键公式回顾：

• 布局选择代价：$C_{\text{layout}} = \alpha \cdot C_{\text{compute}} + \beta \cdot C_{\text{memory}} + \gamma \cdot C_{\text{transform}}$
• Padding 计算：$d_{\text{padded}} = \lceil \frac{d}{a} \rceil \times a$
• 混合精度优化：$\min \sum_{v \in V} T_v(p_v) + \sum_{(u,v) \in E} C_{uv}(p_u, p_v)$

练习题

练习 6.1（基础）：布局转换开销计算

给定张量 $T \in \mathbb{R}^{32 \times 256 \times 56 \times 56}$，数据类型为 FP16，内存带宽为 900 GB/s。计算从 NCHW 转换到 NHWC 布局的理论时间下界。

提示：考虑读写各一次的内存访问量。

练习 6.2（基础）：对齐计算

某 NPU 要求所有维度都对齐到 16 的倍数。给定输入张量维度 [7, 129, 55, 31]，计算：

1. Padded 后的维度
2. 内存浪费比例

提示：使用向上取整公式。

练习 6.3（基础）：缓存命中率分析

在 L1 缓存为 48KB 的 GPU 上执行 3×3 卷积，输入通道数为 512，批次大小为 1。比较 NCHW 和 NHWC 布局下，计算一个输出像素时的缓存命中率。假设数据类型为 FP16。

提示：计算卷积窗口的数据量是否能装入 L1 缓存。

练习 6.4（挑战）：布局传播算法

给定如下计算图，每个算子在不同布局下的计算时间（ms）和转换代价为 2ms。使用动态规划求解最优布局分配。

输入 → Conv1 → ReLU → Conv2 → 输出

| 算子 | NCHW | NHWC |

提示：构建状态转移方程。

练习 6.5（挑战）：混合精度优化

某模型包含三个连续的矩阵乘法：$Y = ABC$，其中 $A \in \mathbb{R}^{1024 \times 512}$，$B \in \mathbb{R}^{512 \times 256}$，$C \in \mathbb{R}^{256 \times 128}$。不同精度下的计算时间（ms）和转换时间如下：

| 操作 | FP32 | FP16 | INT8 |

转换时间：FP32↔FP16: 1ms，FP16↔INT8: 2ms，FP32↔INT8: 3ms

求最优精度分配方案。

提示：枚举所有可能的精度组合。

练习 6.6（挑战）：稀疏数据布局

设计一个数据结构存储稀疏率为 90% 的 $1000 \times 1000$ 矩阵，要求：

1. 支持高效的行访问
2. 支持高效的列访问
3. 最小化存储开销

分析你的设计在不同访问模式下的时间复杂度。

提示：考虑 CSR 和 CSC 的混合方案。

CSR格式：

- values[]: 非零值数组（100,000 个元素）
- col_indices[]: 列索引（100,000 个元素）
- row_ptr[]: 行指针（1,001 个元素）

CSC格式：

- values[]: 共享 CSR 的值数组
- row_indices[]: 行索引（100,000 个元素）
- col_ptr[]: 列指针（1,001 个元素）

练习 6.7（开放）：自动驾驶多传感器布局

设计一个统一的内存布局方案，处理自动驾驶中的多模态数据：

• 6 个摄像头图像（1920×1080×3，30 FPS）
• 1 个 LiDAR 点云（~100,000 点/帧，10 FPS）
• 5 个毫米波雷达（~1,000 点/帧，20 FPS）

要求最小化数据搬移，支持时间同步。

提示：考虑环形缓冲区和时间戳对齐。

Level 1（高频）：摄像头环形缓冲区

- 6 × 1920×1080×3 × 2字节(FP16) = 74.6 MB/帧
- 缓存 100ms（3帧）= 224 MB
- 布局：NHWC（便于 CNN 处理）

Level 2（中频）：雷达环形缓冲区

- 5 × 1000 × 16字节 = 80 KB/帧
- 缓存 200ms（4帧）= 320 KB
- 布局：结构体数组（x,y,z,velocity,intensity）

Level 3（低频）：LiDAR 环形缓冲区

- 100,000 × 16字节 = 1.6 MB/帧
- 缓存 300ms（3帧）= 4.8 MB
- 布局：八叉树 + 原始点云

常见陷阱与错误（Gotchas）

1. 布局转换的隐式开销

陷阱：只考虑单个算子的最优布局，忽略全局转换开销。

案例：ResNet 中，虽然深度可分离卷积在 NHWC 下更快，但频繁的布局转换反而降低整体性能。

解决方案：使用全图优化，设置转换代价阈值。

2. 对齐导致的内存爆炸

陷阱：盲目对齐所有张量到硬件要求的最大值。

案例：将所有维度对齐到 256，小张量内存使用增加 100 倍。

解决方案：

• 分级对齐策略
• 只对性能关键路径应用严格对齐

3. 混合精度的数值不稳定

陷阱：在梯度累积时使用 FP16，导致下溢。

案例：大模型训练中，梯度值 < 1e-7 在 FP16 下变为 0。

解决方案：

• 梯度使用 FP32 累积
• 动态损失缩放

4. 稀疏格式选择不当

陷阱：对所有稀疏矩阵使用同一格式。

案例：对行稀疏矩阵使用 CSC 格式，性能下降 10 倍。

解决方案：根据访问模式和稀疏模式选择格式。

5. 缓存行冲突

陷阱：Padding 大小恰好是缓存大小的倍数。

案例：每行 padding 到 2048 字节，在 32KB 8-way 缓存上产生冲突。

解决方案：使用素数大小的 padding 或添加随机偏移。

最佳实践检查清单

设计阶段

• [ ] 分析目标硬件的内存层次结构
• [ ] 确定主要算子的访存模式
• [ ] 评估不同布局的理论性能上界
• [ ] 考虑多精度训练/推理需求

实现阶段

• [ ] 实现布局无关的算子接口
• [ ] 添加布局转换的 profiling 代码
• [ ] 支持运行时布局选择
• [ ] 实现内存池避免频繁分配

优化阶段

• [ ] 使用真实硬件测试不同布局
• [ ] 分析缓存命中率和带宽利用率
• [ ] 识别并消除不必要的转换
• [ ] 平衡内存使用和计算效率

验证阶段

• [ ] 对比优化前后的端到端性能
• [ ] 检查内存使用峰值
• [ ] 验证数值精度（特别是混合精度）
• [ ] 测试边界情况（很小/很大的张量）

部署阶段

• [ ] 根据部署硬件微调布局策略
• [ ] 添加布局选择的配置选项
• [ ] 文档化布局决策原理
• [ ] 提供性能调优指南