昇腾图引擎GE的算子图编译优化与自动微分切图策略和整图下沉执行机制深度技术解读：从CANN开源仓库看架构原理与部署实践

前言

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 IR

GE 的输入通常是框架导出的计算图。在 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 计算图 IRgraph=ge.Graph()x=graph.placeholder(shape=[64,3,224,224],dtype="float16")w=graph.const(value=weight_tensor,name="conv1_w")conv=graph.op("Conv2D",inputs=[x,w],attrs={"stride":[1,1]})bn=graph.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 数量，但可能引入新的同步点或导致并行度下降。

Conv+BN 是深度学习图中最常见的融合模式。在分离执行时，Conv2D 先输出一个特征图写入 HBM，BatchNorm 再从 HBM 读取该特征图并做归一化。融合后，Conv2D 的输出直接作为 BatchNorm 的输入在 Vector 单元上完成缩放和偏移，中间张量不会离开片上缓存。这个融合模式在 ResNet 和 VGG 类骨干网络中几乎都会被触发。

MatMul+Add 融合则更多出现在全连接层和注意力投影中。分离执行时，MatMul 的结果需要写回 HBM，Add 再读取并做偏置相加。融合后，Add 的偏置加载可以与 MatMul 的累加过程重叠，减少一次完整的 HBM 读写。注意力机制中的 Q/K/V 投影和输出投影也常以类似方式融合。

融合边界与调度约束之间经常出现冲突。例如，一条融合规则要求两个算子必须位于同一个流上，但原图中它们分别属于不同的并行分支。如果强行合并，可能会破坏原有的并行执行结构，导致整体延迟增加。GE 的处理方式是在模式匹配阶段就把调度约束纳入收益估算，拒绝那些会破坏关键路径并行的融合。另一个常见冲突是融合后算子形状超出当前算子实现的限制。GE 会把图拆成可融合子图和不可融合子图，分别走不同的代码生成路径。

布局约束也是融合边界的重要组成部分。两个算子即使计算逻辑上可融合，如果它们的输入输出布局不一致，融合后可能需要额外插入布局转换节点。当转换开销超过融合收益时，GE 会保留原图。这种取舍说明融合优化不是单纯越多越好，而是要与整体调度目标一致。

# 定义 Conv+BN 融合规则pattern=ge.FusionPattern(anchor="Conv2D",followers=["BatchNorm"],constraint=lambdanode:node["data_format"]=="NCHW"andnode["training"]==False)# Conv+BN 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 节点forward=ge.build_graph(model,inputs)checkpoint_nodes=["layer_1_relu","layer_3_relu"]backward=ge.grad(forward,wrt=params,checkpoints=checkpoint_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 接口，可以导出任务图执行时间线，帮助开发者定位哪些任务之间存在不必要的同步或空闲间隙。

# 配置整图下沉执行session=ge.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 的图编译优化前后的差异。

这个对比表说明，GE 的优化是面向运行时性能的设计，不是面向编译速度的优化。使用 GE 后，训练迭代延迟和显存压力通常会有所改善，但首次编译时间会比逐算子方案更长。在部署推理服务时，这种取舍是可接受的，因为编译只发生一次，而推理会反复执行。在训练场景下，如果模型结构频繁变化，过长的编译时间会影响实验迭代效率，这时需要合理设置 GE 的优化级别。

在实际工程中，选择优化策略时需要考虑模型的规模、迭代频率、部署模式。小型模型可能因图规模有限而无法充分展现融合收益，大型模型则更容易从融合和下沉中获益。离线推理任务对编译时间并不敏感，可以启用最高优化等级；在线服务需要考虑冷启动时间，需要在首次请求延迟和持续吞吐之间做取舍。训练任务如果每个 step 的图结构固定，编译开销可以被大量迭代摊平；如果图结构动态变化，过高的优化等级反而会成为瓶颈。

优化级别的设置需要在编译深度和运行收益之间做平衡。GE 通常提供多个优化等级，低等级只做基础优化，编译速度快但运行时收益有限；高等级会启用更多融合和下沉策略，编译时间更长但执行效率更高。对于开发阶段频繁修改的模型，可以选择较低的优化等级以缩短迭代周期；对于最终部署的模型，则应该启用完整优化以获得最佳运行时性能。

结尾

GE 作为 CANN 的图编译核心，其融合优化和整图下沉执行是昇腾硬件发挥算力的关键上游决策。图编译前端把框架计算图转换为统一的 GE IR，让后续的常量折叠、CSE、算子融合、自动微分等优化能够统一处理。算子融合引擎在规则匹配、收益估算和边界冲突之间做权衡，把 Conv+BN、MatMul+Add 等常见模式合并为高效的融合算子。自动微分模块生成反向图，并通过 checkpointing 和内存复用策略控制训练显存。整图下沉执行把 Host 从繁重的逐算子调度中解放出来，让 Device 侧的 Runtime 自主管理任务流和事件同步。

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