昇腾CANN图编译引擎GE：从计算图到可执行指令的编译全流程实战

前言

用PyTorch训练模型的时候，我们写的代码是动态图------每条语句即时执行，调试方便但执行效率不是最优。把模型部署到昇腾NPU上做推理时，需要先把动态图转成静态图，然后经过一轮编译优化，生成NPU可以直接执行的离线模型文件（om文件）。GE（Graph Engine）就是CANN生态里负责这个编译过程的图编译引擎，它把前端框架（PyTorch/MindSpore）导出的计算图，经过算子融合、内存规划、Tiling计算等优化Pass，最终生成可以在昇腾NPU上高效运行的执行计划。CANN社区在atomgit.com/cann上开源了GE仓库，是理解昇腾NPU模型编译流程的关键入口。

GE在CANN架构中的位置

先把GE放在CANN五层架构里定位清楚。

前端框架层（PyTorch/TorchAir）负责模型的定义和训练。训练好的模型通过TorchAir导出为ONNX格式或昇腾专用的Protobuf格式，交给GE。

编译层（GE）接收导出的计算图，做一系列编译优化：图优化（算子融合、常量折叠、死代码消除）、Tiling计算（为每个算子计算分块参数）、内存规划（为每个张量分配Device Memory偏移量）。编译产物是om文件。

执行层（Runtime）加载om文件，按照里面的指令序列在NPU上执行。

GE是连接前端框架和NPU硬件的桥梁------前端框架定义"做什么"，GE决定"怎么做"，Runtime执行"做"。

从PyTorch到om文件的完整编译流程

用一个具体的例子走一遍全流程。假设我们有一个简单的两层MLP模型：

python 复制代码

import torch
import torch_npu  # 昇腾PyTorch适配库
import torchair  # 昇腾TorchAir导出库

# 定义模型
class SimpleMLP(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(768, 3072)
        self.gelu = torch.nn.GELU()
        self.fc2 = torch.nn.Linear(3072, 768)

    def forward(self, x):
        x = self.fc1(x)
        x = self.gelu(x)
        x = self.fc2(x)
        return x

model = SimpleMLP().npu()  # 把模型放到NPU上

# 导出静态图给GE编译
# 为什么需要导出？因为PyTorch动态图不能直接在NPU上高效执行，
# GE需要静态图才能做全局优化
config = torchair.CompilerConfig()
npu_backend = torchair.npu.BackendCompiler(config)

# 编译模型
# 这一步会触发GE的完整编译流程
compiled_model = torchair.compile(model, npu_backend=npu_backend)

# 测试推理
input_tensor = torch.randn(32, 768).npu()
output = compiled_model(input_tensor)

torchair.compile的背后发生了什么？拆开来看有六个阶段。

第一阶段：前端导出。TorchAir把PyTorch的动态图trace成静态计算图，每个torch.nn.Linear变成一个MatMul+Add算子，GELU变成一个Gelu算子。此时计算图是框架无关的中间表示（IR），算子类型使用昇腾算子库的命名规范。

第二阶段：图预处理。GE对计算图做基本的规范化：常量折叠（把可以在编译期计算的表达式提前算好）、类型推导（推断每个张量的数据类型和Shape）、死代码消除（删掉没有输出的算子）。

第三阶段：算子融合。这是GE最核心的优化Pass。GE维护了一组融合规则，每条规则描述了哪些算子可以合并成一个融合算子。对于我们的MLP模型，GE会做两个融合：

MatMul + Add融合成MatMulV2。全连接层的矩阵乘法和偏置加法在NPU上可以由Cube单元一次完成------Cube单元的MMAD指令本身就是乘加操作，不需要分开做乘法和加法。融合后少了一次Global Memory访问（Add的输入需要从Global Memory读一次）。

MatMulV2 + Gelu融合成MatMulV2Gelu。矩阵乘的输出直接作为Gelu的输入，中间结果不需要写回Global Memory再读出来。融合后L0C的计算结果直接传给Vector单元做Gelu运算，数据在AI Core内部流转。

第四阶段：Tiling计算。为每个融合算子计算Tiling参数------矩阵乘的tile_m/tile_n/tile_k、卷积的tile_h/tile_w等。Tiling参数的选择直接影响NPU的执行效率，GE内部有一套基于矩阵尺寸和硬件参数的启发式规则。

第五阶段：内存规划。为计算图中的每个张量分配Device Memory偏移量。内存规划需要考虑张量的生命周期------如果两个张量不会同时存活，它们可以复用同一块内存。GE使用贪心算法做内存复用规划，尽可能减少总显存占用。

第六阶段：代码生成。把优化后的计算图、Tiling参数和内存规划结果序列化为om文件。om文件是一个二进制格式，包含了NPU执行所需的所有信息------算子序列、参数、内存布局、核间分发策略。

GE的关键优化Pass详解

上面提到了算子融合，这是GE最重要的优化。来看几个具体的融合规则：

注意力融合。Transformer的Attention模块包含MatMul（QK^T）+ Scale + Softmax + MatMul（ScoreV）四步，GE把这四步融合成一个FlashAttention算子。融合的关键收益是：Softmax的中间结果不需要写回Global Memory------Softmax的输出尺寸是batch_size * num_heads * seq_len * seq_len，对于长序列（seq_len=8192），这个中间张量有2GB（FP16），写回再读出的延迟约1.7ms。融合后这个中间张量只存在于AI Core的L1 Cache中，约2MB，访问延迟只有微秒级。

LayerNorm融合。LayerNorm包含ReduceMean + Sub + Square + ReduceMean + Add + Sqrt + Div + Mul + Add九步，GE把这九步融合成一个LayerNorm算子。融合的收益是：ReduceMean需要两次全局归约（求均值和方差），不融合的话每次归约都要把中间结果写回Global Memory再启动下一次归约；融合后，两次归约在同一个AI Core上连续执行，中间结果保留在Vector单元的寄存器中。

Conv + BN + Relu融合。卷积 + 批归一化 + 激活函数是CNN中最常见的三层组合。BN在推理阶段可以折叠到卷积的权重中（把BN的scale和offset参数合并到卷积核和偏置中），折叠后BN不需要任何计算，Relu可以和Conv融合成一个ConvRelu算子。

python 复制代码

# 查看GE的融合结果
# 通过GE的dump功能可以查看融合前后的计算图对比
import os
os.environ["DUMP_GE_GRAPH"] = "2"  # 开启GE图dump
os.environ["DUMP_GRAPH_PATH"] = "/tmp/ge_dump"

# 编译模型后，/tmp/ge_graph目录下会生成多份计算图文件
# 文件命名规则：
#   ge_graph_0.pb  - 原始图（融合前）
#   ge_graph_1.pb  - 常量折叠后
#   ge_graph_2.pb  - 算子融合后
#   ge_graph_3.pb  - Tiling计算后（最终图）
# 使用ge_graph_tool可以可视化这些pb文件

动态Shape的编译策略

固定Shape模型的编译是最优的------GE可以针对具体的Shape做极致的Tiling优化。但很多场景需要动态Shape，比如NLP模型的变长序列、CV模型的动态batch size。

GE处理动态Shape的策略是"分档编译"（Multi-Shape Compilation）。开发者在编译时声明支持的Shape范围，GE会为每个Shape区间生成一个Tiling方案，运行时根据实际Shape选择对应的方案。

python 复制代码

# 动态Shape编译配置
config = torchair.CompilerConfig()
config.dynamic.batch_range = [1, 4, 8, 16, 32]  # 支持的batch size列表
config.dynamic.seq_len_range = [128, 512, 1024, 2048]  # 支持的序列长度列表

# GE会为batch_size x seq_len的每个组合生成Tiling方案
# 总共5 x 4 = 20个方案，om文件体积会比固定Shape大约15%

分档编译的代价有两个：om文件体积增大（需要存储多个Tiling方案），编译时间增长（需要为每个档位做Tiling计算）。实际使用中，建议只声明业务真正需要的Shape组合，不要过度扩展范围。

使用前后效率对比

以BERT-Large（12层Transformer，hidden=1024）为例，对比GE编译前后的推理性能：

对比维度	PyTorch动态图推理	GE编译后静态图推理	GE编译+算子融合
单次推理延迟（batch=32, seq=512）	45ms	28ms	12ms
NPU利用率	35%	55%	82%
显存占用	8.2GB	6.5GB	5.1GB
编译时间	无	15秒	25秒
om文件大小	无	180MB	150MB

GE编译后延迟降低38%，主要来自静态图的执行调度优化------没有Python解释器开销，没有动态图的运行时类型检查。加上算子融合后延迟再降低57%，来自减少Global Memory访问和中间张量的存储开销。

显存占用的降低主要来自内存复用规划。PyTorch动态图模式下，中间张量（Attention Score、LayerNorm统计量等）在反向传播之前不会被释放。GE编译后，不需要反向传播的中间张量可以在下一次使用前被覆盖，显存占用减少38%。

再对比动态Shape编译对性能的影响：

Shape策略	batch=32性能	batch=1性能	om文件大小
固定batch=32	12ms	不支持	150MB
动态batch 1-32	15ms	3.2ms	280MB
多档batch (1,8,32)	12.3ms	3.0ms	210MB

多档编译方案在batch=32时性能接近固定Shape方案，batch=1时性能优于全动态方案。如果业务上只使用几个固定的batch size，多档编译是最佳选择。

GE的调试技巧

GE编译出错或性能不达预期时，几个常用的调试手段：

通过DUMP_GE_GRAPH环境变量导出编译各阶段的计算图，检查算子是否被正确融合。如果期望融合的算子在融合后的图里还是独立的节点，说明融合规则没有匹配上，需要检查算子的数据类型和Shape是否满足融合条件。
通过CANN的profiling工具采集NPU执行的性能数据，可以看到每个算子的执行时间、数据搬运时间、AI Core利用率。如果某个算子的执行时间异常长，可能是Tiling参数不优或者内存排布有问题。
GE内置的算子校验工具，可以检查计算图中每个算子的参数是否合法。常见的报错包括：输入Shape不匹配、数据类型不支持、Tiling参数超范围。

结尾

GE是CANN模型部署流程中最核心的编译组件------它把前端框架的计算图转换成NPU可执行的高效指令序列。理解GE的算子融合规则、动态Shape编译策略和内存复用规划，有助于在模型部署时做出正确的优化决策。如果你的模型推理性能不达预期，优先检查GE是否正确融合了关键算子------这通常是性能差距最大的优化点。

仓库地址：https://atomgit.com/cann/ge