GE 自定义算子开发指南【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge1. 概述自定义算子允许开发者将自有 kernelAscend C、Triton 或其他设备侧实现接入 GE 图编译和执行流程而无需修改 GE 框架代码。适用场景框架前端PyTorch / TensorFlow引入了 GE 不认识的算子需要在 NPU 上执行需要自定义 kernel 实现如融合算子、特定优化并接入 GE 编译和下沉流程需要在 GE 图中嵌入第三方 kernel binary与 GE 内置算子的区别维度内置算子自定义算子注册方式GE 内部 REG_OP 引擎注册REG_OP REG_AUTO_MAPPING_OP以 .so 交付件加载编译产物随 GE 发布独立 .so通过 ASCEND_CUSTOM_OPP_PATH 加载执行方式GE 引擎调度DNNEngine 等框架回调用户实现的 Execute / Compile更新方式随 GE 版本更新独立于 GE可随时替换 .so2. 核心概念2.1 接口体系所有自定义算子实现类继承自BaseCustomOp按需组合以下能力接口BaseCustomOp公共基类 ├── EagerExecuteOp → Execute(ctx) 运行时执行 ├── ShapeInferOp → InferShape(ctx) Shape 推导 │ → InferDataType(ctx) DataType 推导 ├── CompilableOp → Compile(ctx) 在线编译 ├── PortableOp → Serialize(buf) 序列化到 OM │ → Deserialize(buf) 从 OM 反序列化 └── ArgsUpdater → UpdateHostArgs(ctx) I/O 地址刷新每个接口对应一个独立的回调时机GE 在图编译或执行的特定阶段调用对应回调。2.2 注册机制自定义算子需要两个注册REG_OP— 定义算子的输入、输出和属性 proto供构图侧创建节点时使用REG_OP(MyCustomOp) .INPUT(x, TensorType({DT_FLOAT, DT_FLOAT16})) .INPUT(y, TensorType({DT_FLOAT, DT_FLOAT16})) .OUTPUT(z, TensorType({DT_FLOAT, DT_FLOAT16})) .OP_END_FACTORY_REG(MyCustomOp);REG_AUTO_MAPPING_OP— 将 C 实现类映射到 op typeGE 按 op type 名查找并实例化class MyCustomOp : public EagerExecuteOp { /* ... */ }; REG_AUTO_MAPPING_OP(MyCustomOp);命名约束REG_OP 中的 op type 名、REG_AUTO_MAPPING_OP 的类名、构图侧使用的 op type 必须三者一致。2.3 交付件与 OPP 包自定义算子以.so交付件形式被 GE 加载通过ASCEND_CUSTOM_OPP_PATH环境变量指定路径。推荐的 OPP 包目录结构custom_opp_root/ ├── op_graph/ │ ├── include/ │ │ └── my_custom_op.h // REG_OP proto 头文件 │ └── lib/ │ └── os/arch/ │ └── libcust_opapi.so // 交付件 .so └── framework/ // 框架侧交付件按需 └── tensorflow/ └── npu_supported_ops.json配置方式export ASCEND_CUSTOM_OPP_PATHcustom_opp_root:$ASCEND_CUSTOM_OPP_PATH3. 快速开始以下 5 步实现一个最小自定义算子以 Add 为例Step 1: 定义算子 proto// my_add.h REG_OP(AddCustom) .INPUT(x, TensorType({DT_FLOAT, DT_FLOAT16})) .INPUT(y, TensorType({DT_FLOAT, DT_FLOAT16})) .OUTPUT(z, TensorType({DT_FLOAT, DT_FLOAT16})) .OP_END_FACTORY_REG(AddCustom);Step 2: 实现执行逻辑class AddCustom : public EagerExecuteOp { graphStatus Execute(gert::EagerOpExecutionContext *ctx) override { auto *x ctx-GetInputTensor(0); auto *y ctx-GetInputTensor(1); auto *z ctx-MallocOutputTensor(0, x-GetShape(), x-GetFormat(), x-GetDataType()); // ... launch kernel with ctx-GetStream() ... return GRAPH_SUCCESS; } };Step 3: 注册REG_AUTO_MAPPING_OP(AddCustom);Step 4: 编译为 .socmake -S . -B build cmake --build build -j$(nproc)Step 5: 配置并运行export ASCEND_CUSTOM_OPP_PATH$(pwd)/build:$ASCEND_CUSTOM_OPP_PATH # 运行你的构图/执行程序完整可运行代码参见examples/custom_op/ascendc_add_custom/。4. 场景选择指南4.1 决策表场景推荐接口组合典型链路参考样例动态图在线执行EagerExecuteOpShapeInferOp(可选)构图 → 直接执行ascendc_add_custom在线编译 在线执行EagerExecuteOpCompilableOpShapeInferOp(可选)构图 → Compile → 执行—离线 OM 模型下沉EagerExecuteOpCompilableOpShapeInferOpPortableOpAIR → ATC → OM → ACLcompilable_add_custom数据依赖 shapeEagerExecuteOpShapeInferOp构图 → 执行 shape buffer 回写data_dependent_shape_custom4.2 场景 A动态图在线执行最简单的场景。kernel 已预编译好Ascend C.asc同库编译、Triton.npubin等只需在 Execute 中加载并 launch。构图 → GE 回调 Execute → 获取输入 → 分配输出 → launch kernel → 返回关键点只需实现EagerExecuteOpShapeInferOp 可选如果输出 shape 与输入相同且 REG_OP 已声明可省略kernel binary 随 .so 一起编译或在 Execute 中从文件加载4.3 场景 B在线编译 在线执行kernel 源码需要在 GE 编译阶段通过 RTCRuntime Compilation编译为 device binary。构图 → GE 回调 CompileRTC 编译 kernel → GE 回调 Execute加载 binarylaunch kernel关键点Compile 回调中可通过ctx-GetInputTensor(i)获取输入的 shape、dtype、format 等完整元信息建议按 shape 生成 binary key支持多 shape / 多 kernel 管理编译产物缓存在实现类的成员变量中供 Execute 使用4.4 场景 C离线 OM 模型下沉在场景 B 基础上需要将编译产物序列化到 OM 模型文件中后续加载 OM 时反序列化恢复。构图 → ATC 离线编译 → 回调 CompileRTC 编译 → 回调 Serializebinary 写入 OM → ACL 加载 OM → 回调 Deserialize从 OM 恢复 binary → 回调 Executelaunch kernel关键点必须实现PortableOpSerialize/Deserialize 的 buffer 格式由用户自定义GE 只透传不解析ShapeInferOp 在 OM 编译阶段被调用用于推导输出 shape/dtype支持多份 binary 的序列化按 key 管理4.5 场景 D数据依赖 shape三类算子输出 shape 依赖运行时输入数据如 Where、NonZero编译期只能确定上界。InferShape → 输出 shape 上界含 unknown 维 Execute → 按上界分配输出 分配 shape buffer → launch kernelkernel 写回实际 shape 到 shape buffer → stream sync → D2H 读回 shape buffer → 更新输出 shape关键点InferShape 中用-1表示 unknown 维Execute 中按最大 shape 分配输出额外分配 shape buffer通过MallocWorkSpacekernel 执行后需自行同步 stream 并读回 shape buffershape buffer 格式[ndim, d0, d1, ...]uint64_t 数组完整代码参见examples/custom_op/data_dependent_shape_custom/。5. 接口详解5.1 EagerExecuteOp::Execute运行时执行回调是最核心的接口。graphStatus Execute(gert::EagerOpExecutionContext *ctx) override;上下文 API方法用途ctx-GetInputTensor(index)获取输入 Tensor含 addr、shape、dtype、formatctx-MallocOutputTensor(index, shape, format, dtype)分配输出 Tensor device 内存ctx-MallocWorkSpace(size)分配 workspace device 内存ctx-GetStream()获取执行流ctx-GetOutputTensor(index)获取已分配的输出 Tensorctx-MakeOutputRefInput(out_idx, in_idx)输出引用输入地址零拷贝典型流程graphStatus Execute(gert::EagerOpExecutionContext *ctx) { auto *x ctx-GetInputTensor(0); auto *y ctx-MallocOutputTensor(0, x-GetShape(), x-GetFormat(), x-GetDataType()); void *stream ctx-GetStream(); LaunchMyKernel(x-GetAddr(), y-GetAddr(), stream); return GRAPH_SUCCESS; }注意事项输出内存由 context 管理生命周期调用者无需释放GetInputTensor返回const Tensor*不可修改输入数据对于预编译 kernel binary 场景需在 Execute 中加载 binary 并通过aclrtLaunchKernelWithHostArgslaunch5.2 ShapeInferOp::InferShape / InferDataType编译期 Shape 和 DataType 推导回调。graphStatus InferShape(gert::InferShapeContext *ctx) override; graphStatus InferDataType(gert::InferDataTypeContext *ctx) override;InferShapeContext API方法用途ctx-GetInputShape(index)获取输入 origin Shapeconst Shape*ctx-GetOutputShape(index)获取可修改的输出 ShapeShape*ctx-GetInputTensor(index)获取输入 Tensor可访问数据InferDataTypeContext API方法用途ctx-GetInputDataType(index)获取输入 DataTypectx-SetOutputDataType(index, dtype)设置输出 DataType典型流程graphStatus InferShape(gert::InferShapeContext *ctx) { auto *in_shape ctx-GetInputShape(0); auto *out_shape ctx-GetOutputShape(0); *out_shape *in_shape; return GRAPH_SUCCESS; } graphStatus InferDataType(gert::InferDataTypeContext *ctx) { return ctx-SetOutputDataType(0, ctx-GetInputDataType(0)); }5.3 CompilableOp::Compile在线编译回调在 GE/ATC 编译阶段被调用。graphStatus Compile(gert::OpCompileContext *ctx) override;OpCompileContext API方法用途ctx-GetInputTensor(index)获取编译期输入 Tensor含 shape、dtype、formatctx-GetOutputTensor(index)获取编译期输出 Tensorctx-GetOption(key, value)获取编译选项ctx-GetPlatformInfos(...)获取目标平台信息芯片型号等典型流程graphStatus Compile(gert::OpCompileContext *ctx) { auto *input ctx-GetInputTensor(0); auto key BuildBinaryKey(input-GetShape()); auto source LoadFile(my_kernel.cpp); aclrtcProg prog; aclrtcCreateProg(prog, source.c_str(), kernel_name, 0, nullptr, nullptr); aclrtcCompileProg(prog, num_options, options); size_t bin_size; aclrtcGetBinDataSize(prog, bin_size); std::vectoruint8_t binary(bin_size); aclrtcGetBinData(prog, binary.data()); device_elves_[key] std::move(binary); aclrtcDestroyProg(prog); return GRAPH_SUCCESS; }5.4 PortableOp::Serialize / Deserialize序列化/反序列化回调用于 OM 模型下沉。graphStatus Serialize(std::vectoruint8_t buffer) override; graphStatus Deserialize(const std::vectoruint8_t buffer) override;关键点buffer 格式完全由用户自定义GE 不解析只透传Serialize 在 ATC 编译阶段被调用将 Compile 产物写入 bufferDeserialize 在 OM 加载执行阶段被调用从 buffer 恢复 Compile 产物多 binary 场景建议序列化格式[count][key_len][key][bin_len][bin]...5.5 ArgsUpdater::UpdateHostArgsI/O 地址变化时的回调用于刷新 kernel args 中的地址引用。graphStatus UpdateHostArgs(gert::UpdateArgsContext *ctx) override;上下文 API方法用途ctx-GetKernelArgs(placement, index)获取指定位置的 kernel args继承自 EagerOpExecutionContext 的所有方法获取输入/输出 Tensor 等适用场景静态图下沉后模型加载时 I/O 地址可能与编译时不同需要刷新 args buffer 中的地址。6. 前端接入6.1 GE 原生构图GE 原生构图时构图侧需要能看到 REG_OP proto 头文件并在图中创建与 REG_AUTO_MAPPING_OP 注册名一致的 op type。auto op ge::OperatorFactory::CreateOperator(my_node, AddCustom); op.SetInput(x, input_x); op.SetInput(y, input_y); graph.AddOp(op);6.2 PyTorch TorchAir 入图PyTorch 场景需要两侧交付件配合Python 侧 GE 侧 ┌─────────────────────┐ ┌──────────────────────┐ │ TORCH_LIBRARY │ │ REG_OP REG_AUTO_ │ │ (定义算子签名) │ converter │ MAPPING_OP │ │ │ ──────────→ │ (交付件 .so) │ │ TORCH_LIBRARY_IMPL │ │ │ │ (注册 PrivateUse1 │ │ │ │ / Meta / XLA key) │ │ │ └─────────────────────┘ └──────────────────────┘Python 侧需要TORCH_LIBRARY定义算子 schemaTORCH_LIBRARY_IMPL注册 PrivateUse1NPU 执行、Metashape 推导、XLA图编译等 keytorchair.register_fx_node_ge_converter将 PyTorch FX 节点映射到 GE 自定义算子torchair.register_fx_node_ge_converter(torch.ops.my_ops.my_op.default) def convert_my_op(x, y, meta_outputsNone): return torchair.ge.custom_op(AddCustom, inputs{x: x, y: y}, outputs[z])完整代码参见examples/custom_op/ascendc_add_custom/。6.3 TensorFlow 入图TensorFlow 场景需要TensorFlow 侧自定义算子.so声明算子可见GE 侧交付件.so只需REG_AUTO_MAPPING_OP无需额外提供REG_OPnpu_supported_ops.json供 TensorFlow Adapter 识别TensorFlow 侧 GE 侧 ┌──────────────────┐ ┌────────────────────────────┐ │ libcustom_ops.so │ │ libcust_opapi.so │ │ (算子声明) │ Adapter │ (REG_AUTO_MAPPING_OP) │ │ │ ────────→ │ │ │ npu_supported_ │ │ GE proto 由 TF 原型自动生成 │ │ ops.json │ │ 无需额外编写 REG_OP │ └──────────────────┘ └────────────────────────────┘与 GE 原生构图的区别TensorFlow 场景下REG_AUTO_MAPPING_OP可以从 TF 算子原型自动生成 GE 算子原型输入、输出、属性开发者无需额外编写REG_OP定义。这进一步降低了 TensorFlow 自定义算子的接入负担。完整代码参见examples/custom_op/triton_add_custom/。6.4 ONNX 入图ONNX 前端与 PyTorch / TensorFlow 不同GE 的 ONNX Parser 不识别的 op type 会直接报错PARAM_INVALID因此需要额外编写一个ONNX 解析插件将 ONNX 节点的属性映射为 GE 算子属性。ONNX 模型 ONNX 解析插件 GE 侧 ┌──────────────────┐ ┌────────────────────────┐ ┌──────────────────────┐ │ NodeProto │ │ REGISTER_CUSTOM_OP │ │ REG_OP REG_AUTO_ │ │ (op_type, │ dlopen │ (OriginOpType 映射 │ │ MAPPING_OP │ │ attributes) │ ────────→ │ ParseParamsFn 解析) │ ───→ │ (交付件 .so) │ └──────────────────┘ └────────────────────────┘ └──────────────────────┘ONNX 算子类型标识格式domain::version::OpType如ai.onnx::11::Conv、com.example::1::MyOp。Parser 会按此格式在 OpRegistry 中查找映射找不到则解析失败。ONNX 解析插件需要使用REGISTER_CUSTOM_OP注册 ONNX op type 到 GE op type 的映射实现ParseParamsFn从 ONNXNodeProto中提取属性并设置到 GEOperator#include register/register.h #include proto/onnx/ge_onnx.pb.h // 属性解析函数将 ONNX NodeProto 属性映射到 GE Operator 属性 static domi::Status ParseMyOp(const google::protobuf::Message *op_src, ge::Operator op_dest) { auto *node dynamic_castconst ge::onnx::NodeProto *(op_src); for (const auto attr : node-attribute()) { if (attr.name() alpha) { op_dest.SetAttr(alpha, attr.f()); } } return domi::SUCCESS; } // 注册将 ONNX op type 映射到 GE op type REGISTER_CUSTOM_OP(MyCustomOp) .FrameworkType(domi::ONNX) .OriginOpType({ai.onnx::11::MyOp, ai.onnx::13::MyOp}) .ParseParamsFn(ParseMyOp);REGISTER_CUSTOM_OP 关键方法方法用途.FrameworkType(domi::ONNX)声明为 ONNX 框架插件.OriginOpType(domain::ver::Type)指定 ONNX op type支持列表覆盖多版本.ParseParamsFn(fn)属性解析函数Message* → Operator.ParseParamsByOperatorFn(fn)自动映射路径的属性解析Operator → Operator.ParseOpToGraphFn(fn)将单个算子展开为子图可选.ParseSubgraphPostFn(fn)含子图算子的后处理可选交付件组织ONNX 场景需要两个 .so交付件职责放置路径ONNX 解析插件 .soREGISTER_CUSTOM_OPParseParamsFn负责 ONNX → GE 的属性映射ASCEND_CUSTOM_OPP_PATH下的插件目录GE 自定义算子 .soREG_OPREG_AUTO_MAPPING_OP 执行逻辑op_graph/lib/os/arch/注意事项ONNX 解析插件 .so 必须在 ATC 编译前通过ASCEND_CUSTOM_OPP_PATH加载否则 Parser 无法识别自定义 op typeOriginOpType中的 domain 为空时默认使用ai.onnx自定义 domain 需显式指定如果 ONNX 算子可映射为已有 GE 算子的组合可通过ParseOpToGraphFn将其展开为子图无需编写 GE 侧自定义算子 .so7. 构建与部署7.1 CMake 构建配置关键构建要素# 头文件路径必须包含 target_include_directories(my_op PRIVATE ${ASCEND_HOME_PATH}/include ${ASCEND_HOME_PATH}/include/graph ${ASCEND_HOME_PATH}/include/register ${ASCEND_HOME_PATH}/include/external ) # 链接库按需选择 target_link_libraries(my_op PRIVATE ascendcl acl_rt ge_compiler ge_executor ) # 输出 .so add_library(my_op SHARED custom_op.cpp)PyTorch 场景额外依赖需链接torch_npu、c10、lowering等并包含 torch_npu 头文件路径。Ascend C kernel 同库编译使用find_package(ASC REQUIRED)并将.asc文件加入add_library源文件列表。7.2 OPP 包目录结构构建产物需按 OPP 包规范安装output/ # ASCEND_CUSTOM_OPP_PATH 指向此目录 ├── op_graph/ │ ├── include/my_op.h # proto 头文件供构图侧使用 │ └── lib/linux/x86_64/ │ ├── libcust_opapi.so # 交付件 │ └── my_kernel.cpp # kernel 源码RTC 场景需要 └── framework/tensorflow/ └── npu_supported_ops.json # TensorFlow 场景需要7.3 ATC 离线编译流程离线 OM 下沉的完整链路# 1. 构图并导出 AIR ./graph_build_program # 生成 model.air # 2. ATC 离线编译 atc --modelmodel.air --framework1 --outputmodel --soc_versionAscend910B1 # 3. ACL 加载执行 ./model_exec_program model.om前提ASCEND_CUSTOM_OPP_PATH已指向包含交付件的 OPP 包根目录。ATC 编译时会自动加载 .so 并回调 Compile 和 Serialize。8. 调试与常见问题8.1 开发检查项算子类型名、注册类名、构图侧 op type 三者一致ASCEND_CUSTOM_OPP_PATH指向 OPP 包根目录而非 .so 所在目录kernel binary / 源码路径不依赖临时工作目录PyTorch 场景同时检查 Python 侧和 GE 侧交付件是否都已加载TensorFlow 场景同时检查 TF 侧 .so、GE 侧 .so 和 npu_supported_ops.jsonGE 侧无需额外编写 REG_OP离线 OM 场景确认soc_version与实际环境一致ONNX 场景确认解析插件 .so 和 GE 交付件 .so 都已加载且OriginOpType格式正确8.2 常见错误排查现象可能原因排查方法GE 找不到自定义算子ASCEND_CUSTOM_OPP_PATH未配置或路径错误检查环境变量确认op_graph/lib/os/arch/下有 .soExecute 返回 null input输入索引与 REG_OP 定义不匹配核对 REG_OP 的 INPUT 顺序和 Execute 中的 GetInputTensor 索引编译期 shape 为空Compile 阶段输入 shape 未就绪检查是否在 Compile 中正确读取 GetInputTensorOM 加载后 Execute 失败Deserialize 未正确恢复 binary检查 Serialize/Deserialize 的 buffer 格式是否对称PyTorch 图模式报错converter 未注册或 op type 名不匹配确认register_fx_node_ge_converter中的 GE op type 与 REG_AUTO_MAPPING_OP 一致TensorFlow 算子不可见npu_supported_ops.json 缺失或格式错误检查 JSON 文件是否在framework/tensorflow/下ONNX 解析报 The type is not supportedONNX 解析插件未加载或 OriginOpType 不匹配确认解析插件 .so 已通过ASCEND_CUSTOM_OPP_PATH加载且OriginOpType格式为domain::version::OpType【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考