前言GEGraph Engine是 CANN 的核心组件负责把计算图转成 NPU 可以执行的计划。理解 GE 的优化逻辑才能写出高性能的模型。一、计算图的表示PyTorch 或 ONNX 的模型在 GE 里被表示成ComputeGraph结构// GE 内部的计算图结构简化classComputeGraph{std::vectorNode*nodes_;// 算子节点std::vectorEdge*edges_;// 数据边std::mapstring,Tensor*tensors_;// 中间 tensor};classNode{string op_type_;// 算子类型Conv, MatMul, ReLU...std::vectorTensor*inputs_;// 输入 tensorstd::vectorTensor*outputs_;// 输出 tensorOpDesc*op_desc_;// 算子属性};导出计算图用torch.jit.trace把 PyTorch 模型转成计算图importtorchimporttorch_npu modelMyModel().npu().eval)tracedtorch.jit.trace(model,torch.randn(1,3,224,224).npu())# 导出 GE 计算图torch.npu.save_ge_graph(traced,model_graph.txt)用 Netron 打开model_graph.txt可以看到完整的计算图结构。二、常量折叠Constant Folding常量折叠是最基础的优化把编译时能算出来的结果提前算好。优化前importtorchclassModel(torch.nn.Module):def__init__(self):super().__init__()self.scaletorch.tensor(2.0)# 常量self.biastorch.tensor(1.0)# 常量defforward(self,x):returnx*self.scaleself.bias# 乘法和加法在运行时执行计算图input ──→ [Mul] ──→ [Add] ──→ output ↑ ↑ scale biasGE 优化后如果scale和bias在编译时已知GE 会把Mul Add合并成一个算子// 优化后的伪代码outputx*2.01.0;// 编译时算好系数计算图input ──→ [Scale] ──→ outputScale算子的参数在编译时已经确定运行时只需要一次乘加操作。触发条件常量折叠的条件所有输入都是常量不是动态 tensor算子没有副作用不会修改全局状态# 不会触发常量折叠输入是动态的defforward(self,x):scaletorch.tensor(2.0)returnx*scale# x 是动态输入不能提前算# 会触发常量折叠defforward(self,x):atorch.tensor(2.0)btorch.tensor(3.0)cab# 编译时算出 c 5.0returnx*c三、公共子表达式消除CSE如果计算图里有重复的计算GE 会消除重复的子图。优化前defforward(self,x):ax1bx1# 重复计算returnab计算图x ──→ [Add] ──→ a ──┐ │ ↑ ├─→ [Add] ──→ output └─→ [Add] ──→ b ─┘ ↑ 1两个Add算子完全相同浪费计算资源。GE 优化后defforward(self,x):ax1returnaa# 复用 a计算图x ──→ [Add] ──→ a ──┬─→ [Add] ──→ output ↑ │ 1 └─→ (复用)只执行一次Add结果被复用。CSE 的限制# 不会触发 CSE算子有随机性defforward(self,x):atorch.randn_like(x)# 每次结果不同btorch.randn_like(x)returnab# 不能复用随机算子、Dropout 等有状态的算子不会触发 CSE。四、算子融合Operator Fusion算子融合是 GE 最重要的优化把多个小算子合并成一个大算子减少显存读写。Conv BN ReLU 融合# 原始模型classModel(torch.nn.Module):def__init__(self):self.convtorch.nn.Conv2d(3,64,3)self.bntorch.nn.BatchNorm2d(64)self.relutorch.nn.ReLU()defforward(self,x):xself.conv(x)# 写显存xself.bn(x)# 读显存、写显存xself.relu(x)# 读显存、写显存returnx三次显存读写中间结果conv_out和bn_out都要存到显存。融合后GE 会把这三个算子合成一个ConvBNReLU算子// 融合算子的伪代码voidConvBNReLU(Tensor input,Tensor weight,Tensor bn_weight,Tensor bn_bias,Tensor output){// 整个计算在 UB 里完成不写回显存for(inti0;ioutput_size;i){floatconv_outconv_compute(input,weight,i);floatbn_out(conv_out-mean)/std*bn_weightbn_bias;output[i]relu(bn_out);}}只需要一次显存写入最终输出中间结果在 UB 里流转。融合条件GE 支持的融合模板融合模式条件性能提升Conv BN ReLUBN 在 eval 模式40%MatMul Add ReLUAdd 是 bias15%FlashAttentionQ/K/V 来自同一输入3xAttention 部分LayerNorm Dropout ResidualDropout 比例固定20%五、内存规划Memory PlanningGE 会分析每个 tensor 的生命周期让不重叠的 tensor 复用同一块显存。生命周期分析defforward(self,x):aop1(x)# a 的生命周期创建到 op3 使用bop2(a)# b 的生命周期创建到 op4 使用cop3(a)# a 在这里最后使用dop4(b)# b 在这里最后使用returnd时间线a: |----------| (op1 → op3) b: |----------| (op2 → op4)a和b的生命周期有重叠不能复用同一块内存。复用示例defforward(self,x):aop1(x)bop2(a)# a 在这里最后使用可以释放cop3(b)# c 可以复用 a 的内存returnc时间线a: |-----| b: |----------| c: |-----| (复用 a 的内存)查看内存规划结果exportGE_MEMORY_PLANNING_LOG1atc--modelmodel.onnx--outputmodel日志会显示每个 tensor 的内存偏移和大小以及内存复用情况。六、算子调度Kernel Selection同一个算子有多种实现GE 会选择最优的。Cube vs Vector矩阵乘有两种实现// Cube 实现适合大矩阵voidMatMul_Cube(Tensor A,Tensor B,Tensor C){// 用 Cube Unit 硬件加速cube_gemm(A,B,C);}// Vector 实现适合小矩阵或特殊 shapevoidMatMul_Vector(Tensor A,Tensor B,Tensor C){// 用 Vector Unit 软件实现for(inti0;iM;i){for(intj0;jN;j){C[i][j]0;for(intk0;kK;k){C[i][j]A[i][k]*B[k][j];}}}}GE 会根据矩阵大小选择矩阵大小选择原因M, N, K 16CubeCube 效率高M 16 或 N 16VectorCube 对小矩阵效率低K 不是 16 的倍数VectorCube 要求数据对齐强制选择实现# 强制使用 Cube 实现torch.npu.set_op_impl(matmul,implcube)outputtorch.matmul(a,b)# 强制使用 Vector 实现torch.npu.set_op_impl(matmul,implvector)outputtorch.matmul(a,b)七、查看优化后的计算图导出优化前后的图importtorchimporttorch_npu modelMyModel().npu().eval)# 导出优化前的图torch.npu.save_ge_graph(model,before_opt.txt)# 执行一次推理触发优化model(torch.randn(1,3,224,224).npu())# 导出优化后的图torch.npu.save_ge_graph(model,after_opt.txt)用 Netron 对比两个文件可以看到优化前后的差异节点数量减少融合、消除边的数量减少复用、消除内存占用降低参考资源GE 优化原理https://www.hiascend.com/document/detail/zh/CANN/算子融合规则https://www.hiascend.com/document/detail/zh/CANN/Netron 可视化工具https://netron.app/ATC 编译参数说明https://www.hiascend.com/document/detail/zh/CANN/总结GE 的图优化分四个阶段常量折叠消除编译时能算的结果、CSE 消除重复计算、算子融合减少显存读写、内存规划让 tensor 复用空间。理解这些优化之后写模型时可以有意配合把常量提取出来、避免重复计算、用标准算子组合触发融合。用 Netron 查看优化前后的计算图能直观感受 GE 做了什么。调优的时候节点数量和显存占用是两个关键指标——节点越少、显存越小模型性能越好。