CANN ge:图编译引擎在 CANN 架构中的角色
个人主页ujainu文章目录前言ge 的定位不只是翻译器为什么需要图编译三阶段流水线拆解Stage 1图准备——把图洗干净Stage 2图优化——把图压紧实Stage 3图编译——把图落成铁三个关键机制算子融合减少搬运才是真加速内存复用一张 7B 模型省掉 40% 显存模型下沉让 NPU 自己管执行在 CANN 架构中的位置与 metadef 的关系词汇表与编译器与 runtime 的关系编译器与执行器与 TorchAir / TFA 的关系前端接入路径与其他 CANN 仓库的关系graph-autofusion算子自动融合框架AscendCL统一推理接口Ascend C算子开发语言一条完整的编译链路前言同一个 ResNet-50 模型同样的昇腾 NPU 硬件逐算子 eager 执行跑一遍 18ms经过 ge 图优化后只需要 3.2ms。算子没换、硬件没变差距从哪来三个字全局视野。逐算子执行就像每做一道菜洗一次锅——每个动作单独看都没问题全局看全是浪费。geGraph Engine做的就是在执行前把整个计算图看一遍该删的删掉、该合的合并、该复用的复用。它是昇腾 CANN 五层架构中第 3 层编译层的核心组件专门负责把前端框架交过来的计算图变成 NPU 真正能高效跑的执行计划。仓库地址https://atomgit.com/cann/gege 的定位不只是翻译器很多人把 ge 当成框架适配层——PyTorch 出图、ge 接图、runtime 跑图。这个理解只对了一半。ge 更准确的定位是计算图编译器。它不是被动接收图然后原样下发而是对图做编译级优化常量折叠、死节点消除、算子融合、内存复用规划、多流并行编排。这些优化单独拿出一个都不新鲜但 ge 的核心价值在于把它们串成一条三阶段流水线让优化之间不互相打架。为什么需要图编译如果 NPU 的执行模式像 CPU 那样逐指令发射逐算子调度就够了。但昇腾 NPU 的硬件特性决定了它需要全局编译Cube/Vector 双引擎并行NPU 内部有 Cube矩阵乘和 Vector向量运算两套计算单元两者可以同时工作。逐算子调度时Cube 算完一个 MatMul 闲着等 Vector 算完 Add而图编译可以在编译期分析数据依赖把 Cube 和 Vector 的任务编排到不同流水级让双引擎持续满载。HBM 访存瓶颈NPU 的算力远超 HBM 带宽一次 HBM 读写的时间可以完成多次片上计算。逐算子模式下每个中间张量都写回 HBM图编译则通过算子融合让中间数据留在片上存储L1/L0减少搬运次数。Host-Device 调度开销每次 Host 向 Device 下发一条算子指令都要经过驱动调用和 PCIe 传输。模型越大、算子越多这个开销越不可忽视。图编译把整图编译成一条 Task 链一次性下发消除逐算子的 Host-Device 往返。// 逐算子执行每个算子一次 Host→Device 调度for(autoop:graph.ops()){host_launch(op);// 每次都有驱动调用 PCIe 传输开销}// 图编译后整图一次下发Device 自主执行autotask_chainge_compile(graph);// 编译期完成所有优化device_execute(task_chain);// 一次下发NPU 按预编排顺序跑完这就解释了为什么 CANN 需要一个图编译引擎——不是有更好而是没有就不行。NPU 的硬件红利只有通过全局编译才能兑现。三阶段流水线拆解ge 把整个编译流程分为三个阶段每个阶段负责不同的优化维度原始计算图来自 PyTorch / TensorFlow / ONNX / PB │ ▼ ┌─────────────────────────────┐ │ Stage 1: 图准备 │ │ · 形状与数据类型推导 │ │ · 常量折叠 │ │ · 死边消除删掉无用计算 │ │ · 图合法性校验 │ └──────────────┬──────────────┘ ▼ ┌─────────────────────────────┐ │ Stage 2: 图优化 │ │ · 算子融合多算子→单算子 │ │ · 图切分SubGraph 划分 │ │ · 流水编排重新排列执行序│ └──────────────┬──────────────┘ ▼ ┌─────────────────────────────┐ │ Stage 3: 图编译 │ │ · 整图内存复用规划 │ │ · 连续内存分配 │ │ · Task 生成与下发 │ │ · 模型下沉Device端执行 │ └──────────────┬──────────────┘ ▼ 可执行模型交给 runtime准备阶段做删优化阶段做合编译阶段做排。为什么一定要分三步而不是一步到位因为融合改变了图的拓扑结构而内存复用规划需要基于融合后的图做全局分配。如果合并成一步要么丢掉融合机会要么内存规划不准确。Stage 1图准备——把图洗干净图准备阶段的核心任务是让后续优化基于一个干净、确定的图工作。前端框架导出的计算图经常携带大量冗余信息训练时留下的梯度节点、只走一个分支的 Switch/Select、固定 shape 下的动态逻辑……这些如果不清理后续优化就无法展开。形状与数据类型推导是准备阶段的基础工作。ge 从输入张量的 shape 和 dtype 出发沿数据流方向逐算子推导每个节点的输出形状。这个信息是后续几乎所有优化——融合规则匹配、内存规划、Task 切分——的前提// ge 形状推导示意伪代码StatusInferShape(ComputeGraphgraph){for(autonode:graph.GetNodesInTopoOrder()){autoop_descnode.GetOpDesc();// 根据算子类型和输入 shape 推导输出 shapeautoretShapeInference::Infer(op_desc);if(ret!SUCCESS){// 无法静态推导的 shape 标记为动态走动态图分支op_desc-SetDynamicShape(true);}}returnSUCCESS;}常量折叠把编译期就能算出来的常量表达式预计算好。例如Add(Const(1), Const(2))直接替换为Const(3)消灭两个节点和一条边。在训练转推理的场景下BatchNorm 的 running_mean/running_var 是固定常量可以折叠进 Conv 的权重从而省掉 BN 算子# 折叠前Conv → BN → ...BN 在运行时计算conv_outconv(input,weight)bn_outbatchnorm(conv_out,running_mean,running_var,gamma,beta)# 折叠后Conv(fused_weight, fused_bias) → ...BN 参数融入 Conv 权重# fused_weight weight * gamma / sqrt(running_var eps)# fused_bias beta - running_mean * gamma / sqrt(running_var eps)fused_outconv(input,fused_weight,fused_bias)死边消除从图的输出节点反向追溯删掉所有不可达的节点和边。训练图中的梯度子图、Dropout推理时是 Identity等都是典型的清除对象。Stage 2图优化——把图压紧实经过准备阶段清洗后的图已经是确定且无冗余的接下来进入优化阶段。这个阶段的核心目标是减少执行时的算子数量和数据搬运次数。算子融合是最核心的优化手段。ge 内建了一套融合规则库通过模式匹配识别可融合的算子组合将其替换为单个融合算子。经典模式包括 Conv→BN→ReLU、MatMul→Add→GELU、Reduce→Reshape 等。融合的收益不在计算量——计算量不变——而在省掉中间张量的 HBM 写入和读取。图切分针对异构执行场景某些算子必须在 Host 侧执行如文件 I/O、动态 shape 控制某些算子适合在 Device 侧执行。ge 根据算子的执行属性和 SoC 能力将整图切分为多个 SubGraph每个 SubGraph 内部连续执行SubGraph 之间做必要的 Host-Device 同步。流水编排在保持语义等价的前提下重新排列算子执行顺序目标是最大化 NPU 双引擎的并行度。例如把不依赖当前 Cube 结果的 Vector 算子提前与当前 Cube 算子形成流水并行# 编排前串行执行Cube 和 Vector 交替空闲 Cube[MatMul1] → Vector[Add1] → Cube[MatMul2] → Vector[Add2] # 编排后流水并行Cube 和 Vector 同时工作 时间轴 ─────────────────────────────────────► Cube: [MatMul1]──────[MatMul2]──────[MatMul3] Vector: [Add1]──────[Add2]──────[Add3]Stage 3图编译——把图落成铁优化阶段输出的是一个优化后的逻辑图但它还不能直接在 NPU 上执行。图编译阶段的任务是把逻辑图变成物理可执行的 Task 序列。整图内存复用规划是编译阶段最关键的步骤。ge 对融合后的完整图做全局生命周期分析找出所有张量的起止时间点然后求解一个最优的内存分配方案——生命周期不重叠的张量共享同一块显存。这本质上是一个区间图着色问题ge 采用启发式算法在编译期求解// ge 内存复用规划示意伪代码MemoryPlanPlanMemory(ComputeGraphgraph){// Step 1: 计算每个张量的生命周期 [start_op, end_op]autolifetimesAnalyzeLifetime(graph);// Step 2: 生命周期不重叠的张量分到同一组autogroupsIntervalColoring(lifetimes);// Step 3: 每组分配一块共享内存大小 组内最大张量MemoryPlan plan;for(autogroup:groups){size_t block_sizeMax(group.tensor_sizes);plan.AllocBlock(block_size,group.tensor_ids);}returnplan;}连续内存分配进一步优化将多个小张量拼接成一块连续内存减少内存碎片和分配次数。对于动态 batch 场景ge 支持预分配一块足够大的内存池运行时从池中切分避免频繁调用驱动层内存分配接口。Task 生成将图中的每个算子节点翻译为 NPU 可执行的 Task 描述符包含算子类型、输入输出内存地址由内存规划步骤确定、执行流编号等信息。这些 Task 按拓扑序排列构成一条可顺序执行的 Task 链。模型下沉是编译阶段的收尾工作——把整条 Task 链打包下沉到 Device 端NPU 自主循环执行Host 不再逐帧干预。三个关键机制从三阶段流水线中有三个机制值得单独展开。算子融合减少搬运才是真加速ge 的算子融合不是简单的把两个算子绑在一起跑而是将多个算子合并成一个融合算子中间数据不落回内存。以经典的 Conv→BN→ReLU 为例# 融合前3 次算子调用中间结果写回内存output1conv(input,weight)output2batchnorm(output1)output3relu(output2)# 融合后1 次算子调用中间结果留在 NPU 片上存储outputfused_conv_bn_relu(input,weight,bn_param)融合的关键收益不在计算量在于省掉了两次 HBM 写入。NPU 的计算能力远超访存带宽瓶颈往往在搬运而非计算。ge 会自动识别可融合的算子模式并执行融合这个过程不依赖用户干预。但 ge 自身的融合能力有限——更复杂的融合场景由 graph-autofusion 算子自动融合框架处理。ge 负责标准模式匹配graph-autofusion 负责 JIT 编译级的动态融合两者配合覆盖了从静态到动态的完整融合需求。内存复用一张 7B 模型省掉 40% 显存大模型推理中显存往往比算力更紧张。ge 的内存复用机制在图编译阶段做全局分析两个张量的生命周期不重叠就分配同一块显存。# ge 内部的生命周期分析示意伪代码# tensor_a 在第 3 层计算后不再使用# tensor_b 在第 5 层才开始使用# → tensor_b 可以复用 tensor_a 的显存# 不做内存复用每层独立分配峰值显存 Σ(所有层)# 做内存复用峰值显存 max(同时在用层的显存总和)对于 7B 参数的模型这种全局复用通常能减少 40% 左右的峰值显存占用。这个优化发生在 Stage 3 图编译阶段因为它需要知道融合后的完整图拓扑才能准确判断生命周期。模型下沉让 NPU 自己管执行传统推理流程中Host CPU 每帧都下发算子调度指令开销随模型规模线性增长。ge 的模型下沉机制将整个计算图编译成一条预定义的 Task 链一次性下发到 Device 端NPU 自己循环执行Host 不再逐帧干预。# 不下沉Host 逐帧调度 for frame in frames: for op in graph: host_schedule(op) # 每次 Host→Device 往返有延迟 # 下沉后Device 自主执行 device_execute(precompiled_task_chain) # 一次下发NPU 自己跑这在大模型推理场景下尤其关键——token 生成阶段的延迟对用户体验影响直接。下沉后每次 token 生成的 Host 开销从数百微秒降到接近零循环推理的帧间抖动也大幅降低。在 CANN 架构中的位置ge 处于 CANN 五层架构的第 3 层编译层向上承接框架适配向下对接运行时执行第1层 AscendCL统一接口 │ ▼ 第2层 Framework Adaptor框架适配 │ TorchAir / TFA / Triton GE Backend ▼ 第3层 ge图编译引擎 ← 你在这里 │ ├── 上游metadef算子原语定义、图 IR 数据结构 │ ├── 横向graph-autofusion算子自动融合框架 │ └── 下游runtime任务调度与执行与 metadef 的关系词汇表与编译器metadef 为 ge 提供算子原语定义和图 IR 的基础数据结构——可以理解为 ge 的词汇表。metadef 定义了每一个算子的名称、输入输出规格、属性列表和 shape 推导规则。ge 在编译过程中对图的每一步操作都依赖 metadef 的定义形状推导需要查阅算子的 InferShape 函数融合规则匹配需要识别算子的类型和属性合法性校验需要检查输入输出的数据类型是否合规。// metadef 中的算子定义示意伪代码REG_OP(MatMul).Input(x1,TensorType({DT_FLOAT,DT_FLOAT16})).Input(x2,TensorType({DT_FLOAT,DT_FLOAT16})).Output(y,TensorType({DT_FLOAT,DT_FLOAT16})).Attr(transpose_a,AttrType::BOOL,false).Attr(transpose_b,AttrType::BOOL,false).InferShape(MatMulInferShape)// shape 推导函数.Register();没有 metadefge 无法识别图中的算子类型、属性和数据格式整个编译流水线无从启动。与 runtime 的关系编译器与执行器runtime 接收 ge 编译后的 Task 链负责真正的设备侧执行调度。ge 和 runtime 的分工是经典的编译时 vs 运行时分离ge 在编译时做所有可以静态确定的决策融合、内存规划、执行序编排runtime 在运行时处理动态逻辑流调度、事件同步、异常处理。// ge 编译输出 → runtime 执行autocompiled_modelge::Compile(graph,options);// compiled_model 包含Task 链 内存分配方案 流编排信息// runtime 加载并执行autosessionruntime::LoadModel(compiled_model);autooutputsession.Execute(input_tensors);这种分离的好处是运行时的执行路径尽可能短——不需要再遍历图、不需要再判断融合、不需要再分配内存只需要按 Task 链顺序执行即可。与 TorchAir / TFA 的关系前端接入路径框架侧TorchAir 是 PyTorch 图模式接入 ge 的路径TFA 对应 TensorFlowtriton-inference-server-ge-backend 则是 Triton 推理服务在昇腾 NPU 上的部署方案。三条路径最终都汇入 ge 的编译流水线。TorchAir 的核心工作是把 PyTorch 的动态图通过torch.compile或torch.export导出为静态计算图然后转换成 ge 可接受的 Graph IR 格式。这个转换过程需要处理 PyTorch 算子到昇腾算子的映射、动态 shape 的处理、以及 Python 控制流的图化# TorchAir 接入 ge 的典型用法importtorchimporttorch_npu# 昇腾 PyTorch 扩展importtorchairastng# 配置 ge 编译选项configtng.CompilerConfig()config.experimental.fusion_enableTrue# 开启算子融合config.experimental.memory_optimizationTrue# 开启内存复用# 将 PyTorch 模型接入 ge 编译npu_backendtng.get_npu_backend(compiler_configconfig)modeltorch.compile(model,backendnpu_backend)# 推理时自动走 ge 编译 runtime 执行outputmodel(input_tensor)TFATensorFlow Adapter做的是同样的事只是源框架从 PyTorch 换成了 TensorFlow。三条接入路径在 ge 内部汇聚到同一个编译流水线前端框架的差异在 ge 入口处就被抹平了。与其他 CANN 仓库的关系ge 不是孤立工作的它与多个 CANN 仓库有上下游依赖。graph-autofusion算子自动融合框架graph-autofusion 是 CANN 的算子自动融合框架负责 ge 标准融合规则无法覆盖的复杂融合场景。ge 内建的融合规则是白名单式的——只匹配预定义的算子组合模式。当模型中出现新的、未被规则库收录的可融合模式时ge 的融合就到头了。graph-autofusion 通过 JIT 编译技术在运行时动态分析算子间的数据流和计算特征自动生成融合算子的 Ascend C 实现代码。ge 和 graph-autofusion 的配合模式是ge 先用规则匹配做一轮静态融合剩余未融合的算子交由 graph-autofusion 尝试动态融合。AscendCL统一推理接口AscendCLAscend Computing Language是 CANN 的统一接口层位于五层架构最顶部。用户通过 AscendCL 的 API 加载 ge 编译后的离线模型、管理输入输出内存、执行推理。AscendCL 本身不参与编译但它是 ge 编译产物的消费者——ge 输出的.om离线模型文件由 AscendCL 的aclmdlLoadFromFile接口加载执行// AscendCL 加载 ge 编译的离线模型aclError ret;retaclmdlLoadFromFile(resnet50_npu.om,model_id);retaclmdlExecute(model_id,input_dataset,output_dataset);Ascend C算子开发语言Ascend C 是昇腾 NPU 的算子开发语言ge 编译流水线中使用的融合算子和自定义算子最终都用 Ascend C 编写。ge 的算子融合规则库中的融合算子、graph-autofusion 自动生成的融合算子其底层实现都是 Ascend C 代码。ge 在编译期并不关心算子的具体实现语言只需要知道算子的输入输出规格和 shape 推导规则这些由 metadef 定义。但到了运行时Task 链中的每个算子最终都会调用对应的 Ascend C kernel 完成计算。一条完整的编译链路把前面的内容串起来看一个 PyTorch 模型经过 ge 编译的完整过程# 第 1 步PyTorch 模型导出python export.py--modelresnet50--outputresnet50.onnx# 第 2 步ATC 工具调用 ge 进行离线编译# ATC 是 CANN 的模型转换工具底层调用 ge 的编译流水线atc--modelresnet50.onnx\--framework5\--outputresnet50_npu\--soc_versionAscend910ATC 内部会依次调用 ge 的三阶段流水线Stage 1 对 ONNX 图做形状推导和常量折叠Stage 2 执行 Conv→BN→ReLU 等标准融合并切分子图Stage 3 做内存复用规划和 Task 生成。最终输出一个 NPU 可直接加载执行的离线模型.om文件。如果想观察 ge 每个阶段的编译细节可以开启编译日志# 开启 ge 编译详细日志exportASCEND_GLOBAL_LOG_LEVEL0# DEBUG 级别exportASCEND_MODULE_LOG_LEVEL_SETGE:0# ge 模块详细日志atc--modelresnet50.onnx\--framework5\--outputresnet50_npu\--soc_versionAscend910# 日志中可以看到# - Stage 1: 常量折叠消除了多少节点、shape 推导结果# - Stage 2: 哪些融合规则命中、融合前后算子数量变化# - Stage 3: 内存复用方案、Task 数量这条链路的可调参数很多——融合策略开关、内存复用模式重计算 vs 重存储、流数量配置等。具体参数需要根据模型结构和硬件配置针对性调整但默认配置对大多数 CV 和 Transformer 模型已经能给出不错的基线性能。如果你手上有现成的 PyTorch 模型想快速验证 ge 的优化效果直接用 ATC 转换后对比推理耗时是最直接的方式。想深入了解某个子阶段的优化策略ge 仓库的源码和编译日志是最可靠的一手资料https://atomgit.com/cann/ge