MNN 新增算子系列三计算实现完全指南——几何拆解与 CPU 后端的双路径实战【免费下载链接】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 算子的组合零后端成本还是CPU 后端实现为算子编写原生执行代码。读完本文你将掌握GeometryComputer的注册与编排机制、GeometryComputerUtils六大工具函数的真实用法含仓库内CosineSimilarity完整拆解案例、CPUExecution的onResize/onExecute生命周期、以及tools/script/register.py驱动的编译验证流程。前两步Schema 定义、形状计算与本步的衔接关系可参考总览文档 SKILL.md。3.0 第一步决策选择实现方式优先级铁律几何计算 CPU 后端实现。几何计算的本质是把新算子拆解为已有 MNN 算子的组合从而无需为 CPU / Metal / OpenCL / Vulkan / CUDA 等每个后端单独写实现——被拆解出的子算子全部走各自后端已有的成熟 kernel。这是 MNN 中最经济的算子落地方式。判断是否可以用几何计算类别代表算子拆解方式✅ 可以几何拆解CosineSimilarityMul Reduce(Sum) Sqrt Div✅ 可以几何拆解LayerNormMean Sub Variance Rsqrt Mul Add✅ 可以几何拆解SoftmaxExp ReduceSum Div✅ 可以几何拆解各种 element-wise 组合算子若干 Unary/Binary 串联❌ 不适合NMS、Sort需要复杂循环/条件分支❌ 不适合Im2Col、Gather 的特殊变体需要特殊内存访问模式❌ 不适合性能敏感算子几何拆解效率远低于原生实现判断结果可以拆解→ 走 3.A几何计算不可拆解→ 走 3.BCPU 后端实现一个必须警惕的坑源自 SKILL.md 的核心原则拆解后所有变量、字段、宏名指代的必须是代码里真正装的那个量——blockSize是块内元素数还是每行块数、step是元素数还是字节数。名字骗人不会报错编译、对拍、单测都能过但后续基于它写的阈值、门控、分档会一起错。3.A 路径一几何计算实现3.A.1 创建几何计算文件在source/geometry/下创建GeometryMyCustomOp.cpp骨架如下完整可编译模板核心为onCompute纯虚函数// // GeometryMyCustomOp.cpp // MNN // #include geometry/GeometryComputer.hpp #include core/OpCommonUtils.hpp #include geometry/GeometryComputerUtils.hpp namespace MNN { class GeometryMyCustomOp : public GeometryComputer { public: virtual bool onCompute(const Op* op, const std::vectorTensor* inputs, const std::vectorTensor* outputs, Context context, CommandBuffer res) const override { auto input inputs[0]; auto output outputs[0]; // 拆解为已有算子组合 // 示例 1创建中间 Tensor // std::shared_ptrTensor tmpTensor; // tmpTensor.reset(Tensor::createDevicefloat({batch, channel, height, width})); // auto des TensorUtils::getDescribe(tmpTensor.get()); // des-dimensionFormat MNN_DATA_FORMAT_NCHW; // res.extras.emplace_back(tmpTensor); // 示例 2调用 Binary 算子 // auto cmd GeometryComputerUtils::makeBinary(BinaryOpOperation_MUL, // inputA, inputB, outputTensor); // res.command.emplace_back(std::move(cmd)); // 示例 3调用 Unary 算子 // auto cmd GeometryComputerUtils::makeUnary(UnaryOpOperation_SQRT, // inputTensor, outputTensor); // res.command.emplace_back(std::move(cmd)); // 示例 4调用 Reduce 算子 // auto cmd GeometryComputerUtils::makeReduce(ReductionType_SUM, // inputTensor, outputTensor); // res.command.emplace_back(std::move(cmd)); // 示例 5虚拟 Tensor内存重映射零拷贝 reshape/transpose // auto outputDes TensorUtils::getDescribe(output); // outputDes-memoryType Tensor::InsideDescribe::MEMORY_VIRTUAL; // Tensor::InsideDescribe::Region desReg; // desReg.size[0] dim0; // desReg.size[1] dim1; // desReg.size[2] dim2; // desReg.dst.offset 0; // desReg.dst.stride[0] dim1 * dim2; // desReg.dst.stride[1] dim2; // desReg.dst.stride[2] 1; // desReg.src.offset 0; // desReg.src.stride[0] src_stride0; // desReg.src.stride[1] src_stride1; // desReg.src.stride[2] 1; // desReg.origin input; // outputDes-regions.emplace_back(std::move(desReg)); return true; } }; static void _create() { std::shared_ptrGeometryComputer comp(new GeometryMyCustomOp); GeometryComputer::registerGeometryComputer(comp, {OpType_MyCustomOp}); } REGISTER_GEOMETRY(GeometryMyCustomOp, _create); } // namespace MNN几何计算的源码级原理注册与命令缓冲注册机制_create()中调用GeometryComputer::registerGeometryComputer(comp, {OpType_MyCustomOp})把算子类型绑定到几何实现REGISTER_GEOMETRY(GeometryMyCustomOp, _create)宏定义见 GeometryComputer.hpp展开为一个extern函数___GeometryMyCustomOp___create__()供自动生成的注册文件统一调用。registerGeometryComputer与运行时查找入口GeometryComputer::search(opType, compileType)同见 GeometryComputer.hpp。Context 的职责onCompute的第 4 个参数Context封装了后端分配器、前向类型、精度模式与 GPU 模式见 GeometryComputer.hpp其中allocConst(op, shape, type)用于创建并缓存常量 Tensor——同一个op再次调用时会命中缓存避免重复分配。CommandBuffer 的职责res包含command真正要执行的子算子命令队列与extras生命周期由几何计算托管的中间 Tensor。每次makeBinary/makeUnary/makeReduce产出的命令都要emplace_back进res.command所有新建的中间 Tensor 都要emplace_back进res.extras否则内存无法正确托管。常用几何计算工具函数源码出处函数作用声明位置GeometryComputerUtils::makeBinary(op, a, b, out)二元运算加减乘除等GeometryComputerUtils.hppGeometryComputerUtils::makeUnary(op, in, out)一元运算sqrt, exp, log 等GeometryComputerUtils.hppGeometryComputerUtils::makeReduce(type, in, out)规约运算sum, mean, max 等GeometryComputerUtils.hppGeometryComputerUtils::makeMatMul(a, b, out, ...)矩阵乘法可带 Bias、转置标志GeometryComputerUtils.hppcontext.allocConst(op, shape, type)创建并缓存常量 TensorGeometryComputer.hppTensor::createDevicefloat(shape)创建设备中间 TensorTensor 接口MEMORY_VIRTUALregions零拷贝内存重映射reshape/transposeTensor::InsideDescribe::RegionGeometryComputerUtils还额外提供makeLayerNorm、makePool等高阶封装见 GeometryComputerUtils.hpp如果你要拆解的目标恰好是 LayerNorm、Pooling 这类算子可以直接复用不必从 Binary/Unary 手动拼。仓库实战案例GeometryCosineSimilarity 完整拆解文档开头的判断表把CosineSimilarity Mul Reduce(Sum) Sqrt Div列为经典可拆解算子仓库中恰好存在完整实现 GeometryCosineSimilarity.cpp可作为几何计算的最佳范本其执行流水线为用MEMORY_VIRTUALRegion将两个输入 reshape 为{outside, channel, inside}的三维视图L43-L91零拷贝完成布局转换makeBinary(BinaryOpOperation_MUL, ...)分别计算input0*input0、input0*input1、input1*input1L100-L127makeReduce(ReductionType_SUM, ...)沿 channel 维axis1对三个乘积求和L134-L158用context.allocConst(op, {1}, halide_type_offloat())创建eps 1e-8f常量并做ADDL180-L186——这正是allocConst按op缓存常量的典型用法makeUnary(UnaryOpOperation_SQRT, ...)开平方makeBinary(BinaryOpOperation_REALDIV, ...)做最终除法L195-L209最后再把中间结果通过虚拟 Region 映射到输出 TensorL212-L228。一个完整算子就这样被拆成 11 条子命令全程没有写一行真实数值计算。写你自己的GeometryMyCustomOp.cpp时对照此文件逐段替换即可。3.A.2 运行注册脚本并编译python3 tools/script/register.py cd build cmake .. -DMNN_BUILD_TESTON make -j$(nproc)register.py位于 tools/script/register.py会扫描source/geometry/下所有REGISTER_GEOMETRY宏自动把extern void ___GeometryXXX___create__()声明与调用写入 GeometryOPRegister.cpp。该文件头部注释明确写着 This file is generated by Shell for ops register——不要手工编辑它任何改动都会在下次运行 register.py 时被覆盖。同理CPU 算子由它写入 CPU 注册文件。几何计算完成后跳到步骤 4单元测试。3.B 路径二CPU 后端实现当算子无法拆解、或性能敏感需要原生实现时走 CPU 后端路径。核心是继承Execution基类定义见 Execution.hpp实现两个生命周期方法onResize形状变化时申请缓存与onExecute实际计算。3.B.1 创建头文件在source/backend/cpu/下创建CPUMyCustomOp.hpp// // CPUMyCustomOp.hpp // MNN // #ifndef CPUMyCustomOp_hpp #define CPUMyCustomOp_hpp #include core/Execution.hpp namespace MNN { class CPUMyCustomOp : public Execution { public: CPUMyCustomOp(Backend* backend, const Op* op); virtual ~CPUMyCustomOp() default; virtual ErrorCode onResize(const std::vectorTensor* inputs, const std::vectorTensor* outputs) override; virtual ErrorCode onExecute(const std::vectorTensor* inputs, const std::vectorTensor* outputs) override; private: // 算子参数 // int mAxis 0; // 临时缓存 // Tensor mCache; }; } // namespace MNN #endif // CPUMyCustomOp_hpp3.B.2 创建实现文件在source/backend/cpu/下创建CPUMyCustomOp.cpp// // CPUMyCustomOp.cpp // MNN // #include CPUMyCustomOp.hpp #include core/Macro.h #include core/TensorUtils.hpp #include backend/cpu/CPUBackend.hpp namespace MNN { CPUMyCustomOp::CPUMyCustomOp(Backend* backend, const Op* op) : Execution(backend) { // 从 op 中读取参数 // auto param op-main_as_MyCustomOpParam(); // mAxis param-axis(); } ErrorCode CPUMyCustomOp::onResize(const std::vectorTensor* inputs, const std::vectorTensor* outputs) { // 在此申请临时缓存仅在输入形状变化时调用 // backend()-onAcquireBuffer(mCache, Backend::DYNAMIC); // backend()-onReleaseBuffer(mCache, Backend::DYNAMIC); return NO_ERROR; } ErrorCode CPUMyCustomOp::onExecute(const std::vectorTensor* inputs, const std::vectorTensor* outputs) { // 1. 获取输入输出数据指针 auto input inputs[0]; auto output outputs[0]; auto inputPtr input-hostfloat(); auto outputPtr output-hostfloat(); MNN_ASSERT(inputPtr ! nullptr); MNN_ASSERT(outputPtr ! nullptr); // 2. 获取形状信息 int totalSize 1; for (int i 0; i input-dimensions(); i) { totalSize * input-length(i); } // 3. 实际计算简单单线程版本先确保正确性 for (int i 0; i totalSize; i) { outputPtr[i] inputPtr[i]; // ← 替换为实际计算逻辑 } return NO_ERROR; } // 注册 Creator 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, op); } }; REGISTER_CPU_OP_CREATOR(CPUMyCustomOpCreator, OpType_MyCustomOp); } // namespace MNNCPU 后端的源码级原理执行生命周期与注册宏两阶段执行模型MNN 的Execution抽象Execution.hpp刻意把准备与计算分离——onResize只在输入形状变化时被调度器调用用于申请/调整临时缓存onExecute在每次推理时执行实际计算。CPU 算子的临时缓存通过backend()-onAcquireBuffer(buf, Backend::DYNAMIC)/onReleaseBuffer(buf, Backend::DYNAMIC)向后端统一内存池申请DYNAMIC语义表示该内存可在不同算子间复用不要持有跨onExecute的指针。数据访问input-hostfloat()返回 CPU 可直访的连续内存指针按elementSize()长度遍历即可dimensions()/length(i)/elementSize()分别给出维度数、第 i 维大小与总元素数。注意hostint32_t()适用于 int32 输入索引类算子如 Gather常用。注册宏REGISTER_CPU_OP_CREATOR(Creator, OpType)定义于 CPUBackend.hpp把Creator挂到OpType_MyCustomOp与 CPU 后端的映射表上onCreate返回的执行对象在每次推理会话中由后端调度器创建。MNN_ASSERT宏core/Macro.h用于调试期输入合法性校验。返回值语义NO_ERROR表示成功INPUT_DATA_ERROR表示输入数据非法另有NOT_SUPPORT算子在当前后端不可执行等错误码定义于 include/MNN/ErrorCode.hpp。3.B.3 关键 API 速查表API说明input-hostfloat()获取 float 数据指针input-hostint32_t()获取 int32 数据指针input-dimensions()维度数input-length(i)第 i 维大小input-elementSize()总元素数backend()-onAcquireBuffer(buf, Backend::DYNAMIC)申请临时缓存backend()-onReleaseBuffer(buf, Backend::DYNAMIC)释放临时缓存可被复用NO_ERROR返回成功INPUT_DATA_ERROR返回输入错误REGISTER_CPU_OP_CREATOR(Creator, OpType)注册 CPU 算子3.B.4 运行注册脚本并编译python3 tools/script/register.py cd build cmake .. -DMNN_BUILD_TESTON make -j$(nproc)步骤 3 测试标准与验收清单测试方法# 1. register.py 运行成功 python3 tools/script/register.py # 2. 编译通过 cd build cmake .. -DMNN_BUILD_TESTON make -j$(nproc) # 应该无编译和链接错误通过标准register.py 运行无错误编译无错误、无链接错误如果是几何计算GeometryOPRegister.cpp 中包含新算子grep MyCustomOp source/geometry/GeometryOPRegister.cpp可见extern声明与___GeometryMyCustomOp___create__();调用如果是 CPU 实现CPU 注册文件中包含新算子编译通过 ≠ 逻辑正确步骤 3 只保证代码能编译。逻辑正确性在步骤 4单元测试中验证——届时你会把算子的 MNN 输出与参考实现numpy 等对拍。因此在 3.B 的onExecute里先写单线程朴素版本、确保正确性是刻意为之的工程策略性能优化留到步骤 5step5-optimize.md。常见错误排查表错误原因修复编译错误OpType_MyCustomOp undeclaredSchema 头文件未更新重新执行步骤 1 的generate.sh见 step1-schema.md链接错误undefined referenceregister.py 未运行重新运行python3 tools/script/register.py编译错误头文件未找到include 路径错误检查#include的路径下一步步骤 3 通过后进入 step4-test.md步骤 4单元测试。只有测试对拍通过整个算子的正确性才算真正闭环若你的算子走的是 CPU 原生路径且需要覆盖多后端再继续阅读 step5-optimize.md 进行 Metal / OpenCL / Vulkan / CUDA 扩展与性能优化。【免费下载链接】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），仅供参考
