深入剖析昇腾CANN开源仓库ge的架构设计、计算图编译优化机制与工程实践，计算图编译·算子融合·内存优化的内在逻辑与实战剖析

前言

在实际的大模型训练和推理场景中，开发者经常会遇到这样的困惑：为什么明明写的是一个简单的矩阵乘法加激活函数的代码，在昇腾NPU上跑出来的性能却跟预期差了好几倍？为什么同样一个模型，用PyTorch写出来的版本和用昇腾图引擎优化后的版本，显存占用相差接近一半？这些看似玄学的问题，背后其实都跟一个关键组件有关——昇腾CANN异构计算架构中的图引擎（Graph Engine），也就是社区开源仓库ge所承担的核心职责。

就是把ge仓库的内部结构、工作原理、以及它在实际工程中如何发挥作用这些东西，掰开了揉碎了讲清楚。我不会用一个简单的"ge是编译器"来概括它，因为这种概括既不准确也容易误导人。我会从它的架构定位开始，沿着计算图的构建、传递、编译优化这条主线，一层层揭开它的面纱。过程中会穿插真实的代码片段，解释为什么要这样设计，以及它在性能上带来了什么样的改变。

第一章：ge在昇腾CANN五层架构中的位置

在深入ge本身之前，有必要先把它的位置坐标说清楚。昇腾CANN并不是一个单独的编译器，而是一套完整的异构计算架构，涵盖了从应用层到底层硬件的五层结构。从上往下看，第一层是昇腾计算语言层AscendCL，提供给应用开发者的编程接口；第二层是昇腾计算服务层，包含AOL算子库和AOE调优引擎；第三层是昇腾计算编译层，这里就是ge的所在位置，具体承担Graph Compiler图编译器的职责；第四层是昇腾计算执行层，包括Runtime运行时和Graph Executor图执行器；最底层是昇腾计算基础层，包括各种驱动和底层服务。

ge作为Graph Compiler，扮演的是整个编译链路的枢纽角色。它的上游对接各种前端框架（PyTorch、TensorFlow等）传入的计算图表示，下游则生成可以在Runtime上执行的低级算子指令序列。可以把它理解为一个翻译官，前端框架用各自的语言描述"我要做什么计算"，ge负责把这个意图翻译成昇腾NPU能听懂并高效执行的"机器指令"。

这种定位决定了ge不是一个简单的语法解析器，它必须在翻译的过程中完成大量的优化工作——这些优化工作才是它价值的核心所在。接下来的几章会逐一展开它具体做了什么。

第二章：计算图的核心抽象与结构

理解ge的工作原理，首先要理解它操作的对象——计算图。在深度学习框架中，一个神经网络模型会被表达为一个有向无环图（DAG），图中的节点代表算子（比如MatMul、ReLU、Softmax），边代表张量数据的流向。ge的职责就是接收这样一个图，然后对它进行各种分析和改造，最终生成一个更高效的等价图。

在ge的内部实现中，一个计算图由几个核心概念组成。图本身是一个容器，里面包含了所有的节点和边，以及一些全局的元数据信息。每个节点对应一个算子，节点上附带着这个算子的属性参数。边则连接两个节点，代表张量从源节点的输出流向目标节点的输入，边上会标注张量的形状和数据类型信息。

对于开发者来说，理解这个抽象最大的意义在于：ge对图的所有优化，本质上都是对这个DAG结构的修改。它可能在某个地方把两个连续的节点合并成一个新节点，也可能把一个节点拆分成多个节点并行执行，还可能在某些边上插入额外的同步或拷贝节点来保证数据的正确性。所有这些操作，都是为了得到一个执行效率更高的图。

在实际的代码中，ge对图的表示会涉及一些关键数据结构。图对象会维护一个节点列表和一个边列表，每个节点会有一个唯一的标识符、算子类型、以及输入输出边的引用关系。这些数据结构的设计直接影响后续遍历和修改图的效率，所以在图构建的早期阶段，ge就投入了不少精力在数据结构的合理设计上。

第三章：图优化流程与关键机制

ge对计算图的处理流程可以分为几个主要阶段。首先是图的解析和校验阶段，这个阶段ge会读取来自前端框架的图描述，验证图的结构合法性——比如检查是否存在孤立节点、是否存在环形依赖、输入输出的形状是否匹配等。如果图结构本身就有问题，后面的优化和执行都会出问题，所以这个阶段的校验是必不可少的。

接下来是一些基础优化，这些优化不依赖于具体的硬件特性，属于通用的图级别优化。常见的内容包括常量折叠——如果某个节点的输入全是常量，ge就会在编译时把这个节点的结果直接算出来，而不用等到运行时再计算。另一个基础优化是公共子表达式消除，如果图中出现了两个完全相同的计算子图，ge会保留其中一个的结果让另一个直接引用，省掉重复计算。还有一个是无用节点消除，如果某个节点计算出来的结果没有被任何后续节点使用，这个节点就是无用的，ge会把它从图中删掉。

在基础优化之后，ge会进入跟硬件相关的优化阶段，这一阶段才是它真正体现价值的地方。首先是算子融合，这是ge最核心的优化手段之一。在深度学习中，很多算子经常连续出现，比如一个矩阵乘法后面跟着一个ReLU激活函数。在原始的图里，这两个算子是分开执行的，每执行完一个算子就要把结果写回显存，然后下一个算子再从显存把数据读出来。这个读写过程会消耗大量的时间和带宽，特别是对于大尺寸的张量来说，中间的读写延迟往往比实际计算本身还要高。算子融合的思路就是把这两个算子合并成一个融合算子，在昇腾NPU上只启动一次kernel执行，在这个kernel内部连续完成乘法和激活的计算，中间结果不需要写回显存，直接在芯片内部传递。这就像把两段流水线接起来变成一条更长的流水线，省掉了中间的装卸环节。

除了算子融合，ge还会做一些内存优化。在深度学习计算中，显存占用是一个永恒的话题。ge会在编译时分析每个张量的生命周期，找出那些"用过之后就不再需要"的张量，然后复用它们的显存空间。这个过程叫做显存复用优化，它可以让同样大小的显存容纳下更大的模型或者更大的batch size。这个设计背后的逻辑其实很朴素：显存是有限的资源，如果能让多个张量在时间上错开复用同一块物理显存，就能有效降低峰值显存占用。ge通过数据流分析来追踪每个张量的生存期，在这个基础上实现了一个比较高效的显存复用算法。

还有一个重要的优化方向是调度优化。ge会分析图中各个节点之间的依赖关系，然后生成一个并行执行策略。简单来说，如果两个算子之间没有数据依赖关系，它们就可以并行执行。ge会把可以并行执行的算子打包在一起，通过昇腾NPU的多核并行能力同时执行，从而提升整体的吞吐量。这个过程涉及到对计算图的拓扑排序以及对并行度的估算，需要在并行收益和调度开销之间做权衡。

第四章：代码结构与核心模块

ge仓库的代码结构设计体现了模块化的思路，整个仓库按功能划分成了若干相对独立的子模块，每个子模块负责一块具体的职责。这种设计的好处是职责边界清晰，便于理解和维护，同时也为后续的功能扩展留出了空间。

核心模块之一是图的表示层。这一层的代码负责定义计算图的基本数据结构，包括前面提到的节点、边、张量等概念的具体实现。图的表示层还提供了一组图的遍历接口，供其他模块调用。比如深度优先遍历可以用于检测图中的环路，广度优先遍历可以用于按层级组织图的节点拓扑。这些遍历接口是后续各种分析和优化算法的基础，所以在实现的时候特别注重了效率，避免在遍历过程中产生过多的临时对象。

核心模块之二是图的分析器。这个模块负责对输入的计算图进行各种静态分析，包括数据流分析、控制流分析、以及依赖关系分析。数据流分析会追踪每个张量从产生到消费的完整路径，帮助其他模块判断哪些优化是安全的——比如只有当一个张量在融合后不会被其他节点使用时，合并两个算子才是安全的。控制流分析则处理条件分支和循环结构的建模，确保在有条件跳转的情况下，优化不会破坏语义正确性。依赖关系分析会生成一个依赖图，标注每个算子必须等待哪些前置算子完成才能开始执行，这个依赖图是调度优化的输入。

核心模块之三是优化器。这个模块是ge的核心引擎，它实现了各种图优化算法。优化器按照预定的顺序依次调用各个优化 pass，每个 pass 实现一种特定的优化策略。pass之间通过注册机制组织在一起，可以通过配置文件来选择启用哪些 pass以及它们的执行顺序。这种可插拔的设计让ge可以灵活地支持不同的优化场景——在开发调试阶段可以禁用部分优化来方便调试，在生产环境则可以启用所有优化来追求极致性能。

核心模块之四是代码生成器。这个模块负责把优化后的计算图转换成低级算子指令序列，输出给下游的Runtime执行。代码生成器会根据目标硬件的特性（比如昇腾NPU的核数、内存层次结构等）来选择合适的指令编排方式。生成的指令序列包含了算子的输入输出地址、每个核要处理的数据分片、以及必要的同步信号。代码生成器的输出质量直接影响Runtime的执行效率，所以在实现的时候需要充分考虑数据布局、内存对齐、并行度分配等细节问题。

在代码层面，ge的实现大量使用了面向对象的设计模式。各个核心类之间通过接口来交互，这种设计保证了模块之间的松耦合。比如图的表示层只关心数据结构本身，不关心这些数据会被如何使用；优化器只通过接口操作图，不直接依赖具体的数据结构实现。这种分离设计让各个模块可以独立演进，降低了维护成本。

第五章：工程实践中的使用方式

对于大多数开发者来说，ge不是一个需要直接打交道的组件——它隐身在框架和Runtime之间，自动完成自己的工作。但对于想要深入优化模型性能的高级用户或者框架开发者来说，理解ge的工作原理以及如何跟它配合，仍然是有价值的。

一种典型的使用场景是在自定义模型或者新算子的开发过程中。当开发者写了一个新的融合算子，需要把它集成到昇腾的算子体系中时，ge的行为就直接影响了这个算子最终的执行效率。了解ge的融合规则和触发条件，可以帮助开发者更好地设计算子的接口和内部实现，从而更容易被ge识别并应用优化。比如一个融合算子如果暴露了太多细粒度的中间张量，ge可能就不容易判断它们是否可以被融合，这时候把接口设计得更"粗粒度"一些，效果往往会更好。

另一种场景是在性能调优过程中。当模型的执行效率不达预期时，开发者需要诊断瓶颈在哪里。如果ge负责的图优化阶段没有正确识别出某个融合机会，问题的根源可能在于图的结构或者某些属性参数设置不当。ge提供了一些调试手段来输出中间态的计算图，开发者可以对比优化前后的图结构变化，看看是否所有预期的优化都被应用了。这个诊断过程需要对ge的内部机制有一定了解，但从实际效果来看，这种投入往往是值得的——找到并修复一个融合遗漏问题，对模型整体吞吐量的提升可能是成倍的。

在实际项目中，ge还跟框架适配层有紧密的协作。不同的深度学习框架对计算图的表示方式各不相同，PyTorch有TorchScript的图表示，TensorFlow有GraphDef或者FusedGraph。ge通过框架适配器把这些不同格式的图描述转换成它内部统一使用的数据结构。这个转换过程不是简单的翻译，ge会在这个阶段做一些预处理优化，比如把框架特定的算子映射到昇腾对应的算子实现、把框架中的分组卷积等特殊结构转换成ge更擅长的融合模式。这个设计让ge的优化引擎可以专注于图的优化本身，而不用关心前端框架的差异。

从工程实现的角度看，ge跟上下游模块的交互采用了流水线式的架构设计。上游框架适配器把图输入ge，ge在内部经过一系列优化处理后，把结果输出给下游的Runtime。这个流水线式设计让每个环节可以独立演进，同时也便于在出现问题时做分阶段的诊断。在实际部署中，这个流水线会在模型加载的时候执行一次，生成优化后的执行计划缓存在本地，后续每次推理都直接使用缓存的执行计划，省掉了重复编译的开销。

第六章：典型优化案例与代码示例

为了更具体地说明ge在实际场景中如何工作，这章通过几个典型案例来展示它的优化过程和效果。每个案例都会给出优化前后的代码对比，以及ge在这个过程中做了什么具体的改动。

案例一：连续矩阵乘法与激活函数的融合

考虑一个最简单的深度学习前向计算场景，输入x经过一个线性变换后加上偏置，然后通过ReLU激活函数。在一个没有经过ge优化的场景中，这个计算在图级别会被表示为三个连续的节点：矩阵乘法（MatMul）、逐元素加法（Add）、ReLU激活（ReLU）。这三个节点是顺序依赖关系，每次都要把中间结果写回显存再读出来。

# 优化前的计算表达a=torch.matmul(x,w)# WHY: 三个算子各自独立执行，中间结果必须写回显存b=a+b# WHY: 跨算子的数据搬运消耗大量时间和带宽c=torch.relu(b)# WHY: 每次写回和读出都是一次显存的完整读写操作

ge在处理这个图的时候，会识别出这三个算子构成了一个典型的融合模式。它会把它们合并成一个单一的融合算子，在昇腾NPU上只启动一次kernel。这个融合算子在编译时生成，内部的计算逻辑就是把原来三个算子的计算合并在一起。对于MatMul+Add+ReLU这种组合，昇腾NPU上的融合算子只需要一次kernel启动，数据在芯片内部通过寄存器或者高速缓存传递，完全省掉了中间结果的显存访问。

# ge优化后的计算表达# ge识别出MatMul→Add→ReLU的融合模式后，生成融合算子# 在昇腾NPU上只启动一次kernel，数据始终保留在芯片内部# 这就是为什么融合后性能提升如此显著的原因z=fused_linear_relu(x,w,b)# WHY: 一次kernel启动替代原来的三次，数据不离开NPU

案例二：显存复用对大模型训练的影响

在大模型训练场景中，显存的消耗主要来自模型参数、梯度和优化器状态这三个部分，以及前向计算过程中产生的中间激活值。对于一个参数量巨大的Transformer模型，中间激活值占用的显存可能超过模型参数本身。ge通过显存复用优化来应对这个问题。

# 显存优化前的张量分配策略# 每个中间张量都分配独立的显存空间，互不干扰但显存利用率低y1=layer1_forward(x)# WHY: 分配new_size显存，layer2开始后y1的空间其实已不再被使用y2=layer2_forward(y1)# WHY: 分配new_size显存，此时y1的物理空间已空闲y3=layer3_forward(y2)# WHY: 分配new_size显存，y2的物理空间也已空闲# 峰值显存 = 3 × new_size，加上模型参数和梯度，总显存占用很大

ge在编译时分析每个张量的生存期，发现y1在layer2计算完成后就没有任何节点再使用它了，因此y1所占用的显存空间可以被layer2的输出y2复用。同样的逻辑也适用于y2和y3。这种复用策略可以让中间激活值的峰值显存占用降低一半甚至更多。

# ge优化后的显存复用策略# ge通过分析张量生存期，在编译时规划显存复用# y1完成即复用，y2完成即复用，峰值显存从3×new_size降到2×new_size# 这就是为什么启用ge显存优化后，同样硬件能跑更大batch的原因

案例三：并行调度优化对吞吐量的提升

当一个计算图包含多个没有依赖关系的计算分支时，ge会生成一个并行调度策略把这些分支分配到昇腾NPU的不同核上同时执行。这对于多分支的模型结构（比如Vision Transformer中的多个patch embedding分支）尤其有效。

# ge优化前的串行调度# 四个分支依次执行，总耗时是各分支耗时的总和p1=patch_embed_1(x)# WHY: 等待前一个分支完成才开始p2=patch_embed_2(x)# WHY: p1的结果作为输入，必须等p1做完p3=patch_embed_3(x)# WHY: p2做完才能开始p4=patch_embed_4(x)# WHY: 四个串行，总耗时是各分支耗时的累加result=concat([p1,p2,p3,p4])

ge通过依赖分析发现这四个patch embedding分支之间完全没有数据依赖，可以安全地并行执行。它生成一个并行调度计划，把这四个分支分配到昇腾NPU的四个核上同时执行。由于每个核只处理自己那份数据，总耗时从四个分支的串行时间总和变成最长单个分支的耗时。

# ge优化后的并行调度# ge生成并行调度计划，四个分支同时在昇腾NPU上执行# 每个核独立处理一个分支的embedding计算，总耗时接近最长路径# 这就是为什么ge的并行调度优化能显著提升多分支模型的吞吐量p1,p2,p3,p4=parallel_patch_embed(x)# WHY: 并行分配到多个核同时执行result=concat([p1,p2,p3,p4])

第七章：使用前后的效率对比

在并行调度方面，未优化的版本按照拓扑顺序串行执行所有算子，即使两个算子之间完全没有依赖关系也要等前一个完全结束才能开始下一个。ge的并行调度优化通过分析依赖图，把可以并行执行的算子分配到昇腾NPU的多核上同时执行。这个优化对于多分支模型或者含有大量独立计算的模型特别有效。加速比取决于昇腾NPU的可用核数以及可以并行化的算子比例——核数越多、模型中独立算子的比例越高，并行收益就越大。

GE（Graph Engine）是面向昇腾的图编译器和执行器，提供了计算图优化、多流并行、内存复用和模型下沉等技术手段，加速模型执行效率，减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力，并同时支持 onnx、pb 等主流模型格式的解析与编译。

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