CANN/ge: GE 图拆分特性分析
2026/5/9 22:35:33 网站建设 项目流程

GE 图拆分(Graph Split)特性分析

【免费下载链接】geGE(Graph Engine)是面向昇腾的图编译器和执行器,提供了计算图优化、多流并行、内存复用和模型下沉等技术手段,加速模型执行效率,减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力,并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge

1. 特性背景

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 命令行工具编译模型时,可以指定多档 Shape,GE 自动将图拆分为"公共入口 + Case 分支子图"结构:

atc --model=resnet.onnx \ --dynamic_batch_size=1,4,8,16 \ --output=resnet_dyn

或指定更灵活的动态维度组合:

atc --model=bert.onnx \ --input_shape="input:1,-1,128" \ --dynamic_dims="1,32;1,64;1,128" \ --output=bert_dyn

编译结果是一个 OM 模型,运行时根据实际输入的 Shape 自动选择对应分支执行。

2.2 动静态混合场景

模型中部分算子的 Shape 依赖运行时计算结果(如NonMaxSuppression的输出数量取决于输入内容和阈值),而其他算子 Shape 静态。GE 自动将这类图拆分为静态子图和动态子图,静态子图享受预编译 kernel 和静态内存规划的性能优势,动态子图走运行时 tiling 和动态调度路径。

2.3 流水线并行场景(Stage Partition)

对于大型模型,用户可通过ATTR_STAGE_LEVEL属性将算子标记到不同流水线阶段。GE 按阶段拆分子图,每个阶段独立编译执行,阶段间通过同步点协调。

2.4 JIT 增量编译场景

在线模式下(如通过 TorchAir),GE 支持增量式图拆分。随着符号推理逐步完成,已确定 Shape 的子图先编译执行,未确定部分保留等待后续输入信息,实现逐层"剥洋葱"式的编译-执行交替。

3. 编译流水线中的位置

图拆分在GraphManager::OptimizeSubgraph()中按以下顺序执行,位于编译流水线的中段:

StagePartition → EnginePlacer1 → HostcpuEngineUpdatePass → DynamicShapePartition + EnginePlacer2 → CompositeEnginePartition + 子图优化 + Merge → AtomicEnginePartition + 子图优化 + Merge

对应代码入口compiler/graph/manager/graph_manager.ccOptimizeSubgraph()方法。

每一步的职责:

阶段执行器职责
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,128

4.2 运行时配置选项

选项默认值含义
ge.exec.static_model_ops_lower_limit4(ffts+ 场景为 6)静态子图最少算子数阈值,低于此值的静态子图降级为动态。设为 -1 可将所有子图合并为动态图
ge.topoSortingMode默认设为3启用稳定 RDFS 排序,会改变 cluster 合并策略
ge.tiling_schedule_optimize0设为1启用 tiling 下沉(在 AICPU 上执行 tiling)
ge.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中的BasePartitionercompiler/graph/partition/base_cluster.h中的BaseCluster构成。所有具体拆分策略都继承此框架。

5.1.1 拆分流水线

BasePartitioner::PartitionImpl()定义了统一的拆分流程:

InitClusters → MergeClusters → ProcessUniqueClusters → BuildPartitionFrame → CombinePartitionFrame → BuildPartitionSubgraph
  1. InitClusters:为每个节点创建一个独立的 cluster,按策略分类(DATA / KNOWN_SHAPE / UNKNOWN_SHAPE / NETOUTPUT 等)。
  2. MergeClusters:根据特定规则合并相邻 cluster,减少子图数量。
  3. ProcessUniqueClusters:去重、清理合并后的 cluster 集合。
  4. BuildPartitionFrame:为每个 cluster 在根图中创建一个PartitionedCall节点,并将 cluster 内节点移入对应子图。
  5. CombinePartitionFrame:在PartitionedCall节点之间建立数据边。
  6. BuildPartitionSubgraph:在子图内部添加 InnerData / InnerNetOutput 节点,完成 IO 连接。
5.1.2 Cluster 数据结构

BaseCluster的核心字段:

  • type_index_:类型索引,标识 cluster 类别(DATA=0, NETOUTPUT=1, INPUT_NODE=2, STAGE=3, KNOWN_SHAPE=4, UNKNOWN_SHAPE=5)
  • min_/max_:cluster 内节点的拓扑序范围,用于合并判断
  • in_clusters_/out_clusters_:入/出 cluster 的邻接关系
  • nodes_:cluster 包含的节点集合
  • subgraph_:cluster 对应的子图ComputeGraph
  • partition_node_:cluster 在根图中对应的PartitionedCall节点

关键合并操作:

  • Merge()— 无条件合并,将另一个 cluster 的所有节点和邻接关系吸收进来
  • TryMerge()— 仅在不会形成环时合并(通过前向可达性检查)
  • MergeAllPathFrom()— 合并两个 cluster 之间所有路径上的 cluster(双向 BFS 找路径交集)
5.1.3 属性传递机制

PartitionNodeAttrNameManager管理需要在拆分前后传递的节点属性注册表。通过REGISTER_PARTITION_ATTR_NAME宏注册属性,拆分时自动在PartitionedCall节点和子图内部节点间复制这些属性,确保拆分后语义一致。

5.2 动静 Shape 拆分:DynamicShapePartitioner

compiler/graph/partition/dynamic_shape_partition.h中的DynamicShapePartitioner是图拆分的核心策略实现,负责将计算图按动态/静态 Shape 划分为不同子图。

5.2.1 节点分类规则

MarkUnknownShapeNodes()方法按以下规则判断节点是否属于动态 Shape:

  1. 动态 Shape 算子:Tensor Shape 中存在 -1(未知维度)或 -2(未知 rank)
  2. Force Unknown 标记:节点被设置了_force_unknown_shape=true
  3. Tiling 依赖不支持下沉:节点存在动态 tiling 依赖,但不支持在 AICPU 上执行 tiling
  4. 地址刷新不支持:节点_is_support_addr_refresh=false
  5. Host CPU 引擎:节点属于DNN_VM_HOST_CPU引擎
  6. 子图传播:若节点的子图(控制流子图)中存在动态 Shape 算子,则该节点也归为动态
5.2.2 Cluster 合并策略

DynamicShapeCluster继承BaseCluster,按类型分为KNOWN_SHAPE(type_index=4)和UNKNOWN_SHAPE(type_index=5)。

MergeClustersNormal()的合并顺序:

  1. 动态路径吸附:遍历所有UNKNOWN_SHAPEcluster,若两个动态 cluster 之间存在路径,则将路径上所有 cluster 合并为动态。这保证了动态链路的连续性。
  2. 静态单路径合并:遍历KNOWN_SHAPEcluster,若两个静态 cluster 之间仅存在唯一路径(无环),则合并。
  3. 小 cluster 降级:节点数低于ge.exec.static_model_ops_lower_limit阈值的静态 cluster 降级为动态,避免产生过小的静态子图碎片。
  4. 控制流合并:属于同一ATTR_NAME_CONTROL_FLOW_GROUP的控制流节点(如 StreamActive、StreamSwitch)合并到同一 cluster。
  5. RefVariable 合并:引用类型的 Variable 节点与其消费者合并到同一 cluster。
5.2.3 重拆分机制

初次拆分后,DynamicDataFlowPartitionerPass检查数据流算子(Stack/StackPush/StackPop/StackClose)是否横跨动静子图。如果是,将这些算子强制标记为_force_unknown_shape=true,然后调用ReDynamicShapePartitioner()重新拆分。这个迭代过程确保数据流状态在执行语义上的一致性。

5.2.4 整图动态判定

IsGraphNeedUnknownShapePartition()判断整图是否需要走动态拆分流程。如果图中没有任何动态 Shape 节点,则设置_dynamic_shape_partitioned=false,整图走静态编译路径。如果图中静态节点数很少(低于ge.host_scheduling_max_threshold),则整图直接走 Host 调度模式。

5.3 引擎级拆分:EnginePartitioner

compiler/graph/partition/engine_partitioner.h中的EnginePartitioner负责按引擎归属拆分子图。它在动静拆分之后执行,进一步将子图按 AI Core、AI CPU、DVPP 等不同硬件引擎切分。

5.3.1 拆分流程
  1. Initialize:通过EnginePlacer为每个节点分配引擎,创建初始 cluster(每个节点一个,携带引擎名和 stream label)。
  2. MarkClusters:遍历 cluster 对,若两个 cluster 具有相同引擎 + 相同 stream label + 之间无第二条路径,则合并。
  3. SplitSubGraphs:为每个合并后的 cluster 创建ComputeGraph子图,在不同引擎子图之间插入PlaceHolder/End节点对。
  4. 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 流水线阶段拆分:StagePartitioner

compiler/graph/partition/stage_partitioner.h中的StagePartitioner按节点上的ATTR_STAGE_LEVEL属性将计算图拆分为多个流水线阶段。

拆分逻辑:

  1. SplitStageLevel():收集带有ATTR_STAGE_LEVEL属性的节点,将属性向上游传播。
  2. SplitTailStage():将未标记阶段的节点归入最后一个阶段。
  3. StagePartition():使用GraphUtils::BuildSubgraphWithNodes()将每个阶段的节点封装为子图,阶段间通过PartitionedCall节点连接。

每个阶段的父节点会被设置_force_unknown_shape=true,确保阶段间的同步在运行时处理。

5.5 多档 Clone:MultiBatchClonePass

compiler/graph/passes/multi_batch/multi_batch_clone_pass.h中的MultiBatchClonePass处理动态分档场景。它将原始图 Clone 为多份,每份对应一个档位的 Shape,然后通过 Case 节点在运行时选择分支。

5.5.1 构建流程
  1. CollectIoNodes:收集原始图的输入/输出节点,解析用户指定的动态 Shape 参数。
  2. CreateRootGraph:创建根图,包含:
    • Shape 索引节点(Data 或GetDynamicDims
    • 带 max shape 的输入 Data 节点
    • Case 节点
  3. CreateSubgraphs:将原始图 Clone N 份(N 为档位数),每份使用对应档位的静态 Shape,作为 Case 节点的分支子图。
  4. 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 拆分入子图:SplitVariableIntoSubgraphPass

compiler/graph/passes/variable_optimize/split_variable_into_subgraph_pass.h中的SplitVariableIntoSubgraphPass处理 Variable/RefData 节点与控制流子图(If/Case/PartitionedCall/While)的交互。对于需要被子图内部访问的 Variable 节点,将其拷贝到子图内部,确保子图可以独立访问权重数据。对于 While 节点,由于循环语义的特殊性,改为添加控制边而非拷贝。

5.7 JIT 增量拆分:BinaryPartitioner

api/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.cc),PartitionedCall节点可以被"展平"回父图。ExpandPartitionedCallToParentGraph()方法:

  1. 在 PartitionedCall 前后各插入一个 NoOp 节点用于控制依赖。
  2. 将子图内部的 InnerData 节点替换为父图的输入数据边。
  3. 将子图内部的 InnerNetOutput 节点的控制边连接到后置 NoOp 节点。
  4. 将子图所有节点移入父图,更新节点和边的归属关系。

这种展平策略让运行时可以灵活选择是否保持子图隔离。

6.2 静态子图执行:DavinciModelKernel

runtime/v2/kernel/known_subgraph/davinci_model_kernel.h中的DavinciModelKernel负责静态子图的执行。静态子图被编译为DavinciModel,包含预编译的 kernel 二进制和静态内存规划结果。运行时直接加载执行,无需运行时 tiling 计算。

6.3 PartitionedCall Lowering

runtime/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数据流算子重拆分 pass
compiler/graph/partition/optimizer/dynamic_data_flow_engine_reassign_pass.h/.cc数据流引擎重分配 pass
compiler/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编译流水线编排,OptimizeSubgraph
compiler/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静态子图执行 kernel
runtime/v2/engine/gelocal/partitioned_call_converter.ccPartitionedCall lowering

【免费下载链接】geGE(Graph Engine)是面向昇腾的图编译器和执行器,提供了计算图优化、多流并行、内存复用和模型下沉等技术手段,加速模型执行效率,减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力,并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询