GE 静态执行器Known Shape Executor特性分析【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 特性背景静态执行器是 GE 运行时中负责加载和执行静态shape图或者子图的核心组件。它以DavinciModel类为载体提供两种使用模式独立静态模型模式整个模型编译为单个 OM 文件通过ModelManager直接加载执行V2 Runtime Kernel 集成模式在新一代 V2 运行时中静态子图以 Kernel 的形式嵌入执行图通过 Kernel 注册机制管理生命周期创建、参数刷新、执行、工作区更新实现与动态子图的协同调度2. 用户使用场景2.1 场景一离线模型推理ACL 模式用户通过 ATC 工具将模型转换为 OM 文件后在推理侧通过 ACL API 加载并执行模型文件(.om) → aclmdlLoadFromFile → aclmdlExecute → 输出结果底层执行路径ModelManager::LoadModelOffline→DavinciModel::Init→DavinciModel::NnExecute2.2 场景二通过 GeSession 接口执行用户通过 GE V2 API 提供的GeSession类加载和执行计算图这是 GE 运行时的高层编程接口2.3 场景三V2 Runtime Kernel 集成在新一代 V2 运行时中静态子图以 Kernel 的形式集成到执行图中通过注册的 Kernel 函数完成生命周期管理DavinciModelCreate创建并初始化 DavinciModelDavinciModelUpdateArgs刷新输入输出地址DavinciModelExecute触发设备侧执行DavinciModelUpdateWorkspaces更新工作内存基址DavinciModelGetRunAddress获取运行时内存地址用于 Kernel 间地址依赖3. 对外接口3.1 DavinciModel 核心接口DavinciModel是静态执行器的核心类封装了模型加载、内存管理、任务分发和执行的全流程加载阶段Init(ModelParam, outer_fm_mem)模型初始化入口完成内存分配、资源创建、任务下沉Task SinkSetKnownNode(bool)标记是否为混合执行模式下的已知形状子图节点执行阶段NnExecute(stream, async_mode, input_tensor, output_tensor)模型执行入口UpdateKnownNodeArgs(inputs, outputs)为已知形状子图刷新输入输出地址混合执行模式专用CopyModelData拷贝输入数据到设备侧CopyOutputData拷贝输出数据回用户缓冲区资源管理GetRtModelHandle()获取底层 RT 模型句柄GetLogicalMemAllocation()获取逻辑内存分配表UpdateHbmFmMemBases更新 HBM 特征图内存基址3.2 V2 Kernel 注册接口REGISTER_KERNEL(DavinciModelCreate) // 创建模型 REGISTER_KERNEL(DavinciModelCreateV2) // V2 版本创建 REGISTER_KERNEL(DavinciModelUpdateArgs) // 更新参数 REGISTER_KERNEL(DavinciModelExecute) // 执行模型 REGISTER_KERNEL(DavinciModelUpdateWorkspaces) // 更新工作区 REGISTER_KERNEL(DavinciModelGetRunAddress) // 获取运行时地址4. 架构设计4.1 类层次结构Executor (base/common/model/executor.h) ├── ModelExecutor (runtime/v1/graph/execute/model_executor.h) │ └── 委托给 ModelManager - DavinciModel │ DavinciModel (runtime/v1/graph/load/model_manager/davinci_model.h) ├── 独立模型: known_node_ false └── 已知形状子图: known_node_ true └── V2 runtime kernel5. 核心实现5.1 模型加载流程DavinciModel::Init模型加载是静态执行器最复杂的阶段涉及内存分配、资源创建、任务分发等多个子步骤。DavinciModel::Init() │ ├── InitRuntimeParams() │ 从 GeModel 中提取运行时参数内存布局、任务定义、流配置等 │ ├── InitWeightMem() │ 加载权重数据到设备侧 HBM │ ├── InitFixedFeatureMap() │ 设置固定不可刷新特征图内存 │ 这部分内存在模型生命周期内地址不变 │ ├── InitFeatureMapAndP2PMem() │ 设置可刷新特征图内存 │ 支持分段内存sub memory管理 │ ├── PreProcessFileConstants() │ 处理外部权重文件FileConstant 节点 │ 支持权重合并Combined Weights优化 │ ├── InitIoNodes() │ 初始化 Data 和 NetOutput 节点 │ 配置零拷贝Zero-Copy内存映射 │ ├── InitRuntimeResource() │ 创建 RT 模型句柄rtModel_t │ 创建执行流aclrtStream、事件Event、标签Label │ ├── TransAllVarData() │ 将变量Variable数据传输到设备侧 │ ├── InitNodes() │ 初始化所有计算节点 │ 加载 TBE Kernel 句柄、注册算子实现空间 │ └── DoTaskSink() ★ 关键步骤任务下沉到设备 │ ├── BindModelStream() │ 将逻辑流绑定到物理 RT Stream │ ├── InitTaskInfo() │ 从 ModelTaskDef 创建 TaskInfo 对象 │ 初始化 ModelArgsManager参数管理器 │ ├── LoadWithQueue() │ 如果配置了队列调度设置队列执行路径 │ ├── DistributeTask() │ 通过 rtPersistentTaskLaunch 将任务分发到设备侧 │ 所有 Kernel 的 Launch 参数在此阶段预置到设备 │ ├── UpdateStaticModelArgsByFm() │ 用特征图地址初始化参数刷新表 │ └── aclmdlRIBuildEnd() 标记 RT 模型构建完成Task Sink 的设计意义将编译期生成的所有任务Task预先分发到设备侧执行时只需触发rtModelExecute无需每次执行都进行 Kernel Launch。这是静态执行器高性能的核心保障——消除了 Host 侧的 Kernel 启动开销。5.2 模型执行流程DavinciModel::NnExecuteDavinciModel::NnExecute(stream, async_mode, input_tensor, output_tensor) │ ├── InitModelStream(stream) │ 设置执行流 │ ├── CopyModelData(input_tensor, output_tensor) │ │ │ ├── UpdateAllNodeArgs() │ │ 更新所有 Kernel 的 Launch 参数 │ │ 包括输入输出地址、形状信息等 │ │ │ └── CopyInputForNoZeroCopy() │ 对于非零拷贝的输入执行 H2D 数据拷贝 │ ├── rtModelExecute(rt_model_handle_, rt_model_stream_, 0U) ★ 设备侧执行 │ 或 rtModelExecuteSync()MDC 场景带超时控制 │ ├── rtStreamSynchronizeWithTimeout() │ 等待执行完成内建流场景 │ ├── CopyOutputData(output_tensor) │ 将输出数据拷贝回用户缓冲区 │ 零拷贝模式下跳过此步骤 │ └── UpdateOutputTensorShape() 更新输出张量形状动态 Shape 场景5.3 地址刷新机制核心创新在混合执行模式下已知形状子图的输入输出地址在每次迭代时都可能变化。静态执行器通过逻辑内存分配表 活跃地址映射的机制实现高效的地址刷新。5.3.1 数据结构DavinciModel维护以下关键数据结构logical_mem_allocations_逻辑内存分配表记录每个逻辑内存区域的类型INPUT/OUTPUT/FEATURE_MAP、大小、命中次数等元信息allocation_ids_to_active_base_addr_活跃地址映射表将 allocation_id 映射到当前执行的实际设备地址refreshable_input_index_and_allocation_ids_可刷新输入索引到 allocation_id 的映射refreshable_output_index_and_allocation_ids_可刷新输出索引到 allocation_id 的映射refreshable_fm_index_and_allocation_ids_可刷新特征图索引到 allocation_id 的映射5.3.2 刷新流程DavinciModel::UpdateKnownNodeArgs(inputs, outputs) │ ├── ConstructActiveMemBaseAddrsForKnownNode(ret_up, inputs, outputs) │ │ │ ├── 更新 FM 地址 │ │ 遍历 refreshable_fm_index_and_allocation_ids_ │ │ 将 runtime_param_.fm_memory_infos 中的地址写入活跃地址表 │ │ │ ├── 更新输入地址 │ │ 遍历 refreshable_input_index_and_allocation_ids_ │ │ 将用户传入的 inputs[i] 设备地址写入活跃地址表 │ │ 首次执行时包含非冻结输入后续使用 zero_copy_no_frozen │ │ │ └── 更新输出地址 │ 遍历 refreshable_output_index_and_allocation_ids_ │ 将用户传入的 outputs[i] 设备地址写入活跃地址表 │ └── args_manager_.UpdateForExecute(ret_up, rt_model_stream_) 将更新后的活跃地址表拷贝到设备侧 通过 UpdateModelParam Kernel 实现高效刷新 ret_up 决定刷新策略全量刷新 vs 增量刷新设计精妙之处增量刷新策略ret_up变量记录需要刷新的最大策略级别args_manager_根据此值决定只拷贝发生变化的地址最小化 H2D 带宽消耗冻结输入优化对于地址不变的输入Frozen Inputs在首次执行后从刷新列表中排除避免不必要的地址更新零拷贝支持用户提供的 I/O 缓冲区直接映射到设备 Kernel 参数无需中间拷贝5.4 known_node_ 标志的影响known_node_是区分独立静态模型和混合执行模式下已知形状子图的关键标志。设置此标志后通过SetKnownNode(true)DavinciModel的行为发生以下变化行为known_node_ falseknown_node_ trueSession ID 获取使用runtime_param_.graph_id使用runtime_param_.root_graph_id地址刷新使用UpdateAllNodeArgs使用UpdateKnownNodeArgs特征图基址固定不可刷新可刷新feature_base_refreshable_ true错误追踪清理执行跳过变量初始化标准路径特殊路径内存分段可能合并为单段保持分段结构5.5 ModelArgsManager 参数管理ModelArgsManager是静态执行器中管理 Kernel 启动参数的核心组件负责初始化阶段从ModelTaskDef中解析所有任务的参数布局建立逻辑地址到设备参数的映射执行阶段根据活跃地址表更新设备侧参数通过UpdateModelParamKernel 实现高效刷新策略管理维护id_to_policy映射支持全量刷新all-one-time和增量刷新两种策略5.6 内存管理策略静态执行器采用分层内存管理策略设备侧内存布局 │ ├── 权重内存weights_mem_base_ │ 模型权重数据加载后地址固定 │ ├── 固定特征图内存fixed_mem_base_ │ 不可刷新的特征图内存 │ 在模型生命周期内地址不变 │ ├── 可刷新特征图内存mem_base_ │ 支持运行时地址刷新 │ 分段场景下为首个 refreshable FM 段的地址 │ ├── 零拷贝 I/O 内存 │ 用户提供的输入输出缓冲区 │ 地址通过 args_manager_ 刷新 │ └── 变量内存var_mem_base_ 模型变量如 BatchNorm 的 running mean/var6. 编译器侧支持6.1 动静 Shape 图划分DynamicShapePartitioner负责将计算图划分为 KNOWN_SHAPE 和 UNKNOWN_SHAPE 集群DynamicShapePartitioner::Partition() │ ├── MarkUnknownShapeNodes() │ 标记所有包含未知维度-1或未知秩-2的节点 │ ├── InitClusters() │ 为每个节点创建 Cluster │ 类型包括DATA / KNOWN_SHAPE / UNKNOWN_SHAPE / NETOUTPUT │ ├── MergeClusters() │ │ │ ├── MergeClustersUnknownShape() │ │ 如果两个 UNKNOWN_SHAPE 集群相连合并它们 │ │ 合并路径上的所有 KNOWN_SHAPE 集群也会被吞并 │ │ │ ├── MergeClustersNormal() │ │ 如果两个 KNOWN_SHAPE 集群之间只有一条路径合并它们 │ │ │ └── MergeClustersInputData() │ 合并所有 INPUT_DATA 集群 │ └── PruneUniqueClusters() 去重合并后的集群合并规则的关键约束UNKNOWN_SHAPE 集群具有传染性——如果两个未知形状节点之间存在路径路径上的所有已知形状节点都会被标记为未知形状。这是因为已知形状节点的输出可能作为未知形状节点的输入需要统一管理。6.2 已知形状图编译GraphBuilder::BuildForKnownShapeGraph()负责编译 KNOWN_SHAPE 集群生成完整的ModelTaskDef包含所有任务的定义计算精确的内存分配方案MemAllocation生成零拷贝偏移信息ZeroCopyOffset输出GeModel对象包含编译后的图信息和任务定义7. V2 Runtime Kernel 集成V2 运行时将静态子图作为 Kernel 集成到执行图中提供了更细粒度的控制7.1 DavinciModelCreate创建并初始化 DavinciModel 实例从输入获取GeModel对象创建DavinciModel实例设置known_node_true设置 Session ID、Root Graph ID、Step ID 等上下文信息初始化权重内存和特征图内存调用DavinciModel::Init()完成加载将 DavinciModel 指针输出到下游 Kernel7.2 DavinciModelUpdateArgs每次执行前刷新输入输出地址从 KernelContext 获取输入输出 Tensor 的设备地址构造vectoruint64_t地址列表调用DavinciModel::UpdateKnownNodeArgs()刷新地址7.3 DavinciModelExecute触发设备侧执行先调用DavinciModelUpdateArgs刷新地址调用rtModelExecute触发执行调用CopyOutputData拷贝输出数据7.4 DavinciModelUpdateWorkspaces更新工作内存基址从 KernelContext 获取 workspace 地址和内存类型调用DavinciModel::UpdateHbmFmMemBases()更新 HBM 内存调用DavinciModel::UpdateExMemBase()更新其他类型内存7.5 DavinciModelGetRunAddress获取运行时内存地址用于后续 Kernel 的地址依赖根据MemoryBaseTypeOffset内存类型 偏移量查询实际运行地址支持 Weight、FileConstant 等内存类型将地址写入输出 Tensor8. 关键设计决策分析8.1 为什么选择 Task Sink rtModelExecute 架构替代方案对比方案Host 侧开销设备侧开销灵活性每次执行 Kernel Launch高每次都要 Launch低高Task Sink rtModelExecute低仅地址刷新低中全图编译为单个 Kernel最低最低低GE 选择 Task Sink rtModelExecute 的原因是编译期确定性KNOWN_SHAPE 子图的所有参数在编译期已知可以安全地预分发任务执行效率消除 Host 侧 Kernel Launch 开销rtModelExecute 仅触发预置的任务链地址刷新灵活性通过ModelArgsManager实现高效的地址刷新支持动态 I/O 地址8.2 为什么引入 known_node_ 标志而不是创建新类在混合执行模式下DavinciModel通过known_node_标志切换行为而非创建独立的子类。这种设计的选择优点代码复用最大化加载、任务分发、执行等核心逻辑完全共享维护成本低只需在差异点做条件分支内存布局一致两种模式使用相同的内存管理结构代价类职责不够单一一个类同时承担独立模型和子图两种角色条件分支增加复杂度代码中散布if (known_node_)判断代码中的注释也反映了这一点——UpdateKnownNodeArgs方法上的// todo 临时方案注释表明未来可能通过重构来改善这一设计。8.3 地址刷新机制的设计权衡地址刷新机制是静态执行器的核心创新其设计面临以下权衡全量刷新 vs 增量刷新全量刷新简单可靠但 H2D 带宽浪费增量刷新通过ret_up策略级别控制只刷新变化的地址但实现复杂GE 选择了增量刷新策略通过active_mem_base_id_to_plicy映射表记录每个 allocation 的刷新策略级别在UpdateForExecute时根据ret_up决定实际拷贝的数据量。零拷贝 vs 中间拷贝零拷贝用户缓冲区直接映射地址刷新即可无数据拷贝开销中间拷贝GE 内部管理 I/O 缓冲区需要额外的 H2D/D2H 拷贝GE 优先使用零拷贝模式仅在用户缓冲区不满足对齐要求或设备不可达时回退到中间拷贝。9. 性能优化要点9.1 Task Sink 预分发所有 Kernel 任务在模型加载时预分发到设备侧执行时无需 Kernel Launch。这是静态执行器相比动态执行器最大的性能优势来源。9.2 增量地址刷新通过ModelArgsManager的增量刷新策略最小化每次执行时的 H2D 数据传输量。对于 Frozen Inputs地址不变的输入首次执行后不再参与刷新。9.3 零拷贝 I/O用户提供的设备缓冲区直接映射到 Kernel 参数避免中间拷贝。在训练场景和大批量推理场景中收益显著。9.4 流复用通过ReusableStreamAllocator实现 Stream 复用减少 Stream 创建和销毁的开销。在多模型并发加载场景下尤为重要。9.5 Shrink 优化模型加载完成后调用Shrink()释放 Host 侧的GeModel对象减少内存占用。因为所有必要信息已分发到设备侧Host 侧的图结构不再需要。10. 文件清单运行时核心文件路径功能runtime/v1/graph/load/model_manager/davinci_model.hDavinciModel 类定义runtime/v1/graph/load/model_manager/davinci_model.ccDavinciModel 实现约 9281 行runtime/v1/graph/load/model_manager/model_manager.hModelManager 单例定义runtime/v1/graph/load/model_manager/model_manager.ccModelManager 实现runtime/v1/graph/load/model_manager/model_args_manager.h参数管理器定义V2 Kernel文件路径功能runtime/v2/kernel/known_subgraph/davinci_model_kernel.ccV2 Kernel 集成Create/Execute/UpdateArgs编译器文件路径功能compiler/graph/partition/dynamic_shape_partition.hDynamicShapeCluster/Partitioner 定义compiler/graph/partition/dynamic_shape_partition.cc已知/未知形状图划分逻辑基础接口文件路径功能base/common/model/executor.hExecutor 抽象接口【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考