CANN/GE语言无关自定义算子接入

[RFC] 语言无关自定义算子接入 GE【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/geSummary摘要本文提出一种语言无关的自定义算子接入 GE 的机制。通过定义统一的算子接入接口将自定义算子的集成过程与具体的算子编成语言Ascend C、Triton、PPTO 等解耦并提供渐进式的开发体验——从只支持运行时执行到参与编译时优化逐步获得更高的性能收益。Motivation动机当前 GE 对自定义算子接入的支持有 2 个关键痛点只支持 Ascend C 语言开发的自定义算子接入。随着算子多样化编成语言的发展如 Triton 在易用性上有较高吸引力用户希望其他语言开发的算子也能接入 GE。入图交付件多且零散易用性有待提升。开发者需要同时维护 proto 定义、执行逻辑、编译逻辑等多个文件缺乏统一的交付件组织方式。Proposed Design设计方案架构视图通过统一的开发界面对接不同的编成语言。自定义算子以.so交付件形式加载到 GE参与图编译和执行的全流程。渐进式能力模型自定义算子入图分为 3 个阶段开发工作量与性能收益逐步递增阶段核心能力新增交付件性能收益阶段 1Executehost 调度 kernel1 个 .so可运行有 host 调度开销阶段 2.1Execute下沉调度无新增静态 shape 下消除 host 调度开销阶段 2.2 InferShape Compile无新增shape 推导、内存复用、在线编译阶段 3 Serialize / Deserialize无新增离线 OM 部署伪代码开发示例以下以 Add 算子为例展示各阶段的开发者体验。阶段 1动态 shape host 调度只需实现Execute完成 kernel 的加载和 launchclass AddCustom : public EagerExecuteOp { public: graphStatus Execute(gert::EagerOpExecutionContext *ctx) override { // 1. 获取输入 auto *x ctx-GetInputTensor(0); auto *y ctx-GetInputTensor(1); // 2. 分配输出 auto *z ctx-MallocOutputTensor(0, x-GetShape(), x-GetFormat(), x-GetDataType()); // 3. 加载 kernel binary预编译的 npubin / Ascend C binary auto bin_data LoadBinary(add_kernel.npubin); auto func_handle GetKernelFunction(bin_data, add_kernel); // 4. 构造 args 并 launch int64_t n x-GetShapeSize(); int32_t block_num CeilDiv(n, BLOCK_SIZE); struct Args { const void *in0, *in1; void *out; int32_t n, gx, gy, gz; } args {x-GetAddr(), y-GetAddr(), z-GetAddr(), (int32_t)n, block_num, 1, 1}; aclrtLaunchKernelWithHostArgs(func_handle, block_num, ctx-GetStream(), nullptr, args, sizeof(args), nullptr, 0); return GRAPH_SUCCESS; } }; 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); REG_AUTO_MAPPING_OP(AddCustom);效果算子可在 GE 图中运行支持动态 shape但每个推理 step 有 host 侧调度开销。阶段 2静态 shape 下沉在阶段 1 基础上补充ShapeInferOp和CompilableOpclass AddCustom : public EagerExecuteOp, public ShapeInferOp, public CompilableOp { // Execute 同阶段 1省略... graphStatus InferShape(gert::InferShapeContext *ctx) override { *ctx-GetOutputShape(0) *ctx-GetInputShape(0); return GRAPH_SUCCESS; } graphStatus InferDataType(gert::InferDataTypeContext *ctx) override { return ctx-SetOutputDataType(0, ctx-GetInputDataType(0)); } graphStatus Compile(gert::OpCompileContext *ctx) override { auto *input ctx-GetInputTensor(0); auto key BuildKey(input-GetShape()); auto source LoadFile(add_kernel.cpp); aclrtcProg prog; aclrtcCreateProg(prog, source.c_str(), add_kernel, 0, nullptr, nullptr); aclrtcCompileProg(prog, 1, options); size_t bin_size; aclrtcGetBinDataSize(prog, bin_size); device_elves_[key].resize(bin_size); aclrtcGetBinData(prog, device_elves_[key].data()); aclrtcDestroyProg(prog); return GRAPH_SUCCESS; } private: std::mapstd::string, std::vectoruint8_t device_elves_; };效果阶段 2.1无新增交付件静态 shape 下 kernel 下沉调度消除 host 开销阶段 2.2参与 shape 推导和内存复用Compile 阶段完成算子在线编译阶段 3离线 OM 支持在阶段 2 基础上补充PortableOpclass AddCustom : public EagerExecuteOp, public ShapeInferOp, public CompilableOp, public PortableOp { // Execute / InferShape / Compile 同阶段 2省略... graphStatus Serialize(std::vectoruint8_t buffer) override { // 将 device_elves_ 序列化到 buffer格式自定义GE 只透传 return SerializeBinaryMap(device_elves_, buffer); } graphStatus Deserialize(const std::vectoruint8_t buffer) override { // 从 buffer 恢复 device_elves_ return DeserializeBinaryMap(buffer, device_elves_); } };效果编译产物随 OM 文件保存和恢复支持AIR → ATC → OM → ACL离线部署链路。语言公共层封装效果上述基础设施层代码约 60-80 行。各编成语言可构建公共层进一步封装以 Triton 为例// 使用 Triton 公共层后同一个 Add 算子只需 ~10 行 TRITON_CUSTOM_OP(AddCustom) .Kernel(add_kernel) // 声明 kernel 名称 .Binary(add_kernel.npubin) // 声明 binary 路径 .Inputs({x, y}) // 声明输入 .Outputs({z}) // 声明输出 .InferShapeSameAsInput(0) // 输出 shape 第 0 个输入 shape .InferDataTypeSameAsInput(0) // 输出 dtype 第 0 个输入 dtype .TilingStrategy(TilingStrategy::ElementWise) // 自动计算 block_num .Build();封装前后对比重复逻辑基础设施层手动语言公共层自动binary 加载手动aclrtBinaryLoadFromData声明.Binary()路径args 构造手动拼装 packed struct根据 kernel 签名自动生成block_num 计算手动CeilDiv(n, BLOCK_SIZE).TilingStrategy(ElementWise)REG_OP 定义手动编写 proto.Inputs()/.Outputs()自动生成InferShape手动实现.InferShapeSameAsInput(0)基础设施定位与语言公共层层次职责维护方GE 基础设施层统一接入接口、注册机制、编译/执行回调、序列化协议GE 团队语言公共层封装特定语言的 boilerplatebinary 加载、args 构造等各语言 SDK 团队算子开发者只需实现 kernel 逻辑 少量声明算子开发者前端接入前端额外交付件接入方式GE 原生无REG_OP OperatorFactoryPyTorch TorchAirTORCH_LIBRARY converterFX 节点映射到 GE op typeTensorFlowlibcustom_ops.so npu_supported_ops.jsonTF Adapter 构图转换REG_AUTO_MAPPING_OP 自动生成 GE protoONNXREGISTER_CUSTOM_OP 解析插件NodeProto 属性映射到 GE OperatorOpen Questions待讨论问题语言公共层的标准化程度各语言的公共层是否应该由 GE 统一提供模板/SDK还是由各语言团队独立维护多版本兼容当 GE 基础设施层接口演进时如何保证旧版本 .so 交付件在新版 GE 上仍可加载是否需要引入算子版本字段编译期并行安全CustomGraphOptimizer并行回调Compile当前要求算子实现自行保证线程安全。是否应由框架层提供锁机制序列化格式标准化当前PortableOp的 buffer 格式完全由用户自定义。是否需要 GE 提供标准的序列化辅助工具ONNX 自定义 domain 支持当前 ONNX 解析插件需要为每个 domain::version::OpType 显式注册。是否支持通配符或自动发现机制Timeline时间线阶段状态说明阶段 1动态 shape host 调度✅ 已完成参见examples/custom_op/triton_add_custom阶段 2.1静态 shape 下沉 only✅ 已完成同上 sample 验证下沉效果阶段 2.2下沉全量收益✅ 已完成shape 推导、内存复用、在线编译已支持阶段 3离线 OM 支持✅ 已完成参见examples/custom_op/compilable_add_custom语言公共层 规划中各语言 SDK 团队按需构建References开发指南custom_op_development_guide.md架构设计custom_op_architecture.md样例代码examples/custom_op/【免费下载链接】geGEGraph Engine是面向昇腾的图编译器和执行器提供了计算图优化、多流并行、内存复用和模型下沉等技术手段加速模型执行效率减少模型内存占用。 GE 提供对 PyTorch、TensorFlow 前端的友好接入能力并同时支持 onnx、pb 等主流模型格式的解析与编译。项目地址: https://gitcode.com/cann/ge创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考