【端智能】MNN CPU性能优化年度小结

原创 霞影 淘系技术  4月28日

简介

‍

2020年5月，MNN发布了1.0.0版本，作为移动端/服务端/PC均适用的推理引擎，在通用性与高性能方面处于业界领先水平。

这之后，我们并没有停止脚步，而是继续在 CPU/GPU/NPU 持续发力。其中，CPU 作为最普适易用的计算资源，我们在多种架构，多种精度模式下继续进行优化，使 MNN 整体的CPU性能进一步提升：AVX2 架构下性能提升 70%-100%，ARM / ARMv8.2 架构下提升约 10-20%，新增了对 Intel Xeon 处理器的 AVX512 指令集的支持，新建对中低端移动手机的 BF16 优化，新增低比特量化计算支持（Int7 / OAQ / Winograd-Aware）。

▐  成果概述如下

【注：GPU 部分优化成果参见《谈谈MNN GPU性能优化策略》一文】

架构改造

所谓磨刀不误砍柴工，为了支持更多指令集架构与更多精度选择，我们先进行了架构改造，以便于减少后续优化成本。

▐  几何计算

几何计算是我们为了适配各类硬件所提出一种简化算子的方案，其核心思想是从原始算子里面分离坐标映射部分，降解为更基础的算子。

通过几何计算机制，一些原本在CPU上优化得不充分的算子，如 Slice-TF / Pack / Unpack / StridedSlice 等自然消除，相应的业务模型也得到了加速。

但由于各硬件的内存设计不一样，目前几何计算在处理坐标映射时只能比较保守地生成内存拷贝算子而不是直接引用，在处理复杂算子如卷积、反卷积等会因引入额外的性能损耗而无法拆解。

对于CPU，不论是各架构还是各类精度（以下简称分支），内存设计是完全一样的，我们可以在几何计算的基础上，对各分支做复杂算子实现的统一。

▐  算子统一

卷积、反卷积、矩阵乘等算子会有多种实现，以卷积为例有如下实现：

• Slide-Window
  
• Im2Col+GEMM
  
• Winograd-GEMM
  

它们在不同的输入条件下各有优势，MNN 会通过 cost 计算，选择最优实现，保证各类模型都能达到最优性能。

如果针对各分支分别写这些实现，工作量无疑是巨大的。

因此，基于CPU内存同构的特性，我们把繁杂而并不耗时的内存管理、地址计算、预处理、分块并行等复杂逻辑统一实现，每个分支实现核心计算函数。

经过这番改造后，我们算子实现上的优化，可以较低成本同步到各个分支上。

算子优化

▐  矩阵乘实现改进

MNN卷积最初的实现为滑窗算法，为了方便进行SIMD优化，采用NC4HW4 的内存布局，即：

维度上c上对齐到4然后拆分，c的坐标j变换为

Im2Col-GEMM / Winograd-GEMM 算法的实现中，我们会根据这个NC4HW4布局，实现

的矩阵乘汇编，其中X为分块大小，受寄存器数N的限制2X+4≤N，在ARMv8 上为 14 ，然后每次拷贝/计算出 大小的计算，经过矩阵乘计算后填到对应的输出位置。

这样实现默认在C方向上强制做了4对齐，两个输入每次都需要以SIMD单位去读，理论读取内存量为：

可并发的FMLA指令数为X，有如下缺陷：

• 可并行的FMLA指令不足，在中低端CPU上不能完全隐藏掉计算延迟
  
• 内存读取量不够小
  
• 不利于其他SIMD架构实现（比如AVX2并行单位为8）
  

为此，我们修正了矩阵乘算法，采用e_p，1_，h_p的分块模式，e_p，h_p由各架构的核心函数确定，并增加一个把 NC4HW4 转成e_p，1的Pack过程。这样理论读取内存量为： ，可并发的FMLA指令数为（n为SIMD一次计算数）

具体到 ARMv8 架构-浮点分支上，e_p=12,h_p=8,n=4 ，读取内存量 ，并发数为12∗8/4=24>14

虽然额外增加了Pack过程，但总体而言减少了内存读写量，增加了并发数，能达到理论极限性能。

▐  卷积实现改进

在卷积算子的实现中，上一节增加的 Pack 过程可以与 Im2Col 过程合并，仅在Im2Col过程中增加一些转置指令，不影响内存读写的瓶颈耗时。

首先我们看原始的实现，以 w=4,h=2,c=4,n=1,k_x=3,k_y=1 为例，我们需要先基于NC4HW4布局进行Im2Col，再把Im2Col的结果转换到c,e_p的内存布局。

为了对Im2Col和Pack进行合并，我们注意到在w和c的方向上，内存的搬运是有连续性的，因此可以先计算出一个包含一系列矩形框的“骨架”（最多含有 ），每个矩阵框包含源地址、目标偏移、矩形大小等信息，然后再根据这些矩形框，逐个连续从输入源取值并转置到目标位置上。

示例流程如下图：

经过这个优化之后，我们和Google 提出的 的 Indirect Convolution Algorithm （集成在 XNNPack）进行了对比测试，由下表数据可见，在汇编都进行了深度优化的基础上，MNN 的 Im2Col-Pack Fuse - GEMM 卷积算法比XNNPack略优，此外，在非 1x1 卷积（矩阵乘即可，不需要额外处理）较多的模型，综合了 Winograd 算法的 MNN 则明显快于 Tflite-XNNPack 。

A76 上测试数据：

硬件相关优化

▐  ARM 调优

在算子优化的基础上，我们对ARM汇编做了进一步的精细调优，包括但不限于：

• 矩阵乘/packing等广泛存在的存储指令ld* 优化为stp，读取指令选取合适的优化为ldp。
  
• 部分合适的代码优化为读写双pipeline并行，
  
• 部分简单指令例如 sub add指令hide到耗时长指令之内
  

▐  x64-AVX2

主流PC/服务器均支持AVX2的SIMD计算，它有16个256位的向量寄存器，一次可计算8个浮点数。

基于之前的算子优化，我们将 AVX2 的矩阵乘定为24x1x4，并编写汇编深度调优，提升了70%-100%的性能，使AVX2上MNN多数模型运行性能接近或超过理论浮点峰值（注：部分模型用Winograd算法加速，为了和训练框架的统计保证一致，模型计算量仍按滑窗计算）。

▐  x64-AVX512

Intel出的Xeon处理器增加了AVX512指令集，支持512位的向量寄存器，一次可计算16个浮点数，更新的AVX512-VNNI指令集追加了uint8-int8的dot指令，可以一次性计算32个整型乘加。

dot指令-vpdpbusd

FOR j := 0 to 7
    tmp1.word := Signed(ZeroExtend16(a.byte[4*j]) * SignExtend16(b.byte[4*j]))
    tmp2.word := Signed(ZeroExtend16(a.byte[4*j+1]) * SignExtend16(b.byte[4*j+1]))
    tmp3.word := Signed(ZeroExtend16(a.byte[4*j+2]) * SignExtend16(b.byte[4*j+2]))
    tmp4.word := Signed(ZeroExtend16(a.byte[4*j+3]) * SignExtend16(b.byte[4*j+3]))
    dst.dword[j] := src.dword[j] + tmp1 + tmp2 + tmp3 + tmp4
ENDFOR
dst[MAX:256] := 0

在Intel同事的支持下，我们借助AVX512的16单位浮点乘加指令与dot量化计算指令，在AVX2优化的基础上，使浮点矩阵乘加速 60%（1024x1024x1024矩阵由24ms下降到15ms)，量化计算加速200%（单测用例31.53ms下降到10.15ms）。

低精度计算

▐  半精度浮点

采用16位浮点替代浮点计算，一般对精度影响不大，不但能加速，还可以减少一半的内存占用。

2018年之后的新手机大都支持 ARMv8.2 指令集，它支持了FP16的SIMD计算，可以达到浮点的两倍性能。

MNN在2020年即已支持这一功能，在算子统一的基础之上，我们重新设计并实现了ARMv8.2的矩阵分块方案与相应汇编代码，取得了进一步的性能提升。此外，为了满足部分应用使用32位架构部署的需求（减少包大小），我们也把FP16计算相关汇编同步编写了一份32位的汇编，支持了32位架构的FP16加速。

对于老的机器，不支持FP16计算的情况，我们实现了介于 Int8 - FP32 的一种优化方案 —— BF16，以减少内存，提升性能。这个方案的原理是以少量的转换指令（BF16与FP32可以简单的移位指令互转），减少读写内存的损耗。

实测在低端机红米4上，BF16 有 10%-20% 的加速效果，在中端机荣耀10上，大部分模型也有一定性能优势。

红米4 单核性能

荣耀10 单核性能

打开 MNN_SUPPORT_BF16 编译宏，设置 precision = Precision_Low 即可使用。

▐  低比特量化计算

为了进一步提升量化模型计算的性能，我们需要在训练时加一些限定条件，使在特定架构下能用相关指令加速。

MNN支持了OAQ量化计算，7Bit下的Winograd量化计算与AVX2优化，在在PAI平台相应的模型压缩能力支持下，可以在原有量化模型的基础上提升 30% 左右的性能。

简述如下表：

小结

▐  纵向对比

经过这一年优化之后，MNN在ARM/ARMv8.2上有10%-20%左右性能提升，而在 AVX2 上有70%-100%的提升。

Mate30 单核耗时/ms

Mac Retina, 15-inch, Mid 2015 单核耗时/ms

▐  横向对比

过去一年业界推理框架发展很快，有不少的优秀开源框架涌现，TNN / Mindspore-lite / Bolt 都宣称自己的性能业界第一，老将Tflite在集成 XNNPack 后性能也有了质的飞跃。MNN经过这一年优化，在通用性（算子数、训练框架支持，异构计算支持）有明显优势的基础上，性能保持业界领先，如下是和业界开源框架TNN/Bolt/Tengine/Tflite/Mindspore-lite 对比，可见 MNN 整体CPU性能最优：

Mate30-A76-FP32单核耗时/ms

Mate30-A76-FP16单核耗时/ms

小米6-A73-FP32单核耗时/ms

相关开源框架版本选取如下：

• Bolt : 测试时使用master分支, 版本号>r1.1,  commit-id: e951118fb28f5fead39c6e56cba16caef4583a99
  
• TNN : 测试时使用master分支: commit-id: 98003d6a9ce3c917ccf7522bf24de3a4b43e3824
  
• Mindspore-lite: 取release 1.1版本. git clone  [https://github.com/mindspore-ai/mindspore](https://github.com/mindspore-ai/mindspore) -b r1.1
  
• TFlite : 测试时使用master分支 commit-id：ee6fa8155e966654f158dad7da06f0b708e9a650
  

随着新架构、新指令、新算法的涌现，性能优化之路永无止境。MNN将紧随时代潮流，以通用性为基础，适配新硬件特性、新的压缩算法，将推理性能优化到极致，持续帮算法同学解决推理部署的问题。