作者：高洋
转载链接：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算法有效的原理
这时候，我们令w = h = 4，根据论文中的推导方法，就可以给出一组A，B和G：
A=
B=
G=
我们记作
此时Y是个2×2的矩阵，我们就令这个Winograd的变量b=2，此时g是3×3，stride=1，记作winograd_b2f3s1。
二维图像卷积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转换，如下图所示：
多通道卷积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的方式：
对输入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倍左右，十分接近理论上计算量的倍数。下表将两组数据放在一起比较一下：
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专栏。
