CANN/ge：AscendIR图引擎中间表示设计

AscendIR — 图引擎的中间表示设计【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geGE 编译器消费的核心数据结构其设计质量直接决定了编译优化的能力和上限1. 整体架构四层对象模型AscendIR 的核心对象模型由四个层次构成ComputeGraph → Node → OpDesc → GeTensorDesc ↕ AnchorComputeGraph图的容器管理节点集合、输入/输出节点、子图、拓扑排序Node图中的算子节点持有 OpDesc 和一组 Anchor锚点OpDesc算子描述符定义算子的名称、类型、输入/输出张量描述、属性、推导函数等GeTensorDesc张量描述包含形状、数据类型、格式NCHW/NHWC 等、内存布局信息Anchor描述节点间的连接关系分为 DataAnchor 和 ControlAnchorAscendIR 的一个关键设计特征是图中不存在独立的 Edge边对象。边的关系完全由锚点之间的双向引用维护。2. 锚点系统连接关系的内嵌表达2.1 设计细节AscendIR 的 Anchor 类继承体系如下Anchor基类 ├── DataAnchor │ ├── InDataAnchor 输入数据锚点 │ └── OutDataAnchor 输出数据锚点 └── ControlAnchor ├── InControlAnchor 输入控制锚点 └── OutControlAnchor 输出控制锚点每个 Node 在初始化时NodeImpl::Init根据其 OpDesc 中定义的输入/输出数量创建对应数量的 InDataAnchor 和 OutDataAnchor。此外每个节点固定拥有一对InControlAnchor 和 OutControlAnchor索引为 -1。连接关系的维护方式每个 AnchorImpl 内部持有一个vectorweak_ptrAnchor peer_anchors_。当 OutDataAnchor 调用LinkTo(InDataAnchor)时两端互相将对方加入自己的 peer_anchors_ 列表。这是一个双向邻接表设计。关键约束InDataAnchor 只能有一个 peer单输入LinkFrom方法会检查 peer_anchors_ 是否为空OutDataAnchor 可以有多个 peer扇出一个输出可以连接多个下游输入ControlAnchor 可任意连接用于表达执行顺序依赖连接时支持跨类型连接OutDataAnchor 可以连接 InControlAnchor数据→控制依赖OutControlAnchor 也可以连接 InDataAnchor控制→数据依赖2.2 锚点方案的设计考量AscendIR 选择将连接关系内嵌于节点的锚点系统而非使用独立的边对象。对比独立边对象方案如果引入独立的 Edge 对象会增加一层间接性——遍历邻居需要 Edge → Anchor → Node 的两步跳转边的生命周期管理也更为复杂删边时需要同步更新两端的引用序列化时边的顺序维护是一个额外负担。锚点方案的优势O(1) 邻居访问从 InDataAnchor 直接获得唯一的 peer OutDataAnchorGetPeerOutAnchor从 OutDataAnchor 直接遍历所有 peer InDataAnchorGetPeerInDataAnchors无需全局查找。原子性连接/断开LinkTo和Unlink操作同时修改两端保证一致性。Insert和ReplacePeer方法支持在已有连接中插入新节点或替换对端。内存效率weak_ptr 避免循环引用anchor 本身是节点的组成部分而非独立对象减少了内存分配次数。图变换友好GE 编译器的大量 Pass融合、常量折叠、死代码消除需要频繁修改图结构。锚点系统使得断开旧连接、建立新连接操作非常局部化不需要全局重构。使用weak_ptr意味着每次访问都需要lock()操作但在图编译的语境下这一开销远小于简化图变换带来的收益。3. 图结构ComputeGraph 的设计3.1 核心数据结构ComputeGraphImpl 的核心成员std::listNodePtr nodes_节点列表使用 list 而非 vector支持频繁的中间插入和删除std::vectorNodePtr input_nodes_输入节点集合std::vectorpairNodePtr, int32_t output_nodes_info_输出节点及其输出索引std::vectorComputeGraphPtr sub_graph_子图集合std::mapstring, ComputeGraphPtr names_to_subgraph_名称→子图的映射weak_ptrComputeGraph parent_graph_/weak_ptrNode parent_node_父图和父节点引用3.2 节点管理添加节点AddNode创建 Node 对象时同时创建其所有 Anchor调用Node::InitInit 根据 OpDesc 的输入/输出数量为每个端口创建对应的 DataAnchor另外始终创建一对 ControlAnchor。节点被 push_back 到 nodes_ 列表末尾。删除节点RemoveNode这是一个复合操作——首先删除该节点关联的所有 Const 输入节点然后从 input_nodes_ 和 output_nodes_info_ 中移除再调用IsolateNode断开所有连边将上游输出直接连接到下游输入bypass 被删除的节点最后从 nodes_ 列表中移除。融合节点FuseNodeKeepTopo用于算子融合场景将多个原始节点替换为一组融合算子。插入位置选择在原始节点中 topo id 最小的位置同时继承原始节点的流标签、SuperKernel 属性和核数配置。3.3 子图管理AscendIR 支持嵌套子图这是实现控制流If/While/Case的基础子图通过AddSubGraph添加到父图的sub_graph_向量子图通过parent_graph_和parent_node_回指父图GetAllNodes方法会递归遍历所有子图节点通过 OpDesc 中的 subgraph_instance_names 建立关联子图中的 DATA 节点通过ATTR_NAME_PARENT_NODE_INDEX属性与父节点的输入端口对应子图中的 NETOUTPUT 节点汇总子图的输出子图不是一个独立的图而是挂在父节点上的附属结构。父节点如 If/While的 OpDesc 记录了子图实例名称subgraph_instance_names_Graph 通过名称映射找到对应的子图对象。3.4 拓扑排序AscendIR 提供了四种拓扑排序策略策略算法适用场景BFS广度优先训练场景默认DFS深度优先推理场景默认RDFS反向深度优先从输出节点回溯StableRDFS稳定反向DFS尽量保持原有顺序排序策略的选择通过OPTION_TOPOSORTING_MODE配置项控制训练默认 BFS推理默认 DFS。内存优先排序MemoryPriority当启用时拓扑排序会考虑节点的输出张量大小优先处理输出较大的节点通过NodeOutInfo结构比较输出大小和扇出数。这是为了优化内存复用——让大张量尽早释放。DelayTopoSort会将某些链式节点输入来自 variable/const 的链延迟到其消费者附近以提高内存局部性。排序后每个节点的id被重新设置为其在排序结果中的位置索引。排序还会检测环——如果排序后的节点数量不等于总节点数说明图中存在环路。4. 算子描述符OpDesc 的设计4.1 OpDesc 的双重身份OpDesc 是 AscendIR 中信息密度最高的对象同时承担了两种职责静态描述算子的输入/输出张量描述、名称映射、属性编译状态载体id、stream_id、input_offset、output_offset、workspace 等编译过程中逐步填充的信息OpDescImpl 的关键成员vectorGeTensorDescPtr inputs_desc_/outputs_desc_输入/输出张量描述mapstring, uint32_t input_name_idx_/output_name_idx_名称到索引的映射AttrStore attrs_属性存储基于 protobuf 的 AnyValueOpMetadata meta_data_元数据name, type, id, stream_id, offsets 等IRMetaData ir_meta_嵌套在 meta_data_ 中IR 注册信息输入/输出名称、子图 IR 名称等推导函数infer_func_、infer_format_func_、verifier_func_、infer_data_slice_func_4.2 输入/输出的双重索引OpDesc 对输入/输出提供了两种访问方式按索引和按名称。这在处理动态输入如可变数量的输入时尤为重要。动态输入的处理通过AddDynamicInputDesc实现为动态输入端口按顺序创建名为name0, name1, ...的输入描述。AddInputDescMiddle和AddInputDescForward支持在中间位置或头部插入新的输入描述同时更新名称到索引的映射关系。不同前端框架对同一算子的输入命名不同。通过UpdateInputName方法OpDesc 可以在运行时替换名称映射从 Factory 获取的标准名称覆盖 Parser 设置的名称而不改变实际的输入索引和描述。4.3 属性系统AscendIR 的属性系统基于AttrStore底层使用 protobuf 的mapstring, AttrDef存储。支持多种属性类型int、float、string、bool、list、tensor 等通过AttrUtils工具类提供类型安全的访问接口。属性可以附加在任何对象上——ComputeGraph、OpDesc、GeTensorDesc 都支持属性。这意味着编译器可以通过属性传递编译状态信息如内存偏移、流分配结果、融合标记等而无需修改 IR 的核心数据结构。4.4 IR 元数据与算子注册的解耦IRMetaData记录了算子的 IR 层面信息输入/输出的 IR 名称和类型固定/动态/可选、属性名称、子图 IR 名称和类型。这些信息在算子注册时由算子原型OpProto提供在创建 OpDesc 时被复制到实例中。这一层抽象使得 OpDesc 可以在不依赖算子定义仓库的情况下独立存在。算子定义在独立的 .so 文件中通过OpsProtoManager动态加载GE 只需要在编译时查找已注册的算子信息。5. 算子注册与工厂模式5.1 注册机制AscendIR 的算子注册采用自动注册模式算子定义仓库(.so) ──dlopen──→ OperatorFactoryImpl全局注册表 ↓ CreateOperator按类型查找 creator ↓ OpDesc实例化的算子描述OperatorFactoryImpl维护了一系列全局静态的shared_ptrmapstring, FuncTypeoperator_creators_v2_算子创建函数按类型名查找operator_infershape_funcs_形状推导函数operator_inferformat_funcs_格式推导函数operator_verify_funcs_校验函数算子在独立仓库中通过REG_OP宏注册编译为.so文件在运行时由OpsProtoManager加载。OpsProtoManager::Initialize从配置的ge.opsProtoLibPath路径加载所有算子.so调用其中的注册函数。5.2 算子定义独立于 GE 的设计考量算子定义被拆分到独立仓库如 ops-math、ops-transformer主要原因包括多场景共享同一个算子定义需要在多个场景使用——aclnn单算子直调、GE入图编译、调试工具。独立迭代算子的新增和修改频率远高于图引擎。将算子定义独立出来可以在不重新编译 GE 的情况下添加新算子只需替换 .so。插件化架构昇腾 NPU 产品线有不同的算子支持集合。通过动态加载算子 .so不同硬件配置可以加载不同的算子集合。离线编译支持GE 的编译功能可以在无设备的环境下运行离线编译算子 .so 的加载是延迟的不需要在编译时全部加载。算子注册信息的查找是运行时的 map 查找但相较于图编译的整体耗时这一开销可以忽略。5.3 注册的覆写机制注册覆写is_register_overridable允许在加载新的算子 .so 时同名算子的注册函数覆写已有注册。这在不重启进程的情况下更新算子定义时非常有用如热补丁场景。加载完成后覆写机制关闭防止运行时意外修改。6. 张量描述GeTensorDesc6.1 设计GeTensorDesc 描述一个张量的元信息ShapeGeShape张量的维度信息DataType数据类型支持从 FP32 到各种低精度格式如 FP8/FP4/INT4/INT2 等Format数据排布格式NCHW、NHWC、ND 等包含昇腾特有的 5D/6D 格式Origin信息origin_shape、origin_format、origin_data_type记录优化前的原始信息扩展元数据通过 ext_meta_size、weight_size、reuse_input、data_offset、device_type 等GeTensorDesc 使用了 Pimpl 模式通过GeTensorDescImpl使其可以在不影响 API 的情况下修改内部实现。6.2 序列化GeTensorDesc 的序列化通过GeTensorSerializeUtils::GeTensorDescAsProto实现将各字段写入 protobuf 的TensorDescriptor。序列化时需要处理扩展元数据直接映射到 proto 字段属性通过AttrGroupSerialize序列化到 proto 的 attr maporigin 信息通过特殊属性键存储在 attr map 中DataType 映射通过kDataTypeMap静态表实现包含 30 种数据类型的映射7. 图的序列化7.1 架构AscendIR 的序列化基于 protobuf核心 proto 定义在ge_ir.proto中ModelDef ├── GraphDef │ ├── OpDef[] 节点 │ │ ├── name, type, id, stream_id │ │ ├── input[] 前驱节点引用格式为 node_name:output_index │ │ ├── input_desc[], output_desc[] 张量描述 │ │ └── attr{} 属性映射 │ └── attr{} └── attr{}ModelSerializeImp负责序列化的核心逻辑SerializeOpDesc将 OpDesc 序列化为 OpDef protoSerializeEdge遍历节点的所有 InDataAnchor 和 InControlAnchor将前驱节点信息编码为 节点名:输出索引 格式的字符串反序列化时先创建所有节点再根据 input 字符串重建锚点连接7.2 序列化的设计选择边的序列化采用节点名称引用方式node_name:output_index而非序列化锚点本身。这意味着序列化后的数据是自包含的——不需要额外的边表反序列化时需要一次全局的名称查找来重建连接控制边通过索引 -1 标识node_name:-1这一设计在序列化效率和重建效率之间取得了平衡。对于 GE 的典型图规模数千到数万节点名称查找的开销是可接受的。7.3 属性序列化属性序列化采用类型分发模式AttrSerializerRegistry注册了每种属性类型int、float、string、tensor、graph 等的序列化器。每种类型有独立的 Serializer 类如 IntSerializer、FloatSerializer、TensorSerializer实现Serialize(AnyValue → AttrDef)和Deserialize(AttrDef → AnyValue)方法。这种设计支持扩展——新的属性类型只需注册新的序列化器不需要修改序列化框架。8. 图工具类graph/utils/目录提供了丰富的图操作工具GraphUtils图级别的操作AddEdge、RemoveEdge、ReplaceEdgeSrc、CopyGraph、DumpGraph 等NodeUtils节点级别的操作GetInDataNodes、MoveOutputEdges、IsAnchorStatusSet 等OpDescUtilsOpDesc 操作SetSubgraphInstanceName、ClearInDesc 等TensorUtils张量操作CalcTensorMemSize 等AnchorUtils锚点操作GetStatus、SetStatus、GetIdxTypeUtils类型转换DataType ↔ 字符串、Format ↔ 字符串这些工具类大多以静态方法的形式提供采用函数式风格不持有状态。这是 AscendIR 设计中的一个惯用模式核心对象Graph/Node/OpDesc/Anchor提供数据访问接口工具类提供复合操作。9. AscendIR 设计的核心设计原则总结锚点优于边连接关系内嵌于节点避免独立边的生命周期管理复杂性静态图 属性扩展核心结构是静态的 DAG动态信息通过属性系统附加算子与引擎解耦算子定义通过动态加载注册GE 只消费注册结果Pimpl 隔离所有核心类使用 Impl 类分离接口和实现支持 ABI 稳定weak_ptr 防环Anchor 和 Node 对父图的引用使用 weak_ptr避免循环引用导致的内存泄漏名称 索引双寻址OpDesc 的输入输出同时支持按名称和按索引访问适配不同前端框架的习惯理解了 AscendIR 的图结构、锚点系统、算子注册和序列化机制后读者已经掌握了 IR 中哪些信息是编译器需要消费的——节点拓扑、张量描述、属性、推导函数。下一个模块将揭示这些信息进入编译器后经过哪些阶段处理从图合法性校验、形状推导、算子融合到内存规划和流分配。【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考