GE 自定义算子架构设计

GE 自定义算子架构设计

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

1. 简介

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 对自定义算子接入有两个核心诉求：

1. 语言无关：算子开发不拘泥于特定编程语言（Ascend C、Triton、PPTO 等），通过统一的接入接口将算子集成过程与具体编成语言解耦。
2. 渐进式开发体验：开发者可按需逐步补充能力（执行 → 下沉 → 编译优化 → 离线 OM），每完成一个阶段即可获得对应的性能收益，而非一次性完成全部接入工作。

2.2 设计目标

2.3 三阶段演进路线

2.4 基础设施定位与语言公共层

当前的接口体系（BaseCustomOp+ 5 个能力接口）定位为基础设施层，而非面向终端算子开发者的最终 API。各算子编成语言（Ascend C、Triton、PPTO 等）可在此基础设施之上构建语言公共层，封装重复的 boilerplate 逻辑，使单个算子的入图开发量降至最低。

分层职责：

语言公共层可封装的典型 boilerplate：

设计意义：通过将基础设施层与语言公共层分离，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 函数 .TilingParse<MyCompileInfo>(my_parse_func) // tiling 解析函数 .InputsDataDependency({0}) // 声明哪些输入需要数据 .PrivateAttr("attr_name", int64_t(42)); // 私有属性

老版本面向AICore 标准引擎流水线设计：算子提供 tiling 函数和 shape 推导函数，框架的 FE 引擎和 DavinciModel 代劳编译、序列化、调度和地址刷新。

核心差异："谁控制"

新老版本的关键差异不是"有没有"某项能力，而是谁控制这项能力：

其他接口层面的差异：

设计演进动机

1. 老版本面向标准引擎流水线设计：算子必须走 TBE/AscendC → tiling → kernel launch 的标准路径，框架全权代劳编译、序列化和调度。这对内置算子很高效，但无法接纳非标准 kernel（如 Triton npubin、第三方预编译 binary）。

2. 新版本将控制权交给算子开发者：通过虚函数接口暴露执行、编译、序列化和地址刷新的控制权，使任意 kernel binary 都能接入 GE。代价是开发者需要自行实现这些环节（但可通过语言公共层封装）。

3. 两者互补而非替代：老版本适合标准 AICore 算子（框架全托管，开发量小），新版本适合非标准 kernel（用户全控制，灵活性高）。两套机制在 GE 中共存，各自服务于不同的算子接入场景。

3. 系统架构

3.1 架构视图

3.2 与内置算子的关系

自定义算子通过独立的自定义引擎DNN_VM_CUSTOM与内置算子区分：

关键设计决策：自定义算子在流分配阶段与 AiCore 引擎节点合并处理（engine_partitioner.cc），确保自定义算子能正确参与多流并行调度，而不是被隔离到独立流上。

4. 核心组件设计

4.1 接口体系

设计原则：

• 能力组合：开发者按需继承，不强制实现所有接口。GE 通过dynamic_cast在运行时检测算子具备哪些能力。
• 正交设计：每个接口对应独立的回调时机，接口之间无耦合。
• Context 隔离：每个回调接收专用的 Context 对象，只暴露算子需要的信息。

能力检测机制：

auto *base = CustomOpFactory::CreateOrGetCustomOp(op_type); auto *compilable = dynamic_cast<CompilableOp*>(base); if (compilable != nullptr) { compilable->Compile(ctx); }

4.2 注册与工厂机制

实例化策略：CustomOpFactoryImpl对每个 op type 采用懒加载单例模式。同一 op type 的所有图节点共享同一个BaseCustomOp实例。

设计约束：

• 实现类的成员变量是跨节点共享的（如device_elves_map）
• Compile回调可能被并行调用（CustomGraphOptimizer使用线程池），实现需保证线程安全
• 注册是一次性的，重复注册同一 op type 会返回GRAPH_FAILED

4.3 SO 加载机制

加载限制（PluginManager 强制）：

4.4 自定义引擎（DNN_VM_CUSTOM）

自定义算子通过独立的引擎组件接入 GE 编译流程：

5. 关键流程

5.1 SO 加载与注册流程

5.2 编译期流程

5.3 执行期流程

V1 静态执行器（Known Shape）：

V2 动态执行器（Unknown Shape / RT2.0）：

5.4 序列化/反序列化流程

序列化互斥约束：一张图中不能同时包含实现了PortableOp和未实现PortableOp的自定义算子。

6. 前端接入架构

6.1 GE 原生构图

最简单的路径。构图侧直接引用 REG_OP proto 头文件，通过OperatorFactory::CreateOperator创建节点。

6.2 PyTorch + TorchAir

6.3 TensorFlow

与 GE 原生构图不同，TensorFlow 场景下REG_AUTO_MAPPING_OP可从 TF 算子原型自动生成 GE 算子原型，开发者无需额外编写REG_OP。

6.4 ONNX

ONNX 解析插件通过REGISTER_CUSTOM_OP注册，dlopen 时自动收集到OpRegistry。插件需实现ParseParamsFn，将 ONNXNodeProto的属性映射为 GEOperator的属性。

7. 设计约束与不变量

8. 特性交叉分析

9. 非功能需求

9.1 性能

9.2 兼容性

• OM 前向兼容：新版本 GE 可加载旧版本 OM 中的自定义算子序列化数据
• 接口兼容：BaseCustomOp及其子接口均为纯虚函数，新增接口不影响已有实现

9.3 可维护性

• 自定义算子代码完全在 GE 仓外维护，通过 .so 动态加载
• 接口变更需同步更新custom_op.h头文件，属于 GE 公共 API
• Context 类的扩展需保持 POD 布局兼容

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