GE 图拆分Graph Split特性分析【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 特性背景1.1 问题域昇腾 AI 处理器上执行深度学习推理/训练时计算图中的算子特性并非整齐划一。实际业务场景中一张图里常常同时包含以下异构元素静态 Shape 算子输入输出 Shape 在编译期完全确定如卷积、全连接可以预先分配内存、静态调度 kernel执行效率最高。动态 Shape 算子Shape 中存在未知维度-1 或 -2编译期无法确定精确内存布局需要运行时动态计算 tiling 参数、动态分配 workspace。Host 侧算子必须在 Host CPU 上执行的算子如某些控制流操作无法下沉到 Device 侧。不同引擎算子分属 AI Core、AI CPU、DVPP 等不同硬件引擎各自有独立的编译和调度路径。如果将整张图交给单一执行器处理要么所有算子退化为动态执行模式牺牲静态算子的性能要么无法处理动态算子功能不可用。因此GE 需要一种机制在编译期将整图按执行语义拆分为多个子图让每个子图进入最适合的执行路径。1.2 设计目标图拆分模块的核心目标是回答每个节点/路径应进入哪个执行器的问题。它在编译流水线中处于图优化之后、算子编译和内存规划之前是连接高层图优化与底层执行调度的桥梁。拆分结果的优劣直接影响后续所有阶段的正确性和性能。2. 用户使用场景2.1 动态分档场景用户通过 atc 命令行工具编译模型时可以指定多档 ShapeGE 自动将图拆分为公共入口 Case 分支子图结构atc --modelresnet.onnx \ --dynamic_batch_size1,4,8,16 \ --outputresnet_dyn或指定更灵活的动态维度组合atc --modelbert.onnx \ --input_shapeinput:1,-1,128 \ --dynamic_dims1,32;1,64;1,128 \ --outputbert_dyn编译结果是一个 OM 模型运行时根据实际输入的 Shape 自动选择对应分支执行。2.2 动静态混合场景模型中部分算子的 Shape 依赖运行时计算结果如NonMaxSuppression的输出数量取决于输入内容和阈值而其他算子 Shape 静态。GE 自动将这类图拆分为静态子图和动态子图静态子图享受预编译 kernel 和静态内存规划的性能优势动态子图走运行时 tiling 和动态调度路径。2.3 流水线并行场景Stage Partition对于大型模型用户可通过ATTR_STAGE_LEVEL属性将算子标记到不同流水线阶段。GE 按阶段拆分子图每个阶段独立编译执行阶段间通过同步点协调。2.4 JIT 增量编译场景在线模式下如通过 TorchAirGE 支持增量式图拆分。随着符号推理逐步完成已确定 Shape 的子图先编译执行未确定部分保留等待后续输入信息实现逐层剥洋葱式的编译-执行交替。3. 编译流水线中的位置图拆分在GraphManager::OptimizeSubgraph()中按以下顺序执行位于编译流水线的中段StagePartition → EnginePlacer1 → HostcpuEngineUpdatePass → DynamicShapePartition EnginePlacer2 → CompositeEnginePartition 子图优化 Merge → AtomicEnginePartition 子图优化 Merge对应代码入口compiler/graph/manager/graph_manager.cc的OptimizeSubgraph()方法。每一步的职责阶段执行器职责StagePartitionStagePartitioner按流水线阶段拆分EnginePlacer1EnginePlacer为所有节点分配初始引擎HostcpuEngineUpdatePassEnginePlacer提前标记 Host CPU 引擎节点DynamicShapePartitionDynamicShapePartitioner按动静 Shape 拆分生成 PartitionedCall 子图EnginePlacer2EnginePlacer拆分后重新分配引擎CompositeEnginePartitionEnginePartitioner按组合引擎拆分子图并优化AtomicEnginePartitionEnginePartitioner按原子引擎拆分子图并优化4. 对外接口4.1 atc 命令行选项通过api/atc/main_impl.cc暴露三个动态 Shape 选项互斥选项含义示例--dynamic_batch_size动态批大小多档以逗号分隔1,4,8,16--dynamic_image_size动态图像尺寸不同组以分号分隔组内维度以逗号分隔224,224;256,256;512,512--dynamic_dims通用动态维度不同档以分号分隔1,32;1,64;1,1284.2 运行时配置选项选项默认值含义ge.exec.static_model_ops_lower_limit4ffts 场景为 6静态子图最少算子数阈值低于此值的静态子图降级为动态。设为 -1 可将所有子图合并为动态图ge.topoSortingMode默认设为3启用稳定 RDFS 排序会改变 cluster 合并策略ge.tiling_schedule_optimize0设为1启用 tiling 下沉在 AICPU 上执行 tilingge.host_scheduling_max_threshold0静态图节点数低于此阈值时整体走动态执行4.3 图属性接口拆分结果通过图/节点属性传递给下游模块属性作用域含义_dynamic_shape_partitioned图级标识该图是否经过动态 Shape 拆分_force_unknown_shape节点级强制将节点归入动态子图_is_unknown_shape节点级标记节点的动态/静态属性ATTR_STAGE_LEVEL节点级流水线阶段编号ATTR_NAME_MEMORY_DISCONTIGUOUS_ALLOCATION图级启用非连续内存分配动态子图4.4 Python API通过api/python/ge/ge_api_c_wrapper/c_graph.cc暴露子图操作接口GeApiWrapper_Graph_GetAllSubgraphs()— 获取所有子图GeApiWrapper_Graph_GetSubGraph()— 按名称获取子图GeApiWrapper_Graph_AddSubGraph()— 添加子图GeApiWrapper_Graph_RemoveSubgraph()— 删除子图5. 具体实现5.1 基础框架BasePartitioner BaseCluster图拆分的基础框架由compiler/graph/partition/base_partitioner.h中的BasePartitioner和compiler/graph/partition/base_cluster.h中的BaseCluster构成。所有具体拆分策略都继承此框架。5.1.1 拆分流水线BasePartitioner::PartitionImpl()定义了统一的拆分流程InitClusters → MergeClusters → ProcessUniqueClusters → BuildPartitionFrame → CombinePartitionFrame → BuildPartitionSubgraphInitClusters为每个节点创建一个独立的 cluster按策略分类DATA / KNOWN_SHAPE / UNKNOWN_SHAPE / NETOUTPUT 等。MergeClusters根据特定规则合并相邻 cluster减少子图数量。ProcessUniqueClusters去重、清理合并后的 cluster 集合。BuildPartitionFrame为每个 cluster 在根图中创建一个PartitionedCall节点并将 cluster 内节点移入对应子图。CombinePartitionFrame在PartitionedCall节点之间建立数据边。BuildPartitionSubgraph在子图内部添加 InnerData / InnerNetOutput 节点完成 IO 连接。5.1.2 Cluster 数据结构BaseCluster的核心字段type_index_类型索引标识 cluster 类别DATA0, NETOUTPUT1, INPUT_NODE2, STAGE3, KNOWN_SHAPE4, UNKNOWN_SHAPE5min_/max_cluster 内节点的拓扑序范围用于合并判断in_clusters_/out_clusters_入/出 cluster 的邻接关系nodes_cluster 包含的节点集合subgraph_cluster 对应的子图ComputeGraphpartition_node_cluster 在根图中对应的PartitionedCall节点关键合并操作Merge()— 无条件合并将另一个 cluster 的所有节点和邻接关系吸收进来TryMerge()— 仅在不会形成环时合并通过前向可达性检查MergeAllPathFrom()— 合并两个 cluster 之间所有路径上的 cluster双向 BFS 找路径交集5.1.3 属性传递机制PartitionNodeAttrNameManager管理需要在拆分前后传递的节点属性注册表。通过REGISTER_PARTITION_ATTR_NAME宏注册属性拆分时自动在PartitionedCall节点和子图内部节点间复制这些属性确保拆分后语义一致。5.2 动静 Shape 拆分DynamicShapePartitionercompiler/graph/partition/dynamic_shape_partition.h中的DynamicShapePartitioner是图拆分的核心策略实现负责将计算图按动态/静态 Shape 划分为不同子图。5.2.1 节点分类规则MarkUnknownShapeNodes()方法按以下规则判断节点是否属于动态 Shape动态 Shape 算子Tensor Shape 中存在 -1未知维度或 -2未知 rankForce Unknown 标记节点被设置了_force_unknown_shapetrueTiling 依赖不支持下沉节点存在动态 tiling 依赖但不支持在 AICPU 上执行 tiling地址刷新不支持节点_is_support_addr_refreshfalseHost CPU 引擎节点属于DNN_VM_HOST_CPU引擎子图传播若节点的子图控制流子图中存在动态 Shape 算子则该节点也归为动态5.2.2 Cluster 合并策略DynamicShapeCluster继承BaseCluster按类型分为KNOWN_SHAPEtype_index4和UNKNOWN_SHAPEtype_index5。MergeClustersNormal()的合并顺序动态路径吸附遍历所有UNKNOWN_SHAPEcluster若两个动态 cluster 之间存在路径则将路径上所有 cluster 合并为动态。这保证了动态链路的连续性。静态单路径合并遍历KNOWN_SHAPEcluster若两个静态 cluster 之间仅存在唯一路径无环则合并。小 cluster 降级节点数低于ge.exec.static_model_ops_lower_limit阈值的静态 cluster 降级为动态避免产生过小的静态子图碎片。控制流合并属于同一ATTR_NAME_CONTROL_FLOW_GROUP的控制流节点如 StreamActive、StreamSwitch合并到同一 cluster。RefVariable 合并引用类型的 Variable 节点与其消费者合并到同一 cluster。5.2.3 重拆分机制初次拆分后DynamicDataFlowPartitionerPass检查数据流算子Stack/StackPush/StackPop/StackClose是否横跨动静子图。如果是将这些算子强制标记为_force_unknown_shapetrue然后调用ReDynamicShapePartitioner()重新拆分。这个迭代过程确保数据流状态在执行语义上的一致性。5.2.4 整图动态判定IsGraphNeedUnknownShapePartition()判断整图是否需要走动态拆分流程。如果图中没有任何动态 Shape 节点则设置_dynamic_shape_partitionedfalse整图走静态编译路径。如果图中静态节点数很少低于ge.host_scheduling_max_threshold则整图直接走 Host 调度模式。5.3 引擎级拆分EnginePartitionercompiler/graph/partition/engine_partitioner.h中的EnginePartitioner负责按引擎归属拆分子图。它在动静拆分之后执行进一步将子图按 AI Core、AI CPU、DVPP 等不同硬件引擎切分。5.3.1 拆分流程Initialize通过EnginePlacer为每个节点分配引擎创建初始 cluster每个节点一个携带引擎名和 stream label。MarkClusters遍历 cluster 对若两个 cluster 具有相同引擎 相同 stream label 之间无第二条路径则合并。SplitSubGraphs为每个合并后的 cluster 创建ComputeGraph子图在不同引擎子图之间插入PlaceHolder/End节点对。SortSubGraphs对子图拓扑排序将 Data 节点合并到统一的输入子图。5.3.2 PlaceHolder / End 机制与DynamicShapePartitioner使用PartitionedCall不同EnginePartitioner使用PlaceHolder/End节点对作为跨子图的数据传递桥梁End 节点位于源子图中标记子图的输出边界携带源节点和输出索引信息。PlaceHolder 节点位于目标子图中标记子图的输入边界通过peer_index属性与对应的 End 节点配对。子图优化完成后MergeAfterSubGraphOptimization()方法移除所有 PlaceHolder/End 节点对将子图重新合并为完整的计算图。5.3.3 两种拆分模式EnginePartitioner支持两种拆分模式CompositeEnginePartitioning按组合引擎拆分粒度较粗用于大型引擎级别的分离。AtomicEnginePartitioning按原子引擎拆分粒度更细用于更精确的引擎隔离。5.4 流水线阶段拆分StagePartitionercompiler/graph/partition/stage_partitioner.h中的StagePartitioner按节点上的ATTR_STAGE_LEVEL属性将计算图拆分为多个流水线阶段。拆分逻辑SplitStageLevel()收集带有ATTR_STAGE_LEVEL属性的节点将属性向上游传播。SplitTailStage()将未标记阶段的节点归入最后一个阶段。StagePartition()使用GraphUtils::BuildSubgraphWithNodes()将每个阶段的节点封装为子图阶段间通过PartitionedCall节点连接。每个阶段的父节点会被设置_force_unknown_shapetrue确保阶段间的同步在运行时处理。5.5 多档 CloneMultiBatchClonePasscompiler/graph/passes/multi_batch/multi_batch_clone_pass.h中的MultiBatchClonePass处理动态分档场景。它将原始图 Clone 为多份每份对应一个档位的 Shape然后通过 Case 节点在运行时选择分支。5.5.1 构建流程CollectIoNodes收集原始图的输入/输出节点解析用户指定的动态 Shape 参数。CreateRootGraph创建根图包含Shape 索引节点Data 或GetDynamicDims带 max shape 的输入 Data 节点Case 节点CreateSubgraphs将原始图 Clone N 份N 为档位数每份使用对应档位的静态 Shape作为 Case 节点的分支子图。PruneDirectOutput清理冗余的直连输出。5.5.2 Scope 子图创建compiler/graph/passes/multi_batch/create_subgraph_with_scope_pass.h中的CreateSubGraphWithScopePass用于多维度场景。它将具有相同ATTR_NAME_OP_MULTI_DIMS_INPUT_DIMS属性的节点封装到新的PartitionedCall子图中实现按 scope 粒度的子图划分。5.6 Variable 拆分入子图SplitVariableIntoSubgraphPasscompiler/graph/passes/variable_optimize/split_variable_into_subgraph_pass.h中的SplitVariableIntoSubgraphPass处理 Variable/RefData 节点与控制流子图If/Case/PartitionedCall/While的交互。对于需要被子图内部访问的 Variable 节点将其拷贝到子图内部确保子图可以独立访问权重数据。对于 While 节点由于循环语义的特殊性改为添加控制边而非拷贝。5.7 JIT 增量拆分BinaryPartitionerapi/session/jit_execution/utils/partitioner/binary_partitioner.h中的BinaryPartitioner用于在线 JIT 编译场景将图按符号推理结果拆分为已完成推理和未完成推理两部分。5.7.1 拆分逻辑Partition()方法接收一组已完成符号推理的节点将图拆分为sliced_graph包含已推理节点可以立即编译执行。remaining_graph包含未推理节点保留等待后续输入。CheckNodesContainsCycle()验证已推理节点集合不依赖未推理节点的输入确保拆分合法。BinaryGraphBuilder负责构建两个子图建立 IO 映射BinaryGraphIOLinkage处理输入节点替换和去重。5.7.2 执行点管理api/session/jit_execution/exe_points/execution_order.h中的ExecutionOrder管理一系列ExecutionPoint执行点每个执行点对应一个已编译的子图切片。它通过AddNewSlice()方法在每次有新的符号推理完成时调用BinaryPartitioner::Partition()创建新的切片。6. 运行时执行模型6.1 PartitionedCall 子图展开在运行时 lowering 阶段runtime/v2/lowering/graph_converter.ccPartitionedCall节点可以被展平回父图。ExpandPartitionedCallToParentGraph()方法在 PartitionedCall 前后各插入一个 NoOp 节点用于控制依赖。将子图内部的 InnerData 节点替换为父图的输入数据边。将子图内部的 InnerNetOutput 节点的控制边连接到后置 NoOp 节点。将子图所有节点移入父图更新节点和边的归属关系。这种展平策略让运行时可以灵活选择是否保持子图隔离。6.2 静态子图执行DavinciModelKernelruntime/v2/kernel/known_subgraph/davinci_model_kernel.h中的DavinciModelKernel负责静态子图的执行。静态子图被编译为DavinciModel包含预编译的 kernel 二进制和静态内存规划结果。运行时直接加载执行无需运行时 tiling 计算。6.3 PartitionedCall Loweringruntime/v2/engine/gelocal/partitioned_call_converter.cc注册了PartitionedCall节点的 lowering 转换器。它将 PartitionedCall 的输入/输出转换为运行时的数据搬运操作处理子图与父图之间的数据传递。6.4 Stage 同步对于流水线阶段拆分ExpandLastSyncExeNodesToMainGraph()和ExpandFirstExeNodesToMainGraph()方法处理阶段间的同步点展开确保前一个阶段的最后执行节点和后一个阶段的首执行节点之间的依赖关系正确建立。7. 拆分规则总览8. 关键文件索引文件核心内容docs/architecture/constraints/graph_split.md图拆分模块设计约束文档compiler/graph/partition/base_partitioner.h/.cc拆分框架基类定义拆分流水线compiler/graph/partition/base_cluster.h/.ccCluster 基类节点合并与子图构建compiler/graph/partition/dynamic_shape_partition.h/.cc动静 Shape 拆分策略实现compiler/graph/partition/engine_partitioner.h/.cc引擎级拆分PlaceHolder/End 机制compiler/graph/partition/stage_partitioner.h/.cc流水线阶段拆分compiler/graph/partition/engine_place.h/.cc引擎分配器compiler/graph/partition/optimizer/dynamic_data_flow_partitioner_pass.h/.cc数据流算子重拆分 passcompiler/graph/partition/optimizer/dynamic_data_flow_engine_reassign_pass.h/.cc数据流引擎重分配 passcompiler/graph/passes/multi_batch/multi_batch_clone_pass.h/.cc多档 Case 分支构建compiler/graph/passes/multi_batch/create_subgraph_with_scope_pass.h/.ccScope 级子图创建compiler/graph/passes/variable_optimize/split_variable_into_subgraph_pass.h/.ccVariable 拆分入子图compiler/graph/manager/graph_manager.cc编译流水线编排OptimizeSubgraphcompiler/graph/build/graph_builder.cc消费拆分结果进行图构建api/session/jit_execution/utils/partitioner/binary_partitioner.h/.ccJIT 二分拆分器api/session/jit_execution/exe_points/execution_order.h/.ccJIT 执行点管理api/atc/main_impl.ccatc 命令行选项处理runtime/v2/lowering/graph_converter.cc运行时 PartitionedCall 展开runtime/v2/kernel/known_subgraph/davinci_model_kernel.h/.cc静态子图执行 kernelruntime/v2/engine/gelocal/partitioned_call_converter.ccPartitionedCall lowering【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考