Stream Allocator Feature Analysis【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1 Feature BackgroundComputation tasks on Ascend AI processors are organized and scheduled through Streams. A stream is a device-side execution queue. Tasks within the same stream execute strictly in order, while tasks between different streams can execute in parallel. The quality of stream allocation directly affects model execution efficiency. Allocating too few streams cannot fully utilize hardware parallel capability, while allocating too many streams brings excessive synchronization overhead (Event/Notify) and resource occupation.During the process of compiling AscendIR to executable models (OM files), the GE graph compiler needs to complete stream allocation decisions at compilation time. This decision involves three core questions:Which operators can execute in parallel?This requires deciding based on engine type, data dependency relationships, user annotations, and other information.How to synchronize between operators executing in parallel?Events/Notify need to be inserted between different streams to ensure data consistency.How to split when physical stream capacity is limited?A logical stream has an upper limit on the number of tasks it can carry. When exceeded, it needs to be split into multiple physical streams.The stream allocation feature is designed to systematically solve these problems.Applicable ScenariosThe stream allocation feature applies to the following typical scenarios:ScenarioDescriptionStatic Shape Model CompilationModel input shape is known at compilation time. GE can perform fine-grained multi-stream allocation based on complete graph topologyDynamic Shape Model CompilationModel input shape is determined at runtime. GE needs to adopt a more conservative stream allocation strategyMixed Engine ModelModel contains operators from different engines such as AI Core, HCCL (collective communication), AI CPU, DVPP. Stream allocation needs to be based on engine characteristicsTraining Scenario AllReduce ParallelismGradient aggregation (AllReduce) executes in parallel with backward computation to accelerate trainingUser-defined Stream AllocationUser specifies particular operators to be allocated to particular streams through StreamLabel attribute2 Overall ArchitectureThe stream allocation feature spans both compiler and runtime phases, forming a complete pipeline of logical stream allocation → synchronization event insertion → physical stream splitting → runtime stream creation.Module ResponsibilitiesModuleDirectoryResponsibilityStreamAllocatorcompiler/graph/build/stream/Compilation phase stream allocation main entry, coordinating logical stream allocation, synchronization insertion, physical stream splittingLogicalStreamAllocatorcompiler/graph/build/stream/Logical stream allocation for static shape, based on Pass chain architectureDynamicStreamAllocatorcompiler/graph/build/stream/Stream allocation for dynamic shape, with simpler strategyStreamUtilscompiler/graph/build/stream/Common utility functions for stream allocationgert::StreamAllocatorinc/framework/runtime/Runtime V2 path stream creation interfacege::ReusableStreamAllocatorruntime/v1/Runtime V1 path stream reuse pool3 External Interfaces3.1 Compilation Phase APICompilation phase stream allocation is part of the graph compilation pipeline and is not directly exposed to end users. However, after compilation completes, users can query stream allocation results through the following interfaces.GetStreamAllocationSummaryGet stream allocation summary information, including logical stream, physical stream, and attached stream allocation status.Header File:ge/ge_graph_compile_summary.hLibrary File:libge_compiler.soFunction Prototype:Status GetStreamAllocationSummary( std::shared_ptrStreamAllocationSummary stream_allocation) const;The returnedStreamAllocationSummaryobject provides the following query interfaces:InterfaceDescriptionGetAllLogicalStreamInfos()Get allocation information of all logical streamsGetUsrStreamLabels()Get user stream label listGetPhysicalStreamNums()Get physical stream countGetAttachedStreamIds()Get attached stream ID listGetHcclFollowedStreamNums()Get HCCL followed stream countIsAssignedByStreamPass()Determine if assigned by StreamPassLogicalStreamAllocationInfoDetailed information for each logical stream, including:InterfaceDescriptionGetLogicalStreamId()Logical stream IDGetUsrStreamLabel()User stream labelGetAttachedStreamIds()Attached stream IDGetPhysicalStreamNum()Physical stream countGetHcclFollowedStreamNum()HCCL followed stream countGetAllNodes()All nodes on this stream3.2 Runtime APIgert::StreamAllocator (V2 Path)Runtime stream creation interface, creating and managing device streams on demand.Header File:framework/runtime/stream_allocator.hCore Interface:namespace gert { class StreamAllocator { // Supports up to 2024 streams static constexpr size_t kMaxStreamNum 2024U; StreamAllocator(int32_t priority RT_STREAM_PRIORITY_DEFAULT, uint32_t flags RT_STREAM_DEFAULT); ~StreamAllocator(); // Acquire streams on demand, returns continuous vector, auto-creates if insufficient TypedContinuousVectorrtStream_t *AcquireStreams(size_t stream_num) const; }; }This interface is called during model loading phase. Based on the stream count determined at compilation time, it creates all required device streams at once. The implementation usesContinuousVectorto pre-allocate maximum capacity (2024 streams), marking actual used stream count throughSetSize, avoiding frequent memory allocation.ge::ReusableStreamAllocator (V1 Path)Runtime stream reuse pool, used to reuse device streams across models, reducing stream creation/destruction overhead.Header File:runtime/v1/graph/load/model_manager/reusable_stream_allocator.hCore Interface: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); }; }ReusableStreamAllocatormaintains a stream pool keyed bypriority, stream_flag, sorted by task_num. When a new model loads, it first searches the existing stream pool for reusable streams, avoiding repeated calls tortStreamCreate. Each stream tracks the models using it throughrt_model_idset, ensuring it does not reuse streams from its own model.3.3 User Configurable OptionsUsers can influence stream allocation behavior through the following methods:Configuration ItemScopeDescriptionSINGLE_STREAM_ENABLEStatic ShapeEnable single-stream mode, all operators execute on one streamAC_PARALLEL_ENABLEDynamic ShapeValues are 0, 1 or empty, controls whether AI CPU and AI Core execute in parallelEVENTStatic ShapeWhen set to notify, use Notify instead of Event for synchronizationSTREAM_LABEL(Node Attribute)All ScenariosOperator-level stream label, operators with same label are allocated to same streamUSER_STREAM_LABEL(Node Attribute)All ScenariosUser-level stream label, highest priorityPARALLEL_GROUP(Node Attribute)Static ShapeParallel group identifier, operators in same group are allocated to independent streamsATTACHED_STREAM_INFO(Node Attribute)Static ShapeAttached stream information, one node can produce multiple streams4 Specific Implementation4.1 Static Shape Logical Stream AllocationLogical stream allocation under static shape uses aPass chain architecture, where each Pass is responsible for one type of stream allocation rule, executing in order by priority. The design philosophy of this architecture is separation of concerns. Each stream allocation logic is independently encapsulated as a Pass. Adding new stream allocation rules only requires adding new Passes, without modifying existing logic.4.1.1 Pass Chain DetailsUpdateForMdeGroupPass: Allocates new streams to nodes based onNewStreamIdattribute. This is the highest priority stream allocation rule, used to support independent stream requirements for specific operators in MDE (Multi-Data Execution) scenarios.AssignByLabelPass: Allocates streams based onStreamLabelattribute. Subgraphs with the same StreamLabel are allocated to the same stream, different StreamLabels get new streams. This allows upper-layer compilation optimizations (such as fusion Pass) to guide stream allocation by setting StreamLabel.IndependentStreamPass: Allocates independent streams to subgraphs of independent engines (such as HCCL). Operators of independent engines need to exclusively occupy a stream and cannot reuse with other engines. Within the same independent engine, subgraphs with the same StreamLabel share streams.AssignByDependencyPass: The most core stream allocation Pass. It allocates and reuses streams based on data dependency relationships between engine subgraphs. This Pass works as follows:Traverse all subgraphs without stream allocationCheck if predecessor subgraphs have reusable streamsReuse if possible, otherwise allocate new streamStream reuse requires three conditions: same scheduler_id, not independent engine/tagged stream, no engine conflictNodeStreamUpdatePass: Maps subgraph-level stream allocation results to node level. Each node gets the stream_id of its belonging subgraph. Specifically, nodes withATTR_NAME_RTS_LABEL_NODEattribute are allocated to parent stream (instead of subgraph stream), used to support control flow scenarios.UpdateForParallelGroupPass: Reallocates streams to nodes based onPARALLEL_GROUPattribute. Nodes in the same parallel group are allocated to the same new stream. For HCOM operators, if the parallel group name is -1 and has only one input, it tries to reuse the input nodes stream.AllReduceParallelPass: Whenhcom_parallelis enabled, allocates successor non-HCOM nodes of AllReduce operators to new streams, enabling AllReduce to execute in parallel with backward computation. This is a key optimization for training acceleration.UpdateForSkippedEnginePass: Optimizes node stream allocation in skipped engine subgraphs. For patterns likeNodeA(stream1) → Const(stream2) → NodeB(stream1), changes Const nodes stream to stream1, reducing unnecessary synchronization events between two streams.OptimizeIneffectiveMultiStreamPass: Topology optimization Pass, eliminating nominally multi-stream but actually no parallel benefit situations. If a node connects to another stream on all input/output directions, and no other nodes exist between input/output nodes on that stream, move the current node to that stream, reducing synchronization overhead.4.1.2 Attached Stream AllocationAttached Stream is an additional stream produced by a node besides the main stream. Some operators (such as SuperKernel) require multiple streams to execute different computation tasks. Attached stream allocation occurs after main stream allocation completes.AssignAttachedStreamPassgets attached stream information throughATTR_NAME_ATTACHED_STREAM_INFOorATTR_NAME_ATTACHED_STREAM_INFO_LISTattribute, including stream count and reuse_key. Attached streams with the same reuse_key share the same stream, avoiding unnecessary stream creation.After attached stream allocation completes, total stream count main stream count attached stream count.4.2 Dynamic Shape Stream AllocationStream allocation strategy under dynamic shape is more conservative compared to static shape. By default, only one stream is allocated (single-stream mode). Multi-stream is only enabled when configuration allows. This is because dynamic shape graph structure is incomplete at compilation time, preventing precise dependency analysis.Key Differences from Static Shape:DifferenceStatic ShapeDynamic ShapeDefault ModeMulti-streamSingle-streamAllocation GranularityPass chain processing, fine-grained rulesStream allocation by engine, simple rulesStream Reuse StrategyComplex reuse judgment based on dependencyPredecessor/successor subgraph reuseNode-level ConstraintsFewData, Variable, NetOutput, FILECONSTANT forced on main streamAttached StreamSupportedSupported through independent interfaceAssignAttachedResourceSynchronization MechanismEvent Notify dual modeEvent only4.3 Synchronization Event ManagementAfter stream allocation completes, synchronization events are needed between operators on different streams to ensure correct execution order. Synchronization event management is the most complex part of the stream allocation feature.4.3.1 Event TypesTypeDescriptionApplicable ScenariokEventNormal Event, Send/Recv pairingDefault modekNotifyNotify, supports finer-grained synchronizationEnabled throughEVENTnotifyconfiguration4.3.2 Event Insertion RulesThe system inserts synchronization events in the following scenarios:The core logic of event insertion is inStreamAllocator::InsertOneEventInTwoNodes. Traverse all data edges and control edges of the entire graph. When two adjacent nodes belong to different streams, insert a pair of Send/Recv events between them.4.3.3 Event OptimizationAfter inserting events, the system eliminates redundant events through three optimizations:OptimizeBySendEvents: Within the same stream, if the event of Send node A already ensures that Recv node C on stream B executes after A, then no additional event is needed between A and C.OptimizeByRecvEvents: Similarly, eliminate redundancy in the receiving direction.OptimizeByStreamActivate: Optimize cross-stream events throughStreamActivemechanism. When a node on stream A activates stream B throughStreamActive, no additional Event is needed from stream A to stream B, becauseStreamActiveitself implies synchronization semantics. This optimization judges throughIsRecvNodeActivatedBySendNodemethod, tracing back along the activation chain to check for activation relationship.4.3.4 Event ReuseIn multi-dims scenarios, only one dim executes at any moment, so their Events can be reused.ReuseEventForMultiDimsmethod independently numbers Events for each dim, then takes the maximum value as the final Event count. For example: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 Final event_num max(4, 5, 3) 5Additionally, operators can explicitly declare Event reuse relationship throughATTR_NAME_EVENT_MULTIPLEXINGattribute. The system replaces corresponding event IDs based on this.4.3.5 Event Continuity GuaranteeRTS (Runtime Service) requires Event IDs to be allocated continuously starting from 0. Therefore, after all optimization and reuse processing completes, the system remaps Event IDs throughRefreshContinuousEventsmethod to ensure continuity. This logic also applies to Notify.4.4 Physical Stream SplittingLogical stream allocation does not consider task count limits, but physical streams have an upper limit on task count. Physical stream splitting phase is responsible for splitting logical streams that exceed the limit into multiple physical streams.Split Trigger Conditions(StreamAllocator::NeedSpiltNewStream):Not the first node of the streamCurrent stream task count reserved count limitNode has no subgraph (non-control flow node)Not the first node ofStreamActiveNot control flow label nodes likeLabelSet/LabelGotoEx/LabelSwitchByIndexItems to Handle During Splitting:Update nodes stream_id to new physical stream IDInsert synchronization events between nodes before and after the split pointMaintainsplit_stream_id_to_logic_stream_id_mappingHandleContinuousStreamLabel: Nodes with the same label must be split to the same streamHuge Stream: When single-stream mode task count exceeds normal stream limit, the system tries to use Huge Stream, which has higher task capacity.4.5 Stream Activation MechanismStream Activate is the stream scheduling mechanism on the Ascend device side. When a stream needs to wake up another stream, it sends an activation signal throughStreamActiveoperator.Establishment Process of Activation Relationship:Label Activation (SetActiveStreamsByLabel): Traverse all nodes withATTR_NAME_ACTIVE_LABEL_LISTattribute, map labels to actual stream IDs, write toATTR_NAME_ACTIVE_STREAM_LISTattribute.Subgraph Activation (SetActiveStreamsForSubgraphs): For While/For loop subgraphs, set activation stream list for the firstStreamActivenode, ensuring all streams within the subgraph are correctly activated.Switch Node Activation (UpdateActiveStreamsForSwitchNode): AStreamActivenode is inserted afterStreamSwitchnode, activating corresponding streams based on conditional branches.Loop Activation (SetActiveStreamsForLoop): Handle FpBp loops in training scenarios. TheStreamActivenode needs to activate all streams not specifically activated, ensuring all streams are correctly started at the beginning of each iteration.Activation Update After Stream Splitting: Physical stream splitting produces new stream IDs. The system needs to update activation lists of allStreamActivenodes, adding newly split streams to the activation scope.4.6 Synchronization Event Node GenerationIn the final stage of stream allocation, the system needs to convert event information recorded in data structures (node_to_send_events_,node_to_recv_events_, etc.) into actual nodes in the graph (Send/Recv operators).GenerateSyncEventNodesmethod traverses event mappings of all nodes, creates corresponding Send and Recv nodes for each event, and inserts them into correct positions in the graph through control edges. These nodes are converted to device-side Event Record/Wait tasks in subsequent Task generation phase.4.7 Runtime Stream CreationAfter compilation determines stream count, runtime creates corresponding device streams during model loading phase.V2 Path (gert::StreamAllocator)gert::StreamAllocatorpre-allocates a continuous vector that can hold up to 2024rtStream_t. On firstAcquireStreamscall, it creates the specified number of streams in sequence (callingrtStreamCreateWithFlags). Subsequent calls reuse already created streams. Stream destruction completes uniformly in the destructor.V1 Path (ge::ReusableStreamAllocator)V1 path supports cross-model stream reuse.ReusableStreamAllocatormaintains a stream pool keyed bypriority, stream_flag. Each stream records itstask_numand the model ID set using it. When a new model requests streams:Search stream pool for matchingpriority, stream_flagFilter streams not belonging to current modelSort bytask_num, preferentially reuse streams with similar task countCreate new stream if no available stream5 Key Design Decisions5.1 Why Do Static and Dynamic Use Different Stream Allocation Strategies?Static Shape graph topology is completely known at compilation time. GE can precisely analyze all data dependency relationships and perform fine-grained stream reuse. Dynamic Shape graph structure is incomplete at compilation time (some subgraphs only expand at runtime), preventing precise dependency analysis. Therefore, Dynamic Shape adopts a more conservative strategy: default single-stream, only enabling multi-stream when user explicitly configures, with stream allocation rules primarily at engine granularity rather than dependency analysis.5.2 Why Use Pass Chain Instead of Single Allocation Algorithm?The core advantage of Pass chain architecture isextensibilityandmaintainability. Each stream allocation rule (label allocation, engine allocation, dependency allocation, AllReduce parallelism, etc.) is encapsulated as an independent Pass, each maintaining its own state without interference. Adding new stream allocation rules only requires adding new Passes to the chain, without modifying existing logic. If using a single algorithm, all rules would be interwoven, drastically reducing code readability and maintainability.5.3 Why Need Graph Structure Stability Principle?Stream allocation depends on topological sort and topo ID continuity. If graph structure changes during stream allocation (adding or deleting nodes), topo ID becomes discontinuous, affecting subsequent memory reuse (which allocates memory blocks based on topo order) and physical stream splitting (which calculates task count based on topo order). Therefore, stream allocation phase strictly prohibits graph modification operations. All nodes needing insertion (such as Send/Recv, StreamActive) are inserted uniformly after stream allocation completes.5.4 Why Must Attached Stream Be Allocated After Main Stream Allocation Completes?Attached stream IDs are incremented based on main stream IDs. If attached streams and main streams are allocated together, stream IDs become discontinuous, increasing synchronization management complexity. Allocating attached streams all at once after main stream allocation completes ensures stream ID continuity and predictability.6 Constraints and LimitationsConstraintDescriptionSingle-stream mode does not support StreamLabelSingle-stream mode has only one stream, StreamLabel causes conflictEvent ID must be continuousRTS has Event ID continuity validation, must ensure throughRefreshContinuousEventsGraph structure immutableGraph structure cannot change during stream allocation, topo ID must be continuousNotify count upper limitMaximum 1024 Notify supportedSingle node attached stream countOne node supports at most one attached stream (reuse through reuse_key)Multi-thread safetyStreamAllocatorsupports multi-threading but needs to protect shared resources;ScalableAllocatordoes not support multi-thread concurrencyDynamic graph While constraintNetOutput of While operators static body subgraph in dynamic graph must be on stream 0Multi-dim scenario label constraintWhen adding StreamLabel in multi-dim scenarios, need to add dim information to distinguish【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考