GE 图引擎：计算图优化的三阶段流水线

在深度学习框架的世界里，计算图（Computational Graph）是所有神经网络模型的通用表示形式。无论是 PyTorch 的动态图、TensorFlow 的静态图，还是 ONNX 的中间表示，本质上都是对张量计算流程的有向图描述。然而，当这些计算图从训练框架走向昇腾 AI 处理器时，面临的核心挑战是：如何让高层抽象的计算图真正适配底层硬件的并行计算能力

程序员清洒

79人浏览 · 2026-05-22 11:28:14

程序员清洒 · 2026-05-22 11:28:14 发布

在这里插入图片描述

前言

这个问题的答案，就是 GE（Graph Engine，图引擎）存在的根本意义。作为昇腾异构计算架构 CANN 的核心组件之一，GE 承担着从多框架计算图到昇腾达芬奇架构可执行指令之间的桥梁作用。它不是简单的格式转换工具，而是一个完整的计算图编译优化系统——通过图准备、图优化、图编译三阶段流水线，将原始计算图逐步转化为经过深度优化的底层执行方案。

理解 GE 的设计，就是理解昇腾如何把"软件定义的灵活性"与"硬件加速的高性能"统一起来的关键。本文将深入 GE 的架构内部，剖析三阶段流水线的每一个环节，揭示"编译器的编译器"这一核心定位背后的技术逻辑。

一、为什么昇腾需要独立的图引擎

1.1 计算图的多态性与硬件专用性的矛盾

深度学习框架产生的计算图具有高度的多态性。同一个 Transformer 模型，在 PyTorch 中是以 torch.fx.Graph 形式存在，在 TensorFlow 中是以 tf.Graph 形式存在，导出为 ONNX 后又是一套基于 onnx.NodeProto 的表示。这三套表示在算子定义、数据布局、控制流表达上均有差异。

而昇腾 AI 处理器的底层执行模型是高度专用的。昇腾达芬奇架构的 AI Core 以三维立方体计算单元（Cube Unit）为核心，配合向量计算单元（Vector Unit）和标量计算单元（Scalar Unit），要求计算任务必须以特定的任务块（Task Block）形式提交，每个任务块包含明确的寄存器配置、内存地址和同步原语。

从多态的计算图表示到专用的硬件执行模型，中间隔着巨大的语义鸿沟。直接让每个框架各自适配硬件，不仅重复劳动，更无法做跨框架的全局优化。这正是 GE 出现的根本原因：提供一个统一的计算图中间表示，并在此基础上进行与硬件深度耦合的编译优化。

1.2 GE 在 CANN 架构中的位置

GE 在 CANN 软件栈中处于"框架对接层"与"运行时层"之间，其上下游关系如下：

PyTorch/TensorFlow/ONNX/PB
         │
    ┌────▼────┐
    │  Adapter │  ← TorchAir / TFA / ONNX Parser
    └────┬────┘
         │  标准计算图（FE Graph）
    ┌────▼────────────────────┐
    │      GE 图引擎           │
    │  ┌──────────────────┐   │
    │  │ 图准备（Preparation）│   │
    │  ├──────────────────┤   │
    │  │ 图优化（Optimization）│   │
    │  ├──────────────────┤   │
    │  │ 图编译（Compilation） │   │
    │  └──────────────────┘   │
    └────┬────────────────────┘
         │  优化后的计算图（Model/OM）
    ┌────▼────┐
    │ Runtime  │  ← 加载模型、分配内存、下发任务
    └────┬────┘
         │  任务块（Task）
    ┌────▼────┐
    │  昇腾    │
    │ 达芬奇   │
    │ 架构     │
    └─────────┘

GE 的核心职责可以概括为：接收来自多框架的标准计算图，输出经过深度优化的、可直接由 Runtime 加载执行的离线模型文件。

二、三阶段流水线总览

GE 对计算图的处理是一条严格的三个阶段流水线，每个阶段有清晰定义的输入、输出和不变式。

输入：标准计算图（FE Graph，基于 onnx.proto 格式）
  │
  ▼
┌─────────────────────────────────────────────┐
│         阶段一：图准备（Preparation）          │
│  输入：原始计算图（可能含控制流、动态形状）    │
│  输出：静态化、标准化的计算图（Prepare Graph） │
│  核心任务：格式转换、形状推导、校验、控制流展开 │
└──────────────────┬──────────────────────────┘
                   │
                   ▼
┌─────────────────────────────────────────────┐
│       阶段二：图优化（Optimization）          │
│  输入：Prepare Graph                          │
│  输出：优化后的计算图（Fusion Graph）          │
│  核心任务：算子融合、常量折叠、内存规划、        │
│            多流并行分析、模型下沉预处理          │
└──────────────────┬──────────────────────────┘
                   │
                   ▼
┌─────────────────────────────────────────────┐
│        阶段三：图编译（Compilation）           │
│  输入：Fusion Graph                           │
│  输出：离线模型文件（.om 文件）                │
│  核心任务：算子编译（AKE/LLVM）、指令生成、    │
│            内存复用决策固化、任务块调度方案生成  │
└──────────────────┬──────────────────────────┘
                   │
                   ▼
输出：.om 离线模型文件（可由 Runtime 直接加载执行）

这三个阶段是顺序执行的，前一个阶段的输出是后一个阶段的输入，阶段之间通过标准化的图结构传递数据，保证了解耦性和可测试性。

三、阶段一：图准备（Graph Preparation）

3.1 核心任务概述

图准备阶段是 GE 流水线的入口，负责将来自各框架适配器的原始计算图转化为 GE 内部统一的、静态化的标准计算图。这个阶段解决的核心问题是：消除输入图的多样性与不确定性，为后续优化阶段提供一个确定性的优化基础。

3.2 格式转换与解析

GE 通过适配器层接收不同来源的计算图：

PyTorch 路径：通过 TorchAir 适配器，将 torch.fx.Graph 转换为 GE 图。TorchAir 是 CANN 提供的 PyTorch 对接插件，通过 torchair.inference.inference_adapter() 接口实现无缝集成。
TensorFlow 路径：通过 TFA（TensorFlow Adapter）将 tf.Graph 或 SavedModel 转换为 GE 图。
ONNX 路径：通过 ONNX Parser 直接解析 .onnx 文件。
PB 路径：通过 Protobuf Parser 解析 .pb 格式的序列化图。

转换后的统一内部表示是基于 ONNX Protobuf 格式扩展的 ge::ComputeGraph，其核心数据结构如下：

// ge/core/compute_graph.h（简化版）
class ComputeGraph {
public:
    std::string name_;                    // 图名称
    std::vector<NodePtr> nodes_;          // 所有计算节点
    std::vector<EdgePtr> edges_;          // 所有数据流边
    GraphId_t graph_id_;                  // 图唯一标识
    std::shared_ptr<ComputeGraph> parent_; // 父图（用于子图）
    std::vector<ComputeGraphPtr> subs_;   // 子图列表
    // ...
};

逐行解释：

第 3 行：name_ 存储计算图的名称，用于调试和日志追踪。
第 4 行：nodes_ 是图中所有计算节点的容器，每个 NodePtr 指向一个 ge::Node 对象，代表一个算子（如 Conv2D、MatMul）。
第 5 行：edges_ 是数据流边的容器，每个 EdgePtr 描述一个从源节点的输出端口到目标节点的输入端口的数据依赖关系。
第 6 行：graph_id_ 是图的全局唯一标识符，用于在多子图场景下区分不同的计算图。
第 7 行：parent_ 指向父图，用于支持控制流子图（如 If/Else、While 循环）的层次化表示。
第 8 行：subs_ 存储所有直接子图，与 parent_ 共同构成图的层次结构。

3.3 形状推导（Shape Inference）

动态形状是图准备阶段最棘手的问题之一。许多算子（如 Resize、Squeeze）的输出形状依赖于输入张量的运行时值，而 GE 的优化阶段需要确定的形状信息才能进行内存规划和算子融合。

GE 的形状推导引擎通过以下机制处理动态形状：

常量形状传播：对于形状完全由输入决定的算子，直接通过 ONNX shape inference 规则推导。
符号形状约束：对于含动态维度的张量，引入符号变量（Symbolic Dimension），建立符号之间的等式/不等式约束，在编译期尽可能缩小形状范围。
分段静态化：对于无法在编译期确定形状的场景，将计算图按照动态形状发生点切分为多个子图，每个子图内部形状静态，运行时通过"模型下沉"机制动态下发子图。

核心代码路径在 ge/python/ge/graph_utils.py 中的形状推导调用链：

# ge/python/ge/graph_utils.py（简化版）
def infer_shape(graph):
    """对计算图进行形状推导"""
    changed = True
    while changed:                     # 第 4 行
        changed = False
        for node in graph.nodes:       # 第 6 行
            if node.op_type in SHAPE_INFER_MAP:  # 第 7 行
                infer_func = SHAPE_INFER_MAP[node.op_type]
                new_shapes = infer_func(node)     # 第 9 行
                if new_shapes != node.output_shapes:
                    node.output_shapes = new_shapes
                    changed = True                # 第 12 行
    return graph

逐行解释：

第 4 行：使用迭代固定点算法（fixed-point iteration），反复遍历图中节点直到形状信息不再变化。
第 6 行：遍历计算图中的每个节点。
第 7 行：检查当前节点的算子类型是否有注册的形状推导函数。SHAPE_INFER_MAP 是一个全局字典，将算子类型映射到对应的形状推导实现。
第 9 行：调用该算子的形状推导函数，传入节点对象（包含输入形状信息），返回推导后的输出形状列表。
第 10-12 行：如果推导出的新形状与节点当前记录的输出形状不同，则更新节点形状，并将 changed 标记为 True，触发下一轮迭代。

3.4 图校验（Graph Validation）

在形状推导完成后，GE 对计算图进行完整性校验，包括：

拓扑校验：检测图中是否存在环（非法循环依赖）。GE 使用基于 DFS 的拓扑排序算法，若发现后向边则报错。
算子校验：检查每个算子的输入/输出个数、数据类型、格式（NCHW/NHWC/ND）是否合法，是否与昇腾 AI 处理器的支持范围匹配。
内存可行性预判：基于当前形状信息，估算图中所有张量的峰值内存占用，若超过设备内存上限则提前报错，避免编译到一半失败。

图准备阶段的输出是一个静态化、标准化、通过校验的 ComputeGraph 对象，流入图优化阶段。

四、阶段二：图优化（Graph Optimization）

图优化阶段是 GE 最核心、最复杂的部分，直接决定了最终模型的执行性能。这个阶段的目标是：在不改变计算语义的前提下，对计算图进行等价变换，最大化硬件利用率。

4.1 算子融合（Operator Fusion）

算子融合是深度学习编译器最经典的优化手段。GE 的融合策略针对昇腾达芬奇架构的特点进行了深度定制。

4.1.1 融合模式

GE 支持多种融合模式，每种模式对应不同的硬件执行策略：

融合模式	描述	典型场景
`fusion_mode::Basic`	基础融合：相邻算子直接融合	Conv2D + BatchNorm + ReLU
`fusion_mode::Scope`	范围融合：指定子图整体融合	Transformer 中的 Self-Attention 子图
`fusion_mode::Dynamic`	动态融合：支持动态形状的融合	动态输入分辨率的检测网络
`fusion_mode::Compress`	压缩融合：权重压缩后的融合	量化模型中的权重量化+算子融合

4.1.2 融合引擎的实现

GE 的融合决策引擎位于 ge/core/graph_optimizer/fusion_pass_manager.cpp，核心逻辑是基于模式匹配的融合 pass 调度：

// ge/core/graph_optimizer/fusion_pass_manager.cpp（简化版）
Status FusionPassManager::RunFusionPasses(const ComputeGraphPtr &graph) {
    auto passes = GetFusionPassesByPriority();  // 第 3 行
    for (const auto &pass : passes) {           // 第 4 行
        bool changed = true;
        while (changed) {                       // 第 6 行
            changed = false;
            std::vector<NodePtr> matched_nodes;
            pass->Match(graph, matched_nodes);  // 第 9 行
            if (!matched_nodes.empty()) {
                auto ret = pass->Fuse(graph, matched_nodes);  // 第 11 行
                if (ret == SUCCESS) {
                    changed = true;             // 第 13 行
                }
            }
        }
    }
    return SUCCESS;
}

逐行解释：

第 3 行：获取所有已注册的融合 pass，按优先级排序。优先级决定了融合的顺序（例如先融合 Conv+BN，再融合融合后的节点与 ReLU）。
第 4 行：遍历每个融合 pass。
第 6 行：对每个 pass 使用迭代应用策略——一次融合可能产生新的可融合模式，需要反复匹配。
第 9 行：在当前计算图中匹配该 pass 对应的算子模式。Match 方法使用图模式匹配算法（类似子图同构）找到所有符合条件的节点组合。
第 11 行：对匹配到的节点组合执行融合变换，将多个算子节点合并为一个融合算子节点，同时重新连接数据流边。
第 13 行：如果融合成功，设置 changed = true，触发对该 pass 的下一轮匹配。

4.2 常量折叠（Constant Folding）

常量折叠是另一个重要的图优化 pass。其核心思想是：在编译期尽可能预计算所有可计算的常量表达式，避免运行时重复计算。

典型的常量折叠场景包括：

权重+偏置的预融合：如果模型中存在 Constant 节点与 Add 节点的组合，且 Constant 的值在编译期已知，则直接预计算加法结果，消除运行时的 Add 算子。
Shape/Size 算子的静态求值：对于输入形状已知的 Shape 或 Size 算子，直接在编译期计算出结果常量，替换为 Constant 节点。
类型转换的预计算：如果 Cast 算子的输入是常量，则直接在编译期完成类型转换。