爱笑的小姐姐 · 2022年02月14日

“远超”理论浮点峰值

作者:高洋
转载链接:https://zhuanlan.zhihu.com/p/465739282

本文作为上一篇文章(后面简称前文)
高洋:M1 Max初相遇,快快乐乐写卷积

的姊妹篇,旨在介绍前文中提到的可以“远超”理论浮点峰值的卷积算法。

既然“远超”两字带引号,说明并不是真的在物理上超越CPU峰值性能,而是通过数学公式相等变换,将计算量大的算法,转换成计算量小的算法,等效地超越物理峰值。

做信号处理,图像处理的同学,应该对卷积定理非常熟悉,即通过FFT和逆FFT的运算,将计算量是O(nm)的一维或者二维卷积操作,降低到O(max(n×logn, m×logm)),从而大幅提高性能。这个卷积定理可以通过数学方法证明,所以不是近似计算,只不过通过计算机浮点运算方式去算的时候,由于字长有限,计算算法不同,会导致数值上的误差。这些误差是否能接受,要具体问题具体分析。

FFT比较大的问题,在于它是复数域上的算法,对于CNN中的多通道卷积,FFT和逆FFT中间的点乘,会转换成复数矩阵乘法。如果能有一种类似FFT的在实数域上的变换,那么计算量将会减少很多。这个想法在2015年Andrew Lavin和Scott Gray发布的论文Fast Algorithms for Convolutional Neural Networks提出的Winograd卷积算法解决。

Winograd算法已经有几十年的历史,之前一直是为了解决一些特定点数FFT的快速算法问题。Andrew和Scott的论文将Winograd推导到CNN中的卷积计算领域,我们这里忽略文章中的数学推导,只分析一下Winograd算法有效的原理。然后通过有效性的分析,我们再设计在arm上的实现方式。这里仍然使用M1 Max芯片,卷积核大小还是3×3,stride=1,输入输出Tensor还是nchw格式。

Winograd算法有效的原理

image.png
这时候,我们令w = h = 4,根据论文中的推导方法,就可以给出一组A,B和G:

A=
image.png
B=
image.png
G=
image.png
我们记作
image.png
此时Y是个2×2的矩阵,我们就令这个Winograd的变量b=2,此时g是3×3,stride=1,记作winograd_b2f3s1。
image.png
二维图像卷积Winograd变换示意

如上图所示,计算整个二维图像的Winograd卷积,就通过滑动窗口从左到右,从上到下以步长b=2滑过整个图像,对每个4×4分块做Winograd变换变成2×2的矩阵写到结果矩阵中。此时,为了处理边界,4×4分块的个数就是[(in_w - 1) / 2] × [(in_h - 1) / 2],记作in_blk,其中“/”除号表示整除。

我们在实际计算CNN中多通道卷积的时候,不是直接对每个二维图像卷积做Winograd,那样会有非常多冗余操作。事实上,我们可以将输入Tensor一次性做好B转换,filter也提前做好G转换,然后做一次批量的矩阵乘法,再将结果统一做一次A转换,如下图所示:
image.png
多通道卷积Winograd算法流程

这个流程,最左边一列,是对filter的G变换,可以在初始化时候提前转好;中间一列,是每次计算卷积时候,对输入Tensor做的B变换;可以看出filter和输入Tensor变换后的结果,经过转置,正好是16批次的矩阵乘法,其中m = out_c,n = in_blk,k = in_c;矩阵乘法的结果做转置,再对每16个一组的分块做A变换,变成2×2的结果分块;in_blk组2×2分块做好边界处理,正好构成输出Tensor的一个feature map,完成整个计算过程。

我们分析一下整个流程:

[1] 首先filter的G变换由于可以提前算好,因此不算在计算量里;

[2] 输入Tensor的B变换,做了in_c × in_blk组4×4的矩阵乘法,所以总计算量就是2 × 4 × 4 × 4 × in_blk × in_c;

[3] 16批次的矩阵乘法,计算量是16 × 2 × in_blk × in_c × out_c;

[4] 对矩阵乘法结果做A变换,计算量是2 × 4 × 4 × 2 × in_blk × out_c;

我们简单地把[2]定义为Winograd前处理,[3]定义为批量矩阵乘,[4]定义为Winograd后处理。忽略掉这三个步骤计算量里面乘积的常数,会发现前处理的时间复杂度是O(in_blk × in_c),批量矩阵乘时间复杂度是O(in_blk × in_c × out_c),后处理的时间复杂度是O(in_blk × out_c)。

通过观察我们发现三个步骤计算复杂度有一个公共的乘积in_blk,可以约掉;我们知道CNN里面的卷积计算,in_c和out_c在很多情况下都有可能非常大,只要这两个参数足够大,那么这三个步骤的主要计算量,就集中在了批量矩阵乘法中。因此我们经常用矩阵乘法的计算量来近似整个Winograd算法的计算量。

在b2f3s1的这个case中:通用卷积算法的计算量是2 × 3 × 3 × in_c × out_c × out_h × out_w;Winograd算法的计算量是16 × 2 × in_blk × in_c × out_c;其中out_h × out_w在这两个参数足够大的时候,大致等于4 × in_blk,所以用通用卷积算法的计算量除以Winograd算法的计算量,就可以得到36 / 16这个比值,大约是2.25倍。

Winograd的变换多种多样:b,f,s都可以改变,比如b4f3s1,比值就是4;b6f3s1,比值就是5.125;b4f5s1,比值就是6.25... 一般来讲,f和s固定,b越大,这个比值就越大,但是也有上限:比如f3s1,b = out_h = out_w时候,Winograd计算量变成in_h × in_w × 2 × 1 × in_c × out_c,假设in_h = out_h + 2以及in_w = out_w + 2足够大,那么计算量比例就接近18 / 2 = 9。只是在这种case下,in_blk变为1,作为矩阵乘法的一维,相当于退化成矩阵向量乘法,大部分情况下完全无法发挥硬件理论峰值性能,背离了我们探索这种算法的本意。所以对于f3s1这种case,我们常用b=2,4或6这几种大小,就可以平衡利弊,取得不错的效果。

实现方法

本节只介绍b2f3s1的实现,相关代码同样放在了前文的项目里:
https://link.zhihu.com/?target=https://github.com/pigirons/conv3x3_m1

实现的大方向,仍然是采用分块转换并调用sgemm kernel的方式:
image.png
对输入Tensor做分块B变换示意

上图展示了对输入Tensor做分块B变换的逻辑示意,对于连续4组4×4分块左上点在同一行的情况,可以使用arm提供的vld2指令,直接将4个分块的起点拼成一个向量,右边挨着的第二个点也拼成向量,参与向量化的B变换计算。由于B矩阵只有0,1和-1三个值,整个计算过程可以直接展开,跳过0元素,对1和-1直接调用加减法,大幅减少计算负担,核心代码如下(使用arm neon intrinsic实现):

#define CONV_WGB2_F3S1_SRC_TRANS_WMAX(SRC, LDS, DST, LDD) \
{ \
    float32x4x2_t vi[4]; \
    float32x4_t vr[8]; \
    vi[0] = vld2q_f32((SRC) + 0 * (LDS) + 0); \
    vi[1] = vld2q_f32((SRC) + 0 * (LDS) + 2); \
    vi[2] = vld2q_f32((SRC) + 2 * (LDS) + 0); \
    vi[3] = vld2q_f32((SRC) + 2 * (LDS) + 2); \
    vr[0] = vsubq_f32(vi[0].val[0], vi[2].val[0]); \
    vr[1] = vsubq_f32(vi[0].val[1], vi[2].val[1]); \
    vr[2] = vsubq_f32(vi[1].val[0], vi[3].val[0]); \
    vr[3] = vsubq_f32(vi[1].val[1], vi[3].val[1]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 0 * (LDD), vr[4]); \
    vst1q_f32((DST) + 1 * (LDD), vr[5]); \
    vst1q_f32((DST) + 2 * (LDD), vr[6]); \
    vst1q_f32((DST) + 3 * (LDD), vr[7]); \
    vi[0] = vld2q_f32((SRC) + 1 * (LDS) + 0); \
    vi[1] = vld2q_f32((SRC) + 1 * (LDS) + 2); \
    vr[0] = vaddq_f32(vi[0].val[0], vi[2].val[0]); \
    vr[1] = vaddq_f32(vi[0].val[1], vi[2].val[1]); \
    vr[2] = vaddq_f32(vi[1].val[0], vi[3].val[0]); \
    vr[3] = vaddq_f32(vi[1].val[1], vi[3].val[1]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 4 * (LDD), vr[4]); \
    vst1q_f32((DST) + 5 * (LDD), vr[5]); \
    vst1q_f32((DST) + 6 * (LDD), vr[6]); \
    vst1q_f32((DST) + 7 * (LDD), vr[7]); \
    vr[0] = vsubq_f32(vi[2].val[0], vi[0].val[0]); \
    vr[1] = vsubq_f32(vi[2].val[1], vi[0].val[1]); \
    vr[2] = vsubq_f32(vi[3].val[0], vi[1].val[0]); \
    vr[3] = vsubq_f32(vi[3].val[1], vi[1].val[1]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 8 * (LDD), vr[4]); \
    vst1q_f32((DST) + 9 * (LDD), vr[5]); \
    vst1q_f32((DST) + 10 * (LDD), vr[6]); \
    vst1q_f32((DST) + 11 * (LDD), vr[7]); \
    vi[2] = vld2q_f32((SRC) + 3 * (LDS) + 0); \
    vi[3] = vld2q_f32((SRC) + 3 * (LDS) + 2); \
    vr[0] = vsubq_f32(vi[0].val[0], vi[2].val[0]); \
    vr[1] = vsubq_f32(vi[0].val[1], vi[2].val[1]); \
    vr[2] = vsubq_f32(vi[1].val[0], vi[3].val[0]); \
    vr[3] = vsubq_f32(vi[1].val[1], vi[3].val[1]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 12 * (LDD), vr[4]); \
    vst1q_f32((DST) + 13 * (LDD), vr[5]); \
    vst1q_f32((DST) + 14 * (LDD), vr[6]); \
    vst1q_f32((DST) + 15 * (LDD), vr[7]); \
}

对于左上点不在一行的,就只能通过上图示意的步骤,借助一个临时buffer做转置,然后再做B变换写入批量矩阵中,代码如下:

#define CONV_WGB2_F3S1_SRC_TRANS(SRC0, SRC1, SRC2, SRC3, LDS, DST, LDD) \
{ \
    float32x4_t vi[8]; \
    float32x4_t vr[8]; \
    vi[0] = vld1q_f32((SRC0) + 0 * (LDS)); \
    vi[1] = vld1q_f32((SRC1) + 0 * (LDS)); \
    vi[2] = vld1q_f32((SRC2) + 0 * (LDS)); \
    vi[3] = vld1q_f32((SRC3) + 0 * (LDS)); \
    CONV_TRANSPOSE_4X4(vi[0], vi[1], vi[2], vi[3]); \
    vi[4] = vld1q_f32((SRC0) + 2 * (LDS)); \
    vi[5] = vld1q_f32((SRC1) + 2 * (LDS)); \
    vi[6] = vld1q_f32((SRC2) + 2 * (LDS)); \
    vi[7] = vld1q_f32((SRC3) + 2 * (LDS)); \
    CONV_TRANSPOSE_4X4(vi[4], vi[5], vi[6], vi[7]); \
    vr[0] = vsubq_f32(vi[0], vi[4]); \
    vr[1] = vsubq_f32(vi[1], vi[5]); \
    vr[2] = vsubq_f32(vi[2], vi[6]); \
    vr[3] = vsubq_f32(vi[3], vi[7]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 0 * (LDD), vr[4]); \
    vst1q_f32((DST) + 1 * (LDD), vr[5]); \
    vst1q_f32((DST) + 2 * (LDD), vr[6]); \
    vst1q_f32((DST) + 3 * (LDD), vr[7]); \
    vi[0] = vld1q_f32((SRC0) + 1 * (LDS)); \
    vi[1] = vld1q_f32((SRC1) + 1 * (LDS)); \
    vi[2] = vld1q_f32((SRC2) + 1 * (LDS)); \
    vi[3] = vld1q_f32((SRC3) + 1 * (LDS)); \
    CONV_TRANSPOSE_4X4(vi[0], vi[1], vi[2], vi[3]); \
    vr[0] = vaddq_f32(vi[0], vi[4]); \
    vr[1] = vaddq_f32(vi[1], vi[5]); \
    vr[2] = vaddq_f32(vi[2], vi[6]); \
    vr[3] = vaddq_f32(vi[3], vi[7]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 4 * (LDD), vr[4]); \
    vst1q_f32((DST) + 5 * (LDD), vr[5]); \
    vst1q_f32((DST) + 6 * (LDD), vr[6]); \
    vst1q_f32((DST) + 7 * (LDD), vr[7]); \
    vr[0] = vsubq_f32(vi[4], vi[0]); \
    vr[1] = vsubq_f32(vi[5], vi[1]); \
    vr[2] = vsubq_f32(vi[6], vi[2]); \
    vr[3] = vsubq_f32(vi[7], vi[3]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 8 * (LDD), vr[4]); \
    vst1q_f32((DST) + 9 * (LDD), vr[5]); \
    vst1q_f32((DST) + 10 * (LDD), vr[6]); \
    vst1q_f32((DST) + 11 * (LDD), vr[7]); \
    vi[4] = vld1q_f32((SRC0) + 3 * (LDS)); \
    vi[5] = vld1q_f32((SRC1) + 3 * (LDS)); \
    vi[6] = vld1q_f32((SRC2) + 3 * (LDS)); \
    vi[7] = vld1q_f32((SRC3) + 3 * (LDS)); \
    CONV_TRANSPOSE_4X4(vi[4], vi[5], vi[6], vi[7]); \
    vr[0] = vsubq_f32(vi[0], vi[4]); \
    vr[1] = vsubq_f32(vi[1], vi[5]); \
    vr[2] = vsubq_f32(vi[2], vi[6]); \
    vr[3] = vsubq_f32(vi[3], vi[7]); \
    vr[4] = vsubq_f32(vr[0], vr[2]); \
    vr[5] = vaddq_f32(vr[1], vr[2]); \
    vr[6] = vsubq_f32(vr[2], vr[1]); \
    vr[7] = vsubq_f32(vr[1], vr[3]); \
    vst1q_f32((DST) + 12 * (LDD), vr[4]); \
    vst1q_f32((DST) + 13 * (LDD), vr[5]); \
    vst1q_f32((DST) + 14 * (LDD), vr[6]); \
    vst1q_f32((DST) + 15 * (LDD), vr[7]); \
}

输入Tensor的B变换做完以后,就可以直接做批量的矩阵乘法了,连续调用16次sgemm kernel,得到结果矩阵;最后,用类似的方法对结果矩阵做A变换,只是顺序颠倒过来:先对连续4个4×4分块做向量化的A变换,A矩阵同样只有0,1和-1;然后做转置,转置结果是4个2×2分块,每个2×2分块会在同一个向量寄存器中,写入输出Tensor的时候分两次,每次写入2个值,具体实现参考conv_winograd_b2f3s1_m1.cpp即可。

实现过程稍显复杂,我们做一些性能测试,还是使用VGG16的几个重要卷积:

conv = 224×224×64×64,padding = 1,time = 20.517ms,perf = 180.312003 GFLOPS
conv = 112×112×128×128,padding = 1,time = 20.110ms,perf = 183.954548 GFLOPS
conv = 56×56×256×256,padding = 1,time = 20.653ms,perf = 179.118314 GFLOPS
conv = 28×28×512×512,padding = 1,time = 17.141ms,perf = 215.822382 GFLOPS
conv = 14×14×512×512,padding = 1,time = 4.770ms,perf = 193.885390 GFLOPS

这里的GFLOPS,是等效GFLOPS,不是实际值,可以发现,性能相比前文介绍的tile gemm方法大致快了1倍左右,十分接近理论上计算量的倍数。下表将两组数据放在一起比较一下:
image.png

Winograd算法一定是更好的算法吗

通过上面表格的数据,我们发现Winograd算法在VGG16模型上比通用类卷积算法性能高出很多,这一点具有普遍性吗?这一节的标题我使用“更好的算法”,而不是“更快的算法”。说明影响算法好坏的因素不止速度这一点,还包括算法精度,内存占用等其他指标,我们逐一分析。

[1] 通过Winograd算法有效性的原理,我们会发现,in_c和out_c是最关键的因素。它们俩必须都足够大的时候,Winograd的计算量近似估计才成立,不然只要有一个很小(比如16以内),批量矩阵乘法的计算量就无法舍入前处理或者后处理的计算量。这时候要么实际加速比大幅降低,要么甚至不会有加速,比通用算法更慢。最极端的case就是depthwise卷积,相当于进行多组in_c = out_c = 1的通用卷积,所以Winograd一定会特别慢。

[2] in_blk这个值很有意思,它与输入输出Tensor的feature map大小有关系,也与b的大小有关系。当in_blk很小的时候,会对矩阵乘法的效率造成影响,相当于将矩阵乘法的维度n变得很小。同时,in_blk很小往往是feature map太小,这时候如果再用一些b比较大的分块方式,就会造成较多的冗余计算,影响性能。可以看出来VGG16最后两组卷积,由于28>14,前者的性能也会高于后者,且这两个原因都有作用。

[3] 对于1×1的卷积,我们发现Winograd并不会带来计算量的减少,同时由于多了转换的开销,性能反而一定不如通用算法。所以1×1卷积我们一般直接调用矩阵乘法去计算,而不用Winograd算法。

[4] 转换后filter的大小,从原先的3×3×in_c×out_c,增加到了16×in_c×out_c;如果b变大,这个增加量更大,且是平方级别增长。这对内存吃紧的设备是很大的问题。

[5] 通过对不同大小b的Winograd算法与通用算法对比diff,发现b越大,精度越低,但一般不会超过10-5相对误差,多数情况下问题不大。但是从fp32降低精度到fp16以后,这个误差就不会太小了,很多时候b=4时候的fp16算法误差已经会造成模型推理精度掉点了。这也是制约Winograd算法将b做大的一个debuff。

一点思考

回顾一下前文的tile gemm算法,可以发现,想让这个算法适配任意filter size,任意stride,甚至是dilation卷积(任意的hole大小),修改起来都十分容易。我们甚至可以写一个统一的支持任意参数的tile gemm算法,并且其性能一定比同参数的im2col + gemm的方法要快。所以我们完全可以把tile gemm算法作为实现卷积的基础算法,再在适当的case下调用Winograd算法。

同时我们发现,tile gemm和Winograd两种算法在框架上比较接近,都是前处理+矩阵乘法+后处理的模式,且主要的计算量就集中在中间的矩阵乘法上。

对于这些特点,我有时候会思考AI加速器的设计。业界曾经有过关于加速器是直接为某个具体filter size的卷积设计最优电路好,还是专为矩阵乘法设计,通过矩阵乘法计算卷积好。也曾经有过脉动阵列设计多大规模会更好的讨论。如果设计一种电路,由多个小型矩阵乘阵列组成,同时设计多个针对不同算法(比如tile gemm或者Winograd)的前处理和后处理逻辑的电路,然后让它们在软件控制下高效组合调度,不同的卷积参数使用不同的算法以期达到更平衡的性能,支持更多复杂结构网络的高效推理,是不是比单单做大阵列规模,甚至直接优化3×3卷积更好?

当然我不是做硬件芯片设计的,不清楚这里面的复杂性和收益到底怎样,只是随便提提一些想法。

后记

这个系列的文章写完了,不会再有后续了。看在高叔叔这么勤奋发文的份上,大家继续疯狂star吧:
https://github.com/openppl-public

推荐阅读

更多嵌入式AI技术相关内容请关注嵌入式AI专栏。
推荐阅读
关注数
18838
内容数
1372
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息