图优化：布局优化

可视化演示

展示NCHW与NHWC两种4D张量数据布局的内存排列差异与效率、布局推断传播与转换消除三阶段优化流程、以及布局优化对transpose消除和内存访问效率的提升。

布局基础概念

小算第一次管理仓库时，面对四面货架不知所措。他决定先制定“摆放规则”：每个四维箱子（批次 N、品类 C、行 H、列 W）打开后，从里到外一层层摆。先摆完最里层的行，再摆这一行的下一行，摆完一个品类再摆下一个品类。这种从内到外的顺序，就是计算机里的“行优先”（row-major）：最内层的维度变化最快，相邻元素在内存里紧挨着。

在 AI 编译器眼中，张量就是多维数组，而内存是一条连续的一维空间。如何把多维下标映射到一维偏移量，完全由“步长”（stride）决定。行优先下，最后一维步长为 1，其余维度的步长等于它右侧所有维度尺寸的乘积。这个公式虽然简单，却是后续所有布局讨论的基石。

有了步长，访问任意元素都变成一次乘法加加法。比如 NCHW 布局的 (N,C,H,W) 张量，元素 (n,c,h,w) 的偏移是 n×(C×H×W) + c×(H×W) + h×W + w。注意，如果只改步长而不搬数据，你得到的是一个“视图”；真正把数据按新顺序重排，才产生新的内存排列——后者才是布局转换。

小算把这张“维度-跨度对照表”贴在仓库柱子上，每个师傅来取货都能自己算出位置。这张表，就是张量的 strides 元数据。理解了它，后面所有布局话题都有了抓手。

常见布局格式

切配师傅崔姐喜欢“按品类取料”——她要一口气把 32 个通道同一位置的数据抓齐，如果货架上这些数据隔了十万八千里，她就得来回跑。炒菜师傅贺哥则喜欢“按位置取料”——他卷积时要连续滑动窗口，希望同一位置的通道挨在一起。于是仓库衍生出两种主流摆放方案。

NCHW 把批次作大区、通道作中区、高宽作小区，同一通道的数据连续存放；NHWC 把通道变成最内层，同一空间位置 (h,w) 的所有通道紧挨着。两种布局数学上完全等价，但内存访问模式截然不同。下面两个公式分别给出各自的偏移量计算方式。

GPU 的卷积核通常对 NHWC 更友好——它把 C 放在最内层，配合 SIMD 一条指令可以加载多个通道的同一像素；而 CPU 上许多算子对 NCHW 支持更成熟，因为通道维连续访问利于缓存。还有更激进的分块布局（如 NCHW_VECT_C），把通道切成 4 或 8 的向量块，让向量化更进一步。没有绝对好坏，只有合不合适。

小算在仓库里划了 A 区和 B 区，分别按两种规则摆货。他没有私心——哪个区给哪个师傅用，全看师傅的需求。这就是编译器里的“布局分配”：先了解每一位师傅的偏好，再决定货架的摆法。

布局传播机制

仓库有条不成文的规矩：上一道工序的摆放方式，最好就是下一道工序的取料方式。如果每一站都顺手，整条流水线就丝滑；如果哪一站不顺手，搬运工就要插进来重新码货。小算要做的，就是沿着流水线把所有“顺手”关系串起来。

编译器把这套“顺手”翻译成“布局约束”：每个算子在实现时，能高效处理的布局是固定的。编译器需要给图中每条边分配一个布局，使它同时满足上下游两个算子的要求。这个分配过程就叫布局传播（layout assignment）。形式化地说，对图中任意一条数据边，两端的布局必须一致：

实际编译器通常从输出端开始回溯：输出布局是确定的（比如 Softmax 要求 NHWC），往前推，Softmax 的输入最好也是 NHWC；再往前，全连接层的输入输出布局关系又形成新的约束。遇到一个有明确要求的算子，约束就“卡”住一种布局；卡不住的地方，就放一个转换节点。

小算在流水线示意图上画箭头，从最后一道菜往前推，逐站确认“这里顺不顺”。一旦发现某站不顺，他就在图上画个红圈——那里要安排搬运工。布局传播的本质，就是提前把所有红圈找出来。

转换代价分析

搬运工老张最怕的不是箱子重，而是“整批重码”：把 NCHW 摆的货全部挪成 NHWC，等于把仓库 A 区的货全部搬到 B 区再换顺序。耗时永远和货物总量成正比，而且搬完还要对账。

布局转换（layout convert）在计算图里就是一个节点，内部实现是 transpose + copy。transpose 改变逻辑维度顺序，copy 把数据按新顺序写回。对一个 (N,C,H,W) 张量，一次完整的 NCHW→NHWC 转换的代价大约是两倍的张量元素数——读一遍、写一遍。

更麻烦的是，转换会切断算子的融合机会。编译器本来可以让 Conv+BN+ReLU 融合成一个算子，中间结果留在寄存器或缓存里；一旦插入转换，中间结果被迫写回内存，融合链断掉，性能雪上加霜。所以优化时宁可调整整个图的布局，也不轻易插入转换。

小算把老张的“重码成本”折算成表格，每次规划布局时先问一句：这里要让老张干活吗？老张一歇，流水线就快了。真正优秀的调度，是让搬运工从开工到收工都不用出手。

编译器实践

每家中央厨房的调度制度不同。XLA 厨房派了专门的“布局分配专员”；PyTorch 厨房给师傅们发了“备选工具包”；TVM 厨房则由老师傅凭经验现场拍板。三种制度各有各的智慧，也有共同的追求。

XLA 在 HLO 优化阶段有一个专门的 LayoutAssignment pass，它会根据每个算子的后端实现，为每个值分配布局，并把用户指定的布局约束尽量推满全图。PyTorch 通过 channels_last 内存格式让模型按 NHWC 运行，torch.compile 的 Inductor 后端会自动处理布局的选择与转换。TVM 则使用 ALT（AutoTVM-based Layout Transformation），按算子类型挑选最优输入布局。它们的共同目标可以用一个优化式表达：

实践中，布局分配还会考虑“哪个算子最贵就优先迁就谁”。比如卷积占全图 90% 算力，那全图都往卷积的偏好布局上靠，即使其他小算子多几次转换也划算。这背后是一个带权重的全局优化决策。

小算听完三家厨房的汇报，得出结论：好制度不是消灭搬运工，而是让搬运工出现在最便宜的位置。编译器们比来比去，比的都是这句话。

布局优化算法

小算终于要设计自己的调度算法了。他有两把刷子：一把叫“贪心”，眼睛只盯着当前这一步；一把叫“动态规划”（DP），心里装着全局。两把刷子配合使用，才能既快又稳。

贪心法从输出端开始，为每个算子挑一个“局部代价最小”的布局；速度极快，但可能因为只看到眼前而错过更优的全局解。DP 法则枚举每个算子的若干候选布局，用状态转移方程累计最小代价，在候选集有限时能保证找到全局最优。

真正的编译器通常把两者结合：先用启发式规则缩小候选布局集合，再用 DP 或类似搜索求解。成本模型（cost model）负责估算每个算子在某种布局下的执行时间——它本身可能是个小型回归模型或经验表，这也是布局优化能持续迭代的关键。

小算把 DP 表格贴满了一面墙，每个格子都是一条决策路径。他发现，所谓“最优布局”，不过是在每一站都问一次“从哪条路走过来最省”，然后把答案记下来。算法之美，就是把麻烦留给表格，把清爽留给流水线。

性能影响分析

货架怎么摆，直接影响师傅的脚程。如果炒菜师傅要的四种调料在四面墙上，他炒一道菜要跑四趟；如果集成一个调料架，一趟全拿齐。速度差异肉眼可见，布局对 AI 模型的影响也是同样的道理。

布局对性能的影响主要有三处。第一是访存局部性：NHWC 让同一空间位置的通道数据连续，卷积滑动窗口时读的全是相邻地址，缓存命中率大幅提升。第二是 SIMD 向量化：通道在最内层时，一条指令可以加载多个通道的同一像素，计算密度飙升。第三是内存带宽：消除了转置的读改写，梯度回传时的 gather/scatter 也变得更规整。向量化利用率是对第二点的量化：

实测数据最有说服力：ResNet-50 在 GPU 上用 NHWC 配合 Tensor Core，比 NCHW 快 1.3~2 倍；BERT 的矩阵乘在某些加速卡上，布局错了甚至会触发全量拷贝，单算子慢 10 倍。布局优化的收益不是理论上的“一点点”，而是量级上的。

小算把师傅们的脚程数据画成热力图，发现所有慢的站点都贴着“布局不匹配”的标签。他立刻明白：调度主管的最高荣誉，不是把货架擦得多亮，而是让每位师傅少跑路。

工程实践调优

再完美的调度方案也需要日报来校正。小算每天早上第一件事，是看流水线的监控屏：哪站堵了，搬运工老张在哪儿忙，一目了然。没有监控，再好的理论也只是纸上谈兵。

工程中定位布局问题，靠的是 profiler。PyTorch 的 torch.profiler 能列出每个算子的耗时，如果看到大量 copy_、permute、contiguous 节点，基本可以断定布局转换在拖后腿。Nsight 等工具还能进一步看到访存指令与缓存未命中。定位之后，再用 memory_format=torch.channels_last 显式指定模型输入布局，或调整图中关键算子的输入格式。调优效果看这个指标：

常见问题清单：模型用 nn.Conv2d 默认 NCHW，但输入数据是 NHWC——第一个算子就插转换；自定义算子不支持 channels_last，导致下游全部转回 NCHW；多分支网络在汇合处布局冲突，产生隐性转换。解决思路统一：把“谁需要什么布局”的表格列出来，让布局从源头一致。

小算的日报本上，红圈越来越少。他给新人培训时总说：布局优化不是魔法，是“提前规划 + 持续观测”。把这句话抄在仓库门口，流水线就再也没堵过。

AI 应用场景

Python 演示

# 布局分配模拟器：从输出端回溯，为计算图选择布局# 目标是尽量减少布局转换（搬运）的次数 class Op:    """一个算子：记录名字、期望的输入输出布局"""    def __init__(self, name, in_layout, out_layout):        self.name = name        self.in_layout = in_layout        self.out_layout = out_layout def assign_layouts(ops):    """从最后一个算子开始，向前为每条边选定布局"""    n = len(ops)    edges = [None] * (n + 1)          # 每条边的最终布局    edges[n] = ops[-1].out_layout     # 输出端布局固定    converts = 0    for i in range(n - 1, -1, -1):        op = ops[i]        if edges[i + 1] == op.out_layout:            edges[i] = op.in_layout        else:            converts += 1             # 需要插一次搬运            edges[i] = op.in_layout    return edges, converts def naive_layouts(ops):    """朴素方案：每个算子都坚持用自己最舒服的布局"""    n = len(ops)    edges = [None] * (n + 1)    converts = 0    for i, op in enumerate(ops):        if edges[i] is not None and edges[i] != op.in_layout:            converts += 1             # 上游不配合，插转换        edges[i] = op.in_layout        edges[i + 1] = op.out_layout    return edges, converts def main():    # 计算图：Conv(NCHW->NHWC) -> ReLU(NHWC->NHWC) -> Softmax(NHWC->NHWC)    ops = [        Op("Conv",    "NCHW", "NHWC"),        Op("ReLU",    "NHWC", "NHWC"),        Op("Softmax", "NHWC", "NHWC"),    ]    e1, c1 = assign_layouts(ops)    e2, c2 = naive_layouts(ops)    print("回溯传播方案：", c1, "次转换，边布局", e1)    print("朴素方案：", c2, "次转换，边布局", e2) if __name__ == "__main__":    main()

1. class Op:定义一个算子，记录它的名字和期望的输入/输出布局，就像仓库里每位师傅的取料偏好。
2. edges = [None] * (n + 1)图中 n 个算子会形成 n+1 条数据边，这里给每条边准备一个布局位。
3. edges[n] = ops[-1].out_layout从最后一道工序开始，输出端的布局是事先定好的（如同成品菜的装盘方式固定）。
4. if edges[i + 1] == op.out_layout:检查下游能不能接住当前算子的输出布局，能接住就说明这一站“顺手”。
5. converts += 1接不住就记一次“搬运”，对应图中插入一次布局转换节点。
6. edges[i] = op.in_layout把当前算子的输入布局继续传给上游，实现从后往前的约束传播。
7. def naive_layouts(ops):对照用的朴素方案：每个算子只坚持自己最舒服的布局，不管上下游死活。

常见误区

面试题