GE 自定义算子架构设计

GE 自定义算子架构设计【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 简介1.1 目的本文档描述 GE 自定义算子接入机制的架构设计面向 GE 内部开发者和架构师。文档覆盖自定义算子的接口体系、注册与加载机制、编译期和运行期的内部流程以及与 GE 各子系统的交互方式。面向外部算子开发者的使用指南参见custom_op_development_guide.md。1.2 范围本文档覆盖自定义算子的接口设计与注册机制SO 交付件加载与生命周期管理自定义算子在 GE 编译器和执行器中的调度路径前端PyTorch / TensorFlow / ONNX接入架构设计约束与特性交叉分析不覆盖具体算子的 kernel 实现Ascend C / Triton 等GE 内置算子的引擎调度细节2. 总体概述2.1 设计动机GE 对自定义算子接入有两个核心诉求语言无关算子开发不拘泥于特定编程语言Ascend C、Triton、PPTO 等通过统一的接入接口将算子集成过程与具体编成语言解耦。渐进式开发体验开发者可按需逐步补充能力执行 → 下沉 → 编译优化 → 离线 OM每完成一个阶段即可获得对应的性能收益而非一次性完成全部接入工作。2.2 设计目标目标说明语言无关接口层不假设 kernel 的实现语言只关心 kernel binary 的加载和 launch渐进式能力按需组合能力接口从最小可运行到全量下沉逐步演进交付件类组织一个 .so 包含一个或多个算子的全部接入逻辑便于分发和维护与内置算子共存自定义算子与内置算子在同一张图中混合执行共享流分配和内存规划基础设施定位接口层作为基础设施各编成语言可构建公共层进一步降低开发难度2.3 三阶段演进路线阶段核心能力新增交付件性能收益状态阶段 1Executehost 调度 kernel1 个 .so可运行有 host 调度开销已完成阶段 2.1Execute下沉调度无新增静态 shape 下消除 host 调度开销已完成阶段 2.2 InferShape Compile无新增shape 推导、内存复用、算子在线编译已完成阶段 3 Serialize / Deserialize无新增离线 OM 部署已完成2.4 基础设施定位与语言公共层当前的接口体系BaseCustomOp 5 个能力接口定位为基础设施层而非面向终端算子开发者的最终 API。各算子编成语言Ascend C、Triton、PPTO 等可在此基础设施之上构建语言公共层封装重复的 boilerplate 逻辑使单个算子的入图开发量降至最低。分层职责层次职责维护方GE 基础设施层提供统一的算子接入接口、注册机制、编译/执行回调、序列化协议GE 团队语言公共层封装特定编成语言的重复逻辑binary 加载、args 构造、kernel launch 等提供简化的注册宏或 API各语言 SDK 团队算子开发者只需实现 kernel 逻辑 少量声明输入输出规格、tiling 参数等算子开发者语言公共层可封装的典型 boilerplate重复逻辑当前基础设施层封装后语言公共层kernel binary 加载开发者手动调用aclrtBinaryLoadFromDataaclrtBinaryGetFunction公共层自动加载开发者只需声明 kernel 名称args 构造开发者手动拼装 packed struct处理对齐和指针公共层根据 kernel 签名自动生成 argsblock_num 计算开发者手动根据 shape 和 BLOCK_SIZE 计算 grid公共层提供 tiling 辅助函数或自动计算REG_OP REG_AUTO_MAPPING_OP开发者手动编写 proto 定义和注册宏公共层提供声明式 API自动生成 proto 和注册代码ShapeInferOp 实现开发者手动实现 InferShape / InferDataType公共层提供常见推导模式same-as-input、broadcast 等的模板设计意义通过将基础设施层与语言公共层分离GE 保持了接口稳定性和语言无关性同时各语言 SDK 可以独立演进、竞争优化最终让算子开发者只需关注 kernel 本身。2.5 与老版本自定义算子机制的对比在新版本BaseCustomOp机制之前GE 通过IMPL_OP宏 OpImplRegisterV2链式注册来接入自定义算子。理解两者的差异有助于把握新版本的设计意图。老版本机制概述老版本采用无状态函数指针组合模式算子开发者注册一组函数指针框架全权调度// 老版本注册方式 IMPL_OP(MyOp) .InferShape(my_infer_shape_func) // shape 推导函数指针 .InferShapeRange(my_shape_range_func) // 动态 shape 范围推导 .InferDataType(my_infer_datatype_func) // dtype 推导函数指针 .InferFormat(my_infer_format_func) // 格式推导函数指针 .Tiling(my_tiling_func) // AICore tiling 函数 .TilingParseMyCompileInfo(my_parse_func) // tiling 解析函数 .InputsDataDependency({0}) // 声明哪些输入需要数据 .PrivateAttr(attr_name, int64_t(42)); // 私有属性老版本面向AICore 标准引擎流水线设计算子提供 tiling 函数和 shape 推导函数框架的 FE 引擎和 DavinciModel 代劳编译、序列化、调度和地址刷新。核心差异谁控制新老版本的关键差异不是有没有某项能力而是谁控制这项能力能力老版本IMPL_OP新版本BaseCustomOp执行框架内嵌引擎调度 tiling kernel launch用户自定义EagerExecuteOp::Execute()在线编译框架内嵌FE 引擎调度 TBE/AscendC 编译器用户自定义CompilableOp::Compile()序列化框架内嵌FE 引擎自动序列化 tiling data / compile info用户自定义PortableOp::Serialize/Deserialize地址刷新框架内嵌DavinciModel SinkTask 自动处理用户自定义ArgsUpdater::UpdateHostArgs()其他接口层面的差异维度老版本IMPL_OP新版本BaseCustomOp注册风格链式函数指针虚函数继承 工厂注册状态管理无状态函数指针有状态对象成员变量Shape 推导InferShapeKernelFunc函数指针ShapeInferOp::InferShape()虚方法Shape 范围推导InferShapeRangeKernelFunc独立注册复用InferShape-1 表示 unknown 维DataType 推导InferDataTypeKernelFunc函数指针ShapeInferOp::InferDataType()虚方法格式推导InferFormatKernelFunc独立注册暂时不支持设计中数据依赖声明显式.InputsDataDependency({indices})隐式通过 context 访问私有属性.PrivateAttr(name, value)不支持Tiling专用.Tiling()/.TilingParse()接口弱化tiling的概念由子类自主处理设计面向AICore 内置算子标准引擎流水线外部自定义算子任意 kernel binarykernel 语言Ascend C / TBE框架编译任意用户自行编译或 RTC设计演进动机老版本面向标准引擎流水线设计算子必须走 TBE/AscendC → tiling → kernel launch 的标准路径框架全权代劳编译、序列化和调度。这对内置算子很高效但无法接纳非标准 kernel如 Triton npubin、第三方预编译 binary。新版本将控制权交给算子开发者通过虚函数接口暴露执行、编译、序列化和地址刷新的控制权使任意 kernel binary 都能接入 GE。代价是开发者需要自行实现这些环节但可通过语言公共层封装。两者互补而非替代老版本适合标准 AICore 算子框架全托管开发量小新版本适合非标准 kernel用户全控制灵活性高。两套机制在 GE 中共存各自服务于不同的算子接入场景。3. 系统架构3.1 架构视图3.2 与内置算子的关系自定义算子通过独立的自定义引擎DNN_VM_CUSTOM与内置算子区分维度内置算子自定义算子引擎DNNEngineAiCore / VectorCore / AICPUDNN_VM_CUSTOM算子注册算子仓 REG_OP 引擎 OpsKernelInfoStoreREG_OP REG_AUTO_MAPPING_OP CustomOpsKernelInfoStoreKernel 构建TBE / AICPU KernelBuilderCustomOpsKernelBuilder生成 MODEL_TASK_CUSTOM_KERNEL编译优化GE 图优化 pass 引擎内部优化CustomGraphOptimizer回调 Compile执行调度引擎 TaskInfoTBE / AICPUCustomTaskInfo回调 Execute流分配独立引擎流与 AiCore 合并分配流关键设计决策自定义算子在流分配阶段与 AiCore 引擎节点合并处理engine_partitioner.cc确保自定义算子能正确参与多流并行调度而不是被隔离到独立流上。4. 核心组件设计4.1 接口体系设计原则能力组合开发者按需继承不强制实现所有接口。GE 通过dynamic_cast在运行时检测算子具备哪些能力。正交设计每个接口对应独立的回调时机接口之间无耦合。Context 隔离每个回调接收专用的 Context 对象只暴露算子需要的信息。能力检测机制auto *base CustomOpFactory::CreateOrGetCustomOp(op_type); auto *compilable dynamic_castCompilableOp*(base); if (compilable ! nullptr) { compilable-Compile(ctx); }4.2 注册与工厂机制实例化策略CustomOpFactoryImpl对每个 op type 采用懒加载单例模式。同一 op type 的所有图节点共享同一个BaseCustomOp实例。设计约束实现类的成员变量是跨节点共享的如device_elves_mapCompile回调可能被并行调用CustomGraphOptimizer使用线程池实现需保证线程安全注册是一次性的重复注册同一 op type 会返回GRAPH_FAILED4.3 SO 加载机制加载限制PluginManager 强制限制项上限.so 文件数量64 个单个 .so 大小800 MB总加载大小1000 MB4.4 自定义引擎DNN_VM_CUSTOM自定义算子通过独立的引擎组件接入 GE 编译流程组件职责关键文件CustomOpsKernelInfoStore初始化时查询已注册 op type生成 OpInfocompiler/engines/custom_engine/custom_ops_kernel_info_store.ccCustomGraphOptimizer并行遍历自定义算子节点回调 Compilecompiler/engines/custom_engine/custom_graph_optimizer.ccCustomOpsKernelBuilder生成 MODEL_TASK_CUSTOM_KERNEL TaskDefcompiler/engines/custom_engine/custom_ops_kernel_builder.cc5. 关键流程5.1 SO 加载与注册流程5.2 编译期流程5.3 执行期流程V1 静态执行器Known ShapeV2 动态执行器Unknown Shape / RT2.05.4 序列化/反序列化流程序列化互斥约束一张图中不能同时包含实现了PortableOp和未实现PortableOp的自定义算子。6. 前端接入架构6.1 GE 原生构图最简单的路径。构图侧直接引用 REG_OP proto 头文件通过OperatorFactory::CreateOperator创建节点。6.2 PyTorch TorchAir6.3 TensorFlow与 GE 原生构图不同TensorFlow 场景下REG_AUTO_MAPPING_OP可从 TF 算子原型自动生成 GE 算子原型开发者无需额外编写REG_OP。6.4 ONNXONNX 解析插件通过REGISTER_CUSTOM_OP注册dlopen 时自动收集到OpRegistry。插件需实现ParseParamsFn将 ONNXNodeProto的属性映射为 GEOperator的属性。7. 设计约束与不变量约束说明影响单例实例 per op typeCreateOrGetCustomOp为每个 op type 创建唯一实例成员变量跨节点共享Compile并行调用需线程安全dynamic_cast 能力检测GE 通过dynamic_cast判断算子支持哪些接口未实现的接口自动跳过不影响其他流程注册一次性RegisterCustomOpCreator拒绝重复注册同一进程中不可注册同名 op type序列化互斥图中不可混合可序列化和不可序列化的自定义算子OM 下沉场景所有自定义算子必须实现PortableOpSO 加载限制最多 64 个 .so单个 ≤ 800MB总计 ≤ 1000MB大量自定义算子需合并到少量 .so 中共享 AiCore 流自定义算子在流分配时与 AiCore 合并自定义算子可参与 AiCore 的多流并行地址刷新实现ArgsUpdater的算子走预留内存分配路径零拷贝场景需实现ArgsUpdater8. 特性交叉分析场景适用性分析说明静态 Shape适用阶段 2.1 下沉调度消除 host 开销。CustomTaskInfo::Distribute()在 DavinciModel SinkTask 流程中调用 Execute。输出 tensor size 参与逻辑内存复用。动态 Shape适用阶段 1 核心场景。V2 执行器通过LoweringCustomNode生成FindCustomOpExecuteCustomOpkernel运行时 host 调度执行。动态 Shape 静态子图适用静态子图走 V1DavinciModelKernel路径自定义算子通过CustomTaskInfo::Distribute()执行与纯静态 shape 场景一致。离线场景atc 编译适用阶段 3 覆盖。ATC 编译时回调 Compile SerializeOM 加载时回调 Deserialize Execute。在线场景框架适配适用PyTorch/TorchAir 和 TensorFlow 通过各自 Adapter 映射为 GE 自定义算子节点走统一编译执行流程。9. 非功能需求9.1 性能阶段编译期影响执行期影响阶段 1host 调度CustomOpsKernelInfoStore::Initialize遍历已注册 op typeO(n)每个自定义算子节点有 host 侧 Execute 调用开销阶段 2.1下沉无额外编译开销消除 host 调度开销kernel 直接在下沉流中执行阶段 2.2全量CustomGraphOptimizer并行回调 Compile增加编译时间shape 推导和内存复用带来执行期收益阶段 3离线 OMSerialize 增加 OM 保存时间Deserialize 增加 OM 加载时间执行期无额外开销9.2 兼容性OM 前向兼容新版本 GE 可加载旧版本 OM 中的自定义算子序列化数据接口兼容BaseCustomOp及其子接口均为纯虚函数新增接口不影响已有实现9.3 可维护性自定义算子代码完全在 GE 仓外维护通过 .so 动态加载接口变更需同步更新custom_op.h头文件属于 GE 公共 APIContext 类的扩展需保持 POD 布局兼容【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考