前言

做深度学习推理部署的时候,很多人会遇到一个困惑:明明模型已经导出了,明明昇腾NPU就在那里跑着,但整个流程到底是怎样串起来的,从框架层到硬件层中间隔了多少个组件,每个组件各自负责什么,这些东西很难在脑子里形成一个清晰的图景。尤其是 CANN 生态对外部开发者来说一直是个黑箱,很多人只知道自己用了某个工具链,却不清楚这个工具链背后到底在做什么。ge 仓库的出现在很大程度上填补了这个空白——它是昇腾面向图模式的完整编译器和执行器实现,直接把从模型文件到设备可执行二进制这条链路给开源了出来。

本文把焦点放在 ge 本身,试图从架构设计者的角度去拆解这个项目的内部构造。我们会从它的定位出发,弄清楚它在 CANN 五层架构里处于什么位置、和周围的组件是什么关系;接下来深入编译器前端、图优化、算子编译、调度下沉这些核心模块,理解 GE 到底怎么把一个高层次的计算图变成可以在昇腾芯片上直接跑的高效程序;再从执行器的视角看看模型是怎么被加载、怎么被驱动、怎么和底层 Runtime 协作的;在收尾阶段通过一些实际的代码示例来感受 GE 的使用方式,以及它在性能层面带来了哪些实质性的改变。整篇文章不会回避技术细节,但会尽量把每一步背后的设计意图说清楚。

一、GE在CANN生态中的坐标与角色

理解 GE 的第一步,是把它放回整个 CANN 生态里去看。CANN 是昇腾的异构计算架构,覆盖了从应用层到底层硬件的完整软件栈。在官方文档里,这个架构被划分为五层,从上到下依次是:昇腾计算语言层(AscendCL)、昇腾计算服务层、昇腾计算编译层、昇腾计算执行层,以及最底层的昇腾计算基础和硬件层。GE 恰好横跨了编译层和执行层这两层——它既包含了图编译器(GE Compiler),也包含了图执行器(GE Executor),是连接高层框架和底层 Runtime 的关键枢纽。

具体来说,当用户在 PyTorch 或者 TensorFlow 里写好一个模型之后,如果要走图模式在昇腾 NPU 上运行,那么模型会先经过前端适配层被转换成一种叫做 AscendIR(简称 AIR)的中间表示。AscendIR 是一种静态计算图,以算子和张量为基本元素,表达模型的数据流和依赖关系。这个 IR 在此之后被送进 GE Compiler,在这里完成图级优化、算子融合、内存规划、流分配等一系列编译动作,最终生成一种叫做 OM(Offline Model)的离线模型文件。OM 文件是一种序列化的二进制产物,它包含了算子的调度信息、内存布局、权重数据等,执行时由 GE Executor 加载并在昇腾设备上驱动。

这就是所谓的离线编译路径。用户也可以选择不导模型文件,而是让适配层在框架运行时直接驱动 GE,这种模式叫做在线场景,GE 在这里作为框架的后端被实时调用。两种路径的核心区别在于:离线场景下编译和执行完全分开,编译在 Host 侧完成,产物可以独立部署;在线场景下编译和执行绑定在一起,适合开发调试阶段。不管哪种路径,GE Compiler 和 GE Executor 都是最核心的两个子系统,理解它们各自的职责边界是理解整个系统架构的关键。

从仓库分类来看,ge 属于 CANN 开源社区中的编译与运行时仓库,和它同属这一类的还有 metadef、runtime 以及 driver。ge 和这些仓库之间存在紧密的依赖关系:metadef 提供了图层面的数据结构和定义,runtime 提供了设备端的最底层执行能力,ge 则在这两者之上构建了图编译和图执行的完整逻辑。这种分层设计让各个仓库的职责足够清晰,也允许它们各自独立演进。

二、前端适配:从框架到AscendIR的统一入口

GE 的前端适配层负责把来自不同框架的模型转换成统一的 AscendIR 图表示。这件事的重要性在于:PyTorch 有自己的 AtenIR,TensorFlow 有自己的 GraphDef,ONNX 有自己的 protobuf 结构,这些东西从语义上都能表达同一个神经网络,但格式和抽象层次各不相同。如果 GE 需要分别对接每一种框架,就要分别写一套解析器,这是不可维护的。所以 GE 的策略是把这些五花八门的 IR 全部先翻译成 AscendIR,接下来后续的编译流程只跟 AscendIR 打交道。

目前官方维护了两个主要的适配组件。TorchAir 是 PyTorch 的适配器,它把 PyTorch 的AtenIR 转换成 AscendIR。这个转换发生在模型被 torch.compile 或者 torch.export 触发的时候,TorchAir 在框架内部截获计算图,把它吐给 GE。TF Adapter(TensorFlow Adapter)则负责把 TensorFlow 的 GraphDef 转成 AscendIR,走的是类似的思路。适配层和 GE 之间通过一套公开的接口来通信,这套接口定义了 AscendIR 的基本结构和行为规范。适配层可以独立于 GE 演进,只要接口保持兼容,两边就可以各自迭代。

对于 ONNX 和 PB(TensorFlow SavedModel)这类静态模型文件,适配的工作不是由框架适配层来完成,而是由 GE 内置的 Parser 来处理的。GE 的 parser 目录包含了多个子模块,分别对应不同的模型格式:tensorflow 目录处理 PB 文件,onnx 目录处理 ONNX 模型,caffe 目录处理 Caffe 的 prototxt,mindspore 目录处理 MindSpore 的模型。这些 Parser 读入文件后构造出对应的 AscendIR 图,接下来把图送到 GE Compiler 继续处理。Parser 和 Adapter 处理的都是同一个问题——如何把形形色色的模型格式变成 GE 认识的统一 IR——只是输入的来源不同。

这种统一入口的设计带来了一个额外的好处:如果某个框架新增了一个算子支持,开发者只需要在 AscendIR 的算子定义里补充相应的信息,接下来让 GE 的编译器能够正确处理它就行。适配层和 Parser 不需要为每一个新算子都写一遍转换逻辑,因为它们本质上只是做格式翻译,不负责算子语义的实现。算子的真正实现在算子仓里(比如 ops-math、ops-nn、ops-transformer 这些),GE 通过调用算子仓提供的接口来完成具体的计算。

三、AscendIR:图编译器的心脏

AscendIR 是 GE 整个编译流程的核心中间表示。理解它的设计思路,是理解 GE Compiler 为什么长成现在这个样子的大前提。

从抽象层次上看,AscendIR 和 ONNX、StableHLO 属于同一类东西——都是高层图表示(HLO),以算子和张量为基本构件。GE 的编译器在 AscendIR 这个层级做的是图级优化,包括算子融合、公共子表达式消除、死代码消除、常量折叠这些通用的编译器变换,也包括 GE 自研的一些融合类优化,比如基于 Pattern 的手写融合规则和基于算子分类的自动融合。

AscendIR 表达的是一张静态计算图。这个“静态”有两层含义需要区分清楚:第一层是指图结构在编译期固定下来,不会像动态图那样在运行过程中随时增删节点;第二层是指 shape 信息可以是静态的也可以是动态的——如果所有张量的维度在编译期就能确定,那就是静态 Shape 图;如果有些维度要等到运行时才能知道,那就是动态 Shape 图。GE 不支持真正意义上的动态图(即图结构可以在运行时改变),但支持动态 Shape 图。这个设计取舍是有道理的:静态图使得整图视角的优化成为可能,编译器可以在编译期就拿到完整的依赖关系,从而做出更激进的优化决策,比如精确的内存规划、跨算子的融合、流的并行分配等等。

AscendIR 的图结构由几个核心元素构成。Graph 是最顶层的容器,它承载了图中的所有节点、边、输入输出的描述信息。Node 代表一个算子,包含算子类型、输入输出张量的引用、以及附加的属性信息。Tensor 是算子的输入输出数据实体,带有 shape、dtype、format 等元数据。Attribute 是算子的附加配置信息,在构图时就确定好了,比如某些算子的模式参数。Data Edge 表示数据流向,从生产者节点指向消费者节点。Control Edge 表示纯依赖关系,没有数据传递,只约束执行顺序。

在 GE 的实现层面,节点之间的边并不是用一个独立的 Edge 对象来表示的,而是通过“锚点(Anchor)”来描述这种连边关系。DataAnchor 对应数据边,CtrlAnchor 对应控制边。每个锚点维护它对端的锚点引用,通过这种方式来表达节点之间的连接。这个设计的好处是把边和节点的引用关系统一起来管理,在图的遍历和变换操作中可以更高效地处理依赖关系。

四、GE Compiler:编译器的内部构造

GE Compiler 是把 AscendIR 变成可执行模型的核心引擎。它的内部工作分成几个主要阶段,每个阶段处理不同的优化任务。

第一个阶段是图级优化。这一步对整张图进行结构层面的改造,目的是减少冗余计算、压缩无效操作、合并可以合并的算子。通用编译器优化里最常见的几招这里都有:公共子表达式消除(CSE)会把重复出现的等价计算只保留一份,避免重复执行;常量折叠在编译期就把那些值恒定的表达式给算出来,省得运行时再跑一遍;死代码消除删掉那些对最终输出没有任何影响的节点和边。融合类优化是另一个重头戏,GE 支持两种融合路径:一种是基于 Pattern 的手写融合,开发者针对特定模型结构(比如 Transformer 里的 Attention 块)写好融合规则,编译器遇到匹配的 Pattern 就直接把它们合并成一个算子执行,这种方式可控可预测,适合对关键路径做高质量的特殊化处理;另一种是基于算子分类的自动融合,编译器根据算子的计算公式和输入输出依赖关系自动分析哪些算子可以合并,接下来通过代码生成技术生成融合算子再做在线编译,这种方式覆盖的融合空间更大,不需要人工去写 Pattern。

第二个阶段是算子级优化。GE 在做图优化的时候会利用图上已经推导出的 shape 信息,对单个算子进行在线编译。什么意思呢?传统的算子实现往往是针对某个通用的 shape 范围写的,但实际模型里往往出现的是特定的 shape,比如 batch size 是 32、hidden dimension 是 768 这种。GE 可以根据具体的 shape 定制化编译出专门针对这个 shape 的算子实现,从而在指令选择、寄存器分配等层面做出更激进的优化,达到更好的执行性能。

第三个阶段是流分配和内存规划。流分配(Stream Planning)分析计算图中哪些算子之间没有数据依赖、可以并行执行,接下来把它们分配到不同的执行流上,让昇腾芯片的多个计算单元同时开工。内存规划(Memory Planning)在静态 shape 的场景下以整图视角来规划张量的内存布局,通过复用同一块物理内存来存放不同生命周期的中间张量,从而降低峰值内存占用。这两个优化直接决定了模型在硬件上的执行效率和资源利用率。

在收尾阶段一个阶段是模型序列化(Model Serialization)。编译好的模型被序列化成 OM 文件,这个文件包含了算子的调度指令、权重数据、内存布局信息等。OM 文件是离线部署的核心产物——用户把这个文件扔到昇腾设备上,GE Executor 直接加载它就能跑,不需要再走一遍编译流程。

五、GE Executor:设备端执行的守门人

GE Executor 负责模型在昇腾设备上的运行。它做的事情分两大块:模型加载和模型执行。

模型加载阶段,GE Executor 把编译阶段生成的 OM 文件读取进来,提取出其中的算子二进制、权重数据、调度信息等资源,把它们映射到昇腾设备的内存和计算单元上。对于下沉(Sink)模式的模型,加载阶段还会把整个算子执行序列预先传输到设备端,这样执行的时候就不需要主机侧逐个下发指令了——设备端自己驱动整个执行流,主机侧只需要触发一次 launch 就够了。这个设计对于 Host-bound 的模型特别有价值,因为减少 Host 和 Device 之间的通信次数是提升端到端吞吐量的关键。

模型执行阶段,GE Executor 驱动整个图的执行逻辑,按照 AscendIR 里定义的算子顺序和依赖关系来调度算子。这里的调度和 GE Compiler 阶段规划好的流分配方案紧密配合——哪些算子在哪条流上、哪些地方需要流间同步、哪些地方需要等待数据就绪,这些信息都已经在编译阶段确定好了,执行器只需要忠实地执行这套计划。

除了基本的执行控制,GE Executor 还负责一些运行时层面的事情,比如异常处理、执行 profiling、以及和底层 Runtime 的交互。这里需要厘清一个容易混淆的地方:GE Executor 和 CANN 的 Runtime 是两个不同层次的东西。CANN Runtime 位于更底层,它提供的是最基础的设备管理、内存分配、核函数启动这些能力;GE Executor 则在 Runtime 之上,它操作的是经过编译的模型,执行的是图的语义,而不是单个算子。

六、插件与扩展:从自定义算子到自定义Pass

GE 的设计者显然预料到了用户会有定制化需求,所以留了插件和扩展的接口。开发者可以在不修改 GE 主体代码的情况下,通过两种主要方式来扩展它的能力。

第一种是 Ascend C 自定义算子入图。如果现有的算子库里没有你需要的算子,可以用 Ascend C 写一个自定义算子,接下来通过 GE 的接口把它注册到图里,让它参与到后续的编译和执行流程中。Ascend C 是 CANN 生态里的算子开发语言,它的语法风格接近 C++,但专门为昇腾的达芬奇架构做了适配。

第二种是自定义 Pass。如果想在图级别做一些额外的优化,比如针对特定业务场景的特殊融合规则或者图结构变换,可以写一个自定义的 Pass 插入到 GE 的编译流水线里。GE 的编译流水线是阶段化的,每个阶段之间有清晰的接口,这让插入自定义 Pass 变得相对简单——你只需要实现 GE 定义的 Pass 接口,接下来在编译配置里把你的 Pass 注册进去就行了。

GE 的 examples 目录里提供了融合 Pass 的示例代码,可以作为开发的起点。从示例里可以看到,一个自定义 Pass 通常需要定义两个核心函数:一个用来遍历图里的节点、识别融合机会,另一个用来执行实际的融合操作、把多个节点合并成更高效的算子组合。

七、GE的目录结构与模块映射

看过 GE 的源码仓库之后,理解它的目录结构能帮助把之前讲到的各个组件落到实处。api 目录包含 GE 对外暴露的编程接口,这些接口是上层框架适配器和 Parser 调用 GE 功能的入口。base 目录放的是基础工具方法和通用模块定义,属于那种被很多地方引用的公共代码。compiler 目录是 GE Compiler 图编译模块的所在地,所有的图优化、算子编译、模型序列化逻辑都收敛在这里。parser 目录包含了业界前端框架 IR 转 AscendIR 的实现,目前支持 tensorflow、onnx、caffe 和 mindspore 四种格式,每个子目录对应一个框架的 Parser。runtime 目录是 GE Executor 图执行模块的地盘,模型加载和图执行的逻辑在这里实现。graph_metadef 目录定义了图相关的数据结构,包括 Node、Tensor、Anchor 这些核心概念的底层表示。inc 目录是头文件集中存放的地方。examples 目录放了各种使用样例,包括 ResNet50 图像分类和 Qwen 大语言模型推理的完整示例。docs 目录则包含了构建文档、快速开始指南和架构说明等参考资料。

还有一个值得注意的模块是 dflow,它的中文名字叫 DataFlow 执行器,提供异构模型描述和串接执行能力。按照项目规划,dflow 未来会和 GE 解耦、独立成一个单独的仓库。它的存在说明 GE 的团队在架构层面是有所预判的——DataFlow 执行器和图编译器虽然是上下游关系,但它们的演进节奏和使用场景有差异,拆开来各自维护能更好地适应不同的发展需求。

八、生态集成:谁在用GE作为后端

GE 不是孤零零地存在着的,它在整个 CANN 生态里扮演的是底层基础设施的角色,有多个上层的框架和工具已经把它集成进去作为推理或图模式的后端。

TorchAir 把 GE 接入 PyTorch 的图模式。当用户用 torch.compile 编译模型的时候,TorchAir 把生成的 AtenIR 转成 AscendIR,接下来走 GE 的编译流程。如果用的是昇腾的特定分支或者插件,用户实际上感知到的只是换了一个后端设备,但底层跑的就是 GE 的编译器生成的模型。TFA(TensorFlow Adapter)做了类似的事情,把 TensorFlow 的 GraphDef 转成 AscendIR,让 TensorFlow 模型可以在 GE 上跑。这两个适配器都维护在 CANN 的独立仓库里,和 GE 通过公开接口交互,保持了各自的独立演进节奏。

JittorInfer 是基于昇腾芯片的大模型 C++ 推理框架,它把 GE 作为底层的图执行引擎。Triton Inference Server 的 GE Backend 则是把 GE 接入到了 Triton 的推理服务体系里——Triton 是一个流行的开源推理服务框架,支持多模型并发和动态 batching,接入 GE Backend 之后,Triton 就能够把模型编译和执行的任务委派给 GE,自己只做请求调度和结果返回。这些集成都不是 GE 本身提供的功能,而是由对应的上游项目来维护的,但它们的存在说明 GE 作为底层基础设施已经被生态广泛认可。

九、代码示例:从模型文件到OM的离线编译流程

为了让 GE 的工作方式更具体,这里用 atc 工具对 ONNX 模型进行离线编译的例子来说明编译的完整流程。

importacl

# 初始化 ACL,这是调用 GE/ATC 能力的通用前提
acl.init()

# 指定模型文件路径,这里以 ONNX 导出的 ResNet50 为例
model_path = "resnet50.onnx"

# 调用 atc 接口对模型进行编译
# 这个接口会先让 Parser 把 ONNX 转成 AscendIR,
# 接下来启动 GE Compiler 做图优化、算子编译、流分配、内存规划,
# 最终输出 OM 文件存到指定路径
atc_result = acl.atc(
    model=model_path,
    framework=5,          # 5 表示 ONNX 格式
    output="resnet50",     # 输出 OM 文件名(不含后缀)
    soc_version="Ascend910"
)

atc 这个接口把模型文件到 OM 的整个链路封装成了一步调用,这样做是为了降低离线编译的使用门槛。用户不需要关心 Parser 和 GE Compiler 各自的内部流程,只需要给一个模型文件、告诉系统目标芯片型号,剩下的全部由 atc 内部搞定。framework 参数的数值对应了不同的模型格式,ONNX 对应的就是 5,其他框架有各自对应的编号。

再看一个通过 TorchAir 在 PyTorch 在线场景下使用 GE 的例子。这个场景不需要手动调用 atc,GE 是作为 PyTorch 的设备后端直接被框架驱动的。

import torch
import torchair

# 选一台昇腾 NPU 作为 torch 的设备
device = torch.device("npu:0")

# 构造一个简单的线性层作为示例
model = torch.nn.Linear(512, 256).to(device)

# 用 TorchAir 的方式来编译这个模型
# npu_config 里的 backend 指定为 "ge",
# 这样 TorchAir 会把 AtenIR 转换成 AscendIR,
# 接下来走 GE Compiler 做在线编译,
# 在收尾阶段由 GE Executor 驱动执行
compiled_model = torchair.compile(
    model,
    config={"backend": "ge", "npu_id": 0}
)

# 运行推理,输入数据会自动从 Host 传到 Device,执行完再传回来
x = torch.randn(8, 512, device=device)
y = compiled_model(x)

把 backend 指定为 “ge” 而不是直接用某个设备编号,是因为 GE 在这里代表的是整个图执行引擎的运行时,而不仅仅是裸的硬件设备。GE 会在内部决定怎么把算子映射到具体的计算单元、怎么规划内存、怎么调度流。如果直接用 npu:0,默认会走 aclnn 的原生 API 调用路径,那是另一套执行模式,两者在优化策略和执行方式上有本质差别。

在收尾阶段来看一个通过 GE 的 C++ API 直接操作图和节点的例子,这个更接近 GE 内部的工作方式。

#include "graph/ge_graph.h"
#include "graph/operator.h"
#include "graph/tensor.h"

// 创建一个新的空图,这就是 AscendIR 的载体
ge::Graph graph = ge::Graph::Make("my_inference_graph");

// 向图里添加一个 MatMul 算子节点
ge::Operator matmul("MatMul");
matmul.set_input("x", ge::Tensor());  // 输入占位,具体数据运行时再喂
matmul.set_input("y", ge::Tensor());

// 再添加一个 ReLU 激活节点,它消费 MatMul 的输出
ge::Operator relu("Relu");
relu.set_input("x", matmul.get_output("y"));

// 把两个节点加到图里
graph.AddOp(matmul);
graph.AddOp(relu);

// 用 GE Compiler 对这张图进行编译
ge::GraphOptimizer optimizer;
ge::GraphCompileOptions compile_opts;
compile_opts.output_path = "./my_model.om";

ge::Status status = optimizer.Compile(graph, compile_opts);

这个 API 的设计风格参考了计算图编程的通用范式——先有图再有节点,节点之间通过算子名和数据依赖关系来建立连接。GE Compiler 在处理这张图的时候会遍历所有节点,根据算子类型去算子仓里查对应的实现,接下来执行图优化和算子编译。这里的 MatMul 和 Relu 并不是 GE 自己实现的,而是引用了 ops-nn 算子仓里的定义,GE 只负责把它们编织成一张可执行的图。

十、使用前后的效率对比

在实际部署场景里,GE 带来的改变是系统性的,不仅仅是某一个环节的提速。以下从几个关键维度来对比使用 GE 前后的差异。

场景维度 使用前(无GE方案) 使用后(GE图编译方案)
模型加载方式 逐算子调用,Runtime 每次都要解析图结构 整图下沉,OM 文件一次性加载到设备端
内存管理策略 每个算子独立分配和释放显存 整图视角的内存复用规划,峰值显存占用明显降低
算子执行效率 通用 shape 的预编译算子 基于具体 shape 的在线编译算子,单算子性能更优
多算子并行度 依赖 Runtime 基础调度能力 流分配优化后的并行执行,硬件利用率提升
框架接入方式 需手动适配各框架,算子映射复杂 统一 AscendIR 接口,适配层按框架维护,接入成本低
离线部署能力 需要保留完整的框架运行时 OM 文件独立可执行,部署包体积小、不依赖框架
优化空间 单算子层面的优化 图级别的融合优化、跨算子优化、内存规划联合优化
执行路径 Host 侧逐指令下发,通信开销大 Sink 模式下设备侧自主执行,Host-Device 交互大幅减少

从对比里可以看出,GE 的价值不是某一个点的突破,而是整条链路的协同优化:编译阶段通过图级优化和在线编译提升单算子效率,执行阶段通过下沉调度和流分配减少 Host-Device 通信和硬件空闲时间,内存规划通过整图视角的复用把峰值显存压下去。这几个环节的优化相互支撑,单独拿出来看每一个都不算惊人,但叠加在一起之后带来的整体收益是比较可观的。


仓库链接:https://atomgit.com/cann/ge

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐