MNN ARM CPU 算子性能优化。涵盖计算拆解、函数复用、多线程、数据排布、ARM 汇编编写。采用"先正确,再加速"原则,基于性能基准测试驱动优化。
触发条件:当用户请求优化某个算子/内核在 ARM CPU 上的性能时触发。常见表述包括:"优化xxx的ARM性能"、"加速xxx算子"、"写xxx的NEON实现"、"用SME2实现xxx"等。
本 SKILL 指导 AI Agent 对 MNN 的 ARM CPU 后端进行性能优化。遵循 "先正确,再加速" 原则,每次优化都要保证结果不变。
MNN 的 CoreFunctions 中已经包含了。这些函数已经针对不同指令集(NEON/FP16/SDOT/I8MM/SME2)编写了专门的汇编内核,性能远超任何 C++ 循环或 Vec4 包装。
优化的核心思路是:将算子的计算逻辑拆解为这些已有函数的组合调用,而不是自己写循环。
| 函数 | 作用 | 替代了什么 |
|---|---|---|
gcore->MNNPackedMatMul | 矩阵乘 C = A × B(已有 NEON/FP16/SME2 汇编) | 任何双重循环的矩阵乘 |
gcore->MNNPackedMatMulRemain | 矩阵乘余数处理 | MatMul 的尾部处理 |
gcore->MNNComputeMatMulForE_1 | 矩阵向量乘 y = A × x(E=1 时专用) | 循环实现的 MatVec |
gcore->MNNComputeMatMulForH_1 | 向量矩阵乘 y = x × B(H=1 时专用) | 循环实现的 VecMat |
MNNScaleAndAddBiasScalar | y = x * scale + bias | 循环乘标量/加标量 |
gcore->MNNScaleAndAddBias | 按通道 scale + bias | 循环乘/加 |
MNNExp | 批量 exp(x) | 循环调用 expf() |
MNNSiLu / MNNSiLuLowp | 批量 SiLU 激活 | 循环 x*sigmoid(x) |
MNNSoftmax | Softmax(含 Flash Attention 支持) | 循环 exp + sum + div |
MNNNorm | LayerNorm / RMSNorm | 循环求范数 |
gcore->MNNPackCUnit / MNNUnpackCUnit | NC4HW4 Pack/Unpack | 循环数据重排 |
gcore->MNNPackC4ForMatMul_A | MatMul 的 A 矩阵 Pack | 循环重排 A |
gcore->MNNPackForMatMul_B | MatMul 的 B 矩阵 Pack | 循环重排 B |
gcore->MNNConvRunForLineDepthwise | Depthwise 卷积 | 循环卷积 |
MNNMatrixAdd / MNNMatrixSub | 矩阵加减 | 循环加减 |
MNN_CONCURRENCY_BEGIN/END | 多线程并行 | 单线程循环 |
1. 函数语义完全匹配
在用 MNN 函数替换循环前,必须阅读函数签名和实现,确认数学语义完全一致。常见陷阱:
不要只看函数名就假设可以替换。 如果语义不匹配,应寻找其他函数、调整数据,实在不行才退回手动实现并注释说明。
2. in-place 安全性
部分 MNN 函数不支持 dst == src(in-place 调用),因为内部实现会先写 dst 再读 src:
| 函数 | in-place (dst==src) | 说明 |
|---|---|---|
MNNScaleAndAddBiasScalar | ✅ 安全 | 逐元素操作 |
MNNSiLu / MNNSiLuLowp | ❌ 不安全 | 内部先写 dst 再读 src |
MNNExp | ❌ 不安全 | 同上 |
MNNNorm | ✅ 安全 | 只读 src,只写 dst |
gcore->MNNComputeMatMulForE_1 | ✅ 安全 | 输出独立于输入 |
当函数不支持 in-place 时,需要 scratch buffer(可复用 onResize 预分配的 buffer)。 判断方法:阅读函数实现或写小测试验证
dst==src时结果是否正确。
// ❌ 错误1:已有 MNN 函数时用 Vec4 替代(MNN 函数有汇编优化,Vec4 只是 intrinsic 包装)
using Vec4 = Math::Vec<float, 4>;
for (int i = 0; i + 3 < size; i += 4) {
Vec4 v = Vec4::load(data + i);
v = v * Vec4(scale);
Vec4::save(data + i, v);
}
// ✅ 正确:直接调用已有函数
MNNScaleAndAddBiasScalar(data, data, 0.0f, scale, size);
// ❌ 错误2:不考虑数据规模,总是用重量级函数
// MNNPackedMatMul 需要 Pack/Unpack,对小矩阵开销可能大于计算
// ✅ 正确:根据数据规模选择策略
// 大规模 → 用 MNN 函数(Pack 开销可摊薄)
// 小规模 → 保持朴素循环(编译器自动向量化,零额外开销)
// MatVec(一个维度为1) → 用 MNNComputeMatMulForE_1(无需 Pack)
// ❌ 错误3:用循环实现 MatVec(S^T @ q)
for (int j = 0; j < dv; ++j) {
float sum = 0;
for (int i = 0; i < dk; ++i)
sum += S[i*dv+j] * q[i];
out[j] = sum;
}
// ✅ 正确:直接调用
gcore->MNNComputeMatMulForE_1(q, S, out, nullptr, &matParam, 0);
// ❌ 错误4:用循环实现 exp
for (int i = 0; i < size; ++i) dst[i] = expf(src[i]);
// ✅ 正确:
MNNExp(dst, src, offset, size);
// ❌ 错误5:用 std::vector 分配临时缓存