MNN 自定义算子开发指南从 Schema 描述、模型转换到多后端实现的完整流程【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN在 MNN 中扩展一个新的算子需要打通「模型转换 → Schema 描述 → 维度计算 → 后端执行」四个环节。本篇基于 MNN 官方贡献文档 docs/contribute/op.md 整理覆盖自定义算子从MNN.fbsSchema 定义、TensorFlow / TFLite / Caffe / ONNX 前端转换类编写到 CPU、Metal、Vulkan、OpenCL、OpenGL、QNN 各后端Execution注册的完整链路并结合仓库源码说明自动注册脚本 tools/script/register.py 的工作机制读完即可独立完成一个新算子的端到端落地。整体结构四步走的路径选择在动手添加算子前官方建议先检查目标框架已支持的算子列表避免重复开发./MNNConvert -f CAFFE --OP ./MNNConvert -f TF --OP ./MNNConvert -f ONNX --OP ./MNNConvert -f TORCH --OPMNNConvert由 tools/converter 目录编译产出。MNN 的算子落地分为「模型转换」与「算子实现」两条主线各有多个可选步骤模型转换二选一训练框架导出的 Op 与 MNN 的 Op一一对应在前端转换器中直接添加转换类使用组合器由已有 MNN 算子组合实现可参考 tools/converter/source/optimizer/onnxextra 目录其中包含OnnxClip、OnnxSoftmax、OnnxGather等由基础算子拼装而成的组合实现以及统一调度入口OnnxExtraManager。算子实现按序执行添加 Schema 描述必须添加维度Shape计算若输出维度与输入一致可跳过添加几何计算实现可选——一旦实现几何计算后续各后端就无须再为该算子单独实现添加各后端算子实现可选按需选择后端。一句话总结优先级优先转换然后组合然后几何计算最后各后端实现。组合与几何计算可以覆盖所有后端是成本最低的落地方式只有性能关键或几何无法表达的算子才需要下沉到具体后端。第一步添加 Schema 描述若目标算子不在 MNN 的算子列表中必须先修改模型描述FlatBuffers Schema然后用 generate 脚本重新生成头文件。所有 Schema 源文件位于 schema/default 目录生成产物写入 schema/current如MNN_generated.h。1. 添加算子类型在 schema/default/MNN.fbs 的OpType枚举中追加算子名称enum OpType : int { AbsVal, QuantizedAdd, ... MyCustomOp }从源码结构看OpType枚举中存在若干保留区段追加算子时应避开已有编号Raster/ConvertTensor等占用 128–156 的显式编号段Plugin 256之后是训练类算子257 起299 起是用户自定义/Transformer 类算子Attention、FmhaV2、RoPE、FusedLinear等Extra 512之后是量化算子513 起600 起是控制流算子While、If、LayerNorm、GridSample。普通业务算子一般直接追加到连续段末尾即可。2. 添加算子参数描述如果算子不带参数可跳过此步。首先在 schema/default/MNN.fbs 的OpParameter联合类型中追加参数表名称union OpParameter { QuantizedAdd, ArgMax, AsString, ... MyCustomOpParam }然后按算子来源选择对应的 fbs 文件添加参数表来自 Caffe 的算子写入CaffeOp.fbs来自 TensorFlow 的写入TensorflowOp.fbs两者均在 schema/default 下table MyCustomOpParam { padX:int; padY:int; kernelX:int; kernelY:int; strideX:int; strideY:int; dataType:DataTypeDT_FLOAT; }3. 重新生成描述头文件修改完模型描述后调用 schema/generate.sh 重新生成头文件。该脚本会先检查3rd_party/flatbuffers/tmp/flatc是否存在不存在则自动 cmake 构建flatc再对default目录下全部 fbs 执行flatc -c -b --gen-object-api --reflect-names把.h与二进制描述文件输出到schema/current/。第二步添加模型转换用户可根据自己使用的框架选择对应的前端转换模块添加算子转换支持。添加完模型转换后需要重新 cmake 编译。目前 MNN 支持 TensorFlow、TensorFlow Lite、Caffe、ONNX 和 TorchScript 五种格式源码分别位于 tools/converter/source/tensorflow、tools/converter/source/tflite、tools/converter/source/caffe、tools/converter/source/onnx 与 tools/converter/source/torch。TensorFlow 模型转换在 tools/converter/source/tensorflow 下添加MyCustomOpTf.cpp。可以直接声明转换类也可以利用宏简化代码。直接声明示例class MyCustomOpTf : public tfOpConverter { public: virtual void run(MNN::OpT *dstOp, TmpNode *srcNode, TmpGraph *tempGraph); MyCustomOpTf() {} virtual ~MyCustomOpTf() {} virtual MNN::OpType opType(); virtual MNN::OpParameter type(); }等效宏定义示例DECLARE_OP_CONVERTER(MyCustomOpTf);需要实现run、析构、opType和type函数。其中run用于解析模型的 proto 文件得到参数然后赋值给 flatbuffer 自定义参数参数srcNode保存输入输出节点信息可以根据输入输出节点在tempGraph中找到TmpNode调用find_attr_value(const tensorflow::NodeDef node, const char* key, tensorflow::AttrValue value)获得对应属性值。仓库中大量现成实现可参考例如 tools/converter/source/tensorflow/AsStringTf.cpp 中多次调用find_attr_value(srcNode-tfNode, T, value)提取属性。注册转换类REGISTER_CONVERTER(MyCustomOpTf, MyCustomOp);TensorFlow Lite 模型转换添加转换类在 tools/converter/source/tflite 下添加MyCustomOpTflite.cpp宏定义示例DECLARE_OP_CONVERTER(MyCustomOpTflite);需要实现函数MyCustomOpTflite::opType(int quantizedModel); MyCustomOpTflite::type(int quantizedModel); MyCustomOpTflite::run(MNN::OpT *dstOp, const std::unique_ptrtflite::OperatorT tfliteOp, const std::vectorstd::unique_ptrtflite::TensorT tfliteTensors, const std::vectorstd::unique_ptrtflite::BufferT tfliteModelBuffer, const std::vectorstd::unique_ptrtflite::OperatorCodeT tfliteOpSet, int quantizedModel)其中run函数相比 TensorFlow 版本多一个quantizedModel参数为 true 表示量化模型需转为相应的量化 Op为 false 则转为浮点 Op。run中还必须设置输入、输出 tensor 的 index// set input output index dstOp-inputIndexes.resize(1); dstOp-outputIndexes.resize(1); dstOp-inputIndexes[0] tfliteOp-inputs[0]; dstOp-outputIndexes[0] tfliteOp-outputs[0];注册转换类using namespace tflite; REGISTER_CONVERTER(MyCustomOpTflite, BuiltinOperator_OPName);Caffe 模型转换添加转换类在 tools/converter/source/caffe 下添加MyCustomOp.cpp。类声明示例class MyCustomOp : public OpConverter { public: virtual void run(MNN::OpT* dstOp, const caffe::LayerParameter parameters, const caffe::LayerParameter weight); MyCustomOp() {} virtual ~MyCustomOp() {} virtual MNN::OpType opType(); virtual MNN::OpParameter type(); };实现run、opType、type函数在run中解析 Caffe 参数得到具体参数。其中parameters保存 Op 的参数信息weight保存卷积、BN 等数据参数。注册转换类static OpConverterRegisterMyCustomOp a(MyCustomOp);ONNX 模型转换添加转换类在 tools/converter/source/onnx 下添加MyCustomOpOnnx.cpp类声明示例DECLARE_OP_CONVERTER(MyCustomOpOnnx);需要实现函数MNN::OpType MyCustomOpOnnx::opType(); MNN::OpParameter MyCustomOpOnnx::type(); void MyCustomOpOnnx::run(MNN::OpT* dstOp, const onnx::NodeProto* onnxNode, std::vectorconst onnx::TensorProto* initializers);run函数中onnxNode即 ONNX 原始节点信息权重等数据信息需从initializers中获取。注册转换类REGISTER_CONVERTER(MyCustomOpOnnx, MyCustomOp);第三步添加维度计算如果该 Op 的输出 Tensor 大小与第 1 个输入 Tensor 一致且不需要分析 FLOPS可以跳过这步。添加计算类在 source/shape 目录下添加ShapeMyCustomOp.cppclass MyCustomOpSizeComputer : public SizeComputer { public: virtual bool onComputeSize(const MNN::Op* op, const std::vectorTensor* inputs, const std::vectorTensor* outputs) const override { // set tensor-buffer.type // .dimensions // .dim[x].extent // .dim[x].stride // .dim[x].flag return true; } virtual float onComputeFlops(const MNN::Op* op, const std::vectorTensor* inputs, const std::vectorTensor* outputs) const { return flops_for_calc_output_from_input; } };onComputeSize根据输入 tensor 的维度信息计算输出 tensor 的维度信息并设置输出 tensor 的数据类型计算完成返回true若输入维度未知返回false引擎会据此判断形状是否可解析。onComputeFlops根据输入、输出 tensor 的维度信息返回总计算量供性能分析使用。注册计算类REGISTER_SHAPE(MyCustomOpSizeComputer, OpType_MyCustomOp);添加完形状计算代码后需要在仓库根目录运行python3 tools/script/register.py并重新 cmake。第四步添加各后端实现添加完算子实现后同样需要在仓库根目录运行python3 tools/script/register.py并重新 cmake。register.py 的自注册机制tools/script/register.py 是 MNN 的算子自动注册核心它扫描各目录源码中的注册宏汇总生成统一的注册入口文件。从源码实现看它做了如下几件事扫描 source/shape识别REGISTER_SHAPE/REGISTER_SHAPE_OLD/REGISTER_SHAPE_INPUTS及 RENDER / TRANSFORMER_FUSE 变体生成 source/shape/ShapeRegister.cpp 中的registerShapeOps()扫描 source/backend/cpu识别REGISTER_CPU_OP_CREATOR系列宏生成 source/backend/cpu/CPUOPRegister.cpp 中的registerCPUOps()扫描 source/geometry识别REGISTER_GEOMETRY生成 source/geometry/GeometryOPRegister.cpp 中的registerGeometryOps()分别扫描 CoreML、NNAPI、OpenCL 后端目录生成对应的CoreMLOPRegister.cpp、NNAPIOPRegister.cpp、OpenCLOPRegister.cpp。生成的注册函数内部是extern void ___CreatorName__OpType_X__();声明加调用序列每个注册宏会展开为一个静态构造函数形式的自注册函数。因此只要新增了带注册宏的源文件运行一次 register.py 并重新 cmake 即可让新算子被引擎发现这是理解 MNN「免手工登记」扩展机制的关键。CPU 实现在 source/backend/cpu 目录下添加CPUMyCustomOp.hpp、CPUMyCustomOp.cpp。实现类声明class CPUMyCustomOp : public Execution { public: // 若执行onExecute需要使用缓存在此函数中申请若无可不声明 virtual ErrorCode onResize(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; // 具体的Op执行函数 virtual ErrorCode onExecute(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; };实现onResize与onExecuteonResize中调用backend()-onAcquireBuffer(mCache, Backend::DYNAMIC)申请缓存、backend()-onReleaseBuffer(mCache, Backend::DYNAMIC)回收缓存释放后的内存可被复用onExecute中做必要的输入检查有利于提前发现问题执行完毕正确返回NO_ERROR。注册实现类class CPUMyCustomOpCreator : public CPUBackend::Creator { public: virtual Execution *onCreate(const std::vectorTensor * inputs, const std::vectorTensor * outputs, const MNN::Op *op, Backend *backend) const override { return new CPUMyCustomOp(backend); } }; REGISTER_CPU_OP_CREATOR(CPUMyCustomOpCreator, OpType_MyCustomOp);Metal 实现在 source/backend/metal 目录下添加MetalMyCustomOp.hpp与MetalMyCustomOp.cpp。实现类声明class MetalMyCustomOp : public Execution { public: virtual ErrorCode onResize(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; virtual void onEncode(const std::vectorTensor * inputs, const std::vectorTensor * outputs, idMTLComputeCommandEncoder encoder) override; };实现onResize与onEncode尽量将申请内存和计算 group size 的操作放在onResize中onEncode时使用传入的 encoder 编排计算任务不要自行创建 command buffer 或 encoder。内存使用——这是 Metal 后端与 CPU 最大的区别CPU Tensor 数据存储在 host 指针中而 Metal 数据指针存放在deviceId中deviceId上存储的是idMTLBuffer。由于内存复用机制各 Tensor 可能共用同一块内存以 offset 偏移auto buffer (__bridge idMTLBuffer)(void *)tensor-deviceId(); auto offset TensorUtils::getDescribe(tensor)-extra.offset;Metal Op 的特定参数可通过idMTLBuffer存储。buffer 数据类型可以与 tensor 不同甚至可以混合多种数据类型只需保证创建时指定了正确长度即可auto buffer [context newDeviceBuffer:2 * sizeof(int) 2 * sizeof(__fp16) access:CPUWriteOnly]; ((__fp16 *)buffer.contents)[0] mAlpha / mLocalSize; // alpha ((__fp16 *)buffer.contents)[1] mBeta; // beta ((int *)buffer.contents)[1] mLocalSize; // local size ((int *)buffer.contents)[2] inputs[0]-channel(); // channel创建 buffer 时需指定访问控制权限共三种CPUReadWrite数据在 CPU/GPU 间共享存储一般用于 device bufferCPUWriteOnly数据通过 CPU 写入后不再读取一般用于参数 bufferCPUTransparent数据只在 GPU 中一般用于 heap buffer。MNNMetalContext在创建 buffer 上提供两套相近接口区别在于数据生命周期device 占用的内存在单次推理过程中不会被复用heap 占用的内存在调用-[MNNMetalContext releaseHeapBuffer:]之后可被其他 Op 复用。一般而言 heap 只与CPUTransparent一起使用heap 实际只在 iOS 10 上有效iOS 9- 会回退到 device。Metal 内存布局与 CPU-FP32-Neon 一致Tensor 的 dimensionFormat 为 NC4HW4 时使用 C4NHW4 排布否则按默认线性布局。注册实现类class MetalMyCustomOpCreator : public MetalBackend::Creator { public: virtual Execution *onCreate(const std::vectorTensor * inputs, const MNN::Op *op, Backend *backend) const { return new MetalMyCustomOp(backend); } }; REGISTER_METAL_OP_CREATOR(MetalMyCustomOpCreator, OpType_MyCustomOp);工程更新进入source/backend/metal目录执行 python3 MetalCodeGen.py . 更新自注册文件然后重新运行 CMake或手动在 Xcode 工程中新加文件。Vulkan 实现Vulkan 后端包含两种张量存储类型buffer 与 image开发者可在编译时通过宏MNN_VULKAN_IMAGE选择需要的存储类型添加算子时也要考虑选择何种存储类型并在相应目录下开发。以 image 类型为例主要流程如下目录为source/backend/vulkan/image/。实现 Execution执行脚本 source/backend/vulkan/image/compiler/VulkanCodeGen.py该脚本会向source/backend/vulkan/image/execution中添加VulkanMyOp.hpp与VulkanMyOp.cpp的模板代码实现构造函数从 CPU 中读取常量参数并写入 GPU创建算子所需 pipeline确定要使用的 shader 及 Macro → 设置 descriptorTypes即 shader 中用到的显存对象类型 → 调用getPipeline接口实现onEncode申请显存资源并更新 descriptorSet将 shader 中需要读写的显存对象写入→ 添加 memoryBarrier → 把 pipeline 绑到 cmdBuffer 与 descriptorSet → command dispatch注册算子并添加创建类class VulkanMyCustomOpCreator : public VulkanBackend::Creator { public: virtual Execution* onCreate(const std::vectorTensor* inputs, const MNN::Op* op, Backend* backend) const override { return new VulkanMyCustomOp(op, backend); } }; static bool gResistor []() { VulkanBackend::addCreator(OpType_MyCustomOp, new VulkanMyCustomOpCreator); return true; }();实现 shader 及编译编写 Compute Shader 文件myOp.comp添加至目录source/backend/vulkan/image/execution/glsl将算子中用到的宏加入source/backend/vulkan/image/execution/glsl/macro.json执行脚本 source/backend/vulkan/image/compiler/makeshader.py该脚本将编译myOp.comp并更新AllShader.cpp、AllShader.h与VulkanShaderMap.cpp。MNN Vulkan 当前使用 glslangValidatorglslang 版本 12.2.0commit idd1517d64cfca91f573af1bf7341dc3a5113349c0编译所有 compute shader。如需自行编译后得到的二进制结果与 MNN 仓库中现有编译结果一致必须确保环境中 glslang 版本与 MNN 所使用的一致。OpenCL 实现添加 Kernel在source/backend/opencl/execution/cl目录添加具体的 kernel*.cl。目前 feature map 均使用image2d实现可参考目录下已有实现。然后执行opencl_codegen.py生成 kernel 映射。实现类声明在source/backend/opencl/execution/下添加MyCustomOp.h与MyCustomOp.cpptemplate typename T class MyCustomOp : public Execution { public: virtual ErrorCode onResize(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; virtual ErrorCode onExecute(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; };实现onResize可选与onExecute执行完毕返回NO_ERROR。注册实现类OpenCLCreatorRegisterTypedCreatorMyCustomOpcl_data_t __my_custom_op(OpType_MyCustomOp);从 tools/script/register.py 的 OpenCL 分支可以印证它分别扫描 buffer 与 image 两个 execution 子目录识别REGISTER_OPENCL_OP_CREATOR系列宏按__BUFFER__/__IMAGE__后缀生成OpenCLOPRegister.cpp并通过MNN_OPENCL_BUFFER_CLOSED宏控制 buffer 路径开关。OpenGL 实现添加 Shader在source/backend/opengl/glsl下添加具体的 shader*.glsl不用加文件头feature map 均采用image3d表示可参考目录下已有实现。然后在source/backend/opengl目录执行 makeshader.py。添加 Executor在source/backend/opengl/execution/下添加GLMyCustomOp.h与GLMyCustomOp.cppclass GLMyCustomOp : public Execution { public: GLMyCustomOp(const std::vectorTensor * inputs, const Op *op, Backend *bn); virtual ~GLMyCustomOp(); virtual ErrorCode onExecute(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; virtual ErrorCode onResize(const std::vectorTensor * inputs, const std::vectorTensor * outputs) override; private: std::shared_ptrGLProgram mProgram; };实现onResize可选与onExecute执行完毕返回NO_ERROR。注册实现类GLCreatorRegisterTypedCreatorGLMyCustomOp __my_custom_op(OpType_MyCustomOp);QNN 实现QNN 后端通过调用高通 QNN 的官方算子库qti.aisw实现 MNN 的计算功能。为 QNN 添加算子时需要充分理解 MNN 算子格式与 QNN 算子格式读取 MNN 算子的尺寸与参数再设定 QNN 算子的尺寸与参数。关于 QNN 官方算子库可参考官方文档中的算子定义与算子在 HTP 上的限制说明docs.qualcomm.com上的 80-63442-50 系列文档。添加 Executor在source/backend/qnn/execution目录下添加QnnCustomOp.cpp与QnnCustomOp.hpp实现QnnCustomOp类及QnnCustomOpCreator类的骨架代码实现QnnCustomOp::onEncode函数分两种情况单个 QNN 算子即可实现MNN 算子的计算功能一般包含添加 inputs、添加 scalar params可参考source/backend/qnn/execution/QNNArgmax.cpp、添加 tensor params可参考source/backend/qnn/execution/QNNPool.cpp、添加 outputs、调用 QNN 的节点入图 API 等步骤拼接多个 QNN 算子实现 MNN 算子的计算功能可参考source/backend/qnn/execution/QNNScale.cpp。注册算子在QnnCustomOp.cpp中添加REGISTER_QNN_OP_CREATOR(QnnCustomOpCreator, OpType_CustomOp)。该宏定义于 source/backend/qnn/backend/QNNBackend.hpp在source/backend/qnn/backend/QNNUtils.hpp中添加函数声明extern void ___QnnCustomOpCreator__OpType_CustomOp__();并在QNNUtils.cpp的registerQNNOps函数中追加___QnnCustomOpCreator__OpType_CustomOp__();registerQNNOps的声明见 source/backend/qnn/backend/QNNUtils.hpp。编译与验证流程小结把整个流程按执行顺序汇总每一步的产物与触发条件如下阶段关键动作触发命令 / 脚本0. 查重查询已支持算子./MNNConvert -f {CAFFE\|TF\|ONNX\|TORCH} --OP1. Schema修改 schema/default/MNN.fbsOpType / OpParameter / 参数表sh schema/generate.sh2. 模型转换在前端目录添加转换类并REGISTER_CONVERTER重新 cmake3. 维度计算source/shape 添加 SizeComputer 并REGISTER_SHAPEpython3 tools/script/register.py 重新 cmake4. 几何实现可选source/geometry 添加几何计算并REGISTER_GEOMETRYpython3 tools/script/register.py 重新 cmake5. 后端实现可选各 backend 目录添加 Execution 与 Creatorpython3 tools/script/register.py 重新 cmakeMetal 额外执行MetalCodeGen.pyVulkan/OpenGL 额外执行对应 shader 编译脚本补充两点实践提示注意脚本路径贡献文档中写作tools/scripts/register.py但当前仓库中该脚本实际位于 tools/script/register.py以仓库实际路径为准测试验证算子行为验证可以放在 test/op 目录下的算子测试中扩展框架由 test/MNNTestSuite.cpp 与 test/CommonOpCreator.hpp 提供形状解析逻辑则由 test/core 下的用例覆盖。几何实现一旦就位即可自动服务所有后端这也是官方「几何计算优先于后端实现」建议的根本原因。遵循「优先转换 → 组合 → 几何 → 后端」的顺序逐级下沉能以最少的代码量让新算子在 MNN 全后端可用而当性能关键路径确实需要时再按各后端的onResize/onExecute或onEncode范式补齐具体实现即可。【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
