前言GE 是 CANN 异构计算架构中负责图编译的核心组件。上层框架无论是 PyTorch 还是 TensorFlow在调用昇腾NPU执行训练或推理任务时都需要把计算图转换成昇腾硬件能够直接执行的任务图。这个转换过程发生在 GE 内部它接收框架侧的计算图描述构建统一的中间表示再经过多轮优化、自动微分、算子融合以及整图下沉执行调度最终输出可在 Device 侧运行的指令序列。理解 GE 的内部机制是定位昇腾 NPU 上训练性能瓶颈的关键上游环节。GE 在 CANN 五层架构中位于编译层。它向上通过 Framework Adaptor 对接 PyTorch 和 TensorFlow向下把编译好的任务图交给 Runtime 执行。与单纯的算子库不同GE 不实现具体的矩阵乘或卷积计算逻辑而是决定这些算子以何种顺序、何种组合、何种内存布局被调度到 NPU 上。因此GE 的优化方向直接影响 Cube 单元和 Vector 单元的利用率也影响 HBM 带宽的占用情况。与传统的算子逐层调度相比GE 的工作集中在编译期。它需要在图级别做出决策哪些算子可以合并、哪些中间张量可以消除、反向图如何切分、内存如何复用、Host 与 Device 之间的交互频次如何压缩。这些决策一旦确定在执行阶段很难再做调整。GE 的编译结果直接决定了训练任务能否充分利用昇腾 NPU 的并行计算能力以及能否缓解 HBM 带宽压力。图编译前端从框架计算图到GE IRGE 的输入通常是框架导出的计算图。在 PyTorch 适配路径中TorchAir 会把 PyTorch 的 FX 图或 ATen 算子序列转换为 GE 能够识别的图描述在 TensorFlow 适配路径中GE 则直接解析 TensorFlow 的 GraphDef 或 FunctionDef。无论来源如何GE 都会把这张图重新构建成自己的中间表示称为 GE IR。IR 的节点表示算子边表示张量依赖每个节点携带算子类型、输入输出形状、数据类型、属性字典以及布局信息。IR 构建阶段的一项重要工作是把数据流图和算子图统一起来。数据流图强调张量之间的依赖关系适合描述神经网络的前向传播算子图则强调计算操作的执行顺序更适合后续的编译优化。GE 在转换时会建立双向映射数据流图的边对应算子图的输入输出算子图的节点对应数据流图的计算节点。这种映射让 GE 既能做高层的数据流分析也能做低层的算子调度分析。IR 节点除了算子信息外还保存了形状推导所需的全部元数据。当图中存在广播、切片、reshape 等会改变张量形状的算子时GE 会在编译期完成形状推导并把推导结果写回 IR。形状推导的准确性直接影响后续内存分配和算子融合。如果某个节点的输出形状无法确定GE 会把它标记为动态形状节点并为其生成运行时的形状计算代码。布局转换是前端阶段另一个需要处理的问题。PyTorch 默认使用 NCHW 布局而某些昇腾算子实现可能在 NHWC 或 NC1HWC0 布局上效率更高。GE 会在 IR 中插入布局转换节点把数据从框架布局转换到硬件偏好的布局。布局转换本身是一次内存操作过多插入会增加带宽压力因此 GE 会尽量把布局转换合并到相邻算子中避免独立的 transpose 节点。GE IR 的设计目标是在保持框架语义的同时为硬件优化提供足够的信息。每个算子节点除了记录计算语义还记录了精度模式、内存格式、执行优先级等属性。这些属性在后续的优化阶段会被逐步填充和改写。例如当一个 Conv 算子被决定采用 NHWC 布局时IR 中对应的布局属性会从 NCHW 更新为 NHWC相邻的格式转换节点则会被删除或合并。这种基于属性的等价变换让 GE 能够在不破坏框架语义的前提下对图进行深度重写。常量折叠是 GE 在前端阶段执行的基础优化之一。当图中某个节点的输入全部是常量时GE 会在编译期直接算出该节点的输出并用常量节点替换原计算节点。折叠的好处是减少运行时的计算量也减少需要传输到 Device 的张量。但如果常量节点数量过多折叠后的图体积可能膨胀反而增加模型加载时间。GE 通常会在常量折叠后做一次图规模检查避免过度展开。公共子表达式消除CSE是另一项标准编译器优化。GE 会扫描图中是否存在结构相同的子图如果两个子图具有相同的算子类型、输入和属性并且它们的输出没有被副作用依赖GE 会把它们合并为一个节点让下游节点共享结果。CSE 在 BERT 类模型中效果明显因为自注意力层内部存在大量重复的矩阵变换和转置操作。消除这些重复计算后Device 端的 HBM 读写次数也会减少。前端阶段还需要处理动态形状和标量依赖。当输入形状在编译期无法完全确定时GE 会为相关节点保留动态维度符号并在后续阶段生成形状推导代码。动态形状的处理会增加编译复杂度因为某些融合规则只适用于静态形状。GE 的做法是把图拆成静态子图和动态子图分别应用不同的优化策略。IR 构建完成后GE 会执行一次合法性校验。校验包括算子类型是否被支持、输入输出形状是否匹配、数据类型是否兼容、属性取值是否在合法范围内。校验失败时GE 会向上层框架返回错误信息而不是把问题推迟到执行阶段。这种早期失败的策略可以节省大量调试时间因为 Device 侧的报错信息通常比 Host 侧更难解读。# 构建 GE 计算图 IRgraphge.Graph()xgraph.placeholder(shape[64,3,224,224],dtypefloat16)wgraph.const(valueweight_tensor,nameconv1_w)convgraph.op(Conv2D,inputs[x,w],attrs{stride:[1,1]})bngraph.op(BatchNorm,inputs[conv])# GE builds IR in host memory before launching any NPU kernel; all shape and dtype metadata must be kept for later fusion and memory planning passes.算子融合引擎规则、收益与边界GE 的算子融合引擎负责把多个独立的算子合并成一个融合算子。融合后的算子由单个 NPU 任务完成中间结果可以保留在片上高速缓存中不必写回 HBM。这是降低内存带宽压力、提升计算密度的主要手段。融合规则的定义采用模式匹配方式。每条规则包含一个锚点算子、一组跟随算子、一组约束条件以及一个收益估算函数。GE 在遍历图时从锚点算子出发检查其下游节点是否匹配跟随算子再检查约束条件是否满足。约束条件包括数据布局、数据类型、形状兼容性、算子属性是否允许融合等。例如Conv2D 与 BatchNorm 的融合要求卷积输出与 BN 的输入维度一致并且 BN 的均值、方差、缩放、偏移在编译期已知。收益估算模型决定了一条融合规则是否值得应用。GE 会综合考虑融合后的计算量减少、中间张量消除带来的内存收益、新增算子实现复杂度、以及融合后算子是否已经存在于算子库。如果收益估算为负GE 会放弃融合保持原图结构。收益估算不能只看单点开销还要看融合对后续调度可能造成的影响。某些融合虽然减少了 kernel 数量但可能引入新的同步点或导致并行度下降。ConvBN 是深度学习图中最常见的融合模式。在分离执行时Conv2D 先输出一个特征图写入 HBMBatchNorm 再从 HBM 读取该特征图并做归一化。融合后Conv2D 的输出直接作为 BatchNorm 的输入在 Vector 单元上完成缩放和偏移中间张量不会离开片上缓存。这个融合模式在 ResNet 和 VGG 类骨干网络中几乎都会被触发。MatMulAdd 融合则更多出现在全连接层和注意力投影中。分离执行时MatMul 的结果需要写回 HBMAdd 再读取并做偏置相加。融合后Add 的偏置加载可以与 MatMul 的累加过程重叠减少一次完整的 HBM 读写。注意力机制中的 Q/K/V 投影和输出投影也常以类似方式融合。融合边界与调度约束之间经常出现冲突。例如一条融合规则要求两个算子必须位于同一个流上但原图中它们分别属于不同的并行分支。如果强行合并可能会破坏原有的并行执行结构导致整体延迟增加。GE 的处理方式是在模式匹配阶段就把调度约束纳入收益估算拒绝那些会破坏关键路径并行的融合。另一个常见冲突是融合后算子形状超出当前算子实现的限制。GE 会把图拆成可融合子图和不可融合子图分别走不同的代码生成路径。布局约束也是融合边界的重要组成部分。两个算子即使计算逻辑上可融合如果它们的输入输出布局不一致融合后可能需要额外插入布局转换节点。当转换开销超过融合收益时GE 会保留原图。这种取舍说明融合优化不是单纯越多越好而是要与整体调度目标一致。# 定义 ConvBN 融合规则patternge.FusionPattern(anchorConv2D,followers[BatchNorm],constraintlambdanode:node[data_format]NCHWandnode[training]False)# ConvBN fusion only applies when the BN statistics are frozen so the normalization can be folded into the conv weights and bias without runtime recomputation.自动微分与反向图切分GE 支持在图级别自动生成反向传播图。给定前向图GE 会为每个需要梯度的前向节点构造对应的反向节点并建立前向张量到反向张量的依赖链。自动微分的输入包括前向图的拓扑结构也涵盖哪些张量需要求导、哪些参数是可训练参数、以及是否需要保留前向中间结果。反向图生成面临的主要问题是内存开销。训练深度网络时前向传播的中间结果通常需要保留到反向传播使用。如果每个前向节点都保留全部输出显存占用会迅速增长。GE 采用的策略之一是 checkpointing只在部分关键节点保留中间结果其他节点在反向传播时重新计算。checkpointing 以计算换内存适用于显存受限但计算冗余度较低的场景。checkpointing 节点的选择需要权衡计算与内存。保留过多中间结果会浪费 HBM保留过少则会在反向阶段引入大量重算。GE 提供了基于层级的 checkpointing 策略默认在模型层边界保留激活。用户也可以手动指定保留节点以适配特定模型的内存预算。对于 Transformer 类大模型checkpointing 通常是控制显存占用的必要手段。反向图中的内存分配策略与正向图不同。正向图可以按拓扑顺序一次性分配内存反向图则需要考虑梯度张量的生命周期。GE 会在反向图构建完成后做一次内存复用分析把生命周期不重叠的梯度张量映射到同一块 HBM 地址。这种复用能够降低训练时的峰值显存占用。在多步骤梯度累积场景中反向图需要支持重切分。梯度累积要求每个 micro-batch 的梯度独立累加这意味着 GE 需要为每个 micro-batch 生成独立的反向子图或者在同一张反向图中插入梯度累加节点。GE 的做法通常是后者在反向图末端插入梯度累加算子让不同 micro-batch 的梯度在同一个张量上累加。这样可以在不多次启动完整反向图的情况下完成梯度累积减少 Host 对 Device 的调度次数。反向图切分还涉及数据并行和模型并行的边界。在分布式训练中某些梯度需要在不同 NPU 之间做 AllReduce而另一些梯度只在本地使用。GE 会在反向图中插入通信节点并把它们与计算节点一起调度。通信节点的插入位置会直接影响梯度同步的延迟GE 会尽量把通信节点下沉到 Device 侧减少 Host 介入。# 在反向图中插入 checkpoint 节点forwardge.build_graph(model,inputs)checkpoint_nodes[layer_1_relu,layer_3_relu]backwardge.grad(forward,wrtparams,checkpointscheckpoint_nodes)# Checkpoints trade extra forward recomputation for reduced HBM residency because NPU on-chip L0 buffer is too small to hold all intermediate activations of large models.整图下沉执行GE 的优化最终要落到执行层面。在昇腾 NPU 上执行方式有两种极端逐算子调度和整图下沉。逐算子调度意味着每个算子都由 Host 单独发射到 Device每次发射都涉及一次 Host-Device 交互。这种方式灵活适合调试和动态图但交互开销会累积。整图下沉则把整张图一次性编译成 Device 侧的执行计划Host 只需要在图边界做一次启动图内部的算子调度由 Device 上的运行时完成。Graph Sink 是 GE 实现整图下沉的核心机制。在 Graph Sink 模式下GE 会把优化后的图转换成一张任务图任务图中的节点对应 NPU 上的具体任务边对应任务之间的数据依赖和同步关系。任务图生成后GE 会把它下发到 Device 的 Runtime由 Runtime 负责任务的流分配和事件同步。Host 在每次迭代时只需要把输入数据放到指定地址再触发一次图执行。流分配决定了图中哪些任务可以并行执行。昇腾 NPU 支持多个硬件流每个流上的任务按顺序执行不同流上的任务可以并发。GE 在编译期会分析任务之间的依赖关系把无依赖的任务分配到不同流把有依赖的任务分配到同一流或插入事件同步。流分配策略会影响硬件利用率和执行延迟分配过多流会增加同步开销分配过少流会限制并行度。事件同步在图执行中承担依赖管理职责。当任务 A 的输出是任务 B 的输入时GE 会在两个任务之间插入事件。任务 A 完成后触发事件任务 B 在事件满足后开始执行。事件同步由 Device 硬件直接支持不需要 Host 参与。这使得 Graph Sink 模式下的 Host-Device 交互只发生在图级别而不是每个算子级别。Graph Sink 并非在所有场景下都适用。如果图中包含大量控制流或动态形状Device 侧可能无法预先确定执行路径这时 GE 会退回到逐算子或子图下沉模式。子图下沉是 Graph Sink 和逐算子调度的折中GE 把图拆成多个子图每个子图内部做整图下沉子图之间仍由 Host 调度。这种混合模式在动态图和静态图混合的模型中比较常见。调试整图下沉模式比调试逐算子模式更复杂。因为 Device 侧自主管理任务调度Host 侧很难单步跟踪每个算子的执行。GE 提供了图级 Profiling 接口可以导出任务图执行时间线帮助开发者定位哪些任务之间存在不必要的同步或空闲间隙。# 配置整图下沉执行sessionge.Session()session.set_option(graph_sink,True)session.set_option(stream_num,4)session.compile(graph)session.run(graph)# Graph sink pushes the entire execution plan to the NPU so the host only emits one launch call per graph, cutting host-device round trips below the kernel-launch latency threshold.效率对比GE 的优化措施在不同维度上带来的效果并不一致。有些维度提升明显有些维度则保持原有水平甚至略有增加。下面的对比表从硬件利用率、内存占用、编译时间、Host-Device 交互四个维度描述使用 GE 的图编译优化前后的差异。维度使用前使用后差异来源硬件利用率逐算子调度导致 kernel 之间空闲间隔较多Cube 和 Vector 单元难以持续满载算子融合与整图下沉让任务在 Device 侧连续调度计算单元空闲时间被压缩融合减少中间读写Graph Sink 减少调度间隙内存占用每个独立算子都需要为输出张量分配 HBM 空间峰值显存由中间结果数量决定融合消除部分中间张量checkpointing 策略减少常驻激活内存融合算子复用片上缓存内存复用分析降低 HBM 占用编译时间框架侧直接把图交给 Device 运行时几乎没有编译开销使用 GE 后需要多轮优化、自动微分、任务图生成编译时间并未降低图编译优化本身是一次性开销且会随图规模增长Host-Device 交互每个算子启动都需要 Host 发射 kernel交互次数与算子数量成正比整图下沉把多次交互压缩为一次图级启动交互次数明显下降Graph Sink 让 Device 内部自主完成算子调度这个对比表说明GE 的优化是面向运行时性能的设计不是面向编译速度的优化。使用 GE 后训练迭代延迟和显存压力通常会有所改善但首次编译时间会比逐算子方案更长。在部署推理服务时这种取舍是可接受的因为编译只发生一次而推理会反复执行。在训练场景下如果模型结构频繁变化过长的编译时间会影响实验迭代效率这时需要合理设置 GE 的优化级别。在实际工程中选择优化策略时需要考虑模型的规模、迭代频率、部署模式。小型模型可能因图规模有限而无法充分展现融合收益大型模型则更容易从融合和下沉中获益。离线推理任务对编译时间并不敏感可以启用最高优化等级在线服务需要考虑冷启动时间需要在首次请求延迟和持续吞吐之间做取舍。训练任务如果每个 step 的图结构固定编译开销可以被大量迭代摊平如果图结构动态变化过高的优化等级反而会成为瓶颈。优化级别的设置需要在编译深度和运行收益之间做平衡。GE 通常提供多个优化等级低等级只做基础优化编译速度快但运行时收益有限高等级会启用更多融合和下沉策略编译时间更长但执行效率更高。对于开发阶段频繁修改的模型可以选择较低的优化等级以缩短迭代周期对于最终部署的模型则应该启用完整优化以获得最佳运行时性能。结尾GE 作为 CANN 的图编译核心其融合优化和整图下沉执行是昇腾硬件发挥算力的关键上游决策。图编译前端把框架计算图转换为统一的 GE IR让后续的常量折叠、CSE、算子融合、自动微分等优化能够统一处理。算子融合引擎在规则匹配、收益估算和边界冲突之间做权衡把 ConvBN、MatMulAdd 等常见模式合并为高效的融合算子。自动微分模块生成反向图并通过 checkpointing 和内存复用策略控制训练显存。整图下沉执行把 Host 从繁重的逐算子调度中解放出来让 Device 侧的 Runtime 自主管理任务流和事件同步。仓库地址https://atomgit.com/cann/ge