CANN/GE 流分配特性分析

Stream Allocator流分配特性分析【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1 特性背景昇腾 AI 处理器上的计算任务通过流Stream来组织和调度。流是设备侧的执行队列——同一条流内的任务严格按序执行不同流之间的任务可以并行执行。流分配的质量直接影响模型的执行效率分配的流太少无法充分利用硬件并行能力分配的流太多又会带来过多的同步开销Event/Notify和资源占用。GE 图编译器在将 AscendIR 编译为可执行模型OM 文件的过程中需要在编译期完成流的分配决策。这一决策涉及三个核心问题哪些算子可以并行执行需要根据引擎类型、数据依赖关系、用户标注等信息决定。并行执行的算子之间如何同步不同流之间需要插入 Event/Notify 来保证数据一致性。物理流的容量有限时如何拆分一条逻辑流承载的 task 数量有上限超出时需要拆分为多条物理流。流分配特性正是为系统性地解决这些问题而设计的。适用场景流分配特性适用于以下典型场景场景说明静态 Shape 模型编译模型的输入 shape 在编译期已知GE 可以基于完整的图拓扑进行精细的多流分配动态 Shape 模型编译模型的输入 shape 在运行时才确定GE 需要采用更保守的分流策略混合引擎模型模型中同时包含 AI Core、HCCL集合通信、AI CPU、DVPP 等不同引擎的算子需要按引擎特性分流训练场景的 AllReduce 并行梯度聚合AllReduce与反向计算并行执行以加速训练用户自定义分流用户通过 StreamLabel 属性指定特定算子分配到特定流2 总体架构流分配特性横跨编译器和运行时两个阶段形成逻辑流分配 → 同步事件插入 → 物理流拆分 → 运行时流创建的完整流水线。模块分工模块所在目录职责StreamAllocatorcompiler/graph/build/stream/编译期流分配的总入口协调逻辑流分配、同步插入、物理流拆分LogicalStreamAllocatorcompiler/graph/build/stream/静态 Shape 下的逻辑流分配基于 Pass 链式架构DynamicStreamAllocatorcompiler/graph/build/stream/动态 Shape 下的流分配策略更简洁StreamUtilscompiler/graph/build/stream/流分配的公共工具函数gert::StreamAllocatorinc/framework/runtime/运行时 V2 路径的流创建接口ge::ReusableStreamAllocatorruntime/v1/运行时 V1 路径的流复用池3 对外接口3.1 编译期 API编译期流分配作为图编译流水线的一部分不直接暴露给终端用户。但编译完成后用户可通过以下接口查询流分配结果。GetStreamAllocationSummary获取流分配概要信息包括逻辑流、物理流、附着流的分配情况。头文件ge/ge_graph_compile_summary.h库文件libge_compiler.so函数原型Status GetStreamAllocationSummary( std::shared_ptrStreamAllocationSummary stream_allocation) const;返回的StreamAllocationSummary对象提供以下查询接口接口说明GetAllLogicalStreamInfos()获取所有逻辑流的分配信息GetUsrStreamLabels()获取用户流标签列表GetPhysicalStreamNums()获取物理流数量GetAttachedStreamIds()获取附着流 ID 列表GetHcclFollowedStreamNums()获取 HCCL 后续流数量IsAssignedByStreamPass()判断是否由 StreamPass 分配LogicalStreamAllocationInfo每条逻辑流的详细信息包括接口说明GetLogicalStreamId()逻辑流 IDGetUsrStreamLabel()用户流标签GetAttachedStreamIds()附着流 IDGetPhysicalStreamNum()物理流数量GetHcclFollowedStreamNum()HCCL 后续流数量GetAllNodes()该流上所有节点3.2 运行时 APIgert::StreamAllocatorV2 路径运行时流创建接口按需创建和管理设备流。头文件framework/runtime/stream_allocator.h核心接口namespace gert { class StreamAllocator { // 最多支持 2024 条流 static constexpr size_t kMaxStreamNum 2024U; StreamAllocator(int32_t priority RT_STREAM_PRIORITY_DEFAULT, uint32_t flags RT_STREAM_DEFAULT); ~StreamAllocator(); // 按需获取流返回连续向量不足部分自动创建 TypedContinuousVectorrtStream_t *AcquireStreams(size_t stream_num) const; }; }该接口在模型加载阶段被调用根据编译期确定的流数量一次性创建所需的全部设备流。实现上使用ContinuousVector预分配最大容量2024 条流通过SetSize标记实际使用的流数量避免频繁的内存分配。ge::ReusableStreamAllocatorV1 路径运行时流复用池用于跨模型复用设备流减少流创建/销毁开销。头文件runtime/v1/graph/load/model_manager/reusable_stream_allocator.h核心接口namespace ge { class ReusableStreamAllocator { static ReusableStreamAllocator *Create(); Status GetOrCreateRtStream(aclrtStream stream, uint32_t rt_model_id, int32_t priority, uint32_t stream_flag, uint32_t task_num 0U); Status DestroyStream(aclrtStream stream, bool is_force_destroy false); }; }ReusableStreamAllocator维护一个以priority, stream_flag为键的流池按 task_num 排序。新模型加载时优先从已有流池中查找可复用的流避免重复调用rtStreamCreate。每个流通过rt_model_id集合追踪使用它的模型确保不会复用自身模型的流。3.3 用户可配置项用户可通过以下方式影响流分配行为配置项影响范围说明SINGLE_STREAM_ENABLE静态 Shape开启单流模式所有算子在一条流上执行AC_PARALLEL_ENABLE动态 Shape取值为 0、1 或空控制 AI CPU 与 AI Core 是否并行EVENT静态 Shape设为 notify 时使用 Notify 替代 Event 进行同步STREAM_LABEL节点属性所有场景算子级别的流标签相同标签的算子分配到同一条流USER_STREAM_LABEL节点属性所有场景用户级流标签优先级最高PARALLEL_GROUP节点属性静态 Shape并行组标识同组算子分配到独立流ATTACHED_STREAM_INFO节点属性静态 Shape附着流信息一个节点可产生多条流4 具体实现4.1 静态 Shape 逻辑流分配静态 Shape 下的逻辑流分配采用Pass 链式架构每个 Pass 负责一类分流规则按优先级依次执行。这套架构的设计哲学是关注点分离——每种分流逻辑独立封装为 Pass新增分流规则只需添加新 Pass无需修改已有逻辑。4.1.1 Pass 链详解UpdateForMdeGroupPass根据NewStreamId属性为节点分配新流。这是最高优先级的分流规则用于支持 MDEMulti-Data Execution场景下特定算子的独立流需求。AssignByLabelPass根据StreamLabel属性分流。相同 StreamLabel 的子图分配到同一条流不同 StreamLabel 分配新流。这让上层编译优化如融合 Pass可以通过设置 StreamLabel 来指导流分配。IndependentStreamPass为独立引擎如 HCCL的子图分配独立流。独立引擎的算子需要独占一条流不能与其他引擎复用。同一个独立引擎内相同 StreamLabel 的子图共享流。AssignByDependencyPass最核心的分流 Pass根据引擎子图之间的数据依赖关系进行流分配和流复用。该 Pass 的工作方式是遍历所有未分流的子图查找前驱子图中是否存在可复用的流若可复用则复用否则分配新流流复用需要满足三个条件scheduler_id 相同、不是独立引擎/带标签的流、无引擎冲突NodeStreamUpdatePass将子图级别的流分配结果映射到节点级别。每个节点获得其所属子图的 stream_id。特殊地带有ATTR_NAME_RTS_LABEL_NODE属性的节点会被分配到父流而非子图流用于支持控制流场景。UpdateForParallelGroupPass根据PARALLEL_GROUP属性为节点重新分配流。同一并行组的节点分配到同一条新流。对于 HCOM 算子如果并行组名为 -1 且只有一个输入则尝试复用输入节点的流。AllReduceParallelPass当开启hcom_parallel时将 AllReduce 算子的后继非 HCOM 节点分配到新流使 AllReduce 与反向计算可以并行执行。这是训练加速的关键优化。UpdateForSkippedEnginePass优化跳过引擎skipped engine子图中的节点流分配。对于NodeA(stream1) → Const(stream2) → NodeB(stream1)这类模式将 Const 节点的流改为 stream1从而减少两个流之间不必要的同步事件。OptimizeIneffectiveMultiStreamPass拓扑优化 Pass消除名义上多流但实际不产生并行收益的情况。如果某个节点在所有输入输出方向上都与另一条流相连且在该流上输入输出节点之间没有其他节点则将当前节点移到那条流上从而减少同步开销。4.1.2 附着流分配附着流Attached Stream是一个节点产生的除主流之外的额外流。某些算子如 SuperKernel需要多条流来执行不同的计算任务。附着流分配在主流分配完成后进行。AssignAttachedStreamPass通过ATTR_NAME_ATTACHED_STREAM_INFO或ATTR_NAME_ATTACHED_STREAM_INFO_LIST属性获取附着流信息包括流数量和复用键reuse_key。相同 reuse_key 的附着流共享同一条流避免不必要的流创建。附着流分流完成后总的流数量 主流数量 附着流数量。4.2 动态 Shape 流分配动态 Shape 下的流分配策略相比静态 Shape 更为保守——默认只分配一条流单流模式只有在配置开启多流时才启用多流。这是因为动态 Shape 的图结构在编译期不完整无法进行精确的依赖分析。与静态 Shape 的关键差异差异点静态 Shape动态 Shape默认模式多流单流分流粒度Pass 链式处理规则精细按引擎分流规则简洁流复用策略基于依赖关系的复杂复用判断前驱/后继子图复用节点级约束较少Data、Variable、NetOutput、FILECONSTANT 等强制在主流附着流支持通过独立接口AssignAttachedResource支持同步机制Event Notify 双模式仅 Event4.3 同步事件管理流分配完成后不同流上的算子之间需要同步事件来保证执行顺序正确性。同步事件的管理是流分配特性中最复杂的部分。4.3.1 事件类型类型说明适用场景kEvent普通 EventSend/Recv 配对默认模式kNotifyNotify支持更细粒度的同步通过EVENTnotify配置开启4.3.2 事件插入规则系统在以下场景插入同步事件事件插入的核心逻辑在StreamAllocator::InsertOneEventInTwoNodes中遍历整图的所有数据边和控制边当相邻两个节点属于不同流时在两个节点之间插入一对 Send/Recv 事件。4.3.3 事件优化插入事件后系统会通过三重优化消除冗余事件OptimizeBySendEvents在同一条流内如果 Send 节点 A 的事件已经确保了流 B 上的 Recv 节点 C 在 A 之后执行那么 A 和 C 之间不需要额外的事件。OptimizeByRecvEvents类似地在接收方向上消除冗余。OptimizeByStreamActivate通过StreamActive机制优化跨流事件。当流 A 上的节点通过StreamActive激活了流 B则流 A 到流 B 之间不需要额外的 Event——因为StreamActive本身就隐含了同步语义。该优化通过IsRecvNodeActivatedBySendNode方法判断沿着激活链回溯检查是否存在激活关系。4.3.4 事件复用在多档位multi-dims场景下不同档位在同一时刻只有一个会执行因此它们的 Event 可以复用。ReuseEventForMultiDims方法为每个档位独立编号 Event然后取最大值作为最终 Event 数量。例如dim0: event 0, 1, 2, 3 → 0, 1, 2, 3 dim1: event 4, 5, 6, 7, 8 → 0, 1, 2, 3, 4 dim2: event 9, 10, 11 → 0, 1, 2 最终 event_num max(4, 5, 3) 5此外算子可以通过ATTR_NAME_EVENT_MULTIPLEXING属性显式声明 Event 的复用关系系统据此替换对应的事件 ID。4.3.5 事件连续性保证RTSRuntime Service要求 Event ID 必须从 0 开始连续分配。因此在所有优化和复用处理完成后系统通过RefreshContinuousEvents方法重新映射 Event ID确保连续性。该逻辑同样适用于 Notify。4.4 物理流拆分逻辑流分配不考虑 task 数量限制但物理流承载的 task 数量有上限。物理流拆分阶段负责将超出限制的逻辑流拆分为多条物理流。拆分触发条件StreamAllocator::NeedSpiltNewStream不是流的首个节点当前流上的 task 数量 预留数量 上限节点没有子图非控制流节点不是StreamActive的首节点不是LabelSet/LabelGotoEx/LabelSwitchByIndex等控制流标签节点拆分时需处理的事项更新节点的 stream_id 为新的物理流 ID在拆分点的前后节点之间插入同步事件维护split_stream_id_to_logic_stream_id_映射处理ContinuousStreamLabel相同标签的节点必须拆分到同一条流Huge Stream当单流模式的 task 数量超过普通流上限时系统尝试使用 Huge Stream大流大流具有更高的 task 容量。4.5 流激活机制流激活Stream Activate是昇腾设备侧的流调度机制。当一条流需要唤醒另一条流时通过StreamActive算子发送激活信号。激活关系的建立过程标签激活SetActiveStreamsByLabel遍历所有带有ATTR_NAME_ACTIVE_LABEL_LIST属性的节点将标签映射到实际的流 ID写入ATTR_NAME_ACTIVE_STREAM_LIST属性。子图激活SetActiveStreamsForSubgraphs为 While/For 等循环子图的首个StreamActive节点设置激活流列表确保子图内所有流都被正确激活。Switch 节点激活UpdateActiveStreamsForSwitchNodeStreamSwitch节点后面会插入StreamActive节点根据条件分支激活对应的流。循环激活SetActiveStreamsForLoop处理训练场景的 FpBp 循环StreamActive节点需要激活所有未被特定激活的流确保每次迭代开始时所有流都被正确启动。流拆分后的激活更新物理流拆分会产生新的流 ID系统需要更新所有StreamActive节点的激活列表将拆分后的新流加入激活范围。4.6 同步事件节点生成在流分配的最后阶段系统需要将记录在数据结构中的事件信息node_to_send_events_、node_to_recv_events_等转化为图中的实际节点Send/Recv 算子。GenerateSyncEventNodes方法遍历所有节点的事件映射为每个事件创建对应的 Send 和Recv 节点并通过控制边将它们插入到图中的正确位置。这些节点在后续的 Task 生成阶段会被转换为设备侧的 Event Record/Wait 任务。4.7 运行时流创建编译期确定了流数量后运行时在模型加载阶段创建对应的设备流。V2 路径gert::StreamAllocatorgert::StreamAllocator预分配一个最多容纳 2024 条rtStream_t的连续向量。首次调用AcquireStreams时依次创建指定数量的流调用rtStreamCreateWithFlags后续调用则复用已创建的流。流的销毁在析构函数中统一完成。V1 路径ge::ReusableStreamAllocatorV1 路径支持跨模型复用流。ReusableStreamAllocator维护以priority, stream_flag为键的流池每个流记录其task_num和使用它的模型 ID 集合。当新模型请求流时从流池中查找匹配的priority, stream_flag筛选不属于当前模型的流按task_num排序优先选择 task 数相近的流复用若无可用流则创建新流5 关键设计决策5.1 为什么静态和动态采用不同的分流策略静态 Shape 的图拓扑在编译期完全已知GE 可以精确分析所有数据依赖关系并进行细粒度的流复用。而动态 Shape 的图结构在编译期不完整部分子图是运行时才展开的无法进行精确的依赖分析。因此动态 Shape 采用更保守的策略默认单流仅在用户明确配置时开启多流且分流规则以引擎粒度为主而非依赖分析。5.2 为什么使用 Pass 链而非单一分配算法Pass 链架构的核心优势是可扩展性和可维护性。每种分流规则标签分流、引擎分流、依赖分流、AllReduce 并行等封装为独立 Pass各自维护状态互不干扰。新增分流规则只需添加新 Pass 到链中无需修改已有逻辑。如果采用单一算法所有规则交织在一起代码的可读性和可维护性会急剧下降。5.3 为什么需要图结构稳定性原则流分配依赖于拓扑排序和 topo ID 的连续性。如果在流分配过程中图结构发生变化增删节点会导致 topo ID 不连续进而影响后续的内存复用依赖 topo 顺序分配内存块和物理流拆分依赖 topo 顺序计算 task 数量。因此流分配阶段严格禁止改图操作——所有需要插入的节点如 Send/Recv、StreamActive都在流分配完成后统一插入。5.4 为什么附着流要在主流分配完成后才能分配附着流的 ID 是在主流 ID 的基础上递增的。如果附着流和主流混合分配会导致流 ID 不连续增加同步管理的复杂度。主流分配完成后一次性分配附着流可以确保流 ID 的连续性和可预测性。6 约束与限制约束说明单流模式不支持 StreamLabel单流模式只有一条流StreamLabel 会导致冲突Event ID 必须连续RTS 对 Event ID 有连续性校验必须通过RefreshContinuousEvents保证图结构不可变流分配期间图结构不能改变topo ID 必须连续Notify 数量上限最大支持 1024 个 Notify单节点附着流数量一个节点最多支持一个附着流通过 reuse_key 复用多线程安全StreamAllocator支持多线程但需保护共享资源ScalableAllocator不支持多线程并发动态图 While 约束动态图上 While 算子的静态 body 子图的 NetOutput 必须在 stream 0分档场景标签约束分档场景下添加 StreamLabel 时需要加上档位信息以区分【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考