AI 编译器前端优化全景:9 个 Pass、两类优化、一张地图
这一篇是 AISystem「AI 编译器前端优化」章节的收尾总结。原始材料只有 8 页 slide:一张分层图 + 一张反复高亮的分组目录。目录本身信息量不大,但它给出的分类骨架很值钱——前端优化被切成"神经网络相关"和"代码/计算层面"两大块,共 9 个 pass。
下面我按这个骨架把整章串起来:每个 pass 一句话讲清"为什么需要它"、它吃什么信息、收益怎么量、什么时候会翻车。细节推导不重复,指向本系列已单独成篇的文章。
文章目录
一、这份总结材料的全部内容
8 页 slide,四类内容:
| 页码 | 内容 |
|---|---|
| 1 | 封面:AI 编译器系列·之前端优化 Summary |
| 2 | “Where are we?” 标题 |
| 3 | 分层架构图:Python → 框架 → Graph Optimizer → Ops Optimizer → Backend → 硬件 |
| 4 | 目录:8 项前端优化 |
| 5 | 目录高亮:左列 4 项标为 神经网络相关的优化 |
| 6 | 目录新增 算子替换 - OP Replace,左列变 5 项 |
| 7 | 右列 4 项标为 代码/计算层面的优化 |
| 8 | 版权页 |
也就是说,真正的信息只有两条:
- 前端有 9 个优化(4 + 5);
- 它们被分成两类:神经网络相关 与 代码/计算层面。
剩下的都是"为什么这么分、彼此什么关系、先后怎么排"。这些 slide 里没有,是这篇要补的。
二、先定位:这 9 个 pass 落在哪一层
把第 3 页那张分层图重画一遍(这是整章最重要的一张图):
图中那三块带虚线框的左侧清单,正是三层各自的技术清单——这张表建议直接记住:
| 层 | 输入/输出 IR | 对应技术(图左侧清单) | 回答的问题 |
|---|---|---|---|
| Graph Optimizer | Graph IR | Tensor 表达、NN API 表达、优化器、计算图转换 | 图长什么样:算子怎么连、哪些能合并/删除 |
| Ops Optimizer | Tensor IR | 自动微分、图算融合、剪枝量化、并行切分 | 张量怎么算:单算子内部怎么分解、怎么并行 |
| Backend | → LLVM IR / GE IR | Schedule opt、SoC 存储、kernel tiling、DSL 描述 | 指令怎么生成:循环怎么排、数据搬几次、寄存器够不够 |
关键认知:前端(前两层)产出的是 Tensor IR,后端才把它降解成 LLVM IR / 厂商 IR(如 GE IR)。所以前端优化的共同特征是——它们都在"还不知道生成什么指令"的时候做决定。
大白话:前端管"这张图该怎么长",后端管"长成这样之后,怎么变成机器码"。前端做的每一个决定,都会变成后端的约束。
为什么必须有前端这一层:因为框架之间算子语义不一致(padding 约定、广播规则、Conv 的 4 种实现),不统一成 Graph IR,后面所有优化都得为每个框架写一遍。
三、两类优化:判据其实只有一条
判据一句话:这个优化需不需要知道"这是个神经网络"?
| 神经网络相关的优化 | 代码/计算层面的优化 | |
|---|---|---|
| 依据 | 领域知识:算子语义、张量形状、数据布局、AI 硬件特性 | 通用编译原理:数据流、格、不动点、生命周期 |
| 例子 | 知道 Conv + BN + ReLU 可以折成一个 Conv | 知道 x * 1 可以变成 x(跟是不是 AI 无关) |
| 换个领域还能用吗 | 换到视频编码器基本用不上 | 换到 GCC/LLVM 一模一样能用 |
| 翻车风险 | 语义风险:改了可能算错(padding、广播、布局) | 正确性风险:别名、溢出、浮点结合律 |
| 可验证性 | 靠数值对齐(对拍输出) | 靠形式化(数据流分析的正确性定理) |
大白话:右列是"编译器通用手艺",左列是"AI 专属手艺"。这也解释了为什么传统编译器工程师转 AI 编译,右边那 4 项几乎是白送,左边那 5 项要从头学。
四、九个 pass 逐个过一遍
4.1 图层 — Graph IR
一句话:把各框架的计算图统一成一种统一的图中间表示,并做一轮规范化。
做四件事:
| 动作 | 内容 |
|---|---|
| 统一表达 | 把 TensorFlow 的 Conv2D、PyTorch 的 conv2d、ONNX 的 Conv 映射成同一个图节点 |
| 形状与类型推断 | 每个节点算出 output shape / dtype,为后续所有 pass 提供前提(这是 90% 优化判断的输入) |
| 控制流表达 | If / While / Loop 的控制边与数据边;动态 shape 的符号维度 |
| 规范化(canonicalization) | 把等价写法归一到一种(如把 Pad+Conv 统一成带 padding 属性的 Conv) |
为什么它排第一:后面 8 个 pass 全都读这张图。图的表达质量,决定了后面能优化到什么程度。一个典型例子:如果图 IR 没有保留 padding 语义、只留一个 Pad 算子,融合就会多出一个不该存在的节点。
风险:语义保真。不同框架的广播规则、padding 是 SAME 还是 VALID、axis 约定,映射错一个就是静默的数值错误。
4.2 算子融合 — OP Fusion
一句话:把多个算子合成一个 kernel,省掉中间张量的往返显存。
- 收益公式:省下 ≈ 2 × \approx 2 \times ≈2× 中间张量字节数(写一次 + 读一次)+ 若干次 kernel launch。
- 纵向融合(生产者-消费者,如
Conv + BN + ReLU)与横向融合(同输入的多消费者)。 - 前提:数据依赖清晰、无副作用、融合后寄存器/共享内存放得下。
- 典型实现:TVM 的算子融合基于支配树与
Group;TF 的 Grappler 有remapper/arithmetic_optimizer。
翻车点:融合不是越多越好。把太多 elementwise 塞进一个大 kernel,会让寄存器压力上升、occupancy 下降,实际可能更慢。所以融合要有"收益模型",不是"能融就融"。
详解见本系列《AI 编译器算子融合》一篇(含 Conv+BN+ReLU 的代数折叠公式
γ
/
σ
2
+
ϵ
\gamma/\sqrt{\sigma^2+\epsilon}
γ/σ2+ϵ)。
4.3 算子替换 — OP Replace
一句话:用一个语义等价但更高效的算子,替换原算子。
这是"神经网络相关"里最容易被忽略、但在实际落地中最常见的一项。典型替换:
| 原算子/模式 | 替换为 | 动机 |
|---|---|---|
Conv + BatchNorm(推理) | 单个 Conv(权重已折叠) | 省一个算子 + 一次访存 |
Conv + Scale | 单个 Conv | 同上(BN 拆出来的 Scale 也要折) |
Dropout(推理期) | 删除 / 恒等 | 训练专用,推理无用 |
Conv(通用布局) | CONV_NCHW16c 这类布局专用算子 | 直接吃 NPU 的分形布局,省转换 |
| 大 kernel 卷积 | 分解成多个小 kernel(如 7×7 → 1×7 + 7×1) | 减少乘加次数(可分离卷积) |
MatMul(批量场景) | BatchMatMul | 硬件有专门指令 |
高精度 Gelu | 近似实现(tanh 近似等) | 换性能,但要显式接受精度损失 |
它和"融合"的边界:融合是多个算子 → 一个(改变图结构);替换是一个算子 → 另一个(不改变拓扑)。但两者常常一起用:先把
Conv+BN替换成单个Conv(权重折叠),再把这个Conv和后面的ReLU融合。
风险:替换必须逐位对齐(至少数值误差在容差内),尤其是 padding、dilation、groups 三个属性最容易在替换中被吃掉。布局专用算子的替换尤其危险——它隐含了"数据已经是某个布局"这个前提,前提不成立就直接算错。
4.4 布局转换 — Layout Transform
一句话:把张量从一种逻辑/物理布局搬到另一种,让算子吃到对硬件最友好的内存排布。
- 核心概念:
NCHW/NHWC/NC1HWC0(5HD)、C0=16(FP16)/32(INT8)、FRACTAL Z/NZ 分形。 - 本质:布局就是 stride。
transpose只改 stride(view),真正搬数据才叫转换(contiguous)。 - 代价:一次转换 = 整张量读 + 写。所以原则是"少转换、长直通":在入口/出口各转一次,中间让 layout-tolerant 算子直接吃目标布局。
- 权重预转换:
KCRS → OIHW16i16o这类转换在离线做,不进运行时。
收益不在"转换本身",而在"转换之后剩下的所有算子都能高效跑"。所以它是一次投入、长期收益的优化。
详解见本系列《布局转换原理》与《AI 编译器布局转换算法》两篇(后者含 NHWC → NC1HWC0 的两行变换代码)。
4.5 内存分配 — Memory Allocation
一句话:给每个张量安排一块显存,并让生命周期不重叠的张量复用同一块。
- 判据是生命周期,不是"静态/动态":持久(权重/梯度/优化器状态/常量)、瞬态(激活/workspace)、外部持有(输入输出)。
- 两个省内存的招:Inplace(原地算)与 Memory Sharing(不同时活着的张量共用一块 buffer)。
- 算法本质:
tag + 内存池,把"生命周期重叠"建成干涉图,然后着色——这和寄存器分配是同一个问题。 - 训练显存的真相:混合精度 Adam 大约 16 字节/参数(2 权重 + 2 梯度 + 4 master + 4 动量 + 4 方差),7B 模型光这一项就 112 GB。
翻车点:复用会引入 WAW 假依赖——本来能并行的两条分支,因为共用一块内存被串行化。所以"省内存"和"并行度"是要权衡的。
详解见本系列《AI 编译器内存分配算法》一篇。
4.6 常量折叠 — Constant Fold
一句话:编译期就把纯常量表达式算完,用结果替换表达式。
- 触发条件很硬:所有子表达式都是常量。所以它依赖常量传播,而且必须迭代到不动点(折一个 → 可能带出下一个)。
- 三类可折节点:全常量输入、只看形状的算子(
Shape/Size)、链式传导。 - 收益是算子数、kernel launch、中间张量,不是模型体积(权重是输入常量,不是折出来的)。
- 不能折的情况:有副作用/随机/读变量的算子;结果超过实现的大小上限;浮点语义不允许。
翻车点:把
RandomUniform折了 → 每批采样全一样;把ReadVariableOp折了 → 训练时权重更新被冻结。
详解见本系列《常量折叠原理》一篇(含 clang -O0 vs -O2 的实测对比)。
4.7 公共子表达式消除 — CSE
一句话:同一个值只算一次,重复的子图合并成一份。
- 前提三条:值不变、纯计算、暂存值不变。
- 局部(LVN 值编号)→ 全局(LCM 缓式代码移动,四集合公式,核心是 must 分析)。
- 在计算图上退化成"哈希合并":DAG + 天然 SSA,不需要 φ、不需要支配性检查。
- 哈希键必须完备:算子类型 + 全部属性(padding/strides/dtype/axis)+ 输入顺序 + 输出个数 + 常量内容。
翻车点:少一个属性进哈希键,就会把
SAME和VALID的卷积合并成一个——这是正确性事故,不是性能问题。
详解见本系列《公共子表达式消除原理》一篇。
4.8 死代码消除 — DCE
一句话:没人用的东西,删掉。
它是前三个 pass 的收尾工:常量折叠删了消费者 → 原来的常量输入变孤儿;CSE 合并后 → 被合并的那份变孤儿;算子替换后 → 旧算子(如 BN)变孤儿。
四个容易踩的点:
| 点 | 说明 |
|---|---|
| 连锁删除 | 删一个节点后,它的输入可能也变孤儿 → 必须迭代,不能扫一遍 |
| 副作用节点不能删 | Assign、Print、Assert、以及带控制边的节点,即使无数据消费者也要留 |
| 控制依赖 | TF 的 ^ 边、NPU 图的 ctrl 边,删节点要一起维护 |
| 随机算子 | RandomUniform 的结果没人用可以删(它没有外部可见副作用),但删了会改变随机数消耗序列,可能影响可复现性 |
DCE 的收益常被低估:在 NPU 上,一个多余算子 = 一次 kernel launch + 一次显存占用;批量推理时这是实打实的延迟。
4.9 代数简化 — Simplify Algebraic
一句话:用代数恒等式把"合法但多余"的运算删掉或重排。
分类记这一组就够:
| 类型 | 例子 | 注意 |
|---|---|---|
| 单位元 | x*1 → x、x+0 → x、x/1 → x | 安全 |
| 零元 | x*0 → 0 | 安全(但 0*NaN 语义不同,浮点下要谨慎) |
| 自反 | x-x → 0、x/x → 1 | ⚠️ x/x 在 x=0 时是 NaN 或异常,不能随便折 |
| 结合/分配重排 | x*a + x*b → x*(a+b) | ⚠️ 浮点结合律不成立,会改变结果精度 |
| 矩阵恒等 | MatMul(I, x) → x、Transpose(Transpose(x)) → x、Reshape(Reshape(x)) → Reshape(x) | 安全,且能解锁后续融合 |
| 结构简化 | Concat(Split(x)) → x、Slice(Concat) 合并 | 常见于框架自动生成的图 |
| 幂运算 | pow(x, 2) → x*x | 快很多,但要注意 x 为负/整型溢出 |
这一项和常量折叠的逻辑完全不同:常量折叠是"算出来",代数简化是"变成另一种写法"。两者都是在"代码/计算层面",也都不依赖领域知识。
风险集中在浮点语义:(a+b)+c ≠ a+(b+c) 在浮点下不成立。所以很多编译器把"重排"归到 -ffast-math 一档,默认不开——因为会破坏可复现性。
五、它们不是并列的,是有顺序的
这 9 个 pass 在真实编译器里排成一条循环流水线,跑到不动点:
为什么是这个顺序(每一条都有硬理由):
| 顺序 | 理由 |
|---|---|
| 化简在融合前 | 先删掉多余运算,融合时少塞东西进 kernel,寄存器压力更小 |
| 替换在融合前 | Conv+BN 先折成单个 Conv,后面只需融 Conv+ReLU(融两次变成融一次) |
| 融合在 CSE 前 | 融合会改变图结构,先融后合能合并更多;反过来先 CSE 可能白合(融合后又被吃掉) |
| CSE 在 DCE 前 | CSE 产生孤儿,正好交给 DCE 一次清完 |
| 布局在融合后 | 布局要跟最终算子形态走;过早转换,中间节点被融掉了就白转 |
| 内存分配最后 | 它依赖前面确定的图结构 + 布局 + 生命周期,图一变全得重算 |
大白话:化简 → 改形 → 合并 → 清理 → 落布局 → 分内存,然后回头再看能不能再化简。少一次迭代,就会漏掉一整层优化机会。
六、怎么判断"前端优化做得好不好"
不要只看"算子数减少了几个",按这六个指标看:
| 指标 | 怎么看 | 陷阱 |
|---|---|---|
| 算子数 / kernel 数 | 打印优化后的图,数节点 | 算子少了但每个都变慢 → 等于没优化 |
| 中间张量字节总量 | 统计所有跨算子边界的数据量 | 这才是融合的真实收益来源 |
| 峰值显存 | 跑一遍 profile | 复用可能引入假依赖,省内存换并行度 |
| 布局转换次数 | 数 Transpose/ConvertLayout 节点 | NPU 上每次都是整张量读写 |
| 图编译时间 | pass 总耗时 | 优化过度(如超深融合搜索)会让编译变慢,拖垮迭代 |
| 数值一致性 | 与未优化版本对拍,看 max abs diff / 相对误差 | 代数简化 + 浮点重排会累积误差 |
七、五个常见误区
-
“算子融合越多越好”
错。融合增加寄存器压力、降低 occupancy,超大 elementwise kernel 常常更慢。融合要有收益模型。 -
“常量折叠能减小模型体积”
错。它减的是算子数和中间张量;权重不是被折叠出来的,只有当某个权重只被折掉的算子使用时才变成死常量被删。 -
“布局转换是免费的(只是改 stride)”
错。改 stride 的transpose是免费的,但真正搬数据(contiguous)不是。在 NPU 上一次不必要的布局转换可能吃掉整层算子的收益。 -
“CSE 一定能提速”
不一定。后端(LLVM GVN、Inductor)可能已经做了。图级 CSE 的主要收益是省中间张量和 kernel launch,而不是省那几条指令。 -
“代数简化是纯数学,随便改”
错。x/x → 1(x=0时错)、(a+b)+c → a+(b+c)(浮点不成立)、x*0 → 0(x=NaN时不成立)。默认应保守,激进重排要显式开关。
八、前端的边界在哪
前端这 9 项做完,图的"形状"就定死了:多少个算子、怎么连、什么布局、每块内存多大。剩下的是后端的事:
| 前端决定 | 后端决定 |
|---|---|
| 有几个 kernel | 每个 kernel 内部循环怎么排(tiling) |
| 每个算子吃什么布局 | 数据怎么搬进共享内存/寄存器 |
| 每块 buffer 多大 | 具体分配到哪个地址、寄存器怎么分配 |
| 哪些算子被融合 | 融合后的 kernel 怎么调度、怎么切分到多核 |
所以前端优化的天花板是"图的结构"。结构定错了,后端再强也救不回来——这也是为什么"图层/融合/布局"这三项在 AI 编译器里权重最高。
九、小结
- 前端优化 = 9 个 pass:图层、算子融合、算子替换、布局转换、内存分配(神经网络相关)+ 常量折叠、CSE、DCE、代数简化(代码/计算层面)。
- 分类判据只有一条:需不需要"这是个神经网络"的领域知识。右边 4 项换到 GCC 里一模一样能用。
- 前端在 Graph IR / Tensor IR 上工作,产出 Tensor IR 交给后端降解成 LLVM IR 或厂商 IR(如 GE IR)。前端决定"图的形状",后端决定"指令的形状"。
- 顺序是硬约束:化简 → 替换 → 融合 → CSE → DCE → 布局 → 内存,然后迭代到不动点。
- 收益要按"中间张量字节数 + kernel 数"算,不是按"少了几个节点"算。
- 每个 pass 都有明确的翻车方式:融合过度、哈希键不全、布局转换滥用、浮点重排、假依赖——它们全是正确性/性能的双刃剑。
十、本系列对应文章
| 章节 | 文章 | 状态 |
|---|---|---|
| AI 编译器基本架构 | 《AI 编译器基本架构》 | 已发布 |
| 图层 | 本篇第 4.1 节 | 概览 |
| 算子融合 | 《AI 编译器算子融合》 | 已发布 |
| 算子替换 | 本篇第 4.3 节 | 概览 |
| 布局转换 | 《布局转换原理》《AI 编译器布局转换算法》 | 已发布 |
| 内存分配 | 《AI 编译器内存分配算法》 | 已发布 |
| 常量折叠 | 《常量折叠原理》 | 已发布 |
| 公共子表达式消除 | 《公共子表达式消除原理》 | 已发布 |
| 死代码消除 | 本篇第 4.8 节 | 概览 |
| 代数简化 | 本篇第 4.9 节 | 概览 |
| 挑战与未来 | 《AI 编译器挑战与未来》 | 已发布 |
参考资料
- 原始材料:AISystem
03Compiler/03Frontend/11Summary.pdf(8 页,ZOMI,Ascend & MindSpore) - 分层架构图:该 PDF 第 3 页「Where are we?」
- 分类与高亮:该 PDF 第 5–7 页「Talk Overview of Frontend Optimizer」
- 数据流分析与不动点:Dragon Book 第 9 章
- 布局与分形:
03Compiler/03Frontend/05LayoutAlgo.pdf - 常量折叠 / CSE / 融合 / 内存分配:本系列前序各篇
转载自 CSDN-专业IT技术社区
原文链接:https://blog.csdn.net/u011436427/article/details/170196196



