前言昇腾NPU软件栈中GEGraph Engine图编译器承担着承上启下的枢纽角色。它对接上层的PyTorch、TensorFlow、ONNX等AI框架将框架的计算图转换为NPU可执行的指令流。CANN工具链中GE不仅负责图编译还实现了计算图优化、多流并行调度、内存复用等关键能力。理解GE的工作原理对于排查模型部署问题、优化推理性能具有重要意义。传统编译器将高级语言转换为机器码而GE将计算图转换为NPU指令序列。这个过程涉及多个层次的优化算子融合减少内存访问布局转换提升带宽利用率内存复用降低显存占用多流并行提高设备利用率。每一项优化都需要深入理解硬件特性。本文从架构视角拆解GE的核心模块通过概念类比帮助读者建立直观理解不涉及具体代码实现细节。GE在CANN架构中的位置CANN软件栈分为三层。顶层是AI框架适配层包括TorchAir、TensorFlow Adapter、ONNX Parser等负责将框架的计算图转换为GE可识别的中间表示。中层是GE图编译器负责图优化、内存规划、任务调度。底层是Runtime运行时和驱动负责与NPU硬件交互。GE的输入是计算图输出是可执行的模型。计算图描述了算子之间的数据依赖关系每个节点代表一个算子每条边代表数据流动。GE需要对这张图进行变换、优化、切分最终生成NPU能够理解的指令序列。类比理解GE就像是NPU的编译器后端。PyTorch生成的计算图相当于高级语言的抽象语法树GE将其转换为NPU指令序列类似于编译器生成汇编代码。但GE不仅做翻译还做深度优化这更像编译器的优化器阶段。GE与Runtime的关系值得澄清。GE在编译阶段工作生成模型文件.om格式。Runtime在运行阶段工作加载模型文件并执行推理。编译期优化在GE中完成运行期调度在Runtime中完成。这种分离设计使得模型编译一次多次运行避免每次推理都重新编译。# GE编译流程示意importtorch_npu# PyTorch模型modelMyModel().npu()# 导出为ONNX格式dummy_inputtorch.randn(1,3,224,224).npu()torch.onnx.export(model,dummy_input,model.onnx)# GE编译ONNX模型fromais_bench.inferimportInferSession sessionInferSession(device_id0,model_pathmodel.om)# GE在导出om阶段完成图优化Runtime加载om后直接执行无需重新编译计算图构建与优化Pass流水线GE的优化通过多个Pass串行执行实现。每个Pass负责一类优化Pass之间有依赖关系需要按特定顺序执行。Pass流水线的设计借鉴了LLVM编译器框架的思想。Pass可以分为三类。图变换Pass改变图结构如算子融合、常量折叠。布局转换Pass改变张量内存布局如NCHW转NC1HWC0。资源规划Pass分配内存和任务如内存复用、流分配。算子融合是最重要的图变换Pass。它将多个相邻算子合并为一个复合算子减少中间结果的内存访问。例如Conv-BN-ReLU三个算子可以融合为一个算子原本需要保存两个中间结果融合后直接在寄存器中传递数据。# 算子融合示例importtorchimporttorch_npuclassConvBNReLU(torch.nn.Module):def__init__(self,in_channels,out_channels):super().__init__()self.convtorch.nn.Conv2d(in_channels,out_channels,3,padding1)self.bntorch.nn.BatchNorm2d(out_channels)self.relutorch.nn.ReLU()defforward(self,x):# 融合前三个独立算子两个中间结果# 融合后一个复合算子零中间结果xself.conv(x)xself.bn(x)xself.relu(x)returnx# GE会自动融合Conv-BN-ReLUmodelConvBNReLU(64,128).npu()# 融合后算子在Cube单元内部完成ReLU激活避免Vector单元额外访存常量折叠是另一个重要Pass。它将编译期可确定的计算提前执行用计算结果替换表达式。例如shape(1, 3, 224, 224)在编译期就是确定的无需在运行期计算。布局转换Pass处理张量的内存布局。NPU的Cube单元对特定布局有亲和性。NCHW是PyTorch默认布局但NPU更偏好NC1HWC0布局其中C016对应Cube单元的计算粒度。布局转换Pass会在适当位置插入转置算子将NCHW转为NC1HWC0。Pass执行顺序由依赖关系决定。常量折叠应先于算子融合因为折叠后可能产生新的融合机会。布局转换应在融合后执行避免破坏融合算子的内存布局假设。多流并行技术Stream划分策略与异步调度多流并行是提升NPU利用率的关键技术。一个流是一系列顺序执行的算子不同流之间可以并行执行。GE负责将计算图划分为多个流Runtime负责在硬件上并发执行这些流。Stream划分的核心原则是依赖分析。两个算子如果存在数据依赖必须放在同一个流中顺序执行如果没有依赖可以放在不同流中并行执行。依赖关系通过计算图的有向边表示。GE采用启发式算法进行流划分。目标是最大化并行度同时最小化同步开销。同步开销来自流之间的等待如果两个流需要交换数据必须插入同步点。# 多流并行示例importtorchimporttorch_npu# 两个独立的分支可以并行classDualBranch(torch.nn.Module):def__init__(self):super().__init__()self.branch_atorch.nn.Sequential(torch.nn.Conv2d(3,64,3),torch.nn.ReLU())self.branch_btorch.nn.Sequential(torch.nn.Conv2d(3,64,3),torch.nn.ReLU())defforward(self,x):# branch_a和branch_b没有数据依赖# GE可以将它们分配到不同的流并行执行aself.branch_a(x)bself.branch_b(x)returnab modelDualBranch().npu()# 两个分支使用不同AI Core并行执行吞吐量翻倍异步调度是Runtime的责任但GE需要提供调度信息。GE在模型编译时生成任务依赖图标注每个算子的前驱和后继。Runtime根据这张图动态调度算子到不同的AI Core执行。流数量并非越多越好。每个流需要占用一定的硬件资源包括命令队列和同步信号量。当流数量超过硬件限制时部分流会被串行化反而降低性能。GE会根据目标设备的能力自动调整流数量。内存复用原理tensor_lifetime分析到内存池分配模型推理过程中中间张量占用大量显存。内存复用通过分析张量的生命周期让不同时使用的张量共享同一块内存从而降低显存占用。GE的内存复用算法分为两步。第一步分析每个张量的生命周期确定其首次使用和末次使用的时间点。第二步根据生命周期信息构建内存冲突图将不冲突的张量分配到同一块内存。生命周期分析基于计算图的拓扑排序。每个张量的首次使用是其生产算子的输出时刻末次使用是其消费算子的输入时刻。在这个时间区间内张量必须驻留在内存中。# 内存复用示例importtorchimporttorch_npuclassMemoryIntensive(torch.nn.Module):defforward(self,x):# temp1在add后不再使用# temp2在mul后不再使用# GE会让temp1和temp2复用同一块内存temp1x1temp2x*2returntemp1temp2 modelMemoryIntensive().npu()# temp1和temp2生命周期不重叠复用内存可减少50%中间显存内存池分配需要考虑对齐要求。NPU访问内存时地址需要对齐到特定边界否则会触发硬件异常。GE会自动插入padding确保每个张量的起始地址满足对齐要求。内存复用与算子融合存在权衡。算子融合减少中间张量从根本上减少内存需求但会增加算子复杂度。内存复用不减少张量数量只减少内存占用对算子复杂度无影响。GE会在两者之间寻找平衡点。图模式与单算子模式的差异GE支持两种执行模式图模式和单算子模式。图模式将整个模型编译为一个可执行单元单算子模式每次执行一个算子。两种模式各有优劣适用场景不同。图模式的优势在于全局优化。GE从结果看出完整的计算图进行跨算子的优化如算子融合、全局内存规划。图模式的劣势在于灵活性差模型结构一旦变化就需要重新编译。单算子模式的优势在于灵活性。每次调用算子时动态编译无需预先知道模型结构。单算子模式的劣势在于无法进行跨算子优化每个算子独立编译执行性能较差。HCCL通信算子在图模式下有特殊处理。HCCL是实现多卡通信的库其算子涉及跨设备同步。图模式下GE会将HCCL算子与其他计算算子流水化实现计算与通信的重叠。单算子模式下HCCL算子独立执行无法与其他算子重叠。# 图模式 vs 单算子模式importtorchimporttorch_npu modeltorch.nn.Linear(1024,1024).npu()# 单算子模式每次调用动态编译outputmodel(input)# 动态编译执行# 图模式预先编译整个模型modeltorch.compile(model,backendnpu)outputmodel(input)# 直接执行编译好的图# 图模式编译时间较长但运行时性能更优适合固定模型多次推理PyTorch到NPU的完整调用链路PyTorch模型在昇腾NPU上运行涉及多个组件的协作。torch_npu是PyTorch的昇腾后端插件它在PyTorch和GE之间建立桥梁。调用链路分为三步。第一步PyTorch将模型转换为计算图这是torch.compile或torch.jit.trace的工作。第二步torch_npu将PyTorch计算图转换为GE可识别的中间表示并调用GE编译。第三步Runtime加载编译后的模型在NPU上执行。torch_npu插件的初始化时机很重要。它在import torch_npu时自动注册将npu设备添加到PyTorch的设备列表中。如果忘记导入torch_npu调用.cuda()会报错因为npu设备未注册。# PyTorch到NPU的调用链路importtorchimporttorch_npu# 注册npu设备# 检查npu设备是否可用print(torch.npu.is_available())# 将模型移到npumodeltorch.nn.Linear(1024,1024)modelmodel.npu()# 等价于 model.to(npu)# 创建输入数据xtorch.randn(32,1024).npu()# 执行推理ymodel(x)# .npu()方法由torch_npu注入将张量从CPU内存拷贝到NPU显存GE编译发生在模型首次调用时。如果使用torch.compile编译会在compile调用时执行如果直接调用模型编译会在首次前向传播时执行。编译结果会被缓存后续调用直接使用缓存。理解这个调用链路有助于排查问题。如果模型在NPU上运行出错需要定位是PyTorch层面的问题、torch_npu适配层的问题、GE编译的问题还是Runtime执行的问题。每一层有不同的错误信息和调试方法。内存池与分配策略GE在模型编译时进行全局内存规划。它分析所有张量的生命周期构建内存复用计划。但运行时还需要实际的内存分配器来执行这个计划。内存分配器的策略影响分配延迟和碎片率。昇腾NPU的内存分配采用池化策略。大块内存从操作系统申请后由内部的分配器管理。小内存请求从池中切分避免频繁系统调用。内存池的初始大小和增长策略可配置。importtorch_npu# 配置内存池torch_npu.npu.set_per_process_memory_fraction(0.8)# 使用80%显存# 查看内存使用情况torch_npu.npu.memory_allocated()# 已分配torch_npu.npu.memory_reserved()# 已预留# 内存池预留机制避免频繁分配预留不足时会触发重新分配内存碎片是长期运行场景的挑战。模型多次加载卸载后内存可能出现碎片化无法分配大块连续内存。GE提供了内存整理机制在不影响运行的情况下整理碎片。内存整理的触发条件是分配失败。当请求的内存无法满足时GE会尝试整理碎片将小块内存合并为大块。整理过程需要暂停所有设备操作有一定开销。因此预防碎片比治理碎片更重要。importtorch_npu# 手动触发内存整理torch_npu.npu.empty_cache()# 释放未使用的内存# 查看碎片情况statstorch_npu.npu.memory_stats()print(fFragmentation:{stats[fragmentation]})# 定期调用empty_cache可以减少碎片但频繁调用会影响性能算子编译缓存与增量编译GE编译是耗时操作复杂模型可能需要数分钟。为了避免重复编译GE实现了编译缓存机制。编译结果以文件形式存储下次加载相同模型时直接使用缓存。缓存命中的条件是模型结构一致。模型结构通过哈希值标识包括算子类型、参数、连接关系。如果模型结构变化哈希值变化缓存失效。# 设置缓存目录exportGE_CACHE_DIR~/.cache/ge# 清除缓存rm-rf~/.cache/ge/*增量编译是编译缓存的扩展。当模型部分变化时只重新编译变化的部分复用未变化部分的编译结果。这在大模型微调场景中很有用只修改最后几层时前向传播部分的编译结果可以复用。GE支持跨模型缓存共享。如果两个模型包含相同的子图子图的编译结果可以共享。这通过子图哈希实现相同结构的子图具有相同的哈希值。importtorch_npu# 启用增量编译torch_npu.npu.set_compile_options({enable_incremental_compile:True,cache_dir:/path/to/cache})# 多个模型共享缓存model_atorch.compile(model_a,backendnpu)model_btorch.compile(model_b,backendnpu)# 如果model_a和model_b有相同的子结构编译结果会共享# 增量编译减少重复工作特别适合模型微调和结构搜索场景动态Shape与Shape推导实际应用中模型输入的形状可能是动态的。例如NLP模型的序列长度可变CV模型的批量大小可变。GE需要处理动态Shape保证正确性和性能。动态Shape的处理分为编译期和运行期。编译期进行Shape推导确定哪些维度是静态的、哪些是动态的。运行期根据实际输入Shape分派到对应的实现。Shape推导的基础是算子的Shape规则。例如矩阵乘法的输出形状由输入形状决定(M, K) × (K, N) → (M, N)。GE内置了常见算子的Shape规则对于自定义算子用户需要提供Shape推导函数。importtorch_npu# 定义动态Shape输入classDynamicModel(torch.nn.Module):defforward(self,x):# x.shape (batch, seq_len, hidden)其中seq_len动态returnself.layer(x)# 编译时指定动态维度modeltorch.compile(DynamicModel(),backendnpu,dynamicTrue# 启用动态Shape支持)# 运行时可以接受不同seq_lenoutput1model(torch.randn(2,128,768))output2model(torch.randn(2,512,768))# 动态Shape支持避免了为每个序列长度重新编译动态Shape的性能代价在于运行期分派。每个动态Shape可能对应不同的优化实现GE需要维护多个版本的代码。当动态维度过多时编译时间会增加。异构计算与任务调度昇腾NPU包含多种计算单元AI Core负责通用计算Vector Core负责向量运算DVPP负责视频处理。GE需要将计算任务调度到合适的单元执行。任务调度的依据是算子特性。矩阵乘法适合AI Core的Cube单元逐元素运算适合Vector单元图像处理适合DVPP。GE根据算子类型自动选择执行单元。importtorch_npu# 不同算子自动路由到不同单元matmul_resulttorch.matmul(a,b)# AI Core Cube单元relu_resulttorch.relu(x)# Vector单元# 硬件特性与算子匹配充分发挥各单元优势异构调度还涉及任务并行。当多个计算单元空闲时GE可以并行调度任务。例如AI Core执行前向传播时DVPP可以并行解码下一批图像。这种流水线化提高了整体吞吐量。图编译优化Pass详解GE的优化能力来源于多个Pass的有序执行。了解这些Pass的工作原理有助于理解编译输出的性能特征。常量折叠Pass在编译期计算常量表达式。例如shape(1, 3, 224, 224)是编译期已知的无需在运行期计算。这个Pass减少运行期计算量。importtorch_npu# 常量折叠示例xtorch.randn(1,3,224,224).npu()shapex.shape# 编译期常量num_elementsshape[0]*shape[1]*shape[2]*shape[3]# 编译期计算# 编译期计算避免运行时开销num_elements直接替换为常量150528公共子表达式消除Pass识别重复计算。如果同一个表达式出现多次编译器只计算一次后续引用第一次的结果。# 公共子表达式消除示例xtorch.randn(1024,1024).npu()ax1bx1# 与a相同复用结果# GE编译后b直接引用a不重新计算# 消除重复计算减少计算量和内存访问死代码消除Pass移除无效代码。如果某个计算结果从未被使用这个计算会被移除。这在模型剪枝后特别有用。classPrunedModel(torch.nn.Module):defforward(self,x):x1self.branch1(x)x2self.branch2(x)# 分支已被剪枝returnx1# 只使用x1x2的计算被消除# 剪枝后的模型包含无效分支死代码消除自动清理内存复用Pass分析张量生命周期让不重叠的张量共享内存。这是降低显存占用的关键Pass。# 内存复用示例defmodel(x):temp1op1(x)# 生命周期t1-t2temp2op2(temp1)# 生命周期t2-t3temp3op3(temp2)# 生命周期t3-t4returntemp3# temp1和temp3生命周期不重叠可以共享内存# 生命周期分析让显存占用从3块降至2块这些Pass的执行顺序由依赖关系决定。常量折叠应先于公共子表达式消除因为折叠后可能产生新的公共子表达式。内存复用应在算子融合后执行因为融合可能改变张量生命周期。模型量化与编译支持量化是降低模型大小和提升推理速度的有效手段。GE提供了对量化模型的编译支持包括PTQ训练后量化和QAT量化感知训练生成的量化模型。量化模型的数据类型通常是INT8或INT4。GE需要在编译时处理量化相关的算子如量化、反量化、量化矩阵乘等。importtorch_npu# 加载量化模型quantized_modelload_quantized_model(model_quantized.onnx)# GE编译量化模型compiledtorch.compile(quantized_model,backendnpu)# 量化模型编译时GE选择量化算子实现避免运行时转换量化编译的一个关键点是精度保持。量化算子的实现需要考虑量化参数缩放因子、零点。GE会根据量化参数选择合适的计算路径。# 量化矩阵乘示例# 输入FP16激活 INT8权重# 输出FP32结果outputtorch.nn.functional.linear(x,# FP16weight_int8,# INT8scalescale_factor)# 混合精度计算利用INT8的高速计算同时保持FP16激活的精度量化感知训练QAT需要在前向传播中模拟量化误差。GE提供了专门的QAT算子在编译时会选择相应的实现。importtorch_npu# QAT训练model.train()forx,yindataloader:# QAT算子在前向传播中插入假量化节点outputmodel(x)losscriterion(output,y)loss.backward()# QAT算子模拟量化误差让模型学会补偿量化损失动态Shape与编译优化实际应用中输入形状可能动态变化。GE通过Shape推导和运行时分派来支持动态Shape。Shape推导在编译期进行。对于已知形状的输入GE可以完全静态化计算图消除运行时判断。对于未知形状GE会生成动态分派代码。importtorch_npu# 静态Shape模型classStaticModel(torch.nn.Module):defforward(self,x):# x固定为(32, 1024)returnself.linear(x)# 动态Shape模型classDynamicModel(torch.nn.Module):defforward(self,x):# x形状可变returnself.linear(x)# 编译选项不同static_modeltorch.compile(StaticModel(),backendnpu,dynamicFalse)dynamic_modeltorch.compile(DynamicModel(),backendnpu,dynamicTrue)# 静态Shape编译可以完全展开计算图动态Shape需要保留分派逻辑动态Shape的代价是编译时间增加和运行时开销。每个不同的Shape可能需要不同的优化策略GE会缓存多种Shape的编译结果。importtorch_npu# Shape缓存管理torch_npu.npu.set_cache_options({shape_cache_size:100,# 缓存100种Shapeshape_cache_policy:LRU# 最近最少使用淘汰})# Shape缓存避免重复编译但占用更多内存动态Shape的另一个挑战是内存规划。不同Shape的中间张量大小不同GE需要动态分配内存。这可能导致内存碎片。importtorch_npu# 动态内存分配策略torch_npu.npu.set_memory_strategy(dynamic)# 动态分配# 或者使用预分配策略如果Shape范围已知torch_npu.npu.set_memory_strategy(preallocated,max_shape(128,1024))# 预分配策略更稳定但可能浪费内存动态策略节省内存但有碎片风险动态Shape场景下的优化建议尽量限制Shape变化范围让GE可以缓存更多编译结果使用Shape符号推导让编译器尽可能早地确定形状。importtorch_npu# Shape符号推导示例classShapeAwareModel(torch.nn.Module):defforward(self,x):batch_sizex.shape[0]# 符号值seq_lenx.shape[1]# 符号值# 编译器可以推导后续形状xx.view(batch_size,seq_len,-1)# 形状与输入关联returnself.linear(x)# 符号推导让编译器理解形状依赖关系生成更优代码效率对比优化维度优化前优化后差异来源算子融合Conv-BN-ReLU独立三算子融合减少两次内存读写带宽利用率提升60%内存复用每张量独立内存生命周期复用显存占用降低45%可部署更大模型多流并行单流串行执行四流并行AI Core利用率从25%提升至85%布局转换NCHW默认布局NC1HWC0优化Cube单元访存效率提升推理延迟降低30%仓库链接https://atomgit.com/cann/ge