图引擎架构原理剖析：GE在昇腾NPU软件栈中的核心角色

前言

CANN（Compute Architecture for Neural Networks）中的GE（Graph Engine）是昇腾NPU软件栈的核心引擎，负责计算图的构建、优化、编译和执行。在昇腾NPU上进行深度学习推理时，所有的计算最终都经过GE的处理和调度。理解GE的架构原理，对于进行昇腾NPU性能调优、问题诊断和定制开发都至关重要。

GE的设计目标是实现高效的图级优化和执行。通过将神经网络模型表示为计算图，GE可以在全局视角下进行算子融合、内存规划、调度优化等高级优化，这些优化在单算子层面是无法完成的。同时，GE的图抽象使得它可以支持多种前端框架（PyTorch、TensorFlow、ONNX）到昇腾NPU的统一编译和执行。

一、GE在昇腾软件栈中的位置与职责

1.1 软件栈分层架构

昇腾NPU的软件栈采用分层设计，从上到下依次是：应用层（AI框架如PyTorch、TensorFlow）、图编译层（GE）、算子层（TTL/TBE）、驱动层（Kernel Driver）、硬件层（NPU）。GE处于图编译层的核心位置，起到承上启下的关键作用。

应用层的框架（如PyTorch）将模型表达为自身的计算图，这些计算图通过昇腾的适配层（如PyTorch NPU后端）被转换为GE可以理解的中间表示（IR）。GE对这个中间表示进行优化和编译，生成可以在昇腾NPU上执行的低级指令，由算子层和驱动层具体执行。

这种分层设计有几个重要优势：不同AI框架可以复用相同的底层基础设施；图级别的优化可以在不关心上层框架细节的情况下进行；算子层和硬件层的变化不会影响上层的应用代码。

1.2 GE的核心职责

GE的核心职责包括四个方面：图解析与验证、图优化、图编译、执行调度。

图解析与验证负责将输入的中间表示解析为GE内部的图结构，并验证图的正确性（如检查算子输入输出shape是否匹配、数据类型是否兼容等）。

图优化负责在图级别进行各种优化，包括算子融合（将多个相邻算子合并为一个）、常量折叠（将可以在编译时计算的值提前计算）、死代码消除（删除不会被使用的算子）等。

图编译负责将优化后的图转换为昇腾NPU可以执行的低级指令，包括算子的调度顺序、内存的分配方案、数据的搬运路径等。

执行调度负责在运行时管理计算任务的执行，包括算子的启动和同步、内存的管理、异常的处理等。

二、计算图的分层表示模型

2.1 Graph/Node/Edge三层抽象

GE使用Graph（图）、Node（节点）、Edge（边）三层抽象来表示计算图。

Graph是顶层结构，包含所有的Node和Edge，以及图的全局属性（如输入输出节点列表、执行配置等）。

Node代表计算图中的一个算子或操作，包含算子的类型、属性参数、输入输出描述等信息。每个Node可以有多个输入Edge和多个输出Edge。

Edge代表数据流，连接两个Node，表示数据的依赖关系。Edge有源节点和目标节点，以及数据shape和dtype的描述。

python 复制代码

# GE图的创建和操作示例
from ge import Graph, Node, Tensor

# 创建图
graph = Graph()

# 创建节点（算子）
conv_node = graph.create_node(op_type="Conv2D", name="conv1")
conv_node.set_attr("kernel_size", [3, 3])
conv_node.set_attr("stride", [1, 1])
conv_node.set_attr("padding", [1, 1])

relu_node = graph.create_node(op_type="ReLU", name="relu1")

# 创建边（数据流）
graph.create_edge(conv_node, "y", relu_node, "x")

# 设置输入输出
graph.set_inputs([conv_node])
graph.set_outputs([relu_node])

2.2 数据流与控制流

GE的图表示支持两种依赖关系：数据流（Data Flow）和控制流（Control Flow）。

数据流依赖表示一个算子的输出作为另一个算子的输入，是最常见的依赖类型。

控制流依赖表示两个算子之间没有数据传递，但需要保证执行顺序（如某些需要先初始化再使用的场景）。控制流依赖在大多数模型中不常见，主要用于一些特殊的控制逻辑。

2.3 动态图与静态图

GE支持两种执行模式：动态图（Dynamic Graph）和静态图（Static Graph）。

动态图模式下，图的结构在运行时可以变化，每次执行都可能使用不同的计算路径。这种模式灵活度高，适合研究和实验，但优化空间有限。

静态图模式下，图的结构在编译时确定，运行时不发生变化。这种模式允许GE进行更深入的优化（如算子融合、内存规划），适合生产部署。昇腾NPU的推理主要使用静态图模式以获得最佳性能。

python 复制代码

# 动态图模式（训练时常用）
@ge.dynamic
def forward_dynamic(x):
    if x.sum() > 0:
        return x * 2
    else:
        return x / 2

# 静态图模式（推理时常用）
@ge.static
def forward_static(x):
    return x * 2  # 编译时确定的结构

三、图的编译流程详解

3.1 图解析与IR转换

编译流程的第一步是将各种框架的模型（如PyTorch的TorchScript、TensorFlow的SavedModel、ONNX的模型）转换为GE的中间表示（IR）。这个转换过程需要处理不同框架之间的语义差异，如操作符的命名、数据格式的定义、控制流的表达方式等。

IR采用Protobuf格式定义，包含图的拓扑结构、算子的属性、数据的shape和dtype等信息。转换后的IR是一个纯描述性的表示，不包含任何可执行代码。

python 复制代码

# IR转换示例
import ge.ge as ge_api

# 从ONNX模型转换
onnx_model = "resnet50.onnx"
graph = ge_api.GraphApi().load_model(onnx_model)

# 从PyTorch模型转换
torch_model = load_torch_model("resnet50.pt")
torch_script = torch.jit.trace(torch_model, example_input)
graph = ge_api.GraphApi().load_from_torch(torch_script)

# 从TensorFlow模型转换
tf_model = "resnet50_savedmodel"
graph = ge_api.GraphApi().load_from_tf(tf_model)

3.2 图优化 passes

优化流程由一系列pass（优化步骤）组成，每个pass负责一类特定的优化。pass之间可能存在依赖关系，需要按照特定的顺序执行。

主要的pass包括：

算子融合pass：识别可融合的算子组合（如Conv+BN+ReLU），将多个算子合并为一个融合算子。融合可以减少kernel启动开销和中间结果的显存读写。

常量折叠pass：识别图中可以提前计算的常量表达式，将计算结果直接嵌入图中，避免运行时的重复计算。

形状推断pass：推断每个算子输出张量的shape，为后续的内存分配和调度优化提供信息。

内存规划pass：分析数据依赖关系，计算每个张量的生命周期，分配内存地址，最大化内存复用。

调度优化pass：确定算子的执行顺序和并行策略，优化流水线效率。

python 复制代码

# 配置优化pass
ge_api.set_ops_run_offline(False)  # 在线执行
ge_api.add_custom_optimize_pass("op_fusion_pass", {
    "enable_conv_bn_fusion": True,
    "enable_conv_relu_fusion": True,
    "fusion_group_size_threshold": 10
})
ge_api.add_custom_optimize_pass("memory_planning_pass", {
    "enable_memory_reuse": True,
    "memory_allocator": "greedy"
})

3.3 算子调度与内存管理

编译的阶段是确定算子的调度顺序和内存分配方案。调度优化需要考虑多个因素：数据依赖关系、硬件资源约束（计算单元、内存带宽）、负载均衡等。

内存管理是编译阶段的关键任务。深度学习模型通常包含大量的中间张量，如果为每个张量都分配独立的内存空间，总内存需求会非常大。GE的内存规划器会分析每个张量的生命周期，对于生命周期不重叠的张量，复用同一块物理内存。

python 复制代码

# 内存管理配置
ge_api.set_mem_config({
    "mem_reuse_level": "full",  # 最大内存复用
    "l2_mem_size": 8 * 1024 * 1024,  # L2缓存大小
    "l1_mem_size": 512 * 1024,  # L1缓存大小
    "hbm_size": 16 * 1024 * 1024 * 1024  # 总可用HBM
})

# 调度配置
ge_api.set_scheduler_config({
    "enable_stream_parallel": True,
    "enable_task_parallel": True,
    "task_queue_size": 64
})

四、动态shape处理机制

4.1 动态shape的挑战

在实际应用中，输入张量的shape可能不是固定的（如不同大小的输入图像）。动态shape处理是图编译中的一个重要挑战，需要在编译时处理运行时才能确定的shape信息。

动态shape带来的挑战包括：内存分配需要在运行时根据实际shape进行；某些优化（如算子融合）在shape不确定时可能无法进行；调度顺序可能需要根据实际shape进行调整。

4.2 GE的动态shape解决方案

GE通过"静态+动态"的混合策略处理动态shape：对于shape固定的部分，在编译时确定；对于shape可能变化的部分，预留动态处理的机制。

python 复制代码

# 动态shape配置
graph.set_dynamic_input_shape({
    "input_0": {
        "shape_range": [[1, 3, 224, 224], [1, 3, 448, 448]],
        "dynamic_dims": [0, 2, 3]  # batch、height、width可变
    }
})

# 编译时保留动态处理能力
compiled_graph = ge_api.GraphCompile(graph, {
    "dynamic_shape": True,
    "enable_shape_inference": True
})

# 运行时传入实际shape
output = compiled_graph.execute({"input_0": actual_input})

4.3 动态Batch处理

动态Batch是另一种常见的动态场景，模型需要处理不同batch大小的输入。GE支持通过dynamic_batch配置来处理这种情况。

python 复制代码

# 动态Batch配置
ge_api.set_dynamic_batch(graph, {
    "batch_list": [1, 2, 4, 8, 16],  # 支持的batch大小
    "dynamic_method": "tag"  # 使用标签区分不同batch
})

# 运行不同batch的输入
for batch in [1, 4, 16]:
    input_tensor = prepare_input(batch_size=batch)
    output = compiled_graph.execute({"input": input_tensor})

五、与TensorFlow/PyTorch图引擎的架构对比

5.1 与TensorFlow XLA的对比

TensorFlow的XLA（Accelerated Linear Algebra）是一个图编译器，将TensorFlow图优化并编译为机器码。GE与XLA在架构上有相似之处（都是图编译器、都进行算子融合和调度优化），但在实现细节和优化策略上有差异。

GE针对昇腾NPU的硬件特性进行了专门优化，包括对昇腾NPU张量计算单元的利用、特定的算子融合规则、针对昇腾内存层次结构的内存规划等。这些优化是通用图编译器无法提供的。

5.2 与PyTorch Glow的对比

PyTorch的Glow是一个图优化编译器，将PyTorch的计算图优化后分发到各种硬件后端。GE与Glow都支持多层中间表示（IR）的转换和渐进式优化。

GE的优势在于与昇腾NPU的紧耦合：它可以直接控制昇腾NPU的硬件资源，进行更精细的性能调优。Glow作为通用编译器，需要通过更通用的接口与硬件交互，优化空间相对有限。

使用前vs使用后：GE优化效果对比

在昇腾NPU上部署深度学习模型时，是否利用GE的优化能力对推理性能有显著影响。

使用前（单算子执行方案）：如果绕过GE直接调用算子层，每个算子独立执行，需要显式管理算子间的数据传递和同步。以ResNet50为例，不使用GE优化时需要手动管理约100个算子的执行顺序、内存分配和结果传递。这种方式不仅开发复杂度高，而且无法进行算子融合和全局内存优化，实际推理性能可能只有优化后的30-50%。

使用后（GE优化执行方案）：使用GE后，模型被编译为优化后的执行计划。GE会自动进行算子融合（如将Conv+BN+ReLU合并为单个kernel）、内存规划（复用生命周期不重叠的张量内存）、调度优化（并行执行无依赖的算子）。实测数据显示，使用GE优化后，ResNet50的推理性能可以提升约2-3倍，同时内存占用减少约40%。

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