skills/arm-cpu-optimize/SKILL.md
触发条件:当用户请求优化某个算子/内核在 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 分配临时缓存
std::vector<float> temp(size); // 每次调用都 malloc/free
// ✅ 正确:在 onResize 中用 Backend 的内存池
mTemp.reset(Tensor::createDevice<float>({size}));
backend()->onAcquireBuffer(mTemp.get(), Backend::DYNAMIC);
backend()->onReleaseBuffer(mTemp.get(), Backend::DYNAMIC);
// ❌ 错误6:用裸 NEON intrinsic 写循环(仅在没有 MNN 函数且性能敏感时才考虑,且更推荐写 .S 汇编)
#include "core/SimdHeader.h"
float32x4_t vsum = vdupq_n_f32(0.0f);
// ❌ 错误7:不验证 in-place 安全性就用 dst==src 调用
// 部分 MNN 函数不支持 in-place,会静默产生错误结果
// ✅ 正确:查阅上方 in-place 安全性表,不确定时使用 scratch buffer
// ❌ 错误8:不验证函数语义就替换(函数名匹配 ≠ 数学语义匹配)
// ✅ 正确:替换前阅读函数签名/实现,确认参数含义和计算逻辑完全一致
强烈不建议的做法:
| 避免使用 | 原因 | 应该用什么 |
|---|---|---|
#include "core/SimdHeader.h" | 裸 NEON intrinsic 性能不一定最优且绑定平台 | 优先用 MNN 已有函数 或 编写 .S 汇编文件 |
std::vector<float> 在 onExecute 中 | 每次运行都 malloc/free 开销巨大 | Tensor + Backend 内存池 用在 onResize |
#ifdef MNN_USE_NEON ... #else ... #endif | 增加代码分支,维护困难 | 封装到底层函数,上层统一调用 |
有条件允许 —— Vec4 循环:
| 场景 | 是否允许 | 说明 |
|---|---|---|
| 已有 MNN 函数能覆盖 | ❌ | 必须用 MNN 函数,禁止用 Vec4 替代 |
| 没有对应 MNN 函数,且是性能热点 | ✅ | Vec4 作为中间方案,优于朴素循环 |
| 没有对应 MNN 函数,且计算量极小 | ❌ | 保持朴素循环,编译器自动向量化即可 |
Vec4(
#include "math/Vec.hpp")本质是 intrinsic 的跨平台包装,性能不如专门调优的汇编函数,但远优于朴素标量循环。当 MNN 没有对应的已有函数且评估手写汇编成本过高时,Vec4 是合理的优化手段。
看到一段循环代码 →
├─ 是双重循环的乘加?
│ ├─ 有一个维度为 1(MatVec)? → 用 MNNComputeMatMulForE_1(无需 Pack)
│ ├─ 数据规模大(值得 Pack 开销)? → 用 MNNPackedMatMul
│ └─ 数据规模小(Pack 开销 > 计算)? → 保持朴素循环
├─ 是单循环乘标量/加标量? → 用 MNNScaleAndAddBiasScalar
├─ 是循环调用 expf/silu/sigmoid? → 用 MNNExp/MNNSiLu
├─ 是循环做卷积? → 用 MNNConvRunForLineDepthwise
├─ 是循环做数据重排? → 用 MNNPackCUnit/MNNUnpackCUnit
├─ 是循环做 softmax? → 用 MNNSoftmax
├─ 是循环做范数/归一化? → 用 MNNNorm(先验证语义匹配!)
├─ 以上都不匹配?
│ ├─ 能拆解为已有函数的组合? → 组合调用(如外积 = MatMul 的特例)
│ ├─ 计算量极小? → 保持朴素循环(编译器自动向量化)
│ ├─ 是性能热点且可向量化? → 用 Vec4 循环作为中间方案
│ └─ Vec4 也不够且是核心热点? → 写新的 .S 汇编 kernel
│
⚠️ 替换前:1) 验证函数语义匹配 2) 确认 in-place 安全性
核心限制:
schema/private/和source/internal/目录不应对 AI 公开或被随意修改。
MNN 使用 NC4HW4 数据格式作为默认 CPU 布局,pack 大小由
CoreFunctions::pack决定(FP32=4, FP16=8)
参考学习:在开始优化前,强烈建议阅读
source/backend/cpu/CPUAttention.cpp,学习它如何调用gcore->MNNComputeMatMulForE_1、MNNScaleAndAddBiasScalar、MNNSoftmax等函数。你的优化代码应该通过调用 CoreFunctions 来实现高性能。
报告文件:优化完成后请将性能报告写入
<算子名>_optimization.md,而不是仅在终端打印。
从低到高,每级都向下兼容:
| 指令集 | 编译宏/检测 | Pack 大小 | 典型场景 | 代表芯片 |
|---|---|---|---|---|
| ARMv7 NEON | __arm__ | 4 (FP32) | 基础 SIMD | Cortex-A7/A15 |
| ARMv8 NEON | __aarch64__ | 4 (FP32) | 标准 64 位 | Cortex-A53/A72 |
| ARMv8.2 FP16 | MNN_ARM82 / supportFp16arith | 8 (FP16) | 半精度加速 | A55/A76/A78 |
| ARMv8.2 SDOT | supportSDot | - | INT8 点积加速 | A75+, A55r1+ |
| ARMv8.6 I8MM | MNN_ARM86 / supportI8mm | - | INT8 矩阵乘加速 | A78C/X2/X3 |
| ARMv9 SME2 | MNN_SME2 / supportSME2 | 可变 | 矩阵扩展引擎 | X4/X925 |
所有 CPU 函数指针都注册在 CoreFunctions 中(source/backend/cpu/compute/CommonOptFunction.h)。运行时根据 CPU 能力选择最优实现:
struct CoreFunctions {
// CPU 特性检测
bool supportFp16arith;
bool supportSDot;
bool supportI8mm;
bool supportSME2;
// Pack 参数
int pack; // FP32=4, FP16=8
int bytes; // FP32=4, FP16=2
// 关键函数指针
void(*MNNPackedMatMul)(...); // 矩阵乘主核心
void(*MNNPackedMatMulRemain)(...); // 矩阵乘余数处理
void(*MNNPackC4ForMatMul_A)(...); // 输入数据 Pack
void(*MNNPackForMatMul_B)(...); // 权重数据 Pack
void(*MNNGetMatMulPackMode)(&eP, &lP, &hP); // 获取 Pack 参数
// ...
};
| 排布 | 说明 | 使用场景 |
|---|---|---|
| NC4HW4 / NC8HW8 | 通道方向 Pack4/8,SIMD 友好 | 卷积、Pooling 等 |
| [e/eP, l/lP, eP, lP] | MatMul 的 A 矩阵 Pack | GEMM 优化 |
| [h/hP, l/lP, lP, hP] | MatMul 的 B 矩阵 Pack | 权重重排 |
| NCHW | 标准线性布局 | 形状计算、非 Pack 算子 |
source/backend/cpu/
├── CPUXxx.cpp/.hpp ← 算子主逻辑(调度、多线程)
├── compute/
│ ├── CommonOptFunction.h ← CoreFunctions 定义
│ ├── CommonOptFunction.cpp ← 默认 C++ 实现
│ ├── ConvOpt.h ← 卷积相关函数声明
│ └── ...
├── arm/
│ ├── arm64/
│ │ ├── MNNPackedMatMul.S ← FP32 NEON 矩阵乘
│ │ ├── MNNPackedMatMulFP16.S ← FP16 矩阵乘
│ │ ├── MNNGemmInt8AddBiasScale_*.S ← INT8 GEMM
│ │ ├── MNNPackedMatMul_int8.S ← SDOT 矩阵乘
│ │ ├── MNNPackedMatMulRemain_int8.S ← SDOT 余数
│ │ └── ...
│ └── arm32/
│ ├── MNNGemmInt8*.S ← 32 位 INT8 GEMM
│ └── ...
└── x86_64/
└── ... ← x86 SSE/AVX 实现
┌──────────────────────────────────────────────────────────┐
│ 步骤 0: 计算拆解 (step0-decompose.md) ★ 最关键的步骤 │
│ 输入: 算子源码 │
│ 输出: 每个计算逻辑对应的 MNN 已有函数映射表 │
│ 测试: 所有可复用的函数已识别,不存在遗漏 │
├──────────────────────────────────────────────────────────┤
│ 步骤 1: 建立性能基准 (step1-benchmark.md) │
│ 输入: 待优化的算子名和参数 │
│ 输出: test/speed/ 下的基准测试 + 基线数据 │
│ 测试: 基准测试能稳定运行,数据可复现 │
├──────────────────────────────────────────────────────────┤
│ 步骤 2: 分析瓶颈与制定方案 (step2-analyze.md) │
│ 输入: 步骤 0 的映射表 + 步骤 1 的基线数据 │
│ 输出: 优化方案(哪些换函数、哪些改排布、哪些需要新汇编) │
│ 测试: 方案可行性论证 │
├──────────────────────────────────────────────────────────┤
│ 步骤 3: C++ 优化(函数替换 + 多线程 + 排布) │
│ 输入: 步骤 2 的优化方案 │
│ 输出: 用 MNN 函数替换循环 + 多线程 + 内存池 │
│ 测试: 正确性验证 + 性能对比 │
├──────────────────────────────────────────────────────────┤
│ 步骤 4: ARM 汇编优化 (step4-asm.md)(仅在必要时) │
│ 输入: 步骤 3 中无法用已有函数覆盖的热点 │
│ 输出: ARM NEON/FP16/SDOT/I8MM/SME2 汇编实现 │
│ 测试: 正确性验证 + 性能对比 + 多指令集覆盖 │
├──────────────────────────────────────────────────────────┤
│ 步骤 5: 集成与回归测试 (step5-integrate.md) │
│ 输入: 步骤 3 或 4 通过 │
│ 输出: 全量回归测试 │
│ 测试: 所有相关算子测试通过 + 性能报告 │
└──────────────────────────────────────────────────────────┘
说明:步骤 4(汇编)往往不是必须的。如果通过步骤 3(利用现有 MNN 工具、重排版数据并应用多线程)已经达到了优异的性能提升,便可以直接进入集成测试阶段(步骤 5)。
| 情况 | 处理方式 |
|---|---|
| step3 某个替换导致性能下降 | 回退该替换,保持原实现,重新评估原因(如小数据规模开销大) |
| step3 实施时发现 step0 映射有误 | 修正映射表,选择更适合的函数或退回到手动循环 |
| 无法在 ARM 设备上实测性能 | 明确标记"待实测",给出测试指令,并在力所能及的平台上做验证 |
| step4 汇编写出来但正确性不过 | 回退到安全 C++/Vec4 实现,确保可用性为先 |
| 文件 | 参考价值 | 优先级 |
|---|---|---|
source/backend/cpu/CPUAttention.cpp | 推荐阅读:多线程 + CoreFunctions 调用的标杆实现 | ★★★ |
source/backend/cpu/compute/CommonOptFunction.h | 必备:CoreFunctions 定义及可用函数签名 | ★★★ |
source/backend/cpu/compute/DenseConvolutionTiledExecutor.cpp | MatMul Tiling + Pack + 多线程参考 | ★★ |
test/speed/MatMulSpeed.cpp | 性能基准测试参考模板 | ★★ |
source/backend/cpu/arm/arm64/MNNPackedMatMul.S | 汇编编写参考案例 | ★ |
现在请打开 skills/arm-cpu-optimize/step0-decompose.md,开始步骤 0。