MNN MathOp 参考:用 Express 表达式系统编写元素级、归约与矩阵运算
MNN MathOp 参考用 Express 表达式系统编写元素级、归约与矩阵运算【免费下载链接】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/MNNMNN 的 Express表达式API 中MNN::Express命名空间下的 MathOp 一族函数是构建前向计算图的核心积木从_Add这类逐元素二元运算到_ReduceMean这类跨维归约再到_MatMul、_BatchMatMul这类线性代数运算均以VARPVariable Pointer为输入输出惰性构建出可被 MNN 调度器优化的Expr图。本文以 MathOp API 文档 为主线完整梳理各运算的签名、参数约束与返回类型并结合 MathOp.hpp 头文件声明与 express/MathOp.cpp 实现源码说明这些函数在底层是如何拼装成 OpT 算子的以及_checkNC4HW4等源码级细节对实际使用的约束。1. MathOp 在 Express 体系中的位置MathOp 的全部函数声明在 include/MNN/expr/MathOp.hpp按功能被源码注释划分为四个区块//BinaryOPs二元运算_Add、_Subtract、比较、位运算等//UnaryOPs一元运算三角/双曲函数、激活函数、_Sqrt等//ReduceOPs归约运算_ReduceSum系列及其Mutable变体//OtherOPs_Cast、_MatMul、_ArgMax、_ScatterNd、_OneHot等其余运算。头文件中的函数签名与文档一一对应但存在两处文档未覆盖的补充MNN_PUBLIC VARP _MulSilu(VARP x, VARP y);与MNN_PUBLIC VARP _Silu(VARP x);二者分别对应 express/MathOp.cpp 中的BinaryOpOperation_MUL_SILU和UnaryOpOperation_SILU是服务 LLM 场景的 SiLUSwish激活及其乘加融合变体。使用方式与其他 MathOp 相同包含MNN/expr/ExprCreator.hpp内部引入 MathOp 声明以using namespace MNN::Express;后直接调用。从源码结构看每个 MathOp 函数都只是“构图动作”——它并不执行计算而是创建一个带参数的Expr最终由Executor/Session在实际 run 时才触发后端CPU/OpenCL/Metal 等的 Execution 实现。这一点对理解“返回值何时可读数据”很关键。2. 二元运算BinaryOPs二元运算统一签名为VARP _Xxx(VARP x, VARP y);两个参数支持Halide_Type_Int、Halide_Type_Float、Halide_Type_Int64、Halide_Type_Uint8比较/位运算类按文档限定为 Int/Float返回值类型与x相同。完整清单及语义如下API计算含义返回类型对应源码枚举_Addx y与 x 相同BinaryOpOperation_ADD_Subtractx - y与 x 相同BinaryOpOperation_SUB_Multiplyx * y与 x 相同BinaryOpOperation_MUL_Dividex / y真除与 x 相同BinaryOpOperation_REALDIV_FloorDivx // y整除与 x 相同BinaryOpOperation_FLOORDIV_Powx 的 y 次幂与 x 相同BinaryOpOperation_POW_Minimum/_Maximum逐元素取小/取大与 x 相同MINIMUM/MAXIMUM_Modx % yC 风格余数与 x 相同MOD_FloorMod向零取整除法后的余数与 x 相同FLOORMOD_SquaredDifference(x - y) * (x - y)与 x 相同SquaredDifference_Atan2y / x 的反正切保留符号与 x 相同ATAN2_Greater/_GreaterEqualx y / x yboolGREATER/GREATER_EQUAL_Less/_LessEqualx y / x yboolLESS/LESS_EQUAL_Equal/_NotEqualx y / x ! yboolEQUAL/NOTEQUAL_LogicalOrx 逻辑或 ytrue/falseLOGICALOR_BitwiseAnd/_BitwiseOr/_BitwiseXorx y、x | y、x ^ y与 x 相同BITWISE_AND/BITWISE_OR/BITWISE_XOR2.1 _BiasAdd一个特殊的加法MNN_PUBLIC VARP _BiasAdd(VARP value, VARP bias);_BiasAdd是“bias 限定为一维”的加法特例bias必须是 1-D 变量其长度需与value的通道维度匹配且支持广播value可以是任意维数。与普通_Add的差别在于量化场景下 bias 允许使用与 value 不同的量化类型。从 express/MathOp.cpp 的实现看它直接转发为return _Add(value, bias);即 MNN 通过BinaryOpOperation_ADD加广播机制实现而不是单独的算子。2.2 底层实现_Binary 助手与 NC4HW4 自动转换所有二元运算最终收敛到 express/MathOp.cpp 中的静态函数_Binary(VARP x, VARP y, BinaryOpOperation operation)其做三件事对 x、y 调用_checkNC4HW4若任一输入的内存序为NC4HW4MNN 默认的高效布局先经_Convert(x, NCHW)转为NCHW再参与运算——这意味着用 NC4HW4 输入直接做二元运算时图上会多一个 Convert 节点用FlatBufferBuilder默认容量MNN_DEFAULT_FLATBUFFER_SIZE 32序列化为BinaryOp参数构造OpType_BinaryOp的Expr以{x, y}两个输入创建Variable操作类型编码进BinaryOpT::opType。因此上表第三列的枚举值是这些 API 在模型文件中实际落地的 op 参数阅读 schema/current/MNN_generated.h 中BinaryOp定义可进一步核对。3. 一元运算UnaryOPs一元运算统一签名为VARP _Xxx(VARP x);绝大多数要求输入为Halide_Type_Int或Halide_Type_Float返回与 x 同类型的变量。按功能分组如下符号与取整API语义_Sign符号函数x0 得 0x0 得 -1x0 得 1_Abs逐元素绝对值。例MNN.const((-1.0, -2.0, 3.0), (3,))经_Abs得(1.0, 2.0, 3.0)_Negative逐元素取负。例(-1.0, -2.0, 3.0)得(1.0, 2.0, -3.0)_Floor返回不大于 x 的最大整数_Ceil返回不小于 x 的最小整数_Round返回四舍五入的整数文档限定输入为Halide_Type_Float_Square逐元素平方幂、指数与对数API语义_Sqrt平方根_Rsqrt平方根的倒数1/sqrt(x)_Exp自然指数 e^x_Log自然对数_Log1pln(1 x)_Expm1e^x - 1文档限定Halide_Type_Float_Reciprocal倒数 1/x三角与反三角、双曲与反双曲API语义文档标注类型约束_Sin/_Cos/_Tan正弦 / 余弦 / 正切_Sin、_Cos、_Tan标注为Halide_Type_Float_Asin/_Acos/_Atan反正弦 / 反余弦 / 反正切_Atan标注为Halide_Type_Float_Sinh/_Cosh双曲正弦 / 双曲余弦Int 或 Float_Asinh/_Acosh/_Atanh反双曲正弦 / 反双曲余弦 / 反双曲正切_Atanh标注为Halide_Type_Float误差函数族与激活函数API语义_Erf/_Erfc/_Erfinv高斯误差函数 / 互补误差函数 / erf 的逆函数_Gelu高斯误差线性单元激活文档限定Halide_Type_Float_Tanh双曲正切输入域 [-inf, inf]输出域 [-1, 1]_Sigmoid逻辑斯蒂sigmoid激活_HardswishHardswish 神经网络激活express/MathOp.cpp 中每个一元运算都是一行转发例如VARP _Abs(VARP x) { return _Unary(x, UnaryOpOperation_ABS); }静态助手_Unary将运算编码为OpType_UnaryOpUnaryOpT::opType落图。头文件中的函数顺序_Sign、_Abs、_Negative…_Expm1、_Hardswish与文档条目顺序一致可以据此交叉核对枚举值。一个值得注意的量化细节test/op/UnaryTest.cpp 提供了_UnaryInt8辅助函数——对 INT8 输入部分一元运算_Square、_Sqrt、_Rsqrt等通过预计算查表 bufferUnaryOpT::tableInt8实现逐值映射这也解释了为何一元 API 允许Halide_Type_Int输入量化路径下以 LUT 近似完成。4. 归约运算ReduceOPs归约族全部采用同一签名形态MNN_PUBLIC VARP _ReduceXxx(VARP input_variable, INTS axis {}, bool keepDims false);其中input_variable被归约的变量应为数值类型axis要归约的维度集合取值范围为[-rank(input_variable), rank(input_variable))为空默认时归约所有维度返回单元素变量keepDims为 true 时保留长度为 1 的归约维度否则每归约一个维度 rank 减 1返回与input_variable同类型的归约后变量。具体成员函数及语义API归约方式_ReduceSum/_ReduceSumMutable逐维求和_ReduceMean/_ReduceMeanMutable逐维求均值_ReduceMax/_ReduceMaxMutable逐维取最大_ReduceMin/_ReduceMinMutable逐维取最小_ReduceProd/_ReduceProdMutable逐维求积_ReduceAny/_ReduceAnyMutable逐维“逻辑或”_ReduceAll/_ReduceAllMutable逐维“逻辑与”Mutable变体与非 Mutable 变体的差别在 MathOp.hpp 中可以直接确认axis参数类型为VARP而非INTS即归约轴本身可以是运行时可变的张量。对应 express/MathOp.cpp 中的两个静态助手_Reduce把静态INTS维度写入ReductionParam::dim_ReduceMutable不写 dim 字段而是把axis作为第二个图输入{x, dim}传入让 Shape 推理在运行期解析轴。源码对Mutable路径还留有// TODO: Remove Copy注释从源码结构看这是该路径当前实现的已知开销点。express/MathOp.cpp 还展示了 MathOp 的组合式用法——_ReduceVariance由现成 API 复合而成VARP _ReduceVariance(VARP input_variable, INTS axis, bool keepdims) { auto mean _ReduceMean(input_variable, axis, true); // to use broadcast of subtract auto variance _ReduceMean(_Square(_Subtract(input_variable, mean)), axis, keepdims); return variance; }这正是 Express API 的典型范式方差没有独立算子而是“均值 → 广播相减 → 平方 → 再求均值”四个图节点拼接读者可以按同样方式自行组合任意数学表达式。test/op/ReductionTest.cpp 提供了可直接参照的验证用例_Input({4}, NCHW)写入{-1.0, -2.0, 3.0, 4.0}调用_ReduceSum(input)axis 留空归约全维断言输出为{4.0}ReduceSumMultiTest则覆盖{4, 10, 1, 4}输入按指定轴归约并针对低精度模式precision 2放宽误差阈值至 0.1。5. Eltwise 运算与 Int8 量化版本5.1 带系数的逐元素运算MNN_PUBLIC VARP _Prod(VARP a, VARP b, std::vectorfloat coeff); MNN_PUBLIC VARP _Sum(VARP a, VARP b, std::vectorfloat coeff); MNN_PUBLIC VARP _Max(VARP a, VARP b, std::vectorfloat coeff); MNN_PUBLIC VARP _Sub(VARP a, VARP b, std::vectorfloat coeff);四个函数要求 a、b 为Halide_Type_Floatcoeff是 blob-wise按输入张量整体的系数分别对应EltwiseType_PROD / SUM / MAXIMUM / SUB。文档中_Sub的返回描述写作“下标元素”但从 express/MathOp.cpp 的实现return _Eltwise(a, b, EltwiseType_SUB, coeff);可以确认其实际语义是逐元素求差 a - b文档该处描述有误以源码为准。5.2 Int8 量化 EltwiseMNN_PUBLIC VARP _EltwiseProdInt8(VARP x, VARP y, std::vectorint8_t x_weight, std::vectorint32_t x_bias, std::vectorfloat x_scale, std::vectorfloat x_tensorScale, std::vectorint8_t y_weight, std::vectorint32_t y_bias, std::vectorfloat y_scale, std::vectorfloat y_tensorScale, std::vectorint8_t output_weight, std::vectorint32_t output_bias, std::vectorfloat output_scale, std::vectorfloat output_tensorScale);_EltwiseProdInt8、_EltwiseSumInt8、_EltwiseSubInt8、_EltwiseMaxInt8四个函数签名完全同构分别在 Eltwise 层对 x、y 做累计乘积、求和、求差值、求最大值。参数含义x、y一个变量Halide_Type_Float或Halide_Type_Int类型之一Int8 路径下为量化张量x_weight/y_weight/output_weightx、y、输出数据的权值int8 反量化查找表x_bias/y_bias/output_bias对应的偏差int32 零点x_scale/y_scale/output_scale比例因子x_tensorScale/y_tensorScale/output_tensorScale张量级比例因子。返回 VARP 类型变量。源码中四个 API 统一转发到静态函数_EltwiseInt8(...)它把三组量化参数分别装入QuantizedFloatParamTinputQuan0/inputQuan1/outputQuan构造OpType_EltwiseInt8EltwiseInt8T落图注意 express/MathOp.cpp 中param_x-tensorScale y_tensorScale;一行把 x 的 tensorScale 赋成了 y 的值从源码结构看这属于实现层细节使用该 API 时 x 的tensorScale参数实际未按传入值生效。6. 类型转换、矩阵运算与统计定位6.1 _Cast// 模板版本目标类型由 T 推导 VARP _Cast(VARP x) { return _Cast(x, halide_type_ofT()); } MNN_PUBLIC VARP _Cast(VARP x, halide_type_t dtype);将变量强制转换为目标类型返回与 x 形状相同、与 dtype 类型相同的变量。文档声明支持Halide_Type_Int、Halide_Type_Float、Halide_Type_Int64、Halide_Type_Uint8但从 express/MathOp.cpp 的_convertDataType映射看实际被翻译为有效DataType的只有 float、uint8、int8、int32 四种其余类型含 int64会落到DataType_DT_INVALID使用时应以这四种为目标类型。6.2 _MatMul 与 _BatchMatMulMNN_PUBLIC VARP _MatMul(VARP a, VARP b, bool tranposeA false, bool tranposeB false);计算矩阵 a * 矩阵b。输入必须是二维矩阵且 a 的内部维数tranposeA为 true 时先转置必须匹配 b 的外部维数tranposeB为 true 时先转置。参数tranposeA/tranposeB默认 false。实现上直接构造OpType_MatMulMatMulT写入transposeA/transposeB字段可对照 test/expr/MatMulTest.cpp 查看测试用法。MNN_PUBLIC VARP _BatchMatMul(VARP x, VARP y, bool adj_x false, bool adj_y false);批量矩阵乘x、y 为形状[..., r_x, c_x]与[..., r_y, c_y]的二阶及以上张量逐切片相乘并归入同一批输出输出形状为[..., r_o, c_o]其中r_o c_x if adj_x else r_x、c_o r_y if adj_y else c_y计算公式为output[..., :, :] matrix(x[..., :, :]) * matrix(y[..., :, :])。adj_x/adj_y为 true 时对相应切片先做伴随转置并共轭默认 false。6.3 _NormalizeMNN_PUBLIC VARP _Normalize(VARP x, int32_t acrossSpatial, int32_t channelShared, float eps, std::vectorfloat scale);将 x 转换为指定标准化格式五个参数分别写入NormalizeT的acrossSpatial、channelShared、eps、scale字段文档对前三个参数的说明为占位式描述实际含义以OpType_Normalize的 Normalize 参数表定义为准可参考 source/backend/cpu 下 CPUNormalize 的实现核对语义。6.4 _ArgMax / _ArgMinMNN_PUBLIC VARP _ArgMax(VARP input, int axis 0); MNN_PUBLIC VARP _ArgMin(VARP input, int axis 0);分别返回张量沿指定坐标轴上最大值 / 最小值的索引返回索引值。axis取值范围[-rank(input), rank(input))向量输入用axis 0默认axis 0。输入支持Halide_Type_Float或Halide_Type_Int。源码实现express/MathOp.cpp值得注意两点二者共用ArgMaxT参数结构体_ArgMin只是把 op type 置为OpType_ArgMin同时会先对输入做_checkNC4HW4转换。topK、softmaxThreshold等字段在此 API 中被显式置 0。test/op/ArgMaxTest.cpp 覆盖了对应的取值验证。7. 索引操作与数据构造类 API7.1 _UnravelIndex 与 _Scatter 族MNN_PUBLIC VARP _UnravelIndex(VARP indices, VARP dims);返回indices中的元素在维度为dims的数组中的索引值默认按元组形式返回。MNN_PUBLIC VARP _ScatterNd(VARP indices, VARP updates, VARP shape); MNN_PUBLIC VARP _ScatterNd(VARP indices, VARP updates, VARP shape, VARP input); MNN_PUBLIC VARP _ScatterNd(VARP indices, VARP updates, VARP shape, int reduction); MNN_PUBLIC VARP _ScatterNd(VARP indices, VARP updates, VARP shape, VARP input, int reduction);四个重载语义一致根据声明的索引通过对声明形状张量零张量或传入的input内的单个切片或值做分散更新来构造输出张量。参数indices保存指向输出张量的索引、updates保存索引处的值、shape输出张量规定形状、input可选的输入张量数据、reduction归约方式默认 -1 表示直接写。从 express/MathOp.cpp 看reduction ! -1时会被解释为BinaryOpOperation枚举写入BinaryOpT即同一索引位置多次命中时可选择按何种二元运算合并。MNN_PUBLIC VARP _ScatterElements(VARP data, VARP indices, VARP updates, int reduction -1); MNN_PUBLIC VARP _ScatterElements(VARP data, VARP indices, VARP updates, VARP axis, int reduction -1);根据updates和indices更新data的值并返回第二个重载额外提供axis表示沿行还是列操作reduction同样默认 -1。7.2 _OneHot、_BroadcastTo、_LinSpace、_RandomUnifomMNN_PUBLIC VARP _OneHot(VARP indices, VARP depth, VARP onValue, VARP offValue, int axis -1);独热编码分类问题中将存在的数据类别置为onValue常为 1其余置为offValue常为 0。depth是定位填充维度的标量axis是要填充的轴默认 -1。MNN_PUBLIC VARP _BroadcastTo(VARP a, VARP shape);利用广播规则把a扩展到期望输出形状shape。MNN_PUBLIC VARP _LinSpace(VARP start, VARP stop, VARP num);创建等差数列start为区间最小值stop为区间最大值num为数据量分割份数。MNN_PUBLIC VARP _RandomUnifom(VARP shape, halide_type_t dtype, float low 0.0f, float high 1.0f, int seed0 0, int seed1 0);生成指定形状与类型的均匀分布随机数low/high为区间边界默认 0~1seed0/seed1为随机因子默认 0。函数名中的 “Unifom” 拼写与 MathOp.hpp 声明一致属源码既有拼写调用时请保持原样。7.3 _CumSum、_CumProd、_Svd、_HistogramMNN_PUBLIC VARP _CumSum(VARP x, int axis, bool exclusive false, bool reverse false); MNN_PUBLIC VARP _CumProd(VARP x, int axis);分别计算 x 沿axis坐标轴的累加值与累计乘积返回与 x 同类型的变量。exclusive为 true 时当前元素不参与自身累加结果左移一位reverse控制是否逆向累计。test/op/CumTest.cpp 给出了完整验证对_Input({2, 2, 2}, NCHW)写入{1..8}_CumProd(input, 0)期望{1., 2., 3., 4., 5., 12., 21., 32.}axis1、axis2 各有对应期望序列CumSumTest对_CumSum(input, 0)期望{1., 2., 3., 4., 6., 8., 10., 12.}。MNN_PUBLIC VARPS _Svd(VARP x);奇异值分解。注意返回类型是VARPS变量数组express/MathOp.cpp 中_Svd创建带 3 个输出的Expr并拆分为 U、Σ、V 三个Variable返回这是 MathOp 中唯一的“一函数多输出”API。MNN_PUBLIC VARP _Histogram(VARP x, int bin, int min, int max, int channel -1);直方图统计bin为统计区间个数min/max为统计范围的最小/最大值channel为通道维默认 -1。从 express/MathOp.cpp 结构看其参数被复用进ArgMaxT字段bin→outMaxVal、min→softmaxThreshold、max→topK、channel→axisop type 为OpType_Histogram读者阅读模型文件时可按此对应关系解读。8. 实战用 MathOp 构建并运行一个表达式综合 test/op/ReductionTest.cpp 与 test/op/CumTest.cpp 的写法一个最小可用示例如下依赖 MNN 的 Express 库输入/常量创建来自 include/MNN/expr/NeuralNetWorkOp.hpp 的_Input/_Const#include MNN/expr/Expr.hpp #include MNN/expr/ExprCreator.hpp using namespace MNN::Express; // 1) 构造 4x3 的输入并写数据 auto input _Input({4, 3}, NCHW, halide_type_offloat()); input-setName(input_tensor); const float data[] {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f}; memcpy(input-writeMapfloat(), data, sizeof(data)); input-unMap(); // 2) 组合 MathOp 构图每行均值 全体求和 转 int8 auto rowMean _ReduceMean(input, {0}, false); // 沿 batch 维归约 auto total _ReduceSum(input); // axis 留空归约全维 auto asInt8 _Cast(rowMean, halide_type_ofint8_t()); // 显式指定目标类型 // 3) 取结果表达式图在此触发推理 printf(total %f\n, *total-readMapfloat()); auto meanPtr rowMean-readMapfloat();要点提示_ReduceMean(input, {0}, true)中keepDims true时结果形状保留被归约维为 1便于后续直接广播源码中_ReduceVariance正是依赖这一点做“减均值”若输入来自 NC4HW4 布局的中间张量二元/归约运算前 MNN 会自动插入_Convert到 NCHW_checkNC4HW4一般无需手工处理各 API 的正确性可对照仓库测试用例自行验证一元运算见 test/op/UnaryTest.cpp归约见 test/op/ReductionTest.cpp累计运算见 test/op/CumTest.cpp索引类见 test/op/ArgMaxTest.cpp。9. 文档与源码差异小结撰写或调用 MathOp 接口时建议以头文件 include/MNN/expr/MathOp.hpp 为最终签名依据并注意以下几处文档与源码的出入Mutable 归约的 axis 类型MathOp 文档 将_ReduceSumMutable等写作INTS axis而头文件实际为VARP axis运行时可变轴二者调用方式不同_Sub的语义文档返回描述为“下标元素”源码实现为EltwiseType_SUB逐元素求差以后者为准_Cast的目标类型文档列出 Int64但源码_convertDataType仅映射 float/uint8/int8/int32 为有效DataType拼写与描述小误_RandomUnifom的 “Unifom” 为源码既有拼写_ReduceMinMutable文档描述误写为“最大值”源码为ReductionType_MINIMUM头文件额外 API_Silu与_MulSilu未在文档中列出但已在头文件与实现中就绪可直接用于 LLM 推理图的 SiLU 节点构建。掌握上述清单后你可以仅依赖 Express API 组合出加、减、乘、除、归约、矩阵乘、索引更新等任意数学子图交由 MNN 调度到 CPU、OpenCL、Metal 等后端执行无需手写 OpT 或 FlatBuffer 代码。【免费下载链接】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),仅供参考